Skip to content

Commit

Permalink
chore: Change arrow scalar ids usage (#4347)
Browse files Browse the repository at this point in the history
* Update google-cloud-storage

Signed-off-by: Christopher Camenares <ccamenares@gmail.com>

* test tighter library restriction

Signed-off-by: Christopher Camenares <ccamenares@gmail.com>

* fix lint

Signed-off-by: Christopher Camenares <ccamenares@gmail.com>

* bump <4 again

Signed-off-by: Christopher Camenares <ccamenares@gmail.com>

---------

Signed-off-by: Christopher Camenares <ccamenares@gmail.com>
  • Loading branch information
camenares authored Jul 16, 2024
1 parent a8bc696 commit 5a16364
Show file tree
Hide file tree
Showing 8 changed files with 532 additions and 85 deletions.
15 changes: 12 additions & 3 deletions sdk/python/feast/infra/offline_stores/bigquery.py
Original file line number Diff line number Diff line change
Expand Up @@ -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())
Expand Down Expand Up @@ -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)
Expand Down
Loading

0 comments on commit 5a16364

Please sign in to comment.