fix: Handle complex Spark data types in SparkSource (#3154)

rebstastic · adchia · commit 0e135734e313 · 2022-08-31T13:57:49.000-07:00
* Make sure data types are strings

Signed-off-by: Petra Rebernjak &lt;petra@noom.com&gt;

* Lint

Signed-off-by: Petra Rebernjak &lt;petra@noom.com&gt;

Signed-off-by: Petra Rebernjak &lt;petra@noom.com&gt;
diff --git a/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/spark_source.py b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/spark_source.py
@@ -159,10 +159,7 @@ def get_table_column_names_and_types(
             store_config=config.offline_store
         )
         df = spark_session.sql(f"SELECT * FROM {self.get_table_query_string()}")
-        return (
-            (fields["name"], fields["type"])
-            for fields in df.schema.jsonValue()["fields"]
-        )
+        return ((field.name, field.dataType.simpleString()) for field in df.schema)
 
     def get_table_query_string(self) -> str:
         """Returns a string that can directly be used to reference this table in SQL"""
