Repository: scikit-learn/scikit-learn
Pull request: ENH add parameter subsample to KBinsDiscretizer
Reference Issues/PRs

Takes over the stalled #20397 and #17221 and closes #12492

What does this implement/fix? Explain your changes.
Adds a subsampling strategy to KBinsDiscretizer

Problem statement:
KBinsDiscretizer(strategy='quantile') should subsample the training data when n_samples > 1e5
For large `n_samples` the univariate CDFs estimated on a large enough subsample should be accurate enough and it should protect us from the n log(n) complexity of sorting each column.

We already do that in `QuantileTransformer`.

Maybe we should also do that for `strategy='kmeans'`.

Hints from issue discussion:
Shall I take this ?

Sure go ahead !
Hmmm... What obligation do we have to backwards compatibility?

> Hmmm... What obligation do we have to backwards compatibility?

I proposed to use the traditional `FutureWarning` strategy in the review of #12501.
@amueller @ogrisel @TomDLT  is this still needed? If so, I could work on this.

Commit message: remove subsample=None

Patch to review:
diff --git a/doc/whats_new/v1.1.rst b/doc/whats_new/v1.1.rst
index 7141080afbc06..24fe7f968d305 100644
--- a/doc/whats_new/v1.1.rst
+++ b/doc/whats_new/v1.1.rst
@@ -88,6 +88,14 @@ Changelog
   Setting a transformer to "passthrough" will pass the features unchanged.
   :pr:`20860` by :user:`Shubhraneel Pal <shubhraneel>`.
 
+:mod:`sklearn.preprocessing`
+............................
+
+- |Enhancement| Adds a `subsample` parameter to :class:`preprocessing.KBinsDiscretizer`.
+  This allow users to specify a maximum number of samples to be used while fitting
+  the model. The option is only available when `strategy` is set to `quantile`.
+  :pr:`21445` by :user:`Felipe Bidu <fbidu>` and :user:`Amanda Dsouza <amy12xx>`.
+
 :mod:`sklearn.utils`
 ....................
 
diff --git a/sklearn/preprocessing/_discretization.py b/sklearn/preprocessing/_discretization.py
index 79a667911d354..f77410550c116 100644
--- a/sklearn/preprocessing/_discretization.py
+++ b/sklearn/preprocessing/_discretization.py
@@ -15,7 +15,9 @@
 from ..base import BaseEstimator, TransformerMixin
 from ..utils.validation import check_array
 from ..utils.validation import check_is_fitted
+from ..utils.validation import check_random_state
 from ..utils.validation import _check_feature_names_in
+from ..utils import _safe_indexing
 
 
 class KBinsDiscretizer(TransformerMixin, BaseEstimator):
@@ -63,6 +65,21 @@ class KBinsDiscretizer(TransformerMixin, BaseEstimator):
 
         .. versionadded:: 0.24
 
+    subsample : int or None (default='warn')
+        Maximum number of samples, used to fit the model, for computational
+        efficiency. Used when `strategy="quantile"`.
+        subsample=None means that all the training samples are used when
+        computing the quantiles that determine the binning thresholds.
+        Since quantile computation relies on sorting each column of X and
+        that sorting has an n log(n) time complexity, it is recommended
+        to use subsampling on datasets with a very large number of samples.
+
+    random_state : int, RandomState instance or None (default=None)
+        Determines random number generation for subsampling.
+        Please see `subsample` for more details.
+        Pass an int for reproducible results across multiple function calls.
+        See :term:`Glossary <random_state>`.
+
     Attributes
     ----------
     bin_edges_ : ndarray of ndarray of shape (n_features,)
@@ -136,11 +153,22 @@ class KBinsDiscretizer(TransformerMixin, BaseEstimator):
            [ 0.5,  3.5, -1.5,  1.5]])
     """
 
-    def __init__(self, n_bins=5, *, encode="onehot", strategy="quantile", dtype=None):
+    def __init__(
+        self,
+        n_bins=5,
+        *,
+        encode="onehot",
+        strategy="quantile",
+        dtype=None,
+        subsample="warn",
+        random_state=None,
+    ):
         self.n_bins = n_bins
         self.encode = encode
         self.strategy = strategy
         self.dtype = dtype
+        self.subsample = subsample
+        self.random_state = random_state
 
     def fit(self, X, y=None):
         """
@@ -174,6 +202,44 @@ def fit(self, X, y=None):
                 " instead."
             )
 
