diff --git a/docs/examples.md b/docs/examples.md index 0a3066c..e3be79c 100644 --- a/docs/examples.md +++ b/docs/examples.md @@ -87,7 +87,39 @@ 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 fsspec +import xarray as xr +from xarray_sql import XarrayContext + +# A real GOES-16 ABI cloud-and-moisture file from NOAA's public bucket: +# (y, x) image bands alongside dozens of scalar metadata variables. +url = ( + 'https://noaa-goes16.s3.amazonaws.com/ABI-L2-MCMIPM/2024/001/00/' + 'OR_ABI-L2-MCMIPM1-M6_G16_s20240010000281_e20240010000350_c20240010000426.nc' +) +ds = xr.open_dataset(fsspec.open_local(f'simplecache::{url}')).chunk( + {'y': 250, 'x': 250} +) + +ctx = XarrayContext() +ctx.from_dataset('goes', ds) + +# The gridded bands and the scalar metadata are separate tables. +ctx.sql('SELECT COUNT(*) AS n FROM goes.y_x').to_pandas()['n'][0] # -> 250000 +ctx.sql('SELECT * FROM goes.scalar').to_pandas().shape # -> (1, 89) +``` + +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 diff --git a/tests/test_df.py b/tests/test_df.py index 4ebb675..a2b443f 100644 --- a/tests/test_df.py +++ b/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]}) diff --git a/tests/test_sql.py b/tests/test_sql.py index 74f4c66..857a7bc 100644 --- a/tests/test_sql.py +++ b/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( diff --git a/xarray_sql/df.py b/xarray_sql/df.py index db9d37e..f55b6ad 100644 --- a/xarray_sql/df.py +++ b/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 = { diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py index 3881d98..403eaee 100644 --- a/xarray_sql/sql.py +++ b/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