diff --git a/lib/ncdata/dataset_like.py b/lib/ncdata/dataset_like.py index 3702f95..b92ffa9 100644 --- a/lib/ncdata/dataset_like.py +++ b/lib/ncdata/dataset_like.py @@ -81,6 +81,11 @@ def __setattr__(self, attr, value): else: self.setncattr(attr, value) + def set_auto_chartostring(self, value): + if bool(value): + msg = "Cannot enable 'auto_chartostring' for Nc4Dataselike or Nc4VariableLike." + raise ValueError() + class Nc4DatasetLike(_Nc4DatalikeWithNcattrs): """ diff --git a/lib/ncdata/netcdf4.py b/lib/ncdata/netcdf4.py index b70a928..fd34b87 100644 --- a/lib/ncdata/netcdf4.py +++ b/lib/ncdata/netcdf4.py @@ -10,13 +10,17 @@ from typing import Dict, Optional, Union import dask.array as da +import dask.config +import dask.utils import netCDF4 as nc import numpy as np from . import NcData, NcDimension, NcVariable -__all__ = ["from_nc4", "to_nc4"] +__all__ = ["from_nc4", "to_nc4", "ASSUMED_TYPICAL_STRINGLENGTH"] +#: A generally safe 'guess' at the size of variable-length strings, for chunking decisions. +ASSUMED_TYPICAL_STRINGLENGTH = 500 # The variable arguments which are 'claimed' by the existing to_nc4 code, and # therefore not valid to appear in the 'var_kwargs' parameter. @@ -60,9 +64,12 @@ def _to_nc4_group( ) raise ValueError(msg) + nc4_dtype = var.dtype + data_is_strs = nc4_dtype.kind == "O" + nc4var = nc4object.createVariable( varname=varname, - datatype=var.dtype, + datatype=str if data_is_strs else var.dtype, dimensions=var.dimensions, fill_value=fill_value, **kwargs, @@ -78,6 +85,10 @@ def _to_nc4_group( nc4var.setncattr(attrname, attrval) data = var.data + if data_is_strs: + # cast the data : this means using string-repr of everything + data = data.astype("U") # convert to string-arrays, if not already + if hasattr(data, "compute"): da.store(data, nc4var) else: @@ -209,6 +220,37 @@ def to_nc4( nc4ds.close() +def fix_varlen_string_chunks(chunks: list[int | str], shape: list[int]): + """Choose chunks for a variable containing variable-length strings. + + Given the shape and requested chunks (a list of user-provided numbers and ['auto']), + replace all 'auto's in the input chunks list with fixed numbers. + + Aim at the current (configured) dask default chunk-size. + Assume that variable strings may be of size ASSUMED_TYPICAL_STRINGLENGTH. + """ + chunksize = dask.config.get("array.chunk-size") + chunk_bytes = dask.utils.parse_bytes(chunksize) + flex_dims = [ + i_dim for i_dim, n_dim in enumerate(shape) if chunks[i_dim] == "auto" + ] + given_multiple = np.prod( + [1] + + [dim for i_dim, dim in enumerate(chunks) if i_dim not in flex_dims] + ) + remaining_multiple = max( + 1, chunk_bytes // ASSUMED_TYPICAL_STRINGLENGTH // given_multiple + ) + for i_flexdim in flex_dims[::-1]: + # Replace all 'auto's with numbers: fill trailing dims first, leave 1s in remaining (leading) dims + this_multiple = remaining_multiple + this_dim = shape[i_flexdim] + if this_multiple > this_dim: + this_multiple = this_dim + remaining_multiple //= this_multiple + chunks[i_flexdim] = this_multiple + + def _from_nc4_group(nc4ds: Union[nc.Dataset, nc.Group], dim_chunks) -> NcData: """ Inner routine for :func:`from_nc4`. @@ -230,10 +272,15 @@ def _from_nc4_group(nc4ds: Union[nc.Dataset, nc.Group], dim_chunks) -> NcData: ncdata.dimensions[dimname] = NcDimension(dimname, size, unlimited) for varname, nc4var in nc4ds.variables.items(): + ncdata_var_dtype = nc4var.dtype + data_is_str = ncdata_var_dtype is str + if data_is_str: + # For variable-length strings, actual exchanged data is an object array. + ncdata_var_dtype = "O" var = NcVariable( name=varname, dimensions=nc4var.dimensions, - dtype=nc4var.dtype, + dtype=ncdata_var_dtype, group=ncdata, ) ncdata.variables[varname] = var @@ -255,12 +302,17 @@ def _from_nc4_group(nc4ds: Union[nc.Dataset, nc.Group], dim_chunks) -> NcData: proxy = _NetCDFDataProxy( shape=shape, - dtype=var.dtype, + dtype=ncdata_var_dtype, filepath=parent_ds.filepath(), variable_name=varname, group_names_path=group_names_path, ) chunks = [dim_chunks.get(name, "auto") for name in var.dimensions] + if data_is_str: + # payload is variable length strings + # so data sizes are unknown and standard "auto" chunking fails: make our own assumptions + calculate + fix_varlen_string_chunks(chunks, shape) + var.data = da.from_array( proxy, chunks=chunks, asarray=True, meta=np.ndarray ) diff --git a/tests/data_testcase_schemas.py b/tests/data_testcase_schemas.py index 7622da3..ce68460 100644 --- a/tests/data_testcase_schemas.py +++ b/tests/data_testcase_schemas.py @@ -525,6 +525,7 @@ def standard_testcase(request, session_testdir): # """ "small_rotPole_precipitation", "small_FC_167", + "test_monotonic_coordinate", ], # Xarray can save ~anything "save": [r"test_monotonic_coordinate"], diff --git a/tests/integration/test_iris_xarray_roundtrips.py b/tests/integration/test_iris_xarray_roundtrips.py index d0eb2bf..ae4ffd5 100644 --- a/tests/integration/test_iris_xarray_roundtrips.py +++ b/tests/integration/test_iris_xarray_roundtrips.py @@ -79,9 +79,19 @@ def test_roundtrip_ixi(standard_testcase, use_irislock, adjust_chunks): "testdata____ugrid__21_triangle_example", # Problem with units on time bounds "label_and_climate__small_FC_167", - # Broken UGRID files now won't load in Iris >= 3.10 - "unstructured_grid__mesh_C12", - "_unstructured_grid__theta_nodal_xios", + # # Broken UGRID files now won't load in Iris >= 3.10 + # "unstructured_grid__mesh_C12", + # "_unstructured_grid__theta_nodal_xios", + # **No** mesh files will now roundtrip, since xarray 2026.04.0. + # This is because xarray now normalises the array type when loading + # from ncdata, removing masks and converting to nanarrays (thus making ints + # into floats). + # Previously we could create xarray Datasets with masked integer variables + # as data -- though that was alwaysobviously a bit tricksy. + # Iris can't accept the new form as it insists on integer types for + # mesh connectivities :-( + # TODO: do something to get this working again? + "unstructured", ] ) if any(key in standard_testcase.name for key in exclude_case_keys): @@ -177,7 +187,11 @@ def test_roundtrip_ixi(standard_testcase, use_irislock, adjust_chunks): # FOR NOW: compare with experimental ncdata comparison. # I know this is a bit circular, but it is useful for debugging, for now ... result = dataset_differences( - from_iris(iris_cubes), from_iris(iris_xr_cubes) + from_iris(iris_cubes), + from_iris(iris_xr_cubes), + check_attrs_order=False, + check_dims_order=False, + check_vars_order=False, ) assert result == [] diff --git a/tests/unit/netcdf/test_fix_varlen_string_chunks.py b/tests/unit/netcdf/test_fix_varlen_string_chunks.py new file mode 100644 index 0000000..ad6ee6b --- /dev/null +++ b/tests/unit/netcdf/test_fix_varlen_string_chunks.py @@ -0,0 +1,85 @@ +""" +Tests for :func:`ncdata.netcdf.fix_varlen_string_chunks`. + +""" + +import dask.config +import pytest + +from ncdata.netcdf4 import fix_varlen_string_chunks + + +class TestFixVarlenStringChunks: + @staticmethod + def run(chunks, shape): + chunks_list = list(chunks) + fix_varlen_string_chunks(chunks=chunks_list, shape=shape) + return chunks_list + + def test_noauto(self): + shape = (1, 3, 5, 4) + chunks = (1, 2, 3, 4) + result = self.run(chunks, shape) + assert result == list(chunks) + + def test_fullshape_auto(self): + shape = (1, 3, 5, 4) + chunks = ["auto"] * 4 + result = self.run(chunks, shape) + assert result == list(shape) + + def test_scalar(self): + shape = () + chunks = [] + result = self.run(chunks, shape) + assert result == list(shape) + + def test_auto_last(self): + shape = (1, 3, 5, 14) + chunks = (1, 2, 3, "auto") + result = self.run(chunks, shape) + assert result == [1, 2, 3, 14] + + def test_auto_first(self): + shape = (21, 33, 45, 14) + chunks = ("auto", 2, 3, 4) + result = self.run(chunks, shape) + assert result == [21, 2, 3, 4] + + def test_auto_penult(self): + shape = (1, 3, 5, 14) + chunks = (1, 2, "auto", 4) + result = self.run(chunks, shape) + assert result == [1, 2, 5, 4] + + @pytest.fixture() + def varstr_len_100(self, mocker): + # Allow 100 characters(bytes) in variable-length strings + mocker.patch("ncdata.netcdf4.ASSUMED_TYPICAL_STRINGLENGTH", 100) + yield + + def test_sizelimit_fullauto(self, varstr_len_100): + shape = (10, 12) + chunks = ("auto", "auto") + + # set chunksize to match 40 strings + with dask.config.set({"array.chunk-size": "4000b"}): + result = self.run(chunks, shape) + assert result == [3, 12] + + def test_sizelimit_multispread(self, varstr_len_100): + shape = (10, 12, 3, 3) + chunks = ("auto", "auto", "auto", "auto") + # set chunksize to match 40 strings + with dask.config.set({"array.chunk-size": "4000b"}): + result = self.run(chunks, shape) + assert result == [1, 4, 3, 3] + + def test_sizelimit_mixed(self, varstr_len_100): + shape = (5, 10, 4, 6) + chunks = ("auto", "auto", 2, "auto") + + # set chunksize to match 40 strings + with dask.config.set({"array.chunk-size": "4000b"}): + result = self.run(chunks, shape) + assert result == [1, 3, 2, 6] diff --git a/tests/unit/netcdf/test_from_nc4.py b/tests/unit/netcdf/test_from_nc4.py index d19703d..7bce80d 100644 --- a/tests/unit/netcdf/test_from_nc4.py +++ b/tests/unit/netcdf/test_from_nc4.py @@ -11,6 +11,8 @@ from pathlib import Path +import dask.array as da +import dask.config import netCDF4 as nc import numpy as np import pytest @@ -111,3 +113,62 @@ def test_target_types(sourcetype, tmp_path): diffs = dataset_differences(ncdata, ncdata_expected) assert diffs == [] + + +class TestVarStrs: + def test_load_vlenstrs_basic(self, tmp_path): + varstr_test_spec = { + "dims": [dict(name="x", size=3)], + "vars": [ + dict( + name="var_0", + dims=["x"], + dtype=str, + data=np.array(["one", "two", "three"], dtype="U10"), + ), + ], + } + filepath = tmp_path / "testinput_basic.nc" + ncds = file_and_ncdata_from_spec(filepath, varstr_test_spec) + var = ncds.variables["var_0"] + assert var.dtype == "O" + data = var.data + assert isinstance(data, da.Array) + assert data.shape == (3,) + assert data.dtype == "O" + values = data.compute() + assert values.shape == (3,) + assert values.dtype == "O" + expect = np.array(["one", "two", "three"], dtype="O") + assert np.all(values == expect) + + def test_load_large_chunks(self, tmp_path, mocker): + # NB the input string array to create the file, via make_testcase_dataset, is + # NOT an object-array, as that's not how netCDF4 creates a 'str' type variable. + nparray_onestr = np.array(["this"], dtype="U10").reshape((1, 1)) + darr_onestr = da.from_array(nparray_onestr, chunks=1) + # expand to get a big array of (100 x 100) strings. + big_string_array, _ = da.broadcast_arrays( + darr_onestr, da.zeros((100, 100)) + ) + varstr_test_spec = { + "dims": [dict(name="x", size=100), dict(name="y", size=100)], + "vars": [ + dict( + name="var_1", + dims=["y", "x"], + dtype=str, + data=big_string_array, + ), + ], + } + filepath = tmp_path / "testinput_largearr.nc" + with dask.config.set({"array.chunk-size": "4000b"}): + ncds = file_and_ncdata_from_spec(filepath, varstr_test_spec) + + var = ncds.variables["var_1"] + assert var.dtype == "O" + data = var.data + assert isinstance(data, da.Array) + assert data.shape == (100, 100) + assert data.chunksize == (1, 8) diff --git a/tests/unit/netcdf/test_to_nc4.py b/tests/unit/netcdf/test_to_nc4.py index 85c8589..f5f64ca 100644 --- a/tests/unit/netcdf/test_to_nc4.py +++ b/tests/unit/netcdf/test_to_nc4.py @@ -16,7 +16,7 @@ import numpy as np import pytest -from ncdata import NcData +from ncdata import NcData, NcDimension, NcVariable from ncdata.netcdf4 import from_nc4, to_nc4 from ncdata.utils import dataset_differences @@ -167,3 +167,43 @@ def test_var_kwargs__bad_kwarg(tmp_path): ) with pytest.raises(ValueError, match=expected_msg_regex): to_nc4(ncdata, output_path, var_kwargs=var_kwargs) + + +class TestVarStrs: + def test_basic_varstr(self, tmp_path): + string_data_array = np.array(["one", "two", "three"], dtype="O") + varstr_var = NcVariable("vx", ["x"], data=string_data_array) + ds = NcData(dimensions=[NcDimension("x", 3)], variables=[varstr_var]) + test_filepath = tmp_path / "test_save_basic_varstr.nc" + to_nc4(ds, test_filepath) + with nc.Dataset(test_filepath) as ds: + var = ds.variables["vx"] + assert var.shape == (3,) + assert var.dtype is str + values_arr = var[:] + + assert values_arr.dtype == "O" + expected = np.array(["one", "two", "three"], dtype="O") + assert np.all(values_arr == expected) + + def test_odd_objectarray(self, tmp_path): + # Generally objects content is uniformly *treated* as strings on output + # - the array elements all have 'str()' applied + string_data_array = np.array( + ["one", {"this": "yes", "that": 0}, None, 7], dtype="O" + ) + varstr_var = NcVariable("vx", ["x"], data=string_data_array) + ds = NcData(dimensions=[NcDimension("x", 4)], variables=[varstr_var]) + test_filepath = tmp_path / "test_save_objarr.nc" + to_nc4(ds, test_filepath) + with nc.Dataset(test_filepath) as ds: + var = ds.variables["vx"] + assert var.shape == (4,) + assert var.dtype is str + values_arr = var[:] + + assert values_arr.dtype == "O" + expected = np.array( + ["one", "{'this': 'yes', 'that': 0}", "None", "7"], dtype="O" + ) + assert np.all(values_arr == expected)