Repository: scikit-learn/scikit-learn

Issue:
Apparent mismatch between possible arguments for `average` in the base stochastic gradient class
### Describe the bug

Raised in https://github.com/scikit-learn/scikit-learn/pull/28373#discussion_r1482869890.

The  `average` parameter in [`BaseSGD`](https://github.com/scikit-learn/scikit-learn/blob/main/sklearn/linear_model/_stochastic_gradient.py#L82-L334) (which propagates to `SGDRegressor`, `SGDClassifier` and `SGDOneClassSVM`)  [is constrained](https://github.com/scikit-learn/scikit-learn/blob/main/sklearn/linear_model/_stochastic_gradient.py#L93) to non-negative integers or boolean values: `"average": [Interval(Integral, 0, None, closed="left"), bool, np.bool_]`.

This seems to be at odds with `average=0` seemingly meaning `average=True` which contradicts the typical truth-evaluation of `0`.


### Steps/Code to Reproduce

```
```

### Expected Results

```
```

### Actual Results

```
```

### Versions

This is in the main branch right now.



Discussion:
github-actions[bot]: ## ✔️ Linting Passed
All linting checks passed. Your pull request is in excellent shape! ☀️

<sub> _Generated for commit: [3941875](https://github.com/scikit-learn/scikit-learn/pull/28582/commits/394187597c3d409b7c9a2b117e83a4f7d109e2fe). Link to the linter CI: [here](https://github.com/scikit-learn/scikit-learn/actions/runs/8172162249)_ </sub>

Patch:
diff --git a/doc/whats_new/v1.4.rst b/doc/whats_new/v1.4.rst
index aca929fea..449d266e1 100644
--- a/doc/whats_new/v1.4.rst
+++ b/doc/whats_new/v1.4.rst
@@ -44,6 +44,14 @@ Changelog
   `get_feature_names_out` returns the correct number of feature names.
   :pr:`28306` by :user:`Brendan Lu <brendanlu>`.
 
+:mod:`sklearn.utils`
+....................
+
+- |Fix| :func:`~utils._safe_indexing` now works correctly for polars DataFrame when
+  `axis=0` and supports indexing polars Series.
+  :pr:`28521` by :user:`Yao Xiao <Charlie-XIAO>`.
+
+
 .. _changes_1_4_1:
 
 Version 1.4.1
diff --git a/doc/whats_new/v1.5.rst b/doc/whats_new/v1.5.rst
index 2b97df541..df65e9c16 100644
--- a/doc/whats_new/v1.5.rst
+++ b/doc/whats_new/v1.5.rst
@@ -200,6 +200,13 @@ Changelog
   will now always be `None` when `tol` is set, as `n_nonzero_coefs` is ignored in
   this case. :pr:`28557` by :user:`Lucy Liu <lucyleeow>`.
 
+- |API| Passing `average=0` to disable averaging is deprecated in
+  :class:`linear_model.PassiveAggressiveClassifier`,
+  :class:`linear_model.PassiveAggressiveRegressor`,
+  :class:`linear_model.SGDClassifier`, :class:`linear_model.SGDRegressor` and
+  :class:`linear_model.SGDOneClassSVM`. Pass `average=False` instead.
+  :pr:`28582` by :user:`Jérémie du Boisberranger <jeremiedbb>`.
+
 :mod:`sklearn.manifold`
 .......................
 
diff --git a/examples/miscellaneous/plot_outlier_detection_bench.py b/examples/miscellaneous/plot_outlier_detection_bench.py
index 8e1de1669..7af47fe28 100644
--- a/examples/miscellaneous/plot_outlier_detection_bench.py
+++ b/examples/miscellaneous/plot_outlier_detection_bench.py
@@ -6,9 +6,11 @@ Evaluation of outlier detection estimators
 This example compares two outlier detection algorithms, namely
 :ref:`local_outlier_factor` (LOF) and :ref:`isolation_forest` (IForest), on
 real-world datasets available in :class:`sklearn.datasets`. The goal is to show
-that different algorithms perform well on different datasets.
+that different algorithms perform well on different datasets and contrast their
+training speed and sensitivity to hyperparameters.
 
-The algorithms are trained in an outlier detection context:
+The algorithms are trained (without labels) on the whole dataset assumed to
+contain outliers.
 
 1. The ROC curves are computed using knowledge of the ground-truth labels
 and displayed using :class:`~sklearn.metrics.RocCurveDisplay`.
@@ -314,6 +316,12 @@ _ = plt.tight_layout(pad=2.0)  # spacing between subplots
 # datasets. The score for IForest is slightly better for the SA dataset and LOF
 # performs considerably better on the Ames housing dataset than IForest.
 #
+# Recall however that Isolation Forest tends to train much faster than LOF on
+# datasets with a large number of samples. LOF needs to compute pairwise
+# distances to find nearest neighbors, which has a quadratic complexity with respect
+# to the number of observations. This can make this method prohibitive on large
+# datasets.
+#
 # Ablation study
 # ==============
 #
diff --git a/sklearn/linear_model/_stochastic_gradient.py b/sklearn/linear_model/_stochastic_gradient.py
index 22aee7d4e..67187bbdb 100644
--- a/sklearn/linear_model/_stochastic_gradient.py
+++ b/sklearn/linear_model/_stochastic_gradient.py
@@ -603,6 +603,17 @@ class BaseSGDClassifier(LinearClassifierMixin, BaseSGD, metaclass=ABCMeta):
             reset=first_call,
         )
 
+        if first_call:
+            # TODO(1.7) remove 0 from average parameter constraint
+            if not isinstance(self.average, (bool, np.bool_)) and self.average == 0:
+                warnings.warn(
+                    (
+                        "Passing average=0 to disable averaging is deprecated and will"
+                        " be removed in 1.7. Please use average=False instead."
+                    ),
+                    FutureWarning,
+                )
+
         n_samples, n_features = X.shape
 
         _check_partial_fit_first_call(self, classes)
@@ -678,6 +689,16 @@ class BaseSGDClassifier(LinearClassifierMixin, BaseSGD, metaclass=ABCMeta):
             # delete the attribute otherwise _partial_fit thinks it's not the first call
             delattr(self, "classes_")
 
+        # TODO(1.7) remove 0 from average parameter constraint
+        if not isinstance(self.average, (bool, np.bool_)) and self.average == 0:
+            warnings.warn(
+                (
+                    "Passing average=0 to disable averaging is deprecated and will be "
+                    "removed in 1.7. Please use average=False instead."
+                ),
+                FutureWarning,
+            )
+
         # labels can be encoded as float, int, or string literals
         # np.unique sorts in asc order; largest class id is positive class
         y = self._validate_data(y=y)
@@ -1465,6 +1486,17 @@ class BaseSGDRegressor(RegressorMixin, BaseSGD):
         )
         y = y.astype(X.dtype, copy=False)
 
