From 1e0ae7a86908900df63de2bed73cca86f6887ff9 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sat, 26 Sep 2026 21:18:38 -0700 Subject: [PATCH 1/2] Fix chunked reconstruction losing windows over nanosecond times Polars reads a pd.Timestamp or datetime literal as microseconds. The chunked round-trip builds its window predicates from such literals, so against a nanosecond time column (pandas' and xarray's default): - a single-step window, sent as a value list, matched nothing, and the spill path silently returned NaN for it (e.g. the last window of 3 steps in chunks of 2); - on a Polars frame, is_in (a join internally) failed on mismatched key types; - range bounds were truncated to microseconds, so windows over times closer than that returned rows they had not asked for. Window values are now Polars literals of the column's own type, built from their integer nanosecond counts, and expressed in the column's zone when it has one. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01AMTHTEAoyKzUJLvFmg5G6t --- tests/test_lazy_roundtrip.py | 22 +++++++++++++ xarray_sql/lazyscan.py | 64 ++++++++++++++++++++++++------------ 2 files changed, 65 insertions(+), 21 deletions(-) diff --git a/tests/test_lazy_roundtrip.py b/tests/test_lazy_roundtrip.py index d8148a5..e1efabe 100644 --- a/tests/test_lazy_roundtrip.py +++ b/tests/test_lazy_roundtrip.py @@ -419,3 +419,25 @@ def test_timedelta_coordinates_returned_as_text(chunks): ) xr.testing.assert_identical(out.compute(), ds) + + +@pytest.mark.parametrize("chunk", [1, 2, 3]) +@pytest.mark.parametrize("via", ["polars", "spill"]) +def test_chunked_windows_over_nanosecond_times(chunk, via): + # Four times a nanosecond apart, in windows that leave a partial or + # single-step last window (sent as a value list, not a range). + pl = pytest.importorskip("polars") + times = pd.to_datetime(["2021-01-01"] * 4) + pd.to_timedelta( + [0, 1, 2, 3], unit="ns" + ) + ds = xr.Dataset( + {"v": ("time", [1.0, 2.0, 3.0, 4.0])}, coords={"time": times} + ) + table = pa.table({"time": pa.array(times.values), "v": ds.v.values}) + result = pl.from_arrow(table).lazy() if via == "polars" else table + + out = xql.to_dataset( + result, template=ds, chunks={"time": chunk}, spill=via == "spill" + ) + + xr.testing.assert_identical(out.compute(), ds) diff --git a/xarray_sql/lazyscan.py b/xarray_sql/lazyscan.py index 91de0a2..3d6514f 100644 --- a/xarray_sql/lazyscan.py +++ b/xarray_sql/lazyscan.py @@ -75,20 +75,6 @@ def _plain(value: Any) -> Any: return value -def _zoned(value: Any, zone: str | None) -> Any: - """A plain literal comparable with a column in time zone *zone*. - - Window values come from numpy, which has no time zones: a - ``datetime64`` is a UTC instant. Engines refuse to compare a naive - literal with a zone-aware column (or compare it as local time), so - the literal is labeled UTC and expressed in the column's zone. - """ - plain = _plain(value) - if zone and isinstance(plain, pd.Timestamp) and plain.tzinfo is None: - return plain.tz_localize("UTC").tz_convert(zone) - return plain - - class LazyResultHandle(Protocol): """A re-executable query result (see module docstring).""" @@ -285,6 +271,41 @@ def run() -> None: self._run(run) +def _polars_times(values: np.ndarray, dtype: Any) -> Any: + """Window times as a Polars Series of the column's own type. + + Built from the numpy values, so nanoseconds survive: Polars reads a + ``pd.Timestamp`` or ``datetime`` literal as microseconds, which never + equals a value in a nanosecond column (``is_in`` then matches + nothing, or fails as a join on mismatched key types). A naive window + time is a UTC instant; it is expressed in the column's zone, if any. + """ + import polars as pl + + series = pl.Series(np.asarray(values, dtype="datetime64[ns]")) + if dtype.time_zone: + series = series.dt.replace_time_zone("UTC").dt.convert_time_zone( + dtype.time_zone + ) + return series.cast(dtype) + + +def _polars_value(value: Any, dtype: Any) -> Any: + """One window bound as a Polars literal comparable with *dtype*.""" + import polars as pl + + if isinstance(dtype, pl.Datetime) and isinstance(value, np.datetime64): + # From the integer count, which is exact; a Python datetime is not. + nanos = int(value.astype("datetime64[ns]").astype("int64")) + literal = pl.lit(nanos, dtype=pl.Int64).cast(pl.Datetime("ns")) + if dtype.time_zone: + literal = literal.dt.replace_time_zone("UTC").dt.convert_time_zone( + dtype.time_zone + ) + return literal.cast(dtype) + return _plain(value) + + class PolarsHandle: """Handle over a ``polars.LazyFrame``. @@ -317,10 +338,12 @@ def fetch( schema = self._lf.collect_schema() exprs = [] for dim, (kind, a, b) in specs.items(): - zone = getattr(schema.get(dim), "time_zone", None) + dtype = schema.get(dim) if kind == "range": exprs.append( - pl.col(dim).is_between(_zoned(a, zone), _zoned(b, zone)) + pl.col(dim).is_between( + _polars_value(a, dtype), _polars_value(b, dtype) + ) ) elif getattr(a, "dtype", None) is not None and a.dtype.kind == "f": # Upstream Polars translates float ``is_in`` literals @@ -331,14 +354,13 @@ def fetch( # of values. exprs.append( pl.any_horizontal( - [ - pl.col(dim).is_between(*(_zoned(v, zone),) * 2) - for v in a - ] + [pl.col(dim).is_between(*(_plain(v),) * 2) for v in a] ) ) + elif isinstance(dtype, pl.Datetime): + exprs.append(pl.col(dim).is_in(_polars_times(a, dtype))) else: - exprs.append(pl.col(dim).is_in([_zoned(v, zone) for v in a])) + exprs.append(pl.col(dim).is_in([_plain(v) for v in a])) lf = self._lf.filter(*exprs) if exprs else self._lf out = _collect_streaming(lf.select([pl.col(n) for n in columns])) return cast(list[pa.RecordBatch], out.to_arrow().to_batches()) From 9da28061e03cda87fd3a1972b6e6904126a5b0e3 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Tue, 29 Sep 2026 09:53:27 -0700 Subject: [PATCH 2/2] Keep nanosecond durations exact in Polars window filters too Durations had the same problem as times: a `timedelta64` window value became a `pd.Timedelta`, which Polars reads as microseconds, so a single-step window over a `timedelta64[ns]` coordinate matched nothing (or failed as a join). Duration columns now get literals of their own type from integer nanosecond counts, like datetime columns. The test now covers durations and zone-aware times as well as naive ones, so a single-step window over a zone-aware column is exercised. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01AMTHTEAoyKzUJLvFmg5G6t --- tests/test_lazy_roundtrip.py | 21 +++++++------- xarray_sql/lazyscan.py | 55 ++++++++++++++++++++++-------------- 2 files changed, 45 insertions(+), 31 deletions(-) diff --git a/tests/test_lazy_roundtrip.py b/tests/test_lazy_roundtrip.py index e1efabe..e2a7db5 100644 --- a/tests/test_lazy_roundtrip.py +++ b/tests/test_lazy_roundtrip.py @@ -423,21 +423,22 @@ def test_timedelta_coordinates_returned_as_text(chunks): @pytest.mark.parametrize("chunk", [1, 2, 3]) @pytest.mark.parametrize("via", ["polars", "spill"]) -def test_chunked_windows_over_nanosecond_times(chunk, via): - # Four times a nanosecond apart, in windows that leave a partial or +@pytest.mark.parametrize("kind", ["naive", "zoned", "duration"]) +def test_chunked_windows_over_nanosecond_values(chunk, via, kind): + # Four values a nanosecond apart, in windows that leave a partial or # single-step last window (sent as a value list, not a range). pl = pytest.importorskip("polars") - times = pd.to_datetime(["2021-01-01"] * 4) + pd.to_timedelta( - [0, 1, 2, 3], unit="ns" - ) - ds = xr.Dataset( - {"v": ("time", [1.0, 2.0, 3.0, 4.0])}, coords={"time": times} - ) - table = pa.table({"time": pa.array(times.values), "v": ds.v.values}) + steps = pd.to_timedelta([0, 1, 2, 3], unit="ns") + coord = steps if kind == "duration" else pd.Timestamp("2021-01-01") + steps + ds = xr.Dataset({"v": ("t", [1.0, 2.0, 3.0, 4.0])}, coords={"t": coord}) + column = pa.array(coord.values) + if kind == "zoned": + column = column.cast(pa.timestamp("ns", tz="Asia/Tokyo")) + table = pa.table({"t": column, "v": ds.v.values}) result = pl.from_arrow(table).lazy() if via == "polars" else table out = xql.to_dataset( - result, template=ds, chunks={"time": chunk}, spill=via == "spill" + result, template=ds, chunks={"t": chunk}, spill=via == "spill" ) xr.testing.assert_identical(out.compute(), ds) diff --git a/xarray_sql/lazyscan.py b/xarray_sql/lazyscan.py index 3d6514f..e8f36fc 100644 --- a/xarray_sql/lazyscan.py +++ b/xarray_sql/lazyscan.py @@ -271,38 +271,51 @@ def run() -> None: self._run(run) +def _nanosecond(dtype: Any) -> str | None: + """The numpy ``[ns]`` type matching a Polars time column, else ``None``.""" + import polars as pl + + if isinstance(dtype, pl.Datetime): + return "datetime64[ns]" + if isinstance(dtype, pl.Duration): + return "timedelta64[ns]" + return None + + +def _in_zone(expr: Any, dtype: Any) -> Any: + """*expr*, a naive UTC time, in *dtype*'s zone if it has one.""" + zone = getattr(dtype, "time_zone", None) + if zone: + return expr.dt.replace_time_zone("UTC").dt.convert_time_zone(zone) + return expr + + def _polars_times(values: np.ndarray, dtype: Any) -> Any: - """Window times as a Polars Series of the column's own type. + """Window times or durations as a Polars Series of the column's type. Built from the numpy values, so nanoseconds survive: Polars reads a - ``pd.Timestamp`` or ``datetime`` literal as microseconds, which never - equals a value in a nanosecond column (``is_in`` then matches - nothing, or fails as a join on mismatched key types). A naive window - time is a UTC instant; it is expressed in the column's zone, if any. + ``pd.Timestamp``, ``pd.Timedelta``, or ``datetime`` literal as + microseconds, which never equals a value in a nanosecond column + (``is_in`` then matches nothing, or fails as a join on mismatched + key types). A naive window time is a UTC instant; it is expressed in + the column's zone, if any. """ import polars as pl - series = pl.Series(np.asarray(values, dtype="datetime64[ns]")) - if dtype.time_zone: - series = series.dt.replace_time_zone("UTC").dt.convert_time_zone( - dtype.time_zone - ) - return series.cast(dtype) + series = pl.Series(np.asarray(values, dtype=_nanosecond(dtype))) + return _in_zone(series, dtype).cast(dtype) def _polars_value(value: Any, dtype: Any) -> Any: """One window bound as a Polars literal comparable with *dtype*.""" import polars as pl - if isinstance(dtype, pl.Datetime) and isinstance(value, np.datetime64): - # From the integer count, which is exact; a Python datetime is not. - nanos = int(value.astype("datetime64[ns]").astype("int64")) - literal = pl.lit(nanos, dtype=pl.Int64).cast(pl.Datetime("ns")) - if dtype.time_zone: - literal = literal.dt.replace_time_zone("UTC").dt.convert_time_zone( - dtype.time_zone - ) - return literal.cast(dtype) + unit = _nanosecond(dtype) + if unit and isinstance(value, (np.datetime64, np.timedelta64)): + # From the integer count, which is exact; a Python object is not. + nanos = int(value.astype(unit).astype("int64")) + literal = pl.lit(nanos, dtype=pl.Int64).cast(type(dtype)("ns")) + return _in_zone(literal, dtype).cast(dtype) return _plain(value) @@ -357,7 +370,7 @@ def fetch( [pl.col(dim).is_between(*(_plain(v),) * 2) for v in a] ) ) - elif isinstance(dtype, pl.Datetime): + elif _nanosecond(dtype): exprs.append(pl.col(dim).is_in(_polars_times(a, dtype))) else: exprs.append(pl.col(dim).is_in([_plain(v) for v in a]))