polars-bigquery 0.2.0__py3-none-any.whl

polars_bigquery/__init__.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+from polars_bigquery._read_bigquery import (
+    Client,
+)
+from polars_bigquery.core.version import __version__
+
+__all__ = [
+    "Client",
+    "__version__",
+]

polars_bigquery/_read_bigquery.py

@@ -0,0 +1,146 @@
+from __future__ import annotations
+
+from collections.abc import Iterator
+from typing import Any
+
+import arrow_bigquery
+import arrow_bigquery.api.resources
+import polars as pl
+import polars.io.plugins
+
+import polars_bigquery.core.schema
+import polars_bigquery.core.version
+from polars_bigquery.core import bigquery_rest, compiler
+
+
+def _get_user_agent(user_agent: str | None) -> str:
+    ua = f"polars-bigquery/{polars_bigquery.core.version.__version__}"
+
+    if user_agent:
+        return f"{ua} {user_agent}"
+    else:
+        return ua
+
+
+class Client:
+    """Client for reading data from BigQuery into Polars.
+
+    Wraps an arrow_bigquery Client to keep connections open and reuse credentials across operations.
+    """
+
+    def __init__(
+        self,
+        *,
+        quota_project_id: str,
+        credentials_provider: pl.CredentialProviderGCP | None = None,
+        user_agent: str | None = None,
+    ) -> None:
+        if credentials_provider is None:
+            credentials_provider = pl.CredentialProviderGCP(
+                quota_project_id=quota_project_id
+            )
+        self._credentials_provider = credentials_provider
+        self._user_agent = _get_user_agent(user_agent)
+        self._quota_project_id = quota_project_id
+        self._rest_client = bigquery_rest.BigQueryRestClient(
+            quota_project_id=quota_project_id,
+            credentials_provider=credentials_provider,
+            user_agent=self._user_agent,
+        )
+        self._arrow_client = arrow_bigquery.Client(
+            quota_project_id=quota_project_id,
+            credentials_provider=credentials_provider,
+            user_agent=self._user_agent,
+        )
+
+    def close(self) -> None:
+        """Close underlying REST session resources."""
+        self._rest_client.close()
+
+    def __enter__(self) -> Client:  # noqa: PYI034
+        return self
+
+    def __exit__(self, *exc_info: object) -> None:
+        self.close()
+
+    @property
+    def credentials_provider(self) -> pl.CredentialProviderGCP:
+        return self._credentials_provider
+
+    @property
+    def quota_project_id(self) -> str:
+        return self._quota_project_id
+
+    def read_table(
+        self,
+        table: Any,
+        *,
+        maintain_order: bool = False,
+    ) -> pl.DataFrame:
+        table_ref = arrow_bigquery.api.resources.parse_table_id(table)
+        arrow_stream_exporter = self._arrow_client.read_table(
+            table_ref,
+            maintain_order=maintain_order,
+        )
+        return pl.DataFrame(arrow_stream_exporter)
+
+    def read_query(
+        self,
+        query: str,
+        *,
+        maintain_order: bool = False,
+    ) -> pl.DataFrame:
+        table = self._rest_client.run_query(query)
+        table_ref = arrow_bigquery.api.resources.parse_table_id(table)
+        arrow_stream_exporter = self._arrow_client.read_table(
+            table_ref,
+            maintain_order=maintain_order,
+        )
+        return pl.DataFrame(arrow_stream_exporter)
+
+    def scan_table(
+        self,
+        table: Any,
+    ) -> pl.LazyFrame:
+        table_ref = arrow_bigquery.api.resources.parse_table_id(table)
+        table_metadata = self._rest_client.get_table_metadata(table_ref)
+        schema = polars_bigquery.core.schema.extract_polars_schema(table_metadata)
+
+        def source_generator(
+            with_columns: list[str] | None,
+            predicate: pl.Expr | None,
+            n_rows: int | None,
+            batch_size: int | None,
+        ) -> Iterator[pl.DataFrame]:
+            columns = with_columns if with_columns is not None else list(schema.keys())
+            predicate_cols = (
+                predicate.meta.root_names() if predicate is not None else []
+            )
+            selected_fields = list(dict.fromkeys([*columns, *predicate_cols]))
+
+            arrow_stream_exporter = self._arrow_client.read_table(
+                table_ref,
+                maintain_order=False,
+                selected_fields=selected_fields,
+                row_restriction=compiler.predicate_to_row_restriction(
+                    predicate=predicate
+                )
+                if predicate is not None
+                else "",
+            )
+            lazyframe = pl.scan_arrow_c_stream(arrow_stream_exporter)
+
+            if predicate is not None:
+                lazyframe = lazyframe.filter(predicate)
+
+            lazyframe = lazyframe.select(columns)
+
+            if n_rows is not None:
+                lazyframe = lazyframe.limit(n_rows)
+
+            return lazyframe.collect_batches(chunk_size=batch_size)
+
+        return polars.io.plugins.register_io_source(
+            io_source=source_generator,
+            schema=schema,
+        )

