Add test for passing polars dataframes into from_estimator

patrickkwang · patrickkwang · commit bcbc83a9c7fa · 2024-03-28T08:57:27.000-04:00
diff --git a/sklearn/inspection/_plot/tests/test_boundary_decision_display.py b/sklearn/inspection/_plot/tests/test_boundary_decision_display.py
@@ -468,8 +468,8 @@ def test_string_target(pyplot):
     )
 
 
-def test_dataframe_support(pyplot):
-    """Check that passing a dataframe at fit and to the Display does not
+def test_pandas_dataframe_support(pyplot):
+    """Check that passing a pandas dataframe at fit and to the Display does not
     raise warnings.
 
     Non-regression test for:
@@ -485,6 +485,23 @@ def test_dataframe_support(pyplot):
         DecisionBoundaryDisplay.from_estimator(estimator, df, response_method="predict")
 
 
+def test_polars_dataframe_support(pyplot):
+    """Check that passing a polars dataframe at fit and to the Display does not
+    raise warnings.
+
+    Non-regression test for:
+    https://github.com/scikit-learn/scikit-learn/issues/23311
+    """
+    pl = pytest.importorskip("polars")
+    df = pl.DataFrame({"col_x": X[:, 0], "col_y": X[:, 1]})
+    estimator = LogisticRegression().fit(df, y)
+
+    with warnings.catch_warnings():
+        # no warnings linked to feature names validation should be raised
+        warnings.simplefilter("error", UserWarning)
+        DecisionBoundaryDisplay.from_estimator(estimator, df, response_method="predict")
+
+
 @pytest.mark.parametrize("response_method", ["predict_proba", "decision_function"])
 def test_class_of_interest_binary(pyplot, response_method):
     """Check the behaviour of passing `class_of_interest` for plotting the output of