+        n_samples, n_features = X.shape
+
+        if self.strategy == "quantile" and self.subsample is not None:
+            if self.subsample == "warn":
+                if X.shape[0] > 2e5:
+                    warnings.warn(
+                        "In version 1.2 onwards, subsample=2e5 "
+                        "will be used by default. Set subsample explicitly to "
+                        "silence this warning in the mean time. Set "
+                        "subsample=None to disable subsampling explicitly.",
+                        FutureWarning,
+                    )
+            elif isinstance(self.subsample, numbers.Integral):
+                if self.subsample <= 0:
+                    raise ValueError(
+                        f"Invalid value for `subsample`: {self.subsample!r}. The "
+                        "number of subsamples must be at least one."
+                    )
+
+                rng = check_random_state(self.random_state)
+                if n_samples > self.subsample:
+                    subsample_idx = rng.choice(
+                        n_samples, size=self.subsample, replace=False
+                    )
+                    X = _safe_indexing(X, subsample_idx)
+            else:
+                raise ValueError(
+                    f"Invalid value for `subsample`: {self.subsample!r}. "
+                    "`subsample` must be int or None."
+                )
+        elif self.strategy != "quantile" and isinstance(
+            self.subsample, numbers.Integral
+        ):
+            raise ValueError(
+                f"Invalid parameter for `strategy`: {self.strategy}. "
+                '`subsample` must be used with `strategy`="quantile".'
+            )
+
         valid_encode = ("onehot", "onehot-dense", "ordinal")
         if self.encode not in valid_encode:
             raise ValueError(
diff --git a/sklearn/preprocessing/tests/test_discretization.py b/sklearn/preprocessing/tests/test_discretization.py
index a053332619e39..879cacc33bc0c 100644
--- a/sklearn/preprocessing/tests/test_discretization.py
+++ b/sklearn/preprocessing/tests/test_discretization.py
@@ -3,6 +3,7 @@
 import scipy.sparse as sp
 import warnings
 
+from sklearn import clone
 from sklearn.preprocessing import KBinsDiscretizer
 from sklearn.preprocessing import OneHotEncoder
 from sklearn.utils._testing import (
@@ -37,16 +38,16 @@ def test_valid_n_bins():
 def test_invalid_n_bins():
     est = KBinsDiscretizer(n_bins=1)
     err_msg = (
-        "KBinsDiscretizer received an invalid "
-        "number of bins. Received 1, expected at least 2."
+        "KBinsDiscretizer received an invalid number of bins. Received 1, expected at"
+        " least 2."
     )
     with pytest.raises(ValueError, match=err_msg):
         est.fit_transform(X)
 
     est = KBinsDiscretizer(n_bins=1.1)
     err_msg = (
-        "KBinsDiscretizer received an invalid "
-        "n_bins type. Received float, expected int."
+        "KBinsDiscretizer received an invalid n_bins type. Received float, expected"
+        " int."
     )
     with pytest.raises(ValueError, match=err_msg):
         est.fit_transform(X)
@@ -357,3 +358,74 @@ def test_32_equal_64(input_dtype, encode):
     Xt_64 = kbd_64.transform(X_input)
 
     assert_allclose_dense_sparse(Xt_32, Xt_64)
+
+
+@pytest.mark.parametrize("subsample", [None, "warn"])
+def test_kbinsdiscretizer_subsample_default(subsample):
+    # Since the size of X is small (< 2e5), subsampling will not take place.
+    X = np.array([-2, 1.5, -4, -1]).reshape(-1, 1)
+    kbd_default = KBinsDiscretizer(n_bins=10, encode="ordinal", strategy="quantile")
+    kbd_default.fit(X)
+
+    kbd_with_subsampling = clone(kbd_default)
+    kbd_with_subsampling.set_params(subsample=subsample)
+    kbd_with_subsampling.fit(X)
+
+    for bin_kbd_default, bin_kbd_with_subsampling in zip(
+        kbd_default.bin_edges_[0], kbd_with_subsampling.bin_edges_[0]
+    ):
+        np.testing.assert_allclose(bin_kbd_default, bin_kbd_with_subsampling)
+    assert kbd_default.bin_edges_.shape == kbd_with_subsampling.bin_edges_.shape
+
+
+def test_kbinsdiscretizer_subsample_strategy_other():
+    X = np.array([-2, 1.5, -4, -1]).reshape(-1, 1)
+    kbd = KBinsDiscretizer(n_bins=10, encode="ordinal", strategy="uniform", subsample=3)
+
+    err_msg = '`subsample` must be used with `strategy`="quantile".'
+    with pytest.raises(ValueError, match=err_msg):
+        kbd.fit(X)
+
+
+def test_kbinsdiscretizer_subsample_value_other():
+    X = np.array([-2, 1.5, -4, -1]).reshape(-1, 1)
+    kbd = KBinsDiscretizer(
+        n_bins=10, encode="ordinal", strategy="quantile", subsample="full"
+    )
+
+    with pytest.raises(ValueError, match="`subsample` must be int or None."):
+        kbd.fit(X)
+
+
+def test_kbinsdiscretizer_subsample_warn():
+    X = np.random.rand(200001, 1).reshape(-1, 1)
+    kbd = KBinsDiscretizer(n_bins=100, encode="ordinal", strategy="quantile")
+
+    msg = "In version 1.2 onwards, subsample=2e5 will be used by default."
+    with pytest.warns(FutureWarning, match=msg):
+        kbd.fit(X)
+
+
+@pytest.mark.parametrize("subsample", [0, int(2e5)])
+def test_kbinsdiscretizer_subsample_values(subsample):
+    X = np.random.rand(220000, 1).reshape(-1, 1)
+    kbd_default = KBinsDiscretizer(n_bins=10, encode="ordinal", strategy="quantile")
+
+    kbd_with_subsampling = clone(kbd_default)
+    kbd_with_subsampling.set_params(subsample=subsample)
+
+    if subsample == 0:
+        with pytest.raises(
+            ValueError, match="number of subsamples must be at least one."
+        ):
+            kbd_with_subsampling.fit(X)
+    else:
+        msg = "In version 1.2 onwards, subsample=2e5 will be used by default."
+        with pytest.warns(FutureWarning, match=msg):
+            kbd_default.fit(X)
+
+        kbd_with_subsampling.fit(X)
+        assert not np.all(
+            kbd_default.bin_edges_[0] == kbd_with_subsampling.bin_edges_[0]
+        )
+        assert kbd_default.bin_edges_.shape == kbd_with_subsampling.bin_edges_.shape


Code review comments:
diff_hunk: @@ -174,6 +202,44 @@ def fit(self, X, y=None):
                 " instead."
             )
 
+        n_samples, n_features = X.shape
+
+        if self.strategy == "quantile" and self.subsample is not None:
+            if self.subsample == "warn":
+                if X.shape[0] > 2e5:
+                    warnings.warn(
+                        "In version 1.2 onwards, subsample=2e5 "
+                        "will be used by default. Set subsample explicitly to "
+                        "silence this warning in the mean time. Set "
+                        "subsample=None to disable subsampling explicitly.",
+                        FutureWarning,
+                    )
+            elif isinstance(self.subsample, numbers.Integral):
+                if self.subsample <= 0:
+                    raise ValueError(
+                        f"Invalid value for `subsample`: {self.subsample!r}. The "
+                        "number of subsamples must be at least one."
+                    )
+
+                rng = check_random_state(self.random_state)
+                if n_samples > self.subsample:
+                    subsample_idx = rng.choice(
+                        n_samples, size=self.subsample, replace=False
+                    )
+                    X = _safe_indexing(X, subsample_idx)
+            else:
+                raise ValueError(
+                    f"Invalid value for `subsample`: {self.subsample!r}. "
+                    "`subsample` must be int or None."
+                )
+        elif self.strategy != "quantile" and isinstance(
+            self.subsample, numbers.Integral
+        ):
+            raise ValueError(
+                f"Invalid parameter for `strategy`: {self.strategy}. "
+                '`subsample` must be used with `strategy`="quantile".'
path: sklearn/preprocessing/_discretization.py
text: @user1:
```suggestion
                '`subsample` must be used with `strategy="quantile"`.'
```

diff_hunk: @@ -174,6 +202,44 @@ def fit(self, X, y=None):
                 " instead."
             )
 
+        n_samples, n_features = X.shape
+
+        if self.strategy == "quantile" and self.subsample is not None:
+            if self.subsample == "warn":
+                if X.shape[0] > 2e5:
+                    warnings.warn(
+                        "In version 1.2 onwards, subsample=2e5 "
+                        "will be used by default. Set subsample explicitly to "
+                        "silence this warning in the mean time. Set "
+                        "subsample=None to disable subsampling explicitly.",
+                        FutureWarning,
+                    )
+            elif isinstance(self.subsample, numbers.Integral):
+                if self.subsample <= 0:
+                    raise ValueError(
+                        f"Invalid value for `subsample`: {self.subsample!r}. The "
+                        "number of subsamples must be at least one."
+                    )
+
+                rng = check_random_state(self.random_state)
+                if n_samples > self.subsample:
+                    subsample_idx = rng.choice(
+                        n_samples, size=self.subsample, replace=False
+                    )
+                    X = _safe_indexing(X, subsample_idx)
+            else:
+                raise ValueError(
+                    f"Invalid value for `subsample`: {self.subsample!r}. "
+                    "`subsample` must be int or None."
+                )
path: sklearn/preprocessing/_discretization.py
text: @user1:
We can take advantage of the `check_scalar` to raise the right error message:

```suggestion
            else:
                self.subsample = check_scalar(
                    self.subsample, "subsample", numbers.Integral, min_val=1
                )
                rng = check_random_state(self.random_state)
                if n_samples > self.subsample:
                    subsample_idx = rng.choice(
                        n_samples, size=self.subsample, replace=False
                    )
                    X = _safe_indexing(X, subsample_idx)
```

`check_scalar` can be imported as `from ..utils.validation import check_scalar`

diff_hunk: @@ -63,6 +65,21 @@ class KBinsDiscretizer(TransformerMixin, BaseEstimator):
 
         .. versionadded:: 0.24
 
+    subsample : int or None (default='warn')
+        Maximum number of samples, used to fit the model, for computational
+        efficiency. Used when `strategy="quantile"`.
+        subsample=None means that all the training samples are used when
+        computing the quantiles that determine the binning thresholds.
+        Since quantile computation relies on sorting each column of X and
+        that sorting has an n log(n) time complexity, it is recommended
+        to use subsampling on datasets with a very large number of samples.
+
+    random_state : int, RandomState instance or None (default=None)
+        Determines random number generation for subsampling.
+        Please see `subsample` for more details.
+        Pass an int for reproducible results across multiple function calls.
+        See :term:`Glossary <random_state>`.
path: sklearn/preprocessing/_discretization.py
text: @user1:
This follow the same description template:
```suggestion
        Determines random number generation for subsampling.
        Pass an int for reproducible results across multiple function calls.
        See the `subsample` parameter for more details.
        See :term:`Glossary <random_state>`.
```

diff_hunk: @@ -174,6 +202,44 @@ def fit(self, X, y=None):
                 " instead."
             )
 
