Copilot commented on code in PR #49679: URL: https://github.com/apache/arrow/pull/49679#discussion_r3551401401
########## cpp/src/arrow/compute/kernels/vector_search_sorted.cc: ########## @@ -0,0 +1,1166 @@ +// 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/compute/api_vector.h" + +#include <algorithm> +#include <memory> +#include <optional> +#include <ranges> +#include <type_traits> +#include <utility> + +#include "arrow/array/array_primitive.h" +#include "arrow/array/array_run_end.h" +#include "arrow/array/concatenate.h" +#include "arrow/array/util.h" +#include "arrow/buffer_builder.h" +#include "arrow/chunk_resolver.h" +#include "arrow/compute/function.h" +#include "arrow/compute/kernels/codegen_internal.h" +#include "arrow/compute/kernels/vector_sort_internal.h" +#include "arrow/compute/registry.h" +#include "arrow/compute/registry_internal.h" +#include "arrow/type_traits.h" +#include "arrow/util/checked_cast.h" +#include "arrow/util/logging_internal.h" +#include "arrow/util/ree_util.h" +#include "arrow/util/unreachable.h" + +namespace arrow { + +using internal::checked_cast; + +namespace compute::internal { +namespace { + +/// Return the static default options instance used by the meta-function. +const SearchSortedOptions* GetDefaultSearchSortedOptions() { + static const auto kDefaultSearchSortedOptions = SearchSortedOptions::Defaults(); + return &kDefaultSearchSortedOptions; +} + +const FunctionDoc search_sorted_doc( + "Find insertion indices for sorted input", + ("Return the index where each needle should be inserted in a sorted input array\n" + "to maintain ascending order.\n" + "\n" + "With side='left', returns the first suitable index (lower bound).\n" + "With side='right', returns the last suitable index (upper bound).\n" + "\n" + "The searched values may be provided as an array or chunked array and must\n" + "already be sorted in ascending order. Null values in the searched array are\n" + "supported when clustered entirely at the start or\n" + "entirely at the end. Non-null needles are matched only against the non-null\n" + "portion of the searched array. Needles may be a scalar, array, or chunked\n" + "array. Null needles emit nulls in the output."), + {"values", "needles"}, "SearchSortedOptions"); + +// This file implements search_sorted as a normalization pipeline around one +// typed binary-search core. +// +// The searched values are first validated, unwrapped to their logical type, +// and adapted to a uniform accessor interface. Plain arrays and chunked arrays +// expose logical element access directly. Run-end encoded (REE) arrays expose a +// search domain over physical runs while still translating insertion positions +// back to logical indices. The typed accessors are intentionally thin: shared +// offset, length, run-resolution, and logical-index bookkeeping lives in +// non-templated helpers so that the binary search is specialized by physical +// value type without cloning the surrounding normalization logic. +// +// Values null handling is normalized before any search happens. Nulls are only +// accepted when clustered entirely at the start or entirely at the end of the +// sorted values. The implementation computes the contiguous non-null logical +// window once and then searches only within that window. For REE values this +// requires logical null counting, because nullness lives in the values child +// rather than in a top-level validity bitmap. +// +// Needles are normalized before accessor dispatch where possible. Null scalar +// needles return a null scalar result immediately. Non-null scalar needles are +// materialized as length-1 arrays so the main implementation only needs one +// array-oriented emit path. REE needles are handled separately: the kernel +// searches each physical REE value once, rebuilds a temporary REE UInt64 +// result with the same logical run ends, and then run-end decodes it back to +// the dense public output shape. Plain array needles are visited element by +// element through one callback interface that propagates logical nulls. +// +// The actual comparison/search step is shared across all normalized inputs. +// After dispatching to the logical/physical Arrow representation, the kernel +// runs a lower-bound or upper-bound binary search depending on +// `SearchSortedOptions::side`, then maps the found position back to the caller- +// visible logical insertion index. +// +// Output materialization is centralized in a UInt64 builder with an optional +// validity bitmap. Non-null-only needles only build the values buffer, while +// nullable needles also emit the null bitmap. Logical null detection uses +// `ComputeLogicalNullCount()` so plain arrays, chunked arrays, and REE inputs +// all participate in the same output-nullability decision. +// +// High-level flow: +// +// values datum +// | +// +--> ValidateSortedValuesInput +// | +// +--> LogicalType / FindNonNullValuesRange +// | +// +--> VisitValuesAccessor +// | +// +--> PlainArrayAccessor +// | +// +--> RunEndEncodedValuesAccessor +// | +// +--> ChunkedArrayAccessor +// | +// `--> ChunkedRunEndEncodedValuesAccessor +// +// needles datum +// | +// +--> ValidateNeedleInput +// | +// +--> scalar null +// | `--> return null scalar +// | +// +--> scalar value +// | `--> MakeArrayFromScalar(length=1) +// | +// +--> REE needles +// | +--> search physical runs once +// | +--> rebuild temporary REE uint64 result +// | `--> RunEndDecode back to dense output +// | +// `--> DatumHasNulls / VisitNeedleRuns +// | +// +--> plain array -> one logical element per slot +// | +// `--> chunked input -> recurse chunk by chunk +// +// normalized values accessor + normalized needle runs +// | +// `--> FindInsertionPoint<T> +// | +// +--> side = left -> lower_bound semantics +// | +// `--> side = right -> upper_bound semantics +// +// result materialization +// | +// +--> no needle nulls +// | `--> InsertionIndexBuilder<false> +// | `--> fill uint64 buffer directly +// | +// `--> nullable needles +// `--> InsertionIndexBuilder<true> +// +--> AppendNulls for null runs +// `--> bulk fill repeated indices and validity bits +// +// A rough map of the file: +// +// [validation + type helpers] +// | +// [value accessors] +// | +// [needle visitors] +// | +// [typed search + output helpers] +// | +// [meta-function dispatch] +// + +#define VISIT_SEARCH_SORTED_PHYSICAL_TYPES(VISIT) \ + VISIT(BooleanType) \ + VISIT(Int8Type) \ + VISIT(Int16Type) \ + VISIT(Int32Type) \ + VISIT(Int64Type) \ + VISIT(UInt8Type) \ + VISIT(UInt16Type) \ + VISIT(UInt32Type) \ + VISIT(UInt64Type) \ + VISIT(FloatType) \ + VISIT(DoubleType) \ + VISIT(BinaryType) \ + VISIT(LargeBinaryType) \ + VISIT(BinaryViewType) Review Comment: The kernel dispatch only visits a subset of sortable physical types (no FixedSizeBinary/Decimal128/Decimal256), but the docs added in `docs/source/cpp/compute.rst` describe the input types as "Boolean, Numeric, Temporal, Binary- and String-like", which typically includes fixed-size binary and decimals in Arrow. Either narrow the documented supported types or extend the dispatch list to include the missing sortable physical types so behavior matches the documentation and other sort-related kernels. ########## cpp/src/arrow/compute/kernels/vector_search_sorted.cc: ########## @@ -0,0 +1,1166 @@ +// 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/compute/api_vector.h" + +#include <algorithm> +#include <memory> +#include <optional> +#include <ranges> +#include <type_traits> +#include <utility> + +#include "arrow/array/array_primitive.h" +#include "arrow/array/array_run_end.h" +#include "arrow/array/concatenate.h" +#include "arrow/array/util.h" +#include "arrow/buffer_builder.h" +#include "arrow/chunk_resolver.h" +#include "arrow/compute/function.h" +#include "arrow/compute/kernels/codegen_internal.h" +#include "arrow/compute/kernels/vector_sort_internal.h" +#include "arrow/compute/registry.h" +#include "arrow/compute/registry_internal.h" +#include "arrow/type_traits.h" +#include "arrow/util/checked_cast.h" +#include "arrow/util/logging_internal.h" +#include "arrow/util/ree_util.h" +#include "arrow/util/unreachable.h" + +namespace arrow { + +using internal::checked_cast; + +namespace compute::internal { +namespace { + +/// Return the static default options instance used by the meta-function. +const SearchSortedOptions* GetDefaultSearchSortedOptions() { + static const auto kDefaultSearchSortedOptions = SearchSortedOptions::Defaults(); + return &kDefaultSearchSortedOptions; +} + +const FunctionDoc search_sorted_doc( + "Find insertion indices for sorted input", + ("Return the index where each needle should be inserted in a sorted input array\n" + "to maintain ascending order.\n" + "\n" + "With side='left', returns the first suitable index (lower bound).\n" + "With side='right', returns the last suitable index (upper bound).\n" + "\n" + "The searched values may be provided as an array or chunked array and must\n" + "already be sorted in ascending order. Null values in the searched array are\n" + "supported when clustered entirely at the start or\n" + "entirely at the end. Non-null needles are matched only against the non-null\n" + "portion of the searched array. Needles may be a scalar, array, or chunked\n" + "array. Null needles emit nulls in the output."), + {"values", "needles"}, "SearchSortedOptions"); + +// This file implements search_sorted as a normalization pipeline around one +// typed binary-search core. +// +// The searched values are first validated, unwrapped to their logical type, +// and adapted to a uniform accessor interface. Plain arrays and chunked arrays +// expose logical element access directly. Run-end encoded (REE) arrays expose a +// search domain over physical runs while still translating insertion positions +// back to logical indices. The typed accessors are intentionally thin: shared +// offset, length, run-resolution, and logical-index bookkeeping lives in +// non-templated helpers so that the binary search is specialized by physical +// value type without cloning the surrounding normalization logic. +// +// Values null handling is normalized before any search happens. Nulls are only +// accepted when clustered entirely at the start or entirely at the end of the +// sorted values. The implementation computes the contiguous non-null logical +// window once and then searches only within that window. For REE values this +// requires logical null counting, because nullness lives in the values child +// rather than in a top-level validity bitmap. +// +// Needles are normalized before accessor dispatch where possible. Null scalar +// needles return a null scalar result immediately. Non-null scalar needles are +// materialized as length-1 arrays so the main implementation only needs one +// array-oriented emit path. REE needles are handled separately: the kernel +// searches each physical REE value once, rebuilds a temporary REE UInt64 +// result with the same logical run ends, and then run-end decodes it back to +// the dense public output shape. Plain array needles are visited element by +// element through one callback interface that propagates logical nulls. +// +// The actual comparison/search step is shared across all normalized inputs. +// After dispatching to the logical/physical Arrow representation, the kernel +// runs a lower-bound or upper-bound binary search depending on +// `SearchSortedOptions::side`, then maps the found position back to the caller- +// visible logical insertion index. +// +// Output materialization is centralized in a UInt64 builder with an optional +// validity bitmap. Non-null-only needles only build the values buffer, while +// nullable needles also emit the null bitmap. Logical null detection uses +// `ComputeLogicalNullCount()` so plain arrays, chunked arrays, and REE inputs +// all participate in the same output-nullability decision. +// +// High-level flow: +// +// values datum +// | +// +--> ValidateSortedValuesInput +// | +// +--> LogicalType / FindNonNullValuesRange +// | +// +--> VisitValuesAccessor +// | +// +--> PlainArrayAccessor +// | +// +--> RunEndEncodedValuesAccessor +// | +// +--> ChunkedArrayAccessor +// | +// `--> ChunkedRunEndEncodedValuesAccessor +// +// needles datum +// | +// +--> ValidateNeedleInput +// | +// +--> scalar null +// | `--> return null scalar +// | +// +--> scalar value +// | `--> MakeArrayFromScalar(length=1) +// | +// +--> REE needles +// | +--> search physical runs once +// | +--> rebuild temporary REE uint64 result +// | `--> RunEndDecode back to dense output +// | +// `--> DatumHasNulls / VisitNeedleRuns +// | +// +--> plain array -> one logical element per slot +// | +// `--> chunked input -> recurse chunk by chunk +// +// normalized values accessor + normalized needle runs +// | +// `--> FindInsertionPoint<T> +// | +// +--> side = left -> lower_bound semantics +// | +// `--> side = right -> upper_bound semantics +// +// result materialization +// | +// +--> no needle nulls +// | `--> InsertionIndexBuilder<false> +// | `--> fill uint64 buffer directly +// | +// `--> nullable needles +// `--> InsertionIndexBuilder<true> +// +--> AppendNulls for null runs +// `--> bulk fill repeated indices and validity bits +// +// A rough map of the file: +// +// [validation + type helpers] +// | +// [value accessors] +// | +// [needle visitors] +// | +// [typed search + output helpers] +// | +// [meta-function dispatch] +// + +#define VISIT_SEARCH_SORTED_PHYSICAL_TYPES(VISIT) \ + VISIT(BooleanType) \ + VISIT(Int8Type) \ + VISIT(Int16Type) \ + VISIT(Int32Type) \ + VISIT(Int64Type) \ + VISIT(UInt8Type) \ + VISIT(UInt16Type) \ + VISIT(UInt32Type) \ + VISIT(UInt64Type) \ + VISIT(FloatType) \ + VISIT(DoubleType) \ + VISIT(BinaryType) \ + VISIT(LargeBinaryType) \ + VISIT(BinaryViewType) + +template <typename ArrowType> +using SearchValue = typename GetViewType<ArrowType>::T; + +struct NonNullValuesRange { + int64_t offset = 0; + int64_t length = 0; + + /// Return whether the range spans the full searched values input. + bool is_identity(int64_t full_length) const { + return (offset == 0) && (length == full_length); + } +}; + +// Convert ArrayData to its physical representation so that typed accessors +// can be constructed with a physical ArrowType (e.g. Date32 → Int32). +// For REE arrays, only the values child type is converted; the REE wrapper +// type stays unchanged. +inline std::shared_ptr<ArrayData> ToPhysicalData( + const std::shared_ptr<ArrayData>& data, + const std::shared_ptr<DataType>& physical_type) { + if (data->type->id() == Type::RUN_END_ENCODED) { + auto result = data->Copy(); + auto values_copy = result->child_data[1]->Copy(); + values_copy->type = physical_type; + result->child_data[1] = std::move(values_copy); + return result; + } + auto result = data->Copy(); + result->type = physical_type; + return result; +} + +/// Read a run-end value from any supported run-end integer representation. +inline int64_t GetRunEndValue(const ArraySpan& run_ends, int64_t physical_index) { + switch (run_ends.type->id()) { + case Type::INT16: + return run_ends.GetValues<int16_t>(1)[physical_index]; + case Type::INT32: + return run_ends.GetValues<int32_t>(1)[physical_index]; + case Type::INT64: + return run_ends.GetValues<int64_t>(1)[physical_index]; + default: + DCHECK(false) << "Unexpected run-end type for search_sorted values: " + << run_ends.type->ToString(); + return 0; + } +} + +/// Comparator implementing Arrow's ascending-order semantics for supported types. +template <typename ArrowType> +struct SearchSortedCompare { + using ValueType = SearchValue<ArrowType>; + + int operator()(const ValueType& left, const ValueType& right) const { + return CompareTypeValues<ArrowType>(left, right, SortOrder::Ascending, + NullPlacement::AtEnd); + } +}; + +class SearchWindow { + public: + SearchWindow(int64_t length, NonNullValuesRange non_null_range) + : offset_(non_null_range.offset), + length_(non_null_range.is_identity(length) ? length : non_null_range.length) {} + + int64_t length() const { return length_; } + + uint64_t LogicalInsertionIndex(int64_t index) const { + return static_cast<uint64_t>(index); + } + + protected: + int64_t physical_offset() const { return offset_; } + + private: + int64_t offset_ = 0; + int64_t length_; +}; + +/// Access logical values from a plain Arrow array. +template <typename ArrowType> +class PlainArrayAccessor : public SearchWindow { + public: + using ArrayType = typename TypeTraits<ArrowType>::ArrayType; + using ValueType = SearchValue<ArrowType>; + + /// Build a typed accessor over a plain array payload, optionally restricted + /// to a non-null subrange. + explicit PlainArrayAccessor(const std::shared_ptr<ArrayData>& array_data, + NonNullValuesRange non_null_range = {}) + : SearchWindow(array_data->length, non_null_range), array_(array_data) {} + + /// Return the logical value at the given position within the search window. + ValueType Value(int64_t index) const { + return GetViewType<ArrowType>::LogicalValue( + array_.GetView(physical_offset() + index)); + } + + private: + ArrayType array_; +}; + +class RunEndEncodedValuesAccessorBase { + public: + explicit RunEndEncodedValuesAccessorBase(const RunEndEncodedArray& array, + NonNullValuesRange logical_non_null_range = {}) + : array_(array), + values_(array.values()), + array_span_(*array.data()), + physical_range_(::arrow::ree_util::FindPhysicalRange(array_span_, array.offset(), + array.length())) { + InitSearchWindow(logical_non_null_range); + } + + int64_t length() const { return search_length_; } + + int64_t NullCount() const { + return values_->Slice(physical_range_.first, physical_range_.second)->null_count(); + } + + uint64_t LogicalInsertionIndex(int64_t index) const { + DCHECK_GE(index, 0); + DCHECK_LE(index, search_length_); + return PhysicalLogicalInsertionIndex(search_offset_ + index) - + PhysicalLogicalInsertionIndex(search_offset_); + } + + int64_t logical_length() const { return array_.length(); } + + protected: + int64_t PhysicalIndex(int64_t index) const { + return physical_range_.first + search_offset_ + index; + } + + private: + /// Initialize the search window from a logical non-null range, converting to + /// physical-run coordinates. + void InitSearchWindow(const NonNullValuesRange& logical_non_null_range) { + search_length_ = physical_range_.second; + if (!logical_non_null_range.is_identity(array_.length())) { + const auto null_runs = NullCount(); + search_offset_ = logical_non_null_range.offset > 0 ? null_runs : 0; + search_length_ = physical_range_.second - null_runs; + } + } + + /// Compute a logical insertion index from an absolute physical run index + /// (before search-window offsetting). + uint64_t PhysicalLogicalInsertionIndex(int64_t physical_index) const { + if (physical_index == 0) { + return 0; + } + if (physical_index == physical_range_.second) { + return static_cast<uint64_t>(array_.length()); + } + return static_cast<uint64_t>( + LogicalRunEnd(physical_range_.first + physical_index - 1)); + } + + /// Return the logical run end corresponding to a physical run index. + int64_t LogicalRunEnd(int64_t physical_index) const { + // The run-end value is an absolute (cumulative) logical position in the + // full array. Subtract array_.offset() to get a position relative to the + // current slice. Clamp to 0, when the slice offset falls in the middle of + // a physical run the first runend after the slice start is always positive, + // but defensive clamping guards against edge cases where a run-end lands + // exactly at (or before) the slice offset. + const int64_t logical_run_end = std::max<int64_t>( + GetRunEndValue(::arrow::ree_util::RunEndsArray(array_span_), physical_index) - + array_.offset(), + 0); + // The physical range returned by FindPhysicalRange may include a trailing + // run that extends beyond the logical slice. Clamp to array_.length() so + // the result stays within the slice boundary. + return std::min(logical_run_end, array_.length()); + } + + const RunEndEncodedArray& array_; + std::shared_ptr<Array> values_; + ArraySpan array_span_; + std::pair<int64_t, int64_t> physical_range_; + int64_t search_offset_ = 0; + int64_t search_length_ = 0; +}; + +/// Access logical values from a run-end encoded Arrow array. +template <typename ArrowType> +class RunEndEncodedValuesAccessor : public RunEndEncodedValuesAccessorBase { + public: + using ArrayType = typename TypeTraits<ArrowType>::ArrayType; + using ValueType = SearchValue<ArrowType>; + + /// Build a typed accessor over a run-end encoded payload, optionally restricted + /// to a logical non-null subrange. + explicit RunEndEncodedValuesAccessor(const RunEndEncodedArray& array, + NonNullValuesRange logical_non_null_range = {}) + : RunEndEncodedValuesAccessorBase(array, logical_non_null_range), + values_(array.values()->data()) {} + + /// Return the logical value at the given physical run position within the + /// search window. + ValueType Value(int64_t index) const { + return GetViewType<ArrowType>::LogicalValue(values_.GetView(PhysicalIndex(index))); + } + + ArrayType values_; +}; + +/// Access logical values from a chunked Arrow array without combining chunks. +template <typename ArrowType> +class ChunkedArrayAccessor : public SearchWindow { + public: + using ArrayType = typename TypeTraits<ArrowType>::ArrayType; + using ValueType = SearchValue<ArrowType>; + + /// Build an accessor that resolves logical indices across chunk boundaries + /// without concatenating the input, optionally restricted to a subrange. + explicit ChunkedArrayAccessor(const ChunkedArray& chunked_array, + NonNullValuesRange non_null_range = {}) + : SearchWindow(chunked_array.length(), non_null_range), + chunked_array_(chunked_array), + resolver_(chunked_array.chunks()) { + chunks_.reserve(static_cast<size_t>(chunked_array.num_chunks())); + for (const auto& chunk : chunked_array.chunks()) { + DCHECK_NE(chunk->type_id(), Type::RUN_END_ENCODED); + chunks_.emplace_back(chunk->data()); + } + } + + /// Resolve a logical index within the search window to its chunk-local + /// storage and return that value. + ValueType Value(int64_t index) const { + const auto location = resolver_.Resolve(physical_offset() + index); + DCHECK_LT(location.chunk_index, chunked_array_.num_chunks()); + return GetViewType<ArrowType>::LogicalValue( + chunks_[location.chunk_index].GetView(location.index_in_chunk)); + } + + private: + const ChunkedArray& chunked_array_; + ChunkResolver resolver_; + std::vector<ArrayType> chunks_; +}; + +struct ChunkedRunEndEncodedMetadata { + std::vector<int64_t> logical_offsets; + std::vector<int64_t> run_offsets; + int64_t total_run_count = 0; +}; + +inline int64_t GetPhysicalRunCount(const RunEndEncodedArray& array) { + ArraySpan span(*array.data()); + return ::arrow::ree_util::FindPhysicalRange(span, array.offset(), array.length()) + .second; +} + +inline ChunkedRunEndEncodedMetadata MakeChunkedRunEndEncodedMetadata( + const ChunkedArray& chunked_array) { + ChunkedRunEndEncodedMetadata metadata; + metadata.logical_offsets.reserve(static_cast<size_t>(chunked_array.num_chunks())); + metadata.run_offsets.reserve(static_cast<size_t>(chunked_array.num_chunks()) + 1); + metadata.run_offsets.push_back(0); + + int64_t logical_offset = 0; + for (const auto& chunk : chunked_array.chunks()) { + if (chunk->length() == 0) { + continue; + } + DCHECK_EQ(chunk->type_id(), Type::RUN_END_ENCODED); + const auto& ree_chunk = checked_cast<const RunEndEncodedArray&>(*chunk); + metadata.logical_offsets.push_back(logical_offset); + logical_offset += chunk->length(); + metadata.total_run_count += GetPhysicalRunCount(ree_chunk); + metadata.run_offsets.push_back(metadata.total_run_count); + } + + DCHECK_EQ(logical_offset, chunked_array.length()); + return metadata; +} + +class ChunkedRunEndEncodedValuesAccessorBase { + public: + ChunkedRunEndEncodedValuesAccessorBase(int64_t logical_length, + ChunkedRunEndEncodedMetadata metadata) + : logical_length_(logical_length), + total_run_count_(metadata.total_run_count), + run_resolver_(std::move(metadata.run_offsets)), + logical_offsets_(std::move(metadata.logical_offsets)) {} + + int64_t length() const { return search_length_; } + + int64_t logical_length() const { return logical_length_; } + + protected: + void FinalizeSearchWindow(const NonNullValuesRange& logical_non_null_range, + int64_t null_runs) { + search_length_ = total_run_count_; + if (!logical_non_null_range.is_identity(logical_length_)) { + offset_ = logical_non_null_range.offset > 0 ? null_runs : 0; + search_length_ = total_run_count_ - null_runs; + } + } + + int64_t SearchOffset() const { return offset_; } + + int64_t PhysicalRunIndex(int64_t index) const { return offset_ + index; } + + std::pair<size_t, int64_t> ResolveRunIndex(int64_t index) const { + DCHECK_LT(index, total_run_count_); + const auto location = run_resolver_.Resolve(index); + return {static_cast<size_t>(location.chunk_index), location.index_in_chunk}; + } + + int64_t total_run_count() const { return total_run_count_; } + + int64_t logical_offset(size_t chunk_index) const { + return logical_offsets_[chunk_index]; + } + + private: + int64_t logical_length_; + int64_t total_run_count_ = 0; + ChunkResolver run_resolver_; + std::vector<int64_t> logical_offsets_; + int64_t offset_ = 0; + int64_t search_length_ = 0; +}; + +template <typename ArrowType> +class ChunkedRunEndEncodedValuesAccessor : public ChunkedRunEndEncodedValuesAccessorBase { + public: + using ValueType = SearchValue<ArrowType>; + + /// Flatten a chunked REE input into a logical sequence of physical runs while + /// preserving enough offset information to map search results back to logical + /// array positions, optionally restricted to a non-null subrange. + explicit ChunkedRunEndEncodedValuesAccessor( + const ChunkedArray& chunked_array, NonNullValuesRange logical_non_null_range = {}) + : ChunkedRunEndEncodedValuesAccessorBase( + chunked_array.length(), MakeChunkedRunEndEncodedMetadata(chunked_array)) { + const auto chunk_count = chunked_array.num_chunks(); + accessors_.reserve(static_cast<size_t>(chunk_count)); + for (const auto& chunk : chunked_array.chunks()) { + if (chunk->length() != 0) { + DCHECK_EQ(chunk->type_id(), Type::RUN_END_ENCODED); + const auto& ree_chunk = checked_cast<const RunEndEncodedArray&>(*chunk); + accessors_.emplace_back( + ree_chunk, NonNullValuesRange{.offset = 0, .length = ree_chunk.length()}); + } + } + FinalizeSearchWindow(logical_non_null_range, NullCount()); + } + + /// Translate a binary-search position within the search window back to a + /// logical insertion index relative to the window start. + uint64_t LogicalInsertionIndex(int64_t index) const { + DCHECK_GE(index, 0); + DCHECK_LE(index, length()); + + const auto phys_index = PhysicalRunIndex(index); + return RawLogicalInsertionIndex(phys_index) - + RawLogicalInsertionIndex(SearchOffset()); + } + + /// Resolve a global physical-run index within the search window to the + /// owning chunk accessor. + ValueType Value(int64_t index) const { + const auto [chunk_index, local_index] = ResolveRunIndex(PhysicalRunIndex(index)); + return accessors_[chunk_index].Value(local_index); + } + + /// Count leading null physical runs across chunks. Validation guarantees that + /// any null runs are clustered entirely at one end of the logical values. + int64_t NullCount() const { + int64_t null_run_count = 0; + for (const auto& accessor : accessors_) { + const auto local_null_run_count = accessor.NullCount(); + null_run_count += local_null_run_count; + if (local_null_run_count != accessor.length()) { + break; + } + } + return null_run_count; + } + + private: + uint64_t RawLogicalInsertionIndex(int64_t phys_index) const { + DCHECK_GE(phys_index, 0); + DCHECK_LE(phys_index, total_run_count()); + if (phys_index == 0) return 0; + if (phys_index == total_run_count()) return logical_length(); + const auto [chunk_index, local_index] = ResolveRunIndex(phys_index); + return static_cast<uint64_t>(logical_offset(chunk_index)) + + accessors_[chunk_index].LogicalInsertionIndex(local_index); + } + + std::vector<RunEndEncodedValuesAccessor<ArrowType>> accessors_; +}; + +/// Validate the supplied null counts and produce the logical non-null window +/// that will actually participate in binary search. +inline Result<NonNullValuesRange> MakeNonNullValuesRange(int64_t full_length, + int64_t null_count, + int64_t leading_null_count, + int64_t trailing_null_count) { + NonNullValuesRange non_null_values_range{.offset = 0, .length = full_length}; + + if (leading_null_count == full_length) { + non_null_values_range.length = 0; + return non_null_values_range; + } + + if (leading_null_count > 0) { + non_null_values_range.offset = leading_null_count; + non_null_values_range.length = full_length - leading_null_count; + return non_null_values_range; + } + + non_null_values_range.length = full_length - trailing_null_count; + return non_null_values_range; +} + +/// Build the searchable non-null window once the side containing clustered +/// nulls is already known. +inline Result<NonNullValuesRange> MakeNonNullValuesRangeFromNullPlacement( + int64_t full_length, int64_t null_count, bool has_leading_nulls) { + return MakeNonNullValuesRange(full_length, null_count, + has_leading_nulls ? null_count : 0, + has_leading_nulls ? 0 : null_count); +} + +/// Return whether a logical position in a chunked array is null. +inline bool IsNull(const ChunkedArray& values, int64_t index) { + DCHECK_GE(index, 0); + DCHECK_LT(index, values.length()); + + ChunkResolver resolver(values.chunks()); + const auto location = resolver.Resolve(index); + return values.chunk(static_cast<int>(location.chunk_index)) + ->IsNull(location.index_in_chunk); +} + +/// Infer the non-null search window from total null count plus a predicate that +/// can test whether any logical position is null. +template <typename IsNullAt> +inline Result<NonNullValuesRange> FindNonNullValuesRangeFromNullCount( + int64_t length, int64_t null_count, IsNullAt&& is_null_at) { + DCHECK_GT(null_count, 0); + DCHECK_LE(null_count, length); + + const bool has_leading_nulls = is_null_at(0); + return MakeNonNullValuesRangeFromNullPlacement(length, null_count, has_leading_nulls); +} Review Comment: `FindNonNullValuesRangeFromNullCount` assumes that all nulls in `values` are clustered at either the start or the end based only on `is_null_at(0)`, but it never verifies that the computed non-null window is actually null-free. If nulls appear in the middle, the search window can include null slots and the kernel will read/compare undefined values, yielding incorrect results silently. Consider validating that the inferred non-null window contains zero nulls and return an Invalid status otherwise. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
