From 83f35de2ebe4672d0d57290840d0153ac8a016e8 Mon Sep 17 00:00:00 2001 From: Felipe Oliveira Carvalho Date: Tue, 24 Sep 2024 11:36:24 -0300 Subject: [PATCH] GH-43953: [C++] Add tests based on random data and benchmarks to ChunkResolver::ResolveMany (#43954) ### Rationale for this change Improve tests and add benchmarks. I wrote the tests and benchmarks while trying to improve the performance of `ResolveMany` and failing at it. ### What changes are included in this PR? Tests, benchmarks, and changes that don't really affect performance but might unlock more optimization opportunities in the future. ### Are these changes tested? Yes. * GitHub Issue: #43953 Lead-authored-by: Felipe Oliveira Carvalho Co-authored-by: Antoine Pitrou Signed-off-by: Antoine Pitrou --- cpp/src/arrow/CMakeLists.txt | 1 + cpp/src/arrow/chunk_resolver.cc | 104 ++++++++------ cpp/src/arrow/chunk_resolver.h | 57 ++++---- cpp/src/arrow/chunk_resolver_benchmark.cc | 166 ++++++++++++++++++++++ cpp/src/arrow/chunked_array.cc | 4 +- cpp/src/arrow/chunked_array_test.cc | 68 +++++++++ 6 files changed, 326 insertions(+), 74 deletions(-) create mode 100644 cpp/src/arrow/chunk_resolver_benchmark.cc diff --git a/cpp/src/arrow/CMakeLists.txt b/cpp/src/arrow/CMakeLists.txt index e77a02d0c0800..c911f0f4e9481 100644 --- a/cpp/src/arrow/CMakeLists.txt +++ b/cpp/src/arrow/CMakeLists.txt @@ -1218,6 +1218,7 @@ add_arrow_test(sparse_tensor_test) add_arrow_test(stl_test SOURCES stl_iterator_test.cc stl_test.cc) add_arrow_benchmark(builder_benchmark) +add_arrow_benchmark(chunk_resolver_benchmark) add_arrow_benchmark(compare_benchmark) add_arrow_benchmark(memory_pool_benchmark) add_arrow_benchmark(type_benchmark) diff --git a/cpp/src/arrow/chunk_resolver.cc b/cpp/src/arrow/chunk_resolver.cc index 854127480744e..bda6b17810299 100644 --- a/cpp/src/arrow/chunk_resolver.cc +++ b/cpp/src/arrow/chunk_resolver.cc @@ -55,43 +55,57 @@ inline std::vector MakeChunksOffsets(const std::vector& chunks) { return offsets; } +template +inline TypedChunkLocation ResolveOneInline(uint32_t num_offsets, + const uint64_t* offsets, + IndexType typed_logical_index, + int32_t num_chunks, + int32_t chunk_hint) { + const auto index = static_cast(typed_logical_index); + // use or update chunk_hint + if (index >= offsets[chunk_hint] && + (chunk_hint == num_chunks || index < offsets[chunk_hint + 1])) { + // hint is correct! + } else { + // lo < hi is guaranteed by `num_offsets = chunks.size() + 1` + auto chunk_index = + ChunkResolver::Bisect(index, offsets, /*lo=*/0, /*hi=*/num_offsets); + chunk_hint = static_cast(chunk_index); + } + // chunk_index is in [0, chunks.size()] no matter what the value + // of logical_index is, so it's always safe to dereference offsets + // as it contains chunks.size()+1 values. + auto loc = TypedChunkLocation( + /*chunk_index=*/chunk_hint, + /*index_in_chunk=*/typed_logical_index - + static_cast(offsets[chunk_hint])); +#if defined(ARROW_VALGRIND) || defined(ADDRESS_SANITIZER) + // Make it more likely that Valgrind/ASAN can catch an invalid memory + // access by poisoning the index-in-chunk value when the logical + // index is out-of-bounds. + if (static_cast(loc.chunk_index) == num_chunks) { + loc.index_in_chunk = std::numeric_limits::max(); + } +#endif + return loc; +} + /// \pre all the pre-conditions of ChunkResolver::ResolveMany() /// \pre num_offsets - 1 <= std::numeric_limits::max() template -void ResolveManyInline(size_t num_offsets, const int64_t* signed_offsets, +void ResolveManyInline(uint32_t num_offsets, const int64_t* signed_offsets, int64_t n_indices, const IndexType* logical_index_vec, TypedChunkLocation* out_chunk_location_vec, - IndexType chunk_hint) { + int32_t chunk_hint) { auto* offsets = reinterpret_cast(signed_offsets); - const auto num_chunks = static_cast(num_offsets - 1); + const auto num_chunks = static_cast(num_offsets - 1); // chunk_hint in [0, num_offsets) per the precondition. for (int64_t i = 0; i < n_indices; i++) { - auto typed_logical_index = logical_index_vec[i]; - const auto index = static_cast(typed_logical_index); - // use or update chunk_hint - if (index >= offsets[chunk_hint] && - (chunk_hint == num_chunks || index < offsets[chunk_hint + 1])) { - // hint is correct! - } else { - // lo < hi is guaranteed by `num_offsets = chunks.size() + 1` - auto chunk_index = - ChunkResolver::Bisect(index, offsets, /*lo=*/0, /*hi=*/num_offsets); - chunk_hint = static_cast(chunk_index); - } - out_chunk_location_vec[i].chunk_index = chunk_hint; - // chunk_index is in [0, chunks.size()] no matter what the - // value of logical_index is, so it's always safe to dereference - // offset_ as it contains chunks.size()+1 values. - out_chunk_location_vec[i].index_in_chunk = - typed_logical_index - static_cast(offsets[chunk_hint]); -#if defined(ARROW_VALGRIND) || defined(ADDRESS_SANITIZER) - // Make it more likely that Valgrind/ASAN can catch an invalid memory - // access by poisoning the index-in-chunk value when the logical - // index is out-of-bounds. - if (chunk_hint == num_chunks) { - out_chunk_location_vec[i].index_in_chunk = std::numeric_limits::max(); - } -#endif + const auto typed_logical_index = logical_index_vec[i]; + const auto loc = ResolveOneInline(num_offsets, offsets, typed_logical_index, + num_chunks, chunk_hint); + out_chunk_location_vec[i] = loc; + chunk_hint = static_cast(loc.chunk_index); } } @@ -127,30 +141,30 @@ ChunkResolver& ChunkResolver::operator=(const ChunkResolver& other) noexcept { void ChunkResolver::ResolveManyImpl(int64_t n_indices, const uint8_t* logical_index_vec, TypedChunkLocation* out_chunk_location_vec, - uint8_t chunk_hint) const { - ResolveManyInline(offsets_.size(), offsets_.data(), n_indices, logical_index_vec, - out_chunk_location_vec, chunk_hint); -} - -void ChunkResolver::ResolveManyImpl(int64_t n_indices, const uint32_t* logical_index_vec, - TypedChunkLocation* out_chunk_location_vec, - uint32_t chunk_hint) const { - ResolveManyInline(offsets_.size(), offsets_.data(), n_indices, logical_index_vec, - out_chunk_location_vec, chunk_hint); + int32_t chunk_hint) const { + ResolveManyInline(static_cast(offsets_.size()), offsets_.data(), n_indices, + logical_index_vec, out_chunk_location_vec, chunk_hint); } void ChunkResolver::ResolveManyImpl(int64_t n_indices, const uint16_t* logical_index_vec, TypedChunkLocation* out_chunk_location_vec, - uint16_t chunk_hint) const { - ResolveManyInline(offsets_.size(), offsets_.data(), n_indices, logical_index_vec, - out_chunk_location_vec, chunk_hint); + int32_t chunk_hint) const { + ResolveManyInline(static_cast(offsets_.size()), offsets_.data(), n_indices, + logical_index_vec, out_chunk_location_vec, chunk_hint); +} + +void ChunkResolver::ResolveManyImpl(int64_t n_indices, const uint32_t* logical_index_vec, + TypedChunkLocation* out_chunk_location_vec, + int32_t chunk_hint) const { + ResolveManyInline(static_cast(offsets_.size()), offsets_.data(), n_indices, + logical_index_vec, out_chunk_location_vec, chunk_hint); } void ChunkResolver::ResolveManyImpl(int64_t n_indices, const uint64_t* logical_index_vec, TypedChunkLocation* out_chunk_location_vec, - uint64_t chunk_hint) const { - ResolveManyInline(offsets_.size(), offsets_.data(), n_indices, logical_index_vec, - out_chunk_location_vec, chunk_hint); + int32_t chunk_hint) const { + ResolveManyInline(static_cast(offsets_.size()), offsets_.data(), n_indices, + logical_index_vec, out_chunk_location_vec, chunk_hint); } } // namespace arrow::internal diff --git a/cpp/src/arrow/chunk_resolver.h b/cpp/src/arrow/chunk_resolver.h index 83fda62387fe1..4a5e27c05361f 100644 --- a/cpp/src/arrow/chunk_resolver.h +++ b/cpp/src/arrow/chunk_resolver.h @@ -72,7 +72,7 @@ struct ARROW_EXPORT ChunkResolver { /// \brief Cache of the index of the last resolved chunk. /// /// \invariant `cached_chunk_ in [0, chunks.size()]` - mutable std::atomic cached_chunk_; + mutable std::atomic cached_chunk_; public: explicit ChunkResolver(const ArrayVector& chunks) noexcept; @@ -92,6 +92,8 @@ struct ARROW_EXPORT ChunkResolver { for (size_t i = 1; i < offsets_.size(); i++) { assert(offsets_[i] >= offsets_[i - 1]); } + assert(offsets_.size() - 1 <= + static_cast(std::numeric_limits::max())); #endif } @@ -102,7 +104,7 @@ struct ARROW_EXPORT ChunkResolver { ChunkResolver& operator=(const ChunkResolver& other) noexcept; int64_t logical_array_length() const { return offsets_.back(); } - int64_t num_chunks() const { return static_cast(offsets_.size()) - 1; } + int32_t num_chunks() const { return static_cast(offsets_.size() - 1); } int64_t chunk_length(int64_t chunk_index) const { return offsets_[chunk_index + 1] - offsets_[chunk_index]; @@ -140,9 +142,9 @@ struct ARROW_EXPORT ChunkResolver { /// bounds, or with chunk_index == chunks.size() if logical index is /// `>= chunked_array.length()`. inline ChunkLocation ResolveWithHint(int64_t index, ChunkLocation hint) const { - assert(hint.chunk_index < static_cast(offsets_.size())); - const auto chunk_index = - ResolveChunkIndex(index, hint.chunk_index); + assert(hint.chunk_index < static_cast(offsets_.size())); + const auto chunk_index = ResolveChunkIndex( + index, static_cast(hint.chunk_index)); return ChunkLocation{chunk_index, index - offsets_[chunk_index]}; } @@ -169,13 +171,12 @@ struct ARROW_EXPORT ChunkResolver { [[nodiscard]] bool ResolveMany(int64_t n_indices, const IndexType* logical_index_vec, TypedChunkLocation* out_chunk_location_vec, IndexType chunk_hint = 0) const { - if constexpr (sizeof(IndexType) < sizeof(uint64_t)) { + if constexpr (sizeof(IndexType) < sizeof(uint32_t)) { // The max value returned by Bisect is `offsets.size() - 1` (= chunks.size()). - constexpr uint64_t kMaxIndexTypeValue = std::numeric_limits::max(); + constexpr int64_t kMaxIndexTypeValue = std::numeric_limits::max(); // A ChunkedArray with enough empty chunks can make the index of a chunk // exceed the logical index and thus the maximum value of IndexType. - const bool chunk_index_fits_on_type = - static_cast(offsets_.size() - 1) <= kMaxIndexTypeValue; + const bool chunk_index_fits_on_type = num_chunks() <= kMaxIndexTypeValue; if (ARROW_PREDICT_FALSE(!chunk_index_fits_on_type)) { return false; } @@ -194,17 +195,18 @@ struct ARROW_EXPORT ChunkResolver { using U = std::make_unsigned_t; ResolveManyImpl(n_indices, reinterpret_cast(logical_index_vec), reinterpret_cast*>(out_chunk_location_vec), - static_cast(chunk_hint)); + static_cast(chunk_hint)); } else { static_assert(std::is_unsigned_v); - ResolveManyImpl(n_indices, logical_index_vec, out_chunk_location_vec, chunk_hint); + ResolveManyImpl(n_indices, logical_index_vec, out_chunk_location_vec, + static_cast(chunk_hint)); } return true; } private: template - inline int64_t ResolveChunkIndex(int64_t index, int64_t cached_chunk) const { + inline int64_t ResolveChunkIndex(int64_t index, int32_t cached_chunk) const { // It is common for algorithms sequentially processing arrays to make consecutive // accesses at a relatively small distance from each other, hence often falling in the // same chunk. @@ -212,16 +214,17 @@ struct ARROW_EXPORT ChunkResolver { // This is guaranteed when merging (assuming each side of the merge uses its // own resolver), and is the most common case in recursive invocations of // partitioning. - const auto num_offsets = static_cast(offsets_.size()); + const auto num_offsets = static_cast(offsets_.size()); const int64_t* offsets = offsets_.data(); if (ARROW_PREDICT_TRUE(index >= offsets[cached_chunk]) && - (cached_chunk + 1 == num_offsets || index < offsets[cached_chunk + 1])) { + (static_cast(cached_chunk + 1) == num_offsets || + index < offsets[cached_chunk + 1])) { return cached_chunk; } // lo < hi is guaranteed by `num_offsets = chunks.size() + 1` const auto chunk_index = Bisect(index, offsets, /*lo=*/0, /*hi=*/num_offsets); if constexpr (StoreCachedChunk) { - assert(chunk_index < static_cast(offsets_.size())); + assert(static_cast(chunk_index) < static_cast(offsets_.size())); cached_chunk_.store(chunk_index, std::memory_order_relaxed); } return chunk_index; @@ -230,13 +233,13 @@ struct ARROW_EXPORT ChunkResolver { /// \pre all the pre-conditions of ChunkResolver::ResolveMany() /// \pre num_offsets - 1 <= std::numeric_limits::max() void ResolveManyImpl(int64_t, const uint8_t*, TypedChunkLocation*, - uint8_t) const; + int32_t) const; void ResolveManyImpl(int64_t, const uint16_t*, TypedChunkLocation*, - uint16_t) const; + int32_t) const; void ResolveManyImpl(int64_t, const uint32_t*, TypedChunkLocation*, - uint32_t) const; + int32_t) const; void ResolveManyImpl(int64_t, const uint64_t*, TypedChunkLocation*, - uint64_t) const; + int32_t) const; public: /// \brief Find the index of the chunk that contains the logical index. @@ -249,15 +252,15 @@ struct ARROW_EXPORT ChunkResolver { /// \pre index >= 0 (otherwise, when index is negative, hi-1 is returned) /// \pre lo < hi /// \pre lo >= 0 && hi <= offsets_.size() - static inline int64_t Bisect(int64_t index, const int64_t* offsets, int64_t lo, - int64_t hi) { + static inline int32_t Bisect(int64_t index, const int64_t* offsets, int32_t lo, + int32_t hi) { return Bisect(static_cast(index), - reinterpret_cast(offsets), static_cast(lo), - static_cast(hi)); + reinterpret_cast(offsets), static_cast(lo), + static_cast(hi)); } - static inline int64_t Bisect(uint64_t index, const uint64_t* offsets, uint64_t lo, - uint64_t hi) { + static inline int32_t Bisect(uint64_t index, const uint64_t* offsets, uint32_t lo, + uint32_t hi) { // Similar to std::upper_bound(), but slightly different as our offsets // array always starts with 0. auto n = hi - lo; @@ -265,8 +268,8 @@ struct ARROW_EXPORT ChunkResolver { // (lo < hi is guaranteed by the precondition). assert(n > 1 && "lo < hi is a precondition of Bisect"); do { - const uint64_t m = n >> 1; - const uint64_t mid = lo + m; + const uint32_t m = n >> 1; + const uint32_t mid = lo + m; if (index >= offsets[mid]) { lo = mid; n -= m; diff --git a/cpp/src/arrow/chunk_resolver_benchmark.cc b/cpp/src/arrow/chunk_resolver_benchmark.cc new file mode 100644 index 0000000000000..0756de3fbe930 --- /dev/null +++ b/cpp/src/arrow/chunk_resolver_benchmark.cc @@ -0,0 +1,166 @@ +// 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 "benchmark/benchmark.h" + +#include +#include +#include +#include + +#include "arrow/chunk_resolver.h" +#include "arrow/util/benchmark_util.h" +#include "arrow/util/pcg_random.h" + +namespace arrow { + +using internal::ChunkResolver; +using internal::TypedChunkLocation; + +namespace { + +int64_t constexpr kChunkedArrayLength = std::numeric_limits::max(); + +struct ResolveManyBenchmark { + benchmark::State& state; + random::pcg64 rng; + // Values from the state.range(i) + int64_t chunked_array_length; + int32_t num_chunks; + int64_t num_logical_indices; + + explicit ResolveManyBenchmark(benchmark::State& state) + : state(state), + rng(42), + chunked_array_length(state.range(0)), + num_chunks(static_cast(state.range(1))), + num_logical_indices(state.range(2)) {} + + std::vector GenChunkedArrayOffsets() { + std::uniform_int_distribution offset_gen(1, chunked_array_length); + std::vector offsets; + offsets.reserve(num_chunks + 1); + offsets.push_back(0); + while (offsets.size() < static_cast(num_chunks)) { + offsets.push_back(offset_gen(rng)); + } + offsets.push_back(chunked_array_length); + std::sort(offsets.begin() + 1, offsets.end()); + return offsets; + } + + template + std::vector GenRandomIndices(IndexType max_index, bool sorted) { + std::uniform_int_distribution index_gen(0, max_index); + std::vector indices; + indices.reserve(num_logical_indices); + while (indices.size() < static_cast(num_logical_indices)) { + indices.push_back(index_gen(rng)); + } + if (sorted) { + std::sort(indices.begin(), indices.end()); + } + return indices; + } + + template + void Bench(bool sorted) { + if constexpr (sizeof(IndexType) < 8) { + constexpr uint64_t kLimitIndex = std::numeric_limits::max(); + ARROW_CHECK_LE(static_cast(chunked_array_length), kLimitIndex); + } + const auto max_random_index = static_cast(chunked_array_length); + auto offsets = GenChunkedArrayOffsets(); + auto logical_indices = GenRandomIndices(max_random_index, sorted); + ChunkResolver resolver(std::move(offsets)); + std::vector> chunk_location_vec(num_logical_indices); + BENCHMARK_UNUSED bool all_succeeded = true; + for (auto _ : state) { + const bool success = resolver.ResolveMany( + num_logical_indices, logical_indices.data(), chunk_location_vec.data()); + all_succeeded &= success; + } + ARROW_CHECK(all_succeeded); + state.SetItemsProcessed(state.iterations() * num_logical_indices); + } +}; + +template +void ResolveManySetArgs(benchmark::internal::Benchmark* bench) { + constexpr int32_t kNonAligned = 3; + const int64_t kNumIndicesFew = (kChunkedArrayLength >> 7) - kNonAligned; + const int64_t kNumIndicesMany = (kChunkedArrayLength >> 1) - kNonAligned; + + bench->ArgNames({"chunked_array_length", "num_chunks", "num_indices"}); + + switch (sizeof(IndexType)) { + case 1: + // Unexpected. See comments below. + case 2: + case 4: + case 8: + bench->Args({kChunkedArrayLength, /*num_chunks*/ 10000, kNumIndicesFew}); + bench->Args({kChunkedArrayLength, /*num_chunks*/ 100, kNumIndicesFew}); + bench->Args({kChunkedArrayLength, /*num_chunks*/ 10000, kNumIndicesMany}); + bench->Args({kChunkedArrayLength, /*num_chunks*/ 100, kNumIndicesMany}); + break; + } +} + +template +void ResolveManyBench(benchmark::State& state, bool sorted) { + ResolveManyBenchmark{state}.Bench(sorted); +} + +void ResolveManyUInt16Random(benchmark::State& state) { + ResolveManyBench(state, false); +} + +void ResolveManyUInt32Random(benchmark::State& state) { + ResolveManyBench(state, false); +} + +void ResolveManyUInt64Random(benchmark::State& state) { + ResolveManyBench(state, false); +} + +void ResolveManyUInt16Sorted(benchmark::State& state) { + ResolveManyBench(state, true); +} + +void ResolveManyUInt32Sorted(benchmark::State& state) { + ResolveManyBench(state, true); +} + +void ResolveManyUInt64Sorted(benchmark::State& state) { + ResolveManyBench(state, true); +} + +} // namespace + +// We don't benchmark with uint8_t because it's too fast -- any meaningful +// array of logical indices will contain all the possible values. + +BENCHMARK(ResolveManyUInt16Random)->Apply(ResolveManySetArgs); +BENCHMARK(ResolveManyUInt32Random)->Apply(ResolveManySetArgs); +BENCHMARK(ResolveManyUInt64Random)->Apply(ResolveManySetArgs); + +BENCHMARK(ResolveManyUInt16Sorted)->Apply(ResolveManySetArgs); +BENCHMARK(ResolveManyUInt32Sorted)->Apply(ResolveManySetArgs); +BENCHMARK(ResolveManyUInt64Sorted)->Apply(ResolveManySetArgs); + +} // namespace arrow diff --git a/cpp/src/arrow/chunked_array.cc b/cpp/src/arrow/chunked_array.cc index dd6aa51534fcb..2ea2f514b9cab 100644 --- a/cpp/src/arrow/chunked_array.cc +++ b/cpp/src/arrow/chunked_array.cc @@ -51,11 +51,11 @@ ChunkedArray::ChunkedArray(ArrayVector chunks, std::shared_ptr type) null_count_(0), chunk_resolver_{chunks_} { if (type_ == nullptr) { - ARROW_CHECK_GT(chunks_.size(), 0) + ARROW_CHECK_GT(chunks_.size(), static_cast(0)) << "cannot construct ChunkedArray from empty vector and omitted type"; type_ = chunks_[0]->type(); } - + ARROW_CHECK_LE(chunks.size(), static_cast(std::numeric_limits::max())); for (const auto& chunk : chunks_) { length_ += chunk->length(); null_count_ += chunk->null_count(); diff --git a/cpp/src/arrow/chunked_array_test.cc b/cpp/src/arrow/chunked_array_test.cc index bf9d4af7c7bb0..f98dde689c237 100644 --- a/cpp/src/arrow/chunked_array_test.cc +++ b/cpp/src/arrow/chunked_array_test.cc @@ -32,6 +32,7 @@ #include "arrow/type.h" #include "arrow/util/endian.h" #include "arrow/util/key_value_metadata.h" +#include "arrow/util/pcg_random.h" namespace arrow { @@ -373,10 +374,27 @@ TEST(TestChunkResolver, Resolve) { ASSERT_EQ(resolver.Resolve(10).chunk_index, 3); } +template +std::vector GenChunkedArrayOffsets(RNG& rng, int32_t num_chunks, + int64_t chunked_array_len) { + std::uniform_int_distribution offset_gen(1, chunked_array_len - 1); + std::vector offsets; + offsets.reserve(num_chunks + 1); + offsets.push_back(0); + while (offsets.size() < static_cast(num_chunks)) { + offsets.push_back(offset_gen(rng)); + } + offsets.push_back(chunked_array_len); + std::sort(offsets.begin() + 1, offsets.end()); + return offsets; +} + template class TestChunkResolverMany : public ::testing::Test { public: using IndexType = T; + static constexpr int32_t kMaxInt32 = std::numeric_limits::max(); + static constexpr uint64_t kMaxValidIndex = std::numeric_limits::max(); Result> ResolveMany( const ChunkResolver& resolver, const std::vector& logical_index_vec) { @@ -510,6 +528,55 @@ class TestChunkResolverMany : public ::testing::Test { ASSERT_OK(ResolveMany(resolver_with_empty, logical_index_vec)); } } + + void TestRandomInput(int32_t num_chunks, int64_t chunked_array_len) { + random::pcg64 rng(42); + + // Generate random chunk offsets... + auto offsets = GenChunkedArrayOffsets(rng, num_chunks, chunked_array_len); + ASSERT_EQ(offsets.size(), static_cast(num_chunks) + 1); + // ...and ensure there is at least one empty chunk. + std::uniform_int_distribution chunk_index_gen( + 1, static_cast(num_chunks - 1)); + auto chunk_index = chunk_index_gen(rng); + offsets[chunk_index] = offsets[chunk_index - 1]; + + // Generate random query array of logical indices... + const auto num_logical_indices = 3 * static_cast(num_chunks) / 2; + std::vector logical_index_vec; + logical_index_vec.reserve(num_logical_indices); + std::uniform_int_distribution logical_index_gen(1, kMaxValidIndex); + for (int64_t i = 0; i < num_logical_indices; i++) { + const auto index = static_cast(logical_index_gen(rng)); + logical_index_vec.push_back(index); + } + // ...and sprinkle some extreme logical index values. + std::uniform_int_distribution position_gen(0, logical_index_vec.size() - 1); + for (int i = 0; i < 2; i++) { + auto max_valid_index = + std::min(kMaxValidIndex, static_cast(chunked_array_len)); + // zero and last valid logical index + logical_index_vec[position_gen(rng)] = 0; + logical_index_vec[position_gen(rng)] = static_cast(max_valid_index - 1); + // out of bounds indices + logical_index_vec[position_gen(rng)] = static_cast(max_valid_index); + if (max_valid_index < kMaxValidIndex) { + logical_index_vec[position_gen(rng)] = + static_cast(max_valid_index + 1); + } + } + + ChunkResolver resolver(std::move(offsets)); + CheckResolveMany(resolver, logical_index_vec); + } + + void TestRandomInput() { + const int64_t num_chunks = static_cast( + std::min(kMaxValidIndex - 1, static_cast(1) << 16)); + const int64_t avg_chunk_length = 20; + const int64_t chunked_array_len = num_chunks * 2 * avg_chunk_length; + TestRandomInput(num_chunks, chunked_array_len); + } }; TYPED_TEST_SUITE(TestChunkResolverMany, IndexTypes); @@ -517,5 +584,6 @@ TYPED_TEST_SUITE(TestChunkResolverMany, IndexTypes); TYPED_TEST(TestChunkResolverMany, Basics) { this->TestBasics(); } TYPED_TEST(TestChunkResolverMany, OutOfBounds) { this->TestOutOfBounds(); } TYPED_TEST(TestChunkResolverMany, Overflow) { this->TestOverflow(); } +TYPED_TEST(TestChunkResolverMany, RandomInput) { this->TestRandomInput(); } } // namespace arrow