updates with_feature-group def and adds test

Author: patrickmcarlos

src/sagemaker/feature_store/dataset_builder.py

@@ -43,6 +43,32 @@ class TableType(Enum):
     DATA_FRAME = "DataFrame"
 
 
+@attr.s
+class JoinTypeEnum(Enum):
+    """Enum of Join types.
+    The Join comparator can be "INNER_JOIN", "LEFT_JOIN", "RIGHT_JOIN", "FULL_JOIN"
+    """
+
+    INNER_JOIN = "JOIN"
+    LEFT_JOIN = "LEFT JOIN"
+    RIGHT_JOIN = "RIGHT JOIN"
+    FULL_JOIN = "FULL JOIN"
+
+
+@attr.s
+class JoinComparatorEnum(Enum):
+    """Enum of Join comparators.
+    The Join comparator can be "EQUALS", "GREATER_THAN", "LESS_THAN",
+    "GREATER_THAN_OR_EQUAL_TO", or "LESS_THAN_OR_EQUAL_TO"
+    """
+
+    EQUALS = "="
+    GREATER_THAN = ">"
+    GREATER_THAN_OR_EQUAL_TO = ">="
+    LESS_THAN = "<"
+    LESS_THAN_OR_EQUAL_TO = "<="
+
+
 @attr.s
 class FeatureGroupToBeMerged:
     """FeatureGroup metadata which will be used for SQL join.
@@ -68,6 +94,13 @@ class FeatureGroupToBeMerged:
             be used as target join key (default: None).
         table_type (TableType): A TableType representing the type of table if it is Feature Group or
             Panda Data Frame (default: None).
+        feature_name_in_target (str): A string representing the feature in the target feature group
+            that will be compared to the target feature in the base feature group
+        join_comparator (JoinComparatorEnum): A JoinComparatorEnum representing the comparator used
+            when joining the target feature in the base feature group and the feature in the target
+            feature group (default: None).
+        join_type (JoinTypeEnum): A JoinTypeEnum representing the type of join between the base and
+            target feature groups. (default: None).
     """
 
     features: List[str] = attr.ib()
@@ -80,12 +113,18 @@ class FeatureGroupToBeMerged:
     event_time_identifier_feature: FeatureDefinition = attr.ib()
     target_feature_name_in_base: str = attr.ib(default=None)
     table_type: TableType = attr.ib(default=None)
+    feature_name_in_target: str = attr.ib(default=None)
+    join_comparator: JoinComparatorEnum = attr.ib(default=None)
+    join_type: JoinTypeEnum = attr.ib(default=None)
 
 
 def construct_feature_group_to_be_merged(
-    feature_group: FeatureGroup,
+    target_feature_group: FeatureGroup,
     included_feature_names: List[str],
     target_feature_name_in_base: str = None,
+    feature_name_in_target: str = None,
+    join_comparator: JoinComparatorEnum = None,
+    join_type: JoinTypeEnum = None
 ) -> FeatureGroupToBeMerged:
     """Construct a FeatureGroupToBeMerged object by provided parameters.
 
@@ -101,12 +140,12 @@ def construct_feature_group_to_be_merged(
     Raises:
         ValueError: Invalid feature name(s) in included_feature_names.
     """
-    feature_group_metadata = feature_group.describe()
+    feature_group_metadata = target_feature_group.describe()
     data_catalog_config = feature_group_metadata.get("OfflineStoreConfig", {}).get(
         "DataCatalogConfig", None
     )
     if not data_catalog_config:
-        raise RuntimeError(f"No metastore is configured with FeatureGroup {feature_group.name}.")
+        raise RuntimeError(f"No metastore is configured with FeatureGroup {target_feature_group.name}.")
 
     record_identifier_feature_name = feature_group_metadata.get("RecordIdentifierFeatureName", None)
     feature_definitions = feature_group_metadata.get("FeatureDefinitions", [])
