apache / arrow-nanoarrow

Helpers for Arrow C Data & Arrow C Stream interfaces
https://arrow.apache.org/nanoarrow
Apache License 2.0
151 stars 34 forks source link

feat(python): Create string/binary arrays from iterables #430

Closed paleolimbot closed 2 months ago

paleolimbot commented 2 months ago

This PR adds support for building string and binary arrays via iterable.

It also cleans up a few parts of #426 that resulted in the wheel builds failing for (at least) PyPy 3.8 and 3.9. We can circle back to the performance of building from iterables (and whether or not pack_into() is essential) when all the wheels are building reliably.

import nanoarrow as na

strings = ["pizza", "yogurt", "noodles", "peanut butter sandwiches"]

na.Array(strings, na.string())
#> nanoarrow.Array<string>[4]
#> 'pizza'
#> 'yogurt'
#> 'noodles'
#> 'peanut butter sandwiches'

na.Array((s.encode() for s in strings), na.binary())
#> nanoarrow.Array<binary>[4]
#> b'pizza'
#> b'yogurt'
#> b'noodles'
#> b'peanut butter sandwiches'

The "build from iterable" code is now sufficiently complicated that it should be separated out. I did an initial attempt at that for this PR; however, it scrambles things up a bit and is complicated by the interdependence between the functions that sanitize arguments (e.g., c_schema(), c_array()) and the functions that build from iterable.

Currently faster for strings and slightly slower for bytes than pyarrow.

from itertools import cycle, islice
import nanoarrow as na
import pyarrow as pa

strings = ["pizza", "yogurt", "noodles", "peanut butter sandwiches"]
binary = [s.encode() for s in strings]

def many_strings():
    return islice(cycle(strings), int(1e6))

def many_strings_with_nulls():
    return islice(cycle(strings + [None]), int(1e6))

def many_bytes():
    return islice(cycle(binary), int(1e6))

def many_bytes_with_nulls():
    return islice(cycle(binary + [None]), int(1e6))

%timeit pa.array(many_strings(), pa.string())
#> 23.4 ms ± 488 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit na.c_array(many_strings(), na.string())
#> 14.3 ms ± 112 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit pa.array(many_strings_with_nulls(), pa.string())
#> 21.4 ms ± 340 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit na.c_array(many_strings_with_nulls(), na.string())
#> 17.1 ms ± 340 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit pa.array(many_bytes(), pa.binary())
#> 19.7 ms ± 283 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit na.c_array(many_bytes(), na.binary())
#> 16.3 ms ± 136 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit pa.array(many_bytes_with_nulls(), pa.binary())
#> 17.6 ms ± 37.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit na.c_array(many_bytes_with_nulls(), na.binary())
#> 19 ms ± 378 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)