From 91ca9a26ca4ed518f61ebebf2f7e842f0ec80176 Mon Sep 17 00:00:00 2001 From: Pavel Borobov Date: Mon, 9 Aug 2021 17:40:28 -0700 Subject: [PATCH 1/5] Add support for database schema def in RedshiftSource --- protos/feast/core/DataSource.proto | 3 ++ .../feast/infra/offline_stores/redshift.py | 18 ++++++--- .../infra/offline_stores/redshift_source.py | 38 ++++++++++++++++--- sdk/python/feast/infra/utils/aws_utils.py | 12 +++++- 4 files changed, 59 insertions(+), 12 deletions(-) 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/infra/offline_stores/redshift.py b/sdk/python/feast/infra/offline_stores/redshift.py index c4c5f4f06b8..b0badfe5eaf 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 """ @@ -237,6 +240,7 @@ def to_arrow(self) -> pa.Table: self._config.offline_store.iam_role, query, self._drop_columns, + self._config.offline_store.temp_schema_name, ) def to_s3(self) -> str: @@ -254,13 +258,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 +297,22 @@ 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..9f5e3877b89 100644 --- a/sdk/python/feast/infra/offline_stores/redshift_source.py +++ b/sdk/python/feast/infra/offline_stores/redshift_source.py @@ -13,6 +13,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 +26,7 @@ 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): @@ -95,7 +96,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})" @@ -153,9 +155,19 @@ 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 +197,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 table(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 +224,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 +240,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..ca7362f36d7 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 @@ -325,6 +330,7 @@ def unload_redshift_query_to_pa( iam_role: str, query: str, drop_columns: Optional[List[str]] = None, + temp_schema_name: Optional[str] = None, ) -> pa.Table: """ Unload Redshift Query results to S3 and get the results in PyArrow Table format """ bucket, key = get_bucket_and_key(s3_path) @@ -356,6 +362,7 @@ def unload_redshift_query_to_df( iam_role: str, query: str, drop_columns: Optional[List[str]] = None, + schema: Optional[List[str]] = None, ) -> pd.DataFrame: """ Unload Redshift Query results to S3 and get the results in Pandas DataFrame format """ table = unload_redshift_query_to_pa( @@ -368,5 +375,6 @@ def unload_redshift_query_to_df( iam_role, query, drop_columns, + schema, ) return table.to_pandas() From 90f8fa70e0877318543af798057715fa08b758ca Mon Sep 17 00:00:00 2001 From: Pavel Borobov Date: Mon, 9 Aug 2021 18:02:50 -0700 Subject: [PATCH 2/5] Add edge case for undefined table specification in multiple schemas --- sdk/python/feast/errors.py | 6 +++++ .../infra/offline_stores/redshift_source.py | 27 ++++++++++++++----- 2 files changed, 26 insertions(+), 7 deletions(-) diff --git a/sdk/python/feast/errors.py b/sdk/python/feast/errors.py index 1202d4df49b..51388907f20 100644 --- a/sdk/python/feast/errors.py +++ b/sdk/python/feast/errors.py @@ -9,6 +9,12 @@ def __init__(self, path): f"Unable to find table at '{path}'. Please check that table exists." ) +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_source.py b/sdk/python/feast/infra/offline_stores/redshift_source.py index 9f5e3877b89..a4168549779 100644 --- a/sdk/python/feast/infra/offline_stores/redshift_source.py +++ b/sdk/python/feast/infra/offline_stores/redshift_source.py @@ -2,7 +2,7 @@ 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 @@ -48,6 +48,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 @@ -61,6 +62,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): """ @@ -119,12 +124,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 @@ -134,6 +143,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( From db28e2d45c174f14d06d27f7a4159c6749a0cc29 Mon Sep 17 00:00:00 2001 From: Pavel Borobov Date: Mon, 9 Aug 2021 18:05:47 -0700 Subject: [PATCH 3/5] Remove temp_schema_name parameter from few places --- sdk/python/feast/infra/offline_stores/redshift.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/sdk/python/feast/infra/offline_stores/redshift.py b/sdk/python/feast/infra/offline_stores/redshift.py index b0badfe5eaf..7bafff8ce07 100644 --- a/sdk/python/feast/infra/offline_stores/redshift.py +++ b/sdk/python/feast/infra/offline_stores/redshift.py @@ -239,8 +239,7 @@ def to_arrow(self) -> pa.Table: self._s3_path, self._config.offline_store.iam_role, query, - self._drop_columns, - self._config.offline_store.temp_schema_name, + self._drop_columns ) def to_s3(self) -> str: From 4fa934b0f2209e842a641fb6abe3cbb7c645eead Mon Sep 17 00:00:00 2001 From: Pavel Borobov Date: Mon, 9 Aug 2021 18:06:37 -0700 Subject: [PATCH 4/5] Remove temp_schema_name parameter from few places --- sdk/python/feast/infra/utils/aws_utils.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/sdk/python/feast/infra/utils/aws_utils.py b/sdk/python/feast/infra/utils/aws_utils.py index ca7362f36d7..6dcb8321d83 100644 --- a/sdk/python/feast/infra/utils/aws_utils.py +++ b/sdk/python/feast/infra/utils/aws_utils.py @@ -330,7 +330,6 @@ def unload_redshift_query_to_pa( iam_role: str, query: str, drop_columns: Optional[List[str]] = None, - temp_schema_name: Optional[str] = None, ) -> pa.Table: """ Unload Redshift Query results to S3 and get the results in PyArrow Table format """ bucket, key = get_bucket_and_key(s3_path) @@ -362,7 +361,6 @@ def unload_redshift_query_to_df( iam_role: str, query: str, drop_columns: Optional[List[str]] = None, - schema: Optional[List[str]] = None, ) -> pd.DataFrame: """ Unload Redshift Query results to S3 and get the results in Pandas DataFrame format """ table = unload_redshift_query_to_pa( @@ -375,6 +373,5 @@ def unload_redshift_query_to_df( iam_role, query, drop_columns, - schema, ) return table.to_pandas() From 50eb4ec36124d850a6ec60f27a7b3fae4ad6dc71 Mon Sep 17 00:00:00 2001 From: Pavel Borobov Date: Tue, 10 Aug 2021 15:48:22 -0700 Subject: [PATCH 5/5] Fix linter issues; Fix copypaste typo --- sdk/python/feast/errors.py | 2 + .../feast/infra/offline_stores/redshift.py | 13 +++--- .../infra/offline_stores/redshift_source.py | 44 +++++++++++-------- sdk/python/feast/infra/utils/aws_utils.py | 8 ++-- 4 files changed, 40 insertions(+), 27 deletions(-) diff --git a/sdk/python/feast/errors.py b/sdk/python/feast/errors.py index 51388907f20..8524bd10883 100644 --- a/sdk/python/feast/errors.py +++ b/sdk/python/feast/errors.py @@ -9,6 +9,7 @@ def __init__(self, path): f"Unable to find table at '{path}'. Please check that table exists." ) + class UndefinedDatasourceSchemaException(Exception): def __init__(self, schemas_found, path): super().__init__( @@ -16,6 +17,7 @@ def __init__(self, schemas_found, 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 7bafff8ce07..fd8d91d9d06 100644 --- a/sdk/python/feast/infra/offline_stores/redshift.py +++ b/sdk/python/feast/infra/offline_stores/redshift.py @@ -239,7 +239,7 @@ def to_arrow(self) -> pa.Table: self._s3_path, self._config.offline_store.iam_role, query, - self._drop_columns + self._drop_columns, ) def to_s3(self) -> str: @@ -260,8 +260,8 @@ def to_s3(self) -> str: 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: - schema_prefix = f'{schema}.' if schema is not None else '' - full_table_name = f'{schema_prefix}{table_name}' + 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: @@ -302,7 +302,7 @@ def _upload_entity_df_and_get_entity_schema( 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}' + full_table_name = f"{config.offline_store.temp_schema_name}.{table_name}" aws_utils.execute_redshift_statement( redshift_client, config.offline_store.cluster_id, @@ -311,7 +311,10 @@ def _upload_entity_df_and_get_entity_schema( f"CREATE TABLE {full_table_name} AS ({entity_df})", ) limited_entity_df = RedshiftRetrievalJob( - f"SELECT * FROM {full_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 a4168549779..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, UndefinedDatasourceSchemaException +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 @@ -26,7 +30,9 @@ def __init__( date_partition_column, ) - self._redshift_options = RedshiftOptions(table=table, schema=schema, query=query) + self._redshift_options = RedshiftOptions( + table=table, schema=schema, query=query + ) @staticmethod def from_proto(data_source: DataSourceProto): @@ -65,7 +71,7 @@ def query(self): @property def schema(self): return self._redshift_options.schema - + @property def redshift_options(self): """ @@ -101,7 +107,7 @@ 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: - schema_prefix = f'{self.schema}.' if self.schema is not None else '' + schema_prefix = f"{self.schema}." if self.schema is not None else "" return f'"{schema_prefix}{self.table}"' else: return f"({self.query})" @@ -125,14 +131,14 @@ def get_table_column_names_and_types( if self.table is not None: try: desribe_table_req = { - 'ClusterIdentifier': config.offline_store.cluster_id, - 'Database': config.offline_store.database, - 'DbUser': config.offline_store.user, - 'Table': self.table, + "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 - + desribe_table_req["Schema"] = self.schema + table = client.describe_table(**desribe_table_req) except ClientError as e: if e.response["Error"]["Code"] == "ValidationException": @@ -143,9 +149,9 @@ 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"]} + unique_schemas = {col["schemaName"] for col in table["ColumnList"]} if len(unique_schemas) > 1: - raise UndefinedDatasourceSchemaException(unique_schemas, self.table) + raise UndefinedDatasourceSchemaException(unique_schemas, self.table) columns = table["ColumnList"] else: @@ -168,14 +174,16 @@ class RedshiftOptions: DataSource Redshift options used to source features from Redshift query """ - def __init__(self, table: Optional[str], query: Optional[str], schema: 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. + 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 @@ -216,9 +224,9 @@ def schema(self): Returns the schema name of this Redshift table schema """ return self._schema - + @schema.setter - def table(self, schema_name): + def schema(self, schema_name): """ Sets the schema ref of this Redshift table schema """ @@ -237,9 +245,9 @@ def from_proto(cls, redshift_options_proto: DataSourceProto.RedshiftOptions): """ redshift_options = cls( - table=redshift_options_proto.table, + table=redshift_options_proto.table, query=redshift_options_proto.query, - schema=redshift_options_proto.schema + schema=redshift_options_proto.schema, ) return redshift_options diff --git a/sdk/python/feast/infra/utils/aws_utils.py b/sdk/python/feast/infra/utils/aws_utils.py index 6dcb8321d83..949c33c7719 100644 --- a/sdk/python/feast/infra/utils/aws_utils.py +++ b/sdk/python/feast/infra/utils/aws_utils.py @@ -205,8 +205,8 @@ 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}' + 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 {full_table_name}({column_query_list}); " + f"COPY {full_table_name} FROM '{s3_path}' IAM_ROLE '{iam_role}' FORMAT AS PARQUET" @@ -230,7 +230,7 @@ def temporarily_upload_df_to_redshift( iam_role: str, table_name: str, df: pd.DataFrame, - schema_name: Optional[str] = None + schema_name: Optional[str] = None, ) -> Iterator[None]: """Uploads a Pandas DataFrame to Redshift as a new table with cleanup logic. @@ -253,7 +253,7 @@ def temporarily_upload_df_to_redshift( iam_role, table_name, df, - schema_name + schema_name, ) yield