pythonpyarrow

pyarrow chunkedarray get items at given indices


Say I have

In [3]: import pyarrow as pa

In [4]: ca = pa.chunked_array([[1,2,3], [4,5,6]])

I'd like to extract elements [1, 4, 2] and end up with

<pyarrow.lib.Int64Array object at 0x7f6eb43c2d40>
[
  2,
  5,
  3
]

as if I was doing NumPy-style indexing


Solution

  • Use the take function

    import pyarrow as pa
    ca = pa.chunked_array([[1, 2, 3], [4, 5, 6]])
    subset = ca.take([1, 2, 3])