Skip to content

Commit f5f0160

Browse files
committed
fix : conflicts
2 parents 7b7518b + 10b7ad5 commit f5f0160

5 files changed

Lines changed: 146 additions & 11 deletions

File tree

docs/examples.md

Lines changed: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,39 @@ If you omit `table_names`, each table is named by joining its dimension names
8787
with underscores, e.g. `era5.time_latitude_longitude` and
8888
`era5.time_level_latitude_longitude`.
8989

90-
A runnable version of this example lives at
90+
## GOES satellite imagery (scalar variables)
91+
92+
Real-world stores often mix gridded data with scalar (0-dimensional) metadata.
93+
GOES satellite imagery, for example, pairs `(y, x)` image bands with dozens of
94+
scalar variables such as `goes_imager_projection`. `from_dataset` groups all the
95+
scalars into a single one-row table named `scalar`:
96+
97+
```python
98+
import fsspec
99+
import xarray as xr
100+
from xarray_sql import XarrayContext
101+
102+
# A real GOES-16 ABI cloud-and-moisture file from NOAA's public bucket:
103+
# (y, x) image bands alongside dozens of scalar metadata variables.
104+
url = (
105+
'https://noaa-goes16.s3.amazonaws.com/ABI-L2-MCMIPM/2024/001/00/'
106+
'OR_ABI-L2-MCMIPM1-M6_G16_s20240010000281_e20240010000350_c20240010000426.nc'
107+
)
108+
ds = xr.open_dataset(fsspec.open_local(f'simplecache::{url}')).chunk(
109+
{'y': 250, 'x': 250}
110+
)
111+
112+
ctx = XarrayContext()
113+
ctx.from_dataset('goes', ds)
114+
115+
# The gridded bands and the scalar metadata are separate tables.
116+
ctx.sql('SELECT COUNT(*) AS n FROM goes.y_x').to_pandas()['n'][0] # -> 250000
117+
ctx.sql('SELECT * FROM goes.scalar').to_pandas().shape # -> (1, 89)
118+
```
119+
120+
Override the default name like any other group with `table_names={(): 'metadata'}`.
121+
122+
A runnable version of the ERA5 example lives at
91123
[`perf_tests/era5_temp_profile.py`](../perf_tests/era5_temp_profile.py).
92124

93125
[arco-era5]: https://github.com/google-research/arco-era5

tests/test_df.py

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import numpy as np
44
import pandas as pd
55
import pyarrow as pa
6+
import pytest
67
import xarray as xr
78

