Skip to content

Commit 561d6f4

Browse files
woopwoop
authored
Add YAML export to Python SDK (#782)
* Add feature set export method * Fix equality in feature set and remove YAML path export * Remove spec from yaml export * Conditionally delete meta from feature set when empty
1 parent 5d3d81f commit 561d6f4

3 files changed

Lines changed: 61 additions & 4 deletions

File tree

‎sdk/go/client.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ func (fc *GrpcClient) GetOnlineFeatures(ctx context.Context, req *OnlineFeatures
5454
// collect unqiue entity refs from entity rows
5555
entityRefs := make(map[string]struct{})
5656
for _, entityRows := range req.Entities {
57-
for ref, _ := range entityRows {
57+
for ref := range entityRows {
5858
entityRefs[ref] = struct{}{}
5959
}
6060
}

‎sdk/python/feast/feature_set.py‎

Lines changed: 38 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,10 @@
1717

1818
import pandas as pd
1919
import pyarrow as pa
20+
import yaml
2021
from google.protobuf import json_format
2122
from google.protobuf.duration_pb2 import Duration
22-
from google.protobuf.json_format import MessageToJson
23+
from google.protobuf.json_format import MessageToDict, MessageToJson
2324
from google.protobuf.message import Message
2425
from pandas.api.types import is_datetime64_ns_dtype
2526
from pyarrow.lib import TimestampType
@@ -79,12 +80,18 @@ def __eq__(self, other):
7980
if key not in other.fields.keys() or self.fields[key] != other.fields[key]:
8081
return False
8182

83+
if self.fields[key] != other.fields[key]:
84+
return False
85+
8286
if (
8387
self.name != other.name
8488
or self.project != other.project
8589
or self.max_age != other.max_age
8690
):
8791
return False
92+
93+
if self.source != other.source:
94+
return False
8895
return True
8996

9097
def __str__(self):
@@ -783,13 +790,18 @@ def from_proto(cls, feature_set_proto: FeatureSetProto):
783790
entities=[
784791
Entity.from_proto(entity) for entity in feature_set_proto.spec.entities
785792
],
786-
max_age=feature_set_proto.spec.max_age,
793+
max_age=(
794+
None
795+
if feature_set_proto.spec.max_age.seconds == 0
796+
and feature_set_proto.spec.max_age.nanos == 0
797+
else feature_set_proto.spec.max_age
798+
),
787799
source=(
788800
None
789801
if feature_set_proto.spec.source.type == 0
790802
else Source.from_proto(feature_set_proto.spec.source)
791803
),
792-
project=feature_set_proto.spec.project
804+
project=None
793805
if len(feature_set_proto.spec.project) == 0
794806
else feature_set_proto.spec.project,
795807
)
@@ -828,6 +840,29 @@ def to_proto(self) -> FeatureSetProto:
828840

829841
return FeatureSetProto(spec=spec, meta=meta)
830842

843+
def to_dict(self) -> Dict:
844+
"""
845+
Converts feature set to dict
846+
847+
:return: Dictionary object representation of feature set
848+
"""
849+
feature_set_dict = MessageToDict(self.to_proto())
850+
851+
# Remove meta when empty for more readable exports
852+
if feature_set_dict["meta"] == {}:
853+
del feature_set_dict["meta"]
854+
855+
return feature_set_dict
856+
857+
def to_yaml(self):
858+
"""
859+
Converts a feature set to a YAML string.
860+
861+
:return: Feature set string returned in YAML format
862+
"""
863+
feature_set_dict = self.to_dict()
864+
return yaml.dump(feature_set_dict, allow_unicode=True, sort_keys=False)
865+
831866

832867
class FeatureSetRef:
833868
"""

‎sdk/python/tests/test_feature_set.py‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
14+
1415
import pathlib
1516
from concurrent import futures
1617
from datetime import datetime
@@ -244,6 +245,27 @@ def test_export_tfx_schema(self):
244245
for actual, expected in zip(actual_schema.feature, expected_schema.feature):
245246
assert actual.SerializeToString() == expected.SerializeToString()
246247

248+
def test_feature_set_import_export_yaml(self):
249+
250+
test_feature_set = FeatureSet(
251+
name="bikeshare",
252+
entities=[Entity(name="station_id", dtype=ValueType.INT64)],
253+
features=[
254+
Feature(name="name", dtype=ValueType.STRING),
255+
Feature(name="longitude", dtype=ValueType.FLOAT),
256+
Feature(name="location", dtype=ValueType.STRING),
257+
],
258+
)
259+
260+
# Create a string YAML representation of the feature set
261+
string_yaml = test_feature_set.to_yaml()
262+
263+
# Create a new feature set object from the YAML string
264+
actual_feature_set_from_string = FeatureSet.from_yaml(string_yaml)
265+
266+
# Ensure equality is upheld to original feature set
267+
assert test_feature_set == actual_feature_set_from_string
268+
247269

248270
def make_tfx_schema_domain_info_inline(schema):
249271
# Copy top-level domain info defined in the schema to inline definition.

0 commit comments

Comments
 (0)