Skip to content

Commit 8c190d4

Browse files
alxmrsclaude
andcommitted
Test and document shared CTEs and views in prepare
Co-Authored-By: Claude Opus 5.5 <[email protected]> Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h
1 parent e46f822 commit 8c190d4

3 files changed

Lines changed: 124 additions & 7 deletions

File tree

‎docs/performance.md‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -204,6 +204,15 @@ planning, as it is over tables of more than about a megabyte, a call
204204
plans the query again and runs as fast as `ctx.sql`. So does a call of
205205
a plan over a lazy table (Dask or a lazily opened file).
206206

207+
DataFusion plans a CTE or a view in full at each place a query reads
208+
it. A model whose layers each read the layer before more than once
209+
grows with every layer: a one-layer GPT's loss plans as 9,882 operators,
210+
repeats of only 87 distinct subplans. `prepare` runs each CTE or view read more
211+
than once just once per call, and keeps its rows in a table for the
212+
rest of the query to read. That GPT's loss takes 16 ms per call, and
213+
2.7 s with `ctx.sql`. A query with `WITH RECURSIVE`, or a subquery with
214+
CTEs of its own, is planned as it's written.
215+
207216
## Round-trip faster with ORDER BY
208217

209218
Results that arrive **grid-ordered** — sorted by the dimension columns,

‎tests/test_prepare.py‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -262,3 +262,108 @@ def test_a_quoted_table_replaced_in_sql_counts_too():
262262
ctx.sql('CREATE OR REPLACE TABLE "Mixed" AS VALUES (1), (2)').collect()
263263
with pytest.raises(ValueError, match="registered again"):
264264
query()
265+
266+
267+
def _planned(query, a, b, *setup):
268+
"""The rows of ``query`` planned from scratch over ``a`` and ``b``."""
269+
fresh = Context()
270+
fresh.from_dataset("a", a)
271+
fresh.from_dataset("b", b)
272+
for statement in setup:
273+
fresh.sql(statement)
274+
return fresh.sql(query).to_arrow_table()
275+
276+
277+
def _calls(query, *setup):
278+
"""A prepared ``query``'s rows on new values of ``a``, against planning
279+
it."""
280+
rng = np.random.default_rng(1)
281+
b = _b(rng.normal(size=(4, 3)))
282+
ctx = Context()
283+
ctx.from_dataset("a", _a(rng.normal(size=(4, 3))))
284+
ctx.from_dataset("b", b)
285+
for statement in setup:
286+
ctx.sql(statement)
287+
prepared = ctx.prepare(query)
288+
for _ in range(3):
289+
a = _a(rng.normal(size=(4, 3)))
290+
yield prepared(a=a), _planned(query, a, b, *setup)
291+
292+
293+
# Each reads `t` twice, so a call runs `t` once and stores its rows.
294+
@pytest.mark.parametrize(
295+
"query",
296+
[
297+
# A CTE read twice.
298+
(
299+
"WITH t AS (SELECT i, j, 2 * v AS v2 FROM a) "
300+
"SELECT x.i, x.j, x.v2 * y.v2 AS p FROM t x JOIN t y "
301+
"ON x.i = y.i AND x.j = y.j ORDER BY x.i, x.j"
302+
),
303+
# Read twice, by a CTE read twice itself.
304+
(
305+
"WITH t AS (SELECT i, j, v FROM a), "
306+
"u AS (SELECT x.i, x.j, x.v + y.v AS v FROM t x JOIN t y "
307+
"ON x.i = y.i AND x.j = y.j), "
308+
"s AS (SELECT u.i, SUM(u.v * b.w) AS s FROM u JOIN b "
309+
"ON u.i = b.i AND u.j = b.j GROUP BY u.i) "
310+
"SELECT p.i, p.s + q.s AS s FROM s p JOIN s q ON p.i = q.i ORDER BY p.i"
311+
),
312+
# Rows whose number changes from call to call.
313+
(
314+
"WITH t AS (SELECT i, j, v FROM a WHERE v > 0) "
315+
"SELECT (SELECT COUNT(*) FROM t) AS n, (SELECT SUM(v) FROM t) AS s"
316+
),
317+
# Rows with nulls, and with strings, which aren't written in place.
318+
(
319+
"WITH t AS (SELECT i, j, CASE WHEN v > 0 THEN v END AS v, "
320+
"CASE WHEN v > 0 THEN 'up' ELSE 'down' END AS d FROM a) "
321+
"SELECT x.i, x.j, x.v, y.d FROM t x JOIN t y "
322+
"ON x.i = y.i AND x.j = y.j ORDER BY x.i, x.j"
323+
),
324+
# Columns named where the CTE is.
325+
(
326+
"WITH t(r, c, val) AS (SELECT i, j, v FROM a) "
327+
"SELECT x.r, x.c, x.val + y.val AS z FROM t x JOIN t y "
328+
"ON x.r = y.r AND x.c = y.c ORDER BY x.r, x.c"
329+
),
330+
],
331+
)
332+
def test_a_cte_read_more_than_once_gives_the_rows_planning_gives(query):
333+
for rows, want in _calls(query):
334+
assert rows.equals(want)
335+
336+
337+
def test_a_view_read_more_than_once_gives_the_rows_planning_gives():
338+
view = (
339+
"CREATE VIEW t AS WITH p AS (SELECT i, j, 2 * v AS v FROM a) "
340+
"SELECT x.i, x.j, x.v + y.v AS v FROM p x JOIN p y "
341+
"ON x.i = y.i AND x.j = y.j"
342+
)
343+
query = (
344+
"SELECT x.i, SUM(x.v * y.v) AS s FROM t x JOIN t y "
345+
"ON x.i = y.i AND x.j = y.j GROUP BY x.i ORDER BY x.i"
346+
)
347+
for rows, want in _calls(query, view):
348+
assert rows.equals(want)
349+
350+
351+
def test_a_cte_named_like_a_view_hides_it():
352+
view = "CREATE VIEW t AS SELECT i, j, -v AS v FROM a"
353+
query = (
354+
"WITH t AS (SELECT i, j, v FROM a) "
355+
"SELECT x.i, x.j, x.v + y.v AS v FROM t x JOIN t y "
356+
"ON x.i = y.i AND x.j = y.j ORDER BY x.i, x.j"
357+
)
358+
for rows, want in _calls(query, view):
359+
assert rows.equals(want)
360+
361+
362+
def test_a_recursive_cte_runs_again():
363+
query = (
364+
"WITH RECURSIVE r(n) AS (SELECT 1 UNION ALL SELECT n + 1 FROM r "
365+
"WHERE n < 4) SELECT SUM(n) AS s FROM r"
366+
)
367+
prepared = Context().prepare(query)
368+
assert prepared().to_pydict() == {"s": [10]}
369+
assert prepared().to_pydict() == {"s": [10]}

