pyarrow-bigquery 0.1.0__py3-none-any.whl

pyarrow/bigquery/__init__.py

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

pyarrow/bigquery/read.py

@@ -0,0 +1,232 @@
+from __future__ import annotations
+
+import logging
+import os
+import time
+import tempfile
+import multiprocessing
+import threading
+import shutil
+
+from google.cloud import bigquery_storage
+from google.cloud import bigquery
+from google.cloud.exceptions import NotFound
+
+import pyarrow as pa
+import pyarrow.feather as fa
+
+from . import some_itertools
+
+
+logger = logging.getLogger(__name__)
+
+
+def _bq_table_exists(project: str, location: str):
+    client = bigquery.Client(project=project)
+
+    try:
+        client.get_table(location)
+        logger.debug(f"Table {location} already exists")
+    except NotFound as e:
+        logger.debug("Table {location} is not found")
+        raise e
+
+
+def _bq_read_create_strems(
+    read_client: bigquery_storage.BigQueryReadClient,
+    parent: str,
+    location: str,
+    selected_fields: list | None,
+    row_restrictions: str | None,
+    max_stream_count: int,
+) -> tuple[list[str], pa.Schema]:
+    project, dataset, table = location.split(".")
+
+    read_session = bigquery_storage.ReadSession(
+        table=f"projects/{project}/datasets/{dataset}/tables/{table}",
+        data_format=bigquery_storage.DataFormat.ARROW,
+        read_options={
+            "selected_fields": selected_fields,
+            "row_restriction": row_restrictions,
+        },
+    )
+
+    read_session = read_client.create_read_session(
+        parent=f"projects/{parent}",
+        read_session=read_session,
+        max_stream_count=max_stream_count,
+    )
+
+    schema_buffer = pa.py_buffer(read_session.arrow_schema.serialized_schema)
+    schema = pa.ipc.read_schema(schema_buffer)
+
+    return read_session.streams, schema
+
+
+def _stream_worker(read_client, read_streams, table_schema, batch_size, queue_results, temp_dir):
+    batches = []
+
+    for stream in read_streams:
+        t = time.time()
+
+        for message in read_client.read_rows(stream.name):
+            record_batch = pa.ipc.read_record_batch(message.arrow_record_batch.serialized_record_batch, table_schema)
+
+            batches.append(record_batch)
+
+            if sum(b.num_rows for b in batches) >= batch_size:
+                table = pa.Table.from_batches(batches)
+
+                element = tempfile.mktemp(dir=temp_dir)
+                fa.write_feather(table[:batch_size], element)
+                queue_results.put(element)
+
+                batches = table[batch_size:].to_batches()
+
+        logger.debug(f"Stream {stream.name} done in {time.time()-t:.2f} seconds")
+
+    if batches:
+        table = pa.Table.from_batches(batches)
+        element = tempfile.mktemp(dir=temp_dir)
+        fa.write_feather(table, element)
+        queue_results.put(element)
+
+    queue_results.put(None)
+
+
+def reader(
+    source: str,
+    *,
+    project: str | None = None,
+    columns: list | None = None,
+    row_restrictions: str | None = None,
+    worker_count: int = multiprocessing.cpu_count(),
+    worker_type: type[threading.Thread] | type[multiprocessing.Process] = threading.Thread,
+    batch_size: int = 100,
+):
+    t0 = time.time()
+    project_id, *_ = source.split(".")
+
+    if not project:
+        project = project_id
+
+    queue_results = multiprocessing.Queue()
+    read_client = bigquery_storage.BigQueryReadClient()
+
+    _bq_table_exists(project, source)
+
+    streams, streams_schema = _bq_read_create_strems(
+        read_client=read_client,
+        parent=project,
+        location=source,
+        selected_fields=columns,
+        row_restrictions=row_restrictions,
+        max_stream_count=worker_count * 3,
+    )
+    workers_done = 0
+
+    assert streams, "No streams to read, Table might be empty"
+
+    logger.debug(f"Number of workers: {worker_count}, number of streams: {len(streams)}")
+
+    actual_worker_count = min(worker_count, len(streams))
+
+    logger.debug(f"Actual worker count: {actual_worker_count}")
+
+    temp_dir = tempfile.mkdtemp()
+
+    try:
+        for streams in some_itertools.to_split(streams, actual_worker_count):
+            e = worker_type(
+                target=_stream_worker,
+                args=(
+                    read_client,
+                    streams,
+                    streams_schema,
+                    batch_size,
+                    queue_results,
+                    temp_dir,
+                ),
+            )
+            e.start()
+
+        while True:
+            element = queue_results.get()
+
+            if not element:
+                workers_done += 1
+
+                if workers_done == actual_worker_count:
+                    break
+            else:
+                table = fa.read_table(element)
+                os.remove(element)
+                yield table
+    finally:
+        t = time.time()
+        shutil.rmtree(temp_dir, ignore_errors=True)
+        logger.debug(f"Time to cleanup temp directory: {time.time()-t:.2f}")
+        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,
+    *,
+    project: str | None = None,
+    columns: list | None = None,
+    row_restrictions: str | None = None,
+    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(
+            source=source,
+            project=project,
+            columns=columns,
+            row_restrictions=row_restrictions,
+            worker_count=worker_count,
+            worker_type=worker_type,
+            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
+        )
+    )

