Function test_get_columns

python/pyarrow/tests/interchange/test_interchange_spec.py:216–237

(uint, int, float, np_float_str, use_batch)

)
@pytest.mark.parametrize("use_batch", [False, True])
def test_get_columns(uint, int, float, np_float_str, use_batch):
    arr = [[1, 2, 3], [4, 5]]
    arr_float = np.array([1, 2, 3, 4, 5], dtype=np.dtype(np_float_str))
    table = pa.table(
        {
            "a": pa.chunked_array(arr, type=uint),
            "b": pa.chunked_array(arr, type=int),
            "c": pa.array(arr_float, type=float)
        }
    )
    if use_batch:
        table = table.combine_chunks().to_batches()[0]
    df = table.__dataframe__()
    for col in df.get_columns():
        assert col.size() == 5
        assert col.num_chunks() == 1

    # 0 = DtypeKind.INT, 1 = DtypeKind.UINT, 2 = DtypeKind.FLOAT,
    # see DtypeKind class in column.py
    assert df.get_column(0).dtype[0] == 1  # UINT
    assert df.get_column(1).dtype[0] == 0  # INT
    assert df.get_column(2).dtype[0] == 2  # FLOAT


@pytest.mark.parametrize(