Add update_schema method to FeatureSet

- Update Field, Feature and Entity class with fields from presence_constraints, shape_type and domain_info
feast-dev · davidheryanto · Jan 21, 2020 · Jan 21, 2020 · Jan 21, 2020 · Jan 21, 2020
commit d33c546951104798e5e8e9d0ec2f108ad29010ec
@@ -11,11 +11,12 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from tensorflow_metadata.proto.v0 import schema_pb2
 
-from feast.value_type import ValueType
 from feast.core.FeatureSet_pb2 import EntitySpec as EntityProto
-from feast.types import Value_pb2 as ValueTypeProto
 from feast.field import Field
+from feast.types import Value_pb2 as ValueTypeProto
+from feast.value_type import ValueType
 
 
 class Entity(Field):
@@ -29,17 +30,42 @@ def to_proto(self) -> EntityProto:
             Returns EntitySpec object
         """
         value_type = ValueTypeProto.ValueType.Enum.Value(self.dtype.name)
-        return EntityProto(name=self.name, value_type=value_type)
+        return EntityProto(
+            name=self.name,
+            value_type=value_type,
+            presence=self.presence,
+            group_presence=self.group_presence,
+            shape=self.shape,
+            value_count=self.value_count,
+            domain=self.domain,
+            int_domain=self.int_domain,
+            float_domain=self.float_domain,
+            string_domain=self.string_domain,
+            bool_domain=self.bool_domain,
+            struct_domain=self.struct_domain,
+            natural_language_domain=self.natural_language_domain,
+            image_domain=self.image_domain,
+            mid_domain=self.mid_domain,
+            url_domain=self.url_domain,
+            time_domain=self.time_domain,
+            time_of_day_domain=self.time_of_day_domain,
+        )
 
     @classmethod
-    def from_proto(cls, entity_proto: EntityProto):
+    def from_proto(cls, entity_proto: EntityProto, schema: schema_pb2.Schema = None):
         """
         Creates a Feast Entity object from its Protocol Buffer representation
 
         Args:
             entity_proto: EntitySpec protobuf object
