cr

achals · achals · commit 9182d3d6391e · 2022-03-22T11:34:40.000-07:00
Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;
diff --git a/sdk/python/feast/data_source.py b/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"):
diff --git a/sdk/python/tests/unit/test_data_sources.py b/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