89
from xarray_sql.df import (
@@ -58,6 +59,37 @@ def test_explode_data_equal_one_last(air):
5859
assert air.isel(iselection).equals(ds)
5960

6061

62+
def test_block_slices_scalar_dataset_yields_single_block():
63+
# A dimensionless dataset (e.g. scalar metadata variables) has exactly
64+
# one block: the whole, empty selection.
65+
ds = xr.Dataset({"projection": ((), 0)})
66+
assert list(block_slices(ds)) == [{}]
67+
68+
69+
def test_block_slices_scalar_ignores_irrelevant_chunks():
70+
ds = xr.Dataset({"projection": ((), 0)})
71+
assert list(block_slices(ds, chunks={"time": 4})) == [{}]
72+
73+
74+
def test_block_slices_filters_chunk_keys_to_dataset_dims(air_small):
75+
# A chunk key for a dimension the dataset doesn't have is ignored,
76+
# rather than raising.
77+
base = list(block_slices(air_small, chunks={"time": 4, "lat": 3, "lon": 4}))
78+
extra = list(
79+
block_slices(
80+
air_small, chunks={"time": 4, "lat": 3, "lon": 4, "absent": 2}
81+
)
82+
)
83+
assert len(extra) == len(base)
84+
85+
86+
def test_block_slices_dimensional_unchunked_raises():
87+
# A dataset with dimensions but no chunking is still a user error.
88+
ds = xr.Dataset({"v": (["x"], np.arange(3))}, coords={"x": np.arange(3)})
89+
with pytest.raises(AssertionError):
90+
list(block_slices(ds))
91+
92+
6193
def test_from_map_basic():
6294
def make_df(x):
6395
return pd.DataFrame({"value": [x, x * 2], "index": [0, 1]})

tests/test_sql.py

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -359,6 +359,14 @@ def test_multiple_dim_groups(self):
359359
def test_empty_dataset(self):
360360
assert _group_vars_by_dims(xr.Dataset()) == {}
361361

362+
def test_includes_scalar_group(self):
363+
"""Scalar (0-dim) variables group under the empty dims tuple."""
364+
ds = xr.Dataset(
365+
{"band": (["y", "x"], np.zeros((2, 3))), "projection": ((), 0)}
366+
)
367+
groups = _group_vars_by_dims(ds)
368+
assert groups == {("y", "x"): ["band"], (): ["projection"]}
369+
362370
def test_ignores_coords(self):
363371
"""Coordinate variables shouldn't be returned as groups."""
364372
ds = xr.Dataset(
@@ -404,6 +412,49 @@ def test_registers_multiple_tables(self, mixed_ds):
404412
assert len(surface) == 2 * 3 * 4
405413
assert len(upper) == 2 * 3 * 4 * 2
406414

415+
@pytest.fixture
416+
def scalar_and_array_ds(self):
417+
"""A gridded variable plus a scalar metadata variable (GOES-like)."""
418+
np.random.seed(0)
419+
return xr.Dataset(
420+
{
421+
"temperature_2m": (
422+
["time", "lat", "lon"],
423+
np.random.rand(2, 3, 4),
424+
),
425+
"projection": ((), 0),
426+
},
427+
coords={
428+
"time": pd.date_range("2020-01-01", periods=2),
429+
"lat": np.linspace(-90, 90, 3),
430+
"lon": np.linspace(-180, 180, 4),
431+
},
432+
).chunk({"time": 1})
433+
434+
def test_registers_scalar_var_as_single_row_table(
435+
self, scalar_and_array_ds
436+
):
437+
ctx = XarrayContext()
438+
ctx.from_dataset("goes", scalar_and_array_ds)
439+
surface = ctx.sql("SELECT * FROM goes.time_lat_lon").to_pandas()
440+
scalar = ctx.sql("SELECT * FROM goes.scalar").to_pandas()
441+
assert "temperature_2m" in surface.columns
442+
assert "projection" in scalar.columns
443+
assert len(scalar) == 1
444+
445+
def test_scalar_group_in_catalog(self, scalar_and_array_ds):
446+
ctx = XarrayContext()
447+
ctx.from_dataset("goes", scalar_and_array_ds)
448+
tables = set(ctx.catalog().schema("goes").table_names())
449+
assert tables == {"time_lat_lon", "scalar"}
450+
451+
def test_scalar_table_name_override(self, scalar_and_array_ds):
452+
ctx = XarrayContext()
453+
ctx.from_dataset("goes", scalar_and_array_ds, table_names={(): "meta"})
454+
result = ctx.sql("SELECT * FROM goes.meta").to_pandas()
455+
assert "projection" in result.columns
456+
assert len(result) == 1
457+
407458
def test_table_names_override(self, mixed_ds):
408459
ctx = XarrayContext()
409460
ctx.from_dataset(

xarray_sql/df.py

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,13 +30,26 @@ def _get_chunk_slicer(
3030
def block_slices(ds: xr.Dataset, chunks: Chunks = None) -> Iterator[Block]:
3131
"""Compute block slices for a chunked Dataset."""
3232
if chunks is not None:
33+
# Only chunk dimensions this dataset actually has. A sub-dataset
34+
# (e.g. one dimension group of a heterogeneous Dataset) need not
35+
# contain every dimension named in the chunks spec.
36+
chunks = {dim: size for dim, size in chunks.items() if dim in ds.sizes}
37+
if chunks:
3338
for_chunking = ds.copy(data=None, deep=False).chunk(chunks)
3439
chunks = for_chunking.chunks
3540
del for_chunking
3641
else:
3742
chunks = ds.chunks
3843

39-
assert chunks, "Dataset `ds` must be chunked or `chunks` must be provided."
44+
if not chunks:
45+
# No chunkable dimensions. A dimensionless dataset (e.g. scalar
46+
# metadata variables) is a single block; a dataset that has
47+
# dimensions but no chunking is a user error.
48+
assert not ds.sizes, (
49+
"Dataset `ds` must be chunked or `chunks` must be provided."
50+
)
51+
yield {}
52+
return
4053

4154
# chunks is Dict[str, Tuple[int, ...]] from xarray
4255
chunk_bounds = {

xarray_sql/sql.py

Lines changed: 16 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -90,20 +90,22 @@ def from_dataset(
9090
groups = _group_vars_by_dims(input_table)
9191

9292
if len(groups) <= 1:
93+
self._registered_datasets[name] = input_table
9394
return self._from_dataset(name, input_table, chunks)
9495

9596
table_names = table_names or {}
9697
schema = Schema.memory_schema(self)
9798
self.catalog().register_schema(name, schema)
9899

99100
for dims, var_names in groups.items():
100-
sub_name = table_names.get(dims, "_".join(dims))
101+
# Scalar variables group under empty dims, where "_".join(()) is
102+
# the empty string; fall back to a valid default table name.
103+
sub_name = table_names.get(dims, "_".join(dims) or "scalar")
101104
sub_ds = input_table[var_names]
102-
schema.register_table(sub_name, read_xarray_table(sub_ds, chunks))
103-
# Track the fully-qualified name for XarrayDataFrame metadata
104-
# recovery on round-trip.
105+
self._from_dataset(sub_name, sub_ds, chunks, schema=schema)
106+
# Track the fully-qualified name so XarrayDataFrame metadata
107+
# recovery can find this Dataset on round-trip.
105108
self._registered_datasets[f"{name}.{sub_name}"] = sub_ds
106-
self._maybe_register_cftime_udf(sub_ds)
107109

108110
return self
109111

@@ -112,12 +114,17 @@ def _from_dataset(
112114
table_name: str,
113115
input_table: xr.Dataset,
114116
chunks: Chunks = None,
117+
schema: Schema | None = None,
115118
):
116-
"""Register a uniform-dimension Dataset as a single SQL table."""
119+
"""Register a Dataset as a single SQL table.
117120
118-
table = read_xarray_table(input_table, chunks)
119-
self.register_table(table_name, table)
120-
self._registered_datasets[table_name] = input_table
121+
Registers a top-level table by default, or a table inside ``schema``
122+
(a SQL namespace) when one is given.
123+
"""
124+
register = (
125+
self.register_table if schema is None else schema.register_table
126+
)
127+
register(table_name, read_xarray_table(input_table, chunks))
121128
self._maybe_register_cftime_udf(input_table)
122129
return self
123130

0 commit comments

Comments
 (0)