diff --git a/tests/test_sql.py b/tests/test_sql.py index f83af95d..56645401 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,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) diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py index 229dbc5a..632f9bca 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 @@ -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.