‎xarray_sql/sql.py‎

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -12,14 +12,14 @@
1212
from types import ModuleType
1313

1414
from . import cftime as cft
15+
from ._native import share as _share_sql
1516
from .df import (
1617
Chunks,
1718
TableNames,
1819
group_vars_by_dims,
1920
resolve_table_names,
2021
shared_coord_arrays,
2122
)
22-
from ._native import share as _share_sql
2323
from .ds import XarrayDataFrame
2424
from .reader import _in_memory_partitions, read_xarray_table
2525

@@ -525,20 +525,22 @@ def write(self, rows: pa.Table) -> bool:
525525
schema and number of rows, in numeric columns without nulls."""
526526
if rows.schema != self._schema or rows.num_rows != self._rows:
527527
return False
528-
if any(
529-
array is None or column.null_count
530-
for array, column in zip(self._arrays, rows.columns)
528+
arrays = [array for array in self._arrays if array is not None]
529+
if len(arrays) < len(self._arrays) or any(
530+
column.null_count for column in rows.columns
531531
):
532532
return False
533-
for array, column in zip(self._arrays, rows.columns):
533+
for array, column in zip(arrays, rows.columns):
534534
array[...] = column.to_numpy()
535535
return True
536536

537537

538538
def _writable(column: pa.ChunkedArray) -> np.ndarray | None:
539539
"""A copy of ``column`` that pyarrow can share memory with, if it's
540540
numeric and has no nulls."""
541-
numeric = pa.types.is_integer(column.type) or pa.types.is_floating(column.type)
541+
numeric = pa.types.is_integer(column.type) or pa.types.is_floating(
542+
column.type
543+
)
542544
if not numeric or column.null_count:
543545
return None
544546
return np.array(column.to_numpy(), copy=True)
@@ -563,9 +565,10 @@ def _share(
563565
).to_pydict()
564566
except Exception: # noqa: BLE001 - no information schema to read
565567
views = {"table_name": [], "definition": []}
566-
return _share_sql(
568+
shared: tuple[list[tuple[str, str]], str] | None = _share_sql(
567569
query, dict(zip(views["table_name"], views["definition"])), prefix
568570
)
571+
return shared
569572

570573

571574
def _copy(ctx: SessionContext, plan: ExecutionPlan) -> ExecutionPlan:

0 commit comments

Comments
 (0)