pyarrow/bigquery/some_itertools.py

@@ -0,0 +1,12 @@
+def to_chunks(table, chunk_size):
+    for i in range(0, len(table), chunk_size):
+        yield table[i : i + chunk_size]
+
+
+def to_split(table, split_number):
+    k, m = divmod(len(table), split_number)
+    return (table[i * k + min(i, m) : (i + 1) * k + min(i + 1, m)] for i in range(split_number))
+
+
+def first(iterable, condition):
+    return next((x for x in iterable if condition(x)), None)

pyarrow/bigquery/write/__init__.py

@@ -0,0 +1,238 @@
+from __future__ import annotations
+
+import os
+import time
+import datetime
+import multiprocessing
+import threading
+import shutil
+import tempfile
+import logging
+import collections
+
+
+from google.cloud import bigquery
+from google.cloud import bigquery_storage_v1
+from google.cloud.bigquery_storage_v1.writer import AppendRowsStream
+from google.protobuf import descriptor_pb2
+from google.api_core.exceptions import Unknown, NotFound
+from google.api_core import retry
+
+import pyarrow as pa
+import pyarrow.feather as fa
+
+from . import pa_to_bq
+from . import pa_to_pb
+from . import upload
+from .. import some_itertools
+
+logger = logging.getLogger(__name__)
+
+
+Stream = collections.namedtuple("Stream", ["append_rows_stream", "write_stream"])
+
+
+def _bq_create_table(*, project, location, schema, expire, overwrite):
+    client = bigquery.Client(project=project)
+
+    if overwrite:
+        client.delete_table(location, not_found_ok=True)
+
+    bq_schema = pa_to_bq.generate(schema)
+
+    table = bigquery.Table(location, schema=bq_schema)
+
+    client.create_table(table)
+
+    if expire:
+        table.expires = datetime.datetime.now() + datetime.timedelta(seconds=expire)
+        client.update_table(table, ["expires"])
+
+    logger.debug(f"Created BigQuery table '{location}'")
+
+
+def _bq_write_create_stream(write_client: bigquery_storage_v1.BigQueryWriteClient, parent, protobuf_definition):
+    write_stream = write_client.create_write_stream(
+        parent=parent,
+        write_stream=bigquery_storage_v1.types.WriteStream(type=bigquery_storage_v1.types.WriteStream.Type.PENDING),
+        retry=retry.Retry(predicate=retry.if_exception_type(Unknown, NotFound)),
+    )
+
+    proto_schema = bigquery_storage_v1.types.ProtoSchema()
+    proto_descriptor = descriptor_pb2.DescriptorProto()
+    protobuf_definition.CopyToProto(proto_descriptor)
+    proto_schema.proto_descriptor = proto_descriptor
+
+    proto_data = bigquery_storage_v1.types.AppendRowsRequest.ProtoData()
+    proto_data.writer_schema = proto_schema
+
+    request_template = bigquery_storage_v1.types.AppendRowsRequest()
+    request_template.write_stream = write_stream.name
+    request_template.proto_rows = proto_data
+
+    append_rows_stream = AppendRowsStream(write_client, request_template)
+
+    return Stream(
+        write_stream=write_stream,
+        append_rows_stream=append_rows_stream,
+    )
+
+
+def _bq_storage_close_stream(write_client, stream, parent):
+    stream.append_rows_stream.close()
+    write_client.finalize_write_stream(name=stream.write_stream.name)
+
+    batch_commit_write_streams_request = bigquery_storage_v1.types.BatchCommitWriteStreamsRequest()
+    batch_commit_write_streams_request.parent = parent
+    batch_commit_write_streams_request.write_streams = [stream.write_stream.name]
+
+    write_client.batch_commit_write_streams(batch_commit_write_streams_request)
+
+    logger.debug(f"Stream '{stream.write_stream.name}' closed")
+
+
+def _stream_worker(
+    write_client: bigquery_storage_v1.BigQueryWriteClient,
+    parent: str,
+    schema_protobuf,
+    queue_results,
+):
+    stream = _bq_write_create_stream(write_client, parent, schema_protobuf)
+
+    offset = 0
+
+    while True:
+        element = queue_results.get()
+        if element is None:
+            break
+
+        table = fa.read_table(element)
+
+        upload.upload_data(stream, table, schema_protobuf, offset)
+
+        os.remove(element)
+
+        offset += table.num_rows
+
+    _bq_storage_close_stream(write_client, stream, parent)
+
+
+class writer:
+    """Method to handle"""
+
+    def __init__(
+        self,
+        schema: pa.Schema,
+        where: str,
+        *,
+        project: str | None = None,
+        table_create: bool = True,
+        table_expire: int | None = None,
+        table_overwrite: bool = False,
+        worker_count: int = multiprocessing.cpu_count(),
+        worker_type: type[threading.Thread] | type[multiprocessing.Process] = threading.Thread,
+        batch_size: int = 100,
+    ):
+        self.project = project
+        self.where = where
+        self.schema = schema
+
+        self.table_create = table_create
+        self.table_expire = table_expire
+        self.table_overwrite = table_overwrite
+
+        self.worker_count = worker_count
+        self.worker_type = worker_type
+
+        self.batch_size = batch_size
+
+        project_id, dataset_id, table_id = where.split(".")
+
+        self.parent = f"projects/{project_id}/datasets/{dataset_id}/tables/{table_id}"
+
+        if not self.project:
+            self.project = project_id
+
+    def __enter__(self):
+        self.t0 = time.time()
+        self.temp_dir = tempfile.mkdtemp()
+        self.schema_protobuf = pa_to_pb.generate(self.schema)
+
+        if self.table_create:
+            _bq_create_table(
+                project=self.project,
+                location=self.where,
+                schema=self.schema,
+                expire=self.table_expire,
+                overwrite=self.table_overwrite,
+            )
+
+        self.queue_results = multiprocessing.Queue()
+        self.workers = []
+        write_client = bigquery_storage_v1.BigQueryWriteClient()
+
+        for _ in range(self.worker_count):
+            worker = self.worker_type(
+                target=_stream_worker,
+                args=(
+                    write_client,
+                    self.parent,
+                    self.schema_protobuf,
+                    self.queue_results,
+                ),
+            )
+            worker.start()
+            self.workers.append(worker)
+
+        return self
+
+    def write_table(self, table):
+        for table_chunk in some_itertools.to_chunks(table, self.batch_size):
+            element = tempfile.mktemp(dir=self.temp_dir)
+            fa.write_feather(table_chunk, element)
+            self.queue_results.put(element)
+
+    def write_batch(self, batch):
+        element = tempfile.mktemp(dir=self.temp_dir)
+        fa.write_feather(pa.Table.from_batches([batch]), element)
+        self.queue_results.put(element)
+
+    def __exit__(self, *_, **__):
+        for _ in range(self.worker_count):
+            self.queue_results.put(None)
+
+        for w in self.workers:
+            w.join()
+
+        shutil.rmtree(self.temp_dir, ignore_errors=True)
+
+        logger.debug(f"Time taken: {time.time() - self.t0}")
+
+
+def write_table(
+    table: pa.Table,
+    where: str,
+    *,
+    project: str | None = None,
+    table_create: bool = True,
+    table_expire: int | None = None,
+    table_overwrite: bool = False,
+    worker_count: int = multiprocessing.cpu_count(),
+    worker_type: type[threading.Thread] | type[multiprocessing.Process] = threading.Thread,
+    batch_size: int = 100,
+):
+    assert table.num_rows > 0, "Table is empty"
+
+    with writer(
+        where=where,
+        schema=table.schema,
+        project=project,
+        table_create=table_create,
+        table_expire=table_expire,
+        table_overwrite=table_overwrite,
+        worker_count=worker_count,
+        worker_type=worker_type,
+        batch_size=batch_size,
+    ) as w:
+        for table_chunk in some_itertools.to_split(table, w.worker_count):
+            w.write_table(table_chunk)

