iter

Author: glemaitre

imblearn/base.py

@@ -11,9 +11,13 @@
 from sklearn.base import BaseEstimator
 from sklearn.preprocessing import label_binarize
 
+from .dask._support import is_dask_collection
 from .utils import check_sampling_strategy, check_target_type
-from .utils._validation import ArraysTransformer
-from .utils._validation import _deprecate_positional_args
+from .utils._validation import (
+    ArraysTransformer,
+    _deprecate_positional_args,
+    get_classes_counts,
+)
 from .utils.wrapper import check_classification_targets
 
 
@@ -45,9 +49,13 @@ def fit(self, X, y):
         self : object
             Return the instance itself.
         """
-        X, y, _ = self._check_X_y(X, y)
+        dask_collection = any([is_dask_collection(arr) for arr in (X, y)])
+        if (not dask_collection or
+                (dask_collection and self.validate_if_dask_collection)):
+            X, y, _ = self._check_X_y(X, y)
+        self._classes_counts = get_classes_counts(y)
         self.sampling_strategy_ = check_sampling_strategy(
-            self.sampling_strategy, y, self._sampling_type
+            self.sampling_strategy, self._classes_counts, self._sampling_type
         )
         return self
 
@@ -72,12 +80,19 @@ def fit_resample(self, X, y):
         y_resampled : array-like of shape (n_samples_new,)
             The corresponding label of `X_resampled`.
         """
-        check_classification_targets(y)
         arrays_transformer = ArraysTransformer(X, y)
-        X, y, binarize_y = self._check_X_y(X, y)
+        dask_collection = any([is_dask_collection(arr) for arr in (X, y)])
+        if (not dask_collection or
+                (dask_collection and self.validate_if_dask_collection)):
+            check_classification_targets(y)
+            X, y, binarize_y = self._check_X_y(X, y)
+        else:
+            X, y = arrays_transformer.to_dask_array(X, y)
+            binarize_y = False
 
+        self._classes_counts = get_classes_counts(y)
         self.sampling_strategy_ = check_sampling_strategy(
-            self.sampling_strategy, y, self._sampling_type
+            self.sampling_strategy, self._classes_counts, self._sampling_type
         )
 
         output = self._fit_resample(X, y)
@@ -125,8 +140,13 @@ class BaseSampler(SamplerMixin):
     instead.
     """
 
-    def __init__(self, sampling_strategy="auto"):
+    def __init__(
+        self,
+        sampling_strategy="auto",
+        validate_if_dask_collection=False,
+    ):
         self.sampling_strategy = sampling_strategy
+        self.validate_if_dask_collection = validate_if_dask_collection
 
     def _check_X_y(self, X, y, accept_sparse=None):
         if accept_sparse is None:

imblearn/dask/_support.py

@@ -1,13 +1,9 @@
-_REGISTERED_DASK_CONTAINER = []
+def is_dask_collection(container):
+    try:
+        # to keep dask as an optional depency, keep the statement in a
+        # try/except statement
+        from dask import is_dask_collection
 
-try:
-    from dask import array, dataframe
-    _REGISTERED_DASK_CONTAINER += [
-        array.Array, dataframe.Series, dataframe.DataFrame,
-    ]
-except ImportError:
-    pass
-
-
-def is_dask_container(container):
-    return isinstance(container, tuple(_REGISTERED_DASK_CONTAINER))
+        return is_dask_collection(container)
+    except ImportError:
+        return False

imblearn/under_sampling/_prototype_selection/_random_under_sampler.py

@@ -10,16 +10,20 @@
 from sklearn.utils import _safe_indexing
 
 from ..base import BaseUnderSampler
-from ...dask._support import is_dask_container
+from ...dask._support import is_dask_collection
 from ...utils import check_target_type
 from ...utils import Substitution
-from ...utils._docstring import _random_state_docstring
+from ...utils._docstring import (
+    _random_state_docstring,
+    _validate_if_dask_collection_docstring
+)
 from ...utils._validation import _deprecate_positional_args
 
 
 @Substitution(
     sampling_strategy=BaseUnderSampler._sampling_strategy_docstring,
     random_state=_random_state_docstring,
+    validate_if_dask_collection=_validate_if_dask_collection_docstring,
 )
 class RandomUnderSampler(BaseUnderSampler):
     """Class to perform random under-sampling.
@@ -38,6 +42,8 @@ class RandomUnderSampler(BaseUnderSampler):
     replacement : bool, default=False
         Whether the sample is with or without replacement.
 
+    {validate_if_dask_collection}
+
     Attributes
     ----------
     sample_indices_ : ndarray of shape (n_new_samples,)
@@ -74,22 +80,23 @@ class RandomUnderSampler(BaseUnderSampler):
 
     @_deprecate_positional_args
     def __init__(
-        self, *, sampling_strategy="auto", random_state=None, replacement=False
+        self,
+        *,
+        sampling_strategy="auto",
+        random_state=None,
+        replacement=False,
+        validate_if_dask_collection=False,
     ):
-        super().__init__(sampling_strategy=sampling_strategy)
+        super().__init__(
+            sampling_strategy=sampling_strategy,
+            validate_if_dask_collection=validate_if_dask_collection,
+        )
         self.random_state = random_state
         self.replacement = replacement
 
     def _check_X_y(self, X, y):
-        if is_dask_container(y) and hasattr(y, "to_dask_array"):
-            y = y.to_dask_array()
-            y.compute_chunk_sizes()
-        y, binarize_y, self._uniques = check_target_type(
-            y,
-            indicate_one_vs_all=True,
-            return_unique=True,
-        )
-        if not any([is_dask_container(arr) for arr in (X, y)]):
+        y, binarize_y = check_target_type(y, indicate_one_vs_all=True)
+        if not any([is_dask_collection(arr) for arr in (X, y)]):
             X, y = self._validate_data(
                 X,
                 y,
@@ -98,24 +105,23 @@ def _check_X_y(self, X, y):
                 dtype=None,
                 force_all_finite=False,
             )
-        elif is_dask_container(X) and hasattr(X, "to_dask_array"):
-            X = X.to_dask_array()
-            X.compute_chunk_sizes()
         return X, y, binarize_y
 
     @staticmethod
     def _find_target_class_indices(y, target_class):
         target_class_indices = np.flatnonzero(y == target_class)
-        if is_dask_container(y):
-            return target_class_indices.compute()
+        if is_dask_collection(y):
+            from dask import compute
+
+            return compute(target_class_indices)[0]
         return target_class_indices
 
     def _fit_resample(self, X, y):
         random_state = check_random_state(self.random_state)
 
         idx_under = []
 
-        for target_class in self._uniques:
+        for target_class in self._classes_counts:
             target_class_indices = self._find_target_class_indices(
                 y, target_class
             )

imblearn/utils/_docstring.py

@@ -41,3 +41,10 @@ def __call__(self, obj):
         `Glossary <https://scikit-learn.org/stable/glossary.html#term-n-jobs>`_
         for more details.
     """.rstrip()
+
+_validate_if_dask_collection_docstring = \
+  """validate_if_dask_collection : bool, default=False
+        Whether or not `X` and `y` should be validated. This parameter applies
+        only when `X` and `y` are Dask collections where validation might be
+        potentially costly.
+    """.rstrip()