Skip to content

Commit 7c5025d

Browse files
khorshuhengkhorshuheng
committed
Use dictionary instead of class to avoid mandatory pyspark dependencies for Feast SDK
Signed-off-by: Khor Shu Heng <khor.heng@gojek.com>
1 parent f29f326 commit 7c5025d

1 file changed

Lines changed: 39 additions & 55 deletions

File tree

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

Lines changed: 39 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -2,16 +2,9 @@
22
import json
33
import os
44
import subprocess
5-
from typing import List
5+
from typing import Dict, List
66
from urllib.parse import urlparse
77

8-
from feast.pyspark.historical_feature_retrieval_job import (
9-
EntitySource,
10-
FeatureTable,
11-
FeatureTableSource,
12-
FileDestination,
13-
)
14-
158

169
class SparkJobFailure(Exception):
1710
"""
@@ -68,12 +61,9 @@ def __init__(self, job_id: str, process: subprocess.Popen, output_file_uri: str)
6861
This is the returned historical feature retrieval job result for StandaloneClusterLauncher.
6962
7063
Args:
71-
job_id (str):
72-
Historical feature retrieval job id.
73-
process (subprocess.Popen):
74-
Pyspark driver process, spawned by the launcher.
75-
output_file_uri (str):
76-
Uri to the historical feature retrieval job output file.
64+
job_id (str): Historical feature retrieval job id.
65+
process (subprocess.Popen): Pyspark driver process, spawned by the launcher.
66+
output_file_uri (str): Uri to the historical feature retrieval job output file.
7767
"""
7868
self.job_id = job_id
7969
self._process = process
@@ -108,12 +98,10 @@ def __init__(self, job_id, operation, output_file_uri):
10898
This is the returned historical feature retrieval job result for DataprocClusterLauncher.
10999
110100
Args:
111-
job_id (str):
112-
Historical feature retrieval job id.
113-
operation (google.api.core.operation.Operation):
114-
A Future for the spark job result, returned by the dataproc client.
115-
output_file_uri (str):
116-
Uri to the historical feature retrieval job output file.
101+
job_id (str): Historical feature retrieval job id.
102+
operation (google.api.core.operation.Operation): A Future for the spark job result,
103+
returned by the dataproc client.
104+
output_file_uri (str): Uri to the historical feature retrieval job output file.
117105
"""
118106
self.job_id = job_id
119107
self._operation = operation
@@ -135,35 +123,29 @@ class JobLauncher(abc.ABC):
135123
Submits spark jobs to a spark cluster. Currently supports only historical feature retrieval jobs.
136124
"""
137125

