[FSTORE-1906] Improve online schema validation (#728)

Author: bubriks

python/hsfs/core/feature_group_engine.py

@@ -107,9 +107,11 @@ def save(
         )
 
         if (
-            feature_group.online_enabled
-            and not feature_group.embedding_index
-            and validation_options.get("online_schema_validation", True)
+            not validation_options
+            or (
+                validation_options.get("online_schema_validation", True) # for backwards compatibility
+                and validation_options.get("schema_validation", True)
+            )
         ):
             # validate df schema
             dataframe_features = DataFrameValidator().validate_schema(
@@ -190,9 +192,11 @@ def insert(
         )
 
         if (
-            feature_group.online_enabled
-            and not feature_group.embedding_index
-            and validation_options.get("online_schema_validation", True)
+            not validation_options
+            or (
+                validation_options.get("online_schema_validation", True) # for backwards compatibility
+                and validation_options.get("schema_validation", True)
+            )
         ):
             # validate df schema
             dataframe_features = DataFrameValidator().validate_schema(

python/hsfs/core/schema_validation.py

@@ -60,13 +60,6 @@ def _raise_validation_error(self, errors):
 
     def validate_schema(self, feature_group, df, df_features):
         """Common validation rules"""
-        if feature_group.online_enabled is False:
-            logger.warning("Feature group is not online enabled. Skipping validation")
-            return df_features
-        if feature_group._embedding_index is not None:
-            logger.warning("Feature group is embedding type. Skipping validation")
-            return df_features
-
         validator = self.get_validator(df)
         if validator is None:
             # If no validator is found for this type, skip validation and return df_features
@@ -251,33 +244,57 @@ def _validate_df_specifics(self, feature_group, df):
         is_pk_null = False
         is_string_length_exceeded = False
 
-        # Check for null values in primary key columns
-        for pk in feature_group.primary_key:
-            if df.filter(df[pk].isNull()).count() > 0:
+        # Collect string columns and primary key columns
+        string_cols = [
+            f.name for f in df.schema.fields if isinstance(f.dataType, StringType)
+        ]
+        pk_cols = list(feature_group.primary_key) if feature_group.primary_key else []
+
+        # Build a single aggregation expression list:
+        # - sum(when(col.isNull(),1).otherwise(0)) for pk null counts
+        # - max(length(col)) for string max lengths
+        agg_exprs = []
+        for pk in pk_cols:
+            agg_exprs.append(
+                sf.sum(sf.when(sf.col(pk).isNull(), 1).otherwise(0)).alias(
+                    f"__nulls_{pk}"
+                )
+            )
+        for col in string_cols:
+            agg_exprs.append(sf.max(sf.length(sf.col(col))).alias(f"__maxlen_{col}"))
+
+        # If there is nothing to compute, return early
+        if not agg_exprs:
+            return errors, column_lengths, is_pk_null, is_string_length_exceeded
+
+        # Execute a single aggregation job
+        agg_row = df.agg(*agg_exprs).collect()[0]
+
+        # Evaluate primary key nulls
+        for pk in pk_cols:
+            null_count = agg_row[f"__nulls_{pk}"]
+            if null_count is not None and int(null_count) > 0:
                 errors[pk] = f"Primary key column {pk} contains null values."
                 is_pk_null = True
 
-        # Check string lengths for string columns
-        for field in df.schema.fields:
-            if isinstance(field.dataType, StringType):
-                col = field.name
-                # Compute max length - PySpark specific way
-                currentmax_row = df.select(sf.max(sf.length(col))).collect()[0][0]
-                currentmax = 0 if currentmax_row is None else currentmax_row
-
-                col_max_len = (
-                    self.get_online_varchar_length(
-                        self.get_feature_from_list(col, feature_group.features)
-                    )
-                    if feature_group.features
-                    else 100
+        # Evaluate string length violations
+        for col in string_cols:
+            currentmax = agg_row[f"__maxlen_{col}"]
+            currentmax = 0 if currentmax is None else int(currentmax)
+
+            col_max_len = (
+                self.get_online_varchar_length(
+                    self.get_feature_from_list(col, feature_group.features)
                 )
+                if feature_group.features
+                else 100
+            )
 
-                if col_max_len is not None and currentmax > col_max_len:
-                    errors[col] = (
-                        f"String length exceeded. Column {col} has string values longer than maximum column limit of {col_max_len} characters."
-                    )
-                    column_lengths[col] = currentmax
-                    is_string_length_exceeded = True
+            if col_max_len is not None and currentmax > col_max_len:
+                errors[col] = (
+                    f"String length exceeded. Column {col} has string values longer than maximum column limit of {col_max_len} characters."
+                )
+                column_lengths[col] = currentmax
+                is_string_length_exceeded = True
 
         return errors, column_lengths, is_pk_null, is_string_length_exceeded

python/hsfs/feature_group.py

@@ -2911,7 +2911,7 @@ def save(
                 * key `run_validation` boolean value, set to `False` to skip validation temporarily on ingestion.
                 * key `save_report` boolean value, set to `False` to skip upload of the validation report to Hopsworks.
                 * key `ge_validate_kwargs` a dictionary containing kwargs for the validate method of Great Expectations.
-                * key `online_schema_validation` boolean value, set to `True` to validate the schema for online ingestion.
+                * key `schema_validation` boolean value, set to `True` to validate the schema.
             wait: Wait for job and online ingestion to finish before returning, defaults to `False`.
                 Shortcut for write_options `{"wait_for_job": False, "wait_for_online_ingestion": False}`.
 
@@ -3120,7 +3120,7 @@ def insert(
                 * key `ge_validate_kwargs` a dictionary containing kwargs for the validate method of Great Expectations.
                 * key `fetch_expectation_suite` a boolean value, by default `True`, to control whether the expectation
                    suite of the feature group should be fetched before every insert.
-                * key `online_schema_validation` boolean value, set to `True` to validate the schema for online ingestion.
+                * key `schema_validation` boolean value, set to `True` to validate the schema.
             wait: Wait for job and online ingestion to finish before returning, defaults to `False`.
                 Shortcut for write_options `{"wait_for_job": False, "wait_for_online_ingestion": False}`.
             transformation_context: `Dict[str, Any]` A dictionary mapping variable names to objects that will be provided as contextual information to the transformation function at runtime.

python/tests/core/test_feature_group_engine.py

@@ -148,7 +148,24 @@ def test_save_ge_report(self, mocker):
         # Assert
         assert mock_engine_get_instance.return_value.save_dataframe.call_count == 0
 
-    def test_insert(self, mocker):
+    @pytest.mark.parametrize(
+        "online_enabled,validation_options,should_validate_schema",
+        [
+            # Online enabled
+            (True, None, True),
+            (True, {}, True),
+            (True, {"online_schema_validation": False}, False),
+            (True, {"schema_validation": False}, False),
+            # Not enabled
+            (False, None, True),
+            (False, {}, True),
+            (False, {"online_schema_validation": False}, False),
+            (False, {"schema_validation": False}, False),
+        ],
+    )
+    def test_insert(
+        self, online_enabled, validation_options, should_validate_schema, mocker
+    ):
         # Arrange
         feature_store_id = 99
 
@@ -162,6 +179,9 @@ def test_insert(self, mocker):
         )
         mocker.patch("hsfs.core.great_expectation_engine.GreatExpectationEngine")
         mock_fg_api = mocker.patch("hsfs.core.feature_group_api.FeatureGroupApi")
+        mock_validate_schema = mocker.patch(
+            "hsfs.core.schema_validation.DataFrameValidator.validate_schema"
+        )
 
         fg_engine = feature_group_engine.FeatureGroupEngine(
             feature_store_id=feature_store_id
@@ -174,6 +194,7 @@ def test_insert(self, mocker):
             primary_key=[],
             foreign_key=[],
             partition_key=[],
+            online_enabled=online_enabled,
         )
 
         # Act
@@ -184,11 +205,59 @@ def test_insert(self, mocker):
             operation=None,
             storage=None,
             write_options=None,
+            validation_options=validation_options,
         )
 
         # Assert
-        assert mock_fg_api.return_value.delete_content.call_count == 0
         assert mock_engine_get_instance.return_value.save_dataframe.call_count == 1
+        assert mock_fg_api.return_value.delete_content.call_count == 0
+        assert mock_validate_schema.called == should_validate_schema
+
+    def test_insert_storage(self, mocker):
+        # Arrange
+        feature_store_id = 99
+
+        mocker.patch("hsfs.engine.get_type")
+        mock_engine_get_instance = mocker.patch("hsfs.engine.get_instance")
+        mocker.patch(
+            "hsfs.core.feature_group_engine.FeatureGroupEngine.save_feature_group_metadata"
+        )
+        mocker.patch(
+            "hsfs.core.feature_group_engine.FeatureGroupEngine._verify_schema_compatibility"
+        )
+        mocker.patch("hsfs.core.great_expectation_engine.GreatExpectationEngine")
+        mock_fg_api = mocker.patch("hsfs.core.feature_group_api.FeatureGroupApi")
+
+        fg_engine = feature_group_engine.FeatureGroupEngine(
+            feature_store_id=feature_store_id
+        )
+
+        fg = feature_group.FeatureGroup(
+            name="test",
+            version=1,
+            featurestore_id=feature_store_id,
+            primary_key=[],
+            foreign_key=[],
+            partition_key=[],
+        )
+
+        # Act
+        with pytest.raises(exceptions.FeatureStoreException) as e_info:
+            fg_engine.insert(
+                feature_group=fg,
+                feature_dataframe=None,
+                overwrite=None,
+                operation=None,
+                storage="online",
+                write_options=None,
+            )
+
+        # Assert
+        assert mock_fg_api.return_value.delete_content.call_count == 0
+        assert mock_engine_get_instance.return_value.save_dataframe.call_count == 0
+        assert (
+            str(e_info.value) == "Online storage is not enabled for this feature group."
+        )
 
     def test_insert_transformation_functions(self, mocker):
         # Arrange
@@ -343,52 +412,6 @@ def test_insert_ge_report(self, mocker):
         assert mock_fg_api.return_value.delete_content.call_count == 0
         assert mock_engine_get_instance.return_value.save_dataframe.call_count == 0
 
-    def test_insert_storage(self, mocker):
-        # Arrange
-        feature_store_id = 99
-
-        mocker.patch("hsfs.engine.get_type")
-        mock_engine_get_instance = mocker.patch("hsfs.engine.get_instance")
-        mocker.patch(
-            "hsfs.core.feature_group_engine.FeatureGroupEngine.save_feature_group_metadata"
-        )
-        mocker.patch(
-            "hsfs.core.feature_group_engine.FeatureGroupEngine._verify_schema_compatibility"
-        )
-        mocker.patch("hsfs.core.great_expectation_engine.GreatExpectationEngine")
-        mock_fg_api = mocker.patch("hsfs.core.feature_group_api.FeatureGroupApi")
-
-        fg_engine = feature_group_engine.FeatureGroupEngine(
-            feature_store_id=feature_store_id
-        )
-
-        fg = feature_group.FeatureGroup(
-            name="test",
-            version=1,
-            featurestore_id=feature_store_id,
-            primary_key=[],
-            foreign_key=[],
-            partition_key=[],
-        )
-
-        # Act
-        with pytest.raises(exceptions.FeatureStoreException) as e_info:
-            fg_engine.insert(
-                feature_group=fg,
-                feature_dataframe=None,
-                overwrite=None,
-                operation=None,
-                storage="online",
-                write_options=None,
-            )
-
-        # Assert
-        assert mock_fg_api.return_value.delete_content.call_count == 0
-        assert mock_engine_get_instance.return_value.save_dataframe.call_count == 0
-        assert (
-            str(e_info.value) == "Online storage is not enabled for this feature group."
-        )
-
     def test_insert_overwrite(self, mocker):
         # Arrange
         feature_store_id = 99