Skip to content

Commit 86d5c00

Browse files
alxmrsclaude
andauthored
Split a query across partitions only where enough rows reach it (#289)
Part of workstream 1, per-query overhead, in #285. ## What `Context()`, given no `config`, now sets DataFusion's `datafusion.execution.use_row_number_estimates_to_optimize_partitioning` to `true`. Everything else stays as `SessionContext()` sets it, including the information schema. A `config` given is used as it is. ## Why DataFusion decides whether to split an input across partitions from its row count. When the count is exact, as it is for a table, DataFusion splits only inputs larger than a batch. When the count is an estimate, as it is after a join, a filter or an aggregate, DataFusion splits the input however few rows the estimate says it has, unless this setting is on. apache/datafusion#26087 describes this, and apache/datafusion#18595 asks to stop repartitioning small datasets. The ML demos run many queries of a few thousand rows each. On 12 cores, every one of them was split into 12 partitions and merged back: | Measured (DataFusion 55.0.0, i7-8700, quiet machine) | before | after | one partition | |---|---:|---:|---:| | ddx's 67 step plans of one GPT SGD step, executing only | 151 ms | 69 ms | 61 ms | | GPT SGD step (`02_nanogpt.py`, later steps) | 0.655 s | 0.415 s | | | GPT forward loss: operators, seconds | 13,756, 3.6 s | 9,882, 3.1 s | | | Attention, 512 tokens, forward and gradient, both spellings | 0.28 / 0.86 / 0.41 / 0.97 s | 0.27 / 0.87 / 0.42 / 0.98 s | | Every check against JAX passes with the setting on. Lazy Datasets of 4 and 16 million rows plan the same `RepartitionExec`s, and take the same time, with and without it: a group-by, a filtered group-by and a self-join. ## Risk A low estimate on an input that's really large would leave it in one partition. A filter's estimate is a fraction of its input's rows, so a large table still gives a large estimate. Anyone who wants DataFusion's default can pass `SessionConfig()`. ## Tests - `Context()` runs a small join and aggregate with no `RepartitionExec`. - A given config is used as it is: with 4 target partitions, the same query repartitions. - `Context()` still has the information schema. A bare `SessionConfig()` turns it off, and ddx reads view definitions from it. The ML demos' README numbers will be re-recorded in #269, or after it, once this merges. 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h <!-- This is an auto-generated comment: release notes by coderabbit.ai --> <!-- coderabbit-ai-generated:v1 --> ## Summary by CodeRabbit * **New Features** * SQL contexts now expose the information schema by default, including table listings. * By default, small-input joins use row-number estimates when determining partitioning and omit repartitioning when it is unnecessary. Configuring a target partition count can still result in repartitioning. * You can provide a custom session configuration when creating a context; supplied settings are used unchanged. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Claude Opus 5.5 <[email protected]>
1 parent c70b1b2 commit 86d5c00

2 files changed

Lines changed: 57 additions & 3 deletions

File tree

‎tests/test_sql.py‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import pandas as pd
55
import pytest
66
import xarray as xr
7+
from datafusion import SessionConfig
78

89
import xarray_sql
910
from xarray_sql import XarrayContext
@@ -13,6 +14,40 @@ def test_context_is_the_xarray_context():
1314
assert xarray_sql.Context is XarrayContext
1415

1516

17+
SMALL_JOIN = """
18+
SELECT a.j, SUM(a.v * b.v) AS s FROM a JOIN b ON a.i = b.i AND a.j = b.j
19+
GROUP BY a.j ORDER BY a.j
20+
"""
21+
22+
23+
def _small_join_plan(ctx):
24+
ds = xr.Dataset(
25+
{"v": (("i", "j"), np.ones((40, 25)))},
26+
coords={"i": np.arange(40), "j": np.arange(25)},
27+
)
28+
ctx.from_dataset("a", ds)
29+
ctx.from_dataset("b", ds)
30+
assert ctx.sql(SMALL_JOIN).to_pydict()["s"] == [40.0] * 25
31+
return ctx.sql(SMALL_JOIN).execution_plan().display_indent()
32+
33+
34+
def test_a_small_query_is_not_split_across_partitions():
35+
assert "RepartitionExec" not in _small_join_plan(xarray_sql.Context())
36+
37+
38+
def test_a_config_given_is_used_as_it_is():
39+
config = SessionConfig().with_target_partitions(4)
40+
plan = _small_join_plan(xarray_sql.Context(config))
41+
assert "RepartitionExec: partitioning=RoundRobinBatch(4)" in plan
42+
43+
44+
def test_a_context_has_the_information_schema():
45+
rows = xarray_sql.Context().sql(
46+
"SELECT table_name FROM information_schema.tables"
47+
)
48+
assert rows.count() > 0
49+
50+
1651
def test_sanity(air_dataset_small):
1752
ctx = XarrayContext()
1853
ctx.from_dataset("air", air_dataset_small)

‎xarray_sql/sql.py‎

Lines changed: 22 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
import pyarrow as pa
44
import xarray as xr
5-
from datafusion import SessionContext
5+
from datafusion import SessionConfig, SessionContext
66
from datafusion.catalog import Schema
77
from types import ModuleType
88

@@ -32,8 +32,27 @@ class Context(SessionContext):
3232
Also available under its original name, ``XarrayContext``.
3333
"""
3434

35-
def __init__(self, *args, **kwargs):
36-
super().__init__(*args, **kwargs)
35+
def __init__(self, config: SessionConfig | None = None, *args, **kwargs):
36+
if config is None:
37+
# Split a query's work across partitions only where enough rows
38+
# reach it, by DataFusion's estimate. DataFusion otherwise splits
39+
# every input whose row count is an estimate, as it is after a
40+
# join, a filter or an aggregate, however few rows that is, which
41+
# makes a query of a few thousand rows two to three times slower
42+
# to run. A config given is used as it is.
43+
#
44+
# A bare SessionConfig turns off the information schema, which
45+
# SessionContext() turns on; otherwise the two are the same.
46+
config = (
47+
SessionConfig()
48+
.with_information_schema(True)
49+
.set(
50+
"datafusion.execution."
51+
"use_row_number_estimates_to_optimize_partitioning",
52+
"true",
53+
)
54+
)
55+
super().__init__(config, *args, **kwargs)
3756
# Track registered xarray Datasets so XarrayDataFrame can recover
3857
# defaults (dimension_columns) and metadata (var/dataset attrs,
3958
# non-dim coords, dim-coord dtype) that the forward pivot drops.

0 commit comments

Comments
 (0)