Skip to content

Commit 013fc57

Browse files
fix: Keep the BatchFeatureView type and settings when copying it (#6989)
* fix: Keep the BatchFeatureView type and settings when copying it Signed-off-by: LuisFigueroaG <[email protected]> * fix: Add type hints to BatchFeatureView.__copy__ and its test Signed-off-by: LuisFigueroaG <[email protected]> --------- Signed-off-by: LuisFigueroaG <[email protected]>
1 parent 015cc84 commit 013fc57

2 files changed

Lines changed: 77 additions & 1 deletion

File tree

‎sdk/python/feast/batch_feature_view.py‎

Lines changed: 35 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
1+
import copy
12
import functools
23
import warnings
34
from datetime import datetime, timedelta
4-
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
5+
from typing import Any, Callable, Dict, List, Optional, Tuple, Union, cast
56

67
import dill
78

@@ -161,6 +162,39 @@ def __init__(
161162
version=version,
162163
)
163164

165+
def __copy__(self) -> "BatchFeatureView":
166+
fv = BatchFeatureView(
167+
name=self.name,
168+
mode=cast(Union[TransformationMode, str], self.mode),
169+
source=cast(List["BatchFeatureView"], self.source_views)
170+
if self.source_views
171+
else self.batch_source,
172+
sink_source=self.batch_source if self.source_views else None,
173+
schema=self.schema,
174+
ttl=self.ttl,
175+
tags=self.tags,
176+
online=self.online,
177+
offline=self.offline,
178+
description=self.description,
179+
owner=self.owner,
180+
org=self.org,
181+
udf=self.udf,
182+
udf_string=self.udf_string,
183+
feature_transformation=self.feature_transformation,
184+
batch_engine=self.batch_engine,
185+
aggregations=self.aggregations,
186+
enable_validation=self.enable_validation,
187+
version=self.version,
188+
)
189+
fv.enabled = self.enabled
190+
fv.state = self.state
191+
fv.entities = self.entities
192+
fv.features = copy.copy(self.features)
193+
fv.entity_columns = copy.copy(self.entity_columns)
194+
fv.projection = copy.copy(self.projection)
195+
fv._raw_feature_transformation_proto = self._raw_feature_transformation_proto
196+
return fv
197+
164198
def get_feature_transformation(self) -> Optional[Transformation]:
165199
if not self.udf:
166200
return None

‎sdk/python/tests/unit/test_feature_views.py‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
1+
import copy
12
from datetime import timedelta
23

34
import pandas as pd
45
import pytest
56
from typeguard import TypeCheckError
67

8+
from feast.aggregation import Aggregation
79
from feast.batch_feature_view import BatchFeatureView
810
from feast.data_format import AvroFormat
911
from feast.data_source import KafkaSource
@@ -404,6 +406,46 @@ def transform_udf(df: pd.DataFrame) -> pd.DataFrame:
404406
assert second_deserialized.feature_transformation is not None
405407

406408

409+
def test_batch_feature_view_copy() -> None:
410+
def transform_udf(df: pd.DataFrame) -> pd.DataFrame:
411+
return df
412+
413+
bfv = BatchFeatureView(
414+
name="test_batch_feature_view",
415+
entities=[Entity(name="test_entity", join_keys=["entity_id"])],
416+
schema=[
417+
Field(name="entity_id", dtype=String),
418+
Field(name="input_feature", dtype=Int64),
419+
],
420+
source=FileSource(
421+
name="test-source",
422+
path="test_data.parquet",
423+
timestamp_field="event_timestamp",
424+
),
425+
udf=transform_udf,
426+
mode="pandas",
427+
batch_engine={"spark.sql.shuffle.partitions": "4"},
428+
aggregations=[
429+
Aggregation(
430+
column="input_feature", function="sum", time_window=timedelta(days=1)
431+
)
432+
],
433+
)
434+
435+
bfv_copy = copy.copy(bfv)
436+
assert isinstance(bfv_copy, BatchFeatureView)
437+
assert bfv_copy == bfv
438+
assert bfv_copy.mode == "pandas"
439+
assert bfv_copy.udf is transform_udf
440+
assert bfv_copy.feature_transformation is bfv.feature_transformation
441+
assert bfv_copy.batch_engine == {"spark.sql.shuffle.partitions": "4"}
442+
assert bfv_copy.aggregations == bfv.aggregations
443+
444+
projected = bfv[["input_feature"]]
445+
assert isinstance(projected, BatchFeatureView)
446+
assert projected.to_proto().spec.HasField("feature_transformation")
447+
448+
407449
def test_transformation_mode_serialization():
408450
"""
409451
Test that transformation mode is properly serialized to proto and deserialized back.

0 commit comments

Comments
 (0)