GH-39984: [Python] Add ChunkedArray import/export to/from C

[paleolimbot]

Rationale for this change

ChunkedArrays have an unambiguous representation as a stream of arrays. #39455 added the ability to import/export in C++...this PR wires up the new functions in pyarrow.

What changes are included in this PR?

• Added __arrow_c_stream__() and _import_from_c_capsule() to the ChunkedArray

Are these changes tested?

Yes! Tests were added.

Are there any user-facing changes?

Yes! But I'm not sure where the protocol methods are documented.

import pyarrow as pa
import nanoarrow as na
chunked = pa.chunked_array([pa.array([0, 1, 2]), pa.array([3, 4, 5])])
[na.c_array_view(item) for item in na.c_array_stream(chunked)]

[<nanoarrow.c_lib.CArrayView>
 - storage_type: 'int64'
 - length: 3
 - offset: 0
 - null_count: 0
 - buffers[2]:
   - <bool validity[0 b] >
   - <int64 data[24 b] 0 1 2>
 - dictionary: NULL
 - children[0]:,
 <nanoarrow.c_lib.CArrayView>
 - storage_type: 'int64'
 - length: 3
 - offset: 0
 - null_count: 0
 - buffers[2]:
   - <bool validity[0 b] >
   - <int64 data[24 b] 3 4 5>
 - dictionary: NULL
 - children[0]:]

stream_capsule = chunked.__arrow_c_stream__()
chunked2 = chunked._import_from_c_capsule(stream_capsule)
chunked2

<pyarrow.lib.ChunkedArray object at 0x105bb70b0>
[
  [
    0,
    1,
    2
  ],
  [
    3,
    4,
    5
  ]
]

• Closes: [Python] Add ChunkedArray import/export to/from C #39984

[github-actions]

⚠️ GitHub issue #39984 has been automatically assigned in GitHub to PR creator.

[pitrou]

Not strictly necessary, but it would be nicer (both for memory consumption and for latency) to cast each batch when required, rather than all the table up front.

You could simply use RecordBatchReader.from_batches with a generator that casts each batch in turn. Something like:

        batches = table.to_batches()
        if requested_schema is not None:
            out_schema = Schema._import_from_c_capsule(requested_schema)
            if self.schema != out_schema:
                batches = (batch.cast(out_schema) for batch in batches)

        return RecordBatchReader.from_batches(batches)

(or you can fold the functionality directly in PyRecordBatchReader)

[pitrou]

This is probably not required.

[pitrou]

Same remark as in Table.__arrow_c_stream__.

[paleolimbot]

Is this as much of a concern as with a Table? I can't think of any clean way to do this lazily on a per-chunk basis, although I'm happy to remove the feature if it's that bad of an idea.

A quick check suggests casting up front is faster...perhaps because it obliterates the chunks:

import pyarrow as pa
import numpy as np

n = int(1e6)
n_chunks = 1000
per_chunk = n // n_chunks
chunks = [np.random.random(per_chunk) for i in range(n_chunks)]

chunked = pa.chunked_array(chunks)

def roundtrip_chunked():
    stream_capsule = chunked.__arrow_c_stream__()
    chunked._import_from_c_capsule(stream_capsule)

%timeit roundtrip_chunked()
#> 3.72 ms ± 15.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
    
def roundtrip_chunked_cast():
    stream_capsule = chunked.cast(pa.float()).__arrow_c_stream__()
    chunked._import_from_c_capsule(stream_capsule)

%timeit roundtrip_chunked_cast()
#> 1.52 ms ± 1.79 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

[pitrou]

Well, I gave a possible solution above.

Also, your benchmark is bizarre: where is the casting in roundtrip_chunked?

[paleolimbot]

Nowhere! A roundtrip without casting is almost certainly faster than a roundtrip with casting, and it is already twice as slow as the cast + roundtrip. I'm sure there's room to make a better benchmark (I'm also running a C++ debug build), but I'm personally convinced that the cast + export solution is not so bad that it should not be attempted.

Well, I gave a possible solution above.

For RecordBatches? I don't think we have a way to do that for a stream of Array in Arrow C++ or in pyarrow?

I'm happy to remove the feature as well and leave to be implemented properly later...I didn't anticipate it being controversial.

[pitrou]

A roundtrip without casting is almost certainly faster than a roundtrip with casting, and it is already twice as slow as the cast + roundtrip.

If it doesn't make sense then there's something wrong.

[pitrou]

I'm happy to remove the feature as well and leave to be implemented properly later...I didn't anticipate it being controversial.

It's not controversial. The implementation is. I'm sure for simple benchmarks with a small dataset, an otherwise idle machine, and enough RAM to hold multiple copies of the dataset, casting everything at once can seem slightly faster because it saves some overhead. That doesn't make it a viable strategy in the general case.

[paleolimbot]

Got it! I dropped the cast and we can circle back and do it right. I opened a PR to do it properly for the batch-wise RecordBatchReader export and that might serve as a template for how this should work, too.

[pitrou]

This should be fixed instead of working around it.

[pitrou]

@paleolimbot I think it would be nice to add a test_nanoarrow.py to exercise nanoarrow integration.

[paleolimbot]

Not strictly necessary, but it would be nicer (both for memory consumption and for latency) to cast each batch when required, rather than all the table up front.

I removed table casting for now...I'm still game to try in #40066 , maybe with a CastingRecordBatchReader : public RecordBatchReader that would solve the issue for streams, too.

I think it would be nice to add a test_nanoarrow.py to exercise nanoarrow integration.

I opened #40065 ...probably best done when nanoarrow's Python is a little more stable (i.e., when there is a nanoarrow.Array and nanoarrow.ArrayStream, neither of which currently exist).

[paleolimbot]

Is this as much of a concern as with a Table? I can't think of any clean way to do this lazily on a per-chunk basis, although I'm happy to remove the feature if it's that bad of an idea.

A quick check suggests casting up front is faster...perhaps because it obliterates the chunks:

import pyarrow as pa
import numpy as np

n = int(1e6)
n_chunks = 1000
per_chunk = n // n_chunks
chunks = [np.random.random(per_chunk) for i in range(n_chunks)]

chunked = pa.chunked_array(chunks)

def roundtrip_chunked():
    stream_capsule = chunked.__arrow_c_stream__()
    chunked._import_from_c_capsule(stream_capsule)

%timeit roundtrip_chunked()
#> 3.72 ms ± 15.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
    
def roundtrip_chunked_cast():
    stream_capsule = chunked.cast(pa.float()).__arrow_c_stream__()
    chunked._import_from_c_capsule(stream_capsule)

%timeit roundtrip_chunked_cast()
#> 1.52 ms ± 1.79 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

[pitrou]

Thank you @paleolimbot !

[conbench-apache-arrow]

After merging your PR, Conbench analyzed the 6 benchmarking runs that have been run so far on merge-commit 91bf1c9.

There was 1 benchmark result with an error:

• Commit Run on ursa-i9-9960x at 2024-02-14 11:17:50Z
  • tpch (R) with engine=arrow, format=parquet, language=R, memory_map=False, query_id=TPCH-19, scale_factor=10

There were no benchmark performance regressions. 🎉

The full Conbench report has more details.