Function test_buffer

python/pyarrow/tests/interchange/test_interchange_spec.py:244–274

(int, use_batch)

)
@pytest.mark.parametrize("use_batch", [False, True])
def test_buffer(int, use_batch):
    arr = [0, 1, -1]
    table = pa.table({"a": pa.array(arr, type=int)})
    if use_batch:
        table = table.to_batches()[0]
    df = table.__dataframe__()
    col = df.get_column(0)
    buf = col.get_buffers()

    dataBuf, dataDtype = buf["data"]

    assert dataBuf.bufsize > 0
    assert dataBuf.ptr != 0
    device, _ = dataBuf.__dlpack_device__()

    # 0 = DtypeKind.INT
    # see DtypeKind class in column.py
    assert dataDtype[0] == 0

    if device == 1:  # CPU-only as we're going to directly read memory here
        bitwidth = dataDtype[1]
        ctype = {
            8: ctypes.c_int8,
            16: ctypes.c_int16,
            32: ctypes.c_int32,
            64: ctypes.c_int64,
        }[bitwidth]

        for idx, truth in enumerate(arr):
            val = ctype.from_address(dataBuf.ptr + idx * (bitwidth // 8)).value
            assert val == truth, f"Buffer at index {idx} mismatch"


@pytest.mark.parametrize(