GH-49677 [Python][C++][Compute] Add search sorted compute kernel (#49679)
### Rationale for this change
Add the implemenation of the search sorted compute kernel based on the numpy function: https://numpy.org/doc/stable/reference/generated/numpy.searchsorted.html
### What changes are included in this PR?
Implementation of the C++ kernel + Python API.
Tests in C++ and Python
### Are these changes tested?
Yes
### Are there any user-facing changes?
No breaking change
* GitHub Issue: #49677
Lead-authored-by: Alexis Placet <2400067+Alex-PLACET@users.noreply.github.com>
Co-authored-by: Alexis Placet <alexis.placet.dev@pm.me>
Co-authored-by: Antoine Pitrou <antoine@python.org>
Co-authored-by: Copilot <copilot@github.com>
Signed-off-by: Antoine Pitrou <antoine@python.org>
diff --git a/cpp/src/arrow/CMakeLists.txt b/cpp/src/arrow/CMakeLists.txt
index f6671fe..d2b3eac 100644
--- a/cpp/src/arrow/CMakeLists.txt
+++ b/cpp/src/arrow/CMakeLists.txt
@@ -871,6 +871,7 @@
compute/kernels/vector_rank.cc
compute/kernels/vector_replace.cc
compute/kernels/vector_run_end_encode.cc
+ compute/kernels/vector_search_sorted.cc
compute/kernels/vector_select_k.cc
compute/kernels/vector_sort.cc
compute/kernels/vector_statistics.cc
diff --git a/cpp/src/arrow/array/array_binary.cc b/cpp/src/arrow/array/array_binary.cc
index 1266819..e0f1938 100644
--- a/cpp/src/arrow/array/array_binary.cc
+++ b/cpp/src/arrow/array/array_binary.cc
@@ -91,7 +91,7 @@
Status LargeStringArray::ValidateUTF8() const { return internal::ValidateUTF8(*data_); }
BinaryViewArray::BinaryViewArray(std::shared_ptr<ArrayData> data) {
- ARROW_CHECK_EQ(data->type->id(), Type::BINARY_VIEW);
+ ARROW_CHECK(is_binary_view_like(data->type->id()));
SetData(std::move(data));
}
diff --git a/cpp/src/arrow/compute/api_vector.cc b/cpp/src/arrow/compute/api_vector.cc
index 57b2eda..048b6a8 100644
--- a/cpp/src/arrow/compute/api_vector.cc
+++ b/cpp/src/arrow/compute/api_vector.cc
@@ -51,6 +51,7 @@
using compute::NullPlacement;
using compute::RankOptions;
using compute::RankQuantileOptions;
+using compute::SearchSortedOptions;
template <>
struct EnumTraits<FilterOptions::NullSelectionBehavior>
@@ -82,6 +83,22 @@
return "<INVALID>";
}
};
+
+template <>
+struct EnumTraits<SearchSortedOptions::Side>
+ : BasicEnumTraits<SearchSortedOptions::Side, SearchSortedOptions::Left,
+ SearchSortedOptions::Right> {
+ static std::string name() { return "SearchSortedOptions::Side"; }
+ static std::string value_name(SearchSortedOptions::Side value) {
+ switch (value) {
+ case SearchSortedOptions::Left:
+ return "Left";
+ case SearchSortedOptions::Right:
+ return "Right";
+ }
+ return "<INVALID>";
+ }
+};
template <>
struct EnumTraits<RankOptions::Tiebreaker>
: BasicEnumTraits<RankOptions::Tiebreaker, RankOptions::Min, RankOptions::Max,
@@ -125,6 +142,8 @@
static auto kArraySortOptionsType = GetFunctionOptionsType<ArraySortOptions>(
DataMember("order", &ArraySortOptions::order),
DataMember("null_placement", &ArraySortOptions::null_placement));
+static auto kSearchSortedOptionsType = GetFunctionOptionsType<SearchSortedOptions>(
+ DataMember("side", &SearchSortedOptions::side));
static auto kSortOptionsType = GetFunctionOptionsType<SortOptions>(
CoercedDataMember("sort_keys", &SortOptions::sort_keys, &SortOptions::GetSortKeys));
static auto kPartitionNthOptionsType = GetFunctionOptionsType<PartitionNthOptions>(
@@ -182,12 +201,15 @@
null_placement(null_placement) {}
constexpr char ArraySortOptions::kTypeName[];
+SearchSortedOptions::SearchSortedOptions(SearchSortedOptions::Side side)
+ : FunctionOptions(internal::kSearchSortedOptionsType), side(side) {}
+constexpr char SearchSortedOptions::kTypeName[];
+
ARROW_SUPPRESS_DEPRECATION_WARNING
SortOptions::SortOptions(std::vector<SortKey> sort_keys)
: FunctionOptions(internal::kSortOptionsType),
sort_keys(std::move(sort_keys)),
null_placement(std::nullopt) {}
-
SortOptions::SortOptions(std::vector<SortKey> sort_keys,
std::optional<NullPlacement> null_placement)
: FunctionOptions(internal::kSortOptionsType),
@@ -277,6 +299,7 @@
DCHECK_OK(registry->AddFunctionOptionsType(kDictionaryEncodeOptionsType));
DCHECK_OK(registry->AddFunctionOptionsType(kRunEndEncodeOptionsType));
DCHECK_OK(registry->AddFunctionOptionsType(kArraySortOptionsType));
+ DCHECK_OK(registry->AddFunctionOptionsType(kSearchSortedOptionsType));
DCHECK_OK(registry->AddFunctionOptionsType(kSortOptionsType));
DCHECK_OK(registry->AddFunctionOptionsType(kPartitionNthOptionsType));
DCHECK_OK(registry->AddFunctionOptionsType(kSelectKOptionsType));
@@ -318,6 +341,11 @@
return result.make_array();
}
+Result<Datum> SearchSorted(const Datum& values, const Datum& needles,
+ const SearchSortedOptions& options, ExecContext* ctx) {
+ return CallFunction("search_sorted", {values, needles}, &options, ctx);
+}
+
Result<Datum> ReplaceWithMask(const Datum& values, const Datum& mask,
const Datum& replacements, ExecContext* ctx) {
return CallFunction("replace_with_mask", {values, mask, replacements}, ctx);
diff --git a/cpp/src/arrow/compute/api_vector.h b/cpp/src/arrow/compute/api_vector.h
index 2de7137..b55ed9e 100644
--- a/cpp/src/arrow/compute/api_vector.h
+++ b/cpp/src/arrow/compute/api_vector.h
@@ -103,6 +103,21 @@
NullPlacement null_placement;
};
+class ARROW_EXPORT SearchSortedOptions : public FunctionOptions {
+ public:
+ enum Side {
+ Left,
+ Right,
+ };
+
+ explicit SearchSortedOptions(Side side = Side::Left);
+ static constexpr const char kTypeName[] = "SearchSortedOptions";
+ static SearchSortedOptions Defaults() { return SearchSortedOptions(); }
+
+ /// Whether to return the leftmost or rightmost insertion point.
+ Side side;
+};
+
class ARROW_EXPORT SortOptions : public FunctionOptions {
public:
explicit SortOptions(std::vector<SortKey> sort_keys = {});
@@ -598,6 +613,29 @@
const SelectKOptions& options,
ExecContext* ctx = NULLPTR);
+/// \brief Find insertion indices that preserve sorted order.
+///
+/// The `values` datum must be a plain array, chunked array, or run-end encoded
+/// array (including chunked run-end encoded) sorted in ascending order.
+/// `needles` may be a scalar, plain array, chunked array, or run-end encoded
+/// array (including chunked run-end encoded) whose logical value type matches
+/// `values`.
+///
+/// Nulls in `values` are supported when clustered entirely at the start or the
+/// end of the sorted array. Non-null needles are matched only against the
+/// non-null portion of `values`. Null needles yield null outputs.
+///
+/// \param[in] values sorted array to search within
+/// \param[in] needles scalar or array-like values to search for
+/// \param[in] options selects left or right insertion semantics
+/// \param[in] ctx the function execution context, optional
+/// \return insertion indices as uint64 scalar or array
+ARROW_EXPORT
+Result<Datum> SearchSorted(
+ const Datum& values, const Datum& needles,
+ const SearchSortedOptions& options = SearchSortedOptions::Defaults(),
+ ExecContext* ctx = NULLPTR);
+
/// \brief Return the indices that would sort an array.
///
/// Perform an indirect sort of array. The output array will contain
diff --git a/cpp/src/arrow/compute/initialize.cc b/cpp/src/arrow/compute/initialize.cc
index d88835d..ec531e8 100644
--- a/cpp/src/arrow/compute/initialize.cc
+++ b/cpp/src/arrow/compute/initialize.cc
@@ -48,6 +48,7 @@
internal::RegisterVectorNested(registry);
internal::RegisterVectorRank(registry);
internal::RegisterVectorReplace(registry);
+ internal::RegisterVectorSearchSorted(registry);
internal::RegisterVectorSelectK(registry);
internal::RegisterVectorSort(registry);
internal::RegisterVectorRunEndEncode(registry);
diff --git a/cpp/src/arrow/compute/kernels/CMakeLists.txt b/cpp/src/arrow/compute/kernels/CMakeLists.txt
index 15955b5..d07356a 100644
--- a/cpp/src/arrow/compute/kernels/CMakeLists.txt
+++ b/cpp/src/arrow/compute/kernels/CMakeLists.txt
@@ -121,6 +121,13 @@
arrow_compute_kernels_testing
arrow_compute_testing)
+add_arrow_compute_test(vector_search_sorted_test
+ SOURCES
+ vector_search_sorted_test.cc
+ EXTRA_LINK_LIBS
+ arrow_compute_kernels_testing
+ arrow_compute_testing)
+
add_arrow_compute_test(vector_selection_test
SOURCES
vector_selection_test.cc
@@ -141,6 +148,7 @@
add_arrow_compute_benchmark(vector_partition_benchmark)
add_arrow_compute_benchmark(vector_topk_benchmark)
add_arrow_compute_benchmark(vector_replace_benchmark)
+add_arrow_compute_benchmark(vector_search_sorted_benchmark)
add_arrow_compute_benchmark(vector_selection_benchmark)
# ----------------------------------------------------------------------
diff --git a/cpp/src/arrow/compute/kernels/meson.build b/cpp/src/arrow/compute/kernels/meson.build
index fb68244..0a8c16f 100644
--- a/cpp/src/arrow/compute/kernels/meson.build
+++ b/cpp/src/arrow/compute/kernels/meson.build
@@ -132,6 +132,7 @@
'vector_partition_benchmark',
'vector_topk_benchmark',
'vector_replace_benchmark',
+ 'vector_search_sorted_benchmark',
'vector_selection_benchmark',
]
diff --git a/cpp/src/arrow/compute/kernels/vector_search_sorted.cc b/cpp/src/arrow/compute/kernels/vector_search_sorted.cc
new file mode 100644
index 0000000..e2e90ef
--- /dev/null
+++ b/cpp/src/arrow/compute/kernels/vector_search_sorted.cc
@@ -0,0 +1,881 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/compute/api_vector.h"
+
+#include <algorithm>
+#include <memory>
+#include <optional>
+#include <ranges>
+#include <type_traits>
+#include <utility>
+
+#include "arrow/array/array_primitive.h"
+#include "arrow/array/array_run_end.h"
+#include "arrow/array/concatenate.h"
+#include "arrow/array/util.h"
+#include "arrow/buffer_builder.h"
+#include "arrow/chunk_resolver.h"
+#include "arrow/compute/function.h"
+#include "arrow/compute/kernels/codegen_internal.h"
+#include "arrow/compute/kernels/vector_sort_internal.h"
+#include "arrow/compute/registry.h"
+#include "arrow/compute/registry_internal.h"
+#include "arrow/type_traits.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/float16.h"
+#include "arrow/util/logging_internal.h"
+#include "arrow/util/ree_util.h"
+#include "arrow/util/unreachable.h"
+
+namespace arrow {
+
+using internal::checked_cast;
+using util::Float16;
+
+namespace compute::internal {
+namespace {
+
+/// Return the static default options instance used by the meta-function.
+const SearchSortedOptions* GetDefaultSearchSortedOptions() {
+ static const auto kDefaultSearchSortedOptions = SearchSortedOptions::Defaults();
+ return &kDefaultSearchSortedOptions;
+}
+
+const FunctionDoc search_sorted_doc(
+ "Find insertion indices for sorted input",
+ ("Return the index where each needle should be inserted in a sorted input array\n"
+ "to maintain ascending order.\n"
+ "\n"
+ "With side='left', returns the first suitable index (lower bound).\n"
+ "With side='right', returns the last suitable index (upper bound).\n"
+ "\n"
+ "The searched values may be provided as an array or chunked array and must\n"
+ "already be sorted in ascending order. Null values in the searched array are\n"
+ "supported when clustered entirely at the start or\n"
+ "entirely at the end. Non-null needles are matched only against the non-null\n"
+ "portion of the searched array. Needles may be a scalar, array, or chunked\n"
+ "array. Null needles emit nulls in the output."),
+ {"values", "needles"}, "SearchSortedOptions");
+
+#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(HalfFloatType) \
+ 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);
+ }
+};
+
+// The three first members are ordered by "nullness"
+enum class NullGeometry : int { NoNulls, AllNans, AllNulls, Empty, AtStart, AtEnd };
+
+inline bool IsNanPrimitive(const Array& array, int64_t index) {
+ switch (array.type_id()) {
+ case Type::FLOAT:
+ return std::isnan(checked_cast<const FloatArray&>(array).Value(index));
+ case Type::DOUBLE:
+ return std::isnan(checked_cast<const DoubleArray&>(array).Value(index));
+ case Type::HALF_FLOAT:
+ return Float16::FromBits(checked_cast<const HalfFloatArray&>(array).Value(index))
+ .is_nan();
+ default:
+ return false;
+ }
+}
+
+/// Detect the NullGeometry of a primitive or run-end-encoded array
+NullGeometry DetectNullGeometry(const Array& array) {
+ if (array.length() == 0) {
+ return NullGeometry::Empty;
+ }
+ if (array.type_id() == Type::RUN_END_ENCODED) {
+ const auto& ree_array = checked_cast<const RunEndEncodedArray&>(array);
+ auto range = ::arrow::ree_util::FindPhysicalRange(*array.data(), array.offset(),
+ array.length());
+ return DetectNullGeometry(*ree_array.values()->Slice(range.first, range.second));
+ }
+ bool null_at_start = array.IsNull(0);
+ bool null_at_end = array.IsNull(array.length() - 1);
+ if (null_at_start && !null_at_end) {
+ return NullGeometry::AtStart;
+ }
+ if (!null_at_start && null_at_end) {
+ return NullGeometry::AtEnd;
+ }
+ if (null_at_start && null_at_end) {
+ return NullGeometry::AllNulls;
+ }
+
+ // No nulls, look at NaNs
+ bool nan_at_start = IsNanPrimitive(array, 0);
+ bool nan_at_end = IsNanPrimitive(array, array.length() - 1);
+ if (nan_at_start && !nan_at_end) {
+ return NullGeometry::AtStart;
+ }
+ if (!nan_at_start && nan_at_end) {
+ return NullGeometry::AtEnd;
+ }
+ if (nan_at_start && nan_at_end) {
+ return NullGeometry::AllNans;
+ }
+ return NullGeometry::NoNulls;
+}
+
+/// Detect the NullGeometry of a chunked array
+NullGeometry DetectNullGeometry(const ChunkedArray& chunked_array) {
+ auto previous_geometry = NullGeometry::Empty;
+
+ for (const auto& chunk : chunked_array.chunks()) {
+ auto null_geometry = DetectNullGeometry(*chunk);
+ if (null_geometry == NullGeometry::Empty) {
+ continue;
+ }
+ if (null_geometry == NullGeometry::AtStart || null_geometry == NullGeometry::AtEnd) {
+ // We can conclude from this chunk alone
+ return null_geometry;
+ }
+ // If the previous chunk had different results, we can compare and decide
+ if (previous_geometry != NullGeometry::Empty && previous_geometry != null_geometry) {
+ return previous_geometry < null_geometry ? NullGeometry::AtEnd
+ : NullGeometry::AtStart;
+ }
+ previous_geometry = null_geometry;
+ }
+ return previous_geometry;
+}
+
+/// Validate the supplied null counts and produce the logical non-null window
+/// that will actually participate in binary search.
+NonNullValuesRange MakeNonNullValuesRange(int64_t full_length, int64_t null_count,
+ int64_t leading_null_count,
+ int64_t trailing_null_count) {
+ if (leading_null_count > 0) {
+ return {.offset = leading_null_count, .length = full_length - leading_null_count};
+ } else {
+ return {.offset = 0, .length = full_length - trailing_null_count};
+ }
+}
+
+/// Build the searchable non-null window once the side containing clustered
+/// nulls is already known.
+NonNullValuesRange MakeNonNullValuesRangeFromNullPlacement(int64_t full_length,
+ int64_t null_count,
+ NullPlacement null_placement) {
+ return MakeNonNullValuesRange(
+ full_length, null_count, null_placement == NullPlacement::AtStart ? null_count : 0,
+ null_placement == NullPlacement::AtStart ? 0 : null_count);
+}
+
+// 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.
+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) {
+ const auto& ree_type = checked_cast<const RunEndEncodedType&>(*data->type);
+ auto result = data->Copy();
+ auto values_copy = result->child_data[1]->Copy();
+ values_copy->type = physical_type;
+ result->type = run_end_encoded(ree_type.run_end_type(), physical_type);
+ result->child_data[1] = std::move(values_copy);
+ return result;
+ }
+ auto result = data->Copy();
+ result->type = physical_type;
+ return result;
+}
+
+/// Read a run-end value from any supported run-end integer representation.
+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;
+ }
+}
+
+class SearchWindow {
+ public:
+ explicit SearchWindow(NonNullValuesRange non_null_range)
+ : offset_(non_null_range.offset), length_(non_null_range.length) {}
+
+ int64_t length() const { return length_; }
+
+ int64_t LogicalInsertionIndex(int64_t index) const { return index + physical_offset(); }
+
+ protected:
+ int64_t physical_offset() const { return offset_; }
+
+ private:
+ int64_t offset_ = 0;
+ int64_t length_;
+};
+
+/// Access logical values from a plain Arrow array.
+template <typename ArrowType>
+class PlainArrayAccessor : public SearchWindow {
+ public:
+ using ArrayType = typename TypeTraits<ArrowType>::ArrayType;
+ using ValueType = SearchValue<ArrowType>;
+
+ PlainArrayAccessor(const std::shared_ptr<ArrayData>& array_data,
+ NullPlacement null_placement)
+ : PlainArrayAccessor(array_data, MakeNonNullValuesRangeFromNullPlacement(
+ array_data->length, array_data->GetNullCount(),
+ null_placement)) {}
+
+ PlainArrayAccessor(const std::shared_ptr<ArrayData>& array_data,
+ NonNullValuesRange non_null_range)
+ : SearchWindow(non_null_range), array_(array_data) {}
+
+ /// Return the logical value at the given position within the search window.
+ ValueType Value(int64_t index) const {
+ return GetViewType<ArrowType>::LogicalValue(
+ array_.GetView(physical_offset() + index));
+ }
+
+ private:
+ ArrayType array_;
+};
+
+class RunEndEncodedValuesAccessorBase {
+ public:
+ explicit RunEndEncodedValuesAccessorBase(RunEndEncodedArray array)
+ : array_(std::move(array)),
+ array_span_(*array_.data()),
+ run_ends_span_(::arrow::ree_util::RunEndsArray(array_span_)),
+ physical_range_(::arrow::ree_util::FindPhysicalRange(array_span_, array_.offset(),
+ array_.length())) {
+ values_ = array_.values()->Slice(physical_range_.first, physical_range_.second);
+ }
+
+ protected:
+ int64_t PhysicalIndex(int64_t index) const {
+ return physical_range_.first + /*search_offset_ + */ index;
+ }
+
+ RunEndEncodedArray array_;
+ std::shared_ptr<Array> values_;
+ ArraySpan array_span_;
+ ArraySpan run_ends_span_;
+ std::pair<int64_t, int64_t> physical_range_;
+};
+
+/// Access logical values from a run-end encoded Arrow array.
+template <typename ArrowType>
+class RunEndEncodedValuesAccessor : public RunEndEncodedValuesAccessorBase {
+ public:
+ using ArrayType = typename TypeTraits<ArrowType>::ArrayType;
+ using ValueType = SearchValue<ArrowType>;
+
+ RunEndEncodedValuesAccessor(RunEndEncodedArray array, NullPlacement null_placement)
+ : RunEndEncodedValuesAccessorBase(std::move(array)),
+ physical_accessor_(values_->data(), null_placement) {}
+
+ RunEndEncodedValuesAccessor(const std::shared_ptr<ArrayData>& array_data,
+ NullPlacement null_placement)
+ : RunEndEncodedValuesAccessor(RunEndEncodedArray(array_data), null_placement) {}
+
+ /// Return the logical value at the given physical run position within the
+ /// search window.
+ ValueType Value(int64_t index) const { return physical_accessor_.Value(index); }
+
+ int64_t length() const { return physical_accessor_.length(); }
+
+ int64_t LogicalInsertionIndex(int64_t index) const {
+ auto physical_index = physical_accessor_.LogicalInsertionIndex(index);
+
+ DCHECK_GE(physical_index, 0);
+ DCHECK_LE(physical_index, physical_range_.second);
+ if (physical_index == 0) {
+ return 0;
+ } else if (physical_index == physical_range_.second) {
+ return array_.length();
+ } else {
+ auto run_end =
+ GetRunEndValue(run_ends_span_, physical_index + physical_range_.first - 1);
+ DCHECK_GE(run_end, array_.offset());
+ DCHECK_LE(run_end, array_.offset() + array_.length());
+ return run_end - array_.offset();
+ }
+ }
+
+ protected:
+ PlainArrayAccessor<ArrowType> physical_accessor_;
+};
+
+/// Return the logical type of a datum, unwrapping run-end encoding when present.
+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;
+}
+
+/// Reject nested run-end encoded values. TODO: Support this case in the future if there
+/// is demand for it.
+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();
+}
+
+/// Validate the searched values input shape and supported encoding.
+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.
+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();
+}
+
+/// Find the insertion point into a dense array
+template <typename ArrowType, typename Accessor>
+int64_t FindInsertionDense(const Accessor& array, const SearchValue<ArrowType>& needle,
+ SearchSortedOptions::Side side, NullPlacement null_placement) {
+ // When looking for the Left side, we want equal values to be considered greater
+ // than the needle (1), otherwise smaller (-1).
+ const int on_equality = (side == SearchSortedOptions::Left) ? 1 : -1;
+ int64_t first = 0;
+ int64_t count = array.length();
+
+ auto compare = [&](auto left, auto right) {
+ // The same comparison function as used for sorting, taking account null_placement
+ // when NaNs are involved.
+ // XXX Instead of detecting NaN-ness during each comparison, we could
+ // take advantage of null_placement to single out the range of NaNs that's at
+ // the beginning or end of the array.
+ return CompareTypeValues<ArrowType>(left, right, SortOrder::Ascending, null_placement,
+ /*on_equality=*/on_equality);
+ };
+
+ while (count > 0) {
+ const int64_t step = count / 2;
+ const int64_t it = first + step;
+ const bool advance = compare(array.Value(it), needle) < 0;
+ if (advance) {
+ first = it + 1;
+ count -= step + 1;
+ } else {
+ count = step;
+ }
+ }
+ return first;
+}
+
+/// Find the insertion chunk in an array of vector chunks. The chunk index is returned.
+template <typename ArrowType, typename Accessor>
+int64_t FindInsertionChunk(const std::vector<Accessor>& chunks,
+ const SearchValue<ArrowType>& needle,
+ SearchSortedOptions::Side side, NullPlacement null_placement) {
+ // When looking for the Left side, we want equal values to be considered greater
+ // than the needle (1), otherwise smaller (-1).
+ const int on_equality = (side == SearchSortedOptions::Left) ? 1 : -1;
+ int64_t first = 0;
+ int64_t count = static_cast<int64_t>(chunks.size());
+
+ auto compare = [&](auto left, auto right) {
+ // The same comparison function as used for sorting, taking account null_placement
+ // when NaNs are involved.
+ // XXX Instead of detecting NaN-ness during each comparison, we could
+ // take advantage of null_placement to single out the range of NaNs that's at
+ // the beginning or end of the sorted_values.
+ return CompareTypeValues<ArrowType>(left, right, SortOrder::Ascending, null_placement,
+ on_equality);
+ };
+
+ while (count > 0) {
+ const int64_t step = count / 2;
+ const int64_t it = first + step;
+ bool advance;
+ const auto& chunk = chunks[it];
+ if (chunk.length() == 0) {
+ // If nulls are clustered at the start, advance towards the end.
+ advance = (null_placement == NullPlacement::AtStart);
+ } else {
+ auto chunk_first = chunk.Value(0);
+ auto chunk_last = chunk.Value(chunk.length() - 1);
+ if (compare(chunk_first, needle) > 0) {
+ // First chunk value too large => go left
+ advance = false;
+ } else if (compare(chunk_last, needle) < 0) {
+ // Last chunk value too small => go right
+ advance = true;
+ } else {
+ // Insertion point is in this chunk
+ first = it;
+ break;
+ }
+ }
+ if (advance) {
+ first = it + 1;
+ count -= step + 1;
+ } else {
+ count = step;
+ }
+ }
+ return first;
+}
+
+/// Find the insertion point into a chunked array.
+template <typename ArrowType, typename Accessor>
+ChunkLocation FindInsertionChunked(const std::vector<Accessor>& chunks,
+ const SearchValue<ArrowType>& needle,
+ SearchSortedOptions::Side side,
+ NullPlacement null_placement) {
+ // A naive implementation would search directly in the chunked array,
+ // with each indexed access taking O(log n) time.
+ // It is much faster to first narrow down the search to a single chunk
+ // (by using a binary search among chunk boundaries, see FindInsertionChunk)
+ // and then do a dense binary search (FindInsertionDense).
+ DCHECK_GT(chunks.size(), 0);
+ int64_t chunk_index =
+ FindInsertionChunk<ArrowType>(chunks, needle, side, null_placement);
+ int64_t index_in_chunk;
+ if (chunk_index == static_cast<int64_t>(chunks.size())) {
+ // Inserting at the right of the last chunk
+ --chunk_index;
+ index_in_chunk = chunks.back().length();
+ } else {
+ index_in_chunk =
+ FindInsertionDense<ArrowType>(chunks[chunk_index], needle, side, null_placement);
+ }
+ return {chunk_index, chunks[chunk_index].LogicalInsertionIndex(index_in_chunk)};
+}
+
+template <typename ArrowType, typename Accessor>
+class ChunkedSearchSorted {
+ public:
+ ChunkedSearchSorted(const ArrayVector& chunks, SearchSortedOptions::Side side,
+ NullPlacement null_placement)
+ : side_(side), null_placement_(null_placement) {
+ // Initialize accessors from non-empty chunks
+ chunk_accessors_.reserve(chunks.size());
+ chunk_offsets_.reserve(chunks.size());
+ int64_t offset = 0;
+ for (const auto& chunk : chunks) {
+ if (chunk->length() > 0) {
+ auto accessor = Accessor(chunk->data(), null_placement);
+ chunk_accessors_.push_back(std::move(accessor));
+ chunk_offsets_.push_back(offset);
+ offset += chunk->length();
+ }
+ }
+ }
+
+ int64_t FindLogicalInsertionIndex(const SearchValue<ArrowType>& needle) const {
+ if (chunk_accessors_.empty()) {
+ return 0;
+ }
+ ChunkLocation location =
+ FindInsertionChunked<ArrowType>(chunk_accessors_, needle, side_, null_placement_);
+ DCHECK_LT(location.chunk_index, static_cast<int64_t>(chunk_offsets_.size()));
+ return chunk_offsets_[location.chunk_index] + location.index_in_chunk;
+ }
+
+ protected:
+ SearchSortedOptions::Side side_;
+ NullPlacement null_placement_;
+ std::vector<Accessor> chunk_accessors_;
+ std::vector<int64_t> chunk_offsets_;
+};
+
+template <typename ArrowType>
+using VisitedNeedle = std::optional<SearchValue<ArrowType>>;
+
+/// 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 std::optional<SearchValue<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 or plain-array needles through a uniform callback interface
+/// of logical elements.
+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));
+ }
+
+ 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) {}
+
+ /// Reserve the final output size up front so append operations can use the
+ /// builders' unchecked fast path.
+ 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();
+ }
+
+ /// Append a null output slot for a null needle.
+ Status AppendNull() {
+ DCHECK(nullable_);
+ indices_builder_.UnsafeAppend(uint64_t{0});
+ null_bitmap_builder_.UnsafeAppend(false);
+ ++null_count_;
+ return Status::OK();
+ }
+
+ /// Append one computed insertion index for a non-null needle.
+ Status AppendValue(uint64_t insertion_index) {
+ indices_builder_.UnsafeAppend(insertion_index);
+ if (nullable_) {
+ null_bitmap_builder_.UnsafeAppend(true);
+ }
+ return Status::OK();
+ }
+
+ /// Finish building the output UInt64 array, attaching the null bitmap only
+ /// when nullable output was requested.
+ Result<std::shared_ptr<Array>> Finish() && {
+ DCHECK_EQ(indices_builder_.length(), expected_length_);
+ ARROW_ASSIGN_OR_RAISE(auto indices, indices_builder_.Finish());
+
+ std::shared_ptr<Buffer> null_bitmap;
+ if (nullable_) {
+ DCHECK_EQ(null_bitmap_builder_.length(), expected_length_);
+ ARROW_ASSIGN_OR_RAISE(null_bitmap, null_bitmap_builder_.Finish());
+ }
+
+ return MakeArray(ArrayData::Make(uint64(), expected_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 null_count_ = 0;
+};
+
+Result<Datum> ComputeRunEndEncodedNeedleInsertionIndices(
+ const Datum& values, const RunEndEncodedArray& needles,
+ SearchSortedOptions::Side side, ExecContext* ctx) {
+ ExecContext* exec_ctx = ctx != NULLPTR ? ctx : default_exec_context();
+
+ // Search each physical REE value once, then rebuild the run-end encoded shape
+ // and decode back to the dense logical result expected by the public API.
+ 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);
+}
+
+template <typename ArrowType, typename ValuesAccessor>
+Result<Datum> ComputeInsertionIndicesWithAccessor(
+ const Datum& sorted_values, const Datum& needles, SearchSortedOptions::Side side,
+ NullPlacement null_placement, uint64_t insertion_offset, ExecContext* ctx) {
+ // Only emit a null bitmap if necessary
+ const bool has_nulls = needles.ComputeLogicalNullCount() > 0;
+ InsertionIndexBuilder output(ctx->memory_pool(), has_nulls);
+ ARROW_RETURN_NOT_OK(output.Init(needles.length()));
+
+ // Array and ChunkedArray follow the same path, an Array having just a single chunk.
+ // The trivial case with one chunk does not add overhead, so it's not worth
+ // the maintenance hassle to have a separate path for Array.
+ ChunkedSearchSorted<ArrowType, ValuesAccessor> search_sorted(sorted_values.chunks(),
+ side, null_placement);
+
+ auto emit_search_result = [&](const VisitedNeedle<ArrowType>& needle) -> Status {
+ if (!needle.has_value()) {
+ return output.AppendNull();
+ }
+ const auto insertion_index = search_sorted.FindLogicalInsertionIndex(*needle);
+ return output.AppendValue(static_cast<uint64_t>(insertion_index));
+ };
+
+ RETURN_NOT_OK(VisitNeedleRuns<ArrowType>(needles, emit_search_result));
+
+ return std::move(output).Finish().As<Datum>();
+}
+
+template <typename ArrowType>
+Result<Datum> ComputeInsertionIndices(const Datum& sorted_values, const Datum& needles,
+ SearchSortedOptions::Side side,
+ NullPlacement null_placement,
+ uint64_t insertion_offset, ExecContext* ctx) {
+ auto physical_type = TypeTraits<ArrowType>::type_singleton();
+ DCHECK_NE(physical_type->id(), Type::RUN_END_ENCODED);
+
+ if (sorted_values.type()->id() == Type::RUN_END_ENCODED) {
+ return ComputeInsertionIndicesWithAccessor<ArrowType,
+ RunEndEncodedValuesAccessor<ArrowType>>(
+ sorted_values, needles, side, null_placement, insertion_offset, ctx);
+ } else {
+ return ComputeInsertionIndicesWithAccessor<ArrowType, PlainArrayAccessor<ArrowType>>(
+ sorted_values, needles, side, null_placement, insertion_offset, ctx);
+ }
+}
+
+/// 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());
+ }
+
+ // Chunked needles are handled at the top level so the typed dispatch
+ // below only ever sees non-chunked (scalar /array) needles.
+ if (args[1].is_chunked_array()) {
+ return ExecuteChunkedNeedles(args[0], *args[1].chunked_array(),
+ static_cast<const SearchSortedOptions&>(*options),
+ ctx);
+ }
+
+ auto null_placement = DetectNullPlacement(args[0]);
+ ARROW_ASSIGN_OR_RAISE(auto non_null_values_range,
+ FindNonNullValuesRange(args[0], null_placement));
+ auto result = DispatchByType(args[0], non_null_values_range, args[1],
+ static_cast<const SearchSortedOptions&>(*options),
+ null_placement, ctx);
+ return result;
+ }
+
+ private:
+ /// Process each needle chunk independently and concatenate the results.
+ Result<Datum> ExecuteChunkedNeedles(const Datum& values, const ChunkedArray& needles,
+ const SearchSortedOptions& options,
+ ExecContext* ctx) const {
+ if (needles.num_chunks() == 0) {
+ return MakeEmptyArray(uint64(), ctx->memory_pool()).As<Datum>();
+ }
+ ArrayVector result_chunks;
+ result_chunks.reserve(static_cast<size_t>(needles.num_chunks()));
+ for (const auto& chunk : needles.chunks()) {
+ ARROW_ASSIGN_OR_RAISE(auto chunk_result,
+ ExecuteImpl({values, Datum(chunk)}, &options, ctx));
+ result_chunks.push_back(chunk_result.make_array());
+ }
+ ARROW_ASSIGN_OR_RAISE(auto out, Concatenate(result_chunks, ctx->memory_pool()));
+ return Datum(std::move(out));
+ }
+
+ /// Compute the non-null search window on the logical view of the values
+ /// input, regardless of its physical storage.
+ Result<NonNullValuesRange> FindNonNullValuesRange(const Datum& values,
+ NullPlacement null_placement) const {
+ const int64_t null_count = values.ComputeLogicalNullCount();
+ return MakeNonNullValuesRangeFromNullPlacement(values.length(), null_count,
+ null_placement);
+ }
+
+ NullPlacement DetectNullPlacement(const Datum& values) const {
+ const auto null_geometry = values.is_chunked_array()
+ ? DetectNullGeometry(*values.chunked_array())
+ : DetectNullGeometry(*values.make_array());
+ switch (null_geometry) {
+ case NullGeometry::AtStart:
+ return NullPlacement::AtStart;
+ case NullGeometry::AtEnd:
+ return NullPlacement::AtEnd;
+ default:
+ // Shouldn't matter as there are either no nulls or only nulls
+ DCHECK(values.null_count() == 0 || values.null_count() == values.length());
+ return NullPlacement::AtEnd;
+ }
+ }
+
+ /// 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,
+ NullPlacement null_placement, 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();
+ }
+
+ auto physical_type = GetPhysicalType(logical_type_ptr);
+ switch (physical_type->id()) {
+#define VISIT(TYPE) \
+ case TYPE::type_id: \
+ return DispatchHaystack<TYPE>(values, non_null_values_range, needles, options.side, \
+ null_placement, ctx);
+ VISIT_SEARCH_SORTED_PHYSICAL_TYPES(VISIT)
+#undef VISIT
+ default:
+ break;
+ }
+ return Status::NotImplemented("search_sorted is not implemented for type ",
+ logical_type_ptr->ToString());
+ }
+
+ /// Dispatch the physical representation of the searched values.
+ template <typename ArrowType>
+ Result<Datum> DispatchHaystack(const Datum& values,
+ const NonNullValuesRange& non_null_values_range,
+ const Datum& needles, SearchSortedOptions::Side side,
+ NullPlacement null_placement, ExecContext* ctx) const {
+ 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,
+ DispatchHaystack<ArrowType>(values, non_null_values_range, Datum(scalar_arr),
+ side, null_placement, ctx));
+ ARROW_ASSIGN_OR_RAISE(auto result_scalar, result.make_array()->GetScalar(0));
+ return Datum(std::move(result_scalar));
+ }
+
+ // XXX This doesn't need to be in the type-specialized DispatchHaystack
+ if (needles.type()->id() == Type::RUN_END_ENCODED) {
+ return ComputeRunEndEncodedNeedleInsertionIndices(
+ values, RunEndEncodedArray(needles.array()), side, ctx);
+ }
+
+ return ComputeInsertionIndices<ArrowType>(
+ values, needles, side, null_placement,
+ static_cast<uint64_t>(non_null_values_range.offset), ctx);
+ }
+};
+
+} // namespace
+
+/// Register the search_sorted vector kernel in the global compute registry.
+void RegisterVectorSearchSorted(FunctionRegistry* registry) {
+ DCHECK_OK(registry->AddFunction(std::make_shared<SearchSortedMetaFunction>()));
+}
+
+} // namespace compute::internal
+} // namespace arrow
diff --git a/cpp/src/arrow/compute/kernels/vector_search_sorted_benchmark.cc b/cpp/src/arrow/compute/kernels/vector_search_sorted_benchmark.cc
new file mode 100644
index 0000000..1caf79a
--- /dev/null
+++ b/cpp/src/arrow/compute/kernels/vector_search_sorted_benchmark.cc
@@ -0,0 +1,275 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "benchmark/benchmark.h"
+
+#include <algorithm>
+#include <cmath>
+#include <concepts>
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <string_view>
+#include <vector>
+
+#include "arrow/array.h"
+#include "arrow/array/builder_binary.h"
+#include "arrow/array/builder_primitive.h"
+#include "arrow/compute/api_vector.h"
+#include "arrow/datum.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/random.h"
+#include "arrow/util/benchmark_util.h"
+
+namespace arrow {
+namespace compute {
+
+constexpr auto kSeed = 0x5EA4C42;
+constexpr int32_t kStringMinLength = 8;
+constexpr int32_t kStringMaxLength = 20;
+
+struct DenseBenchmarkArgs {
+ int64_t logical_length_;
+ double null_probability_;
+ SearchSortedOptions::Side side_;
+
+ explicit DenseBenchmarkArgs(const benchmark::State& state)
+ : logical_length_(state.range(0)), null_probability_(state.range(1) / 100.0) {
+ side_ = std::array{SearchSortedOptions::Left, SearchSortedOptions::Right}.at(
+ state.range(2));
+ }
+
+ static void SetArgs(benchmark::internal::Benchmark* bench) {
+ bench->Unit(benchmark::kMicrosecond);
+ bench->ArgNames({"length", "null_probability", "side"});
+ for (const auto size : std::vector<int64_t>{kL1Size, kL2Size}) {
+ for (const double null_probability : {0.0, 0.5, 0.9}) {
+ for (const bool right_side : {false, true}) {
+ bench->Args(
+ {size / 4, static_cast<int64_t>(null_probability * 100), right_side});
+ }
+ }
+ }
+ }
+};
+
+struct DenseBenchmarkBase {
+ double null_probability_;
+ random::RandomArrayGenerator rand_;
+
+ explicit DenseBenchmarkBase(const DenseBenchmarkArgs& args)
+ : null_probability_(args.null_probability_), rand_(kSeed) {}
+};
+
+struct Int64Benchmark : public DenseBenchmarkBase {
+ using DenseBenchmarkBase::DenseBenchmarkBase;
+
+ std::shared_ptr<Array> BuildSortedValues(int64_t length) {
+ auto values = std::static_pointer_cast<Int64Array>(RandomArray(length));
+ std::vector<int64_t> data(values->raw_values(),
+ values->raw_values() + values->length());
+ std::ranges::sort(data);
+
+ Int64Builder builder;
+ ABORT_NOT_OK(builder.AppendValues(data));
+ return builder.Finish().ValueOrDie();
+ }
+
+ std::shared_ptr<Array> BuildNeedles(int64_t length) {
+ // NOTE it's not required to have the same number of needles as values;
+ // it just makes the benchmarks simpler.
+ return RandomArray(length);
+ }
+
+ protected:
+ std::shared_ptr<Array> RandomArray(int64_t length) {
+ return rand_.Int64(length, 0, /*max=*/length * 4, null_probability_);
+ }
+};
+
+struct StringBenchmark : public DenseBenchmarkBase {
+ using DenseBenchmarkBase::DenseBenchmarkBase;
+
+ std::shared_ptr<Array> BuildSortedValues(int64_t length) {
+ auto values = std::static_pointer_cast<StringArray>(
+ rand_.String(length, kStringMinLength, kStringMaxLength, null_probability_));
+
+ std::vector<std::string_view> data;
+ data.reserve(static_cast<size_t>(values->length()));
+ for (int64_t index = 0; index < values->length(); ++index) {
+ data.push_back(values->GetView(index));
+ }
+ std::ranges::sort(data);
+
+ StringBuilder builder;
+ ABORT_NOT_OK(builder.Reserve(values->length()));
+ ABORT_NOT_OK(builder.ReserveData(values->total_values_length()));
+ for (const auto view : data) {
+ builder.UnsafeAppend(view);
+ }
+ return std::static_pointer_cast<StringArray>(builder.Finish().ValueOrDie());
+ }
+
+ std::shared_ptr<Array> BuildNeedles(int64_t length) {
+ return rand_.String(length, kStringMinLength, kStringMaxLength, null_probability_);
+ }
+};
+
+template <typename DenseBenchmark, bool kREENeedles>
+ requires std::derived_from<DenseBenchmark, DenseBenchmarkBase>
+struct REEBenchmark {
+ static constexpr int64_t kAverageRunLength = 50;
+ int64_t logical_length_;
+ int64_t physical_length_;
+ DenseBenchmark dense_benchmark_;
+
+ explicit REEBenchmark(const DenseBenchmarkArgs& args) : dense_benchmark_(args) {}
+
+ std::shared_ptr<Array> BuildSortedValues(int64_t logical_length) {
+ return Encode(dense_benchmark_.BuildSortedValues(logical_length / kAverageRunLength),
+ logical_length);
+ }
+
+ std::shared_ptr<Array> BuildNeedles(int64_t logical_length) {
+ if constexpr (kREENeedles) {
+ return Encode(dense_benchmark_.BuildNeedles(logical_length / kAverageRunLength),
+ logical_length);
+ } else {
+ return dense_benchmark_.BuildNeedles(logical_length);
+ }
+ }
+
+ protected:
+ std::shared_ptr<Array> Encode(std::shared_ptr<Array> values, int64_t logical_length) {
+ return dense_benchmark_.rand_.RunEndEncoded(values, logical_length);
+ }
+};
+
+template <typename DenseBenchmark>
+using REEValuesDenseNeedlesBenchmark =
+ REEBenchmark<DenseBenchmark, /*kREENeedles=*/false>;
+
+template <typename DenseBenchmark>
+using REEValuesREENeedlesBenchmark = REEBenchmark<DenseBenchmark, /*kREENeedles=*/true>;
+
+struct NoOpChunker {
+ Datum operator()(std::shared_ptr<Array> array) { return array; }
+};
+
+struct StaticChunker {
+ static constexpr int kNumChunks = 8;
+
+ Datum operator()(std::shared_ptr<Array> array) {
+ ArrayVector chunks;
+ int64_t chunk_start = 0;
+ for (int64_t i = 0; i < kNumChunks; ++i) {
+ const int64_t chunk_end = static_cast<int64_t>(
+ ceil(static_cast<double>(i + 1) / kNumChunks * array->length()));
+ chunks.push_back(
+ array->SliceSafe(chunk_start, chunk_end - chunk_start).ValueOrDie());
+ chunk_start = chunk_end;
+ }
+ ARROW_CHECK_EQ(chunk_start, array->length());
+ auto chunked_array = ChunkedArray::Make(std::move(chunks)).ValueOrDie();
+ ARROW_CHECK_EQ(chunked_array->length(), array->length());
+ return chunked_array;
+ }
+};
+
+void SetBenchmarkCounters(benchmark::State& state, const Datum& values,
+ const Datum& needles, SearchSortedOptions::Side side) {
+ const auto needles_length = needles.length();
+ state.SetItemsProcessed(state.iterations() * needles_length);
+}
+
+void RunSearchSortedBenchmark(benchmark::State& state, const Datum& values,
+ const Datum& needles, SearchSortedOptions::Side side) {
+ const SearchSortedOptions options(side);
+ for (auto _ : state) {
+ auto result = SearchSorted(values, needles, options);
+ ABORT_NOT_OK(result.status());
+ benchmark::DoNotOptimize(result.ValueUnsafe());
+ }
+ SetBenchmarkCounters(state, values, needles, side);
+}
+
+template <typename Benchmark, typename Chunker>
+void RunSearchSortedBenchmark(benchmark::State& state, DenseBenchmarkArgs args,
+ Benchmark benchmark, SearchSortedOptions::Side side) {
+ Chunker chunker;
+ RunSearchSortedBenchmark(state,
+ chunker(benchmark.BuildSortedValues(args.logical_length_)),
+ chunker(benchmark.BuildNeedles(args.logical_length_)), side);
+}
+
+template <typename Benchmark, typename Chunker>
+void RunSearchSortedBenchmark(benchmark::State& state, DenseBenchmarkArgs args) {
+ RunSearchSortedBenchmark<Benchmark, Chunker>(state, args, Benchmark(args), args.side_);
+}
+
+static void SearchSortedDenseInt64Array(benchmark::State& state) {
+ RunSearchSortedBenchmark<Int64Benchmark, NoOpChunker>(state, DenseBenchmarkArgs(state));
+}
+
+static void SearchSortedDenseStringArray(benchmark::State& state) {
+ RunSearchSortedBenchmark<StringBenchmark, NoOpChunker>(state,
+ DenseBenchmarkArgs(state));
+}
+
+static void SearchSortedDenseInt64ChunkedArray(benchmark::State& state) {
+ RunSearchSortedBenchmark<Int64Benchmark, StaticChunker>(state,
+ DenseBenchmarkArgs(state));
+}
+
+static void SearchSortedDenseStringChunkedArray(benchmark::State& state) {
+ RunSearchSortedBenchmark<StringBenchmark, StaticChunker>(state,
+ DenseBenchmarkArgs(state));
+}
+
+static void SearchSortedREEInt64Array(benchmark::State& state) {
+ RunSearchSortedBenchmark<REEValuesDenseNeedlesBenchmark<Int64Benchmark>, NoOpChunker>(
+ state, DenseBenchmarkArgs(state));
+}
+
+static void SearchSortedREEInt64ChunkedArray(benchmark::State& state) {
+ RunSearchSortedBenchmark<REEValuesDenseNeedlesBenchmark<Int64Benchmark>, StaticChunker>(
+ state, DenseBenchmarkArgs(state));
+}
+
+static void SearchSortedREEInt64ArrayREENeedles(benchmark::State& state) {
+ RunSearchSortedBenchmark<REEValuesREENeedlesBenchmark<Int64Benchmark>, NoOpChunker>(
+ state, DenseBenchmarkArgs(state));
+}
+
+static void SearchSortedREEInt64ChunkedArrayREENeedles(benchmark::State& state) {
+ RunSearchSortedBenchmark<REEValuesREENeedlesBenchmark<Int64Benchmark>, StaticChunker>(
+ state, DenseBenchmarkArgs(state));
+}
+
+BENCHMARK(SearchSortedDenseInt64Array)->Apply(DenseBenchmarkArgs::SetArgs);
+BENCHMARK(SearchSortedDenseStringArray)->Apply(DenseBenchmarkArgs::SetArgs);
+BENCHMARK(SearchSortedDenseInt64ChunkedArray)->Apply(DenseBenchmarkArgs::SetArgs);
+BENCHMARK(SearchSortedDenseStringChunkedArray)->Apply(DenseBenchmarkArgs::SetArgs);
+
+BENCHMARK(SearchSortedREEInt64Array)->Apply(DenseBenchmarkArgs::SetArgs);
+BENCHMARK(SearchSortedREEInt64ChunkedArray)->Apply(DenseBenchmarkArgs::SetArgs);
+
+BENCHMARK(SearchSortedREEInt64ArrayREENeedles)->Apply(DenseBenchmarkArgs::SetArgs);
+BENCHMARK(SearchSortedREEInt64ChunkedArrayREENeedles)->Apply(DenseBenchmarkArgs::SetArgs);
+
+} // namespace compute
+} // namespace arrow
diff --git a/cpp/src/arrow/compute/kernels/vector_search_sorted_test.cc b/cpp/src/arrow/compute/kernels/vector_search_sorted_test.cc
new file mode 100644
index 0000000..00144db
--- /dev/null
+++ b/cpp/src/arrow/compute/kernels/vector_search_sorted_test.cc
@@ -0,0 +1,672 @@
+// 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 <memory>
+#include <ostream>
+#include <string>
+#include <vector>
+
+#include <gtest/gtest.h>
+
+#include "arrow/array/concatenate.h"
+#include "arrow/compute/api.h"
+#include "arrow/compute/kernels/test_util_internal.h"
+#include "arrow/testing/gtest_util.h"
+
+namespace arrow {
+
+using internal::checked_cast;
+
+namespace compute {
+namespace {
+
+Result<std::shared_ptr<Array>> REEFromJSON(const std::shared_ptr<DataType>& ree_type,
+ const std::string& json) {
+ auto ree_type_ptr = checked_cast<const RunEndEncodedType*>(ree_type.get());
+ auto array = ArrayFromJSON(ree_type_ptr->value_type(), json);
+ ARROW_ASSIGN_OR_RAISE(
+ auto datum, RunEndEncode(array, RunEndEncodeOptions{ree_type_ptr->run_end_type()}));
+ return datum.make_array();
+}
+
+void CheckSearchSorted(const Datum& values, const Datum& needles,
+ SearchSortedOptions::Side side, const std::string& expected_json) {
+ ASSERT_OK_AND_ASSIGN(auto result,
+ SearchSorted(values, needles, SearchSortedOptions(side)));
+ ASSERT_TRUE(result.is_array());
+ ASSERT_OK(result.make_array()->ValidateFull());
+
+ AssertArraysEqual(*ArrayFromJSON(uint64(), expected_json), *result.make_array());
+}
+
+void CheckSearchSorted(const Datum& values, const Datum& needles,
+ const std::string& expected_left_json,
+ const std::string& expected_right_json) {
+ CheckSearchSorted(values, needles, SearchSortedOptions::Left, expected_left_json);
+ CheckSearchSorted(values, needles, SearchSortedOptions::Right, expected_right_json);
+}
+
+void CheckSimpleSearchSorted(const std::shared_ptr<DataType>& type,
+ const std::string& values_json,
+ const std::string& needles_json,
+ const std::string& expected_left_json,
+ const std::string& expected_right_json) {
+ auto values = ArrayFromJSON(type, values_json);
+ auto needles = ArrayFromJSON(type, needles_json);
+
+ CheckSearchSorted(Datum(values), Datum(needles), expected_left_json,
+ expected_right_json);
+}
+
+void CheckScalarSearchSorted(const Datum& values, const std::shared_ptr<Array>& needles,
+ const std::string& expected_left_json,
+ const std::string& expected_right_json) {
+ auto expected_left = ArrayFromJSON(uint64(), expected_left_json);
+ auto expected_right = ArrayFromJSON(uint64(), expected_right_json);
+
+ ASSERT_EQ(needles->length(), expected_left->length());
+ ASSERT_EQ(needles->length(), expected_right->length());
+
+ for (int64_t index = 0; index < needles->length(); ++index) {
+ ASSERT_OK_AND_ASSIGN(auto needle, needles->GetScalar(index));
+ ASSERT_OK_AND_ASSIGN(auto left,
+ SearchSorted(values, Datum(needle),
+ SearchSortedOptions(SearchSortedOptions::Left)));
+ ASSERT_OK_AND_ASSIGN(auto right,
+ SearchSorted(values, Datum(needle),
+ SearchSortedOptions(SearchSortedOptions::Right)));
+
+ ASSERT_TRUE(left.is_scalar());
+ ASSERT_TRUE(right.is_scalar());
+
+ ASSERT_OK_AND_ASSIGN(auto expected_left_scalar, expected_left->GetScalar(index));
+ ASSERT_OK_AND_ASSIGN(auto expected_right_scalar, expected_right->GetScalar(index));
+ AssertScalarsEqual(*expected_left_scalar, *left.scalar());
+ AssertScalarsEqual(*expected_right_scalar, *right.scalar());
+ }
+}
+
+void CheckSimpleScalarSearchSorted(const std::shared_ptr<DataType>& type,
+ const std::string& values_json,
+ const std::string& needles_json,
+ const std::string& expected_left_json,
+ const std::string& expected_right_json) {
+ auto values = ArrayFromJSON(type, values_json);
+ auto needles = ArrayFromJSON(type, needles_json);
+ CheckScalarSearchSorted(Datum(values), needles, expected_left_json,
+ expected_right_json);
+}
+
+void CheckSimpleSearchSortedAndScalar(const std::shared_ptr<DataType>& type,
+ const std::string& values_json,
+ const std::string& needles_json,
+ const std::string& expected_left_json,
+ const std::string& expected_right_json) {
+ auto values = ArrayFromJSON(type, values_json);
+ auto needles = ArrayFromJSON(type, needles_json);
+
+ CheckSearchSorted(Datum(values), Datum(needles), expected_left_json,
+ expected_right_json);
+ CheckScalarSearchSorted(Datum(values), needles, expected_left_json,
+ expected_right_json);
+}
+
+void CheckChunkedSearchSortedAndConcatenated(const std::shared_ptr<ChunkedArray>& values,
+ const std::shared_ptr<ChunkedArray>& needles,
+ const std::string& expected_left_json,
+ const std::string& expected_right_json) {
+ CheckSearchSorted(Datum(values), Datum(needles), expected_left_json,
+ expected_right_json);
+
+ ASSERT_OK_AND_ASSIGN(auto concatenated_values, Concatenate(values->chunks()));
+ ASSERT_OK_AND_ASSIGN(auto concatenated_needles, Concatenate(needles->chunks()));
+
+ CheckSearchSorted(Datum(concatenated_values), Datum(concatenated_needles),
+ expected_left_json, expected_right_json);
+}
+
+const std::vector<std::shared_ptr<DataType>> kSupportedFloatTypes{float16(), float32(),
+ float64()};
+
+struct SearchSortedSmokeCase {
+ std::string name;
+ std::shared_ptr<DataType> type;
+ std::string values_json;
+ std::string needles_json;
+ std::string expected_left_json;
+ std::string expected_right_json;
+
+ // Define a custom print since the default Googletest print trips Valgrind
+ friend std::ostream& operator<<(std::ostream& os, const SearchSortedSmokeCase& param) {
+ os << "SearchSortedSmokeCase{\"" << param.name << "\"}";
+ return os;
+ }
+};
+
+std::vector<SearchSortedSmokeCase> SupportedTypeSmokeCases() {
+ return {
+ {"Boolean", boolean(), "[false, false, false, true, true]", "[false, true]",
+ "[0, 3]", "[3, 5]"},
+ {
+ "Int8",
+ int8(),
+ "[1, 3, 3, 5, 8]",
+ "[0, 3, 9]",
+ "[0, 1, 5]",
+ "[0, 3, 5]",
+ },
+ {
+ "Int16",
+ int16(),
+ "[1, 3, 3, 5, 8]",
+ "[0, 3, 9]",
+ "[0, 1, 5]",
+ "[0, 3, 5]",
+ },
+ {
+ "Int32",
+ int32(),
+ "[1, 3, 3, 5, 8]",
+ "[0, 3, 9]",
+ "[0, 1, 5]",
+ "[0, 3, 5]",
+ },
+ {
+ "Int64",
+ int64(),
+ "[1, 3, 3, 5, 8]",
+ "[0, 3, 9]",
+ "[0, 1, 5]",
+ "[0, 3, 5]",
+ },
+ {
+ "UInt8",
+ uint8(),
+ "[1, 3, 3, 5, 8]",
+ "[0, 3, 9]",
+ "[0, 1, 5]",
+ "[0, 3, 5]",
+ },
+ {
+ "UInt16",
+ uint16(),
+ "[1, 3, 3, 5, 8]",
+ "[0, 3, 9]",
+ "[0, 1, 5]",
+ "[0, 3, 5]",
+ },
+ {
+ "UInt32",
+ uint32(),
+ "[1, 3, 3, 5, 8]",
+ "[0, 3, 9]",
+ "[0, 1, 5]",
+ "[0, 3, 5]",
+ },
+ {
+ "UInt64",
+ uint64(),
+ "[1, 3, 3, 5, 8]",
+ "[0, 3, 9]",
+ "[0, 1, 5]",
+ "[0, 3, 5]",
+ },
+ {"Float16", float16(), "[1.0, 3.0, 3.0, 5.0, 8.0]", "[0.0, 3.0, 9.0]", "[0, 1, 5]",
+ "[0, 3, 5]"},
+ {"Float32", float32(), "[1.0, 3.0, 3.0, 5.0, 8.0]", "[0.0, 3.0, 9.0]", "[0, 1, 5]",
+ "[0, 3, 5]"},
+ {"Float64", float64(), "[1.0, 3.0, 3.0, 5.0, 8.0]", "[0.0, 3.0, 9.0]", "[0, 1, 5]",
+ "[0, 3, 5]"},
+ {
+ "Date32",
+ date32(),
+ "[1, 3, 3, 5, 8]",
+ "[0, 3, 9]",
+ "[0, 1, 5]",
+ "[0, 3, 5]",
+ },
+ {
+ "Date64",
+ date64(),
+ "[86400000, 259200000, 259200000, 432000000, 691200000]",
+ "[0, 259200000, 777600000]",
+ "[0, 1, 5]",
+ "[0, 3, 5]",
+ },
+ {"Time32", time32(TimeUnit::SECOND), "[1, 3, 3, 5, 8]", "[0, 3, 9]", "[0, 1, 5]",
+ "[0, 3, 5]"},
+ {"Time64", time64(TimeUnit::NANO), "[1, 3, 3, 5, 8]", "[0, 3, 9]", "[0, 1, 5]",
+ "[0, 3, 5]"},
+ {"Timestamp", timestamp(TimeUnit::SECOND),
+ R"(["1970-01-02", "1970-01-04", "1970-01-04", "1970-01-06", "1970-01-09"])",
+ R"(["1970-01-01", "1970-01-04", "1970-01-10"])", "[0, 1, 5]", "[0, 3, 5]"},
+ {"Duration", duration(TimeUnit::NANO), "[1, 3, 3, 5, 8]", "[0, 3, 9]", "[0, 1, 5]",
+ "[0, 3, 5]"},
+ {"Binary", binary(), R"(["aa", "bb", "bb", "dd", "ff"])",
+ R"(["a", "c", "bb", "z"])", "[0, 3, 1, 5]", "[0, 3, 3, 5]"},
+ {"String", utf8(), R"(["aa", "bb", "bb", "dd", "ff"])", R"(["a", "c", "bb", "z"])",
+ "[0, 3, 1, 5]", "[0, 3, 3, 5]"},
+ {"LargeBinary", large_binary(), R"(["aa", "bb", "bb", "dd", "ff"])",
+ R"(["a", "c", "bb", "z"])", "[0, 3, 1, 5]", "[0, 3, 3, 5]"},
+ {"LargeString", large_utf8(), R"(["aa", "bb", "bb", "dd", "ff"])",
+ R"(["a", "c", "bb", "z"])", "[0, 3, 1, 5]", "[0, 3, 3, 5]"},
+ {"BinaryView", binary_view(), R"(["aa", "bb", "bb", "dd", "ff"])",
+ R"(["a", "c", "bb", "z"])", "[0, 3, 1, 5]", "[0, 3, 3, 5]"},
+ {"StringView", utf8_view(), R"(["aa", "bb", "bb", "dd", "ff"])",
+ R"(["a", "c", "bb", "z"])", "[0, 3, 1, 5]", "[0, 3, 3, 5]"},
+ };
+}
+
+class SearchSortedSupportedTypesTest
+ : public ::testing::TestWithParam<SearchSortedSmokeCase> {};
+
+TEST(SearchSorted, BasicLeftRight) {
+ CheckSimpleSearchSorted(int64(), "[100, 200, 200, 300, 300]", "[50, 200, 250, 400]",
+ "[0, 1, 3, 5]", "[0, 3, 3, 5]");
+}
+
+TEST(SearchSorted, ScalarNeedle) {
+ auto values = ArrayFromJSON(int32(), "[1, 3, 5, 7]");
+
+ ASSERT_OK_AND_ASSIGN(
+ auto result, SearchSorted(Datum(values), Datum(std::make_shared<Int32Scalar>(5)),
+ SearchSortedOptions(SearchSortedOptions::Right)));
+
+ ASSERT_TRUE(result.is_scalar());
+ ASSERT_EQ(checked_cast<const UInt64Scalar&>(*result.scalar()).value, 3);
+}
+
+TEST(SearchSorted, ScalarStringNeedle) {
+ auto values = ArrayFromJSON(utf8(), R"(["aa", "bb", "bb", "cc"])");
+
+ ASSERT_OK_AND_ASSIGN(
+ auto result,
+ SearchSorted(Datum(values), Datum(std::make_shared<StringScalar>("bb")),
+ SearchSortedOptions(SearchSortedOptions::Right)));
+
+ ASSERT_TRUE(result.is_scalar());
+ ASSERT_EQ(checked_cast<const UInt64Scalar&>(*result.scalar()).value, 3);
+}
+
+TEST(SearchSorted, EmptyHaystack) {
+ CheckSimpleSearchSorted(int16(), "[]", "[1, 2, 3]", "[0, 0, 0]", "[0, 0, 0]");
+}
+
+TEST(SearchSorted, ValuesWithLeadingNulls) {
+ CheckSimpleSearchSorted(int32(), "[null, 200, 300, 300]", "[50, 200, 250, 400]",
+ "[1, 1, 2, 4]", "[1, 2, 2, 4]");
+}
+
+TEST(SearchSorted, ValuesWithTrailingNulls) {
+ CheckSimpleSearchSorted(int32(), "[200, 300, 300, null, null]", "[50, 200, 250, 400]",
+ "[0, 0, 1, 3]", "[0, 1, 1, 3]");
+}
+
+TEST(SearchSorted, FloatValuesWithNaNs) {
+ for (const auto& type : kSupportedFloatTypes) {
+ // Nulls last, NaN in needles but not in haystack
+ CheckSimpleSearchSortedAndScalar(type, "[1.0, 3.0, 3.0, 5.0]",
+ "[3.0, 0.0, NaN, 4.0, 6.0]", "[1, 0, 4, 3, 4]",
+ "[3, 0, 4, 3, 4]");
+ CheckSimpleSearchSortedAndScalar(type, "[1.0, 3.0, 3.0, 5.0, null]",
+ "[3.0, 0.0, NaN, 4.0, 6.0]", "[1, 0, 4, 3, 4]",
+ "[3, 0, 4, 3, 4]");
+ // Nulls last, NaN in needles and in haystack
+ CheckSimpleSearchSortedAndScalar(type, "[1.0, 3.0, 3.0, 5.0, NaN, NaN]",
+ "[3.0, 0.0, NaN, 4.0]", "[1, 0, 4, 3]",
+ "[3, 0, 6, 3]");
+ CheckSimpleSearchSortedAndScalar(type, "[1.0, 3.0, 3.0, 5.0, NaN, NaN, null]",
+ "[3.0, 0.0, NaN, 4.0]", "[1, 0, 4, 3]",
+ "[3, 0, 6, 3]");
+ // Nulls first, NaN in needles but not in haystack
+ CheckSimpleSearchSortedAndScalar(type, "[null, 1.0, 3.0, 3.0, 5.0]",
+ "[3.0, 0.0, NaN, 4.0, 6.0]", "[2, 1, 1, 4, 5]",
+ "[4, 1, 1, 4, 5]");
+ // Nulls first, NaN in needles and in haystack
+ CheckSimpleSearchSortedAndScalar(type, "[NaN, NaN, NaN, 1.0, 3.0, 3.0, 5.0]",
+ "[3.0, 0.0, NaN, 4.0, 6.0]", "[4, 3, 0, 6, 7]",
+ "[6, 3, 3, 6, 7]");
+ CheckSimpleSearchSortedAndScalar(type, "[null, NaN, NaN, 1.0, 3.0, 3.0, 5.0]",
+ "[3.0, 0.0, NaN, 4.0, 6.0]", "[4, 3, 1, 6, 7]",
+ "[6, 3, 3, 6, 7]");
+ }
+}
+
+TEST(SearchSorted, NullNeedlesEmitNull) {
+ CheckSimpleSearchSorted(int32(), "[null, 200, 300, 300]", "[null, 50, 200, null, 400]",
+ "[null, 1, 1, null, 4]", "[null, 1, 2, null, 4]");
+
+ auto values = ArrayFromJSON(int32(), "[null, 200, 300, 300]");
+
+ ASSERT_OK_AND_ASSIGN(auto scalar_result,
+ SearchSorted(Datum(values), Datum(std::make_shared<Int32Scalar>()),
+ SearchSortedOptions(SearchSortedOptions::Left)));
+ ASSERT_TRUE(scalar_result.is_scalar());
+ ASSERT_FALSE(scalar_result.scalar()->is_valid);
+ ASSERT_TRUE(scalar_result.scalar()->type->Equals(uint64()));
+}
+
+TEST(SearchSorted, ChunkedValues) {
+ auto values = std::make_shared<ChunkedArray>(ArrayVector{
+ ArrayFromJSON(int32(), "[10, 10]"),
+ ArrayFromJSON(int32(), "[10, 30, 50]"),
+ });
+ auto needles = ArrayFromJSON(int32(), "[10, 20, 60]");
+ CheckSearchSorted(Datum(values), Datum(needles), "[0, 3, 5]", "[3, 3, 5]");
+
+ // Zero chunks
+ values = std::make_shared<ChunkedArray>(ArrayVector{}, int32());
+ CheckSearchSorted(Datum(values), Datum(needles), "[0, 0, 0]", "[0, 0, 0]");
+}
+
+TEST(SearchSorted, ChunkedNeedles) {
+ auto values = ArrayFromJSON(int32(), "[1, 1, 3, 5, 8]");
+ auto needles = std::make_shared<ChunkedArray>(ArrayVector{
+ ArrayFromJSON(int32(), "[null, 0, 1]"),
+ ArrayFromJSON(int32(), "[4, null, 9]"),
+ });
+ CheckSearchSorted(Datum(values), Datum(needles), "[null, 0, 0, 3, null, 5]",
+ "[null, 0, 2, 3, null, 5]");
+
+ // Zero chunks
+ needles = std::make_shared<ChunkedArray>(ArrayVector{}, int32());
+ CheckSearchSorted(Datum(values), Datum(needles), "[]", "[]");
+}
+
+TEST(SearchSorted, ChunkedValuesChunkedNeedles) {
+ auto values = std::make_shared<ChunkedArray>(ArrayVector{
+ ArrayFromJSON(int32(), "[1, 1]"),
+ ArrayFromJSON(int32(), "[3]"),
+ ArrayFromJSON(int32(), "[5, 8]"),
+ });
+ auto needles = std::make_shared<ChunkedArray>(ArrayVector{
+ ArrayFromJSON(int32(), "[null, 0, 1]"),
+ ArrayFromJSON(int32(), "[4]"),
+ ArrayFromJSON(int32(), "[null, 9]"),
+ });
+
+ CheckChunkedSearchSortedAndConcatenated(values, needles, "[null, 0, 0, 3, null, 5]",
+ "[null, 0, 2, 3, null, 5]");
+}
+
+TEST(SearchSorted, ChunkedFloatValuesWithNaNs) {
+ for (const auto& type : kSupportedFloatTypes) {
+ auto needles = ArrayFromJSON(type, "[3.0, 0.0, NaN, 4.0, 6.0]");
+ {
+ // Nulls last
+ auto values_without_nans =
+ ChunkedArrayFromJSON(type, {"[]", "[1.0, 3.0]", "[3.0, 5.0]"});
+ // (stress NullPlacement detection by chunking in different ways)
+ auto values_with_nans1 =
+ ChunkedArrayFromJSON(type, {"[]", "[1.0, 3.0]", "[3.0, 5.0, NaN, NaN]"});
+ auto values_with_nans2 =
+ ChunkedArrayFromJSON(type, {"[]", "[1.0, 3.0]", "[3.0, 5.0]", "[NaN, NaN]"});
+ auto values_with_nans_and_nulls1 = ChunkedArrayFromJSON(
+ type, {"[]", "[1.0, 3.0]", "[3.0, 5.0, NaN, NaN, null, null]"});
+ auto values_with_nans_and_nulls2 = ChunkedArrayFromJSON(
+ type, {"[]", "[1.0, 3.0]", "[3.0, 5.0]", "[NaN, NaN]", "[null, null]"});
+ auto values_with_nans_and_nulls3 = ChunkedArrayFromJSON(
+ type, {"[]", "[1.0, 3.0]", "[3.0, 5.0]", "[NaN]", "[NaN, null, null]"});
+ CheckSearchSorted(Datum(values_without_nans), Datum(needles), "[1, 0, 4, 3, 4]",
+ "[3, 0, 4, 3, 4]");
+ CheckSearchSorted(Datum(values_with_nans1), Datum(needles), "[1, 0, 4, 3, 4]",
+ "[3, 0, 6, 3, 4]");
+ CheckSearchSorted(Datum(values_with_nans2), Datum(needles), "[1, 0, 4, 3, 4]",
+ "[3, 0, 6, 3, 4]");
+ CheckSearchSorted(Datum(values_with_nans_and_nulls1), Datum(needles),
+ "[1, 0, 4, 3, 4]", "[3, 0, 6, 3, 4]");
+ CheckSearchSorted(Datum(values_with_nans_and_nulls2), Datum(needles),
+ "[1, 0, 4, 3, 4]", "[3, 0, 6, 3, 4]");
+ CheckSearchSorted(Datum(values_with_nans_and_nulls3), Datum(needles),
+ "[1, 0, 4, 3, 4]", "[3, 0, 6, 3, 4]");
+ }
+ {
+ // Nulls first
+ auto values_with_nans1 =
+ ChunkedArrayFromJSON(type, {"[]", "[NaN, NaN, 1.0, 3.0]", "[3.0, 5.0]"});
+ auto values_with_nans2 =
+ ChunkedArrayFromJSON(type, {"[]", "[NaN, NaN]", "[1.0, 3.0]", "[3.0, 5.0]"});
+ auto values_with_nans_and_nulls1 = ChunkedArrayFromJSON(
+ type, {"[]", "[null, null, NaN, NaN, 1.0, 3.0]", "[3.0, 5.0]"});
+ auto values_with_nans_and_nulls2 = ChunkedArrayFromJSON(
+ type, {"[]", "[null, null, NaN]", "[NaN, 1.0, 3.0]", "[3.0, 5.0]"});
+ auto values_with_nans_and_nulls3 = ChunkedArrayFromJSON(
+ type, {"[]", "[null, null]", "[NaN]", "[NaN, 1.0, 3.0]", "[3.0, 5.0]"});
+ CheckSearchSorted(Datum(values_with_nans1), Datum(needles), "[3, 2, 0, 5, 6]",
+ "[5, 2, 2, 5, 6]");
+ CheckSearchSorted(Datum(values_with_nans2), Datum(needles), "[3, 2, 0, 5, 6]",
+ "[5, 2, 2, 5, 6]");
+ CheckSearchSorted(Datum(values_with_nans_and_nulls1), Datum(needles),
+ "[5, 4, 2, 7, 8]", "[7, 4, 4, 7, 8]");
+ CheckSearchSorted(Datum(values_with_nans_and_nulls2), Datum(needles),
+ "[5, 4, 2, 7, 8]", "[7, 4, 4, 7, 8]");
+ CheckSearchSorted(Datum(values_with_nans_and_nulls3), Datum(needles),
+ "[5, 4, 2, 7, 8]", "[7, 4, 4, 7, 8]");
+ }
+ }
+}
+
+TEST(SearchSorted, ChunkedRunEndEncodedValues) {
+ auto values_type = run_end_encoded(int16(), int32());
+ ASSERT_OK_AND_ASSIGN(auto left_chunk, REEFromJSON(values_type, "[10, 10, 10]"));
+ ASSERT_OK_AND_ASSIGN(auto right_chunk, REEFromJSON(values_type, "[30, 30, 50]"));
+ auto values = std::make_shared<ChunkedArray>(ArrayVector{left_chunk, right_chunk});
+ auto needles = ArrayFromJSON(int32(), "[5, 10, 20, 30, 40, 50, 60]");
+
+ CheckSearchSorted(Datum(values), Datum(needles), "[0, 0, 3, 3, 5, 5, 6]",
+ "[0, 3, 3, 5, 5, 6, 6]");
+}
+
+TEST(SearchSorted, SlicedChunkedRunEndEncodedValues) {
+ auto values_type = run_end_encoded(int16(), int32());
+ ASSERT_OK_AND_ASSIGN(auto left_chunk, REEFromJSON(values_type, "[10, 10, 10]"));
+ ASSERT_OK_AND_ASSIGN(auto right_chunk, REEFromJSON(values_type, "[30, 30, 50]"));
+ auto values = std::make_shared<ChunkedArray>(
+ ArrayVector{left_chunk->Slice(1, 2), right_chunk->Slice(0, 2)});
+ auto needles = ArrayFromJSON(int32(), "[5, 10, 20, 30, 40, 50]");
+
+ CheckSearchSorted(Datum(values), Datum(needles), "[0, 0, 2, 2, 4, 4]",
+ "[0, 2, 2, 4, 4, 4]");
+}
+
+TEST(SearchSorted, ChunkedRunEndEncodedNeedles) {
+ auto values = ArrayFromJSON(int32(), "[1, 1, 3, 5, 8]");
+ auto needles_type = run_end_encoded(int32(), int32());
+ ASSERT_OK_AND_ASSIGN(auto left_chunk, REEFromJSON(needles_type, "[0, 0, 1, 1]"));
+ ASSERT_OK_AND_ASSIGN(auto right_chunk, REEFromJSON(needles_type, "[4, 4, 9]"));
+ auto needles = std::make_shared<ChunkedArray>(ArrayVector{left_chunk, right_chunk});
+
+ CheckSearchSorted(Datum(values), Datum(needles), SearchSortedOptions::Right,
+ "[0, 0, 2, 2, 3, 3, 5]");
+}
+
+TEST(SearchSorted, ChunkedRunEndEncodedValuesLeadingNullsAcrossEmptyChunks) {
+ auto values_type = run_end_encoded(int16(), int32());
+ ASSERT_OK_AND_ASSIGN(auto empty_chunk, REEFromJSON(values_type, "[]"));
+ ASSERT_OK_AND_ASSIGN(auto null_chunk, REEFromJSON(values_type, "[null, null]"));
+ ASSERT_OK_AND_ASSIGN(auto data_chunk, REEFromJSON(values_type, "[2, 4, 4]"));
+ auto values = std::make_shared<ChunkedArray>(
+ ArrayVector{empty_chunk, null_chunk, empty_chunk, data_chunk});
+ auto needles = ArrayFromJSON(int32(), "[1, 4, 8]");
+
+ CheckSearchSorted(Datum(values), Datum(needles), "[2, 3, 5]", "[2, 5, 5]");
+}
+
+TEST(SearchSorted, ChunkedRunEndEncodedValuesTrailingNullsAcrossEmptyChunks) {
+ auto values_type = run_end_encoded(int16(), int32());
+ ASSERT_OK_AND_ASSIGN(auto data_chunk, REEFromJSON(values_type, "[2, 4, 4]"));
+ ASSERT_OK_AND_ASSIGN(auto empty_chunk, REEFromJSON(values_type, "[]"));
+ ASSERT_OK_AND_ASSIGN(auto null_chunk, REEFromJSON(values_type, "[null, null]"));
+ auto values =
+ std::make_shared<ChunkedArray>(ArrayVector{data_chunk, empty_chunk, null_chunk});
+ auto needles = ArrayFromJSON(int32(), "[1, 4, 8]");
+
+ CheckSearchSorted(Datum(values), Datum(needles), "[0, 1, 3]", "[0, 3, 3]");
+}
+
+TEST(SearchSorted, ChunkedValuesLeadingNullsAcrossEmptyChunks) {
+ auto values = ChunkedArrayFromJSON(int32(), {"[]", "[null, null]", "[]", "[2, 4, 4]"});
+ auto needles = ArrayFromJSON(int32(), "[1, 4, 8]");
+
+ CheckSearchSorted(Datum(values), Datum(needles), "[2, 3, 5]", "[2, 5, 5]");
+}
+
+TEST(SearchSorted, ChunkedValuesTrailingNullsAcrossEmptyChunks) {
+ auto values = ChunkedArrayFromJSON(int32(), {"[]", "[2, 4, 4]", "[]", "[null, null]"});
+ auto needles = ArrayFromJSON(int32(), "[1, 4, 8]");
+
+ CheckSearchSorted(Datum(values), Datum(needles), "[0, 1, 3]", "[0, 3, 3]");
+}
+
+TEST(SearchSorted, RunEndEncodedDate32Values) {
+ auto values_type = run_end_encoded(int16(), date32());
+ ASSERT_OK_AND_ASSIGN(auto values, REEFromJSON(values_type, "[1, 1, 3]"));
+ auto needles = ArrayFromJSON(date32(), "[2]");
+
+ CheckSearchSorted(Datum(values), Datum(needles), "[2]", "[2]");
+}
+
+TEST(SearchSorted, RunEndEncodedNulls) {
+ auto values_type = run_end_encoded(int16(), int32());
+ ASSERT_OK_AND_ASSIGN(auto ree_values,
+ REEFromJSON(values_type, "[null, null, 2, 4, 4]"));
+ auto needles_type = run_end_encoded(int16(), int32());
+ ASSERT_OK_AND_ASSIGN(auto ree_needles,
+ REEFromJSON(needles_type, "[null, null, 1, 4, 4, null, 8]"));
+
+ CheckSearchSorted(Datum(ree_values), Datum(ree_needles),
+ "[null, null, 2, 3, 3, null, 5]", "[null, null, 2, 5, 5, null, 5]");
+}
+
+TEST(SearchSorted, RunEndEncodedValuesWithTrailingNulls) {
+ auto values_type = run_end_encoded(int16(), int32());
+ ASSERT_OK_AND_ASSIGN(auto ree_values,
+ REEFromJSON(values_type, "[2, 4, 4, null, null]"));
+ auto needles = ArrayFromJSON(int32(), "[1, 4, 8]");
+
+ CheckSearchSorted(Datum(ree_values), Datum(needles), "[0, 1, 3]", "[0, 3, 3]");
+}
+
+TEST(SearchSorted, SlicedRunEndEncodedValuesIgnoreNullRunsOutsideSlice) {
+ auto values_type = run_end_encoded(int16(), int32());
+ ASSERT_OK_AND_ASSIGN(auto ree_values,
+ REEFromJSON(values_type, "[null, null, 2, 4, 4, null]"));
+ auto sliced = ree_values->Slice(2, 3);
+ auto needles = ArrayFromJSON(int32(), "[1, 4, 8]");
+
+ CheckSearchSorted(Datum(sliced), Datum(needles), "[0, 1, 3]", "[0, 3, 3]");
+}
+
+TEST(SearchSorted, RunEndEncodedNeedlesWithNullRuns) {
+ auto values = ArrayFromJSON(int32(), "[1, 1, 3, 5, 8]");
+ auto needles_type = run_end_encoded(int32(), int32());
+ ASSERT_OK_AND_ASSIGN(
+ auto ree_needles,
+ REEFromJSON(needles_type, "[null, null, 0, 0, 0, 1, 1, 4, 4, 4, null, 9, 9]"));
+
+ CheckSearchSorted(Datum(values), Datum(ree_needles),
+ "[null, null, 0, 0, 0, 0, 0, 3, 3, 3, null, 5, 5]",
+ "[null, null, 0, 0, 0, 2, 2, 3, 3, 3, null, 5, 5]");
+}
+
+TEST(SearchSorted, RejectMismatchedTypes) {
+ auto values = ArrayFromJSON(int32(), "[1, 2, 3]");
+ auto needles = ArrayFromJSON(int64(), "[2]");
+
+ ASSERT_RAISES(TypeError, SearchSorted(Datum(values), Datum(needles)));
+}
+
+TEST(SearchSorted, RunEndEncodedValues) {
+ auto values_type = run_end_encoded(int16(), int32());
+ ASSERT_OK_AND_ASSIGN(auto ree_values, REEFromJSON(values_type, "[1, 1, 1, 3, 3, 5]"));
+ auto needles = ArrayFromJSON(int32(), "[0, 1, 2, 3, 4, 5, 6]");
+
+ CheckSearchSorted(Datum(ree_values), Datum(needles), "[0, 0, 3, 3, 5, 5, 6]",
+ "[0, 3, 3, 5, 5, 6, 6]");
+}
+
+TEST(SearchSorted, RunEndEncodedNeedles) {
+ auto values = ArrayFromJSON(int32(), "[1, 1, 3, 5, 8]");
+ auto needles_type = run_end_encoded(int32(), int32());
+ ASSERT_OK_AND_ASSIGN(auto ree_needles,
+ REEFromJSON(needles_type, "[0, 0, 1, 1, 4, 4, 9]"));
+
+ CheckSearchSorted(Datum(values), Datum(ree_needles), SearchSortedOptions::Right,
+ "[0, 0, 2, 2, 3, 3, 5]");
+}
+
+TEST(SearchSorted, SlicedRunEndEncodedNeedles) {
+ auto values = ArrayFromJSON(int32(), "[1, 1, 3, 5, 8]");
+ auto needles_type = run_end_encoded(int32(), int32());
+ ASSERT_OK_AND_ASSIGN(auto ree_needles,
+ REEFromJSON(needles_type, "[null, 0, 0, 1, 1, 4, 4, 9, null]"));
+ auto sliced = ree_needles->Slice(1, 7);
+
+ CheckSearchSorted(Datum(values), Datum(sliced), "[0, 0, 0, 0, 3, 3, 5]",
+ "[0, 0, 2, 2, 3, 3, 5]");
+}
+
+TEST(SearchSorted, SlicedRunEndEncodedValues) {
+ auto values_type = run_end_encoded(int32(), int32());
+ ASSERT_OK_AND_ASSIGN(auto ree_values,
+ REEFromJSON(values_type, "[10, 10, 20, 20, 20, 40, 40, 90]"));
+ auto sliced = ree_values->Slice(1, 5);
+ auto needles = ArrayFromJSON(int32(), "[5, 10, 20, 30, 40, 80, 90, 100]");
+ CheckSearchSorted(Datum(sliced), Datum(needles), "[0, 0, 1, 4, 4, 5, 5, 5]",
+ "[0, 1, 4, 4, 5, 5, 5, 5]");
+
+ sliced = ree_values->Slice(3, 4);
+ CheckSearchSorted(Datum(sliced), Datum(needles), "[0, 0, 0, 2, 2, 4, 4, 4]",
+ "[0, 0, 2, 2, 4, 4, 4, 4]");
+}
+
+TEST(SearchSorted, SlicedRunEndEncodedValuesWithLeadingNulls) {
+ auto values_type = run_end_encoded(int32(), int32());
+ ASSERT_OK_AND_ASSIGN(
+ auto ree_values,
+ REEFromJSON(values_type, "[null, null, 10, 20, 20, 20, 40, 40, 90]"));
+ auto sliced = ree_values->Slice(1, 5);
+ auto needles = ArrayFromJSON(int32(), "[5, 10, 20, 30, 40, 50]");
+ CheckSearchSorted(Datum(sliced), Datum(needles), "[1, 1, 2, 5, 5, 5]",
+ "[1, 2, 5, 5, 5, 5]");
+
+ sliced = ree_values->Slice(3, 4);
+ CheckSearchSorted(Datum(sliced), Datum(needles), "[0, 0, 0, 3, 3, 4]",
+ "[0, 0, 3, 3, 4, 4]");
+}
+
+TEST_P(SearchSortedSupportedTypesTest, ArraySmoke) {
+ const auto& param = GetParam();
+ CheckSimpleSearchSorted(param.type, param.values_json, param.needles_json,
+ param.expected_left_json, param.expected_right_json);
+}
+
+TEST_P(SearchSortedSupportedTypesTest, ScalarSmoke) {
+ const auto& param = GetParam();
+ CheckSimpleScalarSearchSorted(param.type, param.values_json, param.needles_json,
+ param.expected_left_json, param.expected_right_json);
+}
+
+INSTANTIATE_TEST_SUITE_P(SupportedTypes, SearchSortedSupportedTypesTest,
+ ::testing::ValuesIn(SupportedTypeSmokeCases()),
+ [](const ::testing::TestParamInfo<SearchSortedSmokeCase>& info) {
+ return info.param.name;
+ });
+
+} // namespace
+} // namespace compute
+} // namespace arrow
diff --git a/cpp/src/arrow/compute/kernels/vector_sort_internal.h b/cpp/src/arrow/compute/kernels/vector_sort_internal.h
index c19acc5..8bb10ef 100644
--- a/cpp/src/arrow/compute/kernels/vector_sort_internal.h
+++ b/cpp/src/arrow/compute/kernels/vector_sort_internal.h
@@ -31,6 +31,8 @@
#include "arrow/table.h"
#include "arrow/type.h"
#include "arrow/type_traits.h"
+#include "arrow/util/float16.h"
+#include "arrow/util/math_internal.h"
namespace arrow::compute::internal {
@@ -77,19 +79,24 @@
template <typename TypeClass>
constexpr bool has_null_like_values() {
- return is_physical_floating(TypeClass::type_id);
+ return is_floating(TypeClass::type_id);
+}
+
+template <typename TypeClass, typename Value>
+bool is_null_like(Value value) {
+ return ::arrow::internal::IsNan(value);
}
// Compare two values, taking NaNs into account
template <typename Type, typename Value>
int CompareTypeValues(Value&& left, Value&& right, SortOrder order,
- NullPlacement null_placement) {
+ NullPlacement null_placement, int on_equality_result = 0) {
if constexpr (has_null_like_values<Type>()) {
- const bool is_nan_left = std::isnan(left);
- const bool is_nan_right = std::isnan(right);
+ const bool is_nan_left = is_null_like<Type>(left);
+ const bool is_nan_right = is_null_like<Type>(right);
if (is_nan_left && is_nan_right) {
- return 0;
+ return on_equality_result;
} else if (is_nan_left) {
return null_placement == NullPlacement::AtStart ? -1 : 1;
} else if (is_nan_right) {
@@ -98,7 +105,7 @@
}
int compared;
if (left == right) {
- compared = 0;
+ compared = on_equality_result;
} else if (left > right) {
compared = 1;
} else {
@@ -252,17 +259,22 @@
template <typename ArrayType, typename Partitioner>
NanPartition PartitionNans(std::span<uint64_t> indices, const ArrayType& values,
int64_t offset, NullPlacement null_placement) {
- if constexpr (has_null_like_values<typename ArrayType::TypeClass>()) {
+ using TypeClass = typename ArrayType::TypeClass;
+ if constexpr (has_null_like_values<TypeClass>()) {
Partitioner partitioner;
if (null_placement == NullPlacement::AtStart) {
auto non_null_like_tail = partitioner(indices, [&values, &offset](uint64_t ind) {
- return std::isnan(values.GetView(static_cast<int64_t>(ind) - offset));
+ const auto value = GetViewType<TypeClass>::LogicalValue(
+ values.GetView(static_cast<int64_t>(ind) - offset));
+ return is_null_like<TypeClass>(value);
});
return NanPartition{.non_null_like_range = non_null_like_tail,
.nan_range = {indices.data(), non_null_like_tail.data()}};
} else {
auto nan_tail = partitioner(indices, [&values, &offset](uint64_t ind) {
- return !std::isnan(values.GetView(static_cast<int64_t>(ind) - offset));
+ const auto value = GetViewType<TypeClass>::LogicalValue(
+ values.GetView(static_cast<int64_t>(ind) - offset));
+ return !is_null_like<TypeClass>(value);
});
return NanPartition{.non_null_like_range = {indices.data(), nan_tail.data()},
.nan_range = nan_tail};
diff --git a/cpp/src/arrow/compute/registry_internal.h b/cpp/src/arrow/compute/registry_internal.h
index 5b9d7f8..d457f6f 100644
--- a/cpp/src/arrow/compute/registry_internal.h
+++ b/cpp/src/arrow/compute/registry_internal.h
@@ -50,6 +50,7 @@
void RegisterVectorNested(FunctionRegistry* registry);
void RegisterVectorRank(FunctionRegistry* registry);
void RegisterVectorReplace(FunctionRegistry* registry);
+void RegisterVectorSearchSorted(FunctionRegistry* registry);
void RegisterVectorSelectK(FunctionRegistry* registry);
void RegisterVectorSelection(FunctionRegistry* registry);
void RegisterVectorSort(FunctionRegistry* registry);
diff --git a/cpp/src/arrow/meson.build b/cpp/src/arrow/meson.build
index fea26ef..418165e 100644
--- a/cpp/src/arrow/meson.build
+++ b/cpp/src/arrow/meson.build
@@ -610,6 +610,7 @@
'compute/kernels/vector_rank.cc',
'compute/kernels/vector_replace.cc',
'compute/kernels/vector_run_end_encode.cc',
+ 'compute/kernels/vector_search_sorted.cc',
'compute/kernels/vector_select_k.cc',
'compute/kernels/vector_sort.cc',
'compute/kernels/vector_statistics.cc',
diff --git a/cpp/src/arrow/testing/random.cc b/cpp/src/arrow/testing/random.cc
index ce73b3f..db7a60e 100644
--- a/cpp/src/arrow/testing/random.cc
+++ b/cpp/src/arrow/testing/random.cc
@@ -966,24 +966,42 @@
}
std::shared_ptr<Array> RandomArrayGenerator::RunEndEncoded(
- std::shared_ptr<DataType> value_type, int64_t logical_size, double null_probability) {
- Int32Builder run_ends_builder;
- pcg32 rng(seed());
+ std::shared_ptr<DataType> value_type, int64_t logical_size, double null_probability,
+ int64_t average_run_length) {
+ const auto physical_size = std::max<int64_t>(1, logical_size / average_run_length);
+ std::shared_ptr<Array> values =
+ ArrayOf(std::move(value_type), physical_size, null_probability);
+ return RunEndEncoded(values, logical_size);
+}
- DCHECK_LE(logical_size, std::numeric_limits<int32_t>::max());
-
- std::uniform_int_distribution<int64_t> distribution(1, 100);
- int64_t current_end = 0;
- while (current_end < logical_size) {
- current_end += distribution(rng);
- current_end = std::min(current_end, logical_size);
- ARROW_CHECK_OK(run_ends_builder.Append(static_cast<int32_t>(current_end)));
+std::shared_ptr<Array> RandomArrayGenerator::RunEndEncoded(
+ const std::shared_ptr<Array>& values, int64_t logical_size) {
+ if (logical_size == 0) {
+ return MakeEmptyArray(run_end_encoded(int32(), values->type())).ValueOrDie();
}
- std::shared_ptr<Array> run_ends = *run_ends_builder.Finish();
- std::shared_ptr<Array> values =
- ArrayOf(std::move(value_type), run_ends->length(), null_probability);
+ ARROW_CHECK_GT(values->length(), 0);
+ ARROW_CHECK_LE(logical_size, std::numeric_limits<int32_t>::max());
+ ARROW_CHECK_GE(logical_size, values->length());
+ int64_t current_end = 0;
+ Int32Builder run_ends_builder;
+ ARROW_CHECK_OK(run_ends_builder.Reserve(values->length()));
+
+ // Ideally, we would generate random run-ends, but we need to make both unique
+ // and monotonic, so we use the same (approximate) run length instead.
+ for (int64_t i = 0; i < values->length(); ++i) {
+ auto remaining_logical_size = logical_size - current_end;
+ auto remaining_physical_size = values->length() - i;
+ // Runs must never be empty.
+ auto run_length = static_cast<int64_t>(
+ ceil(static_cast<double>(remaining_logical_size) / remaining_physical_size));
+ current_end += run_length;
+ run_ends_builder.UnsafeAppend(static_cast<int32_t>(current_end));
+ }
+ DCHECK_LE(current_end, logical_size);
+
+ std::shared_ptr<Array> run_ends = *run_ends_builder.Finish();
return RunEndEncodedArray::Make(logical_size, run_ends, values).ValueOrDie();
}
diff --git a/cpp/src/arrow/testing/random.h b/cpp/src/arrow/testing/random.h
index bc21307..a953547 100644
--- a/cpp/src/arrow/testing/random.h
+++ b/cpp/src/arrow/testing/random.h
@@ -560,7 +560,20 @@
/// \return a generated Array
std::shared_ptr<Array> RunEndEncoded(std::shared_ptr<DataType> value_type,
int64_t logical_size,
- double null_probability = 0.0);
+ double null_probability = 0.0,
+ int64_t average_run_length = 50);
+
+ /// \brief Generate a random RunEndEncodedArray
+ ///
+ /// \param[in] values The underlying physical values
+ /// \param[in] logical_size The logical length of the generated array
+ ///
+ /// `logical_size` must be at least as large as the length of `values`,
+ /// and `values` must be non-empty if `logical_size` is non-zero.
+ ///
+ /// \return a generated Array
+ std::shared_ptr<Array> RunEndEncoded(const std::shared_ptr<Array>& values,
+ int64_t logical_size);
/// \brief Generate a random SparseUnionArray
///
diff --git a/cpp/src/arrow/testing/random_test.cc b/cpp/src/arrow/testing/random_test.cc
index 279fb6d..e1fddbd 100644
--- a/cpp/src/arrow/testing/random_test.cc
+++ b/cpp/src/arrow/testing/random_test.cc
@@ -647,20 +647,48 @@
TEST(RandomRunEndEncoded, Basics) {
random::RandomArrayGenerator rng(42);
for (const double null_probability : {0.0, 0.1, 1.0}) {
- SCOPED_TRACE("null_probability = " + std::to_string(null_probability));
- auto array = rng.ArrayOf(run_end_encoded(int32(), int16()), 12345, null_probability);
+ ARROW_SCOPED_TRACE("null_probability = ", null_probability);
+ for (const int64_t logical_length : {12345}) {
+ auto check_run_end_encoded = [&](const std::shared_ptr<Array>& array,
+ int64_t average_run_length) {
+ ASSERT_OK(array->ValidateFull());
+ ASSERT_EQ(array->length(), logical_length);
+ const auto& ree_array = checked_cast<const RunEndEncodedArray&>(*array);
+ ASSERT_EQ(*ree_array.type(), *run_end_encoded(int32(), int16()));
+ const int64_t physical_length = ree_array.run_ends()->length();
+ ASSERT_EQ(ree_array.values()->length(), physical_length);
+ const auto actual_average_run_length =
+ static_cast<double>(logical_length) / physical_length;
+ ASSERT_GE(actual_average_run_length, average_run_length * 0.9);
+ ASSERT_LE(actual_average_run_length, average_run_length * 1.1);
+ if (null_probability == 0.0) {
+ ASSERT_EQ(ree_array.values()->null_count(), 0);
+ }
+ if (null_probability == 1.0) {
+ ASSERT_EQ(ree_array.values()->null_count(), physical_length);
+ }
+ };
+
+ auto array = rng.ArrayOf(run_end_encoded(int32(), int16()), logical_length,
+ null_probability);
+ // 50 is the default value in RandomArrayGenerator::RunEndEncoded
+ check_run_end_encoded(array, /*average_run_length=*/50);
+
+ for (const int64_t physical_length : {100, 1000}) {
+ auto values =
+ rng.Int16(physical_length, /*min=*/0, /*max=*/16384, null_probability);
+ auto array = rng.RunEndEncoded(values, logical_length);
+ check_run_end_encoded(array,
+ /*average_run_length=*/logical_length / physical_length);
+ }
+ }
+ }
+ // Small-sized REE arrays should be generated adequately
+ for (const int64_t logical_length : {0, 1, 10, 55}) {
+ auto array = rng.ArrayOf(run_end_encoded(int32(), int16()), logical_length,
+ /*null_probability=*/0.2);
ASSERT_OK(array->ValidateFull());
- ASSERT_EQ(array->length(), 12345);
- const auto& ree_array = checked_cast<const RunEndEncodedArray&>(*array);
- ASSERT_EQ(*ree_array.type(), *run_end_encoded(int32(), int16()));
- const int64_t physical_length = ree_array.run_ends()->length();
- ASSERT_EQ(ree_array.values()->length(), physical_length);
- if (null_probability == 0.0) {
- ASSERT_EQ(ree_array.values()->null_count(), 0);
- }
- if (null_probability == 1.0) {
- ASSERT_EQ(ree_array.values()->null_count(), physical_length);
- }
+ ASSERT_EQ(array->length(), logical_length);
}
}
diff --git a/cpp/src/arrow/util/math_internal.h b/cpp/src/arrow/util/math_internal.h
index a57083c..0375fc5 100644
--- a/cpp/src/arrow/util/math_internal.h
+++ b/cpp/src/arrow/util/math_internal.h
@@ -21,6 +21,7 @@
#include <cmath>
#include <initializer_list>
+#include "arrow/util/float16.h"
#include "arrow/util/macros.h"
#include "arrow/util/visibility.h"
@@ -75,4 +76,8 @@
static_assert(ReversePow2(4) == 2);
static_assert(ReversePow2(2) == 1);
+inline bool IsNan(float v) { return std::isnan(v); }
+inline bool IsNan(double v) { return std::isnan(v); }
+inline bool IsNan(::arrow::util::Float16 v) { return v.is_nan(); }
+
} // namespace arrow::internal
diff --git a/docs/source/cpp/compute.rst b/docs/source/cpp/compute.rst
index 1e067c5..f24fea2 100644
--- a/docs/source/cpp/compute.rst
+++ b/docs/source/cpp/compute.rst
@@ -1879,6 +1879,8 @@
+-----------------------+------------+---------------------------------------------------------+-------------------+-------------------------------+----------------+
| sort_indices | Unary | Boolean, Numeric, Temporal, Binary- and String-like | UInt64 | :struct:`SortOptions` | \(1) \(6) |
+-----------------------+------------+---------------------------------------------------------+-------------------+-------------------------------+----------------+
+| search_sorted | Binary | Boolean, Numeric, Temporal, Binary- and String-like | UInt64 | :struct:`SearchSortedOptions` | \(8) |
++-----------------------+------------+---------------------------------------------------------+-------------------+-------------------------------+----------------+
* \(1) The output is an array of indices into the input, that define a
@@ -1907,6 +1909,13 @@
* \(7) The output is an array of indices into the input, that define a
non-stable sort of the input.
+* \(8) The first argument must be sorted in ascending order. If it contains
+ nulls, they must be clustered entirely at the start or the end, and non-null
+ needles are matched only against the non-null portion. The second argument
+ may be a scalar, array, or run-end encoded array. Null needles yield null
+ outputs. Both arguments must have the same logical type. A scalar needle
+ yields a UInt64 scalar; otherwise the result is a UInt64 array.
+
.. _cpp-compute-vector-structural-transforms:
Structural transforms
diff --git a/python/pyarrow/_compute.pyx b/python/pyarrow/_compute.pyx
index da237d9..1c0c779 100644
--- a/python/pyarrow/_compute.pyx
+++ b/python/pyarrow/_compute.pyx
@@ -2080,6 +2080,14 @@
_raise_invalid_function_option(null_placement, "null placement")
+cdef CSearchSortedSide unwrap_search_sorted_side(side) except *:
+ if side == "left":
+ return CSearchSortedSide_Left
+ elif side == "right":
+ return CSearchSortedSide_Right
+ _raise_invalid_function_option(side, "search sorted side")
+
+
cdef class _PartitionNthOptions(FunctionOptions):
def _set_options(self, pivot, null_placement):
self.wrapped.reset(new CPartitionNthOptions(
@@ -2249,6 +2257,27 @@
self._set_options(order, null_placement)
+cdef class _SearchSortedOptions(FunctionOptions):
+ def _set_options(self, side):
+ self.wrapped.reset(new CSearchSortedOptions(
+ unwrap_search_sorted_side(side)))
+
+
+class SearchSortedOptions(_SearchSortedOptions):
+ """
+ Options for the `search_sorted` function.
+
+ Parameters
+ ----------
+ side : str, default "left"
+ Whether to return the leftmost or rightmost insertion point.
+ Accepted values are "left", "right".
+ """
+
+ def __init__(self, side="left"):
+ self._set_options(side)
+
+
cdef class _SortOptions(FunctionOptions):
def _set_options(self, sort_keys, null_placement):
if null_placement is None:
diff --git a/python/pyarrow/_compute_docstrings.py b/python/pyarrow/_compute_docstrings.py
index 079f00d..e5189f0 100644
--- a/python/pyarrow/_compute_docstrings.py
+++ b/python/pyarrow/_compute_docstrings.py
@@ -42,6 +42,41 @@
]
"""
+function_doc_additions["search_sorted"] = """
+ Examples
+ --------
+ >>> import pyarrow as pa
+ >>> import pyarrow.compute as pc
+ >>> values = pa.array([1, 1, 3, 5, 8])
+ >>> pc.search_sorted(values, pa.array([0, 1, 4, 9]))
+ <pyarrow.lib.UInt64Array object at ...>
+ [
+ 0,
+ 0,
+ 3,
+ 5
+ ]
+ >>> with_nulls = pa.array([None, 200, 300, 300], type=pa.int64())
+ >>> pc.search_sorted(
+ ... with_nulls, pa.array([50, 200, None, 400], type=pa.int64())
+ ... )
+ <pyarrow.lib.UInt64Array object at ...>
+ [
+ 1,
+ 1,
+ null,
+ 4
+ ]
+ >>> pc.search_sorted(values, pa.array([0, 1, 4, 9]), side="right")
+ <pyarrow.lib.UInt64Array object at ...>
+ [
+ 0,
+ 2,
+ 3,
+ 5
+ ]
+ """
+
function_doc_additions["mode"] = """
Examples
--------
diff --git a/python/pyarrow/compute.py b/python/pyarrow/compute.py
index 477ad7d..d34eccd 100644
--- a/python/pyarrow/compute.py
+++ b/python/pyarrow/compute.py
@@ -68,6 +68,7 @@
RoundToMultipleOptions,
ScalarAggregateOptions,
ScatterOptions,
+ SearchSortedOptions,
SelectKOptions,
SetLookupOptions,
SkewOptions,
diff --git a/python/pyarrow/includes/libarrow.pxd b/python/pyarrow/includes/libarrow.pxd
index 16784b5..455d11d 100644
--- a/python/pyarrow/includes/libarrow.pxd
+++ b/python/pyarrow/includes/libarrow.pxd
@@ -2831,6 +2831,16 @@
CSortOrder order
CNullPlacement null_placement
+ cdef enum CSearchSortedSide \
+ "arrow::compute::SearchSortedOptions::Side":
+ CSearchSortedSide_Left "arrow::compute::SearchSortedOptions::Left"
+ CSearchSortedSide_Right "arrow::compute::SearchSortedOptions::Right"
+
+ cdef cppclass CSearchSortedOptions \
+ "arrow::compute::SearchSortedOptions"(CFunctionOptions):
+ CSearchSortedOptions(CSearchSortedSide side)
+ CSearchSortedSide side
+
cdef cppclass CSortKey" arrow::compute::SortKey":
CSortKey(CFieldRef target, CSortOrder order)
CSortKey(CFieldRef target, CSortOrder order, CNullPlacement null_placement)
diff --git a/python/pyarrow/tests/test_compute.py b/python/pyarrow/tests/test_compute.py
index 68fe9a9..83e36a7 100644
--- a/python/pyarrow/tests/test_compute.py
+++ b/python/pyarrow/tests/test_compute.py
@@ -3267,6 +3267,96 @@
pc.array_sort_indices(arr, order="nonscending")
+def test_search_sorted():
+ values = pa.array([1, 1, 3, 5, 8])
+ needles = pa.array([0, 1, 3, 4, 5, 8, 9])
+
+ expected_left = pa.array([0, 0, 2, 3, 3, 4, 5], type=pa.uint64())
+ expected_right = pa.array([0, 2, 3, 3, 4, 5, 5], type=pa.uint64())
+
+ assert pc.search_sorted(values, needles).equals(expected_left)
+ assert pc.search_sorted(values, needles, side="left").equals(expected_left)
+ assert pc.search_sorted(values, needles, "right").equals(expected_right)
+ assert pc.search_sorted(
+ values, needles, options=pc.SearchSortedOptions(side="right")
+ ).equals(expected_right)
+
+ assert pc.search_sorted(values, pa.scalar(5, type=pa.int64())).as_py() == 3
+ assert pc.search_sorted(
+ values, pa.scalar(5, type=pa.int64()), side="right"
+ ).as_py() == 4
+
+
+def test_search_sorted_null_values():
+ needles = pa.array([50, 200, 250, 400], type=pa.int64())
+
+ values = pa.array([None, 200, 300, 300], type=pa.int64())
+ expected_left = pa.array([1, 1, 2, 4], type=pa.uint64())
+ expected_right = pa.array([1, 2, 2, 4], type=pa.uint64())
+ assert pc.search_sorted(values, needles, side="left").equals(expected_left)
+ assert pc.search_sorted(values, needles, side="right").equals(expected_right)
+
+ values = pa.array([200, 300, 300, None, None], type=pa.int64())
+ expected_left = pa.array([0, 0, 1, 3], type=pa.uint64())
+ expected_right = pa.array([0, 1, 1, 3], type=pa.uint64())
+ assert pc.search_sorted(values, needles, side="left").equals(expected_left)
+ assert pc.search_sorted(values, needles, side="right").equals(expected_right)
+
+
+def test_search_sorted_null_needles_emit_null():
+ values = pa.array([None, 200, 300, 300], type=pa.int64())
+ needles = pa.array([None, 50, 200, None, 400], type=pa.int64())
+
+ expected_left = pa.array([None, 1, 1, None, 4], type=pa.uint64())
+ expected_right = pa.array([None, 1, 2, None, 4], type=pa.uint64())
+
+ assert pc.search_sorted(values, needles, side="left").equals(expected_left)
+ assert pc.search_sorted(values, needles, side="right").equals(expected_right)
+
+ scalar_result = pc.search_sorted(values, pa.scalar(None, type=pa.int64()))
+ assert scalar_result.as_py() is None
+
+
+def test_search_sorted_run_end_encoded():
+ run_ends = pa.array([2, 3, 4, 5], type=pa.int16())
+ encoded_values = pa.array([1, 3, 5, 8], type=pa.int64())
+ values = pa.RunEndEncodedArray.from_arrays(run_ends, encoded_values)
+ needles = pa.array([0, 1, 3, 4, 5, 8, 9], type=pa.int64())
+
+ expected_left = pa.array([0, 0, 2, 3, 3, 4, 5], type=pa.uint64())
+ assert pc.search_sorted(values, needles).equals(expected_left)
+
+ ree_needles = pa.RunEndEncodedArray.from_arrays(
+ pa.array([2, 4, 6], type=pa.int16()),
+ pa.array([1, 4, 9], type=pa.int64())
+ )
+ expected_right = pa.array([2, 2, 3, 3, 5, 5], type=pa.uint64())
+ assert pc.search_sorted(values, ree_needles, side="right").equals(
+ expected_right
+ )
+
+
+def test_search_sorted_run_end_encoded_nulls():
+ values = pa.RunEndEncodedArray.from_arrays(
+ pa.array([2, 3, 5], type=pa.int16()),
+ pa.array([None, 2, 4], type=pa.int64())
+ )
+ needles = pa.RunEndEncodedArray.from_arrays(
+ pa.array([2, 3, 5, 6], type=pa.int16()),
+ pa.array([None, 1, 4, None], type=pa.int64())
+ )
+
+ expected = pa.array([None, None, 2, 3, 3, None], type=pa.uint64())
+ assert pc.search_sorted(values, needles, side="left").equals(expected)
+
+
+def test_search_sorted_errors():
+ values = pa.array([1, 1, 3, 5, 8])
+
+ with pytest.raises(ValueError, match='"middle" is not a valid search sorted side'):
+ pc.search_sorted(values, pa.array([1]), side="middle")
+
+
def test_sort_indices_array():
arr = pa.array([1, 2, None, 0])
result = pc.sort_indices(arr)
