Inferring concatenation order from coordinate data values

[TomNicholas]

If xr.combine_by_coords is used, then indexes must be created for dimension coordinates in order to use them to infer the correct order of concatenation. This requires loading of the data for that coordinate in order to create the index. We can do this, but I think we will then end up with a Dataset which contains some ManifestArrays but also some numpy arrays for the dimension coordinates. This isn't itself a problem - we have that information, we can load it from the files.

The tricky part is that we presumably don't want to write those actual data values back out during the serialization step, we want to only write the chunk manifest for each dimension coordinate. How can we regain access to this information, or carry it around ready for serialization at the end?

[TomNicholas]

My understanding (/ my recollection 😅 ) of the code in xr.combine_by_coords is that it infers concatenation order from the pandas indexes backing any 1D coordinate variables. So as long as those indexes are present and correct, the datasets can be ordered. So perhaps we could just manually create those indexes when we create the dataset (which requires loading data from the files), so the indexes are eager (i.e. how they normally work) but the corresponding array data is virtualized.

This would mean we would need to both read the kerchunk references and actually read data values from the file at dataset construction time. We would want to have a function flag to opt in to this behaviour.

[TomNicholas]

Okay I prematurely closed this, as it was not actually as simple as it seemed.

perhaps we could just manually create those indexes when we create the dataset (which requires loading data from the files), so the indexes are eager (i.e. how they normally work) but the corresponding array data is virtualized.

This would mean we would need to both read the kerchunk references and actually read data values from the file at dataset construction time.

This is what I implemented in #52, but I didn't remember to check this part:

The tricky part is that we presumably don't want to write those actual data values back out during the serialization step, we want to only write the chunk manifest for each dimension coordinate.

Turns out that doesn't work yet - at some point within xr.combine_by_coords the coordinate variable is being recreated from the index as a numpy array, even though the original coordinate variable was a virtual ManifestArray. Hopefully this can be fixed with another change to xarray upstream similar to pydata/xarray#8872.

In the meantime I should remove that section from the docs so as not to commit false advertising! EDIT: Removed in TomNicholas@2c5be3f

cc @norlandrhagen

[benbovy]

So perhaps we could just manually create those indexes when we create the dataset (which requires loading data from the files), so the indexes are eager (i.e. how they normally work) but the corresponding array data is virtualized.

In general I wonder if it would be useful to have some custom Xarray index here. Basically a simple subclass or a very lightweight wrapper around xarray.indexes.PandasIndex where the .create_variables() method is re-implemented such that the indexed coordinate variable doesn't wrap the pandas.Index object but returns a variable still wrapping the original ManifestArray (IIUC).

we presumably don't want to write those actual data values back out during the serialization step

Would such decoupling of the pandas.Index from the indexed coordinate data solve this?

The tricky part is that we still need to refactor Xarray IndexVariable, which is still too tied to pandas.Index and currently prevents implementing lazy custom indexes (see pydata/xarray#8124). Perhaps as a temporary workaround we could use a regular Variable for the indexed coordinate instead of an IndexVariable...

[dcherian]

we presumably don't want to write those actual data values back out during the serialization step,

TWriting would speed things up reading from the consolidated files. Kerchunk calls this "inlining". For e.g. you wouldn't want read ~90k time steps from 90k files to construct a 90K long time coordinate

[TomNicholas]

a variable still wrapping the original ManifestArray (IIUC).

Would such decoupling of the pandas.Index from the indexed coordinate data solve this?

Yeah that sounds like what we need! I basically need the ManifestArray inside the coordinate variable to still exist after xarray concat etc. operations, so that I can write it to disk. So I need the coordinate data to still wrap the original duck array, without having coerced it or dropped it in favour of a numpy array created from the pd.Index.

Perhaps as a temporary workaround we could use a regular Variable for the indexed coordinate instead of an IndexVariable...

I think this is what I was hoping I could achieve with pydata/xarray#8872, but clearly I was missing some subtleties with IndexVariable vs Variable.

In general I wonder if it would be useful to have some custom Xarray index here.

Naively, I don't really understand why this use case requires a custom Xarray index. I am still happy to use PandasIndex for the index class (i.e. not custom), I just want my wrapped coordinate variable data to not be messed with.

[TomNicholas]

TWriting would speed things up reading from the consolidated files. Kerchunk calls this "inlining". For e.g. you wouldn't want read ~90k time steps from 90k files to construct a 90K long time coordinate

@dcherian if we have retained both the index and the ManifestArray, wouldn't the choice of whether or not to "inline" be something that can be deferred to the write step?

[dcherian]

IIUC yes.