+        if first_call:
+            # TODO(1.7) remove 0 from average parameter constraint
+            if not isinstance(self.average, (bool, np.bool_)) and self.average == 0:
+                warnings.warn(
+                    (
+                        "Passing average=0 to disable averaging is deprecated and will"
+                        " be removed in 1.7. Please use average=False instead."
+                    ),
+                    FutureWarning,
+                )
+
         n_samples, n_features = X.shape
 
         sample_weight = _check_sample_weight(sample_weight, X, dtype=X.dtype)
@@ -1542,6 +1574,16 @@ class BaseSGDRegressor(RegressorMixin, BaseSGD):
         intercept_init=None,
         sample_weight=None,
     ):
+        # TODO(1.7) remove 0 from average parameter constraint
+        if not isinstance(self.average, (bool, np.bool_)) and self.average == 0:
+            warnings.warn(
+                (
+                    "Passing average=0 to disable averaging is deprecated and will be "
+                    "removed in 1.7. Please use average=False instead."
+                ),
+                FutureWarning,
+            )
+
         if self.warm_start and getattr(self, "coef_", None) is not None:
             if coef_init is None:
                 coef_init = self.coef_
@@ -2358,6 +2400,17 @@ class SGDOneClassSVM(BaseSGD, OutlierMixin):
             reset=first_call,
         )
 
+        if first_call:
+            # TODO(1.7) remove 0 from average parameter constraint
+            if not isinstance(self.average, (bool, np.bool_)) and self.average == 0:
+                warnings.warn(
+                    (
+                        "Passing average=0 to disable averaging is deprecated and will"
+                        " be removed in 1.7. Please use average=False instead."
+                    ),
+                    FutureWarning,
+                )
+
         n_features = X.shape[1]
 
         # Allocate datastructures from input arguments
