diff --git a/ibis/backends/bigquery/__init__.py b/ibis/backends/bigquery/__init__.py index 1cc1359147bc..e291bdbd83b5 100644 --- a/ibis/backends/bigquery/__init__.py +++ b/ibis/backends/bigquery/__init__.py @@ -24,6 +24,7 @@ import ibis import ibis.backends.sql.compilers as sc import ibis.common.exceptions as com +import ibis.expr.datatypes as dt import ibis.expr.operations as ops import ibis.expr.schema as sch import ibis.expr.types as ir @@ -167,6 +168,17 @@ def _postprocess_arrow( return table_or_batch.rename_columns(names) +def _microsecond_timestamp(dtype: dt.DataType) -> dt.DataType: + """Give a timestamp without a scale BigQuery's microsecond precision. + + BigQuery stores DATETIME and TIMESTAMP values in microseconds, but most + ibis timestamps carry no scale, and polars maps those to nanoseconds. + """ + if dtype.is_timestamp() and dtype.scale is None: + return dtype.copy(scale=6) + return dtype + + class Backend( SupportsTempTables, SQLBackend, CanCreateDatabase, DirectPyArrowExampleLoader ): @@ -995,6 +1007,34 @@ def to_pyarrow( table = _postprocess_arrow(table, list(schema.names)) return expr.__pyarrow_result__(table, schema=schema) + def to_polars( + self, + expr: ir.Expr, + /, + *, + params: Mapping[ir.Scalar, Any] | None = None, + limit: int | str | None = None, + **kwargs: Any, + ) -> pl.DataFrame: + import polars as pl + + from ibis.formats.polars import PolarsData + + table = self.to_pyarrow(expr.as_table(), params=params, limit=limit, **kwargs) + df = pl.from_arrow(table) + if isinstance(expr, ir.Table): + schema = expr.schema() + return PolarsData.convert_table( + df, + sch.Schema( + {name: _microsecond_timestamp(t) for name, t in schema.items()} + ), + ) + elif isinstance(expr, ir.Column): + return PolarsData.convert_column(df, _microsecond_timestamp(expr.type())) + else: + return PolarsData.convert_scalar(df, _microsecond_timestamp(expr.type())) + def to_pyarrow_batches( self, expr: ir.Expr, diff --git a/ibis/backends/bigquery/tests/unit/test_client.py b/ibis/backends/bigquery/tests/unit/test_client.py index a9bc04d2c31c..0c4a711ee314 100644 --- a/ibis/backends/bigquery/tests/unit/test_client.py +++ b/ibis/backends/bigquery/tests/unit/test_client.py @@ -3,7 +3,8 @@ import pytest import sqlglot as sg -from ibis.backends.bigquery import _force_quote_table, client +import ibis +from ibis.backends.bigquery import Backend, _force_quote_table, client @pytest.mark.parametrize( @@ -55,3 +56,37 @@ def test_force_quoting(bq_path_str, expected): table = _force_quote_table(table) assert table.sql("bigquery") == expected + + +def test_to_polars_keeps_microsecond_timestamps(monkeypatch): + import datetime + + pa = pytest.importorskip("pyarrow") + pl = pytest.importorskip("polars") + + value = datetime.datetime(2025, 1, 1, 12, 34, 56, 123456) + arrow = pa.table( + { + "ts": pa.array( + [value.replace(tzinfo=datetime.UTC)], pa.timestamp("us", "UTC") + ), + "naive": pa.array([value], pa.timestamp("us")), + } + ) + con = Backend() + monkeypatch.setattr(con, "to_pyarrow", lambda *_, **__: arrow) + + t = ibis.table({"ts": "timestamp('UTC')", "naive": "timestamp"}, name="t") + + df = con.to_polars(t) + assert df.schema == pl.Schema( + {"ts": pl.Datetime("us", "UTC"), "naive": pl.Datetime("us")} + ) + assert df["naive"][0] == value + + def one_column(expr, *_, **__): + return pa.table({expr.as_table().columns[0]: arrow["naive"]}) + + monkeypatch.setattr(con, "to_pyarrow", one_column) + assert con.to_polars(t.naive).dtype == pl.Datetime("us") + assert con.to_polars(t.naive.max()) == value