Add stream_arrow: read Arrow data through the C Data Interface

[haoxu0]

Adds dx.stream_arrow(data): a stream over Arrow data, handed in through the Arrow C Data Interface.

I hit this training with MLX and wanting a HuggingFace dataset as the input. There is no way to feed one to an mlx-data pipeline today: you convert to JSONL and lose the column projection, or build the stream yourself outside mlx-data and lose the pipeline ops. But a HF dataset is already Arrow — the cache is Arrow IPC and Dataset.data is a pyarrow.Table — so the conversion is the only reason it does not work, and it should not need to exist.

import mlx.data as dx
from datasets import load_dataset

ds = load_dataset("...", split="train")
dset = dx.stream_arrow(ds.data.to_batches(1024))

Parquet, which is what #114 asked about, arrives through the same door:

import pyarrow.parquet as pq

dset = dx.stream_arrow(pq.ParquetFile("train.parquet").iter_batches(1024))

The implementation is straightforward. We reuse the arrow abi.h interface. And implement the ArrowStream similar to CsvStream.

One Sample per record batch, not per row

A batch's column already is a contiguous typed buffer, so it is adopted through the existing Array(type, shape, shared_ptr<void>); Array is unmodified. Emitting one sample per row would mean splitting those buffers and letting Batch reassemble them — over 200,000 rows × 3 columns, 22.8 ms via to_pylist(), 33.5 ms via per-row buffer reads, against 0.02 ms for whole columns.

The consequence, stated in the docstring: the leading dimension is Arrow's batch size, not one this stream chose. Set it at the reader (iter_batches(batch_size=1024)); .batch(n) on top re-batches. This is the part that does not look like the rest of mlx-data, where every stream yields a row, and it is question 1 below.

Ownership is one aliasing shared_ptr per column, pointing at that column's buffer while holding the whole batch, so the batch's release callback fires exactly once, when the last column Array is dropped.

Verification

Built against main at 2f431e9: full cmake build including the Python bindings, 0 errors and no new warnings; 24 Python tests pass (10 new, 14 existing); tests/test_arrow_stream.cpp checks values, slice offset, and that release fires exactly once. That test builds its own ArrowArrayStream out of plain C, so it needs neither Arrow nor a test framework, and it is the only place ownership is checked — counting release calls needs a producer you control, which the Python suite cannot reasonably hand-roll.

It is the first entry in tests/CMakeLists.txt (previously # coming soon), behind a new MLX_BUILD_TESTS that is off by default. Off by default plus a CI that builds through setup.py develop and runs only the Python suite would have meant the test never actually ran here, so linux_build_on_commit now configures it explicitly and runs ctest. Linux only — the test is plain C, so there is nothing platform-specific for the mac job to add.

I did not trust a suite that passed first try, so I reverted the offset arithmetic to void* values = base;: both offset tests fail and the offset-0 case still passes. Arrow's offset is in values, not bytes, and a sliced batch has a non-zero one — ignoring it silently returns the wrong rows, which is the easiest thing to get wrong in this interface.

Open questions

Nulls. Genuinely separate — huggingface datasets can use None because its rows are Python objects; a Sample has no equivalent.