pyarrow/bigquery/write/pa_to_bq.py

@@ -0,0 +1,48 @@
+import pyarrow as pa
+
+from google.cloud import bigquery
+
+from .type_mapping import TYPES_MAPPING
+from ..some_itertools import first
+
+
+def emit(schema):
+    fields = []
+
+    for field in schema:
+        field_mode = "NULLABLE" if field.nullable else "REQUIRED"
+        field_type = field.type
+
+        if pa.types.is_list(field_type):
+            field_mode = "REPEATED"
+            field_type = field.type.value_type
+
+            if pa.types.is_list(field.type.value_type):
+                raise TypeError("Nested lists are not supported")
+
+        if pa.types.is_struct(field_type):
+            fields.append(
+                bigquery.SchemaField(
+                    name=field.name,
+                    field_type="RECORD",
+                    mode=field_mode,
+                    fields=emit(field_type),
+                )
+            )
+
+        elif type_check := first(TYPES_MAPPING, lambda type_check: type_check(field_type)):
+            fields.append(
+                bigquery.SchemaField(
+                    name=field.name,
+                    field_type=TYPES_MAPPING[type_check].bq,
+                    mode=field_mode,
+                )
+            )
+        else:
+            raise TypeError(f"Unsupported type {field_type}")
+
+    return fields
+
+
+def generate(schema):
+    return emit(schema)

