Query level api

Author: xando

src/pyarrow/bigquery/__init__.py

@@ -1,2 +1,2 @@
-from .read import reader, read_table  # noqa
+from .read import reader, read_table, reader_query, read_query  # noqa
 from .write import writer, write_table  # noqa

src/pyarrow/bigquery/read.py

@@ -7,6 +7,7 @@
 import shutil
 
 from google.cloud import bigquery_storage
+from google.cloud import bigquery
 
 import pyarrow as pa
 import pyarrow.feather as fa
@@ -154,6 +155,28 @@ def reader(
         logger.debug(f"Time taken to read: {time.time()-t0:.2f}")
 
 
+def reader_query(
+    project: str,
+    query: str,
+    *,
+    worker_count: int = multiprocessing.cpu_count(),
+    worker_type: type[threading.Thread] | type[multiprocessing.Process] = threading.Thread,
+    batch_size: int = 100,
+):
+    client = bigquery.Client(project=project)
+    job = client.query(query)
+    job.result()
+
+    source = f"{job.destination.project}.{job.destination.dataset_id}.{job.destination.table_id}"
+    return reader(
+        source=source,
+        project=project,
+        worker_count=worker_count,
+        worker_type=worker_type,
+        batch_size=batch_size,
+    )
+
+
 def read_table(
     source: str,
     *,
@@ -175,3 +198,21 @@ def read_table(
             batch_size=batch_size,
         )
     )
+
+def read_query(
+    project: str,
+    query: str,
+    *,
+    worker_count: int = multiprocessing.cpu_count(),
+    worker_type: type[threading.Thread] | type[multiprocessing.Process] = threading.Thread,
+    batch_size: int = 100,
+):
+    return pa.concat_tables(
+        reader_query(
+            project=project,
+            query=query,
+            worker_count=worker_count,
+            worker_type=worker_type,
+            batch_size=batch_size
+        )
+    )

tests/integration/test_upload.py

@@ -6,7 +6,8 @@
 import pyarrow.bigquery as bq
 
 
-LOCATION = f"{os.environ['GCP_PROJECT']}.test.{uuid.uuid4()}"
+PROJECT = os.environ['GCP_PROJECT']
+LOCATION = f"{PROJECT}.test.{uuid.uuid4()}"
 
 
 @pytest.fixture(autouse=True)
@@ -29,6 +30,22 @@ def test_simple():
     assert table_back.sort_by("test").equals(table.sort_by("test"))
 
 
+def test_reader_query():
+    table = pa.Table.from_arrays([[1, 2, 3, 4]], names=["test"])
+
+    bq.write_table(table, LOCATION, table_create=True)
+
+    query = f'SELECT * FROM `{LOCATION}`'
+    table_back1 = pa.concat_tables([t for t in bq.reader_query(project=PROJECT, query=query)])
+
+    table_back2 = bq.read_query(project=PROJECT, query=query)
+
+    assert table_back1.schema == table_back2.schema == table.schema
+
+    assert table_back1.sort_by("test").equals(table.sort_by("test"))
+    assert table_back2.sort_by("test").equals(table.sort_by("test"))
+
+
 def test_context():
     table = pa.Table.from_arrays([[1, 2, 3, 4]], names=["test"])