Alex-PLACET commented on code in PR #49679: URL: https://github.com/apache/arrow/pull/49679#discussion_r3362602904
########## cpp/src/arrow/compute/kernels/vector_search_sorted.cc: ########## @@ -0,0 +1,1186 @@ +// 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_internal.h" +#include "arrow/compute/registry.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] +// + +// Dispatch on physical types rather than logical types to reduce template +// instantiations. Types sharing the same physical layout (e.g. Date32/Int32, +// Timestamp/Int64, String/Binary) share a single code path. +// HalfFloatType is NOT included here; it is handled separately before the +// physical-type switch because its physical type (UInt16) has different +// comparison semantics (Float16 NaN handling). +#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 const std::shared_ptr<Array>* FindFirstNonEmptyChunk(const ChunkedArray& values) { + const auto it = std::ranges::find_if( + values.chunks(), + [](const std::shared_ptr<Array>& chunk) { return chunk->length() != 0; }); Review Comment: Apache Arrow targets C++20 now ########## cpp/src/arrow/compute/kernels/vector_search_sorted.cc: ########## @@ -0,0 +1,1186 @@ +// 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_internal.h" +#include "arrow/compute/registry.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] +// + +// Dispatch on physical types rather than logical types to reduce template +// instantiations. Types sharing the same physical layout (e.g. Date32/Int32, +// Timestamp/Int64, String/Binary) share a single code path. +// HalfFloatType is NOT included here; it is handled separately before the +// physical-type switch because its physical type (UInt16) has different +// comparison semantics (Float16 NaN handling). +#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 const std::shared_ptr<Array>* FindFirstNonEmptyChunk(const ChunkedArray& values) { + const auto it = std::ranges::find_if( + values.chunks(), + [](const std::shared_ptr<Array>& chunk) { return chunk->length() != 0; }); + return it == values.chunks().end() ? nullptr : &*it; +} + +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}); Review Comment: Apache Arrow targets C++20 now -- 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]
