Skip to content

Commit 7cbd232

Browse files
toping4445younggyu-oh
toping4445
authored and
younggyu-oh
committed
fixed bugs, cleaned code, added some methods. test_universal_historical_retrieval - 100% passed
Signed-off-by: Youngkyu OH <toping4445@gmail.com>
1 parent 693e9df commit 7cbd232

13 files changed

Lines changed: 333 additions & 138 deletions

File tree

‎sdk/python/feast/__init__.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,12 @@
55
from importlib_metadata import PackageNotFoundError, version as _version # type: ignore
66

77
from feast.infra.offline_stores.bigquery_source import BigQuerySource
8-
from feast.infra.offline_stores.file_source import FileSource
9-
from feast.infra.offline_stores.redshift_source import RedshiftSource
10-
from feast.infra.offline_stores.snowflake_source import SnowflakeSource
118
from feast.infra.offline_stores.contrib.athena_offline_store.athena_source import (
129
AthenaSource,
1310
)
11+
from feast.infra.offline_stores.file_source import FileSource
12+
from feast.infra.offline_stores.redshift_source import RedshiftSource
13+
from feast.infra.offline_stores.snowflake_source import SnowflakeSource
1414

1515
from .batch_feature_view import BatchFeatureView
1616
from .data_source import KafkaSource, KinesisSource, PushSource, RequestSource

‎sdk/python/feast/batch_feature_view.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
"SnowflakeSource",
1515
"SparkSource",
1616
"TrinoSource",
17+
"AthenaSource",
1718
}
1819

1920

‎sdk/python/feast/infra/offline_stores/contrib/athena_offline_store/athena.py‎