@@ -2448,6 +2501,16 @@ class SGDOneClassSVM(BaseSGD, OutlierMixin):
         offset_init=None,
         sample_weight=None,
     ):
+        # TODO(1.7) remove 0 from average parameter constraint
+        if not isinstance(self.average, (bool, np.bool_)) and self.average == 0:
+            warnings.warn(
+                (
+                    "Passing average=0 to disable averaging is deprecated and will be "
+                    "removed in 1.7. Please use average=False instead."
+                ),
+                FutureWarning,
+            )
+
         if self.warm_start and hasattr(self, "coef_"):
             if coef_init is None:
                 coef_init = self.coef_
diff --git a/sklearn/manifold/_barnes_hut_tsne.pyx b/sklearn/manifold/_barnes_hut_tsne.pyx
index 5c8d0ded2..485a14edc 100644
--- a/sklearn/manifold/_barnes_hut_tsne.pyx
+++ b/sklearn/manifold/_barnes_hut_tsne.pyx
@@ -52,7 +52,6 @@ cdef float compute_gradient(float[:] val_P,
                             float theta,
                             int dof,
                             long start,
-                            long stop,
                             bint compute_error,
                             int num_threads) noexcept nogil:
     # Having created the tree, calculate the gradient
@@ -76,7 +75,7 @@ cdef float compute_gradient(float[:] val_P,
     if take_timing:
         t1 = clock()
     sQ = compute_gradient_negative(pos_reference, neg_f, qt, dof, theta, start,
-                                   stop, num_threads)
+                                   num_threads)
     if take_timing:
         t2 = clock()
         printf("[t-SNE] Computing negative gradient: %e ticks\n", ((float) (t2 - t1)))
@@ -175,16 +174,14 @@ cdef double compute_gradient_negative(float[:, :] pos_reference,
                                       int dof,
                                       float theta,
                                       long start,
-                                      long stop,
                                       int num_threads) noexcept nogil:
-    if stop == -1:
-        stop = pos_reference.shape[0]
     cdef:
         int ax
         int n_dimensions = qt.n_dimensions
         int offset = n_dimensions + 2
         long i, j, idx
-        long n = stop - start
+        long n_samples = pos_reference.shape[0]
+        long n = n_samples - start
         long dta = 0
         long dtb = 0
         float size, dist2s, mult
