Split a query across partitions only where enough rows reach it

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,40 @@ 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)
+    plan = _small_join_plan(xarray_sql.Context(config))
+    assert "RepartitionExec: partitioning=RoundRobinBatch(4)" in plan
+
+
+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)

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
 
@@ -32,8 +32,27 @@ class Context(SessionContext):
     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:
+            # 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 = (
+                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.