From b09f8b7e7956c646eabd276bf557c24a0bf83c74 Mon Sep 17 00:00:00 2001 From: Jacob Klegar Date: Thu, 4 Mar 2021 19:57:05 -0500 Subject: [PATCH 01/11] WIP Pull table_ref from BQ Signed-off-by: Jacob Klegar --- sdk/python/feast/offline_store.py | 60 +++++++++++++++++++++++++++++++ sdk/python/requirements-dev.txt | 2 +- sdk/python/setup.py | 2 +- 3 files changed, 62 insertions(+), 2 deletions(-) create mode 100644 sdk/python/feast/offline_store.py diff --git a/sdk/python/feast/offline_store.py b/sdk/python/feast/offline_store.py new file mode 100644 index 00000000000..9b0d6607934 --- /dev/null +++ b/sdk/python/feast/offline_store.py @@ -0,0 +1,60 @@ +# Copyright 2019 The Feast Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from abc import ABC, abstractmethod +from datetime import datetime +from typing import Optional +import pyarrow + + +class OfflineStore(ABC): + @abstractmethod + def pull_table( + self, + table_ref: str, + event_timestamp_column: str, + start_date: datetime, + end_date: datetime, + ): + pass + + +class BigQueryOfflineStore: + def pull_table( + self, + table_ref: str, + event_timestamp_column: str, + start_date: datetime, + end_date: datetime, + ) -> Optional[pyarrow.Table]: + from google.cloud.bigquery_storage import BigQueryReadClient, types + + project, dataset, table = table_ref.split(".") + client = BigQueryReadClient() + + requested_session = types.ReadSession() + requested_session.table = f"projects/{project}/datasets/{dataset}/tables/{table}" + requested_session.data_format = types.DataFormat.ARROW + requested_session.read_options.row_restriction = f"{event_timestamp_column} BETWEEN TIMESTAMP('{start_date}') AND TIMESTAMP('{end_date}')" + parent = f"projects/{project}" + session = client.create_read_session( + parent=parent, + read_session=requested_session, + max_stream_count=1, + ) + if len(session.streams) == 0: + # will sometimes happen, indicates no data + return None + reader = client.read_rows(session.streams[0].name) + rows = reader.rows(session).to_arrow() + return rows diff --git a/sdk/python/requirements-dev.txt b/sdk/python/requirements-dev.txt index c7a4ad37d91..499501d3c47 100644 --- a/sdk/python/requirements-dev.txt +++ b/sdk/python/requirements-dev.txt @@ -2,7 +2,7 @@ Click==7.* google-api-core==1.22.4 google-auth==1.22.1 google-cloud-bigquery==1.18 -google-cloud-bigquery-storage==0.7.0 +google-cloud-bigquery-storage==2.3.0 google-cloud-dataproc==2.0.2 google-cloud-storage==1.20.0 google-resumable-media>=0.5 diff --git a/sdk/python/setup.py b/sdk/python/setup.py index a58bc4b8b74..b61b4673c91 100644 --- a/sdk/python/setup.py +++ b/sdk/python/setup.py @@ -31,7 +31,7 @@ "google-cloud-storage==1.20.*", "google-cloud-core==1.4.*", "googleapis-common-protos==1.52.*", - "google-cloud-bigquery-storage==0.7.*", + "google-cloud-bigquery-storage==2.3.*", "grpcio==1.31.0", "pandas~=1.0.0", "pandavro==1.5.*", From f2997f3da5dc1782ef8ca2104b6e58d5cf244e72 Mon Sep 17 00:00:00 2001 From: Jacob Klegar Date: Fri, 5 Mar 2021 15:13:26 -0500 Subject: [PATCH 02/11] More additions Signed-off-by: Jacob Klegar --- sdk/python/feast/big_query_source.py | 22 +++++++- sdk/python/feast/feature_view.py | 7 +++ sdk/python/feast/offline_store.py | 76 +++++++++++++++++++++------- 3 files changed, 86 insertions(+), 19 deletions(-) diff --git a/sdk/python/feast/big_query_source.py b/sdk/python/feast/big_query_source.py index 5fe3df1bd23..5836aea992b 100644 --- a/sdk/python/feast/big_query_source.py +++ b/sdk/python/feast/big_query_source.py @@ -27,8 +27,28 @@ def __init__( field_mapping: Optional[Dict[str, str]], query: Optional[str], ): - if (table_ref is None) != (query is None): + if (table_ref is None) == (query is None): raise Exception("Exactly one of table_ref and query should be specified") + if field_mapping is not None: + for value in field_mapping.values(): + if list(field_mapping.values()).count(value) > 1: + raise Exception( + f"Two fields cannot be mapped to the same name {value}" + ) + + if event_timestamp_column in field_mapping.keys(): + raise Exception( + f"The field {event_timestamp_column} is mapped to {field_mapping[event_timestamp_column]}. Please either remove this field mapping or use {field_mapping[event_timestamp_column]} as the event_timestamp_column." + ) + + if ( + created_timestamp_column is not None + and created_timestamp_column in field_mapping.keys() + ): + raise Exception( + f"The field {created_timestamp_column} is mapped to {field_mapping[created_timestamp_column]}. Please either remove this field mapping or use {field_mapping[created_timestamp_column]} as the _timestamp_column." + ) + self.table_ref = table_ref self.event_timestamp_column = event_timestamp_column self.created_timestamp_column = created_timestamp_column diff --git a/sdk/python/feast/feature_view.py b/sdk/python/feast/feature_view.py index 1047678dca4..e32d3dcfbc3 100644 --- a/sdk/python/feast/feature_view.py +++ b/sdk/python/feast/feature_view.py @@ -35,6 +35,13 @@ def __init__( inputs: BigQuerySource, feature_start_time: datetime, ): + cols = [entity.name for entity in entities] + [feat.name for feat in features] + for col in cols: + if inputs.field_mapping is not None and col in inputs.field_mapping.keys(): + raise Exception( + f"The field {col} is mapped to {inputs.field_mapping[col]} for this data source. Please either remove this field mapping or use {inputs.field_mapping[col]} as the Entity or Feature name." + ) + self.name = name self.entities = entities self.features = features diff --git a/sdk/python/feast/offline_store.py b/sdk/python/feast/offline_store.py index 9b0d6607934..cd981f2b39c 100644 --- a/sdk/python/feast/offline_store.py +++ b/sdk/python/feast/offline_store.py @@ -13,48 +13,88 @@ # limitations under the License. from abc import ABC, abstractmethod from datetime import datetime -from typing import Optional +from typing import Dict, List, Optional + import pyarrow +from feast.entity import Entity +from feast.feature import Feature +from feast.feature_view import FeatureView + class OfflineStore(ABC): + """ + OfflineStore is a non-user-facing object used for all interaction between Feast and the service used for offline storage of features. Currently BigQuery is supported. + """ + @abstractmethod def pull_table( - self, - table_ref: str, - event_timestamp_column: str, - start_date: datetime, - end_date: datetime, + self, feature_view: FeatureView, start_date: datetime, end_date: datetime, ): pass class BigQueryOfflineStore: + """ + BigQueryOfflineStore is a non-user-facing object used for all interaction between Feast and BigQuery. + """ + def pull_table( - self, - table_ref: str, - event_timestamp_column: str, - start_date: datetime, - end_date: datetime, + self, feature_view: FeatureView, start_date: datetime, end_date: datetime, ) -> Optional[pyarrow.Table]: from google.cloud.bigquery_storage import BigQueryReadClient, types - project, dataset, table = table_ref.split(".") + project, dataset, table = feature_view.inputs.table_ref.split(".") client = BigQueryReadClient() requested_session = types.ReadSession() - requested_session.table = f"projects/{project}/datasets/{dataset}/tables/{table}" + requested_session.table = ( + f"projects/{project}/datasets/{dataset}/tables/{table}" + ) requested_session.data_format = types.DataFormat.ARROW + + # if we have mapped fields, use the original field names in the call to BigQuery + event_timestamp_column = feature_view.inputs.event_timestamp_column + fields = ( + [entity.name for entity in feature_view.entities] + + [feature.name for feature in feature_view.features] + + [feature_view.inputs.event_timestamp_column] + ) + if feature_view.inputs.field_mapping is not None: + reverse_field_mapping = { + v: k for k, v in feature_view.inputs.field_mapping.items() + } + event_timestamp_column = ( + reverse_field_mapping[event_timestamp_column] + if event_timestamp_column in reverse_field_mapping.keys() + else event_timestamp_column + ) + fields = [ + reverse_field_mapping[col] + if col in reverse_field_mapping.keys() + else col + for col in fields + ] + requested_session.read_options.selected_fields = fields requested_session.read_options.row_restriction = f"{event_timestamp_column} BETWEEN TIMESTAMP('{start_date}') AND TIMESTAMP('{end_date}')" + parent = f"projects/{project}" session = client.create_read_session( - parent=parent, - read_session=requested_session, - max_stream_count=1, + parent=parent, read_session=requested_session, max_stream_count=1, ) + if len(session.streams) == 0: # will sometimes happen, indicates no data return None reader = client.read_rows(session.streams[0].name) - rows = reader.rows(session).to_arrow() - return rows + table = reader.to_arrow(session) + if feature_view.inputs.field_mapping is not None: + cols = table.column_names + mapped_cols = [ + feature_view.inputs.field_mapping[col] + if col in feature_view.inputs.field_mapping.keys() + else col + for col in cols + ] + table = table.rename_columns(mapped_cols) + return table From 8710676373956d916e5333873c7b746541b31dc8 Mon Sep 17 00:00:00 2001 From: Jacob Klegar Date: Mon, 8 Mar 2021 17:22:32 -0500 Subject: [PATCH 03/11] Refactor a bit + address comments Signed-off-by: Jacob Klegar --- sdk/python/feast/big_query_source.py | 16 ++++----- sdk/python/feast/feature_store.py | 47 ++++++++++++++++++++++++-- sdk/python/feast/feature_view.py | 2 +- sdk/python/feast/offline_store.py | 49 +++++----------------------- sdk/python/requirements-dev.txt | 2 +- sdk/python/setup.py | 2 +- 6 files changed, 64 insertions(+), 54 deletions(-) diff --git a/sdk/python/feast/big_query_source.py b/sdk/python/feast/big_query_source.py index 5836aea992b..53c3c842c55 100644 --- a/sdk/python/feast/big_query_source.py +++ b/sdk/python/feast/big_query_source.py @@ -21,23 +21,23 @@ class BigQuerySource: def __init__( self, - table_ref: Optional[str], event_timestamp_column: str, - created_timestamp_column: Optional[str], - field_mapping: Optional[Dict[str, str]], - query: Optional[str], + table_ref: Optional[str] = None, + created_timestamp_column: Optional[str] = None, + field_mapping: Optional[Dict[str, str]] = None, + query: Optional[str] = None, ): if (table_ref is None) == (query is None): - raise Exception("Exactly one of table_ref and query should be specified") + raise ValueError("Exactly one of table_ref and query should be specified") if field_mapping is not None: for value in field_mapping.values(): if list(field_mapping.values()).count(value) > 1: - raise Exception( + raise ValueError( f"Two fields cannot be mapped to the same name {value}" ) if event_timestamp_column in field_mapping.keys(): - raise Exception( + raise ValueError( f"The field {event_timestamp_column} is mapped to {field_mapping[event_timestamp_column]}. Please either remove this field mapping or use {field_mapping[event_timestamp_column]} as the event_timestamp_column." ) @@ -45,7 +45,7 @@ def __init__( created_timestamp_column is not None and created_timestamp_column in field_mapping.keys() ): - raise Exception( + raise ValueError( f"The field {created_timestamp_column} is mapped to {field_mapping[created_timestamp_column]}. Please either remove this field mapping or use {field_mapping[created_timestamp_column]} as the _timestamp_column." ) diff --git a/sdk/python/feast/feature_store.py b/sdk/python/feast/feature_store.py index 62aed28b991..131f63d6eb6 100644 --- a/sdk/python/feast/feature_store.py +++ b/sdk/python/feast/feature_store.py @@ -11,10 +11,15 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +from datetime import datetime from pathlib import Path from typing import Optional +import pyarrow + +from feast.feature_view import FeatureView from feast.infra.provider import Provider, get_provider +from feast.offline_store import OfflineStore, BigQueryOfflineStore from feast.registry import Registry from feast.repo_config import ( LocalOnlineStoreConfig, @@ -32,10 +37,10 @@ class FeatureStore: config: RepoConfig def __init__( - self, repo_path: Optional[str], config: Optional[RepoConfig], + self, repo_path: Optional[str] = None, config: Optional[RepoConfig] = None, ): if repo_path is not None and config is not None: - raise Exception("You cannot specify both repo_path and config") + raise ValueError("You cannot specify both repo_path and config") if config is not None: self.config = config elif repo_path is not None: @@ -55,3 +60,41 @@ def _get_provider(self) -> Provider: def _get_registry(self) -> Registry: return Registry(self.config.metadata_store) + + def _pull_table(self, feature_view: FeatureView, start_date: datetime, end_date: datetime) -> Optional[pyarrow.Table]: + # if we have mapped fields, use the original field names in the call to the offline store + event_timestamp_column = feature_view.inputs.event_timestamp_column + fields = ( + [entity.name for entity in feature_view.entities] + + [feature.name for feature in feature_view.features] + + [feature_view.inputs.event_timestamp_column] + ) + if feature_view.inputs.field_mapping is not None: + reverse_field_mapping = { + v: k for k, v in feature_view.inputs.field_mapping.items() + } + event_timestamp_column = ( + reverse_field_mapping[event_timestamp_column] + if event_timestamp_column in reverse_field_mapping.keys() + else event_timestamp_column + ) + fields = [ + reverse_field_mapping[col] + if col in reverse_field_mapping.keys() + else col + for col in fields + ] + + table = BigQueryOfflineStore.pull_table(feature_view.inputs.table_ref, fields, event_timestamp_column, start_date, end_date) + + # run feature mapping in the forward direction + if table is not None and feature_view.inputs.field_mapping is not None: + cols = table.column_names + mapped_cols = [ + feature_view.inputs.field_mapping[col] + if col in feature_view.inputs.field_mapping.keys() + else col + for col in cols + ] + table = table.rename_columns(mapped_cols) + return table diff --git a/sdk/python/feast/feature_view.py b/sdk/python/feast/feature_view.py index e32d3dcfbc3..44f0106a040 100644 --- a/sdk/python/feast/feature_view.py +++ b/sdk/python/feast/feature_view.py @@ -38,7 +38,7 @@ def __init__( cols = [entity.name for entity in entities] + [feat.name for feat in features] for col in cols: if inputs.field_mapping is not None and col in inputs.field_mapping.keys(): - raise Exception( + raise ValueError( f"The field {col} is mapped to {inputs.field_mapping[col]} for this data source. Please either remove this field mapping or use {inputs.field_mapping[col]} as the Entity or Feature name." ) diff --git a/sdk/python/feast/offline_store.py b/sdk/python/feast/offline_store.py index cd981f2b39c..06d560b99fd 100644 --- a/sdk/python/feast/offline_store.py +++ b/sdk/python/feast/offline_store.py @@ -13,38 +13,36 @@ # limitations under the License. from abc import ABC, abstractmethod from datetime import datetime -from typing import Dict, List, Optional +from typing import List, Optional import pyarrow -from feast.entity import Entity -from feast.feature import Feature -from feast.feature_view import FeatureView - class OfflineStore(ABC): """ OfflineStore is a non-user-facing object used for all interaction between Feast and the service used for offline storage of features. Currently BigQuery is supported. """ + @staticmethod @abstractmethod def pull_table( - self, feature_view: FeatureView, start_date: datetime, end_date: datetime, - ): + table_ref: str, fields: List[str], event_timestamp_column: str, start_date: datetime, end_date: datetime, + ) -> Optional[pyarrow.Table]: pass -class BigQueryOfflineStore: +class BigQueryOfflineStore(OfflineStore): """ BigQueryOfflineStore is a non-user-facing object used for all interaction between Feast and BigQuery. """ + @staticmethod def pull_table( - self, feature_view: FeatureView, start_date: datetime, end_date: datetime, + table_ref: str, fields: List[str], event_timestamp_column: str, start_date: datetime, end_date: datetime, ) -> Optional[pyarrow.Table]: from google.cloud.bigquery_storage import BigQueryReadClient, types - project, dataset, table = feature_view.inputs.table_ref.split(".") + project, dataset, table = table_ref.split(".") client = BigQueryReadClient() requested_session = types.ReadSession() @@ -53,28 +51,6 @@ def pull_table( ) requested_session.data_format = types.DataFormat.ARROW - # if we have mapped fields, use the original field names in the call to BigQuery - event_timestamp_column = feature_view.inputs.event_timestamp_column - fields = ( - [entity.name for entity in feature_view.entities] - + [feature.name for feature in feature_view.features] - + [feature_view.inputs.event_timestamp_column] - ) - if feature_view.inputs.field_mapping is not None: - reverse_field_mapping = { - v: k for k, v in feature_view.inputs.field_mapping.items() - } - event_timestamp_column = ( - reverse_field_mapping[event_timestamp_column] - if event_timestamp_column in reverse_field_mapping.keys() - else event_timestamp_column - ) - fields = [ - reverse_field_mapping[col] - if col in reverse_field_mapping.keys() - else col - for col in fields - ] requested_session.read_options.selected_fields = fields requested_session.read_options.row_restriction = f"{event_timestamp_column} BETWEEN TIMESTAMP('{start_date}') AND TIMESTAMP('{end_date}')" @@ -88,13 +64,4 @@ def pull_table( return None reader = client.read_rows(session.streams[0].name) table = reader.to_arrow(session) - if feature_view.inputs.field_mapping is not None: - cols = table.column_names - mapped_cols = [ - feature_view.inputs.field_mapping[col] - if col in feature_view.inputs.field_mapping.keys() - else col - for col in cols - ] - table = table.rename_columns(mapped_cols) return table diff --git a/sdk/python/requirements-dev.txt b/sdk/python/requirements-dev.txt index 499501d3c47..414244dbccc 100644 --- a/sdk/python/requirements-dev.txt +++ b/sdk/python/requirements-dev.txt @@ -2,7 +2,7 @@ Click==7.* google-api-core==1.22.4 google-auth==1.22.1 google-cloud-bigquery==1.18 -google-cloud-bigquery-storage==2.3.0 +google-cloud-bigquery-storage>=2.3.0 google-cloud-dataproc==2.0.2 google-cloud-storage==1.20.0 google-resumable-media>=0.5 diff --git a/sdk/python/setup.py b/sdk/python/setup.py index b61b4673c91..fd6e139d435 100644 --- a/sdk/python/setup.py +++ b/sdk/python/setup.py @@ -31,7 +31,7 @@ "google-cloud-storage==1.20.*", "google-cloud-core==1.4.*", "googleapis-common-protos==1.52.*", - "google-cloud-bigquery-storage==2.3.*", + "google-cloud-bigquery-storage>=2.3", "grpcio==1.31.0", "pandas~=1.0.0", "pandavro==1.5.*", From c7350a52a16e821abf65a861ce638f04b8b4285a Mon Sep 17 00:00:00 2001 From: Jacob Klegar Date: Mon, 8 Mar 2021 18:01:50 -0500 Subject: [PATCH 04/11] lint Signed-off-by: Jacob Klegar --- sdk/python/feast/feature_store.py | 19 ++++++++++++++++--- sdk/python/feast/offline_store.py | 12 ++++++++++-- 2 files changed, 26 insertions(+), 5 deletions(-) diff --git a/sdk/python/feast/feature_store.py b/sdk/python/feast/feature_store.py index 131f63d6eb6..ad55083f2f7 100644 --- a/sdk/python/feast/feature_store.py +++ b/sdk/python/feast/feature_store.py @@ -19,7 +19,7 @@ from feast.feature_view import FeatureView from feast.infra.provider import Provider, get_provider -from feast.offline_store import OfflineStore, BigQueryOfflineStore +from feast.offline_store import BigQueryOfflineStore from feast.registry import Registry from feast.repo_config import ( LocalOnlineStoreConfig, @@ -61,7 +61,14 @@ def _get_provider(self) -> Provider: def _get_registry(self) -> Registry: return Registry(self.config.metadata_store) - def _pull_table(self, feature_view: FeatureView, start_date: datetime, end_date: datetime) -> Optional[pyarrow.Table]: + def _pull_table( + self, feature_view: FeatureView, start_date: datetime, end_date: datetime + ) -> Optional[pyarrow.Table]: + if feature_view.inputs.table_ref is None: + raise NotImplementedError( + "Ingestion is not yet implemented for query-based sources." + ) + # if we have mapped fields, use the original field names in the call to the offline store event_timestamp_column = feature_view.inputs.event_timestamp_column fields = ( @@ -85,7 +92,13 @@ def _pull_table(self, feature_view: FeatureView, start_date: datetime, end_date: for col in fields ] - table = BigQueryOfflineStore.pull_table(feature_view.inputs.table_ref, fields, event_timestamp_column, start_date, end_date) + table = BigQueryOfflineStore.pull_table( + feature_view.inputs.table_ref, + fields, + event_timestamp_column, + start_date, + end_date, + ) # run feature mapping in the forward direction if table is not None and feature_view.inputs.field_mapping is not None: diff --git a/sdk/python/feast/offline_store.py b/sdk/python/feast/offline_store.py index 06d560b99fd..7a39930bb95 100644 --- a/sdk/python/feast/offline_store.py +++ b/sdk/python/feast/offline_store.py @@ -26,7 +26,11 @@ class OfflineStore(ABC): @staticmethod @abstractmethod def pull_table( - table_ref: str, fields: List[str], event_timestamp_column: str, start_date: datetime, end_date: datetime, + table_ref: str, + fields: List[str], + event_timestamp_column: str, + start_date: datetime, + end_date: datetime, ) -> Optional[pyarrow.Table]: pass @@ -38,7 +42,11 @@ class BigQueryOfflineStore(OfflineStore): @staticmethod def pull_table( - table_ref: str, fields: List[str], event_timestamp_column: str, start_date: datetime, end_date: datetime, + table_ref: str, + fields: List[str], + event_timestamp_column: str, + start_date: datetime, + end_date: datetime, ) -> Optional[pyarrow.Table]: from google.cloud.bigquery_storage import BigQueryReadClient, types From 13c288cf8c284aed6e2c36bd7dde596282c3b700 Mon Sep 17 00:00:00 2001 From: Jacob Klegar Date: Wed, 10 Mar 2021 12:01:20 -0500 Subject: [PATCH 05/11] Rebase and add created_timestamp_column Signed-off-by: Jacob Klegar --- sdk/python/feast/feature_store.py | 2 ++ sdk/python/feast/repo_config.py | 8 ++++---- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/sdk/python/feast/feature_store.py b/sdk/python/feast/feature_store.py index ad55083f2f7..d2527ce61dc 100644 --- a/sdk/python/feast/feature_store.py +++ b/sdk/python/feast/feature_store.py @@ -76,6 +76,8 @@ def _pull_table( + [feature.name for feature in feature_view.features] + [feature_view.inputs.event_timestamp_column] ) + if feature_view.inputs.created_timestamp_column is not None: + fields.append(feature_view.inputs.created_timestamp_column) if feature_view.inputs.field_mapping is not None: reverse_field_mapping = { v: k for k, v in feature_view.inputs.field_mapping.items() diff --git a/sdk/python/feast/repo_config.py b/sdk/python/feast/repo_config.py index aae9c217d19..af2c1904ca4 100644 --- a/sdk/python/feast/repo_config.py +++ b/sdk/python/feast/repo_config.py @@ -19,10 +19,10 @@ class OnlineStoreConfig(NamedTuple): class RepoConfig(NamedTuple): - metadata_store: str - project: str - provider: str - online_store: OnlineStoreConfig + metadata_store: str = "./metadata_store" + project: str = "default" + provider: str = "local" + online_store: OnlineStoreConfig = OnlineStoreConfig() def load_repo_config(repo_path: Path) -> RepoConfig: From 0f251d7bc4437cab5e11092aee9cf15142ab29e7 Mon Sep 17 00:00:00 2001 From: Jacob Klegar Date: Wed, 10 Mar 2021 17:37:32 -0500 Subject: [PATCH 06/11] Filter by timestamp in BQ first Signed-off-by: Jacob Klegar --- sdk/python/feast/feature_store.py | 30 +++++++++++------ sdk/python/feast/offline_store.py | 56 ++++++++++++++++++------------- sdk/python/requirements-dev.txt | 1 - sdk/python/setup.py | 1 - 4 files changed, 53 insertions(+), 35 deletions(-) diff --git a/sdk/python/feast/feature_store.py b/sdk/python/feast/feature_store.py index d2527ce61dc..7a4bc243bd0 100644 --- a/sdk/python/feast/feature_store.py +++ b/sdk/python/feast/feature_store.py @@ -71,13 +71,9 @@ def _pull_table( # if we have mapped fields, use the original field names in the call to the offline store event_timestamp_column = feature_view.inputs.event_timestamp_column - fields = ( - [entity.name for entity in feature_view.entities] - + [feature.name for feature in feature_view.features] - + [feature_view.inputs.event_timestamp_column] - ) - if feature_view.inputs.created_timestamp_column is not None: - fields.append(feature_view.inputs.created_timestamp_column) + entity_names = [entity.name for entity in feature_view.entities] + feature_names = [feature.name for feature in feature_view.features] + created_timestamp_column = feature_view.inputs.created_timestamp_column if feature_view.inputs.field_mapping is not None: reverse_field_mapping = { v: k for k, v in feature_view.inputs.field_mapping.items() @@ -87,17 +83,31 @@ def _pull_table( if event_timestamp_column in reverse_field_mapping.keys() else event_timestamp_column ) - fields = [ + created_timestamp_column = ( + reverse_field_mapping[created_timestamp_column] + if created_timestamp_column is not None + and created_timestamp_column in reverse_field_mapping.keys() + else created_timestamp_column + ) + entity_names = [ + reverse_field_mapping[col] + if col in reverse_field_mapping.keys() + else col + for col in entity_names + ] + feature_names = [ reverse_field_mapping[col] if col in reverse_field_mapping.keys() else col - for col in fields + for col in feature_names ] table = BigQueryOfflineStore.pull_table( feature_view.inputs.table_ref, - fields, + entity_names, + feature_names, event_timestamp_column, + created_timestamp_column, start_date, end_date, ) diff --git a/sdk/python/feast/offline_store.py b/sdk/python/feast/offline_store.py index 7a39930bb95..03b88c3b990 100644 --- a/sdk/python/feast/offline_store.py +++ b/sdk/python/feast/offline_store.py @@ -27,8 +27,10 @@ class OfflineStore(ABC): @abstractmethod def pull_table( table_ref: str, - fields: List[str], + entity_names: List[str], + feature_names: List[str], event_timestamp_column: str, + created_timestamp_column: Optional[str], start_date: datetime, end_date: datetime, ) -> Optional[pyarrow.Table]: @@ -43,33 +45,41 @@ class BigQueryOfflineStore(OfflineStore): @staticmethod def pull_table( table_ref: str, - fields: List[str], + entity_names: List[str], + feature_names: List[str], event_timestamp_column: str, + created_timestamp_column: Optional[str], start_date: datetime, end_date: datetime, - ) -> Optional[pyarrow.Table]: - from google.cloud.bigquery_storage import BigQueryReadClient, types + ) -> pyarrow.Table: - project, dataset, table = table_ref.split(".") - client = BigQueryReadClient() + partition_by_entity_string = ", ".join(entity_names) + if partition_by_entity_string != "": + partition_by_entity_string = "PARTITION BY " + partition_by_entity_string + feature_string = ", ".join(feature_names) + timestamps = [event_timestamp_column] + if created_timestamp_column is not None: + timestamps.append(created_timestamp_column) + timestamp_string = ", ".join(timestamps) + timestamp_desc_string = " DESC, ".join(timestamps) + " DESC" + field_string = ", ".join(entity_names + feature_names + timestamps) - requested_session = types.ReadSession() - requested_session.table = ( - f"projects/{project}/datasets/{dataset}/tables/{table}" + query = f""" + SELECT {field_string} + FROM ( + SELECT {field_string}, + ROW_NUMBER() OVER({partition_by_entity_string} ORDER BY {timestamp_desc_string}) AS _feast_row + FROM `{table_ref}` + WHERE {event_timestamp_column} BETWEEN TIMESTAMP('{start_date}') AND TIMESTAMP('{end_date}') ) - requested_session.data_format = types.DataFormat.ARROW - - requested_session.read_options.selected_fields = fields - requested_session.read_options.row_restriction = f"{event_timestamp_column} BETWEEN TIMESTAMP('{start_date}') AND TIMESTAMP('{end_date}')" + WHERE _feast_row = 1 + """ + return BigQueryOfflineStore._pull_query(query) - parent = f"projects/{project}" - session = client.create_read_session( - parent=parent, read_session=requested_session, max_stream_count=1, - ) + @staticmethod + def _pull_query(query: str) -> pyarrow.Table: + from google.cloud import bigquery - if len(session.streams) == 0: - # will sometimes happen, indicates no data - return None - reader = client.read_rows(session.streams[0].name) - table = reader.to_arrow(session) - return table + client = bigquery.Client() + query_job = client.query(query) + return query_job.to_arrow() diff --git a/sdk/python/requirements-dev.txt b/sdk/python/requirements-dev.txt index 414244dbccc..ce59be00249 100644 --- a/sdk/python/requirements-dev.txt +++ b/sdk/python/requirements-dev.txt @@ -2,7 +2,6 @@ Click==7.* google-api-core==1.22.4 google-auth==1.22.1 google-cloud-bigquery==1.18 -google-cloud-bigquery-storage>=2.3.0 google-cloud-dataproc==2.0.2 google-cloud-storage==1.20.0 google-resumable-media>=0.5 diff --git a/sdk/python/setup.py b/sdk/python/setup.py index fd6e139d435..41ae8a67a61 100644 --- a/sdk/python/setup.py +++ b/sdk/python/setup.py @@ -31,7 +31,6 @@ "google-cloud-storage==1.20.*", "google-cloud-core==1.4.*", "googleapis-common-protos==1.52.*", - "google-cloud-bigquery-storage>=2.3", "grpcio==1.31.0", "pandas~=1.0.0", "pandavro==1.5.*", From 6052d7aeb837635f30606f4a1976242b0be564fc Mon Sep 17 00:00:00 2001 From: Jacob Klegar Date: Wed, 10 Mar 2021 17:43:47 -0500 Subject: [PATCH 07/11] Lint Signed-off-by: Jacob Klegar --- sdk/python/feast/offline_store.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/sdk/python/feast/offline_store.py b/sdk/python/feast/offline_store.py index 03b88c3b990..bc603ba0035 100644 --- a/sdk/python/feast/offline_store.py +++ b/sdk/python/feast/offline_store.py @@ -56,11 +56,9 @@ def pull_table( partition_by_entity_string = ", ".join(entity_names) if partition_by_entity_string != "": partition_by_entity_string = "PARTITION BY " + partition_by_entity_string - feature_string = ", ".join(feature_names) timestamps = [event_timestamp_column] if created_timestamp_column is not None: timestamps.append(created_timestamp_column) - timestamp_string = ", ".join(timestamps) timestamp_desc_string = " DESC, ".join(timestamps) + " DESC" field_string = ", ".join(entity_names + feature_names + timestamps) From 9198c092cb78de89427ee512090b65700690ae67 Mon Sep 17 00:00:00 2001 From: Jacob Klegar Date: Thu, 11 Mar 2021 13:12:05 -0500 Subject: [PATCH 08/11] Move field mapping to standalone functions Signed-off-by: Jacob Klegar --- sdk/python/feast/feature_store.py | 67 ------------------------------- sdk/python/feast/offline_store.py | 63 +++++++++++++++++++++++++++-- 2 files changed, 60 insertions(+), 70 deletions(-) diff --git a/sdk/python/feast/feature_store.py b/sdk/python/feast/feature_store.py index 7a4bc243bd0..efd46cabff3 100644 --- a/sdk/python/feast/feature_store.py +++ b/sdk/python/feast/feature_store.py @@ -11,15 +11,11 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from datetime import datetime from pathlib import Path from typing import Optional -import pyarrow -from feast.feature_view import FeatureView from feast.infra.provider import Provider, get_provider -from feast.offline_store import BigQueryOfflineStore from feast.registry import Registry from feast.repo_config import ( LocalOnlineStoreConfig, @@ -60,66 +56,3 @@ def _get_provider(self) -> Provider: def _get_registry(self) -> Registry: return Registry(self.config.metadata_store) - - def _pull_table( - self, feature_view: FeatureView, start_date: datetime, end_date: datetime - ) -> Optional[pyarrow.Table]: - if feature_view.inputs.table_ref is None: - raise NotImplementedError( - "Ingestion is not yet implemented for query-based sources." - ) - - # if we have mapped fields, use the original field names in the call to the offline store - event_timestamp_column = feature_view.inputs.event_timestamp_column - entity_names = [entity.name for entity in feature_view.entities] - feature_names = [feature.name for feature in feature_view.features] - created_timestamp_column = feature_view.inputs.created_timestamp_column - if feature_view.inputs.field_mapping is not None: - reverse_field_mapping = { - v: k for k, v in feature_view.inputs.field_mapping.items() - } - event_timestamp_column = ( - reverse_field_mapping[event_timestamp_column] - if event_timestamp_column in reverse_field_mapping.keys() - else event_timestamp_column - ) - created_timestamp_column = ( - reverse_field_mapping[created_timestamp_column] - if created_timestamp_column is not None - and created_timestamp_column in reverse_field_mapping.keys() - else created_timestamp_column - ) - entity_names = [ - reverse_field_mapping[col] - if col in reverse_field_mapping.keys() - else col - for col in entity_names - ] - feature_names = [ - reverse_field_mapping[col] - if col in reverse_field_mapping.keys() - else col - for col in feature_names - ] - - table = BigQueryOfflineStore.pull_table( - feature_view.inputs.table_ref, - entity_names, - feature_names, - event_timestamp_column, - created_timestamp_column, - start_date, - end_date, - ) - - # run feature mapping in the forward direction - if table is not None and feature_view.inputs.field_mapping is not None: - cols = table.column_names - mapped_cols = [ - feature_view.inputs.field_mapping[col] - if col in feature_view.inputs.field_mapping.keys() - else col - for col in cols - ] - table = table.rename_columns(mapped_cols) - return table diff --git a/sdk/python/feast/offline_store.py b/sdk/python/feast/offline_store.py index bc603ba0035..585b8c5f5f0 100644 --- a/sdk/python/feast/offline_store.py +++ b/sdk/python/feast/offline_store.py @@ -13,10 +13,12 @@ # limitations under the License. from abc import ABC, abstractmethod from datetime import datetime -from typing import List, Optional +from typing import List, Optional, Tuple import pyarrow +from feast.feature_view import FeatureView + class OfflineStore(ABC): """ @@ -25,7 +27,7 @@ class OfflineStore(ABC): @staticmethod @abstractmethod - def pull_table( + def pull_latest_from_table( table_ref: str, entity_names: List[str], feature_names: List[str], @@ -43,7 +45,7 @@ class BigQueryOfflineStore(OfflineStore): """ @staticmethod - def pull_table( + def pull_latest_from_table( table_ref: str, entity_names: List[str], feature_names: List[str], @@ -81,3 +83,58 @@ def _pull_query(query: str) -> pyarrow.Table: client = bigquery.Client() query_job = client.query(query) return query_job.to_arrow() + + +def run_reverse_field_mapping( + feature_view: FeatureView, +) -> Tuple[List[str], List[str], str, Optional[str]]: + # if we have mapped fields, use the original field names in the call to the offline store + event_timestamp_column = feature_view.inputs.event_timestamp_column + entity_names = [entity.name for entity in feature_view.entities] + feature_names = [feature.name for feature in feature_view.features] + created_timestamp_column = feature_view.inputs.created_timestamp_column + if feature_view.inputs.field_mapping is not None: + reverse_field_mapping = { + v: k for k, v in feature_view.inputs.field_mapping.items() + } + event_timestamp_column = ( + reverse_field_mapping[event_timestamp_column] + if event_timestamp_column in reverse_field_mapping.keys() + else event_timestamp_column + ) + created_timestamp_column = ( + reverse_field_mapping[created_timestamp_column] + if created_timestamp_column is not None + and created_timestamp_column in reverse_field_mapping.keys() + else created_timestamp_column + ) + entity_names = [ + reverse_field_mapping[col] if col in reverse_field_mapping.keys() else col + for col in entity_names + ] + feature_names = [ + reverse_field_mapping[col] if col in reverse_field_mapping.keys() else col + for col in feature_names + ] + return ( + entity_names, + feature_names, + event_timestamp_column, + created_timestamp_column, + ) + + +def run_forward_field_mapping( + table: pyarrow.Table, feature_view: FeatureView +) -> pyarrow.Table: + # run field mapping in the forward direction + if table is not None and feature_view.inputs.field_mapping is not None: + cols = table.column_names + mapped_cols = [ + feature_view.inputs.field_mapping[col] + if col in feature_view.inputs.field_mapping.keys() + else col + for col in cols + ] + table = table.rename_columns(mapped_cols) + return table From b864271629dd8f85d90252285349e826e66232e6 Mon Sep 17 00:00:00 2001 From: Jacob Klegar Date: Thu, 11 Mar 2021 16:04:07 -0500 Subject: [PATCH 09/11] Update docstrings Signed-off-by: Jacob Klegar --- sdk/python/feast/offline_store.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/sdk/python/feast/offline_store.py b/sdk/python/feast/offline_store.py index 585b8c5f5f0..455edce48f8 100644 --- a/sdk/python/feast/offline_store.py +++ b/sdk/python/feast/offline_store.py @@ -40,10 +40,6 @@ def pull_latest_from_table( class BigQueryOfflineStore(OfflineStore): - """ - BigQueryOfflineStore is a non-user-facing object used for all interaction between Feast and BigQuery. - """ - @staticmethod def pull_latest_from_table( table_ref: str, @@ -88,6 +84,14 @@ def _pull_query(query: str) -> pyarrow.Table: def run_reverse_field_mapping( feature_view: FeatureView, ) -> Tuple[List[str], List[str], str, Optional[str]]: + """ + If a field mapping exists, run it in reverse on the entity names, feature names, event timestamp column, and created timestamp column to get the names of the relevant columns in the BigQuery table. + + Args: + feature_view: FeatureView object containing the field mapping as well as the names to reverse-map. + Returns: + Tuple containing the list of reverse-mapped entity names, reverse-mapped feature names, reverse-mapped event timestamp column, and reverse-mapped created timestamp column that will be passed into the query to the offline store. + """ # if we have mapped fields, use the original field names in the call to the offline store event_timestamp_column = feature_view.inputs.event_timestamp_column entity_names = [entity.name for entity in feature_view.entities] From 36beacda9fe854854d7137f50afb2b9fa16e8ffb Mon Sep 17 00:00:00 2001 From: Jacob Klegar Date: Thu, 11 Mar 2021 17:31:59 -0500 Subject: [PATCH 10/11] Address comments Signed-off-by: Jacob Klegar --- sdk/python/feast/offline_store.py | 36 ++++++++++++++++--------------- 1 file changed, 19 insertions(+), 17 deletions(-) diff --git a/sdk/python/feast/offline_store.py b/sdk/python/feast/offline_store.py index 455edce48f8..54e9cbf08bf 100644 --- a/sdk/python/feast/offline_store.py +++ b/sdk/python/feast/offline_store.py @@ -22,19 +22,13 @@ class OfflineStore(ABC): """ - OfflineStore is a non-user-facing object used for all interaction between Feast and the service used for offline storage of features. Currently BigQuery is supported. + OfflineStore is an object used for all interaction between Feast and the service used for offline storage of features. Currently BigQuery is supported. """ @staticmethod @abstractmethod def pull_latest_from_table( - table_ref: str, - entity_names: List[str], - feature_names: List[str], - event_timestamp_column: str, - created_timestamp_column: Optional[str], - start_date: datetime, - end_date: datetime, + feature_view: FeatureView, start_date: datetime, end_date: datetime, ) -> Optional[pyarrow.Table]: pass @@ -42,14 +36,19 @@ def pull_latest_from_table( class BigQueryOfflineStore(OfflineStore): @staticmethod def pull_latest_from_table( - table_ref: str, - entity_names: List[str], - feature_names: List[str], - event_timestamp_column: str, - created_timestamp_column: Optional[str], - start_date: datetime, - end_date: datetime, + feature_view: FeatureView, start_date: datetime, end_date: datetime, ) -> pyarrow.Table: + if feature_view.inputs.table_ref is None: + raise ValueError( + "This function can only be called on a FeatureView with a table_ref" + ) + + ( + entity_names, + feature_names, + event_timestamp_column, + created_timestamp_column, + ) = run_reverse_field_mapping(feature_view) partition_by_entity_string = ", ".join(entity_names) if partition_by_entity_string != "": @@ -65,12 +64,15 @@ def pull_latest_from_table( FROM ( SELECT {field_string}, ROW_NUMBER() OVER({partition_by_entity_string} ORDER BY {timestamp_desc_string}) AS _feast_row - FROM `{table_ref}` + FROM `{feature_view.inputs.table_ref}` WHERE {event_timestamp_column} BETWEEN TIMESTAMP('{start_date}') AND TIMESTAMP('{end_date}') ) WHERE _feast_row = 1 """ - return BigQueryOfflineStore._pull_query(query) + + table = BigQueryOfflineStore._pull_query(query) + table = run_forward_field_mapping(table, feature_view) + return table @staticmethod def _pull_query(query: str) -> pyarrow.Table: From 85465a1b9ee4de05bcc8c080f43e6a7acb500e63 Mon Sep 17 00:00:00 2001 From: Jacob Klegar Date: Thu, 11 Mar 2021 20:27:21 -0500 Subject: [PATCH 11/11] Rebase Signed-off-by: Jacob Klegar --- sdk/python/feast/feature_store.py | 1 - sdk/python/feast/repo_config.py | 8 ++++---- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/sdk/python/feast/feature_store.py b/sdk/python/feast/feature_store.py index efd46cabff3..00c2befc508 100644 --- a/sdk/python/feast/feature_store.py +++ b/sdk/python/feast/feature_store.py @@ -14,7 +14,6 @@ from pathlib import Path from typing import Optional - from feast.infra.provider import Provider, get_provider from feast.registry import Registry from feast.repo_config import ( diff --git a/sdk/python/feast/repo_config.py b/sdk/python/feast/repo_config.py index af2c1904ca4..aae9c217d19 100644 --- a/sdk/python/feast/repo_config.py +++ b/sdk/python/feast/repo_config.py @@ -19,10 +19,10 @@ class OnlineStoreConfig(NamedTuple): class RepoConfig(NamedTuple): - metadata_store: str = "./metadata_store" - project: str = "default" - provider: str = "local" - online_store: OnlineStoreConfig = OnlineStoreConfig() + metadata_store: str + project: str + provider: str + online_store: OnlineStoreConfig def load_repo_config(repo_path: Path) -> RepoConfig: