From 1cadb98d619fd868ddb5dacaea63cc57ed7a4188 Mon Sep 17 00:00:00 2001 From: Akshay Shah Date: Mon, 12 Oct 2026 12:04:53 -0500 Subject: [PATCH] fix(mlops): generate valid DatasetBuilder SQL for String event times Two v3 regressions in DatasetBuilder query generation: - Event time casts compared feature_type to FeatureTypeEnum.STRING, but FeatureDefinition is now a pydantic shape that stores feature_type as a plain str ('String'), so the comparison was always False and every event-time range filter and point-in-time join wrapped ISO-8601 String event times in from_unixtime(), which Athena rejects for varchar. Compare the normalized type value instead. - The deleted-records CTE joined event time conditions without a leading space, so with_event_time_range() without as_of() produced 'WHERE is_deletedAND ...'. (v2 started the clause with a newline.) --- .../mlops/feature_store/dataset_builder.py | 19 +++++-- .../feature_store/test_dataset_builder.py | 56 +++++++++++++++++++ 2 files changed, 71 insertions(+), 4 deletions(-) diff --git a/sagemaker-mlops/src/sagemaker/mlops/feature_store/dataset_builder.py b/sagemaker-mlops/src/sagemaker/mlops/feature_store/dataset_builder.py index 2805450998..02558bfee4 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/feature_store/dataset_builder.py +++ b/sagemaker-mlops/src/sagemaker/mlops/feature_store/dataset_builder.py @@ -43,6 +43,17 @@ } +def _is_string_feature_type(feature_type) -> bool: + """Whether a feature type (a FeatureTypeEnum or its plain str value) is String. + + FeatureDefinition is a pydantic shape, so feature_type is stored as a plain str even when + a FeatureTypeEnum is passed in, and never compares equal to FeatureTypeEnum.STRING. + """ + if isinstance(feature_type, FeatureTypeEnum): + feature_type = feature_type.value + return feature_type == FeatureTypeEnum.STRING.value + + class TableType(Enum): """Kind of table a dataset builder can read from.""" @@ -750,7 +761,7 @@ def _construct_deleted_query(self, fg: FeatureGroupToBeMerged, suffix: str) -> s conds = self._construct_event_time_conditions( f"origin_{suffix}", fg.event_time_identifier_feature ) - event_conds = "".join(f"AND {c}\n" for c in conds) + event_conds = "".join(f" AND {c}\n" for c in conds) return ( f"deleted_{suffix} AS (\n" @@ -803,7 +814,7 @@ def _construct_event_time_conditions( ) -> List[str]: cast_fn = ( "from_iso8601_timestamp" - if event_time_feature.feature_type == FeatureTypeEnum.STRING + if _is_string_feature_type(event_time_feature.feature_type) else "from_unixtime" ) conditions = [] @@ -833,12 +844,12 @@ def _construct_join_condition(self, fg: FeatureGroupToBeMerged, suffix: str) -> if self._point_in_time_accurate_join: base_cast = ( "from_iso8601_timestamp" - if self._event_time_identifier_feature_type == FeatureTypeEnum.STRING + if _is_string_feature_type(self._event_time_identifier_feature_type) else "from_unixtime" ) fg_cast = ( "from_iso8601_timestamp" - if fg.event_time_identifier_feature.feature_type == FeatureTypeEnum.STRING + if _is_string_feature_type(fg.event_time_identifier_feature.feature_type) else "from_unixtime" ) join += ( diff --git a/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_dataset_builder.py b/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_dataset_builder.py index e1361c3093..6102a67207 100644 --- a/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_dataset_builder.py +++ b/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_dataset_builder.py @@ -601,3 +601,59 @@ def test_to_csv_from_feature_group_skips_register_when_disabled( builder._to_csv_from_feature_group() mock_register.assert_not_called() + + +class TestDatasetBuilderStringEventTimeQuery: + """Query generation for feature groups whose event time feature is a String.""" + + @staticmethod + def _feature_group(name): + from sagemaker.core.shapes import ( + DataCatalogConfig, + FeatureDefinition as FeatureDefinitionShape, + OfflineStoreConfig, + S3StorageConfig, + ) + + return FeatureGroup.model_construct( + feature_group_name=name, + record_identifier_feature_name="id", + event_time_feature_name="et", + feature_definitions=[ + FeatureDefinitionShape(feature_name="id", feature_type="String"), + FeatureDefinitionShape(feature_name="et", feature_type="String"), + FeatureDefinitionShape(feature_name="v", feature_type="Fractional"), + ], + offline_store_config=OfflineStoreConfig( + s3_storage_config=S3StorageConfig(s3_uri="s3://bucket/prefix"), + data_catalog_config=DataCatalogConfig( + table_name=f"{name}_table", catalog="AwsDataCatalog", database="db" + ), + ), + ) + + def test_event_time_range_and_point_in_time_join_query(self): + base, other = self._feature_group("base"), self._feature_group("other") + + with patch( + "sagemaker.mlops.feature_store.dataset_builder.FeatureGroup.get", + side_effect=lambda feature_group_name: self._feature_group(feature_group_name), + ): + builder = DatasetBuilder.create( + base=base, output_path="s3://bucket/out", session=Mock() + ) + builder.with_feature_group(other).point_in_time_accurate_join().with_event_time_range( + datetime.datetime(2024, 1, 1), datetime.datetime(2024, 2, 1) + ) + with patch.object( + DatasetBuilder, "_run_query", return_value={"QueryExecution": {"Query": "q"}} + ) as mock_run_query: + builder._to_csv_from_feature_group() + + query = mock_run_query.call_args[0][0] + # String (ISO-8601) event times must be parsed, not treated as unix timestamps. + assert 'from_iso8601_timestamp(origin_base."et")' in query + assert 'from_iso8601_timestamp(fg_0."et")' in query + assert 'from_unixtime(origin_base."et")' not in query + # The deleted-records CTE needs a separator before the event time conditions. + assert "is_deletedAND" not in query