Concatenating datasets with variable length strings materializes them in memory

[tomwhite]

I want to concatenate a collection of datasets with a variable that contains variable length strings and save them to Zarr. The strings are mostly very short (1-2 characters), but a few can be much longer (thousands of characters), and can be of arbitrary length, so not suitable for a fixed size representation.

Zarr's VLenUTF8 is ideal for this. It is recommended on http://xarray.pydata.org/en/latest/user-guide/io.html#zarr, where it says "To store variable length strings, convert them to object arrays first with dtype=object".

The datasets to be concatenated are produced in parallel, and the size of each is not known in advance.

Here's the MVCE:

import numpy as np
import xarray as xr

# Create a dataset of strings with object dtype, and save to Zarr (this uses VLenUTF8 encoding)
ds1 = xr.Dataset(dict(foo=("x", np.array(["a", "bbbbb"], dtype="O"))))
ds1.to_zarr("ds1.zarr")

# Create another dataset
ds2 = xr.Dataset(dict(foo=("x", np.array(["cc", "dddddddd", "e"], dtype="O"))))
ds2.to_zarr("ds2.zarr")

# Load the datasets from Zarr, and concatenate them
datasets = [xr.open_zarr(f, concat_characters=False) for f in ("ds1.zarr", "ds2.zarr")]

ds = xr.concat(datasets, dim="x", data_vars="minimal")
for var in ds.variables:  # avoid round-tripping problem in https://github.com/pydata/xarray/issues/3476
    ds[var].encoding.clear()

ds = ds.chunk({"x": 2}) # rechunk since Zarr requires fixed chunk sizes

# Save as a single Zarr file
ds.to_zarr("ds.zarr")

This produces a warning, because it has loaded the entire ds.foo variable into memory to determine its type. Obviously for large datasets this is not scalable.

SerializationWarning: variable None has data in the form of a dask array with dtype=object, which means it is being loaded into memory to determine a data type that can be safely stored on disk. To avoid this, coerce this variable to a fixed-size dtype with astype() before saving it.

So, my question is: how can I concatenate datasets with variable length strings without materalizing them in memory?

[tomwhite]

As a workaround I tried to coerce the variable to the vlen dtype that Xarray uses internally, but this had no effect:

VLEN_DTYPE = xr.coding.strings.create_vlen_dtype(str)
ds["foo"] = ds["foo"].astype(VLEN_DTYPE)

I think this is because VLEN_DTYPE and NumPy's object dtype compare equal so Dask treats the astype as a no op:

>>> import xarray as xr
>>> VLEN_DTYPE = xr.coding.strings.create_vlen_dtype(str)
>>> np.dtype("O") == VLEN_DTYPE
True