Skip to content

Commit 03c8317

Browse files
felixwang9817felixwang9817
authored
Delete tables (#1916)
* Ensure that extra table uploaded to Redshift is always removed Signed-off-by: Felix Wang <wangfelix98@gmail.com> * Ensure that extra table uploaded to Bigquery is always removed Signed-off-by: Felix Wang <wangfelix98@gmail.com> * Change drop to drop if exists Signed-off-by: Felix Wang <wangfelix98@gmail.com>
1 parent 189ffb0 commit 03c8317

2 files changed

Lines changed: 92 additions & 68 deletions

File tree

‎sdk/python/feast/infra/offline_stores/bigquery.py‎

Lines changed: 81 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
1+
import contextlib
12
import uuid
23
from datetime import date, datetime, timedelta
3-
from typing import Dict, List, Optional, Union
4+
from typing import Callable, ContextManager, Dict, Iterator, List, Optional, Union
45

56
import numpy as np
67
import pandas as pd
@@ -122,38 +123,47 @@ def get_historical_features(
122123
client, client.project, config.offline_store.dataset
123124
)
124125

125-
entity_schema = _upload_entity_df_and_get_entity_schema(
126-
client=client, table_name=table_reference, entity_df=entity_df,
127-
)
126+
@contextlib.contextmanager
127+
def query_generator() -> Iterator[str]:
128+
entity_schema = _upload_entity_df_and_get_entity_schema(
129+
client=client, table_name=table_reference, entity_df=entity_df,
130+
)
128131

129-
entity_df_event_timestamp_col = offline_utils.infer_event_timestamp_from_entity_df(
130-
entity_schema
131-
)
132+
entity_df_event_timestamp_col = offline_utils.infer_event_timestamp_from_entity_df(
133+
entity_schema
134+
)
132135

133-
expected_join_keys = offline_utils.get_expected_join_keys(
134-
project, feature_views, registry
135-
)
136+
expected_join_keys = offline_utils.get_expected_join_keys(
137+
project, feature_views, registry
138+
)
136139

137-
offline_utils.assert_expected_columns_in_entity_df(
138-
entity_schema, expected_join_keys, entity_df_event_timestamp_col
139-
)
140+
offline_utils.assert_expected_columns_in_entity_df(
141+
entity_schema, expected_join_keys, entity_df_event_timestamp_col
142+
)
140143

141-
# Build a query context containing all information required to template the BigQuery SQL query
142-
query_context = offline_utils.get_feature_view_query_context(
143-
feature_refs, feature_views, registry, project,
144-
)
144+
# Build a query context containing all information required to template the BigQuery SQL query
145+
query_context = offline_utils.get_feature_view_query_context(
146+
feature_refs, feature_views, registry, project,
147+
)
145148

146-
# Generate the BigQuery SQL query from the query context
147-
query = offline_utils.build_point_in_time_query(
148-
query_context,
149-
left_table_query_string=table_reference,
150-
entity_df_event_timestamp_col=entity_df_event_timestamp_col,
151-
query_template=MULTIPLE_FEATURE_VIEW_POINT_IN_TIME_JOIN,
152-
full_feature_names=full_feature_names,
153-
)
149+
# Generate the BigQuery SQL query from the query context
150+
query = offline_utils.build_point_in_time_query(
151+
query_context,
152+
left_table_query_string=table_reference,
153+
entity_df_event_timestamp_col=entity_df_event_timestamp_col,
154+
query_template=MULTIPLE_FEATURE_VIEW_POINT_IN_TIME_JOIN,
155+
full_feature_names=full_feature_names,
156+
)
157+
158+
try:
159+
yield query
160+
finally:
161+
# Asynchronously clean up the uploaded Bigquery table, which will expire
162+
# if cleanup fails
163+
client.delete_table(table=table_reference, not_found_ok=True)
154164

155165
return BigQueryRetrievalJob(
156-
query=query,
166+
query=query_generator,
157167
client=client,
158168
config=config,
159169
full_feature_names=full_feature_names,
@@ -166,13 +176,22 @@ def get_historical_features(
166176
class BigQueryRetrievalJob(RetrievalJob):
167177
def __init__(
168178
self,
169-
query: str,
179+
query: Union[str, Callable[[], ContextManager[str]]],
170180
client: bigquery.Client,
171181
config: RepoConfig,
172182
full_feature_names: bool,
173183
on_demand_feature_views: Optional[List[OnDemandFeatureView]],
174184
):
175-
self.query = query
185+
if not isinstance(query, str):
186+
self._query_generator = query
187+
else:
188+
189+
@contextlib.contextmanager
190+
def query_generator() -> Iterator[str]:
191+
assert isinstance(query, str)
192+
yield query
193+
194+
self._query_generator = query_generator
176195
self.client = client
177196
self.config = config
178197
self._full_feature_names = full_feature_names
@@ -187,15 +206,16 @@ def on_demand_feature_views(self) -> Optional[List[OnDemandFeatureView]]:
187206
return self._on_demand_feature_views
188207

189208
def _to_df_internal(self) -> pd.DataFrame:
190-
# TODO: Ideally only start this job when the user runs "get_historical_features", not when they run to_df()
191-
df = self.client.query(self.query).to_dataframe(create_bqstorage_client=True)
192-
return df
209+
with self._query_generator() as query:
210+
df = self.client.query(query).to_dataframe(create_bqstorage_client=True)
211+
return df
193212

194213
def to_sql(self) -> str:
195214
"""
196215
Returns the SQL query that will be executed in BigQuery to build the historical feature table.
197216
"""
198-
return self.query
217+
with self._query_generator() as query:
218+
return query
199219

200220
def to_bigquery(
201221
self,
@@ -215,36 +235,39 @@ def to_bigquery(
215235
Returns:
216236
Returns the destination table name or returns None if job_config.dry_run is True.
217237
"""
238+
with self._query_generator() as query:
239+
if not job_config:
240+
today = date.today().strftime("%Y%m%d")
241+
rand_id = str(uuid.uuid4())[:7]
242+
path = f"{self.client.project}.{self.config.offline_store.dataset}.historical_{today}_{rand_id}"
243+
job_config = bigquery.QueryJobConfig(destination=path)
244+
245+
if not job_config.dry_run and self.on_demand_feature_views is not None:
246+
job = _write_pyarrow_table_to_bq(
247+
self.client, self.to_arrow(), job_config.destination
248+
)
249+
job.result()
250+
print(f"Done writing to '{job_config.destination}'.")
251+
return str(job_config.destination)
252+
253+
bq_job = self.client.query(query, job_config=job_config)
254+
255+
if job_config.dry_run:
256+
print(
257+
"This query will process {} bytes.".format(
258+
bq_job.total_bytes_processed
259+
)
260+
)
261+
return None
262+
263+
block_until_done(client=self.client, bq_job=bq_job, timeout=timeout)
218264

219-
if not job_config:
220-
today = date.today().strftime("%Y%m%d")
221-
rand_id = str(uuid.uuid4())[:7]
222-
path = f"{self.client.project}.{self.config.offline_store.dataset}.historical_{today}_{rand_id}"
223-
job_config = bigquery.QueryJobConfig(destination=path)
224-
225-
if not job_config.dry_run and self.on_demand_feature_views is not None:
226-
job = _write_pyarrow_table_to_bq(
227-
self.client, self.to_arrow(), job_config.destination
228-
)
229-
job.result()
230265
print(f"Done writing to '{job_config.destination}'.")
231266
return str(job_config.destination)
232267

233-
bq_job = self.client.query(self.query, job_config=job_config)
234-
235-
if job_config.dry_run:
236-
print(
237-
"This query will process {} bytes.".format(bq_job.total_bytes_processed)
238-
)
239-
return None
240-
241-
block_until_done(client=self.client, bq_job=bq_job, timeout=timeout)
242-
243-
print(f"Done writing to '{job_config.destination}'.")
244-
return str(job_config.destination)
245-
246268
def _to_arrow_internal(self) -> pyarrow.Table:
247-
return self.client.query(self.query).to_arrow()
269+
with self._query_generator() as query:
270+
return self.client.query(query).to_arrow()
248271

249272

250273
def block_until_done(
@@ -325,13 +348,13 @@ def _upload_entity_df_and_get_entity_schema(
325348
limited_entity_df = (
326349
client.query(f"SELECT * FROM {table_name} LIMIT 1").result().to_dataframe()
327350
)
351+
328352
entity_schema = dict(zip(limited_entity_df.columns, limited_entity_df.dtypes))
329353
elif isinstance(entity_df, pd.DataFrame):
330354
# Drop the index so that we dont have unnecessary columns
331355
entity_df.reset_index(drop=True, inplace=True)
332356
job = _write_df_to_bq(client, entity_df, table_name)
333357
block_until_done(client, job)
334-
335358
entity_schema = dict(zip(entity_df.columns, entity_df.dtypes))
336359
else:
337360
raise InvalidEntityType(type(entity_df))

‎sdk/python/feast/infra/offline_stores/redshift.py‎

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -153,16 +153,17 @@ def query_generator() -> Iterator[str]:
153153
full_feature_names=full_feature_names,
154154
)
155155

156-
yield query
157-
158-
# Clean up the uploaded Redshift table
159-
aws_utils.execute_redshift_statement(
160-
redshift_client,
161-
config.offline_store.cluster_id,
162-
config.offline_store.database,
163-
config.offline_store.user,
164-
f"DROP TABLE {table_name}",
165-
)
156+
try:
157+
yield query
158+
finally:
159+
# Always clean up the uploaded Redshift table
160+
aws_utils.execute_redshift_statement(
161+
redshift_client,
162+
config.offline_store.cluster_id,
163+
config.offline_store.database,
164+
config.offline_store.user,
165+
f"DROP TABLE IF EXISTS {table_name}",
166+
)
166167

167168
return RedshiftRetrievalJob(
168169
query=query_generator,

0 commit comments

Comments
 (0)