apache / apache/arrow

[C++] TableBatchReader does not propagate device_type to record batches

Open
#44,049 0 comments 0 reactions 0 assignees View on GitHub
Component: C++ Type: bug
Dominant language
C++
Stars
17.1k
Forks
4.3k
Avg merge
3d 13h
Merged PRs (30d)
88

Description

### Describe the bug, including details regarding any error messages, version, and platform.

When creating a RecordBatchReader from a Table object, the reader always creates record batches with the default device_type (CPU).

This is a pyarrow repro using this PR on a CUDA device: https://github.com/apache/arrow/pull/43974
```
>>> import pyarrow as pa
>>> import pyarrow.cuda
>>> cuda_ctx = pyarrow.cuda.Context(0)
>>> schema = pa.schema([pa.field('c0', pa.int32()), pa.field('c1', pa.int32())])
>>> cpu_arrays = [pa.array([1, 2, 3, 4, 5], schema.field(0).type),
... pa.array([-10, -5, 0, None, 10], schema.field(1).type)]
>>> cuda_arrays = [arr.copy_to(cuda_ctx.memory_manager) for arr in cpu_arrays]
>>> cuda_chunked_array = pa.chunked_array(cuda_arrays)
>>> cuda_table = pa.table([cuda_chunked_array, cuda_chunked_array], schema=schema)

>>> print(cuda_table.is_cpu)
False

>>> for batch in cuda_table.to_batches():
... print(batch.device_type)
...
DeviceAllocationType.CPU
DeviceAllocationType.CPU
```

### Component(s)

C++

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.