+        n_samples, n_features = X.shape
+
+        if self.strategy == "quantile" and self.subsample is not None:
+            if self.subsample == "warn":
+                if X.shape[0] > 2e5:
+                    warnings.warn(
+                        "In version 1.2 onwards, subsample=2e5 "
+                        "will be used by default. Set subsample explicitly to "
+                        "silence this warning in the mean time. Set "
+                        "subsample=None to disable subsampling explicitly.",
+                        FutureWarning,
+                    )
+            elif isinstance(self.subsample, numbers.Integral):
+                if self.subsample <= 0:
+                    raise ValueError(
+                        f"Invalid value for `subsample`: {self.subsample!r}. The "
+                        "number of subsamples must be at least one."
+                    )
path: sklearn/preprocessing/_discretization.py
text: @user1:
We are changing validation checks to use [`check_scalar`](https://github.com/scikit-learn/scikit-learn/blob/d1285a9232338c205d4641281efe7b13e48892c2/sklearn/utils/validation.py#L1239) when possible (see https://github.com/scikit-learn/scikit-learn/issues/20724).

You should be able to use it for this scalar validation check and potentially for other ones (if any).

diff_hunk: @@ -357,3 +358,74 @@ def test_32_equal_64(input_dtype, encode):
     Xt_64 = kbd_64.transform(X_input)
 
     assert_allclose_dense_sparse(Xt_32, Xt_64)
+
+
+@pytest.mark.parametrize("subsample", [None, "warn"])
+def test_kbinsdiscretizer_subsample_default(subsample):
+    # Since the size of X is small (< 2e5), subsampling will not take place.
+    X = np.array([-2, 1.5, -4, -1]).reshape(-1, 1)
+    kbd_default = KBinsDiscretizer(n_bins=10, encode="ordinal", strategy="quantile")
+    kbd_default.fit(X)
+
+    kbd_with_subsampling = clone(kbd_default)
+    kbd_with_subsampling.set_params(subsample=subsample)
+    kbd_with_subsampling.fit(X)
+
+    for bin_kbd_default, bin_kbd_with_subsampling in zip(
+        kbd_default.bin_edges_[0], kbd_with_subsampling.bin_edges_[0]
+    ):
+        np.testing.assert_allclose(bin_kbd_default, bin_kbd_with_subsampling)
+    assert kbd_default.bin_edges_.shape == kbd_with_subsampling.bin_edges_.shape
+
+
+def test_kbinsdiscretizer_subsample_strategy_other():
+    X = np.array([-2, 1.5, -4, -1]).reshape(-1, 1)
+    kbd = KBinsDiscretizer(n_bins=10, encode="ordinal", strategy="uniform", subsample=3)
+
+    err_msg = '`subsample` must be used with `strategy`="quantile".'
+    with pytest.raises(ValueError, match=err_msg):
+        kbd.fit(X)
+
+
+def test_kbinsdiscretizer_subsample_value_other():
+    X = np.array([-2, 1.5, -4, -1]).reshape(-1, 1)
+    kbd = KBinsDiscretizer(
+        n_bins=10, encode="ordinal", strategy="quantile", subsample="full"
+    )
+
+    with pytest.raises(ValueError, match="`subsample` must be int or None."):
path: sklearn/preprocessing/tests/test_discretization.py
text: @user1:
By using `check_scalar` the code will change with:

```suggestion
   msg = (
       "subsample must be an instance of <class 'numbers.Integral'>, not "
       "<class 'str'>."
   )
    with pytest.raises(TypeError, match=msg):
```

diff_hunk: @@ -357,3 +358,74 @@ def test_32_equal_64(input_dtype, encode):
     Xt_64 = kbd_64.transform(X_input)
 
     assert_allclose_dense_sparse(Xt_32, Xt_64)
+
+
+@pytest.mark.parametrize("subsample", [None, "warn"])
+def test_kbinsdiscretizer_subsample_default(subsample):
+    # Since the size of X is small (< 2e5), subsampling will not take place.
+    X = np.array([-2, 1.5, -4, -1]).reshape(-1, 1)
+    kbd_default = KBinsDiscretizer(n_bins=10, encode="ordinal", strategy="quantile")
+    kbd_default.fit(X)
+
+    kbd_with_subsampling = clone(kbd_default)
+    kbd_with_subsampling.set_params(subsample=subsample)
+    kbd_with_subsampling.fit(X)
+
+    for bin_kbd_default, bin_kbd_with_subsampling in zip(
+        kbd_default.bin_edges_[0], kbd_with_subsampling.bin_edges_[0]
+    ):
+        np.testing.assert_allclose(bin_kbd_default, bin_kbd_with_subsampling)
+    assert kbd_default.bin_edges_.shape == kbd_with_subsampling.bin_edges_.shape
+
+
+def test_kbinsdiscretizer_subsample_strategy_other():
+    X = np.array([-2, 1.5, -4, -1]).reshape(-1, 1)
+    kbd = KBinsDiscretizer(n_bins=10, encode="ordinal", strategy="uniform", subsample=3)
+
+    err_msg = '`subsample` must be used with `strategy`="quantile".'
+    with pytest.raises(ValueError, match=err_msg):
+        kbd.fit(X)
+
+
+def test_kbinsdiscretizer_subsample_value_other():
+    X = np.array([-2, 1.5, -4, -1]).reshape(-1, 1)
+    kbd = KBinsDiscretizer(
+        n_bins=10, encode="ordinal", strategy="quantile", subsample="full"
+    )
+
+    with pytest.raises(ValueError, match="`subsample` must be int or None."):
+        kbd.fit(X)
+
+
+def test_kbinsdiscretizer_subsample_warn():
+    X = np.random.rand(200001, 1).reshape(-1, 1)
+    kbd = KBinsDiscretizer(n_bins=100, encode="ordinal", strategy="quantile")
+
+    msg = "In version 1.2 onwards, subsample=2e5 will be used by default."
+    with pytest.warns(FutureWarning, match=msg):
+        kbd.fit(X)
+
+
+@pytest.mark.parametrize("subsample", [0, int(2e5)])
+def test_kbinsdiscretizer_subsample_values(subsample):
+    X = np.random.rand(220000, 1).reshape(-1, 1)
+    kbd_default = KBinsDiscretizer(n_bins=10, encode="ordinal", strategy="quantile")
+
+    kbd_with_subsampling = clone(kbd_default)
+    kbd_with_subsampling.set_params(subsample=subsample)
+
+    if subsample == 0:
+        with pytest.raises(
+            ValueError, match="number of subsamples must be at least one."
path: sklearn/preprocessing/tests/test_discretization.py
text: @user1:
Here, we also need to change the message matched.

```suggestion
            ValueError, match="subsample == 0, must be >= 1."
```

diff_hunk: @@ -63,6 +65,21 @@ class KBinsDiscretizer(TransformerMixin, BaseEstimator):
 
         .. versionadded:: 0.24
 
+    subsample : int or None (default='warn')
+        Maximum number of samples, used to fit the model, for computational
+        efficiency. Used when `strategy="quantile"`.
+        subsample=None means that all the training samples are used when
+        computing the quantiles that determine the binning thresholds.
+        Since quantile computation relies on sorting each column of X and
+        that sorting has an n log(n) time complexity, it is recommended
+        to use subsampling on datasets with a very large number of samples.
+
+    random_state : int, RandomState instance or None (default=None)
path: sklearn/preprocessing/_discretization.py
text: @user1:
```suggestion
    random_state : int, RandomState instance or None, default=None
```

diff_hunk: @@ -357,3 +358,74 @@ def test_32_equal_64(input_dtype, encode):
     Xt_64 = kbd_64.transform(X_input)
 
     assert_allclose_dense_sparse(Xt_32, Xt_64)
+
+
+@pytest.mark.parametrize("subsample", [None, "warn"])
+def test_kbinsdiscretizer_subsample_default(subsample):
+    # Since the size of X is small (< 2e5), subsampling will not take place.
+    X = np.array([-2, 1.5, -4, -1]).reshape(-1, 1)
+    kbd_default = KBinsDiscretizer(n_bins=10, encode="ordinal", strategy="quantile")
+    kbd_default.fit(X)
+
+    kbd_with_subsampling = clone(kbd_default)
+    kbd_with_subsampling.set_params(subsample=subsample)
+    kbd_with_subsampling.fit(X)
+
+    for bin_kbd_default, bin_kbd_with_subsampling in zip(
+        kbd_default.bin_edges_[0], kbd_with_subsampling.bin_edges_[0]
+    ):
+        np.testing.assert_allclose(bin_kbd_default, bin_kbd_with_subsampling)
+    assert kbd_default.bin_edges_.shape == kbd_with_subsampling.bin_edges_.shape
+
+
+def test_kbinsdiscretizer_subsample_strategy_other():
+    X = np.array([-2, 1.5, -4, -1]).reshape(-1, 1)
+    kbd = KBinsDiscretizer(n_bins=10, encode="ordinal", strategy="uniform", subsample=3)
+
+    err_msg = '`subsample` must be used with `strategy`="quantile".'
+    with pytest.raises(ValueError, match=err_msg):
+        kbd.fit(X)
+
+
+def test_kbinsdiscretizer_subsample_value_other():
+    X = np.array([-2, 1.5, -4, -1]).reshape(-1, 1)
+    kbd = KBinsDiscretizer(
+        n_bins=10, encode="ordinal", strategy="quantile", subsample="full"
+    )
+
+    with pytest.raises(ValueError, match="`subsample` must be int or None."):
+        kbd.fit(X)
+
+
+def test_kbinsdiscretizer_subsample_warn():
+    X = np.random.rand(200001, 1).reshape(-1, 1)
+    kbd = KBinsDiscretizer(n_bins=100, encode="ordinal", strategy="quantile")
+
+    msg = "In version 1.2 onwards, subsample=2e5 will be used by default."
+    with pytest.warns(FutureWarning, match=msg):
+        kbd.fit(X)
path: sklearn/preprocessing/tests/test_discretization.py
text: @user1:
```suggestion
# TODO: Remove in 1.2
def test_kbinsdiscretizer_subsample_warn():
    X = np.random.rand(200001, 1).reshape(-1, 1)
    kbd = KBinsDiscretizer(n_bins=100, encode="ordinal", strategy="quantile")

    msg = "In version 1.2 onwards, subsample=2e5 will be used by default."
    with pytest.warns(FutureWarning, match=msg):
        kbd.fit(X)
```

diff_hunk: @@ -357,3 +358,74 @@ def test_32_equal_64(input_dtype, encode):
     Xt_64 = kbd_64.transform(X_input)
 
     assert_allclose_dense_sparse(Xt_32, Xt_64)
+
+
+@pytest.mark.parametrize("subsample", [None, "warn"])
+def test_kbinsdiscretizer_subsample_default(subsample):
+    # Since the size of X is small (< 2e5), subsampling will not take place.
+    X = np.array([-2, 1.5, -4, -1]).reshape(-1, 1)
+    kbd_default = KBinsDiscretizer(n_bins=10, encode="ordinal", strategy="quantile")
+    kbd_default.fit(X)
+
+    kbd_with_subsampling = clone(kbd_default)
+    kbd_with_subsampling.set_params(subsample=subsample)
+    kbd_with_subsampling.fit(X)
+
+    for bin_kbd_default, bin_kbd_with_subsampling in zip(
+        kbd_default.bin_edges_[0], kbd_with_subsampling.bin_edges_[0]
+    ):
+        np.testing.assert_allclose(bin_kbd_default, bin_kbd_with_subsampling)
+    assert kbd_default.bin_edges_.shape == kbd_with_subsampling.bin_edges_.shape
+
+
+def test_kbinsdiscretizer_subsample_strategy_other():
path: sklearn/preprocessing/tests/test_discretization.py
text: @user1:
```suggestion
def test_kbinsdiscretizer_subsample_invalid_strategy():
```

diff_hunk: @@ -357,3 +358,74 @@ def test_32_equal_64(input_dtype, encode):
     Xt_64 = kbd_64.transform(X_input)
 
     assert_allclose_dense_sparse(Xt_32, Xt_64)
+
+
+@pytest.mark.parametrize("subsample", [None, "warn"])
+def test_kbinsdiscretizer_subsample_default(subsample):
+    # Since the size of X is small (< 2e5), subsampling will not take place.
+    X = np.array([-2, 1.5, -4, -1]).reshape(-1, 1)
+    kbd_default = KBinsDiscretizer(n_bins=10, encode="ordinal", strategy="quantile")
+    kbd_default.fit(X)
+
+    kbd_with_subsampling = clone(kbd_default)
+    kbd_with_subsampling.set_params(subsample=subsample)
+    kbd_with_subsampling.fit(X)
+
+    for bin_kbd_default, bin_kbd_with_subsampling in zip(
+        kbd_default.bin_edges_[0], kbd_with_subsampling.bin_edges_[0]
+    ):
+        np.testing.assert_allclose(bin_kbd_default, bin_kbd_with_subsampling)
+    assert kbd_default.bin_edges_.shape == kbd_with_subsampling.bin_edges_.shape
+
+
+def test_kbinsdiscretizer_subsample_strategy_other():
+    X = np.array([-2, 1.5, -4, -1]).reshape(-1, 1)
+    kbd = KBinsDiscretizer(n_bins=10, encode="ordinal", strategy="uniform", subsample=3)
+
+    err_msg = '`subsample` must be used with `strategy`="quantile".'
+    with pytest.raises(ValueError, match=err_msg):
+        kbd.fit(X)
+
+
+def test_kbinsdiscretizer_subsample_value_other():
path: sklearn/preprocessing/tests/test_discretization.py
text: @user1:
```suggestion
def test_kbinsdiscretizer_subsample_invalid_type():
```

diff_hunk: @@ -357,3 +358,74 @@ def test_32_equal_64(input_dtype, encode):
     Xt_64 = kbd_64.transform(X_input)
 
     assert_allclose_dense_sparse(Xt_32, Xt_64)
+
+
+@pytest.mark.parametrize("subsample", [None, "warn"])
+def test_kbinsdiscretizer_subsample_default(subsample):
+    # Since the size of X is small (< 2e5), subsampling will not take place.
+    X = np.array([-2, 1.5, -4, -1]).reshape(-1, 1)
+    kbd_default = KBinsDiscretizer(n_bins=10, encode="ordinal", strategy="quantile")
+    kbd_default.fit(X)
+
+    kbd_with_subsampling = clone(kbd_default)
+    kbd_with_subsampling.set_params(subsample=subsample)
+    kbd_with_subsampling.fit(X)
+
+    for bin_kbd_default, bin_kbd_with_subsampling in zip(
+        kbd_default.bin_edges_[0], kbd_with_subsampling.bin_edges_[0]
+    ):
+        np.testing.assert_allclose(bin_kbd_default, bin_kbd_with_subsampling)
+    assert kbd_default.bin_edges_.shape == kbd_with_subsampling.bin_edges_.shape
+
+
+def test_kbinsdiscretizer_subsample_strategy_other():
+    X = np.array([-2, 1.5, -4, -1]).reshape(-1, 1)
+    kbd = KBinsDiscretizer(n_bins=10, encode="ordinal", strategy="uniform", subsample=3)
+
+    err_msg = '`subsample` must be used with `strategy`="quantile".'
+    with pytest.raises(ValueError, match=err_msg):
+        kbd.fit(X)
+
+
+def test_kbinsdiscretizer_subsample_value_other():
+    X = np.array([-2, 1.5, -4, -1]).reshape(-1, 1)
+    kbd = KBinsDiscretizer(
+        n_bins=10, encode="ordinal", strategy="quantile", subsample="full"
+    )
+
+    with pytest.raises(ValueError, match="`subsample` must be int or None."):
+        kbd.fit(X)
+
+
+def test_kbinsdiscretizer_subsample_warn():
+    X = np.random.rand(200001, 1).reshape(-1, 1)
+    kbd = KBinsDiscretizer(n_bins=100, encode="ordinal", strategy="quantile")
+
+    msg = "In version 1.2 onwards, subsample=2e5 will be used by default."
+    with pytest.warns(FutureWarning, match=msg):
+        kbd.fit(X)
+
+
+@pytest.mark.parametrize("subsample", [0, int(2e5)])
+def test_kbinsdiscretizer_subsample_values(subsample):
+    X = np.random.rand(220000, 1).reshape(-1, 1)
+    kbd_default = KBinsDiscretizer(n_bins=10, encode="ordinal", strategy="quantile")
+
+    kbd_with_subsampling = clone(kbd_default)
+    kbd_with_subsampling.set_params(subsample=subsample)
+
+    if subsample == 0:
+        with pytest.raises(
+            ValueError, match="number of subsamples must be at least one."
+        ):
+            kbd_with_subsampling.fit(X)
+    else:
+        msg = "In version 1.2 onwards, subsample=2e5 will be used by default."
+        with pytest.warns(FutureWarning, match=msg):
+            kbd_default.fit(X)
path: sklearn/preprocessing/tests/test_discretization.py
text: @user1:
```suggestion
        # TODO: Remove in 1.2
        msg = "In version 1.2 onwards, subsample=2e5 will be used by default."
        with pytest.warns(FutureWarning, match=msg):
            kbd_default.fit(X)
```

diff_hunk: @@ -63,6 +65,21 @@ class KBinsDiscretizer(TransformerMixin, BaseEstimator):
 
         .. versionadded:: 0.24
 
+    subsample : int or None (default='warn')
+        Maximum number of samples, used to fit the model, for computational
+        efficiency. Used when `strategy="quantile"`.
+        subsample=None means that all the training samples are used when
+        computing the quantiles that determine the binning thresholds.
+        Since quantile computation relies on sorting each column of X and
+        that sorting has an n log(n) time complexity, it is recommended
+        to use subsampling on datasets with a very large number of samples.
path: sklearn/preprocessing/_discretization.py
text: @user1:
This renders maths, which is nice for big-O notation.
```suggestion
        Since quantile computation relies on sorting each column of X and
        that sorting has an :math:`\mathcal{O}(n \log(n))` time complexity,
        it is recommended to use subsampling on datasets with a very large 
        number of samples.
```

@author:
I get a commit error with the math notation:

flake8...................................................................Failed
- hook id: flake8
- exit code: 1

sklearn/preprocessing/_discretization.py:75:36: W605 invalid escape sequence '\m'
sklearn/preprocessing/_discretization.py:75:50: W605 invalid escape sequence '\l'

@user2:
we can maybe just keep the code notation
```
`O(n log(n))`
```

This is more readable when reading the raw docstring and the rendering would be ok in the HTML

@user1:
My bad, this syntax might solely work for reStructuredText files and might need characters to be escaped otherwise -- which apart after having been HMTL-rendered is not nice to read.

@user2:
we might need to declare the string as a raw string:

```
r"""
The docstring with some math inside.
"""

diff_hunk: @@ -174,6 +202,44 @@ def fit(self, X, y=None):
                 " instead."
             )
 
+        n_samples, n_features = X.shape
+
+        if self.strategy == "quantile" and self.subsample is not None:
+            if self.subsample == "warn":
+                if X.shape[0] > 2e5:
path: sklearn/preprocessing/_discretization.py
text: @user1:
```suggestion
                if n_samples > 2e5:
```