@@ -204,7 +201,7 @@ cdef double compute_gradient_negative(float[:, :] pos_reference,
         force = <float *> malloc(sizeof(float) * n_dimensions)
         neg_force = <float *> malloc(sizeof(float) * n_dimensions)
 
-        for i in prange(start, stop, schedule='static'):
+        for i in prange(start, n_samples, schedule='static'):
             # Clear the arrays
             for ax in range(n_dimensions):
                 force[ax] = 0.0
@@ -292,7 +289,7 @@ def gradient(float[:] val_P,
         printf("[t-SNE] Computing gradient\n%s", EMPTY_STRING)
 
     C = compute_gradient(val_P, pos_output, neighbors, indptr, forces,
-                         qt, theta, dof, skip_num_points, -1, compute_error,
+                         qt, theta, dof, skip_num_points, compute_error,
                          num_threads)
 
     if verbose > 10:
diff --git a/sklearn/preprocessing/_label.py b/sklearn/preprocessing/_label.py
index 48533c7ec..301dc19bb 100644
--- a/sklearn/preprocessing/_label.py
+++ b/sklearn/preprocessing/_label.py
@@ -827,7 +827,7 @@ class MultiLabelBinarizer(TransformerMixin, BaseEstimator, auto_wrap_output_keys
         class_mapping[:] = tmp
         self.classes_, inverse = np.unique(class_mapping, return_inverse=True)
         # ensure yt.indices keeps its current dtype
-        yt.indices = np.array(inverse[yt.indices], dtype=yt.indices.dtype, copy=False)
+        yt.indices = np.asarray(inverse[yt.indices], dtype=yt.indices.dtype)
 
         if not self.sparse_output:
             yt = yt.toarray()
diff --git a/sklearn/utils/__init__.py b/sklearn/utils/__init__.py
index 354d82400..028615f18 100644
--- a/sklearn/utils/__init__.py
+++ b/sklearn/utils/__init__.py
@@ -5,6 +5,7 @@ The :mod:`sklearn.utils` module includes various utilities.
 import numbers
 import platform
 import struct
+import sys
 import timeit
 import warnings
 from collections.abc import Sequence
@@ -29,7 +30,7 @@ from .murmurhash import murmurhash3_32
 from .validation import (
     _is_arraylike_not_scalar,
     _is_pandas_df,
-    _is_polars_df,
+    _is_polars_df_or_series,
     _use_interchange_protocol,
     as_float_array,
     assert_all_finite,
@@ -126,12 +127,29 @@ def _polars_indexing(X, key, key_dtype, axis):
     """Indexing X with polars interchange protocol."""
     # Polars behavior is more consistent with lists
     if isinstance(key, np.ndarray):
+        # Convert each element of the array to a Python scalar
         key = key.tolist()
+    elif not (np.isscalar(key) or isinstance(key, slice)):
+        key = list(key)
 
     if axis == 1:
+        # Here we are certain to have a polars DataFrame; which can be indexed with
+        # integer and string scalar, and list of integer, string and boolean
         return X[:, key]
-    else:
-        return X[key]
+
+    if key_dtype == "bool":
+        # Boolean mask can be indexed in the same way for Series and DataFrame (axis=0)
+        return X.filter(key)
+
+    # Integer scalar and list of integer can be indexed in the same way for Series and
+    # DataFrame (axis=0)
+    X_indexed = X[key]
+    if np.isscalar(key) and len(X.shape) == 2:
+        # `X_indexed` is a DataFrame with a single row; we return a Series to be
+        # consistent with pandas
+        pl = sys.modules["polars"]
+        return pl.Series(X_indexed.row(0))
+    return X_indexed
 
 
 def _determine_key_type(key, accept_slice=True):
@@ -272,11 +290,11 @@ def _safe_indexing(X, indices, *, axis=0):
     if axis == 1 and isinstance(X, list):
         raise ValueError("axis=1 is not supported for lists")
 
-    if axis == 1 and hasattr(X, "ndim") and X.ndim != 2:
+    if axis == 1 and hasattr(X, "shape") and len(X.shape) != 2:
         raise ValueError(
-            "'X' should be a 2D NumPy array, 2D sparse matrix or pandas "
+            "'X' should be a 2D NumPy array, 2D sparse matrix or "
             "dataframe when indexing the columns (i.e. 'axis=1'). "
-            "Got {} instead with {} dimension(s).".format(type(X), X.ndim)
+            "Got {} instead with {} dimension(s).".format(type(X), len(X.shape))
         )
 
     if (
@@ -289,10 +307,10 @@ def _safe_indexing(X, indices, *, axis=0):
         )
 
     if hasattr(X, "iloc"):
-        # TODO: we should probably use _is_pandas_df(X) instead but this would
-        # require updating some tests such as test_train_test_split_mock_pandas.
+        # TODO: we should probably use _is_pandas_df_or_series(X) instead but this
+        # would require updating some tests such as test_train_test_split_mock_pandas.
         return _pandas_indexing(X, indices, indices_dtype, axis=axis)
-    elif _is_polars_df(X):
+    elif _is_polars_df_or_series(X):
         return _polars_indexing(X, indices, indices_dtype, axis=axis)
     elif hasattr(X, "shape"):
         return _array_indexing(X, indices, indices_dtype, axis=axis)
diff --git a/sklearn/utils/fixes.py b/sklearn/utils/fixes.py
index acbcc7793..8eca047b1 100644
--- a/sklearn/utils/fixes.py
+++ b/sklearn/utils/fixes.py
@@ -275,14 +275,17 @@ except ImportError:
     from scipy.integrate import trapz as trapezoid  # type: ignore  # noqa
 
 
-# TODO: Remove when Pandas > 2.2 is the minimum supported version
+# TODO: Adapt when Pandas > 2.2 is the minimum supported version
 def pd_fillna(pd, frame):
     pd_version = parse_version(pd.__version__).base_version
     if parse_version(pd_version) < parse_version("2.2"):
         frame = frame.fillna(value=np.nan)
     else:
+        infer_objects_kwargs = (
+            {} if parse_version(pd_version) >= parse_version("3") else {"copy": False}
+        )
         with pd.option_context("future.no_silent_downcasting", True):
-            frame = frame.fillna(value=np.nan).infer_objects(copy=False)
+            frame = frame.fillna(value=np.nan).infer_objects(**infer_objects_kwargs)
     return frame
 
 
diff --git a/sklearn/utils/validation.py b/sklearn/utils/validation.py
index c60b38ca6..3bfd835c1 100644
--- a/sklearn/utils/validation.py
+++ b/sklearn/utils/validation.py
@@ -2175,6 +2175,15 @@ def _is_pandas_df(X):
     return isinstance(X, pd.DataFrame)
 
 
+def _is_polars_df_or_series(X):
+    """Return True if the X is a polars dataframe or series."""
+    try:
+        pl = sys.modules["polars"]
+    except KeyError:
+        return False
+    return isinstance(X, (pl.DataFrame, pl.Series))
+
+
 def _is_polars_df(X):
     """Return True if the X is a polars dataframe."""
     try:


Test patch:
diff --git a/sklearn/linear_model/tests/test_passive_aggressive.py b/sklearn/linear_model/tests/test_passive_aggressive.py
index bcfd58b1e..0bcb19eb9 100644
--- a/sklearn/linear_model/tests/test_passive_aggressive.py
+++ b/sklearn/linear_model/tests/test_passive_aggressive.py
@@ -266,3 +266,13 @@ def test_regressor_undefined_methods():
     reg = PassiveAggressiveRegressor(max_iter=100)
     with pytest.raises(AttributeError):
         reg.transform(X)
+
+
+# TODO(1.7): remove
+@pytest.mark.parametrize(
+    "Estimator", [PassiveAggressiveClassifier, PassiveAggressiveRegressor]
+)
+def test_passive_aggressive_deprecated_average(Estimator):
+    est = Estimator(average=0)
+    with pytest.warns(FutureWarning, match="average=0"):
+        est.fit(X, y)
diff --git a/sklearn/linear_model/tests/test_sgd.py b/sklearn/linear_model/tests/test_sgd.py
index d68eaa6d9..46e153c5c 100644
--- a/sklearn/linear_model/tests/test_sgd.py
+++ b/sklearn/linear_model/tests/test_sgd.py
@@ -360,7 +360,7 @@ def test_late_onset_averaging_reached(klass):
         shuffle=False,
     )
     clf2 = klass(
-        average=0,
+        average=False,
         learning_rate="constant",
         loss="squared_error",
         eta0=eta0,
@@ -1546,7 +1546,12 @@ def test_late_onset_averaging_reached_oneclass(klass):
     )
     # 1 pass over the training set with no averaging
     clf2 = klass(
-        average=0, learning_rate="constant", eta0=eta0, nu=nu, max_iter=1, shuffle=False
+        average=False,
+        learning_rate="constant",
+        eta0=eta0,
+        nu=nu,
+        max_iter=1,
+        shuffle=False,
     )
 
     clf1.fit(X)
@@ -2222,3 +2227,11 @@ def test_loss_attribute_deprecation(Estimator):
 
     with pytest.warns(FutureWarning, match="`loss_function_` was deprecated"):
         est.loss_function_
+
+
+# TODO(1.7): remove
+@pytest.mark.parametrize("Estimator", [SGDClassifier, SGDRegressor, SGDOneClassSVM])
+def test_passive_aggressive_deprecated_average(Estimator):
+    est = Estimator(average=0)
+    with pytest.warns(FutureWarning, match="average=0"):
+        est.fit(X, Y)
diff --git a/sklearn/utils/_testing.py b/sklearn/utils/_testing.py
index c78765cb1..d21e4bf1d 100644
--- a/sklearn/utils/_testing.py
+++ b/sklearn/utils/_testing.py
@@ -748,7 +748,9 @@ def _convert_container(
     container : array-like
         The container to convert.
     constructor_name : {"list", "tuple", "array", "sparse", "dataframe", \
-            "series", "index", "slice", "sparse_csr", "sparse_csc"}
+            "series", "index", "slice", "sparse_csr", "sparse_csc", \
+            "sparse_csr_array", "sparse_csc_array", "pyarrow", "polars", \
+            "polars_series"}
         The type of the returned container.
     columns_name : index or array-like, default=None
         For pandas container supporting `columns_names`, it will affect
@@ -808,6 +810,9 @@ def _convert_container(
     elif constructor_name == "series":
         pd = pytest.importorskip("pandas", minversion=minversion)
         return pd.Series(container, dtype=dtype)
+    elif constructor_name == "polars_series":
+        pl = pytest.importorskip("polars", minversion=minversion)
+        return pl.Series(values=container)
     elif constructor_name == "index":
         pd = pytest.importorskip("pandas", minversion=minversion)
         return pd.Index(container, dtype=dtype)
diff --git a/sklearn/utils/tests/test_testing.py b/sklearn/utils/tests/test_testing.py
index c219cb8c0..e176ef160 100644
--- a/sklearn/utils/tests/test_testing.py
+++ b/sklearn/utils/tests/test_testing.py
@@ -684,8 +684,8 @@ def test_convert_container(
 ):
     """Check that we convert the container to the right type of array with the
     right data type."""
-    if constructor_name in ("dataframe", "series", "index"):
-        # delay the import of pandas within the function to only skip this test
+    if constructor_name in ("dataframe", "polars", "series", "polars_series", "index"):
+        # delay the import of pandas/polars within the function to only skip this test
         # instead of the whole file
         container_type = container_type()
     container = [0, 1]
diff --git a/sklearn/utils/tests/test_utils.py b/sklearn/utils/tests/test_utils.py
index d87a4a57c..fa91b2868 100644
--- a/sklearn/utils/tests/test_utils.py
+++ b/sklearn/utils/tests/test_utils.py
@@ -11,7 +11,6 @@ from sklearn.externals._packaging.version import parse as parse_version
 from sklearn.utils import (
     _determine_key_type,
     _get_column_indices,
-    _is_polars_df,
     _message_with_time,
     _print_elapsed_time,
     _safe_assign,
@@ -33,6 +32,7 @@ from sklearn.utils._testing import (
     assert_no_warnings,
 )
 from sklearn.utils.fixes import CSC_CONTAINERS, CSR_CONTAINERS
+from sklearn.utils.validation import _is_polars_df
 
 # toy array
 X_toy = np.arange(9).reshape((3, 3))
@@ -249,7 +249,9 @@ def test_determine_key_type_slice_error():
         _determine_key_type(slice(0, 2, 1), accept_slice=False)
 
 
-@pytest.mark.parametrize("array_type", ["list", "array", "sparse", "dataframe"])
+@pytest.mark.parametrize(
+    "array_type", ["list", "array", "sparse", "dataframe", "polars"]
+)
 @pytest.mark.parametrize("indices_type", ["list", "tuple", "array", "series", "slice"])
 def test_safe_indexing_2d_container_axis_0(array_type, indices_type):
     indices = [1, 2]
@@ -263,7 +265,7 @@ def test_safe_indexing_2d_container_axis_0(array_type, indices_type):
     )
 
 
-@pytest.mark.parametrize("array_type", ["list", "array", "series"])
+@pytest.mark.parametrize("array_type", ["list", "array", "series", "polars_series"])
 @pytest.mark.parametrize("indices_type", ["list", "tuple", "array", "series", "slice"])
 def test_safe_indexing_1d_container(array_type, indices_type):
     indices = [1, 2]
@@ -275,7 +277,7 @@ def test_safe_indexing_1d_container(array_type, indices_type):
     assert_allclose_dense_sparse(subset, _convert_container([2, 3], array_type))
 
 
-@pytest.mark.parametrize("array_type", ["array", "sparse", "dataframe"])
+@pytest.mark.parametrize("array_type", ["array", "sparse", "dataframe", "polars"])
 @pytest.mark.parametrize("indices_type", ["list", "tuple", "array", "series", "slice"])
 @pytest.mark.parametrize("indices", [[1, 2], ["col_1", "col_2"]])
 def test_safe_indexing_2d_container_axis_1(array_type, indices_type, indices):
@@ -291,7 +293,7 @@ def test_safe_indexing_2d_container_axis_1(array_type, indices_type, indices):
     )
     indices_converted = _convert_container(indices_converted, indices_type)
 
-    if isinstance(indices[0], str) and array_type != "dataframe":
+    if isinstance(indices[0], str) and array_type not in ("dataframe", "polars"):
         err_msg = (
             "Specifying the columns using strings is only supported for dataframes"
         )
@@ -306,7 +308,7 @@ def test_safe_indexing_2d_container_axis_1(array_type, indices_type, indices):
 
 @pytest.mark.parametrize("array_read_only", [True, False])
 @pytest.mark.parametrize("indices_read_only", [True, False])
-@pytest.mark.parametrize("array_type", ["array", "sparse", "dataframe"])
+@pytest.mark.parametrize("array_type", ["array", "sparse", "dataframe", "polars"])
 @pytest.mark.parametrize("indices_type", ["array", "series"])
 @pytest.mark.parametrize(
     "axis, expected_array", [(0, [[4, 5, 6], [7, 8, 9]]), (1, [[2, 3], [5, 6], [8, 9]])]
@@ -326,7 +328,7 @@ def test_safe_indexing_2d_read_only_axis_1(
     assert_allclose_dense_sparse(subset, _convert_container(expected_array, array_type))
 
 
-@pytest.mark.parametrize("array_type", ["list", "array", "series"])
+@pytest.mark.parametrize("array_type", ["list", "array", "series", "polars_series"])
 @pytest.mark.parametrize("indices_type", ["list", "tuple", "array", "series"])
 def test_safe_indexing_1d_container_mask(array_type, indices_type):
     indices = [False] + [True] * 2 + [False] * 6
@@ -336,7 +338,7 @@ def test_safe_indexing_1d_container_mask(array_type, indices_type):
     assert_allclose_dense_sparse(subset, _convert_container([2, 3], array_type))
 
 
-@pytest.mark.parametrize("array_type", ["array", "sparse", "dataframe"])
+@pytest.mark.parametrize("array_type", ["array", "sparse", "dataframe", "polars"])
 @pytest.mark.parametrize("indices_type", ["list", "tuple", "array", "series"])
 @pytest.mark.parametrize(
     "axis, expected_subset",
@@ -363,6 +365,7 @@ def test_safe_indexing_2d_mask(array_type, indices_type, axis, expected_subset):
         ("array", "array"),
         ("sparse", "sparse"),
         ("dataframe", "series"),
+        ("polars", "polars_series"),
     ],
 )
 def test_safe_indexing_2d_scalar_axis_0(array_type, expected_output_type):
@@ -373,7 +376,7 @@ def test_safe_indexing_2d_scalar_axis_0(array_type, expected_output_type):
     assert_allclose_dense_sparse(subset, expected_array)
 
 
-@pytest.mark.parametrize("array_type", ["list", "array", "series"])
+@pytest.mark.parametrize("array_type", ["list", "array", "series", "polars_series"])
 def test_safe_indexing_1d_scalar(array_type):
     array = _convert_container([1, 2, 3, 4, 5, 6, 7, 8, 9], array_type)
     indices = 2
@@ -383,7 +386,12 @@ def test_safe_indexing_1d_scalar(array_type):
 
 @pytest.mark.parametrize(
     "array_type, expected_output_type",
-    [("array", "array"), ("sparse", "sparse"), ("dataframe", "series")],
+    [
+        ("array", "array"),
+        ("sparse", "sparse"),
+        ("dataframe", "series"),
+        ("polars", "polars_series"),
+    ],
 )
 @pytest.mark.parametrize("indices", [2, "col_2"])
 def test_safe_indexing_2d_scalar_axis_1(array_type, expected_output_type, indices):
@@ -392,7 +400,7 @@ def test_safe_indexing_2d_scalar_axis_1(array_type, expected_output_type, indice
         [[1, 2, 3], [4, 5, 6], [7, 8, 9]], array_type, columns_name
     )
 
-    if isinstance(indices, str) and array_type != "dataframe":
+    if isinstance(indices, str) and array_type not in ("dataframe", "polars"):
         err_msg = (
             "Specifying the columns using strings is only supported for dataframes"
         )
@@ -429,7 +437,7 @@ def test_safe_indexing_error_axis(axis):
         _safe_indexing(X_toy, [0, 1], axis=axis)
 
 
-@pytest.mark.parametrize("X_constructor", ["array", "series"])
+@pytest.mark.parametrize("X_constructor", ["array", "series", "polars_series"])
 def test_safe_indexing_1d_array_error(X_constructor):
     # check that we are raising an error if the array-like passed is 1D and
     # we try to index on the 2nd dimension
@@ -439,8 +447,11 @@ def test_safe_indexing_1d_array_error(X_constructor):
     elif X_constructor == "series":
         pd = pytest.importorskip("pandas")
         X_constructor = pd.Series(X)
+    elif X_constructor == "polars_series":
+        pl = pytest.importorskip("polars")
+        X_constructor = pl.Series(values=X)
 
-    err_msg = "'X' should be a 2D NumPy array, 2D sparse matrix or pandas"
+    err_msg = "'X' should be a 2D NumPy array, 2D sparse matrix or dataframe"
     with pytest.raises(ValueError, match=err_msg):
         _safe_indexing(X_constructor, [0, 1], axis=1)