fix value error messages in data drift calculations

src/evidently/calculations/data_drift.py

@@ -66,12 +66,17 @@ def get_one_column_drift(
     dataset_columns: DatasetColumns,
     column_type: Optional[str] = None,
 ) -> ColumnDataDriftMetrics:
+    if column_name not in current_data:
+        raise ValueError(f"Cannot find column '{column_name}' in current dataset")
+
+    if column_name not in reference_data:
+        raise ValueError(f"Cannot find column '{column_name}' in reference dataset")
 
     if column_type is None:
         column_type = recognize_column_type(dataset=reference_data, column_name=column_name, columns=dataset_columns)
 
     if column_type not in ("cat", "num"):
-        raise ValueError(f"Cannot calculate drift metric for column {column_name} with type {column_type}")
+        raise ValueError(f"Cannot calculate drift metric for column '{column_name}' with type {column_type}")
 
     if column_name == dataset_columns.utility_columns.target and column_type == "num":
         stattest = options.num_target_stattest_func
@@ -86,22 +91,26 @@ def get_one_column_drift(
     current_column = current_data[column_name]
     reference_column = reference_data[column_name]
 
-    if column_type == "num":
-        if not pd.api.types.is_numeric_dtype(reference_column):
-            raise ValueError(f"Column {column_name} in reference dataset should contain numerical values only.")
-
-        if not pd.api.types.is_numeric_dtype(current_column):
-            raise ValueError(f"Column {column_name} in current dataset should contain numerical values only.")
-
+    # clean and check the column in reference dataset
     reference_column = reference_column.replace([-np.inf, np.inf], np.nan).dropna()
 
     if reference_column.empty:
-        raise ValueError(f"Column '{column_name}' in reference dataset has no values for drift calculation.")
+        raise ValueError(
+            f"An empty column '{column_name}' was provided for drift calculation in the reference dataset."
+        )
 
+    # clean and check the column in current dataset
     current_column = current_column.replace([-np.inf, np.inf], np.nan).dropna()
 
     if current_column.empty:
-        raise ValueError(f"Column '{column_name}' in current dataset has no values for drift calculation.")
+        raise ValueError(f"An empty column '{column_name}' was provided for drift calculation in the current dataset.")
+
+    if column_type == "num":
+        if not pd.api.types.is_numeric_dtype(reference_column):
+            raise ValueError(f"Column '{column_name}' in reference dataset should contain numerical values only.")
+
+        if not pd.api.types.is_numeric_dtype(current_column):
+            raise ValueError(f"Column '{column_name}' in current dataset should contain numerical values only.")
 
     drift_test_function = get_stattest(reference_column, current_column, column_type, stattest)
     drift_result = drift_test_function(reference_column, current_column, column_type, threshold)
@@ -115,9 +124,6 @@ def get_one_column_drift(
     )
 
     if column_type == "num":
-        if not pd.api.types.is_numeric_dtype(reference_column) or not pd.api.types.is_numeric_dtype(current_column):
-            raise ValueError(f"Column {column_name} should only contain numerical values.")
-
         numeric_columns = dataset_columns.num_feature_names
 
         if column_name not in numeric_columns:
@@ -161,6 +167,7 @@ def get_one_column_drift(
         result.current_small_distribution = list(
             reversed(list(map(list, zip(*sorted(current_counts.items(), key=lambda x: str(x[0]))))))
         )
+
     distribution_for_plot = get_distribution_for_column(
         column_name=column_name,
         column_type=column_type,

src/evidently/metrics/data_drift/column_drift_metric.py

@@ -53,10 +53,10 @@ def calculate(self, data: InputData) -> ColumnDriftMetricResults:
             raise ValueError("Reference dataset should be present")
 
         if self.column_name not in data.current_data:
-            raise ValueError(f"Cannot find column {self.column_name} in current dataset")
+            raise ValueError(f"Cannot find column '{self.column_name}' in current dataset")
 
         if self.column_name not in data.reference_data:
-            raise ValueError(f"Cannot find column {self.column_name} in reference dataset")
+            raise ValueError(f"Cannot find column '{self.column_name}' in reference dataset")
 
         dataset_columns = process_columns(data.reference_data, data.column_mapping)
         drift_result = get_one_column_drift(

tests/calculations/data_clean.py → tests/calculations/test_data_clean.py

@@ -3,7 +3,7 @@
 import pytest
 from pytest import approx
 
-from evidently.calculations import replace_infinity_values_to_nan
+from evidently.utils.data_operations import replace_infinity_values_to_nan
 
 
 @pytest.mark.parametrize(

tests/calculations/test_data_drift.py

@@ -0,0 +1,200 @@
+from typing import List
+from typing import Optional
+from typing import Union
+
+import pandas as pd
+import pytest
+
+from evidently import ColumnMapping
+from evidently.calculations.data_drift import ensure_prediction_column_is_string
+from evidently.calculations.data_drift import get_one_column_drift
+from evidently.options import DataDriftOptions
+from evidently.utils.data_operations import process_columns
+
+
+@pytest.mark.parametrize(
+    "prediction_column, current_data, reference_data, threshold, expected_prediction_column",
+    (
+        (None, pd.DataFrame({}), pd.DataFrame({}), 0.0, None),
+        ("preds", pd.DataFrame({"preds": [1, 2, 3]}), pd.DataFrame({"preds": [1, 2, 3]}), 0.0, "preds"),
+        (
+            ["pred_a", "pred_b"],
+            pd.DataFrame({"pred_a": [1, 0, 1], "pred_b": [1, 0, 1]}),
+            pd.DataFrame({"pred_a": [1, 0, 1], "pred_b": [1, 0, 1]}),
+            0.0,
+            "predicted_labels",
+        ),
+        (
+            ["pred_a", "pred_b", "pred_c", "pred_d"],
+            pd.DataFrame(
+                {
+                    "pred_a": [0.5, 0, 0.8],
+                    "pred_b": [0, 0.2, 0.5],
+                    "pred_c": [0.3, 0.2, 0.5],
+                    "pred_d": [0.1, 0.1, 0.9],
+                }
+            ),
+            pd.DataFrame(
+                {
+                    "pred_a": [1, 0, 0, 0],
+                    "pred_b": [0, 1, 0, 0],
+                    "pred_c": [0, 0, 1, 0],
+                    "pred_d": [0, 0, 0, 1],
+                }
+            ),
+            0.3,
+            "predicted_labels",
+        ),
+    ),
+)
+def test_ensure_prediction_column_is_string(
+    prediction_column: Optional[Union[str, List]],
+    current_data: pd.DataFrame,
+    reference_data: pd.DataFrame,
+    threshold: float,
+    expected_prediction_column: Optional[str],
+):
+    result = ensure_prediction_column_is_string(
+        prediction_column=prediction_column,
+        current_data=current_data,
+        reference_data=reference_data,
+        threshold=threshold,
+    )
+    assert result == expected_prediction_column
+
+    # check that string prediction column or a new predicted_labels is in datasets
+    if prediction_column is not None:
+        assert result in current_data
+        assert result in reference_data
+
+
+@pytest.mark.parametrize(
+    "current_data, reference_data, column_name, options, column_type, expected_drift_detected",
+    (
+        (pd.DataFrame({"test": [1, 2, 3]}), pd.DataFrame({"test": [1, 2, 3]}), "test", DataDriftOptions(), None, False),
+        (
+            pd.DataFrame({"test": [1, 2, 3]}),
+            pd.DataFrame({"test": [1, 2, 3]}),
+            "test",
+            DataDriftOptions(),
+            "cat",
+            False,
+        ),
+        (
+            pd.DataFrame({"test": [1, 2, 3], "target": [1, 2, 3]}),
+            pd.DataFrame({"test": [1, 2, 3], "target": [3, 2, 1]}),
+            "test",
+            DataDriftOptions(),
+            None,
+            False,
+        ),
+        (
+            pd.DataFrame({"test": [1, 2, 3], "target": [1, 2, 3]}),
+            pd.DataFrame({"test": [1, 2, 3], "target": [3, 2, 1]}),
+            "test",
+            DataDriftOptions(),
+            "cat",
+            False,
+        ),
+        (
+            pd.DataFrame({"test": [1, 2, 3], "target": [1, 2, 3]}),
+            pd.DataFrame({"test": [4, 5, 6], "target": [1, 2, 3]}),
+            "target",
+            DataDriftOptions(),
+            None,
+            False,
+        ),
+    ),
+)
+def test_get_one_column_drift_success(
+    current_data: pd.DataFrame,
+    reference_data: pd.DataFrame,
+    column_name: str,
+    options: DataDriftOptions,
+    column_type: Optional[str],
+    expected_drift_detected: bool,
+):
+    dataset_columns = process_columns(reference_data, ColumnMapping())
+    result = get_one_column_drift(
+        current_data=current_data,
+        reference_data=reference_data,
+        column_name=column_name,
+        options=options,
+        dataset_columns=dataset_columns,
+        column_type=column_type,
+    )
+    assert result.drift_detected == expected_drift_detected
+
+
+@pytest.mark.parametrize(
+    "current_data, reference_data, column_name, options, column_type, expected_value_error",
+    (
+        (
+            pd.DataFrame({"test": [1, 2, 3]}),
+            pd.DataFrame({"test": [1, 2, 3]}),
+            "feature",
+            DataDriftOptions(),
+            None,
+            "Cannot find column 'feature' in current dataset",
+        ),
+        (
+            pd.DataFrame({"feature": [1, 2, 3]}),
+            pd.DataFrame({"test": [1, 2, 3]}),
+            "feature",
+            DataDriftOptions(),
+            None,
+            "Cannot find column 'feature' in reference dataset",
+        ),
+        (
+            pd.DataFrame({"test": [None, None, None]}),
+            pd.DataFrame({"test": [1, 2, 3]}),
+            "test",
+            DataDriftOptions(),
+            None,
+            "An empty column 'test' was provided for drift calculation in the current dataset.",
+        ),
+        (
+            pd.DataFrame({"test": [1, 2, 3]}),
+            pd.DataFrame({"test": [None, None, None]}),
+            "test",
+            DataDriftOptions(),
+            None,
+            "An empty column 'test' was provided for drift calculation in the reference dataset.",
+        ),
+        (
+            pd.DataFrame({"test": ["a", 2, "c"]}),
+            pd.DataFrame({"test": [1, 2, 3]}),
+            "test",
+            DataDriftOptions(),
+            "num",
+            "Column 'test' in current dataset should contain numerical values only.",
+        ),
+        (
+            pd.DataFrame({"test": [1, 2, 3]}),
+            pd.DataFrame({"test": ["a", "b", 3]}),
+            "test",
+            DataDriftOptions(),
+            "num",
+            "Column 'test' in reference dataset should contain numerical values only.",
+        ),
+    ),
+)
+def test_get_one_column_drift_value_error(
+    current_data: pd.DataFrame,
+    reference_data: pd.DataFrame,
+    column_name: str,
+    options: DataDriftOptions,
+    column_type: Optional[str],
+    expected_value_error: bool,
+):
+    dataset_columns = process_columns(reference_data, ColumnMapping())
+    with pytest.raises(ValueError) as error:
+        get_one_column_drift(
+            current_data=current_data,
+            reference_data=reference_data,
+            column_name=column_name,
+            options=options,
+            dataset_columns=dataset_columns,
+            column_type=column_type,
+        )
+    assert error.value.args[0] == expected_value_error

tests/calculations/data_quality.py → tests/calculations/test_data_quality.py

@@ -2,7 +2,7 @@
 import pandas as pd
 import pytest
 
-from evidently.calculations import get_rows_count
+from evidently.calculations.data_quality import get_rows_count
 
 
 @pytest.mark.parametrize(

tests/metrics/data_drift/test_column_drift_metric.py

@@ -62,6 +62,7 @@ def test_column_drift_metric_success(
             None,
             ColumnDriftMetric(column_name="col"),
         ),
+        # no not-nan values in the column
         (
             pd.DataFrame({"col": [None, np.inf, -np.inf]}),
             pd.DataFrame({"col": [1, 2, 3]}),