Skip to content

Commit

Permalink
Handle scalar dataset variables
Browse files Browse the repository at this point in the history
  • Loading branch information
ghidalgo3 committed Jul 30, 2024
1 parent 179bb2a commit 3b091e8
Show file tree
Hide file tree
Showing 2 changed files with 34 additions and 5 deletions.
26 changes: 26 additions & 0 deletions virtualizarr/tests/test_xarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -419,6 +419,32 @@ def test_open_virtual_dataset_passes_expected_args(
}
mock_read_kerchunk.assert_called_once_with(**args)

@patch("virtualizarr.kerchunk.parse_array_refs")
def test_open_dataset_with_scalars(self, mock_parse_array_refs, tmpdir):
expected_attrs = {"units": "K", "long_name": "scalar"}
mock_parse_array_refs.return_value = (
None,
ZArray(
chunks=(),
shape=(),
compressor={},
dtype=np.dtype("float32"),
fill_value=0.0,
filters=None,
order="C",
zarr_format=2,
),
{**expected_attrs, **{"_ARRAY_DIMENSIONS": []}},
)

ds = xr.Dataset()
ds["scalar"] = xr.Variable(dims=(), data=None, attrs=expected_attrs)
ds.to_netcdf(f"{tmpdir}/scalar.nc")
vds = open_virtual_dataset(f"{tmpdir}/scalar.nc")
assert vds.scalar.dims == ()
assert vds.scalar.data == 0.0
assert vds.scalar.attrs == expected_attrs


class TestRenamePaths:
def test_rename_to_str(self, netcdf4_file):
Expand Down
13 changes: 8 additions & 5 deletions virtualizarr/xarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -333,13 +333,16 @@ def variable_from_kerchunk_refs(

arr_refs = kerchunk.extract_array_refs(refs, var_name)
chunk_dict, zarray, zattrs = kerchunk.parse_array_refs(arr_refs)

manifest = ChunkManifest._from_kerchunk_chunk_dict(chunk_dict)

# we want to remove the _ARRAY_DIMENSIONS from the final variables' .attrs
dims = zattrs.pop("_ARRAY_DIMENSIONS")

varr = virtual_array_class(zarray=zarray, chunkmanifest=manifest)
if chunk_dict:
manifest = ChunkManifest._from_kerchunk_chunk_dict(chunk_dict)
varr = virtual_array_class(zarray=zarray, chunkmanifest=manifest)
else:
# This means we encountered a scalar variable of dimension 0,
# very likely that it actually has no numeric value and its only purpose
# is to communicate dataset attributes.
varr = zarray.fill_value

return xr.Variable(data=varr, dims=dims, attrs=zattrs)

Expand Down

0 comments on commit 3b091e8

Please sign in to comment.