pyarrow/bigquery/write/pa_to_pb.py

@@ -0,0 +1,151 @@
+import logging
+import time
+import random
+import string
+
+from google.protobuf.descriptor_pb2 import DescriptorProto, FieldDescriptorProto
+from google.protobuf import descriptor_pb2, descriptor_pool, message_factory
+
+import pyarrow as pa
+from .type_mapping import TYPES_MAPPING
+from ..some_itertools import first
+
+logger = logging.getLogger(__name__)
+
+
+FORCE_STRING = [
+    pa.types.is_date,
+    pa.types.is_decimal,
+    pa.types.is_time,
+    pa.types.is_timestamp,
+]
+
+GRPC_UPLOAD_LIMIT = 10485760
+
+
+def random_string(length):
+    characters = string.ascii_letters + string.digits
+    return "".join(random.choices(characters, k=length))
+
+
+def emit(schema, message_name):
+    message_descriptor = DescriptorProto()
+    message_descriptor.name = message_name
+
+    for idx, field in enumerate(schema, 1):
+        _optional = field.nullable
+        _repeated = False
+        _type = field.type
+
+        if pa.types.is_list(_type):
+            _repeated = True
+            _type = field.type.value_type
+
+            if pa.types.is_list(field.type.value_type):
+                raise TypeError("Nested lists are not supported")
+
+        if pa.types.is_struct(_type):
+            proto_type = f"{message_name}_{field.name}"
+
+            message_descriptor.nested_type.extend([emit(schema=_type, message_name=proto_type)])
+
+            label = (
+                FieldDescriptorProto.LABEL_REPEATED
+                if _repeated
+                else FieldDescriptorProto.LABEL_OPTIONAL
+                if _optional
+                else FieldDescriptorProto.LABEL_REQUIRED
+            )
+
+            message_descriptor.field.add(
+                name=field.name,
+                number=idx,
+                label=label,
+                type=FieldDescriptorProto.TYPE_MESSAGE,
+                type_name=proto_type,
+            )
+
+        elif type_check := first(TYPES_MAPPING, lambda type_check: type_check(_type)):
+            label = (
+                FieldDescriptorProto.LABEL_REPEATED
+                if _repeated
+                else FieldDescriptorProto.LABEL_OPTIONAL
+                if _optional
+                else FieldDescriptorProto.LABEL_REQUIRED
+            )
+            message_descriptor.field.add(
+                name=field.name,
+                number=idx,
+                label=label,
+                type=TYPES_MAPPING[type_check].pb,
+            )
+
+        else:
+            raise TypeError(f"Unsupported type {_type}")
+
+    return message_descriptor
+
+
+def generate(schema):
+    # NOTE. (I think)
+    #
+    # 1)    Since we are using the same descriptor pool,
+    #       we need to make sure that the message name is unique.
+    # 2)    Above applies to files added to the pool as well.
+
+    message_name = f"Message_{random_string(10)}"
+    file_name = f"{random_string(10)}.proto"
+
+    message_type = emit(schema, message_name=message_name)
+
+    pool = descriptor_pool.Default()
+    pool.AddSerializedFile(
+        descriptor_pb2.FileDescriptorProto(name=file_name, message_type=[message_type]).SerializeToString()
+    )
+
+    return pool.FindMessageTypeByName(message_name)
+
+
+def ammend_schema(schema):
+    def _cast_field(field):
+        if pa.types.is_struct(field.type):
+            new_fields = [_cast_field(sub_field) for sub_field in field.type]
+            return pa.field(field.name, pa.struct(new_fields))
+
+        elif pa.types.is_list(field.type):
+            new_value_field = _cast_field(field.type.value_field)
+            return pa.field(field.name, pa.list_(new_value_field.type))
+
+        elif any(type_check(field.type) for type_check in FORCE_STRING):
+            return pa.field(field.name, pa.string())
+        else:
+            return field
+
+    new_fields = [_cast_field(field) for field in schema]
+    return pa.schema(new_fields)
+
+
+def serialize(pa_table, protobuf_definition):
+    Message = message_factory.GetMessageClass(protobuf_definition)
+
+    pa_table = pa_table.cast(ammend_schema(pa_table.schema))
+
+    rows = []
+    size = 0
+
+    for element in pa_table.to_pylist():
+        t0 = time.time()
+        message = Message(**element)
+        size += message.ByteSize()
+
+        logger.debug(f"Time taken to serialize: {(time.time() - t0):.4f}")
+
+        if size > GRPC_UPLOAD_LIMIT:
+            assert rows, "Row is too large to fit in a single request"
+            yield rows
+            rows = []
+
+        rows.append(message.SerializeToString())
+
+    if rows:
+        yield rows

