Function test_recordbatch_non_cpu

python/pyarrow/tests/test_table.py:3817–3974

(cuda_context, cpu_recordbatch, cuda_recordbatch,
                             cuda_arrays, schema)

def test_recordbatch_non_cpu(cuda_context, cpu_recordbatch, cuda_recordbatch,
                             cuda_arrays, schema):
    verify_cuda_recordbatch(cuda_recordbatch, expected_schema=schema)
    N = cuda_recordbatch.num_rows

    # shape test
    assert cuda_recordbatch.shape == (5, 2)

    # columns() test
    assert len(cuda_recordbatch.columns) == 2

    # add_column(), set_column() test
    for fn in [cuda_recordbatch.add_column, cuda_recordbatch.set_column]:
        col = pa.array([-2, -1, 0, 1, 2], pa.int8()
                       ).copy_to(cuda_context.memory_manager)
        new_batch = fn(2, 'c2', col)
        verify_cuda_recordbatch(
            new_batch, expected_schema=schema.append(pa.field('c2', pa.int8())))
        err_msg = ("Got column on device <DeviceAllocationType.CPU: 1>, "
                   "but expected <DeviceAllocationType.CUDA: 2>.")
        with pytest.raises(TypeError, match=err_msg):
            fn(2, 'c2', [1] * N)

    # remove_column() test
    new_batch = cuda_recordbatch.remove_column(1)
    verify_cuda_recordbatch(new_batch, expected_schema=schema.remove(1))

    # drop_columns() test
    new_batch = cuda_recordbatch.drop_columns(['c1'])
    verify_cuda_recordbatch(new_batch, expected_schema=schema.remove(1))
    empty_batch = cuda_recordbatch.drop_columns(['c0', 'c1'])
    assert len(empty_batch.columns) == 0
    assert empty_batch.device_type == pa.DeviceAllocationType.CUDA

    # select() test
    new_batch = cuda_recordbatch.select(['c0'])
    verify_cuda_recordbatch(new_batch, expected_schema=schema.remove(1))

    # cast() test
    new_schema = pa.schema([pa.field('c0', pa.int64()), pa.field('c1', pa.int64())])
    with pytest.raises(NotImplementedError):
        cuda_recordbatch.cast(new_schema)

    # drop_null() test
    null_col = pa.array([1] * N, mask=[True, False, True, False, True]).copy_to(
        cuda_context.memory_manager)
    cuda_recordbatch_with_nulls = cuda_recordbatch.add_column(2, 'c2', null_col)
    with pytest.raises(NotImplementedError):
        cuda_recordbatch_with_nulls.drop_null()

    # filter() test
    with pytest.raises(NotImplementedError):
        cuda_recordbatch.filter([True] * N)

    # take() test
    with pytest.raises(NotImplementedError):
        cuda_recordbatch.take([0])