Fix inference of BigQuery ARRAY types. (feast-dev#2245)

judahrand · web-flow · commit 7c531774d099 · 2022-02-01T12:11:11.000-08:00
* Support more BigQuery ARRAY types

Signed-off-by: Judah Rand &lt;17158624+judahrand@users.noreply.github.com&gt;

* Correctly infer BigQuery ARRAY types

Signed-off-by: Judah Rand &lt;17158624+judahrand@users.noreply.github.com&gt;
diff --git a/sdk/python/feast/infra/offline_stores/bigquery_source.py b/sdk/python/feast/infra/offline_stores/bigquery_source.py
@@ -1,4 +1,4 @@
-from typing import Callable, Dict, Iterable, Optional, Tuple
+from typing import Callable, Dict, Iterable, List, Optional, Tuple
 
 from feast import type_map
 from feast.data_source import DataSource
@@ -123,18 +123,20 @@ def get_table_column_names_and_types(
 
         client = bigquery.Client()
         if self.table_ref is not None:
-            table_schema = client.get_table(self.table_ref).schema
-            if not isinstance(table_schema[0], bigquery.schema.SchemaField):
+            schema = client.get_table(self.table_ref).schema
+            if not isinstance(schema[0], bigquery.schema.SchemaField):
                 raise TypeError("Could not parse BigQuery table schema.")
-
-            name_type_pairs = [(field.name, field.field_type) for field in table_schema]
         else:
             bq_columns_query = f"SELECT * FROM ({self.query}) LIMIT 1"
             queryRes = client.query(bq_columns_query).result()
-            name_type_pairs = [
-                (schema_field.name, schema_field.field_type)
-                for schema_field in queryRes.schema
-            ]
+            schema = queryRes.schema
+
+        name_type_pairs: List[Tuple[str, str]] = []
+        for field in schema:
+            bq_type_as_str = field.field_type
+            if field.mode == "REPEATED":
+                bq_type_as_str = "ARRAY<" + bq_type_as_str + ">"
+            name_type_pairs.append((field.name, bq_type_as_str))
 
         return name_type_pairs
 
diff --git a/sdk/python/feast/type_map.py b/sdk/python/feast/type_map.py
@@ -442,6 +442,11 @@ def pa_to_feast_value_type(pa_type_as_str: str) -> ValueType:
 
 
 def bq_to_feast_value_type(bq_type_as_str: str) -> ValueType:
+    is_list = False
+    if bq_type_as_str.startswith("ARRAY<"):
+        is_list = True
+        bq_type_as_str = bq_type_as_str[6:-1]
+
     type_map: Dict[str, ValueType] = {
         "DATETIME": ValueType.UNIX_TIMESTAMP,
         "TIMESTAMP": ValueType.UNIX_TIMESTAMP,
@@ -453,15 +458,14 @@ def bq_to_feast_value_type(bq_type_as_str: str) -> ValueType:
         "BYTES": ValueType.BYTES,
         "BOOL": ValueType.BOOL,
         "BOOLEAN": ValueType.BOOL,  # legacy sql data type
-        "ARRAY<INT64>": ValueType.INT64_LIST,
-        "ARRAY<FLOAT64>": ValueType.DOUBLE_LIST,
-        "ARRAY<STRING>": ValueType.STRING_LIST,
-        "ARRAY<BYTES>": ValueType.BYTES_LIST,
-        "ARRAY<BOOL>": ValueType.BOOL_LIST,
         "NULL": ValueType.NULL,
     }
 
-    return type_map[bq_type_as_str]
+    value_type = type_map[bq_type_as_str]
+    if is_list:
+        value_type = ValueType[value_type.name + "_LIST"]
+
+    return value_type
 
 
 def redshift_to_feast_value_type(redshift_type_as_str: str) -> ValueType:
