Skip to content

Commit de8fdb9

Browse files
jklegarjklegar
committed
Add test
Signed-off-by: Jacob Klegar <jacob@tecton.ai>
1 parent 00a55e5 commit de8fdb9

4 files changed

Lines changed: 75 additions & 16 deletions

File tree

‎sdk/python/feast/data_source.py‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -171,8 +171,7 @@ def to_proto(self) -> DataSourceProto.BigQueryOptions:
171171
"""
172172

173173
bigquery_options_proto = DataSourceProto.BigQueryOptions(
174-
table_ref=self.table_ref,
175-
query=self.query,
174+
table_ref=self.table_ref, query=self.query,
176175
)
177176

178177
return bigquery_options_proto
@@ -462,7 +461,9 @@ def from_proto(data_source):
462461
created_timestamp_column=data_source.created_timestamp_column,
463462
date_partition_column=data_source.date_partition_column,
464463
)
465-
elif (data_source.bigquery_options.table_ref or data_source.bigquery_options.query):
464+
elif (
465+
data_source.bigquery_options.table_ref or data_source.bigquery_options.query
466+
):
466467
data_source_obj = BigQuerySource(
467468
field_mapping=data_source.field_mapping,
468469
table_ref=data_source.bigquery_options.table_ref,

‎sdk/python/feast/feature_store.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -225,7 +225,7 @@ def materialize(
225225
) = _run_reverse_field_mapping(feature_view)
226226

227227
offline_store = get_offline_store(self.config)
228-
table = offline_store.pull_latest_from_table(
228+
table = offline_store.pull_latest_from_table_or_query(
229229
feature_view.input,
230230
entity_names,
231231
feature_names,

‎sdk/python/feast/offline_store.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@ class OfflineStore(ABC):
8686

8787
@staticmethod
8888
@abstractmethod
89-
def pull_latest_from_table(
89+
def pull_latest_from_table_or_query(
9090
data_source: DataSource,
9191
entity_names: List[str],
9292
feature_names: List[str],
@@ -115,7 +115,7 @@ def get_historical_features(
115115

116116
class BigQueryOfflineStore(OfflineStore):
117117
@staticmethod
118-
def pull_latest_from_table(
118+
def pull_latest_from_table_or_query(
119119
data_source: DataSource,
120120
entity_names: List[str],
121121
feature_names: List[str],
@@ -126,9 +126,9 @@ def pull_latest_from_table(
126126
) -> pyarrow.Table:
127127
assert isinstance(data_source, BigQuerySource)
128128
if data_source.table_ref:
129-
from_table = f"`{data_source.table_ref}`"
129+
from_expression = f"`{data_source.table_ref}`"
130130
else:
131-
from_table = f"({data_source.query})"
131+
from_expression = f"({data_source.query})"
132132

133133
partition_by_entity_string = ", ".join(entity_names)
134134
if partition_by_entity_string != "":
@@ -144,7 +144,7 @@ def pull_latest_from_table(
144144
FROM (
145145
SELECT {field_string},
146146
ROW_NUMBER() OVER({partition_by_entity_string} ORDER BY {timestamp_desc_string}) AS _feast_row
147-
FROM {from_table}
147+
FROM {from_expression}
148148
WHERE {event_timestamp_column} BETWEEN TIMESTAMP('{start_date}') AND TIMESTAMP('{end_date}')
149149
)
150150
WHERE _feast_row = 1
@@ -286,7 +286,7 @@ def build_point_in_time_query(
286286

287287
class FileOfflineStore(OfflineStore):
288288
@staticmethod
289-
def pull_latest_from_table(
289+
def pull_latest_from_table_or_query(
290290
data_source: DataSource,
291291
entity_names: List[str],
292292
feature_names: List[str],

‎sdk/python/tests/test_bigquery_ingestion.py‎

Lines changed: 64 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ def setup_method(self):
2929
) # 2 weeks in milliseconds
3030
self.client.update_dataset(dataset, ["default_table_expiration_ms"])
3131

32-
def test_bigquery_ingestion_correctness(self):
32+
def test_bigquery_table_to_datastore_correctness(self):
3333
# create dataset
3434
ts = pd.Timestamp.now(tz="UTC").round("ms")
3535
checked_value = (
@@ -45,15 +45,13 @@ def test_bigquery_ingestion_correctness(self):
4545

4646
# load dataset into BigQuery
4747
job_config = bigquery.LoadJobConfig()
48-
table_id = (
49-
f"{self.gcp_project}.{self.bigquery_dataset}.correctness_{int(time.time())}"
50-
)
48+
table_id = f"{self.gcp_project}.{self.bigquery_dataset}.table_correctness_{int(time.time())}"
5149
job = self.client.load_table_from_dataframe(df, table_id, job_config=job_config)
5250
job.result()
5351

5452
# create FeatureView
5553
fv = FeatureView(
56-
name="test_bq_correctness",
54+
name="test_bq_table_correctness",
5755
entities=["driver_id"],
5856
features=[Feature("value", ValueType.FLOAT)],
5957
ttl=timedelta(minutes=5),
@@ -78,7 +76,67 @@ def test_bigquery_ingestion_correctness(self):
7876

7977
# run materialize()
8078
fs.materialize(
81-
["test_bq_correctness"],
79+
[fv.name],
80+
datetime.utcnow() - timedelta(minutes=5),
81+
datetime.utcnow() - timedelta(minutes=0),
82+
)
83+
84+
# check result of materialize()
85+
entity_key = EntityKeyProto(
86+
entity_names=["driver_id"], entity_values=[ValueProto(int64_val=1)]
87+
)
88+
t, val = fs._get_provider().online_read("default", fv, entity_key)
89+
assert abs(val["value"].double_val - checked_value) < 1e-6
90+
91+
def test_bigquery_query_to_datastore_correctness(self):
92+
# create dataset
93+
ts = pd.Timestamp.now(tz="UTC").round("ms")
94+
checked_value = (
95+
random.random()
96+
) # random value so test doesn't still work if no values written to online store
97+
data = {
98+
"id": [1, 2, 1],
99+
"value": [0.1, 0.2, checked_value],
100+
"ts_1": [ts - timedelta(minutes=2), ts, ts],
101+
"created_ts": [ts, ts, ts],
102+
}
103+
df = pd.DataFrame.from_dict(data)
104+
105+
# load dataset into BigQuery
106+
job_config = bigquery.LoadJobConfig()
107+
table_id = f"{self.gcp_project}.{self.bigquery_dataset}.query_correctness_{int(time.time())}"
108+
query = f"SELECT * FROM `{table_id}`"
109+
job = self.client.load_table_from_dataframe(df, table_id, job_config=job_config)
110+
job.result()
111+
112+
# create FeatureView
113+
fv = FeatureView(
114+
name="test_bq_query_correctness",
115+
entities=["driver_id"],
116+
features=[Feature("value", ValueType.FLOAT)],
117+
ttl=timedelta(minutes=5),
118+
input=BigQuerySource(
119+
event_timestamp_column="ts",
120+
created_timestamp_column="created_ts",
121+
field_mapping={"ts_1": "ts", "id": "driver_id"},
122+
date_partition_column="",
123+
query=query,
124+
),
125+
)
126+
config = RepoConfig(
127+
metadata_store="./metadata.db",
128+
project="default",
129+
provider="gcp",
130+
online_store=OnlineStoreConfig(
131+
local=LocalOnlineStoreConfig("online_store.db")
132+
),
133+
)
134+
fs = FeatureStore(config=config)
135+
fs.apply([fv])
136+
137+
# run materialize()
138+
fs.materialize(
139+
[fv.name],
82140
datetime.utcnow() - timedelta(minutes=5),
83141
datetime.utcnow() - timedelta(minutes=0),
84142
)

0 commit comments

Comments
 (0)