forked from feast-dev/feast
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_unit_feature_store.py
More file actions
50 lines (39 loc) · 1.37 KB
/
test_unit_feature_store.py
File metadata and controls
50 lines (39 loc) · 1.37 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
from dataclasses import dataclass
from typing import Dict, List
from feast import FeatureStore
from feast.protos.feast.types.Value_pb2 import Value
@dataclass
class MockFeatureViewProjection:
join_key_map: Dict[str, str]
@dataclass
class MockFeatureView:
name: str
entities: List[str]
projection: MockFeatureViewProjection
def test_get_unique_entities():
entity_values = {
"entity_1": [Value(int64_val=1), Value(int64_val=2), Value(int64_val=1)],
"entity_2": [
Value(string_val="1"),
Value(string_val="2"),
Value(string_val="1"),
],
"entity_3": [Value(int64_val=8), Value(int64_val=9), Value(int64_val=10)],
}
entity_name_to_join_key_map = {"entity_1": "entity_1", "entity_2": "entity_2"}
fv = MockFeatureView(
name="fv_1",
entities=["entity_1", "entity_2"],
projection=MockFeatureViewProjection(join_key_map={}),
)
unique_entities, indexes = FeatureStore._get_unique_entities(
FeatureStore,
table=fv,
join_key_values=entity_values,
entity_name_to_join_key_map=entity_name_to_join_key_map,
)
assert unique_entities == (
{"entity_1": Value(int64_val=1), "entity_2": Value(string_val="1")},
{"entity_1": Value(int64_val=2), "entity_2": Value(string_val="2")},
)
assert indexes == ([0, 2], [1])