Repository navigation
Expand file tree
/
Copy pathlazyscan.py
More file actions
462 lines (373 loc) · 17.8 KB
/
Copy pathlazyscan.py
File metadata and controls
462 lines (373 loc) · 17.8 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
"""Re-executable engine handles behind the lazy chunked round-trip.
The lazy path of ``to_dataset(chunks=...)`` re-executes the engine's
query per accessed chunk, narrowed to that chunk's coordinate window and
columns. That requires the engine result to be *re-executable* — a
handle onto the query, not a one-shot stream of its rows. Each handle
here adapts one engine's native lazy surface to the three operations the
reconstruction needs:
* [schema][xarray_sql.lazyscan.LazyResultHandle.schema] — result column names/types, without
executing the query;
* [distinct][xarray_sql.lazyscan.LazyResultHandle.distinct] — one column's distinct values
(coordinate discovery; the caller sorts);
* [fetch][xarray_sql.lazyscan.LazyResultHandle.fetch] — the result narrowed by per-dimension
windows and projected to the requested columns, as Arrow batches.
Windows are passed as [DimSpec][xarray_sql.lazyscan.DimSpec] values instead of rendered SQL so
each engine can express them with its own *typed* expression API —
strings would re-open every literal-formatting pitfall (timestamps,
floats, quoting) per dialect.
Handles compose with the registration seam: when the wrapped query
scans a Dataset registered through xarray-sql's pushdown machinery, the
per-chunk range filter flows back through the engine into
[XarrayPushdownDataset][xarray_sql.backends.pyarrow.XarrayPushdownDataset], so each
output chunk's access reads only the source chunks it maps onto.
"""
from __future__ import annotations
import weakref
from collections.abc import Iterator
from concurrent.futures import ThreadPoolExecutor
from typing import Any, Literal, Protocol, cast
import numpy as np
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
from .df import _CAN_START_THREADS
DimSpec = tuple[Literal["range", "values"], Any, Any]
"""One dimension's window: ``("range", lo, hi)`` (inclusive bounds; the
requested coordinate positions are contiguous) or ``("values", array,
None)`` (explicit value list, for stepped/fancy indexers).
The payload stays ``Any``: the values are coordinate scalars handed to
the engine's typed expression API, which does the comparing — Python
never orders them, and a concrete union over coordinate dtypes
(timestamps, cftime, numerics, strings) would stay incomplete."""
def _collect_streaming(lf: Any) -> Any:
"""Collect a Polars LazyFrame on the streaming engine, if available.
``collect(engine="streaming")`` needs polars >= 1.25 (the test
extra pins higher); an older installed polars raises TypeError on
the unknown keyword, and the plain in-memory collect is a correct,
if less memory-frugal, fallback.
"""
try:
return lf.collect(engine="streaming")
except TypeError:
return lf.collect()
def _plain(value: Any) -> Any:
"""A plain-Python literal (numpy scalars don't travel to engines)."""
if isinstance(value, np.datetime64):
return pd.Timestamp(value)
if isinstance(value, np.timedelta64):
return pd.Timedelta(value)
if isinstance(value, np.generic):
return value.item()
return value
class LazyResultHandle(Protocol):
"""A re-executable query result (see module docstring)."""
supports_chunked: bool = True
"""Whether fetch() may be driven from consumer worker threads (the
chunked reconstruction). Handles for engines that cannot safely
re-execute under foreign threads set this False; the eager path
remains available."""
def schema(self) -> pa.Schema: ...
def distinct(self, column: str) -> np.ndarray: ...
def fetch(
self, specs: dict[str, DimSpec], columns: list[str]
) -> list[pa.RecordBatch]: ...
def spill_parquet(self, path: str) -> None: ...
# Handles may additionally offer ``stream(columns)``, yielding the
# unfiltered result's Arrow batches incrementally instead of
# materializing it first. The eager round-trip uses it to enforce
# ``max_result_bytes`` while collecting; handles without it refuse
# the budget rather than blow past it after collecting.
class DataFusionHandle:
"""Handle over a ``datafusion.DataFrame``."""
supports_chunked = True
def __init__(self, df: Any) -> None:
self._df = df
def schema(self) -> pa.Schema:
return self._df.schema()
def distinct(self, column: str) -> np.ndarray:
from datafusion import col
dim_only = self._df.select(col(f'"{column}"')).distinct()
batches = [b.to_pyarrow() for b in dim_only.execute_stream()]
if not batches:
return cast(np.ndarray, np.asarray([]))
return cast(
np.ndarray,
np.concatenate(
[b.column(0).to_numpy(zero_copy_only=False) for b in batches]
),
)
def fetch(
self, specs: dict[str, DimSpec], columns: list[str]
) -> list[pa.RecordBatch]:
from datafusion import col, literal
predicate = None
for dim, (kind, a, b) in specs.items():
c = col(f'"{dim}"')
if kind == "range":
p = (c >= literal(a)) & (c <= literal(b))
else:
# DataFusion 52.0.0 exposes no clean ``Expr.in_list``
# from Python; OR-chained equalities constant-fold
# equivalently and stay typed.
p = c == literal(a[0])
for v in a[1:]:
p = p | (c == literal(v))
predicate = p if predicate is None else predicate & p
out = self._df if predicate is None else self._df.filter(predicate)
out = out.select(*(col(f'"{n}"') for n in columns))
return [b.to_pyarrow() for b in out.execute_stream()]
def stream(self, columns: list[str]) -> Iterator[pa.RecordBatch]:
"""Execute once, yielding Arrow batches as the plan produces them."""
from datafusion import col
out = self._df.select(*(col(f'"{n}"') for n in columns))
return (b.to_pyarrow() for b in out.execute_stream())
def spill_parquet(self, path: str) -> None:
with pq.ParquetWriter(path, self.schema()) as writer:
for batch in self._df.execute_stream():
writer.write_batch(batch.to_pyarrow())
class DuckDBHandle:
"""Handle over a ``duckdb.DuckDBPyRelation``.
Relations are lazy relational algebra: ``filter``/``project`` derive
new relations and every materialization runs the query again, so a
single relation can serve any number of per-chunk fetches, each
narrowed to its own window — the *re-executable* property the
module docstring requires of every handle. Predicates are built
with DuckDB's typed expression API, never rendered SQL text.
Every engine call runs on one dedicated thread owned by the handle.
A relation is bound to one connection, and a query over a table
registered through xarray-sql re-enters Python from DuckDB's
execution threads (the Arrow scan callback); driving such queries
directly from several consumer threads at once (dask computing
output chunks of a lazy round-trip) deadlocks between the
connection's serialization, the callback's need for the GIL, and
the consumer pool's own thread management. Funnelling execution
through a single pre-started thread reproduces the topology that is
known safe — one thread inside the engine, every other thread
parked on a GIL-releasing wait.
"""
supports_chunked = False
"""Chunked (lazy) reconstruction is disabled for DuckDB relations.
Windows of a chunked round-trip re-execute the relation from the
consumer's worker threads (dask). A DuckDB query whose source is a
Python-callback Arrow scan (any table registered through xarray-sql)
intermittently deadlocks inside duckdb-python/CPython when other
Python threads start or stop during execution — reproduced on
duckdb 1.4-1.5 / CPython 3.12 / macOS at ~50% of runs, regardless
of ``SET threads=1``, connection serialization, or pool pre-warming.
Until that upstream race is fixed, chunked DuckDB round-trips fail
fast instead of hanging; the eager path (and every other handle
operation) runs on one dedicated thread and is unaffected.
"""
def __init__(self, rel: Any) -> None:
self._rel = rel
self._runner: ThreadPoolExecutor | None = None
if not _CAN_START_THREADS:
# Single-threaded interpreters (Pyodide) have no other
# threads to deadlock against; run engine calls inline.
return
self._runner = ThreadPoolExecutor(max_workers=1)
self._runner.submit(lambda: None).result() # start the thread now
# Stop the dedicated engine thread when the handle dies; it
# would otherwise linger for the life of the process, one per
# discarded handle. The callback is bound to the executor, not
# the handle, so the finalizer holds no reference back to self.
weakref.finalize(
self, self._runner.shutdown, wait=False, cancel_futures=True
)
def _run(self, fn: Any) -> Any:
if self._runner is None:
return fn()
return self._runner.submit(fn).result()
@staticmethod
def _to_arrow_table(rel: Any) -> pa.Table:
if hasattr(rel, "to_arrow_table"):
return rel.to_arrow_table()
return rel.fetch_arrow_table() # duckdb < 1.5
@staticmethod
def _to_arrow_reader(rel: Any) -> pa.RecordBatchReader:
if hasattr(rel, "to_arrow_reader"):
return rel.to_arrow_reader()
return rel.fetch_record_batch() # duckdb < 1.5
def schema(self) -> pa.Schema:
return self._run(
lambda: self._to_arrow_table(self._rel.limit(0)).schema
)
def distinct(self, column: str) -> np.ndarray:
import duckdb
table = self._run(
lambda: self._to_arrow_table(
self._rel.project(duckdb.ColumnExpression(column)).distinct()
)
)
return cast(
np.ndarray,
np.asarray(table.column(0).to_numpy(zero_copy_only=False)),
)
def fetch(
self, specs: dict[str, DimSpec], columns: list[str]
) -> list[pa.RecordBatch]:
import duckdb
predicate = None
for dim, (kind, a, b) in specs.items():
c = duckdb.ColumnExpression(dim)
if kind == "range":
p = (c >= duckdb.ConstantExpression(_plain(a))) & (
c <= duckdb.ConstantExpression(_plain(b))
)
else:
p = c.isin(*(duckdb.ConstantExpression(_plain(v)) for v in a))
predicate = p if predicate is None else predicate & p
rel = self._rel if predicate is None else self._rel.filter(predicate)
rel = rel.project(*(duckdb.ColumnExpression(n) for n in columns))
return cast(
list[pa.RecordBatch],
self._run(lambda: list(self._to_arrow_reader(rel))),
)
def stream(self, columns: list[str]) -> Iterator[pa.RecordBatch]:
"""Execute once, yielding Arrow batches as DuckDB produces them.
The reader is closed on the engine thread, even when the consumer
stops early, so the query is never torn down from another thread.
"""
import duckdb
rel = self._rel.project(*(duckdb.ColumnExpression(n) for n in columns))
reader = self._run(lambda: self._to_arrow_reader(rel))
batches = iter(reader)
try:
while (batch := self._run(lambda: next(batches, None))) is not None:
yield batch
finally:
self._run(reader.close)
def spill_parquet(self, path: str) -> None:
def run() -> None:
reader = self._to_arrow_reader(self._rel)
with pq.ParquetWriter(path, reader.schema) as writer:
for batch in reader:
writer.write_batch(batch)
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 or durations as a Polars Series of the column's type.
Built from the numpy values, so nanoseconds survive: Polars reads a
``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=_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
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)
class PolarsHandle:
"""Handle over a ``polars.LazyFrame``.
Per-window fetches run on the streaming engine, so a window read
never materializes more than the window even when the frame scans
an out-of-core source.
"""
supports_chunked = True
def __init__(self, lf: Any) -> None:
self._lf = lf
def schema(self) -> pa.Schema:
import polars as pl
return pl.DataFrame(schema=self._lf.collect_schema()).to_arrow().schema
def distinct(self, column: str) -> np.ndarray:
import polars as pl
out = _collect_streaming(self._lf.select(pl.col(column).unique()))
return cast(np.ndarray, out.to_series().to_numpy())
def fetch(
self, specs: dict[str, DimSpec], columns: list[str]
) -> list[pa.RecordBatch]:
import polars as pl
schema = self._lf.collect_schema()
exprs = []
for dim, (kind, a, b) in specs.items():
dtype = schema.get(dim)
if kind == "range":
exprs.append(
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
# imprecisely (silently matching nothing); degenerate
# ranges compare exactly. Reproduced on polars 1.42.
# ``any_horizontal`` keeps the disjunction flat — a
# left-deep OR chain plans quadratically in the number
# of values.
exprs.append(
pl.any_horizontal(
[pl.col(dim).is_between(*(_plain(v),) * 2) for v in a]
)
)
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]))
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())
def stream(self, columns: list[str]) -> Iterator[pa.RecordBatch]:
"""Execute once, yielding Arrow batches incrementally.
Unlike [fetch][xarray_sql.lazyscan.PolarsHandle.fetch], whose ``collect`` materializes the whole
result inside the engine before any batch surfaces, this yields
batches as the streaming engine produces them, so a byte budget
can fire before the result is fully in memory. Requires
``LazyFrame.collect_batches`` (polars >= 1.33); older Polars
raises here, before anything is collected.
"""
import polars as pl
lf = self._lf.select([pl.col(n) for n in columns])
if not hasattr(lf, "collect_batches"):
raise ValueError(
"streaming collection of a Polars LazyFrame requires "
"polars >= 1.33 (LazyFrame.collect_batches)."
)
def generate() -> Iterator[pa.RecordBatch]:
for frame in lf.collect_batches(engine="streaming"):
yield from frame.to_arrow().to_batches()
return generate()
def spill_parquet(self, path: str) -> None:
self._lf.sink_parquet(path)
def resolve_lazy_handle(result: Any) -> LazyResultHandle | None:
"""Adapt an engine result to a handle, or ``None`` if it is one-shot.
Recognizes DuckDB relations, Polars lazy *and* eager frames (an
eager frame re-executes trivially over its in-memory data), and
DataFusion DataFrames. ``pyarrow`` tables/readers and bare
``__arrow_c_stream__`` objects are one-shot streams: there is no
query to re-execute, so the lazy path cannot serve them.
"""
root = type(result).__module__.split(".")[0]
if root in ("duckdb", "_duckdb") and hasattr(result, "filter"):
return DuckDBHandle(result)
if root == "polars":
import polars as pl
if isinstance(result, pl.LazyFrame):
return PolarsHandle(result)
if isinstance(result, pl.DataFrame):
return PolarsHandle(result.lazy())
if hasattr(result, "execute_stream") and hasattr(result, "logical_plan"):
return DataFusionHandle(result)
return None