chore: Change arrow scalar ids usage (feast-dev#4347)

* Update google-cloud-storage Signed-off-by: Christopher Camenares <[email protected]> * test tighter library restriction Signed-off-by: Christopher Camenares <[email protected]> * fix lint Signed-off-by: Christopher Camenares <[email protected]> * bump <4 again Signed-off-by: Christopher Camenares <[email protected]> --------- Signed-off-by: Christopher Camenares <[email protected]>
shuchu · Jul 16, 2024 · 5a16364 · 5a16364
1 parent a8bc696
commit 5a16364
Show file tree

Hide file tree

Showing 8 changed files with 532 additions and 85 deletions.
diff --git a/sdk/python/feast/infra/offline_stores/bigquery.py b/sdk/python/feast/infra/offline_stores/bigquery.py
@@ -59,14 +59,23 @@
     from google.auth.exceptions import DefaultCredentialsError
     from google.cloud import bigquery
     from google.cloud.bigquery import Client, SchemaField, Table
-    from google.cloud.bigquery._pandas_helpers import ARROW_SCALAR_IDS_TO_BQ
     from google.cloud.storage import Client as StorageClient
 
 except ImportError as e:
     from feast.errors import FeastExtrasDependencyImportError
 
     raise FeastExtrasDependencyImportError("gcp", str(e))
 
+try:
+    from google.cloud.bigquery._pyarrow_helpers import _ARROW_SCALAR_IDS_TO_BQ
+except ImportError:
+    try:
+        from google.cloud.bigquery._pandas_helpers import (  # type: ignore
+            ARROW_SCALAR_IDS_TO_BQ as _ARROW_SCALAR_IDS_TO_BQ,
+        )
+    except ImportError as e:
+        raise FeastExtrasDependencyImportError("gcp", str(e))
+
 
 def get_http_client_info():
     return http_client_info.ClientInfo(user_agent=get_user_agent())
@@ -794,10 +803,10 @@ def arrow_schema_to_bq_schema(arrow_schema: pyarrow.Schema) -> List[SchemaField]
     for field in arrow_schema:
         if pyarrow.types.is_list(field.type):
             detected_mode = "REPEATED"
-            detected_type = ARROW_SCALAR_IDS_TO_BQ[field.type.value_type.id]
+            detected_type = _ARROW_SCALAR_IDS_TO_BQ[field.type.value_type.id]
         else:
             detected_mode = "NULLABLE"
-            detected_type = ARROW_SCALAR_IDS_TO_BQ[field.type.id]
+            detected_type = _ARROW_SCALAR_IDS_TO_BQ[field.type.id]
 
         bq_schema.append(
             SchemaField(name=field.name, field_type=detected_type, mode=detected_mode)