# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# Licensed under the Apache License, Version 2.0
"""Dataset Builder for FeatureStore."""
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Dict, List, Union
import datetime
import pandas as pd
from sagemaker.core.helper.session_helper import Session
from sagemaker.core.telemetry import Feature, _telemetry_emitter
from sagemaker.mlops.feature_store import FeatureGroup
from sagemaker.mlops.feature_store.feature_definition import FeatureDefinition, FeatureTypeEnum
from sagemaker.mlops.feature_store.feature_utils import (
upload_dataframe_to_s3,
download_csv_from_s3,
run_athena_query,
)
_DEFAULT_CATALOG = "AwsDataCatalog"
_DEFAULT_DATABASE = "sagemaker_featurestore"
_DTYPE_TO_FEATURE_TYPE = {
"object": "String", "string": "String",
"int64": "Integral", "int32": "Integral",
"float64": "Fractional", "float32": "Fractional",
}
_DTYPE_TO_ATHENA_TYPE = {
"object": "STRING", "int64": "INT", "float64": "DOUBLE",
"bool": "BOOLEAN", "datetime64[ns]": "TIMESTAMP",
}
[docs]
class TableType(Enum):
FEATURE_GROUP = "FeatureGroup"
DATA_FRAME = "DataFrame"
[docs]
class JoinTypeEnum(Enum):
INNER_JOIN = "JOIN"
LEFT_JOIN = "LEFT JOIN"
RIGHT_JOIN = "RIGHT JOIN"
FULL_JOIN = "FULL JOIN"
CROSS_JOIN = "CROSS JOIN"
[docs]
class JoinComparatorEnum(Enum):
EQUALS = "="
GREATER_THAN = ">"
GREATER_THAN_OR_EQUAL_TO = ">="
LESS_THAN = "<"
LESS_THAN_OR_EQUAL_TO = "<="
NOT_EQUAL_TO = "<>"
[docs]
@dataclass
class FeatureGroupToBeMerged:
"""FeatureGroup metadata which will be used for SQL join.
This class instantiates a FeatureGroupToBeMerged object that comprises a list of feature names,
a list of feature names which will be included in SQL query, a database, an Athena table name,
a feature name of record identifier, a feature name of event time identifier and a feature name
of base which is the target join key.
Attributes:
features (List[str]): A list of strings representing feature names of this FeatureGroup.
included_feature_names (List[str]): A list of strings representing features to be
included in the SQL join.
projected_feature_names (List[str]): A list of strings representing features to be
included for final projection in output.
catalog (str): A string representing the catalog.
database (str): A string representing the database.
table_name (str): A string representing the Athena table name of this FeatureGroup.
record_identifier_feature_name (str): A string representing the record identifier feature.
event_time_identifier_feature (FeatureDefinition): A FeatureDefinition representing the
event time identifier feature.
target_feature_name_in_base (str): A string representing the feature name in base which will
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 name in the target feature
group that will be compared to the target feature in the base feature group.
If None is provided, the record identifier feature will be used in the
SQL join. (default: None).
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: JoinComparatorEnum.EQUALS).
join_type (JoinTypeEnum): A JoinTypeEnum representing the type of join between
the base and target feature groups. (default: JoinTypeEnum.INNER_JOIN).
"""
features: List[str]
included_feature_names: List[str]
projected_feature_names: List[str]
catalog: str
database: str
table_name: str
record_identifier_feature_name: str
event_time_identifier_feature: FeatureDefinition
target_feature_name_in_base: str = None
table_type: TableType = None
feature_name_in_target: str = None
join_comparator: JoinComparatorEnum = JoinComparatorEnum.EQUALS
join_type: JoinTypeEnum = JoinTypeEnum.INNER_JOIN
[docs]
def construct_feature_group_to_be_merged(
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 = JoinComparatorEnum.EQUALS,
join_type: JoinTypeEnum = JoinTypeEnum.INNER_JOIN,
) -> FeatureGroupToBeMerged:
"""Construct a FeatureGroupToBeMerged object by provided parameters.
Args:
target_feature_group (FeatureGroup): A FeatureGroup object.
included_feature_names (List[str]): A list of strings representing features to be
included in the output.
target_feature_name_in_base (str): A string representing the feature name in base which
will be used as target join key (default: None).
feature_name_in_target (str): A string representing the feature name in the target feature
group that will be compared to the target feature in the base feature group.
If None is provided, the record identifier feature will be used in the
SQL join. (default: None).
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: JoinComparatorEnum.EQUALS).
join_type (JoinTypeEnum): A JoinTypeEnum representing the type of join between
the base and target feature groups. (default: JoinTypeEnum.INNER_JOIN).
Returns:
A FeatureGroupToBeMerged object.
Raises:
RuntimeError: No metastore is configured with the FeatureGroup.
ValueError: Invalid feature name(s) in included_feature_names.
"""
fg = FeatureGroup.get(feature_group_name=target_feature_group.feature_group_name)
if not fg.offline_store_config or not fg.offline_store_config.data_catalog_config:
raise RuntimeError(f"No metastore configured for FeatureGroup {fg.feature_group_name}.")
catalog_config = fg.offline_store_config.data_catalog_config
disable_glue = getattr(catalog_config, "disable_glue_table_creation", False) or False
features = [fd.feature_name for fd in fg.feature_definitions]
record_id = fg.record_identifier_feature_name
event_time_name = fg.event_time_feature_name
event_time_type = next(
(fd.feature_type for fd in fg.feature_definitions if fd.feature_name == event_time_name),
None
)
if feature_name_in_target and feature_name_in_target not in features:
raise ValueError(f"Feature {feature_name_in_target} not found in {fg.feature_group_name}")
for feat in included_feature_names or []:
if feat not in features:
raise ValueError(f"Feature {feat} not found in {fg.feature_group_name}")
if not included_feature_names:
included_feature_names = features.copy()
projected_feature_names = features.copy()
else:
projected_feature_names = included_feature_names.copy()
if record_id not in included_feature_names:
included_feature_names.append(record_id)
if event_time_name not in included_feature_names:
included_feature_names.append(event_time_name)
return FeatureGroupToBeMerged(
features=features,
included_feature_names=included_feature_names,
projected_feature_names=projected_feature_names,
catalog=catalog_config.catalog if disable_glue else _DEFAULT_CATALOG,
database=catalog_config.database,
table_name=catalog_config.table_name,
record_identifier_feature_name=record_id,
event_time_identifier_feature=FeatureDefinition(
feature_name=event_time_name, feature_type=FeatureTypeEnum(event_time_type).value
),
target_feature_name_in_base=target_feature_name_in_base,
table_type=TableType.FEATURE_GROUP,
feature_name_in_target=feature_name_in_target,
join_comparator=join_comparator,
join_type=join_type,
)
[docs]
@dataclass
class DatasetBuilder:
"""DatasetBuilder definition.
This class instantiates a DatasetBuilder object that comprises a base, a list of feature names,
an output path and a KMS key ID.
Attributes:
_sagemaker_session (Session): Session instance to perform boto calls.
_base (Union[FeatureGroup, DataFrame]): A base which can be either a FeatureGroup or a
pandas.DataFrame and will be used to merge other FeatureGroups and generate a Dataset.
_output_path (str): An S3 URI which stores the output .csv file.
_record_identifier_feature_name (str): A string representing the record identifier feature
if base is a DataFrame (default: None).
_event_time_identifier_feature_name (str): A string representing the event time identifier
feature if base is a DataFrame (default: None).
_included_feature_names (List[str]): A list of strings representing features to be
included in the output. If not set, all features will be included in the output.
(default: None).
_kms_key_id (str): A KMS key id. If set, will be used to encrypt the result file
(default: None).
_point_in_time_accurate_join (bool): A boolean representing if point-in-time join
is applied to the resulting dataframe when calling "to_dataframe".
When set to True, users can retrieve data using "row-level time travel"
according to the event times provided to the DatasetBuilder. This requires that the
entity dataframe with event times is submitted as the base in the constructor
(default: False).
_include_duplicated_records (bool): A boolean representing whether the resulting dataframe
when calling "to_dataframe" should include duplicated records (default: False).
_include_deleted_records (bool): A boolean representing whether the resulting
dataframe when calling "to_dataframe" should include deleted records (default: False).
_number_of_recent_records (int): An integer representing how many records will be
returned for each record identifier (default: 1).
_number_of_records (int): An integer representing the number of records that should be
returned in the resulting dataframe when calling "to_dataframe" (default: None).
_write_time_ending_timestamp (datetime.datetime): A datetime that represents the latest
write time for a record to be included in the resulting dataset. Records with a
newer write time will be omitted from the resulting dataset. (default: None).
_event_time_starting_timestamp (datetime.datetime): A datetime that represents the earliest
event time for a record to be included in the resulting dataset. Records
with an older event time will be omitted from the resulting dataset. (default: None).
_event_time_ending_timestamp (datetime.datetime): A datetime that represents the latest
event time for a record to be included in the resulting dataset. Records
with a newer event time will be omitted from the resulting dataset. (default: None).
_feature_groups_to_be_merged (List[FeatureGroupToBeMerged]): A list of
FeatureGroupToBeMerged which will be joined to base (default: []).
_event_time_identifier_feature_type (FeatureTypeEnum): A FeatureTypeEnum representing the
type of event time identifier feature (default: None).
"""
_sagemaker_session: Session
_base: Union[FeatureGroup, pd.DataFrame]
_output_path: str
_record_identifier_feature_name: str = None
_event_time_identifier_feature_name: str = None
_included_feature_names: List[str] = None
_kms_key_id: str = None
_event_time_identifier_feature_type: FeatureTypeEnum = None
_point_in_time_accurate_join: bool = field(default=False, init=False)
_include_duplicated_records: bool = field(default=False, init=False)
_include_deleted_records: bool = field(default=False, init=False)
_number_of_recent_records: int = field(default=None, init=False)
_number_of_records: int = field(default=None, init=False)
_write_time_ending_timestamp: datetime.datetime = field(default=None, init=False)
_event_time_starting_timestamp: datetime.datetime = field(default=None, init=False)
_event_time_ending_timestamp: datetime.datetime = field(default=None, init=False)
_feature_groups_to_be_merged: List[FeatureGroupToBeMerged] = field(default_factory=list, init=False)
[docs]
@classmethod
def create(
cls,
base: Union[FeatureGroup, pd.DataFrame],
output_path: str,
session: Session,
record_identifier_feature_name: str = None,
event_time_identifier_feature_name: str = None,
included_feature_names: List[str] = None,
kms_key_id: str = None,
) -> "DatasetBuilder":
"""Create a DatasetBuilder for generating a Dataset.
Args:
base: A FeatureGroup or DataFrame to use as the base.
output_path: S3 URI for output.
session: SageMaker session.
record_identifier_feature_name: Required if base is DataFrame.
event_time_identifier_feature_name: Required if base is DataFrame.
included_feature_names: Features to include in output.
kms_key_id: KMS key for encryption.
Returns:
DatasetBuilder instance.
"""
if isinstance(base, pd.DataFrame):
if not record_identifier_feature_name or not event_time_identifier_feature_name:
raise ValueError(
"record_identifier_feature_name and event_time_identifier_feature_name "
"are required when base is a DataFrame."
)
return cls(
_sagemaker_session=session,
_base=base,
_output_path=output_path,
_record_identifier_feature_name=record_identifier_feature_name,
_event_time_identifier_feature_name=event_time_identifier_feature_name,
_included_feature_names=included_feature_names,
_kms_key_id=kms_key_id,
)
[docs]
def with_feature_group(
self,
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 = JoinComparatorEnum.EQUALS,
join_type: JoinTypeEnum = JoinTypeEnum.INNER_JOIN,
) -> "DatasetBuilder":
"""Join FeatureGroup with base.
Args:
feature_group (FeatureGroup): A target FeatureGroup which will be joined to base.
target_feature_name_in_base (str): A string representing the feature name in base which
will be used as a join key (default: None).
included_feature_names (List[str]): A list of strings representing features to be
included in the output (default: None).
feature_name_in_target (str): A string representing the feature name in the target
feature group that will be compared to the target feature in the base feature group.
If None is provided, the record identifier feature will be used in the
SQL join. (default: None).
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: JoinComparatorEnum.EQUALS).
join_type (JoinTypeEnum): A JoinTypeEnum representing the type of join between
the base and target feature groups. (default: JoinTypeEnum.INNER_JOIN).
Returns:
This DatasetBuilder object.
"""
self._feature_groups_to_be_merged.append(
construct_feature_group_to_be_merged(
feature_group, included_feature_names, target_feature_name_in_base,
feature_name_in_target, join_comparator, join_type,
)
)
return self
[docs]
def point_in_time_accurate_join(self) -> "DatasetBuilder":
"""Enable point-in-time accurate join.
Returns:
This DatasetBuilder object.
"""
self._point_in_time_accurate_join = True
return self
[docs]
def include_duplicated_records(self) -> "DatasetBuilder":
"""Include duplicated records in dataset.
Returns:
This DatasetBuilder object.
"""
self._include_duplicated_records = True
return self
[docs]
def include_deleted_records(self) -> "DatasetBuilder":
"""Include deleted records in dataset.
Returns:
This DatasetBuilder object.
"""
self._include_deleted_records = True
return self
[docs]
def with_number_of_recent_records_by_record_identifier(self, n: int) -> "DatasetBuilder":
"""Set number_of_recent_records field with provided input.
Args:
n (int): An int that how many recent records will be returned for
each record identifier.
Returns:
This DatasetBuilder object.
"""
self._number_of_recent_records = n
return self
[docs]
def with_number_of_records_from_query_results(self, n: int) -> "DatasetBuilder":
"""Set number_of_records field with provided input.
Args:
n (int): An int that how many records will be returned.
Returns:
This DatasetBuilder object.
"""
self._number_of_records = n
return self
[docs]
def as_of(self, timestamp: datetime.datetime) -> "DatasetBuilder":
"""Set write_time_ending_timestamp field with provided input.
Args:
timestamp (datetime.datetime): A datetime that all records' write time in dataset will
be before it.
Returns:
This DatasetBuilder object.
"""
self._write_time_ending_timestamp = timestamp
return self
[docs]
def with_event_time_range(
self,
starting_timestamp: datetime.datetime = None,
ending_timestamp: datetime.datetime = None,
) -> "DatasetBuilder":
"""Set event_time_starting_timestamp and event_time_ending_timestamp with provided inputs.
Args:
starting_timestamp (datetime.datetime): A datetime that all records' event time in
dataset will be after it (default: None).
ending_timestamp (datetime.datetime): A datetime that all records' event time in dataset
will be before it (default: None).
Returns:
This DatasetBuilder object.
"""
self._event_time_starting_timestamp = starting_timestamp
self._event_time_ending_timestamp = ending_timestamp
return self
[docs]
@_telemetry_emitter(Feature.FEATURE_STORE, "DatasetBuilder.to_csv_file")
def to_csv_file(self) -> tuple[str, str]:
"""Get query string and result in .csv format file.
Returns:
tuple: A tuple containing:
- str: The S3 path of the .csv file
- str: The query string executed
Note:
This method returns a tuple (csv_path, query_string).
To get just the CSV path: csv_path, _ = builder.to_csv_file()
"""
if isinstance(self._base, pd.DataFrame):
return self._to_csv_from_dataframe()
if isinstance(self._base, FeatureGroup):
return self._to_csv_from_feature_group()
raise ValueError("Base must be either a FeatureGroup or a DataFrame.")
[docs]
@_telemetry_emitter(Feature.FEATURE_STORE, "DatasetBuilder.to_dataframe")
def to_dataframe(self) -> tuple[pd.DataFrame, str]:
"""Get query string and result in pandas.DataFrame.
Returns:
tuple: A tuple containing:
- pd.DataFrame: The pandas DataFrame object
- str: The query string executed
Note:
This method returns a tuple (dataframe, query_string).
To get just the DataFrame: df, _ = builder.to_dataframe()
"""
csv_file, query_string = self.to_csv_file()
df = download_csv_from_s3(csv_file, self._sagemaker_session, self._kms_key_id)
if "row_recent" in df.columns:
df = df.drop("row_recent", axis="columns")
return df, query_string
def _to_csv_from_dataframe(self) -> tuple[str, str]:
s3_folder, temp_table_name = upload_dataframe_to_s3(
self._base, self._output_path, self._sagemaker_session, self._kms_key_id
)
self._create_temp_table(temp_table_name, s3_folder)
base_features = list(self._base.columns)
event_time_dtype = str(self._base[self._event_time_identifier_feature_name].dtypes)
self._event_time_identifier_feature_type = FeatureTypeEnum(
_DTYPE_TO_FEATURE_TYPE.get(event_time_dtype, "String")
)
included = self._included_feature_names or base_features
fg_to_merge = FeatureGroupToBeMerged(
features=base_features,
included_feature_names=included,
projected_feature_names=included,
catalog=_DEFAULT_CATALOG,
database=_DEFAULT_DATABASE,
table_name=temp_table_name,
record_identifier_feature_name=self._record_identifier_feature_name,
event_time_identifier_feature=FeatureDefinition(
feature_name=self._event_time_identifier_feature_name,
feature_type=self._event_time_identifier_feature_type,
),
table_type=TableType.DATA_FRAME,
)
query_string = self._construct_query_string(fg_to_merge)
result = self._run_query(query_string, _DEFAULT_CATALOG, _DEFAULT_DATABASE)
return self._extract_result(result)
def _to_csv_from_feature_group(self) -> tuple[str, str]:
base_fg = construct_feature_group_to_be_merged(self._base, self._included_feature_names)
self._record_identifier_feature_name = base_fg.record_identifier_feature_name
self._event_time_identifier_feature_name = base_fg.event_time_identifier_feature.feature_name
self._event_time_identifier_feature_type = base_fg.event_time_identifier_feature.feature_type
query_string = self._construct_query_string(base_fg)
result = self._run_query(query_string, base_fg.catalog, base_fg.database)
return self._extract_result(result)
def _extract_result(self, query_result: dict) -> tuple[str, str]:
execution = query_result.get("QueryExecution", {})
return (
execution.get("ResultConfiguration", {}).get("OutputLocation"),
execution.get("Query"),
)
def _run_query(self, query_string: str, catalog: str, database: str) -> Dict[str, Any]:
return run_athena_query(
session=self._sagemaker_session,
catalog=catalog,
database=database,
query_string=query_string,
output_location=self._output_path,
kms_key=self._kms_key_id,
)
def _create_temp_table(self, temp_table_name: str, s3_folder: str):
columns = ", ".join(
f"{col} {_DTYPE_TO_ATHENA_TYPE.get(str(self._base[col].dtypes), 'STRING')}"
for col in self._base.columns
)
serde = '"separatorChar" = ",", "quoteChar" = "`", "escapeChar" = "\\\\"'
query = (
f"CREATE EXTERNAL TABLE {temp_table_name} ({columns}) "
f"ROW FORMAT SERDE 'org.apache.hadoop.hive.serde2.OpenCSVSerde' "
f"WITH SERDEPROPERTIES ({serde}) LOCATION '{s3_folder}';"
)
self._run_query(query, _DEFAULT_CATALOG, _DEFAULT_DATABASE)
def _construct_query_string(self, base: FeatureGroupToBeMerged) -> str:
base_query = self._construct_table_query(base, "base")
query = f"WITH fg_base AS ({base_query})"
for i, fg in enumerate(self._feature_groups_to_be_merged):
fg_query = self._construct_table_query(fg, str(i))
query += f",\nfg_{i} AS ({fg_query})"
selected = ", ".join(f"fg_base.{f}" for f in base.projected_feature_names)
selected_final = ", ".join(base.projected_feature_names)
for i, fg in enumerate(self._feature_groups_to_be_merged):
selected += ", " + ", ".join(
f'fg_{i}."{f}" as "{f}.{i+1}"' for f in fg.projected_feature_names
)
selected_final += ", " + ", ".join(
f'"{f}.{i+1}"' for f in fg.projected_feature_names
)
query += (
f"\nSELECT {selected_final}\nFROM (\n"
f"SELECT {selected}, row_number() OVER (\n"
f'PARTITION BY fg_base."{base.record_identifier_feature_name}"\n'
f'ORDER BY fg_base."{base.event_time_identifier_feature.feature_name}" DESC'
)
join_strings = []
for i, fg in enumerate(self._feature_groups_to_be_merged):
if not fg.target_feature_name_in_base:
fg.target_feature_name_in_base = self._record_identifier_feature_name
elif fg.target_feature_name_in_base not in base.features:
raise ValueError(f"Feature {fg.target_feature_name_in_base} not found in base")
query += f', fg_{i}."{fg.event_time_identifier_feature.feature_name}" DESC'
join_strings.append(self._construct_join_condition(fg, str(i)))
recent_where = ""
if self._number_of_recent_records is not None and self._number_of_recent_records >= 0:
recent_where = f"WHERE row_recent <= {self._number_of_recent_records}"
query += f"\n) AS row_recent\nFROM fg_base{''.join(join_strings)}\n)\n{recent_where}"
if self._number_of_records is not None and self._number_of_records >= 0:
query += f"\nLIMIT {self._number_of_records}"
return query
def _construct_table_query(self, fg: FeatureGroupToBeMerged, suffix: str) -> str:
included = ", ".join(f'table_{suffix}."{f}"' for f in fg.included_feature_names)
included_with_write = included
if fg.table_type is TableType.FEATURE_GROUP:
included_with_write += f', table_{suffix}."write_time"'
record_id = fg.record_identifier_feature_name
event_time = fg.event_time_identifier_feature.feature_name
if self._include_duplicated_records and self._include_deleted_records:
return (
f"SELECT {included}\n"
f'FROM "{fg.database}"."{fg.table_name}" table_{suffix}\n'
+ self._construct_where_query_string(suffix, fg.event_time_identifier_feature, ["NOT is_deleted"])
)
if fg.table_type is TableType.FEATURE_GROUP and self._include_deleted_records:
rank = f'ORDER BY origin_{suffix}."api_invocation_time" DESC, origin_{suffix}."write_time" DESC\n'
return (
f"SELECT {included}\nFROM (\n"
f"SELECT *, row_number() OVER (\n"
f'PARTITION BY origin_{suffix}."{record_id}", origin_{suffix}."{event_time}"\n'
f"{rank}) AS row_{suffix}\n"
f'FROM "{fg.database}"."{fg.table_name}" origin_{suffix}\n'
f"WHERE NOT is_deleted) AS table_{suffix}\n"
+ self._construct_where_query_string(suffix, fg.event_time_identifier_feature, [f"row_{suffix} = 1"])
)
if fg.table_type is TableType.FEATURE_GROUP:
dedup = self._construct_dedup_query(fg, suffix)
deleted = self._construct_deleted_query(fg, suffix)
rank_cond = (
f'OR (table_{suffix}."{event_time}" = deleted_{suffix}."{event_time}" '
f'AND table_{suffix}."api_invocation_time" > deleted_{suffix}."api_invocation_time")\n'
f'OR (table_{suffix}."{event_time}" = deleted_{suffix}."{event_time}" '
f'AND table_{suffix}."api_invocation_time" = deleted_{suffix}."api_invocation_time" '
f'AND table_{suffix}."write_time" > deleted_{suffix}."write_time")\n'
)
if self._include_duplicated_records:
return (
f"WITH {deleted}\n"
f"SELECT {included}\nFROM (\n"
f"SELECT {included_with_write}\n"
f'FROM "{fg.database}"."{fg.table_name}" table_{suffix}\n'
f"LEFT JOIN deleted_{suffix} ON table_{suffix}.\"{record_id}\" = deleted_{suffix}.\"{record_id}\"\n"
f'WHERE deleted_{suffix}."{record_id}" IS NULL\n'
f"UNION ALL\n"
f"SELECT {included_with_write}\nFROM deleted_{suffix}\n"
f'JOIN "{fg.database}"."{fg.table_name}" table_{suffix}\n'
f'ON table_{suffix}."{record_id}" = deleted_{suffix}."{record_id}"\n'
f'AND (table_{suffix}."{event_time}" > deleted_{suffix}."{event_time}"\n{rank_cond})\n'
f") AS table_{suffix}\n"
+ self._construct_where_query_string(suffix, fg.event_time_identifier_feature, [])
)
return (
f"WITH {dedup},\n{deleted}\n"
f"SELECT {included}\nFROM (\n"
f"SELECT {included_with_write}\nFROM table_{suffix}\n"
f"LEFT JOIN deleted_{suffix} ON table_{suffix}.\"{record_id}\" = deleted_{suffix}.\"{record_id}\"\n"
f'WHERE deleted_{suffix}."{record_id}" IS NULL\n'
f"UNION ALL\n"
f"SELECT {included_with_write}\nFROM deleted_{suffix}\n"
f"JOIN table_{suffix} ON table_{suffix}.\"{record_id}\" = deleted_{suffix}.\"{record_id}\"\n"
f'AND (table_{suffix}."{event_time}" > deleted_{suffix}."{event_time}"\n{rank_cond})\n'
f") AS table_{suffix}\n"
+ self._construct_where_query_string(suffix, fg.event_time_identifier_feature, [])
)
dedup = self._construct_dedup_query(fg, suffix)
return (
f"WITH {dedup}\n"
f"SELECT {included}\nFROM (\n"
f"SELECT {included_with_write}\nFROM table_{suffix}\n"
f") AS table_{suffix}\n"
+ self._construct_where_query_string(suffix, fg.event_time_identifier_feature, [])
)
def _construct_dedup_query(self, fg: FeatureGroupToBeMerged, suffix: str) -> str:
record_id = fg.record_identifier_feature_name
event_time = fg.event_time_identifier_feature.feature_name
rank = ""
is_fg = fg.table_type is TableType.FEATURE_GROUP
if is_fg:
rank = f'ORDER BY origin_{suffix}."api_invocation_time" DESC, origin_{suffix}."write_time" DESC\n'
where_conds = []
if is_fg and self._write_time_ending_timestamp:
where_conds.append(self._construct_write_time_condition(f"origin_{suffix}"))
where_conds.extend(self._construct_event_time_conditions(f"origin_{suffix}", fg.event_time_identifier_feature))
where_str = f"WHERE {' AND '.join(where_conds)}\n" if where_conds else ""
dedup_where = f"WHERE dedup_row_{suffix} = 1\n" if is_fg else ""
return (
f"table_{suffix} AS (\n"
f"SELECT *\nFROM (\n"
f"SELECT *, row_number() OVER (\n"
f'PARTITION BY origin_{suffix}."{record_id}", origin_{suffix}."{event_time}"\n'
f"{rank}) AS dedup_row_{suffix}\n"
f'FROM "{fg.database}"."{fg.table_name}" origin_{suffix}\n'
f"{where_str})\n{dedup_where})"
)
def _construct_deleted_query(self, fg: FeatureGroupToBeMerged, suffix: str) -> str:
record_id = fg.record_identifier_feature_name
event_time = fg.event_time_identifier_feature.feature_name
rank = f'ORDER BY origin_{suffix}."{event_time}" DESC'
if fg.table_type is TableType.FEATURE_GROUP:
rank += f', origin_{suffix}."api_invocation_time" DESC, origin_{suffix}."write_time" DESC\n'
write_cond = ""
if fg.table_type is TableType.FEATURE_GROUP and self._write_time_ending_timestamp:
write_cond = f" AND {self._construct_write_time_condition(f'origin_{suffix}')}\n"
event_conds = ""
if self._event_time_starting_timestamp and self._event_time_ending_timestamp:
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)
return (
f"deleted_{suffix} AS (\n"
f"SELECT *\nFROM (\n"
f"SELECT *, row_number() OVER (\n"
f'PARTITION BY origin_{suffix}."{record_id}"\n'
f"{rank}) AS deleted_row_{suffix}\n"
f'FROM "{fg.database}"."{fg.table_name}" origin_{suffix}\n'
f"WHERE is_deleted{write_cond}{event_conds})\n"
f"WHERE deleted_row_{suffix} = 1\n)"
)
def _construct_where_query_string(
self, suffix: str, event_time_feature: FeatureDefinition, conditions: List[str]
) -> str:
self._validate_options()
if isinstance(self._base, FeatureGroup) and self._write_time_ending_timestamp:
conditions.append(self._construct_write_time_condition(f"table_{suffix}"))
conditions.extend(self._construct_event_time_conditions(f"table_{suffix}", event_time_feature))
return f"WHERE {' AND '.join(conditions)}" if conditions else ""
def _validate_options(self):
is_df_base = isinstance(self._base, pd.DataFrame)
no_joins = len(self._feature_groups_to_be_merged) == 0
if self._number_of_recent_records is not None and self._number_of_recent_records < 0:
raise ValueError("number_of_recent_records must be non-negative.")
if self._number_of_records is not None and self._number_of_records < 0:
raise ValueError("number_of_records must be non-negative.")
if is_df_base and no_joins:
if self._include_deleted_records:
raise ValueError("include_deleted_records() only works for FeatureGroup if no join.")
if self._include_duplicated_records:
raise ValueError("include_duplicated_records() only works for FeatureGroup if no join.")
if self._write_time_ending_timestamp:
raise ValueError("as_of() only works for FeatureGroup if no join.")
if self._point_in_time_accurate_join and no_joins:
raise ValueError("point_in_time_accurate_join() requires at least one join.")
def _construct_event_time_conditions(self, table: str, event_time_feature: FeatureDefinition) -> List[str]:
cast_fn = "from_iso8601_timestamp" if event_time_feature.feature_type == FeatureTypeEnum.STRING else "from_unixtime"
conditions = []
if self._event_time_starting_timestamp:
conditions.append(
f'{cast_fn}({table}."{event_time_feature.feature_name}") >= '
f"from_unixtime({self._event_time_starting_timestamp.timestamp()})"
)
if self._event_time_ending_timestamp:
conditions.append(
f'{cast_fn}({table}."{event_time_feature.feature_name}") <= '
f"from_unixtime({self._event_time_ending_timestamp.timestamp()})"
)
return conditions
def _construct_write_time_condition(self, table: str) -> str:
ts = self._write_time_ending_timestamp.replace(microsecond=0)
return f'{table}."write_time" <= to_timestamp(\'{ts}\', \'yyyy-mm-dd hh24:mi:ss\')'
def _construct_join_condition(self, fg: FeatureGroupToBeMerged, suffix: str) -> str:
target_feature = fg.feature_name_in_target or fg.record_identifier_feature_name
join = (
f"\n{fg.join_type.value} fg_{suffix}\n"
f'ON fg_base."{fg.target_feature_name_in_base}" {fg.join_comparator.value} fg_{suffix}."{target_feature}"'
)
if self._point_in_time_accurate_join:
base_cast = "from_iso8601_timestamp" if self._event_time_identifier_feature_type == FeatureTypeEnum.STRING else "from_unixtime"
fg_cast = "from_iso8601_timestamp" if fg.event_time_identifier_feature.feature_type == FeatureTypeEnum.STRING else "from_unixtime"
join += (
f'\nAND {base_cast}(fg_base."{self._event_time_identifier_feature_name}") >= '
f'{fg_cast}(fg_{suffix}."{fg.event_time_identifier_feature.feature_name}")'
)
return join