Skip to content

Commit 39c9fbe

Browse files
pyalexpyalex
authored
Refactor Spark Job launcher API (#1060)
* simplify launcher API Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com> * uniq destination property with others Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com>
1 parent 065b310 commit 39c9fbe

4 files changed

Lines changed: 152 additions & 195 deletions

File tree

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

Lines changed: 119 additions & 119 deletions
Original file line numberDiff line numberDiff line change
@@ -86,8 +86,113 @@ def __init__(
8686
feature_tables_sources: List[Dict],
8787
entity_source: Dict,
8888
destination: Dict,
89-
**kwargs,
9089
):
90+
"""
91+
Args:
92+
entity_source (Dict): Entity data source configuration.
93+
feature_tables_sources (List[Dict]): List of feature tables data sources configurations.
94+
feature_tables (List[Dict]): List of feature table specification.
95+
The order of the feature table must correspond to that of feature_tables_sources.
96+
destination (Dict): Retrieval job output destination.
97+
98+
Examples:
99+
>>> # Entity source from file
100+
>>> entity_source = {
101+
"file": {
102+
"format": "parquet",
103+
"path": "gs://some-gcs-bucket/customer",
104+
"event_timestamp_column": "event_timestamp",
105+
"options": {
106+
"mergeSchema": "true"
107+
} # Optional. Options to be passed to Spark while reading the dataframe from source.
108+
"field_mapping": {
109+
"id": "customer_id"
110+
} # Optional. Map the columns, where the key is the original column name and the value is the new column name.
111+
112+
}
113+
}
114+
115+
>>> # Entity source from BigQuery
116+
>>> entity_source = {
117+
"bq": {
118+
"project": "gcp_project_id",
119+
"dataset": "bq_dataset",
120+
"table": "customer",
121+
"event_timestamp_column": "event_timestamp",
122+
}
123+
}
124+
125+
>>> feature_tables_sources = [
126+
{
127+
"bq": {
128+
"project": "gcp_project_id",
129+
"dataset": "bq_dataset",
130+
"table": "customer_transactions",
131+
"event_timestamp_column": "event_timestamp",
132+
"created_timestamp_column": "created_timestamp" # This field is mandatory for feature tables.
133+
}
134+
},
135+
136+
{
137+
"file": {
138+
"format": "parquet",
139+
"path": "gs://some-gcs-bucket/customer_profile",
140+
"event_timestamp_column": "event_timestamp",
141+
"created_timestamp_column": "created_timestamp",
142+
"options": {
143+
"mergeSchema": "true"
144+
}
145+
}
146+
},
147+
]
148+
149+
150+
>>> feature_tables = [
151+
{
152+
"name": "customer_transactions",
153+
"entities": [
154+
{
155+
"name": "customer
156+
"type": "int32"
157+
}
158+
],
159+
"features": [
160+
{
161+
"name": "total_transactions"
162+
"type": "double"
163+
},
164+
{
165+
"name": "total_discounts"
166+
"type": "double"
167+
}
168+
],
169+
"max_age": 86400 # In seconds.
170+
},
171+
172+
{
173+
"name": "customer_profile",
174+
"entities": [
175+
{
176+
"name": "customer
177+
"type": "int32"
178+
}
179+
],
180+
"features": [
181+
{
182+
"name": "is_vip"
183+
"type": "bool"
184+
}
185+
],
186+
187+
}
188+
]
189+
190+
>>> destination = {
191+
"format": "parquet",
192+
"path": "gs://some-gcs-bucket/retrieval_output"
193+
}
194+
195+
"""
91196
self._feature_tables = feature_tables
92197
self._feature_tables_sources = feature_tables_sources
93198
self._entity_source = entity_source
@@ -114,6 +219,9 @@ def get_arguments(self) -> List[str]:
114219
json.dumps(self._destination),
115220
]
116221

222+
def get_destination_path(self) -> str:
223+
return self._destination["path"]
224+
117225

