Support heterogeneous datasets with scalar variables in from_dataset

docs/examples.md

@@ -87,7 +87,30 @@ 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
+## GOES satellite imagery (scalar variables)
+
+Real-world stores often mix gridded data with scalar (0-dimensional) metadata.
+GOES satellite imagery, for example, pairs `(y, x)` image bands with dozens of
+scalar variables such as `goes_imager_projection`. `from_dataset` groups all the
+scalars into a single one-row table named `scalar`:
+
+```python
+import xarray as xr
+from xarray_sql import XarrayContext
+
+# A GOES-16 ABI cloud-and-moisture file: (y, x) bands + scalar metadata.
+ds = xr.open_dataset('OR_ABI-L2-MCMIPM1-M6_G16_....nc').chunk({'y': 250, 'x': 250})
+
+ctx = XarrayContext()
+ctx.from_dataset('goes', ds)
+
+ctx.sql('SELECT COUNT(*) FROM goes.y_x')   # the gridded bands
+ctx.sql('SELECT * FROM goes.scalar')       # one row of metadata
+```
+
+Override the default name like any other group with `table_names={(): 'metadata'}`.
+
+A runnable version of the ERA5 example lives at
 [`perf_tests/era5_temp_profile.py`](../perf_tests/era5_temp_profile.py).
 
 [arco-era5]: https://github.com/google-research/arco-era5

tests/test_df.py

@@ -3,6 +3,7 @@
 import numpy as np
 import pandas as pd
 import pyarrow as pa
