Skip to content

Commit 2f9a13e

Browse files
pyalexpyalex
committed
Configurable materialization destination for view in BigQuerySource (#1201)
* configurable materialization destination Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com> * use materialization options in batch ingestion Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com> * fix no default value Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com>
1 parent b7bbc60 commit 2f9a13e

5 files changed

Lines changed: 54 additions & 10 deletions

File tree

‎sdk/python/feast/constants.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,14 @@ class ConfigOptions(metaclass=ConfigMeta):
151151
#: Directory where Spark is installed
152152
SPARK_HOME: Optional[str] = None
153153

154+
#: The project id where the materialized view of BigQuerySource is going to be created
155+
#: by default, use the same project where view is located
156+
SPARK_BQ_MATERIALIZATION_PROJECT: Optional[str] = None
157+
158+
#: The dataset id where the materialized view of BigQuerySource is going to be created
159+
#: by default, use the same dataset where view is located
160+
SPARK_BQ_MATERIALIZATION_DATASET: Optional[str] = None
161+
154162
#: Dataproc cluster to run Feast Spark Jobs in
155163
DATAPROC_CLUSTER_NAME: Optional[str] = None
156164

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

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -123,6 +123,8 @@ class BigQuerySource(Source):
123123
event_timestamp_column (str): Column representing the event timestamp.
124124
created_timestamp_column (str): Column representing the creation timestamp. Required
125125
only if the source corresponds to a feature table.
126+
materialization (Dict[str, str]): Optional. Destination for materialized view,
127+
e.g. dict(project="...", dataset="...).
126128
"""
127129

128130
def __init__(
@@ -133,13 +135,15 @@ def __init__(
133135
event_timestamp_column: str,
134136
created_timestamp_column: Optional[str],
135137
field_mapping: Optional[Dict[str, str]],
138+
materialization: Optional[Dict[str, str]] = None,
136139
):
137140
super().__init__(
138141
event_timestamp_column, created_timestamp_column, field_mapping
139142
)
140143
self.project = project
141144
self.dataset = dataset
142145
self.table = table
146+
self.materialization = materialization
143147

144148
@property
145149
def spark_format(self) -> str:
@@ -151,7 +155,15 @@ def spark_path(self) -> str:
151155

152156
@property
153157
def spark_read_options(self) -> Dict[str, str]:
154-
return {**super().spark_read_options, "viewsEnabled": "true"}
158+
opts = {**super().spark_read_options, "viewsEnabled": "true"}
159+
if self.materialization:
160+
opts.update(
161+
{
162+
"materializationProject": self.materialization["project"],
163+
"materializationDataset": self.materialization["dataset"],
164+
}
165+
)
166+
return opts
155167

156168

157169
def _source_from_dict(dct: Dict) -> Source:
@@ -174,6 +186,7 @@ def _source_from_dict(dct: Dict) -> Source:
174186
field_mapping=dct["bq"].get("field_mapping", {}),
175187
event_timestamp_column=dct["bq"]["event_timestamp_column"],
176188
created_timestamp_column=dct["bq"].get("created_timestamp_column"),
189+
materialization=dct["bq"].get("materialization"),
177190
)
178191

179192

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

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ def resolve_launcher(config: Config) -> JobLauncher:
7171
return _launchers[config.get(opt.SPARK_LAUNCHER)](config)
7272

7373

74-
def _source_to_argument(source: DataSource):
74+
def _source_to_argument(source: DataSource, config: Config):
7575
common_properties = {
7676
"field_mapping": dict(source.field_mapping),
7777
"event_timestamp_column": source.event_timestamp_column,
@@ -94,6 +94,14 @@ def _source_to_argument(source: DataSource):
9494
properties["project"] = project
9595
properties["dataset"] = dataset
9696
properties["table"] = table
97+
if config.exists(opt.SPARK_BQ_MATERIALIZATION_PROJECT) and config.exists(
98+
opt.SPARK_BQ_MATERIALIZATION_DATASET
99+
):
100+
properties["materialization"] = dict(
101+
project=config.get(opt.SPARK_BQ_MATERIALIZATION_PROJECT),
102+
dataset=config.get(opt.SPARK_BQ_MATERIALIZATION_DATASET),
103+
)
104+
97105
return {"bq": properties}
98106

99107
if isinstance(source, KafkaSource):
@@ -141,9 +149,9 @@ def start_historical_feature_retrieval_spark_session(
141149
spark_session = SparkSession.builder.getOrCreate()
142150
return retrieve_historical_features(
143151
spark=spark_session,
144-
entity_source_conf=_source_to_argument(entity_source),
152+
entity_source_conf=_source_to_argument(entity_source, client._config),
145153
feature_tables_sources_conf=[
146-
_source_to_argument(feature_table.batch_source)
154+
_source_to_argument(feature_table.batch_source, client._config)
147155
for feature_table in feature_tables
148156
],
149157
feature_tables_conf=[
@@ -164,14 +172,15 @@ def start_historical_feature_retrieval_job(
164172
launcher = resolve_launcher(client._config)
165173
feature_sources = [
166174
_source_to_argument(
167-
replace_bq_table_with_joined_view(feature_table, entity_source)
175+
replace_bq_table_with_joined_view(feature_table, entity_source),
176+
client._config,
168177
)
169178
for feature_table in feature_tables
170179
]
171180

172181
return launcher.historical_feature_retrieval(
173182
RetrievalJobParameters(
174-
entity_source=_source_to_argument(entity_source),
183+
entity_source=_source_to_argument(entity_source, client._config),
175184
feature_tables_sources=feature_sources,
176185
feature_tables=[
177186
_feature_table_to_argument(client, project, feature_table)
@@ -224,7 +233,7 @@ def start_offline_to_online_ingestion(
224233
return launcher.offline_to_online_ingestion(
225234
BatchIngestionJobParameters(
226235
jar=client._config.get(opt.SPARK_INGESTION_JAR),
227-
source=_source_to_argument(feature_table.batch_source),
236+
source=_source_to_argument(feature_table.batch_source, client._config),
228237
feature_table=_feature_table_to_argument(client, project, feature_table),
229238
start=start,
230239
end=end,
@@ -251,7 +260,7 @@ def get_stream_to_online_ingestion_params(
251260
return StreamIngestionJobParameters(
252261
jar=client._config.get(opt.SPARK_INGESTION_JAR),
253262
extra_jars=extra_jars,
254-
source=_source_to_argument(feature_table.stream_source),
263+
source=_source_to_argument(feature_table.stream_source, client._config),
255264
feature_table=_feature_table_to_argument(client, project, feature_table),
256265
redis_host=client._config.get(opt.REDIS_HOST),
257266
redis_port=client._config.getint(opt.REDIS_PORT),

‎spark/ingestion/src/main/scala/feast/ingestion/IngestionJobConfig.scala‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,14 +59,17 @@ case class FileSource(
5959
override val datePartitionColumn: Option[String] = None
6060
) extends BatchSource
6161

62+
case class BQMaterializationConfig(project: String, dataset: String)
63+
6264
case class BQSource(
6365
project: String,
6466
dataset: String,
6567
table: String,
6668
override val fieldMapping: Map[String, String],
6769
override val eventTimestampColumn: String,
6870
override val createdTimestampColumn: Option[String] = None,
69-
override val datePartitionColumn: Option[String] = None
71+
override val datePartitionColumn: Option[String] = None,
72+
materialization: Option[BQMaterializationConfig] = None
7073
) extends BatchSource
7174

7275
case class KafkaSource(

‎spark/ingestion/src/main/scala/feast/ingestion/sources/bq/BigQueryReader.scala‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,9 +30,20 @@ object BigQueryReader {
3030
start: DateTime,
3131
end: DateTime
3232
): DataFrame = {
33-
sqlContext.read
33+
val reader = sqlContext.read
3434
.format("bigquery")
3535
.option("viewsEnabled", "true")
36+
37+
source.materialization match {
38+
case Some(materializationConfig) =>
39+
reader
40+
.option("materializationProject", materializationConfig.project)
41+
.option("materializationDataset", materializationConfig.dataset)
42+
43+
case _ => ()
44+
}
45+
46+
reader
3647
.load(s"${source.project}.${source.dataset}.${source.table}")
3748
.filter(col(source.eventTimestampColumn) >= new Timestamp(start.getMillis))
3849
.filter(col(source.eventTimestampColumn) < new Timestamp(end.getMillis))

0 commit comments

Comments
 (0)