apache / arrow-nanoarrow

Helpers for Arrow C Data & Arrow C Stream interfaces
https://arrow.apache.org/nanoarrow
Apache License 2.0
149 stars 34 forks source link

feat(python): Implement array from buffer for non-CPU arrays #550

Open paleolimbot opened 5 days ago

paleolimbot commented 5 days ago

Requires building with (e.g.) export NANOARROW_PYTHON_CUDA=/usr/local/cuda and a cupy install:

import cupy as cp
import nanoarrow as na
from nanoarrow import device
from nanoarrow.c_buffer import CBuffer

x_gpu = cp.array([1, 2, 3])

na.c_array_from_buffers(
    na.int64(),
    3,
    [None, CBuffer.from_dlpack(x_gpu)],
    move=True,
    device=device.resolve(device.DeviceType.CUDA.value, 0)
)

#> <nanoarrow.device.CDeviceArray>
#> - device_type: CUDA <2>
#> - device_id: 0
#> - array: <nanoarrow.c_array.CArray int64>
#>   - length: 3
#>   - offset: 0
#>   - null_count: 0
#>   - buffers: (0, 81681973248)
#>   - dictionary: NULL
#>   - children[0]: