Function test_table_non_cpu

python/pyarrow/tests/test_table.py:3880–4066

(cuda_context, cpu_table, cuda_table,
                       cuda_arrays, cuda_recordbatch, schema)

def test_table_non_cpu(cuda_context, cpu_table, cuda_table,
                       cuda_arrays, cuda_recordbatch, schema):
    verify_cuda_table(cuda_table, expected_schema=schema)
    N = cuda_table.num_rows

    # shape test
    assert cuda_table.shape == (10, 2)

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

    # add_column(), set_column() test
    for fn in [cuda_table.add_column, cuda_table.set_column]:
        cpu_col = pa.array([1] * N, pa.int8())
        cuda_col = cpu_col.copy_to(cuda_context.memory_manager)
        new_table = fn(2, 'c2', cuda_col)
        verify_cuda_table(new_table, expected_schema=schema.append(
            pa.field('c2', pa.int8())))
        new_table = fn(2, 'c2', cpu_col)
        assert new_table.is_cpu is False
        assert new_table.column(0).is_cpu is False
        assert new_table.column(1).is_cpu is False
        assert new_table.column(2).is_cpu is True

    # remove_column() test
    new_table = cuda_table.remove_column(1)
    verify_cuda_table(new_table, expected_schema=schema.remove(1))

    # drop_columns() test
    new_table = cuda_table.drop_columns(['c1'])
    verify_cuda_table(new_table, expected_schema=schema.remove(1))
    new_table = cuda_table.drop_columns(['c0', 'c1'])
    assert len(new_table.columns) == 0
    assert new_table.is_cpu

    # select() test
    new_table = cuda_table.select(['c0'])
    verify_cuda_table(new_table, 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_table.cast(new_schema)

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

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

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