Use ARROW as data format in DB API cursor

googleapis · May 14, 2020 · c715a74 · c715a74
1 parent 530a5db
commit c715a74
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 4 deletions.
diff --git a/google/cloud/bigquery/dbapi/_helpers.py b/google/cloud/bigquery/dbapi/_helpers.py
@@ -234,7 +234,9 @@ def to_bq_table_rows(rows_iterable):
     """
 
     def to_table_row(row):
-        values = tuple(row.values())
+        # NOTE: We fetch ARROW values, thus we need to convert them to Python
+        # objects with as_py().
+        values = tuple(value.as_py() for value in row.values())
         keys_to_index = {key: i for i, key in enumerate(row.keys())}
         return table.Row(values, keys_to_index)
 

diff --git a/google/cloud/bigquery/dbapi/cursor.py b/google/cloud/bigquery/dbapi/cursor.py
@@ -292,11 +292,12 @@ def _bqstorage_fetch(self, bqstorage_client):
                 "projects/{}".format(table_reference.project),
                 # a single stream only, as DB API is not well-suited for multithreading
                 requested_streams=1,
+                format_=bigquery_storage_v1beta1.enums.DataFormat.ARROW,
             )
         else:
             requested_session = bigquery_storage_v1.types.ReadSession(
                 table=table_reference.to_bqstorage(),
-                data_format=bigquery_storage_v1.enums.DataFormat.AVRO,
+                data_format=bigquery_storage_v1.enums.DataFormat.ARROW,
             )
             read_session = bqstorage_client.create_read_session(
                 parent="projects/{}".format(table_reference.project),

diff --git a/tests/system.py b/tests/system.py
@@ -1667,7 +1667,7 @@ def test_dbapi_fetch_w_bqstorage_client_small_result_set(self):
     @unittest.skipIf(
         bigquery_storage_v1 is None, "Requires `google-cloud-bigquery-storage`"
     )
-    @unittest.skipIf(fastavro is None, "Requires `fastavro`")
+    @unittest.skipIf(pyarrow is None, "Requires `pyarrow`")
     def test_dbapi_fetch_w_bqstorage_client_large_result_set(self):
         bqstorage_client = bigquery_storage_v1.BigQueryReadClient(
             credentials=Config.CLIENT._credentials
@@ -1717,7 +1717,7 @@ def test_dbapi_fetch_w_bqstorage_client_large_result_set(self):
     @unittest.skipIf(
         bigquery_storage_v1beta1 is None, "Requires `google-cloud-bigquery-storage`"
     )
-    @unittest.skipIf(fastavro is None, "Requires `fastavro`")
+    @unittest.skipIf(pyarrow is None, "Requires `pyarrow`")
     def test_dbapi_fetch_w_bqstorage_client_v1beta1_large_result_set(self):
         bqstorage_client = bigquery_storage_v1beta1.BigQueryStorageClient(
             credentials=Config.CLIENT._credentials