Function test_non_cpu_array

python/pyarrow/tests/test_array.py:4436–4518

()

def test_non_cpu_array():
    cuda = pytest.importorskip("pyarrow.cuda", exc_type=ImportError)
    ctx = cuda.Context(0)

    data = np.arange(4, dtype=np.int32)
    validity = np.array([True, False, True, False], dtype=np.bool_)
    cuda_data_buf = ctx.buffer_from_data(data)
    cuda_validity_buf = ctx.buffer_from_data(validity)
    arr = pa.Array.from_buffers(pa.int32(), 4, [None, cuda_data_buf])
    arr2 = pa.Array.from_buffers(pa.int32(), 4, [None, cuda_data_buf])
    arr_with_nulls = pa.Array.from_buffers(
        pa.int32(), 4, [cuda_validity_buf, cuda_data_buf])

    # Supported
    arr.validate()
    assert arr.offset == 0
    assert arr.buffers() == [None, cuda_data_buf]
    assert arr.device_type == pa.DeviceAllocationType.CUDA
    assert arr.is_cpu is False
    assert len(arr) == 4
    assert arr.slice(2, 2).offset == 2
    assert repr(arr)
    assert str(arr)

    # TODO support DLPack for CUDA
    with pytest.raises(NotImplementedError):
        arr.__dlpack__()
    with pytest.raises(NotImplementedError):
        arr.__dlpack_device__()

    # Not Supported
    with pytest.raises(NotImplementedError):
        arr.diff(arr2)
    with pytest.raises(NotImplementedError):
        arr.cast(pa.int64())
    with pytest.raises(NotImplementedError):
        arr.view(pa.int64())
    with pytest.raises(NotImplementedError):
        arr.sum()
    with pytest.raises(NotImplementedError):
        arr.unique()
    with pytest.raises(NotImplementedError):
        arr.dictionary_encode()
    with pytest.raises(NotImplementedError):
        arr.value_counts()
    with pytest.raises(NotImplementedError):
        arr_with_nulls.null_count
    with pytest.raises(NotImplementedError):
        arr.nbytes
    with pytest.raises(NotImplementedError):
        arr.get_total_buffer_size()
    with pytest.raises(NotImplementedError):
        [i for i in iter(arr)]
    with pytest.raises(NotImplementedError):
        arr == arr2
    with pytest.raises(NotImplementedError):
        arr.is_null()
    with pytest.raises(NotImplementedError):