138-
# entity_source: Source,
139-
# feature_tables_sources: List[Source],
140-
# feature_tables: List[FeatureTable],
141-
# destination
142-
143126
@abc.abstractmethod
144127
def historical_feature_retrieval(
145128
self,
146129
pyspark_script: str,
147-
entity_source: EntitySource,
148-
feature_tables_sources: List[FeatureTableSource],
149-
feature_tables: List[FeatureTable],
150-
destination: FileDestination,
130+
entity_source_conf: Dict,
131+
feature_tables_sources_conf: List[Dict],
132+
feature_tables_conf: List[Dict],
133+
destination_conf: Dict,
151134
job_id: str,
152135
**kwargs,
153136
) -> RetrievalJob:
154137
"""
155138
Submits a historical feature retrieval job to a Spark cluster.
156139
157140
Args:
158-
pyspark_script (str):
159-
Local file path to the pyspark script for historical feature retrieval.
160-
entity_source (EntitySource): Entity data source.
161-
feature_tables_sources (FeatureTableSource): List of feature tables data sources.
162-
feature_tables (List[FeatureTable]): List of feature table specification.
141+
pyspark_script (str): Local file path to the pyspark script for historical feature
142+
retrieval.
143+
entity_source_conf (List[Dict]): Entity data source configuration.
144+
feature_tables_sources_conf (Dict): List of feature tables data sources configurations.
145+
feature_tables_conf (List[Dict]): List of feature table specification.
163146
The order of the feature table must correspond to that of feature_tables_sources.
164-
destination (FileDestination): Retrieval job output destination.
165-
job_id (str):
166-
A job id that is unique for each job submission.
147+
destination_conf (Dict): Retrieval job output destination.
148+
job_id (str): A job id that is unique for each job submission.
167149
168150
Raises:
169151
SparkJobFailure: The spark job submission failed, encountered error
@@ -202,13 +184,14 @@ def spark_submit_script_path(self):
202184
def historical_feature_retrieval(
203185
self,
204186
pyspark_script: str,
205-
entity_source: EntitySource,
206-
feature_tables_sources: List[FeatureTableSource],
207-
feature_tables: List[FeatureTable],
208-
destination: FileDestination,
187+
entity_source_conf: Dict,
188+
feature_tables_sources_conf: List[Dict],
189+
feature_tables_conf: List[Dict],
190+
destination_conf: Dict,
209191
job_id: str,
210192
**kwargs,
211193
) -> RetrievalJob:
194+
212195
submission_cmd = [
213196
self.spark_submit_script_path,
214197
"--master",
@@ -217,17 +200,17 @@ def historical_feature_retrieval(
217200
job_id,
218201
pyspark_script,
219202
"--feature-tables",
220-
json.dumps([ft._asdict() for ft in feature_tables]),
203+
json.dumps(feature_tables_conf),
221204
"--feature-tables-sources",
222-
json.dumps([fts._asdict() for fts in feature_tables_sources]),
205+
json.dumps(feature_tables_sources_conf),
223206
"--entity-source",
224-
json.dumps(entity_source._asdict()),
207+
json.dumps(entity_source_conf),
225208
"--destination",
226-
json.dumps(destination._asdict()),
209+
json.dumps(destination_conf),
227210
]
228211

229212
process = subprocess.Popen(submission_cmd, shell=True)
230-
output_file = destination.path
213+
output_file = destination_conf["path"]
231214
return StandaloneClusterRetrievalJob(job_id, process, output_file)
232215

233216

@@ -288,13 +271,14 @@ def _stage_files(self, pyspark_script: str, job_id: str) -> str:
288271
def historical_feature_retrieval(
289272
self,
290273
pyspark_script: str,
291-
entity_source: EntitySource,
292-
feature_tables_sources: List[FeatureTableSource],
293-
feature_tables: List[FeatureTable],
294-
destination: FileDestination,
274+
entity_source_conf: Dict,
275+
feature_tables_sources_conf: List[Dict],
276+
feature_tables_conf: List[Dict],
277+
destination_conf: Dict,
295278
job_id: str,
296279
**kwargs,
297280
) -> RetrievalJob:
281+
298282
pyspark_gcs = self._stage_files(pyspark_script, job_id)
299283
job = {
300284
"reference": {"job_id": job_id},
@@ -303,18 +287,18 @@ def historical_feature_retrieval(
303287
"main_python_file_uri": pyspark_gcs,
304288
"args": [
305289
"--feature-tables",
306-
json.dumps([ft._asdict() for ft in feature_tables]),
290+
json.dumps(feature_tables_conf),
307291
"--feature-tables-sources",
308-
json.dumps([fts._asdict() for fts in feature_tables_sources]),
292+
json.dumps(feature_tables_sources_conf),
309293
"--entity-source",
310-
json.dumps(entity_source._asdict()),
294+
json.dumps(entity_source_conf),
311295
"--destination",
312-
json.dumps(destination._asdict()),
296+
json.dumps(destination_conf),
313297
],
314298
},
315299
}
316300
operation = self.job_client.submit_job_as_operation(
317301
request={"project_id": self.project_id, "region": self.region, "job": job}
318302
)
319-
output_file = destination.path
303+
output_file = destination_conf["path"]
320304
return DataprocRetrievalJob(job_id, operation, output_file)

0 commit comments

Comments
 (0)