polars_bigquery/core/bigquery_rest.py

@@ -0,0 +1,300 @@
+"""Utilities for running a query in BigQuery and getting the results as a table.
+
+TODO(tswast): Implement this in Rust using jobs.query and the
+JOB_CREATION_OPTIONAL parameter to improve latency in small query results.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import time
+import uuid
+from typing import Any
+
+import arrow_bigquery.api.resources
+import polars as pl
+import requests
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
+
+import polars_bigquery.exceptions
+
+_BIGQUERY_ENDPOINT = "https://bigquery.googleapis.com/bigquery/v2"
+# Keep timeout at least as long as the server-side timeout defined at
+# https://github.com/googleapis/googleapis/blob/4bcbf04e688ffddb8ec1a20a2349df32b9fcbd9c/google/cloud/bigquery/v2/bigquery_grpc_service_config.json#L170-L181
+_CONNECT_TIMEOUT = 60.0
+_READ_TIMEOUT = 250.0
+DEFAULT_TIMEOUT: tuple[float, float] = (_CONNECT_TIMEOUT, _READ_TIMEOUT)
+MAX_POLL_SECONDS: float = 21600.0
+
+
+def create_resilient_session() -> requests.Session:
+    """Create a requests.Session with connection pooling and exponential backoff."""
+    session = requests.Session()
+    retry_strategy = Retry(
+        total=5,
+        backoff_factor=0.5,
+        status_forcelist=(429, 500, 502, 503, 504),
+        allowed_methods=frozenset(["GET", "POST"]),
+        respect_retry_after_header=True,
+    )
+    adapter = HTTPAdapter(
+        max_retries=retry_strategy,
+        pool_connections=10,
+        pool_maxsize=20,
+    )
+    session.mount("https://", adapter)
+    session.mount("http://", adapter)
+    return session
+
+
+def _raise_for_bigquery_error(response: requests.Response) -> None:
+    """Raise BigQueryError with structured message on HTTP error."""
+    try:
+        response.raise_for_status()
+    except requests.exceptions.HTTPError as exc:
+        message = str(exc)
+        with contextlib.suppress(ValueError, KeyError, TypeError):
+            err_json = response.json()
+            if isinstance(err_json, dict) and "error" in err_json:
+                err_obj = err_json["error"]
+                if isinstance(err_obj, dict) and "message" in err_obj:
+                    message = err_obj["message"]
+        raise polars_bigquery.exceptions.BigQueryError(message) from exc
+
+
+def _get_table_metadata_url(
+    *, table_ref: arrow_bigquery.api.resources.BigQueryTableId
+) -> str:
+    return (
+        f"{_BIGQUERY_ENDPOINT}/projects/{table_ref.project_id}"
+        f"/datasets/{table_ref.dataset_id}/tables/{table_ref.table_id}"
+        "?fields=schema,timePartitioning"
+    )
+
+
+def _get_jobs_insert_url(quota_project_id: str) -> str:
+    return f"{_BIGQUERY_ENDPOINT}/projects/{quota_project_id}/jobs"
+
+
+def _get_jobs_insert_body(query: str, quota_project_id: str) -> dict:
+    return {
+        "jobReference": {
+            "projectId": quota_project_id,
+            "jobId": f"polars_bq_{uuid.uuid4().hex}",
+        },
+        "configuration": {
+            "query": {
+                "query": query,
+                "useLegacySql": False,
+            }
+        },
+    }
+
+
+def _get_jobs_get_url(
+    job_id: str, quota_project_id: str, location: str | None = None
+) -> str:
+    url = f"{_BIGQUERY_ENDPOINT}/projects/{quota_project_id}/jobs/{job_id}"
+    if location:
+        url += f"?location={location}"
+    return url
+
+
+def _get_query_results_url(
+    job_id: str, quota_project_id: str, location: str | None = None
+) -> str:
+    # Set maxResults=0 and timeoutMs=10000 to use this endpoint to wait for job completion.
+    url = (
+        f"{_BIGQUERY_ENDPOINT}/projects/{quota_project_id}/queries/{job_id}/results"
+        "?maxResults=0&timeoutMs=10000"
+    )
+    if location:
+        url += f"&location={location}"
+    return url
+
+
+def _wait_for_job(
+    job_id: str,
+    quota_project_id: str,
+    credentials_provider: pl.CredentialProviderGCP,
+    user_agent: str,
+    *,
+    location: str | None = None,
+    session: requests.Session | None = None,
+    timeout: tuple[float, float] = DEFAULT_TIMEOUT,
+    max_poll_seconds: float = MAX_POLL_SECONDS,
+) -> dict:
+    http = session if session is not None else requests
+    start_time = time.monotonic()
+    while True:
+        if time.monotonic() - start_time > max_poll_seconds:
+            raise polars_bigquery.exceptions.BigQueryError(
+                f"Timed out waiting for BigQuery job {job_id} after {max_poll_seconds}s"
+            )
+        headers = _get_headers(
+            quota_project_id=quota_project_id,
+            credentials_provider=credentials_provider,
+            user_agent=user_agent,
+        )
+        response = http.get(
+            _get_jobs_get_url(job_id, quota_project_id, location=location),
+            headers=headers,
+            timeout=timeout,
+        )
+        _raise_for_bigquery_error(response)
+        job = response.json()
+
+        if job["status"]["state"] == "DONE":
+            if "errorResult" in job["status"]:
+                raise polars_bigquery.exceptions.BigQueryError(
+                    job["status"]["errorResult"]["message"]
+                )
+            return job
+
+        # jobs.getQueryResults waits about 10s or until the query finishes.
+        poll_start = time.monotonic()
+        poll_response = http.get(
+            _get_query_results_url(job_id, quota_project_id, location=location),
+            headers=headers,
+            timeout=timeout,
+        )
+        _raise_for_bigquery_error(poll_response)
+        # Avoid busy-spinning if the polling endpoint returns without blocking
+        elapsed = time.monotonic() - poll_start
+        if elapsed < 1.0:
+            time.sleep(1.0 - elapsed)
+
+
+def _get_headers(
+    *,
+    quota_project_id: str,
+    credentials_provider: pl.CredentialProviderGCP,
+    user_agent: str,
+) -> dict:
+    token_data, _ = credentials_provider()
+    token = token_data["bearer_token"]
+    return {
+        "Authorization": f"Bearer {token}",
+        "User-Agent": user_agent,
+        "Content-Type": "application/json",
+        "x-goog-user-project": quota_project_id,
+    }
+
+
+def run_query(
+    query: str,
+    *,
+    quota_project_id: str,
+    credentials_provider: pl.CredentialProviderGCP,
+    user_agent: str,
+    session: requests.Session | None = None,
+    timeout: tuple[float, float] = DEFAULT_TIMEOUT,
+) -> str:
+    """Run a query and return the destination table from the job resource."""
+    http = session if session is not None else requests
+    headers = _get_headers(
+        quota_project_id=quota_project_id,
+        credentials_provider=credentials_provider,
+        user_agent=user_agent,
+    )
+
+    # 1. Insert the job
+    insert_url = _get_jobs_insert_url(quota_project_id)
+    body = _get_jobs_insert_body(query, quota_project_id)
+    response = http.post(insert_url, headers=headers, json=body, timeout=timeout)
+    _raise_for_bigquery_error(response)
+    job_resource = response.json()
+    job_ref = job_resource.get("jobReference", {})
+    job_id = job_ref["jobId"]
+    location = job_ref.get("location")
+
+    # 2. Wait for the job to complete (refreshing credentials on each poll)
+    job = _wait_for_job(
+        job_id,
+        quota_project_id,
+        credentials_provider,
+        user_agent,
+        location=location,
+        session=session,
+        timeout=timeout,
+    )
+
+    # 3. Return the destination table ID
+    dest = job["configuration"]["query"]["destinationTable"]
+    return f"{dest['projectId']}.{dest['datasetId']}.{dest['tableId']}"
+
+
+def get_table_metadata(
+    table_ref: arrow_bigquery.api.resources.BigQueryTableId,
+    *,
+    quota_project_id: str,
+    credentials_provider: pl.CredentialProviderGCP,
+    user_agent: str,
+    session: requests.Session | None = None,
+    timeout: tuple[float, float] = DEFAULT_TIMEOUT,
+) -> dict:
+    http = session if session is not None else requests
+    headers = _get_headers(
+        quota_project_id=quota_project_id,
+        credentials_provider=credentials_provider,
+        user_agent=user_agent,
+    )
+    table_metadata_url = _get_table_metadata_url(table_ref=table_ref)
+    response = http.get(table_metadata_url, headers=headers, timeout=timeout)
+    _raise_for_bigquery_error(response)
+    return response.json()
+
+
+class BigQueryRestClient:
+    """Stateful BigQuery REST client with connection pooling and metadata caching."""
+
+    def __init__(
+        self,
+        *,
+        quota_project_id: str,
+        credentials_provider: pl.CredentialProviderGCP,
+        user_agent: str,
+        session: requests.Session | None = None,
+    ) -> None:
+        self._quota_project_id = quota_project_id
+        self._credentials_provider = credentials_provider
+        self._user_agent = user_agent
+        self._owns_session = session is None
+        self._session = session if session is not None else create_resilient_session()
+
+    @property
+    def session(self) -> requests.Session:
+        return self._session
+
+    def close(self) -> None:
+        """Close the underlying HTTP session if owned by this client."""
+        if self._owns_session:
+            self._session.close()
+
+    def __enter__(self) -> BigQueryRestClient:  # noqa: PYI034
+        return self
+
+    def __exit__(self, *exc_info: object) -> None:
+        self.close()
+
+    def run_query(self, query: str) -> str:
+        return run_query(
+            query,
+            quota_project_id=self._quota_project_id,
+            credentials_provider=self._credentials_provider,
+            user_agent=self._user_agent,
+            session=self._session,
+        )
+
+    def get_table_metadata(
+        self,
+        table_ref: arrow_bigquery.api.resources.BigQueryTableId,
+    ) -> dict[str, Any]:
+        return get_table_metadata(
+            table_ref,
+            quota_project_id=self._quota_project_id,
+            credentials_provider=self._credentials_provider,
+            user_agent=self._user_agent,
+            session=self._session,
+        )

polars_bigquery/core/compiler/__init__.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+import io
+import json
+from typing import Any
+
+import polars as pl
+
+from polars_bigquery.core.compiler import ir, parser, rewriter, sql
+from polars_bigquery.core.compiler.parser import json_to_ir
+from polars_bigquery.core.compiler.rewriter import rewrite_ir
+from polars_bigquery.core.compiler.sql import ir_to_sql
+
+__all__ = [
+    "ir",
+    "ir_to_sql",
+    "json_to_ir",
+    "parser",
+    "predicate_to_row_restriction",
+    "rewrite_ir",
+    "rewriter",
+    "sql",
+]
+
+
+def _json_expr_to_row_restriction(expr_json: dict[str, Any]) -> str | None:
+    """Compile a Polars expression JSON dict into a BigQuery SQL row restriction.
+
+    Returns None if unknown operators are found and can't guarantee a superset of rows.
+    """
+    ir_tree = json_to_ir(expr_json)
+    ir_tree = rewrite_ir(ir_tree)
+    return ir_to_sql(ir_tree)
+
+
+def predicate_to_row_restriction(predicate: pl.Expr) -> str:
+    """Orchestrate compilation of a Polars expression to a BigQuery row restriction SQL string."""
+    predicate_json_file = io.BytesIO()
+    predicate.meta.serialize(predicate_json_file, format="json")
+    predicate_json_file.seek(0)
+    try:
+        predicate_json = json.load(predicate_json_file)
+        row_restriction = _json_expr_to_row_restriction(predicate_json)
+    except (RecursionError, json.JSONDecodeError):
+        row_restriction = None
+    return row_restriction if row_restriction is not None else ""

polars_bigquery/core/compiler/ir/__init__.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+from polars_bigquery.core.compiler.ir import (
+    base,
+    boolean,
+    comparison,
+    list_,
+    numeric,
+    string,
+    temporal,
+)
+from polars_bigquery.core.compiler.ir.base import (
+    BinaryExpr,
+    Column,
+    Expr,
+    Literal,
+    NullLiteral,
+    UnaryExpr,
+    Unsupported,
+)
+from polars_bigquery.core.compiler.ir.boolean import (
+    And,
+    BoolLiteral,
+    IsNotNull,
+    IsNull,
+    Not,
+    Or,
+)
+from polars_bigquery.core.compiler.ir.comparison import (
+    Eq,
+    Gt,
+    GtEq,
+    Lt,
+    LtEq,
+    NotEq,
+)
+from polars_bigquery.core.compiler.ir.list_ import IsIn, ListLiteral
+from polars_bigquery.core.compiler.ir.numeric import (
+    FloatLiteral,
+    IntLiteral,
+    IsFinite,
+    IsInfinite,
+    IsNan,
+    IsNotNan,
+)
+from polars_bigquery.core.compiler.ir.string import (
+    Contains,
+    EndsWith,
+    Lowercase,
+    StartsWith,
+    StringLiteral,
+    Uppercase,
+)
+from polars_bigquery.core.compiler.ir.temporal import (
+    DateLiteral,
+    TimestampLiteral,
+    TimestampUnit,
+)
+
+__all__ = [
+    "And",
+    "BinaryExpr",
+    "BoolLiteral",
+    "Column",
+    "Contains",
+    "DateLiteral",
+    "EndsWith",
+    "Eq",
+    "Expr",
+    "FloatLiteral",
+    "Gt",
+    "GtEq",
+    "IntLiteral",
+    "IsFinite",
+    "IsIn",
+    "IsInfinite",
+    "IsNan",
+    "IsNotNan",
+    "IsNotNull",
+    "IsNull",
+    "ListLiteral",
+    "Literal",
+    "Lowercase",
+    "Lt",
+    "LtEq",
+    "Not",
+    "NotEq",
+    "NullLiteral",
+    "Or",
+    "StartsWith",
+    "StringLiteral",
+    "TimestampLiteral",
+    "TimestampUnit",
+    "UnaryExpr",
+    "Unsupported",
+    "Uppercase",
+    "base",
+    "boolean",
+    "comparison",
+    "list_",
+    "numeric",
+    "string",
+    "temporal",
+]

polars_bigquery/core/compiler/ir/base.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+import abc
+import dataclasses
+
+
+@dataclasses.dataclass(frozen=True)
+class Expr(abc.ABC):
+    """Base class for all predicate Intermediate Representation (IR) nodes."""
+
+    @abc.abstractmethod
+    def children(self) -> tuple[Expr, ...]:
+        """Return the child expressions of this node in order."""
+
+
+@dataclasses.dataclass(frozen=True)
+class UnaryExpr(Expr):
+    """Base class for unary operation IR nodes."""
+
+    expr: Expr
+
+    def children(self) -> tuple[Expr, ...]:
+        return (self.expr,)
+
+
+@dataclasses.dataclass(frozen=True)
+class BinaryExpr(Expr):
+    """Base class for binary operation IR nodes."""
+
+    left: Expr
+    right: Expr
+
+    def children(self) -> tuple[Expr, ...]:
+        return (self.left, self.right)
+
+
+@dataclasses.dataclass(frozen=True)
+class Unsupported(Expr):
+    """IR node representing an unsupported or unknown expression."""
+
+    operands: tuple[Expr, ...] = ()
+
+    def children(self) -> tuple[Expr, ...]:
+        return self.operands
+
+
+@dataclasses.dataclass(frozen=True)
+class Column(Expr):
+    """IR node representing a column reference."""
+
+    name: str
+
+    def children(self) -> tuple[Expr, ...]:
+        return ()
+
+
+@dataclasses.dataclass(frozen=True)
+class Literal(Expr):
+    """Base class for literal IR nodes."""
+
+    def children(self) -> tuple[Expr, ...]:
+        return ()
+
+
+@dataclasses.dataclass(frozen=True)
+class NullLiteral(Literal):
+    """IR node representing a NULL literal."""