+import pytest
 import xarray as xr
 
 from xarray_sql.df import (
@@ -58,6 +59,37 @@ def test_explode_data_equal_one_last(air):
     assert air.isel(iselection).equals(ds)
 
 
+def test_block_slices_scalar_dataset_yields_single_block():
+    # A dimensionless dataset (e.g. scalar metadata variables) has exactly
+    # one block: the whole, empty selection.
+    ds = xr.Dataset({"projection": ((), 0)})
+    assert list(block_slices(ds)) == [{}]
+
+
+def test_block_slices_scalar_ignores_irrelevant_chunks():
+    ds = xr.Dataset({"projection": ((), 0)})
+    assert list(block_slices(ds, chunks={"time": 4})) == [{}]
+
+
+def test_block_slices_filters_chunk_keys_to_dataset_dims(air_small):
+    # A chunk key for a dimension the dataset doesn't have is ignored,
+    # rather than raising.
+    base = list(block_slices(air_small, chunks={"time": 4, "lat": 3, "lon": 4}))
+    extra = list(
+        block_slices(
+            air_small, chunks={"time": 4, "lat": 3, "lon": 4, "absent": 2}
+        )
+    )
+    assert len(extra) == len(base)
+
+
+def test_block_slices_dimensional_unchunked_raises():
+    # A dataset with dimensions but no chunking is still a user error.
+    ds = xr.Dataset({"v": (["x"], np.arange(3))}, coords={"x": np.arange(3)})
+    with pytest.raises(AssertionError):
+        list(block_slices(ds))
+
+
 def test_from_map_basic():
     def make_df(x):
         return pd.DataFrame({"value": [x, x * 2], "index": [0, 1]})

tests/test_sql.py

@@ -359,6 +359,14 @@ def test_multiple_dim_groups(self):
     def test_empty_dataset(self):
         assert _group_vars_by_dims(xr.Dataset()) == {}
 
+    def test_includes_scalar_group(self):
+        """Scalar (0-dim) variables group under the empty dims tuple."""
+        ds = xr.Dataset(
+            {"band": (["y", "x"], np.zeros((2, 3))), "projection": ((), 0)}
+        )
+        groups = _group_vars_by_dims(ds)
+        assert groups == {("y", "x"): ["band"], (): ["projection"]}
+
     def test_ignores_coords(self):
         """Coordinate variables shouldn't be returned as groups."""
         ds = xr.Dataset(
@@ -404,6 +412,49 @@ def test_registers_multiple_tables(self, mixed_ds):
         assert len(surface) == 2 * 3 * 4
         assert len(upper) == 2 * 3 * 4 * 2
 
+    @pytest.fixture
+    def scalar_and_array_ds(self):
+        """A gridded variable plus a scalar metadata variable (GOES-like)."""
+        np.random.seed(0)
+        return xr.Dataset(
+            {
+                "temperature_2m": (
+                    ["time", "lat", "lon"],
+                    np.random.rand(2, 3, 4),
+                ),
+                "projection": ((), 0),
+            },
+            coords={
+                "time": pd.date_range("2020-01-01", periods=2),
+                "lat": np.linspace(-90, 90, 3),
+                "lon": np.linspace(-180, 180, 4),
+            },
+        ).chunk({"time": 1})
+
+    def test_registers_scalar_var_as_single_row_table(
+        self, scalar_and_array_ds
+    ):
+        ctx = XarrayContext()
+        ctx.from_dataset("goes", scalar_and_array_ds)
+        surface = ctx.sql("SELECT * FROM goes.time_lat_lon").to_pandas()
+        scalar = ctx.sql("SELECT * FROM goes.scalar").to_pandas()
+        assert "temperature_2m" in surface.columns
+        assert "projection" in scalar.columns
+        assert len(scalar) == 1
+
+    def test_scalar_group_in_catalog(self, scalar_and_array_ds):
+        ctx = XarrayContext()
+        ctx.from_dataset("goes", scalar_and_array_ds)
+        tables = set(ctx.catalog().schema("goes").table_names())
+        assert tables == {"time_lat_lon", "scalar"}
+
+    def test_scalar_table_name_override(self, scalar_and_array_ds):
+        ctx = XarrayContext()
+        ctx.from_dataset("goes", scalar_and_array_ds, table_names={(): "meta"})
+        result = ctx.sql("SELECT * FROM goes.meta").to_pandas()
+        assert "projection" in result.columns
+        assert len(result) == 1
+
     def test_table_names_override(self, mixed_ds):
         ctx = XarrayContext()
         ctx.from_dataset(

xarray_sql/df.py

@@ -30,13 +30,26 @@ def _get_chunk_slicer(
 def block_slices(ds: xr.Dataset, chunks: Chunks = None) -> Iterator[Block]:
     """Compute block slices for a chunked Dataset."""
     if chunks is not None:
+        # Only chunk dimensions this dataset actually has. A sub-dataset
+        # (e.g. one dimension group of a heterogeneous Dataset) need not
+        # contain every dimension named in the chunks spec.
+        chunks = {dim: size for dim, size in chunks.items() if dim in ds.sizes}
+    if chunks:
         for_chunking = ds.copy(data=None, deep=False).chunk(chunks)
         chunks = for_chunking.chunks
         del for_chunking
     else:
         chunks = ds.chunks
 
-    assert chunks, "Dataset `ds` must be chunked or `chunks` must be provided."
+    if not chunks:
+        # No chunkable dimensions. A dimensionless dataset (e.g. scalar
+        # metadata variables) is a single block; a dataset that has
+        # dimensions but no chunking is a user error.
+        assert not ds.sizes, (
+            "Dataset `ds` must be chunked or `chunks` must be provided."
+        )
+        yield {}
+        return
 
     # chunks is Dict[str, Tuple[int, ...]] from xarray
     chunk_bounds = {

xarray_sql/sql.py

@@ -86,10 +86,12 @@ def from_dataset(
         self.catalog().register_schema(name, schema)
 
         for dims, var_names in groups.items():
-            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)
+            # Scalar variables group under empty dims, where "_".join(()) is
+            # the empty string; fall back to a valid default table name.
+            sub_name = table_names.get(dims, "_".join(dims) or "scalar")
+            self._from_dataset(
+                sub_name, input_table[var_names], chunks, schema=schema
+            )
 
         return self
 
@@ -98,11 +100,17 @@ def _from_dataset(
         table_name: str,
         input_table: xr.Dataset,
         chunks: Chunks = None,
+        schema: Schema | None = None,
     ):
-        """Register a uniform-dimension Dataset as a single SQL table."""
+        """Register a Dataset as a single SQL table.
 
-        table = read_xarray_table(input_table, chunks)
-        self.register_table(table_name, table)
+        Registers a top-level table by default, or a table inside ``schema``
+        (a SQL namespace) when one is given.
+        """
+        register = (
+            self.register_table if schema is None else schema.register_table
+        )
+        register(table_name, read_xarray_table(input_table, chunks))
         self._maybe_register_cftime_udf(input_table)
         return self