Update featureset spec for featureset source (Azure#34668)

* Update featureset spec for featureset source

* Update source_matadata_schema

* Correct typo

* Move checks to constructor and add more unit tests

* Fix pylint errors

---------

Co-authored-by: Anya Li <[email protected]>

Author: anyalee0221

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/_feature_set/featureset_spec_metadata_schema.py

@@ -22,7 +22,7 @@ class FeaturesetSpecMetadataSchema(YamlFileSchema):
     source = fields.Nested(SourceMetadataSchema, required=True)
     feature_transformation_code = fields.Nested(FeatureTransformationCodeMetadataSchema, required=False)
     features = fields.List(NestedField(FeatureSchema), required=True, allow_none=False)
-    index_columns = fields.List(NestedField(DataColumnSchema), required=True, allow_none=False)
+    index_columns = fields.List(NestedField(DataColumnSchema), required=False)
     source_lookback = fields.Nested(DelayMetadataSchema, required=False)
     temporal_join_lookback = fields.Nested(DelayMetadataSchema, required=False)
 

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/_feature_set/source_metadata_schema.py

@@ -18,7 +18,7 @@
 class SourceMetadataSchema(metaclass=PatchedSchemaMeta):
     type = fields.Str(required=True)
     path = fields.Str(required=False)
-    timestamp_column = fields.Nested(TimestampColumnMetadataSchema, required=True)
+    timestamp_column = fields.Nested(TimestampColumnMetadataSchema, required=False)
     source_delay = fields.Nested(DelayMetadataSchema, required=False)
     source_process_code = fields.Nested(SourceProcessCodeSchema, load_only=True, required=False)
     dict = fields.Dict(keys=fields.Str(), values=fields.Str(), data_key="kwargs", load_only=True, required=False)

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_feature_set/featureset_spec_metadata.py

@@ -15,6 +15,7 @@
 from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY
 from azure.ai.ml.entities._feature_store_entity.data_column import DataColumn
 from azure.ai.ml.entities._util import load_from_dict
+from azure.ai.ml.exceptions import ErrorCategory, ErrorTarget, ValidationErrorType, ValidationException
 
 from .delay_metadata import DelayMetadata
 from .feature import Feature
@@ -31,11 +32,29 @@ def __init__(
         source: SourceMetadata,
         feature_transformation_code: Optional[FeatureTransformationCodeMetadata] = None,
         features: List[Feature],
-        index_columns: List[DataColumn],
+        index_columns: Optional[List[DataColumn]] = None,
         source_lookback: Optional[DelayMetadata] = None,
         temporal_join_lookback: Optional[DelayMetadata] = None,
         **_kwargs: Any,
     ):
+        if source.type == "featureset" and index_columns:
+            msg = f"You cannot provide index_columns for {source.type} feature source."
+            raise ValidationException(
+                message=msg,
+                no_personal_data_message=msg,
+                error_type=ValidationErrorType.INVALID_VALUE,
+                target=ErrorTarget.FEATURE_SET,
+                error_category=ErrorCategory.USER_ERROR,
+            )
+        if not index_columns and source.type != "featureset":
+            msg = f"You need to provide index_columns for {source.type} feature source."
+            raise ValidationException(
+                message=msg,
+                no_personal_data_message=msg,
+                error_type=ValidationErrorType.INVALID_VALUE,
+                target=ErrorTarget.FEATURE_SET,
+                error_category=ErrorCategory.USER_ERROR,
+            )
         self.source = source
         self.feature_transformation_code = feature_transformation_code
         self.features = features
@@ -74,6 +93,7 @@ def _load(
         res: FeaturesetSpecMetadata = load_from_dict(
             FeaturesetSpecMetadataSchema, yaml_data, context, "", unknown=INCLUDE, **kwargs
         )
+
         return res
 
     def _to_dict(self) -> Dict:

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_feature_set/source_metadata.py

@@ -18,14 +18,44 @@ def __init__(
         self,
         *,
         type: str,
-        timestamp_column: TimestampColumnMetadata,
+        timestamp_column: Optional[TimestampColumnMetadata] = None,
         path: Optional[str] = None,
         source_delay: Optional[DelayMetadata] = None,
         source_process_code: Optional[SourceProcessCodeMetadata] = None,
         dict: Optional[Dict] = None,
         **kwargs: Any,
     ):
+        if type != "featureset":
+            if not timestamp_column:
+                msg = f"You need to provide timestamp_column for {type} feature source."
+                raise ValidationException(
+                    message=msg,
+                    no_personal_data_message=msg,
+                    error_type=ValidationErrorType.INVALID_VALUE,
+                    target=ErrorTarget.FEATURE_SET,
+                    error_category=ErrorCategory.USER_ERROR,
+                )
+
         if type != "custom":
+            if type == "featureset":
+                if not path:
+                    msg = f"You need to provide path for {type} feature source."
+                    raise ValidationException(
+                        message=msg,
+                        no_personal_data_message=msg,
+                        error_type=ValidationErrorType.INVALID_VALUE,
+                        target=ErrorTarget.FEATURE_SET,
+                        error_category=ErrorCategory.USER_ERROR,
+                    )
+                if timestamp_column:
+                    msg = f"Cannot provide timestamp_column for {type} feature source."
+                    raise ValidationException(
+                        message=msg,
+                        no_personal_data_message=msg,
+                        error_type=ValidationErrorType.INVALID_VALUE,
+                        target=ErrorTarget.FEATURE_SET,
+                        error_category=ErrorCategory.USER_ERROR,
+                    )
             if not (path and not dict and not source_process_code):
                 msg = f"Cannot provide source_process_code or kwargs for {type} feature source."
                 raise ValidationException(

sdk/ml/azure-ai-ml/tests/feature_set/unittests/test_feature_set_spec_schema.py

@@ -36,6 +36,19 @@ def test_feature_set_spec_load(self) -> None:
         assert len(fspec.features) == 3
         assert len(fspec.source.kwargs.keys()) == 3
 
+        spec_path = "./tests/test_configs/feature_set/featureset_source_spec"
+        featureset_spec_contents = read_feature_set_metadata(path=spec_path)
+        featureset_spec_yaml_path = Path(spec_path, "FeatureSetSpec.yaml")
+        fspec = FeaturesetSpecMetadata._load(featureset_spec_contents, featureset_spec_yaml_path)
+
+        assert fspec.feature_transformation_code is not None
+        assert fspec.source is not None
+        assert fspec.source.timestamp_column is None
+        assert len(fspec.features) == 3
+        assert fspec.index_columns is None
+        assert fspec.source.source_delay is None
+        assert fspec.source.timestamp_column is None
+
     def test_feature_set_spec_load_failure(self) -> None:
         spec_path = "./tests/test_configs/feature_set/invalid_spec1"
         featureset_spec_contents = read_feature_set_metadata(path=spec_path)
@@ -48,3 +61,33 @@ def test_feature_set_spec_load_failure(self) -> None:
         featureset_spec_yaml_path = Path(spec_path, "FeatureSetSpec.yaml")
         with pytest.raises(ValidationException):
             FeaturesetSpecMetadata._load(featureset_spec_contents, featureset_spec_yaml_path)
+
+        spec_path = "./tests/test_configs/feature_set/invalid_spec3"
+        featureset_spec_contents = read_feature_set_metadata(path=spec_path)
+        featureset_spec_yaml_path = Path(spec_path, "FeatureSetSpec.yaml")
+        with pytest.raises(ValidationException):
+            FeaturesetSpecMetadata._load(featureset_spec_contents, featureset_spec_yaml_path)
+
+        spec_path = "./tests/test_configs/feature_set/invalid_spec4"
+        featureset_spec_contents = read_feature_set_metadata(path=spec_path)
+        featureset_spec_yaml_path = Path(spec_path, "FeatureSetSpec.yaml")
+        with pytest.raises(ValidationException):
+            FeaturesetSpecMetadata._load(featureset_spec_contents, featureset_spec_yaml_path)
+
+        spec_path = "./tests/test_configs/feature_set/invalid_spec5"
+        featureset_spec_contents = read_feature_set_metadata(path=spec_path)
+        featureset_spec_yaml_path = Path(spec_path, "FeatureSetSpec.yaml")
+        with pytest.raises(ValidationException):
+            FeaturesetSpecMetadata._load(featureset_spec_contents, featureset_spec_yaml_path)
+
+        spec_path = "./tests/test_configs/feature_set/invalid_spec6"
+        featureset_spec_contents = read_feature_set_metadata(path=spec_path)
+        featureset_spec_yaml_path = Path(spec_path, "FeatureSetSpec.yaml")
+        with pytest.raises(ValidationException):
+            FeaturesetSpecMetadata._load(featureset_spec_contents, featureset_spec_yaml_path)
+
+        spec_path = "./tests/test_configs/feature_set/invalid_spec7"
+        featureset_spec_contents = read_feature_set_metadata(path=spec_path)
+        featureset_spec_yaml_path = Path(spec_path, "FeatureSetSpec.yaml")
+        with pytest.raises(ValidationException):
+            FeaturesetSpecMetadata._load(featureset_spec_contents, featureset_spec_yaml_path)

sdk/ml/azure-ai-ml/tests/test_configs/feature_set/featureset_source_spec/FeatureSetSpec.yaml

@@ -0,0 +1,21 @@
+source:
+  type: featureset
+  path: azureml://subscriptions/my_sub/resourcegroups/my_rg/workspaces/my_fs/feturesets/source_feature_set_name/versions/version1
+feature_transformation_code:
+  path: ./code
+  transformer_class: driver_hourly_transform.DriverHourlyTransformer
+features:
+  - name: conv_rate
+    type: double
+  - name: acc_rate
+    type: double
+  - name: avg_daily_trips
+    type: double
+source_lookback:
+    days: 30
+    hours: 0
+    minutes: 0
+temporal_join_lookback:
+    days: 2
+    hours: 0
+    minutes: 0

sdk/ml/azure-ai-ml/tests/test_configs/feature_set/invalid_spec3/FeatureSetSpec.yaml

@@ -0,0 +1,22 @@
+# No path for featureset feature source
+source:
+  type: featureset
+feature_transformation_code:
+  path: ./code
+  transformer_class: driver_hourly_transform.DriverHourlyTransformer
+features:
+  - name: conv_rate
+    type: double
+  - name: acc_rate
+    type: double
+  - name: avg_daily_trips
+    type: double
+source_lookback:
+    days: 30
+    hours: 0
+    minutes: 0
+temporal_join_lookback:
+    days: 2
+    hours: 0
+    minutes: 0
+

sdk/ml/azure-ai-ml/tests/test_configs/feature_set/invalid_spec4/FeatureSetSpec.yaml

@@ -0,0 +1,31 @@
+# No timestamp_column for custom feature source
+source:
+  type: custom
+  kwargs:
+    k1: v1
+    k2: v2
+    k3: v3
+  source_process_code:
+    path: ./source_process_code
+    process_class: source_process.MyDataSourceLoader
+feature_transformation_code:
+  path: ./code
+  transformer_class: driver_hourly_transform.DriverHourlyTransformer
+features:
+  - name: conv_rate
+    type: double
+  - name: acc_rate
+    type: double
+  - name: avg_daily_trips
+    type: double
+index_columns:
+  - name: driver_id
+    type: long
+source_lookback:
+    days: 30
+    hours: 0
+    minutes: 0
+temporal_join_lookback:
+    days: 2
+    hours: 0
+    minutes: 0

sdk/ml/azure-ai-ml/tests/test_configs/feature_set/invalid_spec5/FeatureSetSpec.yaml

@@ -0,0 +1,31 @@
+# No index_columns for custom feature source
+source:
+  type: custom
+  kwargs:
+    k1: v1
+    k2: v2
+    k3: v3
+  timestamp_column:
+    name: timestamp
+    format: "%Y-%m-%d %H:%M:%S"
+  source_process_code:
+    path: ./source_process_code
+    process_class: source_process.MyDataSourceLoader
+feature_transformation_code:
+  path: ./code
+  transformer_class: driver_hourly_transform.DriverHourlyTransformer
+features:
+  - name: conv_rate
+    type: double
+  - name: acc_rate
+    type: double
+  - name: avg_daily_trips
+    type: double
+source_lookback:
+    days: 30
+    hours: 0
+    minutes: 0
+temporal_join_lookback:
+    days: 2
+    hours: 0
+    minutes: 0

sdk/ml/azure-ai-ml/tests/test_configs/feature_set/invalid_spec6/FeatureSetSpec.yaml

@@ -0,0 +1,25 @@
+# timestamp_column is not supported for featureset feature source
+source:
+  type: featureset
+  path: azureml://subscriptions/my_sub/resourcegroups/my_rg/workspaces/my_fs/feturesets/source_feature_set_name/versions/version1
+  timestamp_column:
+    name: timestamp
+    format: "%Y-%m-%d %H:%M:%S"
+feature_transformation_code:
+  path: ./code
+  transformer_class: driver_hourly_transform.DriverHourlyTransformer
+features:
+  - name: conv_rate
+    type: double
+  - name: acc_rate
+    type: double
+  - name: avg_daily_trips
+    type: double
+source_lookback:
+    days: 30
+    hours: 0
+    minutes: 0
+temporal_join_lookback:
+    days: 2
+    hours: 0
+    minutes: 0