ARROW-11989: [C++][Python] Improve ChunkedArray's complexity for the access of elements

cpp/src/arrow/CMakeLists.txt

@@ -158,6 +158,7 @@ set(ARROW_SRCS
     builder.cc
     buffer.cc
     chunked_array.cc
+    chunk_resolver.cc
     compare.cc
     config.cc
     datum.cc

cpp/src/arrow/array/array_base.cc

@@ -165,8 +165,9 @@ struct ScalarFromArraySlotImpl {
 
   Result<std::shared_ptr<Scalar>> Finish() && {
     if (index_ >= array_.length()) {
-      return Status::IndexError("tried to refer to element ", index_,
-                                " but array is only ", array_.length(), " long");
+      return Status::IndexError("index with value of ", index_,
+                                " is out-of-bounds for array of length ",
+                                array_.length());
     }
 
     if (array_.IsNull(index_)) {

cpp/src/arrow/chunk_resolver.cc

@@ -0,0 +1,69 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/chunk_resolver.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+#include "arrow/array.h"
+#include "arrow/record_batch.h"
+
+namespace arrow {
+namespace internal {
+
+namespace {
+template <typename T>
+int64_t GetLength(const T& array) {
+  // General case assumes argument is an Array pointer
+  return array->length();
+}
+
+template <>
+int64_t GetLength<std::shared_ptr<RecordBatch>>(
+    const std::shared_ptr<RecordBatch>& batch) {
+  return batch->num_rows();
+}
+
+template <typename T>
+inline std::vector<int64_t> MakeChunksOffsets(const std::vector<T>& chunks) {
+  std::vector<int64_t> offsets(chunks.size() + 1);
+  int64_t offset = 0;
+  std::transform(chunks.begin(), chunks.end(), offsets.begin(),
+                 [&offset](const T& chunk) {
+                   auto curr_offset = offset;
+                   offset += GetLength(chunk);
+                   return curr_offset;
+                 });
+  offsets[chunks.size()] = offset;
+  return offsets;
+}
+}  // namespace
+
+ChunkResolver::ChunkResolver(const ArrayVector& chunks)
+    : offsets_(MakeChunksOffsets(chunks)), cached_chunk_(0) {}
+
+ChunkResolver::ChunkResolver(const std::vector<const Array*>& chunks)
+    : offsets_(MakeChunksOffsets(chunks)), cached_chunk_(0) {}
+
+ChunkResolver::ChunkResolver(const RecordBatchVector& batches)
+    : offsets_(MakeChunksOffsets(batches)), cached_chunk_(0) {}
+
+}  // namespace internal
+}  // namespace arrow

cpp/src/arrow/chunk_resolver.h

@@ -0,0 +1,104 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <atomic>
+#include <cstdint>
+#include <vector>
+
+#include "arrow/type_fwd.h"
+#include "arrow/util/macros.h"
+
+namespace arrow {
+namespace internal {
+
+struct ChunkLocation {
+  int64_t chunk_index, index_in_chunk;
+};
+
+// An object that resolves an array chunk depending on a logical index
+struct ChunkResolver {
+  explicit ChunkResolver(const ArrayVector& chunks);
+
+  explicit ChunkResolver(const std::vector<const Array*>& chunks);
+
+  explicit ChunkResolver(const RecordBatchVector& batches);
+
+  ChunkResolver(ChunkResolver&& other)
+      : offsets_(std::move(other.offsets_)), cached_chunk_(other.cached_chunk_.load()) {}
+
+  ChunkResolver& operator=(ChunkResolver&& other) {
+    offsets_ = std::move(other.offsets_);
+    cached_chunk_.store(other.cached_chunk_.load());
+    return *this;
+  }
+
+  /// \brief Return a ChunkLocation containing the chunk index and in-chunk value index of
+  /// the chunked array at logical index
+  inline ChunkLocation Resolve(const int64_t index) const {
+    // It is common for the algorithms below to make consecutive accesses at
+    // a relatively small distance from each other, hence often falling in
+    // the same chunk.
+    // This is trivial when merging (assuming each side of the merge uses
+    // its own resolver), but also in the inner recursive invocations of
+    // partitioning.
+    if (offsets_.size() <= 1) {
+      return {0, index};
+    }
+    const auto cached_chunk = cached_chunk_.load();
+    const bool cache_hit =
+        (index >= offsets_[cached_chunk] && index < offsets_[cached_chunk + 1]);
+    if (ARROW_PREDICT_TRUE(cache_hit)) {
+      return {cached_chunk, index - offsets_[cached_chunk]};
+    }
+    auto chunk_index = Bisect(index);
+    cached_chunk_.store(chunk_index);
+    return {chunk_index, index - offsets_[chunk_index]};
+  }
+
+ protected:
+  // Find the chunk index corresponding to a value index using binary search
+  inline int64_t Bisect(const int64_t index) const {
+    // Like std::upper_bound(), but hand-written as it can help the compiler.
+    // Search [lo, lo + n)
+    int64_t lo = 0;
+    auto n = static_cast<int64_t>(offsets_.size());
+    while (n > 1) {
+      const int64_t m = n >> 1;
+      const int64_t mid = lo + m;
+      if (static_cast<int64_t>(index) >= offsets_[mid]) {
+        lo = mid;
+        n -= m;
+      } else {
+        n = m;
+      }
+    }
+    return lo;
+  }
+
+ private:
+  // Collection of starting offsets used for binary search
+  std::vector<int64_t> offsets_;
+
+  // Tracks the most recently used chunk index to allow fast
+  // access for consecutive indices corresponding to the same chunk
+  mutable std::atomic<int64_t> cached_chunk_;
+};
+
+}  // namespace internal
+}  // namespace arrow

cpp/src/arrow/chunked_array.cc

@@ -43,15 +43,17 @@ class MemoryPool;
 // ChunkedArray methods
 
 ChunkedArray::ChunkedArray(ArrayVector chunks, std::shared_ptr<DataType> type)
-    : chunks_(std::move(chunks)), type_(std::move(type)) {
-  length_ = 0;
-  null_count_ = 0;
-
+    : chunks_(std::move(chunks)),
+      type_(std::move(type)),
+      length_(0),
+      null_count_(0),
+      chunk_resolver_{chunks_} {
   if (type_ == nullptr) {
     ARROW_CHECK_GT(chunks_.size(), 0)
         << "cannot construct ChunkedArray from empty vector and omitted type";
     type_ = chunks_[0]->type();
   }
+
   for (const auto& chunk : chunks_) {
     length_ += chunk->length();
     null_count_ += chunk->null_count();
@@ -147,13 +149,12 @@ bool ChunkedArray::ApproxEquals(const ChunkedArray& other,
 }
 
 Result<std::shared_ptr<Scalar>> ChunkedArray::GetScalar(int64_t index) const {
-  for (const auto& chunk : chunks_) {
-    if (index < chunk->length()) {
-      return chunk->GetScalar(index);
-    }
-    index -= chunk->length();
+  const auto loc = chunk_resolver_.Resolve(index);
+  if (loc.chunk_index >= static_cast<int64_t>(chunks_.size())) {
+    return Status::IndexError("index with value of ", index,
+                              " is out-of-bounds for chunked array of length ", length_);
   }
-  return Status::Invalid("index out of bounds");
+  return chunks_[loc.chunk_index]->GetScalar(loc.index_in_chunk);
 }
 
 std::shared_ptr<ChunkedArray> ChunkedArray::Slice(int64_t offset, int64_t length) const {

cpp/src/arrow/chunked_array.h

@@ -23,6 +23,7 @@
 #include <utility>
 #include <vector>
 
+#include "arrow/chunk_resolver.h"
 #include "arrow/compare.h"
 #include "arrow/result.h"
 #include "arrow/status.h"
@@ -177,11 +178,12 @@ class ARROW_EXPORT ChunkedArray {
 
  protected:
   ArrayVector chunks_;
+  std::shared_ptr<DataType> type_;
   int64_t length_;
   int64_t null_count_;
-  std::shared_ptr<DataType> type_;
 
  private:
+  internal::ChunkResolver chunk_resolver_;
   ARROW_DISALLOW_COPY_AND_ASSIGN(ChunkedArray);
 };
 

cpp/src/arrow/chunked_array_test.cc

@@ -278,7 +278,7 @@ TEST_F(TestChunkedArray, GetScalar) {
   check_scalar(carr, 4, **MakeScalar(ty, 3));
   check_scalar(carr, 6, **MakeScalar(ty, 5));
 
-  ASSERT_RAISES(Invalid, carr.GetScalar(7));
+  ASSERT_RAISES(IndexError, carr.GetScalar(7));
 }
 
 }  // namespace arrow