forked from feast-dev/feast
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathsaved_dataset.py
More file actions
185 lines (150 loc) · 6.06 KB
/
Copy pathsaved_dataset.py
File metadata and controls
185 lines (150 loc) · 6.06 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
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
from abc import abstractmethod
from datetime import datetime
from typing import TYPE_CHECKING, Dict, List, Optional, Type, cast
import pandas as pd
import pyarrow
from google.protobuf.json_format import MessageToJson
from feast.data_source import DataSource
from feast.protos.feast.core.SavedDataset_pb2 import SavedDataset as SavedDatasetProto
from feast.protos.feast.core.SavedDataset_pb2 import SavedDatasetMeta, SavedDatasetSpec
from feast.protos.feast.core.SavedDataset_pb2 import (
SavedDatasetStorage as SavedDatasetStorageProto,
)
if TYPE_CHECKING:
from feast.infra.offline_stores.offline_store import RetrievalJob
class _StorageRegistry(type):
classes_by_proto_attr_name: Dict[str, Type["SavedDatasetStorage"]] = {}
def __new__(cls, name, bases, dct):
kls = type.__new__(cls, name, bases, dct)
if dct.get("_proto_attr_name"):
cls.classes_by_proto_attr_name[dct["_proto_attr_name"]] = kls
return kls
class SavedDatasetStorage(metaclass=_StorageRegistry):
_proto_attr_name: str
@staticmethod
def from_proto(storage_proto: SavedDatasetStorageProto) -> "SavedDatasetStorage":
proto_attr_name = cast(str, storage_proto.WhichOneof("kind"))
return _StorageRegistry.classes_by_proto_attr_name[proto_attr_name].from_proto(
storage_proto
)
@abstractmethod
def to_proto(self) -> SavedDatasetStorageProto:
...
@abstractmethod
def to_data_source(self) -> DataSource:
...
class SavedDataset:
name: str
features: List[str]
join_keys: List[str]
full_feature_names: bool
storage: SavedDatasetStorage
tags: Dict[str, str]
created_timestamp: Optional[datetime] = None
last_updated_timestamp: Optional[datetime] = None
min_event_timestamp: Optional[datetime] = None
max_event_timestamp: Optional[datetime] = None
def __init__(
self,
name: str,
features: List[str],
join_keys: List[str],
storage: SavedDatasetStorage,
full_feature_names: bool = False,
tags: Optional[Dict[str, str]] = None,
):
self.name = name
self.features = features
self.join_keys = join_keys
self.storage = storage
self.full_feature_names = full_feature_names
self.tags = tags or {}
def __repr__(self):
items = (f"{k} = {v}" for k, v in self.__dict__.items())
return f"<{self.__class__.__name__}({', '.join(items)})>"
def __str__(self):
return str(MessageToJson(self.to_proto()))
def __hash__(self):
return hash((id(self), self.name))
def __eq__(self, other):
if not isinstance(other, SavedDataset):
raise TypeError(
"Comparisons should only involve FeatureService class objects."
)
if self.name != other.name:
return False
if sorted(self.features) != sorted(other.features):
return False
return True
@staticmethod
def from_proto(saved_dataset_proto: SavedDatasetProto):
"""
Converts a SavedDatasetProto to a SavedDataset object.
Args:
saved_dataset_proto: A protobuf representation of a SavedDataset.
"""
ds = SavedDataset(
name=saved_dataset_proto.spec.name,
features=list(saved_dataset_proto.spec.features),
join_keys=list(saved_dataset_proto.spec.join_keys),
full_feature_names=saved_dataset_proto.spec.full_feature_names,
storage=SavedDatasetStorage.from_proto(saved_dataset_proto.spec.storage),
tags=dict(saved_dataset_proto.spec.tags.items()),
)
if saved_dataset_proto.meta.HasField("created_timestamp"):
ds.created_timestamp = (
saved_dataset_proto.meta.created_timestamp.ToDatetime()
)
if saved_dataset_proto.meta.HasField("last_updated_timestamp"):
ds.last_updated_timestamp = (
saved_dataset_proto.meta.last_updated_timestamp.ToDatetime()
)
if saved_dataset_proto.meta.HasField("min_event_timestamp"):
ds.min_event_timestamp = (
saved_dataset_proto.meta.min_event_timestamp.ToDatetime()
)
if saved_dataset_proto.meta.HasField("max_event_timestamp"):
ds.max_event_timestamp = (
saved_dataset_proto.meta.max_event_timestamp.ToDatetime()
)
return ds
def to_proto(self) -> SavedDatasetProto:
"""
Converts a SavedDataset to its protobuf representation.
Returns:
A SavedDatasetProto protobuf.
"""
meta = SavedDatasetMeta()
if self.created_timestamp:
meta.created_timestamp.FromDatetime(self.created_timestamp)
if self.min_event_timestamp:
meta.min_event_timestamp.FromDatetime(self.min_event_timestamp)
if self.max_event_timestamp:
meta.max_event_timestamp.FromDatetime(self.max_event_timestamp)
spec = SavedDatasetSpec(
name=self.name,
features=self.features,
join_keys=self.join_keys,
full_feature_names=self.full_feature_names,
storage=self.storage.to_proto(),
tags=self.tags,
)
feature_service_proto = SavedDatasetProto(spec=spec, meta=meta)
return feature_service_proto
def with_retrieval_job(self, retrieval_job: "RetrievalJob") -> "SavedDataset":
self._retrieval_job = retrieval_job
return self
def to_df(self) -> pd.DataFrame:
if not self._retrieval_job:
raise RuntimeError(
"To load this dataset use FeatureStore.get_saved_dataset() "
"instead of instantiating it directly."
)
return self._retrieval_job.to_df()
def to_arrow(self) -> pyarrow.Table:
if not self._retrieval_job:
raise RuntimeError(
"To load this dataset use FeatureStore.get_saved_dataset() "
"instead of instantiating it directly."
)
return self._retrieval_job.to_arrow()