Necessary columns can be unordered

plotly · Aug 15, 2023 · 83c2260 · 83c2260
1 parent e2a7cb4
commit 83c2260
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 9 deletions.
diff --git a/packages/python/plotly/plotly/express/_core.py b/packages/python/plotly/plotly/express/_core.py
@@ -1419,15 +1419,13 @@ def build_dataframe(args, constructor):
             else:
                 # Save precious resources by only interchanging columns that are
                 # actually going to be plotted.
-                necessary_columns = [
+                necessary_columns = {
                     i for i in args.values() if isinstance(i, str) and i in columns
-                ]
+                }
                 for field in args:
                     if args[field] is not None and field in array_attrables:
-                        necessary_columns.extend(
-                            [i for i in args[field] if i in columns]
-                        )
-                columns = list(dict.fromkeys(necessary_columns))
+                        necessary_columns.update(i for i in args[field] if i in columns)
+                columns = list(necessary_columns)
                 args["data_frame"] = pd.api.interchange.from_dataframe(
                     args["data_frame"].select_columns_by_name(columns)
                 )

diff --git a/packages/python/plotly/plotly/tests/test_optional/test_px/test_px_input.py b/packages/python/plotly/plotly/tests/test_optional/test_px/test_px_input.py
@@ -8,6 +8,7 @@
 from plotly.express._core import build_dataframe
 from pandas.testing import assert_frame_equal
 
+
 # Fixtures
 # --------
 @pytest.fixture
@@ -292,9 +293,10 @@ def __dataframe__(self):
         ) as mock_from_dataframe:
             build_dataframe(args, go.Scatter)
         mock_from_dataframe.assert_called_once_with(interchange_dataframe_reduced)
-        interchange_dataframe.select_columns_by_name.assert_called_with(
-            ["petal_width", "sepal_length"]
-        )
+        assert set(interchange_dataframe.select_columns_by_name.call_args.args[0]) == {
+            "petal_width",
+            "sepal_length",
+        }
 
         args = dict(data_frame=input_dataframe_reduced, color=None)
         with mock.patch(