pydata / pydata/xarray

Add ability to change underlying array type

Open
#4,234 12 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

topic-arrays
Dominant language
Python
Stars
4.2k
Forks
1.4k
Avg merge
2d 15h
Merged PRs (30d)
14

Description

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

In order to use Xarray with alternative array types like cupy the user needs to be able to specify the underlying array type without digging into internals.

Right now I'm doing something like this.

import xarray as xr
import cupy as cp

ds = xr.tutorial.load_dataset("air_temperature")
ds.air.data = cp.asarray(ds.air.data)

However this will become burdensome when there are many data arrays and feels brittle and prone to errors.

As I see it a conversion could instead be done in a couple of places; on load, or as a utility method.

Currently Xarray supports NumPy and Dask array well. Numpy is the defrault and the way you specify whether a Dask array should be used is to give the chunks kwargs to an open_ function or by calling .chunk() on a DataSet or DataArray.

Side note: There are a few places where the Dask array API bleeds into Xarray in order to have compatibility, the chunk kwarg/method is one, the .compute() method is another. I'm hesitant to do this for other array types, however surfacing the cupy.ndarray.get method could feel natural for cupy users. But for now I think it would be best to take Dask as a special case and try and be generic for everything else.

Describe the solution you'd like

For other array types I would like to propose the addition of an asarray kwarg for the open_ methods and an .asarray() method on DataSet and DataArray. This should take either the array type cupy.ndarray, the asarray method cp.asarray, or preferably either.

This would result in something like the following.

import xarray as xr
import cupy as cp

ds = xr.open_mfdataset("/path/to/files/*.nc", asarray=cp.ndarray)

# or

ds = xr.open_mfdataset("/path/to/files/*.nc")
gds = ds.asarray(cp.ndarray)

These operations would convert all data arrays to cupy arrays. For the case that ds is backed by Dask arrays it would use map_blocks to cast each block to the appropriate array type.

It is still unclear what to do about index variables, which are currently of type pandas.Index. For cupy it may be more appropriate to use a cudf.Index instead to ensure both are on the GPU. However this would add a dependency on cudf and potentially increase complexity here.

Describe alternatives you've considered

Instead of an asarray kwarg/method something like to_cupy/from_cupy could be done. However I feel this makes less sense because the object type is not changing, just that of the underlying data structure.

Another option would be to go more high level with it. For example a gpu kwarg and to_gpu/from_gpu method could be added in the same way. This would abstract things even further and give users a choice about hardware rather than software. This would also be a fine solution but I think it may special case too much and a more generic solution would be better.

Additional context
Related to #4212.

I'm keen to start implementing this. But would like some discussion/feedback before I dive in here.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reviewing the discussion on this issue and related issue #4212; the payload names no implementation files, tests, or entry points. The work is not ready until the API design is agreed, including conversion of eager and Dask-backed arrays and the treatment of index variables.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, pandas, python
Domain
data
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.