[C++] TableBatchReader does not propagate device_type to record batches
- 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
Assessment
This issue has not been assessed yet.