cr

Signed-off-by: Achal Shah <achals@gmail.com>

Author: achals

sdk/python/feast/data_source.py

@@ -560,8 +560,8 @@ def get_table_column_names_and_types(
     def from_proto(data_source: DataSourceProto):
         schema_pb = data_source.push_options.schema
         schema = {}
-        for key in schema_pb.keys():
-            schema[key] = ValueType(schema_pb.get(key))
+        for key, value in schema_pb.items():
+            schema[key] = value
 
         batch_source = None
         if data_source.push_options.HasField("batch_source"):

sdk/python/tests/unit/test_data_sources.py

@@ -31,4 +31,4 @@ def test_push_with_batch():
 
     assert push_source.name == push_source_unproto.name
     assert push_source.schema == push_source_unproto.schema
-    assert push_source.batch_source == push_source_unproto.batch_source
+    assert push_source.batch_source.name == push_source_unproto.batch_source.name