+            schema: Schema from Tensorflow metadata, will be used to reference domain
+                defined at the schema level
 
         Returns:
             Entity object
         """
-        return cls(name=entity_proto.name, dtype=ValueType(entity_proto.value_type))
+        entity = cls(name=entity_proto.name, dtype=ValueType(entity_proto.value_type))
+        entity.update_presence_constraints(entity_proto)
+        entity.update_shape_type(entity_proto)
+        entity.update_domain_info(entity_proto, schema)
+        return entity
@@ -11,11 +11,12 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from tensorflow_metadata.proto.v0 import schema_pb2
 
-from feast.value_type import ValueType
 from feast.core.FeatureSet_pb2 import FeatureSpec as FeatureProto
-from feast.types import Value_pb2 as ValueTypeProto
 from feast.field import Field
+from feast.types import Value_pb2 as ValueTypeProto
+from feast.value_type import ValueType
 
 
 class Feature(Field):
@@ -24,9 +25,43 @@ class Feature(Field):
     def to_proto(self) -> FeatureProto:
         """Converts Feature object to its Protocol Buffer representation"""
         value_type = ValueTypeProto.ValueType.Enum.Value(self.dtype.name)
-        return FeatureProto(name=self.name, value_type=value_type)
+        return FeatureProto(
+            name=self.name,
+            value_type=value_type,
+            presence=self.presence,
+            group_presence=self.group_presence,
+            shape=self.shape,
+            value_count=self.value_count,
+            domain=self.domain,
+            int_domain=self.int_domain,
+            float_domain=self.float_domain,
+            string_domain=self.string_domain,
+            bool_domain=self.bool_domain,
+            struct_domain=self.struct_domain,
+            natural_language_domain=self.natural_language_domain,
+            image_domain=self.image_domain,
+            mid_domain=self.mid_domain,
+            url_domain=self.url_domain,
+            time_domain=self.time_domain,
+            time_of_day_domain=self.time_of_day_domain,
+        )
 
     @classmethod
-    def from_proto(cls, feature_proto: FeatureProto):
-        """Converts Protobuf Feature to its SDK equivalent"""
-        return cls(name=feature_proto.name, dtype=ValueType(feature_proto.value_type))
+    def from_proto(cls, feature_proto: FeatureProto, schema: schema_pb2.Schema = None):
+        """
+
+        Args:
+            feature_proto: FeatureSpec protobuf object
+            schema: Schema from Tensorflow metadata, will be used to reference domain
+                defined at the schema level
+
+        Returns:
+            Feature object
+        """
+        feature = cls(
+            name=feature_proto.name, dtype=ValueType(feature_proto.value_type)
+        )
+        feature.update_presence_constraints(feature_proto)
+        feature.update_shape_type(feature_proto)
+        feature.update_domain_info(feature_proto, schema)
+        return feature
@@ -11,15 +11,20 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
+import warnings
 from collections import OrderedDict
 from typing import Dict
 from typing import List, Optional
-from tensorflow_metadata.proto.v0.schema_pb2 import Schema
 
 import pandas as pd
 import pyarrow as pa
+from google.protobuf import json_format
+from google.protobuf.duration_pb2 import Duration
+from google.protobuf.json_format import MessageToJson
+from pandas.api.types import is_datetime64_ns_dtype
+from pyarrow.lib import TimestampType
+from tensorflow_metadata.proto.v0.schema_pb2 import Schema
+
 from feast.core.FeatureSet_pb2 import FeatureSet as FeatureSetProto
 from feast.core.FeatureSet_pb2 import FeatureSetMeta as FeatureSetMetaProto
 from feast.core.FeatureSet_pb2 import FeatureSetSpec as FeatureSetSpecProto
@@ -30,15 +35,6 @@
 from feast.type_map import DATETIME_COLUMN
 from feast.type_map import pa_to_feast_value_type
 from feast.type_map import python_type_to_feast_value_type
-from google.protobuf import json_format
-from feast.core.FeatureSet_pb2 import FeatureSetSpec as FeatureSetSpecProto
-from feast.core.FeatureSet_pb2 import FeatureSetMeta as FeatureSetMetaProto
-from feast.core.FeatureSet_pb2 import FeatureSet as FeatureSetProto
-from google.protobuf.duration_pb2 import Duration
-from feast.type_map import python_type_to_feast_value_type
-from google.protobuf.json_format import MessageToJson
-from pandas.api.types import is_datetime64_ns_dtype
-from pyarrow.lib import TimestampType
 
 
 class FeatureSet:
@@ -664,11 +660,40 @@ def is_valid(self):
         if len(self.entities) == 0:
             raise ValueError(f"No entities found in feature set {self.name}")
 
-    def update_schema(self):
-        pass
+    def update_schema(self, schema: Schema):
+        """
+        Updates presence_constraints, shape_type and domain_info for all entities
+        and features in the FeatureSet from schema in Tensorflow metadata.
+
+        Args:
+            schema: schema from Tensorflow metadata
+
+        Returns:
+            None
 
-    def get_schema(self) -> Schema:
-        pass
+        """
+        name_to_feature = {f.name: f for f in self.features}
+        name_to_entity = {e.name: e for e in self.entities}
+
+        for feature_from_new_schema in schema.feature:
+
+            if feature_from_new_schema.name in name_to_feature:
+                feature = name_to_feature[feature_from_new_schema.name]
+                feature.update_presence_constraints(feature_from_new_schema)
+                feature.update_shape_type(feature_from_new_schema)
+                feature.update_domain_info(feature_from_new_schema, schema)
+
+            elif feature_from_new_schema.name in name_to_entity:
+                entity = name_to_entity[feature_from_new_schema.name]
+                entity.update_presence_constraints(feature_from_new_schema)
+                entity.update_shape_type(feature_from_new_schema)
+                entity.update_domain_info(feature_from_new_schema, schema)
+
+            else:
+                warnings.warn(
+                    f"The provided schema contains feature name '{feature_from_new_schema.name}' "
+                    f"that does not exist in the FeatureSet '{self.name}' in Feast"
+                )
 
     @classmethod
     def from_yaml(cls, yml: str):