[C++] Implement Take on ChunkedArray for DataFrame use

[asfimport]

Follow up to ARROW-2667

Reporter: Wes McKinney / @wesm
Assignee: Neal Richardson / @nealrichardson

Related issues:

• [C++][R] Move filter and take code from Rcpp to C++ library (Is contained by)
• [C++/Python] Add pandas-like take method to Array (is related to)

_{Note: This issue was originally created as ARROW-5454. Please see the migration documentation for further details.}

[asfimport]

Artem KOZHEVNIKOV / @artemru:
if it were in pure python, we could do something like below (relying on pa.array.take)

import numpy as np
import pyarrow as pa
from pandas.core.sorting import get_group_index_sorter

def take_on_chunked_array(charr, indices):
    indices = np.array(indices, dtype=np.int)
    if indices.max() > len(charr):
        raise IndexError()

    indices[indices < 0] += len(charr)

    if indices.min() < 0:
        raise IndexError()

    lengths = np.fromiter(map(len, charr.chunks), dtype=np.int64)
    cum_lengths = lengths.cumsum()

    bins = np.searchsorted(cum_lengths, indices, side="right")
    limits_idx = np.concatenate([[0], np.bincount(bins).cumsum()])

    sort_idx = get_group_index_sorter(bins, len(cum_lengths))
    del bins

    indices = indices[sort_idx]
    sort_idx = np.argsort(sort_idx, kind="merge")  # inverse sort indices

    cum_lengths -= lengths
    res_array = pa.concat_arrays([charr.chunks[i].take(pa.array(indices[limits_idx[i]:limits_idx[i + 1]] - cum_length))
                                  for i, cum_length in enumerate(cum_lengths)])
    return res_array.take(pa.array(sort_idx))


charr = pa.chunked_array([pa.array([0, 1]), pa.array([2, 3, 4]), pa.array([5, 6, 7, 8])])
take_on_chunked_array(charr, np.array([6, 0, 3])).to_numpy()
pa.concat_arrays(charr.chunks).take(pa.array([6, 0, 3])).to_numpy()

Do we want something similar in C++ ? Should we reuse cpp:Array:Take method and concat_arrays (or we want to avoid an extra copy) ? if we don't reuse array.take we can of cause avoid sorting indices back and forth.

[asfimport]

Wes McKinney / @wesm:
A C++ implementation would be useful. It makes most sense in the context of the "Data Frame" project

https://docs.google.com/document/d/1XHe_j87n2VHGzEbnLe786GHbbcbrzbjgG8D0IXWAeHg/edit?usp=sharing

I plan to dig into this finally after Labor Day

[asfimport]

Artem KOZHEVNIKOV / @artemru:
@wesm, could you have a look on this ? apart a arrow:DataFrame project, in meanwhile this feature can be very useful to work with pyarrow data structures via pandas.ExtenensionArray.

[asfimport]

Wes McKinney / @wesm:
You're welcome to contribute an implementation. It does not have to wait for the data frame project work to begin

[asfimport]

Neal Richardson / @nealrichardson:
See also ARROW-6784