Alex-PLACET commented on code in PR #49679: URL: https://github.com/apache/arrow/pull/49679#discussion_r3492386691
########## cpp/src/arrow/compute/kernels/vector_search_sorted.cc: ########## @@ -0,0 +1,1183 @@ +// 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 <numeric> +#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 { + +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 small pipeline that first normalizes +// Arrow input shapes and then runs one typed binary-search core on logical +// values. +// +// Plain arrays, run-end encoded arrays, chunked arrays, and scalar needles are +// all adapted into a common accessor and run-visitor model so the search logic +// does not care about physical layout. +// +// After validation, the kernel isolates the contiguous non-null window of the +// searched values, because nulls are only supported when clustered at one end. +// That window uses logical null counting for run-end encoded inputs, whose +// nulls live in the values child rather than in a top-level validity bitmap. +// +// Needles then follow one of two paths. Scalars and plain arrays go through a +// shared logical-run visitor: scalars become a single run, plain arrays become +// one-element runs, and chunked inputs recurse chunk by chunk. Run-end encoded +// needles take a simpler physical-run path: search each physical needle once, +// rebuild a temporary run-end encoded uint64 result with the same run ends, +// and run-end decode it back to the dense output shape. +// +// Output materialization is unified behind a typed-buffer builder with an +// optional validity bitmap. Non-null-only needles only build the uint64 values +// buffer, while nullable needles also emit a null bitmap. +// +// High-level flow: +// +// values datum +// | +// +--> ValidateSortedValuesInput +// | +// +--> LogicalType / FindNonNullValuesRange +// | +// +--> VisitValuesAccessor +// | +// +--> PlainArrayAccessor +// | +// +--> RunEndEncodedValuesAccessor +// | +// +--> ChunkedArrayAccessor +// | +// `--> ChunkedRunEndEncodedValuesAccessor +// +// needles datum +// | +// +--> ValidateNeedleInput +// | +// +--> DatumHasNulls +// | +// +--> REE needles +// | +--> search physical runs once +// | +--> rebuild temporary REE uint64 result +// | `--> RunEndDecode back to dense output +// | +// `--> VisitNeedleRuns +// | +// +--> scalar needle -> one logical run +// | +// +--> plain array -> one-element runs +// | +// `--> 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; +} + +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); + } +}; + +/// Access logical values from a plain Arrow array. +template <typename ArrowType> +class PlainArrayAccessor { + public: + using ArrayType = typename TypeTraits<ArrowType>::ArrayType; + using ValueType = SearchValue<ArrowType>; + + /// Build a typed accessor over a plain array payload. + explicit PlainArrayAccessor(const std::shared_ptr<ArrayData>& array_data) + : array_(array_data) {} + + /// Return the logical length of the searched values. + int64_t length() const { return array_.length(); } + + /// Return the logical value at the given logical position. + ValueType Value(int64_t index) const { + return GetViewType<ArrowType>::LogicalValue(array_.GetView(index)); + } + + uint64_t LogicalInsertionIndex(int64_t index) const { + return static_cast<uint64_t>(index); + } + + private: + ArrayType array_; +}; + +/// Access logical values from a run-end encoded Arrow array. +template <typename ArrowType> +class RunEndEncodedValuesAccessor { + public: + using ArrayType = typename TypeTraits<ArrowType>::ArrayType; + using ValueType = SearchValue<ArrowType>; + + /// Build a typed accessor over a run-end encoded payload. + explicit RunEndEncodedValuesAccessor(const RunEndEncodedArray& array) + : array_(array), + values_(array.values()->data()), + array_span_(*array.data()), + physical_range_(::arrow::ree_util::FindPhysicalRange(array_span_, array.offset(), + array.length())) {} + + /// Return the number of physical runs used as the search domain. + int64_t length() const { return physical_range_.second; } + + /// Return the logical value at the given physical run position. + ValueType Value(int64_t index) const { + const auto physical_index = physical_range_.first + index; + return GetViewType<ArrowType>::LogicalValue(values_.GetView(physical_index)); + } + + int64_t LeadingNullRunCount() const { + int64_t null_run_count = 0; + for (int64_t index = 0; index < physical_range_.second; ++index) { + if (!values_.IsNull(physical_range_.first + index)) { + break; + } + ++null_run_count; + } + return null_run_count; + } + + int64_t TrailingNullRunCount() const { + int64_t null_run_count = 0; + for (int64_t index = physical_range_.second; index > 0; --index) { + if (!values_.IsNull(physical_range_.first + index - 1)) { + break; + } + ++null_run_count; + } + return null_run_count; + } + + uint64_t LogicalInsertionIndex(int64_t index) const { + DCHECK_GE(index, 0); + DCHECK_LE(index, physical_range_.second); + + if (index == 0) { + return 0; + } + if (index == physical_range_.second) { + return static_cast<uint64_t>(array_.length()); + } + return static_cast<uint64_t>(LogicalRunEnd(physical_range_.first + index - 1)); + } + + int64_t logical_length() const { return array_.length(); } + + private: + 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_; + ArrayType values_; + ArraySpan array_span_; + std::pair<int64_t, int64_t> physical_range_; +}; + +/// Access logical values from a chunked Arrow array without combining chunks. +template <typename ArrowType> +class ChunkedArrayAccessor { + public: + using ArrayType = typename TypeTraits<ArrowType>::ArrayType; + using ValueType = SearchValue<ArrowType>; + + explicit ChunkedArrayAccessor(const ChunkedArray& chunked_array) + : 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()); + } + } + + int64_t length() const { return chunked_array_.length(); } + + ValueType Value(int64_t index) const { + const auto location = resolver_.Resolve(index); + DCHECK_LT(location.chunk_index, chunked_array_.num_chunks()); + return GetViewType<ArrowType>::LogicalValue( + chunks_[location.chunk_index].GetView(location.index_in_chunk)); + } + + uint64_t LogicalInsertionIndex(int64_t index) const { + return static_cast<uint64_t>(index); + } + + private: + const ChunkedArray& chunked_array_; + ChunkResolver resolver_; + std::vector<ArrayType> chunks_; +}; + +template <typename ArrowType> +class ChunkedRunEndEncodedValuesAccessor { + public: + using ValueType = SearchValue<ArrowType>; + + explicit ChunkedRunEndEncodedValuesAccessor(const ChunkedArray& chunked_array) + : chunked_array_(chunked_array), logical_length_(chunked_array.length()) { + const auto chunk_count = chunked_array_.num_chunks(); + run_offsets_.reserve(static_cast<size_t>(chunk_count)); + logical_offsets_.reserve(static_cast<size_t>(chunk_count)); + accessors_.reserve(static_cast<size_t>(chunk_count)); + + int64_t selected_run_start = 0; + int64_t selected_logical_start = 0; + + 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); + run_offsets_.push_back(selected_run_start); + logical_offsets_.push_back(selected_logical_start); + accessors_.emplace_back(ree_chunk); + + selected_run_start += accessors_.back().length(); + selected_logical_start += chunk->length(); + } + } + + DCHECK_EQ(selected_logical_start, logical_length_); + total_run_count_ = selected_run_start; + } + + int64_t length() const { return total_run_count_; } + + ValueType Value(int64_t index) const { + const auto [chunk_index, local_index] = ResolveRun(index); + return accessors_[chunk_index].Value(local_index); + } + + int64_t LeadingNullRunCount() const { + int64_t null_run_count = 0; + for (const auto& accessor : accessors_) { + const auto local_null_run_count = accessor.LeadingNullRunCount(); + null_run_count += local_null_run_count; + if (local_null_run_count != accessor.length()) { + break; + } + } + return null_run_count; + } + + int64_t TrailingNullRunCount() const { + int64_t null_run_count = 0; + for (auto it = accessors_.rbegin(); it != accessors_.rend(); ++it) { + const auto local_null_run_count = it->TrailingNullRunCount(); + null_run_count += local_null_run_count; + if (local_null_run_count != it->length()) { + break; + } + } + return null_run_count; + } + + uint64_t LogicalInsertionIndex(int64_t index) const { + DCHECK_GE(index, 0); + DCHECK_LE(index, total_run_count_); + + if (index == 0) { + return 0; + } + if (index == total_run_count_) { + return static_cast<uint64_t>(logical_length_); + } + + const auto [chunk_index, local_index] = ResolveRun(index); + return static_cast<uint64_t>(logical_offsets_[chunk_index]) + + accessors_[chunk_index].LogicalInsertionIndex(local_index); + } + + int64_t logical_length() const { return logical_length_; } + + private: + std::pair<size_t, int64_t> ResolveRun(int64_t index) const { + DCHECK_LT(index, total_run_count_); + const auto it = std::upper_bound(run_offsets_.begin(), run_offsets_.end(), index); + DCHECK_NE(it, run_offsets_.begin()); + const auto chunk_index = + static_cast<size_t>(std::distance(run_offsets_.begin(), it) - 1); + return {chunk_index, index - run_offsets_[chunk_index]}; + } + + const ChunkedArray& chunked_array_; + int64_t logical_length_; + int64_t total_run_count_ = 0; + std::vector<int64_t> run_offsets_; + std::vector<int64_t> logical_offsets_; + std::vector<RunEndEncodedValuesAccessor<ArrowType>> accessors_; +}; + +constexpr std::string_view kClusteredNullValuesError = + "search_sorted values with nulls must be clustered at the start or end."; + +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) { + if (leading_null_count != null_count) { + return Status::Invalid(kClusteredNullValuesError); + } + non_null_values_range.offset = leading_null_count; + non_null_values_range.length = full_length - leading_null_count; + return non_null_values_range; + } + + if (trailing_null_count == 0 || trailing_null_count != null_count) { + return Status::Invalid(kClusteredNullValuesError); + } + + non_null_values_range.length = full_length - trailing_null_count; + return non_null_values_range; +} + +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); +} + +inline int64_t GetLogicalNullCount(const ArrayData& values) { + if (!values.MayHaveLogicalNulls()) { + return 0; + } + if (values.type->id() == Type::RUN_END_ENCODED) { + return values.ComputeLogicalNullCount(); + } + return values.GetNullCount(); +} + +inline int64_t GetLogicalNullCount(const ChunkedArray& values) { + if (values.type()->id() != Type::RUN_END_ENCODED) { + return values.null_count(); + } + + auto chunk_null_counts = values.chunks() | std::views::transform([](const auto& chunk) { + return GetLogicalNullCount(*chunk->data()); + }); + return std::reduce(chunk_null_counts.begin(), chunk_null_counts.end(), int64_t{0}); +} + +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(location.chunk_index)->IsNull(location.index_in_chunk); +} + +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); + if (has_leading_nulls) { + if (!is_null_at(null_count - 1) || (null_count < length && is_null_at(null_count))) { + return Status::Invalid(kClusteredNullValuesError); + } + } else { + const auto first_trailing_null_index = length - null_count; + if (!is_null_at(first_trailing_null_index) || !is_null_at(length - 1) || + is_null_at(first_trailing_null_index - 1)) { + return Status::Invalid(kClusteredNullValuesError); + } + } + return MakeNonNullValuesRangeFromNullPlacement(length, null_count, has_leading_nulls); +} + +/// Present a contiguous non-null slice of the searched values through the same +/// accessor interface as the original values container. +template <typename ValuesAccessor> +class NonNullValuesAccessor { + public: + /// Wrap the original accessor with the discovered non-null subrange. + explicit NonNullValuesAccessor(const ValuesAccessor& values, + const NonNullValuesRange& non_null_values_range) + : values_(values), + offset_(non_null_values_range.offset), + length_(non_null_values_range.length), + base_insertion_index_(values_.LogicalInsertionIndex(offset_)) {} + + /// Return the number of accessible non-null values. + int64_t length() const noexcept { return length_; } + + /// Return the value at the given index within the non-null subrange. + auto Value(int64_t index) const { return values_.Value(offset_ + index); } + + uint64_t LogicalInsertionIndex(int64_t index) const { + return values_.LogicalInsertionIndex(offset_ + index) - base_insertion_index_; + } + + private: + const ValuesAccessor& values_; + int64_t offset_; + int64_t length_; + uint64_t base_insertion_index_; +}; + +/// Return the logical type of a datum, unwrapping run-end encoding when present. +inline const DataType& LogicalType(const Datum& datum) { + const auto& type = *datum.type(); + if (type.id() == Type::RUN_END_ENCODED) { + return *checked_cast<const RunEndEncodedType&>(type).value_type(); + } + return type; +} + +/// Return whether a scalar or array needle input contains any logical nulls. +inline bool DatumHasNulls(const Datum& datum) { + if (datum.is_scalar()) { + return !datum.scalar()->is_valid; + } + + if (datum.is_chunked_array()) { + const auto& chunked_array = *datum.chunked_array(); + if (chunked_array.null_count() > 0) { + return true; + } + if (chunked_array.type()->id() != Type::RUN_END_ENCODED) { + return false; + } + return std::ranges::any_of( + chunked_array.chunks(), [](const std::shared_ptr<Array>& chunk) { + const auto& ree_chunk = checked_cast<const RunEndEncodedArray&>(*chunk); + return ree_chunk.values()->null_count() != 0; + }); + } + + const auto& array_data = datum.array(); + const bool has_nulls = array_data->GetNullCount() > 0; + if (array_data->type->id() == Type::RUN_END_ENCODED) { + RunEndEncodedArray run_end_encoded(array_data); + return has_nulls || (run_end_encoded.values()->null_count() != 0); + } + return has_nulls; +} + +/// Reject nested run-end encoded values. TODO: Support this case in the future if there +/// is demand for it. +inline Status ValidateRunEndEncodedLogicalValueType(const DataType& type, + const char* name) { + const auto& ree_type = checked_cast<const RunEndEncodedType&>(type); + if (ree_type.value_type()->id() == Type::RUN_END_ENCODED) { + return Status::TypeError("Nested run-end encoded ", name, " are not supported"); + } + return Status::OK(); +} + +/// Compute the contiguous non-null window of the searched values. +/// +inline Result<NonNullValuesRange> FindNonNullValuesRange(const ArrayData& values) { + NonNullValuesRange non_null_values_range{.offset = 0, .length = values.length}; + + const auto null_count = GetLogicalNullCount(values); + if (null_count == 0) { + return non_null_values_range; + } + + return FindNonNullValuesRangeFromNullCount( + values.length, null_count, [&](int64_t index) { return values.IsNull(index); }); +} + +/// Validate the searched values input shape and supported encoding. +inline Status ValidateSortedValuesInput(const Datum& datum) { + if (!(datum.is_array() || datum.is_chunked_array())) { + return Status::TypeError("search_sorted values must be an array or chunked array"); + } + + const auto& type = *datum.type(); + if (type.id() == Type::RUN_END_ENCODED) { + return ValidateRunEndEncodedLogicalValueType(type, "values"); + } + + return Status::OK(); +} + +/// Validate the needles input shape and supported encoding. +/// Needles can be a scalar, array, or chunked array. Array-like needles must not have +/// nested run-end encoding since that is not currently supported. +inline Status ValidateNeedleInput(const Datum& datum) { + if (!(datum.is_array() || datum.is_chunked_array() || datum.is_scalar())) { + return Status::TypeError( + "search_sorted needles must be a scalar, array, or chunked array"); + } + + if ((datum.is_array() || datum.is_chunked_array()) && + datum.type()->id() == Type::RUN_END_ENCODED) { + return ValidateRunEndEncodedLogicalValueType(*datum.type(), "needles"); + } + return Status::OK(); +} + +inline Result<NonNullValuesRange> FindNonNullValuesRange(const ChunkedArray& values) { + NonNullValuesRange non_null_values_range{.offset = 0, .length = values.length()}; + + const auto null_count = GetLogicalNullCount(values); + if (null_count == 0) { + return non_null_values_range; + } + + return FindNonNullValuesRangeFromNullCount( + values.length(), null_count, [&](int64_t index) { return IsNull(values, index); }); +} + +/// Perform a lower- or upper-bound binary search over already sorted values. +template <SearchSortedOptions::Side side, typename ArrowType, typename Accessor> +uint64_t FindInsertionPointImpl(const Accessor& sorted_values, + const SearchValue<ArrowType>& needle) { + SearchSortedCompare<ArrowType> compare; + int64_t first = 0; + int64_t count = sorted_values.length(); + + // TODO(search_sorted): For fixed-width primitive haystacks, investigate a SIMD-friendly + // batched search path . + while (count > 0) { + const int64_t step = count / 2; + const int64_t it = first + step; + const bool advance = [&] { + if constexpr (side == SearchSortedOptions::Left) { + return compare(sorted_values.Value(it), needle) < 0; + } else { + return compare(needle, sorted_values.Value(it)) >= 0; + } + }(); + if (advance) { + first = it + 1; + count -= step + 1; + } else { + count = step; + } + } + return static_cast<uint64_t>(first); +} + +template <typename ArrowType, typename Accessor> +uint64_t FindInsertionPoint(const Accessor& sorted_values, + const SearchValue<ArrowType>& needle, + SearchSortedOptions::Side side) { + switch (side) { + case SearchSortedOptions::Left: + return FindInsertionPointImpl<SearchSortedOptions::Left, ArrowType>(sorted_values, + needle); + case SearchSortedOptions::Right: + return FindInsertionPointImpl<SearchSortedOptions::Right, ArrowType>(sorted_values, + needle); + } + ::arrow::Unreachable("Invalid SearchSortedOptions::Side value"); + return 0; +} + +template <typename ArrowType, typename Accessor> +uint64_t FindLogicalInsertionIndex(const Accessor& sorted_values, + const SearchValue<ArrowType>& needle, + SearchSortedOptions::Side side, + uint64_t insertion_offset) { + const auto search_index = + static_cast<int64_t>(FindInsertionPoint<ArrowType>(sorted_values, needle, side)); + return sorted_values.LogicalInsertionIndex(search_index) + insertion_offset; +} + +template <typename ArrowType> +using VisitedNeedle = std::optional<SearchValue<ArrowType>>; + +/// Normalize a non-null logical needle into the visitor payload type. +template <typename ArrowType> +VisitedNeedle<ArrowType> MakeVisitedNeedle(const SearchValue<ArrowType>& needle) { + return std::optional<SearchValue<ArrowType>>(needle); +} + +/// Read one logical needle value from a physical array position. +template <typename ArrowType, typename ArrayType> +VisitedNeedle<ArrowType> ReadVisitedNeedle(const ArrayType& array, + int64_t physical_index) { + if (array.IsNull(physical_index)) { + return std::nullopt; + } + const auto needle = GetViewType<ArrowType>::LogicalValue(array.GetView(physical_index)); + return MakeVisitedNeedle<ArrowType>(needle); +} + +/// Visit each plain-array needle as single-element logical runs. +template <typename ArrowType, typename Visitor> +Status VisitArrayNeedleRuns(const std::shared_ptr<ArrayData>& needles_data, + Visitor&& visitor) { + using ArrayType = typename TypeTraits<ArrowType>::ArrayType; + + auto physical_type = TypeTraits<ArrowType>::type_singleton(); + auto physical_data = ToPhysicalData(needles_data, physical_type); + ArrayType array(physical_data); + for (int64_t index = 0; index < array.length(); ++index) { + RETURN_NOT_OK(visitor(ReadVisitedNeedle<ArrowType>(array, index))); + } + return Status::OK(); +} + +/// Visit scalar, plain-array, run-end encoded, or chunked needles through a +/// uniform callback interface of logical run lengths. +template <typename ArrowType, typename Visitor> +Status VisitNeedleRuns(const Datum& needles, Visitor&& visitor) { + if (needles.is_scalar()) { + if (!needles.scalar()->is_valid) { + return visitor(std::optional<SearchValue<ArrowType>>{}); + } + ARROW_ASSIGN_OR_RAISE(auto scalar_array, MakeArrayFromScalar(*needles.scalar(), 1)); + return VisitArrayNeedleRuns<ArrowType>(scalar_array->data(), + std::forward<Visitor>(visitor)); + } + + if (needles.is_chunked_array()) { + for (const auto& chunk : needles.chunked_array()->chunks()) { + ARROW_RETURN_NOT_OK((VisitNeedleRuns<ArrowType>(Datum(chunk), visitor))); + } + return Status::OK(); + } + + const auto& needle_data = needles.array(); + return VisitArrayNeedleRuns<ArrowType>(needle_data, visitor); +} + +/// Build uint64 insertion-index arrays with an optional null bitmap. +class InsertionIndexBuilder { + public: + explicit InsertionIndexBuilder(MemoryPool* pool, bool nullable) + : indices_builder_(pool), null_bitmap_builder_(pool), nullable_(nullable) {} + + Status Init(int64_t length) { + expected_length_ = length; + RETURN_NOT_OK(indices_builder_.Reserve(length)); + if (nullable_) { + RETURN_NOT_OK(null_bitmap_builder_.Reserve(length)); + } + return Status::OK(); + } + + Status AppendNull() { + DCHECK_LE(length_ + 1, expected_length_); + DCHECK(nullable_); + indices_builder_.UnsafeAppend(uint64_t{0}); + null_bitmap_builder_.UnsafeAppend(false); + ++null_count_; + ++length_; + return Status::OK(); + } + + Status AppendValue(uint64_t insertion_index) { + DCHECK_LE(length_ + 1, expected_length_); + indices_builder_.UnsafeAppend(insertion_index); + if (nullable_) { + null_bitmap_builder_.UnsafeAppend(true); + } + ++length_; + return Status::OK(); + } + + Result<std::shared_ptr<Array>> Finish() && { + DCHECK_EQ(length_, expected_length_); + ARROW_ASSIGN_OR_RAISE(auto indices, indices_builder_.Finish()); + + std::shared_ptr<Buffer> null_bitmap; + if (nullable_) { + ARROW_ASSIGN_OR_RAISE(null_bitmap, null_bitmap_builder_.Finish()); + } + + return MakeArray(ArrayData::Make( + uint64(), length_, {std::move(null_bitmap), std::move(indices)}, null_count_)); + } + + private: + TypedBufferBuilder<uint64_t> indices_builder_; + TypedBufferBuilder<bool> null_bitmap_builder_; + bool nullable_; + int64_t expected_length_ = 0; + int64_t length_ = 0; + int64_t null_count_ = 0; +}; + +/// Visit normalized needle runs and emit insertion indices through an output +/// policy object. +template <typename ArrowType, typename ValuesAccessor, typename Output> +Status EmitInsertionIndices(const ValuesAccessor& sorted_values, const Datum& needles, + SearchSortedOptions::Side side, uint64_t insertion_offset, + Output* output) { + auto emit_search_result = [&](const VisitedNeedle<ArrowType>& needle) -> Status { + if (!needle.has_value()) { + return output->AppendNull(); + } + const auto insertion_index = FindLogicalInsertionIndex<ArrowType>( + sorted_values, *needle, side, insertion_offset); + return output->AppendValue(insertion_index); + }; + + return VisitNeedleRuns<ArrowType>(needles, emit_search_result); +} + +inline Result<Datum> ComputeRunEndEncodedNeedleInsertionIndices( + const Datum& values, const RunEndEncodedArray& needles, + SearchSortedOptions::Side side, ExecContext* ctx) { + ExecContext* exec_ctx = ctx != NULLPTR ? ctx : default_exec_context(); + + ARROW_ASSIGN_OR_RAISE(auto physical_results, + SearchSorted(values, Datum(needles.LogicalValues()), + SearchSortedOptions(side), exec_ctx)); + + ARROW_ASSIGN_OR_RAISE(auto logical_run_ends, + needles.LogicalRunEnds(exec_ctx->memory_pool())); + ARROW_ASSIGN_OR_RAISE(auto ree_result, + RunEndEncodedArray::Make(needles.length(), logical_run_ends, + physical_results.make_array())); + return RunEndDecode(Datum(ree_result), exec_ctx); +} + +/// Materialize output for scalar or array needles. +template <typename ArrowType, typename ValuesAccessor> +Result<Datum> ComputeInsertionIndices(const ValuesAccessor& sorted_values, + const Datum& values, const Datum& needles, + SearchSortedOptions::Side side, + uint64_t insertion_offset, ExecContext* ctx) { + if (needles.is_scalar()) { + auto scalar = needles.scalar(); + if (!scalar->is_valid) { + return Datum(std::make_shared<UInt64Scalar>()); + } + + ARROW_ASSIGN_OR_RAISE(auto scalar_arr, MakeArrayFromScalar(*scalar, 1)); + ARROW_ASSIGN_OR_RAISE(auto result, ComputeInsertionIndices<ArrowType>( + sorted_values, values, Datum(scalar_arr), side, + insertion_offset, ctx)); + ARROW_ASSIGN_OR_RAISE(auto result_scalar, result.make_array()->GetScalar(0)); + return Datum(std::move(result_scalar)); + } + + if (needles.type()->id() == Type::RUN_END_ENCODED) { + if (needles.is_array()) { + return ComputeRunEndEncodedNeedleInsertionIndices( + values, RunEndEncodedArray(needles.array()), side, ctx); + } + + std::vector<std::shared_ptr<Array>> decoded_chunks; + decoded_chunks.reserve(static_cast<size_t>(needles.chunked_array()->num_chunks())); + for (const auto& chunk : needles.chunked_array()->chunks()) { + ARROW_ASSIGN_OR_RAISE( + auto decoded_chunk, + ComputeRunEndEncodedNeedleInsertionIndices( + values, checked_cast<const RunEndEncodedArray&>(*chunk), side, ctx)); + decoded_chunks.push_back(decoded_chunk.make_array()); + } + ARROW_ASSIGN_OR_RAISE(auto out, Concatenate(decoded_chunks, ctx->memory_pool())); + return Datum(std::move(out)); + } + + auto has_nulls = DatumHasNulls(needles); + InsertionIndexBuilder output(ctx->memory_pool(), has_nulls); + ARROW_RETURN_NOT_OK(output.Init(needles.length())); + ARROW_RETURN_NOT_OK((EmitInsertionIndices<ArrowType>(sorted_values, needles, side, + insertion_offset, &output))); + ARROW_ASSIGN_OR_RAISE(auto out, std::move(output).Finish()); + return Datum(std::move(out)); +} + +// Main entry point for search_sorted over a single array of sorted values and scalar or +// array needles. Handles null presence in the needles and dispatches to the appropriate +// search implementation. +template <typename ArrowType, typename ValuesAccessor> +Result<Datum> SearchWithAccessor(const ValuesAccessor& values_accessor, + const NonNullValuesRange& non_null_values_range, + const Datum& values, const Datum& needles, + SearchSortedOptions::Side side, ExecContext* ctx) { + if (non_null_values_range.is_identity(values_accessor.length())) { + return ComputeInsertionIndices<ArrowType>(values_accessor, values, needles, side, + /*insertion_offset=*/0, ctx); + } + + NonNullValuesAccessor non_null_values(values_accessor, non_null_values_range); + return ComputeInsertionIndices<ArrowType>( + non_null_values, values, needles, side, + static_cast<uint64_t>(non_null_values_range.offset), ctx); +} + +template <typename ValuesAccessor> +NonNullValuesRange MakePhysicalNonNullValuesRange( + const ValuesAccessor& values_accessor, + const NonNullValuesRange& non_null_values_range) { + const auto leading_null_run_count = + non_null_values_range.offset > 0 ? values_accessor.LeadingNullRunCount() : 0; + const auto trailing_null_run_count = + non_null_values_range.offset > 0 ? 0 : values_accessor.TrailingNullRunCount(); + return NonNullValuesRange{.offset = leading_null_run_count, + .length = values_accessor.length() - leading_null_run_count - + trailing_null_run_count}; +} + +template <typename ArrowType, typename ValuesAccessor> +Result<Datum> SearchWithRunEndEncodedAccessor( + const ValuesAccessor& values_accessor, + const NonNullValuesRange& non_null_values_range, const Datum& values, + const Datum& needles, SearchSortedOptions::Side side, ExecContext* ctx) { + if (non_null_values_range.is_identity(values_accessor.logical_length())) { + return ComputeInsertionIndices<ArrowType>(values_accessor, values, needles, side, + /*insertion_offset=*/0, ctx); + } + + NonNullValuesAccessor non_null_values( + values_accessor, + MakePhysicalNonNullValuesRange(values_accessor, non_null_values_range)); + return ComputeInsertionIndices<ArrowType>( + non_null_values, values, needles, side, + static_cast<uint64_t>(non_null_values_range.offset), ctx); +} + +template <typename ArrowType> +Result<Datum> SearchWithAccessor( + const RunEndEncodedValuesAccessor<ArrowType>& values_accessor, + const NonNullValuesRange& non_null_values_range, const Datum& values, + const Datum& needles, SearchSortedOptions::Side side, ExecContext* ctx) { + return SearchWithRunEndEncodedAccessor<ArrowType>( + values_accessor, non_null_values_range, values, needles, side, ctx); +} + +template <typename ArrowType> +Result<Datum> SearchWithAccessor( + const ChunkedRunEndEncodedValuesAccessor<ArrowType>& values_accessor, + const NonNullValuesRange& non_null_values_range, const Datum& values, + const Datum& needles, SearchSortedOptions::Side side, ExecContext* ctx) { + return SearchWithRunEndEncodedAccessor<ArrowType>( + values_accessor, non_null_values_range, values, needles, side, ctx); +} + +// Meta-function implementation for the search_sorted public compute entrypoint. +template <typename ArrowType, typename Visitor> +Result<Datum> VisitValuesAccessor(const std::shared_ptr<ArrayData>& values_data, + Visitor&& visitor) { + auto physical_type = TypeTraits<ArrowType>::type_singleton(); + auto physical_data = ToPhysicalData(values_data, physical_type); + + if (physical_data->type->id() == Type::RUN_END_ENCODED) { + RunEndEncodedArray ree(physical_data); + RunEndEncodedValuesAccessor<ArrowType> values_accessor(ree); + return visitor(values_accessor); + } + + PlainArrayAccessor<ArrowType> values_accessor(physical_data); + return visitor(values_accessor); +} + +template <typename ArrowType, typename Visitor> +Result<Datum> VisitValuesAccessor(const ChunkedArray& values, Visitor&& visitor) { + auto physical_type = TypeTraits<ArrowType>::type_singleton(); + + if (values.type()->id() == Type::RUN_END_ENCODED) { + // Convert each REE chunk's values child to the physical type. + ArrayVector physical_chunks; + physical_chunks.reserve(static_cast<size_t>(values.num_chunks())); + for (const auto& chunk : values.chunks()) { + physical_chunks.push_back(MakeArray(ToPhysicalData(chunk->data(), physical_type))); + } + ChunkedArray physical_chunked(std::move(physical_chunks)); + ChunkedRunEndEncodedValuesAccessor<ArrowType> values_accessor(physical_chunked); + return visitor(values_accessor); + } + + // For plain chunked arrays, convert each chunk to the physical type. + auto physical_chunks = GetPhysicalChunks(values, physical_type); + ChunkedArray physical_chunked(std::move(physical_chunks), physical_type); + ChunkedArrayAccessor<ArrowType> values_accessor(physical_chunked); + return visitor(values_accessor); +} + +/// Meta-function implementation for the search_sorted public compute entrypoint. +/// Validates input shapes and types, normalizes to logical value accessors, and +/// dispatches to the typed search implementation. +class SearchSortedMetaFunction : public MetaFunction { + public: + /// Construct the registry entry with default options and documentation. + SearchSortedMetaFunction() + : MetaFunction("search_sorted", Arity::Binary(), search_sorted_doc, + GetDefaultSearchSortedOptions()) {} + + /// Validate inputs, normalize options, and dispatch to the typed search implementation. + Result<Datum> ExecuteImpl(const std::vector<Datum>& args, + const FunctionOptions* options, + ExecContext* ctx) const override { + RETURN_NOT_OK(ValidateSortedValuesInput(args[0])); + RETURN_NOT_OK(ValidateNeedleInput(args[1])); + + const auto& values_type = LogicalType(args[0]); + const auto& needles_type = LogicalType(args[1]); + if (!values_type.Equals(needles_type)) { + return Status::TypeError( + "search_sorted arguments must have matching logical types, got ", + values_type.ToString(), " and ", needles_type.ToString()); + } + + ARROW_ASSIGN_OR_RAISE(auto non_null_values_range, FindNonNullValuesRange(args[0])); + auto result = DispatchByType(args[0], non_null_values_range, args[1], + static_cast<const SearchSortedOptions&>(*options), ctx); + return result; + } + + private: + [[nodiscard]] Result<NonNullValuesRange> FindNonNullValuesRange( + const Datum& values) const { + if (values.is_chunked_array()) { + return ::arrow::compute::internal::FindNonNullValuesRange(*values.chunked_array()); + } + return ::arrow::compute::internal::FindNonNullValuesRange(*values.array()); + } + + /// Dispatch the logical value type to the matching template specialization. + /// Resolves logical types to physical types via GetPhysicalType() so that + /// types sharing the same physical layout (e.g. Date32/Int32, String/Binary) + /// share a single code path, reducing template instantiations. + Result<Datum> DispatchByType(const Datum& values, + const NonNullValuesRange& non_null_values_range, + const Datum& needles, const SearchSortedOptions& options, + ExecContext* ctx) const { + // Resolve to logical type first (stripping REE wrapper if present). + auto logical_type_ptr = values.type(); + if (logical_type_ptr->id() == Type::RUN_END_ENCODED) { + logical_type_ptr = + checked_cast<const RunEndEncodedType&>(*logical_type_ptr).value_type(); + } + + // HalfFloatType must keep its logical type because its physical type + // (UInt16) uses different comparison semantics (Float16 NaN handling). + if (logical_type_ptr->id() == Type::HALF_FLOAT) { Review Comment: You are right, I fixed it ########## cpp/src/arrow/compute/kernels/vector_search_sorted.cc: ########## @@ -0,0 +1,1183 @@ +// 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 <numeric> +#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 { + +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 small pipeline that first normalizes +// Arrow input shapes and then runs one typed binary-search core on logical +// values. +// +// Plain arrays, run-end encoded arrays, chunked arrays, and scalar needles are +// all adapted into a common accessor and run-visitor model so the search logic +// does not care about physical layout. +// +// After validation, the kernel isolates the contiguous non-null window of the +// searched values, because nulls are only supported when clustered at one end. +// That window uses logical null counting for run-end encoded inputs, whose +// nulls live in the values child rather than in a top-level validity bitmap. +// +// Needles then follow one of two paths. Scalars and plain arrays go through a +// shared logical-run visitor: scalars become a single run, plain arrays become +// one-element runs, and chunked inputs recurse chunk by chunk. Run-end encoded +// needles take a simpler physical-run path: search each physical needle once, +// rebuild a temporary run-end encoded uint64 result with the same run ends, +// and run-end decode it back to the dense output shape. +// +// Output materialization is unified behind a typed-buffer builder with an +// optional validity bitmap. Non-null-only needles only build the uint64 values +// buffer, while nullable needles also emit a null bitmap. +// +// High-level flow: +// +// values datum +// | +// +--> ValidateSortedValuesInput +// | +// +--> LogicalType / FindNonNullValuesRange +// | +// +--> VisitValuesAccessor +// | +// +--> PlainArrayAccessor +// | +// +--> RunEndEncodedValuesAccessor +// | +// +--> ChunkedArrayAccessor +// | +// `--> ChunkedRunEndEncodedValuesAccessor +// +// needles datum +// | +// +--> ValidateNeedleInput +// | +// +--> DatumHasNulls +// | +// +--> REE needles +// | +--> search physical runs once +// | +--> rebuild temporary REE uint64 result +// | `--> RunEndDecode back to dense output +// | +// `--> VisitNeedleRuns +// | +// +--> scalar needle -> one logical run +// | +// +--> plain array -> one-element runs +// | +// `--> 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; +} + +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); + } +}; + +/// Access logical values from a plain Arrow array. +template <typename ArrowType> +class PlainArrayAccessor { + public: + using ArrayType = typename TypeTraits<ArrowType>::ArrayType; + using ValueType = SearchValue<ArrowType>; + + /// Build a typed accessor over a plain array payload. + explicit PlainArrayAccessor(const std::shared_ptr<ArrayData>& array_data) + : array_(array_data) {} + + /// Return the logical length of the searched values. + int64_t length() const { return array_.length(); } + + /// Return the logical value at the given logical position. + ValueType Value(int64_t index) const { + return GetViewType<ArrowType>::LogicalValue(array_.GetView(index)); + } + + uint64_t LogicalInsertionIndex(int64_t index) const { + return static_cast<uint64_t>(index); + } + + private: + ArrayType array_; +}; + +/// Access logical values from a run-end encoded Arrow array. +template <typename ArrowType> +class RunEndEncodedValuesAccessor { + public: + using ArrayType = typename TypeTraits<ArrowType>::ArrayType; + using ValueType = SearchValue<ArrowType>; + + /// Build a typed accessor over a run-end encoded payload. + explicit RunEndEncodedValuesAccessor(const RunEndEncodedArray& array) + : array_(array), + values_(array.values()->data()), + array_span_(*array.data()), + physical_range_(::arrow::ree_util::FindPhysicalRange(array_span_, array.offset(), + array.length())) {} + + /// Return the number of physical runs used as the search domain. + int64_t length() const { return physical_range_.second; } + + /// Return the logical value at the given physical run position. + ValueType Value(int64_t index) const { + const auto physical_index = physical_range_.first + index; + return GetViewType<ArrowType>::LogicalValue(values_.GetView(physical_index)); + } + + int64_t LeadingNullRunCount() const { + int64_t null_run_count = 0; + for (int64_t index = 0; index < physical_range_.second; ++index) { + if (!values_.IsNull(physical_range_.first + index)) { + break; + } + ++null_run_count; + } + return null_run_count; + } + + int64_t TrailingNullRunCount() const { + int64_t null_run_count = 0; + for (int64_t index = physical_range_.second; index > 0; --index) { + if (!values_.IsNull(physical_range_.first + index - 1)) { + break; + } + ++null_run_count; + } + return null_run_count; + } + + uint64_t LogicalInsertionIndex(int64_t index) const { + DCHECK_GE(index, 0); + DCHECK_LE(index, physical_range_.second); + + if (index == 0) { + return 0; + } + if (index == physical_range_.second) { + return static_cast<uint64_t>(array_.length()); + } + return static_cast<uint64_t>(LogicalRunEnd(physical_range_.first + index - 1)); + } + + int64_t logical_length() const { return array_.length(); } + + private: + 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_; + ArrayType values_; + ArraySpan array_span_; + std::pair<int64_t, int64_t> physical_range_; +}; + +/// Access logical values from a chunked Arrow array without combining chunks. +template <typename ArrowType> +class ChunkedArrayAccessor { + public: + using ArrayType = typename TypeTraits<ArrowType>::ArrayType; + using ValueType = SearchValue<ArrowType>; + + explicit ChunkedArrayAccessor(const ChunkedArray& chunked_array) + : 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()); + } + } + + int64_t length() const { return chunked_array_.length(); } + + ValueType Value(int64_t index) const { + const auto location = resolver_.Resolve(index); + DCHECK_LT(location.chunk_index, chunked_array_.num_chunks()); + return GetViewType<ArrowType>::LogicalValue( + chunks_[location.chunk_index].GetView(location.index_in_chunk)); + } + + uint64_t LogicalInsertionIndex(int64_t index) const { + return static_cast<uint64_t>(index); + } + + private: + const ChunkedArray& chunked_array_; + ChunkResolver resolver_; + std::vector<ArrayType> chunks_; +}; + +template <typename ArrowType> +class ChunkedRunEndEncodedValuesAccessor { + public: + using ValueType = SearchValue<ArrowType>; + + explicit ChunkedRunEndEncodedValuesAccessor(const ChunkedArray& chunked_array) + : chunked_array_(chunked_array), logical_length_(chunked_array.length()) { + const auto chunk_count = chunked_array_.num_chunks(); + run_offsets_.reserve(static_cast<size_t>(chunk_count)); + logical_offsets_.reserve(static_cast<size_t>(chunk_count)); + accessors_.reserve(static_cast<size_t>(chunk_count)); + + int64_t selected_run_start = 0; + int64_t selected_logical_start = 0; + + 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); + run_offsets_.push_back(selected_run_start); + logical_offsets_.push_back(selected_logical_start); + accessors_.emplace_back(ree_chunk); + + selected_run_start += accessors_.back().length(); + selected_logical_start += chunk->length(); + } + } + + DCHECK_EQ(selected_logical_start, logical_length_); + total_run_count_ = selected_run_start; + } + + int64_t length() const { return total_run_count_; } + + ValueType Value(int64_t index) const { + const auto [chunk_index, local_index] = ResolveRun(index); + return accessors_[chunk_index].Value(local_index); + } + + int64_t LeadingNullRunCount() const { + int64_t null_run_count = 0; + for (const auto& accessor : accessors_) { + const auto local_null_run_count = accessor.LeadingNullRunCount(); + null_run_count += local_null_run_count; + if (local_null_run_count != accessor.length()) { + break; + } + } + return null_run_count; + } + + int64_t TrailingNullRunCount() const { + int64_t null_run_count = 0; + for (auto it = accessors_.rbegin(); it != accessors_.rend(); ++it) { + const auto local_null_run_count = it->TrailingNullRunCount(); + null_run_count += local_null_run_count; + if (local_null_run_count != it->length()) { + break; + } + } + return null_run_count; + } + + uint64_t LogicalInsertionIndex(int64_t index) const { + DCHECK_GE(index, 0); + DCHECK_LE(index, total_run_count_); + + if (index == 0) { + return 0; + } + if (index == total_run_count_) { + return static_cast<uint64_t>(logical_length_); + } + + const auto [chunk_index, local_index] = ResolveRun(index); + return static_cast<uint64_t>(logical_offsets_[chunk_index]) + + accessors_[chunk_index].LogicalInsertionIndex(local_index); + } + + int64_t logical_length() const { return logical_length_; } + + private: + std::pair<size_t, int64_t> ResolveRun(int64_t index) const { + DCHECK_LT(index, total_run_count_); + const auto it = std::upper_bound(run_offsets_.begin(), run_offsets_.end(), index); + DCHECK_NE(it, run_offsets_.begin()); + const auto chunk_index = + static_cast<size_t>(std::distance(run_offsets_.begin(), it) - 1); + return {chunk_index, index - run_offsets_[chunk_index]}; + } + + const ChunkedArray& chunked_array_; + int64_t logical_length_; + int64_t total_run_count_ = 0; + std::vector<int64_t> run_offsets_; + std::vector<int64_t> logical_offsets_; + std::vector<RunEndEncodedValuesAccessor<ArrowType>> accessors_; +}; + +constexpr std::string_view kClusteredNullValuesError = + "search_sorted values with nulls must be clustered at the start or end."; + +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) { + if (leading_null_count != null_count) { + return Status::Invalid(kClusteredNullValuesError); + } + non_null_values_range.offset = leading_null_count; + non_null_values_range.length = full_length - leading_null_count; + return non_null_values_range; + } + + if (trailing_null_count == 0 || trailing_null_count != null_count) { + return Status::Invalid(kClusteredNullValuesError); + } + + non_null_values_range.length = full_length - trailing_null_count; + return non_null_values_range; +} + +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); +} + +inline int64_t GetLogicalNullCount(const ArrayData& values) { + if (!values.MayHaveLogicalNulls()) { + return 0; + } + if (values.type->id() == Type::RUN_END_ENCODED) { + return values.ComputeLogicalNullCount(); + } + return values.GetNullCount(); Review Comment: Fixed -- 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]
