Skip to content

Commit c4df8c9

Browse files
khorshuhengkhorshuheng
authored
Support TFRecord as one of the output formats for historical feature retrieval (#1222)
* Support TFRecord as one of the output formats for historical feature retrieval Signed-off-by: Khor Shu Heng <khor.heng@gojek.com> * Fix style Signed-off-by: Khor Shu Heng <khor.heng@gojek.com> * Python style fix Signed-off-by: Khor Shu Heng <khor.heng@gojek.com> * Update tfrecord jar to be Spark 3.0 compatible Signed-off-by: Khor Shu Heng <khor.heng@gojek.com> * Add ability to download extra packages for EMR historical retrieval job Signed-off-by: Khor Shu Heng <khor.heng@gojek.com> * Add ability to download extra packages for k8s historical retrieval job Signed-off-by: Khor Shu Heng <khor.heng@gojek.com> * e2e tests for tfrecord output Signed-off-by: Khor Shu Heng <khor.heng@gojek.com> * Patch copy module so that regex can be deep copied Signed-off-by: Khor Shu Heng <khor.heng@gojek.com> * Use separate fixture for tfrecord feast client Signed-off-by: Khor Shu Heng <khor.heng@gojek.com> * Style fix Signed-off-by: Khor Shu Heng <khor.heng@gojek.com> * Fix aws launcher spark submit arguments Signed-off-by: Khor Shu Heng <khor.heng@gojek.com> * Fix spark submit argument sequence Signed-off-by: Khor Shu Heng <khor.heng@gojek.com> * Update docker image for k8s launcher Signed-off-by: Khor Shu Heng <khor.heng@gojek.com> * Revert image change Signed-off-by: Khor Shu Heng <khor.heng@gojek.com> Co-authored-by: Khor Shu Heng <khor.heng@gojek.com>
1 parent a99005c commit c4df8c9

13 files changed

Lines changed: 254 additions & 18 deletions

File tree

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,15 @@ def get_class_name(self) -> Optional[str]:
104104
"""
105105
return None
106106

107+
def get_extra_packages(self) -> List[str]:
108+
"""
109+
Getter for extra maven packages to be included on driver and executor
110+
classpath if applicable.
111+
Returns:
112+
List[str]: List of maven packages
113+
"""
114+
return []
115+
107116
@abc.abstractmethod
108117
def get_arguments(self) -> List[str]:
109118
"""
@@ -122,6 +131,7 @@ def __init__(
122131
feature_tables_sources: List[Dict],
123132
entity_source: Dict,
124133
destination: Dict,
134+
extra_packages: Optional[List[str]] = None,
125135
):
126136
"""
127137
Args:
@@ -130,6 +140,8 @@ def __init__(
130140
feature_tables (List[Dict]): List of feature table specification.
131141
The order of the feature table must correspond to that of feature_tables_sources.
132142
destination (Dict): Retrieval job output destination.
143+
extra_packages (Optional[List[str]): Extra maven packages to be included on Spark driver
144+
and executors classpath.
133145
134146
Examples:
135147
>>> # Entity source from file
@@ -233,6 +245,7 @@ def __init__(
233245
self._feature_tables_sources = feature_tables_sources
234246
self._entity_source = entity_source
235247
self._destination = destination
248+
self._extra_packages = extra_packages if extra_packages else []
236249

237250
def get_name(self) -> str:
238251
all_feature_tables_names = [ft["name"] for ft in self._feature_tables]
@@ -246,6 +259,9 @@ def get_main_file_path(self) -> str:
246259
os.path.dirname(__file__), "historical_feature_retrieval_job.py"
247260
)
248261

262+
def get_extra_packages(self) -> List[str]:
263+
return self._extra_packages
264+
249265
def get_arguments(self) -> List[str]:
250266
def json_b64_encode(obj) -> str:
251267
return b64encode(json.dumps(obj).encode("utf8")).decode("ascii")

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77

88
from pyspark.sql import DataFrame, SparkSession, Window
99
from pyspark.sql.functions import col, expr, monotonically_increasing_id, row_number
10+
from pyspark.sql.types import LongType
1011

1112
EVENT_TIMESTAMP_ALIAS = "event_timestamp"
1213
CREATED_TIMESTAMP_ALIAS = "created_timestamp"
@@ -728,7 +729,15 @@ def start_job(
728729
result = retrieve_historical_features(
729730
spark, entity_source_conf, feature_tables_sources_conf, feature_tables_conf
730731
)
732+
731733
destination = FileDestination(**destination_conf)
734+
if destination.format == "tfrecord":
735+
entity_source = _source_from_dict(entity_source_conf)
736+
result = result.withColumn(
737+
entity_source.event_timestamp_column,
738+
col(entity_source.event_timestamp_column).cast(LongType()),
739+
)
740+
732741
result.write.format(destination.format).mode("overwrite").save(destination.path)
733742

734743

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -200,6 +200,10 @@ def start_historical_feature_retrieval_job(
200200
for feature_table in feature_tables
201201
]
202202

203+
extra_packages = []
204+
if output_format == "tfrecord":
205+
extra_packages.append("com.linkedin.sparktfrecord:spark-tfrecord_2.12:0.3.0")
206+
203207
return launcher.historical_feature_retrieval(
204208
RetrievalJobParameters(
205209
entity_source=_source_to_argument(entity_source, client._config),
@@ -209,6 +213,7 @@ def start_historical_feature_retrieval_job(
209213
for feature_table in feature_tables
210214
],
211215
destination={"format": output_format, "path": output_path},
216+
extra_packages=extra_packages,
212217
)
213218
)
214219

‎sdk/python/feast/pyspark/launchers/aws/emr.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -231,6 +231,7 @@ def historical_feature_retrieval(
231231
pyspark_script_path,
232232
args=job_params.get_arguments(),
233233
output_file_uri=job_params.get_destination_path(),
234+
packages=job_params.get_extra_packages(),
234235
)
235236

236237
job_ref = self._submit_emr_job(step)

‎sdk/python/feast/pyspark/launchers/aws/emr_utils.py‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -281,7 +281,10 @@ def _cancel_job(emr_client, job: EmrJobRef):
281281

282282

283283
def _historical_retrieval_step(
284-
pyspark_script_path: str, args: List[str], output_file_uri: str,
284+
pyspark_script_path: str,
285+
args: List[str],
286+
output_file_uri: str,
287+
packages: List[str] = None,
285288
) -> Dict[str, Any]:
286289

287290
return {
@@ -297,7 +300,10 @@ def _historical_retrieval_step(
297300
"Value": output_file_uri,
298301
},
299302
],
300-
"Args": ["spark-submit", pyspark_script_path] + args,
303+
"Args": ["spark-submit"]
304+
+ (["--packages", ",".join(packages)] if packages else [])
305+
+ [pyspark_script_path]
306+
+ args,
301307
"Jar": "command-runner.jar",
302308
},
303309
}

‎sdk/python/feast/pyspark/launchers/gcloud/dataproc.py‎

Lines changed: 20 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -274,12 +274,20 @@ def dataproc_submit(
274274
"labels": {self.JOB_TYPE_LABEL_KEY: job_params.get_job_type().name.lower()},
275275
}
276276

277+
maven_package_properties = {
278+
"spark.jars.packages": ",".join(job_params.get_extra_packages())
279+
}
280+
common_properties = {
281+
"spark.executor.instances": self.executor_instances,
282+
"spark.executor.cores": self.executor_cores,
283+
"spark.executor.memory": self.executor_memory,
284+
}
277285
# Add job hash to labels only for the stream ingestion job
278286
if isinstance(job_params, StreamIngestionJobParameters):
279287
job_config["labels"][self.JOB_HASH_LABEL_KEY] = job_params.get_job_hash()
280288

281289
if job_params.get_class_name():
282-
properties = {
290+
scala_job_properties = {
283291
"spark.yarn.user.classpath.first": "true",
284292
"spark.executor.instances": self.executor_instances,
285293
"spark.executor.cores": self.executor_cores,
@@ -288,15 +296,18 @@ def dataproc_submit(
288296
"spark.pyspark.python": "python3.7",
289297
}
290298

291-
properties.update(extra_properties)
292-
293299
job_config.update(
294300
{
295301
"spark_job": {
296302
"jar_file_uris": [main_file_uri] + self.EXTERNAL_JARS,
297303
"main_class": job_params.get_class_name(),
298304
"args": job_params.get_arguments(),
299-
"properties": properties,
305+
"properties": {
306+
**scala_job_properties,
307+
**common_properties,
308+
**maven_package_properties,
309+
**extra_properties,
310+
},
300311
}
301312
}
302313
)
@@ -307,7 +318,11 @@ def dataproc_submit(
307318
"main_python_file_uri": main_file_uri,
308319
"jar_file_uris": self.EXTERNAL_JARS,
309320
"args": job_params.get_arguments(),
310-
"properties": extra_properties if extra_properties else {},
321+
"properties": {
322+
**common_properties,
323+
**maven_package_properties,
324+
**extra_properties,
325+
},
311326
}
312327
}
313328
)

‎sdk/python/feast/pyspark/launchers/standalone/local.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -272,7 +272,7 @@ def spark_submit(
272272
"--conf",
273273
"spark.sql.session.timeZone=UTC", # ignore local timezone
274274
"--packages",
275-
BQ_SPARK_PACKAGE,
275+
",".join([BQ_SPARK_PACKAGE] + job_params.get_extra_packages()),
276276
"--jars",
277277
"https://storage.googleapis.com/hadoop-lib/gcs/gcs-connector-hadoop2-latest.jar,"
278278
"https://repo1.maven.org/maven2/org/apache/hadoop/hadoop-aws/2.7.3/hadoop-aws-2.7.3.jar,"

‎sdk/python/tests/test_historical_feature_retrieval.py‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
from feast import Client, Entity, Feature, FeatureTable, FileSource, ValueType
2929
from feast.core import CoreService_pb2_grpc as Core
3030
from feast.data_format import ParquetFormat
31+
from feast.pyspark.abc import SparkJobStatus
3132
from tests.feast_core_server import CoreServicer
3233

3334

@@ -107,6 +108,26 @@ def client_with_local_spark(tmpdir):
107108
)
108109

109110

111+
@pytest.fixture()
112+
def client_with_tfrecord_output(tmpdir):
113+
import pyspark
114+
115+
spark_staging_location = f"file://{os.path.join(tmpdir, 'staging')}"
116+
historical_feature_output_location = (
117+
f"file://{os.path.join(tmpdir, 'historical_feature_retrieval_tfrecord_output')}"
118+
)
119+
120+
return Client(
121+
core_url=f"localhost:{free_port}",
122+
spark_launcher="standalone",
123+
spark_standalone_master="local",
124+
spark_home=os.path.dirname(pyspark.__file__),
125+
spark_staging_location=spark_staging_location,
126+
historical_feature_output_location=historical_feature_output_location,
127+
historical_feature_output_format="tfrecord",
128+
)
129+
130+
110131
@pytest.fixture()
111132
def driver_entity(client):
112133
return client.apply(Entity("driver_id", "description", ValueType.INT32))
@@ -466,3 +487,39 @@ def test_historical_feature_retrieval_with_pandas_dataframe_input(
466487
by=["customer_id", "driver_id", "event_timestamp"]
467488
).reset_index(drop=True),
468489
)
490+
491+
492+
@pytest.mark.usefixtures(
493+
"driver_entity",
494+
"customer_entity",
495+
"bookings_feature_table",
496+
"transactions_feature_table",
497+
)
498+
def test_historical_feature_retrieval_with_tfrecord_output(
499+
client_with_tfrecord_output,
500+
):
501+
502+
customer_driver_pairs_pandas_df = pd.DataFrame(
503+
np.array(
504+
[
505+
[1001, 8001, datetime(year=2020, month=9, day=1, tzinfo=utc)],
506+
[2001, 8001, datetime(year=2020, month=9, day=2, tzinfo=utc)],
507+
[2001, 8002, datetime(year=2020, month=9, day=1, tzinfo=utc)],
508+
[1001, 8001, datetime(year=2020, month=9, day=2, tzinfo=utc)],
509+
[1001, 8001, datetime(year=2020, month=9, day=3, tzinfo=utc)],
510+
[1001, 8001, datetime(year=2020, month=9, day=4, tzinfo=utc)],
511+
]
512+
),
513+
columns=["customer_id", "driver_id", "event_timestamp"],
514+
)
515+
customer_driver_pairs_pandas_df = customer_driver_pairs_pandas_df.astype(
516+
{"customer_id": "int32", "driver_id": "int32"}
517+
)
518+
519+
job_output = client_with_tfrecord_output.get_historical_features(
520+
["transactions:total_transactions", "bookings:total_completed_bookings"],
521+
customer_driver_pairs_pandas_df,
522+
)
523+
524+
job_output.get_output_file_uri()
525+
assert job_output.get_status() == SparkJobStatus.COMPLETED

‎tests/e2e/conftest.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ def pytest_runtest_setup(item):
4343
global_staging_path,
4444
ingestion_job_jar,
4545
local_staging_path,
46+
tfrecord_feast_client,
4647
)
4748

4849
if not os.environ.get("DISABLE_SERVICE_FIXTURES"):

‎tests/e2e/fixtures/client.py‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,82 @@ def feast_client(
103103
return c
104104

105105

106+
@pytest.fixture
107+
def tfrecord_feast_client(
108+
pytestconfig,
109+
feast_core: Tuple[str, int],
110+
local_staging_path,
111+
feast_jobservice: Optional[Tuple[str, int]],
112+
enable_auth,
113+
):
114+
if feast_jobservice is None:
115+
job_service_env = dict()
116+
else:
117+
job_service_env = dict(
118+
job_service_url=f"{feast_jobservice[0]}:{feast_jobservice[1]}"
119+
)
120+
121+
if pytestconfig.getoption("env") == "local":
122+
import pyspark
123+
124+
return Client(
125+
core_url=f"{feast_core[0]}:{feast_core[1]}",
126+
spark_launcher="standalone",
127+
spark_standalone_master="local",
128+
spark_home=os.getenv("SPARK_HOME") or os.path.dirname(pyspark.__file__),
129+
spark_staging_location=os.path.join(local_staging_path, "spark"),
130+
historical_feature_output_format="tfrecord",
131+
historical_feature_output_location=os.path.join(
132+
local_staging_path, "historical_output"
133+
),
134+
**job_service_env,
135+
)
136+
137+
elif pytestconfig.getoption("env") == "gcloud":
138+
c = Client(
139+
core_url=f"{feast_core[0]}:{feast_core[1]}",
140+
spark_launcher="dataproc",
141+
dataproc_cluster_name=pytestconfig.getoption("dataproc_cluster_name"),
142+
dataproc_project=pytestconfig.getoption("dataproc_project"),
143+
dataproc_region=pytestconfig.getoption("dataproc_region"),
144+
spark_staging_location=os.path.join(local_staging_path, "dataproc"),
145+
historical_feature_output_format="tfrecord",
146+
historical_feature_output_location=os.path.join(
147+
local_staging_path, "historical_output"
148+
),
149+
ingestion_drop_invalid_rows=True,
150+
**job_service_env,
151+
)
152+
elif pytestconfig.getoption("env") == "aws":
153+
return Client(
154+
core_url=f"{feast_core[0]}:{feast_core[1]}",
155+
spark_launcher="emr",
156+
emr_cluster_id=pytestconfig.getoption("emr_cluster_id"),
157+
emr_region=pytestconfig.getoption("emr_region"),
158+
spark_staging_location=os.path.join(local_staging_path, "emr"),
159+
emr_log_location=os.path.join(local_staging_path, "emr_logs"),
160+
historical_feature_output_format="tfrecord",
161+
historical_feature_output_location=os.path.join(
162+
local_staging_path, "historical_output"
163+
),
164+
)
165+
elif pytestconfig.getoption("env") == "k8s":
166+
return Client(
167+
core_url=f"{feast_core[0]}:{feast_core[1]}",
168+
spark_launcher="k8s",
169+
spark_staging_location=os.path.join(local_staging_path, "k8s"),
170+
historical_feature_output_format="tfrecord",
171+
historical_feature_output_location=os.path.join(
172+
local_staging_path, "historical_output"
173+
),
174+
)
175+
else:
176+
raise KeyError(f"Unknown environment {pytestconfig.getoption('env')}")
177+
178+
c.set_project(pytestconfig.getoption("feast_project"))
179+
return c
180+
181+
106182
@pytest.fixture(scope="session")
107183
def global_staging_path(pytestconfig):
108184
if pytestconfig.getoption("env") == "local" and not pytestconfig.getoption(

0 commit comments

Comments
 (0)