pydata / pydata/xarray

Please expose __cuda_array_interface__ via the xarray.__array__() function if present

Open
#6,847 5 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement
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?

When using an array type with GPU support, such as CuPy arrays, Numba device arrays or Numba mapped arrays (shared), __cuda_array_interface__ is not exposed by the xarray.__array() function.

I'm using large NetCDF files which I wish to process against reference dataframes and use GPU acceleration to do this.

For example, Numba mapped array:

>>> points = np.random.randn(2, 3)
>>> map_points = nb.cuda.mapped_array_like(points)
>>> map_points.__array_interface__
{'data': (140399865758208, False),
 'strides': None,
 'descr': [('', '<f8')],
 'typestr': '<f8',
 'shape': (2, 3),
 'version': 3}

>>> map_points.__cuda_array_interface__
{'shape': (2, 3),
 'strides': None,
 'data': (140399865758208, False),
 'typestr': '<f8',
 'stream': None,
 'version': 3}

When copied to xarray:

>>> data = xr.DataArray(map_points, dims=("x", "y"), coords={"x": [10, 20]})
>>> data
xarray.DataArray x: 2y: 3
array([[0., 0., 0.],
       [0., 0., 0.]])
Coordinates:
x (x) int64 10 20
Attributes: (0)

Array interface confirms same address for the base (CPU) array as above, i.e. Zero Copy

>>> data.__array__().__array_interface__
{'data': (140399865758208, False),
 'strides': None,
 'descr': [('', '<f8')],
 'typestr': '<f8',
 'shape': (2, 3),
 'version': 3}

However the __cuda_array_interface__ is not exposed

>>> data.__array__().__cuda_array_interface__
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Input In [23], in <cell line: 1>()
----> 1 data.__array__().__cuda_array_interface__

AttributeError: 'numpy.ndarray' object has no attribute '__cuda_array_interface__'
Describe the solution you'd like

Expose __cuda_array_interface__ via the xarray.__array() function so it is available to CuPy and Numba CUDA functions.

Describe alternatives you've considered

As a workaround, I'm not using xarray for NetCDF files. Instead I'm converting them into an dictionary of arrays which provides me with the GPU interfaces.

Additional context

No response

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 at the xarray.array entry point and trace how wrapped array types are converted, comparing the shown NumPy and CUDA array-interface behavior. Done means a CuPy- or Numba-style input preserves cuda_array_interface through array, with regression coverage for that behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, python
Domain
data, performance
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
42/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.