commit 93b127327561aa74ae91fe5e690ced59e1276f11
author Dewey Dunnington <dewey@dunnington.ca> Tue Apr 23 20:58:52 2024 -0300
committer GitHub <noreply@github.com> Tue Apr 23 20:58:52 2024 -0300
tree c6af5d13333c0aad0e687720af1f36a54b1fa04b
parent c677d4d396e75d362a626db6c56207ef4ee4befa

feat(python): Iterate over array buffers (#433)

The idea with this change is to support efficient buffer access for
chunked/streaming input (e.g., make a numpy array). The efficient
implementation is compact but I am not sure it is easy to guess for
anybody not familiar with nanoarrow internals:

```python
with c_array_stream(obj, schema) as stream:
        for array in stream:
            view = array.view()
```

I'm not sure that `iter_chun_data()` is the best name here, but one
would use it like:

```python
import nanoarrow as na

array = na.Array([1, 2, 3], na.int32())

for view in array.iter_chunk_data():
    print(view.offset, view.length, list(view.buffers))
#> 0 3 [nanoarrow.c_lib.CBufferView(bool[0 b] ), nanoarrow.c_lib.CBufferView(int32[12 b] 1 2 3)]
```

This would replace `iter_buffers()` which is a little dangerous to use
(since one might assume the whole buffer represents the array, where we
really need the offset everywhere one might access a buffer). It also
cleans up some of the `ArrayViewIterator` terminology (since an earlier
version of this used the `ArrayViewIterator` instead of the simpler
approach it now uses).

The below benchmark is engineered to find the point where a this
iterator would be slower than `pa.ChunkedArray.to_numpy()` (for a
million doubles in this specific example, PyArrow becomes faster between
100 and 1000 chunks).

```python
import nanoarrow as na
from nanoarrow import c_lib
import pyarrow as pa
import numpy as np

n = int(1e6)
chunk_size = int(1e4)
num_chunks = n // chunk_size
n = chunk_size * num_chunks

chunks = [na.c_array(np.random.random(chunk_size)) for i in range(num_chunks)]
array = na.Array(c_lib.CArrayStream.from_array_list(chunks, na.c_schema(na.float64())))

def make():
    out = np.empty(len(array), dtype=np.float64)

    cursor = 0
    for view in array.iter_chunk_data():
        offset = view.offset
        length = view.length
        data = np.array(view.buffer(1), copy=False)[offset:(offset + length)]
        out[cursor:(cursor + length)] = np.array(data, copy=False)
        cursor += length

    return out


%timeit make()
#> 749 µs ± 37.8 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

chunked = pa.chunked_array([pa.array(item) for item in chunks])
%timeit chunked.to_numpy()
#> 2.07 ms ± 17.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

# With 1000 chunks of size 1000, the number would be
# iter_array_view()
#> 3.02 ms ± 238 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
# chunked.to_numpy()
#> 2.07 ms ± 16.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

np.testing.assert_equal(make(), chunked.to_numpy())
```

---------

Co-authored-by: Dane Pitkin <48041712+danepitkin@users.noreply.github.com>