118226
class RetrievalJob(SparkJob):
119227
"""
@@ -150,7 +258,6 @@ def __init__(
150258
start: datetime,
151259
end: datetime,
152260
jar: str,
153-
**kwargs,
154261
):
155262
self._feature_table = feature_table
156263
self._source = source
@@ -198,139 +305,32 @@ class JobLauncher(abc.ABC):
198305

199306
@abc.abstractmethod
200307
def historical_feature_retrieval(
201-
self,
202-
entity_source_conf: Dict,
203-
feature_tables_sources_conf: List[Dict],
204-
feature_tables_conf: List[Dict],
205-
destination_conf: Dict,
206-
**kwargs,
308+
self, retrieval_job_params: RetrievalJobParameters
207309
) -> RetrievalJob:
208310
"""
209311
Submits a historical feature retrieval job to a Spark cluster.
210312
211-
Args:
212-
entity_source_conf (Dict): Entity data source configuration.
213-
feature_tables_sources_conf (List[Dict]): List of feature tables data sources configurations.
214-
feature_tables_conf (List[Dict]): List of feature table specification.
215-
The order of the feature table must correspond to that of feature_tables_sources.
216-
destination_conf (Dict): Retrieval job output destination.
217-
218313
Raises:
219314
SparkJobFailure: The spark job submission failed, encountered error
220315
during execution, or timeout.
221316
222-
Examples:
223-
>>> # Entity source from file
224-
>>> entity_source_conf = {
225-
"file": {
226-
"format": "parquet",
227-
"path": "gs://some-gcs-bucket/customer",
228-
"event_timestamp_column": "event_timestamp",
229-
"options": {
230-
"mergeSchema": "true"
231-
} # Optional. Options to be passed to Spark while reading the dataframe from source.
232-
"field_mapping": {
233-
"id": "customer_id"
234-
} # Optional. Map the columns, where the key is the original column name and the value is the new column name.
235-
236-
}
237-
}
238-
239-
>>> # Entity source from BigQuery
240-
>>> entity_source_conf = {
241-
"bq": {
242-
"project": "gcp_project_id",
243-
"dataset": "bq_dataset",
244-
"table": "customer",
245-
"event_timestamp_column": "event_timestamp",
246-
}
247-
}
248-
249-
>>> feature_table_sources_conf = [
250-
{
251-
"bq": {
252-
"project": "gcp_project_id",
253-
"dataset": "bq_dataset",
254-
"table": "customer_transactions",
255-
"event_timestamp_column": "event_timestamp",
256-
"created_timestamp_column": "created_timestamp" # This field is mandatory for feature tables.
257-
}
258-
},
259-
260-
{
261-
"file": {
262-
"format": "parquet",
263-
"path": "gs://some-gcs-bucket/customer_profile",
264-
"event_timestamp_column": "event_timestamp",
265-
"created_timestamp_column": "created_timestamp",
266-
"options": {
267-
"mergeSchema": "true"
268-
}
269-
}
270-
},
271-
]
272-
273-
274-
>>> feature_tables_conf = [
275-
{
276-
"name": "customer_transactions",
277-
"entities": [
278-
{
279-
"name": "customer
280-
"type": "int32"
281-
}
282-
],
283-
"features": [
284-
{
285-
"name": "total_transactions"
286-
"type": "double"
287-
},
288-
{
289-
"name": "total_discounts"
290-
"type": "double"
291-
}
292-
],
293-
"max_age": 86400 # In seconds.
294-
},
295-
296-
{
297-
"name": "customer_profile",
298-
"entities": [
299-
{
300-
"name": "customer
301-
"type": "int32"
302-
}
303-
],
304-
"features": [
305-
{
306-
"name": "is_vip"
307-
"type": "bool"
308-
}
309-
],
310-
311-
}
312-
]
313-
314-
>>> destination_conf = {
315-
"format": "parquet",
316-
"path": "gs://some-gcs-bucket/retrieval_output"
317-
}
318-
319317
Returns:
320-
str: file uri to the result file.
318+
RetrievalJob: wrapper around remote job that returns file uri to the result file.
321319
"""
322320
raise NotImplementedError
323321

324322
@abc.abstractmethod
325323
def offline_to_online_ingestion(
326-
self,
327-
jar_path: str,
328-
source_conf: Dict,
329-
feature_table_conf: Dict,
330-
start: datetime,
331-
end: datetime,
324+
self, ingestion_job_params: IngestionJobParameters
332325
) -> IngestionJob:
333326
"""
334327
Submits a batch ingestion job to a Spark cluster.
328+
329+
Raises:
330+
SparkJobFailure: The spark job submission failed, encountered error
331+
during execution, or timeout.
332+
333+
Returns:
334+
IngestionJob: wrapper around remote job that can be used to check when job completed.
335335
"""
336336
raise NotImplementedError

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

Lines changed: 26 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,13 @@
1717
)
1818
from feast.data_source import BigQuerySource, DataSource, FileSource
1919
from feast.feature_table import FeatureTable
20-
from feast.pyspark.abc import IngestionJob, JobLauncher, RetrievalJob
20+
from feast.pyspark.abc import (
21+
IngestionJob,
22+
IngestionJobParameters,
23+
JobLauncher,
24+
RetrievalJob,
25+
RetrievalJobParameters,
26+
)
2127
from feast.staging.storage_client import get_staging_client
2228
from feast.value_type import ValueType
2329

@@ -129,16 +135,18 @@ def start_historical_feature_retrieval_job(
129135
) -> RetrievalJob:
130136
launcher = resolve_launcher(client._config)
131137
return launcher.historical_feature_retrieval(
132-
entity_source_conf=_source_to_argument(entity_source),
133-
feature_tables_sources_conf=[
134-
_source_to_argument(feature_table.batch_source)
135-
for feature_table in feature_tables
136-
],
137-
feature_tables_conf=[
138-
_feature_table_to_argument(client, feature_table)
139-
for feature_table in feature_tables
140-
],
141-
destination_conf={"format": output_format, "path": output_path},
138+
RetrievalJobParameters(
139+
entity_source=_source_to_argument(entity_source),
140+
feature_tables_sources=[
141+
_source_to_argument(feature_table.batch_source)
142+
for feature_table in feature_tables
143+
],
144+
feature_tables=[
145+
_feature_table_to_argument(client, feature_table)
146+
for feature_table in feature_tables
147+
],
148+
destination={"format": output_format, "path": output_path},
149+
)
142150
)
143151

144152

@@ -163,9 +171,11 @@ def start_offline_to_online_ingestion(
163171
local_jar_path = _download_jar(client._config.get(CONFIG_SPARK_INGESTION_JOB_JAR))
164172

165173
return launcher.offline_to_online_ingestion(
166-
jar_path=local_jar_path,
167-
source_conf=_source_to_argument(feature_table.batch_source),
168-
feature_table_conf=_feature_table_to_argument(client, feature_table),
169-
start=start,
170-
end=end,
174+
IngestionJobParameters(
175+
jar=local_jar_path,
176+
source=_source_to_argument(feature_table.batch_source),
177+
feature_table=_feature_table_to_argument(client, feature_table),
178+
start=start,
179+
end=end,
180+
)
171181
)

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

Lines changed: 4 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
import os
22
import uuid
3-
from datetime import datetime
4-
from typing import Dict, List, cast
3+
from typing import cast
54
from urllib.parse import urlparse
65

76
from google.api_core.operation import Operation
@@ -144,37 +143,13 @@ def dataproc_submit(self, job_params: SparkJobParameters) -> Operation:
144143
)
145144

146145
def historical_feature_retrieval(
147-
self,
148-
entity_source_conf: Dict,
149-
feature_tables_sources_conf: List[Dict],
150-
feature_tables_conf: List[Dict],
151-
destination_conf: Dict,
152-
**kwargs,
146+
self, job_params: RetrievalJobParameters
153147
) -> RetrievalJob:
154-
job_params = RetrievalJobParameters(
155-
feature_tables=feature_tables_conf,
156-
feature_tables_sources=feature_tables_sources_conf,
157-
entity_source=entity_source_conf,
158-
destination=destination_conf,
159-
)
160-
161148
return DataprocRetrievalJob(
162-
self.dataproc_submit(job_params), destination_conf["path"]
149+
self.dataproc_submit(job_params), job_params.get_destination_path()
163150
)
164151

165152
def offline_to_online_ingestion(
166-
self,
167-
jar_path: str,
168-
source_conf: Dict,
169-
feature_table_conf: Dict,
170-
start: datetime,
171-
end: datetime,
153+
self, job_params: IngestionJobParameters
172154
) -> IngestionJob:
173-
job_params = IngestionJobParameters(
174-
feature_table=feature_table_conf,
175-
source=source_conf,
176-
start=start,
177-
end=end,
178-
jar=jar_path,
179-
)
180155
return DataprocIngestionJob(self.dataproc_submit(job_params))

0 commit comments

Comments
 (0)