Repository: deepchecks/deepchecks

Issue:
Improve metric calculation with many missing labels
Currently, missing label filtering in metric-related checks is done at the moment of calculation, after sampling was already done. This could result in bad metrics when a large percent of labels are missing. 
Better practice to implement - first filter missing labels, then sample. 


Discussion:

Resolving patch:
diff --git a/deepchecks/base/dataset.py b/deepchecks/base/dataset.py
--- a/deepchecks/base/dataset.py
+++ b/deepchecks/base/dataset.py
@@ -420,19 +420,25 @@ def copy(self: TDataset, new_data: pd.DataFrame) -> TDataset:
                    convert_datetime=False, max_categorical_ratio=self._max_categorical_ratio,
                    max_categories=self._max_categories, label_type=self.label_type)
 
-    def sample(self: TDataset, n_samples: int, replace: bool = False, random_state: t.Optional[int] = None) -> TDataset:
+    def sample(self: TDataset, n_samples: int, replace: bool = False, random_state: t.Optional[int] = None,
+               drop_na_label: bool = False) -> TDataset:
         """Create a copy of the dataset object, with the internal dataframe being a sample of the original dataframe.
 
         Args:
             n_samples (int): Number of samples to draw.
             replace (bool, default False): Whether to sample with replacement.
             random_state (int, default None): Random state.
-
+            drop_na_label (bool, default False): Whether to take sample only from rows with exiting label.
         Returns:
             Dataset: instance of the Dataset with sampled internal dataframe.
         """
-        n_samples = min(n_samples, len(self))
-        return self.copy(self._data.sample(n_samples, replace=replace, random_state=random_state))
+        if drop_na_label and self.label_name:
+            valid_idx = self.label_col.notna()
+            data_to_sample = self.data[valid_idx]
+        else:
+            data_to_sample = self.data
+        n_samples = min(n_samples, len(data_to_sample))
+        return self.copy(data_to_sample.sample(n_samples, replace=replace, random_state=random_state))
 
     @property
     def n_samples(self) -> int:
diff --git a/deepchecks/checks/performance/model_error_analysis.py b/deepchecks/checks/performance/model_error_analysis.py
--- a/deepchecks/checks/performance/model_error_analysis.py
+++ b/deepchecks/checks/performance/model_error_analysis.py
@@ -143,8 +143,8 @@ def run(self, train_dataset: Dataset, test_dataset: Dataset, model=None) -> Chec
 
         cat_features = train_dataset.cat_features
 
-        train_dataset = train_dataset.sample(self.n_samples, random_state=self.random_state)
-        test_dataset = test_dataset.sample(self.n_samples, random_state=self.random_state)
+        train_dataset = train_dataset.sample(self.n_samples, random_state=self.random_state, drop_na_label=True)
+        test_dataset = test_dataset.sample(self.n_samples, random_state=self.random_state, drop_na_label=True)
 
         # Create scoring function, used to calculate the per sample model error
         if task_type == ModelType.REGRESSION:


Test patch:
diff --git a/tests/base/dataset_test.py b/tests/base/dataset_test.py
--- a/tests/base/dataset_test.py
+++ b/tests/base/dataset_test.py
@@ -16,7 +16,7 @@
 from sklearn.datasets import load_iris
 from hamcrest import (
     assert_that, instance_of, equal_to, is_,
-    calling, raises, not_none, has_property, all_of, contains_exactly, has_item
+    calling, raises, not_none, has_property, all_of, contains_exactly, has_item, has_length
 )
 
 from deepchecks import Dataset
@@ -891,3 +891,25 @@ def test_label_numpy_multi_4d_array(iris):
     # Act & Assert
     assert_that(calling(Dataset).with_args(iris, label=label),
                 raises(DeepchecksValueError, 'Label must be either column vector or row vector'))
+
+
+def test_sample_with_nan_labels(iris):
+    # Arrange
+    iris = iris.copy()
+    iris.loc[iris['target'] != 2, 'target'] = None
+    dataset = Dataset(iris, label='target')
+    # Act
+    sample = dataset.sample(10000)
+    # Assert
+    assert_that(sample, has_length(150))
+
+
+def test_sample_drop_nan_labels(iris):
+    # Arrange
+    iris = iris.copy()
+    iris.loc[iris['target'] != 2, 'target'] = None
+    dataset = Dataset(iris, label='target')
+    # Act
+    sample = dataset.sample(10000, drop_na_label=True)
+    # Assert
+    assert_that(sample, has_length(50))