pyarrow/bigquery/write/type_mapping.py

@@ -0,0 +1,21 @@
+import collections
+
+import pyarrow as pa
+from google.protobuf.descriptor_pb2 import FieldDescriptorProto
+from google.cloud import bigquery
+
+
+TypeMapping = collections.namedtuple("TypeMapping", ["bq", "pb"])
+
+
+TYPES_MAPPING = {
+    pa.types.is_binary: TypeMapping(bigquery.SqlTypeNames.BYTES, FieldDescriptorProto.TYPE_BYTES),
+    pa.types.is_boolean: TypeMapping(bigquery.SqlTypeNames.BOOLEAN, FieldDescriptorProto.TYPE_BOOL),
+    pa.types.is_date: TypeMapping(bigquery.SqlTypeNames.DATE, FieldDescriptorProto.TYPE_STRING),
+    pa.types.is_decimal: TypeMapping(bigquery.SqlTypeNames.DECIMAL, FieldDescriptorProto.TYPE_STRING),
+    pa.types.is_floating: TypeMapping(bigquery.SqlTypeNames.FLOAT64, FieldDescriptorProto.TYPE_DOUBLE),
+    pa.types.is_integer: TypeMapping(bigquery.SqlTypeNames.INT64, FieldDescriptorProto.TYPE_INT64),
+    pa.types.is_string: TypeMapping(bigquery.SqlTypeNames.STRING, FieldDescriptorProto.TYPE_STRING),
+    pa.types.is_time: TypeMapping(bigquery.SqlTypeNames.TIME, FieldDescriptorProto.TYPE_STRING),
+    pa.types.is_timestamp: TypeMapping(bigquery.SqlTypeNames.TIMESTAMP, FieldDescriptorProto.TYPE_STRING),
+}

