Skip to content

Commit 4141cd0

Browse files
khorshuhengkhorshuheng
authored
Add unit test for historical retrieval with panda dataframe (#1073)
Signed-off-by: Khor Shu Heng <khor.heng@gojek.com> Co-authored-by: Khor Shu Heng <khor.heng@gojek.com>
1 parent 429fec0 commit 4141cd0

4 files changed

Lines changed: 143 additions & 43 deletions

File tree

‎sdk/python/feast/client.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -893,8 +893,9 @@ def get_historical_features(
893893
feature_tables = self._get_feature_tables_from_feature_refs(
894894
feature_refs, project
895895
)
896-
output_location = self._config.get(
897-
CONFIG_SPARK_HISTORICAL_FEATURE_OUTPUT_LOCATION
896+
output_location = os.path.join(
897+
self._config.get(CONFIG_SPARK_HISTORICAL_FEATURE_OUTPUT_LOCATION),
898+
str(uuid.uuid4()),
898899
)
899900
output_format = self._config.get(CONFIG_SPARK_HISTORICAL_FEATURE_OUTPUT_FORMAT)
900901

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,7 @@ def get_output_file_uri(self, timeout_sec: int = None):
124124
with self._process as p:
125125
try:
126126
p.wait(timeout_sec)
127+
return self._output_file_uri
127128
except Exception:
128129
p.kill()
129130
raise SparkJobFailure("Timeout waiting for subprocess to return")

‎sdk/python/feast/staging/storage_client.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -270,7 +270,7 @@ def list_files(self, bucket: str, path: str) -> List[str]:
270270
raise NotImplementedError("list files not implemented for Local file")
271271

272272
def upload_file(self, local_path: str, bucket: str, remote_path: str):
273-
dest_fpath = "/" + remote_path
273+
dest_fpath = remote_path if remote_path.startswith("/") else "/" + remote_path
274274
os.makedirs(os.path.dirname(dest_fpath), exist_ok=True)
275275
shutil.copy(local_path, dest_fpath)
276276

‎sdk/python/tests/test_historical_feature_retrieval.py‎

Lines changed: 138 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,14 @@
66
from contextlib import closing
77
from datetime import datetime
88
from typing import List, Tuple
9+
from urllib.parse import urlparse
910

1011
import grpc
12+
import numpy as np
13+
import pandas as pd
1114
import pytest
1215
from google.protobuf.duration_pb2 import Duration
16+
from pandas.util.testing import assert_frame_equal
1317
from pyspark.sql import DataFrame, SparkSession
1418
from pyspark.sql.types import (
1519
BooleanType,
@@ -19,6 +23,7 @@
1923
StructType,
2024
TimestampType,
2125
)
26+
from pytz import utc
2227

2328
from feast import Client, Entity, Feature, FeatureTable, FileSource, ValueType
2429
from feast.core import CoreService_pb2_grpc as Core
@@ -82,6 +87,26 @@ def client(server):
8287
return Client(core_url=f"localhost:{free_port}")
8388

8489

90+
@pytest.yield_fixture()
91+
def client_with_local_spark(tmpdir):
92+
import pyspark
93+
94+
spark_staging_location = f"file://{os.path.join(tmpdir, 'staging')}"
95+
historical_feature_output_location = (
96+
f"file://{os.path.join(tmpdir, 'historical_feature_retrieval_output')}"
97+
)
98+
99+
return Client(
100+
core_url=f"localhost:{free_port}",
101+
spark_launcher="standalone",
102+
spark_standalone_master="local",
103+
spark_home=os.path.dirname(pyspark.__file__),
104+
spark_staging_location=spark_staging_location,
105+
historical_feature_output_location=historical_feature_output_location,
106+
historical_feature_output_format="parquet",
107+
)
108+
109+
85110
@pytest.fixture()
86111
def driver_entity(client):
87112
return client.apply_entity(Entity("driver_id", "description", ValueType.INT32))
@@ -116,36 +141,36 @@ def transactions_feature_table(spark, client):
116141
df_data = [
117142
(
118143
1001,
119-
datetime(year=2020, month=9, day=1),
120-
datetime(year=2020, month=9, day=1),
144+
datetime(year=2020, month=9, day=1, tzinfo=utc),
145+
datetime(year=2020, month=9, day=1, tzinfo=utc),
121146
50.0,
122147
True,
123148
),
124149
(
125150
1001,
126-
datetime(year=2020, month=9, day=1),
127-
datetime(year=2020, month=9, day=2),
151+
datetime(year=2020, month=9, day=1, tzinfo=utc),
152+
datetime(year=2020, month=9, day=2, tzinfo=utc),
128153
100.0,
129154
True,
130155
),
131156
(
132157
2001,
133-
datetime(year=2020, month=9, day=1),
134-
datetime(year=2020, month=9, day=1),
158+
datetime(year=2020, month=9, day=1, tzinfo=utc),
159+
datetime(year=2020, month=9, day=1, tzinfo=utc),
135160
400.0,
136161
False,
137162
),
138163
(
139164
1001,
140-
datetime(year=2020, month=9, day=2),
141-
datetime(year=2020, month=9, day=1),
165+
datetime(year=2020, month=9, day=2, tzinfo=utc),
166+
datetime(year=2020, month=9, day=1, tzinfo=utc),
142167
200.0,
143168
False,
144169
),
145170
(
146171
1001,
147-
datetime(year=2020, month=9, day=4),
148-
datetime(year=2020, month=9, day=1),
172+
datetime(year=2020, month=9, day=4, tzinfo=utc),
173+
datetime(year=2020, month=9, day=1, tzinfo=utc),
149174
300.0,
150175
False,
151176
),
@@ -180,20 +205,20 @@ def bookings_feature_table(spark, client):
180205
df_data = [
181206
(
182207
8001,
183-
datetime(year=2020, month=9, day=1),
184-
datetime(year=2020, month=9, day=1),
208+
datetime(year=2020, month=9, day=1, tzinfo=utc),
209+
datetime(year=2020, month=9, day=1, tzinfo=utc),
185210
100,
186211
),
187212
(
188213
8001,
189-
datetime(year=2020, month=9, day=2),
190-
datetime(year=2020, month=9, day=2),
214+
datetime(year=2020, month=9, day=2, tzinfo=utc),
215+
datetime(year=2020, month=9, day=2, tzinfo=utc),
191216
150,
192217
),
193218
(
194219
8002,
195-
datetime(year=2020, month=9, day=2),
196-
datetime(year=2020, month=9, day=2),
220+
datetime(year=2020, month=9, day=2, tzinfo=utc),
221+
datetime(year=2020, month=9, day=2, tzinfo=utc),
197222
200,
198223
),
199224
]
@@ -225,20 +250,20 @@ def bookings_feature_table_with_mapping(spark, client):
225250
df_data = [
226251
(
227252
8001,
228-
datetime(year=2020, month=9, day=1),
229-
datetime(year=2020, month=9, day=1),
253+
datetime(year=2020, month=9, day=1, tzinfo=utc),
254+
datetime(year=2020, month=9, day=1, tzinfo=utc),
230255
100,
231256
),
232257
(
233258
8001,
234-
datetime(year=2020, month=9, day=2),
235-
datetime(year=2020, month=9, day=2),
259+
datetime(year=2020, month=9, day=2, tzinfo=utc),
260+
datetime(year=2020, month=9, day=2, tzinfo=utc),
236261
150,
237262
),
238263
(
239264
8002,
240-
datetime(year=2020, month=9, day=2),
241-
datetime(year=2020, month=9, day=2),
265+
datetime(year=2020, month=9, day=2, tzinfo=utc),
266+
datetime(year=2020, month=9, day=2, tzinfo=utc),
242267
200,
243268
),
244269
]
@@ -273,12 +298,12 @@ def test_historical_feature_retrieval_from_local_spark_session(
273298
]
274299
)
275300
df_data = [
276-
(1001, 8001, datetime(year=2020, month=9, day=1),),
277-
(2001, 8001, datetime(year=2020, month=9, day=2),),
278-
(2001, 8002, datetime(year=2020, month=9, day=1),),
279-
(1001, 8001, datetime(year=2020, month=9, day=2),),
280-
(1001, 8001, datetime(year=2020, month=9, day=3),),
281-
(1001, 8001, datetime(year=2020, month=9, day=4),),
301+
(1001, 8001, datetime(year=2020, month=9, day=1, tzinfo=utc)),
302+
(2001, 8001, datetime(year=2020, month=9, day=2, tzinfo=utc)),
303+
(2001, 8002, datetime(year=2020, month=9, day=1, tzinfo=utc)),
304+
(1001, 8001, datetime(year=2020, month=9, day=2, tzinfo=utc)),
305+
(1001, 8001, datetime(year=2020, month=9, day=3, tzinfo=utc)),
306+
(1001, 8001, datetime(year=2020, month=9, day=4, tzinfo=utc)),
282307
]
283308
temp_dir, file_uri = create_temp_parquet_file(
284309
spark, "customer_driver_pair", schema, df_data
@@ -300,12 +325,12 @@ def test_historical_feature_retrieval_from_local_spark_session(
300325
]
301326
)
302327
expected_joined_df_data = [
303-
(1001, 8001, datetime(year=2020, month=9, day=1), 100.0, 100),
304-
(2001, 8001, datetime(year=2020, month=9, day=2), 400.0, 150),
305-
(2001, 8002, datetime(year=2020, month=9, day=1), 400.0, None),
306-
(1001, 8001, datetime(year=2020, month=9, day=2), 200.0, 150),
307-
(1001, 8001, datetime(year=2020, month=9, day=3), 200.0, 150),
308-
(1001, 8001, datetime(year=2020, month=9, day=4), 300.0, None),
328+
(1001, 8001, datetime(year=2020, month=9, day=1, tzinfo=utc), 100.0, 100),
329+
(2001, 8001, datetime(year=2020, month=9, day=2, tzinfo=utc), 400.0, 150),
330+
(2001, 8002, datetime(year=2020, month=9, day=1, tzinfo=utc), 400.0, None),
331+
(1001, 8001, datetime(year=2020, month=9, day=2, tzinfo=utc), 200.0, 150),
332+
(1001, 8001, datetime(year=2020, month=9, day=3, tzinfo=utc), 200.0, 150),
333+
(1001, 8001, datetime(year=2020, month=9, day=4, tzinfo=utc), 300.0, None),
309334
]
310335
expected_joined_df = spark.createDataFrame(
311336
spark.sparkContext.parallelize(expected_joined_df_data),
@@ -325,9 +350,9 @@ def test_historical_feature_retrieval_with_field_mappings_from_local_spark_sessi
325350
]
326351
)
327352
df_data = [
328-
(8001, datetime(year=2020, month=9, day=1)),
329-
(8001, datetime(year=2020, month=9, day=2)),
330-
(8002, datetime(year=2020, month=9, day=1)),
353+
(8001, datetime(year=2020, month=9, day=1, tzinfo=utc)),
354+
(8001, datetime(year=2020, month=9, day=2, tzinfo=utc)),
355+
(8002, datetime(year=2020, month=9, day=1, tzinfo=utc)),
331356
]
332357
temp_dir, file_uri = create_temp_parquet_file(spark, "drivers", schema, df_data)
333358
entity_source = FileSource(
@@ -344,13 +369,86 @@ def test_historical_feature_retrieval_with_field_mappings_from_local_spark_sessi
344369
]
345370
)
346371
expected_joined_df_data = [
347-
(8001, datetime(year=2020, month=9, day=1), 100),
348-
(8001, datetime(year=2020, month=9, day=2), 150),
349-
(8002, datetime(year=2020, month=9, day=1), None),
372+
(8001, datetime(year=2020, month=9, day=1, tzinfo=utc), 100),
373+
(8001, datetime(year=2020, month=9, day=2, tzinfo=utc), 150),
374+
(8002, datetime(year=2020, month=9, day=1, tzinfo=utc), None),
350375
]
351376
expected_joined_df = spark.createDataFrame(
352377
spark.sparkContext.parallelize(expected_joined_df_data),
353378
expected_joined_df_schema,
354379
)
355380
assert_dataframe_equal(joined_df, expected_joined_df)
356381
shutil.rmtree(temp_dir)
382+
383+
384+
@pytest.mark.usefixtures(
385+
"driver_entity",
386+
"customer_entity",
387+
"bookings_feature_table",
388+
"transactions_feature_table",
389+
)
390+
def test_historical_feature_retrieval_with_pandas_dataframe_input(
391+
client_with_local_spark,
392+
):
393+
394+
customer_driver_pairs_pandas_df = pd.DataFrame(
395+
np.array(
396+
[
397+
[1001, 8001, datetime(year=2020, month=9, day=1, tzinfo=utc)],
398+
[2001, 8001, datetime(year=2020, month=9, day=2, tzinfo=utc)],
399+
[2001, 8002, datetime(year=2020, month=9, day=1, tzinfo=utc)],
400+
[1001, 8001, datetime(year=2020, month=9, day=2, tzinfo=utc)],
401+
[1001, 8001, datetime(year=2020, month=9, day=3, tzinfo=utc)],
402+
[1001, 8001, datetime(year=2020, month=9, day=4, tzinfo=utc)],
403+
]
404+
),
405+
columns=["customer_id", "driver_id", "event_timestamp"],
406+
)
407+
customer_driver_pairs_pandas_df = customer_driver_pairs_pandas_df.astype(
408+
{"customer_id": "int32", "driver_id": "int32"}
409+
)
410+
411+
job_output = client_with_local_spark.get_historical_features(
412+
["transactions:total_transactions", "bookings:total_completed_bookings"],
413+
customer_driver_pairs_pandas_df,
414+
)
415+
416+
output_dir = job_output.get_output_file_uri()
417+
joined_df = pd.read_parquet(urlparse(output_dir).path)
418+
419+
expected_joined_df = pd.DataFrame(
420+
np.array(
421+
[
422+
[1001, 8001, datetime(year=2020, month=9, day=1), 100.0, 100],
423+
[2001, 8001, datetime(year=2020, month=9, day=2), 400.0, 150],
424+
[2001, 8002, datetime(year=2020, month=9, day=1), 400.0, None],
425+
[1001, 8001, datetime(year=2020, month=9, day=2), 200.0, 150],
426+
[1001, 8001, datetime(year=2020, month=9, day=3), 200.0, 150],
427+
[1001, 8001, datetime(year=2020, month=9, day=4), 300.0, None],
428+
]
429+
),
430+
columns=[
431+
"customer_id",
432+
"driver_id",
433+
"event_timestamp",
434+
"transactions__total_transactions",
435+
"bookings__total_completed_bookings",
436+
],
437+
)
438+
expected_joined_df = expected_joined_df.astype(
439+
{
440+
"customer_id": "int32",
441+
"driver_id": "int32",
442+
"transactions__total_transactions": "float64",
443+
"bookings__total_completed_bookings": "float64",
444+
}
445+
)
446+
447+
assert_frame_equal(
448+
joined_df.sort_values(
449+
by=["customer_id", "driver_id", "event_timestamp"]
450+
).reset_index(drop=True),
451+
expected_joined_df.sort_values(
452+
by=["customer_id", "driver_id", "event_timestamp"]
453+
).reset_index(drop=True),
454+
)

0 commit comments

Comments
 (0)