Xarray integration

[cosmicBboy]

Is your feature request related to a problem? Please describe.

xarray is a project that provides a dict-like data container abstraction for ndimensional arrays. It shares some commonalities with pandas, but there many key differences (e.g. coords and attrs).

After chatting with @jhamman about this approach, we decided it would be appropriate to park xarray-schema within the pandera codebase. This issue tracks the planned integration of xarray-schema into the pandera codebase.

Describe the solution you'd like

A good start for this integration is to add a pandera.xarray module exposing the schema and schema component classes specific to xarray:

import numpy as np
import xarray as xr
from pandera.xarray import DataArraySchema, DatasetSchema

da = xr.DataArray(np.ones(4, dtype='i4'), dims=['x'], name='foo')

schema = DataArraySchema(dtype=np.integer, name='foo', shape=(4, ), dims=['x'])

schema.validate(da)

TODO

• Abstract out validation logic to support non-pandas dataframes, e.g. spark, dask, etc #381
• move xarray-schema codebase into pandera.xarray
• TBD adjustments to xarray schema/schema components to fit pandera's validation model.

Describe alternatives you've considered

The main alternative to this integration is to keep xarray-schema as a separate project that's interoperable with pandera. However, given that pandera plans on expanding its scope to validate data containers beyond pandas, it would benefit this project to maintain schema interfaces for multiple (not just pandas-like) data container libraries.

Additional context

• Abstract out validation logic to support non-pandas dataframes, e.g. spark, dask, etc #381

[leroyvn]

Hi, I've been looking for an xarray validation library for a while now and I was wondering: is integration in Pandera still planned? Thanks!

[avcopan]

In reply to @cosmicBboy's request on Discord, here is a simple xarray.DataArray schema that I would like to validate using Pandera:

class RateConstant(DataArrayModel):
    # coords:
    pressure  # sequence of floats: [0.1, 1.0, 10.0, ...]
    temperature  # sequence of floats: [500, 600, 700, ...]

What I want to be able to do is:

1. Validate that the xarray has coordinates names "pressure" and "temperature" (actually checking their types is low priority). Example:

# create the DataArray:
ktp = xarray.DataArray(
    data=[[1e1, 1e2, 1e3, 1e4], [1e5, 1e6, 1e7, 1e8], [1e9, 1e10, 1e11, 1e12]],
    coords={
        "pressure": [1, 10, numpy.inf],
        "temperature": [1000, 1500, 2000, 2500],
    }
)

# validate the DataArray:
ktp = RateConstant.validate(ktp)

2. Use RateConstant.pressure and RateConstant.temperature to access these coordinate names from the schema (to avoid typos and allow me to easily update the code that depends on the schema, if needed). Example:

>>> print(RateConstant.temperature)
"temperature"

Note that this is for a simple DataArray, not a Dataset.

[cosmicBboy]

@avcopan thank you! just to follow up, I'm not super familiar with xarray, so it would it help a lot if you could write out the xarray-native assertions that you would write for the following validations:

Validate that the xarray has coordinates names "pressure" and "temperature" (actually checking their types is low priority). Example:
Use RateConstant.pressure and RateConstant.temperature to access these coordinate names from the schema (to avoid typos and allow me to easily update the code that depends on the schema, if needed). Example:

I can then start to map those over to the pandera api

[avcopan]

Oh, sure!

With the above example, you would just do:

assert "pressure" in ktp.coords
assert "temperature" in ktp.coords

[avcopan]

<DataArray object>.coords acts essentially like a dictionary mapping coordinate keys (e.g. "temperature" and "pressure") to their values.

[bmwilly]

Hey, this would be great! I think it would be useful for lots of people and projects.

I'm curious if this is still planned? Especially since xarray-schema seems to be fairly dead (last commit ~3 years ago).

Thanks!

[cosmicBboy]

hi @bmwilly ! xarray support is still on the table. Can't promise a concrete timeline, but basically I have a couple of things on my plate:

• supporting pandas 3.0: Support pandas 3.0 #2214
• kicking off the narwhals effort: [🐋 narwhals] create stub modules, classes, and functions for narwhals integration #2081. This will take some time too, but will be a parallel effort to replace and standardize the backends for all the existing supported dataframe libraries.

I've been starting to use agents to help out with development and maintenance efforts (with babysitting and review of course), so can probably start moving on this sooner rather than later (at least iterating on an initial implementation).

Just so I get an intuition for the design of this @bmwilly would you say that the https://github.com/xarray-contrib/xarray-schema is close to what you want in terms of developer experience? E.g. I see it exposes a lot of different schema components:

from xarray_schema.components import (
    DTypeSchema,
    DimsSchema,
    ShapeSchema,
    NameSchema,
    ChunksSchema,
    ArrayTypeSchema,
    AttrSchema,
    AttrsSchema
)

My feels a little heavy to me, I'd much prefer limiting components and reducing them to the equivalents of the DataFrame, Series, and Index and all other properties are kwargs for a small number of central objects. Do you have thoughts on this?

[bmwilly]

@cosmicBboy thanks for the quick response and for soliciting feedback :)

would you say that the https://github.com/xarray-contrib/xarray-schema is close to what you want

I agree with you and think they expose too many low-level components. Something analogous to pandera's pandas support would be preferred — i.e., schemas only for xr.Dataset, xr.DataArray, and xr.DataTree.

(Note I'm answering this as someone who hasn't really used xarray-schema so take the above with a grain of salt).

[leroyvn]

In case that's useful, I forked xarray-schema (xarray-validate) to add some features that were missing for my use case (e.g. error collection). The many-component approach turned out to be a bit heavy but useful for unit testing. I preferred to expose a dict-based interface, more concise but more difficult to debug in case of an incorrect schema. I don't know Pandera's architecture, so I'm not sure how relevant that could be.