Lines changed: 76 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -25,23 +25,21 @@
2525
from feast import OnDemandFeatureView
2626
from feast.data_source import DataSource
2727
from feast.errors import InvalidEntityType
28-
from feast.feature_logging import LoggingConfig, LoggingSource, LoggingDestination
28+
from feast.feature_logging import LoggingConfig, LoggingDestination, LoggingSource
2929
from feast.feature_view import DUMMY_ENTITY_ID, DUMMY_ENTITY_VAL, FeatureView
30+
from feast.infra.offline_stores import offline_utils
31+
from feast.infra.offline_stores.contrib.athena_offline_store.athena_source import (
32+
AthenaLoggingDestination,
33+
AthenaSource,
34+
SavedDatasetAthenaStorage,
35+
)
3036
from feast.infra.offline_stores.offline_store import (
3137
OfflineStore,
3238
RetrievalJob,
3339
RetrievalMetadata,
3440
)
35-
36-
from feast.infra.offline_stores.contrib.athena_offline_store.athena_source import (
37-
AthenaSource,
38-
AthenaLoggingDestination,
39-
SavedDatasetAthenaStorage,
40-
)
4141
from feast.infra.utils import aws_utils
42-
from feast.infra.offline_stores import offline_utils
43-
44-
from feast.registry import Registry
42+
from feast.registry import Registry, BaseRegistry
4543
from feast.repo_config import FeastConfigBaseModel, RepoConfig
4644
from feast.saved_dataset import SavedDatasetStorage
4745
from feast.usage import log_exceptions_and_usage
@@ -82,7 +80,7 @@ def pull_latest_from_table_or_query(
8280
assert isinstance(data_source, AthenaSource)
8381
assert isinstance(config.offline_store, AthenaOfflineStoreConfig)
8482

85-
from_expression = data_source.get_table_query_string()
83+
from_expression = data_source.get_table_query_string(config)
8684

8785
partition_by_join_key_string = ", ".join(join_key_columns)
8886
if partition_by_join_key_string != "":
@@ -99,9 +97,7 @@ def pull_latest_from_table_or_query(
9997

10098
date_partition_column = data_source.date_partition_column
10199

102-
athena_client = aws_utils.get_athena_data_client(
103-
config.offline_store.region
104-
)
100+
athena_client = aws_utils.get_athena_data_client(config.offline_store.region)
105101
s3_resource = aws_utils.get_s3_resource(config.offline_store.region)
106102

107103
start_date = start_date.astimezone(tz=utc)
@@ -142,15 +138,13 @@ def pull_all_from_table_or_query(
142138
end_date: datetime,
143139
) -> RetrievalJob:
144140
assert isinstance(data_source, AthenaSource)
145-
from_expression = data_source.get_table_query_string()
141+
from_expression = data_source.get_table_query_string(config)
146142

147143
field_string = ", ".join(
148144
join_key_columns + feature_name_columns + [timestamp_field]
149145
)
150146

151-
athena_client = aws_utils.get_athena_data_client(
152-
config.offline_store.region
153-
)
147+
athena_client = aws_utils.get_athena_data_client(config.offline_store.region)
154148
s3_resource = aws_utils.get_s3_resource(config.offline_store.region)
155149

156150
date_partition_column = data_source.date_partition_column
@@ -186,9 +180,7 @@ def get_historical_features(
186180
) -> RetrievalJob:
187181
assert isinstance(config.offline_store, AthenaOfflineStoreConfig)
188182

189-
athena_client = aws_utils.get_athena_data_client(
190-
config.offline_store.region
191-
)
183+
athena_client = aws_utils.get_athena_data_client(config.offline_store.region)
192184
s3_resource = aws_utils.get_s3_resource(config.offline_store.region)
193185

194186
# get pandas dataframe consisting of 1 row (LIMIT 1) and generate the schema out of it
@@ -197,23 +189,24 @@ def get_historical_features(
197189
)
198190

199191
# find timestamp column of entity df.(default = "event_timestamp"). Exception occurs if there are more than two timestamp columns.
200-
entity_df_event_timestamp_col = offline_utils.infer_event_timestamp_from_entity_df(
201-
entity_schema
192+
entity_df_event_timestamp_col = (
193+
offline_utils.infer_event_timestamp_from_entity_df(entity_schema)
202194
)
203195

204196
# get min,max of event_timestamp.
205197
entity_df_event_timestamp_range = _get_entity_df_event_timestamp_range(
206-
entity_df, entity_df_event_timestamp_col, athena_client, config,
198+
entity_df,
199+
entity_df_event_timestamp_col,
200+
athena_client,
201+
config,
207202
)
208203

209204
@contextlib.contextmanager
210205
def query_generator() -> Iterator[str]:
211206

212207
table_name = offline_utils.get_temp_entity_table_name()
213208

214-
_upload_entity_df(
215-
entity_df, athena_client, config, s3_resource, table_name
216-
)
209+
_upload_entity_df(entity_df, athena_client, config, s3_resource, table_name)
217210

218211
expected_join_keys = offline_utils.get_expected_join_keys(
219212
project, feature_views, registry
@@ -232,7 +225,6 @@ def query_generator() -> Iterator[str]:
232225
entity_df_event_timestamp_range,
233226
)
234227

235-
236228
# Generate the Athena SQL query from the query context
237229
query = offline_utils.build_point_in_time_query(
238230
query_context,
@@ -247,17 +239,20 @@ def query_generator() -> Iterator[str]:
247239
yield query
248240
finally:
249241

250-
#Always clean up the temp Athena table
242+
# Always clean up the temp Athena table
251243
aws_utils.execute_athena_query(
252244
athena_client,
253245
config.offline_store.data_source,
254246
config.offline_store.database,
255247
f"DROP TABLE IF EXISTS {config.offline_store.database}.{table_name}",
256248
)
257249

258-
bucket = config.offline_store.s3_staging_location.replace("s3://", "").split("/", 1)[0]
259-
aws_utils.delete_s3_directory(s3_resource,bucket, "entity_df/"+table_name+"/")
260-
250+
bucket = config.offline_store.s3_staging_location.replace(
251+
"s3://", ""
252+
).split("/", 1)[0]
253+
aws_utils.delete_s3_directory(
254+
s3_resource, bucket, "entity_df/" + table_name + "/"
255+
)
261256

262257
return AthenaRetrievalJob(
263258
query=query_generator,
@@ -276,21 +271,18 @@ def query_generator() -> Iterator[str]:
276271
),
277272
)
278273

279-
280274
@staticmethod
281275
def write_logged_features(
282276
config: RepoConfig,
283277
data: Union[pyarrow.Table, Path],
284278
source: LoggingSource,
285279
logging_config: LoggingConfig,
286-
registry: Registry,
280+
registry: BaseRegistry,
287281
):
288282
destination = logging_config.destination
289283
assert isinstance(destination, AthenaLoggingDestination)
290284

291-
athena_client = aws_utils.get_athena_data_client(
292-
config.offline_store.region
293-
)
285+
athena_client = aws_utils.get_athena_data_client(config.offline_store.region)
294286
s3_resource = aws_utils.get_s3_resource(config.offline_store.region)
295287
if isinstance(data, Path):
296288
s3_path = f"{config.offline_store.s3_staging_location}/logged_features/{uuid.uuid4()}"
@@ -299,7 +291,7 @@ def write_logged_features(
299291

300292
aws_utils.upload_arrow_table_to_athena(
301293
table=data,
302-
athena_data_client=athena_client,
294+
athena_client=athena_client,
303295
data_source=config.offline_store.data_source,
304296
database=config.offline_store.database,
305297
s3_resource=s3_resource,
@@ -332,7 +324,6 @@ def __init__(
332324
on_demand_feature_views (optional): A list of on demand transforms to apply at retrieval time
333325
"""
334326

335-
336327
if not isinstance(query, str):
337328
self._query_generator = query
338329
else:
@@ -352,7 +343,6 @@ def query_generator() -> Iterator[str]:
352343
)
353344
self._metadata = metadata
354345

355-
356346
@property
357347
def full_feature_names(self) -> bool:
358348
return self._full_feature_names
@@ -362,9 +352,15 @@ def on_demand_feature_views(self) -> Optional[List[OnDemandFeatureView]]:
362352
return self._on_demand_feature_views
363353

364354
def get_temp_s3_path(self) -> str:
365-
return self._config.offline_store.s3_staging_location + "/unload/" + str(uuid.uuid4())
355+
return (
356+
self._config.offline_store.s3_staging_location
357+
+ "/unload/"
358+
+ str(uuid.uuid4())
359+
)
366360

367-
def get_temp_table_dml_header(self, temp_table_name:str, temp_external_location:str) -> str:
361+
def get_temp_table_dml_header(
362+
self, temp_table_name: str, temp_external_location: str
363+
) -> str:
368364
temp_table_dml_header = f"""
369365
CREATE TABLE {temp_table_name}
370366
WITH (
@@ -387,7 +383,8 @@ def _to_df_internal(self) -> pd.DataFrame:
387383
self._config.offline_store.database,
388384
self._s3_resource,
389385
temp_external_location,
390-
self.get_temp_table_dml_header(temp_table_name, temp_external_location) + query,
386+
self.get_temp_table_dml_header(temp_table_name, temp_external_location)
387+
+ query,
391388
temp_table_name,
392389
)
393390

@@ -402,7 +399,8 @@ def _to_arrow_internal(self) -> pa.Table:
402399
self._config.offline_store.database,
403400
self._s3_resource,
404401
temp_external_location,
405-
self.get_temp_table_dml_header(temp_table_name, temp_external_location) + query,
402+
self.get_temp_table_dml_header(temp_table_name, temp_external_location)
403+
+ query,
406404
temp_table_name,
407405
)
408406

@@ -412,7 +410,33 @@ def metadata(self) -> Optional[RetrievalMetadata]:
412410

413411
def persist(self, storage: SavedDatasetStorage):
414412
assert isinstance(storage, SavedDatasetAthenaStorage)
415-
# self.to_athena(table_name=storage.athena_options.table)
413+
self.to_athena(table_name=storage.athena_options.table)
414+
415+
@log_exceptions_and_usage
416+
def to_athena(self, table_name: str) -> None:
417+
418+
if self.on_demand_feature_views:
419+
transformed_df = self.to_df()
420+
421+
_upload_entity_df(
422+
transformed_df,
423+
self._athena_client,
424+
self._config,
425+
self._s3_resource,
426+
table_name,
427+
)
428+
429+
return
430+
431+
with self._query_generator() as query:
432+
query = f'CREATE TABLE "{table_name}" AS ({query});\n'
433+
434+
aws_utils.execute_athena_query(
435+
self._athena_client,
436+
self._config.offline_store.data_source,
437+
self._config.offline_store.database,
438+
query,
439+
)
416440

417441

418442
def _upload_entity_df(
@@ -496,12 +520,14 @@ def _get_entity_df_event_timestamp_range(
496520
f"SELECT MIN({entity_df_event_timestamp_col}) AS min, MAX({entity_df_event_timestamp_col}) AS max "
497521
f"FROM ({entity_df})",
498522
)
499-
res = aws_utils.get_athena_query_result(athena_client, statement_id)[
500-
"Records"
501-
][0]
523+
res = aws_utils.get_athena_query_result(athena_client, statement_id)
502524
entity_df_event_timestamp_range = (
503-
res.parse(res[0]["stringValue"]),
504-
res.parse(res[1]["stringValue"]),
525+
datetime.strptime(
526+
res["Rows"][1]["Data"][0]["VarCharValue"], "%Y-%m-%d %H:%M:%S.%f"
527+
),
528+
datetime.strptime(
529+
res["Rows"][1]["Data"][1]["VarCharValue"], "%Y-%m-%d %H:%M:%S.%f"
530+
),
505531
)
506532
else:
507533
raise InvalidEntityType(type(entity_df))

0 commit comments

Comments
 (0)