pyarrow/bigquery/write/upload.py

@@ -0,0 +1,30 @@
+import tenacity
+
+from google.cloud.bigquery_storage_v1 import types
+from google.api_core.exceptions import Unknown
+
+from . import pa_to_pb
+
+
+@tenacity.retry(
+        stop=tenacity.stop_after_attempt(5),
+        retry=tenacity.retry_if_exception_type(Unknown))
+def _send(stream, serialized_rows, offset):
+    proto_rows = types.ProtoRows()
+    proto_rows.serialized_rows.extend(serialized_rows)
+
+    proto_data = types.AppendRowsRequest.ProtoData()
+    proto_data.rows = proto_rows
+
+    request = types.AppendRowsRequest()
+    request.offset = offset
+    request.proto_rows = proto_data
+
+    stream.append_rows_stream.send(request).result()
+
+
+def upload_data(stream, pa_table, protobuf_definition, offset):
+    local_offset = 0
+    for serialized_rows in pa_to_pb.serialize(pa_table, protobuf_definition):
+        _send(stream, serialized_rows, offset + local_offset)
+        local_offset += len(serialized_rows)

pyarrow_bigquery-0.1.0.dist-info/METADATA

@@ -0,0 +1,262 @@
+Metadata-Version: 2.1
+Name: pyarrow-bigquery
+Version: 0.1.0
+Summary: A simple library to **write to** and **download from** BigQuery tables as PyArrow tables.
+Author-email: Sebastian Pawluś <sebastian.pawlus@gmail.com>
+License: MIT
+Keywords: pyarrow,bigquery
+Description-Content-Type: text/markdown
+Requires-Dist: google-cloud-bigquery <5,>=3
+Requires-Dist: google-cloud-bigquery-storage <3,>=2
+Requires-Dist: pyarrow <17,>=16
+Requires-Dist: tenacity
+
+# pyarrow-bigquery
+
+A simple library to **write to** and **download from** BigQuery tables as PyArrow tables.
+
+## Installation 
+
+```bash 
+pip install pyarrow-bigquery
+```
+
+## Quick Start
+
+This guide will help you quickly get started with `pyarrow-bigquery`, a library that allows you to **read** from and **write** to Google BigQuery using PyArrow.
+
+### Reading from BigQuery
+
+`pyarrow-bigquery` exposes two methods to read BigQuery tables as PyArrow tables. Depending on your use case or the size of the table, you might want to use one method over the other.
+
+#### Read the Whole Table
+
+When the table is small enough to fit in memory, you can read it directly using `bq.read_table`.
+
+```python
+import pyarrow.bigquery as bq
+
+table = bq.read_table("gcp_project.dataset.small_table")
+
+print(table.num_rows)
+```
+
+#### Read with Batches
+
+If the target table is larger than memory or you have other reasons not to fetch the whole table at once, you can use the `bq.reader` iterator method along with the `batch_size` parameter to limit how much data is fetched per iteration.
+
+```python
+import pyarrow.bigquery as bq
+
+for table in bq.reader("gcp_project.dataset.big_table", batch_size=100):
+    print(table.num_rows)
+```
+
+### Writing to BigQuery
+
+Similarly, the package exposes two methods to write to BigQuery. Depending on your use case or the size of the table, you might want to use one method over the other.
+
+#### Write the Whole Table
+
+When you want to write a complete table at once, you can use the `bq.write_table` method.
+
+```python
+import pyarrow as pa
+import pyarrow.bigquery as bq
+
+table = pa.Table.from_arrays([[1, 2, 3, 4]], names=['integers'])
+
+bq.write_table(table, 'gcp_project.dataset.table')
+```
+
+#### Write in Batches (Smaller Chunks)
+
+If you need to write data in smaller chunks, you can use the `bq.writer` method with the `schema` parameter to define the table structure.
+
+```python
+import pyarrow as pa
+import pyarrow.bigquery as bq
+
+schema = pa.schema([
+    ("integers", pa.int64())
+])
+
+with bq.writer("gcp_project.dataset.table", schema=schema) as w:
+    w.write_batch(record_batch)
+    w.write_table(table)
+```
+
+## API Reference
+
+### `pyarrow.bigquery.write_table`
+
+Write a PyArrow Table to a BigQuery Table. No return value.
+
+**Parameters:**
+
+- `table`: `pa.Table`  
+  PyArrow table.
+
+- `where`: `str`  
+  Destination location in BigQuery catalog.
+
+- `project`: `str`, *default* `None`  
+  BigQuery execution project, also the billing project. If not provided, it will be extracted from `where`.
+
+- `table_create`: `bool`, *default* `True`  
+  Specifies if the BigQuery table should be created.
+
+- `table_expire`: `None | int`, *default* `None`  
+  Amount of seconds after which the created table will expire. Used only if `table_create` is `True`. Set to `None` to disable expiration.
+
+- `table_overwrite`: `bool`, *default* `False`  
+  If the table already exists, destroy it and create a new one.
+
+- `worker_type`: `threading.Thread | multiprocessing.Process`, *default* `threading.Thread`  
+  Worker backend for fetching data.
+
+- `worker_count`: `int`, *default* `os.cpu_count()`  
+  Number of threads or processes to use for fetching data from BigQuery.
+
+- `batch_size`: `int`, *default* `100`  
+  Batch size for fetched rows.
+
+```python
+bq.write_table(table, 'gcp_project.dataset.table')
+```
+
+### `pyarrow.bigquery.writer`
+
+Context manager version of the write method. Useful when the PyArrow table is larger than memory size or the table is available in chunks.
+
+**Parameters:**
+
+- `schema`: `pa.Schema`  
+  PyArrow schema.
+
+- `where`: `str`  
+  Destination location in BigQuery catalog.
+
+- `project`: `str`, *default* `None`  
+  BigQuery execution project, also the billing project. If not provided, it will be extracted from `where`.
+
+- `table_create`: `bool`, *default* `True`  
+  Specifies if the BigQuery table should be created.
+
+- `table_expire`: `None | int`, *default* `None`  
+  Amount of seconds after which the created table will expire. Used only if `table_create` is `True`. Set to `None` to disable expiration.
+
+- `table_overwrite`: `bool`, *default* `False`  
+  If the table already exists, destroy it and create a new one.
+
+- `worker_type`: `threading.Thread | multiprocessing.Process`, *default* `threading.Thread`  
+  Worker backend for writing data.
+
+- `worker_count`: `int`, *default* `os.cpu_count()`  
+  Number of threads or processes to use for writing data to BigQuery.
+
+- `batch_size`: `int`, *default* `100`  
+  Batch size used for writes. Table will be automatically split to this value.
+
+Depending on the use case, you might want to use one of the methods below to write your data to a BigQuery table, using either `pa.Table` or `pa.RecordBatch`.
+
+#### `pyarrow.bigquery.writer.write_table`
+
+Context manager method to write a table.
+
+**Parameters:**
+
+- `table`: `pa.Table`  
+  PyArrow table.
+
+```python
+import pyarrow as pa
+import pyarrow.bigquery as bq
+
+schema = pa.schema([("value", pa.list_(pa.int64()))])
+
+with bq.writer("gcp_project.dataset.table", schema=schema) as w:
+    for a in range(1000):
+        w.write_table(pa.Table.from_pylist([{'value': [a] * 10}]))
+```
+
+#### `pyarrow.bigquery.writer.write_batch`
+
+Context manager method to write a record batch.
+
+**Parameters:**
+
+- `batch`: `pa.RecordBatch`  
+  PyArrow record batch.
+
+```python
+import pyarrow as pa
+import pyarrow.bigquery as bq
+
+schema = pa.schema([("value", pa.list_(pa.int64()))])
+
+with bq.writer("gcp_project.dataset.table", schema=schema) as w:
+    for a in range(1000):
+        w.write_batch(pa.RecordBatch.from_pylist([{'value': [1] * 10}]))
+```
+
+### `pyarrow.bigquery.read_table`
+
+**Parameters:**
+
+- `source`: `str`  
+  BigQuery table location.
+
+- `project`: `str`, *default* `None`  
+  BigQuery execution project, also the billing project. If not provided, it will be extracted from `source`.
+
+- `columns`: `str`, *default* `None`  
+  Columns to download. When not provided, all available columns will be downloaded.
+
+- `row_restrictions`: `str`, *default* `None`  
+  Row level filtering executed on the BigQuery side. More in [BigQuery documentation](https://cloud.google.com/bigquery/docs/reference/storage/rpc/google.cloud.bigquery.storage.v1beta1).
+
+- `worker_type`: `threading.Thread | multiprocessing.Process`, *default* `threading.Thread`  
+  Worker backend for fetching data.
+
+- `worker_count`: `int`, *default* `os.cpu_count()`  
+  Number of threads or processes to use for fetching data from BigQuery.
+
+- `batch_size`: `int`, *default* `100`  
+  Batch size used for fetching. Table will be automatically split to this value.
+
+### `pyarrow.bigquery.reader`
+
+**Parameters:**
+
+- `source`: `str`  
+  BigQuery table location.
+
+- `project`: `str`, *default* `None`  
+  BigQuery execution project, also the billing project. If not provided, it will be extracted from `source`.
+
+- `columns`: `str`, *default* `None`  
+  Columns to download. When not provided, all available columns will be downloaded.
+
+- `row_restrictions`: `str`, *default* `None`  
+  Row level filtering executed on the BigQuery side. More in [BigQuery documentation](https://cloud.google.com/bigquery/docs/reference/storage/rpc/google.cloud.bigquery.storage.v1beta1).
+
+- `worker_type`: `threading.Thread | multiprocessing.Process`, *default* `threading.Thread`  
+  Worker backend for fetching data.
+
+- `worker_count`: `int`, *default* `os.cpu_count()`  
+  Number of threads or processes to use for fetching data from BigQuery.
+
+- `batch_size`: `int`, *default* `100`  
+  Batch size used for fetching. Table will be automatically split to this value.
+
+```python
+import pyarrow as pa
+import pyarrow.bigquery as bq
+
+parts = []
+for part in bq.reader("gcp_project.dataset.table"):
+    parts.append(part)
+
+table = pa.concat_tables(parts)
+```

pyarrow_bigquery-0.1.0.dist-info/RECORD

@@ -0,0 +1,12 @@
+pyarrow/bigquery/__init__.py,sha256=BZBKPSDDcVkf3s6d91hx4OHoeJm927BN4T33rY_xMqs,118
+pyarrow/bigquery/read.py,sha256=MgKhvkeKqq0ha6XSN-Fe59EgPdGOHVLZ-PTCWpmqMwA,6524
+pyarrow/bigquery/some_itertools.py,sha256=XCbKojLfSYqV6TDAtCV6uQXxxuuU9Jlk02fjxCv8UVI,394
+pyarrow/bigquery/write/__init__.py,sha256=u3gntAxTAMRCDj6ogTyp3OicnAjsgK7LKMgVIPyQBOs,7098
+pyarrow/bigquery/write/pa_to_bq.py,sha256=NADdA7bywjpMvS7kCLnTCOdDyTzWiQQm6wFFeqYm5dw,1309
+pyarrow/bigquery/write/pa_to_pb.py,sha256=6UUWpGudPs8aVEhdrsYTg6OJ9PdI8UUxFi55WM2nfl0,4427
+pyarrow/bigquery/write/type_mapping.py,sha256=O8ellbgF6_yEET4E9gQvtSKz4Q0750BMwRctTli8aXI,1138
+pyarrow/bigquery/write/upload.py,sha256=uVmKNUAtoFneACAgFWanefNrJNjSE4D6kQdYw4B8pDk,912
+pyarrow_bigquery-0.1.0.dist-info/METADATA,sha256=7kCiT33XF3QLKbNWcD8TzGn_8burhpBy8AcO0ju3jDI,8204
+pyarrow_bigquery-0.1.0.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+pyarrow_bigquery-0.1.0.dist-info/top_level.txt,sha256=fPLFY23J70iLX3TKZtbNM2WS9DlDdIA5d9WX0dloJVY,8
+pyarrow_bigquery-0.1.0.dist-info/RECORD,,

pyarrow_bigquery-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: bdist_wheel (0.43.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

pyarrow_bigquery-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+pyarrow