From 7ddadec77f6a9b18eef0e3f48a2f1c0b8fd2863e Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sun, 25 Oct 2026 17:56:18 -0700 Subject: [PATCH 1/3] Split a query across partitions only where enough rows reach it Without a config, Context now sets DataFusion's use_row_number_estimates_to_optimize_partitioning, so a small input after a join, filter or aggregate is not repartitioned. A ddx training step of the GPT demo, 67 small plans, runs 2.2 times faster, and a whole SGD step takes 0.415 s instead of 0.655 s. Large inputs plan as before. A config given is used as it is. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h --- tests/test_sql.py | 34 ++++++++++++++++++++++++++++++++++ xarray_sql/sql.py | 25 ++++++++++++++++++++++--- 2 files changed, 56 insertions(+), 3 deletions(-) diff --git a/tests/test_sql.py b/tests/test_sql.py index f83af95d..ce79e82e 100644 --- a/tests/test_sql.py +++ b/tests/test_sql.py @@ -4,6 +4,7 @@ import pandas as pd import pytest import xarray as xr +from datafusion import SessionConfig import xarray_sql from xarray_sql import XarrayContext @@ -13,6 +14,39 @@ def test_context_is_the_xarray_context(): assert xarray_sql.Context is XarrayContext +SMALL_JOIN = """ +SELECT a.j, SUM(a.v * b.v) AS s FROM a JOIN b ON a.i = b.i AND a.j = b.j +GROUP BY a.j ORDER BY a.j +""" + + +def _small_join_plan(ctx): + ds = xr.Dataset( + {"v": (("i", "j"), np.ones((40, 25)))}, + coords={"i": np.arange(40), "j": np.arange(25)}, + ) + ctx.from_dataset("a", ds) + ctx.from_dataset("b", ds) + assert ctx.sql(SMALL_JOIN).to_pydict()["s"] == [40.0] * 25 + return ctx.sql(SMALL_JOIN).execution_plan().display_indent() + + +def test_a_small_query_is_not_split_across_partitions(): + assert "RepartitionExec" not in _small_join_plan(xarray_sql.Context()) + + +def test_a_config_given_is_used_as_it_is(): + config = SessionConfig().with_target_partitions(4) + assert "RepartitionExec" in _small_join_plan(xarray_sql.Context(config)) + + +def test_a_context_has_the_information_schema(): + rows = xarray_sql.Context().sql( + "SELECT table_name FROM information_schema.tables" + ) + assert rows.count() > 0 + + def test_sanity(air_dataset_small): ctx = XarrayContext() ctx.from_dataset("air", air_dataset_small) diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py index 229dbc5a..bc29e509 100644 --- a/xarray_sql/sql.py +++ b/xarray_sql/sql.py @@ -2,7 +2,7 @@ import pyarrow as pa import xarray as xr -from datafusion import SessionContext +from datafusion import SessionConfig, SessionContext from datafusion.catalog import Schema from types import ModuleType @@ -29,11 +29,30 @@ class Context(SessionContext): """A datafusion `SessionContext` that also supports `xarray.Dataset`s. + Without a `config`, a context splits a query's work across partitions + only where enough rows reach it, by DataFusion's estimate. DataFusion + otherwise splits every input whose row count is an estimate, as it is + after a join, a filter or an aggregate, however few rows that is. + Splitting small inputs makes a query of a few thousand rows two to + three times slower to run. A `config` given is used as it is. + Also available under its original name, ``XarrayContext``. """ - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) + def __init__(self, config: SessionConfig | None = None, *args, **kwargs): + if config is None: + # A bare SessionConfig turns off the information schema, which + # SessionContext() turns on; otherwise the two are the same. + config = ( + SessionConfig() + .with_information_schema(True) + .set( + "datafusion.execution." + "use_row_number_estimates_to_optimize_partitioning", + "true", + ) + ) + super().__init__(config, *args, **kwargs) # Track registered xarray Datasets so XarrayDataFrame can recover # defaults (dimension_columns) and metadata (var/dataset attrs, # non-dim coords, dim-coord dtype) that the forward pivot drops. From aa8097fb999fd4ed9754210e98f5fbc49610b436 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sun, 25 Oct 2026 18:13:24 -0700 Subject: [PATCH 2/3] Address review: check a given config's partition count Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h --- tests/test_sql.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_sql.py b/tests/test_sql.py index ce79e82e..56645401 100644 --- a/tests/test_sql.py +++ b/tests/test_sql.py @@ -37,7 +37,8 @@ def test_a_small_query_is_not_split_across_partitions(): def test_a_config_given_is_used_as_it_is(): config = SessionConfig().with_target_partitions(4) - assert "RepartitionExec" in _small_join_plan(xarray_sql.Context(config)) + plan = _small_join_plan(xarray_sql.Context(config)) + assert "RepartitionExec: partitioning=RoundRobinBatch(4)" in plan def test_a_context_has_the_information_schema(): From 26ab9b8b8157bdc1b75438ec2921be6978fd9469 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sun, 25 Oct 2026 18:57:30 -0700 Subject: [PATCH 3/3] Address review: explain the default config in a comment, not the docstring Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h --- xarray_sql/sql.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py index bc29e509..632f9bca 100644 --- a/xarray_sql/sql.py +++ b/xarray_sql/sql.py @@ -29,18 +29,18 @@ class Context(SessionContext): """A datafusion `SessionContext` that also supports `xarray.Dataset`s. - Without a `config`, a context splits a query's work across partitions - only where enough rows reach it, by DataFusion's estimate. DataFusion - otherwise splits every input whose row count is an estimate, as it is - after a join, a filter or an aggregate, however few rows that is. - Splitting small inputs makes a query of a few thousand rows two to - three times slower to run. A `config` given is used as it is. - Also available under its original name, ``XarrayContext``. """ def __init__(self, config: SessionConfig | None = None, *args, **kwargs): if config is None: + # Split a query's work across partitions only where enough rows + # reach it, by DataFusion's estimate. DataFusion otherwise splits + # every input whose row count is an estimate, as it is after a + # join, a filter or an aggregate, however few rows that is, which + # makes a query of a few thousand rows two to three times slower + # to run. A config given is used as it is. + # # A bare SessionConfig turns off the information schema, which # SessionContext() turns on; otherwise the two are the same. config = (