@@ -126,10 +165,15 @@ def construct_feature_group_to_be_merged(
     catalog = data_catalog_config.get("Catalog", None) if disable_glue else _DEFAULT_CATALOG
     features = [feature.get("FeatureName", None) for feature in feature_definitions]
 
+    if (feature_name_in_target is not None and feature_name_in_target not in features):
+        raise ValueError(
+            f"Feature {feature_name_in_target} not found in FeatureGroup {target_feature_group.name}"
+        )
+
     for included_feature in included_feature_names or []:
         if included_feature not in features:
             raise ValueError(
-                f"Feature {included_feature} not found in FeatureGroup {feature_group.name}"
+                f"Feature {included_feature} not found in FeatureGroup {target_feature_group.name}"
             )
     if not included_feature_names:
         included_feature_names = features
@@ -151,6 +195,9 @@ def construct_feature_group_to_be_merged(
         FeatureDefinition(event_time_identifier_feature_name, event_time_identifier_feature_type),
         target_feature_name_in_base,
         TableType.FEATURE_GROUP,
+        feature_name_in_target,
+        join_comparator,
+        join_type
     )
 
 
@@ -227,6 +274,9 @@ def with_feature_group(
         feature_group: FeatureGroup,
         target_feature_name_in_base: str = None,
         included_feature_names: List[str] = None,
+        feature_name_in_target: str = None,
+        join_comparator: JoinComparatorEnum = None,
+        join_type: JoinTypeEnum = None
     ):
         """Join FeatureGroup with base.
 
@@ -241,7 +291,11 @@ def with_feature_group(
         """
         self._feature_groups_to_be_merged.append(
             construct_feature_group_to_be_merged(
-                feature_group, included_feature_names, target_feature_name_in_base
+                feature_group, included_feature_names,
+                target_feature_name_in_base,
+                feature_name_in_target,
+                join_comparator,
+                join_type
             )
         )
         return self
@@ -905,10 +959,22 @@ def _construct_join_condition(self, feature_group: FeatureGroupToBeMerged, suffi
         Returns:
             The JOIN query string.
         """
+
+        join_type = (feature_group.join_type if feature_group.join_type is not None
+                     else JoinTypeEnum.INNER_JOIN)
+
+        join_comparator = (feature_group.join_comparator
+                           if feature_group.join_comparator is not None
+                           else JoinComparatorEnum.EQUALS)
+
+        feature_name_in_target = (feature_group.feature_name_in_target
+                                  if feature_group.feature_name_in_target is not None
+                                  else feature_group.record_identifier_feature_name)
+
         join_condition_string = (
-            f"\nJOIN fg_{suffix}\n"
-            + f'ON fg_base."{feature_group.target_feature_name_in_base}" = '
-            + f'fg_{suffix}."{feature_group.record_identifier_feature_name}"'
+            f"\n{join_type.value} fg_{suffix}\n"
+            + f'ON fg_base."{feature_group.target_feature_name_in_base}" {join_comparator.value} '
+            + f'fg_{suffix}."{feature_name_in_target}"'
         )
         base_timestamp_cast_function_name = "from_unixtime"
         if self._event_time_identifier_feature_type == FeatureTypeEnum.STRING:

tests/unit/sagemaker/feature_store/test_dataset_builder.py

@@ -23,11 +23,13 @@
     DatasetBuilder,
     FeatureGroupToBeMerged,
     TableType,
+    JoinComparatorEnum,
+    JoinTypeEnum
 )
 from sagemaker.feature_store.feature_group import (
     FeatureDefinition,
     FeatureGroup,
-    FeatureTypeEnum,
+    FeatureTypeEnum
 )
 
 
@@ -528,6 +530,149 @@ def test_construct_query_string(sagemaker_session_mock):
         + ")\n"
     )
 
+# Tests the optional feature_name_in_target, join_comparator and join_type parameters
+def test_with_feature_group_with_optional_params_query_string(sagemaker_session_mock):
+    base_feature_group = FeatureGroup(name="base_feature_group", sagemaker_session=sagemaker_session_mock)
+    target_feature_group = FeatureGroup(name="target_feature_group", sagemaker_session=sagemaker_session_mock)
+
+    dataset_builder = DatasetBuilder(
+        sagemaker_session=sagemaker_session_mock,
+        base=base_feature_group,
+        output_path="file/to/path",
+        record_identifier_feature_name="target-feature",
+    )
+
+    dataset_builder._event_time_identifier_feature_name = "target-feature"
+
+    sagemaker_session_mock.describe_feature_group.return_value = {
+        "OfflineStoreConfig": {"DataCatalogConfig": {"TableName": "table-name", "Database": "database"}},
+        "RecordIdentifierFeatureName": "feature-1",
+        "EventTimeFeatureName": "feature-2",
+        "FeatureDefinitions": [
+            {"FeatureName": "feature-1", "FeatureType": "String"},
+            {"FeatureName": "feature-2", "FeatureType": "String"},
+        ],
+    }
+
+    dataset_builder.with_feature_group(target_feature_group, 
+        "target-feature", 
+        ["feature-1", "feature-2"], 
+        "feature-2", 
+        JoinComparatorEnum.GREATER_THAN, 
+        JoinTypeEnum.FULL_JOIN)
+
+    query_string = dataset_builder._construct_query_string(BASE)
+
+    assert (
+        query_string
+        == "WITH fg_base AS (WITH table_base AS (\n"
+        + "SELECT *\n"
+        + "FROM (\n"
+        + "SELECT *, row_number() OVER (\n"
+        + 'PARTITION BY origin_base."target-feature", origin_base."other-feature"\n'
+        + 'ORDER BY origin_base."api_invocation_time" DESC, origin_base."write_time" DESC\n'
+        + ") AS dedup_row_base\n"
+        + 'FROM "database"."base-table" origin_base\n'
+        + ")\n"
+        + "WHERE dedup_row_base = 1\n"
+        + "),\n"
+        + "deleted_base AS (\n"
+        + "SELECT *\n"
+        + "FROM (\n"
+        + "SELECT *, row_number() OVER (\n"
+        + 'PARTITION BY origin_base."target-feature"\n'
+        + 'ORDER BY origin_base."other-feature" DESC, origin_base."api_invocation_time" '
+        + 'DESC, origin_base."write_time" DESC\n'
+        + ") AS deleted_row_base\n"
+        + 'FROM "database"."base-table" origin_base\n'
+        + "WHERE is_deleted\n"
+        + ")\n"
+        + "WHERE deleted_row_base = 1\n"
+        + ")\n"
+        + 'SELECT table_base."target-feature", table_base."other-feature"\n'
+        + "FROM (\n"
+        + 'SELECT table_base."target-feature", table_base."other-feature", '
+        + 'table_base."write_time"\n'
+        + "FROM table_base\n"
+        + "LEFT JOIN deleted_base\n"
+        + 'ON table_base."target-feature" = deleted_base."target-feature"\n'
+        + 'WHERE deleted_base."target-feature" IS NULL\n'
+        + "UNION ALL\n"
+        + 'SELECT table_base."target-feature", table_base."other-feature", '
+        + 'table_base."write_time"\n'
+        + "FROM deleted_base\n"
+        + "JOIN table_base\n"
+        + 'ON table_base."target-feature" = deleted_base."target-feature"\n'
+        + "AND (\n"
+        + 'table_base."other-feature" > deleted_base."other-feature"\n'
+        + 'OR (table_base."other-feature" = deleted_base."other-feature" AND '
+        + 'table_base."api_invocation_time" > deleted_base."api_invocation_time")\n'
+        + 'OR (table_base."other-feature" = deleted_base."other-feature" AND '
+        + 'table_base."api_invocation_time" = deleted_base."api_invocation_time" AND '
+        + 'table_base."write_time" > deleted_base."write_time")\n'
+        + ")\n"
+        + ") AS table_base\n"
+        + "),\n"
+        + "fg_0 AS (WITH table_0 AS (\n"
+        + "SELECT *\n"
+        + "FROM (\n"
+        + "SELECT *, row_number() OVER (\n"
+        + 'PARTITION BY origin_0."feature-1", origin_0."feature-2"\n'
+        + 'ORDER BY origin_0."api_invocation_time" DESC, origin_0."write_time" DESC\n'
+        + ") AS dedup_row_0\n"
+        + 'FROM "database"."table-name" origin_0\n'
+        + ")\n"
+        + "WHERE dedup_row_0 = 1\n"
+        + "),\n"
+        + "deleted_0 AS (\n"
+        + "SELECT *\n"
+        + "FROM (\n"
+        + "SELECT *, row_number() OVER (\n"
+        + 'PARTITION BY origin_0."feature-1"\n'
+        + 'ORDER BY origin_0."feature-2" DESC, origin_0."api_invocation_time" DESC, '
+        + 'origin_0."write_time" DESC\n'
+        + ") AS deleted_row_0\n"
+        + 'FROM "database"."table-name" origin_0\n'
+        + "WHERE is_deleted\n"
+        + ")\n"
+        + "WHERE deleted_row_0 = 1\n"
+        + ")\n"
+        + 'SELECT table_0."feature-1", table_0."feature-2"\n'
+        + "FROM (\n"
+        + 'SELECT table_0."feature-1", table_0."feature-2", table_0."write_time"\n'
+        + "FROM table_0\n"
+        + "LEFT JOIN deleted_0\n"
+        + 'ON table_0."feature-1" = deleted_0."feature-1"\n'
+        + 'WHERE deleted_0."feature-1" IS NULL\n'
+        + "UNION ALL\n"
+        + 'SELECT table_0."feature-1", table_0."feature-2", table_0."write_time"\n'
+        + "FROM deleted_0\n"
+        + "JOIN table_0\n"
+        + 'ON table_0."feature-1" = deleted_0."feature-1"\n'
+        + "AND (\n"
+        + 'table_0."feature-2" > deleted_0."feature-2"\n'
+        + 'OR (table_0."feature-2" = deleted_0."feature-2" AND '
+        + 'table_0."api_invocation_time" > deleted_0."api_invocation_time")\n'
+        + 'OR (table_0."feature-2" = deleted_0."feature-2" AND '
+        + 'table_0."api_invocation_time" = deleted_0."api_invocation_time" AND '
+        + 'table_0."write_time" > deleted_0."write_time")\n'
+        + ")\n"
+        + ") AS table_0\n"
+        + ")\n"
+        + 'SELECT target-feature, other-feature, "feature-1.1", "feature-2.1"\n'
+        + "FROM (\n"
+        + 'SELECT fg_base.target-feature, fg_base.other-feature, fg_0."feature-1" as '
+        + '"feature-1.1", fg_0."feature-2" as "feature-2.1", row_number() OVER (\n'
+        + 'PARTITION BY fg_base."target-feature"\n'
+        + 'ORDER BY fg_base."other-feature" DESC, fg_0."feature-2" DESC\n'
+        + ") AS row_recent\n"
+        + "FROM fg_base\n"
+        + "FULL JOIN fg_0\n"
+        + 'ON fg_base."target-feature" > fg_0."feature-2"\n'
+        + ")\n"
+    )
+
+
 
 def test_create_temp_table(sagemaker_session_mock):
     dataframe = pd.DataFrame({"feature-1": [420, 380, 390], "feature-2": [50, 40, 45]})