diff --git a/protos/feast/core/DataSource.proto b/protos/feast/core/DataSource.proto index 099ba32d929..55a112969ac 100644 --- a/protos/feast/core/DataSource.proto +++ b/protos/feast/core/DataSource.proto @@ -118,6 +118,9 @@ message DataSource { // SQL query that returns a table containing feature data. Must contain an event_timestamp column, and respective // entity columns string query = 2; + + // Redshift table schema name + string schema = 3; } // Defines configuration for custom third-party data sources. diff --git a/sdk/python/feast/errors.py b/sdk/python/feast/errors.py index 1202d4df49b..8524bd10883 100644 --- a/sdk/python/feast/errors.py +++ b/sdk/python/feast/errors.py @@ -10,6 +10,14 @@ def __init__(self, path): ) +class UndefinedDatasourceSchemaException(Exception): + def __init__(self, schemas_found, path): + super().__init__( + f"Unable to identify correct schema location from table at '{path}'. \ + Found schemas: [{', '.join(schemas_found)}]" + ) + + class FeastObjectNotFoundException(Exception): pass diff --git a/sdk/python/feast/infra/offline_stores/redshift.py b/sdk/python/feast/infra/offline_stores/redshift.py index c4c5f4f06b8..fd8d91d9d06 100644 --- a/sdk/python/feast/infra/offline_stores/redshift.py +++ b/sdk/python/feast/infra/offline_stores/redshift.py @@ -38,6 +38,9 @@ class RedshiftOfflineStoreConfig(FeastConfigBaseModel): database: StrictStr """ Redshift database name """ + temp_schema_name: StrictStr + """ Redshift schema name to offload temporary tables """ + s3_staging_location: StrictStr """ S3 path for importing & exporting data to Redshift """ @@ -254,13 +257,15 @@ def to_s3(self) -> str: ) return self._s3_path - def to_redshift(self, table_name: str) -> None: + def to_redshift(self, table_name: str, schema: Optional[str] = None) -> None: """ Save dataset as a new Redshift table """ with self._query_generator() as query: - query = f'CREATE TABLE "{table_name}" AS ({query});\n' + schema_prefix = f"{schema}." if schema is not None else "" + full_table_name = f"{schema_prefix}{table_name}" + query = f'CREATE TABLE "{full_table_name}" AS ({query});\n' if self._drop_columns is not None: for column in self._drop_columns: - query += f"ALTER TABLE {table_name} DROP COLUMN {column};\n" + query += f"ALTER TABLE {full_table_name} DROP COLUMN {column};\n" aws_utils.execute_redshift_statement( self._redshift_client, @@ -291,20 +296,25 @@ def _upload_entity_df_and_get_entity_schema( config.offline_store.iam_role, table_name, entity_df, + config.offline_store.temp_schema_name, ) return dict(zip(entity_df.columns, entity_df.dtypes)) elif isinstance(entity_df, str): # If the entity_df is a string (SQL query), create a Redshift table out of it, # get pandas dataframe consisting of 1 row (LIMIT 1) and generate the schema out of it + full_table_name = f"{config.offline_store.temp_schema_name}.{table_name}" aws_utils.execute_redshift_statement( redshift_client, config.offline_store.cluster_id, config.offline_store.database, config.offline_store.user, - f"CREATE TABLE {table_name} AS ({entity_df})", + f"CREATE TABLE {full_table_name} AS ({entity_df})", ) limited_entity_df = RedshiftRetrievalJob( - f"SELECT * FROM {table_name} LIMIT 1", redshift_client, s3_resource, config + f"SELECT * FROM {full_table_name} LIMIT 1", + redshift_client, + s3_resource, + config, ).to_df() return dict(zip(limited_entity_df.columns, limited_entity_df.dtypes)) else: diff --git a/sdk/python/feast/infra/offline_stores/redshift_source.py b/sdk/python/feast/infra/offline_stores/redshift_source.py index 6e3a8e0f293..b28b0362271 100644 --- a/sdk/python/feast/infra/offline_stores/redshift_source.py +++ b/sdk/python/feast/infra/offline_stores/redshift_source.py @@ -2,7 +2,11 @@ from feast import type_map from feast.data_source import DataSource -from feast.errors import DataSourceNotFoundException, RedshiftCredentialsError +from feast.errors import ( + DataSourceNotFoundException, + RedshiftCredentialsError, + UndefinedDatasourceSchemaException, +) from feast.protos.feast.core.DataSource_pb2 import DataSource as DataSourceProto from feast.repo_config import RepoConfig from feast.value_type import ValueType @@ -13,6 +17,7 @@ def __init__( self, event_timestamp_column: Optional[str] = "", table: Optional[str] = None, + schema: Optional[str] = None, created_timestamp_column: Optional[str] = "", field_mapping: Optional[Dict[str, str]] = None, date_partition_column: Optional[str] = "", @@ -25,7 +30,9 @@ def __init__( date_partition_column, ) - self._redshift_options = RedshiftOptions(table=table, query=query) + self._redshift_options = RedshiftOptions( + table=table, schema=schema, query=query + ) @staticmethod def from_proto(data_source: DataSourceProto): @@ -47,6 +54,7 @@ def __eq__(self, other): return ( self.redshift_options.table == other.redshift_options.table and self.redshift_options.query == other.redshift_options.query + and self.redshift_options.schema == other.redshift_options.schema and self.event_timestamp_column == other.event_timestamp_column and self.created_timestamp_column == other.created_timestamp_column and self.field_mapping == other.field_mapping @@ -60,6 +68,10 @@ def table(self): def query(self): return self._redshift_options.query + @property + def schema(self): + return self._redshift_options.schema + @property def redshift_options(self): """ @@ -95,7 +107,8 @@ def validate(self, config: RepoConfig): def get_table_query_string(self) -> str: """Returns a string that can directly be used to reference this table in SQL""" if self.table: - return f'"{self.table}"' + schema_prefix = f"{self.schema}." if self.schema is not None else "" + return f'"{schema_prefix}{self.table}"' else: return f"({self.query})" @@ -117,12 +130,16 @@ def get_table_column_names_and_types( if self.table is not None: try: - table = client.describe_table( - ClusterIdentifier=config.offline_store.cluster_id, - Database=config.offline_store.database, - DbUser=config.offline_store.user, - Table=self.table, - ) + desribe_table_req = { + "ClusterIdentifier": config.offline_store.cluster_id, + "Database": config.offline_store.database, + "DbUser": config.offline_store.user, + "Table": self.table, + } + if self.schema is not None: + desribe_table_req["Schema"] = self.schema + + table = client.describe_table(**desribe_table_req) except ClientError as e: if e.response["Error"]["Code"] == "ValidationException": raise RedshiftCredentialsError() from e @@ -132,6 +149,10 @@ def get_table_column_names_and_types( if len(table["ColumnList"]) == 0: raise DataSourceNotFoundException(self.table) + unique_schemas = {col["schemaName"] for col in table["ColumnList"]} + if len(unique_schemas) > 1: + raise UndefinedDatasourceSchemaException(unique_schemas, self.table) + columns = table["ColumnList"] else: statement_id = aws_utils.execute_redshift_statement( @@ -153,9 +174,21 @@ class RedshiftOptions: DataSource Redshift options used to source features from Redshift query """ - def __init__(self, table: Optional[str], query: Optional[str]): + def __init__( + self, table: Optional[str], query: Optional[str], schema: Optional[str] + ): + """Redshift options to encapsulate logic for parsing and working with 2 kinds of source creation + table + schema or query + + Args: + table (Optional[str]): Redshift table to be looked for in redshift cluster to form datasource + query (Optional[str]): Query to run to gather datasource + schema (Optional[str]): Schema in redshift cluster to lookup a table. + Has to be provided in case of tables with same name. + """ self._table = table self._query = query + self._schema = schema @property def query(self): @@ -185,6 +218,20 @@ def table(self, table_name): """ self._table = table_name + @property + def schema(self): + """ + Returns the schema name of this Redshift table schema + """ + return self._schema + + @schema.setter + def schema(self, schema_name): + """ + Sets the schema ref of this Redshift table schema + """ + self._schema = schema_name + @classmethod def from_proto(cls, redshift_options_proto: DataSourceProto.RedshiftOptions): """ @@ -198,7 +245,9 @@ def from_proto(cls, redshift_options_proto: DataSourceProto.RedshiftOptions): """ redshift_options = cls( - table=redshift_options_proto.table, query=redshift_options_proto.query, + table=redshift_options_proto.table, + query=redshift_options_proto.query, + schema=redshift_options_proto.schema, ) return redshift_options @@ -212,7 +261,7 @@ def to_proto(self) -> DataSourceProto.RedshiftOptions: """ redshift_options_proto = DataSourceProto.RedshiftOptions( - table=self.table, query=self.query, + table=self.table, query=self.query, schema=self.schema ) return redshift_options_proto diff --git a/sdk/python/feast/infra/utils/aws_utils.py b/sdk/python/feast/infra/utils/aws_utils.py index aea460cfb84..949c33c7719 100644 --- a/sdk/python/feast/infra/utils/aws_utils.py +++ b/sdk/python/feast/infra/utils/aws_utils.py @@ -146,6 +146,7 @@ def upload_df_to_redshift( iam_role: str, table_name: str, df: pd.DataFrame, + schema_name: Optional[str] = None, ) -> None: """Uploads a Pandas DataFrame to Redshift as a new table. @@ -204,9 +205,11 @@ def upload_df_to_redshift( # Create the table with the desired schema and # copy the Parquet file contents to the Redshift table + schema_prefix = f"{schema_name}." if schema_name is not None else "" + full_table_name = f"{schema_prefix}{table_name}" create_and_copy_query = ( - f"CREATE TABLE {table_name}({column_query_list}); " - + f"COPY {table_name} FROM '{s3_path}' IAM_ROLE '{iam_role}' FORMAT AS PARQUET" + f"CREATE TABLE {full_table_name}({column_query_list}); " + + f"COPY {full_table_name} FROM '{s3_path}' IAM_ROLE '{iam_role}' FORMAT AS PARQUET" ) execute_redshift_statement( redshift_data_client, cluster_id, database, user, create_and_copy_query @@ -227,6 +230,7 @@ def temporarily_upload_df_to_redshift( iam_role: str, table_name: str, df: pd.DataFrame, + schema_name: Optional[str] = None, ) -> Iterator[None]: """Uploads a Pandas DataFrame to Redshift as a new table with cleanup logic. @@ -249,6 +253,7 @@ def temporarily_upload_df_to_redshift( iam_role, table_name, df, + schema_name, ) yield