From aa4ded72fff8ade4bcab453fdb9b2542d57b6409 Mon Sep 17 00:00:00 2001 From: Evan Park Date: Tue, 28 Apr 2026 20:03:36 -0700 Subject: [PATCH 1/8] added dim grouping and table registration for differing dims in xarraycontext --- xarray_sql/sql.py | 60 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 60 insertions(+) diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py index f30f5b8..b485d8c 100644 --- a/xarray_sql/sql.py +++ b/xarray_sql/sql.py @@ -1,5 +1,6 @@ import xarray as xr from datafusion import SessionContext +from collections import defaultdict from . import cftime as cft from .df import Chunks @@ -48,6 +49,7 @@ def from_dataset( Returns: self, to allow chaining. """ + table = read_xarray_table(input_table, chunks) self.register_table(table_name, table) @@ -61,3 +63,61 @@ def from_dataset( break # One UDF per context is enough. return self + + def from_dataset_multi( + self, + dataset_name: str, + input_table: xr.Dataset, + dim_aliases: dict[tuple[str, ...], str] | None = None, + chunks: Chunks = None, + ): + """Register an xarray Dataset with mixed dimensions as multiple SQL tables. + + When a Dataset contains variables with different dimensions (e.g. some + variables on a 3D grid and others on a 4D grid), this method splits them + into separate tables and registers each one. Each table is named + ``.__...`` by default, or using the override + provided in ``dim_aliases``:: + + ctx.from_dataset_multi('era5', ds, chunks={'time': 24}) + # registers: 'era5.time_lat_lon' and 'era5.time_lat_lon_level' + ctx.sql('SELECT AVG(temperature_2m) FROM "era5.time_lat_lon"') + + Args: + dataset_name: A name for the dataset, used as a prefix for all + registered table names. + input_table: An xarray Dataset. Variables may have differing dimensions. + dim_aliases: Optional mapping from dimension tuples to custom table + name suffixes. For example, + ``{('time', 'lat', 'lon'): 'surface'}`` registers the table as + ``era5.surface`` instead of ``era5.time_lat_lon``. + chunks: Xarray-like chunks specification. If not provided, uses + the Dataset's existing chunks. + + Returns: + self, to allow chaining. + """ + + dim_aliases = dim_aliases or {} + groups = _group_vars_by_dims(input_table) + + for dims, var_names in groups.items(): + suffix = dim_aliases.get(dims, "_".join(dims)) + table_name = f"{dataset_name}_{suffix}" + sub_ds = input_table[var_names] + self.from_dataset(table_name, sub_ds, chunks) + + return self + + +def _group_vars_by_dims(ds: xr.Dataset) -> dict[tuple[str, ...], list[str]]: + """Group variables in the dataset based on shared dims. + + ("time", "lat", "lon"): ["temperature_2m", "wind_speed"], + ("time", "lat", "lon", "level"): ["pressure", "humidity"] + """ + groups = defaultdict(list) + for var_name, var in ds.data_vars.items(): + dims = var.dims + groups[dims].append(var_name) + return groups From c57f9e495e36be2b615f087c22ac41e64fdfb193 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sun, 17 May 2026 13:46:21 -0700 Subject: [PATCH 2/8] Added tests for new multi-table structure. --- tests/test_sql.py | 107 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 107 insertions(+) diff --git a/tests/test_sql.py b/tests/test_sql.py index 01929f0..0fba987 100644 --- a/tests/test_sql.py +++ b/tests/test_sql.py @@ -6,6 +6,7 @@ import xarray as xr from xarray_sql import XarrayContext +from xarray_sql.sql import _group_vars_by_dims def test_sanity(air_dataset_small): @@ -322,3 +323,109 @@ def test_gregorian_like_no_cftime_udf(self): ctx.sql( "SELECT COUNT(*) FROM rasm WHERE time >= cftime('1980-01-01')" ).collect() + + +class TestGroupVarsByDims: + """Unit tests for the private _group_vars_by_dims helper.""" + + def test_single_dim_group(self): + ds = xr.Dataset( + { + "a": (["x", "y"], np.zeros((2, 3))), + "b": (["x", "y"], np.ones((2, 3))), + } + ) + groups = _group_vars_by_dims(ds) + assert groups == {("x", "y"): ["a", "b"]} + + def test_multiple_dim_groups(self): + ds = xr.Dataset( + { + "surface": (["time", "lat", "lon"], np.zeros((2, 3, 4))), + "upper": ( + ["time", "lat", "lon", "level"], + np.zeros((2, 3, 4, 5)), + ), + } + ) + groups = _group_vars_by_dims(ds) + assert set(groups.keys()) == { + ("time", "lat", "lon"), + ("time", "lat", "lon", "level"), + } + assert groups[("time", "lat", "lon")] == ["surface"] + assert groups[("time", "lat", "lon", "level")] == ["upper"] + + def test_empty_dataset(self): + assert _group_vars_by_dims(xr.Dataset()) == {} + + def test_ignores_coords(self): + """Coordinate variables shouldn't be returned as groups.""" + ds = xr.Dataset( + {"v": (["x"], np.arange(3))}, + coords={"x": np.arange(3), "label": ("x", ["a", "b", "c"])}, + ) + groups = _group_vars_by_dims(ds) + assert groups == {("x",): ["v"]} + + +class TestFromDatasetMultiDims: + """from_dataset should split datasets with mixed dims into multiple tables.""" + + @pytest.fixture + def mixed_ds(self): + np.random.seed(0) + return xr.Dataset( + { + "temperature_2m": ( + ["time", "lat", "lon"], + np.random.rand(2, 3, 4), + ), + "pressure": ( + ["time", "lat", "lon", "level"], + np.random.rand(2, 3, 4, 2), + ), + }, + coords={ + "time": pd.date_range("2020-01-01", periods=2), + "lat": np.linspace(-90, 90, 3), + "lon": np.linspace(-180, 180, 4), + "level": [500, 1000], + }, + ).chunk({"time": 1}) + + def test_registers_multiple_tables(self, mixed_ds): + ctx = XarrayContext() + ctx.from_dataset("era5", mixed_ds) + surface = ctx.sql('SELECT * FROM "era5_time_lat_lon"').to_pandas() + upper = ctx.sql('SELECT * FROM "era5_time_lat_lon_level"').to_pandas() + assert "temperature_2m" in surface.columns + assert "pressure" in upper.columns + assert len(surface) == 2 * 3 * 4 + assert len(upper) == 2 * 3 * 4 * 2 + + def test_dim_group_aliases(self, mixed_ds): + ctx = XarrayContext() + ctx.from_dataset( + "era5", + mixed_ds, + dim_group_aliases={("time", "lat", "lon"): "surface"}, + ) + result = ctx.sql('SELECT * FROM "era5_surface"').to_pandas() + assert "temperature_2m" in result.columns + # Non-aliased group falls back to default suffix. + upper = ctx.sql('SELECT * FROM "era5_time_lat_lon_level"').to_pandas() + assert "pressure" in upper.columns + + def test_uniform_dims_uses_table_name_directly(self, mixed_ds): + """A single dim group should register under the bare table_name.""" + ds = mixed_ds[["temperature_2m"]] + ctx = XarrayContext() + ctx.from_dataset("surface", ds) + result = ctx.sql("SELECT * FROM surface").to_pandas() + assert "temperature_2m" in result.columns + + def test_dim_group_aliases_is_keyword_only(self, mixed_ds): + ctx = XarrayContext() + with pytest.raises(TypeError): + ctx.from_dataset("era5", mixed_ds, {("time",): "x"}) From 849c55c73c28da848ca42e437767cb4e2339f1b8 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sun, 17 May 2026 13:49:39 -0700 Subject: [PATCH 3/8] Primary interface now supports registering multiple tables. --- xarray_sql/sql.py | 100 ++++++++++++++++++++++------------------------ 1 file changed, 48 insertions(+), 52 deletions(-) diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py index b485d8c..4c49c10 100644 --- a/xarray_sql/sql.py +++ b/xarray_sql/sql.py @@ -14,9 +14,30 @@ def from_dataset( self, table_name: str, input_table: xr.Dataset, + *, + dim_group_aliases: dict[tuple[str, ...], str] | None = None, chunks: Chunks = None, ): - """Register an xarray Dataset as a queryable SQL table. + """Register an xarray Dataset as one or more queryable SQL tables. + + When all data variables share the same dimensions, the dataset is + registered as a single table named ``table_name``. When variables + have differing dimensions (e.g. some on a 3D grid and others on a + 4D grid), the dataset is split into one table per dimension group, + each named ``___...`` by default:: + + ctx.from_dataset('era5', ds, chunks={'time': 24}) + # registers: 'era5_time_lat_lon' and 'era5_time_lat_lon_level' + ctx.sql('SELECT AVG(temperature_2m) FROM "era5_time_lat_lon"') + + Use ``dim_group_aliases`` to override the suffix for specific + dimension tuples:: + + ctx.from_dataset( + 'era5', ds, + dim_group_aliases={('time', 'lat', 'lon'): 'surface'}, + ) + # registers: 'era5_surface' and 'era5_time_lat_lon_level' For datasets with non-Gregorian cftime coordinates (e.g. 360_day, julian), a ``cftime()`` scalar UDF is automatically registered so @@ -25,9 +46,6 @@ def from_dataset( ctx.from_dataset("ds360", ds, chunks={"time": 6}) ctx.sql("SELECT * FROM ds360 WHERE time >= cftime('2000-07-01')") - The UDF converts a date string to the int64 offset used to store - that calendar's time axis. - .. note:: Only one ``cftime()`` UDF is registered per context, using the @@ -40,15 +58,38 @@ def from_dataset( for each calendar. Args: - table_name: The SQL table name to register the dataset under. - input_table: An xarray Dataset. All data_vars must share the - same dimensions. + table_name: The SQL table name (or table-name prefix, when the + dataset has mixed dimensions) to register the dataset under. + input_table: An xarray Dataset. + dim_group_aliases: Optional mapping from dimension tuples to + custom table-name suffixes, used when the dataset has + variables with differing dimensions. chunks: Xarray-like chunks specification. If not provided, uses the Dataset's existing chunks. Returns: self, to allow chaining. """ + groups = _group_vars_by_dims(input_table) + + if len(groups) <= 1: + return self._from_dataset(table_name, input_table, chunks) + + dim_group_aliases = dim_group_aliases or {} + for dims, var_names in groups.items(): + suffix = dim_group_aliases.get(dims, "_".join(dims)) + sub_name = f"{table_name}_{suffix}" + self._from_dataset(sub_name, input_table[var_names], chunks) + + return self + + def _from_dataset( + self, + table_name: str, + input_table: xr.Dataset, + chunks: Chunks = None, + ): + """Register a uniform-dimension Dataset as a single SQL table.""" table = read_xarray_table(input_table, chunks) self.register_table(table_name, table) @@ -64,51 +105,6 @@ def from_dataset( return self - def from_dataset_multi( - self, - dataset_name: str, - input_table: xr.Dataset, - dim_aliases: dict[tuple[str, ...], str] | None = None, - chunks: Chunks = None, - ): - """Register an xarray Dataset with mixed dimensions as multiple SQL tables. - - When a Dataset contains variables with different dimensions (e.g. some - variables on a 3D grid and others on a 4D grid), this method splits them - into separate tables and registers each one. Each table is named - ``.__...`` by default, or using the override - provided in ``dim_aliases``:: - - ctx.from_dataset_multi('era5', ds, chunks={'time': 24}) - # registers: 'era5.time_lat_lon' and 'era5.time_lat_lon_level' - ctx.sql('SELECT AVG(temperature_2m) FROM "era5.time_lat_lon"') - - Args: - dataset_name: A name for the dataset, used as a prefix for all - registered table names. - input_table: An xarray Dataset. Variables may have differing dimensions. - dim_aliases: Optional mapping from dimension tuples to custom table - name suffixes. For example, - ``{('time', 'lat', 'lon'): 'surface'}`` registers the table as - ``era5.surface`` instead of ``era5.time_lat_lon``. - chunks: Xarray-like chunks specification. If not provided, uses - the Dataset's existing chunks. - - Returns: - self, to allow chaining. - """ - - dim_aliases = dim_aliases or {} - groups = _group_vars_by_dims(input_table) - - for dims, var_names in groups.items(): - suffix = dim_aliases.get(dims, "_".join(dims)) - table_name = f"{dataset_name}_{suffix}" - sub_ds = input_table[var_names] - self.from_dataset(table_name, sub_ds, chunks) - - return self - def _group_vars_by_dims(ds: xr.Dataset) -> dict[tuple[str, ...], list[str]]: """Group variables in the dataset based on shared dims. From 1ab6c86d7e834ef0da56960ecc370deda15b0f79 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sun, 17 May 2026 14:07:50 -0700 Subject: [PATCH 4/8] fix: cargo.lock has incorrect xql version. --- Cargo.lock | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index 41ef28b..1d53519 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3375,7 +3375,7 @@ checksum = "ea2f10b9bb0928dfb1b42b65e1f9e36f7f54dbdf08457afefb38afcdec4fa2bb" [[package]] name = "xarray_sql" -version = "0.2.1" +version = "0.2.2" dependencies = [ "arrow", "async-stream", From 7b35cb0c235023e6e1b132bfb5a96ead5d708ca8 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sun, 17 May 2026 14:12:02 -0700 Subject: [PATCH 5/8] Better interface: making use of catalog hierarchy of schemas and tables. --- tests/test_sql.py | 18 +++++++++++++----- xarray_sql/sql.py | 46 +++++++++++++++++++++++++++------------------- 2 files changed, 40 insertions(+), 24 deletions(-) diff --git a/tests/test_sql.py b/tests/test_sql.py index 0fba987..6d228a9 100644 --- a/tests/test_sql.py +++ b/tests/test_sql.py @@ -397,8 +397,8 @@ def mixed_ds(self): def test_registers_multiple_tables(self, mixed_ds): ctx = XarrayContext() ctx.from_dataset("era5", mixed_ds) - surface = ctx.sql('SELECT * FROM "era5_time_lat_lon"').to_pandas() - upper = ctx.sql('SELECT * FROM "era5_time_lat_lon_level"').to_pandas() + surface = ctx.sql("SELECT * FROM era5.time_lat_lon").to_pandas() + upper = ctx.sql("SELECT * FROM era5.time_lat_lon_level").to_pandas() assert "temperature_2m" in surface.columns assert "pressure" in upper.columns assert len(surface) == 2 * 3 * 4 @@ -411,12 +411,20 @@ def test_dim_group_aliases(self, mixed_ds): mixed_ds, dim_group_aliases={("time", "lat", "lon"): "surface"}, ) - result = ctx.sql('SELECT * FROM "era5_surface"').to_pandas() + result = ctx.sql("SELECT * FROM era5.surface").to_pandas() assert "temperature_2m" in result.columns - # Non-aliased group falls back to default suffix. - upper = ctx.sql('SELECT * FROM "era5_time_lat_lon_level"').to_pandas() + # Non-aliased group falls back to the default joined-dim table name. + upper = ctx.sql("SELECT * FROM era5.time_lat_lon_level").to_pandas() assert "pressure" in upper.columns + def test_schema_registered_in_catalog(self, mixed_ds): + """Mixed-dim datasets should create a SQL schema under the catalog.""" + ctx = XarrayContext() + ctx.from_dataset("era5", mixed_ds) + assert "era5" in ctx.catalog().schema_names() + tables = set(ctx.catalog().schema("era5").table_names()) + assert tables == {"time_lat_lon", "time_lat_lon_level"} + def test_uniform_dims_uses_table_name_directly(self, mixed_ds): """A single dim group should register under the bare table_name.""" ds = mixed_ds[["temperature_2m"]] diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py index 4c49c10..2fc313d 100644 --- a/xarray_sql/sql.py +++ b/xarray_sql/sql.py @@ -1,5 +1,6 @@ import xarray as xr from datafusion import SessionContext +from datafusion.catalog import Schema from collections import defaultdict from . import cftime as cft @@ -23,12 +24,14 @@ def from_dataset( When all data variables share the same dimensions, the dataset is registered as a single table named ``table_name``. When variables have differing dimensions (e.g. some on a 3D grid and others on a - 4D grid), the dataset is split into one table per dimension group, - each named ``___...`` by default:: + 4D grid), the dataset is split into one table per dimension group. + The tables are registered under a SQL schema (namespace) named + ``table_name`` and named ``__...`` by default:: ctx.from_dataset('era5', ds, chunks={'time': 24}) - # registers: 'era5_time_lat_lon' and 'era5_time_lat_lon_level' - ctx.sql('SELECT AVG(temperature_2m) FROM "era5_time_lat_lon"') + # registers tables: 'era5.time_lat_lon' and + # 'era5.time_lat_lon_level' + ctx.sql('SELECT AVG(temperature_2m) FROM era5.time_lat_lon') Use ``dim_group_aliases`` to override the suffix for specific dimension tuples:: @@ -37,7 +40,7 @@ def from_dataset( 'era5', ds, dim_group_aliases={('time', 'lat', 'lon'): 'surface'}, ) - # registers: 'era5_surface' and 'era5_time_lat_lon_level' + ctx.sql('SELECT * FROM era5.surface') For datasets with non-Gregorian cftime coordinates (e.g. 360_day, julian), a ``cftime()`` scalar UDF is automatically registered so @@ -58,12 +61,13 @@ def from_dataset( for each calendar. Args: - table_name: The SQL table name (or table-name prefix, when the - dataset has mixed dimensions) to register the dataset under. + table_name: The SQL table name. For datasets with mixed + dimensions, this becomes the name of a SQL schema + (namespace) containing one table per dimension group. input_table: An xarray Dataset. dim_group_aliases: Optional mapping from dimension tuples to - custom table-name suffixes, used when the dataset has - variables with differing dimensions. + custom table names within the schema, used when the dataset + has variables with differing dimensions. chunks: Xarray-like chunks specification. If not provided, uses the Dataset's existing chunks. @@ -76,10 +80,14 @@ def from_dataset( return self._from_dataset(table_name, input_table, chunks) dim_group_aliases = dim_group_aliases or {} + schema = Schema.memory_schema(self) + self.catalog().register_schema(table_name, schema) + for dims, var_names in groups.items(): - suffix = dim_group_aliases.get(dims, "_".join(dims)) - sub_name = f"{table_name}_{suffix}" - self._from_dataset(sub_name, input_table[var_names], chunks) + sub_name = dim_group_aliases.get(dims, "_".join(dims)) + sub_ds = input_table[var_names] + schema.register_table(sub_name, read_xarray_table(sub_ds, chunks)) + self._maybe_register_cftime_udf(sub_ds) return self @@ -93,18 +101,18 @@ def _from_dataset( table = read_xarray_table(input_table, chunks) self.register_table(table_name, table) + self._maybe_register_cftime_udf(input_table) + return self - # Auto-register a cftime() UDF for non-Gregorian cftime coordinates - # so users can write: WHERE time > cftime('0500-01-01') - for coord_name in input_table.dims: - if cft.is_cftime_index(input_table, coord_name): - units, cal = cft.encoding(input_table, coord_name) + def _maybe_register_cftime_udf(self, ds: xr.Dataset) -> None: + """Auto-register a cftime() UDF for non-Gregorian cftime coordinates.""" + for coord_name in ds.dims: + if cft.is_cftime_index(ds, coord_name): + units, cal = cft.encoding(ds, coord_name) if not cft.is_gregorian_like(cal): self.register_udf(cft.make_cftime_udf(units, cal)) break # One UDF per context is enough. - return self - def _group_vars_by_dims(ds: xr.Dataset) -> dict[tuple[str, ...], list[str]]: """Group variables in the dataset based on shared dims. From 3d596677d4d8b46d3863d18e673bf8b5f59432ae Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sun, 17 May 2026 14:46:39 -0700 Subject: [PATCH 6/8] More accurate names for the top level API. --- tests/test_sql.py | 10 +++++----- xarray_sql/sql.py | 40 +++++++++++++++++++++------------------- 2 files changed, 26 insertions(+), 24 deletions(-) diff --git a/tests/test_sql.py b/tests/test_sql.py index 6d228a9..74f4c66 100644 --- a/tests/test_sql.py +++ b/tests/test_sql.py @@ -404,12 +404,12 @@ def test_registers_multiple_tables(self, mixed_ds): assert len(surface) == 2 * 3 * 4 assert len(upper) == 2 * 3 * 4 * 2 - def test_dim_group_aliases(self, mixed_ds): + def test_table_names_override(self, mixed_ds): ctx = XarrayContext() ctx.from_dataset( "era5", mixed_ds, - dim_group_aliases={("time", "lat", "lon"): "surface"}, + table_names={("time", "lat", "lon"): "surface"}, ) result = ctx.sql("SELECT * FROM era5.surface").to_pandas() assert "temperature_2m" in result.columns @@ -425,15 +425,15 @@ def test_schema_registered_in_catalog(self, mixed_ds): tables = set(ctx.catalog().schema("era5").table_names()) assert tables == {"time_lat_lon", "time_lat_lon_level"} - def test_uniform_dims_uses_table_name_directly(self, mixed_ds): - """A single dim group should register under the bare table_name.""" + def test_uniform_dims_uses_name_directly(self, mixed_ds): + """A single dim group should register under the bare name.""" ds = mixed_ds[["temperature_2m"]] ctx = XarrayContext() ctx.from_dataset("surface", ds) result = ctx.sql("SELECT * FROM surface").to_pandas() assert "temperature_2m" in result.columns - def test_dim_group_aliases_is_keyword_only(self, mixed_ds): + def test_table_names_is_keyword_only(self, mixed_ds): ctx = XarrayContext() with pytest.raises(TypeError): ctx.from_dataset("era5", mixed_ds, {("time",): "x"}) diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py index 2fc313d..3881d98 100644 --- a/xarray_sql/sql.py +++ b/xarray_sql/sql.py @@ -13,32 +13,32 @@ class XarrayContext(SessionContext): def from_dataset( self, - table_name: str, + name: str, input_table: xr.Dataset, *, - dim_group_aliases: dict[tuple[str, ...], str] | None = None, + table_names: dict[tuple[str, ...], str] | None = None, chunks: Chunks = None, ): """Register an xarray Dataset as one or more queryable SQL tables. When all data variables share the same dimensions, the dataset is - registered as a single table named ``table_name``. When variables - have differing dimensions (e.g. some on a 3D grid and others on a - 4D grid), the dataset is split into one table per dimension group. + registered as a single table named ``name``. When variables have + differing dimensions (e.g. some on a 3D grid and others on a 4D + grid), the dataset is split into one table per dimension group. The tables are registered under a SQL schema (namespace) named - ``table_name`` and named ``__...`` by default:: + ``name`` and named ``__...`` by default:: ctx.from_dataset('era5', ds, chunks={'time': 24}) # registers tables: 'era5.time_lat_lon' and # 'era5.time_lat_lon_level' ctx.sql('SELECT AVG(temperature_2m) FROM era5.time_lat_lon') - Use ``dim_group_aliases`` to override the suffix for specific - dimension tuples:: + Use ``table_names`` to override the name for specific dimension + tuples:: ctx.from_dataset( 'era5', ds, - dim_group_aliases={('time', 'lat', 'lon'): 'surface'}, + table_names={('time', 'lat', 'lon'): 'surface'}, ) ctx.sql('SELECT * FROM era5.surface') @@ -61,13 +61,15 @@ def from_dataset( for each calendar. Args: - table_name: The SQL table name. For datasets with mixed - dimensions, this becomes the name of a SQL schema - (namespace) containing one table per dimension group. + name: The SQL identifier under which the dataset is registered. + For datasets with uniform dimensions, this is the table + name. For datasets with mixed dimensions, this is the name + of a SQL schema (namespace) containing one table per + dimension group. input_table: An xarray Dataset. - dim_group_aliases: Optional mapping from dimension tuples to - custom table names within the schema, used when the dataset - has variables with differing dimensions. + table_names: Optional mapping from dimension tuples to custom + table names within the schema, used when the dataset has + variables with differing dimensions. chunks: Xarray-like chunks specification. If not provided, uses the Dataset's existing chunks. @@ -77,14 +79,14 @@ def from_dataset( groups = _group_vars_by_dims(input_table) if len(groups) <= 1: - return self._from_dataset(table_name, input_table, chunks) + return self._from_dataset(name, input_table, chunks) - dim_group_aliases = dim_group_aliases or {} + table_names = table_names or {} schema = Schema.memory_schema(self) - self.catalog().register_schema(table_name, schema) + self.catalog().register_schema(name, schema) for dims, var_names in groups.items(): - sub_name = dim_group_aliases.get(dims, "_".join(dims)) + sub_name = table_names.get(dims, "_".join(dims)) sub_ds = input_table[var_names] schema.register_table(sub_name, read_xarray_table(sub_ds, chunks)) self._maybe_register_cftime_udf(sub_ds) From 6cebb32b92b1b2891f30731efb6079ea9cc01dac Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sun, 17 May 2026 14:47:29 -0700 Subject: [PATCH 7/8] BFD: Putting ARCO-ERA5 example at the front of the project!! --- README.md | 80 +++++++++++++++++++++++------------------------- docs/examples.md | 59 +++++++++++++++++++++++++++++++++++ 2 files changed, 97 insertions(+), 42 deletions(-) diff --git a/README.md b/README.md index 5b03ea5..791da43 100644 --- a/README.md +++ b/README.md @@ -19,52 +19,48 @@ This is an experiment to provide a SQL interface for array datasets. import xarray as xr import xarray_sql as xql -ds = xr.tutorial.open_dataset('air_temperature') +# Anonymous read from the public GCS bucket — no auth required. +url = 'gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3' +full = xr.open_zarr(url, chunks=None, storage_options={'token': 'anon'}) + +# A few hours over the NYC area for January 2020, two pressure levels. +ds = full[['2m_temperature', 'temperature']].sel( + time=slice('2020-01-01', '2020-01-01T05'), + latitude=slice(40, 39), + longitude=slice(286, 287), # ERA5 uses 0-360 longitudes + level=[500, 850], +).chunk({'time': 3}) -# The same as a dask-sql Context; i.e. an Apache DataFusion Context. ctx = xql.XarrayContext() -ctx.from_dataset('air', ds, chunks=dict(time=24)) # the dataset needs to be chunked! -# data is only materialized when we make a query. - -result = ctx.sql(''' - SELECT - "lat", "lon", AVG("air") as air_avg - FROM - "air" - GROUP BY - "lat", "lon" -''') -# DataFrame() -# +------+-------+--------------------+ -# | lat | lon | air_avg | -# +------+-------+--------------------+ -# | 75.0 | 205.0 | 259.88662671232834 | -# | 75.0 | 207.5 | 259.48268150684896 | -# | 75.0 | 230.0 | 258.9192123287667 | -# | 75.0 | 275.0 | 257.07574315068456 | -# | 75.0 | 322.5 | 250.11792123287654 | -# | 75.0 | 325.0 | 250.81590068493134 | -# | 72.5 | 205.0 | 262.74933904109537 | -# | 72.5 | 207.5 | 262.5384315068488 | -# | 72.5 | 230.0 | 260.82879452054743 | -# | 72.5 | 275.0 | 257.3063321917804 | -# +------+-------+--------------------+ -# Data truncated. - -# The full query is only made when we call `collect()`, or, in this case, -# `to_pandas()`. -df = result.to_pandas() -df.head() -# lat lon air_avg -# 0 75.0 232.5 258.836188 -# 1 75.0 247.5 257.716171 -# 2 75.0 262.5 257.347959 -# 3 75.0 277.5 257.671308 -# 4 72.5 232.5 260.654401 +ctx.from_dataset('era5', ds, table_names={ + ('time', 'latitude', 'longitude'): 'surface', + ('time', 'level', 'latitude', 'longitude'): 'atmosphere', +}) +# Tables in the 'era5' schema: 'surface', 'atmosphere' + +ctx.sql(''' + SELECT AVG("2m_temperature") - 273.15 AS avg_c + FROM era5.surface +''').to_pandas() +# avg_c +# 0 8.640069 + +ctx.sql(''' + SELECT level, AVG(temperature) - 273.15 AS avg_c + FROM era5.atmosphere + GROUP BY level + ORDER BY level +''').to_pandas() +# level avg_c +# 0 500 -20.215307 +# 1 850 -2.379477 ``` -Succinctly, we "pivot" Xarray Datasets (with consistent dimensions) to treat them like tables so we can run -SQL queries against them. +> A runnable version of this example lives at +> [`perf_tests/era5_nyc_avg_temp.py`](perf_tests/era5_nyc_avg_temp.py). + +Succinctly, we "pivot" Xarray Datasets to treat them like tables so we can run +SQL queries against them. ## Why build this? diff --git a/docs/examples.md b/docs/examples.md index cff0f33..b8bae01 100644 --- a/docs/examples.md +++ b/docs/examples.md @@ -21,3 +21,62 @@ result = ctx.sql(''' df = result.to_pandas() df.head() ``` + +## Mixed-dimension datasets: ARCO-ERA5 + +When a Dataset has variables with differing dimensions (e.g. surface fields on +`(time, latitude, longitude)` and atmospheric fields on +`(time, level, latitude, longitude)`), `from_dataset` splits them into one +table per dimension group, registered together under a SQL schema named after +the first argument. [ARCO-ERA5][arco-era5] is a good example: 262 of its +variables are surface fields and 11 are atmospheric. + +Use the `table_names` kwarg to give each dimension group a friendly name: + +```python +import xarray as xr +import xarray_sql as xql + +# Open ARCO-ERA5 directly from GCS (anonymous read). +url = 'gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3' +full = xr.open_zarr(url, chunks=None, storage_options={'token': 'anon'}) + +# Small slice for the example: a few hours over the NYC area, two pressure levels, +# one variable from each dimension group. +ds = full[['2m_temperature', 'temperature']].sel( + time=slice('2020-01-01', '2020-01-01T05'), + latitude=slice(40, 39), + longitude=slice(286, 287), # ERA5 uses 0–360° longitudes + level=[500, 850], +).chunk({'time': 3}) + +ctx = xql.XarrayContext() +ctx.from_dataset('era5', ds, table_names={ + ('time', 'latitude', 'longitude'): 'surface', + ('time', 'level', 'latitude', 'longitude'): 'atmosphere', +}) +# Registers two tables under a SQL schema named 'era5': 'surface' and 'atmosphere'. + +# Query the surface table. +ctx.sql(''' + SELECT AVG("2m_temperature") - 273.15 AS avg_c + FROM era5.surface +''').to_pandas() + +# Query the atmospheric table. +ctx.sql(''' + SELECT level, AVG(temperature) - 273.15 AS avg_c + FROM era5.atmosphere + GROUP BY level + ORDER BY level +''').to_pandas() +``` + +If you omit `table_names`, each table is named by joining its dimension names +with underscores, e.g. `era5.time_latitude_longitude` and +`era5.time_level_latitude_longitude`. + +A runnable version of this example lives at +[`perf_tests/era5_nyc_avg_temp.py`](../perf_tests/era5_nyc_avg_temp.py). + +[arco-era5]: https://github.com/google-research/arco-era5 From 4c825a1e3fa21661feeb260776b04924ae27501f Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sun, 17 May 2026 15:57:17 -0700 Subject: [PATCH 8/8] SQL-forward README and examples. --- README.md | 86 +++++++++++++++++++++------ docs/examples.md | 37 ++++++++---- perf_tests/era5_temp_profile.py | 101 ++++++++++++++++++++++++++++++++ 3 files changed, 193 insertions(+), 31 deletions(-) create mode 100644 perf_tests/era5_temp_profile.py diff --git a/README.md b/README.md index 791da43..58afc7c 100644 --- a/README.md +++ b/README.md @@ -19,45 +19,95 @@ This is an experiment to provide a SQL interface for array datasets. import xarray as xr import xarray_sql as xql -# Anonymous read from the public GCS bucket — no auth required. -url = 'gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3' -full = xr.open_zarr(url, chunks=None, storage_options={'token': 'anon'}) - -# A few hours over the NYC area for January 2020, two pressure levels. -ds = full[['2m_temperature', 'temperature']].sel( - time=slice('2020-01-01', '2020-01-01T05'), - latitude=slice(40, 39), - longitude=slice(286, 287), # ERA5 uses 0-360 longitudes - level=[500, 850], -).chunk({'time': 3}) + +# Open a year of ARCO-ERA5 — all 273 variables. Selecting a year up front +# keeps Dask's partition setup cheap before any chunks are read from GCS. +ds = ( + xr.open_zarr('gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3', + chunks=None, + storage_options={'token': 'anon'}) # Anonymous read from the public GCS bucket — no auth required. + .sel(time='2020') + .chunk({'time': 1}) +) ctx = xql.XarrayContext() ctx.from_dataset('era5', ds, table_names={ ('time', 'latitude', 'longitude'): 'surface', ('time', 'level', 'latitude', 'longitude'): 'atmosphere', }) -# Tables in the 'era5' schema: 'surface', 'atmosphere' +# Registration: ~0.5s for a full year of hourly ERA5, all variables. + + +# Heads up: ARCO-ERA5 has 262 surface + 11 atmospheric variables. The library +# pushes column projection down to Zarr, so SELECT only fetches what you ask +# for — but `SELECT * FROM era5.surface` would try to pull every variable +# across the year (terabytes from GCS). +# ---> Always SELECT specific columns. <--- +# Average 2m-temperature over NYC on the morning of 2020-01-01. The library +# pushes WHERE clauses on dimension columns down to partition pruning. ctx.sql(''' SELECT AVG("2m_temperature") - 273.15 AS avg_c FROM era5.surface + WHERE time BETWEEN TIMESTAMP '2020-01-01' + AND TIMESTAMP '2020-01-01 05:00:00' + AND latitude BETWEEN 39 AND 40 + AND longitude BETWEEN 286 AND 287 -- ERA5 uses 0-360 longitudes ''').to_pandas() # avg_c # 0 8.640069 +# Average temperature per pressure level, globally. ctx.sql(''' SELECT level, AVG(temperature) - 273.15 AS avg_c FROM era5.atmosphere + WHERE time BETWEEN TIMESTAMP '2020-01-01' + AND TIMESTAMP '2020-01-01 05:00:00' GROUP BY level - ORDER BY level + ORDER BY level DESC ''').to_pandas() -# level avg_c -# 0 500 -20.215307 -# 1 850 -2.379477 +# level avg_c +# 0 1000 6.621012 ← surface +# 1 975 5.185638 +# 2 950 4.028429 +# 3 925 3.082812 +# 4 900 2.210917 +# 5 875 1.395018 +# 6 850 0.634267 +# 7 825 -0.210372 +# 8 800 -1.181075 +# 9 775 -2.306465 +# 10 750 -3.535534 +# 11 700 -6.241685 +# 12 650 -9.236364 +# 13 600 -12.580938 +# 14 550 -16.335386 +# 15 500 -20.643604 +# 16 450 -25.573401 +# 17 400 -31.156920 +# 18 350 -37.400552 +# 19 300 -43.852607 +# 20 250 -49.322132 +# 21 225 -51.569113 +# 22 200 -53.693248 +# 23 175 -55.890484 +# 24 150 -58.382290 +# 25 125 -61.091916 +# 26 100 -63.624885 ← tropopause +# 27 70 -63.182300 +# 28 50 -60.124845 +# 29 30 -55.986327 +# 30 20 -52.433089 +# 31 10 -44.140750 +# 32 7 -38.707350 +# 33 5 -32.621999 +# 34 3 -21.509175 +# 35 2 -13.355764 +# 36 1 -9.020513 ← top of atmosphere ``` -> A runnable version of this example lives at -> [`perf_tests/era5_nyc_avg_temp.py`](perf_tests/era5_nyc_avg_temp.py). +_(A runnable version of this example lives at +[`perf_tests/era5_temp_profile.py`](perf_tests/era5_temp_profile.py).)_ Succinctly, we "pivot" Xarray Datasets to treat them like tables so we can run SQL queries against them. diff --git a/docs/examples.md b/docs/examples.md index b8bae01..0a3066c 100644 --- a/docs/examples.md +++ b/docs/examples.md @@ -31,7 +31,9 @@ table per dimension group, registered together under a SQL schema named after the first argument. [ARCO-ERA5][arco-era5] is a good example: 262 of its variables are surface fields and 11 are atmospheric. -Use the `table_names` kwarg to give each dimension group a friendly name: +Open a year of ARCO-ERA5 and let SQL `WHERE` clauses do the filtering — the +library prunes time partitions and pushes dimension-column filters down. Use +the `table_names` kwarg to give each dimension group a friendly name: ```python import xarray as xr @@ -41,14 +43,16 @@ import xarray_sql as xql url = 'gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3' full = xr.open_zarr(url, chunks=None, storage_options={'token': 'anon'}) -# Small slice for the example: a few hours over the NYC area, two pressure levels, -# one variable from each dimension group. -ds = full[['2m_temperature', 'temperature']].sel( - time=slice('2020-01-01', '2020-01-01T05'), - latitude=slice(40, 39), - longitude=slice(286, 287), # ERA5 uses 0–360° longitudes - level=[500, 850], -).chunk({'time': 3}) +# A full year of hourly ERA5 — all 273 variables. No spatial slicing on the +# xarray side; SQL WHERE clauses below express the filters. `chunks={'time': 1}` +# aligns Dask chunks to native Zarr chunks of shape (1, 37, 721, 1440) so +# chunk reads from GCS happen concurrently. +# +# Heads up: 262 of those variables are surface and 11 are atmospheric. The +# library pushes column projection down, so SELECT only fetches what you ask +# for — but `SELECT * FROM era5.surface` would try to pull every variable +# across the year (terabytes from GCS). Always SELECT specific columns. +ds = full.sel(time='2020').chunk({'time': 1}) ctx = xql.XarrayContext() ctx.from_dataset('era5', ds, table_names={ @@ -57,18 +61,25 @@ ctx.from_dataset('era5', ds, table_names={ }) # Registers two tables under a SQL schema named 'era5': 'surface' and 'atmosphere'. -# Query the surface table. +# Average 2m-temperature over the NYC area on the morning of 2020-01-01. ctx.sql(''' SELECT AVG("2m_temperature") - 273.15 AS avg_c FROM era5.surface + WHERE time BETWEEN TIMESTAMP '2020-01-01' + AND TIMESTAMP '2020-01-01 05:00:00' + AND latitude BETWEEN 39 AND 40 + AND longitude BETWEEN 286 AND 287 ''').to_pandas() -# Query the atmospheric table. +# Average temperature per pressure level, globally — the standard +# atmospheric temperature profile. Scans ~230M rows. ctx.sql(''' SELECT level, AVG(temperature) - 273.15 AS avg_c FROM era5.atmosphere + WHERE time BETWEEN TIMESTAMP '2020-01-01' + AND TIMESTAMP '2020-01-01 05:00:00' GROUP BY level - ORDER BY level + ORDER BY level DESC -- surface (1000 hPa) first ''').to_pandas() ``` @@ -77,6 +88,6 @@ with underscores, e.g. `era5.time_latitude_longitude` and `era5.time_level_latitude_longitude`. A runnable version of this example lives at -[`perf_tests/era5_nyc_avg_temp.py`](../perf_tests/era5_nyc_avg_temp.py). +[`perf_tests/era5_temp_profile.py`](../perf_tests/era5_temp_profile.py). [arco-era5]: https://github.com/google-research/arco-era5 diff --git a/perf_tests/era5_temp_profile.py b/perf_tests/era5_temp_profile.py new file mode 100644 index 0000000..28275a9 --- /dev/null +++ b/perf_tests/era5_temp_profile.py @@ -0,0 +1,101 @@ +#!/usr/bin/env python3 +"""Surface and global-atmospheric temperatures on 2020-01-01, in SQL. + +Two queries against ARCO-ERA5 on the morning of January 1, 2020: + + * **Surface (local).** Average 2m-temperature over a small grid covering the + New York City area for the first six hours. + * **Atmosphere (global).** Average temperature per pressure level, computed + over the entire planet for the same six hours — a classic atmospheric + temperature profile (surface around 1000 hPa is warmest, tropopause near + 100 hPa is coldest). + +Both queries express their filters entirely in SQL: ``xr.open_zarr`` is given +a single calendar year and no spatial slicing. The library's table provider +prunes time partitions for ``WHERE time …`` filters, and pushes ``WHERE +latitude/longitude …`` down to dimension columns. + +ARCO-ERA5's atmospheric variables are stored in native Zarr chunks of shape +``(1, 37, 721, 1440)`` — about 150 MB per hour. We align Dask chunks to that +shape with ``chunks={'time': 1}`` so chunks fetch from GCS concurrently. The +global atmospheric query scans ~230M rows after pruning. + +The Zarr is read anonymously from the public GCS bucket — no auth required. +""" + +import time + +import xarray as xr + +import xarray_sql as xql + + +URL = "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3" + + +def main() -> None: + full = xr.open_zarr(URL, chunks=None, storage_options={"token": "anon"}) + + # Open a full calendar year — all 273 variables. No spatial slicing on + # the xarray side; SQL WHERE clauses below express the filters. + # + # Heads up: the library pushes column projection down to Zarr, so SELECT + # only fetches what you ask for — but `SELECT * FROM era5.surface` would + # try to read every variable across the year (terabytes from GCS). + # Always SELECT specific columns. + ds = full.sel(time="2020").chunk({"time": 1}) + print( + "ARCO-ERA5 opened: year 2020, " + f"{ds.sizes['time']:,} hourly time steps, " + f"{len(ds.data_vars)} variables (no spatial pre-slicing)." + ) + + ctx = xql.XarrayContext() + t0 = time.perf_counter() + ctx.from_dataset( + "era5", + ds, + table_names={ + ("time", "latitude", "longitude"): "surface", + ("time", "level", "latitude", "longitude"): "atmosphere", + }, + ) + print(f"Registration: {time.perf_counter() - t0:.2f}s") + ctx.sql("SELECT 1").to_pandas() # warm the planner + + print("\nAverage 2m-temperature over NYC, 2020-01-01 00:00-05:00 UTC (°C):") + t0 = time.perf_counter() + surface = ctx.sql( + """ + SELECT AVG("2m_temperature") - 273.15 AS avg_c + FROM era5.surface + WHERE time BETWEEN TIMESTAMP '2020-01-01' + AND TIMESTAMP '2020-01-01 05:00:00' + AND latitude BETWEEN 39 AND 40 + AND longitude BETWEEN 286 AND 287 -- ERA5 uses 0-360 longitudes + """ + ).to_pandas() + print(surface) + print(f" ({time.perf_counter() - t0:.2f}s)") + + print( + "\nAverage temperature per pressure level, globally, " + "2020-01-01 00:00-05:00 UTC (°C):" + ) + t0 = time.perf_counter() + profile = ctx.sql( + """ + SELECT level, AVG(temperature) - 273.15 AS avg_c + FROM era5.atmosphere + WHERE time BETWEEN TIMESTAMP '2020-01-01' + AND TIMESTAMP '2020-01-01 05:00:00' + GROUP BY level + ORDER BY level DESC -- surface (1000 hPa) first, top of atmosphere last + """ + ).to_pandas() + print(profile.to_string(index=False)) + print(f" ({time.perf_counter() - t0:.2f}s, ~230M rows scanned)") + + +if __name__ == "__main__": + main()