Skip to content

Commit e8b24bb

Browse files
khorshuhengkhorshuheng
authored
Feast SDK integration for historical feature retrieval using Spark (#1054)
* Feast SDK integration for historical feature retrieval using Spark Signed-off-by: Khor Shu Heng <khor.heng@gojek.com> * Downgrade pyspark dependencies, add more tests Signed-off-by: Khor Shu Heng <khor.heng@gojek.com> * Don't expose the historical feature output config directly to the user Signed-off-by: Khor Shu Heng <khor.heng@gojek.com> Co-authored-by: Khor Shu Heng <khor.heng@gojek.com>
1 parent 5e9a717 commit e8b24bb

17 files changed

Lines changed: 1305 additions & 591 deletions

‎sdk/python/feast/client.py‎

Lines changed: 119 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,8 @@
1414
import logging
1515
import multiprocessing
1616
import shutil
17+
import uuid
18+
from itertools import groupby
1719
from typing import Any, Dict, List, Optional, Union
1820

1921
import grpc
@@ -30,6 +32,8 @@
3032
CONFIG_SERVING_ENABLE_SSL_KEY,
3133
CONFIG_SERVING_SERVER_SSL_CERT_KEY,
3234
CONFIG_SERVING_URL_KEY,
35+
CONFIG_SPARK_HISTORICAL_FEATURE_OUTPUT_FORMAT,
36+
CONFIG_SPARK_HISTORICAL_FEATURE_OUTPUT_LOCATION,
3337
FEAST_DEFAULT_OPTIONS,
3438
)
3539
from feast.core.CoreService_pb2 import (
@@ -70,6 +74,11 @@
7074
_write_partitioned_table_from_source,
7175
)
7276
from feast.online_response import OnlineResponse, _infer_online_entity_rows
77+
from feast.pyspark.abc import RetrievalJob
78+
from feast.pyspark.launcher import (
79+
start_historical_feature_retrieval_job,
80+
start_historical_feature_retrieval_spark_session,
81+
)
7382
from feast.serving.ServingService_pb2 import (
7483
GetFeastServingInfoRequest,
7584
GetOnlineFeaturesRequestV2,
@@ -723,7 +732,6 @@ def get_online_features(
723732
) -> OnlineResponse:
724733
"""
725734
Retrieves the latest online feature data from Feast Serving.
726-
727735
Args:
728736
feature_refs: List of feature references that will be returned for each entity.
729737
Each feature reference should have the following format:
@@ -733,12 +741,10 @@ def get_online_features(
733741
entity_rows: A list of dictionaries where each key-value is an entity-name, entity-value pair.
734742
project: Optionally specify the the project override. If specified, uses given project for retrieval.
735743
Overrides the projects specified in Feature References if also are specified.
736-
737744
Returns:
738745
GetOnlineFeaturesResponse containing the feature data in records.
739746
Each EntityRow provided will yield one record, which contains
740747
data fields with data value and field status metadata (if included).
741-
742748
Examples:
743749
>>> from feast import Client
744750
>>>
@@ -767,3 +773,113 @@ def get_online_features(
767773

768774
response = OnlineResponse(response)
769775
return response
776+
777+
def get_historical_features(
778+
self,
779+
feature_refs: List[str],
780+
entity_source: Union[FileSource, BigQuerySource],
781+
project: str = None,
782+
) -> RetrievalJob:
783+
"""
784+
Launch a historical feature retrieval job.
785+
786+
Args:
787+
feature_refs: List of feature references that will be returned for each entity.
788+
Each feature reference should have the following format:
789+
"feature_table:feature" where "feature_table" & "feature" refer to
790+
the feature and feature table names respectively.
791+
entity_source (Union[FileSource, BigQuerySource]): Source for the entity rows.
792+
The user needs to make sure that the source is accessible from the Spark cluster
793+
that will be used for the retrieval job.
794+
project: Specifies the project that contains the feature tables
795+
which the requested features belong to.
796+
797+
Returns:
798+
Returns a retrieval job object that can be used to monitor retrieval
799+
progress asynchronously, and can be used to materialize the
800+
results.
801+
802+
Examples:
803+
>>> from feast import Client
804+
>>> from datetime import datetime
805+
>>> feast_client = Client(core_url="localhost:6565")
806+
>>> feature_refs = ["bookings:bookings_7d", "bookings:booking_14d"]
807+
>>> entity_source = FileSource("event_timestamp", "parquet", "gs://some-bucket/customer")
808+
>>> feature_retrieval_job = feast_client.get_historical_features(
809+
>>> feature_refs, entity_source, project="my_project")
810+
>>> output_file_uri = feature_retrieval_job.get_output_file_uri()
811+
"gs://some-bucket/output/
812+
"""
813+
feature_tables = self._get_feature_tables_from_feature_refs(
814+
feature_refs, project
815+
)
816+
output_location = self._config.get(
817+
CONFIG_SPARK_HISTORICAL_FEATURE_OUTPUT_LOCATION
818+
)
819+
output_format = self._config.get(CONFIG_SPARK_HISTORICAL_FEATURE_OUTPUT_FORMAT)
820+
job_id = f"historical-feature-{str(uuid.uuid4())}"
821+
822+
return start_historical_feature_retrieval_job(
823+
self, entity_source, feature_tables, output_format, output_location, job_id
824+
)
825+
826+
def get_historical_features_df(
827+
self,
828+
feature_refs: List[str],
829+
entity_source: Union[FileSource, BigQuerySource],
830+
project: str = None,
831+
):
832+
"""
833+
Launch a historical feature retrieval job.
834+
835+
Args:
836+
feature_refs: List of feature references that will be returned for each entity.
837+
Each feature reference should have the following format:
838+
"feature_table:feature" where "feature_table" & "feature" refer to
839+
the feature and feature table names respectively.
840+
entity_source (Union[FileSource, BigQuerySource]): Source for the entity rows.
841+
The user needs to make sure that the source is accessible from the Spark cluster
842+
that will be used for the retrieval job.
843+
project: Specifies the project that contains the feature tables
844+
which the requested features belong to.
845+
846+
Returns:
847+
Returns the historical feature retrieval result in the form of Spark dataframe.
848+
849+
Examples:
850+
>>> from feast import Client
851+
>>> from datetime import datetime
852+
>>> from pyspark.sql import SparkSession
853+
>>> spark = SparkSession.builder.getOrCreate()
854+
>>> feast_client = Client(core_url="localhost:6565")
855+
>>> feature_refs = ["bookings:bookings_7d", "bookings:booking_14d"]
856+
>>> entity_source = FileSource("event_timestamp", "parquet", "gs://some-bucket/customer")
857+
>>> df = feast_client.get_historical_features(
858+
>>> feature_refs, entity_source, project="my_project")
859+
"""
860+
feature_tables = self._get_feature_tables_from_feature_refs(
861+
feature_refs, project
862+
)
863+
return start_historical_feature_retrieval_spark_session(
864+
self, entity_source, feature_tables
865+
)
866+
867+
def _get_feature_tables_from_feature_refs(
868+
self, feature_refs: List[str], project: Optional[str]
869+
):
870+
feature_refs_grouped_by_table = [
871+
(feature_table_name, list(grouped_feature_refs))
872+
for feature_table_name, grouped_feature_refs in groupby(
873+
feature_refs, lambda x: x.split(":")[0]
874+
)
875+
]
876+
877+
feature_tables = []
878+
for feature_table_name, grouped_feature_refs in feature_refs_grouped_by_table:
879+
feature_table = self.get_feature_table(feature_table_name, project)
880+
feature_names = [f.split(":")[-1] for f in grouped_feature_refs]
881+
feature_table.features = [
882+
f for f in feature_table.features if f.name in feature_names
883+
]
884+
feature_tables.append(feature_table)
885+
return feature_tables

‎sdk/python/feast/constants.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,20 @@ class AuthProvider(Enum):
6464
CONFIG_TIMEOUT_KEY = "timeout"
6565
CONFIG_MAX_WAIT_INTERVAL_KEY = "max_wait_interval"
6666

67+
# Spark Job Config
68+
CONFIG_SPARK_LAUNCHER = "spark_launcher" # standalone, dataproc, emr
69+
70+
CONFIG_SPARK_STANDALONE_MASTER = "spark_standalone_master"
71+
72+
CONFIG_SPARK_DATAPROC_CLUSTER_NAME = "dataproc_cluster_name"
73+
CONFIG_SPARK_DATAPROC_PROJECT = "dataproc_project"
74+
CONFIG_SPARK_DATAPROC_REGION = "dataproc_region"
75+
CONFIG_SPARK_DATAPROC_STAGING_LOCATION = "dataproc_staging_location"
76+
77+
CONFIG_SPARK_HISTORICAL_FEATURE_OUTPUT_FORMAT = "historical_feature_output_format"
78+
CONFIG_SPARK_HISTORICAL_FEATURE_OUTPUT_LOCATION = "historical_feature_output_location"
79+
80+
6781
# Configuration option default values
6882
FEAST_DEFAULT_OPTIONS = {
6983
# Default Feast project to use

‎sdk/python/feast/pyspark/abc.py‎

Lines changed: 193 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,193 @@
1+
import abc
2+
from typing import Dict, List
3+
4+
5+
class SparkJobFailure(Exception):
6+
"""
7+
Job submission failed, encountered error during execution, or timeout
8+
"""
9+
10+
pass
11+
12+
13+
class SparkJob(abc.ABC):
14+
"""
15+
Base class for all spark jobs
16+
"""
17+
18+
@abc.abstractmethod
19+
def get_id(self) -> str:
20+
"""
21+
Getter for the job id. The job id must be unique for each spark job submission.
22+
23+
Returns:
24+
str: Job id.
25+
"""
26+
raise NotImplementedError
27+
28+
29+
class RetrievalJob(SparkJob):
30+
"""
31+
Container for the historical feature retrieval job result
32+
"""
33+
34+
@abc.abstractmethod
35+
def get_output_file_uri(self, timeout_sec=None):
36+
"""
37+
Get output file uri to the result file. This method will block until the
38+
job succeeded, or if the job didn't execute successfully within timeout.
39+
40+
Args:
41+
timeout_sec (int):
42+
Max no of seconds to wait until job is done. If "timeout_sec"
43+
is exceeded or if the job fails, an exception will be raised.
44+
45+
Raises:
46+
SparkJobFailure:
47+
The spark job submission failed, encountered error during execution,
48+
or timeout.
49+
50+
Returns:
51+
str: file uri to the result file.
52+
"""
53+
raise NotImplementedError
54+
55+
56+
class IngestionJob(SparkJob):
57+
pass
58+
59+
60+
class JobLauncher(abc.ABC):
61+
"""
62+
Submits spark jobs to a spark cluster. Currently supports only historical feature retrieval jobs.
63+
"""
64+
65+
@abc.abstractmethod
66+
def historical_feature_retrieval(
67+
self,
68+
pyspark_script: str,
69+
entity_source_conf: Dict,
70+
feature_tables_sources_conf: List[Dict],
71+
feature_tables_conf: List[Dict],
72+
destination_conf: Dict,
73+
job_id: str,
74+
**kwargs,
75+
) -> RetrievalJob:
76+
"""
77+
Submits a historical feature retrieval job to a Spark cluster.
78+
79+
Args:
80+
pyspark_script (str): Local file path to the pyspark script for historical feature
81+
retrieval.
82+
entity_source_conf (Dict): Entity data source configuration.
83+
feature_tables_sources_conf (List[Dict]): List of feature tables data sources configurations.
84+
feature_tables_conf (List[Dict]): List of feature table specification.
85+
The order of the feature table must correspond to that of feature_tables_sources.
86+
destination_conf (Dict): Retrieval job output destination.
87+
job_id (str): A job id that is unique for each job submission.
88+
89+
Raises:
90+
SparkJobFailure: The spark job submission failed, encountered error
91+
during execution, or timeout.
92+
93+
Examples:
94+
>>> # Entity source from file
95+
>>> entity_source_conf = {
96+
"file": {
97+
"format": "parquet",
98+
"path": "gs://some-gcs-bucket/customer",
99+
"event_timestamp_column": "event_timestamp",
100+
"options": {
101+
"mergeSchema": "true"
102+
} # Optional. Options to be passed to Spark while reading the dataframe from source.
103+
"field_mapping": {
104+
"id": "customer_id"
105+
} # Optional. Map the columns, where the key is the original column name and the value is the new column name.
106+
107+
}
108+
}
109+
110+
>>> # Entity source from BigQuery
111+
>>> entity_source_conf = {
112+
"bq": {
113+
"project": "gcp_project_id",
114+
"dataset": "bq_dataset",
115+
"table": "customer",
116+
"event_timestamp_column": "event_timestamp",
117+
}
118+
}
119+
120+
>>> feature_table_sources_conf = [
121+
{
122+
"bq": {
123+
"project": "gcp_project_id",
124+
"dataset": "bq_dataset",
125+
"table": "customer_transactions",
126+
"event_timestamp_column": "event_timestamp",
127+
"created_timestamp_column": "created_timestamp" # This field is mandatory for feature tables.
128+
}
129+
},
130+
131+
{
132+
"file": {
133+
"format": "parquet",
134+
"path": "gs://some-gcs-bucket/customer_profile",
135+
"event_timestamp_column": "event_timestamp",
136+
"created_timestamp_column": "created_timestamp",
137+
"options": {
138+
"mergeSchema": "true"
139+
}
140+
}
141+
},
142+
]
143+
144+
145+
>>> feature_tables_conf = [
146+
{
147+
"name": "customer_transactions",
148+
"entities": [
149+
{
150+
"name": "customer
151+
"type": "int32"
152+
}
153+
],
154+
"features": [
155+
{
156+
"name": "total_transactions"
157+
"type": "double"
158+
},
159+
{
160+
"name": "total_discounts"
161+
"type": "double"
162+
}
163+
],
164+
"max_age": 86400 # In seconds.
165+
},
166+
167+
{
168+
"name": "customer_profile",
169+
"entities": [
170+
{
171+
"name": "customer
172+
"type": "int32"
173+
}
174+
],
175+
"features": [
176+
{
177+
"name": "is_vip"
178+
"type": "bool"
179+
}
180+
],
181+
182+
}
183+
]
184+
185+
>>> destination_conf = {
186+
"format": "parquet",
187+
"path": "gs://some-gcs-bucket/retrieval_output"
188+
}
189+
190+
Returns:
191+
str: file uri to the result file.
192+
"""
193+
raise NotImplementedError

0 commit comments

Comments
 (0)