From 764023f5ec79ff0eff7554e4975800823861f063 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Fri, 21 Mar 2025 17:38:13 +0800 Subject: [PATCH 01/71] Special form v2 initial implementation --- cpp/src/arrow/CMakeLists.txt | 1 + cpp/src/arrow/compute/CMakeLists.txt | 8 + cpp/src/arrow/compute/exec.cc | 53 +- cpp/src/arrow/compute/exec.h | 48 +- cpp/src/arrow/compute/exec_internal.h | 3 + cpp/src/arrow/compute/expression.cc | 269 +++- cpp/src/arrow/compute/expression.h | 24 +- cpp/src/arrow/compute/kernel.h | 2 + cpp/src/arrow/compute/special_form.cc | 1001 +++++++++++++++ cpp/src/arrow/compute/special_form.h | 53 + .../arrow/compute/special_form_benchmark.cc | 283 +++++ cpp/src/arrow/compute/special_form_test.cc | 1086 +++++++++++++++++ cpp/src/arrow/compute/type_fwd.h | 3 + 13 files changed, 2765 insertions(+), 69 deletions(-) create mode 100644 cpp/src/arrow/compute/special_form.cc create mode 100644 cpp/src/arrow/compute/special_form.h create mode 100644 cpp/src/arrow/compute/special_form_benchmark.cc create mode 100644 cpp/src/arrow/compute/special_form_test.cc diff --git a/cpp/src/arrow/CMakeLists.txt b/cpp/src/arrow/CMakeLists.txt index 2e5c67e07b6e..a6975dcf4af8 100644 --- a/cpp/src/arrow/CMakeLists.txt +++ b/cpp/src/arrow/CMakeLists.txt @@ -728,6 +728,7 @@ set(ARROW_COMPUTE_SRCS compute/kernel.cc compute/ordering.cc compute/registry.cc + compute/special_form.cc compute/kernels/chunked_internal.cc compute/kernels/codegen_internal.cc compute/kernels/scalar_cast_boolean.cc diff --git a/cpp/src/arrow/compute/CMakeLists.txt b/cpp/src/arrow/compute/CMakeLists.txt index 6c530a76e18f..16c8a4a69f57 100644 --- a/cpp/src/arrow/compute/CMakeLists.txt +++ b/cpp/src/arrow/compute/CMakeLists.txt @@ -174,8 +174,16 @@ add_arrow_compute_test(row_test EXTRA_LINK_LIBS arrow_compute_testing) +add_arrow_compute_test(special_form_test + SOURCES + special_form_test.cc + EXTRA_LINK_LIBS + arrow_compute_testing) + add_arrow_compute_benchmark(function_benchmark) +add_arrow_compute_benchmark(special_form_benchmark PREFIX "arrow-compute") + add_subdirectory(kernels) add_subdirectory(row) diff --git a/cpp/src/arrow/compute/exec.cc b/cpp/src/arrow/compute/exec.cc index 1be398fdae9e..f3c855eda0b5 100644 --- a/cpp/src/arrow/compute/exec.cc +++ b/cpp/src/arrow/compute/exec.cc @@ -27,6 +27,7 @@ #include "arrow/array/array_base.h" #include "arrow/array/array_primitive.h" +#include "arrow/array/builder_primitive.h" #include "arrow/array/data.h" #include "arrow/array/util.h" #include "arrow/buffer.h" @@ -50,6 +51,7 @@ #include "arrow/util/logging_internal.h" #include "arrow/util/thread_pool.h" #include "arrow/util/vector.h" +#include "arrow/visit_data_inline.h" namespace arrow { @@ -367,6 +369,7 @@ Status ExecSpanIterator::Init(const ExecBatch& batch, int64_t max_chunksize, value_offsets_.clear(); value_offsets_.resize(args_->size(), 0); max_chunksize_ = std::min(length_, max_chunksize); + selection_vector_ = batch.selection_vector.get(); return Status::OK(); } @@ -440,6 +443,15 @@ bool ExecSpanIterator::Next(ExecSpan* span) { if (have_all_scalars_ && promote_if_all_scalars_) { PromoteExecSpanScalars(span); + } else { + if (selection_vector_) { + if (have_chunked_arrays_) { + span->selection_vector = SelectionVectorSpan(selection_vector_->indices(), 0); + } else { + span->selection_vector = SelectionVectorSpan(selection_vector_->indices(), + selection_vector_->length()); + } + } } initialized_ = true; @@ -452,6 +464,19 @@ bool ExecSpanIterator::Next(ExecSpan* span) { int64_t iteration_size = std::min(length_ - position_, max_chunksize_); if (have_chunked_arrays_) { iteration_size = GetNextChunkSpan(iteration_size, span); + if (selection_vector_) { + auto indices_begin = + span->selection_vector.indices() + span->selection_vector.length(); + auto indices_end = selection_vector_->indices() + selection_vector_->length(); + DCHECK_LE(indices_begin, indices_end); + auto chunk_row_id_end = position_ + iteration_size; + int64_t num_indices = 0; + while (indices_begin + num_indices < indices_end && + *(indices_begin + num_indices) < chunk_row_id_end) { + ++num_indices; + } + span->selection_vector.SetSlice(span->selection_vector.length(), num_indices); + } } // Now, adjust the span @@ -1352,11 +1377,31 @@ SelectionVector::SelectionVector(std::shared_ptr data) SelectionVector::SelectionVector(const Array& arr) : SelectionVector(arr.data()) {} -int32_t SelectionVector::length() const { return static_cast(data_->length); } - Result> SelectionVector::FromMask( - const BooleanArray& arr) { - return Status::NotImplemented("FromMask"); + const BooleanArray& arr, MemoryPool* pool) { + Int32Builder builder(pool); + RETURN_NOT_OK(builder.Reserve(arr.true_count())); + ArraySpan span(*arr.data()); + int32_t i = 0; + VisitArraySpanInline( + span, + [&](bool mask) { + if (mask) { + builder.UnsafeAppend(i); + } + ++i; + }, + [&]() { ++i; }); + ARROW_ASSIGN_OR_RAISE(auto indices, builder.Finish()); + return std::make_shared(indices->data()); +} + +int64_t SelectionVector::length() const { return data_->length; } + +void SelectionVectorSpan::SetSlice(int64_t offset, int64_t length) { + DCHECK_NE(indices_, nullptr); + offset_ += offset; + length_ = length; } Result CallFunction(const std::string& func_name, const std::vector& args, diff --git a/cpp/src/arrow/compute/exec.h b/cpp/src/arrow/compute/exec.h index dae7e1ea6868..b6a0943d4a5c 100644 --- a/cpp/src/arrow/compute/exec.h +++ b/cpp/src/arrow/compute/exec.h @@ -141,16 +141,38 @@ class ARROW_EXPORT SelectionVector { explicit SelectionVector(const Array& arr); /// \brief Create SelectionVector from boolean mask - static Result> FromMask(const BooleanArray& arr); + static Result> FromMask( + const BooleanArray& arr, MemoryPool* pool = default_memory_pool()); + std::shared_ptr data() const { return data_; } const int32_t* indices() const { return indices_; } - int32_t length() const; + int64_t length() const; private: std::shared_ptr data_; const int32_t* indices_; }; +class ARROW_EXPORT SelectionVectorSpan { + public: + explicit SelectionVectorSpan(const int32_t* indices = NULLPTR, int64_t length = 0, + int64_t offset = 0) + : indices_(indices), length_(length), offset_(offset) {} + + void SetSlice(int64_t offset, int64_t length); + + const int32_t* indices() const { return indices_ + offset_; } + + int64_t length() const { return length_; } + + int64_t offset() const { return offset_; } + + private: + const int32_t* indices_; + int64_t length_; + int64_t offset_; +}; + /// An index to represent that a batch does not belong to an ordered stream constexpr int64_t kUnsequencedIndex = -1; @@ -173,8 +195,11 @@ constexpr int64_t kUnsequencedIndex = -1; struct ARROW_EXPORT ExecBatch { ExecBatch() = default; - ExecBatch(std::vector values, int64_t length) - : values(std::move(values)), length(length) {} + ExecBatch(std::vector values, int64_t length, + std::shared_ptr selection_vector = NULLPTR) + : values(std::move(values)), + length(length), + selection_vector(std::move(selection_vector)) {} explicit ExecBatch(const RecordBatch& batch); @@ -196,13 +221,6 @@ struct ARROW_EXPORT ExecBatch { /// exec function for processing. std::vector values; - /// A deferred filter represented as an array of indices into the values. - /// - /// For example, the filter [true, true, false, true] would be represented as - /// the selection vector [0, 1, 3]. When the selection vector is set, - /// ExecBatch::length is equal to the length of this array. - std::shared_ptr selection_vector; - /// A predicate Expression guaranteed to evaluate to true for all rows in this batch. Expression guarantee = literal(true); @@ -218,6 +236,13 @@ struct ARROW_EXPORT ExecBatch { /// whether any values are Scalar. int64_t length = 0; + /// A deferred filter represented as an array of indices into the values. + /// + /// For example, the filter [true, true, false, true] would be represented as + /// the selection vector [0, 1, 3]. When the selection vector is set, + /// ExecBatch::length is equal to the length of this array. + std::shared_ptr selection_vector; + /// \brief index of this batch in a sorted stream of batches /// /// This index must be strictly monotonic starting at 0 without gaps or @@ -422,6 +447,7 @@ struct ARROW_EXPORT ExecSpan { int64_t length = 0; std::vector values; + SelectionVectorSpan selection_vector; }; /// \defgroup compute-call-function One-shot calls to compute functions diff --git a/cpp/src/arrow/compute/exec_internal.h b/cpp/src/arrow/compute/exec_internal.h index 7e4f364a9288..b92cd257f586 100644 --- a/cpp/src/arrow/compute/exec_internal.h +++ b/cpp/src/arrow/compute/exec_internal.h @@ -42,6 +42,8 @@ namespace detail { /// \brief Break std::vector into a sequence of non-owning /// ExecSpan for kernel execution. The lifetime of the Datum vector /// must be longer than the lifetime of this object +// TODO: Can this struct sense the presence of selection vector and not split ExecBatch if +// so? class ARROW_EXPORT ExecSpanIterator { public: ExecSpanIterator() = default; @@ -83,6 +85,7 @@ class ARROW_EXPORT ExecSpanIterator { bool have_all_scalars_ = false; bool promote_if_all_scalars_ = true; const std::vector* args_; + SelectionVector* selection_vector_ = NULLPTR; std::vector chunk_indexes_; std::vector value_positions_; diff --git a/cpp/src/arrow/compute/expression.cc b/cpp/src/arrow/compute/expression.cc index 3c2ec1004022..ad95579f9fe6 100644 --- a/cpp/src/arrow/compute/expression.cc +++ b/cpp/src/arrow/compute/expression.cc @@ -30,6 +30,7 @@ #include "arrow/compute/exec_internal.h" #include "arrow/compute/expression_internal.h" #include "arrow/compute/function_internal.h" +#include "arrow/compute/special_form.h" #include "arrow/compute/util.h" #include "arrow/io/memory.h" #ifdef ARROW_IPC @@ -59,6 +60,13 @@ void Expression::Call::ComputeHash() { } } +void Expression::Special::ComputeHash() { + hash = std::hash{}(special_form->name); + for (const auto& arg : arguments) { + arrow::internal::hash_combine(hash, arg.hash()); + } +} + Expression::Expression(Call call) { call.ComputeHash(); impl_ = std::make_shared(std::move(call)); @@ -70,6 +78,9 @@ Expression::Expression(Datum literal) Expression::Expression(Parameter parameter) : impl_(std::make_shared(std::move(parameter))) {} +Expression::Expression(Special special) + : impl_(std::make_shared(std::move(special))) {} + Expression literal(Datum lit) { return Expression(std::move(lit)); } Expression field_ref(FieldRef ref) { @@ -85,6 +96,13 @@ Expression call(std::string function, std::vector arguments, return Expression(std::move(call)); } +Expression if_else_special(Expression cond, Expression if_true, Expression if_false) { + Expression::Special special; + special.special_form = GetIfElseSpecialForm(); + special.arguments = {std::move(cond), std::move(if_true), std::move(if_false)}; + return Expression(std::move(special)); +} + const Datum* Expression::literal() const { if (impl_ == nullptr) return nullptr; @@ -110,6 +128,12 @@ const Expression::Call* Expression::call() const { return std::get_if(impl_.get()); } +const Expression::Special* Expression::special() const { + if (impl_ == nullptr) return nullptr; + + return std::get_if(impl_.get()); +} + const DataType* Expression::type() const { if (impl_ == nullptr) return nullptr; @@ -121,9 +145,27 @@ const DataType* Expression::type() const { return parameter->type.type; } + if (const Special* special = this->special()) { + return special->type.type; + } + return CallNotNull(*this)->type.type; } +bool Expression::selection_vector_aware() const { + DCHECK(IsBound()); + + if (literal() || field_ref()) { + return true; + } + + if (auto special = this->special()) { + return special->selection_vector_aware; + } + + return CallNotNull(*this)->selection_vector_aware; +} + namespace { std::string PrintDatum(const Datum& datum) { @@ -169,6 +211,15 @@ std::string Expression::ToString() const { return ref->ToString(); } + if (auto sp = special()) { + std::string out = sp->special_form->name + "_special("; + for (const auto& arg : sp->arguments) { + out += arg.ToString() + ", "; + } + out.resize(out.size() - 2); + return out + ")"; + } + auto call = CallNotNull(*this); auto binary = [&](std::string op) { return "(" + call->arguments[0].ToString() + " " + op + " " + @@ -239,6 +290,12 @@ bool Expression::Equals(const Expression& other) const { return ref->Equals(*other.field_ref()); } + // TODO + // if (auto special = this->special()) { + if (this->special()) { + return true; + } + auto call = CallNotNull(*this); auto other_call = CallNotNull(other); @@ -276,6 +333,10 @@ size_t Expression::hash() const { return ref->hash(); } + if (auto special = this->special()) { + return special->hash; + } + return CallNotNull(*this)->hash; } @@ -300,6 +361,8 @@ bool Expression::IsScalarExpression() const { if (field_ref()) return true; + if (special()) return true; + auto call = CallNotNull(*this); for (const Expression& arg : call->arguments) { @@ -360,6 +423,8 @@ bool Expression::IsSatisfiable() const { if (field_ref()) return true; + if (special()) return true; + auto call = CallNotNull(*this); // invert(true_unless_null(x)) is always false or null by definition @@ -536,73 +601,141 @@ inline std::vector GetTypesWithSmallestLiteralRepresentation( return types; } -// Produce a bound Expression from unbound Call and bound arguments. Result BindNonRecursive(Expression::Call call, bool insert_implicit_casts, - compute::ExecContext* exec_context) { - DCHECK(std::all_of(call.arguments.begin(), call.arguments.end(), + compute::ExecContext* exec_context); +Result BindNonRecursive(Expression::Special special, + bool insert_implicit_casts, + compute::ExecContext* exec_context); + +template +Status DispatchForBind(DispatchExactFn&& dispatch_exact, DispatchBestFn&& dispatch_best, + FinishBind&& finish_bind, std::vector& arguments, + bool insert_implicit_casts, ExecContext* exec_context) { + DCHECK(std::all_of(arguments.begin(), arguments.end(), [](const Expression& argument) { return argument.IsBound(); })); - std::vector types = GetTypes(call.arguments); - ARROW_ASSIGN_OR_RAISE(call.function, GetFunction(call, exec_context)); + std::vector types = GetTypes(arguments); // First try and bind exactly - Result maybe_exact_match = call.function->DispatchExact(types); + auto maybe_exact_match = dispatch_exact(types); if (maybe_exact_match.ok()) { - call.kernel = *maybe_exact_match; - } else { - if (!insert_implicit_casts) { - return maybe_exact_match.status(); + auto output = std::move(*maybe_exact_match); + if (finish_bind(types, std::move(output)).ok()) { + return Status::OK(); } + } - // If exact binding fails, and we are allowed to cast, then prefer casting literals - // first. Since DispatchBest generally prefers up-casting the best way to do this is - // first down-cast the literals as much as possible - types = GetTypesWithSmallestLiteralRepresentation(call.arguments); - ARROW_ASSIGN_OR_RAISE(call.kernel, call.function->DispatchBest(&types)); + if (!insert_implicit_casts) { + return maybe_exact_match.status(); + } - for (size_t i = 0; i < types.size(); ++i) { - if (types[i] == call.arguments[i].type()) continue; + // If exact binding fails, and we are allowed to cast, then prefer casting literals + // first. Since DispatchBest generally prefers up-casting the best way to do this is + // first down-cast the literals as much as possible + types = GetTypesWithSmallestLiteralRepresentation(arguments); + ARROW_ASSIGN_OR_RAISE(auto output, dispatch_best(&types)); + // ARROW_ASSIGN_OR_RAISE(call.kernel, call.function->DispatchBest(&types)); - if (const Datum* lit = call.arguments[i].literal()) { - ARROW_ASSIGN_OR_RAISE(Datum new_lit, - compute::Cast(*lit, types[i].GetSharedPtr())); - call.arguments[i] = literal(std::move(new_lit)); - continue; - } + for (size_t i = 0; i < types.size(); ++i) { + if (types[i] == arguments[i].type()) continue; + + if (const Datum* lit = arguments[i].literal()) { + ARROW_ASSIGN_OR_RAISE(Datum new_lit, compute::Cast(*lit, types[i].GetSharedPtr())); + arguments[i] = literal(std::move(new_lit)); + continue; + } + + // construct an implicit cast Expression with which to replace this argument + Expression::Call implicit_cast; + implicit_cast.function_name = "cast"; + implicit_cast.arguments = {std::move(arguments[i])}; + + // TODO(wesm): Use TypeHolder in options + implicit_cast.options = std::make_shared( + compute::CastOptions::Safe(types[i].GetSharedPtr())); - // construct an implicit cast Expression with which to replace this argument - Expression::Call implicit_cast; - implicit_cast.function_name = "cast"; - implicit_cast.arguments = {std::move(call.arguments[i])}; + ARROW_ASSIGN_OR_RAISE( + arguments[i], BindNonRecursive(std::move(implicit_cast), + /*insert_implicit_casts=*/false, exec_context)); + } + + return finish_bind(types, std::move(output)); +} + +// Produce a bound Expression from unbound Call and bound arguments. +Result BindNonRecursive(Expression::Call call, bool insert_implicit_casts, + compute::ExecContext* exec_context) { + ARROW_ASSIGN_OR_RAISE(call.function, GetFunction(call, exec_context)); - // TODO(wesm): Use TypeHolder in options - implicit_cast.options = std::make_shared( - compute::CastOptions::Safe(types[i].GetSharedPtr())); + auto FinishBind = [&](const std::vector& types, const Kernel* kernel) { + call.kernel = kernel; + compute::KernelContext kernel_context(exec_context, call.kernel); + if (call.kernel->init) { + const FunctionOptions* options = + call.options ? call.options.get() : call.function->default_options(); ARROW_ASSIGN_OR_RAISE( - call.arguments[i], - BindNonRecursive(std::move(implicit_cast), - /*insert_implicit_casts=*/false, exec_context)); + call.kernel_state, + call.kernel->init(&kernel_context, {call.kernel, types, options})); + + kernel_context.SetState(call.kernel_state.get()); } - } - compute::KernelContext kernel_context(exec_context, call.kernel); - if (call.kernel->init) { - const FunctionOptions* options = - call.options ? call.options.get() : call.function->default_options(); - ARROW_ASSIGN_OR_RAISE( - call.kernel_state, - call.kernel->init(&kernel_context, {call.kernel, types, options})); + call.selection_vector_aware = + call.kernel->selection_vector_aware && + std::all_of(call.arguments.begin(), call.arguments.end(), + [](const Expression& arg) { return arg.selection_vector_aware(); }); - kernel_context.SetState(call.kernel_state.get()); - } + ARROW_ASSIGN_OR_RAISE( + call.type, call.kernel->signature->out_type().Resolve(&kernel_context, types)); + return Status::OK(); + }; - ARROW_ASSIGN_OR_RAISE( - call.type, call.kernel->signature->out_type().Resolve(&kernel_context, types)); + RETURN_NOT_OK(DispatchForBind( + [&](const std::vector& types) -> Result { + return call.function->DispatchExact(types); + }, + [&](std::vector* types) -> Result { + return call.function->DispatchBest(types); + }, + FinishBind, call.arguments, insert_implicit_casts, exec_context)); return Expression(std::move(call)); } +// Produce a bound Expression from unbound Special and bound arguments. +Result BindNonRecursive(Expression::Special special, + bool insert_implicit_casts, + compute::ExecContext* exec_context) { + DCHECK(std::all_of(special.arguments.begin(), special.arguments.end(), + [](const Expression& argument) { return argument.IsBound(); })); + + auto FinishBind = [&](const std::vector& types, + std::unique_ptr&& special_exec) { + special.special_exec = std::move(special_exec); + + // Selection vector awareness-es of the subexpressions is fully taken over by the + // special form so not recursive. + special.selection_vector_aware = special.special_form->selection_vector_aware; + + ARROW_ASSIGN_OR_RAISE(special.type, + special.special_exec->Bind(special.arguments, exec_context)); + + return Status::OK(); + }; + + RETURN_NOT_OK(DispatchForBind( + [&](const std::vector& types) -> Result> { + return special.special_form->DispatchExact(types, exec_context); + }, + [&](std::vector* types) -> Result> { + return special.special_form->DispatchBest(types, exec_context); + }, + FinishBind, special.arguments, insert_implicit_casts, exec_context)); + + return Expression(std::move(special)); +} + template Result BindImpl(Expression expr, const TypeOrSchema& in, compute::ExecContext* exec_context) { @@ -624,12 +757,19 @@ Result BindImpl(Expression expr, const TypeOrSchema& in, return Expression{std::move(param)}; } + if (expr.special()) { + auto special = *expr.special(); + for (auto& argument : special.arguments) { + ARROW_ASSIGN_OR_RAISE(argument, BindImpl(std::move(argument), in, exec_context)); + } + return BindNonRecursive(special, /*insert_implicit_casts=*/true, exec_context); + } + auto call = *CallNotNull(expr); for (auto& argument : call.arguments) { ARROW_ASSIGN_OR_RAISE(argument, BindImpl(std::move(argument), in, exec_context)); } - return BindNonRecursive(std::move(call), - /*insert_implicit_casts=*/true, exec_context); + return BindNonRecursive(call, /*insert_implicit_casts=*/true, exec_context); } } // namespace @@ -735,6 +875,8 @@ Result ExecuteScalarExpression(const Expression& expr, const ExecBatch& i "ExecuteScalarExpression cannot Execute non-scalar expression ", expr.ToString()); } + DCHECK(!input.selection_vector || expr.selection_vector_aware()); + if (auto lit = expr.literal()) return *lit; if (auto param = expr.parameter()) { @@ -758,6 +900,10 @@ Result ExecuteScalarExpression(const Expression& expr, const ExecBatch& i return field; } + if (auto special = expr.special()) { + return special->special_exec->Execute(input, exec_context); + } + auto call = CallNotNull(expr); std::vector arguments(call->arguments.size()); @@ -789,7 +935,8 @@ Result ExecuteScalarExpression(const Expression& expr, const ExecBatch& i RETURN_NOT_OK(executor->Init(&kernel_context, {kernel, types, options})); compute::detail::DatumAccumulator listener; - RETURN_NOT_OK(executor->Execute(ExecBatch(arguments, input_length), &listener)); + RETURN_NOT_OK(executor->Execute( + ExecBatch(arguments, input_length, input.selection_vector), &listener)); const auto out = executor->WrapResults(arguments, listener.values()); #ifndef NDEBUG DCHECK_OK(executor->CheckResultType(out, call->function_name.c_str())); @@ -817,12 +964,21 @@ std::vector FieldsInExpression(const Expression& expr) { return {*ref}; } - std::vector fields; - for (const Expression& arg : CallNotNull(expr)->arguments) { - auto argument_fields = FieldsInExpression(arg); - std::move(argument_fields.begin(), argument_fields.end(), std::back_inserter(fields)); + const auto& fields = [](const auto& expr) { + std::vector fields; + for (const Expression& arg : expr->arguments) { + auto argument_fields = FieldsInExpression(arg); + std::move(argument_fields.begin(), argument_fields.end(), + std::back_inserter(fields)); + } + return fields; + }; + + if (auto sp = expr.special()) { + return fields(sp); } - return fields; + + return fields(CallNotNull(expr)); } bool ExpressionHasFieldRefs(const Expression& expr) { @@ -830,6 +986,13 @@ bool ExpressionHasFieldRefs(const Expression& expr) { if (expr.field_ref()) return true; + if (auto sp = expr.special()) { + for (const Expression& arg : sp->arguments) { + if (ExpressionHasFieldRefs(arg)) return true; + } + return false; + } + for (const Expression& arg : CallNotNull(expr)->arguments) { if (ExpressionHasFieldRefs(arg)) return true; } diff --git a/cpp/src/arrow/compute/expression.h b/cpp/src/arrow/compute/expression.h index b8ce50675c8c..61cebfb8ddee 100644 --- a/cpp/src/arrow/compute/expression.h +++ b/cpp/src/arrow/compute/expression.h @@ -56,6 +56,21 @@ class ARROW_EXPORT Expression { const Kernel* kernel = NULLPTR; std::shared_ptr kernel_state; TypeHolder type; + bool selection_vector_aware; + + void ComputeHash(); + }; + + struct Special { + std::shared_ptr special_form; + std::vector arguments; + // Cached hash value + size_t hash; + + // post-Bind properties: + std::shared_ptr special_exec; + TypeHolder type; + bool selection_vector_aware; void ComputeHash(); }; @@ -112,12 +127,16 @@ class ARROW_EXPORT Expression { const Datum* literal() const; /// Access a FieldRef or return nullptr if this expression is not a field_ref const FieldRef* field_ref() const; + /// Access a FieldRef or return nullptr if this expression is not a field_ref + const Special* special() const; /// The type to which this expression will evaluate const DataType* type() const; // XXX someday // NullGeneralization::type nullable() const; + bool selection_vector_aware() const; + struct Parameter { FieldRef ref; @@ -131,11 +150,12 @@ class ARROW_EXPORT Expression { explicit Expression(Call call); explicit Expression(Datum literal); explicit Expression(Parameter parameter); + explicit Expression(Special special); static bool Identical(const Expression& l, const Expression& r); private: - using Impl = std::variant; + using Impl = std::variant; std::shared_ptr impl_; }; @@ -169,6 +189,8 @@ Expression call(std::string function, std::vector arguments, std::make_shared(std::move(options))); } +Expression if_else_special(Expression cond, Expression if_true, Expression if_false); + /// Assemble a list of all fields referenced by an Expression at any depth. ARROW_EXPORT std::vector FieldsInExpression(const Expression&); diff --git a/cpp/src/arrow/compute/kernel.h b/cpp/src/arrow/compute/kernel.h index 0d4f9d6ff436..e1d799a93946 100644 --- a/cpp/src/arrow/compute/kernel.h +++ b/cpp/src/arrow/compute/kernel.h @@ -542,6 +542,8 @@ struct ARROW_EXPORT Kernel { /// so that the most optimized kernel supported on a host's processor can be chosen. SimdLevel::type simd_level = SimdLevel::NONE; + bool selection_vector_aware = false; + // Additional kernel-specific data std::shared_ptr data; }; diff --git a/cpp/src/arrow/compute/special_form.cc b/cpp/src/arrow/compute/special_form.cc new file mode 100644 index 000000000000..6c0595dead86 --- /dev/null +++ b/cpp/src/arrow/compute/special_form.cc @@ -0,0 +1,1001 @@ +// 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/special_form.h" + +#include "arrow/array/builder_primitive.h" +#include "arrow/chunk_resolver.h" +#include "arrow/compute/api_vector.h" +#include "arrow/compute/exec.h" +#include "arrow/compute/expression.h" +#include "arrow/compute/expression_internal.h" +#include "arrow/compute/registry.h" +#include "arrow/util/logging_internal.h" + +namespace arrow::compute { + +namespace { + +// TODO: Clean free functions. + +struct BodyMask; + +struct BranchMask : public std::enable_shared_from_this { + virtual ~BranchMask() = default; + + Result> ApplyCondSparse( + const Expression& expr, const ExecBatch& input, ExecContext* exec_context) const { + ARROW_ASSIGN_OR_RAISE(auto datum, ApplySparse(expr, input, exec_context)); + return MakeBodyMaskFromSparseDatum(datum, exec_context); + } + + Result> ApplyCondDense( + const Expression& expr, const ExecBatch& input, ExecContext* exec_context) const { + ARROW_ASSIGN_OR_RAISE(auto datum, ApplyDense(expr, input, exec_context)); + return MakeBodyMaskFromDenseDatum(datum, exec_context); + } + + virtual bool empty() const = 0; + + protected: + virtual Result ApplySparse(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const = 0; + + virtual Result ApplyDense(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const = 0; + + virtual Result> GetSelectionVector() const = 0; + + virtual Result> MakeBodyMaskFromScalar( + const BooleanScalar& scalar, ExecContext* exec_context) const = 0; + + virtual Result> MakeBodyMaskFromSparseBitmap( + const std::shared_ptr& bitmap, ExecContext* exec_context) const = 0; + + virtual Result> MakeBodyMaskFromDenseBitmap( + const std::shared_ptr& bitmap, ExecContext* exec_context) const = 0; + + virtual Result> MakeBodyMaskFromSparseBitmap( + const std::shared_ptr& bitmap, ExecContext* exec_context) const = 0; + + virtual Result> MakeBodyMaskFromDenseBitmap( + const std::shared_ptr& bitmap, ExecContext* exec_context) const = 0; + + private: + Result> MakeBodyMaskFromSparseDatum( + const Datum& datum, ExecContext* exec_context) const { + DCHECK(datum.type()->id() == Type::BOOL); + if (datum.is_scalar()) { + return MakeBodyMaskFromScalar(datum.scalar_as(), exec_context); + } + if (datum.is_array()) { + return MakeBodyMaskFromSparseBitmap(datum.array_as(), exec_context); + } + DCHECK(datum.is_chunked_array()); + return MakeBodyMaskFromSparseBitmap(datum.chunked_array(), exec_context); + } + + Result> MakeBodyMaskFromDenseDatum( + const Datum& datum, ExecContext* exec_context) const { + DCHECK(datum.type()->id() == Type::BOOL); + if (datum.is_scalar()) { + return MakeBodyMaskFromScalar(datum.scalar_as(), exec_context); + } + if (datum.is_array()) { + return MakeBodyMaskFromDenseBitmap(datum.array_as(), exec_context); + } + DCHECK(datum.is_chunked_array()); + return MakeBodyMaskFromDenseBitmap(datum.chunked_array(), exec_context); + } + + friend struct NestedBodyMask; +}; + +struct BodyMask : public std::enable_shared_from_this { + virtual ~BodyMask() = default; + + virtual bool empty() const = 0; + + virtual Result ApplySparse(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const = 0; + + virtual Result ApplyDense(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const = 0; + + virtual Result> GetSelectionVector() const = 0; + + virtual Result> NextBranchMask() const = 0; +}; + +struct AllPassBranchMask : public BranchMask { + explicit AllPassBranchMask(int64_t length) : length_(length) {} + + bool empty() const override { return false; } + + protected: + Result ApplySparse(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const override { + DCHECK_EQ(input.length, length_); + auto input_with_sel_vec = input; + input_with_sel_vec.selection_vector = nullptr; + return ExecuteScalarExpression(expr, input_with_sel_vec, exec_context); + } + + Result ApplyDense(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const override { + DCHECK_EQ(input.length, length_); + auto input_with_sel_vec = input; + input_with_sel_vec.selection_vector = nullptr; + return ExecuteScalarExpression(expr, input_with_sel_vec, exec_context); + } + + Result> GetSelectionVector() const override { + return nullptr; + } + + Result> MakeBodyMaskFromScalar( + const BooleanScalar& scalar, ExecContext* exec_context) const override; + + Result> MakeBodyMaskFromSparseBitmap( + const std::shared_ptr& bitmap, + ExecContext* exec_context) const override; + + Result> MakeBodyMaskFromDenseBitmap( + const std::shared_ptr& bitmap, + ExecContext* exec_context) const override; + + Result> MakeBodyMaskFromSparseBitmap( + const std::shared_ptr& bitmap, + ExecContext* exec_context) const override; + + Result> MakeBodyMaskFromDenseBitmap( + const std::shared_ptr& bitmap, + ExecContext* exec_context) const override; + + private: + int64_t length_; +}; + +struct AllFailBranchMask : public BranchMask { + AllFailBranchMask() = default; + + bool empty() const override { return true; } + + protected: + Result ApplySparse(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const override { + DCHECK(false); + return Status::Invalid("AllFailBranchMask::ApplySparse should not be called"); + } + + Result ApplyDense(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const override { + DCHECK(false); + return Status::Invalid("AllFailBranchMask::ApplyDense should not be called"); + } + + Result> GetSelectionVector() const override { + DCHECK(false); + return Status::Invalid("AllFailBranchMask::GetSelectionVector should not be called"); + } + + Result> MakeBodyMaskFromScalar( + const BooleanScalar& scalar, ExecContext* exec_context) const override { + DCHECK(false); + return Status::Invalid("AllFailBranchMask::MakeBodyMask should not be called"); + } + + Result> MakeBodyMaskFromSparseBitmap( + const std::shared_ptr& bitmap, + ExecContext* exec_context) const override { + DCHECK(false); + return Status::Invalid( + "AllFailBranchMask::MakeBodyMaskFromSparseBitmap should not be called"); + } + + Result> MakeBodyMaskFromDenseBitmap( + const std::shared_ptr& bitmap, + ExecContext* exec_context) const override { + DCHECK(false); + return Status::Invalid( + "AllFailBranchMask::MakeBodyMaskFromDenseBitmap should not be called"); + } + + Result> MakeBodyMaskFromSparseBitmap( + const std::shared_ptr& bitmap, + ExecContext* exec_context) const override { + DCHECK(false); + return Status::Invalid( + "AllFailBranchMask::MakeBodyMaskFromSparseBitmap should not be called"); + } + + Result> MakeBodyMaskFromDenseBitmap( + const std::shared_ptr& bitmap, + ExecContext* exec_context) const override { + DCHECK(false); + return Status::Invalid( + "AllFailBranchMask::MakeBodyMaskFromDenseBitmap should not be called"); + } +}; + +struct NestedBodyMask : public BodyMask { + explicit NestedBodyMask(std::shared_ptr branch_mask) + : branch_mask_(std::move(branch_mask)) {} + + protected: + Result DelegateApplySparse(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const { + return branch_mask_->ApplySparse(expr, input, exec_context); + } + + Result DelegateApplyDense(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const { + return branch_mask_->ApplyDense(expr, input, exec_context); + } + + Result> DelegateGetSelectionVector() const { + return branch_mask_->GetSelectionVector(); + } + + protected: + std::shared_ptr branch_mask_; +}; + +struct AllNullBodyMask : public NestedBodyMask { + using NestedBodyMask::NestedBodyMask; + + bool empty() const override { return true; } + + Result ApplySparse(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const override { + DCHECK(false); + return Status::Invalid("AllNullBodyMask::ApplySparse should not be called"); + } + + Result ApplyDense(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const override { + DCHECK(false); + return Status::Invalid("AllNullBodyMask::ApplyDense should not be called"); + } + + Result> GetSelectionVector() const override { + DCHECK(false); + return Status::Invalid("AllNullBodyMask::GetSelectionVector should not be called"); + } + + Result> NextBranchMask() const override { + return std::make_shared(); + } +}; + +struct AllPassBodyMask : public NestedBodyMask { + using NestedBodyMask::NestedBodyMask; + + bool empty() const override { return false; } + + Result ApplySparse(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const override { + return DelegateApplySparse(expr, input, exec_context); + } + + Result ApplyDense(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const override { + return DelegateApplyDense(expr, input, exec_context); + } + + Result> GetSelectionVector() const override { + return DelegateGetSelectionVector(); + } + + Result> NextBranchMask() const override { + return std::make_shared(); + } +}; + +struct AllFailBodyMask : public NestedBodyMask { + using NestedBodyMask::NestedBodyMask; + + bool empty() const override { return true; } + + Result ApplySparse(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const override { + DCHECK(false); + return Status::Invalid("AllFailBodyMask::ApplySparse should not be called"); + } + + Result ApplyDense(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const override { + DCHECK(false); + return Status::Invalid("AllFailBodyMask::ApplyDense should not be called"); + } + + Result> GetSelectionVector() const override { + DCHECK(false); + return Status::Invalid("AllFailBodyMask::GetSelectionVector should not be called"); + } + + Result> NextBranchMask() const override { + return branch_mask_; + } +}; + +Result TakeBySelectionVector(const ExecBatch& input, + const Datum& selection_vector, + ExecContext* exec_context) { + std::vector values(input.num_values()); + for (int i = 0; i < input.num_values(); ++i) { + ARROW_ASSIGN_OR_RAISE( + values[i], Take(input[i], selection_vector, TakeOptions{/*boundcheck=*/false}, + exec_context)); + } + return ExecBatch::Make(std::move(values), selection_vector.length()); +} + +struct ConditionalBranchMask : public BranchMask { + ConditionalBranchMask(std::shared_ptr selection_vector, int64_t length) + : selection_vector_(std::move(selection_vector)), length_(length) {} + + bool empty() const override { return selection_vector_->length() == 0; } + + protected: + Result ApplySparse(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const override { + auto sparse_input = input; + sparse_input.selection_vector = selection_vector_; + return ExecuteScalarExpression(expr, sparse_input, exec_context); + } + + Result ApplyDense(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const override { + ARROW_ASSIGN_OR_RAISE( + auto dense_input, + TakeBySelectionVector(input, *selection_vector_->data(), exec_context)); + return ExecuteScalarExpression(expr, dense_input, exec_context); + } + + Result> GetSelectionVector() const override { + return selection_vector_; + } + + Result> MakeBodyMaskFromScalar( + const BooleanScalar& scalar, ExecContext* exec_context) const override; + + Result> MakeBodyMaskFromSparseBitmap( + const std::shared_ptr& bitmap, + ExecContext* exec_context) const override; + + Result> MakeBodyMaskFromDenseBitmap( + const std::shared_ptr& bitmap, + ExecContext* exec_context) const override; + + Result> MakeBodyMaskFromSparseBitmap( + const std::shared_ptr& bitmap, + ExecContext* exec_context) const override; + + Result> MakeBodyMaskFromDenseBitmap( + const std::shared_ptr& bitmap, + ExecContext* exec_context) const override; + + protected: + std::shared_ptr selection_vector_ = nullptr; + int64_t length_ = 0; +}; + +struct ConditionalBodyMask : public BodyMask { + ConditionalBodyMask(std::shared_ptr body, + std::shared_ptr rest, int64_t length) + : body_(std::move(body)), rest_(std::move(rest)), length_(length) {} + + bool empty() const override { return body_->length() == 0; } + + Result ApplySparse(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const override { + auto sparse_input = input; + sparse_input.selection_vector = body_; + return ExecuteScalarExpression(expr, sparse_input, exec_context); + } + + Result ApplyDense(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const override { + ARROW_ASSIGN_OR_RAISE(auto dense_input, + TakeBySelectionVector(input, *body_->data(), exec_context)); + return ExecuteScalarExpression(expr, dense_input, exec_context); + } + + Result> GetSelectionVector() const override { + return body_; + } + + Result> NextBranchMask() const override { + if (!rest_) { + return std::make_shared(); + } + return std::make_shared(rest_, length_); + } + + private: + std::shared_ptr body_; + std::shared_ptr rest_; + int64_t length_; +}; + +Result> BodyMaskFromScalar( + const BooleanScalar& scalar, std::shared_ptr branch_mask, + ExecContext* exec_context) { + if (!scalar.is_valid) { + return std::make_shared(std::move(branch_mask)); + } else if (scalar.value) { + return std::make_shared(std::move(branch_mask)); + } else { + return std::make_shared(std::move(branch_mask)); + } +} + +Result> AllPassBranchMask::MakeBodyMaskFromScalar( + const BooleanScalar& scalar, ExecContext* exec_context) const { + return BodyMaskFromScalar(scalar, shared_from_this(), exec_context); +} + +Result> AllPassBranchMask::MakeBodyMaskFromSparseBitmap( + const std::shared_ptr& bitmap, ExecContext* exec_context) const { + DCHECK_EQ(bitmap->length(), length_); + + Int32Builder body_builder(exec_context->memory_pool()); + Int32Builder rest_builder(exec_context->memory_pool()); + RETURN_NOT_OK(body_builder.Reserve(length_)); + RETURN_NOT_OK(rest_builder.Reserve(length_)); + + ArraySpan span(*bitmap->data()); + int32_t i = 0; + VisitArraySpanInline( + span, + [&](bool mask) { + if (mask) { + body_builder.UnsafeAppend(i); + } else { + rest_builder.UnsafeAppend(i); + } + ++i; + }, + [&]() { ++i; }); + + ARROW_ASSIGN_OR_RAISE(auto body_arr, body_builder.Finish()); + ARROW_ASSIGN_OR_RAISE(auto rest_arr, rest_builder.Finish()); + auto body = std::make_shared(body_arr->data()); + auto rest = std::make_shared(rest_arr->data()); + return std::make_shared(std::move(body), std::move(rest), length_); +} + +Result> AllPassBranchMask::MakeBodyMaskFromDenseBitmap( + const std::shared_ptr& bitmap, ExecContext* exec_context) const { + return MakeBodyMaskFromSparseBitmap(bitmap, exec_context); +} + +Result> AllPassBranchMask::MakeBodyMaskFromSparseBitmap( + const std::shared_ptr& bitmap, ExecContext* exec_context) const { + DCHECK_EQ(bitmap->length(), length_); + + Int32Builder body_builder(exec_context->memory_pool()); + Int32Builder rest_builder(exec_context->memory_pool()); + RETURN_NOT_OK(body_builder.Reserve(length_)); + RETURN_NOT_OK(rest_builder.Reserve(length_)); + + int32_t i = 0; + for (const auto& chunk : bitmap->chunks()) { + DCHECK_EQ(chunk->type()->id(), Type::BOOL); + ArraySpan span(*chunk->data()); + VisitArraySpanInline( + span, + [&](bool mask) { + if (mask) { + body_builder.UnsafeAppend(i); + } else { + rest_builder.UnsafeAppend(i); + } + ++i; + }, + [&]() { ++i; }); + } + + ARROW_ASSIGN_OR_RAISE(auto body_arr, body_builder.Finish()); + ARROW_ASSIGN_OR_RAISE(auto rest_arr, rest_builder.Finish()); + auto body = std::make_shared(body_arr->data()); + auto rest = std::make_shared(rest_arr->data()); + return std::make_shared(std::move(body), std::move(rest), length_); +} + +Result> AllPassBranchMask::MakeBodyMaskFromDenseBitmap( + const std::shared_ptr& bitmap, ExecContext* exec_context) const { + return MakeBodyMaskFromSparseBitmap(bitmap, exec_context); +} + +Result> ConditionalBranchMask::MakeBodyMaskFromScalar( + const BooleanScalar& scalar, ExecContext* exec_context) const { + return BodyMaskFromScalar(scalar, shared_from_this(), exec_context); +} + +Result> +ConditionalBranchMask::MakeBodyMaskFromSparseBitmap( + const std::shared_ptr& bitmap, ExecContext* exec_context) const { + DCHECK_EQ(bitmap->length(), length_); + + Int32Builder body_builder(exec_context->memory_pool()); + Int32Builder rest_builder(exec_context->memory_pool()); + RETURN_NOT_OK(body_builder.Reserve(length_)); + RETURN_NOT_OK(rest_builder.Reserve(length_)); + + for (int64_t i = 0; i < selection_vector_->length(); ++i) { + auto index = selection_vector_->indices()[i]; + if (!bitmap->IsNull(index)) { + if (bitmap->Value(index)) { + body_builder.UnsafeAppend(index); + } else { + rest_builder.UnsafeAppend(index); + } + } + } + + ARROW_ASSIGN_OR_RAISE(auto body_arr, body_builder.Finish()); + ARROW_ASSIGN_OR_RAISE(auto rest_arr, rest_builder.Finish()); + auto body = std::make_shared(body_arr->data()); + auto rest = std::make_shared(rest_arr->data()); + return std::make_shared(std::move(body), std::move(rest), length_); +} + +Result> +ConditionalBranchMask::MakeBodyMaskFromDenseBitmap( + const std::shared_ptr& bitmap, ExecContext* exec_context) const { + DCHECK_EQ(bitmap->length(), selection_vector_->length()); + + Int32Builder body_builder(exec_context->memory_pool()); + Int32Builder rest_builder(exec_context->memory_pool()); + RETURN_NOT_OK(body_builder.Reserve(bitmap->true_count())); + RETURN_NOT_OK(rest_builder.Reserve(bitmap->false_count())); + + for (int64_t i = 0; i < selection_vector_->length(); ++i) { + if (!bitmap->IsNull(i)) { + if (bitmap->Value(i)) { + body_builder.UnsafeAppend(selection_vector_->indices()[i]); + } else { + rest_builder.UnsafeAppend(selection_vector_->indices()[i]); + } + } + } + + ARROW_ASSIGN_OR_RAISE(auto body_arr, body_builder.Finish()); + ARROW_ASSIGN_OR_RAISE(auto rest_arr, rest_builder.Finish()); + auto body = std::make_shared(body_arr->data()); + auto rest = std::make_shared(rest_arr->data()); + return std::make_shared(std::move(body), std::move(rest), length_); +} + +Result> +ConditionalBranchMask::MakeBodyMaskFromSparseBitmap( + const std::shared_ptr& bitmap, ExecContext* exec_context) const { + DCHECK_EQ(bitmap->length(), length_); + + std::vector boolean_arrays(bitmap->num_chunks()); + std::transform(bitmap->chunks().begin(), bitmap->chunks().end(), boolean_arrays.begin(), + [](const auto& chunk) { + DCHECK_EQ(chunk->type()->id(), Type::BOOL); + return checked_cast(chunk.get()); + }); + + Int32Builder body_builder(exec_context->memory_pool()); + Int32Builder rest_builder(exec_context->memory_pool()); + RETURN_NOT_OK(body_builder.Reserve(length_)); + RETURN_NOT_OK(rest_builder.Reserve(length_)); + + ChunkResolver resolver(bitmap->chunks()); + ChunkLocation location; + for (int64_t i = 0; i < selection_vector_->length(); ++i) { + auto index = selection_vector_->indices()[i]; + location = resolver.ResolveWithHint(index, location); + if (boolean_arrays[location.chunk_index]->IsValid(location.index_in_chunk)) { + if (boolean_arrays[location.chunk_index]->Value(location.index_in_chunk)) { + body_builder.UnsafeAppend(index); + } else { + rest_builder.UnsafeAppend(index); + } + } + } + + ARROW_ASSIGN_OR_RAISE(auto body_arr, body_builder.Finish()); + ARROW_ASSIGN_OR_RAISE(auto rest_arr, rest_builder.Finish()); + auto body = std::make_shared(body_arr->data()); + auto rest = std::make_shared(rest_arr->data()); + return std::make_shared(std::move(body), std::move(rest), length_); +} + +Result> +ConditionalBranchMask::MakeBodyMaskFromDenseBitmap( + const std::shared_ptr& bitmap, ExecContext* exec_context) const { + DCHECK_EQ(bitmap->length(), selection_vector_->length()); + + std::vector boolean_arrays(bitmap->num_chunks()); + std::transform(bitmap->chunks().begin(), bitmap->chunks().end(), boolean_arrays.begin(), + [](const auto& chunk) { + DCHECK_EQ(chunk->type()->id(), Type::BOOL); + return checked_cast(chunk.get()); + }); + + Int32Builder body_builder(exec_context->memory_pool()); + Int32Builder rest_builder(exec_context->memory_pool()); + RETURN_NOT_OK(body_builder.Reserve(selection_vector_->length())); + RETURN_NOT_OK(body_builder.Reserve(selection_vector_->length())); + + ChunkResolver resolver(bitmap->chunks()); + ChunkLocation location; + for (int64_t i = 0; i < selection_vector_->length(); ++i) { + location = resolver.ResolveWithHint(i, location); + if (boolean_arrays[location.chunk_index]->IsValid(location.index_in_chunk)) { + if (boolean_arrays[location.chunk_index]->Value(location.index_in_chunk)) { + body_builder.UnsafeAppend(selection_vector_->indices()[i]); + } else { + rest_builder.UnsafeAppend(selection_vector_->indices()[i]); + } + } + } + + ARROW_ASSIGN_OR_RAISE(auto body_arr, body_builder.Finish()); + ARROW_ASSIGN_OR_RAISE(auto rest_arr, rest_builder.Finish()); + auto body = std::make_shared(body_arr->data()); + auto rest = std::make_shared(rest_arr->data()); + return std::make_shared(std::move(body), std::move(rest), length_); +} + +struct Branch { + Expression cond; + Expression body; +}; + +template +struct ConditionalExecutor { + ConditionalExecutor(const std::vector& branches, const TypeHolder& result_type) + : branches(branches), result_type(result_type) {} + + ARROW_DISALLOW_COPY_AND_ASSIGN(ConditionalExecutor); + ARROW_DEFAULT_MOVE_AND_ASSIGN(ConditionalExecutor); + + Result Execute(const ExecBatch& input, ExecContext* exec_context) const&& { + DCHECK(!branches.empty()); + + BranchResults results; + results.Reserve(branches.size()); + ARROW_ASSIGN_OR_RAISE(auto branch_mask, InitBranchMask(input, exec_context)); + for (const auto& branch : branches) { + if (branch_mask->empty()) { + break; + } + ARROW_ASSIGN_OR_RAISE(auto body_mask, + ApplyCond(branch_mask, branch.cond, input, exec_context)); + if (body_mask->empty()) { + ARROW_ASSIGN_OR_RAISE(branch_mask, body_mask->NextBranchMask()); + continue; + } + ARROW_ASSIGN_OR_RAISE(auto body_result, + static_cast(this)->ApplyBody( + body_mask, branch.body, input, exec_context)); + DCHECK(body_result.type()->Equals(*result_type)); + ARROW_ASSIGN_OR_RAISE(auto selection_vector, body_mask->GetSelectionVector()); + results.Emplace(std::move(body_result), std::move(selection_vector)); + ARROW_ASSIGN_OR_RAISE(branch_mask, body_mask->NextBranchMask()); + } + return static_cast(this)->MultiplexBranchResults(input, results, + exec_context); + } + + protected: + struct BranchResults { + void Reserve(int64_t size) { + body_results_.reserve(size); + selection_vectors_.reserve(size); + } + + void Emplace(Datum body_result, std::shared_ptr selection_vector) { + body_results_.emplace_back(std::move(body_result)); + selection_vectors_.emplace_back(std::move(selection_vector)); + } + + bool empty() const { return body_results_.empty(); } + + size_t size() const { return body_results_.size(); } + + const std::vector& body_results() const { return body_results_; } + + const std::vector>& selection_vectors() const { + return selection_vectors_; + } + + private: + std::vector body_results_; + std::vector> selection_vectors_; + }; + + private: + Result> InitBranchMask( + const ExecBatch& input, ExecContext* exec_context) const { + if (input.selection_vector) { + return std::make_shared(input.selection_vector, + input.length); + } + return std::make_shared(input.length); + } + + Result> ApplyCond( + const std::shared_ptr& branch_mask, const Expression& cond, + const ExecBatch& input, ExecContext* exec_context) const { + if (cond.selection_vector_aware()) { + return branch_mask->ApplyCondSparse(cond, input, exec_context); + } + return branch_mask->ApplyCondDense(cond, input, exec_context); + } + + protected: + const std::vector& branches; + const TypeHolder& result_type; +}; + +struct SparseConditionalExecutor : public ConditionalExecutor { + using ConditionalExecutor::ConditionalExecutor; + + Result ApplyBody(const std::shared_ptr& body_mask, + const Expression& body, const ExecBatch& input, + ExecContext* exec_context) const { + return body_mask->ApplySparse(body, input, exec_context); + } + + Result MultiplexBranchResults(const ExecBatch& input, + const BranchResults& results, + ExecContext* exec_context) const { + if (results.empty()) { + return MakeArrayOfNull(result_type.GetSharedPtr(), input.length, + exec_context->memory_pool()); + } + + if (results.size() == 1) { + if (const auto& result = results.body_results()[0]; + results.selection_vectors()[0] == nullptr || + results.selection_vectors()[0]->length() == + (input.selection_vector ? input.selection_vector->length() + : input.length)) { + return result; + } + } + + std::vector choose_args; + choose_args.reserve(results.size() + 1); + ARROW_ASSIGN_OR_RAISE(auto indices, ChooseIndices(results.selection_vectors(), + input.length, exec_context)); + choose_args.emplace_back(std::move(indices)); + choose_args.insert(choose_args.end(), results.body_results().begin(), + results.body_results().end()); + return CallFunction("choose", choose_args, exec_context); + } + + private: + Result ChooseIndices( + const std::vector>& selection_vectors, + int64_t length, ExecContext* exec_context) const { + const int64_t validity_bytes = bit_util::BytesForBits(length); + ARROW_ASSIGN_OR_RAISE( + std::shared_ptr validity_buf, + AllocateResizableBuffer(validity_bytes, exec_context->memory_pool())); + auto validity_data = validity_buf->mutable_data_as(); + std::memset(validity_data, 0, validity_bytes); + + ARROW_ASSIGN_OR_RAISE( + std::shared_ptr indices_buf, + AllocateResizableBuffer(length * sizeof(int32_t), exec_context->memory_pool())); + auto indices_data = indices_buf->mutable_data_as(); + for (int32_t index = 0; index < static_cast(selection_vectors.size()); + ++index) { + DCHECK_NE(selection_vectors[index], nullptr); + DCHECK_GT(selection_vectors[index]->length(), 0); + auto row_ids = selection_vectors[index]->indices(); + for (int64_t i = 0; i < selection_vectors[index]->length(); ++i) { + const int32_t row_id = row_ids[i]; + DCHECK_EQ(bit_util::GetBit(validity_data, row_id), false); + bit_util::SetBitTo(validity_data, row_id, true); + indices_data[row_id] = index; + } + } + + return ArrayData::Make(int32(), length, + {std::move(validity_buf), std::move(indices_buf)}); + } +}; + +struct DenseConditionalExecutor : public ConditionalExecutor { + using ConditionalExecutor::ConditionalExecutor; + + Result ApplyBody(const std::shared_ptr& body_mask, + const Expression& body, const ExecBatch& input, + ExecContext* exec_context) const { + return body_mask->ApplyDense(body, input, exec_context); + } + + Result MultiplexBranchResults(const ExecBatch& input, + const BranchResults& results, + ExecContext* exec_context) const { + if (results.empty()) { + return MakeArrayOfNull(result_type.GetSharedPtr(), input.length, + exec_context->memory_pool()); + } + + if (results.size() == 1) { + if (!results.selection_vectors()[0] || + results.selection_vectors()[0]->length() == input.length) { + return results.body_results()[0]; + } + } + + ARROW_ASSIGN_OR_RAISE( + auto body_results, + ToChunkedArray(results, + [&](const Datum& value, + const std::shared_ptr& selection_vector, + ArrayVector& chunks) -> Status { + DCHECK_NE(selection_vector, nullptr); + DCHECK_GT(selection_vector->length(), 0); + DCHECK(value.is_scalar() || value.is_arraylike()); + if (value.is_scalar()) { + ARROW_ASSIGN_OR_RAISE( + auto arr, MakeArrayFromScalar( + *value.scalar(), selection_vector->length(), + exec_context->memory_pool())); + chunks.push_back(std::move(arr)); + } else if (value.is_array() && value.length() > 0) { + chunks.push_back(value.make_array()); + } else { + DCHECK(value.is_chunked_array()); + for (const auto& chunk : value.chunked_array()->chunks()) { + if (chunk->length() > 0) { + chunks.push_back(chunk); + } + } + } + return Status::OK(); + })); + ARROW_ASSIGN_OR_RAISE( + auto indices, + ToChunkedArray( + results, + [](const Datum&, const std::shared_ptr& selection_vector, + ArrayVector& chunks) -> Status { + DCHECK_NE(selection_vector, nullptr); + DCHECK_GT(selection_vector->length(), 0); + chunks.push_back(MakeArray(selection_vector->data())); + return Status::OK(); + })); + ARROW_ASSIGN_OR_RAISE( + auto result, + Scatter(body_results, indices, ScatterOptions{/*max_index=*/input.length - 1})); + DCHECK(result.is_arraylike()); + if (result.is_chunked_array() && result.chunked_array()->num_chunks() == 1) { + return result.chunked_array()->chunk(0); + } else { + return result; + } + } + + private: + template + Result> ToChunkedArray(const BranchResults& results, + ChunkFunc&& chunk_func) const { + ArrayVector chunks; + chunks.reserve(results.size()); + for (size_t i = 0; i < results.size(); ++i) { + RETURN_NOT_OK( + chunk_func(results.body_results()[i], results.selection_vectors()[i], chunks)); + } + return std::make_shared(std::move(chunks)); + } +}; + +class IfElseSpecialExec : public SpecialExec { + public: + explicit IfElseSpecialExec(std::shared_ptr function, const Kernel* kernel) + : function(std::move(function)), kernel(kernel), branches(2) {} + + Result Bind(const std::vector& arguments, + ExecContext* exec_context) override { + DCHECK(std::all_of(arguments.begin(), arguments.end(), + [](const Expression& argument) { return argument.IsBound(); })); + + DCHECK_EQ(arguments.size(), 3); + const auto& cond = arguments[0]; + const auto& if_true = arguments[1]; + const auto& if_false = arguments[2]; + + { + auto types = GetTypes(arguments); + KernelContext kernel_context(exec_context, kernel); + std::unique_ptr kernel_state; + if (kernel->init) { + const FunctionOptions* options = function->default_options(); + ARROW_ASSIGN_OR_RAISE(kernel_state, + kernel->init(&kernel_context, {kernel, types, options})); + kernel_context.SetState(kernel_state.get()); + } + ARROW_ASSIGN_OR_RAISE( + type, kernel->signature->out_type().Resolve(&kernel_context, types)); + } + + DCHECK_EQ(cond.type()->id(), Type::BOOL); + DCHECK_EQ(type, *if_true.type()); + DCHECK_EQ(type, *if_false.type()); + + all_bodies_selection_vector_aware = + if_true.selection_vector_aware() && if_false.selection_vector_aware(); + + branches[0] = {cond, if_true}; + branches[1] = {literal(true), if_false}; + + return type; + } + + Result Execute(const ExecBatch& input, + ExecContext* exec_context) const override { + if (all_bodies_selection_vector_aware) { + return SparseConditionalExecutor(branches, type).Execute(input, exec_context); + } else { + return DenseConditionalExecutor(branches, type).Execute(input, exec_context); + } + } + + private: + // For Bind. + std::shared_ptr function; + const Kernel* kernel; + + // Post-bind, for Execute. + TypeHolder type; + bool all_bodies_selection_vector_aware; + std::vector branches; +}; + +class IfElseSpecialForm : public SpecialForm { + public: + IfElseSpecialForm() + : SpecialForm(/*name=*/"if_else", /*selection_vector_aware=*/true) {} + + Result> DispatchExact( + const std::vector& types, ExecContext* exec_context) const override { + ARROW_ASSIGN_OR_RAISE(auto function, + exec_context->func_registry()->GetFunction(name)); + ARROW_ASSIGN_OR_RAISE(auto kernel, function->DispatchExact(types)); + return std::make_unique(std::move(function), kernel); + } + + Result> DispatchBest( + std::vector* types, ExecContext* exec_context) const override { + ARROW_ASSIGN_OR_RAISE(auto function, + exec_context->func_registry()->GetFunction(name)); + ARROW_ASSIGN_OR_RAISE(auto kernel, function->DispatchBest(types)); + return std::make_unique(std::move(function), kernel); + } +}; + +} // namespace + +std::shared_ptr GetIfElseSpecialForm() { + static auto instance = std::make_shared(); + return instance; +} + +} // namespace arrow::compute diff --git a/cpp/src/arrow/compute/special_form.h b/cpp/src/arrow/compute/special_form.h new file mode 100644 index 000000000000..8273e8405154 --- /dev/null +++ b/cpp/src/arrow/compute/special_form.h @@ -0,0 +1,53 @@ +// 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/expression.h" + +namespace arrow::compute { + +class ARROW_EXPORT SpecialExec { + public: + virtual ~SpecialExec() = default; + + virtual Result Bind(const std::vector& arguments, + ExecContext* exec_context) = 0; + + virtual Result Execute(const ExecBatch& input, + ExecContext* exec_context) const = 0; +}; + +class ARROW_EXPORT SpecialForm { + public: + explicit SpecialForm(std::string name, bool selection_vector_aware = false) + : name(std::move(name)), selection_vector_aware(selection_vector_aware) {} + + virtual ~SpecialForm() = default; + + virtual Result> DispatchExact( + const std::vector& types, ExecContext* exec_context) const = 0; + + virtual Result> DispatchBest( + std::vector* types, ExecContext* exec_context) const = 0; + + public: + const std::string name; + const bool selection_vector_aware = false; +}; + +std::shared_ptr GetIfElseSpecialForm(); + +} // namespace arrow::compute diff --git a/cpp/src/arrow/compute/special_form_benchmark.cc b/cpp/src/arrow/compute/special_form_benchmark.cc new file mode 100644 index 000000000000..234ba75afa33 --- /dev/null +++ b/cpp/src/arrow/compute/special_form_benchmark.cc @@ -0,0 +1,283 @@ +// 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 "arrow/compute/exec.h" +#include "arrow/compute/expression.h" +#include "arrow/compute/function.h" +#include "arrow/compute/function_internal.h" +#include "arrow/compute/kernels/codegen_internal.h" +#include "arrow/compute/registry.h" +#include "arrow/testing/generator.h" +#include "arrow/testing/random.h" +#include "arrow/util/logging.h" + +namespace arrow { + +namespace compute { + +namespace { + +struct PayloadOptions : public FunctionOptions { + explicit PayloadOptions(int64_t load = 0); + static constexpr char const kTypeName[] = "PayloadOptions"; + static PayloadOptions Defaults() { return PayloadOptions(); } + int64_t load = 0; +}; + +static auto kPayloadOptionsType = internal::GetFunctionOptionsType( + arrow::internal::DataMember("load", &PayloadOptions::load)); + +PayloadOptions::PayloadOptions(int64_t load) + : FunctionOptions(kPayloadOptionsType), load(load) {} + +const PayloadOptions* GetDefaultPayloadOptions() { + static const auto kDefaultPayloadOptions = PayloadOptions::Defaults(); + return &kDefaultPayloadOptions; +} + +using PayloadState = internal::OptionsWrapper; + +Status PayloadExec(KernelContext* ctx, const ExecSpan& span, ExecResult* out) { + ARROW_CHECK_EQ(span.num_values(), 1); + const auto& arg = span[0]; + ARROW_CHECK(arg.is_array()); + + int64_t load = PayloadState::Get(ctx).load; + int64_t load_length = + span.selection_vector.indices() ? span.selection_vector.length() : arg.length(); + for (int64_t i = 0; i < load_length; ++i) { + volatile int64_t j = load; + while (j-- > 0) { + } + } + *out->array_data_mutable() = *arg.array.ToArrayData(); + return Status::OK(); +} + +Status RegisterAuxilaryFunctions() { + auto registry = GetFunctionRegistry(); + + { + if (registry->CanAddFunctionOptionsType(kPayloadOptionsType).ok()) { + RETURN_NOT_OK(registry->AddFunctionOptionsType(kPayloadOptionsType)); + } + } + { + auto register_payload_func = [&](const std::string& name, + bool sv_awareness) -> Status { + auto func = std::make_shared( + name, Arity::Unary(), FunctionDoc::Empty(), GetDefaultPayloadOptions()); + + ScalarKernel kernel({InputType::Any()}, internal::FirstType, PayloadExec, + PayloadState::Init); + kernel.selection_vector_aware = sv_awareness; + kernel.can_write_into_slices = false; + kernel.null_handling = NullHandling::COMPUTED_NO_PREALLOCATE; + kernel.mem_allocation = MemAllocation::NO_PREALLOCATE; + RETURN_NOT_OK(func->AddKernel(kernel)); + if (registry->CanAddFunction(func, /*allow_overwrite=*/false).ok()) { + RETURN_NOT_OK(registry->AddFunction(std::move(func))); + } + return Status::OK(); + }; + + RETURN_NOT_OK(register_payload_func("payload_sv_aware", true)); + RETURN_NOT_OK(register_payload_func("payload_sv_unaware", false)); + } + + return Status::OK(); +} + +Expression if_else_regular(Expression cond, Expression if_true, Expression if_false) { + return call("if_else", {std::move(cond), std::move(if_true), std::move(if_false)}); +} + +Expression sv_suppress(Expression arg) { + return call("payload_sv_unaware", {std::move(arg)}); +} + +Expression heavy(Expression arg) { + return call("payload_sv_aware", {std::move(arg)}, + std::make_shared(/*load=*/1024)); +} + +Expression sv_unaware_if_else_regular(Expression cond, Expression if_true, + Expression if_false) { + return if_else_regular(std::move(cond), sv_suppress(std::move(if_true)), + sv_suppress(std::move(if_false))); +} + +Expression sv_unaware_if_else_special(Expression cond, Expression if_true, + Expression if_false) { + return if_else_special(std::move(cond), sv_suppress(std::move(if_true)), + sv_suppress(std::move(if_false))); +} + +auto kBooleanNull = literal(MakeNullScalar(boolean())); + +void BenchmarkIfElse( + benchmark::State& state, + std::function if_else_func, + Expression cond, Expression if_true, Expression if_false, + const std::shared_ptr& schema, const ExecBatch& batch) { + ARROW_CHECK_OK(RegisterAuxilaryFunctions()); + + auto if_else = if_else_func(std::move(cond), std::move(if_true), std::move(if_false)); + auto bound = if_else.Bind(*schema).ValueOrDie(); + for (auto _ : state) { + ARROW_CHECK_OK(ExecuteScalarExpression(bound, batch).status()); + } + + state.SetItemsProcessed(batch.length * state.iterations()); +} + +} // namespace + +#ifdef BM +# error("BM is defined") +#else +# define BENCHMARK_IF_ELSE_WITH_BASELINE_AND_SV_SUPPRESS(BM, name, ...) \ + BM(name##_regular, if_else_regular, ##__VA_ARGS__); \ + BM(name##_special, if_else_special, ##__VA_ARGS__); \ + BM(name##_regular_sv_unaware, sv_unaware_if_else_regular, ##__VA_ARGS__); \ + BM(name##_special_sv_unaware, sv_unaware_if_else_special, ##__VA_ARGS__); +#endif + +#define BENCHMARK_IF_ELSE(BM, name, if_else, arg_names, args, ...) \ + BENCHMARK_CAPTURE(BM, name, if_else, ##__VA_ARGS__) \ + ->ArgNames(arg_names) \ + ->ArgsProduct(args) + +template +static void BM_IfElseTrivialCond( + benchmark::State& state, + std::function if_else_func, + Expression cond, Args&&...) { + const int64_t num_rows = state.range(0); + + auto schema = arrow::schema({field("i1", int32()), field("i2", int32())}); + + auto i1 = ConstantArrayGenerator::Int32(num_rows, 1); + auto i2 = ConstantArrayGenerator::Int32(num_rows, 0); + ExecBatch batch{std::vector{std::move(i1), std::move(i2)}, num_rows}; + + BenchmarkIfElse(state, std::move(if_else_func), std::move(cond), field_ref("i1"), + field_ref("i2"), schema, batch); +} + +const std::vector kNumRowsArgNames{"num_rows"}; +const std::vector kNumRowsArg = benchmark::CreateRange(1, 64 * 1024, 32); + +#define BM(name, if_else, ...) \ + BENCHMARK_IF_ELSE(BM_IfElseTrivialCond, name, if_else, kNumRowsArgNames, \ + {kNumRowsArg}, ##__VA_ARGS__) +BENCHMARK_IF_ELSE_WITH_BASELINE_AND_SV_SUPPRESS(BM, literal_null, kBooleanNull) +BENCHMARK_IF_ELSE_WITH_BASELINE_AND_SV_SUPPRESS(BM, literal_true, literal(true)) +BENCHMARK_IF_ELSE_WITH_BASELINE_AND_SV_SUPPRESS(BM, literal_false, literal(false)) +#undef BM + +namespace { + +void BenchmarkIfElseWithCondArray( + benchmark::State& state, + std::function if_else_func, + int64_t num_rows, double true_probability, double null_probability) { + random::RandomArrayGenerator rag(42); + auto b = rag.Boolean(num_rows, true_probability, null_probability); + auto schema = arrow::schema({field("b", boolean())}); + + ExecBatch batch{std::vector{std::move(b)}, num_rows}; + + BenchmarkIfElse(state, std::move(if_else_func), field_ref("b"), literal(1), literal(0), + schema, batch); +} + +} // namespace + +template +static void BM_IfElseNumRows( + benchmark::State& state, + std::function if_else_func, + Args&&...) { + const int64_t num_rows = state.range(0); + + BenchmarkIfElseWithCondArray(state, std::move(if_else_func), num_rows, + /*true_probability=*/0.5, /*null_probability=*/0.0); +} + +#define BM(name, if_else, ...) \ + BENCHMARK_IF_ELSE(BM_IfElseNumRows, name, if_else, kNumRowsArgNames, {kNumRowsArg}) +BENCHMARK_IF_ELSE_WITH_BASELINE_AND_SV_SUPPRESS(BM, num_rows) +#undef BM + +template +static void BM_IfElseNullProbability( + benchmark::State& state, + std::function if_else_func, + Args&&...) { + const int64_t num_rows = state.range(0); + const double null_probability = state.range(1) / 100.0; + + BenchmarkIfElseWithCondArray(state, std::move(if_else_func), num_rows, + /*true_probability=*/0.5, null_probability); +} + +const std::vector kNumRowsAndNullProbabilityArgNames{"num_rows", + "null_probability"}; +const std::vector> kNumRowsAndNullProbabilityArgs{ + {4 * 1024, 64 * 1024}, {0, 50, 90, 100}}; + +#define BM(name, if_else, ...) \ + BENCHMARK_IF_ELSE(BM_IfElseNullProbability, name, if_else, \ + kNumRowsAndNullProbabilityArgNames, kNumRowsAndNullProbabilityArgs) +BENCHMARK_IF_ELSE_WITH_BASELINE_AND_SV_SUPPRESS(BM, null_probability); +#undef BM + +template +void BM_IfElseWithOneHeavySide( + benchmark::State& state, + std::function if_else_func, + Args&&...) { + const double heavy_ratio = state.range(0) / 100.0; + constexpr int64_t num_rows = 65536; + + random::RandomArrayGenerator rag(42); + auto b = + rag.Boolean(num_rows, /*true_probability=*/heavy_ratio, /*null_probability=*/0.0); + auto i = ConstantArrayGenerator::Int32(num_rows, 42); + auto schema = arrow::schema({field("b", boolean()), field("i", int32())}); + + ExecBatch batch{std::vector{std::move(b), std::move(i)}, num_rows}; + + BenchmarkIfElse(state, std::move(if_else_func), field_ref("b"), heavy(field_ref("i")), + literal(0), schema, batch); +} + +const std::vector kHeavyRatioArgNames{"heavy_ratio"}; +const std::vector kHeavyRatioArgs{{0, 25, 50, 75, 100}}; + +#define BM(name, if_else, ...) \ + BENCHMARK_IF_ELSE(BM_IfElseWithOneHeavySide, name, if_else, kHeavyRatioArgNames, \ + {kHeavyRatioArgs}) +BENCHMARK_IF_ELSE_WITH_BASELINE_AND_SV_SUPPRESS(BM, one_heavy_side); +#undef BM + +} // namespace compute + +} // namespace arrow diff --git a/cpp/src/arrow/compute/special_form_test.cc b/cpp/src/arrow/compute/special_form_test.cc new file mode 100644 index 000000000000..efe3f7c08f63 --- /dev/null +++ b/cpp/src/arrow/compute/special_form_test.cc @@ -0,0 +1,1086 @@ +// 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/special_form.h" + +#include + +#include "arrow/compute/exec.h" +#include "arrow/compute/expression.h" +#include "arrow/compute/function.h" +#include "arrow/compute/kernels/codegen_internal.h" +#include "arrow/compute/registry.h" +#include "arrow/testing/gtest_util.h" +#include "arrow/util/logging_internal.h" + +namespace arrow::compute { + +namespace { + +void AssertEqualIgnoreShape(const Datum& expected, const Datum& result) { + if (expected.kind() == result.kind()) { + AssertDatumsEqual(expected, result); + return; + } + if (expected.is_scalar()) { + ASSERT_OK_AND_ASSIGN(auto expected_array, + MakeArrayFromScalar(*expected.scalar(), result.length())); + AssertDatumsEqual(expected_array, result); + return; + } + if (result.is_scalar()) { + ASSERT_OK_AND_ASSIGN(auto result_array, + MakeArrayFromScalar(*result.scalar(), expected.length())); + AssertDatumsEqual(expected, result_array); + return; + } +} + +Result ExecuteExpr(const Expression& expr, const std::shared_ptr& schema, + const ExecBatch& batch, + ExecContext* exec_context = default_exec_context()) { + ARROW_ASSIGN_OR_RAISE(auto bound, expr.Bind(*schema, exec_context)); + return ExecuteScalarExpression(bound, batch, exec_context); +} + +#define AssertExprRaisesWithMessage(expr, schema, batch, ENUM, message, ...) \ + { \ + ASSERT_RAISES_WITH_MESSAGE(ENUM, message, \ + ExecuteExpr(expr, schema, batch, ##__VA_ARGS__)); \ + } + +void AssertExprEqualIgnoreShape(const Expression& expr, + const std::shared_ptr& schema, + const ExecBatch& batch, const Datum& expected, + ExecContext* exec_context = default_exec_context()) { + ASSERT_OK_AND_ASSIGN(auto result, ExecuteExpr(expr, schema, batch, exec_context)); + AssertEqualIgnoreShape(expected, result); +} + +void AssertExprEqualExprsIgnoreShape(const Expression& expr, + const std::vector& exprs, + const std::shared_ptr& schema, + const ExecBatch& batch, + ExecContext* exec_context = default_exec_context()) { + ASSERT_OK_AND_ASSIGN(auto expected, ExecuteExpr(expr, schema, batch, exec_context)); + for (const auto& e : exprs) { + ARROW_SCOPED_TRACE(e.ToString()); + AssertExprEqualIgnoreShape(e, schema, batch, expected, exec_context); + } +} + +auto kBooleanNull = literal(MakeNullScalar(boolean())); +auto kIntNull = literal(MakeNullScalar(int32())); + +Expression if_else_regular(Expression cond, Expression if_true, Expression if_false) { + return call("if_else", {std::move(cond), std::move(if_true), std::move(if_false)}); +} +Expression unreachable(Expression arg) { return call("unreachable", {std::move(arg)}); } +Expression sv_aware(Expression arg) { return call("sv_aware", {std::move(arg)}); } +Expression sv_suppress(Expression arg) { return call("sv_suppress", {std::move(arg)}); } +Expression assert_sv_exist(Expression arg) { + return call("assert_sv_exist", {std::move(arg)}); +} +Expression assert_sv_empty(Expression arg) { + return call("assert_sv_empty", {std::move(arg)}); +} + +} // namespace + +TEST(IfElseSpecialForm, Basic) { + { + ARROW_SCOPED_TRACE("if (b != 0) then a / b else b"); + auto cond = call("not_equal", {field_ref("b"), literal(0)}); + auto if_true = call("divide", {field_ref("a"), field_ref("b")}); + auto if_false = field_ref("b"); + auto schema = arrow::schema({field("a", int32()), field("b", int32())}); + auto rb = RecordBatchFromJSON(schema, R"([ + [1, 1], + [2, 1], + [3, 0], + [4, 1], + [5, 1] + ])"); + auto batch = ExecBatch(*rb); + auto if_else_sp = if_else_special(cond, if_true, if_false); + { + auto expected = ArrayFromJSON(int32(), "[1, 2, 0, 4, 5]"); + AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); + } + { + ARROW_SCOPED_TRACE("(if (b != 0) then a / b else b) + 1"); + auto plus_one = call("add", {if_else_sp, literal(1)}); + { + auto expected = ArrayFromJSON(int32(), "[2, 3, 1, 5, 6]"); + AssertExprEqualIgnoreShape(plus_one, schema, batch, expected); + } + { + ARROW_SCOPED_TRACE( + "if ((if (b != 0) then a / b else b) + 1 != 1) then a / b else b"); + auto cond = call("not_equal", {plus_one, literal(1)}); + auto if_true = call("divide", {field_ref("a"), field_ref("b")}); + auto if_false = field_ref("b"); + auto if_else_sp = if_else_special(cond, if_true, if_false); + auto expected = ArrayFromJSON(int32(), "[1, 2, 0, 4, 5]"); + AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); + } + } + } + { + ARROW_SCOPED_TRACE("if (b != 0) then a else a"); + auto cond = call("not_equal", {field_ref("b"), literal(0)}); + auto if_true = field_ref("a"); + auto if_false = field_ref("a"); + auto schema = arrow::schema({field("a", int32()), field("b", int32())}); + auto rb = RecordBatchFromJSON(schema, R"([ + [1, 1], + [2, 1], + [3, 0], + [4, 1], + [5, 1] + ])"); + auto batch = ExecBatch(*rb); + auto if_else_sp = if_else_special(cond, if_true, if_false); + auto expected = ArrayFromJSON(int32(), "[1, 2, 3, 4, 5]"); + AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); + } +} + +TEST(IfElseSpecialForm, ImplicitCast) { + auto schema = arrow::schema({field("i8", int8()), field("i32", int32())}); + for (const auto& if_else_sp : + {if_else_special(literal(true), field_ref("i8"), field_ref("i32")), + if_else_special(literal(true), field_ref("i32"), field_ref("i8")), + // Literal will be downcast. + if_else_special(literal(true), field_ref("i32"), + literal(static_cast(0)))}) { + ARROW_SCOPED_TRACE(if_else_sp.ToString()); + ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema)); + ASSERT_EQ(bound.type()->id(), Type::INT32); + } +} + +class IfElseSpecialFormTest : public ::testing::Test { + protected: + static void SetUpTestSuite() { ASSERT_OK(RegisterAuxilaryFunctions()); } + + protected: + static Status UnreachableExec(KernelContext*, const ExecSpan&, ExecResult*) { + return Status::Invalid("Unreachable"); + } + + static Status IdentityExec(KernelContext*, const ExecSpan& span, ExecResult* out) { + DCHECK_EQ(span.num_values(), 1); + const auto& arg = span[0]; + DCHECK(arg.is_array()); + *out->array_data_mutable() = *arg.array.ToArrayData(); + return Status::OK(); + } + + template + static Status AssertSelectionVectorExec(KernelContext* kernel_ctx, const ExecSpan& span, + ExecResult* out) { + if constexpr (sv_existence) { + if (!span.selection_vector.indices()) { + return Status::Invalid("There is no selection vector"); + } + } else { + if (span.selection_vector.indices()) { + return Status::Invalid("There is a selection vector"); + } + } + return IdentityExec(kernel_ctx, span, out); + } + + static Status RegisterAuxilaryFunctions() { + auto registry = GetFunctionRegistry(); + + { + auto register_unreachable_func = [&](const std::string& name) -> Status { + auto func = + std::make_shared(name, Arity::Unary(), FunctionDoc::Empty()); + + ScalarKernel kernel({InputType::Any()}, internal::FirstType, UnreachableExec); + kernel.selection_vector_aware = true; + kernel.can_write_into_slices = false; + kernel.null_handling = NullHandling::COMPUTED_NO_PREALLOCATE; + kernel.mem_allocation = MemAllocation::NO_PREALLOCATE; + RETURN_NOT_OK(func->AddKernel(kernel)); + RETURN_NOT_OK(registry->AddFunction(std::move(func))); + return Status::OK(); + }; + + RETURN_NOT_OK(register_unreachable_func("unreachable")); + } + + { + auto register_sv_awareness_func = [&](const std::string& name, + bool sv_awareness) -> Status { + auto func = + std::make_shared(name, Arity::Unary(), FunctionDoc::Empty()); + + ScalarKernel kernel({InputType::Any()}, internal::FirstType, IdentityExec); + kernel.selection_vector_aware = sv_awareness; + kernel.can_write_into_slices = false; + kernel.null_handling = NullHandling::COMPUTED_NO_PREALLOCATE; + kernel.mem_allocation = MemAllocation::NO_PREALLOCATE; + RETURN_NOT_OK(func->AddKernel(kernel)); + RETURN_NOT_OK(registry->AddFunction(std::move(func))); + return Status::OK(); + }; + + RETURN_NOT_OK(register_sv_awareness_func("sv_aware", true)); + RETURN_NOT_OK(register_sv_awareness_func("sv_suppress", false)); + } + + { + auto register_assert_sv_func = [&](const std::string& name, + bool sv_existence) -> Status { + auto func = + std::make_shared(name, Arity::Unary(), FunctionDoc::Empty()); + + ArrayKernelExec exec; + if (sv_existence) { + exec = AssertSelectionVectorExec; + } else { + exec = AssertSelectionVectorExec; + } + ScalarKernel kernel({InputType::Any()}, internal::FirstType, std::move(exec)); + kernel.selection_vector_aware = true; + kernel.can_write_into_slices = false; + kernel.null_handling = NullHandling::COMPUTED_NO_PREALLOCATE; + kernel.mem_allocation = MemAllocation::NO_PREALLOCATE; + RETURN_NOT_OK(func->AddKernel(kernel)); + RETURN_NOT_OK(registry->AddFunction(std::move(func))); + return Status::OK(); + }; + + RETURN_NOT_OK(register_assert_sv_func("assert_sv_exist", /*sv_existence=*/true)); + RETURN_NOT_OK(register_assert_sv_func("assert_sv_empty", /*sv_existence=*/false)); + } + + return Status::OK(); + } + + static std::vector SuppressSelectionVectorAwareForIfElse( + const Expression& cond, const Expression& if_true, const Expression& if_false) { + auto suppress_if_else_recursive = + [&](const Expression& expr) -> std::vector { + if (const auto& sp = expr.special(); sp && sp->special_form->name == "if_else") { + const auto& cond = sp->arguments[0]; + const auto& if_true = sp->arguments[1]; + const auto& if_false = sp->arguments[2]; + return SuppressSelectionVectorAwareForIfElse(cond, if_true, if_false); + } else { + return {expr}; + } + }; + auto suppressed_conds = suppress_if_else_recursive(cond); + auto suppressed_if_trues = suppress_if_else_recursive(if_true); + auto suppressed_if_falses = suppress_if_else_recursive(if_false); + std::vector result; + for (const auto& suppressed_cond : suppressed_conds) { + for (const auto& suppressed_if_true : suppressed_if_trues) { + for (const auto& suppressed_if_false : suppressed_if_falses) { + result.emplace_back( + if_else_special(suppressed_cond, suppressed_if_true, suppressed_if_false)); + result.emplace_back(if_else_special(sv_suppress(suppressed_cond), + suppressed_if_true, suppressed_if_false)); + result.emplace_back(if_else_special(suppressed_cond, + sv_suppress(suppressed_if_true), + sv_suppress(suppressed_if_false))); + } + } + } + return result; + } + + static void CheckIfElseIgnoreShape(const Expression& cond, const Expression& if_true, + const Expression& if_false, + const std::shared_ptr& schema, + const ExecBatch& batch, + ExecContext* exec_context = default_exec_context()) { + auto if_else = if_else_regular(cond, if_true, if_false); + auto exprs = SuppressSelectionVectorAwareForIfElse(cond, if_true, if_false); + AssertExprEqualExprsIgnoreShape(if_else, exprs, schema, batch, exec_context); + } +}; + +TEST_F(IfElseSpecialFormTest, AuxilaryFunction) { + auto schema = arrow::schema({field("a", boolean())}); + auto a = field_ref("a"); + auto batch = ExecBatch({*ArrayFromJSON(boolean(), "[null, true, false]")}, 3); + { + ARROW_SCOPED_TRACE("unreachable"); + AssertExprRaisesWithMessage(unreachable(a), schema, batch, Invalid, + "Invalid: Unreachable"); + } + { + ARROW_SCOPED_TRACE("selection vector awareness"); + { + ASSERT_OK_AND_ASSIGN(auto bound, sv_aware(a).Bind(*schema)); + ASSERT_TRUE(bound.selection_vector_aware()); + } + { + ASSERT_OK_AND_ASSIGN(auto bound, sv_suppress(a).Bind(*schema)); + ASSERT_FALSE(bound.selection_vector_aware()); + } + { + ASSERT_OK_AND_ASSIGN(auto bound, sv_aware(sv_aware(a)).Bind(*schema)); + ASSERT_TRUE(bound.selection_vector_aware()); + } + { + ASSERT_OK_AND_ASSIGN(auto bound, sv_aware(sv_suppress(a)).Bind(*schema)); + ASSERT_FALSE(bound.selection_vector_aware()); + } + { + ASSERT_OK_AND_ASSIGN(auto bound, sv_suppress(sv_aware(a)).Bind(*schema)); + ASSERT_FALSE(bound.selection_vector_aware()); + } + } + { + ARROW_SCOPED_TRACE("assert selection vector existence"); + { + ARROW_SCOPED_TRACE("if (a) then a else a"); + auto cond = a; + auto if_true = a; + auto if_false = a; + auto expected = ArrayFromJSON(boolean(), "[null, true, false]"); + { + auto if_else_sp = + if_else_special(cond, assert_sv_exist(if_true), assert_sv_exist(if_false)); + AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); + } + { + auto if_else_sp = + if_else_special(cond, sv_suppress(if_true), assert_sv_exist(if_false)); + AssertExprRaisesWithMessage(if_else_sp, schema, batch, Invalid, + "Invalid: There is no selection vector"); + } + { + auto if_else_sp = + if_else_special(cond, assert_sv_exist(if_true), sv_suppress(if_false)); + AssertExprRaisesWithMessage(if_else_sp, schema, batch, Invalid, + "Invalid: There is no selection vector"); + } + { + auto if_else_sp = + if_else_special(cond, sv_suppress(if_true), assert_sv_empty(if_false)); + AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); + } + { + auto if_else_sp = + if_else_special(cond, assert_sv_empty(if_true), sv_suppress(if_false)); + AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); + } + { + auto if_else_sp = if_else_special(cond, assert_sv_empty(if_true), if_false); + AssertExprRaisesWithMessage(if_else_sp, schema, batch, Invalid, + "Invalid: There is a selection vector"); + } + { + auto if_else_sp = if_else_special(cond, if_true, assert_sv_empty(if_false)); + AssertExprRaisesWithMessage(if_else_sp, schema, batch, Invalid, + "Invalid: There is a selection vector"); + } + } + { + ARROW_SCOPED_TRACE("if (true) then a else false"); + auto cond = literal(true); + auto if_true = a; + auto if_false = literal(false); + auto expected = ArrayFromJSON(boolean(), "[null, true, false]"); + { + auto if_else_sp = if_else_special(cond, assert_sv_exist(if_true), if_false); + AssertExprRaisesWithMessage(if_else_sp, schema, batch, Invalid, + "Invalid: There is no selection vector"); + } + { + auto if_else_sp = if_else_special(cond, assert_sv_empty(if_true), if_false); + AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); + } + } + } +} + +TEST_F(IfElseSpecialFormTest, SelectionVectorAwareness) { + { + ARROW_SCOPED_TRACE("literal"); + ASSERT_TRUE(kBooleanNull.selection_vector_aware()); + ASSERT_TRUE(kIntNull.selection_vector_aware()); + } + { + ARROW_SCOPED_TRACE("field ref"); + auto schema = arrow::schema({field("a", int32())}); + ASSERT_OK_AND_ASSIGN(auto bound, field_ref("a").Bind(*schema)); + ASSERT_TRUE(bound.selection_vector_aware()); + } + { + ARROW_SCOPED_TRACE("if else special"); + auto schema = arrow::schema({}); + auto cond = literal(true); + auto if_true = literal(1); + auto if_false = literal(0); + for (const auto& if_else_sp : + {if_else_special(cond, if_true, if_false), + if_else_special(sv_suppress(cond), if_true, if_false), + if_else_special(cond, sv_suppress(if_true), if_false), + if_else_special(cond, if_true, sv_suppress(if_false))}) { + ARROW_SCOPED_TRACE(if_else_sp.ToString()); + ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema)); + ASSERT_TRUE(bound.selection_vector_aware()); + } + } +} + +TEST_F(IfElseSpecialFormTest, SelectionVectorExistence) { + auto schema = arrow::schema({field("b_null", boolean()), field("b_true", boolean()), + field("b_false", boolean()), field("b", boolean()), + field("i1", int32()), field("i2", int32())}); + auto b_null = field_ref("b_null"); + auto b_true = field_ref("b_true"); + auto b_false = field_ref("b_false"); + auto b = field_ref("b"); + auto i1 = field_ref("i1"); + auto i2 = field_ref("i2"); + auto batch = ExecBatch( + { + *ArrayFromJSON(boolean(), "[null, null, null]"), + *ArrayFromJSON(boolean(), "[true, true, true]"), + *ArrayFromJSON(boolean(), "[false, false, false]"), + *ArrayFromJSON(boolean(), "[null, true, false]"), + *ArrayFromJSON(int32(), "[0, 1, -1]"), + *ArrayFromJSON(int32(), "[0, -1, 1]"), + }, + 3); + + { + ARROW_SCOPED_TRACE("all null condition"); + auto expected = ArrayFromJSON(int32(), "[null, null, null]"); + for (const auto& cond : + {kBooleanNull, b_null, assert_sv_empty(kBooleanNull), assert_sv_empty(b_null)}) { + ARROW_SCOPED_TRACE("cond: " + cond.ToString()); + auto if_else_sp = if_else_special(cond, unreachable(i1), unreachable(i2)); + AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); + } + } + { + ARROW_SCOPED_TRACE("all true condition"); + auto expected = ArrayFromJSON(int32(), "[0, 1, -1]"); + for (const auto& if_else_sp : + {if_else_special(literal(true), i1, unreachable(i2)), + if_else_special(b_true, i1, unreachable(i2)), + if_else_special(assert_sv_empty(literal(true)), assert_sv_empty(i1), + unreachable(i2)), + if_else_special(assert_sv_empty(b_true), assert_sv_exist(i1), unreachable(i2)), + if_else_special(assert_sv_empty(b_true), assert_sv_empty(i1), + sv_suppress(unreachable(i2)))}) { + ARROW_SCOPED_TRACE(if_else_sp.ToString()); + AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); + } + } + { + ARROW_SCOPED_TRACE("all false condition"); + auto expected = ArrayFromJSON(int32(), "[0, -1, 1]"); + for (const auto& if_else_sp : + {if_else_special(literal(false), unreachable(i1), i2), + if_else_special(b_false, unreachable(i1), i2), + if_else_special(assert_sv_empty(literal(false)), unreachable(i1), + assert_sv_empty(i2)), + if_else_special(assert_sv_empty(b_false), unreachable(i1), assert_sv_exist(i2)), + if_else_special(assert_sv_empty(b_false), sv_suppress(unreachable(i1)), + assert_sv_empty(i2))}) { + ARROW_SCOPED_TRACE(if_else_sp.ToString()); + AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); + } + } + { + ARROW_SCOPED_TRACE("even condition"); + auto expected = ArrayFromJSON(int32(), "[null, 1, 1]"); + for (const auto& if_else_sp : + {if_else_special(b, i1, i2), + if_else_special(assert_sv_empty(b), assert_sv_exist(i1), assert_sv_exist(i2)), + if_else_special(sv_suppress(b), assert_sv_exist(i1), assert_sv_exist(i2)), + if_else_special(b, sv_suppress(i1), assert_sv_empty(i2)), + if_else_special(b, assert_sv_empty(i1), sv_suppress(i2))}) { + ARROW_SCOPED_TRACE(if_else_sp.ToString()); + AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); + } + } + { + ARROW_SCOPED_TRACE("literal bodies"); + auto expected = ArrayFromJSON(int32(), "[null, 1, 1]"); + for (const auto& if_else_sp : + {if_else_special(assert_sv_empty(b), assert_sv_empty(literal(1)), + assert_sv_exist(i2)), + if_else_special(assert_sv_empty(b), assert_sv_exist(i1), + assert_sv_empty(literal(1))), + if_else_special(assert_sv_empty(b), assert_sv_empty(literal(1)), + assert_sv_empty(literal(1)))}) { + ARROW_SCOPED_TRACE(if_else_sp.ToString()); + AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); + } + } + { + ARROW_SCOPED_TRACE("nested"); + auto expected = ArrayFromJSON(int32(), "[null, 1, 1]"); + for (const auto& if_else_sp : + {if_else_special(b, if_else_special(b, i1, i2), if_else_special(b, i1, i2)), + // The nested if_else_special will see a selection vector. + if_else_special(b, assert_sv_exist(if_else_special(b, i1, i2)), + assert_sv_exist(if_else_special(b, i1, i2))), + // The arguments of nested if_else_special will see a selection vector. + if_else_special(b, + if_else_special(assert_sv_empty(literal(true)), + assert_sv_exist(i1), unreachable(i2)), + if_else_special(assert_sv_empty(literal(false)), + unreachable(i1), assert_sv_exist(i2))), + if_else_special(b, + if_else_special(assert_sv_exist(b), assert_sv_exist(i1), + assert_sv_exist(i2)), + if_else_special(assert_sv_exist(b), assert_sv_exist(i1), + assert_sv_exist(i2))), + // Selection vector existences with some argument of the nested if_else_special + // being selection vector unaware. + if_else_special( + b, + assert_sv_exist(if_else_special(sv_suppress(literal(true)), + assert_sv_exist(i1), unreachable(i2))), + assert_sv_exist(if_else_special(assert_sv_empty(literal(false)), + unreachable(i1), assert_sv_exist(i2)))), + if_else_special(b, + assert_sv_exist(if_else_special( + sv_suppress(b), assert_sv_exist(i1), assert_sv_exist(i2))), + assert_sv_exist(if_else_special( + assert_sv_exist(b), unreachable(i1), assert_sv_exist(i2)))), + if_else_special( + b, + assert_sv_exist(if_else_special(assert_sv_empty(literal(true)), + sv_suppress(i1), unreachable(i2))), + assert_sv_exist(if_else_special(assert_sv_empty(literal(false)), + unreachable(i1), assert_sv_exist(i2)))), + if_else_special( + b, + assert_sv_exist(if_else_special(assert_sv_exist(b), sv_suppress(i1), + assert_sv_empty(i2))), + assert_sv_exist(if_else_special(assert_sv_exist(b), assert_sv_exist(i1), + assert_sv_exist(i2)))), + if_else_special( + b, + assert_sv_exist(if_else_special(assert_sv_empty(literal(true)), + assert_sv_empty(i1), + sv_suppress(unreachable(i2)))), + assert_sv_exist(if_else_special(assert_sv_empty(literal(false)), + unreachable(i1), assert_sv_exist(i2)))), + if_else_special( + b, + assert_sv_exist(if_else_special(assert_sv_exist(b), assert_sv_empty(i1), + sv_suppress(i2))), + assert_sv_exist(if_else_special(assert_sv_exist(b), assert_sv_exist(i1), + assert_sv_exist(i2)))), + if_else_special( + b, + assert_sv_exist(if_else_special(assert_sv_empty(literal(true)), + assert_sv_exist(i1), unreachable(i2))), + assert_sv_exist(if_else_special(assert_sv_empty(literal(false)), + unreachable(i1), assert_sv_exist(i2)))), + if_else_special( + b, + assert_sv_exist(if_else_special(assert_sv_exist(b), assert_sv_exist(i1), + assert_sv_exist(i2))), + assert_sv_exist(if_else_special(sv_suppress(b), assert_sv_exist(i1), + assert_sv_exist(i2)))), + if_else_special( + b, + assert_sv_exist(if_else_special(assert_sv_empty(literal(true)), + assert_sv_exist(i1), unreachable(i2))), + assert_sv_exist(if_else_special(assert_sv_empty(literal(false)), + sv_suppress(unreachable(i1)), + assert_sv_empty(i2)))), + if_else_special( + b, + assert_sv_exist(if_else_special(assert_sv_exist(b), assert_sv_exist(i1), + assert_sv_exist(i2))), + assert_sv_exist(if_else_special(assert_sv_exist(b), sv_suppress(i1), + assert_sv_empty(i2)))), + if_else_special( + b, + assert_sv_exist(if_else_special(assert_sv_empty(literal(true)), + assert_sv_exist(i1), unreachable(i2))), + assert_sv_exist(if_else_special(assert_sv_empty(literal(false)), + assert_sv_empty(unreachable(i1)), + sv_suppress(i2)))), + if_else_special( + b, + assert_sv_exist(if_else_special(assert_sv_exist(b), assert_sv_exist(i1), + assert_sv_exist(i2))), + assert_sv_exist(if_else_special(assert_sv_exist(b), assert_sv_empty(i1), + sv_suppress(i2))))}) { + ARROW_SCOPED_TRACE(if_else_sp.ToString()); + AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); + } + } +} + +TEST_F(IfElseSpecialFormTest, SelectionVectorExistenceExecChunkSize) { + ExecContext exec_context; + constexpr int64_t num_rows = 8; + auto schema = arrow::schema({field("b", boolean()), field("i", int32())}); + auto b = field_ref("b"); + auto i = field_ref("i"); + auto batch = ExecBatch( + { + *ArrayFromJSON(boolean(), "[true, true, true, true, true, true, true, false]"), + *ArrayFromJSON(int32(), "[42, 42, 42, 42, 42, 42, 42, 42]"), + }, + num_rows); + { + ARROW_SCOPED_TRACE("exec_chunksize >= batch_size"); + for (auto chunksize : {num_rows, num_rows + 1}) { + ARROW_SCOPED_TRACE("exec_chunksize: " + std::to_string(chunksize)); + exec_context.set_exec_chunksize(chunksize); + { + ARROW_SCOPED_TRACE("all literal bodies"); + auto if_else_sp = + if_else_special(b, assert_sv_empty(literal(1)), assert_sv_empty(literal(0))); + ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); + ASSERT_TRUE(bound.selection_vector_aware()); + ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); + } + { + ARROW_SCOPED_TRACE("array true body"); + auto if_else_sp = + if_else_special(b, assert_sv_exist(i), assert_sv_empty(literal(0))); + ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); + ASSERT_TRUE(bound.selection_vector_aware()); + ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); + } + { + ARROW_SCOPED_TRACE("array false body"); + auto if_else_sp = + if_else_special(b, assert_sv_empty(literal(1)), assert_sv_exist(i)); + ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); + ASSERT_TRUE(bound.selection_vector_aware()); + ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); + } + { + ARROW_SCOPED_TRACE("all array bodies"); + auto if_else_sp = if_else_special(b, assert_sv_exist(i), assert_sv_exist(i)); + ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); + ASSERT_TRUE(bound.selection_vector_aware()); + ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); + } + } + } + { + ARROW_SCOPED_TRACE("exec_chunksize < batch_size"); + exec_context.set_exec_chunksize(num_rows - 1); + { + ARROW_SCOPED_TRACE("all literal bodies"); + auto if_else_sp = + if_else_special(b, assert_sv_empty(literal(1)), assert_sv_empty(literal(0))); + ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); + ASSERT_TRUE(bound.selection_vector_aware()); + ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); + } + { + ARROW_SCOPED_TRACE("array true body"); + auto if_else_sp = + if_else_special(b, assert_sv_exist(i), assert_sv_empty(literal(0))); + ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); + ASSERT_TRUE(bound.selection_vector_aware()); + ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); + } + { + ARROW_SCOPED_TRACE("array false body"); + auto if_else_sp = + if_else_special(b, assert_sv_empty(literal(1)), assert_sv_exist(i)); + ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); + ASSERT_TRUE(bound.selection_vector_aware()); + ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); + } + { + ARROW_SCOPED_TRACE("all array bodies"); + auto if_else_sp = if_else_special(b, assert_sv_exist(i), assert_sv_exist(i)); + ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); + ASSERT_TRUE(bound.selection_vector_aware()); + ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); + } + } + { + ARROW_SCOPED_TRACE("nested"); + { + ARROW_SCOPED_TRACE("exec_chunksize == batch_size"); + exec_context.set_exec_chunksize(num_rows); + auto if_else_sp = if_else_special( + b, + assert_sv_exist(if_else_special(assert_sv_exist(b), assert_sv_exist(i), + assert_sv_exist(i))), + assert_sv_exist(if_else_special(assert_sv_exist(b), assert_sv_exist(i), + assert_sv_exist(i)))); + ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); + ASSERT_TRUE(bound.selection_vector_aware()); + ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); + } + { + ARROW_SCOPED_TRACE("exec_chunksize < batch_size"); + exec_context.set_exec_chunksize(num_rows - 2); + auto if_else_sp = if_else_special( + b, + assert_sv_exist(if_else_special(assert_sv_exist(b), assert_sv_exist(i), + assert_sv_exist(i))), + assert_sv_exist(if_else_special( + assert_sv_exist(b), + assert_sv_exist(if_else_special(assert_sv_exist(b), assert_sv_exist(i), + assert_sv_exist(i))), + assert_sv_exist(i)))); + ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); + ASSERT_TRUE(bound.selection_vector_aware()); + ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); + } + } +} + +// TODO: Selection vector existence of chunked array. + +TEST_F(IfElseSpecialFormTest, ResultShape) { + auto schema = arrow::schema({field("b_null", boolean()), field("b_true", boolean()), + field("b_false", boolean())}); + auto b_null = field_ref("b_null"); + auto b_true = field_ref("b_true"); + auto b_false = field_ref("b_false"); + auto batch = ExecBatch( + { + *ArrayFromJSON(boolean(), "[null, null, null]"), + *ArrayFromJSON(boolean(), "[true, true, true]"), + *ArrayFromJSON(boolean(), "[false, false, false]"), + }, + 3); + { + ARROW_SCOPED_TRACE("if (null) then 1 else 0"); + auto expected = ArrayFromJSON(int32(), "[null, null, null]"); + for (const auto& cond : {kBooleanNull, b_null}) { + ARROW_SCOPED_TRACE("cond: " + cond.ToString()); + for (const auto& if_else_sp : + SuppressSelectionVectorAwareForIfElse(cond, literal(1), literal(0))) { + ARROW_SCOPED_TRACE(if_else_sp.ToString()); + ASSERT_OK_AND_ASSIGN(auto result, ExecuteExpr(if_else_sp, schema, batch)); + AssertDatumsEqual(expected, result); + } + } + } + { + ARROW_SCOPED_TRACE("if (true) then 1 else 0"); + auto expected = MakeScalar(1); + for (const auto& cond : {literal(true), b_true}) { + ARROW_SCOPED_TRACE("cond: " + cond.ToString()); + for (const auto& if_else_sp : + SuppressSelectionVectorAwareForIfElse(cond, literal(1), literal(0))) { + ARROW_SCOPED_TRACE(if_else_sp.ToString()); + ASSERT_OK_AND_ASSIGN(auto result, ExecuteExpr(if_else_sp, schema, batch)); + AssertDatumsEqual(expected, result); + } + } + } + { + ARROW_SCOPED_TRACE("if (false) then 1 else 0"); + auto expected = MakeScalar(0); + for (const auto& cond : {literal(false), b_false}) { + ARROW_SCOPED_TRACE("cond: " + cond.ToString()); + for (const auto& if_else_sp : + SuppressSelectionVectorAwareForIfElse(cond, literal(1), literal(0))) { + ARROW_SCOPED_TRACE(if_else_sp.ToString()); + ASSERT_OK_AND_ASSIGN(auto result, ExecuteExpr(if_else_sp, schema, batch)); + AssertDatumsEqual(expected, result); + } + } + } + { + ARROW_SCOPED_TRACE("if (if (null) true then false) then 1 else 0"); + auto expected = ArrayFromJSON(int32(), "[null, null, null]"); + for (const auto& nested_cond : {kBooleanNull, b_null}) { + auto cond = if_else_special(nested_cond, literal(true), literal(false)); + ARROW_SCOPED_TRACE("nested cond: " + cond.ToString()); + for (const auto& if_else_sp : + SuppressSelectionVectorAwareForIfElse(cond, literal(1), literal(0))) { + ARROW_SCOPED_TRACE(if_else_sp.ToString()); + ASSERT_OK_AND_ASSIGN(auto result, ExecuteExpr(if_else_sp, schema, batch)); + AssertDatumsEqual(expected, result); + } + } + } + { + ARROW_SCOPED_TRACE("if (if (true) then true else false) then 1 else 0"); + auto expected = MakeScalar(1); + for (const auto& nested_cond : {literal(true), b_true}) { + auto cond = if_else_special(nested_cond, nested_cond, literal(false)); + ARROW_SCOPED_TRACE("nested cond: " + cond.ToString()); + for (const auto& if_else_sp : + SuppressSelectionVectorAwareForIfElse(cond, literal(1), literal(0))) { + ARROW_SCOPED_TRACE(if_else_sp.ToString()); + ASSERT_OK_AND_ASSIGN(auto result, ExecuteExpr(if_else_sp, schema, batch)); + AssertDatumsEqual(expected, result); + } + } + } + { + ARROW_SCOPED_TRACE("if (if (false) then true else false) then 1 else 0"); + auto expected = MakeScalar(0); + for (const auto& nested_cond : {literal(false), b_false}) { + auto cond = if_else_special(nested_cond, literal(true), nested_cond); + ARROW_SCOPED_TRACE("nested cond: " + cond.ToString()); + for (const auto& if_else_sp : + SuppressSelectionVectorAwareForIfElse(cond, literal(1), literal(0))) { + ARROW_SCOPED_TRACE(if_else_sp.ToString()); + ASSERT_OK_AND_ASSIGN(auto result, ExecuteExpr(if_else_sp, schema, batch)); + AssertDatumsEqual(expected, result); + } + } + } + // TODO: Non-scalar branch bodies. +} + +namespace { + +auto kCanonicalSchema = arrow::schema({field("a", boolean()), field("b", int32())}); + +auto kCanonicalA = field_ref("a"); +auto kCanonicalB = field_ref("b"); + +const auto kCanonicalBooleanCols = {kBooleanNull, literal(true), literal(false), + kCanonicalA}; +const auto kCanonicalIntCols = {kIntNull, literal(42), kCanonicalB}; + +const std::vector kCanonicalBatches = { + ExecBatch(*RecordBatchFromJSON(kCanonicalSchema, R"([ + [null, 0], + [null, null], + [null, 1] + ])")), + ExecBatch(*RecordBatchFromJSON(kCanonicalSchema, R"([ + [true, 0], + [true, null], + [true, 1] + ])")), + ExecBatch(*RecordBatchFromJSON(kCanonicalSchema, R"([ + [false, 0], + [false, null], + [false, 1] + ])")), + ExecBatch(*RecordBatchFromJSON(kCanonicalSchema, R"([ + [false, 0], + [true, null], + [null, 1] + ])")), +}; + +} // namespace + +TEST_F(IfElseSpecialFormTest, Simple) { + const auto& schema = kCanonicalSchema; + const auto& boolean_datums = kCanonicalBooleanCols; + const auto& int_datums = kCanonicalIntCols; + const auto& batches = kCanonicalBatches; + for (const auto& cond : boolean_datums) { + for (const auto& if_true : int_datums) { + for (const auto& if_false : int_datums) { + for (const auto& batch : batches) { + CheckIfElseIgnoreShape(cond, if_true, if_false, schema, batch); + } + } + } + } +} + +TEST_F(IfElseSpecialFormTest, NestedSimple) { + const auto& schema = kCanonicalSchema; + const auto& a = kCanonicalA; + const auto& b = kCanonicalB; + ExecBatch batch(*RecordBatchFromJSON(kCanonicalSchema, R"([ + [false, 0], + [true, null], + [null, 1] + ])")); + for (const auto& cond : { + if_else_special(a, kBooleanNull, a), + if_else_special(a, a, literal(true)), + }) { + for (const auto& if_true : { + if_else_special(a, kIntNull, b), + if_else_special(a, b, literal(42)), + }) { + for (const auto& if_false : { + if_else_special(a, kIntNull, b), + if_else_special(a, b, literal(42)), + }) { + CheckIfElseIgnoreShape(cond, if_true, if_false, schema, batch); + } + } + } +} + +// TODO: Deprecate this test due to slowness. +TEST_F(IfElseSpecialFormTest, NestedConditionComplex) { + const auto& batches = kCanonicalBatches; + const auto& schema = kCanonicalSchema; + const auto& boolean_datums = kCanonicalBooleanCols; + const auto& int_datums = kCanonicalIntCols; + for (const auto& nested_cond : boolean_datums) { + for (const auto& nested_if_true : boolean_datums) { + for (const auto& nested_if_false : boolean_datums) { + auto nested_if_else_sp = + if_else_special(nested_cond, nested_if_true, nested_if_false); + for (const auto& if_true : int_datums) { + for (const auto& if_false : int_datums) { + for (const auto& batch : batches) { + CheckIfElseIgnoreShape(nested_if_else_sp, if_true, if_false, schema, batch); + } + } + } + } + } + } +} + +// TODO: Deprecate this test due to slowness. +TEST_F(IfElseSpecialFormTest, NestedBodyComplex) { + const auto& batches = kCanonicalBatches; + const auto& schema = kCanonicalSchema; + const auto& boolean_datums = kCanonicalBooleanCols; + const auto& int_datums = kCanonicalIntCols; + for (const auto& cond : boolean_datums) { + for (const auto& nested_cond : boolean_datums) { + for (const auto& nested_if_true : int_datums) { + for (const auto& nested_if_false : int_datums) { + auto nested_if_else_sp = + if_else_special(nested_cond, nested_if_true, nested_if_false); + for (const auto& batch : batches) { + CheckIfElseIgnoreShape(cond, nested_if_else_sp, nested_if_else_sp, schema, + batch); + } + } + } + } + } +} + +// TODO: Deprecate this test due to slowness. +TEST_F(IfElseSpecialFormTest, NestedComplex) { + const auto& batches = kCanonicalBatches; + const auto& schema = kCanonicalSchema; + const auto& boolean_datums = kCanonicalBooleanCols; + const auto& int_datums = kCanonicalIntCols; + for (const auto& cond_nested_cond : boolean_datums) { + for (const auto& cond_nested_if_true : boolean_datums) { + for (const auto& cond_nested_if_false : boolean_datums) { + auto cond = + if_else_special(cond_nested_cond, cond_nested_if_true, cond_nested_if_false); + for (const auto& nested_cond : boolean_datums) { + for (const auto& nested_if_true : int_datums) { + for (const auto& nested_if_false : int_datums) { + auto nested_if_else_sp = + if_else_special(nested_cond, nested_if_true, nested_if_false); + for (const auto& batch : batches) { + CheckIfElseIgnoreShape(cond, nested_if_else_sp, nested_if_else_sp, schema, + batch); + } + } + } + } + } + } + } +} + +// TODO: ChunkedArray. + +namespace { +template +Status ConstantKernelExec(KernelContext*, const ExecSpan& span, ExecResult* out) { + DCHECK_EQ(span.num_values(), 1); + DCHECK_EQ(span.length, 1); + DCHECK(out->is_array_span()); + DCHECK_EQ(out->length(), 1); + if constexpr (!selection_vector_aware) { + if (span.selection_vector.length() > 0) { + return Status::Invalid("There is a selection vector"); + } + } + int32_t* out_values = out->array_span_mutable()->GetValues(1); + *out_values = 0; + return Status::OK(); +} + +static Status RegisterConstantFunctions() { + auto registry = GetFunctionRegistry(); + + auto register_test_func = [&](const std::string& name, + bool selection_vector_aware) -> Status { + auto zero = + std::make_shared(name, Arity::Unary(), FunctionDoc::Empty()); + + ArrayKernelExec exec; + if (selection_vector_aware) { + exec = ConstantKernelExec; + } else { + exec = ConstantKernelExec; + } + ScalarKernel kernel({InputType::Any()}, OutputType{int32()}, std::move(exec)); + kernel.selection_vector_aware = selection_vector_aware; + kernel.can_write_into_slices = true; + kernel.null_handling = NullHandling::OUTPUT_NOT_NULL; + kernel.mem_allocation = MemAllocation::PREALLOCATE; + RETURN_NOT_OK(zero->AddKernel(kernel)); + RETURN_NOT_OK(registry->AddFunction(std::move(zero))); + return Status::OK(); + }; + + RETURN_NOT_OK(register_test_func("zero_panic", false)); + RETURN_NOT_OK(register_test_func("zero_calm", true)); + + return Status::OK(); +} + +} // namespace + +TEST(IfElseSpecialForm, Reference) { + ASSERT_OK(RegisterConstantFunctions()); + + auto schema = arrow::schema({field("a", int32()), field("b", int32())}); + std::vector batches = { + ExecBatch(*RecordBatchFromJSON(schema, R"([])")), + ExecBatch(*RecordBatchFromJSON(schema, R"([ + [1, 0], + [1, 0], + [1, 0] + ])")), + ExecBatch(*RecordBatchFromJSON(schema, R"([ + [1, 0], + [null, 0], + [1, null] + ])")), + }; + for (const auto& batch : batches) { + auto expr = call("zero_panic", {literal(42)}); + ASSERT_OK_AND_ASSIGN(auto bound, expr.Bind(*schema)); + ASSERT_OK_AND_ASSIGN(auto result, ExecuteScalarExpression(bound, batch)); + std::cout << result.ToString() << std::endl; + } + // TODO: The result shape of exec_chunksize, chunked input, and preallocate. +} + +} // namespace arrow::compute diff --git a/cpp/src/arrow/compute/type_fwd.h b/cpp/src/arrow/compute/type_fwd.h index 016d97a0dbc2..e316469107a9 100644 --- a/cpp/src/arrow/compute/type_fwd.h +++ b/cpp/src/arrow/compute/type_fwd.h @@ -52,6 +52,9 @@ struct KernelState; class Expression; +class SpecialExec; +class SpecialForm; + ARROW_EXPORT ExecContext* default_exec_context(); ARROW_EXPORT ExecContext* threaded_exec_context(); From 432205487017f4e2ecf7fbf28189e77b8e356f27 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Fri, 1 Aug 2025 09:30:30 +0800 Subject: [PATCH 02/71] WIP --- cpp/src/arrow/compute/kernel.h | 22 +++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/cpp/src/arrow/compute/kernel.h b/cpp/src/arrow/compute/kernel.h index e1d799a93946..7acd169a8f69 100644 --- a/cpp/src/arrow/compute/kernel.h +++ b/cpp/src/arrow/compute/kernel.h @@ -529,6 +529,8 @@ struct ARROW_EXPORT Kernel { static Status InitAll(KernelContext*, const KernelInitArgs&, std::vector>*); + virtual bool selection_vector_aware() const { return false; } + /// \brief Indicates whether execution can benefit from parallelization /// (splitting large chunks into smaller chunks and using multiple /// threads). Some kernels may not support parallel execution at @@ -542,7 +544,7 @@ struct ARROW_EXPORT Kernel { /// so that the most optimized kernel supported on a host's processor can be chosen. SimdLevel::type simd_level = SimdLevel::NONE; - bool selection_vector_aware = false; + // bool selection_vector_aware = false; // Additional kernel-specific data std::shared_ptr data; @@ -557,6 +559,10 @@ struct ARROW_EXPORT Kernel { /// employed this may not be possible. using ArrayKernelExec = Status (*)(KernelContext*, const ExecSpan&, ExecResult*); +using ArrayKernelSelectionVectorAwareExec = Status (*)(KernelContext*, const ExecSpan&, + const SelectionVectorSpan&, + ExecResult*); + /// \brief Kernel data structure for implementations of ScalarFunction. In /// addition to the members found in Kernel, contains the null handling /// and memory pre-allocation preferences. @@ -564,12 +570,16 @@ struct ARROW_EXPORT ScalarKernel : public Kernel { ScalarKernel() = default; ScalarKernel(std::shared_ptr sig, ArrayKernelExec exec, - KernelInit init = NULLPTR) - : Kernel(std::move(sig), init), exec(exec) {} + KernelInit init = NULLPTR, + ArrayKernelSelectionVectorAwareExec sv_exec = NULLPTR) + : Kernel(std::move(sig), init), exec(exec), sv_exec(sv_exec) {} ScalarKernel(std::vector in_types, OutputType out_type, ArrayKernelExec exec, - KernelInit init = NULLPTR) - : Kernel(std::move(in_types), std::move(out_type), std::move(init)), exec(exec) {} + KernelInit init = NULLPTR, + ArrayKernelSelectionVectorAwareExec sv_exec = NULLPTR) + : Kernel(std::move(in_types), std::move(out_type), std::move(init)), + exec(exec), + sv_exec(sv_exec) {} /// \brief Perform a single invocation of this kernel. Depending on the /// implementation, it may only write into preallocated memory, while in some @@ -577,6 +587,8 @@ struct ARROW_EXPORT ScalarKernel : public Kernel { /// through the KernelContext. ArrayKernelExec exec; + ArrayKernelSelectionVectorAwareExec sv_exec; + /// \brief Writing execution results into larger contiguous allocations /// requires that the kernel be able to write into sliced output ArrayData*, /// including sliced output validity bitmaps. Some kernel implementations may From 2ac65af47206593d500680b2fe1d42f8fe7abbb5 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Sat, 2 Aug 2025 14:34:22 +0800 Subject: [PATCH 03/71] Support selective kernel exec and move dense execution to executor --- cpp/src/arrow/compute/exec.cc | 91 ++++++++++++++++++---- cpp/src/arrow/compute/exec.h | 2 +- cpp/src/arrow/compute/expression.cc | 2 +- cpp/src/arrow/compute/kernel.h | 18 ++--- cpp/src/arrow/compute/special_form.cc | 8 +- cpp/src/arrow/compute/special_form_test.cc | 6 +- 6 files changed, 92 insertions(+), 35 deletions(-) diff --git a/cpp/src/arrow/compute/exec.cc b/cpp/src/arrow/compute/exec.cc index f3c855eda0b5..064662525763 100644 --- a/cpp/src/arrow/compute/exec.cc +++ b/cpp/src/arrow/compute/exec.cc @@ -451,6 +451,8 @@ bool ExecSpanIterator::Next(ExecSpan* span) { span->selection_vector = SelectionVectorSpan(selection_vector_->indices(), selection_vector_->length()); } + } else { + span->selection_vector = std::nullopt; } } @@ -466,7 +468,7 @@ bool ExecSpanIterator::Next(ExecSpan* span) { iteration_size = GetNextChunkSpan(iteration_size, span); if (selection_vector_) { auto indices_begin = - span->selection_vector.indices() + span->selection_vector.length(); + span->selection_vector->indices() + span->selection_vector->length(); auto indices_end = selection_vector_->indices() + selection_vector_->length(); DCHECK_LE(indices_begin, indices_end); auto chunk_row_id_end = position_ + iteration_size; @@ -475,7 +477,7 @@ bool ExecSpanIterator::Next(ExecSpan* span) { *(indices_begin + num_indices) < chunk_row_id_end) { ++num_indices; } - span->selection_vector.SetSlice(span->selection_vector.length(), num_indices); + span->selection_vector->SetSlice(span->selection_vector->length(), num_indices); } } @@ -806,8 +808,6 @@ class KernelExecutorImpl : public KernelExecutor { class ScalarExecutor : public KernelExecutorImpl { public: Status Execute(const ExecBatch& batch, ExecListener* listener) override { - RETURN_NOT_OK(span_iterator_.Init(batch, exec_context()->exec_chunksize())); - if (batch.length == 0) { // For zero-length batches, we do nothing except return a zero-length // array of the correct output type @@ -817,6 +817,31 @@ class ScalarExecutor : public KernelExecutorImpl { return EmitResult(result->data(), listener); } + if (!batch.selection_vector || kernel_->selective_exec) { + return ExecuteSparse(batch, listener); + } + + return ExecuteDense(batch, listener); + } + + Datum WrapResults(const std::vector& inputs, + const std::vector& outputs) override { + // If execution yielded multiple chunks (because large arrays were split + // based on the ExecContext parameters, then the result is a ChunkedArray + if (HaveChunkedArray(inputs) || outputs.size() > 1) { + return ToChunkedArray(outputs, output_type_); + } else { + // Outputs have just one element + return outputs[0]; + } + } + + protected: + Status ExecuteSparse(const ExecBatch& batch, ExecListener* listener) { + DCHECK(!batch.selection_vector || kernel_->selective_exec); + + RETURN_NOT_OK(span_iterator_.Init(batch, exec_context()->exec_chunksize())); + // If the executor is configured to produce a single large Array output for // kernels supporting preallocation, then we do so up front and then // iterate over slices of that large array. Otherwise, we preallocate prior @@ -836,19 +861,47 @@ class ScalarExecutor : public KernelExecutorImpl { } } - Datum WrapResults(const std::vector& inputs, - const std::vector& outputs) override { - // If execution yielded multiple chunks (because large arrays were split - // based on the ExecContext parameters, then the result is a ChunkedArray - if (HaveChunkedArray(inputs) || outputs.size() > 1) { - return ToChunkedArray(outputs, output_type_); + Status ExecuteDense(const ExecBatch& batch, ExecListener* listener) { + DCHECK(batch.selection_vector && !kernel_->selective_exec); + + if (CheckIfAllScalar(batch)) { + ExecBatch input = batch; + input.selection_vector = nullptr; + return ExecuteSparse(input, listener); + } + + std::vector values(batch.num_values()); + for (int i = 0; i < batch.num_values(); ++i) { + if (batch[i].is_scalar()) { + // Skip Take for scalars since it is not currently supported. + values[i] = batch[i]; + continue; + } + ARROW_ASSIGN_OR_RAISE(values[i], + Take(batch[i], *batch.selection_vector->data(), + TakeOptions{/*boundcheck=*/false}, exec_context())); + } + ExecBatch input; + ARROW_ASSIGN_OR_RAISE( + input, ExecBatch::Make(std::move(values), batch.selection_vector->length())); + + DatumAccumulator dense_listener; + RETURN_NOT_OK(ExecuteSparse(input, &dense_listener)); + auto dense_results = dense_listener.values(); + + Datum dense_datum; + if (dense_results.size() > 1) { + dense_datum = ToChunkedArray(dense_results, output_type_); } else { - // Outputs have just one element - return outputs[0]; + dense_datum = dense_results[0]; } + + ARROW_ASSIGN_OR_RAISE(auto result, + Scatter(dense_datum, *batch.selection_vector->data(), + ScatterOptions{/*max_index=*/batch.length - 1})); + return listener->OnResult(std::move(result)); } - protected: Status EmitResult(std::shared_ptr out, ExecListener* listener) { if (span_iterator_.have_all_scalars()) { // ARROW-16757 We boxed scalar inputs as ArraySpan, so now we have to @@ -913,7 +966,7 @@ class ScalarExecutor : public KernelExecutorImpl { } else if (kernel_->null_handling == NullHandling::OUTPUT_NOT_NULL) { result_span->null_count = 0; } - RETURN_NOT_OK(kernel_->exec(kernel_ctx_, input, out)); + RETURN_NOT_OK(InvokeKernel(input, out)); // Output type didn't change DCHECK(out->is_array_span()); return Status::OK(); @@ -942,7 +995,7 @@ class ScalarExecutor : public KernelExecutorImpl { out_arr->null_count = 0; } - RETURN_NOT_OK(kernel_->exec(kernel_ctx_, input, &output)); + RETURN_NOT_OK(InvokeKernel(input, &output)); // Output type didn't change DCHECK(output.is_array_data()); @@ -1008,6 +1061,14 @@ class ScalarExecutor : public KernelExecutorImpl { return Status::OK(); } + Status InvokeKernel(const ExecSpan& input, ExecResult* out) { + if (input.selection_vector) { + DCHECK_NE(kernel_->selective_exec, nullptr); + return kernel_->selective_exec(kernel_ctx_, input, out); + } + return kernel_->exec(kernel_ctx_, input, out); + } + // Used to account for the case where we do not preallocate a // validity bitmap because the inputs are all non-null and we're // using NullHandling::INTERSECTION to compute the validity bitmap diff --git a/cpp/src/arrow/compute/exec.h b/cpp/src/arrow/compute/exec.h index b6a0943d4a5c..f2d15b5258f1 100644 --- a/cpp/src/arrow/compute/exec.h +++ b/cpp/src/arrow/compute/exec.h @@ -447,7 +447,7 @@ struct ARROW_EXPORT ExecSpan { int64_t length = 0; std::vector values; - SelectionVectorSpan selection_vector; + std::optional selection_vector; }; /// \defgroup compute-call-function One-shot calls to compute functions diff --git a/cpp/src/arrow/compute/expression.cc b/cpp/src/arrow/compute/expression.cc index ad95579f9fe6..edb97ecf4b78 100644 --- a/cpp/src/arrow/compute/expression.cc +++ b/cpp/src/arrow/compute/expression.cc @@ -875,7 +875,7 @@ Result ExecuteScalarExpression(const Expression& expr, const ExecBatch& i "ExecuteScalarExpression cannot Execute non-scalar expression ", expr.ToString()); } - DCHECK(!input.selection_vector || expr.selection_vector_aware()); + // DCHECK(!input.selection_vector || expr.selection_vector_aware()); if (auto lit = expr.literal()) return *lit; diff --git a/cpp/src/arrow/compute/kernel.h b/cpp/src/arrow/compute/kernel.h index 7acd169a8f69..6b064dddf9ea 100644 --- a/cpp/src/arrow/compute/kernel.h +++ b/cpp/src/arrow/compute/kernel.h @@ -529,8 +529,6 @@ struct ARROW_EXPORT Kernel { static Status InitAll(KernelContext*, const KernelInitArgs&, std::vector>*); - virtual bool selection_vector_aware() const { return false; } - /// \brief Indicates whether execution can benefit from parallelization /// (splitting large chunks into smaller chunks and using multiple /// threads). Some kernels may not support parallel execution at @@ -544,7 +542,7 @@ struct ARROW_EXPORT Kernel { /// so that the most optimized kernel supported on a host's processor can be chosen. SimdLevel::type simd_level = SimdLevel::NONE; - // bool selection_vector_aware = false; + bool selection_vector_aware = false; // Additional kernel-specific data std::shared_ptr data; @@ -559,9 +557,7 @@ struct ARROW_EXPORT Kernel { /// employed this may not be possible. using ArrayKernelExec = Status (*)(KernelContext*, const ExecSpan&, ExecResult*); -using ArrayKernelSelectionVectorAwareExec = Status (*)(KernelContext*, const ExecSpan&, - const SelectionVectorSpan&, - ExecResult*); +using ArrayKernelSelectiveExec = Status (*)(KernelContext*, const ExecSpan&, ExecResult*); /// \brief Kernel data structure for implementations of ScalarFunction. In /// addition to the members found in Kernel, contains the null handling @@ -571,15 +567,15 @@ struct ARROW_EXPORT ScalarKernel : public Kernel { ScalarKernel(std::shared_ptr sig, ArrayKernelExec exec, KernelInit init = NULLPTR, - ArrayKernelSelectionVectorAwareExec sv_exec = NULLPTR) - : Kernel(std::move(sig), init), exec(exec), sv_exec(sv_exec) {} + ArrayKernelSelectiveExec selective_exec = NULLPTR) + : Kernel(std::move(sig), init), exec(exec), selective_exec(selective_exec) {} ScalarKernel(std::vector in_types, OutputType out_type, ArrayKernelExec exec, KernelInit init = NULLPTR, - ArrayKernelSelectionVectorAwareExec sv_exec = NULLPTR) + ArrayKernelSelectiveExec selective_exec = NULLPTR) : Kernel(std::move(in_types), std::move(out_type), std::move(init)), exec(exec), - sv_exec(sv_exec) {} + selective_exec(selective_exec) {} /// \brief Perform a single invocation of this kernel. Depending on the /// implementation, it may only write into preallocated memory, while in some @@ -587,7 +583,7 @@ struct ARROW_EXPORT ScalarKernel : public Kernel { /// through the KernelContext. ArrayKernelExec exec; - ArrayKernelSelectionVectorAwareExec sv_exec; + ArrayKernelSelectiveExec selective_exec; /// \brief Writing execution results into larger contiguous allocations /// requires that the kernel be able to write into sliced output ArrayData*, diff --git a/cpp/src/arrow/compute/special_form.cc b/cpp/src/arrow/compute/special_form.cc index 6c0595dead86..2f97a940535c 100644 --- a/cpp/src/arrow/compute/special_form.cc +++ b/cpp/src/arrow/compute/special_form.cc @@ -951,11 +951,11 @@ class IfElseSpecialExec : public SpecialExec { Result Execute(const ExecBatch& input, ExecContext* exec_context) const override { - if (all_bodies_selection_vector_aware) { + // if (all_bodies_selection_vector_aware) { return SparseConditionalExecutor(branches, type).Execute(input, exec_context); - } else { - return DenseConditionalExecutor(branches, type).Execute(input, exec_context); - } + // } else { + // return DenseConditionalExecutor(branches, type).Execute(input, exec_context); + // } } private: diff --git a/cpp/src/arrow/compute/special_form_test.cc b/cpp/src/arrow/compute/special_form_test.cc index efe3f7c08f63..89392c71a202 100644 --- a/cpp/src/arrow/compute/special_form_test.cc +++ b/cpp/src/arrow/compute/special_form_test.cc @@ -195,11 +195,11 @@ class IfElseSpecialFormTest : public ::testing::Test { static Status AssertSelectionVectorExec(KernelContext* kernel_ctx, const ExecSpan& span, ExecResult* out) { if constexpr (sv_existence) { - if (!span.selection_vector.indices()) { + if (!span.selection_vector->indices()) { return Status::Invalid("There is no selection vector"); } } else { - if (span.selection_vector.indices()) { + if (span.selection_vector->indices()) { return Status::Invalid("There is a selection vector"); } } @@ -1016,7 +1016,7 @@ Status ConstantKernelExec(KernelContext*, const ExecSpan& span, ExecResult* out) DCHECK(out->is_array_span()); DCHECK_EQ(out->length(), 1); if constexpr (!selection_vector_aware) { - if (span.selection_vector.length() > 0) { + if (span.selection_vector->length() > 0) { return Status::Invalid("There is a selection vector"); } } From 66f85ae68f32fec3a82e3f9da16bf1804e988e57 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Sat, 2 Aug 2025 15:50:56 +0800 Subject: [PATCH 04/71] Remove the concept of selection vector awareness --- cpp/src/arrow/compute/expression.cc | 25 --- cpp/src/arrow/compute/expression.h | 4 - cpp/src/arrow/compute/kernel.h | 2 - cpp/src/arrow/compute/special_form.cc | 20 +- cpp/src/arrow/compute/special_form.h | 4 +- cpp/src/arrow/compute/special_form_test.cc | 245 ++++++++------------- 6 files changed, 92 insertions(+), 208 deletions(-) diff --git a/cpp/src/arrow/compute/expression.cc b/cpp/src/arrow/compute/expression.cc index edb97ecf4b78..f44f474b3f00 100644 --- a/cpp/src/arrow/compute/expression.cc +++ b/cpp/src/arrow/compute/expression.cc @@ -152,20 +152,6 @@ const DataType* Expression::type() const { return CallNotNull(*this)->type.type; } -bool Expression::selection_vector_aware() const { - DCHECK(IsBound()); - - if (literal() || field_ref()) { - return true; - } - - if (auto special = this->special()) { - return special->selection_vector_aware; - } - - return CallNotNull(*this)->selection_vector_aware; -} - namespace { std::string PrintDatum(const Datum& datum) { @@ -681,11 +667,6 @@ Result BindNonRecursive(Expression::Call call, bool insert_implicit_ kernel_context.SetState(call.kernel_state.get()); } - call.selection_vector_aware = - call.kernel->selection_vector_aware && - std::all_of(call.arguments.begin(), call.arguments.end(), - [](const Expression& arg) { return arg.selection_vector_aware(); }); - ARROW_ASSIGN_OR_RAISE( call.type, call.kernel->signature->out_type().Resolve(&kernel_context, types)); return Status::OK(); @@ -714,10 +695,6 @@ Result BindNonRecursive(Expression::Special special, std::unique_ptr&& special_exec) { special.special_exec = std::move(special_exec); - // Selection vector awareness-es of the subexpressions is fully taken over by the - // special form so not recursive. - special.selection_vector_aware = special.special_form->selection_vector_aware; - ARROW_ASSIGN_OR_RAISE(special.type, special.special_exec->Bind(special.arguments, exec_context)); @@ -875,8 +852,6 @@ Result ExecuteScalarExpression(const Expression& expr, const ExecBatch& i "ExecuteScalarExpression cannot Execute non-scalar expression ", expr.ToString()); } - // DCHECK(!input.selection_vector || expr.selection_vector_aware()); - if (auto lit = expr.literal()) return *lit; if (auto param = expr.parameter()) { diff --git a/cpp/src/arrow/compute/expression.h b/cpp/src/arrow/compute/expression.h index 61cebfb8ddee..2aa8187e3cd7 100644 --- a/cpp/src/arrow/compute/expression.h +++ b/cpp/src/arrow/compute/expression.h @@ -56,7 +56,6 @@ class ARROW_EXPORT Expression { const Kernel* kernel = NULLPTR; std::shared_ptr kernel_state; TypeHolder type; - bool selection_vector_aware; void ComputeHash(); }; @@ -70,7 +69,6 @@ class ARROW_EXPORT Expression { // post-Bind properties: std::shared_ptr special_exec; TypeHolder type; - bool selection_vector_aware; void ComputeHash(); }; @@ -135,8 +133,6 @@ class ARROW_EXPORT Expression { // XXX someday // NullGeneralization::type nullable() const; - bool selection_vector_aware() const; - struct Parameter { FieldRef ref; diff --git a/cpp/src/arrow/compute/kernel.h b/cpp/src/arrow/compute/kernel.h index 6b064dddf9ea..a4e319361bec 100644 --- a/cpp/src/arrow/compute/kernel.h +++ b/cpp/src/arrow/compute/kernel.h @@ -542,8 +542,6 @@ struct ARROW_EXPORT Kernel { /// so that the most optimized kernel supported on a host's processor can be chosen. SimdLevel::type simd_level = SimdLevel::NONE; - bool selection_vector_aware = false; - // Additional kernel-specific data std::shared_ptr data; }; diff --git a/cpp/src/arrow/compute/special_form.cc b/cpp/src/arrow/compute/special_form.cc index 2f97a940535c..6fdea40cd037 100644 --- a/cpp/src/arrow/compute/special_form.cc +++ b/cpp/src/arrow/compute/special_form.cc @@ -739,10 +739,8 @@ struct ConditionalExecutor { Result> ApplyCond( const std::shared_ptr& branch_mask, const Expression& cond, const ExecBatch& input, ExecContext* exec_context) const { - if (cond.selection_vector_aware()) { - return branch_mask->ApplyCondSparse(cond, input, exec_context); - } - return branch_mask->ApplyCondDense(cond, input, exec_context); + return branch_mask->ApplyCondSparse(cond, input, exec_context); + // return branch_mask->ApplyCondDense(cond, input, exec_context); } protected: @@ -940,9 +938,6 @@ class IfElseSpecialExec : public SpecialExec { DCHECK_EQ(type, *if_true.type()); DCHECK_EQ(type, *if_false.type()); - all_bodies_selection_vector_aware = - if_true.selection_vector_aware() && if_false.selection_vector_aware(); - branches[0] = {cond, if_true}; branches[1] = {literal(true), if_false}; @@ -951,11 +946,8 @@ class IfElseSpecialExec : public SpecialExec { Result Execute(const ExecBatch& input, ExecContext* exec_context) const override { - // if (all_bodies_selection_vector_aware) { - return SparseConditionalExecutor(branches, type).Execute(input, exec_context); - // } else { - // return DenseConditionalExecutor(branches, type).Execute(input, exec_context); - // } + return SparseConditionalExecutor(branches, type).Execute(input, exec_context); + // return DenseConditionalExecutor(branches, type).Execute(input, exec_context); } private: @@ -965,14 +957,12 @@ class IfElseSpecialExec : public SpecialExec { // Post-bind, for Execute. TypeHolder type; - bool all_bodies_selection_vector_aware; std::vector branches; }; class IfElseSpecialForm : public SpecialForm { public: - IfElseSpecialForm() - : SpecialForm(/*name=*/"if_else", /*selection_vector_aware=*/true) {} + IfElseSpecialForm() : SpecialForm(/*name=*/"if_else") {} Result> DispatchExact( const std::vector& types, ExecContext* exec_context) const override { diff --git a/cpp/src/arrow/compute/special_form.h b/cpp/src/arrow/compute/special_form.h index 8273e8405154..9582ce418d32 100644 --- a/cpp/src/arrow/compute/special_form.h +++ b/cpp/src/arrow/compute/special_form.h @@ -32,8 +32,7 @@ class ARROW_EXPORT SpecialExec { class ARROW_EXPORT SpecialForm { public: - explicit SpecialForm(std::string name, bool selection_vector_aware = false) - : name(std::move(name)), selection_vector_aware(selection_vector_aware) {} + explicit SpecialForm(std::string name) : name(std::move(name)) {} virtual ~SpecialForm() = default; @@ -45,7 +44,6 @@ class ARROW_EXPORT SpecialForm { public: const std::string name; - const bool selection_vector_aware = false; }; std::shared_ptr GetIfElseSpecialForm(); diff --git a/cpp/src/arrow/compute/special_form_test.cc b/cpp/src/arrow/compute/special_form_test.cc index 89392c71a202..e62e1fb54869 100644 --- a/cpp/src/arrow/compute/special_form_test.cc +++ b/cpp/src/arrow/compute/special_form_test.cc @@ -90,7 +90,6 @@ Expression if_else_regular(Expression cond, Expression if_true, Expression if_fa return call("if_else", {std::move(cond), std::move(if_true), std::move(if_false)}); } Expression unreachable(Expression arg) { return call("unreachable", {std::move(arg)}); } -Expression sv_aware(Expression arg) { return call("sv_aware", {std::move(arg)}); } Expression sv_suppress(Expression arg) { return call("sv_suppress", {std::move(arg)}); } Expression assert_sv_exist(Expression arg) { return call("assert_sv_exist", {std::move(arg)}); @@ -195,11 +194,11 @@ class IfElseSpecialFormTest : public ::testing::Test { static Status AssertSelectionVectorExec(KernelContext* kernel_ctx, const ExecSpan& span, ExecResult* out) { if constexpr (sv_existence) { - if (!span.selection_vector->indices()) { + if (!span.selection_vector) { return Status::Invalid("There is no selection vector"); } } else { - if (span.selection_vector->indices()) { + if (span.selection_vector) { return Status::Invalid("There is a selection vector"); } } @@ -215,7 +214,6 @@ class IfElseSpecialFormTest : public ::testing::Test { std::make_shared(name, Arity::Unary(), FunctionDoc::Empty()); ScalarKernel kernel({InputType::Any()}, internal::FirstType, UnreachableExec); - kernel.selection_vector_aware = true; kernel.can_write_into_slices = false; kernel.null_handling = NullHandling::COMPUTED_NO_PREALLOCATE; kernel.mem_allocation = MemAllocation::NO_PREALLOCATE; @@ -234,7 +232,6 @@ class IfElseSpecialFormTest : public ::testing::Test { std::make_shared(name, Arity::Unary(), FunctionDoc::Empty()); ScalarKernel kernel({InputType::Any()}, internal::FirstType, IdentityExec); - kernel.selection_vector_aware = sv_awareness; kernel.can_write_into_slices = false; kernel.null_handling = NullHandling::COMPUTED_NO_PREALLOCATE; kernel.mem_allocation = MemAllocation::NO_PREALLOCATE; @@ -243,7 +240,6 @@ class IfElseSpecialFormTest : public ::testing::Test { return Status::OK(); }; - RETURN_NOT_OK(register_sv_awareness_func("sv_aware", true)); RETURN_NOT_OK(register_sv_awareness_func("sv_suppress", false)); } @@ -253,14 +249,13 @@ class IfElseSpecialFormTest : public ::testing::Test { auto func = std::make_shared(name, Arity::Unary(), FunctionDoc::Empty()); - ArrayKernelExec exec; + ArrayKernelExec exec = AssertSelectionVectorExec; + ArrayKernelSelectiveExec selective_exec = nullptr; if (sv_existence) { - exec = AssertSelectionVectorExec; - } else { - exec = AssertSelectionVectorExec; + selective_exec = AssertSelectionVectorExec; } - ScalarKernel kernel({InputType::Any()}, internal::FirstType, std::move(exec)); - kernel.selection_vector_aware = true; + ScalarKernel kernel({InputType::Any()}, internal::FirstType, std::move(exec), + /*init=*/nullptr, std::move(selective_exec)); kernel.can_write_into_slices = false; kernel.null_handling = NullHandling::COMPUTED_NO_PREALLOCATE; kernel.mem_allocation = MemAllocation::NO_PREALLOCATE; @@ -329,29 +324,6 @@ TEST_F(IfElseSpecialFormTest, AuxilaryFunction) { AssertExprRaisesWithMessage(unreachable(a), schema, batch, Invalid, "Invalid: Unreachable"); } - { - ARROW_SCOPED_TRACE("selection vector awareness"); - { - ASSERT_OK_AND_ASSIGN(auto bound, sv_aware(a).Bind(*schema)); - ASSERT_TRUE(bound.selection_vector_aware()); - } - { - ASSERT_OK_AND_ASSIGN(auto bound, sv_suppress(a).Bind(*schema)); - ASSERT_FALSE(bound.selection_vector_aware()); - } - { - ASSERT_OK_AND_ASSIGN(auto bound, sv_aware(sv_aware(a)).Bind(*schema)); - ASSERT_TRUE(bound.selection_vector_aware()); - } - { - ASSERT_OK_AND_ASSIGN(auto bound, sv_aware(sv_suppress(a)).Bind(*schema)); - ASSERT_FALSE(bound.selection_vector_aware()); - } - { - ASSERT_OK_AND_ASSIGN(auto bound, sv_suppress(sv_aware(a)).Bind(*schema)); - ASSERT_FALSE(bound.selection_vector_aware()); - } - } { ARROW_SCOPED_TRACE("assert selection vector existence"); { @@ -368,14 +340,12 @@ TEST_F(IfElseSpecialFormTest, AuxilaryFunction) { { auto if_else_sp = if_else_special(cond, sv_suppress(if_true), assert_sv_exist(if_false)); - AssertExprRaisesWithMessage(if_else_sp, schema, batch, Invalid, - "Invalid: There is no selection vector"); + AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); } { auto if_else_sp = if_else_special(cond, assert_sv_exist(if_true), sv_suppress(if_false)); - AssertExprRaisesWithMessage(if_else_sp, schema, batch, Invalid, - "Invalid: There is no selection vector"); + AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); } { auto if_else_sp = @@ -389,13 +359,11 @@ TEST_F(IfElseSpecialFormTest, AuxilaryFunction) { } { auto if_else_sp = if_else_special(cond, assert_sv_empty(if_true), if_false); - AssertExprRaisesWithMessage(if_else_sp, schema, batch, Invalid, - "Invalid: There is a selection vector"); + AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); } { auto if_else_sp = if_else_special(cond, if_true, assert_sv_empty(if_false)); - AssertExprRaisesWithMessage(if_else_sp, schema, batch, Invalid, - "Invalid: There is a selection vector"); + AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); } } { @@ -406,8 +374,7 @@ TEST_F(IfElseSpecialFormTest, AuxilaryFunction) { auto expected = ArrayFromJSON(boolean(), "[null, true, false]"); { auto if_else_sp = if_else_special(cond, assert_sv_exist(if_true), if_false); - AssertExprRaisesWithMessage(if_else_sp, schema, batch, Invalid, - "Invalid: There is no selection vector"); + AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); } { auto if_else_sp = if_else_special(cond, assert_sv_empty(if_true), if_false); @@ -417,36 +384,6 @@ TEST_F(IfElseSpecialFormTest, AuxilaryFunction) { } } -TEST_F(IfElseSpecialFormTest, SelectionVectorAwareness) { - { - ARROW_SCOPED_TRACE("literal"); - ASSERT_TRUE(kBooleanNull.selection_vector_aware()); - ASSERT_TRUE(kIntNull.selection_vector_aware()); - } - { - ARROW_SCOPED_TRACE("field ref"); - auto schema = arrow::schema({field("a", int32())}); - ASSERT_OK_AND_ASSIGN(auto bound, field_ref("a").Bind(*schema)); - ASSERT_TRUE(bound.selection_vector_aware()); - } - { - ARROW_SCOPED_TRACE("if else special"); - auto schema = arrow::schema({}); - auto cond = literal(true); - auto if_true = literal(1); - auto if_false = literal(0); - for (const auto& if_else_sp : - {if_else_special(cond, if_true, if_false), - if_else_special(sv_suppress(cond), if_true, if_false), - if_else_special(cond, sv_suppress(if_true), if_false), - if_else_special(cond, if_true, sv_suppress(if_false))}) { - ARROW_SCOPED_TRACE(if_else_sp.ToString()); - ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema)); - ASSERT_TRUE(bound.selection_vector_aware()); - } - } -} - TEST_F(IfElseSpecialFormTest, SelectionVectorExistence) { auto schema = arrow::schema({field("b_null", boolean()), field("b_true", boolean()), field("b_false", boolean()), field("b", boolean()), @@ -658,7 +595,6 @@ TEST_F(IfElseSpecialFormTest, SelectionVectorExistenceExecChunkSize) { auto if_else_sp = if_else_special(b, assert_sv_empty(literal(1)), assert_sv_empty(literal(0))); ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); - ASSERT_TRUE(bound.selection_vector_aware()); ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); } { @@ -666,7 +602,6 @@ TEST_F(IfElseSpecialFormTest, SelectionVectorExistenceExecChunkSize) { auto if_else_sp = if_else_special(b, assert_sv_exist(i), assert_sv_empty(literal(0))); ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); - ASSERT_TRUE(bound.selection_vector_aware()); ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); } { @@ -674,14 +609,12 @@ TEST_F(IfElseSpecialFormTest, SelectionVectorExistenceExecChunkSize) { auto if_else_sp = if_else_special(b, assert_sv_empty(literal(1)), assert_sv_exist(i)); ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); - ASSERT_TRUE(bound.selection_vector_aware()); ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); } { ARROW_SCOPED_TRACE("all array bodies"); auto if_else_sp = if_else_special(b, assert_sv_exist(i), assert_sv_exist(i)); ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); - ASSERT_TRUE(bound.selection_vector_aware()); ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); } } @@ -694,7 +627,6 @@ TEST_F(IfElseSpecialFormTest, SelectionVectorExistenceExecChunkSize) { auto if_else_sp = if_else_special(b, assert_sv_empty(literal(1)), assert_sv_empty(literal(0))); ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); - ASSERT_TRUE(bound.selection_vector_aware()); ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); } { @@ -702,7 +634,6 @@ TEST_F(IfElseSpecialFormTest, SelectionVectorExistenceExecChunkSize) { auto if_else_sp = if_else_special(b, assert_sv_exist(i), assert_sv_empty(literal(0))); ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); - ASSERT_TRUE(bound.selection_vector_aware()); ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); } { @@ -710,14 +641,12 @@ TEST_F(IfElseSpecialFormTest, SelectionVectorExistenceExecChunkSize) { auto if_else_sp = if_else_special(b, assert_sv_empty(literal(1)), assert_sv_exist(i)); ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); - ASSERT_TRUE(bound.selection_vector_aware()); ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); } { ARROW_SCOPED_TRACE("all array bodies"); auto if_else_sp = if_else_special(b, assert_sv_exist(i), assert_sv_exist(i)); ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); - ASSERT_TRUE(bound.selection_vector_aware()); ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); } } @@ -733,7 +662,6 @@ TEST_F(IfElseSpecialFormTest, SelectionVectorExistenceExecChunkSize) { assert_sv_exist(if_else_special(assert_sv_exist(b), assert_sv_exist(i), assert_sv_exist(i)))); ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); - ASSERT_TRUE(bound.selection_vector_aware()); ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); } { @@ -749,7 +677,6 @@ TEST_F(IfElseSpecialFormTest, SelectionVectorExistenceExecChunkSize) { assert_sv_exist(i))), assert_sv_exist(i)))); ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); - ASSERT_TRUE(bound.selection_vector_aware()); ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); } } @@ -1008,79 +935,79 @@ TEST_F(IfElseSpecialFormTest, NestedComplex) { // TODO: ChunkedArray. -namespace { -template -Status ConstantKernelExec(KernelContext*, const ExecSpan& span, ExecResult* out) { - DCHECK_EQ(span.num_values(), 1); - DCHECK_EQ(span.length, 1); - DCHECK(out->is_array_span()); - DCHECK_EQ(out->length(), 1); - if constexpr (!selection_vector_aware) { - if (span.selection_vector->length() > 0) { - return Status::Invalid("There is a selection vector"); - } - } - int32_t* out_values = out->array_span_mutable()->GetValues(1); - *out_values = 0; - return Status::OK(); -} - -static Status RegisterConstantFunctions() { - auto registry = GetFunctionRegistry(); - - auto register_test_func = [&](const std::string& name, - bool selection_vector_aware) -> Status { - auto zero = - std::make_shared(name, Arity::Unary(), FunctionDoc::Empty()); - - ArrayKernelExec exec; - if (selection_vector_aware) { - exec = ConstantKernelExec; - } else { - exec = ConstantKernelExec; - } - ScalarKernel kernel({InputType::Any()}, OutputType{int32()}, std::move(exec)); - kernel.selection_vector_aware = selection_vector_aware; - kernel.can_write_into_slices = true; - kernel.null_handling = NullHandling::OUTPUT_NOT_NULL; - kernel.mem_allocation = MemAllocation::PREALLOCATE; - RETURN_NOT_OK(zero->AddKernel(kernel)); - RETURN_NOT_OK(registry->AddFunction(std::move(zero))); - return Status::OK(); - }; - - RETURN_NOT_OK(register_test_func("zero_panic", false)); - RETURN_NOT_OK(register_test_func("zero_calm", true)); - - return Status::OK(); -} - -} // namespace - -TEST(IfElseSpecialForm, Reference) { - ASSERT_OK(RegisterConstantFunctions()); - - auto schema = arrow::schema({field("a", int32()), field("b", int32())}); - std::vector batches = { - ExecBatch(*RecordBatchFromJSON(schema, R"([])")), - ExecBatch(*RecordBatchFromJSON(schema, R"([ - [1, 0], - [1, 0], - [1, 0] - ])")), - ExecBatch(*RecordBatchFromJSON(schema, R"([ - [1, 0], - [null, 0], - [1, null] - ])")), - }; - for (const auto& batch : batches) { - auto expr = call("zero_panic", {literal(42)}); - ASSERT_OK_AND_ASSIGN(auto bound, expr.Bind(*schema)); - ASSERT_OK_AND_ASSIGN(auto result, ExecuteScalarExpression(bound, batch)); - std::cout << result.ToString() << std::endl; - } - // TODO: The result shape of exec_chunksize, chunked input, and preallocate. -} +// namespace { +// template +// Status ConstantKernelExec(KernelContext*, const ExecSpan& span, ExecResult* out) { +// DCHECK_EQ(span.num_values(), 1); +// DCHECK_EQ(span.length, 1); +// DCHECK(out->is_array_span()); +// DCHECK_EQ(out->length(), 1); +// if constexpr (!selection_vector_aware) { +// if (span.selection_vector->length() > 0) { +// return Status::Invalid("There is a selection vector"); +// } +// } +// int32_t* out_values = out->array_span_mutable()->GetValues(1); +// *out_values = 0; +// return Status::OK(); +// } + +// static Status RegisterConstantFunctions() { +// auto registry = GetFunctionRegistry(); + +// auto register_test_func = [&](const std::string& name, +// bool selection_vector_aware) -> Status { +// auto zero = +// std::make_shared(name, Arity::Unary(), FunctionDoc::Empty()); + +// ArrayKernelExec exec; +// if (selection_vector_aware) { +// exec = ConstantKernelExec; +// } else { +// exec = ConstantKernelExec; +// } +// ScalarKernel kernel({InputType::Any()}, OutputType{int32()}, std::move(exec)); +// kernel.selection_vector_aware = selection_vector_aware; +// kernel.can_write_into_slices = true; +// kernel.null_handling = NullHandling::OUTPUT_NOT_NULL; +// kernel.mem_allocation = MemAllocation::PREALLOCATE; +// RETURN_NOT_OK(zero->AddKernel(kernel)); +// RETURN_NOT_OK(registry->AddFunction(std::move(zero))); +// return Status::OK(); +// }; + +// RETURN_NOT_OK(register_test_func("zero_panic", false)); +// RETURN_NOT_OK(register_test_func("zero_calm", true)); + +// return Status::OK(); +// } + +// } // namespace + +// TEST(IfElseSpecialForm, Reference) { +// ASSERT_OK(RegisterConstantFunctions()); + +// auto schema = arrow::schema({field("a", int32()), field("b", int32())}); +// std::vector batches = { +// ExecBatch(*RecordBatchFromJSON(schema, R"([])")), +// ExecBatch(*RecordBatchFromJSON(schema, R"([ +// [1, 0], +// [1, 0], +// [1, 0] +// ])")), +// ExecBatch(*RecordBatchFromJSON(schema, R"([ +// [1, 0], +// [null, 0], +// [1, null] +// ])")), +// }; +// for (const auto& batch : batches) { +// auto expr = call("zero_panic", {literal(42)}); +// ASSERT_OK_AND_ASSIGN(auto bound, expr.Bind(*schema)); +// ASSERT_OK_AND_ASSIGN(auto result, ExecuteScalarExpression(bound, batch)); +// std::cout << result.ToString() << std::endl; +// } +// // TODO: The result shape of exec_chunksize, chunked input, and preallocate. +// } } // namespace arrow::compute From ce6d2ff98a92eb8df91640fca5a8b79cee25c731 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Sun, 3 Aug 2025 00:42:27 +0800 Subject: [PATCH 05/71] Cleanup dense things in special form --- cpp/src/arrow/compute/special_form.cc | 467 ++++---------------------- 1 file changed, 68 insertions(+), 399 deletions(-) diff --git a/cpp/src/arrow/compute/special_form.cc b/cpp/src/arrow/compute/special_form.cc index 6fdea40cd037..b1f0ae9037ca 100644 --- a/cpp/src/arrow/compute/special_form.cc +++ b/cpp/src/arrow/compute/special_form.cc @@ -30,77 +30,35 @@ namespace arrow::compute { namespace { -// TODO: Clean free functions. - struct BodyMask; struct BranchMask : public std::enable_shared_from_this { virtual ~BranchMask() = default; - Result> ApplyCondSparse( - const Expression& expr, const ExecBatch& input, ExecContext* exec_context) const { - ARROW_ASSIGN_OR_RAISE(auto datum, ApplySparse(expr, input, exec_context)); - return MakeBodyMaskFromSparseDatum(datum, exec_context); - } - - Result> ApplyCondDense( - const Expression& expr, const ExecBatch& input, ExecContext* exec_context) const { - ARROW_ASSIGN_OR_RAISE(auto datum, ApplyDense(expr, input, exec_context)); - return MakeBodyMaskFromDenseDatum(datum, exec_context); + Result> ApplyCond(const Expression& expr, + const ExecBatch& input, + ExecContext* exec_context) const { + ARROW_ASSIGN_OR_RAISE(auto datum, DoApplyCond(expr, input, exec_context)); + return MakeBodyMaskFromDatum(datum, exec_context); } virtual bool empty() const = 0; protected: - virtual Result ApplySparse(const Expression& expr, const ExecBatch& input, + virtual Result DoApplyCond(const Expression& expr, const ExecBatch& input, ExecContext* exec_context) const = 0; - virtual Result ApplyDense(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const = 0; - virtual Result> GetSelectionVector() const = 0; - virtual Result> MakeBodyMaskFromScalar( - const BooleanScalar& scalar, ExecContext* exec_context) const = 0; - - virtual Result> MakeBodyMaskFromSparseBitmap( - const std::shared_ptr& bitmap, ExecContext* exec_context) const = 0; - - virtual Result> MakeBodyMaskFromDenseBitmap( + virtual Result> MakeBodyMaskFromBitmap( const std::shared_ptr& bitmap, ExecContext* exec_context) const = 0; - virtual Result> MakeBodyMaskFromSparseBitmap( - const std::shared_ptr& bitmap, ExecContext* exec_context) const = 0; - - virtual Result> MakeBodyMaskFromDenseBitmap( + virtual Result> MakeBodyMaskFromBitmap( const std::shared_ptr& bitmap, ExecContext* exec_context) const = 0; private: - Result> MakeBodyMaskFromSparseDatum( - const Datum& datum, ExecContext* exec_context) const { - DCHECK(datum.type()->id() == Type::BOOL); - if (datum.is_scalar()) { - return MakeBodyMaskFromScalar(datum.scalar_as(), exec_context); - } - if (datum.is_array()) { - return MakeBodyMaskFromSparseBitmap(datum.array_as(), exec_context); - } - DCHECK(datum.is_chunked_array()); - return MakeBodyMaskFromSparseBitmap(datum.chunked_array(), exec_context); - } - - Result> MakeBodyMaskFromDenseDatum( - const Datum& datum, ExecContext* exec_context) const { - DCHECK(datum.type()->id() == Type::BOOL); - if (datum.is_scalar()) { - return MakeBodyMaskFromScalar(datum.scalar_as(), exec_context); - } - if (datum.is_array()) { - return MakeBodyMaskFromDenseBitmap(datum.array_as(), exec_context); - } - DCHECK(datum.is_chunked_array()); - return MakeBodyMaskFromDenseBitmap(datum.chunked_array(), exec_context); - } + Result> MakeBodyMaskFromDatum( + const Datum& datum, ExecContext* exec_context) const; friend struct NestedBodyMask; }; @@ -110,11 +68,8 @@ struct BodyMask : public std::enable_shared_from_this { virtual bool empty() const = 0; - virtual Result ApplySparse(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const = 0; - - virtual Result ApplyDense(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const = 0; + virtual Result ApplyCond(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const = 0; virtual Result> GetSelectionVector() const = 0; @@ -127,7 +82,7 @@ struct AllPassBranchMask : public BranchMask { bool empty() const override { return false; } protected: - Result ApplySparse(const Expression& expr, const ExecBatch& input, + Result DoApplyCond(const Expression& expr, const ExecBatch& input, ExecContext* exec_context) const override { DCHECK_EQ(input.length, length_); auto input_with_sel_vec = input; @@ -135,34 +90,15 @@ struct AllPassBranchMask : public BranchMask { return ExecuteScalarExpression(expr, input_with_sel_vec, exec_context); } - Result ApplyDense(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const override { - DCHECK_EQ(input.length, length_); - auto input_with_sel_vec = input; - input_with_sel_vec.selection_vector = nullptr; - return ExecuteScalarExpression(expr, input_with_sel_vec, exec_context); - } - Result> GetSelectionVector() const override { return nullptr; } - Result> MakeBodyMaskFromScalar( - const BooleanScalar& scalar, ExecContext* exec_context) const override; - - Result> MakeBodyMaskFromSparseBitmap( - const std::shared_ptr& bitmap, - ExecContext* exec_context) const override; - - Result> MakeBodyMaskFromDenseBitmap( + Result> MakeBodyMaskFromBitmap( const std::shared_ptr& bitmap, ExecContext* exec_context) const override; - Result> MakeBodyMaskFromSparseBitmap( - const std::shared_ptr& bitmap, - ExecContext* exec_context) const override; - - Result> MakeBodyMaskFromDenseBitmap( + Result> MakeBodyMaskFromBitmap( const std::shared_ptr& bitmap, ExecContext* exec_context) const override; @@ -176,16 +112,10 @@ struct AllFailBranchMask : public BranchMask { bool empty() const override { return true; } protected: - Result ApplySparse(const Expression& expr, const ExecBatch& input, + Result DoApplyCond(const Expression& expr, const ExecBatch& input, ExecContext* exec_context) const override { DCHECK(false); - return Status::Invalid("AllFailBranchMask::ApplySparse should not be called"); - } - - Result ApplyDense(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const override { - DCHECK(false); - return Status::Invalid("AllFailBranchMask::ApplyDense should not be called"); + return Status::Invalid("AllFailBranchMask::DoApplyCond should not be called"); } Result> GetSelectionVector() const override { @@ -193,42 +123,20 @@ struct AllFailBranchMask : public BranchMask { return Status::Invalid("AllFailBranchMask::GetSelectionVector should not be called"); } - Result> MakeBodyMaskFromScalar( - const BooleanScalar& scalar, ExecContext* exec_context) const override { - DCHECK(false); - return Status::Invalid("AllFailBranchMask::MakeBodyMask should not be called"); - } - - Result> MakeBodyMaskFromSparseBitmap( + Result> MakeBodyMaskFromBitmap( const std::shared_ptr& bitmap, ExecContext* exec_context) const override { DCHECK(false); return Status::Invalid( - "AllFailBranchMask::MakeBodyMaskFromSparseBitmap should not be called"); + "AllFailBranchMask::MakeBodyMaskFromBitmap should not be called"); } - Result> MakeBodyMaskFromDenseBitmap( - const std::shared_ptr& bitmap, - ExecContext* exec_context) const override { - DCHECK(false); - return Status::Invalid( - "AllFailBranchMask::MakeBodyMaskFromDenseBitmap should not be called"); - } - - Result> MakeBodyMaskFromSparseBitmap( + Result> MakeBodyMaskFromBitmap( const std::shared_ptr& bitmap, ExecContext* exec_context) const override { DCHECK(false); return Status::Invalid( - "AllFailBranchMask::MakeBodyMaskFromSparseBitmap should not be called"); - } - - Result> MakeBodyMaskFromDenseBitmap( - const std::shared_ptr& bitmap, - ExecContext* exec_context) const override { - DCHECK(false); - return Status::Invalid( - "AllFailBranchMask::MakeBodyMaskFromDenseBitmap should not be called"); + "AllFailBranchMask::MakeBodyMaskFromBitmap should not be called"); } }; @@ -237,14 +145,9 @@ struct NestedBodyMask : public BodyMask { : branch_mask_(std::move(branch_mask)) {} protected: - Result DelegateApplySparse(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const { - return branch_mask_->ApplySparse(expr, input, exec_context); - } - - Result DelegateApplyDense(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const { - return branch_mask_->ApplyDense(expr, input, exec_context); + Result DelegateApplyCond(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const { + return branch_mask_->DoApplyCond(expr, input, exec_context); } Result> DelegateGetSelectionVector() const { @@ -260,16 +163,10 @@ struct AllNullBodyMask : public NestedBodyMask { bool empty() const override { return true; } - Result ApplySparse(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const override { - DCHECK(false); - return Status::Invalid("AllNullBodyMask::ApplySparse should not be called"); - } - - Result ApplyDense(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const override { + Result ApplyCond(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const override { DCHECK(false); - return Status::Invalid("AllNullBodyMask::ApplyDense should not be called"); + return Status::Invalid("AllNullBodyMask::ApplyCond should not be called"); } Result> GetSelectionVector() const override { @@ -287,14 +184,9 @@ struct AllPassBodyMask : public NestedBodyMask { bool empty() const override { return false; } - Result ApplySparse(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const override { - return DelegateApplySparse(expr, input, exec_context); - } - - Result ApplyDense(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const override { - return DelegateApplyDense(expr, input, exec_context); + Result ApplyCond(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const override { + return DelegateApplyCond(expr, input, exec_context); } Result> GetSelectionVector() const override { @@ -311,16 +203,10 @@ struct AllFailBodyMask : public NestedBodyMask { bool empty() const override { return true; } - Result ApplySparse(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const override { - DCHECK(false); - return Status::Invalid("AllFailBodyMask::ApplySparse should not be called"); - } - - Result ApplyDense(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const override { + Result ApplyCond(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const override { DCHECK(false); - return Status::Invalid("AllFailBodyMask::ApplyDense should not be called"); + return Status::Invalid("AllFailBodyMask::ApplyCond should not be called"); } Result> GetSelectionVector() const override { @@ -333,16 +219,24 @@ struct AllFailBodyMask : public NestedBodyMask { } }; -Result TakeBySelectionVector(const ExecBatch& input, - const Datum& selection_vector, - ExecContext* exec_context) { - std::vector values(input.num_values()); - for (int i = 0; i < input.num_values(); ++i) { - ARROW_ASSIGN_OR_RAISE( - values[i], Take(input[i], selection_vector, TakeOptions{/*boundcheck=*/false}, - exec_context)); +Result> BranchMask::MakeBodyMaskFromDatum( + const Datum& datum, ExecContext* exec_context) const { + DCHECK(datum.type()->id() == Type::BOOL); + if (datum.is_scalar()) { + auto scalar = datum.scalar_as(); + if (!scalar.is_valid) { + return std::make_shared(shared_from_this()); + } else if (scalar.value) { + return std::make_shared(shared_from_this()); + } else { + return std::make_shared(shared_from_this()); + } + } + if (datum.is_array()) { + return MakeBodyMaskFromBitmap(datum.array_as(), exec_context); } - return ExecBatch::Make(std::move(values), selection_vector.length()); + DCHECK(datum.is_chunked_array()); + return MakeBodyMaskFromBitmap(datum.chunked_array(), exec_context); } struct ConditionalBranchMask : public BranchMask { @@ -352,41 +246,22 @@ struct ConditionalBranchMask : public BranchMask { bool empty() const override { return selection_vector_->length() == 0; } protected: - Result ApplySparse(const Expression& expr, const ExecBatch& input, + Result DoApplyCond(const Expression& expr, const ExecBatch& input, ExecContext* exec_context) const override { auto sparse_input = input; sparse_input.selection_vector = selection_vector_; return ExecuteScalarExpression(expr, sparse_input, exec_context); } - Result ApplyDense(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const override { - ARROW_ASSIGN_OR_RAISE( - auto dense_input, - TakeBySelectionVector(input, *selection_vector_->data(), exec_context)); - return ExecuteScalarExpression(expr, dense_input, exec_context); - } - Result> GetSelectionVector() const override { return selection_vector_; } - Result> MakeBodyMaskFromScalar( - const BooleanScalar& scalar, ExecContext* exec_context) const override; - - Result> MakeBodyMaskFromSparseBitmap( - const std::shared_ptr& bitmap, - ExecContext* exec_context) const override; - - Result> MakeBodyMaskFromDenseBitmap( + Result> MakeBodyMaskFromBitmap( const std::shared_ptr& bitmap, ExecContext* exec_context) const override; - Result> MakeBodyMaskFromSparseBitmap( - const std::shared_ptr& bitmap, - ExecContext* exec_context) const override; - - Result> MakeBodyMaskFromDenseBitmap( + Result> MakeBodyMaskFromBitmap( const std::shared_ptr& bitmap, ExecContext* exec_context) const override; @@ -402,20 +277,13 @@ struct ConditionalBodyMask : public BodyMask { bool empty() const override { return body_->length() == 0; } - Result ApplySparse(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const override { + Result ApplyCond(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const override { auto sparse_input = input; sparse_input.selection_vector = body_; return ExecuteScalarExpression(expr, sparse_input, exec_context); } - Result ApplyDense(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const override { - ARROW_ASSIGN_OR_RAISE(auto dense_input, - TakeBySelectionVector(input, *body_->data(), exec_context)); - return ExecuteScalarExpression(expr, dense_input, exec_context); - } - Result> GetSelectionVector() const override { return body_; } @@ -433,24 +301,7 @@ struct ConditionalBodyMask : public BodyMask { int64_t length_; }; -Result> BodyMaskFromScalar( - const BooleanScalar& scalar, std::shared_ptr branch_mask, - ExecContext* exec_context) { - if (!scalar.is_valid) { - return std::make_shared(std::move(branch_mask)); - } else if (scalar.value) { - return std::make_shared(std::move(branch_mask)); - } else { - return std::make_shared(std::move(branch_mask)); - } -} - -Result> AllPassBranchMask::MakeBodyMaskFromScalar( - const BooleanScalar& scalar, ExecContext* exec_context) const { - return BodyMaskFromScalar(scalar, shared_from_this(), exec_context); -} - -Result> AllPassBranchMask::MakeBodyMaskFromSparseBitmap( +Result> AllPassBranchMask::MakeBodyMaskFromBitmap( const std::shared_ptr& bitmap, ExecContext* exec_context) const { DCHECK_EQ(bitmap->length(), length_); @@ -480,12 +331,7 @@ Result> AllPassBranchMask::MakeBodyMaskFromSpars return std::make_shared(std::move(body), std::move(rest), length_); } -Result> AllPassBranchMask::MakeBodyMaskFromDenseBitmap( - const std::shared_ptr& bitmap, ExecContext* exec_context) const { - return MakeBodyMaskFromSparseBitmap(bitmap, exec_context); -} - -Result> AllPassBranchMask::MakeBodyMaskFromSparseBitmap( +Result> AllPassBranchMask::MakeBodyMaskFromBitmap( const std::shared_ptr& bitmap, ExecContext* exec_context) const { DCHECK_EQ(bitmap->length(), length_); @@ -518,18 +364,7 @@ Result> AllPassBranchMask::MakeBodyMaskFromSpars return std::make_shared(std::move(body), std::move(rest), length_); } -Result> AllPassBranchMask::MakeBodyMaskFromDenseBitmap( - const std::shared_ptr& bitmap, ExecContext* exec_context) const { - return MakeBodyMaskFromSparseBitmap(bitmap, exec_context); -} - -Result> ConditionalBranchMask::MakeBodyMaskFromScalar( - const BooleanScalar& scalar, ExecContext* exec_context) const { - return BodyMaskFromScalar(scalar, shared_from_this(), exec_context); -} - -Result> -ConditionalBranchMask::MakeBodyMaskFromSparseBitmap( +Result> ConditionalBranchMask::MakeBodyMaskFromBitmap( const std::shared_ptr& bitmap, ExecContext* exec_context) const { DCHECK_EQ(bitmap->length(), length_); @@ -556,35 +391,7 @@ ConditionalBranchMask::MakeBodyMaskFromSparseBitmap( return std::make_shared(std::move(body), std::move(rest), length_); } -Result> -ConditionalBranchMask::MakeBodyMaskFromDenseBitmap( - const std::shared_ptr& bitmap, ExecContext* exec_context) const { - DCHECK_EQ(bitmap->length(), selection_vector_->length()); - - Int32Builder body_builder(exec_context->memory_pool()); - Int32Builder rest_builder(exec_context->memory_pool()); - RETURN_NOT_OK(body_builder.Reserve(bitmap->true_count())); - RETURN_NOT_OK(rest_builder.Reserve(bitmap->false_count())); - - for (int64_t i = 0; i < selection_vector_->length(); ++i) { - if (!bitmap->IsNull(i)) { - if (bitmap->Value(i)) { - body_builder.UnsafeAppend(selection_vector_->indices()[i]); - } else { - rest_builder.UnsafeAppend(selection_vector_->indices()[i]); - } - } - } - - ARROW_ASSIGN_OR_RAISE(auto body_arr, body_builder.Finish()); - ARROW_ASSIGN_OR_RAISE(auto rest_arr, rest_builder.Finish()); - auto body = std::make_shared(body_arr->data()); - auto rest = std::make_shared(rest_arr->data()); - return std::make_shared(std::move(body), std::move(rest), length_); -} - -Result> -ConditionalBranchMask::MakeBodyMaskFromSparseBitmap( +Result> ConditionalBranchMask::MakeBodyMaskFromBitmap( const std::shared_ptr& bitmap, ExecContext* exec_context) const { DCHECK_EQ(bitmap->length(), length_); @@ -621,56 +428,15 @@ ConditionalBranchMask::MakeBodyMaskFromSparseBitmap( return std::make_shared(std::move(body), std::move(rest), length_); } -Result> -ConditionalBranchMask::MakeBodyMaskFromDenseBitmap( - const std::shared_ptr& bitmap, ExecContext* exec_context) const { - DCHECK_EQ(bitmap->length(), selection_vector_->length()); - - std::vector boolean_arrays(bitmap->num_chunks()); - std::transform(bitmap->chunks().begin(), bitmap->chunks().end(), boolean_arrays.begin(), - [](const auto& chunk) { - DCHECK_EQ(chunk->type()->id(), Type::BOOL); - return checked_cast(chunk.get()); - }); - - Int32Builder body_builder(exec_context->memory_pool()); - Int32Builder rest_builder(exec_context->memory_pool()); - RETURN_NOT_OK(body_builder.Reserve(selection_vector_->length())); - RETURN_NOT_OK(body_builder.Reserve(selection_vector_->length())); - - ChunkResolver resolver(bitmap->chunks()); - ChunkLocation location; - for (int64_t i = 0; i < selection_vector_->length(); ++i) { - location = resolver.ResolveWithHint(i, location); - if (boolean_arrays[location.chunk_index]->IsValid(location.index_in_chunk)) { - if (boolean_arrays[location.chunk_index]->Value(location.index_in_chunk)) { - body_builder.UnsafeAppend(selection_vector_->indices()[i]); - } else { - rest_builder.UnsafeAppend(selection_vector_->indices()[i]); - } - } - } - - ARROW_ASSIGN_OR_RAISE(auto body_arr, body_builder.Finish()); - ARROW_ASSIGN_OR_RAISE(auto rest_arr, rest_builder.Finish()); - auto body = std::make_shared(body_arr->data()); - auto rest = std::make_shared(rest_arr->data()); - return std::make_shared(std::move(body), std::move(rest), length_); -} - struct Branch { Expression cond; Expression body; }; -template struct ConditionalExecutor { ConditionalExecutor(const std::vector& branches, const TypeHolder& result_type) : branches(branches), result_type(result_type) {} - ARROW_DISALLOW_COPY_AND_ASSIGN(ConditionalExecutor); - ARROW_DEFAULT_MOVE_AND_ASSIGN(ConditionalExecutor); - Result Execute(const ExecBatch& input, ExecContext* exec_context) const&& { DCHECK(!branches.empty()); @@ -688,18 +454,16 @@ struct ConditionalExecutor { continue; } ARROW_ASSIGN_OR_RAISE(auto body_result, - static_cast(this)->ApplyBody( - body_mask, branch.body, input, exec_context)); + ApplyBody(body_mask, branch.body, input, exec_context)); DCHECK(body_result.type()->Equals(*result_type)); ARROW_ASSIGN_OR_RAISE(auto selection_vector, body_mask->GetSelectionVector()); results.Emplace(std::move(body_result), std::move(selection_vector)); ARROW_ASSIGN_OR_RAISE(branch_mask, body_mask->NextBranchMask()); } - return static_cast(this)->MultiplexBranchResults(input, results, - exec_context); + return MultiplexBranchResults(input, results, exec_context); } - protected: + private: struct BranchResults { void Reserve(int64_t size) { body_results_.reserve(size); @@ -726,7 +490,6 @@ struct ConditionalExecutor { std::vector> selection_vectors_; }; - private: Result> InitBranchMask( const ExecBatch& input, ExecContext* exec_context) const { if (input.selection_vector) { @@ -739,22 +502,13 @@ struct ConditionalExecutor { Result> ApplyCond( const std::shared_ptr& branch_mask, const Expression& cond, const ExecBatch& input, ExecContext* exec_context) const { - return branch_mask->ApplyCondSparse(cond, input, exec_context); - // return branch_mask->ApplyCondDense(cond, input, exec_context); + return branch_mask->ApplyCond(cond, input, exec_context); } - protected: - const std::vector& branches; - const TypeHolder& result_type; -}; - -struct SparseConditionalExecutor : public ConditionalExecutor { - using ConditionalExecutor::ConditionalExecutor; - Result ApplyBody(const std::shared_ptr& body_mask, const Expression& body, const ExecBatch& input, ExecContext* exec_context) const { - return body_mask->ApplySparse(body, input, exec_context); + return body_mask->ApplyCond(body, input, exec_context); } Result MultiplexBranchResults(const ExecBatch& input, @@ -785,7 +539,6 @@ struct SparseConditionalExecutor : public ConditionalExecutor ChooseIndices( const std::vector>& selection_vectors, int64_t length, ExecContext* exec_context) const { @@ -816,93 +569,10 @@ struct SparseConditionalExecutor : public ConditionalExecutor { - using ConditionalExecutor::ConditionalExecutor; - Result ApplyBody(const std::shared_ptr& body_mask, - const Expression& body, const ExecBatch& input, - ExecContext* exec_context) const { - return body_mask->ApplyDense(body, input, exec_context); - } - - Result MultiplexBranchResults(const ExecBatch& input, - const BranchResults& results, - ExecContext* exec_context) const { - if (results.empty()) { - return MakeArrayOfNull(result_type.GetSharedPtr(), input.length, - exec_context->memory_pool()); - } - - if (results.size() == 1) { - if (!results.selection_vectors()[0] || - results.selection_vectors()[0]->length() == input.length) { - return results.body_results()[0]; - } - } - - ARROW_ASSIGN_OR_RAISE( - auto body_results, - ToChunkedArray(results, - [&](const Datum& value, - const std::shared_ptr& selection_vector, - ArrayVector& chunks) -> Status { - DCHECK_NE(selection_vector, nullptr); - DCHECK_GT(selection_vector->length(), 0); - DCHECK(value.is_scalar() || value.is_arraylike()); - if (value.is_scalar()) { - ARROW_ASSIGN_OR_RAISE( - auto arr, MakeArrayFromScalar( - *value.scalar(), selection_vector->length(), - exec_context->memory_pool())); - chunks.push_back(std::move(arr)); - } else if (value.is_array() && value.length() > 0) { - chunks.push_back(value.make_array()); - } else { - DCHECK(value.is_chunked_array()); - for (const auto& chunk : value.chunked_array()->chunks()) { - if (chunk->length() > 0) { - chunks.push_back(chunk); - } - } - } - return Status::OK(); - })); - ARROW_ASSIGN_OR_RAISE( - auto indices, - ToChunkedArray( - results, - [](const Datum&, const std::shared_ptr& selection_vector, - ArrayVector& chunks) -> Status { - DCHECK_NE(selection_vector, nullptr); - DCHECK_GT(selection_vector->length(), 0); - chunks.push_back(MakeArray(selection_vector->data())); - return Status::OK(); - })); - ARROW_ASSIGN_OR_RAISE( - auto result, - Scatter(body_results, indices, ScatterOptions{/*max_index=*/input.length - 1})); - DCHECK(result.is_arraylike()); - if (result.is_chunked_array() && result.chunked_array()->num_chunks() == 1) { - return result.chunked_array()->chunk(0); - } else { - return result; - } - } - - private: - template - Result> ToChunkedArray(const BranchResults& results, - ChunkFunc&& chunk_func) const { - ArrayVector chunks; - chunks.reserve(results.size()); - for (size_t i = 0; i < results.size(); ++i) { - RETURN_NOT_OK( - chunk_func(results.body_results()[i], results.selection_vectors()[i], chunks)); - } - return std::make_shared(std::move(chunks)); - } + protected: + const std::vector& branches; + const TypeHolder& result_type; }; class IfElseSpecialExec : public SpecialExec { @@ -946,8 +616,7 @@ class IfElseSpecialExec : public SpecialExec { Result Execute(const ExecBatch& input, ExecContext* exec_context) const override { - return SparseConditionalExecutor(branches, type).Execute(input, exec_context); - // return DenseConditionalExecutor(branches, type).Execute(input, exec_context); + return ConditionalExecutor(branches, type).Execute(input, exec_context); } private: From e73056077665231794cee774f2c98d6556a23978 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Mon, 4 Aug 2025 13:57:21 +0800 Subject: [PATCH 06/71] Remove selection vector span from exec span, and let kernel have a designated signature --- cpp/src/arrow/compute/exec.cc | 72 ++++++++++------------ cpp/src/arrow/compute/exec.h | 5 -- cpp/src/arrow/compute/exec_internal.h | 6 +- cpp/src/arrow/compute/kernel.h | 3 +- cpp/src/arrow/compute/special_form_test.cc | 27 ++++---- 5 files changed, 51 insertions(+), 62 deletions(-) diff --git a/cpp/src/arrow/compute/exec.cc b/cpp/src/arrow/compute/exec.cc index 064662525763..cf7ffea5a743 100644 --- a/cpp/src/arrow/compute/exec.cc +++ b/cpp/src/arrow/compute/exec.cc @@ -406,7 +406,7 @@ int64_t ExecSpanIterator::GetNextChunkSpan(int64_t iteration_size, ExecSpan* spa return iteration_size; } -bool ExecSpanIterator::Next(ExecSpan* span) { +bool ExecSpanIterator::Next(ExecSpan* span, SelectionVectorSpan* selection_span) { if (!initialized_) { span->length = 0; @@ -445,14 +445,15 @@ bool ExecSpanIterator::Next(ExecSpan* span) { PromoteExecSpanScalars(span); } else { if (selection_vector_) { + DCHECK_NE(selection_span, nullptr); if (have_chunked_arrays_) { - span->selection_vector = SelectionVectorSpan(selection_vector_->indices(), 0); + *selection_span = SelectionVectorSpan(selection_vector_->indices(), 0); } else { - span->selection_vector = SelectionVectorSpan(selection_vector_->indices(), - selection_vector_->length()); + *selection_span = SelectionVectorSpan(selection_vector_->indices(), + selection_vector_->length()); } } else { - span->selection_vector = std::nullopt; + DCHECK_EQ(selection_span, nullptr); } } @@ -467,8 +468,8 @@ bool ExecSpanIterator::Next(ExecSpan* span) { if (have_chunked_arrays_) { iteration_size = GetNextChunkSpan(iteration_size, span); if (selection_vector_) { - auto indices_begin = - span->selection_vector->indices() + span->selection_vector->length(); + DCHECK_NE(selection_span, nullptr); + auto indices_begin = selection_vector_->indices() + selection_position_; auto indices_end = selection_vector_->indices() + selection_vector_->length(); DCHECK_LE(indices_begin, indices_end); auto chunk_row_id_end = position_ + iteration_size; @@ -477,7 +478,8 @@ bool ExecSpanIterator::Next(ExecSpan* span) { *(indices_begin + num_indices) < chunk_row_id_end) { ++num_indices; } - span->selection_vector->SetSlice(span->selection_vector->length(), num_indices); + selection_span->SetSlice(selection_position_, num_indices); + selection_position_ += num_indices; } } @@ -814,6 +816,7 @@ class ScalarExecutor : public KernelExecutorImpl { ARROW_ASSIGN_OR_RAISE(std::shared_ptr result, MakeArrayOfNull(output_type_.GetSharedPtr(), /*length=*/0, exec_context()->memory_pool())); + RETURN_NOT_OK(span_iterator_.Init(batch, exec_context()->exec_chunksize())); return EmitResult(result->data(), listener); } @@ -920,6 +923,11 @@ class ScalarExecutor : public KernelExecutorImpl { // eventually skip the creation of ArrayData altogether std::shared_ptr preallocation; ExecSpan input; + SelectionVectorSpan selection; + SelectionVectorSpan* selection_ptr = nullptr; + if (span_iterator_.have_selection_vector()) { + selection_ptr = &selection; + } ExecResult output; ArraySpan* output_span = output.array_span_mutable(); @@ -931,10 +939,10 @@ class ScalarExecutor : public KernelExecutorImpl { output_span->SetMembers(*preallocation); output_span->offset = 0; int64_t result_offset = 0; - while (span_iterator_.Next(&input)) { + while (span_iterator_.Next(&input, selection_ptr)) { // Set absolute output span position and length output_span->SetSlice(result_offset, input.length); - RETURN_NOT_OK(ExecuteSingleSpan(input, &output)); + RETURN_NOT_OK(ExecuteSingleSpan(input, selection_ptr, &output)); result_offset = span_iterator_.position(); } @@ -944,10 +952,10 @@ class ScalarExecutor : public KernelExecutorImpl { // Fully preallocating, but not contiguously // We preallocate (maybe) only for the output of processing the current // chunk - while (span_iterator_.Next(&input)) { + while (span_iterator_.Next(&input, selection_ptr)) { ARROW_ASSIGN_OR_RAISE(preallocation, PrepareOutput(input.length)); output_span->SetMembers(*preallocation); - RETURN_NOT_OK(ExecuteSingleSpan(input, &output)); + RETURN_NOT_OK(ExecuteSingleSpan(input, selection_ptr, &output)); // Emit the result for this chunk RETURN_NOT_OK(EmitResult(std::move(preallocation), listener)); } @@ -955,7 +963,8 @@ class ScalarExecutor : public KernelExecutorImpl { } } - Status ExecuteSingleSpan(const ExecSpan& input, ExecResult* out) { + Status ExecuteSingleSpan(const ExecSpan& input, const SelectionVectorSpan* selection, + ExecResult* out) { ArraySpan* result_span = out->array_span_mutable(); if (output_type_.type->id() == Type::NA) { result_span->null_count = result_span->length; @@ -966,7 +975,7 @@ class ScalarExecutor : public KernelExecutorImpl { } else if (kernel_->null_handling == NullHandling::OUTPUT_NOT_NULL) { result_span->null_count = 0; } - RETURN_NOT_OK(InvokeKernel(input, out)); + RETURN_NOT_OK(InvokeKernel(input, selection, out)); // Output type didn't change DCHECK(out->is_array_span()); return Status::OK(); @@ -981,8 +990,13 @@ class ScalarExecutor : public KernelExecutorImpl { // We will eventually delete the Scalar output path per // ARROW-16757. ExecSpan input; + SelectionVectorSpan selection; + SelectionVectorSpan* selection_ptr = nullptr; + if (span_iterator_.have_selection_vector()) { + selection_ptr = &selection; + } ExecResult output; - while (span_iterator_.Next(&input)) { + while (span_iterator_.Next(&input, selection_ptr)) { ARROW_ASSIGN_OR_RAISE(output.value, PrepareOutput(input.length)); DCHECK(output.is_array_data()); @@ -995,7 +1009,7 @@ class ScalarExecutor : public KernelExecutorImpl { out_arr->null_count = 0; } - RETURN_NOT_OK(InvokeKernel(input, &output)); + RETURN_NOT_OK(InvokeKernel(input, selection_ptr, &output)); // Output type didn't change DCHECK(output.is_array_data()); @@ -1061,10 +1075,11 @@ class ScalarExecutor : public KernelExecutorImpl { return Status::OK(); } - Status InvokeKernel(const ExecSpan& input, ExecResult* out) { - if (input.selection_vector) { + Status InvokeKernel(const ExecSpan& input, const SelectionVectorSpan* selection, + ExecResult* out) { + if (selection) { DCHECK_NE(kernel_->selective_exec, nullptr); - return kernel_->selective_exec(kernel_ctx_, input, out); + return kernel_->selective_exec(kernel_ctx_, input, *selection, out); } return kernel_->exec(kernel_ctx_, input, out); } @@ -1438,25 +1453,6 @@ SelectionVector::SelectionVector(std::shared_ptr data) SelectionVector::SelectionVector(const Array& arr) : SelectionVector(arr.data()) {} -Result> SelectionVector::FromMask( - const BooleanArray& arr, MemoryPool* pool) { - Int32Builder builder(pool); - RETURN_NOT_OK(builder.Reserve(arr.true_count())); - ArraySpan span(*arr.data()); - int32_t i = 0; - VisitArraySpanInline( - span, - [&](bool mask) { - if (mask) { - builder.UnsafeAppend(i); - } - ++i; - }, - [&]() { ++i; }); - ARROW_ASSIGN_OR_RAISE(auto indices, builder.Finish()); - return std::make_shared(indices->data()); -} - int64_t SelectionVector::length() const { return data_->length; } void SelectionVectorSpan::SetSlice(int64_t offset, int64_t length) { diff --git a/cpp/src/arrow/compute/exec.h b/cpp/src/arrow/compute/exec.h index f2d15b5258f1..fbf2af64273e 100644 --- a/cpp/src/arrow/compute/exec.h +++ b/cpp/src/arrow/compute/exec.h @@ -140,10 +140,6 @@ class ARROW_EXPORT SelectionVector { explicit SelectionVector(const Array& arr); - /// \brief Create SelectionVector from boolean mask - static Result> FromMask( - const BooleanArray& arr, MemoryPool* pool = default_memory_pool()); - std::shared_ptr data() const { return data_; } const int32_t* indices() const { return indices_; } int64_t length() const; @@ -447,7 +443,6 @@ struct ARROW_EXPORT ExecSpan { int64_t length = 0; std::vector values; - std::optional selection_vector; }; /// \defgroup compute-call-function One-shot calls to compute functions diff --git a/cpp/src/arrow/compute/exec_internal.h b/cpp/src/arrow/compute/exec_internal.h index b92cd257f586..cc9b6e278f21 100644 --- a/cpp/src/arrow/compute/exec_internal.h +++ b/cpp/src/arrow/compute/exec_internal.h @@ -42,8 +42,6 @@ namespace detail { /// \brief Break std::vector into a sequence of non-owning /// ExecSpan for kernel execution. The lifetime of the Datum vector /// must be longer than the lifetime of this object -// TODO: Can this struct sense the presence of selection vector and not split ExecBatch if -// so? class ARROW_EXPORT ExecSpanIterator { public: ExecSpanIterator() = default; @@ -68,12 +66,13 @@ class ARROW_EXPORT ExecSpanIterator { /// with a blank ExecSpan after the first iteration, it will not /// work correctly (maybe we will change this later). Return false /// if the iteration is exhausted - bool Next(ExecSpan* span); + bool Next(ExecSpan* span, SelectionVectorSpan* selection_span = NULLPTR); int64_t length() const { return length_; } int64_t position() const { return position_; } bool have_all_scalars() const { return have_all_scalars_; } + bool have_selection_vector() const { return selection_vector_ != NULLPTR; } private: ExecSpanIterator(const std::vector& args, int64_t length, int64_t max_chunksize); @@ -96,6 +95,7 @@ class ARROW_EXPORT ExecSpanIterator { std::vector value_offsets_; int64_t position_ = 0; int64_t length_ = 0; + int64_t selection_position_ = 0; int64_t max_chunksize_; }; diff --git a/cpp/src/arrow/compute/kernel.h b/cpp/src/arrow/compute/kernel.h index a4e319361bec..928a0d0d111a 100644 --- a/cpp/src/arrow/compute/kernel.h +++ b/cpp/src/arrow/compute/kernel.h @@ -555,7 +555,8 @@ struct ARROW_EXPORT Kernel { /// employed this may not be possible. using ArrayKernelExec = Status (*)(KernelContext*, const ExecSpan&, ExecResult*); -using ArrayKernelSelectiveExec = Status (*)(KernelContext*, const ExecSpan&, ExecResult*); +using ArrayKernelSelectiveExec = Status (*)(KernelContext*, const ExecSpan&, + const SelectionVectorSpan&, ExecResult*); /// \brief Kernel data structure for implementations of ScalarFunction. In /// addition to the members found in Kernel, contains the null handling diff --git a/cpp/src/arrow/compute/special_form_test.cc b/cpp/src/arrow/compute/special_form_test.cc index e62e1fb54869..3a2a48131a54 100644 --- a/cpp/src/arrow/compute/special_form_test.cc +++ b/cpp/src/arrow/compute/special_form_test.cc @@ -190,18 +190,15 @@ class IfElseSpecialFormTest : public ::testing::Test { return Status::OK(); } - template - static Status AssertSelectionVectorExec(KernelContext* kernel_ctx, const ExecSpan& span, - ExecResult* out) { - if constexpr (sv_existence) { - if (!span.selection_vector) { - return Status::Invalid("There is no selection vector"); - } - } else { - if (span.selection_vector) { - return Status::Invalid("There is a selection vector"); - } - } + static Status AssertSelectionVectorNotExistExec(KernelContext* kernel_ctx, + const ExecSpan& span, ExecResult* out) { + return IdentityExec(kernel_ctx, span, out); + } + + static Status AssertSelectionVectorExistExec(KernelContext* kernel_ctx, + const ExecSpan& span, + const SelectionVectorSpan& selection, + ExecResult* out) { return IdentityExec(kernel_ctx, span, out); } @@ -249,10 +246,10 @@ class IfElseSpecialFormTest : public ::testing::Test { auto func = std::make_shared(name, Arity::Unary(), FunctionDoc::Empty()); - ArrayKernelExec exec = AssertSelectionVectorExec; + ArrayKernelExec exec = AssertSelectionVectorNotExistExec; ArrayKernelSelectiveExec selective_exec = nullptr; if (sv_existence) { - selective_exec = AssertSelectionVectorExec; + selective_exec = AssertSelectionVectorExistExec; } ScalarKernel kernel({InputType::Any()}, internal::FirstType, std::move(exec), /*init=*/nullptr, std::move(selective_exec)); @@ -587,7 +584,7 @@ TEST_F(IfElseSpecialFormTest, SelectionVectorExistenceExecChunkSize) { num_rows); { ARROW_SCOPED_TRACE("exec_chunksize >= batch_size"); - for (auto chunksize : {num_rows, num_rows + 1}) { + for (auto chunksize : {num_rows / 2, num_rows, num_rows + 1}) { ARROW_SCOPED_TRACE("exec_chunksize: " + std::to_string(chunksize)); exec_context.set_exec_chunksize(chunksize); { From 976efaa6e00fd47cda371f4d7907da05427a5bc5 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Mon, 4 Aug 2025 13:59:02 +0800 Subject: [PATCH 07/71] Fix selection vector span to forward for split exec span --- cpp/src/arrow/compute/exec.cc | 43 ++++++++----- cpp/src/arrow/compute/exec.h | 11 ++-- cpp/src/arrow/compute/special_form_test.cc | 72 +++++++++++++--------- 3 files changed, 75 insertions(+), 51 deletions(-) diff --git a/cpp/src/arrow/compute/exec.cc b/cpp/src/arrow/compute/exec.cc index cf7ffea5a743..fc225e636b42 100644 --- a/cpp/src/arrow/compute/exec.cc +++ b/cpp/src/arrow/compute/exec.cc @@ -467,20 +467,6 @@ bool ExecSpanIterator::Next(ExecSpan* span, SelectionVectorSpan* selection_span) int64_t iteration_size = std::min(length_ - position_, max_chunksize_); if (have_chunked_arrays_) { iteration_size = GetNextChunkSpan(iteration_size, span); - if (selection_vector_) { - DCHECK_NE(selection_span, nullptr); - auto indices_begin = selection_vector_->indices() + selection_position_; - auto indices_end = selection_vector_->indices() + selection_vector_->length(); - DCHECK_LE(indices_begin, indices_end); - auto chunk_row_id_end = position_ + iteration_size; - int64_t num_indices = 0; - while (indices_begin + num_indices < indices_end && - *(indices_begin + num_indices) < chunk_row_id_end) { - ++num_indices; - } - selection_span->SetSlice(selection_position_, num_indices); - selection_position_ += num_indices; - } } // Now, adjust the span @@ -494,6 +480,23 @@ bool ExecSpanIterator::Next(ExecSpan* span, SelectionVectorSpan* selection_span) } } + // Then the selection span + if (selection_vector_) { + DCHECK_NE(selection_span, nullptr); + auto indices_begin = selection_vector_->indices() + selection_position_; + auto indices_end = selection_vector_->indices() + selection_vector_->length(); + DCHECK_LE(indices_begin, indices_end); + auto chunk_row_id_end = position_ + iteration_size; + int64_t num_indices = 0; + while (indices_begin + num_indices < indices_end && + *(indices_begin + num_indices) < chunk_row_id_end) { + ++num_indices; + } + selection_span->SetSlice(selection_position_, num_indices, + static_cast(position_)); + selection_position_ += num_indices; + } + position_ += iteration_size; DCHECK_LE(position_, length_); return true; @@ -1455,10 +1458,18 @@ SelectionVector::SelectionVector(const Array& arr) : SelectionVector(arr.data()) int64_t SelectionVector::length() const { return data_->length; } -void SelectionVectorSpan::SetSlice(int64_t offset, int64_t length) { +void SelectionVectorSpan::SetSlice(int64_t offset, int64_t length, int32_t backstep) { DCHECK_NE(indices_, nullptr); - offset_ += offset; + offset_ = offset; length_ = length; + backstep_ = backstep; +} + +int32_t SelectionVectorSpan::operator[](int64_t i) const { + DCHECK_GE(i, 0); + DCHECK_LT(i, length_); + DCHECK_GE(indices_[i + offset_], backstep_); + return indices_[i + offset_] - backstep_; } Result CallFunction(const std::string& func_name, const std::vector& args, diff --git a/cpp/src/arrow/compute/exec.h b/cpp/src/arrow/compute/exec.h index fbf2af64273e..393cc0b29961 100644 --- a/cpp/src/arrow/compute/exec.h +++ b/cpp/src/arrow/compute/exec.h @@ -152,21 +152,20 @@ class ARROW_EXPORT SelectionVector { class ARROW_EXPORT SelectionVectorSpan { public: explicit SelectionVectorSpan(const int32_t* indices = NULLPTR, int64_t length = 0, - int64_t offset = 0) - : indices_(indices), length_(length), offset_(offset) {} + int64_t offset = 0, int32_t backstep = 0) + : indices_(indices), length_(length), offset_(offset), backstep_(backstep) {} - void SetSlice(int64_t offset, int64_t length); + void SetSlice(int64_t offset, int64_t length, int32_t backward = 0); - const int32_t* indices() const { return indices_ + offset_; } + int32_t operator[](int64_t i) const; int64_t length() const { return length_; } - int64_t offset() const { return offset_; } - private: const int32_t* indices_; int64_t length_; int64_t offset_; + int32_t backstep_; }; /// An index to represent that a batch does not belong to an ordered stream diff --git a/cpp/src/arrow/compute/special_form_test.cc b/cpp/src/arrow/compute/special_form_test.cc index 3a2a48131a54..38cdf78c4a08 100644 --- a/cpp/src/arrow/compute/special_form_test.cc +++ b/cpp/src/arrow/compute/special_form_test.cc @@ -199,6 +199,10 @@ class IfElseSpecialFormTest : public ::testing::Test { const ExecSpan& span, const SelectionVectorSpan& selection, ExecResult* out) { + for (int32_t i = 0; i < selection.length(); ++i) { + EXPECT_GE(selection[i], 0); + EXPECT_LT(selection[i], span.length); + } return IdentityExec(kernel_ctx, span, out); } @@ -579,12 +583,12 @@ TEST_F(IfElseSpecialFormTest, SelectionVectorExistenceExecChunkSize) { auto batch = ExecBatch( { *ArrayFromJSON(boolean(), "[true, true, true, true, true, true, true, false]"), - *ArrayFromJSON(int32(), "[42, 42, 42, 42, 42, 42, 42, 42]"), + *ArrayFromJSON(int32(), "[1, 2, 3, 4, 5, 6, 7, 8]"), }, num_rows); { ARROW_SCOPED_TRACE("exec_chunksize >= batch_size"); - for (auto chunksize : {num_rows / 2, num_rows, num_rows + 1}) { + for (auto chunksize : {num_rows, num_rows + 1}) { ARROW_SCOPED_TRACE("exec_chunksize: " + std::to_string(chunksize)); exec_context.set_exec_chunksize(chunksize); { @@ -618,33 +622,43 @@ TEST_F(IfElseSpecialFormTest, SelectionVectorExistenceExecChunkSize) { } { ARROW_SCOPED_TRACE("exec_chunksize < batch_size"); - exec_context.set_exec_chunksize(num_rows - 1); - { - ARROW_SCOPED_TRACE("all literal bodies"); - auto if_else_sp = - if_else_special(b, assert_sv_empty(literal(1)), assert_sv_empty(literal(0))); - ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); - ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); - } - { - ARROW_SCOPED_TRACE("array true body"); - auto if_else_sp = - if_else_special(b, assert_sv_exist(i), assert_sv_empty(literal(0))); - ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); - ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); - } - { - ARROW_SCOPED_TRACE("array false body"); - auto if_else_sp = - if_else_special(b, assert_sv_empty(literal(1)), assert_sv_exist(i)); - ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); - ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); - } - { - ARROW_SCOPED_TRACE("all array bodies"); - auto if_else_sp = if_else_special(b, assert_sv_exist(i), assert_sv_exist(i)); - ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); - ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); + for (auto chunksize : {num_rows / 2, num_rows - 1}) { + ARROW_SCOPED_TRACE("exec_chunksize: " + std::to_string(chunksize)); + exec_context.set_exec_chunksize(num_rows - 1); + { + ARROW_SCOPED_TRACE("all literal bodies"); + auto if_else_sp = + if_else_special(b, assert_sv_empty(literal(1)), assert_sv_empty(literal(0))); + ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); + ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); + } + { + ARROW_SCOPED_TRACE("array true body"); + auto if_else_sp = + if_else_special(b, assert_sv_exist(i), assert_sv_empty(literal(0))); + // ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); + // ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); + } + { + ARROW_SCOPED_TRACE("array true body"); + auto if_else_sp = + if_else_special(b, assert_sv_empty(i), assert_sv_empty(literal(0))); + ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); + ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); + } + { + ARROW_SCOPED_TRACE("array false body"); + auto if_else_sp = + if_else_special(b, assert_sv_empty(literal(1)), assert_sv_exist(i)); + ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); + ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); + } + { + ARROW_SCOPED_TRACE("all array bodies"); + auto if_else_sp = if_else_special(b, assert_sv_exist(i), assert_sv_exist(i)); + ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); + ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); + } } } { From f1399e6f960c38b8fb1eb29273238a8cc67d07b0 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Mon, 4 Aug 2025 17:07:35 +0800 Subject: [PATCH 08/71] Rename --- cpp/src/arrow/compute/exec.cc | 1 - cpp/src/arrow/compute/special_form.cc | 25 ++++++++++++------------- 2 files changed, 12 insertions(+), 14 deletions(-) diff --git a/cpp/src/arrow/compute/exec.cc b/cpp/src/arrow/compute/exec.cc index fc225e636b42..55f3dc9c18fd 100644 --- a/cpp/src/arrow/compute/exec.cc +++ b/cpp/src/arrow/compute/exec.cc @@ -27,7 +27,6 @@ #include "arrow/array/array_base.h" #include "arrow/array/array_primitive.h" -#include "arrow/array/builder_primitive.h" #include "arrow/array/data.h" #include "arrow/array/util.h" #include "arrow/buffer.h" diff --git a/cpp/src/arrow/compute/special_form.cc b/cpp/src/arrow/compute/special_form.cc index b1f0ae9037ca..b99f73361ba9 100644 --- a/cpp/src/arrow/compute/special_form.cc +++ b/cpp/src/arrow/compute/special_form.cc @@ -60,7 +60,7 @@ struct BranchMask : public std::enable_shared_from_this { Result> MakeBodyMaskFromDatum( const Datum& datum, ExecContext* exec_context) const; - friend struct NestedBodyMask; + friend struct DelegateBodyMask; }; struct BodyMask : public std::enable_shared_from_this { @@ -140,8 +140,8 @@ struct AllFailBranchMask : public BranchMask { } }; -struct NestedBodyMask : public BodyMask { - explicit NestedBodyMask(std::shared_ptr branch_mask) +struct DelegateBodyMask : public BodyMask { + explicit DelegateBodyMask(std::shared_ptr branch_mask) : branch_mask_(std::move(branch_mask)) {} protected: @@ -158,8 +158,8 @@ struct NestedBodyMask : public BodyMask { std::shared_ptr branch_mask_; }; -struct AllNullBodyMask : public NestedBodyMask { - using NestedBodyMask::NestedBodyMask; +struct AllNullBodyMask : public DelegateBodyMask { + using DelegateBodyMask::DelegateBodyMask; bool empty() const override { return true; } @@ -179,8 +179,8 @@ struct AllNullBodyMask : public NestedBodyMask { } }; -struct AllPassBodyMask : public NestedBodyMask { - using NestedBodyMask::NestedBodyMask; +struct AllPassBodyMask : public DelegateBodyMask { + using DelegateBodyMask::DelegateBodyMask; bool empty() const override { return false; } @@ -198,8 +198,8 @@ struct AllPassBodyMask : public NestedBodyMask { } }; -struct AllFailBodyMask : public NestedBodyMask { - using NestedBodyMask::NestedBodyMask; +struct AllFailBodyMask : public DelegateBodyMask { + using DelegateBodyMask::DelegateBodyMask; bool empty() const override { return true; } @@ -460,7 +460,7 @@ struct ConditionalExecutor { results.Emplace(std::move(body_result), std::move(selection_vector)); ARROW_ASSIGN_OR_RAISE(branch_mask, body_mask->NextBranchMask()); } - return MultiplexBranchResults(input, results, exec_context); + return MultiplexResults(input, results, exec_context); } private: @@ -511,9 +511,8 @@ struct ConditionalExecutor { return body_mask->ApplyCond(body, input, exec_context); } - Result MultiplexBranchResults(const ExecBatch& input, - const BranchResults& results, - ExecContext* exec_context) const { + Result MultiplexResults(const ExecBatch& input, const BranchResults& results, + ExecContext* exec_context) const { if (results.empty()) { return MakeArrayOfNull(result_type.GetSharedPtr(), input.length, exec_context->memory_pool()); From 217493f55bd697ee26e5ebbcce5e85f7e2eaf486 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Thu, 7 Aug 2025 00:01:53 +0800 Subject: [PATCH 09/71] Deep refine special form and executor interfaces --- cpp/src/arrow/compute/expression.cc | 197 ++++++++------------ cpp/src/arrow/compute/expression.h | 2 +- cpp/src/arrow/compute/expression_internal.h | 7 + cpp/src/arrow/compute/special_form.cc | 143 ++++++++------ cpp/src/arrow/compute/special_form.h | 25 +-- cpp/src/arrow/compute/special_form_test.cc | 2 +- cpp/src/arrow/compute/type_fwd.h | 2 +- 7 files changed, 188 insertions(+), 190 deletions(-) diff --git a/cpp/src/arrow/compute/expression.cc b/cpp/src/arrow/compute/expression.cc index f44f474b3f00..6b7defa28923 100644 --- a/cpp/src/arrow/compute/expression.cc +++ b/cpp/src/arrow/compute/expression.cc @@ -61,7 +61,7 @@ void Expression::Call::ComputeHash() { } void Expression::Special::ComputeHash() { - hash = std::hash{}(special_form->name); + hash = std::hash{}(special_form->name()); for (const auto& arg : arguments) { arrow::internal::hash_combine(hash, arg.hash()); } @@ -198,7 +198,7 @@ std::string Expression::ToString() const { } if (auto sp = special()) { - std::string out = sp->special_form->name + "_special("; + std::string out = sp->special_form->name() + "_special("; for (const auto& arg : sp->arguments) { out += arg.ToString() + ", "; } @@ -571,91 +571,16 @@ TypeHolder SmallestTypeFor(const arrow::Datum& value) { } } -inline std::vector GetTypesWithSmallestLiteralRepresentation( - const std::vector& exprs) { - std::vector types(exprs.size()); - for (size_t i = 0; i < exprs.size(); ++i) { - DCHECK(exprs[i].IsBound()); - if (const Datum* literal = exprs[i].literal()) { - if (literal->is_scalar()) { - types[i] = SmallestTypeFor(*literal); - } - } else { - types[i] = exprs[i].type(); - } - } - return types; -} - -Result BindNonRecursive(Expression::Call call, bool insert_implicit_casts, - compute::ExecContext* exec_context); -Result BindNonRecursive(Expression::Special special, - bool insert_implicit_casts, - compute::ExecContext* exec_context); - -template -Status DispatchForBind(DispatchExactFn&& dispatch_exact, DispatchBestFn&& dispatch_best, - FinishBind&& finish_bind, std::vector& arguments, - bool insert_implicit_casts, ExecContext* exec_context) { - DCHECK(std::all_of(arguments.begin(), arguments.end(), - [](const Expression& argument) { return argument.IsBound(); })); - - std::vector types = GetTypes(arguments); - - // First try and bind exactly - auto maybe_exact_match = dispatch_exact(types); - if (maybe_exact_match.ok()) { - auto output = std::move(*maybe_exact_match); - if (finish_bind(types, std::move(output)).ok()) { - return Status::OK(); - } - } - - if (!insert_implicit_casts) { - return maybe_exact_match.status(); - } - - // If exact binding fails, and we are allowed to cast, then prefer casting literals - // first. Since DispatchBest generally prefers up-casting the best way to do this is - // first down-cast the literals as much as possible - types = GetTypesWithSmallestLiteralRepresentation(arguments); - ARROW_ASSIGN_OR_RAISE(auto output, dispatch_best(&types)); - // ARROW_ASSIGN_OR_RAISE(call.kernel, call.function->DispatchBest(&types)); - - for (size_t i = 0; i < types.size(); ++i) { - if (types[i] == arguments[i].type()) continue; - - if (const Datum* lit = arguments[i].literal()) { - ARROW_ASSIGN_OR_RAISE(Datum new_lit, compute::Cast(*lit, types[i].GetSharedPtr())); - arguments[i] = literal(std::move(new_lit)); - continue; - } - - // construct an implicit cast Expression with which to replace this argument - Expression::Call implicit_cast; - implicit_cast.function_name = "cast"; - implicit_cast.arguments = {std::move(arguments[i])}; - - // TODO(wesm): Use TypeHolder in options - implicit_cast.options = std::make_shared( - compute::CastOptions::Safe(types[i].GetSharedPtr())); - - ARROW_ASSIGN_OR_RAISE( - arguments[i], BindNonRecursive(std::move(implicit_cast), - /*insert_implicit_casts=*/false, exec_context)); - } - - return finish_bind(types, std::move(output)); -} - // Produce a bound Expression from unbound Call and bound arguments. Result BindNonRecursive(Expression::Call call, bool insert_implicit_casts, compute::ExecContext* exec_context) { - ARROW_ASSIGN_OR_RAISE(call.function, GetFunction(call, exec_context)); + DCHECK(std::all_of(call.arguments.begin(), call.arguments.end(), + [](const Expression& argument) { return argument.IsBound(); })); - auto FinishBind = [&](const std::vector& types, const Kernel* kernel) { - call.kernel = kernel; + std::vector types = GetTypes(call.arguments); + ARROW_ASSIGN_OR_RAISE(call.function, GetFunction(call, exec_context)); + auto FinishBind = [&] { compute::KernelContext kernel_context(exec_context, call.kernel); if (call.kernel->init) { const FunctionOptions* options = @@ -672,45 +597,29 @@ Result BindNonRecursive(Expression::Call call, bool insert_implicit_ return Status::OK(); }; - RETURN_NOT_OK(DispatchForBind( - [&](const std::vector& types) -> Result { - return call.function->DispatchExact(types); - }, - [&](std::vector* types) -> Result { - return call.function->DispatchBest(types); - }, - FinishBind, call.arguments, insert_implicit_casts, exec_context)); - - return Expression(std::move(call)); -} - -// Produce a bound Expression from unbound Special and bound arguments. -Result BindNonRecursive(Expression::Special special, - bool insert_implicit_casts, - compute::ExecContext* exec_context) { - DCHECK(std::all_of(special.arguments.begin(), special.arguments.end(), - [](const Expression& argument) { return argument.IsBound(); })); - - auto FinishBind = [&](const std::vector& types, - std::unique_ptr&& special_exec) { - special.special_exec = std::move(special_exec); + // First try and bind exactly + Result maybe_exact_match = call.function->DispatchExact(types); + if (maybe_exact_match.ok()) { + call.kernel = *maybe_exact_match; + if (FinishBind().ok()) { + return Expression(std::move(call)); + } + } - ARROW_ASSIGN_OR_RAISE(special.type, - special.special_exec->Bind(special.arguments, exec_context)); + if (!insert_implicit_casts) { + return maybe_exact_match.status(); + } - return Status::OK(); - }; + // If exact binding fails, and we are allowed to cast, then prefer casting literals + // first. Since DispatchBest generally prefers up-casting the best way to do this is + // first down-cast the literals as much as possible + types = GetTypesWithSmallestLiteralRepresentation(call.arguments); + ARROW_ASSIGN_OR_RAISE(call.kernel, call.function->DispatchBest(&types)); - RETURN_NOT_OK(DispatchForBind( - [&](const std::vector& types) -> Result> { - return special.special_form->DispatchExact(types, exec_context); - }, - [&](std::vector* types) -> Result> { - return special.special_form->DispatchBest(types, exec_context); - }, - FinishBind, special.arguments, insert_implicit_casts, exec_context)); + ARROW_RETURN_NOT_OK(ImplicitCastArguments(call.arguments, types, exec_context)); - return Expression(std::move(special)); + RETURN_NOT_OK(FinishBind()); + return Expression(std::move(call)); } template @@ -739,7 +648,10 @@ Result BindImpl(Expression expr, const TypeOrSchema& in, for (auto& argument : special.arguments) { ARROW_ASSIGN_OR_RAISE(argument, BindImpl(std::move(argument), in, exec_context)); } - return BindNonRecursive(special, /*insert_implicit_casts=*/true, exec_context); + ARROW_ASSIGN_OR_RAISE(special.special_executor, + special.special_form->Bind(special.arguments, exec_context)); + special.type = special.special_executor->out_type(); + return Expression(std::move(special)); } auto call = *CallNotNull(expr); @@ -751,6 +663,53 @@ Result BindImpl(Expression expr, const TypeOrSchema& in, } // namespace +std::vector GetTypesWithSmallestLiteralRepresentation( + const std::vector& exprs) { + std::vector types(exprs.size()); + for (size_t i = 0; i < exprs.size(); ++i) { + DCHECK(exprs[i].IsBound()); + if (const Datum* literal = exprs[i].literal()) { + if (literal->is_scalar()) { + types[i] = SmallestTypeFor(*literal); + } + } else { + types[i] = exprs[i].type(); + } + } + return types; +} + +Status ImplicitCastArguments(std::vector& arguments, + const std::vector& types, + ExecContext* exec_context) { + DCHECK_EQ(arguments.size(), types.size()); + + for (size_t i = 0; i < types.size(); ++i) { + if (types[i] == arguments[i].type()) continue; + + if (const Datum* lit = arguments[i].literal()) { + ARROW_ASSIGN_OR_RAISE(Datum new_lit, compute::Cast(*lit, types[i].GetSharedPtr())); + arguments[i] = literal(std::move(new_lit)); + continue; + } + + // construct an implicit cast Expression with which to replace this argument + Expression::Call implicit_cast; + implicit_cast.function_name = "cast"; + implicit_cast.arguments = {std::move(arguments[i])}; + + // TODO(wesm): Use TypeHolder in options + implicit_cast.options = std::make_shared( + compute::CastOptions::Safe(types[i].GetSharedPtr())); + + ARROW_ASSIGN_OR_RAISE( + arguments[i], BindNonRecursive(std::move(implicit_cast), + /*insert_implicit_casts=*/false, exec_context)); + } + + return Status::OK(); +} + Result Expression::Bind(const TypeHolder& in, compute::ExecContext* exec_context) const { return BindImpl(*this, *in.type, exec_context); @@ -876,7 +835,7 @@ Result ExecuteScalarExpression(const Expression& expr, const ExecBatch& i } if (auto special = expr.special()) { - return special->special_exec->Execute(input, exec_context); + return special->special_executor->Execute(input, exec_context); } auto call = CallNotNull(expr); diff --git a/cpp/src/arrow/compute/expression.h b/cpp/src/arrow/compute/expression.h index 2aa8187e3cd7..a4cce3f23dae 100644 --- a/cpp/src/arrow/compute/expression.h +++ b/cpp/src/arrow/compute/expression.h @@ -67,7 +67,7 @@ class ARROW_EXPORT Expression { size_t hash; // post-Bind properties: - std::shared_ptr special_exec; + std::shared_ptr special_executor; TypeHolder type; void ComputeHash(); diff --git a/cpp/src/arrow/compute/expression_internal.h b/cpp/src/arrow/compute/expression_internal.h index 50b08a3d7a0f..b9491ce87b50 100644 --- a/cpp/src/arrow/compute/expression_internal.h +++ b/cpp/src/arrow/compute/expression_internal.h @@ -290,5 +290,12 @@ inline Result> GetFunction( return GetCastFunction(*to_type); } +std::vector GetTypesWithSmallestLiteralRepresentation( + const std::vector& exprs); + +Status ImplicitCastArguments(std::vector& arguments, + const std::vector& types, + ExecContext* exec_context); + } // namespace compute } // namespace arrow diff --git a/cpp/src/arrow/compute/special_form.cc b/cpp/src/arrow/compute/special_form.cc index b99f73361ba9..4fbf239e9310 100644 --- a/cpp/src/arrow/compute/special_form.cc +++ b/cpp/src/arrow/compute/special_form.cc @@ -433,8 +433,8 @@ struct Branch { Expression body; }; -struct ConditionalExecutor { - ConditionalExecutor(const std::vector& branches, const TypeHolder& result_type) +struct ConditionalExec { + ConditionalExec(const std::vector& branches, const TypeHolder& result_type) : branches(branches), result_type(result_type) {} Result Execute(const ExecBatch& input, ExecContext* exec_context) const&& { @@ -574,79 +574,108 @@ struct ConditionalExecutor { const TypeHolder& result_type; }; -class IfElseSpecialExec : public SpecialExec { +class ConditionalSpecialExecutor : public SpecialExecutor { public: - explicit IfElseSpecialExec(std::shared_ptr function, const Kernel* kernel) - : function(std::move(function)), kernel(kernel), branches(2) {} + explicit ConditionalSpecialExecutor(TypeHolder out_type, std::vector branches) + : SpecialExecutor(std::move(out_type)), branches(std::move(branches)) {} - Result Bind(const std::vector& arguments, - ExecContext* exec_context) override { + Result Execute(const ExecBatch& input, + ExecContext* exec_context) const override { + return ConditionalExec(branches, out_type()).Execute(input, exec_context); + } + + private: + std::vector branches; +}; + +template +class KernelBackedSpecialForm : public SpecialForm { + public: + using SpecialForm::SpecialForm; + + ARROW_DISALLOW_COPY_AND_ASSIGN(KernelBackedSpecialForm); + ARROW_DEFAULT_MOVE_AND_ASSIGN(KernelBackedSpecialForm); + + Result> Bind(std::vector& arguments, + ExecContext* exec_context) override { DCHECK(std::all_of(arguments.begin(), arguments.end(), [](const Expression& argument) { return argument.IsBound(); })); - DCHECK_EQ(arguments.size(), 3); - const auto& cond = arguments[0]; - const auto& if_true = arguments[1]; - const auto& if_false = arguments[2]; - - { - auto types = GetTypes(arguments); - KernelContext kernel_context(exec_context, kernel); - std::unique_ptr kernel_state; - if (kernel->init) { - const FunctionOptions* options = function->default_options(); - ARROW_ASSIGN_OR_RAISE(kernel_state, - kernel->init(&kernel_context, {kernel, types, options})); - kernel_context.SetState(kernel_state.get()); - } - ARROW_ASSIGN_OR_RAISE( - type, kernel->signature->out_type().Resolve(&kernel_context, types)); + ARROW_ASSIGN_OR_RAISE(auto function, + exec_context->func_registry()->GetFunction(name())); + const Kernel* kernel = nullptr; + std::vector types = GetTypes(arguments); + + // First try and bind exactly + auto maybe_exact_match = function->DispatchExact(types); + if (maybe_exact_match.ok()) { + kernel = std::move(*maybe_exact_match); + } else { + // If exact binding fails, and we are allowed to cast, then prefer casting literals + // first. Since DispatchBest generally prefers up-casting the best way to do this + // is first down-cast the literals as much as possible + types = GetTypesWithSmallestLiteralRepresentation(arguments); + ARROW_ASSIGN_OR_RAISE(kernel, function->DispatchBest(&types)); + ARROW_RETURN_NOT_OK(ImplicitCastArguments(arguments, types, exec_context)); } - DCHECK_EQ(cond.type()->id(), Type::BOOL); - DCHECK_EQ(type, *if_true.type()); - DCHECK_EQ(type, *if_false.type()); - - branches[0] = {cond, if_true}; - branches[1] = {literal(true), if_false}; + KernelContext kernel_context(exec_context, kernel); + std::unique_ptr kernel_state; + if (kernel->init) { + const FunctionOptions* options = function->default_options(); + ARROW_ASSIGN_OR_RAISE(kernel_state, + kernel->init(&kernel_context, {kernel, types, options})); + kernel_context.SetState(kernel_state.get()); + } + ARROW_ASSIGN_OR_RAISE(auto out_type, + kernel->signature->out_type().Resolve(&kernel_context, types)); - return type; + return static_cast(this)->BindImpl( + std::move(out_type), std::move(kernel), std::move(arguments), exec_context); } +}; - Result Execute(const ExecBatch& input, - ExecContext* exec_context) const override { - return ConditionalExecutor(branches, type).Execute(input, exec_context); - } +template +class ConditionalSpecialForm + : public KernelBackedSpecialForm> { + public: + using KernelBackedSpecialForm>::KernelBackedSpecialForm; - private: - // For Bind. - std::shared_ptr function; - const Kernel* kernel; + ARROW_DISALLOW_COPY_AND_ASSIGN(ConditionalSpecialForm); + ARROW_DEFAULT_MOVE_AND_ASSIGN(ConditionalSpecialForm); - // Post-bind, for Execute. - TypeHolder type; - std::vector branches; + protected: + Result> BindImpl(TypeHolder out_type, + const Kernel* kernel, + std::vector arguments, + ExecContext* exec_context) const { + auto branches = static_cast(this)->GetBranches(arguments); + return std::make_unique(std::move(out_type), + std::move(branches)); + } + + friend class KernelBackedSpecialForm>; }; -class IfElseSpecialForm : public SpecialForm { +class IfElseSpecialForm : public ConditionalSpecialForm { public: - IfElseSpecialForm() : SpecialForm(/*name=*/"if_else") {} + IfElseSpecialForm() : ConditionalSpecialForm("if_else") {} - Result> DispatchExact( - const std::vector& types, ExecContext* exec_context) const override { - ARROW_ASSIGN_OR_RAISE(auto function, - exec_context->func_registry()->GetFunction(name)); - ARROW_ASSIGN_OR_RAISE(auto kernel, function->DispatchExact(types)); - return std::make_unique(std::move(function), kernel); - } + ARROW_DISALLOW_COPY_AND_ASSIGN(IfElseSpecialForm); + ARROW_DEFAULT_MOVE_AND_ASSIGN(IfElseSpecialForm); - Result> DispatchBest( - std::vector* types, ExecContext* exec_context) const override { - ARROW_ASSIGN_OR_RAISE(auto function, - exec_context->func_registry()->GetFunction(name)); - ARROW_ASSIGN_OR_RAISE(auto kernel, function->DispatchBest(types)); - return std::make_unique(std::move(function), kernel); + protected: + std::vector GetBranches(std::vector arguments) const { + DCHECK_EQ(arguments.size(), 3); + auto cond = std::move(arguments[0]); + auto if_true = std::move(arguments[1]); + auto if_false = std::move(arguments[2]); + + return std::vector{{std::move(cond), std::move(if_true)}, + {literal(true), std::move(if_false)}}; } + + friend class ConditionalSpecialForm; }; } // namespace diff --git a/cpp/src/arrow/compute/special_form.h b/cpp/src/arrow/compute/special_form.h index 9582ce418d32..24efe09aa84d 100644 --- a/cpp/src/arrow/compute/special_form.h +++ b/cpp/src/arrow/compute/special_form.h @@ -19,31 +19,34 @@ namespace arrow::compute { -class ARROW_EXPORT SpecialExec { +class ARROW_EXPORT SpecialExecutor { public: - virtual ~SpecialExec() = default; + explicit SpecialExecutor(TypeHolder out_type) : out_type_(std::move(out_type)) {} - virtual Result Bind(const std::vector& arguments, - ExecContext* exec_context) = 0; + virtual ~SpecialExecutor() = default; + + const TypeHolder& out_type() const { return out_type_; } virtual Result Execute(const ExecBatch& input, ExecContext* exec_context) const = 0; + + private: + const TypeHolder out_type_; }; class ARROW_EXPORT SpecialForm { public: - explicit SpecialForm(std::string name) : name(std::move(name)) {} + explicit SpecialForm(std::string name) : name_(std::move(name)) {} virtual ~SpecialForm() = default; - virtual Result> DispatchExact( - const std::vector& types, ExecContext* exec_context) const = 0; + const std::string& name() const { return name_; } - virtual Result> DispatchBest( - std::vector* types, ExecContext* exec_context) const = 0; + virtual Result> Bind( + std::vector& arguments, ExecContext* exec_context) = 0; - public: - const std::string name; + private: + std::string name_; }; std::shared_ptr GetIfElseSpecialForm(); diff --git a/cpp/src/arrow/compute/special_form_test.cc b/cpp/src/arrow/compute/special_form_test.cc index 38cdf78c4a08..d69872463f6f 100644 --- a/cpp/src/arrow/compute/special_form_test.cc +++ b/cpp/src/arrow/compute/special_form_test.cc @@ -276,7 +276,7 @@ class IfElseSpecialFormTest : public ::testing::Test { const Expression& cond, const Expression& if_true, const Expression& if_false) { auto suppress_if_else_recursive = [&](const Expression& expr) -> std::vector { - if (const auto& sp = expr.special(); sp && sp->special_form->name == "if_else") { + if (const auto& sp = expr.special(); sp && sp->special_form->name() == "if_else") { const auto& cond = sp->arguments[0]; const auto& if_true = sp->arguments[1]; const auto& if_false = sp->arguments[2]; diff --git a/cpp/src/arrow/compute/type_fwd.h b/cpp/src/arrow/compute/type_fwd.h index e316469107a9..f88f12a7c235 100644 --- a/cpp/src/arrow/compute/type_fwd.h +++ b/cpp/src/arrow/compute/type_fwd.h @@ -52,8 +52,8 @@ struct KernelState; class Expression; -class SpecialExec; class SpecialForm; +class SpecialExecutor; ARROW_EXPORT ExecContext* default_exec_context(); ARROW_EXPORT ExecContext* threaded_exec_context(); From 1b834504709823cd29ef236e7925a7b1a8aee8fe Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Thu, 7 Aug 2025 02:19:46 +0800 Subject: [PATCH 10/71] Abstract dispatch for bind logic --- cpp/src/arrow/compute/expression.cc | 87 +++++++++++---------- cpp/src/arrow/compute/expression_internal.h | 10 +-- cpp/src/arrow/compute/special_form.cc | 53 ++++++------- 3 files changed, 72 insertions(+), 78 deletions(-) diff --git a/cpp/src/arrow/compute/expression.cc b/cpp/src/arrow/compute/expression.cc index 6b7defa28923..502943e73e2e 100644 --- a/cpp/src/arrow/compute/expression.cc +++ b/cpp/src/arrow/compute/expression.cc @@ -571,17 +571,35 @@ TypeHolder SmallestTypeFor(const arrow::Datum& value) { } } +std::vector GetTypesWithSmallestLiteralRepresentation( + const std::vector& exprs) { + std::vector types(exprs.size()); + for (size_t i = 0; i < exprs.size(); ++i) { + DCHECK(exprs[i].IsBound()); + if (const Datum* literal = exprs[i].literal()) { + if (literal->is_scalar()) { + types[i] = SmallestTypeFor(*literal); + } + } else { + types[i] = exprs[i].type(); + } + } + return types; +} + // Produce a bound Expression from unbound Call and bound arguments. Result BindNonRecursive(Expression::Call call, bool insert_implicit_casts, - compute::ExecContext* exec_context) { + ExecContext* exec_context) { DCHECK(std::all_of(call.arguments.begin(), call.arguments.end(), [](const Expression& argument) { return argument.IsBound(); })); std::vector types = GetTypes(call.arguments); ARROW_ASSIGN_OR_RAISE(call.function, GetFunction(call, exec_context)); - auto FinishBind = [&] { - compute::KernelContext kernel_context(exec_context, call.kernel); + auto finish_bind = [&](const Kernel* kernel, + const std::vector& types) -> Status { + call.kernel = kernel; + KernelContext kernel_context(exec_context, call.kernel); if (call.kernel->init) { const FunctionOptions* options = call.options ? call.options.get() : call.function->default_options(); @@ -597,28 +615,8 @@ Result BindNonRecursive(Expression::Call call, bool insert_implicit_ return Status::OK(); }; - // First try and bind exactly - Result maybe_exact_match = call.function->DispatchExact(types); - if (maybe_exact_match.ok()) { - call.kernel = *maybe_exact_match; - if (FinishBind().ok()) { - return Expression(std::move(call)); - } - } - - if (!insert_implicit_casts) { - return maybe_exact_match.status(); - } - - // If exact binding fails, and we are allowed to cast, then prefer casting literals - // first. Since DispatchBest generally prefers up-casting the best way to do this is - // first down-cast the literals as much as possible - types = GetTypesWithSmallestLiteralRepresentation(call.arguments); - ARROW_ASSIGN_OR_RAISE(call.kernel, call.function->DispatchBest(&types)); - - ARROW_RETURN_NOT_OK(ImplicitCastArguments(call.arguments, types, exec_context)); - - RETURN_NOT_OK(FinishBind()); + RETURN_NOT_OK(DispatchForBind(call.function.get(), call.arguments, + insert_implicit_casts, exec_context, finish_bind)); return Expression(std::move(call)); } @@ -663,26 +661,29 @@ Result BindImpl(Expression expr, const TypeOrSchema& in, } // namespace -std::vector GetTypesWithSmallestLiteralRepresentation( - const std::vector& exprs) { - std::vector types(exprs.size()); - for (size_t i = 0; i < exprs.size(); ++i) { - DCHECK(exprs[i].IsBound()); - if (const Datum* literal = exprs[i].literal()) { - if (literal->is_scalar()) { - types[i] = SmallestTypeFor(*literal); - } - } else { - types[i] = exprs[i].type(); +Status DispatchForBind( + const Function* function, std::vector arguments, + bool insert_implicit_casts, ExecContext* exec_context, + std::function&)> finish_bind) { + std::vector types = GetTypes(arguments); + + // First try and bind exactly + Result maybe_exact_match = function->DispatchExact(types); + if (maybe_exact_match.ok()) { + if (finish_bind(*maybe_exact_match, types).ok()) { + return Status::OK(); } } - return types; -} -Status ImplicitCastArguments(std::vector& arguments, - const std::vector& types, - ExecContext* exec_context) { - DCHECK_EQ(arguments.size(), types.size()); + if (!insert_implicit_casts) { + return maybe_exact_match.status(); + } + + // If exact binding fails, and we are allowed to cast, then prefer casting literals + // first. Since DispatchBest generally prefers up-casting the best way to do this is + // first down-cast the literals as much as possible + types = GetTypesWithSmallestLiteralRepresentation(arguments); + ARROW_ASSIGN_OR_RAISE(auto kernel, function->DispatchBest(&types)); for (size_t i = 0; i < types.size(); ++i) { if (types[i] == arguments[i].type()) continue; @@ -707,7 +708,7 @@ Status ImplicitCastArguments(std::vector& arguments, /*insert_implicit_casts=*/false, exec_context)); } - return Status::OK(); + return finish_bind(kernel, types); } Result Expression::Bind(const TypeHolder& in, diff --git a/cpp/src/arrow/compute/expression_internal.h b/cpp/src/arrow/compute/expression_internal.h index b9491ce87b50..ef983886b087 100644 --- a/cpp/src/arrow/compute/expression_internal.h +++ b/cpp/src/arrow/compute/expression_internal.h @@ -290,12 +290,10 @@ inline Result> GetFunction( return GetCastFunction(*to_type); } -std::vector GetTypesWithSmallestLiteralRepresentation( - const std::vector& exprs); - -Status ImplicitCastArguments(std::vector& arguments, - const std::vector& types, - ExecContext* exec_context); +Status DispatchForBind( + const Function* function, std::vector arguments, + bool insert_implicit_casts, ExecContext* exec_context, + std::function&)> finish_bind); } // namespace compute } // namespace arrow diff --git a/cpp/src/arrow/compute/special_form.cc b/cpp/src/arrow/compute/special_form.cc index 4fbf239e9310..f95d9976dcc2 100644 --- a/cpp/src/arrow/compute/special_form.cc +++ b/cpp/src/arrow/compute/special_form.cc @@ -603,35 +603,30 @@ class KernelBackedSpecialForm : public SpecialForm { ARROW_ASSIGN_OR_RAISE(auto function, exec_context->func_registry()->GetFunction(name())); - const Kernel* kernel = nullptr; - std::vector types = GetTypes(arguments); - - // First try and bind exactly - auto maybe_exact_match = function->DispatchExact(types); - if (maybe_exact_match.ok()) { - kernel = std::move(*maybe_exact_match); - } else { - // If exact binding fails, and we are allowed to cast, then prefer casting literals - // first. Since DispatchBest generally prefers up-casting the best way to do this - // is first down-cast the literals as much as possible - types = GetTypesWithSmallestLiteralRepresentation(arguments); - ARROW_ASSIGN_OR_RAISE(kernel, function->DispatchBest(&types)); - ARROW_RETURN_NOT_OK(ImplicitCastArguments(arguments, types, exec_context)); - } - KernelContext kernel_context(exec_context, kernel); - std::unique_ptr kernel_state; - if (kernel->init) { - const FunctionOptions* options = function->default_options(); - ARROW_ASSIGN_OR_RAISE(kernel_state, - kernel->init(&kernel_context, {kernel, types, options})); - kernel_context.SetState(kernel_state.get()); - } - ARROW_ASSIGN_OR_RAISE(auto out_type, - kernel->signature->out_type().Resolve(&kernel_context, types)); + const Kernel* kernel = nullptr; + TypeHolder out_type; + auto finish_bind = [&](const Kernel* dispatched_kernel, + const std::vector& types) -> Status { + kernel = dispatched_kernel; + KernelContext kernel_context(exec_context, kernel); + std::unique_ptr kernel_state; + if (kernel->init) { + const FunctionOptions* options = function->default_options(); + ARROW_ASSIGN_OR_RAISE(kernel_state, + kernel->init(&kernel_context, {kernel, types, options})); + kernel_context.SetState(kernel_state.get()); + } + ARROW_ASSIGN_OR_RAISE( + out_type, kernel->signature->out_type().Resolve(&kernel_context, types)); + return Status::OK(); + }; - return static_cast(this)->BindImpl( - std::move(out_type), std::move(kernel), std::move(arguments), exec_context); + RETURN_NOT_OK(DispatchForBind(function.get(), arguments, + /*insert_implicit_casts=*/true, exec_context, + finish_bind)); + return static_cast(this)->BindImpl(kernel, std::move(out_type), + std::move(arguments), exec_context); } }; @@ -645,8 +640,8 @@ class ConditionalSpecialForm ARROW_DEFAULT_MOVE_AND_ASSIGN(ConditionalSpecialForm); protected: - Result> BindImpl(TypeHolder out_type, - const Kernel* kernel, + Result> BindImpl(const Kernel* kernel, + TypeHolder out_type, std::vector arguments, ExecContext* exec_context) const { auto branches = static_cast(this)->GetBranches(arguments); From 0e19231c000820c18786f1f581bcdaf43e103a01 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Thu, 7 Aug 2025 10:16:48 +0800 Subject: [PATCH 11/71] Fix build benchmark --- .../arrow/compute/special_form_benchmark.cc | 33 ++++++++++++++----- 1 file changed, 24 insertions(+), 9 deletions(-) diff --git a/cpp/src/arrow/compute/special_form_benchmark.cc b/cpp/src/arrow/compute/special_form_benchmark.cc index 234ba75afa33..11c6380b2725 100644 --- a/cpp/src/arrow/compute/special_form_benchmark.cc +++ b/cpp/src/arrow/compute/special_form_benchmark.cc @@ -53,19 +53,34 @@ const PayloadOptions* GetDefaultPayloadOptions() { using PayloadState = internal::OptionsWrapper; +void DoLoad(int64_t load, int64_t length) { + for (int64_t i = 0; i < length; ++i) { + volatile int64_t j = load; + while (j-- > 0) { + // Do nothing, just burn CPU cycles + } + } +} + Status PayloadExec(KernelContext* ctx, const ExecSpan& span, ExecResult* out) { ARROW_CHECK_EQ(span.num_values(), 1); const auto& arg = span[0]; ARROW_CHECK(arg.is_array()); int64_t load = PayloadState::Get(ctx).load; - int64_t load_length = - span.selection_vector.indices() ? span.selection_vector.length() : arg.length(); - for (int64_t i = 0; i < load_length; ++i) { - volatile int64_t j = load; - while (j-- > 0) { - } - } + DoLoad(load, arg.length()); + *out->array_data_mutable() = *arg.array.ToArrayData(); + return Status::OK(); +} + +Status PayloadSelectiveExec(KernelContext* ctx, const ExecSpan& span, + const SelectionVectorSpan selection_span, ExecResult* out) { + ARROW_CHECK_EQ(span.num_values(), 1); + const auto& arg = span[0]; + ARROW_CHECK(arg.is_array()); + + int64_t load = PayloadState::Get(ctx).load; + DoLoad(load, span.length); *out->array_data_mutable() = *arg.array.ToArrayData(); return Status::OK(); } @@ -85,7 +100,7 @@ Status RegisterAuxilaryFunctions() { name, Arity::Unary(), FunctionDoc::Empty(), GetDefaultPayloadOptions()); ScalarKernel kernel({InputType::Any()}, internal::FirstType, PayloadExec, - PayloadState::Init); + PayloadState::Init, PayloadSelectiveExec); kernel.selection_vector_aware = sv_awareness; kernel.can_write_into_slices = false; kernel.null_handling = NullHandling::COMPUTED_NO_PREALLOCATE; @@ -150,7 +165,7 @@ void BenchmarkIfElse( } // namespace #ifdef BM -# error("BM is defined") +# error ("BM is defined") #else # define BENCHMARK_IF_ELSE_WITH_BASELINE_AND_SV_SUPPRESS(BM, name, ...) \ BM(name##_regular, if_else_regular, ##__VA_ARGS__); \ From a6491d9f190da9f8246d232445616880389b5307 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Wed, 13 Aug 2025 15:11:44 +0800 Subject: [PATCH 12/71] Rename backstep --- cpp/src/arrow/compute/exec.cc | 12 +++--------- cpp/src/arrow/compute/exec.h | 15 ++++++++++----- 2 files changed, 13 insertions(+), 14 deletions(-) diff --git a/cpp/src/arrow/compute/exec.cc b/cpp/src/arrow/compute/exec.cc index 55f3dc9c18fd..cd6a30ac445f 100644 --- a/cpp/src/arrow/compute/exec.cc +++ b/cpp/src/arrow/compute/exec.cc @@ -1457,18 +1457,12 @@ SelectionVector::SelectionVector(const Array& arr) : SelectionVector(arr.data()) int64_t SelectionVector::length() const { return data_->length; } -void SelectionVectorSpan::SetSlice(int64_t offset, int64_t length, int32_t backstep) { +void SelectionVectorSpan::SetSlice(int64_t offset, int64_t length, + int32_t index_back_shift) { DCHECK_NE(indices_, nullptr); offset_ = offset; length_ = length; - backstep_ = backstep; -} - -int32_t SelectionVectorSpan::operator[](int64_t i) const { - DCHECK_GE(i, 0); - DCHECK_LT(i, length_); - DCHECK_GE(indices_[i + offset_], backstep_); - return indices_[i + offset_] - backstep_; + index_back_shift_ = index_back_shift; } Result CallFunction(const std::string& func_name, const std::vector& args, diff --git a/cpp/src/arrow/compute/exec.h b/cpp/src/arrow/compute/exec.h index 393cc0b29961..e96b383ef624 100644 --- a/cpp/src/arrow/compute/exec.h +++ b/cpp/src/arrow/compute/exec.h @@ -152,12 +152,17 @@ class ARROW_EXPORT SelectionVector { class ARROW_EXPORT SelectionVectorSpan { public: explicit SelectionVectorSpan(const int32_t* indices = NULLPTR, int64_t length = 0, - int64_t offset = 0, int32_t backstep = 0) - : indices_(indices), length_(length), offset_(offset), backstep_(backstep) {} + int64_t offset = 0, int32_t index_back_shift = 0) + : indices_(indices), + length_(length), + offset_(offset), + index_back_shift_(index_back_shift) {} - void SetSlice(int64_t offset, int64_t length, int32_t backward = 0); + void SetSlice(int64_t offset, int64_t length, int32_t index_back_shift = 0); - int32_t operator[](int64_t i) const; + int32_t operator[](int64_t i) const { + return indices_[i + offset_] - index_back_shift_; + } int64_t length() const { return length_; } @@ -165,7 +170,7 @@ class ARROW_EXPORT SelectionVectorSpan { const int32_t* indices_; int64_t length_; int64_t offset_; - int32_t backstep_; + int32_t index_back_shift_; }; /// An index to represent that a batch does not belong to an ordered stream From 94871c97783ec0c96fcc427ac0ed69f4d6934101 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Thu, 14 Aug 2025 00:43:12 +0800 Subject: [PATCH 13/71] Add selection vector span test --- cpp/src/arrow/compute/exec_test.cc | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/cpp/src/arrow/compute/exec_test.cc b/cpp/src/arrow/compute/exec_test.cc index 8314ad1d5c3f..8f57b782ab80 100644 --- a/cpp/src/arrow/compute/exec_test.cc +++ b/cpp/src/arrow/compute/exec_test.cc @@ -153,6 +153,24 @@ TEST(SelectionVector, Basics) { ASSERT_EQ(3, sel_vector->indices()[1]); } +TEST(SelectionVectorSpan, Basics) { + auto indices = ArrayFromJSON(int32(), "[0, 3, 7]"); + SelectionVectorSpan sel_span(indices->data()->GetValues(1), + indices->length() - 1, + /*offset=*/1, /*index_back_shift=*/1); + ASSERT_EQ(sel_span[0], 2); + ASSERT_EQ(sel_span[1], 6); + + sel_span.SetSlice(/*offset=*/1, /*length=*/2, /*index_back_shift=*/0); + ASSERT_EQ(sel_span[0], 3); + ASSERT_EQ(sel_span[1], 7); + + sel_span.SetSlice(/*offset=*/0, /*length=*/3); + ASSERT_EQ(sel_span[0], 0); + ASSERT_EQ(sel_span[1], 3); + ASSERT_EQ(sel_span[2], 7); +} + void AssertValidityZeroExtraBits(const uint8_t* data, int64_t length, int64_t offset) { const int64_t bit_extent = ((offset + length + 7) / 8) * 8; for (int64_t i = offset + length; i < bit_extent; ++i) { From f5dc732f56973356b87c3965301ca6a0b4beb08b Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Thu, 14 Aug 2025 11:19:05 +0800 Subject: [PATCH 14/71] Small fix on selection span --- cpp/src/arrow/compute/exec.cc | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/cpp/src/arrow/compute/exec.cc b/cpp/src/arrow/compute/exec.cc index cd6a30ac445f..c75c666904ff 100644 --- a/cpp/src/arrow/compute/exec.cc +++ b/cpp/src/arrow/compute/exec.cc @@ -445,12 +445,7 @@ bool ExecSpanIterator::Next(ExecSpan* span, SelectionVectorSpan* selection_span) } else { if (selection_vector_) { DCHECK_NE(selection_span, nullptr); - if (have_chunked_arrays_) { - *selection_span = SelectionVectorSpan(selection_vector_->indices(), 0); - } else { - *selection_span = SelectionVectorSpan(selection_vector_->indices(), - selection_vector_->length()); - } + *selection_span = SelectionVectorSpan(selection_vector_->indices()); } else { DCHECK_EQ(selection_span, nullptr); } From 3785d7c3ff5bf0558c2fc8f050d9cbfdff9ae3f0 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Thu, 14 Aug 2025 14:25:06 +0800 Subject: [PATCH 15/71] Add selection span iteration test --- cpp/src/arrow/compute/exec.cc | 14 +++++--- cpp/src/arrow/compute/exec_internal.h | 3 ++ cpp/src/arrow/compute/exec_test.cc | 51 +++++++++++++++++++++++++-- 3 files changed, 61 insertions(+), 7 deletions(-) diff --git a/cpp/src/arrow/compute/exec.cc b/cpp/src/arrow/compute/exec.cc index c75c666904ff..9a7928e5b165 100644 --- a/cpp/src/arrow/compute/exec.cc +++ b/cpp/src/arrow/compute/exec.cc @@ -360,6 +360,7 @@ Status ExecSpanIterator::Init(const ExecBatch& batch, int64_t max_chunksize, have_all_scalars_ = CheckIfAllScalar(batch); promote_if_all_scalars_ = promote_if_all_scalars; position_ = 0; + selection_position_ = 0; length_ = batch.length; chunk_indexes_.clear(); chunk_indexes_.resize(args_->size(), 0); @@ -369,6 +370,11 @@ Status ExecSpanIterator::Init(const ExecBatch& batch, int64_t max_chunksize, value_offsets_.resize(args_->size(), 0); max_chunksize_ = std::min(length_, max_chunksize); selection_vector_ = batch.selection_vector.get(); + if (selection_vector_) { + selection_length_ = selection_vector_->length(); + } else { + selection_length_ = 0; + } return Status::OK(); } @@ -440,14 +446,14 @@ bool ExecSpanIterator::Next(ExecSpan* span, SelectionVectorSpan* selection_span) } } - if (have_all_scalars_ && promote_if_all_scalars_) { - PromoteExecSpanScalars(span); + if (have_all_scalars_) { + if (promote_if_all_scalars_) { + PromoteExecSpanScalars(span); + } } else { if (selection_vector_) { DCHECK_NE(selection_span, nullptr); *selection_span = SelectionVectorSpan(selection_vector_->indices()); - } else { - DCHECK_EQ(selection_span, nullptr); } } diff --git a/cpp/src/arrow/compute/exec_internal.h b/cpp/src/arrow/compute/exec_internal.h index cc9b6e278f21..1c4a49f8e484 100644 --- a/cpp/src/arrow/compute/exec_internal.h +++ b/cpp/src/arrow/compute/exec_internal.h @@ -69,7 +69,9 @@ class ARROW_EXPORT ExecSpanIterator { bool Next(ExecSpan* span, SelectionVectorSpan* selection_span = NULLPTR); int64_t length() const { return length_; } + int64_t selection_length() const { return selection_length_; } int64_t position() const { return position_; } + int64_t selection_position() const { return selection_position_; } bool have_all_scalars() const { return have_all_scalars_; } bool have_selection_vector() const { return selection_vector_ != NULLPTR; } @@ -95,6 +97,7 @@ class ARROW_EXPORT ExecSpanIterator { std::vector value_offsets_; int64_t position_ = 0; int64_t length_ = 0; + int64_t selection_length_ = 0; int64_t selection_position_ = 0; int64_t max_chunksize_; }; diff --git a/cpp/src/arrow/compute/exec_test.cc b/cpp/src/arrow/compute/exec_test.cc index 8f57b782ab80..382094a52688 100644 --- a/cpp/src/arrow/compute/exec_test.cc +++ b/cpp/src/arrow/compute/exec_test.cc @@ -750,13 +750,23 @@ class TestExecSpanIterator : public TestComputeInternals { } void CheckIteration(const ExecBatch& input, int chunksize, const std::vector& ex_batch_sizes) { + ASSERT_EQ(input.selection_vector, nullptr); + std::vector ex_selection_sizes(ex_batch_sizes.size(), 0); + return CheckIteration(input, chunksize, ex_batch_sizes, ex_selection_sizes); + } + void CheckIteration(const ExecBatch& input, int chunksize, + const std::vector& ex_batch_sizes, + const std::vector& ex_selection_sizes) { SetupIterator(input, chunksize); ExecSpan batch; - int64_t position = 0; + SelectionVectorSpan selection; + int64_t position = 0, selection_position = 0; for (size_t i = 0; i < ex_batch_sizes.size(); ++i) { ASSERT_EQ(position, iterator_.position()); - ASSERT_TRUE(iterator_.Next(&batch)); + ASSERT_EQ(selection_position, iterator_.selection_position()); + ASSERT_TRUE(iterator_.Next(&batch, &selection)); ASSERT_EQ(ex_batch_sizes[i], batch.length); + ASSERT_EQ(ex_selection_sizes[i], selection.length()); for (size_t j = 0; j < input.values.size(); ++j) { switch (input[j].kind()) { @@ -782,12 +792,22 @@ class TestExecSpanIterator : public TestComputeInternals { break; } } + if (iterator_.have_selection_vector()) { + for (int64_t j = 0; j < selection.length(); ++j) { + ASSERT_EQ(input.selection_vector->indices()[selection_position + j] - position, + selection[j]); + ASSERT_GE(selection[j], 0); + ASSERT_LT(selection[j], batch.length); + } + } position += ex_batch_sizes[i]; + selection_position += ex_selection_sizes[i]; } // Ensure that the iterator is exhausted - ASSERT_FALSE(iterator_.Next(&batch)); + ASSERT_FALSE(iterator_.Next(&batch, &selection)); ASSERT_EQ(iterator_.length(), iterator_.position()); + ASSERT_EQ(iterator_.selection_length(), iterator_.selection_position()); } protected: @@ -899,6 +919,31 @@ TEST_F(TestExecSpanIterator, ZeroLengthInputs) { CheckArgs(input); } +TEST_F(TestExecSpanIterator, SelectionSpanBasic) { + ExecBatch batch( + {Datum(GetInt32Array(30)), Datum(GetInt32Array(30)), + Datum(std::make_shared(5)), Datum(MakeNullScalar(boolean()))}, + 30, std::make_shared(*ArrayFromJSON(int32(), "[1, 2, 7, 29]"))); + + CheckIteration(batch, /*chunksize=*/7, {7, 7, 7, 7, 2}, {2, 1, 0, 0, 1}); + CheckIteration(batch, /*chunksize=*/10, {10, 10, 10}, {3, 0, 1}); + CheckIteration(batch, /*chunksize=*/20, {20, 10}, {3, 1}); + CheckIteration(batch, /*chunksize=*/30, {30}, {4}); +} + +TEST_F(TestExecSpanIterator, SelectionSpanChunked) { + ExecBatch batch( + {Datum(GetInt32Chunked({0, 20, 10})), Datum(GetInt32Chunked({15, 15})), + Datum(GetInt32Array(30)), Datum(std::make_shared(5)), + Datum(MakeNullScalar(boolean()))}, + 30, std::make_shared(*ArrayFromJSON(int32(), "[1, 2, 7, 29]"))); + + CheckIteration(batch, /*chunksize=*/7, {7, 7, 1, 5, 7, 3}, {2, 1, 0, 0, 0, 1}); + CheckIteration(batch, /*chunksize=*/10, {10, 5, 5, 10}, {3, 0, 0, 1}); + CheckIteration(batch, /*chunksize=*/20, {15, 5, 10}, {3, 0, 1}); + CheckIteration(batch, /*chunksize=*/30, {15, 5, 10}, {3, 0, 1}); +} + // ---------------------------------------------------------------------- // Scalar function execution From 3644aef200a3a26fcae4d6ddf1ea3d52b4f2e945 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Thu, 14 Aug 2025 20:37:42 +0800 Subject: [PATCH 16/71] Add selective function execution test part 1 --- cpp/src/arrow/compute/exec.cc | 10 +- cpp/src/arrow/compute/exec_test.cc | 290 ++++++++++++++++++++++++++++- cpp/src/arrow/compute/function.cc | 13 +- cpp/src/arrow/compute/function.h | 4 + cpp/src/arrow/compute/kernel.h | 23 ++- 5 files changed, 319 insertions(+), 21 deletions(-) diff --git a/cpp/src/arrow/compute/exec.cc b/cpp/src/arrow/compute/exec.cc index 9a7928e5b165..73a65e28c0e5 100644 --- a/cpp/src/arrow/compute/exec.cc +++ b/cpp/src/arrow/compute/exec.cc @@ -446,11 +446,11 @@ bool ExecSpanIterator::Next(ExecSpan* span, SelectionVectorSpan* selection_span) } } - if (have_all_scalars_) { - if (promote_if_all_scalars_) { - PromoteExecSpanScalars(span); - } - } else { + if (have_all_scalars_ && promote_if_all_scalars_) { + PromoteExecSpanScalars(span); + } + + if (!have_all_scalars_ || promote_if_all_scalars_) { if (selection_vector_) { DCHECK_NE(selection_span, nullptr); *selection_span = SelectionVectorSpan(selection_vector_->indices()); diff --git a/cpp/src/arrow/compute/exec_test.cc b/cpp/src/arrow/compute/exec_test.cc index 382094a52688..d8940beafda3 100644 --- a/cpp/src/arrow/compute/exec_test.cc +++ b/cpp/src/arrow/compute/exec_test.cc @@ -30,6 +30,7 @@ #include "arrow/chunked_array.h" #include "arrow/compute/exec.h" #include "arrow/compute/exec_internal.h" +#include "arrow/compute/expression.h" #include "arrow/compute/function.h" #include "arrow/compute/function_internal.h" #include "arrow/compute/kernel.h" @@ -39,6 +40,7 @@ #include "arrow/record_batch.h" #include "arrow/scalar.h" #include "arrow/status.h" +#include "arrow/testing/generator.h" #include "arrow/type.h" #include "arrow/util/bit_util.h" #include "arrow/util/bitmap_ops.h" @@ -947,6 +949,20 @@ TEST_F(TestExecSpanIterator, SelectionSpanChunked) { // ---------------------------------------------------------------------- // Scalar function execution +template +auto TrivialSelectiveExec(Exec&& exec) { + return [exec = std::forward(exec)](KernelContext* ctx, const ExecSpan& batch, + const SelectionVectorSpan& selection, + ExecResult* out) { + for (int i = 0; i < selection.length(); ++i) { + auto row_id = selection[i]; + EXPECT_GE(row_id, 0); + EXPECT_LT(row_id, batch.length); + } + return exec(ctx, batch, out); + }; +} + Status ExecCopyArrayData(KernelContext*, const ExecSpan& batch, ExecResult* out) { DCHECK_EQ(1, batch.num_values()); int value_size = batch[0].type()->byte_width(); @@ -959,6 +975,11 @@ Status ExecCopyArrayData(KernelContext*, const ExecSpan& batch, ExecResult* out) return Status::OK(); } +Status SelectiveExecCopyArrayData(KernelContext* ctx, const ExecSpan& batch, + const SelectionVectorSpan& selection, ExecResult* out) { + return TrivialSelectiveExec(ExecCopyArrayData)(ctx, batch, selection, out); +} + Status ExecCopyArraySpan(KernelContext*, const ExecSpan& batch, ExecResult* out) { DCHECK_EQ(1, batch.num_values()); int value_size = batch[0].type()->byte_width(); @@ -970,6 +991,11 @@ Status ExecCopyArraySpan(KernelContext*, const ExecSpan& batch, ExecResult* out) return Status::OK(); } +Status SelectiveExecCopyArraySpan(KernelContext* ctx, const ExecSpan& batch, + const SelectionVectorSpan& selection, ExecResult* out) { + return TrivialSelectiveExec(ExecCopyArraySpan)(ctx, batch, selection, out); +} + Status ExecComputedBitmap(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) { // Propagate nulls not used. Check that the out bitmap isn't the same already // as the input bitmap @@ -986,6 +1012,12 @@ Status ExecComputedBitmap(KernelContext* ctx, const ExecSpan& batch, ExecResult* return ExecCopyArraySpan(ctx, batch, out); } +Status SelectiveExecComputedBitmap(KernelContext* ctx, const ExecSpan& batch, + const SelectionVectorSpan& selection, + ExecResult* out) { + return TrivialSelectiveExec(ExecComputedBitmap)(ctx, batch, selection, out); +} + Status ExecNoPreallocatedData(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) { // Validity preallocated, but not the data @@ -997,6 +1029,12 @@ Status ExecNoPreallocatedData(KernelContext* ctx, const ExecSpan& batch, return ExecCopyArrayData(ctx, batch, out); } +Status SelectiveExecNoPreallocatedData(KernelContext* ctx, const ExecSpan& batch, + const SelectionVectorSpan& selection, + ExecResult* out) { + return TrivialSelectiveExec(ExecNoPreallocatedData)(ctx, batch, selection, out); +} + Status ExecNoPreallocatedAnything(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) { // Neither validity nor data preallocated @@ -1012,6 +1050,12 @@ Status ExecNoPreallocatedAnything(KernelContext* ctx, const ExecSpan& batch, return ExecNoPreallocatedData(ctx, batch, out); } +Status SelectiveExecNoPreallocatedAnything(KernelContext* ctx, const ExecSpan& batch, + const SelectionVectorSpan& selection, + ExecResult* out) { + return TrivialSelectiveExec(ExecNoPreallocatedAnything)(ctx, batch, selection, out); +} + class ExampleOptions : public FunctionOptions { public: explicit ExampleOptions(std::shared_ptr value); @@ -1066,6 +1110,11 @@ Status ExecStateful(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) return Status::OK(); } +Status SelectiveExecStateful(KernelContext* ctx, const ExecSpan& batch, + const SelectionVectorSpan& selection, ExecResult* out) { + return TrivialSelectiveExec(ExecStateful)(ctx, batch, selection, out); +} + Status ExecAddInt32(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) { const int32_t* left_data = batch[0].array.GetValues(1); const int32_t* right_data = batch[1].array.GetValues(1); @@ -1076,6 +1125,11 @@ Status ExecAddInt32(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) return Status::OK(); } +Status SelectiveExecAddInt32(KernelContext* ctx, const ExecSpan& batch, + const SelectionVectorSpan& selection, ExecResult* out) { + return TrivialSelectiveExec(ExecAddInt32)(ctx, batch, selection, out); +} + class TestCallScalarFunction : public TestComputeInternals { protected: static bool initialized_; @@ -1086,9 +1140,13 @@ class TestCallScalarFunction : public TestComputeInternals { if (!initialized_) { initialized_ = true; AddCopyFunctions(); + AddSelectiveCopyFunctions(); AddNoPreallocateFunctions(); + AddSelectiveNoPreallocateFunctions(); AddStatefulFunction(); + AddSelectiveStatefulFunction(); AddScalarFunction(); + AddSelectiveScalarFunction(); } } @@ -1115,6 +1173,34 @@ class TestCallScalarFunction : public TestComputeInternals { ASSERT_OK(registry->AddFunction(func2)); } + void AddSelectiveCopyFunctions() { + auto registry = GetFunctionRegistry(); + + // This function simply copies memory from the input argument into the + // (preallocated) output + auto func = std::make_shared("test_copy_selective", Arity::Unary(), + /*doc=*/FunctionDoc::Empty()); + + // Add a few kernels. Our implementation only accepts arrays + ASSERT_OK(func->AddKernel({uint8()}, uint8(), ExecCopyArraySpan, + SelectiveExecCopyArraySpan)); + ASSERT_OK(func->AddKernel({int32()}, int32(), ExecCopyArraySpan, + SelectiveExecCopyArraySpan)); + ASSERT_OK(func->AddKernel({float64()}, float64(), ExecCopyArraySpan, + SelectiveExecCopyArraySpan)); + ASSERT_OK(registry->AddFunction(func)); + + // A version which doesn't want the executor to call PropagateNulls + auto func2 = + std::make_shared("test_copy_computed_bitmap_selective", + Arity::Unary(), /*doc=*/FunctionDoc::Empty()); + ScalarKernel kernel({uint8()}, uint8(), ExecComputedBitmap, + SelectiveExecComputedBitmap); + kernel.null_handling = NullHandling::COMPUTED_PREALLOCATE; + ASSERT_OK(func2->AddKernel(kernel)); + ASSERT_OK(registry->AddFunction(func2)); + } + void AddNoPreallocateFunctions() { auto registry = GetFunctionRegistry(); @@ -1137,6 +1223,32 @@ class TestCallScalarFunction : public TestComputeInternals { ASSERT_OK(registry->AddFunction(f2)); } + void AddSelectiveNoPreallocateFunctions() { + auto registry = GetFunctionRegistry(); + + // A function that allocates its own output memory. We have cases for both + // non-preallocated data and non-preallocated validity bitmap + auto f1 = + std::make_shared("test_nopre_data_selective", Arity::Unary(), + /*doc=*/FunctionDoc::Empty()); + auto f2 = + std::make_shared("test_nopre_validity_or_data_selective", + Arity::Unary(), /*doc=*/FunctionDoc::Empty()); + + ScalarKernel kernel({uint8()}, uint8(), ExecNoPreallocatedData, + SelectiveExecNoPreallocatedData); + kernel.mem_allocation = MemAllocation::NO_PREALLOCATE; + ASSERT_OK(f1->AddKernel(kernel)); + + kernel.exec = ExecNoPreallocatedAnything; + kernel.selective_exec = SelectiveExecNoPreallocatedAnything; + kernel.null_handling = NullHandling::COMPUTED_NO_PREALLOCATE; + ASSERT_OK(f2->AddKernel(kernel)); + + ASSERT_OK(registry->AddFunction(f1)); + ASSERT_OK(registry->AddFunction(f2)); + } + void AddStatefulFunction() { auto registry = GetFunctionRegistry(); @@ -1150,6 +1262,21 @@ class TestCallScalarFunction : public TestComputeInternals { ASSERT_OK(registry->AddFunction(func)); } + void AddSelectiveStatefulFunction() { + auto registry = GetFunctionRegistry(); + + // This function's behavior depends on a static parameter that is made + // available to the kernel's execution function through its Options object + auto func = + std::make_shared("test_stateful_selective", Arity::Unary(), + /*doc=*/FunctionDoc::Empty()); + + ScalarKernel kernel({int32()}, int32(), ExecStateful, SelectiveExecStateful, + InitStateful); + ASSERT_OK(func->AddKernel(kernel)); + ASSERT_OK(registry->AddFunction(func)); + } + void AddScalarFunction() { auto registry = GetFunctionRegistry(); @@ -1158,6 +1285,17 @@ class TestCallScalarFunction : public TestComputeInternals { ASSERT_OK(func->AddKernel({int32(), int32()}, int32(), ExecAddInt32)); ASSERT_OK(registry->AddFunction(func)); } + + void AddSelectiveScalarFunction() { + auto registry = GetFunctionRegistry(); + + auto func = std::make_shared("test_scalar_add_int32_selective", + Arity::Binary(), + /*doc=*/FunctionDoc::Empty()); + ASSERT_OK(func->AddKernel({int32(), int32()}, int32(), ExecAddInt32, + SelectiveExecAddInt32)); + ASSERT_OK(registry->AddFunction(func)); + } }; bool TestCallScalarFunction::initialized_ = false; @@ -1169,8 +1307,10 @@ class FunctionCaller { virtual Result Call(const std::vector& args, const FunctionOptions* options, ExecContext* ctx = NULLPTR) = 0; - virtual Result Call(const std::vector& args, - ExecContext* ctx = NULLPTR) = 0; + + virtual Result Call(const std::vector& args, ExecContext* ctx = NULLPTR) { + return Call(args, /*options=*/nullptr, ctx); + } }; using FunctionCallerMaker = std::function>( @@ -1193,9 +1333,6 @@ class SimpleFunctionCaller : public FunctionCaller { ExecContext* ctx) override { return CallFunction(func_name, args, options, ctx); } - Result Call(const std::vector& args, ExecContext* ctx) override { - return CallFunction(func_name, args, ctx); - } std::string func_name; }; @@ -1233,11 +1370,90 @@ class ExecFunctionCaller : public FunctionCaller { ARROW_RETURN_NOT_OK(func_exec->Init(options, ctx)); return func_exec->Execute(args); } + + std::shared_ptr func_exec; +}; + +class ExpressionFunctionCaller : public FunctionCaller { + public: + ExpressionFunctionCaller(std::string func_name, std::shared_ptr schema) + : func_name_(std::move(func_name)), schema_(std::move(schema)) {} + + template + static Result> Make(std::string func_name, + std::vector in_types) { + std::vector> fields(in_types.size()); + for (size_t i = 0; i < in_types.size(); ++i) { + fields[i] = field("arg" + std::to_string(i), in_types[i].GetSharedPtr()); + } + auto s = schema(std::move(fields)); + return std::make_shared(std::move(func_name), std::move(s)); + } + + Result Call(const std::vector& args, const FunctionOptions* options, + ExecContext* ctx) override { + bool all_same = false; + auto length = InferBatchLength(args, &all_same); + ARROW_ASSIGN_OR_RAISE(auto selection, GetSelection(length)); + ExecBatch batch(args, length, std::move(selection)); + std::vector expr_args(args.size()); + for (int i = 0; i < static_cast(args.size()); ++i) { + expr_args[i] = field_ref(i); + } + Expression expr = + call(func_name_, std::move(expr_args), options ? options->Copy() : nullptr); + ARROW_ASSIGN_OR_RAISE(auto bound, expr.Bind(*schema_, ctx)); + return ExecuteScalarExpression(bound, batch, ctx); + } + + static Result> Maker(const std::string& func_name, + std::vector in_types) { + return Make(func_name, std::move(in_types)); + } + Result Call(const std::vector& args, ExecContext* ctx) override { - return Call(args, nullptr, ctx); + return Call(args, /*options=*/nullptr, ctx); } - std::shared_ptr func_exec; + protected: + virtual Result> GetSelection(int64_t length) const { + return nullptr; + } + + private: + std::string func_name_; + std::shared_ptr schema_; +}; + +class SelectiveFunctionCaller : public ExpressionFunctionCaller { + public: + SelectiveFunctionCaller(std::string func_name, std::shared_ptr schema) + : ExpressionFunctionCaller(std::move(func_name) + "_selective", std::move(schema)) { + } + + static Result> Maker(const std::string& func_name, + std::vector in_types) { + return ExpressionFunctionCaller::Make(std::move(func_name), + std::move(in_types)); + } +}; + +class SelectAllFunctionCaller : public SelectiveFunctionCaller { + public: + SelectAllFunctionCaller(std::string func_name, std::shared_ptr schema) + : SelectiveFunctionCaller(std::move(func_name), std::move(schema)) {} + + static Result> Maker(const std::string& func_name, + std::vector in_types) { + return ExpressionFunctionCaller::Make(std::move(func_name), + std::move(in_types)); + } + + protected: + Result> GetSelection(int64_t length) const override { + ARROW_ASSIGN_OR_RAISE(auto arr, gen::Step()->Generate(length)); + return std::make_shared(*arr); + } }; class TestCallScalarFunctionArgumentValidation : public TestCallScalarFunction { @@ -1273,6 +1489,18 @@ TEST_F(TestCallScalarFunctionArgumentValidation, ExecCall) { TestCallScalarFunctionArgumentValidation::DoTest(ExecFunctionCaller::Maker); } +TEST_F(TestCallScalarFunctionArgumentValidation, ExpressionCall) { + TestCallScalarFunctionArgumentValidation::DoTest(ExpressionFunctionCaller::Maker); +} + +TEST_F(TestCallScalarFunctionArgumentValidation, SelectiveFunctionCaller) { + TestCallScalarFunctionArgumentValidation::DoTest(SelectiveFunctionCaller::Maker); +} + +TEST_F(TestCallScalarFunctionArgumentValidation, SelectAllFunctionCaller) { + TestCallScalarFunctionArgumentValidation::DoTest(SelectAllFunctionCaller::Maker); +} + class TestCallScalarFunctionPreallocationCases : public TestCallScalarFunction { protected: void DoTest(FunctionCallerMaker caller_maker); @@ -1352,6 +1580,18 @@ TEST_F(TestCallScalarFunctionPreallocationCases, ExecCaller) { TestCallScalarFunctionPreallocationCases::DoTest(ExecFunctionCaller::Maker); } +TEST_F(TestCallScalarFunctionPreallocationCases, ExpressionCaller) { + TestCallScalarFunctionPreallocationCases::DoTest(ExpressionFunctionCaller::Maker); +} + +TEST_F(TestCallScalarFunctionPreallocationCases, SelectiveFunctionCaller) { + TestCallScalarFunctionPreallocationCases::DoTest(SelectiveFunctionCaller::Maker); +} + +TEST_F(TestCallScalarFunctionPreallocationCases, SelectAllFunctionCaller) { + TestCallScalarFunctionPreallocationCases::DoTest(SelectAllFunctionCaller::Maker); +} + class TestCallScalarFunctionBasicNonStandardCases : public TestCallScalarFunction { protected: void DoTest(FunctionCallerMaker caller_maker); @@ -1407,6 +1647,18 @@ TEST_F(TestCallScalarFunctionBasicNonStandardCases, ExecCall) { TestCallScalarFunctionBasicNonStandardCases::DoTest(ExecFunctionCaller::Maker); } +TEST_F(TestCallScalarFunctionBasicNonStandardCases, ExpressionCall) { + TestCallScalarFunctionBasicNonStandardCases::DoTest(ExpressionFunctionCaller::Maker); +} + +TEST_F(TestCallScalarFunctionBasicNonStandardCases, SelectiveFunctionCaller) { + TestCallScalarFunctionBasicNonStandardCases::DoTest(SelectiveFunctionCaller::Maker); +} + +TEST_F(TestCallScalarFunctionBasicNonStandardCases, SelectAllFunctionCaller) { + TestCallScalarFunctionBasicNonStandardCases::DoTest(SelectAllFunctionCaller::Maker); +} + class TestCallScalarFunctionStatefulKernel : public TestCallScalarFunction { protected: void DoTest(FunctionCallerMaker caller_maker); @@ -1433,6 +1685,18 @@ TEST_F(TestCallScalarFunctionStatefulKernel, ExecCall) { TestCallScalarFunctionStatefulKernel::DoTest(ExecFunctionCaller::Maker); } +TEST_F(TestCallScalarFunctionStatefulKernel, ExpressionCall) { + TestCallScalarFunctionStatefulKernel::DoTest(ExpressionFunctionCaller::Maker); +} + +TEST_F(TestCallScalarFunctionStatefulKernel, SelectiveFunctionCaller) { + TestCallScalarFunctionStatefulKernel::DoTest(SelectiveFunctionCaller::Maker); +} + +TEST_F(TestCallScalarFunctionStatefulKernel, SelectAllFunctionCaller) { + TestCallScalarFunctionStatefulKernel::DoTest(SelectAllFunctionCaller::Maker); +} + class TestCallScalarFunctionScalarFunction : public TestCallScalarFunction { protected: void DoTest(FunctionCallerMaker caller_maker); @@ -1459,6 +1723,18 @@ TEST_F(TestCallScalarFunctionScalarFunction, ExecCall) { TestCallScalarFunctionScalarFunction::DoTest(ExecFunctionCaller::Maker); } +TEST_F(TestCallScalarFunctionScalarFunction, ExpressionCall) { + TestCallScalarFunctionScalarFunction::DoTest(ExpressionFunctionCaller::Maker); +} + +TEST_F(TestCallScalarFunctionScalarFunction, SelectiveFunctionCaller) { + TestCallScalarFunctionScalarFunction::DoTest(SelectiveFunctionCaller::Maker); +} + +TEST_F(TestCallScalarFunctionScalarFunction, SelectAllFunctionCaller) { + TestCallScalarFunctionScalarFunction::DoTest(SelectAllFunctionCaller::Maker); +} + TEST(Ordering, IsSuborderOf) { Ordering a{{SortKey{3}, SortKey{1}, SortKey{7}}}; Ordering b{{SortKey{3}, SortKey{1}}}; diff --git a/cpp/src/arrow/compute/function.cc b/cpp/src/arrow/compute/function.cc index b0b12a690f86..96efc989aa85 100644 --- a/cpp/src/arrow/compute/function.cc +++ b/cpp/src/arrow/compute/function.cc @@ -410,7 +410,16 @@ Status Function::Validate() const { } Status ScalarFunction::AddKernel(std::vector in_types, OutputType out_type, - ArrayKernelExec exec, KernelInit init, + ArrayKernelExec exec, ArrayKernelExec exec, + KernelInit init, + std::shared_ptr constraint) { + return AddKernel(std::move(in_types), std::move(out_type), std::move(exec), + /*selective_exec=*/nullptr, std::move(init)); +} + +Status ScalarFunction::AddKernel(std::vector in_types, OutputType out_type, + ArrayKernelExec exec, + ArrayKernelSelectiveExec selective_exec, KernelInit init, std::shared_ptr constraint) { RETURN_NOT_OK(CheckArity(in_types.size())); @@ -419,7 +428,7 @@ Status ScalarFunction::AddKernel(std::vector in_types, OutputType out } auto sig = KernelSignature::Make(std::move(in_types), std::move(out_type), arity_.is_varargs, std::move(constraint)); - kernels_.emplace_back(std::move(sig), exec, init); + kernels_.emplace_back(std::move(sig), exec, selective_exec, init); return Status::OK(); } diff --git a/cpp/src/arrow/compute/function.h b/cpp/src/arrow/compute/function.h index 399081e2a737..b6046e655084 100644 --- a/cpp/src/arrow/compute/function.h +++ b/cpp/src/arrow/compute/function.h @@ -311,6 +311,10 @@ class ARROW_EXPORT ScalarFunction : public detail::FunctionImpl { ArrayKernelExec exec, KernelInit init = NULLPTR, std::shared_ptr constraint = NULLPTR); + Status AddKernel(std::vector in_types, OutputType out_type, + ArrayKernelExec exec, ArrayKernelSelectiveExec selective_exec, + KernelInit init = NULLPTR); + /// \brief Add a kernel (function implementation). Returns error if the /// kernel's signature does not match the function's arity. Status AddKernel(ScalarKernel kernel); diff --git a/cpp/src/arrow/compute/kernel.h b/cpp/src/arrow/compute/kernel.h index 928a0d0d111a..6c7cccbed713 100644 --- a/cpp/src/arrow/compute/kernel.h +++ b/cpp/src/arrow/compute/kernel.h @@ -565,16 +565,25 @@ struct ARROW_EXPORT ScalarKernel : public Kernel { ScalarKernel() = default; ScalarKernel(std::shared_ptr sig, ArrayKernelExec exec, - KernelInit init = NULLPTR, - ArrayKernelSelectiveExec selective_exec = NULLPTR) - : Kernel(std::move(sig), init), exec(exec), selective_exec(selective_exec) {} + ArrayKernelSelectiveExec selective_exec, KernelInit init = NULLPTR) + : Kernel(std::move(sig), std::move(init)), + exec(std::move(exec)), + selective_exec(std::move(selective_exec)) {} ScalarKernel(std::vector in_types, OutputType out_type, ArrayKernelExec exec, - KernelInit init = NULLPTR, - ArrayKernelSelectiveExec selective_exec = NULLPTR) + ArrayKernelSelectiveExec selective_exec, KernelInit init = NULLPTR) : Kernel(std::move(in_types), std::move(out_type), std::move(init)), - exec(exec), - selective_exec(selective_exec) {} + exec(std::move(exec)), + selective_exec(std::move(selective_exec)) {} + + ScalarKernel(std::shared_ptr sig, ArrayKernelExec exec, + KernelInit init = NULLPTR) + : ScalarKernel(std::move(sig), std::move(exec), NULLPTR, std::move(init)) {} + + ScalarKernel(std::vector in_types, OutputType out_type, ArrayKernelExec exec, + KernelInit init = NULLPTR) + : ScalarKernel(std::move(in_types), std::move(out_type), std::move(exec), NULLPTR, + std::move(init)) {} /// \brief Perform a single invocation of this kernel. Depending on the /// implementation, it may only write into preallocated memory, while in some From 96e90761afa27b38556b94c9f591b46a78136138 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Fri, 15 Aug 2025 00:28:52 +0800 Subject: [PATCH 17/71] Add real dense execution cases --- cpp/src/arrow/compute/exec_test.cc | 110 +++++++++++++++++++++-------- 1 file changed, 79 insertions(+), 31 deletions(-) diff --git a/cpp/src/arrow/compute/exec_test.cc b/cpp/src/arrow/compute/exec_test.cc index d8940beafda3..2b91b5f2d310 100644 --- a/cpp/src/arrow/compute/exec_test.cc +++ b/cpp/src/arrow/compute/exec_test.cc @@ -1374,6 +1374,7 @@ class ExecFunctionCaller : public FunctionCaller { std::shared_ptr func_exec; }; +// Call the function via expression with an optional selection vector. class ExpressionFunctionCaller : public FunctionCaller { public: ExpressionFunctionCaller(std::string func_name, std::shared_ptr schema) @@ -1420,39 +1421,65 @@ class ExpressionFunctionCaller : public FunctionCaller { return nullptr; } + Result> MakeFullSelection(int64_t length) const { + ARROW_ASSIGN_OR_RAISE(auto arr, gen::Step()->Generate(length)); + return std::make_shared(*arr); + } + private: std::string func_name_; std::shared_ptr schema_; }; -class SelectiveFunctionCaller : public ExpressionFunctionCaller { +// Call the non-selective function via expression with full selection, triggering the +// dense execution path. +class DenseFunctionCaller : public ExpressionFunctionCaller { public: - SelectiveFunctionCaller(std::string func_name, std::shared_ptr schema) + using ExpressionFunctionCaller::ExpressionFunctionCaller; + + static Result> Maker(const std::string& func_name, + std::vector in_types) { + return ExpressionFunctionCaller::Make(std::move(func_name), + std::move(in_types)); + } + + protected: + Result> GetSelection(int64_t length) const override { + return MakeFullSelection(length); + } +}; + +// Call the selective counterpart of the function via expression with no selection vector, +// triggering the regular sparse execution path. +class RegularSparseFunctionCaller : public ExpressionFunctionCaller { + public: + RegularSparseFunctionCaller(std::string func_name, std::shared_ptr schema) : ExpressionFunctionCaller(std::move(func_name) + "_selective", std::move(schema)) { } static Result> Maker(const std::string& func_name, std::vector in_types) { - return ExpressionFunctionCaller::Make(std::move(func_name), - std::move(in_types)); + return ExpressionFunctionCaller::Make( + std::move(func_name), std::move(in_types)); } }; -class SelectAllFunctionCaller : public SelectiveFunctionCaller { +// Call the selective counterpart of the function via expression with full selection, +// triggering the selective sparse execution path. +class SelectiveSparseFunctionCaller : public RegularSparseFunctionCaller { public: - SelectAllFunctionCaller(std::string func_name, std::shared_ptr schema) - : SelectiveFunctionCaller(std::move(func_name), std::move(schema)) {} + SelectiveSparseFunctionCaller(std::string func_name, std::shared_ptr schema) + : RegularSparseFunctionCaller(std::move(func_name), std::move(schema)) {} static Result> Maker(const std::string& func_name, std::vector in_types) { - return ExpressionFunctionCaller::Make(std::move(func_name), - std::move(in_types)); + return ExpressionFunctionCaller::Make( + std::move(func_name), std::move(in_types)); } protected: Result> GetSelection(int64_t length) const override { - ARROW_ASSIGN_OR_RAISE(auto arr, gen::Step()->Generate(length)); - return std::make_shared(*arr); + return MakeFullSelection(length); } }; @@ -1493,12 +1520,16 @@ TEST_F(TestCallScalarFunctionArgumentValidation, ExpressionCall) { TestCallScalarFunctionArgumentValidation::DoTest(ExpressionFunctionCaller::Maker); } -TEST_F(TestCallScalarFunctionArgumentValidation, SelectiveFunctionCaller) { - TestCallScalarFunctionArgumentValidation::DoTest(SelectiveFunctionCaller::Maker); +TEST_F(TestCallScalarFunctionArgumentValidation, DenseFunctionCaller) { + TestCallScalarFunctionArgumentValidation::DoTest(DenseFunctionCaller::Maker); } -TEST_F(TestCallScalarFunctionArgumentValidation, SelectAllFunctionCaller) { - TestCallScalarFunctionArgumentValidation::DoTest(SelectAllFunctionCaller::Maker); +TEST_F(TestCallScalarFunctionArgumentValidation, RegularSparseFunctionCaller) { + TestCallScalarFunctionArgumentValidation::DoTest(RegularSparseFunctionCaller::Maker); +} + +TEST_F(TestCallScalarFunctionArgumentValidation, SelectiveSparseFunctionCaller) { + TestCallScalarFunctionArgumentValidation::DoTest(SelectiveSparseFunctionCaller::Maker); } class TestCallScalarFunctionPreallocationCases : public TestCallScalarFunction { @@ -1584,12 +1615,16 @@ TEST_F(TestCallScalarFunctionPreallocationCases, ExpressionCaller) { TestCallScalarFunctionPreallocationCases::DoTest(ExpressionFunctionCaller::Maker); } -TEST_F(TestCallScalarFunctionPreallocationCases, SelectiveFunctionCaller) { - TestCallScalarFunctionPreallocationCases::DoTest(SelectiveFunctionCaller::Maker); +TEST_F(TestCallScalarFunctionPreallocationCases, DenseFunctionCaller) { + TestCallScalarFunctionPreallocationCases::DoTest(DenseFunctionCaller::Maker); +} + +TEST_F(TestCallScalarFunctionPreallocationCases, RegularSparseFunctionCaller) { + TestCallScalarFunctionPreallocationCases::DoTest(RegularSparseFunctionCaller::Maker); } -TEST_F(TestCallScalarFunctionPreallocationCases, SelectAllFunctionCaller) { - TestCallScalarFunctionPreallocationCases::DoTest(SelectAllFunctionCaller::Maker); +TEST_F(TestCallScalarFunctionPreallocationCases, SelectiveSparseFunctionCaller) { + TestCallScalarFunctionPreallocationCases::DoTest(SelectiveSparseFunctionCaller::Maker); } class TestCallScalarFunctionBasicNonStandardCases : public TestCallScalarFunction { @@ -1651,12 +1686,17 @@ TEST_F(TestCallScalarFunctionBasicNonStandardCases, ExpressionCall) { TestCallScalarFunctionBasicNonStandardCases::DoTest(ExpressionFunctionCaller::Maker); } -TEST_F(TestCallScalarFunctionBasicNonStandardCases, SelectiveFunctionCaller) { - TestCallScalarFunctionBasicNonStandardCases::DoTest(SelectiveFunctionCaller::Maker); +TEST_F(TestCallScalarFunctionBasicNonStandardCases, DenseFunctionCaller) { + TestCallScalarFunctionBasicNonStandardCases::DoTest(DenseFunctionCaller::Maker); } -TEST_F(TestCallScalarFunctionBasicNonStandardCases, SelectAllFunctionCaller) { - TestCallScalarFunctionBasicNonStandardCases::DoTest(SelectAllFunctionCaller::Maker); +TEST_F(TestCallScalarFunctionBasicNonStandardCases, RegularSparseFunctionCaller) { + TestCallScalarFunctionBasicNonStandardCases::DoTest(RegularSparseFunctionCaller::Maker); +} + +TEST_F(TestCallScalarFunctionBasicNonStandardCases, SelectiveSparseFunctionCaller) { + TestCallScalarFunctionBasicNonStandardCases::DoTest( + SelectiveSparseFunctionCaller::Maker); } class TestCallScalarFunctionStatefulKernel : public TestCallScalarFunction { @@ -1689,12 +1729,16 @@ TEST_F(TestCallScalarFunctionStatefulKernel, ExpressionCall) { TestCallScalarFunctionStatefulKernel::DoTest(ExpressionFunctionCaller::Maker); } -TEST_F(TestCallScalarFunctionStatefulKernel, SelectiveFunctionCaller) { - TestCallScalarFunctionStatefulKernel::DoTest(SelectiveFunctionCaller::Maker); +TEST_F(TestCallScalarFunctionStatefulKernel, DenseFunctionCaller) { + TestCallScalarFunctionStatefulKernel::DoTest(DenseFunctionCaller::Maker); +} + +TEST_F(TestCallScalarFunctionStatefulKernel, RegularSparseFunctionCaller) { + TestCallScalarFunctionStatefulKernel::DoTest(RegularSparseFunctionCaller::Maker); } -TEST_F(TestCallScalarFunctionStatefulKernel, SelectAllFunctionCaller) { - TestCallScalarFunctionStatefulKernel::DoTest(SelectAllFunctionCaller::Maker); +TEST_F(TestCallScalarFunctionStatefulKernel, SelectiveSparseFunctionCaller) { + TestCallScalarFunctionStatefulKernel::DoTest(SelectiveSparseFunctionCaller::Maker); } class TestCallScalarFunctionScalarFunction : public TestCallScalarFunction { @@ -1727,12 +1771,16 @@ TEST_F(TestCallScalarFunctionScalarFunction, ExpressionCall) { TestCallScalarFunctionScalarFunction::DoTest(ExpressionFunctionCaller::Maker); } -TEST_F(TestCallScalarFunctionScalarFunction, SelectiveFunctionCaller) { - TestCallScalarFunctionScalarFunction::DoTest(SelectiveFunctionCaller::Maker); +TEST_F(TestCallScalarFunctionScalarFunction, DenseFunctionCaller) { + TestCallScalarFunctionScalarFunction::DoTest(DenseFunctionCaller::Maker); +} + +TEST_F(TestCallScalarFunctionScalarFunction, RegularSparseFunctionCaller) { + TestCallScalarFunctionScalarFunction::DoTest(RegularSparseFunctionCaller::Maker); } -TEST_F(TestCallScalarFunctionScalarFunction, SelectAllFunctionCaller) { - TestCallScalarFunctionScalarFunction::DoTest(SelectAllFunctionCaller::Maker); +TEST_F(TestCallScalarFunctionScalarFunction, SelectiveSparseFunctionCaller) { + TestCallScalarFunctionScalarFunction::DoTest(SelectiveSparseFunctionCaller::Maker); } TEST(Ordering, IsSuborderOf) { From fa7b95a56d0239cb2bffc5c88a53613540c878a7 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Fri, 15 Aug 2025 10:23:49 +0800 Subject: [PATCH 18/71] Split chunked case for dense caller --- cpp/src/arrow/compute/exec_test.cc | 164 ++++++++++++++++++++++------- 1 file changed, 128 insertions(+), 36 deletions(-) diff --git a/cpp/src/arrow/compute/exec_test.cc b/cpp/src/arrow/compute/exec_test.cc index 2b91b5f2d310..7d93e551989e 100644 --- a/cpp/src/arrow/compute/exec_test.cc +++ b/cpp/src/arrow/compute/exec_test.cc @@ -1431,24 +1431,6 @@ class ExpressionFunctionCaller : public FunctionCaller { std::shared_ptr schema_; }; -// Call the non-selective function via expression with full selection, triggering the -// dense execution path. -class DenseFunctionCaller : public ExpressionFunctionCaller { - public: - using ExpressionFunctionCaller::ExpressionFunctionCaller; - - static Result> Maker(const std::string& func_name, - std::vector in_types) { - return ExpressionFunctionCaller::Make(std::move(func_name), - std::move(in_types)); - } - - protected: - Result> GetSelection(int64_t length) const override { - return MakeFullSelection(length); - } -}; - // Call the selective counterpart of the function via expression with no selection vector, // triggering the regular sparse execution path. class RegularSparseFunctionCaller : public ExpressionFunctionCaller { @@ -1483,6 +1465,24 @@ class SelectiveSparseFunctionCaller : public RegularSparseFunctionCaller { } }; +// Call the non-selective function via expression with full selection, triggering the +// dense execution path. +class DenseFunctionCaller : public ExpressionFunctionCaller { + public: + using ExpressionFunctionCaller::ExpressionFunctionCaller; + + static Result> Maker(const std::string& func_name, + std::vector in_types) { + return ExpressionFunctionCaller::Make(std::move(func_name), + std::move(in_types)); + } + + protected: + Result> GetSelection(int64_t length) const override { + return MakeFullSelection(length); + } +}; + class TestCallScalarFunctionArgumentValidation : public TestCallScalarFunction { protected: void DoTest(FunctionCallerMaker caller_maker); @@ -1520,10 +1520,6 @@ TEST_F(TestCallScalarFunctionArgumentValidation, ExpressionCall) { TestCallScalarFunctionArgumentValidation::DoTest(ExpressionFunctionCaller::Maker); } -TEST_F(TestCallScalarFunctionArgumentValidation, DenseFunctionCaller) { - TestCallScalarFunctionArgumentValidation::DoTest(DenseFunctionCaller::Maker); -} - TEST_F(TestCallScalarFunctionArgumentValidation, RegularSparseFunctionCaller) { TestCallScalarFunctionArgumentValidation::DoTest(RegularSparseFunctionCaller::Maker); } @@ -1532,9 +1528,14 @@ TEST_F(TestCallScalarFunctionArgumentValidation, SelectiveSparseFunctionCaller) TestCallScalarFunctionArgumentValidation::DoTest(SelectiveSparseFunctionCaller::Maker); } +TEST_F(TestCallScalarFunctionArgumentValidation, DenseFunctionCaller) { + TestCallScalarFunctionArgumentValidation::DoTest(DenseFunctionCaller::Maker); +} + class TestCallScalarFunctionPreallocationCases : public TestCallScalarFunction { protected: void DoTest(FunctionCallerMaker caller_maker); + void DoTestIndependentPreallocate(FunctionCallerMaker caller_maker); }; void TestCallScalarFunctionPreallocationCases::DoTest(FunctionCallerMaker caller_maker) { @@ -1580,6 +1581,23 @@ void TestCallScalarFunctionPreallocationCases::DoTest(FunctionCallerMaker caller ASSERT_EQ(1, actual->num_chunks()); AssertChunkedEquivalent(*carr, *actual); } + }; + + ASSERT_OK_AND_ASSIGN(auto test_copy, caller_maker("test_copy", {uint8()})); + CheckFunction(test_copy); + ASSERT_OK_AND_ASSIGN(auto test_copy_computed_bitmap, + caller_maker("test_copy_computed_bitmap", {uint8()})); + CheckFunction(test_copy_computed_bitmap); +} + +void TestCallScalarFunctionPreallocationCases::DoTestIndependentPreallocate( + FunctionCallerMaker caller_maker) { + double null_prob = 0.2; + + auto arr = GetUInt8Array(100, null_prob); + + auto CheckFunction = [&](std::shared_ptr test_copy) { + ResetContexts(); // Preallocate independently for each batch { @@ -1615,10 +1633,6 @@ TEST_F(TestCallScalarFunctionPreallocationCases, ExpressionCaller) { TestCallScalarFunctionPreallocationCases::DoTest(ExpressionFunctionCaller::Maker); } -TEST_F(TestCallScalarFunctionPreallocationCases, DenseFunctionCaller) { - TestCallScalarFunctionPreallocationCases::DoTest(DenseFunctionCaller::Maker); -} - TEST_F(TestCallScalarFunctionPreallocationCases, RegularSparseFunctionCaller) { TestCallScalarFunctionPreallocationCases::DoTest(RegularSparseFunctionCaller::Maker); } @@ -1627,9 +1641,41 @@ TEST_F(TestCallScalarFunctionPreallocationCases, SelectiveSparseFunctionCaller) TestCallScalarFunctionPreallocationCases::DoTest(SelectiveSparseFunctionCaller::Maker); } +TEST_F(TestCallScalarFunctionPreallocationCases, DenseFunctionCaller) { + TestCallScalarFunctionPreallocationCases::DoTest(DenseFunctionCaller::Maker); +} + +TEST_F(TestCallScalarFunctionPreallocationCases, IndependentPreallocateSimpleCaller) { + TestCallScalarFunctionPreallocationCases::DoTestIndependentPreallocate( + SimpleFunctionCaller::Maker); +} + +TEST_F(TestCallScalarFunctionPreallocationCases, IndependentPreallocateExecCaller) { + TestCallScalarFunctionPreallocationCases::DoTestIndependentPreallocate( + ExecFunctionCaller::Maker); +} + +TEST_F(TestCallScalarFunctionPreallocationCases, IndependentPreallocateExpressionCaller) { + TestCallScalarFunctionPreallocationCases::DoTestIndependentPreallocate( + ExpressionFunctionCaller::Maker); +} + +TEST_F(TestCallScalarFunctionPreallocationCases, + IndependentPreallocateRegularSparseFunctionCaller) { + TestCallScalarFunctionPreallocationCases::DoTestIndependentPreallocate( + RegularSparseFunctionCaller::Maker); +} + +TEST_F(TestCallScalarFunctionPreallocationCases, + IndependentPreallocateSelectiveSparseFunctionCaller) { + TestCallScalarFunctionPreallocationCases::DoTestIndependentPreallocate( + SelectiveSparseFunctionCaller::Maker); +} + class TestCallScalarFunctionBasicNonStandardCases : public TestCallScalarFunction { protected: void DoTest(FunctionCallerMaker caller_maker); + void DoTestChunked(FunctionCallerMaker caller_maker); }; void TestCallScalarFunctionBasicNonStandardCases::DoTest( @@ -1653,6 +1699,30 @@ void TestCallScalarFunctionBasicNonStandardCases::DoTest( ASSERT_OK_AND_ASSIGN(Datum result, test_nopre->Call(args)); AssertArraysEqual(*arr, *result.make_array(), true); } + }; + + ASSERT_OK_AND_ASSIGN(auto test_nopre_data, caller_maker("test_nopre_data", {uint8()})); + CheckFunction(test_nopre_data); + ASSERT_OK_AND_ASSIGN(auto test_nopre_validity_or_data, + caller_maker("test_nopre_validity_or_data", {uint8()})); + CheckFunction(test_nopre_validity_or_data); +} + +void TestCallScalarFunctionBasicNonStandardCases::DoTestChunked( + FunctionCallerMaker caller_maker) { + // Test a handful of cases + // + // * Validity bitmap computed by kernel rather than using PropagateNulls + // * Data not pre-allocated + // * Validity bitmap not pre-allocated + + double null_prob = 0.2; + + auto arr = GetUInt8Array(1000, null_prob); + std::vector args = {Datum(arr)}; + + auto CheckFunction = [&](std::shared_ptr test_nopre) { + ResetContexts(); // Split execution into 3 chunks { @@ -1686,15 +1756,37 @@ TEST_F(TestCallScalarFunctionBasicNonStandardCases, ExpressionCall) { TestCallScalarFunctionBasicNonStandardCases::DoTest(ExpressionFunctionCaller::Maker); } +TEST_F(TestCallScalarFunctionBasicNonStandardCases, RegularSparseFunctionCaller) { + TestCallScalarFunctionBasicNonStandardCases::DoTest(RegularSparseFunctionCaller::Maker); +} + +TEST_F(TestCallScalarFunctionBasicNonStandardCases, SelectiveSparseFunctionCaller) { + TestCallScalarFunctionBasicNonStandardCases::DoTest( + SelectiveSparseFunctionCaller::Maker); +} + TEST_F(TestCallScalarFunctionBasicNonStandardCases, DenseFunctionCaller) { TestCallScalarFunctionBasicNonStandardCases::DoTest(DenseFunctionCaller::Maker); } -TEST_F(TestCallScalarFunctionBasicNonStandardCases, RegularSparseFunctionCaller) { +TEST_F(TestCallScalarFunctionBasicNonStandardCases, ChunkedSimpleCall) { + TestCallScalarFunctionBasicNonStandardCases::DoTest(SimpleFunctionCaller::Maker); +} + +TEST_F(TestCallScalarFunctionBasicNonStandardCases, ChunkedExecCall) { + TestCallScalarFunctionBasicNonStandardCases::DoTest(ExecFunctionCaller::Maker); +} + +TEST_F(TestCallScalarFunctionBasicNonStandardCases, ChunkedExpressionCall) { + TestCallScalarFunctionBasicNonStandardCases::DoTest(ExpressionFunctionCaller::Maker); +} + +TEST_F(TestCallScalarFunctionBasicNonStandardCases, ChunkedRegularSparseFunctionCaller) { TestCallScalarFunctionBasicNonStandardCases::DoTest(RegularSparseFunctionCaller::Maker); } -TEST_F(TestCallScalarFunctionBasicNonStandardCases, SelectiveSparseFunctionCaller) { +TEST_F(TestCallScalarFunctionBasicNonStandardCases, + ChunkedSelectiveSparseFunctionCaller) { TestCallScalarFunctionBasicNonStandardCases::DoTest( SelectiveSparseFunctionCaller::Maker); } @@ -1729,10 +1821,6 @@ TEST_F(TestCallScalarFunctionStatefulKernel, ExpressionCall) { TestCallScalarFunctionStatefulKernel::DoTest(ExpressionFunctionCaller::Maker); } -TEST_F(TestCallScalarFunctionStatefulKernel, DenseFunctionCaller) { - TestCallScalarFunctionStatefulKernel::DoTest(DenseFunctionCaller::Maker); -} - TEST_F(TestCallScalarFunctionStatefulKernel, RegularSparseFunctionCaller) { TestCallScalarFunctionStatefulKernel::DoTest(RegularSparseFunctionCaller::Maker); } @@ -1741,6 +1829,10 @@ TEST_F(TestCallScalarFunctionStatefulKernel, SelectiveSparseFunctionCaller) { TestCallScalarFunctionStatefulKernel::DoTest(SelectiveSparseFunctionCaller::Maker); } +TEST_F(TestCallScalarFunctionStatefulKernel, DenseFunctionCaller) { + TestCallScalarFunctionStatefulKernel::DoTest(DenseFunctionCaller::Maker); +} + class TestCallScalarFunctionScalarFunction : public TestCallScalarFunction { protected: void DoTest(FunctionCallerMaker caller_maker); @@ -1771,10 +1863,6 @@ TEST_F(TestCallScalarFunctionScalarFunction, ExpressionCall) { TestCallScalarFunctionScalarFunction::DoTest(ExpressionFunctionCaller::Maker); } -TEST_F(TestCallScalarFunctionScalarFunction, DenseFunctionCaller) { - TestCallScalarFunctionScalarFunction::DoTest(DenseFunctionCaller::Maker); -} - TEST_F(TestCallScalarFunctionScalarFunction, RegularSparseFunctionCaller) { TestCallScalarFunctionScalarFunction::DoTest(RegularSparseFunctionCaller::Maker); } @@ -1783,6 +1871,10 @@ TEST_F(TestCallScalarFunctionScalarFunction, SelectiveSparseFunctionCaller) { TestCallScalarFunctionScalarFunction::DoTest(SelectiveSparseFunctionCaller::Maker); } +TEST_F(TestCallScalarFunctionScalarFunction, DenseFunctionCaller) { + TestCallScalarFunctionScalarFunction::DoTest(DenseFunctionCaller::Maker); +} + TEST(Ordering, IsSuborderOf) { Ordering a{{SortKey{3}, SortKey{1}, SortKey{7}}}; Ordering b{{SortKey{3}, SortKey{1}}}; From 9802cfc1882c1467e8acb09369e5f40f6d0c087f Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Fri, 15 Aug 2025 11:10:54 +0800 Subject: [PATCH 19/71] Refine caller test --- cpp/src/arrow/compute/exec_test.cc | 555 +++++++++++------------------ 1 file changed, 212 insertions(+), 343 deletions(-) diff --git a/cpp/src/arrow/compute/exec_test.cc b/cpp/src/arrow/compute/exec_test.cc index 7d93e551989e..6cd6179b5ab2 100644 --- a/cpp/src/arrow/compute/exec_test.cc +++ b/cpp/src/arrow/compute/exec_test.cc @@ -1304,11 +1304,14 @@ class FunctionCaller { public: virtual ~FunctionCaller() = default; + virtual std::string name() const = 0; + virtual Result Call(const std::vector& args, const FunctionOptions* options, - ExecContext* ctx = NULLPTR) = 0; + ExecContext* ctx = NULLPTR) const = 0; - virtual Result Call(const std::vector& args, ExecContext* ctx = NULLPTR) { + virtual Result Call(const std::vector& args, + ExecContext* ctx = NULLPTR) const { return Call(args, /*options=*/nullptr, ctx); } }; @@ -1320,6 +1323,8 @@ class SimpleFunctionCaller : public FunctionCaller { public: explicit SimpleFunctionCaller(const std::string& func_name) : func_name(func_name) {} + std::string name() const override { return "simple_caller"; } + static Result> Make(const std::string& func_name) { return std::make_shared(func_name); } @@ -1330,7 +1335,7 @@ class SimpleFunctionCaller : public FunctionCaller { } Result Call(const std::vector& args, const FunctionOptions* options, - ExecContext* ctx) override { + ExecContext* ctx) const override { return CallFunction(func_name, args, options, ctx); } @@ -1342,6 +1347,8 @@ class ExecFunctionCaller : public FunctionCaller { explicit ExecFunctionCaller(std::shared_ptr func_exec) : func_exec(std::move(func_exec)) {} + std::string name() const override { return "exec_caller"; } + static Result> Make( const std::string& func_name, const std::vector& args, const FunctionOptions* options = nullptr, @@ -1366,7 +1373,7 @@ class ExecFunctionCaller : public FunctionCaller { } Result Call(const std::vector& args, const FunctionOptions* options, - ExecContext* ctx) override { + ExecContext* ctx) const override { ARROW_RETURN_NOT_OK(func_exec->Init(options, ctx)); return func_exec->Execute(args); } @@ -1380,6 +1387,8 @@ class ExpressionFunctionCaller : public FunctionCaller { ExpressionFunctionCaller(std::string func_name, std::shared_ptr schema) : func_name_(std::move(func_name)), schema_(std::move(schema)) {} + std::string name() const override { return "expression_caller"; } + template static Result> Make(std::string func_name, std::vector in_types) { @@ -1392,7 +1401,7 @@ class ExpressionFunctionCaller : public FunctionCaller { } Result Call(const std::vector& args, const FunctionOptions* options, - ExecContext* ctx) override { + ExecContext* ctx) const override { bool all_same = false; auto length = InferBatchLength(args, &all_same); ARROW_ASSIGN_OR_RAISE(auto selection, GetSelection(length)); @@ -1412,10 +1421,6 @@ class ExpressionFunctionCaller : public FunctionCaller { return Make(func_name, std::move(in_types)); } - Result Call(const std::vector& args, ExecContext* ctx) override { - return Call(args, /*options=*/nullptr, ctx); - } - protected: virtual Result> GetSelection(int64_t length) const { return nullptr; @@ -1439,6 +1444,8 @@ class RegularSparseFunctionCaller : public ExpressionFunctionCaller { : ExpressionFunctionCaller(std::move(func_name) + "_selective", std::move(schema)) { } + std::string name() const override { return "regular_sparse_caller"; } + static Result> Maker(const std::string& func_name, std::vector in_types) { return ExpressionFunctionCaller::Make( @@ -1453,6 +1460,8 @@ class SelectiveSparseFunctionCaller : public RegularSparseFunctionCaller { SelectiveSparseFunctionCaller(std::string func_name, std::shared_ptr schema) : RegularSparseFunctionCaller(std::move(func_name), std::move(schema)) {} + std::string name() const override { return "selective_sparse_caller"; } + static Result> Maker(const std::string& func_name, std::vector in_types) { return ExpressionFunctionCaller::Make( @@ -1471,6 +1480,8 @@ class DenseFunctionCaller : public ExpressionFunctionCaller { public: using ExpressionFunctionCaller::ExpressionFunctionCaller; + std::string name() const override { return "dense_caller"; } + static Result> Maker(const std::string& func_name, std::vector in_types) { return ExpressionFunctionCaller::Make(std::move(func_name), @@ -1483,203 +1494,139 @@ class DenseFunctionCaller : public ExpressionFunctionCaller { } }; -class TestCallScalarFunctionArgumentValidation : public TestCallScalarFunction { - protected: - void DoTest(FunctionCallerMaker caller_maker); -}; - -void TestCallScalarFunctionArgumentValidation::DoTest(FunctionCallerMaker caller_maker) { - ASSERT_OK_AND_ASSIGN(auto test_copy, caller_maker("test_copy", {int32()})); - - // Copy accepts only a single array argument - Datum d1(GetInt32Array(10)); +class TestCallScalarFunctionArgumentValidation : public TestCallScalarFunction {}; - // Too many args - std::vector args = {d1, d1}; - ASSERT_RAISES(Invalid, test_copy->Call(args)); +TEST_F(TestCallScalarFunctionArgumentValidation, Basic) { + for (const auto& caller_maker : + {SimpleFunctionCaller::Maker, ExecFunctionCaller::Maker, + ExpressionFunctionCaller::Maker, RegularSparseFunctionCaller::Maker, + SelectiveSparseFunctionCaller::Maker, DenseFunctionCaller::Maker}) { + ASSERT_OK_AND_ASSIGN(auto test_copy, caller_maker("test_copy", {int32()})); + ARROW_SCOPED_TRACE(test_copy->name()); - // Too few - args = {}; - ASSERT_RAISES(Invalid, test_copy->Call(args)); - - // Cannot do scalar - Datum d1_scalar(std::make_shared(5)); - ASSERT_OK_AND_ASSIGN(auto result, test_copy->Call({d1})); - ASSERT_OK_AND_ASSIGN(result, test_copy->Call({d1_scalar})); -} - -TEST_F(TestCallScalarFunctionArgumentValidation, SimpleCall) { - TestCallScalarFunctionArgumentValidation::DoTest(SimpleFunctionCaller::Maker); -} + ResetContexts(); -TEST_F(TestCallScalarFunctionArgumentValidation, ExecCall) { - TestCallScalarFunctionArgumentValidation::DoTest(ExecFunctionCaller::Maker); -} + // Copy accepts only a single array argument + Datum d1(GetInt32Array(10)); -TEST_F(TestCallScalarFunctionArgumentValidation, ExpressionCall) { - TestCallScalarFunctionArgumentValidation::DoTest(ExpressionFunctionCaller::Maker); -} + // Too many args + std::vector args = {d1, d1}; + ASSERT_RAISES(Invalid, test_copy->Call(args)); -TEST_F(TestCallScalarFunctionArgumentValidation, RegularSparseFunctionCaller) { - TestCallScalarFunctionArgumentValidation::DoTest(RegularSparseFunctionCaller::Maker); -} + // Too few + args = {}; + ASSERT_RAISES(Invalid, test_copy->Call(args)); -TEST_F(TestCallScalarFunctionArgumentValidation, SelectiveSparseFunctionCaller) { - TestCallScalarFunctionArgumentValidation::DoTest(SelectiveSparseFunctionCaller::Maker); -} - -TEST_F(TestCallScalarFunctionArgumentValidation, DenseFunctionCaller) { - TestCallScalarFunctionArgumentValidation::DoTest(DenseFunctionCaller::Maker); + // Cannot do scalar + Datum d1_scalar(std::make_shared(5)); + ASSERT_OK_AND_ASSIGN(auto result, test_copy->Call({d1})); + ASSERT_OK_AND_ASSIGN(result, test_copy->Call({d1_scalar})); + } } -class TestCallScalarFunctionPreallocationCases : public TestCallScalarFunction { - protected: - void DoTest(FunctionCallerMaker caller_maker); - void DoTestIndependentPreallocate(FunctionCallerMaker caller_maker); -}; +class TestCallScalarFunctionPreallocationCases : public TestCallScalarFunction {}; -void TestCallScalarFunctionPreallocationCases::DoTest(FunctionCallerMaker caller_maker) { +TEST_F(TestCallScalarFunctionPreallocationCases, Basic) { double null_prob = 0.2; auto arr = GetUInt8Array(100, null_prob); - auto CheckFunction = [&](std::shared_ptr test_copy) { - ResetContexts(); - - // The default should be a single array output - { - std::vector args = {Datum(arr)}; - ASSERT_OK_AND_ASSIGN(Datum result, test_copy->Call(args)); - ASSERT_EQ(Datum::ARRAY, result.kind()); - AssertArraysEqual(*arr, *result.make_array()); - } - - // Set the exec_chunksize to be smaller, so now we have several invocations - // of the kernel, but still the output is one array - { - std::vector args = {Datum(arr)}; - exec_ctx_->set_exec_chunksize(80); - ASSERT_OK_AND_ASSIGN(Datum result, test_copy->Call(args, exec_ctx_.get())); - AssertArraysEqual(*arr, *result.make_array()); - } + for (const auto& name : {"test_copy", "test_copy_computed_bitmap"}) { + ARROW_SCOPED_TRACE(name); - { - // Chunksize not multiple of 8 - std::vector args = {Datum(arr)}; - exec_ctx_->set_exec_chunksize(11); - ASSERT_OK_AND_ASSIGN(Datum result, test_copy->Call(args, exec_ctx_.get())); - AssertArraysEqual(*arr, *result.make_array()); - } - - // Input is chunked, output has one big chunk - { - auto carr = - std::make_shared(ArrayVector{arr->Slice(0, 10), arr->Slice(10)}); - std::vector args = {Datum(carr)}; - ASSERT_OK_AND_ASSIGN(Datum result, test_copy->Call(args, exec_ctx_.get())); - std::shared_ptr actual = result.chunked_array(); - ASSERT_EQ(1, actual->num_chunks()); - AssertChunkedEquivalent(*carr, *actual); - } - }; + for (const auto& caller_maker : + {SimpleFunctionCaller::Maker, ExecFunctionCaller::Maker, + ExpressionFunctionCaller::Maker, RegularSparseFunctionCaller::Maker, + SelectiveSparseFunctionCaller::Maker, DenseFunctionCaller::Maker}) { + ASSERT_OK_AND_ASSIGN(auto test_copy, caller_maker(name, {uint8()})); + ARROW_SCOPED_TRACE(test_copy->name()); - ASSERT_OK_AND_ASSIGN(auto test_copy, caller_maker("test_copy", {uint8()})); - CheckFunction(test_copy); - ASSERT_OK_AND_ASSIGN(auto test_copy_computed_bitmap, - caller_maker("test_copy_computed_bitmap", {uint8()})); - CheckFunction(test_copy_computed_bitmap); -} + ResetContexts(); -void TestCallScalarFunctionPreallocationCases::DoTestIndependentPreallocate( - FunctionCallerMaker caller_maker) { - double null_prob = 0.2; + // The default should be a single array output + { + std::vector args = {Datum(arr)}; + ASSERT_OK_AND_ASSIGN(Datum result, test_copy->Call(args)); + ASSERT_EQ(Datum::ARRAY, result.kind()); + AssertArraysEqual(*arr, *result.make_array()); + } - auto arr = GetUInt8Array(100, null_prob); + // Set the exec_chunksize to be smaller, so now we have several invocations + // of the kernel, but still the output is one array + { + std::vector args = {Datum(arr)}; + exec_ctx_->set_exec_chunksize(80); + ASSERT_OK_AND_ASSIGN(Datum result, test_copy->Call(args, exec_ctx_.get())); + AssertArraysEqual(*arr, *result.make_array()); + } - auto CheckFunction = [&](std::shared_ptr test_copy) { - ResetContexts(); + { + // Chunksize not multiple of 8 + std::vector args = {Datum(arr)}; + exec_ctx_->set_exec_chunksize(11); + ASSERT_OK_AND_ASSIGN(Datum result, test_copy->Call(args, exec_ctx_.get())); + AssertArraysEqual(*arr, *result.make_array()); + } - // Preallocate independently for each batch - { - std::vector args = {Datum(arr)}; - exec_ctx_->set_preallocate_contiguous(false); - exec_ctx_->set_exec_chunksize(40); - ASSERT_OK_AND_ASSIGN(Datum result, test_copy->Call(args, exec_ctx_.get())); - ASSERT_EQ(Datum::CHUNKED_ARRAY, result.kind()); - const ChunkedArray& carr = *result.chunked_array(); - ASSERT_EQ(3, carr.num_chunks()); - AssertArraysEqual(*arr->Slice(0, 40), *carr.chunk(0)); - AssertArraysEqual(*arr->Slice(40, 40), *carr.chunk(1)); - AssertArraysEqual(*arr->Slice(80), *carr.chunk(2)); + // Input is chunked, output has one big chunk + { + auto carr = std::make_shared( + ArrayVector{arr->Slice(0, 10), arr->Slice(10)}); + std::vector args = {Datum(carr)}; + ASSERT_OK_AND_ASSIGN(Datum result, test_copy->Call(args, exec_ctx_.get())); + std::shared_ptr actual = result.chunked_array(); + ASSERT_EQ(1, actual->num_chunks()); + AssertChunkedEquivalent(*carr, *actual); + } } - }; - - ASSERT_OK_AND_ASSIGN(auto test_copy, caller_maker("test_copy", {uint8()})); - CheckFunction(test_copy); - ASSERT_OK_AND_ASSIGN(auto test_copy_computed_bitmap, - caller_maker("test_copy_computed_bitmap", {uint8()})); - CheckFunction(test_copy_computed_bitmap); -} - -TEST_F(TestCallScalarFunctionPreallocationCases, SimpleCaller) { - TestCallScalarFunctionPreallocationCases::DoTest(SimpleFunctionCaller::Maker); -} - -TEST_F(TestCallScalarFunctionPreallocationCases, ExecCaller) { - TestCallScalarFunctionPreallocationCases::DoTest(ExecFunctionCaller::Maker); -} - -TEST_F(TestCallScalarFunctionPreallocationCases, ExpressionCaller) { - TestCallScalarFunctionPreallocationCases::DoTest(ExpressionFunctionCaller::Maker); -} - -TEST_F(TestCallScalarFunctionPreallocationCases, RegularSparseFunctionCaller) { - TestCallScalarFunctionPreallocationCases::DoTest(RegularSparseFunctionCaller::Maker); -} -TEST_F(TestCallScalarFunctionPreallocationCases, SelectiveSparseFunctionCaller) { - TestCallScalarFunctionPreallocationCases::DoTest(SelectiveSparseFunctionCaller::Maker); -} - -TEST_F(TestCallScalarFunctionPreallocationCases, DenseFunctionCaller) { - TestCallScalarFunctionPreallocationCases::DoTest(DenseFunctionCaller::Maker); -} - -TEST_F(TestCallScalarFunctionPreallocationCases, IndependentPreallocateSimpleCaller) { - TestCallScalarFunctionPreallocationCases::DoTestIndependentPreallocate( - SimpleFunctionCaller::Maker); -} - -TEST_F(TestCallScalarFunctionPreallocationCases, IndependentPreallocateExecCaller) { - TestCallScalarFunctionPreallocationCases::DoTestIndependentPreallocate( - ExecFunctionCaller::Maker); -} - -TEST_F(TestCallScalarFunctionPreallocationCases, IndependentPreallocateExpressionCaller) { - TestCallScalarFunctionPreallocationCases::DoTestIndependentPreallocate( - ExpressionFunctionCaller::Maker); -} - -TEST_F(TestCallScalarFunctionPreallocationCases, - IndependentPreallocateRegularSparseFunctionCaller) { - TestCallScalarFunctionPreallocationCases::DoTestIndependentPreallocate( - RegularSparseFunctionCaller::Maker); -} + for (const auto& caller_maker : + {SimpleFunctionCaller::Maker, ExecFunctionCaller::Maker, + ExpressionFunctionCaller::Maker, RegularSparseFunctionCaller::Maker, + SelectiveSparseFunctionCaller::Maker}) { + ASSERT_OK_AND_ASSIGN(auto test_copy, caller_maker(name, {uint8()})); + ARROW_SCOPED_TRACE(test_copy->name()); + + ResetContexts(); + // Preallocate independently for each batch + { + std::vector args = {Datum(arr)}; + exec_ctx_->set_preallocate_contiguous(false); + exec_ctx_->set_exec_chunksize(40); + ASSERT_OK_AND_ASSIGN(Datum result, test_copy->Call(args, exec_ctx_.get())); + ASSERT_EQ(Datum::CHUNKED_ARRAY, result.kind()); + const ChunkedArray& carr = *result.chunked_array(); + ASSERT_EQ(3, carr.num_chunks()); + AssertArraysEqual(*arr->Slice(0, 40), *carr.chunk(0)); + AssertArraysEqual(*arr->Slice(40, 40), *carr.chunk(1)); + AssertArraysEqual(*arr->Slice(80), *carr.chunk(2)); + } + } -TEST_F(TestCallScalarFunctionPreallocationCases, - IndependentPreallocateSelectiveSparseFunctionCaller) { - TestCallScalarFunctionPreallocationCases::DoTestIndependentPreallocate( - SelectiveSparseFunctionCaller::Maker); + for (const auto& caller_maker : {DenseFunctionCaller::Maker}) { + ASSERT_OK_AND_ASSIGN(auto test_copy, caller_maker(name, {uint8()})); + ARROW_SCOPED_TRACE(test_copy->name()); + + ResetContexts(); + + // Preallocate independently for each batch + { + std::vector args = {Datum(arr)}; + exec_ctx_->set_preallocate_contiguous(false); + exec_ctx_->set_exec_chunksize(40); + ASSERT_OK_AND_ASSIGN(Datum result, test_copy->Call(args, exec_ctx_.get())); + ASSERT_EQ(Datum::CHUNKED_ARRAY, result.kind()); + const ChunkedArray& carr = *result.chunked_array(); + ASSERT_EQ(1, carr.num_chunks()); + AssertArraysEqual(*arr, *carr.chunk(0)); + } + } + } } -class TestCallScalarFunctionBasicNonStandardCases : public TestCallScalarFunction { - protected: - void DoTest(FunctionCallerMaker caller_maker); - void DoTestChunked(FunctionCallerMaker caller_maker); -}; +class TestCallScalarFunctionBasicNonStandardCases : public TestCallScalarFunction {}; -void TestCallScalarFunctionBasicNonStandardCases::DoTest( - FunctionCallerMaker caller_maker) { +TEST_F(TestCallScalarFunctionBasicNonStandardCases, Basic) { // Test a handful of cases // // * Validity bitmap computed by kernel rather than using PropagateNulls @@ -1691,188 +1638,110 @@ void TestCallScalarFunctionBasicNonStandardCases::DoTest( auto arr = GetUInt8Array(1000, null_prob); std::vector args = {Datum(arr)}; - auto CheckFunction = [&](std::shared_ptr test_nopre) { - ResetContexts(); + for (const auto& name : {"test_nopre_data", "test_nopre_validity_or_data"}) { + ARROW_SCOPED_TRACE(name); - // The default should be a single array output - { - ASSERT_OK_AND_ASSIGN(Datum result, test_nopre->Call(args)); - AssertArraysEqual(*arr, *result.make_array(), true); - } - }; - - ASSERT_OK_AND_ASSIGN(auto test_nopre_data, caller_maker("test_nopre_data", {uint8()})); - CheckFunction(test_nopre_data); - ASSERT_OK_AND_ASSIGN(auto test_nopre_validity_or_data, - caller_maker("test_nopre_validity_or_data", {uint8()})); - CheckFunction(test_nopre_validity_or_data); -} + for (const auto& caller_maker : + {SimpleFunctionCaller::Maker, ExecFunctionCaller::Maker, + ExpressionFunctionCaller::Maker, RegularSparseFunctionCaller::Maker, + SelectiveSparseFunctionCaller::Maker, DenseFunctionCaller::Maker}) { + ASSERT_OK_AND_ASSIGN(auto test_nopre, caller_maker(name, {uint8()})); + ARROW_SCOPED_TRACE(test_nopre->name()); -void TestCallScalarFunctionBasicNonStandardCases::DoTestChunked( - FunctionCallerMaker caller_maker) { - // Test a handful of cases - // - // * Validity bitmap computed by kernel rather than using PropagateNulls - // * Data not pre-allocated - // * Validity bitmap not pre-allocated + ResetContexts(); - double null_prob = 0.2; - - auto arr = GetUInt8Array(1000, null_prob); - std::vector args = {Datum(arr)}; - - auto CheckFunction = [&](std::shared_ptr test_nopre) { - ResetContexts(); - - // Split execution into 3 chunks - { - exec_ctx_->set_exec_chunksize(400); - ASSERT_OK_AND_ASSIGN(Datum result, test_nopre->Call(args, exec_ctx_.get())); - ASSERT_EQ(Datum::CHUNKED_ARRAY, result.kind()); - const ChunkedArray& carr = *result.chunked_array(); - ASSERT_EQ(3, carr.num_chunks()); - AssertArraysEqual(*arr->Slice(0, 400), *carr.chunk(0)); - AssertArraysEqual(*arr->Slice(400, 400), *carr.chunk(1)); - AssertArraysEqual(*arr->Slice(800), *carr.chunk(2)); + // The default should be a single array output + { + ASSERT_OK_AND_ASSIGN(Datum result, test_nopre->Call(args)); + AssertArraysEqual(*arr, *result.make_array(), true); + } } - }; - - ASSERT_OK_AND_ASSIGN(auto test_nopre_data, caller_maker("test_nopre_data", {uint8()})); - CheckFunction(test_nopre_data); - ASSERT_OK_AND_ASSIGN(auto test_nopre_validity_or_data, - caller_maker("test_nopre_validity_or_data", {uint8()})); - CheckFunction(test_nopre_validity_or_data); -} - -TEST_F(TestCallScalarFunctionBasicNonStandardCases, SimpleCall) { - TestCallScalarFunctionBasicNonStandardCases::DoTest(SimpleFunctionCaller::Maker); -} - -TEST_F(TestCallScalarFunctionBasicNonStandardCases, ExecCall) { - TestCallScalarFunctionBasicNonStandardCases::DoTest(ExecFunctionCaller::Maker); -} - -TEST_F(TestCallScalarFunctionBasicNonStandardCases, ExpressionCall) { - TestCallScalarFunctionBasicNonStandardCases::DoTest(ExpressionFunctionCaller::Maker); -} - -TEST_F(TestCallScalarFunctionBasicNonStandardCases, RegularSparseFunctionCaller) { - TestCallScalarFunctionBasicNonStandardCases::DoTest(RegularSparseFunctionCaller::Maker); -} - -TEST_F(TestCallScalarFunctionBasicNonStandardCases, SelectiveSparseFunctionCaller) { - TestCallScalarFunctionBasicNonStandardCases::DoTest( - SelectiveSparseFunctionCaller::Maker); -} - -TEST_F(TestCallScalarFunctionBasicNonStandardCases, DenseFunctionCaller) { - TestCallScalarFunctionBasicNonStandardCases::DoTest(DenseFunctionCaller::Maker); -} - -TEST_F(TestCallScalarFunctionBasicNonStandardCases, ChunkedSimpleCall) { - TestCallScalarFunctionBasicNonStandardCases::DoTest(SimpleFunctionCaller::Maker); -} - -TEST_F(TestCallScalarFunctionBasicNonStandardCases, ChunkedExecCall) { - TestCallScalarFunctionBasicNonStandardCases::DoTest(ExecFunctionCaller::Maker); -} -TEST_F(TestCallScalarFunctionBasicNonStandardCases, ChunkedExpressionCall) { - TestCallScalarFunctionBasicNonStandardCases::DoTest(ExpressionFunctionCaller::Maker); -} - -TEST_F(TestCallScalarFunctionBasicNonStandardCases, ChunkedRegularSparseFunctionCaller) { - TestCallScalarFunctionBasicNonStandardCases::DoTest(RegularSparseFunctionCaller::Maker); -} - -TEST_F(TestCallScalarFunctionBasicNonStandardCases, - ChunkedSelectiveSparseFunctionCaller) { - TestCallScalarFunctionBasicNonStandardCases::DoTest( - SelectiveSparseFunctionCaller::Maker); -} - -class TestCallScalarFunctionStatefulKernel : public TestCallScalarFunction { - protected: - void DoTest(FunctionCallerMaker caller_maker); -}; - -void TestCallScalarFunctionStatefulKernel::DoTest(FunctionCallerMaker caller_maker) { - ASSERT_OK_AND_ASSIGN(auto test_stateful, caller_maker("test_stateful", {int32()})); - - auto input = ArrayFromJSON(int32(), "[1, 2, 3, null, 5]"); - auto multiplier = std::make_shared(2); - auto expected = ArrayFromJSON(int32(), "[2, 4, 6, null, 10]"); - - ExampleOptions options(multiplier); - std::vector args = {Datum(input)}; - ASSERT_OK_AND_ASSIGN(Datum result, test_stateful->Call(args, &options)); - AssertArraysEqual(*expected, *result.make_array()); -} - -TEST_F(TestCallScalarFunctionStatefulKernel, Simplecall) { - TestCallScalarFunctionStatefulKernel::DoTest(SimpleFunctionCaller::Maker); -} - -TEST_F(TestCallScalarFunctionStatefulKernel, ExecCall) { - TestCallScalarFunctionStatefulKernel::DoTest(ExecFunctionCaller::Maker); -} - -TEST_F(TestCallScalarFunctionStatefulKernel, ExpressionCall) { - TestCallScalarFunctionStatefulKernel::DoTest(ExpressionFunctionCaller::Maker); -} + for (const auto& caller_maker : + {SimpleFunctionCaller::Maker, ExecFunctionCaller::Maker, + ExpressionFunctionCaller::Maker, RegularSparseFunctionCaller::Maker, + SelectiveSparseFunctionCaller::Maker}) { + ASSERT_OK_AND_ASSIGN(auto test_nopre, caller_maker(name, {uint8()})); + ARROW_SCOPED_TRACE(test_nopre->name()); + + ResetContexts(); + + // Split execution into 3 chunks + { + exec_ctx_->set_exec_chunksize(400); + ASSERT_OK_AND_ASSIGN(Datum result, test_nopre->Call(args, exec_ctx_.get())); + ASSERT_EQ(Datum::CHUNKED_ARRAY, result.kind()); + const ChunkedArray& carr = *result.chunked_array(); + ASSERT_EQ(3, carr.num_chunks()); + AssertArraysEqual(*arr->Slice(0, 400), *carr.chunk(0)); + AssertArraysEqual(*arr->Slice(400, 400), *carr.chunk(1)); + AssertArraysEqual(*arr->Slice(800), *carr.chunk(2)); + } + } -TEST_F(TestCallScalarFunctionStatefulKernel, RegularSparseFunctionCaller) { - TestCallScalarFunctionStatefulKernel::DoTest(RegularSparseFunctionCaller::Maker); -} + for (const auto& caller_maker : {DenseFunctionCaller::Maker}) { + ASSERT_OK_AND_ASSIGN(auto test_nopre, caller_maker(name, {uint8()})); + ARROW_SCOPED_TRACE(test_nopre->name()); -TEST_F(TestCallScalarFunctionStatefulKernel, SelectiveSparseFunctionCaller) { - TestCallScalarFunctionStatefulKernel::DoTest(SelectiveSparseFunctionCaller::Maker); -} + ResetContexts(); -TEST_F(TestCallScalarFunctionStatefulKernel, DenseFunctionCaller) { - TestCallScalarFunctionStatefulKernel::DoTest(DenseFunctionCaller::Maker); + // Split execution into 3 chunks + { + exec_ctx_->set_exec_chunksize(400); + ASSERT_OK_AND_ASSIGN(Datum result, test_nopre->Call(args, exec_ctx_.get())); + ASSERT_EQ(Datum::CHUNKED_ARRAY, result.kind()); + const ChunkedArray& carr = *result.chunked_array(); + ASSERT_EQ(1, carr.num_chunks()); + AssertArraysEqual(*arr, *carr.chunk(0)); + } + } + } } -class TestCallScalarFunctionScalarFunction : public TestCallScalarFunction { - protected: - void DoTest(FunctionCallerMaker caller_maker); -}; +class TestCallScalarFunctionStatefulKernel : public TestCallScalarFunction {}; -void TestCallScalarFunctionScalarFunction::DoTest(FunctionCallerMaker caller_maker) { - ASSERT_OK_AND_ASSIGN(auto test_scalar_add_int32, - caller_maker("test_scalar_add_int32", {int32(), int32()})); +TEST_F(TestCallScalarFunctionStatefulKernel, Basic) { + for (const auto& caller_maker : + {SimpleFunctionCaller::Maker, ExecFunctionCaller::Maker, + ExpressionFunctionCaller::Maker, RegularSparseFunctionCaller::Maker, + SelectiveSparseFunctionCaller::Maker, DenseFunctionCaller::Maker}) { + ASSERT_OK_AND_ASSIGN(auto test_stateful, caller_maker("test_stateful", {int32()})); + ARROW_SCOPED_TRACE(test_stateful->name()); - std::vector args = {Datum(std::make_shared(5)), - Datum(std::make_shared(7))}; - ASSERT_OK_AND_ASSIGN(Datum result, test_scalar_add_int32->Call(args)); - ASSERT_EQ(Datum::SCALAR, result.kind()); + ResetContexts(); - auto expected = std::make_shared(12); - ASSERT_TRUE(expected->Equals(*result.scalar())); -} + auto input = ArrayFromJSON(int32(), "[1, 2, 3, null, 5]"); + auto multiplier = std::make_shared(2); + auto expected = ArrayFromJSON(int32(), "[2, 4, 6, null, 10]"); -TEST_F(TestCallScalarFunctionScalarFunction, SimpleCall) { - TestCallScalarFunctionScalarFunction::DoTest(SimpleFunctionCaller::Maker); + ExampleOptions options(multiplier); + std::vector args = {Datum(input)}; + ASSERT_OK_AND_ASSIGN(Datum result, test_stateful->Call(args, &options)); + AssertArraysEqual(*expected, *result.make_array()); + } } -TEST_F(TestCallScalarFunctionScalarFunction, ExecCall) { - TestCallScalarFunctionScalarFunction::DoTest(ExecFunctionCaller::Maker); -} +class TestCallScalarFunctionScalarFunction : public TestCallScalarFunction {}; -TEST_F(TestCallScalarFunctionScalarFunction, ExpressionCall) { - TestCallScalarFunctionScalarFunction::DoTest(ExpressionFunctionCaller::Maker); -} +TEST_F(TestCallScalarFunctionScalarFunction, Basic) { + for (const auto& caller_maker : + {SimpleFunctionCaller::Maker, ExecFunctionCaller::Maker, + ExpressionFunctionCaller::Maker, RegularSparseFunctionCaller::Maker, + SelectiveSparseFunctionCaller::Maker, DenseFunctionCaller::Maker}) { + ASSERT_OK_AND_ASSIGN(auto test_scalar_add_int32, + caller_maker("test_scalar_add_int32", {int32(), int32()})); + ARROW_SCOPED_TRACE(test_scalar_add_int32->name()); -TEST_F(TestCallScalarFunctionScalarFunction, RegularSparseFunctionCaller) { - TestCallScalarFunctionScalarFunction::DoTest(RegularSparseFunctionCaller::Maker); -} + ResetContexts(); -TEST_F(TestCallScalarFunctionScalarFunction, SelectiveSparseFunctionCaller) { - TestCallScalarFunctionScalarFunction::DoTest(SelectiveSparseFunctionCaller::Maker); -} + std::vector args = {Datum(std::make_shared(5)), + Datum(std::make_shared(7))}; + ASSERT_OK_AND_ASSIGN(Datum result, test_scalar_add_int32->Call(args)); + ASSERT_EQ(Datum::SCALAR, result.kind()); -TEST_F(TestCallScalarFunctionScalarFunction, DenseFunctionCaller) { - TestCallScalarFunctionScalarFunction::DoTest(DenseFunctionCaller::Maker); + auto expected = std::make_shared(12); + ASSERT_TRUE(expected->Equals(*result.scalar())); + } } TEST(Ordering, IsSuborderOf) { From 1cbb45cc589099bc4c1fa7f112bcca499227f514 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Mon, 18 Aug 2025 15:59:37 +0800 Subject: [PATCH 20/71] Change to decent selective kernels --- cpp/src/arrow/compute/exec_test.cc | 146 +++++++++++++++++++++++++---- 1 file changed, 128 insertions(+), 18 deletions(-) diff --git a/cpp/src/arrow/compute/exec_test.cc b/cpp/src/arrow/compute/exec_test.cc index 6cd6179b5ab2..257486c855d9 100644 --- a/cpp/src/arrow/compute/exec_test.cc +++ b/cpp/src/arrow/compute/exec_test.cc @@ -949,18 +949,20 @@ TEST_F(TestExecSpanIterator, SelectionSpanChunked) { // ---------------------------------------------------------------------- // Scalar function execution -template -auto TrivialSelectiveExec(Exec&& exec) { - return [exec = std::forward(exec)](KernelContext* ctx, const ExecSpan& batch, - const SelectionVectorSpan& selection, - ExecResult* out) { - for (int i = 0; i < selection.length(); ++i) { - auto row_id = selection[i]; - EXPECT_GE(row_id, 0); - EXPECT_LT(row_id, batch.length); +template +Status VisitSelectionVector(const SelectionVectorSpan& selection, int64_t length, + OnSelectionFn&& on_selection, + OnNonSelectionFn&& on_non_selection) { + int64_t selected = 0; + for (int64_t i = 0; i < length; ++i) { + if (selected < selection.length() && i == selection[selected]) { + on_selection(i); + ++selected; + } else { + on_non_selection(i); } - return exec(ctx, batch, out); - }; + } + return Status::OK(); } Status ExecCopyArrayData(KernelContext*, const ExecSpan& batch, ExecResult* out) { @@ -975,9 +977,32 @@ Status ExecCopyArrayData(KernelContext*, const ExecSpan& batch, ExecResult* out) return Status::OK(); } +constexpr int8_t kNonSelectionValue = 0xFE; + Status SelectiveExecCopyArrayData(KernelContext* ctx, const ExecSpan& batch, const SelectionVectorSpan& selection, ExecResult* out) { - return TrivialSelectiveExec(ExecCopyArrayData)(ctx, batch, selection, out); + DCHECK_EQ(1, batch.num_values()); + int value_size = batch[0].type()->byte_width(); + + const ArraySpan& arg0 = batch[0].array; + ArrayData* out_arr = out->array_data().get(); + uint8_t* dst_validity = out_arr->buffers[0]->mutable_data(); + int64_t dst_validity_offset = out_arr->offset; + uint8_t* dst = out_arr->buffers[1]->mutable_data() + out_arr->offset * value_size; + const uint8_t* src = arg0.buffers[1].data + arg0.offset * value_size; + return VisitSelectionVector( + selection, batch.length, + [&](int64_t i) { + // Copy the selected value + std::memcpy(dst + i * value_size, src + i * value_size, value_size); + }, + [&](int64_t i) { + // Set the non-selected as valid (regardless of its precomputed validity) and set + // its values with a special value + bit_util::SetBit(dst_validity, dst_validity_offset + i); + std::memset(dst + i * value_size, kNonSelectionValue, value_size); + }); + return Status::OK(); } Status ExecCopyArraySpan(KernelContext*, const ExecSpan& batch, ExecResult* out) { @@ -993,7 +1018,26 @@ Status ExecCopyArraySpan(KernelContext*, const ExecSpan& batch, ExecResult* out) Status SelectiveExecCopyArraySpan(KernelContext* ctx, const ExecSpan& batch, const SelectionVectorSpan& selection, ExecResult* out) { - return TrivialSelectiveExec(ExecCopyArraySpan)(ctx, batch, selection, out); + DCHECK_EQ(1, batch.num_values()); + int value_size = batch[0].type()->byte_width(); + const ArraySpan& arg0 = batch[0].array; + ArraySpan* out_arr = out->array_span_mutable(); + uint8_t* dst_validity = out_arr->buffers[0].data; + int64_t dst_validity_offset = out_arr->offset; + uint8_t* dst = out_arr->buffers[1].data + out_arr->offset * value_size; + const uint8_t* src = arg0.buffers[1].data + arg0.offset * value_size; + return VisitSelectionVector( + selection, batch.length, + [&](int64_t i) { + // Copy the selected value + std::memcpy(dst + i * value_size, src + i * value_size, value_size); + }, + [&](int64_t i) { + // Set the non-selected as valid (regardless of its precomputed validity) and set + // its values with a special value + bit_util::SetBit(dst_validity, dst_validity_offset + i); + std::memset(dst + i * value_size, kNonSelectionValue, value_size); + }); } Status ExecComputedBitmap(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) { @@ -1015,7 +1059,19 @@ Status ExecComputedBitmap(KernelContext* ctx, const ExecSpan& batch, ExecResult* Status SelectiveExecComputedBitmap(KernelContext* ctx, const ExecSpan& batch, const SelectionVectorSpan& selection, ExecResult* out) { - return TrivialSelectiveExec(ExecComputedBitmap)(ctx, batch, selection, out); + // Propagate nulls not used. Check that the out bitmap isn't the same already + // as the input bitmap + const ArraySpan& arg0 = batch[0].array; + ArraySpan* out_arr = out->array_span_mutable(); + if (CountSetBits(arg0.buffers[0].data, arg0.offset, batch.length) > 0) { + // Check that the bitmap has not been already copied over + DCHECK(!BitmapEquals(arg0.buffers[0].data, arg0.offset, out_arr->buffers[0].data, + out_arr->offset, batch.length)); + } + + CopyBitmap(arg0.buffers[0].data, arg0.offset, batch.length, out_arr->buffers[0].data, + out_arr->offset); + return SelectiveExecCopyArraySpan(ctx, batch, selection, out); } Status ExecNoPreallocatedData(KernelContext* ctx, const ExecSpan& batch, @@ -1032,7 +1088,13 @@ Status ExecNoPreallocatedData(KernelContext* ctx, const ExecSpan& batch, Status SelectiveExecNoPreallocatedData(KernelContext* ctx, const ExecSpan& batch, const SelectionVectorSpan& selection, ExecResult* out) { - return TrivialSelectiveExec(ExecNoPreallocatedData)(ctx, batch, selection, out); + // Validity preallocated, but not the data + ArrayData* out_arr = out->array_data().get(); + DCHECK_EQ(0, out_arr->offset); + int value_size = batch[0].type()->byte_width(); + Status s = (ctx->Allocate(out_arr->length * value_size).Value(&out_arr->buffers[1])); + DCHECK_OK(s); + return SelectiveExecCopyArrayData(ctx, batch, selection, out); } Status ExecNoPreallocatedAnything(KernelContext* ctx, const ExecSpan& batch, @@ -1053,7 +1115,17 @@ Status ExecNoPreallocatedAnything(KernelContext* ctx, const ExecSpan& batch, Status SelectiveExecNoPreallocatedAnything(KernelContext* ctx, const ExecSpan& batch, const SelectionVectorSpan& selection, ExecResult* out) { - return TrivialSelectiveExec(ExecNoPreallocatedAnything)(ctx, batch, selection, out); + // Neither validity nor data preallocated + ArrayData* out_arr = out->array_data().get(); + DCHECK_EQ(0, out_arr->offset); + Status s = (ctx->AllocateBitmap(out_arr->length).Value(&out_arr->buffers[0])); + DCHECK_OK(s); + const ArraySpan& arg0 = batch[0].array; + CopyBitmap(arg0.buffers[0].data, arg0.offset, batch.length, + out_arr->buffers[0]->mutable_data(), /*offset=*/0); + + // Reuse the kernel that allocates the data + return SelectiveExecNoPreallocatedData(ctx, batch, selection, out); } class ExampleOptions : public FunctionOptions { @@ -1112,7 +1184,28 @@ Status ExecStateful(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) Status SelectiveExecStateful(KernelContext* ctx, const ExecSpan& batch, const SelectionVectorSpan& selection, ExecResult* out) { - return TrivialSelectiveExec(ExecStateful)(ctx, batch, selection, out); + // We take the value from the state and multiply the data in batch[0] with it + ExampleState* state = static_cast(ctx->state()); + int32_t multiplier = checked_cast(*state->value).value; + + const ArraySpan& arg0 = batch[0].array; + ArraySpan* out_arr = out->array_span_mutable(); + const int32_t* arg0_data = arg0.GetValues(1); + uint8_t* dst_validity = out_arr->buffers[0].data; + int64_t dst_validity_offset = out_arr->offset; + int32_t* dst = out_arr->GetValues(1); + return VisitSelectionVector( + selection, batch.length, + [&](int64_t i) { + // Copy the selected value + dst[i] = arg0_data[i] * multiplier; + }, + [&](int64_t i) { + // Set the non-selected as valid (regardless of its precomputed validity) and set + // its values with a special value + bit_util::SetBit(dst_validity, dst_validity_offset + i); + dst[i] = kNonSelectionValue; + }); } Status ExecAddInt32(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) { @@ -1127,7 +1220,24 @@ Status ExecAddInt32(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) Status SelectiveExecAddInt32(KernelContext* ctx, const ExecSpan& batch, const SelectionVectorSpan& selection, ExecResult* out) { - return TrivialSelectiveExec(ExecAddInt32)(ctx, batch, selection, out); + const int32_t* left_data = batch[0].array.GetValues(1); + const int32_t* right_data = batch[1].array.GetValues(1); + ArraySpan* out_arr = out->array_span_mutable(); + uint8_t* dst_validity = out_arr->buffers[0].data; + int64_t dst_validity_offset = out_arr->offset; + int32_t* out_data = out_arr->GetValues(1); + return VisitSelectionVector( + selection, batch.length, + [&](int64_t i) { + // Copy the selected value + out_data[i] = left_data[i] + right_data[i]; + }, + [&](int64_t i) { + // Set the non-selected as valid (regardless of its precomputed validity) and set + // its values with a special value + bit_util::SetBit(dst_validity, dst_validity_offset + i); + out_data[i] = kNonSelectionValue; + }); } class TestCallScalarFunction : public TestComputeInternals { From 52d062905040c52e322e11d106649919559809b2 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Mon, 18 Aug 2025 19:30:05 +0800 Subject: [PATCH 21/71] Reorg existing tests --- cpp/src/arrow/compute/exec_test.cc | 158 ++++++++++++++++------------- 1 file changed, 89 insertions(+), 69 deletions(-) diff --git a/cpp/src/arrow/compute/exec_test.cc b/cpp/src/arrow/compute/exec_test.cc index 257486c855d9..86b765876b95 100644 --- a/cpp/src/arrow/compute/exec_test.cc +++ b/cpp/src/arrow/compute/exec_test.cc @@ -1494,20 +1494,21 @@ class ExecFunctionCaller : public FunctionCaller { // Call the function via expression with an optional selection vector. class ExpressionFunctionCaller : public FunctionCaller { public: - ExpressionFunctionCaller(std::string func_name, std::shared_ptr schema) - : func_name_(std::move(func_name)), schema_(std::move(schema)) {} + ExpressionFunctionCaller(std::string func_name, const std::vector& in_types) + : func_name_(std::move(func_name)) { + std::vector> fields(in_types.size()); + for (size_t i = 0; i < in_types.size(); ++i) { + fields[i] = field("arg" + std::to_string(i), in_types[i].GetSharedPtr()); + } + schema_ = schema(std::move(fields)); + } std::string name() const override { return "expression_caller"; } template static Result> Make(std::string func_name, std::vector in_types) { - std::vector> fields(in_types.size()); - for (size_t i = 0; i < in_types.size(); ++i) { - fields[i] = field("arg" + std::to_string(i), in_types[i].GetSharedPtr()); - } - auto s = schema(std::move(fields)); - return std::make_shared(std::move(func_name), std::move(s)); + return std::make_shared(std::move(func_name), std::move(in_types)); } Result Call(const std::vector& args, const FunctionOptions* options, @@ -1515,6 +1516,15 @@ class ExpressionFunctionCaller : public FunctionCaller { bool all_same = false; auto length = InferBatchLength(args, &all_same); ARROW_ASSIGN_OR_RAISE(auto selection, GetSelection(length)); + return CallWithSelection(args, options, std::move(selection), ctx); + } + + Result CallWithSelection(const std::vector& args, + const FunctionOptions* options, + std::shared_ptr selection, + ExecContext* ctx) const { + bool all_same = false; + auto length = InferBatchLength(args, &all_same); ExecBatch batch(args, length, std::move(selection)); std::vector expr_args(args.size()); for (int i = 0; i < static_cast(args.size()); ++i) { @@ -1550,9 +1560,9 @@ class ExpressionFunctionCaller : public FunctionCaller { // triggering the regular sparse execution path. class RegularSparseFunctionCaller : public ExpressionFunctionCaller { public: - RegularSparseFunctionCaller(std::string func_name, std::shared_ptr schema) - : ExpressionFunctionCaller(std::move(func_name) + "_selective", std::move(schema)) { - } + RegularSparseFunctionCaller(std::string func_name, + const std::vector& in_types) + : ExpressionFunctionCaller(std::move(func_name) + "_selective", in_types) {} std::string name() const override { return "regular_sparse_caller"; } @@ -1567,8 +1577,9 @@ class RegularSparseFunctionCaller : public ExpressionFunctionCaller { // triggering the selective sparse execution path. class SelectiveSparseFunctionCaller : public RegularSparseFunctionCaller { public: - SelectiveSparseFunctionCaller(std::string func_name, std::shared_ptr schema) - : RegularSparseFunctionCaller(std::move(func_name), std::move(schema)) {} + SelectiveSparseFunctionCaller(std::string func_name, + const std::vector& in_types) + : RegularSparseFunctionCaller(std::move(func_name), in_types) {} std::string name() const override { return "selective_sparse_caller"; } @@ -1613,7 +1624,6 @@ TEST_F(TestCallScalarFunctionArgumentValidation, Basic) { SelectiveSparseFunctionCaller::Maker, DenseFunctionCaller::Maker}) { ASSERT_OK_AND_ASSIGN(auto test_copy, caller_maker("test_copy", {int32()})); ARROW_SCOPED_TRACE(test_copy->name()); - ResetContexts(); // Copy accepts only a single array argument @@ -1634,23 +1644,21 @@ TEST_F(TestCallScalarFunctionArgumentValidation, Basic) { } } -class TestCallScalarFunctionPreallocationCases : public TestCallScalarFunction {}; +class TestCallScalarFunctionPreallocationCases : public TestCallScalarFunction { + protected: + std::shared_ptr GetTestArray() { return GetUInt8Array(100, 0.2); } +}; TEST_F(TestCallScalarFunctionPreallocationCases, Basic) { - double null_prob = 0.2; - - auto arr = GetUInt8Array(100, null_prob); - + auto arr = GetTestArray(); for (const auto& name : {"test_copy", "test_copy_computed_bitmap"}) { ARROW_SCOPED_TRACE(name); - for (const auto& caller_maker : {SimpleFunctionCaller::Maker, ExecFunctionCaller::Maker, ExpressionFunctionCaller::Maker, RegularSparseFunctionCaller::Maker, SelectiveSparseFunctionCaller::Maker, DenseFunctionCaller::Maker}) { ASSERT_OK_AND_ASSIGN(auto test_copy, caller_maker(name, {uint8()})); ARROW_SCOPED_TRACE(test_copy->name()); - ResetContexts(); // The default should be a single array output @@ -1689,18 +1697,24 @@ TEST_F(TestCallScalarFunctionPreallocationCases, Basic) { AssertChunkedEquivalent(*carr, *actual); } } + } +} +TEST_F(TestCallScalarFunctionPreallocationCases, IndependentPreallocate) { + auto arr = GetTestArray(); + std::vector args = {Datum(arr)}; + for (const auto& name : {"test_copy", "test_copy_computed_bitmap"}) { + ARROW_SCOPED_TRACE(name); for (const auto& caller_maker : {SimpleFunctionCaller::Maker, ExecFunctionCaller::Maker, ExpressionFunctionCaller::Maker, RegularSparseFunctionCaller::Maker, SelectiveSparseFunctionCaller::Maker}) { ASSERT_OK_AND_ASSIGN(auto test_copy, caller_maker(name, {uint8()})); ARROW_SCOPED_TRACE(test_copy->name()); - ResetContexts(); + // Preallocate independently for each batch { - std::vector args = {Datum(arr)}; exec_ctx_->set_preallocate_contiguous(false); exec_ctx_->set_exec_chunksize(40); ASSERT_OK_AND_ASSIGN(Datum result, test_copy->Call(args, exec_ctx_.get())); @@ -1712,52 +1726,51 @@ TEST_F(TestCallScalarFunctionPreallocationCases, Basic) { AssertArraysEqual(*arr->Slice(80), *carr.chunk(2)); } } + } +} - for (const auto& caller_maker : {DenseFunctionCaller::Maker}) { - ASSERT_OK_AND_ASSIGN(auto test_copy, caller_maker(name, {uint8()})); - ARROW_SCOPED_TRACE(test_copy->name()); - - ResetContexts(); +TEST_F(TestCallScalarFunctionPreallocationCases, IndependentPreallocateDense) { + auto arr = GetTestArray(); + for (const auto& name : {"test_copy", "test_copy_computed_bitmap"}) { + ARROW_SCOPED_TRACE(name); + ASSERT_OK_AND_ASSIGN(auto test_copy, DenseFunctionCaller::Maker(name, {uint8()})); + ResetContexts(); - // Preallocate independently for each batch - { - std::vector args = {Datum(arr)}; - exec_ctx_->set_preallocate_contiguous(false); - exec_ctx_->set_exec_chunksize(40); - ASSERT_OK_AND_ASSIGN(Datum result, test_copy->Call(args, exec_ctx_.get())); - ASSERT_EQ(Datum::CHUNKED_ARRAY, result.kind()); - const ChunkedArray& carr = *result.chunked_array(); - ASSERT_EQ(1, carr.num_chunks()); - AssertArraysEqual(*arr, *carr.chunk(0)); - } + // Preallocate independently for each batch + { + std::vector args = {Datum(arr)}; + exec_ctx_->set_preallocate_contiguous(false); + exec_ctx_->set_exec_chunksize(40); + ASSERT_OK_AND_ASSIGN(Datum result, test_copy->Call(args, exec_ctx_.get())); + ASSERT_EQ(Datum::CHUNKED_ARRAY, result.kind()); + const ChunkedArray& carr = *result.chunked_array(); + ASSERT_EQ(1, carr.num_chunks()); + AssertArraysEqual(*arr, *carr.chunk(0)); } } } -class TestCallScalarFunctionBasicNonStandardCases : public TestCallScalarFunction {}; +// Test a handful of cases +// +// * Validity bitmap computed by kernel rather than using PropagateNulls +// * Data not pre-allocated +// * Validity bitmap not pre-allocated +class TestCallScalarFunctionBasicNonStandardCases : public TestCallScalarFunction { + protected: + std::shared_ptr GetTestArray() { return GetUInt8Array(1000, 0.2); } +}; TEST_F(TestCallScalarFunctionBasicNonStandardCases, Basic) { - // Test a handful of cases - // - // * Validity bitmap computed by kernel rather than using PropagateNulls - // * Data not pre-allocated - // * Validity bitmap not pre-allocated - - double null_prob = 0.2; - - auto arr = GetUInt8Array(1000, null_prob); + auto arr = GetTestArray(); std::vector args = {Datum(arr)}; - for (const auto& name : {"test_nopre_data", "test_nopre_validity_or_data"}) { ARROW_SCOPED_TRACE(name); - for (const auto& caller_maker : {SimpleFunctionCaller::Maker, ExecFunctionCaller::Maker, ExpressionFunctionCaller::Maker, RegularSparseFunctionCaller::Maker, SelectiveSparseFunctionCaller::Maker, DenseFunctionCaller::Maker}) { ASSERT_OK_AND_ASSIGN(auto test_nopre, caller_maker(name, {uint8()})); ARROW_SCOPED_TRACE(test_nopre->name()); - ResetContexts(); // The default should be a single array output @@ -1766,14 +1779,20 @@ TEST_F(TestCallScalarFunctionBasicNonStandardCases, Basic) { AssertArraysEqual(*arr, *result.make_array(), true); } } + } +} +TEST_F(TestCallScalarFunctionBasicNonStandardCases, SplitExecution) { + auto arr = GetTestArray(); + std::vector args = {Datum(arr)}; + for (const auto& name : {"test_nopre_data", "test_nopre_validity_or_data"}) { + ARROW_SCOPED_TRACE(name); for (const auto& caller_maker : {SimpleFunctionCaller::Maker, ExecFunctionCaller::Maker, ExpressionFunctionCaller::Maker, RegularSparseFunctionCaller::Maker, SelectiveSparseFunctionCaller::Maker}) { ASSERT_OK_AND_ASSIGN(auto test_nopre, caller_maker(name, {uint8()})); ARROW_SCOPED_TRACE(test_nopre->name()); - ResetContexts(); // Split execution into 3 chunks @@ -1788,22 +1807,25 @@ TEST_F(TestCallScalarFunctionBasicNonStandardCases, Basic) { AssertArraysEqual(*arr->Slice(800), *carr.chunk(2)); } } + } +} - for (const auto& caller_maker : {DenseFunctionCaller::Maker}) { - ASSERT_OK_AND_ASSIGN(auto test_nopre, caller_maker(name, {uint8()})); - ARROW_SCOPED_TRACE(test_nopre->name()); - - ResetContexts(); +TEST_F(TestCallScalarFunctionBasicNonStandardCases, SplitExecutionDense) { + auto arr = GetTestArray(); + std::vector args = {Datum(arr)}; + for (const auto& name : {"test_nopre_data", "test_nopre_validity_or_data"}) { + ARROW_SCOPED_TRACE(name); + ASSERT_OK_AND_ASSIGN(auto test_nopre, DenseFunctionCaller::Maker(name, {uint8()})); + ResetContexts(); - // Split execution into 3 chunks - { - exec_ctx_->set_exec_chunksize(400); - ASSERT_OK_AND_ASSIGN(Datum result, test_nopre->Call(args, exec_ctx_.get())); - ASSERT_EQ(Datum::CHUNKED_ARRAY, result.kind()); - const ChunkedArray& carr = *result.chunked_array(); - ASSERT_EQ(1, carr.num_chunks()); - AssertArraysEqual(*arr, *carr.chunk(0)); - } + // Split execution into 3 chunks + { + exec_ctx_->set_exec_chunksize(400); + ASSERT_OK_AND_ASSIGN(Datum result, test_nopre->Call(args, exec_ctx_.get())); + ASSERT_EQ(Datum::CHUNKED_ARRAY, result.kind()); + const ChunkedArray& carr = *result.chunked_array(); + ASSERT_EQ(1, carr.num_chunks()); + AssertArraysEqual(*arr, *carr.chunk(0)); } } } @@ -1817,7 +1839,6 @@ TEST_F(TestCallScalarFunctionStatefulKernel, Basic) { SelectiveSparseFunctionCaller::Maker, DenseFunctionCaller::Maker}) { ASSERT_OK_AND_ASSIGN(auto test_stateful, caller_maker("test_stateful", {int32()})); ARROW_SCOPED_TRACE(test_stateful->name()); - ResetContexts(); auto input = ArrayFromJSON(int32(), "[1, 2, 3, null, 5]"); @@ -1841,7 +1862,6 @@ TEST_F(TestCallScalarFunctionScalarFunction, Basic) { ASSERT_OK_AND_ASSIGN(auto test_scalar_add_int32, caller_maker("test_scalar_add_int32", {int32(), int32()})); ARROW_SCOPED_TRACE(test_scalar_add_int32->name()); - ResetContexts(); std::vector args = {Datum(std::make_shared(5)), From 4f4c81b28f59cbf84e807d9c3450b84547d9a3bb Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Mon, 18 Aug 2025 21:44:55 +0800 Subject: [PATCH 22/71] Add selection argument to function caller --- cpp/src/arrow/compute/exec_test.cc | 44 +++++++++++++++++++----------- 1 file changed, 28 insertions(+), 16 deletions(-) diff --git a/cpp/src/arrow/compute/exec_test.cc b/cpp/src/arrow/compute/exec_test.cc index 86b765876b95..cb62c4c72390 100644 --- a/cpp/src/arrow/compute/exec_test.cc +++ b/cpp/src/arrow/compute/exec_test.cc @@ -1417,9 +1417,16 @@ class FunctionCaller { virtual std::string name() const = 0; virtual Result Call(const std::vector& args, + std::shared_ptr selection, const FunctionOptions* options, ExecContext* ctx = NULLPTR) const = 0; + virtual Result Call(const std::vector& args, + const FunctionOptions* options, + ExecContext* ctx = NULLPTR) const { + return Call(args, nullptr, options, ctx); + } + virtual Result Call(const std::vector& args, ExecContext* ctx = NULLPTR) const { return Call(args, /*options=*/nullptr, ctx); @@ -1444,8 +1451,11 @@ class SimpleFunctionCaller : public FunctionCaller { return Make(func_name); } - Result Call(const std::vector& args, const FunctionOptions* options, - ExecContext* ctx) const override { + Result Call(const std::vector& args, + std::shared_ptr selection, + const FunctionOptions* options, ExecContext* ctx) const override { + ARROW_RETURN_IF(selection != nullptr, + Status::Invalid("Selection vector not supported")); return CallFunction(func_name, args, options, ctx); } @@ -1482,8 +1492,11 @@ class ExecFunctionCaller : public FunctionCaller { return Make(func_name, std::move(in_types)); } - Result Call(const std::vector& args, const FunctionOptions* options, - ExecContext* ctx) const override { + Result Call(const std::vector& args, + std::shared_ptr selection, + const FunctionOptions* options, ExecContext* ctx) const override { + ARROW_RETURN_IF(selection != nullptr, + Status::Invalid("Selection vector not supported")); ARROW_RETURN_NOT_OK(func_exec->Init(options, ctx)); return func_exec->Execute(args); } @@ -1511,18 +1524,9 @@ class ExpressionFunctionCaller : public FunctionCaller { return std::make_shared(std::move(func_name), std::move(in_types)); } - Result Call(const std::vector& args, const FunctionOptions* options, - ExecContext* ctx) const override { - bool all_same = false; - auto length = InferBatchLength(args, &all_same); - ARROW_ASSIGN_OR_RAISE(auto selection, GetSelection(length)); - return CallWithSelection(args, options, std::move(selection), ctx); - } - - Result CallWithSelection(const std::vector& args, - const FunctionOptions* options, - std::shared_ptr selection, - ExecContext* ctx) const { + Result Call(const std::vector& args, + std::shared_ptr selection, + const FunctionOptions* options, ExecContext* ctx) const override { bool all_same = false; auto length = InferBatchLength(args, &all_same); ExecBatch batch(args, length, std::move(selection)); @@ -1536,6 +1540,14 @@ class ExpressionFunctionCaller : public FunctionCaller { return ExecuteScalarExpression(bound, batch, ctx); } + Result Call(const std::vector& args, const FunctionOptions* options, + ExecContext* ctx) const override { + bool all_same = false; + auto length = InferBatchLength(args, &all_same); + ARROW_ASSIGN_OR_RAISE(auto selection, GetSelection(length)); + return Call(args, std::move(selection), options, ctx); + } + static Result> Maker(const std::string& func_name, std::vector in_types) { return Make(func_name, std::move(in_types)); From dabdf5dc285d1474b91534238f6c1bbe9193d266 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Tue, 19 Aug 2025 16:40:11 +0800 Subject: [PATCH 23/71] Add partial for sparse and dense cases --- cpp/src/arrow/compute/exec.cc | 24 +- cpp/src/arrow/compute/exec_test.cc | 455 +++++++++++++++++++++++------ 2 files changed, 379 insertions(+), 100 deletions(-) diff --git a/cpp/src/arrow/compute/exec.cc b/cpp/src/arrow/compute/exec.cc index 73a65e28c0e5..a11d292997b9 100644 --- a/cpp/src/arrow/compute/exec.cc +++ b/cpp/src/arrow/compute/exec.cc @@ -726,7 +726,14 @@ std::shared_ptr ToChunkedArray(const std::vector& values, // Skip empty chunks continue; } - arrays.emplace_back(val.make_array()); + if (val.is_chunked_array()) { + for (const auto& chunk : val.chunked_array()->chunks()) { + arrays.emplace_back(chunk); + } + } else { + DCHECK(val.is_array()); + arrays.emplace_back(val.make_array()); + } } return std::make_shared(std::move(arrays), type.GetSharedPtr()); } @@ -887,23 +894,16 @@ class ScalarExecutor : public KernelExecutorImpl { Take(batch[i], *batch.selection_vector->data(), TakeOptions{/*boundcheck=*/false}, exec_context())); } - ExecBatch input; ARROW_ASSIGN_OR_RAISE( - input, ExecBatch::Make(std::move(values), batch.selection_vector->length())); + ExecBatch input, + ExecBatch::Make(std::move(values), batch.selection_vector->length())); DatumAccumulator dense_listener; RETURN_NOT_OK(ExecuteSparse(input, &dense_listener)); - auto dense_results = dense_listener.values(); - - Datum dense_datum; - if (dense_results.size() > 1) { - dense_datum = ToChunkedArray(dense_results, output_type_); - } else { - dense_datum = dense_results[0]; - } + Datum dense_result = WrapResults(input.values, dense_listener.values()); ARROW_ASSIGN_OR_RAISE(auto result, - Scatter(dense_datum, *batch.selection_vector->data(), + Scatter(dense_result, *batch.selection_vector->data(), ScatterOptions{/*max_index=*/batch.length - 1})); return listener->OnResult(std::move(result)); } diff --git a/cpp/src/arrow/compute/exec_test.cc b/cpp/src/arrow/compute/exec_test.cc index cb62c4c72390..3bc3feabff09 100644 --- a/cpp/src/arrow/compute/exec_test.cc +++ b/cpp/src/arrow/compute/exec_test.cc @@ -921,11 +921,22 @@ TEST_F(TestExecSpanIterator, ZeroLengthInputs) { CheckArgs(input); } +std::shared_ptr MakeSelectionVector(const std::string& json) { + return std::make_shared(*ArrayFromJSON(int32(), json)); +} + +std::shared_ptr MakeSelectionVectorUntil(int64_t length) { + auto res = gen::Step()->Generate(length); + DCHECK_OK(res.status()); + auto arr = res.ValueUnsafe(); + return std::make_shared(*arr); +} + TEST_F(TestExecSpanIterator, SelectionSpanBasic) { ExecBatch batch( {Datum(GetInt32Array(30)), Datum(GetInt32Array(30)), Datum(std::make_shared(5)), Datum(MakeNullScalar(boolean()))}, - 30, std::make_shared(*ArrayFromJSON(int32(), "[1, 2, 7, 29]"))); + 30, MakeSelectionVector("[1, 2, 7, 29]")); CheckIteration(batch, /*chunksize=*/7, {7, 7, 7, 7, 2}, {2, 1, 0, 0, 1}); CheckIteration(batch, /*chunksize=*/10, {10, 10, 10}, {3, 0, 1}); @@ -934,11 +945,10 @@ TEST_F(TestExecSpanIterator, SelectionSpanBasic) { } TEST_F(TestExecSpanIterator, SelectionSpanChunked) { - ExecBatch batch( - {Datum(GetInt32Chunked({0, 20, 10})), Datum(GetInt32Chunked({15, 15})), - Datum(GetInt32Array(30)), Datum(std::make_shared(5)), - Datum(MakeNullScalar(boolean()))}, - 30, std::make_shared(*ArrayFromJSON(int32(), "[1, 2, 7, 29]"))); + ExecBatch batch({Datum(GetInt32Chunked({0, 20, 10})), Datum(GetInt32Chunked({15, 15})), + Datum(GetInt32Array(30)), Datum(std::make_shared(5)), + Datum(MakeNullScalar(boolean()))}, + 30, MakeSelectionVector("[1, 2, 7, 29]")); CheckIteration(batch, /*chunksize=*/7, {7, 7, 1, 5, 7, 3}, {2, 1, 0, 0, 0, 1}); CheckIteration(batch, /*chunksize=*/10, {10, 5, 5, 10}, {3, 0, 0, 1}); @@ -950,9 +960,9 @@ TEST_F(TestExecSpanIterator, SelectionSpanChunked) { // Scalar function execution template -Status VisitSelectionVector(const SelectionVectorSpan& selection, int64_t length, - OnSelectionFn&& on_selection, - OnNonSelectionFn&& on_non_selection) { +void VisitSelectionVector(const SelectionVectorSpan& selection, int64_t length, + OnSelectionFn&& on_selection, + OnNonSelectionFn&& on_non_selection) { int64_t selected = 0; for (int64_t i = 0; i < length; ++i) { if (selected < selection.length() && i == selection[selected]) { @@ -962,7 +972,117 @@ Status VisitSelectionVector(const SelectionVectorSpan& selection, int64_t length on_non_selection(i); } } - return Status::OK(); +} + +constexpr uint8_t kNonSelectionValue = 0xFE; + +void AssertArraysEqualSparseWithSelection(const Array& src, + const SelectionVectorSpan& selection, + const Array& dst) { + ASSERT_EQ(src.length(), dst.length()); + ASSERT_EQ(src.type()->id(), dst.type()->id()); + + int value_size = src.type()->byte_width(); + const uint8_t* src_validity = src.data()->buffers[0]->data(); + const uint8_t* dst_validity = dst.data()->buffers[0]->data(); + const uint8_t* src_data = src.data()->buffers[1]->data(); + const uint8_t* dst_data = dst.data()->buffers[1]->data(); + int64_t src_offset = src.data()->offset; + int64_t dst_offset = dst.data()->offset; + + VisitSelectionVector( + selection, src.length(), + [&](int64_t i) { + // Selected values should match + ASSERT_EQ(bit_util::GetBit(src_validity, src_offset + i), + bit_util::GetBit(dst_validity, dst_offset + i)); + if (bit_util::GetBit(src_validity, src_offset + i)) { + ASSERT_EQ(memcmp(src_data + (src_offset + i) * value_size, + dst_data + (dst_offset + i) * value_size, value_size), + 0); + } + }, + [&](int64_t i) { + // Non-selected values should be the special value in the output + ASSERT_TRUE(bit_util::GetBit(dst_validity, dst_offset + i)); + for (int j = 0; j < value_size; ++j) { + ASSERT_EQ(dst_data[(dst_offset + i) * value_size + j], kNonSelectionValue) + << "at index " << i << " and byte " << j; + } + }); +} + +void AssertArraysEqualDenseWithSelection(const Array& src, + const SelectionVectorSpan& selection, + const Array& dst) { + ASSERT_EQ(src.length(), dst.length()); + ASSERT_EQ(src.type()->id(), dst.type()->id()); + + int value_size = src.type()->byte_width(); + const uint8_t* src_validity = src.data()->buffers[0]->data(); + const uint8_t* dst_validity = dst.data()->buffers[0]->data(); + const uint8_t* src_data = src.data()->buffers[1]->data(); + const uint8_t* dst_data = dst.data()->buffers[1]->data(); + int64_t src_offset = src.data()->offset; + int64_t dst_offset = dst.data()->offset; + + VisitSelectionVector( + selection, src.length(), + [&](int64_t i) { + // Selected values should match + ASSERT_EQ(bit_util::GetBit(src_validity, src_offset + i), + bit_util::GetBit(dst_validity, dst_offset + i)); + if (bit_util::GetBit(src_validity, src_offset + i)) { + ASSERT_EQ(memcmp(src_data + (src_offset + i) * value_size, + dst_data + (dst_offset + i) * value_size, value_size), + 0); + } + }, + [&](int64_t i) { + // Non-selected values should be invalid in the output + ASSERT_FALSE(bit_util::GetBit(dst_validity, dst_offset + i)); + }); +} + +void AssertChunkedExecResultsEqualSparseWithSelection(int64_t exec_chunksize, + const Array& input, + const SelectionVector* selection, + const Datum& result) { + ASSERT_EQ(Datum::CHUNKED_ARRAY, result.kind()); + const ChunkedArray& carr = *result.chunked_array(); + SelectionVectorSpan selection_span(selection->indices(), selection->length()); + ASSERT_EQ(bit_util::CeilDiv(input.length(), exec_chunksize), carr.num_chunks()); + int64_t selection_idx = 0; + for (int i = 0; i < carr.num_chunks(); ++i) { + auto next_selection_idx = selection_idx; + while (next_selection_idx < selection->length() && + selection->indices()[next_selection_idx] < exec_chunksize * (i + 1)) { + ++next_selection_idx; + } + selection_span.SetSlice(selection_idx, next_selection_idx - selection_idx, + static_cast(exec_chunksize * i)); + selection_idx = next_selection_idx; + AssertArraysEqualSparseWithSelection( + *input.Slice(exec_chunksize * i, + std::min(exec_chunksize, input.length() - exec_chunksize * i)), + selection_span, *carr.chunk(i)); + } +} + +void AssertChunkedExecResultsEqualDenseWithSelection(int64_t exec_chunksize, + const Array& input, + const SelectionVector* selection, + const Datum& result) { + SelectionVectorSpan selection_span(selection->indices(), selection->length()); + if (selection_span.length() <= exec_chunksize) { + ASSERT_EQ(Datum::ARRAY, result.kind()); + AssertArraysEqualDenseWithSelection(input, selection_span, *result.make_array()); + } else { + ASSERT_EQ(Datum::CHUNKED_ARRAY, result.kind()); + const ChunkedArray& carr = *result.chunked_array(); + ASSERT_EQ(1, carr.num_chunks()); + AssertArraysEqualDenseWithSelection(input, selection_span, *carr.chunk(0)); + } } Status ExecCopyArrayData(KernelContext*, const ExecSpan& batch, ExecResult* out) { @@ -977,8 +1097,6 @@ Status ExecCopyArrayData(KernelContext*, const ExecSpan& batch, ExecResult* out) return Status::OK(); } -constexpr int8_t kNonSelectionValue = 0xFE; - Status SelectiveExecCopyArrayData(KernelContext* ctx, const ExecSpan& batch, const SelectionVectorSpan& selection, ExecResult* out) { DCHECK_EQ(1, batch.num_values()); @@ -990,7 +1108,7 @@ Status SelectiveExecCopyArrayData(KernelContext* ctx, const ExecSpan& batch, int64_t dst_validity_offset = out_arr->offset; uint8_t* dst = out_arr->buffers[1]->mutable_data() + out_arr->offset * value_size; const uint8_t* src = arg0.buffers[1].data + arg0.offset * value_size; - return VisitSelectionVector( + VisitSelectionVector( selection, batch.length, [&](int64_t i) { // Copy the selected value @@ -1026,7 +1144,7 @@ Status SelectiveExecCopyArraySpan(KernelContext* ctx, const ExecSpan& batch, int64_t dst_validity_offset = out_arr->offset; uint8_t* dst = out_arr->buffers[1].data + out_arr->offset * value_size; const uint8_t* src = arg0.buffers[1].data + arg0.offset * value_size; - return VisitSelectionVector( + VisitSelectionVector( selection, batch.length, [&](int64_t i) { // Copy the selected value @@ -1038,6 +1156,7 @@ Status SelectiveExecCopyArraySpan(KernelContext* ctx, const ExecSpan& batch, bit_util::SetBit(dst_validity, dst_validity_offset + i); std::memset(dst + i * value_size, kNonSelectionValue, value_size); }); + return Status::OK(); } Status ExecComputedBitmap(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) { @@ -1194,7 +1313,7 @@ Status SelectiveExecStateful(KernelContext* ctx, const ExecSpan& batch, uint8_t* dst_validity = out_arr->buffers[0].data; int64_t dst_validity_offset = out_arr->offset; int32_t* dst = out_arr->GetValues(1); - return VisitSelectionVector( + VisitSelectionVector( selection, batch.length, [&](int64_t i) { // Copy the selected value @@ -1206,6 +1325,7 @@ Status SelectiveExecStateful(KernelContext* ctx, const ExecSpan& batch, bit_util::SetBit(dst_validity, dst_validity_offset + i); dst[i] = kNonSelectionValue; }); + return Status::OK(); } Status ExecAddInt32(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) { @@ -1226,7 +1346,7 @@ Status SelectiveExecAddInt32(KernelContext* ctx, const ExecSpan& batch, uint8_t* dst_validity = out_arr->buffers[0].data; int64_t dst_validity_offset = out_arr->offset; int32_t* out_data = out_arr->GetValues(1); - return VisitSelectionVector( + VisitSelectionVector( selection, batch.length, [&](int64_t i) { // Copy the selected value @@ -1238,6 +1358,7 @@ Status SelectiveExecAddInt32(KernelContext* ctx, const ExecSpan& batch, bit_util::SetBit(dst_validity, dst_validity_offset + i); out_data[i] = kNonSelectionValue; }); + return Status::OK(); } class TestCallScalarFunction : public TestComputeInternals { @@ -1418,7 +1539,7 @@ class FunctionCaller { virtual Result Call(const std::vector& args, std::shared_ptr selection, - const FunctionOptions* options, + const FunctionOptions* options = NULLPTR, ExecContext* ctx = NULLPTR) const = 0; virtual Result Call(const std::vector& args, @@ -1570,28 +1691,28 @@ class ExpressionFunctionCaller : public FunctionCaller { // Call the selective counterpart of the function via expression with no selection vector, // triggering the regular sparse execution path. -class RegularSparseFunctionCaller : public ExpressionFunctionCaller { +class NullSelectionSparseFunctionCaller : public ExpressionFunctionCaller { public: - RegularSparseFunctionCaller(std::string func_name, - const std::vector& in_types) + NullSelectionSparseFunctionCaller(std::string func_name, + const std::vector& in_types) : ExpressionFunctionCaller(std::move(func_name) + "_selective", in_types) {} std::string name() const override { return "regular_sparse_caller"; } static Result> Maker(const std::string& func_name, std::vector in_types) { - return ExpressionFunctionCaller::Make( + return ExpressionFunctionCaller::Make( std::move(func_name), std::move(in_types)); } }; // Call the selective counterpart of the function via expression with full selection, // triggering the selective sparse execution path. -class SelectiveSparseFunctionCaller : public RegularSparseFunctionCaller { +class SelectiveSparseFunctionCaller : public NullSelectionSparseFunctionCaller { public: SelectiveSparseFunctionCaller(std::string func_name, const std::vector& in_types) - : RegularSparseFunctionCaller(std::move(func_name), in_types) {} + : NullSelectionSparseFunctionCaller(std::move(func_name), in_types) {} std::string name() const override { return "selective_sparse_caller"; } @@ -1630,10 +1751,8 @@ class DenseFunctionCaller : public ExpressionFunctionCaller { class TestCallScalarFunctionArgumentValidation : public TestCallScalarFunction {}; TEST_F(TestCallScalarFunctionArgumentValidation, Basic) { - for (const auto& caller_maker : - {SimpleFunctionCaller::Maker, ExecFunctionCaller::Maker, - ExpressionFunctionCaller::Maker, RegularSparseFunctionCaller::Maker, - SelectiveSparseFunctionCaller::Maker, DenseFunctionCaller::Maker}) { + for (const auto& caller_maker : {SimpleFunctionCaller::Maker, ExecFunctionCaller::Maker, + NullSelectionSparseFunctionCaller::Maker}) { ASSERT_OK_AND_ASSIGN(auto test_copy, caller_maker("test_copy", {int32()})); ARROW_SCOPED_TRACE(test_copy->name()); ResetContexts(); @@ -1659,6 +1778,75 @@ TEST_F(TestCallScalarFunctionArgumentValidation, Basic) { class TestCallScalarFunctionPreallocationCases : public TestCallScalarFunction { protected: std::shared_ptr GetTestArray() { return GetUInt8Array(100, 0.2); } + + std::vector> GetSelectionVectors() { + return {MakeSelectionVector("[]"), + MakeSelectionVector("[0]"), + MakeSelectionVector("[42]"), + MakeSelectionVector("[99]"), + MakeSelectionVector("[0, 1, 2, 3, 4]"), + MakeSelectionVector("[0, 42, 99]"), + MakeSelectionVectorUntil(40), + MakeSelectionVectorUntil(41), + MakeSelectionVectorUntil(99), + MakeSelectionVectorUntil(100)}; + } + + template + void DoTestBasic(const FunctionCaller* caller, const Array& input, + std::shared_ptr selection, CheckFunc&& check_func) { + // The default should be a single array output + { + std::vector args = {Datum(input)}; + ASSERT_OK_AND_ASSIGN(Datum result, caller->Call(args, selection)); + check_func(result); + } + + // Set the exec_chunksize to be smaller, so now we have several invocations + // of the kernel, but still the output is one array + { + std::vector args = {Datum(input)}; + exec_ctx_->set_exec_chunksize(80); + ASSERT_OK_AND_ASSIGN( + Datum result, + caller->Call(args, selection, /*options=*/nullptr, exec_ctx_.get())); + check_func(result); + } + + { + // Chunksize not multiple of 8 + std::vector args = {Datum(input)}; + exec_ctx_->set_exec_chunksize(11); + ASSERT_OK_AND_ASSIGN( + Datum result, + caller->Call(args, selection, /*options=*/nullptr, exec_ctx_.get())); + check_func(result); + } + } + + template + void DoTestChunked(const FunctionCaller* caller, const ChunkedArray& input, + std::shared_ptr selection, CheckFunc&& check_func) { + // Input is chunked, output has one big chunk + std::vector args = {Datum(input)}; + ASSERT_OK_AND_ASSIGN(Datum result, caller->Call(args, selection, /*options=*/nullptr, + exec_ctx_.get())); + check_func(result); + } + + template + void DoTestIndependentPreallocate(const FunctionCaller* caller, int64_t exec_chunksize, + const Array& input, + std::shared_ptr selection, + CheckFunc&& check_func) { + // Preallocate independently for each batch + std::vector args = {Datum(input)}; + exec_ctx_->set_preallocate_contiguous(false); + exec_ctx_->set_exec_chunksize(exec_chunksize); + ASSERT_OK_AND_ASSIGN(Datum result, caller->Call(args, selection, /*options=*/nullptr, + exec_ctx_.get())); + check_func(result); + } }; TEST_F(TestCallScalarFunctionPreallocationCases, Basic) { @@ -1667,97 +1855,188 @@ TEST_F(TestCallScalarFunctionPreallocationCases, Basic) { ARROW_SCOPED_TRACE(name); for (const auto& caller_maker : {SimpleFunctionCaller::Maker, ExecFunctionCaller::Maker, - ExpressionFunctionCaller::Maker, RegularSparseFunctionCaller::Maker, - SelectiveSparseFunctionCaller::Maker, DenseFunctionCaller::Maker}) { + NullSelectionSparseFunctionCaller::Maker}) { ASSERT_OK_AND_ASSIGN(auto test_copy, caller_maker(name, {uint8()})); ARROW_SCOPED_TRACE(test_copy->name()); ResetContexts(); - // The default should be a single array output - { - std::vector args = {Datum(arr)}; - ASSERT_OK_AND_ASSIGN(Datum result, test_copy->Call(args)); + DoTestBasic(test_copy.get(), *arr, /*selection=*/nullptr, [&](const Datum& result) { ASSERT_EQ(Datum::ARRAY, result.kind()); AssertArraysEqual(*arr, *result.make_array()); - } + }); + } + } +} - // Set the exec_chunksize to be smaller, so now we have several invocations - // of the kernel, but still the output is one array - { - std::vector args = {Datum(arr)}; - exec_ctx_->set_exec_chunksize(80); - ASSERT_OK_AND_ASSIGN(Datum result, test_copy->Call(args, exec_ctx_.get())); - AssertArraysEqual(*arr, *result.make_array()); - } +TEST_F(TestCallScalarFunctionPreallocationCases, BasicSelectiveSparse) { + auto arr = GetTestArray(); + auto selections = GetSelectionVectors(); + for (const auto& name : + {"test_copy_selective", "test_copy_computed_bitmap_selective"}) { + ARROW_SCOPED_TRACE(name); + ASSERT_OK_AND_ASSIGN(auto test_copy, + ExpressionFunctionCaller::Maker(name, {uint8()})); + for (const auto& selection : selections) { + SelectionVectorSpan selection_span(selection->indices(), selection->length()); + ResetContexts(); - { - // Chunksize not multiple of 8 - std::vector args = {Datum(arr)}; - exec_ctx_->set_exec_chunksize(11); - ASSERT_OK_AND_ASSIGN(Datum result, test_copy->Call(args, exec_ctx_.get())); - AssertArraysEqual(*arr, *result.make_array()); - } + DoTestBasic(test_copy.get(), *arr, selection, [&](const Datum& result) { + ASSERT_EQ(Datum::ARRAY, result.kind()); + AssertArraysEqualSparseWithSelection(*arr, selection_span, *result.make_array()); + }); + } + } +} - // Input is chunked, output has one big chunk - { - auto carr = std::make_shared( - ArrayVector{arr->Slice(0, 10), arr->Slice(10)}); - std::vector args = {Datum(carr)}; - ASSERT_OK_AND_ASSIGN(Datum result, test_copy->Call(args, exec_ctx_.get())); +TEST_F(TestCallScalarFunctionPreallocationCases, BasicSelectiveDense) { + auto arr = GetTestArray(); + auto selections = GetSelectionVectors(); + for (const auto& name : {"test_copy", "test_copy_computed_bitmap"}) { + ARROW_SCOPED_TRACE(name); + ASSERT_OK_AND_ASSIGN(auto test_copy, + ExpressionFunctionCaller::Maker(name, {uint8()})); + for (const auto& selection : selections) { + SelectionVectorSpan selection_span(selection->indices(), selection->length()); + ResetContexts(); + + DoTestBasic(test_copy.get(), *arr, selection, [&](const Datum& result) { + ASSERT_EQ(Datum::ARRAY, result.kind()); + AssertArraysEqualDenseWithSelection(*arr, selection_span, *result.make_array()); + }); + } + } +} + +TEST_F(TestCallScalarFunctionPreallocationCases, Chunked) { + auto arr = GetTestArray(); + auto carr = + std::make_shared(ArrayVector{arr->Slice(0, 10), arr->Slice(10)}); + for (const auto& name : {"test_copy", "test_copy_computed_bitmap"}) { + ARROW_SCOPED_TRACE(name); + for (const auto& caller_maker : + {SimpleFunctionCaller::Maker, ExecFunctionCaller::Maker, + NullSelectionSparseFunctionCaller::Maker}) { + ASSERT_OK_AND_ASSIGN(auto test_copy, caller_maker(name, {uint8()})); + ARROW_SCOPED_TRACE(test_copy->name()); + ResetContexts(); + + DoTestChunked(test_copy.get(), *carr, /*selection=*/nullptr, + [&](const Datum& result) { + ASSERT_EQ(Datum::CHUNKED_ARRAY, result.kind()); + std::shared_ptr actual = result.chunked_array(); + ASSERT_EQ(1, actual->num_chunks()); + AssertChunkedEquivalent(*carr, *actual); + }); + } + } +} + +TEST_F(TestCallScalarFunctionPreallocationCases, ChunkedSelectiveSparse) { + auto arr = GetTestArray(); + auto carr = + std::make_shared(ArrayVector{arr->Slice(0, 10), arr->Slice(10)}); + auto selections = GetSelectionVectors(); + for (const auto& name : + {"test_copy_selective", "test_copy_computed_bitmap_selective"}) { + ARROW_SCOPED_TRACE(name); + ASSERT_OK_AND_ASSIGN(auto test_copy, + ExpressionFunctionCaller::Maker(name, {uint8()})); + for (const auto& selection : selections) { + SelectionVectorSpan selection_span(selection->indices(), selection->length()); + ResetContexts(); + + DoTestChunked(test_copy.get(), *carr, selection, [&](const Datum& result) { + ASSERT_EQ(Datum::CHUNKED_ARRAY, result.kind()); std::shared_ptr actual = result.chunked_array(); ASSERT_EQ(1, actual->num_chunks()); - AssertChunkedEquivalent(*carr, *actual); - } + AssertArraysEqualSparseWithSelection(*arr, selection_span, *actual->chunk(0)); + }); + } + } +} + +TEST_F(TestCallScalarFunctionPreallocationCases, ChunkedSelectiveDense) { + auto arr = GetTestArray(); + auto carr = + std::make_shared(ArrayVector{arr->Slice(0, 10), arr->Slice(10)}); + auto selections = GetSelectionVectors(); + for (const auto& name : {"test_copy", "test_copy_computed_bitmap"}) { + ARROW_SCOPED_TRACE(name); + ASSERT_OK_AND_ASSIGN(auto test_copy, + ExpressionFunctionCaller::Maker(name, {uint8()})); + for (const auto& selection : selections) { + SelectionVectorSpan selection_span(selection->indices(), selection->length()); + ResetContexts(); + + DoTestChunked(test_copy.get(), *carr, selection, [&](const Datum& result) { + ASSERT_EQ(Datum::CHUNKED_ARRAY, result.kind()); + std::shared_ptr actual = result.chunked_array(); + ASSERT_EQ(1, actual->num_chunks()); + AssertArraysEqualDenseWithSelection(*arr, selection_span, *actual->chunk(0)); + }); } } } TEST_F(TestCallScalarFunctionPreallocationCases, IndependentPreallocate) { auto arr = GetTestArray(); - std::vector args = {Datum(arr)}; for (const auto& name : {"test_copy", "test_copy_computed_bitmap"}) { ARROW_SCOPED_TRACE(name); for (const auto& caller_maker : {SimpleFunctionCaller::Maker, ExecFunctionCaller::Maker, - ExpressionFunctionCaller::Maker, RegularSparseFunctionCaller::Maker, - SelectiveSparseFunctionCaller::Maker}) { + NullSelectionSparseFunctionCaller::Maker}) { ASSERT_OK_AND_ASSIGN(auto test_copy, caller_maker(name, {uint8()})); ARROW_SCOPED_TRACE(test_copy->name()); ResetContexts(); - // Preallocate independently for each batch - { - exec_ctx_->set_preallocate_contiguous(false); - exec_ctx_->set_exec_chunksize(40); - ASSERT_OK_AND_ASSIGN(Datum result, test_copy->Call(args, exec_ctx_.get())); - ASSERT_EQ(Datum::CHUNKED_ARRAY, result.kind()); - const ChunkedArray& carr = *result.chunked_array(); - ASSERT_EQ(3, carr.num_chunks()); - AssertArraysEqual(*arr->Slice(0, 40), *carr.chunk(0)); - AssertArraysEqual(*arr->Slice(40, 40), *carr.chunk(1)); - AssertArraysEqual(*arr->Slice(80), *carr.chunk(2)); - } + DoTestIndependentPreallocate( + test_copy.get(), /*exec_chunksize=*/40, *arr, /*selection=*/nullptr, + [&](const Datum& result) { + ASSERT_EQ(Datum::CHUNKED_ARRAY, result.kind()); + const ChunkedArray& carr = *result.chunked_array(); + ASSERT_EQ(3, carr.num_chunks()); + AssertArraysEqual(*arr->Slice(0, 40), *carr.chunk(0)); + AssertArraysEqual(*arr->Slice(40, 40), *carr.chunk(1)); + AssertArraysEqual(*arr->Slice(80), *carr.chunk(2)); + }); } } } -TEST_F(TestCallScalarFunctionPreallocationCases, IndependentPreallocateDense) { +TEST_F(TestCallScalarFunctionPreallocationCases, IndependentPreallocateSelectiveSparse) { auto arr = GetTestArray(); + auto selections = GetSelectionVectors(); + for (const auto& name : + {"test_copy_selective", "test_copy_computed_bitmap_selective"}) { + ARROW_SCOPED_TRACE(name); + ASSERT_OK_AND_ASSIGN(auto test_copy, DenseFunctionCaller::Maker(name, {uint8()})); + for (const auto& selection : selections) { + ResetContexts(); + + const int64_t exec_chunksize = 40; + DoTestIndependentPreallocate(test_copy.get(), exec_chunksize, *arr, selection, + [&](const Datum& result) { + AssertChunkedExecResultsEqualSparseWithSelection( + exec_chunksize, *arr, selection.get(), result); + }); + } + } +} + +TEST_F(TestCallScalarFunctionPreallocationCases, IndependentPreallocateSelectiveDense) { + auto arr = GetTestArray(); + auto selections = GetSelectionVectors(); for (const auto& name : {"test_copy", "test_copy_computed_bitmap"}) { ARROW_SCOPED_TRACE(name); ASSERT_OK_AND_ASSIGN(auto test_copy, DenseFunctionCaller::Maker(name, {uint8()})); - ResetContexts(); + for (const auto& selection : selections) { + ResetContexts(); - // Preallocate independently for each batch - { - std::vector args = {Datum(arr)}; - exec_ctx_->set_preallocate_contiguous(false); - exec_ctx_->set_exec_chunksize(40); - ASSERT_OK_AND_ASSIGN(Datum result, test_copy->Call(args, exec_ctx_.get())); - ASSERT_EQ(Datum::CHUNKED_ARRAY, result.kind()); - const ChunkedArray& carr = *result.chunked_array(); - ASSERT_EQ(1, carr.num_chunks()); - AssertArraysEqual(*arr, *carr.chunk(0)); + DoTestIndependentPreallocate(test_copy.get(), /*exec_chunksize=*/40, *arr, + selection, [&](const Datum& result) { + AssertChunkedExecResultsEqualDenseWithSelection( + 40, *arr, selection.get(), result); + }); } } } @@ -1779,7 +2058,7 @@ TEST_F(TestCallScalarFunctionBasicNonStandardCases, Basic) { ARROW_SCOPED_TRACE(name); for (const auto& caller_maker : {SimpleFunctionCaller::Maker, ExecFunctionCaller::Maker, - ExpressionFunctionCaller::Maker, RegularSparseFunctionCaller::Maker, + ExpressionFunctionCaller::Maker, NullSelectionSparseFunctionCaller::Maker, SelectiveSparseFunctionCaller::Maker, DenseFunctionCaller::Maker}) { ASSERT_OK_AND_ASSIGN(auto test_nopre, caller_maker(name, {uint8()})); ARROW_SCOPED_TRACE(test_nopre->name()); @@ -1801,7 +2080,7 @@ TEST_F(TestCallScalarFunctionBasicNonStandardCases, SplitExecution) { ARROW_SCOPED_TRACE(name); for (const auto& caller_maker : {SimpleFunctionCaller::Maker, ExecFunctionCaller::Maker, - ExpressionFunctionCaller::Maker, RegularSparseFunctionCaller::Maker, + ExpressionFunctionCaller::Maker, NullSelectionSparseFunctionCaller::Maker, SelectiveSparseFunctionCaller::Maker}) { ASSERT_OK_AND_ASSIGN(auto test_nopre, caller_maker(name, {uint8()})); ARROW_SCOPED_TRACE(test_nopre->name()); @@ -1847,7 +2126,7 @@ class TestCallScalarFunctionStatefulKernel : public TestCallScalarFunction {}; TEST_F(TestCallScalarFunctionStatefulKernel, Basic) { for (const auto& caller_maker : {SimpleFunctionCaller::Maker, ExecFunctionCaller::Maker, - ExpressionFunctionCaller::Maker, RegularSparseFunctionCaller::Maker, + ExpressionFunctionCaller::Maker, NullSelectionSparseFunctionCaller::Maker, SelectiveSparseFunctionCaller::Maker, DenseFunctionCaller::Maker}) { ASSERT_OK_AND_ASSIGN(auto test_stateful, caller_maker("test_stateful", {int32()})); ARROW_SCOPED_TRACE(test_stateful->name()); @@ -1869,7 +2148,7 @@ class TestCallScalarFunctionScalarFunction : public TestCallScalarFunction {}; TEST_F(TestCallScalarFunctionScalarFunction, Basic) { for (const auto& caller_maker : {SimpleFunctionCaller::Maker, ExecFunctionCaller::Maker, - ExpressionFunctionCaller::Maker, RegularSparseFunctionCaller::Maker, + ExpressionFunctionCaller::Maker, NullSelectionSparseFunctionCaller::Maker, SelectiveSparseFunctionCaller::Maker, DenseFunctionCaller::Maker}) { ASSERT_OK_AND_ASSIGN(auto test_scalar_add_int32, caller_maker("test_scalar_add_int32", {int32(), int32()})); From 14a5523484739f6a3f91deee1debefd4cae8da44 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Tue, 19 Aug 2025 18:19:24 +0800 Subject: [PATCH 24/71] Test finished? --- cpp/src/arrow/compute/exec_test.cc | 429 ++++++++++++++++++----------- 1 file changed, 271 insertions(+), 158 deletions(-) diff --git a/cpp/src/arrow/compute/exec_test.cc b/cpp/src/arrow/compute/exec_test.cc index 3bc3feabff09..483774cf04c1 100644 --- a/cpp/src/arrow/compute/exec_test.cc +++ b/cpp/src/arrow/compute/exec_test.cc @@ -974,7 +974,7 @@ void VisitSelectionVector(const SelectionVectorSpan& selection, int64_t length, } } -constexpr uint8_t kNonSelectionValue = 0xFE; +constexpr uint8_t kNonSelectedByte = 0xFE; void AssertArraysEqualSparseWithSelection(const Array& src, const SelectionVectorSpan& selection, @@ -1003,11 +1003,10 @@ void AssertArraysEqualSparseWithSelection(const Array& src, } }, [&](int64_t i) { - // Non-selected values should be the special value in the output + // Non-selected values should be the valid special value in the output ASSERT_TRUE(bit_util::GetBit(dst_validity, dst_offset + i)); for (int j = 0; j < value_size; ++j) { - ASSERT_EQ(dst_data[(dst_offset + i) * value_size + j], kNonSelectionValue) - << "at index " << i << " and byte " << j; + ASSERT_EQ(dst_data[(dst_offset + i) * value_size + j], kNonSelectedByte); } }); } @@ -1118,7 +1117,7 @@ Status SelectiveExecCopyArrayData(KernelContext* ctx, const ExecSpan& batch, // Set the non-selected as valid (regardless of its precomputed validity) and set // its values with a special value bit_util::SetBit(dst_validity, dst_validity_offset + i); - std::memset(dst + i * value_size, kNonSelectionValue, value_size); + std::memset(dst + i * value_size, kNonSelectedByte, value_size); }); return Status::OK(); } @@ -1154,7 +1153,7 @@ Status SelectiveExecCopyArraySpan(KernelContext* ctx, const ExecSpan& batch, // Set the non-selected as valid (regardless of its precomputed validity) and set // its values with a special value bit_util::SetBit(dst_validity, dst_validity_offset + i); - std::memset(dst + i * value_size, kNonSelectionValue, value_size); + std::memset(dst + i * value_size, kNonSelectedByte, value_size); }); return Status::OK(); } @@ -1323,7 +1322,7 @@ Status SelectiveExecStateful(KernelContext* ctx, const ExecSpan& batch, // Set the non-selected as valid (regardless of its precomputed validity) and set // its values with a special value bit_util::SetBit(dst_validity, dst_validity_offset + i); - dst[i] = kNonSelectionValue; + memset(dst + i, kNonSelectedByte, sizeof(int32_t)); }); return Status::OK(); } @@ -1356,7 +1355,7 @@ Status SelectiveExecAddInt32(KernelContext* ctx, const ExecSpan& batch, // Set the non-selected as valid (regardless of its precomputed validity) and set // its values with a special value bit_util::SetBit(dst_validity, dst_validity_offset + i); - out_data[i] = kNonSelectionValue; + memset(out_data + i, kNonSelectedByte, sizeof(int32_t)); }); return Status::OK(); } @@ -1639,10 +1638,10 @@ class ExpressionFunctionCaller : public FunctionCaller { std::string name() const override { return "expression_caller"; } - template static Result> Make(std::string func_name, std::vector in_types) { - return std::make_shared(std::move(func_name), std::move(in_types)); + return std::make_shared(std::move(func_name), + std::move(in_types)); } Result Call(const std::vector& args, @@ -1663,25 +1662,12 @@ class ExpressionFunctionCaller : public FunctionCaller { Result Call(const std::vector& args, const FunctionOptions* options, ExecContext* ctx) const override { - bool all_same = false; - auto length = InferBatchLength(args, &all_same); - ARROW_ASSIGN_OR_RAISE(auto selection, GetSelection(length)); - return Call(args, std::move(selection), options, ctx); + return Call(args, /*selection=*/nullptr, options, ctx); } static Result> Maker(const std::string& func_name, std::vector in_types) { - return Make(func_name, std::move(in_types)); - } - - protected: - virtual Result> GetSelection(int64_t length) const { - return nullptr; - } - - Result> MakeFullSelection(int64_t length) const { - ARROW_ASSIGN_OR_RAISE(auto arr, gen::Step()->Generate(length)); - return std::make_shared(*arr); + return Make(func_name, std::move(in_types)); } private: @@ -1689,70 +1675,11 @@ class ExpressionFunctionCaller : public FunctionCaller { std::shared_ptr schema_; }; -// Call the selective counterpart of the function via expression with no selection vector, -// triggering the regular sparse execution path. -class NullSelectionSparseFunctionCaller : public ExpressionFunctionCaller { - public: - NullSelectionSparseFunctionCaller(std::string func_name, - const std::vector& in_types) - : ExpressionFunctionCaller(std::move(func_name) + "_selective", in_types) {} - - std::string name() const override { return "regular_sparse_caller"; } - - static Result> Maker(const std::string& func_name, - std::vector in_types) { - return ExpressionFunctionCaller::Make( - std::move(func_name), std::move(in_types)); - } -}; - -// Call the selective counterpart of the function via expression with full selection, -// triggering the selective sparse execution path. -class SelectiveSparseFunctionCaller : public NullSelectionSparseFunctionCaller { - public: - SelectiveSparseFunctionCaller(std::string func_name, - const std::vector& in_types) - : NullSelectionSparseFunctionCaller(std::move(func_name), in_types) {} - - std::string name() const override { return "selective_sparse_caller"; } - - static Result> Maker(const std::string& func_name, - std::vector in_types) { - return ExpressionFunctionCaller::Make( - std::move(func_name), std::move(in_types)); - } - - protected: - Result> GetSelection(int64_t length) const override { - return MakeFullSelection(length); - } -}; - -// Call the non-selective function via expression with full selection, triggering the -// dense execution path. -class DenseFunctionCaller : public ExpressionFunctionCaller { - public: - using ExpressionFunctionCaller::ExpressionFunctionCaller; - - std::string name() const override { return "dense_caller"; } - - static Result> Maker(const std::string& func_name, - std::vector in_types) { - return ExpressionFunctionCaller::Make(std::move(func_name), - std::move(in_types)); - } - - protected: - Result> GetSelection(int64_t length) const override { - return MakeFullSelection(length); - } -}; - class TestCallScalarFunctionArgumentValidation : public TestCallScalarFunction {}; TEST_F(TestCallScalarFunctionArgumentValidation, Basic) { for (const auto& caller_maker : {SimpleFunctionCaller::Maker, ExecFunctionCaller::Maker, - NullSelectionSparseFunctionCaller::Maker}) { + ExpressionFunctionCaller::Maker}) { ASSERT_OK_AND_ASSIGN(auto test_copy, caller_maker("test_copy", {int32()})); ARROW_SCOPED_TRACE(test_copy->name()); ResetContexts(); @@ -1779,7 +1706,7 @@ class TestCallScalarFunctionPreallocationCases : public TestCallScalarFunction { protected: std::shared_ptr GetTestArray() { return GetUInt8Array(100, 0.2); } - std::vector> GetSelectionVectors() { + std::vector> GetTestSelectionVectors() { return {MakeSelectionVector("[]"), MakeSelectionVector("[0]"), MakeSelectionVector("[42]"), @@ -1855,7 +1782,7 @@ TEST_F(TestCallScalarFunctionPreallocationCases, Basic) { ARROW_SCOPED_TRACE(name); for (const auto& caller_maker : {SimpleFunctionCaller::Maker, ExecFunctionCaller::Maker, - NullSelectionSparseFunctionCaller::Maker}) { + ExpressionFunctionCaller::Maker}) { ASSERT_OK_AND_ASSIGN(auto test_copy, caller_maker(name, {uint8()})); ARROW_SCOPED_TRACE(test_copy->name()); ResetContexts(); @@ -1870,7 +1797,7 @@ TEST_F(TestCallScalarFunctionPreallocationCases, Basic) { TEST_F(TestCallScalarFunctionPreallocationCases, BasicSelectiveSparse) { auto arr = GetTestArray(); - auto selections = GetSelectionVectors(); + auto selections = GetTestSelectionVectors(); for (const auto& name : {"test_copy_selective", "test_copy_computed_bitmap_selective"}) { ARROW_SCOPED_TRACE(name); @@ -1890,7 +1817,7 @@ TEST_F(TestCallScalarFunctionPreallocationCases, BasicSelectiveSparse) { TEST_F(TestCallScalarFunctionPreallocationCases, BasicSelectiveDense) { auto arr = GetTestArray(); - auto selections = GetSelectionVectors(); + auto selections = GetTestSelectionVectors(); for (const auto& name : {"test_copy", "test_copy_computed_bitmap"}) { ARROW_SCOPED_TRACE(name); ASSERT_OK_AND_ASSIGN(auto test_copy, @@ -1915,7 +1842,7 @@ TEST_F(TestCallScalarFunctionPreallocationCases, Chunked) { ARROW_SCOPED_TRACE(name); for (const auto& caller_maker : {SimpleFunctionCaller::Maker, ExecFunctionCaller::Maker, - NullSelectionSparseFunctionCaller::Maker}) { + ExpressionFunctionCaller::Maker}) { ASSERT_OK_AND_ASSIGN(auto test_copy, caller_maker(name, {uint8()})); ARROW_SCOPED_TRACE(test_copy->name()); ResetContexts(); @@ -1935,7 +1862,7 @@ TEST_F(TestCallScalarFunctionPreallocationCases, ChunkedSelectiveSparse) { auto arr = GetTestArray(); auto carr = std::make_shared(ArrayVector{arr->Slice(0, 10), arr->Slice(10)}); - auto selections = GetSelectionVectors(); + auto selections = GetTestSelectionVectors(); for (const auto& name : {"test_copy_selective", "test_copy_computed_bitmap_selective"}) { ARROW_SCOPED_TRACE(name); @@ -1959,7 +1886,7 @@ TEST_F(TestCallScalarFunctionPreallocationCases, ChunkedSelectiveDense) { auto arr = GetTestArray(); auto carr = std::make_shared(ArrayVector{arr->Slice(0, 10), arr->Slice(10)}); - auto selections = GetSelectionVectors(); + auto selections = GetTestSelectionVectors(); for (const auto& name : {"test_copy", "test_copy_computed_bitmap"}) { ARROW_SCOPED_TRACE(name); ASSERT_OK_AND_ASSIGN(auto test_copy, @@ -1984,7 +1911,7 @@ TEST_F(TestCallScalarFunctionPreallocationCases, IndependentPreallocate) { ARROW_SCOPED_TRACE(name); for (const auto& caller_maker : {SimpleFunctionCaller::Maker, ExecFunctionCaller::Maker, - NullSelectionSparseFunctionCaller::Maker}) { + ExpressionFunctionCaller::Maker}) { ASSERT_OK_AND_ASSIGN(auto test_copy, caller_maker(name, {uint8()})); ARROW_SCOPED_TRACE(test_copy->name()); ResetContexts(); @@ -2005,15 +1932,16 @@ TEST_F(TestCallScalarFunctionPreallocationCases, IndependentPreallocate) { TEST_F(TestCallScalarFunctionPreallocationCases, IndependentPreallocateSelectiveSparse) { auto arr = GetTestArray(); - auto selections = GetSelectionVectors(); + auto selections = GetTestSelectionVectors(); for (const auto& name : {"test_copy_selective", "test_copy_computed_bitmap_selective"}) { ARROW_SCOPED_TRACE(name); - ASSERT_OK_AND_ASSIGN(auto test_copy, DenseFunctionCaller::Maker(name, {uint8()})); + ASSERT_OK_AND_ASSIGN(auto test_copy, + ExpressionFunctionCaller::Maker(name, {uint8()})); for (const auto& selection : selections) { + const int64_t exec_chunksize = 40; ResetContexts(); - const int64_t exec_chunksize = 40; DoTestIndependentPreallocate(test_copy.get(), exec_chunksize, *arr, selection, [&](const Datum& result) { AssertChunkedExecResultsEqualSparseWithSelection( @@ -2025,17 +1953,19 @@ TEST_F(TestCallScalarFunctionPreallocationCases, IndependentPreallocateSelective TEST_F(TestCallScalarFunctionPreallocationCases, IndependentPreallocateSelectiveDense) { auto arr = GetTestArray(); - auto selections = GetSelectionVectors(); + auto selections = GetTestSelectionVectors(); for (const auto& name : {"test_copy", "test_copy_computed_bitmap"}) { ARROW_SCOPED_TRACE(name); - ASSERT_OK_AND_ASSIGN(auto test_copy, DenseFunctionCaller::Maker(name, {uint8()})); + ASSERT_OK_AND_ASSIGN(auto test_copy, + ExpressionFunctionCaller::Maker(name, {uint8()})); for (const auto& selection : selections) { + const int64_t exec_chunksize = 40; ResetContexts(); - DoTestIndependentPreallocate(test_copy.get(), /*exec_chunksize=*/40, *arr, - selection, [&](const Datum& result) { + DoTestIndependentPreallocate(test_copy.get(), exec_chunksize, *arr, selection, + [&](const Datum& result) { AssertChunkedExecResultsEqualDenseWithSelection( - 40, *arr, selection.get(), result); + exec_chunksize, *arr, selection.get(), result); }); } } @@ -2049,119 +1979,302 @@ TEST_F(TestCallScalarFunctionPreallocationCases, IndependentPreallocateSelective class TestCallScalarFunctionBasicNonStandardCases : public TestCallScalarFunction { protected: std::shared_ptr GetTestArray() { return GetUInt8Array(1000, 0.2); } + + std::vector> GetTestSelectionVectors() { + return {MakeSelectionVector("[]"), MakeSelectionVector("[0]"), + MakeSelectionVector("[999]"), MakeSelectionVectorUntil(400), + MakeSelectionVectorUntil(401), MakeSelectionVectorUntil(1000)}; + } + + template + void DoTestBasic(const FunctionCaller* caller, const Array& input, + std::shared_ptr selection, CheckFunc&& check_func) { + // The default should be a single array output + std::vector args = {Datum(input)}; + ASSERT_OK_AND_ASSIGN(Datum result, caller->Call(args, selection)); + check_func(result); + } + + template + void DoTestSplitExecution(const FunctionCaller* caller, int64_t exec_chunksize, + const Array& input, + std::shared_ptr selection, + CheckFunc&& check_func) { + // Split execution into several chunks + std::vector args = {Datum(input)}; + exec_ctx_->set_exec_chunksize(exec_chunksize); + ASSERT_OK_AND_ASSIGN(Datum result, caller->Call(args, selection, /*options=*/nullptr, + exec_ctx_.get())); + check_func(result); + } }; TEST_F(TestCallScalarFunctionBasicNonStandardCases, Basic) { auto arr = GetTestArray(); - std::vector args = {Datum(arr)}; for (const auto& name : {"test_nopre_data", "test_nopre_validity_or_data"}) { ARROW_SCOPED_TRACE(name); for (const auto& caller_maker : {SimpleFunctionCaller::Maker, ExecFunctionCaller::Maker, - ExpressionFunctionCaller::Maker, NullSelectionSparseFunctionCaller::Maker, - SelectiveSparseFunctionCaller::Maker, DenseFunctionCaller::Maker}) { + ExpressionFunctionCaller::Maker}) { ASSERT_OK_AND_ASSIGN(auto test_nopre, caller_maker(name, {uint8()})); ARROW_SCOPED_TRACE(test_nopre->name()); ResetContexts(); - // The default should be a single array output - { - ASSERT_OK_AND_ASSIGN(Datum result, test_nopre->Call(args)); - AssertArraysEqual(*arr, *result.make_array(), true); - } + DoTestBasic(test_nopre.get(), *arr, /*selection=*/nullptr, + [&](const Datum& result) { + ASSERT_EQ(Datum::ARRAY, result.kind()); + AssertArraysEqual(*arr, *result.make_array(), /*verbose=*/true); + }); + } + } +} + +TEST_F(TestCallScalarFunctionBasicNonStandardCases, BasicSelectiveSparse) { + auto arr = GetTestArray(); + auto selections = GetTestSelectionVectors(); + for (const auto& name : + {"test_nopre_data_selective", "test_nopre_validity_or_data_selective"}) { + ARROW_SCOPED_TRACE(name); + ASSERT_OK_AND_ASSIGN(auto test_nopre, + ExpressionFunctionCaller::Maker(name, {uint8()})); + for (const auto& selection : selections) { + SelectionVectorSpan selection_span(selection->indices(), selection->length()); + ResetContexts(); + + DoTestBasic(test_nopre.get(), *arr, selection, [&](const Datum& result) { + ASSERT_EQ(Datum::ARRAY, result.kind()); + AssertArraysEqualSparseWithSelection(*arr, selection_span, *result.make_array()); + }); + } + } +} + +TEST_F(TestCallScalarFunctionBasicNonStandardCases, BasicSelectiveDense) { + auto arr = GetTestArray(); + auto selections = GetTestSelectionVectors(); + for (const auto& name : {"test_nopre_data", "test_nopre_validity_or_data"}) { + ARROW_SCOPED_TRACE(name); + ASSERT_OK_AND_ASSIGN(auto test_nopre, + ExpressionFunctionCaller::Maker(name, {uint8()})); + for (const auto& selection : selections) { + SelectionVectorSpan selection_span(selection->indices(), selection->length()); + ResetContexts(); + + DoTestBasic(test_nopre.get(), *arr, selection, [&](const Datum& result) { + ASSERT_EQ(Datum::ARRAY, result.kind()); + AssertArraysEqualDenseWithSelection(*arr, selection_span, *result.make_array()); + }); } } } TEST_F(TestCallScalarFunctionBasicNonStandardCases, SplitExecution) { auto arr = GetTestArray(); - std::vector args = {Datum(arr)}; for (const auto& name : {"test_nopre_data", "test_nopre_validity_or_data"}) { ARROW_SCOPED_TRACE(name); for (const auto& caller_maker : {SimpleFunctionCaller::Maker, ExecFunctionCaller::Maker, - ExpressionFunctionCaller::Maker, NullSelectionSparseFunctionCaller::Maker, - SelectiveSparseFunctionCaller::Maker}) { + ExpressionFunctionCaller::Maker}) { ASSERT_OK_AND_ASSIGN(auto test_nopre, caller_maker(name, {uint8()})); ARROW_SCOPED_TRACE(test_nopre->name()); ResetContexts(); - // Split execution into 3 chunks - { - exec_ctx_->set_exec_chunksize(400); - ASSERT_OK_AND_ASSIGN(Datum result, test_nopre->Call(args, exec_ctx_.get())); - ASSERT_EQ(Datum::CHUNKED_ARRAY, result.kind()); - const ChunkedArray& carr = *result.chunked_array(); - ASSERT_EQ(3, carr.num_chunks()); - AssertArraysEqual(*arr->Slice(0, 400), *carr.chunk(0)); - AssertArraysEqual(*arr->Slice(400, 400), *carr.chunk(1)); - AssertArraysEqual(*arr->Slice(800), *carr.chunk(2)); - } + DoTestSplitExecution(test_nopre.get(), /*exec_chunksize=*/400, *arr, + /*selection=*/nullptr, [&](const Datum& result) { + ASSERT_EQ(Datum::CHUNKED_ARRAY, result.kind()); + const ChunkedArray& carr = *result.chunked_array(); + ASSERT_EQ(3, carr.num_chunks()); + AssertArraysEqual(*arr->Slice(0, 400), *carr.chunk(0)); + AssertArraysEqual(*arr->Slice(400, 400), *carr.chunk(1)); + AssertArraysEqual(*arr->Slice(800), *carr.chunk(2)); + }); + } + } +} + +TEST_F(TestCallScalarFunctionBasicNonStandardCases, SplitExecutionSelectiveSparse) { + auto arr = GetTestArray(); + auto selections = GetTestSelectionVectors(); + for (const auto& name : + {"test_nopre_data_selective", "test_nopre_validity_or_data_selective"}) { + ARROW_SCOPED_TRACE(name); + ASSERT_OK_AND_ASSIGN(auto test_nopre, + ExpressionFunctionCaller::Maker(name, {uint8()})); + for (const auto& selection : selections) { + const int64_t exec_chunksize = 400; + ResetContexts(); + + DoTestSplitExecution(test_nopre.get(), exec_chunksize, *arr, selection, + [&](const Datum& result) { + AssertChunkedExecResultsEqualSparseWithSelection( + exec_chunksize, *arr, selection.get(), result); + }); } } } -TEST_F(TestCallScalarFunctionBasicNonStandardCases, SplitExecutionDense) { +TEST_F(TestCallScalarFunctionBasicNonStandardCases, SplitExecutionSelectiveDense) { auto arr = GetTestArray(); - std::vector args = {Datum(arr)}; + auto selections = GetTestSelectionVectors(); for (const auto& name : {"test_nopre_data", "test_nopre_validity_or_data"}) { ARROW_SCOPED_TRACE(name); - ASSERT_OK_AND_ASSIGN(auto test_nopre, DenseFunctionCaller::Maker(name, {uint8()})); - ResetContexts(); + ASSERT_OK_AND_ASSIGN(auto test_nopre, + ExpressionFunctionCaller::Maker(name, {uint8()})); + for (const auto& selection : selections) { + const int64_t exec_chunksize = 400; + ResetContexts(); - // Split execution into 3 chunks - { - exec_ctx_->set_exec_chunksize(400); - ASSERT_OK_AND_ASSIGN(Datum result, test_nopre->Call(args, exec_ctx_.get())); - ASSERT_EQ(Datum::CHUNKED_ARRAY, result.kind()); - const ChunkedArray& carr = *result.chunked_array(); - ASSERT_EQ(1, carr.num_chunks()); - AssertArraysEqual(*arr, *carr.chunk(0)); + DoTestSplitExecution(test_nopre.get(), exec_chunksize, *arr, selection, + [&](const Datum& result) { + AssertChunkedExecResultsEqualDenseWithSelection( + exec_chunksize, *arr, selection.get(), result); + }); } } } -class TestCallScalarFunctionStatefulKernel : public TestCallScalarFunction {}; +class TestCallScalarFunctionStatefulKernel : public TestCallScalarFunction { + protected: + std::shared_ptr GetTestArray() { + return ArrayFromJSON(int32(), "[1, 2, 3, null, 5]"); + } + + static constexpr int32_t kMultiplier = 2; + + std::shared_ptr GetExpected() { + return ArrayFromJSON(int32(), "[2, 4, 6, null, 10]"); + } + + std::vector> GetTestSelectionVectors() { + return {MakeSelectionVector("[]"), MakeSelectionVector("[0]"), + MakeSelectionVector("[4]"), MakeSelectionVectorUntil(2), + MakeSelectionVectorUntil(5)}; + } + + template + void DoTestBasic(const FunctionCaller* caller, const Array& input, + std::shared_ptr multiplier, + std::shared_ptr selection, CheckFunc&& check_func) { + ExampleOptions options(multiplier); + std::vector args = {Datum(input)}; + ASSERT_OK_AND_ASSIGN(Datum result, caller->Call(args, selection, &options)); + check_func(result); + } +}; TEST_F(TestCallScalarFunctionStatefulKernel, Basic) { - for (const auto& caller_maker : - {SimpleFunctionCaller::Maker, ExecFunctionCaller::Maker, - ExpressionFunctionCaller::Maker, NullSelectionSparseFunctionCaller::Maker, - SelectiveSparseFunctionCaller::Maker, DenseFunctionCaller::Maker}) { + auto input = GetTestArray(); + auto multiplier = std::make_shared(kMultiplier); + auto expected = GetExpected(); + for (const auto& caller_maker : {SimpleFunctionCaller::Maker, ExecFunctionCaller::Maker, + ExpressionFunctionCaller::Maker}) { ASSERT_OK_AND_ASSIGN(auto test_stateful, caller_maker("test_stateful", {int32()})); ARROW_SCOPED_TRACE(test_stateful->name()); ResetContexts(); - auto input = ArrayFromJSON(int32(), "[1, 2, 3, null, 5]"); - auto multiplier = std::make_shared(2); - auto expected = ArrayFromJSON(int32(), "[2, 4, 6, null, 10]"); + DoTestBasic( + test_stateful.get(), *input, multiplier, /*selection=*/nullptr, + [&](const Datum& result) { AssertArraysEqual(*expected, *result.make_array()); }); + } +} - ExampleOptions options(multiplier); - std::vector args = {Datum(input)}; - ASSERT_OK_AND_ASSIGN(Datum result, test_stateful->Call(args, &options)); - AssertArraysEqual(*expected, *result.make_array()); +TEST_F(TestCallScalarFunctionStatefulKernel, BasicSelectiveSparse) { + auto input = GetTestArray(); + auto multiplier = std::make_shared(kMultiplier); + auto selections = GetTestSelectionVectors(); + auto expected = GetExpected(); + ASSERT_OK_AND_ASSIGN( + auto caller, ExpressionFunctionCaller::Maker("test_stateful_selective", {int32()})); + for (const auto& selection : selections) { + SelectionVectorSpan selection_span(selection->indices(), selection->length()); + ResetContexts(); + + DoTestBasic(caller.get(), *input, multiplier, selection, [&](const Datum& result) { + ASSERT_EQ(Datum::ARRAY, result.kind()); + AssertArraysEqualSparseWithSelection(*expected, selection_span, + *result.make_array()); + }); + } +} + +TEST_F(TestCallScalarFunctionStatefulKernel, BasicSelectiveDense) { + auto input = GetTestArray(); + auto multiplier = std::make_shared(kMultiplier); + auto selections = GetTestSelectionVectors(); + auto expected = GetExpected(); + ASSERT_OK_AND_ASSIGN(auto caller, + ExpressionFunctionCaller::Maker("test_stateful", {int32()})); + for (const auto& selection : selections) { + SelectionVectorSpan selection_span(selection->indices(), selection->length()); + ResetContexts(); + + DoTestBasic(caller.get(), *input, multiplier, selection, [&](const Datum& result) { + ASSERT_EQ(Datum::ARRAY, result.kind()); + AssertArraysEqualDenseWithSelection(*expected, selection_span, + *result.make_array()); + }); } } -class TestCallScalarFunctionScalarFunction : public TestCallScalarFunction {}; +class TestCallScalarFunctionScalarFunction : public TestCallScalarFunction { + protected: + std::vector GetTestArgs() { + return {Datum(std::make_shared(5)), + Datum(std::make_shared(7))}; + } + + static constexpr int32_t kExpectedResult = 12; + + std::vector> GetTestSelectionVectors() { + return {MakeSelectionVector("[]"), MakeSelectionVector("[0]")}; + } + + void DoTestBasic(const FunctionCaller* caller, const std::vector& args, + std::shared_ptr selection) { + ASSERT_OK_AND_ASSIGN(Datum result, caller->Call(args)); + ASSERT_EQ(Datum::SCALAR, result.kind()); + + auto expected = std::make_shared(kExpectedResult); + ASSERT_TRUE(expected->Equals(*result.scalar())); + } +}; TEST_F(TestCallScalarFunctionScalarFunction, Basic) { - for (const auto& caller_maker : - {SimpleFunctionCaller::Maker, ExecFunctionCaller::Maker, - ExpressionFunctionCaller::Maker, NullSelectionSparseFunctionCaller::Maker, - SelectiveSparseFunctionCaller::Maker, DenseFunctionCaller::Maker}) { + auto args = GetTestArgs(); + for (const auto& caller_maker : {SimpleFunctionCaller::Maker, ExecFunctionCaller::Maker, + ExpressionFunctionCaller::Maker}) { ASSERT_OK_AND_ASSIGN(auto test_scalar_add_int32, caller_maker("test_scalar_add_int32", {int32(), int32()})); ARROW_SCOPED_TRACE(test_scalar_add_int32->name()); ResetContexts(); - std::vector args = {Datum(std::make_shared(5)), - Datum(std::make_shared(7))}; - ASSERT_OK_AND_ASSIGN(Datum result, test_scalar_add_int32->Call(args)); - ASSERT_EQ(Datum::SCALAR, result.kind()); + DoTestBasic(test_scalar_add_int32.get(), args, /*selection=*/nullptr); + } +} - auto expected = std::make_shared(12); - ASSERT_TRUE(expected->Equals(*result.scalar())); +TEST_F(TestCallScalarFunctionScalarFunction, BasicSelectiveSparse) { + auto args = GetTestArgs(); + auto selections = GetTestSelectionVectors(); + ASSERT_OK_AND_ASSIGN(auto test_scalar_add_int32, + ExpressionFunctionCaller::Maker("test_scalar_add_int32_selective", + {int32(), int32()})); + for (const auto& selection : selections) { + ResetContexts(); + + DoTestBasic(test_scalar_add_int32.get(), GetTestArgs(), selection); + } +} + +TEST_F(TestCallScalarFunctionScalarFunction, BasicSelectiveDense) { + auto args = GetTestArgs(); + auto selections = GetTestSelectionVectors(); + ASSERT_OK_AND_ASSIGN( + auto test_scalar_add_int32, + ExpressionFunctionCaller::Maker("test_scalar_add_int32", {int32(), int32()})); + for (const auto& selection : selections) { + ResetContexts(); + + DoTestBasic(test_scalar_add_int32.get(), GetTestArgs(), selection); } } From e2d865fe14386163cc983f821c2d8cf42d5f2d15 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Wed, 20 Aug 2025 14:22:42 +0800 Subject: [PATCH 25/71] More clear naming --- cpp/src/arrow/compute/exec.cc | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/cpp/src/arrow/compute/exec.cc b/cpp/src/arrow/compute/exec.cc index a11d292997b9..e45b831ffd49 100644 --- a/cpp/src/arrow/compute/exec.cc +++ b/cpp/src/arrow/compute/exec.cc @@ -830,11 +830,11 @@ class ScalarExecutor : public KernelExecutorImpl { return EmitResult(result->data(), listener); } - if (!batch.selection_vector || kernel_->selective_exec) { - return ExecuteSparse(batch, listener); + if (batch.selection_vector && !kernel_->selective_exec) { + return ExecuteSelectiveDense(batch, listener); } - return ExecuteDense(batch, listener); + return ExecuteBatch(batch, listener); } Datum WrapResults(const std::vector& inputs, @@ -850,7 +850,7 @@ class ScalarExecutor : public KernelExecutorImpl { } protected: - Status ExecuteSparse(const ExecBatch& batch, ExecListener* listener) { + Status ExecuteBatch(const ExecBatch& batch, ExecListener* listener) { DCHECK(!batch.selection_vector || kernel_->selective_exec); RETURN_NOT_OK(span_iterator_.Init(batch, exec_context()->exec_chunksize())); @@ -874,13 +874,13 @@ class ScalarExecutor : public KernelExecutorImpl { } } - Status ExecuteDense(const ExecBatch& batch, ExecListener* listener) { + Status ExecuteSelectiveDense(const ExecBatch& batch, ExecListener* listener) { DCHECK(batch.selection_vector && !kernel_->selective_exec); if (CheckIfAllScalar(batch)) { ExecBatch input = batch; input.selection_vector = nullptr; - return ExecuteSparse(input, listener); + return ExecuteBatch(input, listener); } std::vector values(batch.num_values()); @@ -899,7 +899,7 @@ class ScalarExecutor : public KernelExecutorImpl { ExecBatch::Make(std::move(values), batch.selection_vector->length())); DatumAccumulator dense_listener; - RETURN_NOT_OK(ExecuteSparse(input, &dense_listener)); + RETURN_NOT_OK(ExecuteBatch(input, &dense_listener)); Datum dense_result = WrapResults(input.values, dense_listener.values()); ARROW_ASSIGN_OR_RAISE(auto result, From cd0d5b8a5b26d5a0cd0f37c9c46887d33fe90d02 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Wed, 3 Sep 2025 12:02:06 +0800 Subject: [PATCH 26/71] Refine binding --- cpp/src/arrow/compute/expression.cc | 122 +++++++++----------- cpp/src/arrow/compute/expression.h | 1 + cpp/src/arrow/compute/expression_internal.h | 7 +- cpp/src/arrow/compute/function.cc | 3 +- cpp/src/arrow/compute/function.h | 3 +- cpp/src/arrow/compute/special_form.cc | 66 +++++------ cpp/src/arrow/compute/special_form.h | 9 +- cpp/src/arrow/compute/special_form_test.cc | 3 +- 8 files changed, 95 insertions(+), 119 deletions(-) diff --git a/cpp/src/arrow/compute/expression.cc b/cpp/src/arrow/compute/expression.cc index 502943e73e2e..4003c4cbe09e 100644 --- a/cpp/src/arrow/compute/expression.cc +++ b/cpp/src/arrow/compute/expression.cc @@ -587,39 +587,6 @@ std::vector GetTypesWithSmallestLiteralRepresentation( return types; } -// Produce a bound Expression from unbound Call and bound arguments. -Result BindNonRecursive(Expression::Call call, bool insert_implicit_casts, - ExecContext* exec_context) { - DCHECK(std::all_of(call.arguments.begin(), call.arguments.end(), - [](const Expression& argument) { return argument.IsBound(); })); - - std::vector types = GetTypes(call.arguments); - ARROW_ASSIGN_OR_RAISE(call.function, GetFunction(call, exec_context)); - - auto finish_bind = [&](const Kernel* kernel, - const std::vector& types) -> Status { - call.kernel = kernel; - KernelContext kernel_context(exec_context, call.kernel); - if (call.kernel->init) { - const FunctionOptions* options = - call.options ? call.options.get() : call.function->default_options(); - ARROW_ASSIGN_OR_RAISE( - call.kernel_state, - call.kernel->init(&kernel_context, {call.kernel, types, options})); - - kernel_context.SetState(call.kernel_state.get()); - } - - ARROW_ASSIGN_OR_RAISE( - call.type, call.kernel->signature->out_type().Resolve(&kernel_context, types)); - return Status::OK(); - }; - - RETURN_NOT_OK(DispatchForBind(call.function.get(), call.arguments, - insert_implicit_casts, exec_context, finish_bind)); - return Expression(std::move(call)); -} - template Result BindImpl(Expression expr, const TypeOrSchema& in, compute::ExecContext* exec_context) { @@ -646,8 +613,9 @@ Result BindImpl(Expression expr, const TypeOrSchema& in, for (auto& argument : special.arguments) { ARROW_ASSIGN_OR_RAISE(argument, BindImpl(std::move(argument), in, exec_context)); } - ARROW_ASSIGN_OR_RAISE(special.special_executor, - special.special_form->Bind(special.arguments, exec_context)); + ARROW_ASSIGN_OR_RAISE( + special.special_executor, + special.special_form->Bind(special.arguments, special.options, exec_context)); special.type = special.special_executor->out_type(); return Expression(std::move(special)); } @@ -661,54 +629,70 @@ Result BindImpl(Expression expr, const TypeOrSchema& in, } // namespace -Status DispatchForBind( - const Function* function, std::vector arguments, - bool insert_implicit_casts, ExecContext* exec_context, - std::function&)> finish_bind) { - std::vector types = GetTypes(arguments); +Result BindNonRecursive(Expression::Call call, bool insert_implicit_casts, + compute::ExecContext* exec_context) { + DCHECK(std::all_of(call.arguments.begin(), call.arguments.end(), + [](const Expression& argument) { return argument.IsBound(); })); + + std::vector types = GetTypes(call.arguments); + ARROW_ASSIGN_OR_RAISE(call.function, GetFunction(call, exec_context)); // First try and bind exactly - Result maybe_exact_match = function->DispatchExact(types); + Result maybe_exact_match = call.function->DispatchExact(types); if (maybe_exact_match.ok()) { - if (finish_bind(*maybe_exact_match, types).ok()) { - return Status::OK(); + call.kernel = *maybe_exact_match; + } else { + if (!insert_implicit_casts) { + return maybe_exact_match.status(); } - } - if (!insert_implicit_casts) { - return maybe_exact_match.status(); - } + // If exact binding fails, and we are allowed to cast, then prefer casting literals + // first. Since DispatchBest generally prefers up-casting the best way to do this is + // first down-cast the literals as much as possible + types = GetTypesWithSmallestLiteralRepresentation(call.arguments); + ARROW_ASSIGN_OR_RAISE(call.kernel, call.function->DispatchBest(&types)); - // If exact binding fails, and we are allowed to cast, then prefer casting literals - // first. Since DispatchBest generally prefers up-casting the best way to do this is - // first down-cast the literals as much as possible - types = GetTypesWithSmallestLiteralRepresentation(arguments); - ARROW_ASSIGN_OR_RAISE(auto kernel, function->DispatchBest(&types)); + for (size_t i = 0; i < types.size(); ++i) { + if (types[i] == call.arguments[i].type()) continue; - for (size_t i = 0; i < types.size(); ++i) { - if (types[i] == arguments[i].type()) continue; + if (const Datum* lit = call.arguments[i].literal()) { + ARROW_ASSIGN_OR_RAISE(Datum new_lit, + compute::Cast(*lit, types[i].GetSharedPtr())); + call.arguments[i] = literal(std::move(new_lit)); + continue; + } - if (const Datum* lit = arguments[i].literal()) { - ARROW_ASSIGN_OR_RAISE(Datum new_lit, compute::Cast(*lit, types[i].GetSharedPtr())); - arguments[i] = literal(std::move(new_lit)); - continue; - } + // construct an implicit cast Expression with which to replace this argument + Expression::Call implicit_cast; + implicit_cast.function_name = "cast"; + implicit_cast.arguments = {std::move(call.arguments[i])}; - // construct an implicit cast Expression with which to replace this argument - Expression::Call implicit_cast; - implicit_cast.function_name = "cast"; - implicit_cast.arguments = {std::move(arguments[i])}; + // TODO(wesm): Use TypeHolder in options + implicit_cast.options = std::make_shared( + compute::CastOptions::Safe(types[i].GetSharedPtr())); - // TODO(wesm): Use TypeHolder in options - implicit_cast.options = std::make_shared( - compute::CastOptions::Safe(types[i].GetSharedPtr())); + ARROW_ASSIGN_OR_RAISE( + call.arguments[i], + BindNonRecursive(std::move(implicit_cast), + /*insert_implicit_casts=*/false, exec_context)); + } + } + compute::KernelContext kernel_context(exec_context, call.kernel); + if (call.kernel->init) { + const FunctionOptions* options = + call.options ? call.options.get() : call.function->default_options(); ARROW_ASSIGN_OR_RAISE( - arguments[i], BindNonRecursive(std::move(implicit_cast), - /*insert_implicit_casts=*/false, exec_context)); + call.kernel_state, + call.kernel->init(&kernel_context, {call.kernel, types, options})); + + kernel_context.SetState(call.kernel_state.get()); } - return finish_bind(kernel, types); + ARROW_ASSIGN_OR_RAISE( + call.type, call.kernel->signature->out_type().Resolve(&kernel_context, types)); + + return Expression(std::move(call)); } Result Expression::Bind(const TypeHolder& in, diff --git a/cpp/src/arrow/compute/expression.h b/cpp/src/arrow/compute/expression.h index a4cce3f23dae..2311f48a44a4 100644 --- a/cpp/src/arrow/compute/expression.h +++ b/cpp/src/arrow/compute/expression.h @@ -63,6 +63,7 @@ class ARROW_EXPORT Expression { struct Special { std::shared_ptr special_form; std::vector arguments; + std::shared_ptr options; // Cached hash value size_t hash; diff --git a/cpp/src/arrow/compute/expression_internal.h b/cpp/src/arrow/compute/expression_internal.h index ef983886b087..ca97e2a1fd88 100644 --- a/cpp/src/arrow/compute/expression_internal.h +++ b/cpp/src/arrow/compute/expression_internal.h @@ -290,10 +290,9 @@ inline Result> GetFunction( return GetCastFunction(*to_type); } -Status DispatchForBind( - const Function* function, std::vector arguments, - bool insert_implicit_casts, ExecContext* exec_context, - std::function&)> finish_bind); +// Produce a bound Expression from unbound Call and bound arguments. +Result BindNonRecursive(Expression::Call call, bool insert_implicit_casts, + ExecContext* exec_context); } // namespace compute } // namespace arrow diff --git a/cpp/src/arrow/compute/function.cc b/cpp/src/arrow/compute/function.cc index 96efc989aa85..fe753da46aec 100644 --- a/cpp/src/arrow/compute/function.cc +++ b/cpp/src/arrow/compute/function.cc @@ -410,8 +410,7 @@ Status Function::Validate() const { } Status ScalarFunction::AddKernel(std::vector in_types, OutputType out_type, - ArrayKernelExec exec, ArrayKernelExec exec, - KernelInit init, + ArrayKernelExec exec, KernelInit init, std::shared_ptr constraint) { return AddKernel(std::move(in_types), std::move(out_type), std::move(exec), /*selective_exec=*/nullptr, std::move(init)); diff --git a/cpp/src/arrow/compute/function.h b/cpp/src/arrow/compute/function.h index b6046e655084..d86c8ab3b5ea 100644 --- a/cpp/src/arrow/compute/function.h +++ b/cpp/src/arrow/compute/function.h @@ -313,7 +313,8 @@ class ARROW_EXPORT ScalarFunction : public detail::FunctionImpl { Status AddKernel(std::vector in_types, OutputType out_type, ArrayKernelExec exec, ArrayKernelSelectiveExec selective_exec, - KernelInit init = NULLPTR); + KernelInit init = NULLPTR, + std::shared_ptr constraint = NULLPTR); /// \brief Add a kernel (function implementation). Returns error if the /// kernel's signature does not match the function's arity. diff --git a/cpp/src/arrow/compute/special_form.cc b/cpp/src/arrow/compute/special_form.cc index f95d9976dcc2..b6196f766441 100644 --- a/cpp/src/arrow/compute/special_form.cc +++ b/cpp/src/arrow/compute/special_form.cc @@ -576,8 +576,10 @@ struct ConditionalExec { class ConditionalSpecialExecutor : public SpecialExecutor { public: - explicit ConditionalSpecialExecutor(TypeHolder out_type, std::vector branches) - : SpecialExecutor(std::move(out_type)), branches(std::move(branches)) {} + explicit ConditionalSpecialExecutor(std::vector branches, TypeHolder out_type, + std::shared_ptr options) + : SpecialExecutor(std::move(out_type), std::move(options)), + branches(std::move(branches)) {} Result Execute(const ExecBatch& input, ExecContext* exec_context) const override { @@ -589,67 +591,51 @@ class ConditionalSpecialExecutor : public SpecialExecutor { }; template -class KernelBackedSpecialForm : public SpecialForm { +class FunctionBackedSpecialForm : public SpecialForm { public: using SpecialForm::SpecialForm; - ARROW_DISALLOW_COPY_AND_ASSIGN(KernelBackedSpecialForm); - ARROW_DEFAULT_MOVE_AND_ASSIGN(KernelBackedSpecialForm); + ARROW_DISALLOW_COPY_AND_ASSIGN(FunctionBackedSpecialForm); + ARROW_DEFAULT_MOVE_AND_ASSIGN(FunctionBackedSpecialForm); Result> Bind(std::vector& arguments, + std::shared_ptr options, ExecContext* exec_context) override { DCHECK(std::all_of(arguments.begin(), arguments.end(), [](const Expression& argument) { return argument.IsBound(); })); - ARROW_ASSIGN_OR_RAISE(auto function, - exec_context->func_registry()->GetFunction(name())); - - const Kernel* kernel = nullptr; - TypeHolder out_type; - auto finish_bind = [&](const Kernel* dispatched_kernel, - const std::vector& types) -> Status { - kernel = dispatched_kernel; - KernelContext kernel_context(exec_context, kernel); - std::unique_ptr kernel_state; - if (kernel->init) { - const FunctionOptions* options = function->default_options(); - ARROW_ASSIGN_OR_RAISE(kernel_state, - kernel->init(&kernel_context, {kernel, types, options})); - kernel_context.SetState(kernel_state.get()); - } - ARROW_ASSIGN_OR_RAISE( - out_type, kernel->signature->out_type().Resolve(&kernel_context, types)); - return Status::OK(); - }; + Expression::Call call; + call.function_name = name(); + call.arguments = std::move(arguments); + call.options = std::move(options); + ARROW_ASSIGN_OR_RAISE( + auto bound, BindNonRecursive(call, /*insert_implicit_casts=*/true, exec_context)); - RETURN_NOT_OK(DispatchForBind(function.get(), arguments, - /*insert_implicit_casts=*/true, exec_context, - finish_bind)); - return static_cast(this)->BindImpl(kernel, std::move(out_type), - std::move(arguments), exec_context); + return static_cast(this)->BindWithBoundCall(CallNotNull(bound), + exec_context); } }; template class ConditionalSpecialForm - : public KernelBackedSpecialForm> { + : public FunctionBackedSpecialForm> { public: - using KernelBackedSpecialForm>::KernelBackedSpecialForm; + using FunctionBackedSpecialForm< + ConditionalSpecialForm>::FunctionBackedSpecialForm; ARROW_DISALLOW_COPY_AND_ASSIGN(ConditionalSpecialForm); ARROW_DEFAULT_MOVE_AND_ASSIGN(ConditionalSpecialForm); protected: - Result> BindImpl(const Kernel* kernel, - TypeHolder out_type, - std::vector arguments, - ExecContext* exec_context) const { - auto branches = static_cast(this)->GetBranches(arguments); - return std::make_unique(std::move(out_type), - std::move(branches)); + Result> BindWithBoundCall( + const Expression::Call* call, ExecContext* exec_context) const { + auto branches = + static_cast(this)->GetBranches(std::move(call->arguments)); + return std::make_unique( + std::move(branches), std::move(call->type), std::move(call->options)); } - friend class KernelBackedSpecialForm>; + friend class FunctionBackedSpecialForm>; }; class IfElseSpecialForm : public ConditionalSpecialForm { diff --git a/cpp/src/arrow/compute/special_form.h b/cpp/src/arrow/compute/special_form.h index 24efe09aa84d..5b8cf1e923fd 100644 --- a/cpp/src/arrow/compute/special_form.h +++ b/cpp/src/arrow/compute/special_form.h @@ -21,17 +21,21 @@ namespace arrow::compute { class ARROW_EXPORT SpecialExecutor { public: - explicit SpecialExecutor(TypeHolder out_type) : out_type_(std::move(out_type)) {} + explicit SpecialExecutor(TypeHolder out_type, + std::shared_ptr options = NULLPTR) + : out_type_(std::move(out_type)), options_(std::move(options)) {} virtual ~SpecialExecutor() = default; const TypeHolder& out_type() const { return out_type_; } + const std::shared_ptr options() const { return options_; } virtual Result Execute(const ExecBatch& input, ExecContext* exec_context) const = 0; private: const TypeHolder out_type_; + const std::shared_ptr options_; }; class ARROW_EXPORT SpecialForm { @@ -43,7 +47,8 @@ class ARROW_EXPORT SpecialForm { const std::string& name() const { return name_; } virtual Result> Bind( - std::vector& arguments, ExecContext* exec_context) = 0; + std::vector& arguments, std::shared_ptr options, + ExecContext* exec_context) = 0; private: std::string name_; diff --git a/cpp/src/arrow/compute/special_form_test.cc b/cpp/src/arrow/compute/special_form_test.cc index d69872463f6f..d20b73754ea1 100644 --- a/cpp/src/arrow/compute/special_form_test.cc +++ b/cpp/src/arrow/compute/special_form_test.cc @@ -256,7 +256,8 @@ class IfElseSpecialFormTest : public ::testing::Test { selective_exec = AssertSelectionVectorExistExec; } ScalarKernel kernel({InputType::Any()}, internal::FirstType, std::move(exec), - /*init=*/nullptr, std::move(selective_exec)); + std::move(selective_exec), + /*init=*/nullptr); kernel.can_write_into_slices = false; kernel.null_handling = NullHandling::COMPUTED_NO_PREALLOCATE; kernel.mem_allocation = MemAllocation::NO_PREALLOCATE; From 69f9d1aa51180f84f9d8f43624f1df98edac8cad Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Wed, 3 Sep 2025 12:13:22 +0800 Subject: [PATCH 27/71] Revert something --- cpp/src/arrow/compute/expression.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/arrow/compute/expression.cc b/cpp/src/arrow/compute/expression.cc index 4003c4cbe09e..3863158725b5 100644 --- a/cpp/src/arrow/compute/expression.cc +++ b/cpp/src/arrow/compute/expression.cc @@ -571,7 +571,7 @@ TypeHolder SmallestTypeFor(const arrow::Datum& value) { } } -std::vector GetTypesWithSmallestLiteralRepresentation( +inline std::vector GetTypesWithSmallestLiteralRepresentation( const std::vector& exprs) { std::vector types(exprs.size()); for (size_t i = 0; i < exprs.size(); ++i) { From b718bb80dcdae1c50624834e177fea45172748f1 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Fri, 5 Sep 2025 11:35:28 +0800 Subject: [PATCH 28/71] Fix rebase error: missed constraint --- cpp/src/arrow/compute/function.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/arrow/compute/function.cc b/cpp/src/arrow/compute/function.cc index fe753da46aec..b9c557aed274 100644 --- a/cpp/src/arrow/compute/function.cc +++ b/cpp/src/arrow/compute/function.cc @@ -413,7 +413,7 @@ Status ScalarFunction::AddKernel(std::vector in_types, OutputType out ArrayKernelExec exec, KernelInit init, std::shared_ptr constraint) { return AddKernel(std::move(in_types), std::move(out_type), std::move(exec), - /*selective_exec=*/nullptr, std::move(init)); + /*selective_exec=*/nullptr, std::move(init), std::move(constraint)); } Status ScalarFunction::AddKernel(std::vector in_types, OutputType out_type, From 624d674cf5af6fae75d755b61a3964d51aa27f7b Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Fri, 5 Sep 2025 13:14:05 +0800 Subject: [PATCH 29/71] Bind special form test --- cpp/src/arrow/compute/expression.cc | 39 +++++--- cpp/src/arrow/compute/expression_internal.h | 6 ++ cpp/src/arrow/compute/expression_test.cc | 100 ++++++++++++++++++++ cpp/src/arrow/compute/special_form.cc | 14 +-- 4 files changed, 138 insertions(+), 21 deletions(-) diff --git a/cpp/src/arrow/compute/expression.cc b/cpp/src/arrow/compute/expression.cc index 3863158725b5..fe13056644e4 100644 --- a/cpp/src/arrow/compute/expression.cc +++ b/cpp/src/arrow/compute/expression.cc @@ -276,10 +276,28 @@ bool Expression::Equals(const Expression& other) const { return ref->Equals(*other.field_ref()); } - // TODO - // if (auto special = this->special()) { - if (this->special()) { - return true; + auto args_and_options_eq = [](const auto* special_or_call, const auto* other) -> bool { + for (size_t i = 0; i < special_or_call->arguments.size(); ++i) { + if (!special_or_call->arguments[i].Equals(other->arguments[i])) { + return false; + } + } + + if (special_or_call->options == other->options) return true; + if (special_or_call->options && other->options) { + return special_or_call->options->Equals(*other->options); + } + return false; + }; + + if (auto special = this->special(); special) { + auto other_special = SpecialNotNull(other); + + if (special->special_form->name() != other_special->special_form->name()) { + return false; + } + + return args_and_options_eq(special, other_special); } auto call = CallNotNull(*this); @@ -290,17 +308,7 @@ bool Expression::Equals(const Expression& other) const { return false; } - for (size_t i = 0; i < call->arguments.size(); ++i) { - if (!call->arguments[i].Equals(other_call->arguments[i])) { - return false; - } - } - - if (call->options == other_call->options) return true; - if (call->options && other_call->options) { - return call->options->Equals(*other_call->options); - } - return false; + return args_and_options_eq(call, other_call); } bool Expression::Identical(const Expression& l, const Expression& r) { @@ -616,6 +624,7 @@ Result BindImpl(Expression expr, const TypeOrSchema& in, ARROW_ASSIGN_OR_RAISE( special.special_executor, special.special_form->Bind(special.arguments, special.options, exec_context)); + DCHECK_NE(special.special_executor, nullptr); special.type = special.special_executor->out_type(); return Expression(std::move(special)); } diff --git a/cpp/src/arrow/compute/expression_internal.h b/cpp/src/arrow/compute/expression_internal.h index ca97e2a1fd88..c7982b30b262 100644 --- a/cpp/src/arrow/compute/expression_internal.h +++ b/cpp/src/arrow/compute/expression_internal.h @@ -44,6 +44,12 @@ inline const Expression::Call* CallNotNull(const Expression& expr) { return call; } +inline const Expression::Special* SpecialNotNull(const Expression& expr) { + auto special = expr.special(); + ARROW_DCHECK_NE(special, nullptr); + return special; +} + inline std::vector GetTypes(const std::vector& exprs) { std::vector types(exprs.size()); for (size_t i = 0; i < exprs.size(); ++i) { diff --git a/cpp/src/arrow/compute/expression_test.cc b/cpp/src/arrow/compute/expression_test.cc index bbab57feebb7..65c3e7a9333e 100644 --- a/cpp/src/arrow/compute/expression_test.cc +++ b/cpp/src/arrow/compute/expression_test.cc @@ -948,6 +948,106 @@ TEST(Expression, BindNestedCall) { EXPECT_TRUE(expr.IsBound()); } +TEST(Expression, BindIfElseSpecialForm) { + { + auto expr = if_else_special(field_ref("bool"), field_ref("i8"), field_ref("i8")); + EXPECT_FALSE(expr.IsBound()); + ExpectBindsTo(expr, no_change, &expr); + EXPECT_TRUE(expr.IsBound()); + EXPECT_TRUE(expr.type()->Equals(*int8())); + } + + // Implicit casts. + { + Expression bound; + ExpectBindsTo(if_else_special(field_ref("bool"), field_ref("i8"), field_ref("i32")), + if_else_special(field_ref("bool"), cast(field_ref("i8"), int32()), + field_ref("i32")), + &bound); + EXPECT_TRUE(bound.IsBound()); + EXPECT_TRUE(bound.type()->Equals(*int32())); + } + { + Expression bound; + ExpectBindsTo(if_else_special(field_ref("bool"), field_ref("i32"), field_ref("i8")), + if_else_special(field_ref("bool"), field_ref("i32"), + cast(field_ref("i8"), int32())), + &bound); + EXPECT_TRUE(bound.IsBound()); + EXPECT_TRUE(bound.type()->Equals(*int32())); + } + + // Nested call. + { + Expression bound; + ExpectBindsTo(if_else_special(equal(field_ref("i8"), field_ref("i8")), + add(field_ref("i8"), literal(42)), + add(field_ref("i32"), literal(42))), + if_else_special(equal(field_ref("i8"), field_ref("i8")), + cast(add(field_ref("i8"), literal(42)), int32()), + add(field_ref("i32"), literal(42))), + &bound); + EXPECT_TRUE(bound.IsBound()); + EXPECT_TRUE(bound.type()->Equals(*int32())); + } + { + Expression bound; + ExpectBindsTo(if_else_special(equal(field_ref("i8"), field_ref("i32")), + add(field_ref("i32"), field_ref("i8")), + add(field_ref("i32"), literal(42))), + if_else_special(equal(cast(field_ref("i8"), int32()), field_ref("i32")), + add(field_ref("i32"), cast(field_ref("i8"), int32())), + add(field_ref("i32"), literal(42))), + &bound); + EXPECT_TRUE(bound.IsBound()); + EXPECT_TRUE(bound.type()->Equals(*int32())); + } + + // Nesting call. + { + Expression bound; + ExpectBindsTo(add(if_else_special(field_ref("bool"), field_ref("i32"), literal(42)), + field_ref("i8")), + add(if_else_special(field_ref("bool"), field_ref("i32"), literal(42)), + cast(field_ref("i8"), int32())), + &bound); + EXPECT_TRUE(bound.IsBound()); + EXPECT_TRUE(bound.type()->Equals(*int32())); + } + { + Expression bound; + ExpectBindsTo( + add(if_else_special(field_ref("bool"), field_ref("i8"), literal(42)), + field_ref("i32")), + add(cast(if_else_special(field_ref("bool"), field_ref("i8"), literal(42)), + int32()), + field_ref("i32")), + &bound); + EXPECT_TRUE(bound.IsBound()); + EXPECT_TRUE(bound.type()->Equals(*int32())); + } + + // Self-nested. + { + Expression bound; + ExpectBindsTo( + if_else_special( + if_else_special(literal(true), literal(true), literal(false)), + if_else_special(field_ref("bool"), field_ref("i8"), field_ref("i32")), + if_else_special(field_ref("bool"), field_ref("i8"), field_ref("i64"))), + if_else_special( + if_else_special(literal(true), literal(true), literal(false)), + cast(if_else_special(field_ref("bool"), cast(field_ref("i8"), int32()), + field_ref("i32")), + int64()), + if_else_special(field_ref("bool"), cast(field_ref("i8"), int64()), + field_ref("i64"))), + &bound); + EXPECT_TRUE(bound.IsBound()); + EXPECT_TRUE(bound.type()->Equals(*int64())); + } +} + TEST(Expression, ExecuteFieldRef) { auto ExpectRefIs = [](FieldRef ref, Datum in, Datum expected) { auto expr = field_ref(ref); diff --git a/cpp/src/arrow/compute/special_form.cc b/cpp/src/arrow/compute/special_form.cc index b6196f766441..d7657fd3f3d5 100644 --- a/cpp/src/arrow/compute/special_form.cc +++ b/cpp/src/arrow/compute/special_form.cc @@ -603,15 +603,17 @@ class FunctionBackedSpecialForm : public SpecialForm { ExecContext* exec_context) override { DCHECK(std::all_of(arguments.begin(), arguments.end(), [](const Expression& argument) { return argument.IsBound(); })); - Expression::Call call; call.function_name = name(); call.arguments = std::move(arguments); call.options = std::move(options); ARROW_ASSIGN_OR_RAISE( auto bound, BindNonRecursive(call, /*insert_implicit_casts=*/true, exec_context)); - - return static_cast(this)->BindWithBoundCall(CallNotNull(bound), + auto bound_call = CallNotNull(bound); + auto bound_call_copy = *bound_call; + arguments = std::move(bound_call->arguments); + options = std::move(bound_call->options); + return static_cast(this)->BindWithBoundCall(std::move(bound_call_copy), exec_context); } }; @@ -628,11 +630,11 @@ class ConditionalSpecialForm protected: Result> BindWithBoundCall( - const Expression::Call* call, ExecContext* exec_context) const { + Expression::Call call, ExecContext* exec_context) const { auto branches = - static_cast(this)->GetBranches(std::move(call->arguments)); + static_cast(this)->GetBranches(std::move(call.arguments)); return std::make_unique( - std::move(branches), std::move(call->type), std::move(call->options)); + std::move(branches), std::move(call.type), std::move(call.options)); } friend class FunctionBackedSpecialForm>; From a7096f7071cc03687464e46f40614911663c5900 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Fri, 5 Sep 2025 14:30:44 +0800 Subject: [PATCH 30/71] Basic expression tests for special form --- cpp/src/arrow/compute/expression_test.cc | 50 +++++++++++++++++++++++- 1 file changed, 49 insertions(+), 1 deletion(-) diff --git a/cpp/src/arrow/compute/expression_test.cc b/cpp/src/arrow/compute/expression_test.cc index 65c3e7a9333e..07fb41e6385e 100644 --- a/cpp/src/arrow/compute/expression_test.cc +++ b/cpp/src/arrow/compute/expression_test.cc @@ -375,6 +375,15 @@ TEST(Expression, Equality) { EXPECT_NE(cast(field_ref("a"), int32()), cast(field_ref("a"), int64())); EXPECT_NE(cast(field_ref("a"), int32()), call("cast", {field_ref("a")}, compute::CastOptions::Unsafe(int32()))); + + EXPECT_EQ(if_else_special(literal(true), field_ref("a"), field_ref("b")), + if_else_special(literal(true), field_ref("a"), field_ref("b"))); + EXPECT_NE(if_else_special(literal(true), field_ref("a"), field_ref("b")), + if_else_special(literal(false), field_ref("a"), field_ref("b"))); + EXPECT_NE(if_else_special(literal(true), field_ref("a"), field_ref("b")), + if_else_special(literal(true), field_ref("b"), field_ref("b"))); + EXPECT_NE(if_else_special(literal(true), field_ref("a"), field_ref("b")), + if_else_special(literal(true), field_ref("a"), field_ref("a"))); } Expression null_literal(const std::shared_ptr& type) { @@ -402,7 +411,17 @@ TEST(Expression, Hash) { // NB: unbound expressions don't check for availability in any registry EXPECT_TRUE(set.emplace(call("widgetify", {})).second); - EXPECT_EQ(set.size(), 8); + EXPECT_TRUE( + set.emplace(if_else_special(field_ref("cond"), field_ref("a"), field_ref("b"))) + .second); + EXPECT_FALSE( + set.emplace(if_else_special(field_ref("cond"), field_ref("a"), field_ref("b"))) + .second); + EXPECT_TRUE( + set.emplace(if_else_special(field_ref("cond"), field_ref("b"), field_ref("a"))) + .second); + + EXPECT_EQ(set.size(), 10); } TEST(Expression, IsScalarExpression) { @@ -422,6 +441,9 @@ TEST(Expression, IsScalarExpression) { // non scalar function EXPECT_FALSE(call("take", {field_ref("a"), literal(arr)}).IsScalarExpression()); + + EXPECT_TRUE(if_else_special(field_ref("cond"), field_ref("a"), field_ref("b")) + .IsScalarExpression()); } TEST(Expression, IsSatisfiable) { @@ -491,6 +513,9 @@ TEST(Expression, IsSatisfiable) { // fill_na) EXPECT_TRUE(Bind(call("is_null", {never_true})).IsSatisfiable()); } + + EXPECT_TRUE(Bind(if_else_special(field_ref("bool"), field_ref("i32"), field_ref("i32"))) + .IsSatisfiable()); } TEST(Expression, FieldsInExpression) { @@ -518,6 +543,20 @@ TEST(Expression, FieldsInExpression) { equal(field_ref("b"), literal(2))), not_(less(field_ref("c"), literal(3)))), {"a", "b", "c"}); + + ExpectFieldsAre(if_else_special(literal(true), field_ref("a"), field_ref("b")), + {"a", "b"}); + ExpectFieldsAre(if_else_special(field_ref("a"), field_ref("b"), field_ref("b")), + {"a", "b", "b"}); + ExpectFieldsAre(if_else_special(field_ref("a"), field_ref("b"), field_ref("c")), + {"a", "b", "c"}); + ExpectFieldsAre( + if_else_special(call("not", {field_ref("a")}), call("not", {field_ref("b")}), + call("not", {field_ref("c")})), + {"a", "b", "c"}); + ExpectFieldsAre( + call("not", {if_else_special(field_ref("a"), field_ref("b"), field_ref("c"))}), + {"a", "b", "c"}); } TEST(Expression, ExpressionHasFieldRefs) { @@ -540,6 +579,15 @@ TEST(Expression, ExpressionHasFieldRefs) { EXPECT_TRUE(ExpressionHasFieldRefs(or_( and_(not_(equal(field_ref("a"), literal(1))), equal(field_ref("b"), literal(2))), not_(less(field_ref("c"), literal(3)))))); + + EXPECT_FALSE( + ExpressionHasFieldRefs(if_else_special(literal(true), literal(1), literal(0)))); + EXPECT_TRUE( + ExpressionHasFieldRefs(if_else_special(field_ref("a"), literal(1), literal(0)))); + EXPECT_TRUE( + ExpressionHasFieldRefs(if_else_special(literal(true), field_ref("a"), literal(0)))); + EXPECT_TRUE( + ExpressionHasFieldRefs(if_else_special(literal(true), literal(0), field_ref("a")))); } TEST(Expression, BindLiteral) { From 410a6366cd44850a448aded1c6b811bbba0e9ec2 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Sun, 7 Sep 2025 23:53:29 +0800 Subject: [PATCH 31/71] Add special form execution test --- cpp/src/arrow/compute/expression_test.cc | 137 ++++++++++++++++++++++- cpp/src/arrow/compute/special_form.cc | 22 ++-- cpp/src/arrow/compute/special_form.h | 4 +- 3 files changed, 146 insertions(+), 17 deletions(-) diff --git a/cpp/src/arrow/compute/expression_test.cc b/cpp/src/arrow/compute/expression_test.cc index 07fb41e6385e..f8af466ae077 100644 --- a/cpp/src/arrow/compute/expression_test.cc +++ b/cpp/src/arrow/compute/expression_test.cc @@ -31,8 +31,10 @@ #include "arrow/compute/expression_internal.h" #include "arrow/compute/function_internal.h" #include "arrow/compute/registry.h" +#include "arrow/compute/special_form.h" #include "arrow/testing/gtest_util.h" #include "arrow/testing/matchers.h" +#include "arrow/util/logging_internal.h" using testing::Eq; using testing::HasSubstr; @@ -47,6 +49,8 @@ using internal::checked_pointer_cast; namespace compute { +namespace { + const std::shared_ptr kBoringSchema = schema({ field("bool", boolean()), field("i8", int8()), @@ -92,6 +96,42 @@ std::string make_range_json(int start, int end) { const auto no_change = std::nullopt; +class EchoSpecialExecutor : public SpecialExecutor { + public: + EchoSpecialExecutor(Expression argument) + : SpecialExecutor(argument.type()), argument_(std::move(argument)) {} + + Result Execute(const ExecBatch& input, + ExecContext* exec_context) const override { + return ExecuteScalarExpression(argument_, input, exec_context); + } + + private: + Expression argument_; +}; + +class EchoSpecialForm : public SpecialForm { + public: + EchoSpecialForm() : SpecialForm("echo") {} + + protected: + Result> Bind( + std::vector& arguments, std::shared_ptr options, + ExecContext* exec_context) const override { + DCHECK_EQ(arguments.size(), 1); + return std::make_unique(arguments[0]); + } +}; + +Expression echo_special(Expression input) { + Expression::Special special; + special.special_form = std::make_shared(); + special.arguments.push_back(std::move(input)); + return Expression(std::move(special)); +} + +} // namespace + TEST(ExpressionUtils, Comparison) { auto cmp_name = [](Datum l, Datum r) { return Comparison::Execute(l, r).Map(Comparison::GetName); @@ -323,6 +363,12 @@ TEST(Expression, ToString) { "round(3.14, {ndigits=0, round_mode=HALF_TO_EVEN})"); EXPECT_EQ(call("random", {}, compute::RandomOptions()).ToString(), "random({initializer=SystemRandom, seed=0})"); + + EXPECT_EQ(echo_special(field_ref("a")).ToString(), "echo_special(a)"); + EXPECT_EQ( + if_else_special(field_ref("cond"), field_ref("if_true"), field_ref("if_false")) + .ToString(), + "if_else_special(cond, if_true, if_false)"); } TEST(Expression, Equality) { @@ -1190,15 +1236,32 @@ Result NaiveExecuteScalarExpression(const Expression& expr, const Datum& return ref->GetOneOrNone(*input.record_batch()); } - auto call = CallNotNull(expr); + auto execute_args = [](const std::vector& args, + const Datum& input) -> Result> { + std::vector results(args.size()); + for (size_t i = 0; i < results.size(); ++i) { + ARROW_ASSIGN_OR_RAISE(results[i], NaiveExecuteScalarExpression(args[i], input)); + } + return results; + }; - std::vector arguments(call->arguments.size()); - for (size_t i = 0; i < arguments.size(); ++i) { - ARROW_ASSIGN_OR_RAISE(arguments[i], - NaiveExecuteScalarExpression(call->arguments[i], input)); + compute::ExecContext exec_context; + + if (auto special = expr.special()) { + auto arguments_cpy = special->arguments; + ARROW_ASSIGN_OR_RAISE( + auto executor, + special->special_form->Bind(arguments_cpy, special->options, &exec_context)); + ARROW_ASSIGN_OR_RAISE(auto arguments, execute_args(special->arguments, input)); + ExecBatch batch{std::move(arguments), input.length()}; + return executor->Execute(batch, &exec_context); } - compute::ExecContext exec_context; + auto call = CallNotNull(expr); + + ARROW_ASSIGN_OR_RAISE(std::vector arguments, + execute_args(call->arguments, input)); + ARROW_ASSIGN_OR_RAISE(auto function, GetFunction(*call, &exec_context)); std::vector types = GetTypes(call->arguments); @@ -1363,6 +1426,68 @@ TEST(Expression, ExecuteDictionaryTransparent) { ])")); } +TEST(Expression, NaiveExecuteSpecialForm) { + ExpectExecute(echo_special(field_ref("i32")), + ArrayFromJSON(struct_({field("i32", int32())}), R"([ + {"i32": 0}, + {"i32": 1}, + {"i32": 2} + ])")); + + ExpectExecute( + if_else_special(field_ref("cond"), field_ref("if_true"), field_ref("if_false")), + ArrayFromJSON(struct_({field("cond", boolean()), field("if_true", int32()), + field("if_false", int32())}), + R"([ + {"cond": null, "if_true": 0, "if_false": 1}, + {"cond": false, "if_true": 2, "if_false": 3}, + {"cond": true, "if_true": 4, "if_false": 5} + ])")); +} + +TEST(Expression, ExecuteSpecialFormNested) { + const int64_t length = 3; + + auto array = ArrayFromJSON(int32(), R"([null, 42, 42])"); + auto scalar = ScalarFromJSON(int32(), R"(42)"); + auto chunked_array = ChunkedArrayFromJSON(int32(), {R"([null, 42])", R"([42])"}); + std::vector inputs = {ExecBatch{{array, array}, length}, + ExecBatch{{array, scalar}, length}, + ExecBatch{{array, chunked_array}, length}, + ExecBatch{{scalar, array}, length}, + ExecBatch{{scalar, scalar}, 1}, + ExecBatch{{scalar, chunked_array}, length}, + ExecBatch{{chunked_array, array}, length}, + ExecBatch{{chunked_array, scalar}, length}, + ExecBatch{{chunked_array, chunked_array}, length}}; + + auto schm = schema({field("a", int32()), field("b", int32())}); + auto a = field_ref("a"); + auto b = field_ref("b"); + std::vector exprs = {add(echo_special(a), b), + add(a, echo_special(b)), + add(echo_special(a), echo_special(b)), + echo_special(add(a, b)), + echo_special(add(echo_special(a), b)), + echo_special(add(a, echo_special(b))), + echo_special(add(echo_special(a), echo_special(b)))}; + + for (const auto& input : inputs) { + ARROW_SCOPED_TRACE("Input: " + input.ToString()); + + ASSERT_OK_AND_ASSIGN(auto bound, add(a, b).Bind(*schm)); + ASSERT_OK_AND_ASSIGN(auto expected, ExecuteScalarExpression(bound, input)); + + for (const auto& expr : exprs) { + ARROW_SCOPED_TRACE("Expr: " + expr.ToString()); + + ASSERT_OK_AND_ASSIGN(auto bound, expr.Bind(*schm)); + ASSERT_OK_AND_ASSIGN(auto result, ExecuteScalarExpression(bound, input)); + AssertDatumsEqual(result, expected, /*verbose=*/true); + } + } +} + void ExpectIdenticalIfUnchanged(Expression modified, Expression original) { if (modified == original) { // no change -> must be identical diff --git a/cpp/src/arrow/compute/special_form.cc b/cpp/src/arrow/compute/special_form.cc index d7657fd3f3d5..8c01bf2d5384 100644 --- a/cpp/src/arrow/compute/special_form.cc +++ b/cpp/src/arrow/compute/special_form.cc @@ -576,14 +576,13 @@ struct ConditionalExec { class ConditionalSpecialExecutor : public SpecialExecutor { public: - explicit ConditionalSpecialExecutor(std::vector branches, TypeHolder out_type, - std::shared_ptr options) - : SpecialExecutor(std::move(out_type), std::move(options)), + ConditionalSpecialExecutor(std::vector branches, TypeHolder out_type) + : SpecialExecutor(std::move(out_type), /*options=*/nullptr), branches(std::move(branches)) {} Result Execute(const ExecBatch& input, ExecContext* exec_context) const override { - return ConditionalExec(branches, out_type()).Execute(input, exec_context); + return ConditionalExec(branches, out_type_).Execute(input, exec_context); } private: @@ -598,9 +597,9 @@ class FunctionBackedSpecialForm : public SpecialForm { ARROW_DISALLOW_COPY_AND_ASSIGN(FunctionBackedSpecialForm); ARROW_DEFAULT_MOVE_AND_ASSIGN(FunctionBackedSpecialForm); - Result> Bind(std::vector& arguments, - std::shared_ptr options, - ExecContext* exec_context) override { + Result> Bind( + std::vector& arguments, std::shared_ptr options, + ExecContext* exec_context) const override { DCHECK(std::all_of(arguments.begin(), arguments.end(), [](const Expression& argument) { return argument.IsBound(); })); Expression::Call call; @@ -631,10 +630,13 @@ class ConditionalSpecialForm protected: Result> BindWithBoundCall( Expression::Call call, ExecContext* exec_context) const { + // Shouldn't have options. This is guaranteed by the call binding. + DCHECK_EQ(call.options, nullptr); + auto branches = static_cast(this)->GetBranches(std::move(call.arguments)); - return std::make_unique( - std::move(branches), std::move(call.type), std::move(call.options)); + return std::make_unique(std::move(branches), + std::move(call.type)); } friend class FunctionBackedSpecialForm>; @@ -649,7 +651,9 @@ class IfElseSpecialForm : public ConditionalSpecialForm { protected: std::vector GetBranches(std::vector arguments) const { + // The arity should be correct. This is guaranteed by the call binding. DCHECK_EQ(arguments.size(), 3); + auto cond = std::move(arguments[0]); auto if_true = std::move(arguments[1]); auto if_false = std::move(arguments[2]); diff --git a/cpp/src/arrow/compute/special_form.h b/cpp/src/arrow/compute/special_form.h index 5b8cf1e923fd..786484ad5263 100644 --- a/cpp/src/arrow/compute/special_form.h +++ b/cpp/src/arrow/compute/special_form.h @@ -33,7 +33,7 @@ class ARROW_EXPORT SpecialExecutor { virtual Result Execute(const ExecBatch& input, ExecContext* exec_context) const = 0; - private: + protected: const TypeHolder out_type_; const std::shared_ptr options_; }; @@ -48,7 +48,7 @@ class ARROW_EXPORT SpecialForm { virtual Result> Bind( std::vector& arguments, std::shared_ptr options, - ExecContext* exec_context) = 0; + ExecContext* exec_context) const = 0; private: std::string name_; From c6d733afd087d8246dcf3f50a50a66a81b8b284a Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Mon, 8 Sep 2025 00:28:02 +0800 Subject: [PATCH 32/71] Add more special form test using echo_special --- cpp/src/arrow/compute/expression_test.cc | 66 +++++++++++++++++++++++- 1 file changed, 64 insertions(+), 2 deletions(-) diff --git a/cpp/src/arrow/compute/expression_test.cc b/cpp/src/arrow/compute/expression_test.cc index f8af466ae077..e002edba045a 100644 --- a/cpp/src/arrow/compute/expression_test.cc +++ b/cpp/src/arrow/compute/expression_test.cc @@ -422,6 +422,19 @@ TEST(Expression, Equality) { EXPECT_NE(cast(field_ref("a"), int32()), call("cast", {field_ref("a")}, compute::CastOptions::Unsafe(int32()))); + EXPECT_EQ(echo_special(literal(42)), echo_special(literal(42))); + EXPECT_NE(echo_special(literal(42)), literal(42)); + EXPECT_NE(echo_special(literal(42)), echo_special(literal(0))); + EXPECT_EQ(echo_special(field_ref("a")), echo_special(field_ref("a"))); + EXPECT_NE(echo_special(field_ref("a")), field_ref("a")); + EXPECT_NE(echo_special(field_ref("a")), echo_special(field_ref("b"))); + EXPECT_EQ(echo_special(add(literal(42), field_ref("a"))), + echo_special(add(literal(42), field_ref("a")))); + EXPECT_NE(echo_special(add(literal(42), field_ref("a"))), + add(literal(42), field_ref("a"))); + EXPECT_NE(echo_special(add(literal(42), field_ref("a"))), + echo_special(add(field_ref("a"), literal(42)))); + EXPECT_EQ(if_else_special(literal(true), field_ref("a"), field_ref("b")), if_else_special(literal(true), field_ref("a"), field_ref("b"))); EXPECT_NE(if_else_special(literal(true), field_ref("a"), field_ref("b")), @@ -430,6 +443,8 @@ TEST(Expression, Equality) { if_else_special(literal(true), field_ref("b"), field_ref("b"))); EXPECT_NE(if_else_special(literal(true), field_ref("a"), field_ref("b")), if_else_special(literal(true), field_ref("a"), field_ref("a"))); + EXPECT_NE(if_else_special(literal(true), field_ref("a"), field_ref("b")), + call("if_else", {literal(true), field_ref("a"), field_ref("b")})); } Expression null_literal(const std::shared_ptr& type) { @@ -457,6 +472,10 @@ TEST(Expression, Hash) { // NB: unbound expressions don't check for availability in any registry EXPECT_TRUE(set.emplace(call("widgetify", {})).second); + EXPECT_TRUE(set.emplace(echo_special(field_ref("a"))).second); + EXPECT_FALSE(set.emplace(echo_special(field_ref("a"))).second) << "already inserted"; + EXPECT_TRUE(set.emplace(echo_special(field_ref("b"))).second); + EXPECT_TRUE( set.emplace(if_else_special(field_ref("cond"), field_ref("a"), field_ref("b"))) .second); @@ -467,7 +486,7 @@ TEST(Expression, Hash) { set.emplace(if_else_special(field_ref("cond"), field_ref("b"), field_ref("a"))) .second); - EXPECT_EQ(set.size(), 10); + EXPECT_EQ(set.size(), 12); } TEST(Expression, IsScalarExpression) { @@ -488,6 +507,8 @@ TEST(Expression, IsScalarExpression) { // non scalar function EXPECT_FALSE(call("take", {field_ref("a"), literal(arr)}).IsScalarExpression()); + EXPECT_TRUE(echo_special(field_ref("a")).IsScalarExpression()); + EXPECT_TRUE(if_else_special(field_ref("cond"), field_ref("a"), field_ref("b")) .IsScalarExpression()); } @@ -560,6 +581,8 @@ TEST(Expression, IsSatisfiable) { EXPECT_TRUE(Bind(call("is_null", {never_true})).IsSatisfiable()); } + EXPECT_TRUE(Bind(echo_special(field_ref("i32"))).IsSatisfiable()); + EXPECT_TRUE(Bind(if_else_special(field_ref("bool"), field_ref("i32"), field_ref("i32"))) .IsSatisfiable()); } @@ -590,6 +613,10 @@ TEST(Expression, FieldsInExpression) { not_(less(field_ref("c"), literal(3)))), {"a", "b", "c"}); + ExpectFieldsAre(echo_special(literal(42)), {}); + ExpectFieldsAre(echo_special(field_ref("a")), {"a"}); + + ExpectFieldsAre(if_else_special(literal(true), literal(1), literal(0)), {}); ExpectFieldsAre(if_else_special(literal(true), field_ref("a"), field_ref("b")), {"a", "b"}); ExpectFieldsAre(if_else_special(field_ref("a"), field_ref("b"), field_ref("b")), @@ -626,6 +653,9 @@ TEST(Expression, ExpressionHasFieldRefs) { and_(not_(equal(field_ref("a"), literal(1))), equal(field_ref("b"), literal(2))), not_(less(field_ref("c"), literal(3)))))); + EXPECT_FALSE(ExpressionHasFieldRefs(echo_special(literal(42)))); + EXPECT_TRUE(ExpressionHasFieldRefs(echo_special(field_ref("a")))); + EXPECT_FALSE( ExpressionHasFieldRefs(if_else_special(literal(true), literal(1), literal(0)))); EXPECT_TRUE( @@ -1042,7 +1072,39 @@ TEST(Expression, BindNestedCall) { EXPECT_TRUE(expr.IsBound()); } -TEST(Expression, BindIfElseSpecialForm) { +TEST(Expression, BindSpecialForm) { + { + auto expr = echo_special(literal(42)); + EXPECT_FALSE(expr.IsBound()); + ExpectBindsTo(expr, no_change, &expr); + EXPECT_TRUE(expr.IsBound()); + EXPECT_TRUE(expr.type()->Equals(*int32())); + } + + { + auto expr = echo_special(field_ref("bool")); + EXPECT_FALSE(expr.IsBound()); + ExpectBindsTo(expr, no_change, &expr); + EXPECT_TRUE(expr.IsBound()); + EXPECT_TRUE(expr.type()->Equals(*boolean())); + } + + { + auto expr = echo_special(add(field_ref("i64"), literal(42))); + EXPECT_FALSE(expr.IsBound()); + ExpectBindsTo(expr, no_change, &expr); + EXPECT_TRUE(expr.IsBound()); + EXPECT_TRUE(expr.type()->Equals(*int64())); + } + + { + auto expr = add(literal(42), echo_special(field_ref("i64"))); + EXPECT_FALSE(expr.IsBound()); + ExpectBindsTo(expr, no_change, &expr); + EXPECT_TRUE(expr.IsBound()); + EXPECT_TRUE(expr.type()->Equals(*int64())); + } + { auto expr = if_else_special(field_ref("bool"), field_ref("i8"), field_ref("i8")); EXPECT_FALSE(expr.IsBound()); From 0c22f9503101fc4f68e515d4b07fb595b64b6fb3 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Fri, 3 Oct 2025 15:37:58 -0700 Subject: [PATCH 33/71] Move if_else special form from compute core to compute lib (#58) --- cpp/src/arrow/CMakeLists.txt | 2 +- cpp/src/arrow/compute/CMakeLists.txt | 10 +- cpp/src/arrow/compute/api.h | 6 + cpp/src/arrow/compute/api_special.h | 26 ++ cpp/src/arrow/compute/expression.cc | 7 - cpp/src/arrow/compute/expression.h | 2 - cpp/src/arrow/compute/expression_test.cc | 241 +----------------- .../arrow/compute/expression_test_internal.h | 90 +++++++ cpp/src/arrow/compute/special/CMakeLists.txt | 26 ++ .../if_else_special.cc} | 14 +- .../if_else_special_benchmark.cc} | 0 .../if_else_special_test.cc} | 199 ++++++++++++++- cpp/src/arrow/compute/special/meson.build | 26 ++ cpp/src/arrow/compute/special_form.h | 2 - 14 files changed, 387 insertions(+), 264 deletions(-) create mode 100644 cpp/src/arrow/compute/api_special.h create mode 100644 cpp/src/arrow/compute/expression_test_internal.h create mode 100644 cpp/src/arrow/compute/special/CMakeLists.txt rename cpp/src/arrow/compute/{special_form.cc => special/if_else_special.cc} (98%) rename cpp/src/arrow/compute/{special_form_benchmark.cc => special/if_else_special_benchmark.cc} (100%) rename cpp/src/arrow/compute/{special_form_test.cc => special/if_else_special_test.cc} (85%) create mode 100644 cpp/src/arrow/compute/special/meson.build diff --git a/cpp/src/arrow/CMakeLists.txt b/cpp/src/arrow/CMakeLists.txt index a6975dcf4af8..5e49f3750ad0 100644 --- a/cpp/src/arrow/CMakeLists.txt +++ b/cpp/src/arrow/CMakeLists.txt @@ -728,7 +728,6 @@ set(ARROW_COMPUTE_SRCS compute/kernel.cc compute/ordering.cc compute/registry.cc - compute/special_form.cc compute/kernels/chunked_internal.cc compute/kernels/codegen_internal.cc compute/kernels/scalar_cast_boolean.cc @@ -803,6 +802,7 @@ if(ARROW_COMPUTE) compute/row/grouper.cc compute/row/row_encoder_internal.cc compute/row/row_internal.cc + compute/special/if_else_special.cc compute/util.cc compute/util_internal.cc) diff --git a/cpp/src/arrow/compute/CMakeLists.txt b/cpp/src/arrow/compute/CMakeLists.txt index 16c8a4a69f57..6434f9edf75f 100644 --- a/cpp/src/arrow/compute/CMakeLists.txt +++ b/cpp/src/arrow/compute/CMakeLists.txt @@ -174,16 +174,10 @@ add_arrow_compute_test(row_test EXTRA_LINK_LIBS arrow_compute_testing) -add_arrow_compute_test(special_form_test - SOURCES - special_form_test.cc - EXTRA_LINK_LIBS - arrow_compute_testing) - add_arrow_compute_benchmark(function_benchmark) -add_arrow_compute_benchmark(special_form_benchmark PREFIX "arrow-compute") - add_subdirectory(kernels) add_subdirectory(row) + +add_subdirectory(special) diff --git a/cpp/src/arrow/compute/api.h b/cpp/src/arrow/compute/api.h index 343e30643cfd..85288de2a106 100644 --- a/cpp/src/arrow/compute/api.h +++ b/cpp/src/arrow/compute/api.h @@ -39,8 +39,14 @@ #include "arrow/compute/registry.h" // IWYU pragma: export #include "arrow/datum.h" // IWYU pragma: export +/// \defgroup expressions Expressions of computation +/// @{ +/// @} + #include "arrow/compute/expression.h" // IWYU pragma: export +#include "arrow/compute/api_special.h" // IWYU pragma: export + /// \defgroup execnode-row Utilities for working with data in a row-major format /// @{ /// @} diff --git a/cpp/src/arrow/compute/api_special.h b/cpp/src/arrow/compute/api_special.h new file mode 100644 index 000000000000..77b0f2b20899 --- /dev/null +++ b/cpp/src/arrow/compute/api_special.h @@ -0,0 +1,26 @@ +// 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/expression.h" + +namespace arrow::compute { + +/// TODO: Doc. +ARROW_EXPORT +Expression if_else_special(Expression cond, Expression if_true, Expression if_false); + +} // namespace arrow::compute \ No newline at end of file diff --git a/cpp/src/arrow/compute/expression.cc b/cpp/src/arrow/compute/expression.cc index fe13056644e4..c884e6d205f2 100644 --- a/cpp/src/arrow/compute/expression.cc +++ b/cpp/src/arrow/compute/expression.cc @@ -96,13 +96,6 @@ Expression call(std::string function, std::vector arguments, return Expression(std::move(call)); } -Expression if_else_special(Expression cond, Expression if_true, Expression if_false) { - Expression::Special special; - special.special_form = GetIfElseSpecialForm(); - special.arguments = {std::move(cond), std::move(if_true), std::move(if_false)}; - return Expression(std::move(special)); -} - const Datum* Expression::literal() const { if (impl_ == nullptr) return nullptr; diff --git a/cpp/src/arrow/compute/expression.h b/cpp/src/arrow/compute/expression.h index 2311f48a44a4..9d2c3eb81e1a 100644 --- a/cpp/src/arrow/compute/expression.h +++ b/cpp/src/arrow/compute/expression.h @@ -186,8 +186,6 @@ Expression call(std::string function, std::vector arguments, std::make_shared(std::move(options))); } -Expression if_else_special(Expression cond, Expression if_true, Expression if_false); - /// Assemble a list of all fields referenced by an Expression at any depth. ARROW_EXPORT std::vector FieldsInExpression(const Expression&); diff --git a/cpp/src/arrow/compute/expression_test.cc b/cpp/src/arrow/compute/expression_test.cc index e002edba045a..c680708c02c9 100644 --- a/cpp/src/arrow/compute/expression_test.cc +++ b/cpp/src/arrow/compute/expression_test.cc @@ -28,11 +28,10 @@ #include #include "arrow/array/builder_primitive.h" -#include "arrow/compute/expression_internal.h" +#include "arrow/compute/expression_test_internal.h" #include "arrow/compute/function_internal.h" #include "arrow/compute/registry.h" #include "arrow/compute/special_form.h" -#include "arrow/testing/gtest_util.h" #include "arrow/testing/matchers.h" #include "arrow/util/logging_internal.h" @@ -49,52 +48,13 @@ using internal::checked_pointer_cast; namespace compute { -namespace { - -const std::shared_ptr kBoringSchema = schema({ - field("bool", boolean()), - field("i8", int8()), - field("i32", int32()), - field("i32_req", int32(), /*nullable=*/false), - field("u32", uint32()), - field("i64", int64()), - field("f32", float32()), - field("f32_req", float32(), /*nullable=*/false), - field("f64", float64()), - field("date64", date64()), - field("str", utf8()), - field("dict_str", dictionary(int32(), utf8())), - field("dict_i32", dictionary(int32(), int32())), - field("ts_ns", timestamp(TimeUnit::NANO)), - field("ts_s", timestamp(TimeUnit::SECOND)), - field("binary", binary()), - field("ts_s_utc", timestamp(TimeUnit::SECOND, "UTC")), -}); - -Expression cast(Expression argument, std::shared_ptr to_type) { - return call("cast", {std::move(argument)}, - compute::CastOptions::Safe(std::move(to_type))); -} - -Expression true_unless_null(Expression argument) { - return call("true_unless_null", {std::move(argument)}); -} - -Expression add(Expression l, Expression r) { - return call("add", {std::move(l), std::move(r)}); -} - -std::string make_range_json(int start, int end) { - std::string result = "["; - for (int i = start; i <= end; ++i) { - if (i > start) result += ","; - result += std::to_string(i); - } - result += "]"; - return result; -} - -const auto no_change = std::nullopt; +using internal::add; +using internal::cast; +using internal::ExpectBindsTo; +using internal::kBoringSchema; +using internal::make_range_json; +using internal::no_change; +using internal::true_unless_null; class EchoSpecialExecutor : public SpecialExecutor { public: @@ -130,8 +90,6 @@ Expression echo_special(Expression input) { return Expression(std::move(special)); } -} // namespace - TEST(ExpressionUtils, Comparison) { auto cmp_name = [](Datum l, Datum r) { return Comparison::Execute(l, r).Map(Comparison::GetName); @@ -365,10 +323,6 @@ TEST(Expression, ToString) { "random({initializer=SystemRandom, seed=0})"); EXPECT_EQ(echo_special(field_ref("a")).ToString(), "echo_special(a)"); - EXPECT_EQ( - if_else_special(field_ref("cond"), field_ref("if_true"), field_ref("if_false")) - .ToString(), - "if_else_special(cond, if_true, if_false)"); } TEST(Expression, Equality) { @@ -434,17 +388,6 @@ TEST(Expression, Equality) { add(literal(42), field_ref("a"))); EXPECT_NE(echo_special(add(literal(42), field_ref("a"))), echo_special(add(field_ref("a"), literal(42)))); - - EXPECT_EQ(if_else_special(literal(true), field_ref("a"), field_ref("b")), - if_else_special(literal(true), field_ref("a"), field_ref("b"))); - EXPECT_NE(if_else_special(literal(true), field_ref("a"), field_ref("b")), - if_else_special(literal(false), field_ref("a"), field_ref("b"))); - EXPECT_NE(if_else_special(literal(true), field_ref("a"), field_ref("b")), - if_else_special(literal(true), field_ref("b"), field_ref("b"))); - EXPECT_NE(if_else_special(literal(true), field_ref("a"), field_ref("b")), - if_else_special(literal(true), field_ref("a"), field_ref("a"))); - EXPECT_NE(if_else_special(literal(true), field_ref("a"), field_ref("b")), - call("if_else", {literal(true), field_ref("a"), field_ref("b")})); } Expression null_literal(const std::shared_ptr& type) { @@ -476,17 +419,7 @@ TEST(Expression, Hash) { EXPECT_FALSE(set.emplace(echo_special(field_ref("a"))).second) << "already inserted"; EXPECT_TRUE(set.emplace(echo_special(field_ref("b"))).second); - EXPECT_TRUE( - set.emplace(if_else_special(field_ref("cond"), field_ref("a"), field_ref("b"))) - .second); - EXPECT_FALSE( - set.emplace(if_else_special(field_ref("cond"), field_ref("a"), field_ref("b"))) - .second); - EXPECT_TRUE( - set.emplace(if_else_special(field_ref("cond"), field_ref("b"), field_ref("a"))) - .second); - - EXPECT_EQ(set.size(), 12); + EXPECT_EQ(set.size(), 10); } TEST(Expression, IsScalarExpression) { @@ -508,9 +441,6 @@ TEST(Expression, IsScalarExpression) { EXPECT_FALSE(call("take", {field_ref("a"), literal(arr)}).IsScalarExpression()); EXPECT_TRUE(echo_special(field_ref("a")).IsScalarExpression()); - - EXPECT_TRUE(if_else_special(field_ref("cond"), field_ref("a"), field_ref("b")) - .IsScalarExpression()); } TEST(Expression, IsSatisfiable) { @@ -582,9 +512,6 @@ TEST(Expression, IsSatisfiable) { } EXPECT_TRUE(Bind(echo_special(field_ref("i32"))).IsSatisfiable()); - - EXPECT_TRUE(Bind(if_else_special(field_ref("bool"), field_ref("i32"), field_ref("i32"))) - .IsSatisfiable()); } TEST(Expression, FieldsInExpression) { @@ -615,21 +542,6 @@ TEST(Expression, FieldsInExpression) { ExpectFieldsAre(echo_special(literal(42)), {}); ExpectFieldsAre(echo_special(field_ref("a")), {"a"}); - - ExpectFieldsAre(if_else_special(literal(true), literal(1), literal(0)), {}); - ExpectFieldsAre(if_else_special(literal(true), field_ref("a"), field_ref("b")), - {"a", "b"}); - ExpectFieldsAre(if_else_special(field_ref("a"), field_ref("b"), field_ref("b")), - {"a", "b", "b"}); - ExpectFieldsAre(if_else_special(field_ref("a"), field_ref("b"), field_ref("c")), - {"a", "b", "c"}); - ExpectFieldsAre( - if_else_special(call("not", {field_ref("a")}), call("not", {field_ref("b")}), - call("not", {field_ref("c")})), - {"a", "b", "c"}); - ExpectFieldsAre( - call("not", {if_else_special(field_ref("a"), field_ref("b"), field_ref("c"))}), - {"a", "b", "c"}); } TEST(Expression, ExpressionHasFieldRefs) { @@ -655,15 +567,6 @@ TEST(Expression, ExpressionHasFieldRefs) { EXPECT_FALSE(ExpressionHasFieldRefs(echo_special(literal(42)))); EXPECT_TRUE(ExpressionHasFieldRefs(echo_special(field_ref("a")))); - - EXPECT_FALSE( - ExpressionHasFieldRefs(if_else_special(literal(true), literal(1), literal(0)))); - EXPECT_TRUE( - ExpressionHasFieldRefs(if_else_special(field_ref("a"), literal(1), literal(0)))); - EXPECT_TRUE( - ExpressionHasFieldRefs(if_else_special(literal(true), field_ref("a"), literal(0)))); - EXPECT_TRUE( - ExpressionHasFieldRefs(if_else_special(literal(true), literal(0), field_ref("a")))); } TEST(Expression, BindLiteral) { @@ -679,24 +582,6 @@ TEST(Expression, BindLiteral) { } } -void ExpectBindsTo(Expression expr, std::optional expected, - Expression* bound_out = nullptr, - const Schema& schema = *kBoringSchema) { - if (!expected) { - expected = expr; - } - - ASSERT_OK_AND_ASSIGN(auto bound, expr.Bind(schema)); - EXPECT_TRUE(bound.IsBound()); - - ASSERT_OK_AND_ASSIGN(expected, expected->Bind(schema)); - EXPECT_EQ(bound, *expected) << " unbound: " << expr.ToString(); - - if (bound_out) { - *bound_out = bound; - } -} - TEST(Expression, BindFieldRef) { // an unbound field_ref does not have the output type set auto expr = field_ref("alpha"); @@ -1104,104 +989,6 @@ TEST(Expression, BindSpecialForm) { EXPECT_TRUE(expr.IsBound()); EXPECT_TRUE(expr.type()->Equals(*int64())); } - - { - auto expr = if_else_special(field_ref("bool"), field_ref("i8"), field_ref("i8")); - EXPECT_FALSE(expr.IsBound()); - ExpectBindsTo(expr, no_change, &expr); - EXPECT_TRUE(expr.IsBound()); - EXPECT_TRUE(expr.type()->Equals(*int8())); - } - - // Implicit casts. - { - Expression bound; - ExpectBindsTo(if_else_special(field_ref("bool"), field_ref("i8"), field_ref("i32")), - if_else_special(field_ref("bool"), cast(field_ref("i8"), int32()), - field_ref("i32")), - &bound); - EXPECT_TRUE(bound.IsBound()); - EXPECT_TRUE(bound.type()->Equals(*int32())); - } - { - Expression bound; - ExpectBindsTo(if_else_special(field_ref("bool"), field_ref("i32"), field_ref("i8")), - if_else_special(field_ref("bool"), field_ref("i32"), - cast(field_ref("i8"), int32())), - &bound); - EXPECT_TRUE(bound.IsBound()); - EXPECT_TRUE(bound.type()->Equals(*int32())); - } - - // Nested call. - { - Expression bound; - ExpectBindsTo(if_else_special(equal(field_ref("i8"), field_ref("i8")), - add(field_ref("i8"), literal(42)), - add(field_ref("i32"), literal(42))), - if_else_special(equal(field_ref("i8"), field_ref("i8")), - cast(add(field_ref("i8"), literal(42)), int32()), - add(field_ref("i32"), literal(42))), - &bound); - EXPECT_TRUE(bound.IsBound()); - EXPECT_TRUE(bound.type()->Equals(*int32())); - } - { - Expression bound; - ExpectBindsTo(if_else_special(equal(field_ref("i8"), field_ref("i32")), - add(field_ref("i32"), field_ref("i8")), - add(field_ref("i32"), literal(42))), - if_else_special(equal(cast(field_ref("i8"), int32()), field_ref("i32")), - add(field_ref("i32"), cast(field_ref("i8"), int32())), - add(field_ref("i32"), literal(42))), - &bound); - EXPECT_TRUE(bound.IsBound()); - EXPECT_TRUE(bound.type()->Equals(*int32())); - } - - // Nesting call. - { - Expression bound; - ExpectBindsTo(add(if_else_special(field_ref("bool"), field_ref("i32"), literal(42)), - field_ref("i8")), - add(if_else_special(field_ref("bool"), field_ref("i32"), literal(42)), - cast(field_ref("i8"), int32())), - &bound); - EXPECT_TRUE(bound.IsBound()); - EXPECT_TRUE(bound.type()->Equals(*int32())); - } - { - Expression bound; - ExpectBindsTo( - add(if_else_special(field_ref("bool"), field_ref("i8"), literal(42)), - field_ref("i32")), - add(cast(if_else_special(field_ref("bool"), field_ref("i8"), literal(42)), - int32()), - field_ref("i32")), - &bound); - EXPECT_TRUE(bound.IsBound()); - EXPECT_TRUE(bound.type()->Equals(*int32())); - } - - // Self-nested. - { - Expression bound; - ExpectBindsTo( - if_else_special( - if_else_special(literal(true), literal(true), literal(false)), - if_else_special(field_ref("bool"), field_ref("i8"), field_ref("i32")), - if_else_special(field_ref("bool"), field_ref("i8"), field_ref("i64"))), - if_else_special( - if_else_special(literal(true), literal(true), literal(false)), - cast(if_else_special(field_ref("bool"), cast(field_ref("i8"), int32()), - field_ref("i32")), - int64()), - if_else_special(field_ref("bool"), cast(field_ref("i8"), int64()), - field_ref("i64"))), - &bound); - EXPECT_TRUE(bound.IsBound()); - EXPECT_TRUE(bound.type()->Equals(*int64())); - } } TEST(Expression, ExecuteFieldRef) { @@ -1495,16 +1282,6 @@ TEST(Expression, NaiveExecuteSpecialForm) { {"i32": 1}, {"i32": 2} ])")); - - ExpectExecute( - if_else_special(field_ref("cond"), field_ref("if_true"), field_ref("if_false")), - ArrayFromJSON(struct_({field("cond", boolean()), field("if_true", int32()), - field("if_false", int32())}), - R"([ - {"cond": null, "if_true": 0, "if_false": 1}, - {"cond": false, "if_true": 2, "if_false": 3}, - {"cond": true, "if_true": 4, "if_false": 5} - ])")); } TEST(Expression, ExecuteSpecialFormNested) { diff --git a/cpp/src/arrow/compute/expression_test_internal.h b/cpp/src/arrow/compute/expression_test_internal.h new file mode 100644 index 000000000000..25d3d6787e16 --- /dev/null +++ b/cpp/src/arrow/compute/expression_test_internal.h @@ -0,0 +1,90 @@ + +// 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 +#include + +#include "arrow/compute/expression_internal.h" +#include "arrow/testing/gtest_util.h" + +namespace arrow::compute::internal { + +const std::shared_ptr kBoringSchema = schema({ + field("bool", boolean()), + field("i8", int8()), + field("i32", int32()), + field("i32_req", int32(), /*nullable=*/false), + field("u32", uint32()), + field("i64", int64()), + field("f32", float32()), + field("f32_req", float32(), /*nullable=*/false), + field("f64", float64()), + field("date64", date64()), + field("str", utf8()), + field("dict_str", dictionary(int32(), utf8())), + field("dict_i32", dictionary(int32(), int32())), + field("ts_ns", timestamp(TimeUnit::NANO)), + field("ts_s", timestamp(TimeUnit::SECOND)), + field("binary", binary()), + field("ts_s_utc", timestamp(TimeUnit::SECOND, "UTC")), +}); + +inline Expression cast(Expression argument, std::shared_ptr to_type) { + return call("cast", {std::move(argument)}, + compute::CastOptions::Safe(std::move(to_type))); +} + +inline Expression true_unless_null(Expression argument) { + return call("true_unless_null", {std::move(argument)}); +} + +inline Expression add(Expression l, Expression r) { + return call("add", {std::move(l), std::move(r)}); +} + +inline std::string make_range_json(int start, int end) { + std::string result = "["; + for (int i = start; i <= end; ++i) { + if (i > start) result += ","; + result += std::to_string(i); + } + result += "]"; + return result; +} + +const auto no_change = std::nullopt; + +inline void ExpectBindsTo(Expression expr, std::optional expected, + Expression* bound_out = nullptr, + const Schema& schema = *kBoringSchema) { + if (!expected) { + expected = expr; + } + + ASSERT_OK_AND_ASSIGN(auto bound, expr.Bind(schema)); + EXPECT_TRUE(bound.IsBound()); + + ASSERT_OK_AND_ASSIGN(expected, expected->Bind(schema)); + EXPECT_EQ(bound, *expected) << " unbound: " << expr.ToString(); + + if (bound_out) { + *bound_out = bound; + } +} + +} // namespace arrow::compute::internal diff --git a/cpp/src/arrow/compute/special/CMakeLists.txt b/cpp/src/arrow/compute/special/CMakeLists.txt new file mode 100644 index 000000000000..1f93a1468baa --- /dev/null +++ b/cpp/src/arrow/compute/special/CMakeLists.txt @@ -0,0 +1,26 @@ +# 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. + +arrow_install_all_headers("arrow/compute/special") + +add_arrow_compute_test(if_else_special_test + SOURCES + if_else_special_test.cc + EXTRA_LINK_LIBS + arrow_compute_testing) + +add_arrow_compute_benchmark(if_else_special_benchmark) diff --git a/cpp/src/arrow/compute/special_form.cc b/cpp/src/arrow/compute/special/if_else_special.cc similarity index 98% rename from cpp/src/arrow/compute/special_form.cc rename to cpp/src/arrow/compute/special/if_else_special.cc index 8c01bf2d5384..09c93d052f2b 100644 --- a/cpp/src/arrow/compute/special_form.cc +++ b/cpp/src/arrow/compute/special/if_else_special.cc @@ -15,8 +15,6 @@ // specific language governing permissions and limitations // under the License. -#include "arrow/compute/special_form.h" - #include "arrow/array/builder_primitive.h" #include "arrow/chunk_resolver.h" #include "arrow/compute/api_vector.h" @@ -24,6 +22,7 @@ #include "arrow/compute/expression.h" #include "arrow/compute/expression_internal.h" #include "arrow/compute/registry.h" +#include "arrow/compute/special_form.h" #include "arrow/util/logging_internal.h" namespace arrow::compute { @@ -665,11 +664,18 @@ class IfElseSpecialForm : public ConditionalSpecialForm { friend class ConditionalSpecialForm; }; -} // namespace - std::shared_ptr GetIfElseSpecialForm() { static auto instance = std::make_shared(); return instance; } +} // namespace + +Expression if_else_special(Expression cond, Expression if_true, Expression if_false) { + Expression::Special special; + special.special_form = GetIfElseSpecialForm(); + special.arguments = {std::move(cond), std::move(if_true), std::move(if_false)}; + return Expression(std::move(special)); +} + } // namespace arrow::compute diff --git a/cpp/src/arrow/compute/special_form_benchmark.cc b/cpp/src/arrow/compute/special/if_else_special_benchmark.cc similarity index 100% rename from cpp/src/arrow/compute/special_form_benchmark.cc rename to cpp/src/arrow/compute/special/if_else_special_benchmark.cc diff --git a/cpp/src/arrow/compute/special_form_test.cc b/cpp/src/arrow/compute/special/if_else_special_test.cc similarity index 85% rename from cpp/src/arrow/compute/special_form_test.cc rename to cpp/src/arrow/compute/special/if_else_special_test.cc index d20b73754ea1..c9ee018f61ae 100644 --- a/cpp/src/arrow/compute/special_form_test.cc +++ b/cpp/src/arrow/compute/special/if_else_special_test.cc @@ -15,22 +15,205 @@ // specific language governing permissions and limitations // under the License. -#include "arrow/compute/special_form.h" - +#include #include +#include -#include "arrow/compute/exec.h" -#include "arrow/compute/expression.h" -#include "arrow/compute/function.h" -#include "arrow/compute/kernels/codegen_internal.h" -#include "arrow/compute/registry.h" -#include "arrow/testing/gtest_util.h" +#include "arrow/compute/api_special.h" +#include "arrow/compute/expression_test_internal.h" +#include "arrow/compute/special_form.h" #include "arrow/util/logging_internal.h" namespace arrow::compute { namespace { +using internal::add; +using internal::cast; +using internal::ExpectBindsTo; +using internal::kBoringSchema; +using internal::no_change; + +TEST(Expression, ToString) { + EXPECT_EQ( + if_else_special(field_ref("cond"), field_ref("if_true"), field_ref("if_false")) + .ToString(), + "if_else_special(cond, if_true, if_false)"); +} + +TEST(Expression, Equality) { + EXPECT_EQ(if_else_special(literal(true), field_ref("a"), field_ref("b")), + if_else_special(literal(true), field_ref("a"), field_ref("b"))); + EXPECT_NE(if_else_special(literal(true), field_ref("a"), field_ref("b")), + if_else_special(literal(false), field_ref("a"), field_ref("b"))); + EXPECT_NE(if_else_special(literal(true), field_ref("a"), field_ref("b")), + if_else_special(literal(true), field_ref("b"), field_ref("b"))); + EXPECT_NE(if_else_special(literal(true), field_ref("a"), field_ref("b")), + if_else_special(literal(true), field_ref("a"), field_ref("a"))); + EXPECT_NE(if_else_special(literal(true), field_ref("a"), field_ref("b")), + call("if_else", {literal(true), field_ref("a"), field_ref("b")})); +} + +TEST(Expression, Hash) { + std::unordered_set set; + + EXPECT_TRUE( + set.emplace(if_else_special(field_ref("cond"), field_ref("a"), field_ref("b"))) + .second); + EXPECT_FALSE( + set.emplace(if_else_special(field_ref("cond"), field_ref("a"), field_ref("b"))) + .second); + EXPECT_TRUE( + set.emplace(if_else_special(field_ref("cond"), field_ref("b"), field_ref("a"))) + .second); + + EXPECT_EQ(set.size(), 2); +} + +TEST(Expression, IsScalarExpression) { + EXPECT_TRUE(if_else_special(field_ref("cond"), field_ref("a"), field_ref("b")) + .IsScalarExpression()); +} + +TEST(Expression, IsSatisfiable) { + auto Bind = [](Expression expr) { return expr.Bind(*kBoringSchema).ValueOrDie(); }; + + EXPECT_TRUE(Bind(if_else_special(field_ref("bool"), field_ref("i32"), field_ref("i32"))) + .IsSatisfiable()); +} + +TEST(Expression, FieldsInExpression) { + auto ExpectFieldsAre = [](Expression expr, std::vector expected) { + EXPECT_THAT(FieldsInExpression(expr), testing::ContainerEq(expected)); + }; + + ExpectFieldsAre(if_else_special(literal(true), literal(1), literal(0)), {}); + ExpectFieldsAre(if_else_special(literal(true), field_ref("a"), field_ref("b")), + {"a", "b"}); + ExpectFieldsAre(if_else_special(field_ref("a"), field_ref("b"), field_ref("b")), + {"a", "b", "b"}); + ExpectFieldsAre(if_else_special(field_ref("a"), field_ref("b"), field_ref("c")), + {"a", "b", "c"}); + ExpectFieldsAre( + if_else_special(call("not", {field_ref("a")}), call("not", {field_ref("b")}), + call("not", {field_ref("c")})), + {"a", "b", "c"}); + ExpectFieldsAre( + call("not", {if_else_special(field_ref("a"), field_ref("b"), field_ref("c"))}), + {"a", "b", "c"}); +} + +TEST(Expression, ExpressionHasFieldRefs) { + EXPECT_FALSE( + ExpressionHasFieldRefs(if_else_special(literal(true), literal(1), literal(0)))); + EXPECT_TRUE( + ExpressionHasFieldRefs(if_else_special(field_ref("a"), literal(1), literal(0)))); + EXPECT_TRUE( + ExpressionHasFieldRefs(if_else_special(literal(true), field_ref("a"), literal(0)))); + EXPECT_TRUE( + ExpressionHasFieldRefs(if_else_special(literal(true), literal(0), field_ref("a")))); +} + +TEST(Expression, BindSpecialForm) { + { + auto expr = if_else_special(field_ref("bool"), field_ref("i8"), field_ref("i8")); + EXPECT_FALSE(expr.IsBound()); + ExpectBindsTo(expr, no_change, &expr); + EXPECT_TRUE(expr.IsBound()); + EXPECT_TRUE(expr.type()->Equals(*int8())); + } + + // Implicit casts. + { + Expression bound; + ExpectBindsTo(if_else_special(field_ref("bool"), field_ref("i8"), field_ref("i32")), + if_else_special(field_ref("bool"), cast(field_ref("i8"), int32()), + field_ref("i32")), + &bound); + EXPECT_TRUE(bound.IsBound()); + EXPECT_TRUE(bound.type()->Equals(*int32())); + } + { + Expression bound; + ExpectBindsTo(if_else_special(field_ref("bool"), field_ref("i32"), field_ref("i8")), + if_else_special(field_ref("bool"), field_ref("i32"), + cast(field_ref("i8"), int32())), + &bound); + EXPECT_TRUE(bound.IsBound()); + EXPECT_TRUE(bound.type()->Equals(*int32())); + } + + // Nested call. + { + Expression bound; + ExpectBindsTo(if_else_special(equal(field_ref("i8"), field_ref("i8")), + add(field_ref("i8"), literal(42)), + add(field_ref("i32"), literal(42))), + if_else_special(equal(field_ref("i8"), field_ref("i8")), + cast(add(field_ref("i8"), literal(42)), int32()), + add(field_ref("i32"), literal(42))), + &bound); + EXPECT_TRUE(bound.IsBound()); + EXPECT_TRUE(bound.type()->Equals(*int32())); + } + { + Expression bound; + ExpectBindsTo(if_else_special(equal(field_ref("i8"), field_ref("i32")), + add(field_ref("i32"), field_ref("i8")), + add(field_ref("i32"), literal(42))), + if_else_special(equal(cast(field_ref("i8"), int32()), field_ref("i32")), + add(field_ref("i32"), cast(field_ref("i8"), int32())), + add(field_ref("i32"), literal(42))), + &bound); + EXPECT_TRUE(bound.IsBound()); + EXPECT_TRUE(bound.type()->Equals(*int32())); + } + + // Nesting call. + { + Expression bound; + ExpectBindsTo(add(if_else_special(field_ref("bool"), field_ref("i32"), literal(42)), + field_ref("i8")), + add(if_else_special(field_ref("bool"), field_ref("i32"), literal(42)), + cast(field_ref("i8"), int32())), + &bound); + EXPECT_TRUE(bound.IsBound()); + EXPECT_TRUE(bound.type()->Equals(*int32())); + } + { + Expression bound; + ExpectBindsTo( + add(if_else_special(field_ref("bool"), field_ref("i8"), literal(42)), + field_ref("i32")), + add(cast(if_else_special(field_ref("bool"), field_ref("i8"), literal(42)), + int32()), + field_ref("i32")), + &bound); + EXPECT_TRUE(bound.IsBound()); + EXPECT_TRUE(bound.type()->Equals(*int32())); + } + + // Self-nested. + { + Expression bound; + ExpectBindsTo( + if_else_special( + if_else_special(literal(true), literal(true), literal(false)), + if_else_special(field_ref("bool"), field_ref("i8"), field_ref("i32")), + if_else_special(field_ref("bool"), field_ref("i8"), field_ref("i64"))), + if_else_special( + if_else_special(literal(true), literal(true), literal(false)), + cast(if_else_special(field_ref("bool"), cast(field_ref("i8"), int32()), + field_ref("i32")), + int64()), + if_else_special(field_ref("bool"), cast(field_ref("i8"), int64()), + field_ref("i64"))), + &bound); + EXPECT_TRUE(bound.IsBound()); + EXPECT_TRUE(bound.type()->Equals(*int64())); + } +} + void AssertEqualIgnoreShape(const Datum& expected, const Datum& result) { if (expected.kind() == result.kind()) { AssertDatumsEqual(expected, result); diff --git a/cpp/src/arrow/compute/special/meson.build b/cpp/src/arrow/compute/special/meson.build new file mode 100644 index 000000000000..df656672c191 --- /dev/null +++ b/cpp/src/arrow/compute/special/meson.build @@ -0,0 +1,26 @@ +# 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. + +# Contains special form implementation. + +if needs_compute + exc = executable( + 'arrow-compute-if-else-special-benchmark', + sources: ['if_else_special_benchmark.cc'], + dependencies: [arrow_compute_dep, arrow_benchmark_dep], + ) +endif diff --git a/cpp/src/arrow/compute/special_form.h b/cpp/src/arrow/compute/special_form.h index 786484ad5263..5588e6a73f2e 100644 --- a/cpp/src/arrow/compute/special_form.h +++ b/cpp/src/arrow/compute/special_form.h @@ -54,6 +54,4 @@ class ARROW_EXPORT SpecialForm { std::string name_; }; -std::shared_ptr GetIfElseSpecialForm(); - } // namespace arrow::compute From 496c2ed30e5912bdf95b2db8fbdb7c772443042d Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Sat, 4 Oct 2025 03:14:06 -0700 Subject: [PATCH 34/71] Make header for conditional stuff so we can test it (#59) --- cpp/src/arrow/CMakeLists.txt | 1 + cpp/src/arrow/compute/expression_internal.h | 2 + cpp/src/arrow/compute/special/CMakeLists.txt | 6 + cpp/src/arrow/compute/special/conditional.cc | 287 ++++++++ .../compute/special/conditional_internal.h | 376 +++++++++++ .../arrow/compute/special/conditional_test.cc | 31 + .../arrow/compute/special/if_else_special.cc | 623 +----------------- .../compute/special/if_else_special_test.cc | 16 +- .../compute/special/special_form_internal.h | 54 ++ cpp/src/arrow/compute/special_form.h | 2 + 10 files changed, 770 insertions(+), 628 deletions(-) create mode 100644 cpp/src/arrow/compute/special/conditional.cc create mode 100644 cpp/src/arrow/compute/special/conditional_internal.h create mode 100644 cpp/src/arrow/compute/special/conditional_test.cc create mode 100644 cpp/src/arrow/compute/special/special_form_internal.h diff --git a/cpp/src/arrow/CMakeLists.txt b/cpp/src/arrow/CMakeLists.txt index 5e49f3750ad0..4042f4e9f8ff 100644 --- a/cpp/src/arrow/CMakeLists.txt +++ b/cpp/src/arrow/CMakeLists.txt @@ -802,6 +802,7 @@ if(ARROW_COMPUTE) compute/row/grouper.cc compute/row/row_encoder_internal.cc compute/row/row_internal.cc + compute/special/conditional.cc compute/special/if_else_special.cc compute/util.cc compute/util_internal.cc) diff --git a/cpp/src/arrow/compute/expression_internal.h b/cpp/src/arrow/compute/expression_internal.h index c7982b30b262..794dcbe76cb3 100644 --- a/cpp/src/arrow/compute/expression_internal.h +++ b/cpp/src/arrow/compute/expression_internal.h @@ -15,6 +15,8 @@ // specific language governing permissions and limitations // under the License. +#pragma once + #include "arrow/compute/expression.h" #include diff --git a/cpp/src/arrow/compute/special/CMakeLists.txt b/cpp/src/arrow/compute/special/CMakeLists.txt index 1f93a1468baa..1ba6d11819b4 100644 --- a/cpp/src/arrow/compute/special/CMakeLists.txt +++ b/cpp/src/arrow/compute/special/CMakeLists.txt @@ -23,4 +23,10 @@ add_arrow_compute_test(if_else_special_test EXTRA_LINK_LIBS arrow_compute_testing) +add_arrow_compute_test(special_internal_test + SOURCES + conditional_test.cc + EXTRA_LINK_LIBS + arrow_compute_testing) + add_arrow_compute_benchmark(if_else_special_benchmark) diff --git a/cpp/src/arrow/compute/special/conditional.cc b/cpp/src/arrow/compute/special/conditional.cc new file mode 100644 index 000000000000..bebafcadae40 --- /dev/null +++ b/cpp/src/arrow/compute/special/conditional.cc @@ -0,0 +1,287 @@ +// 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/special/conditional_internal.h" + +#include "arrow/array/array_primitive.h" +#include "arrow/array/builder_primitive.h" +#include "arrow/visit_data_inline.h" + +namespace arrow::compute::internal { + +Result> BranchMask::MakeBodyMaskFromDatum( + const Datum& datum, ExecContext* exec_context) const { + DCHECK(datum.type()->id() == Type::BOOL); + if (datum.is_scalar()) { + auto scalar = datum.scalar_as(); + if (!scalar.is_valid) { + return std::make_shared(shared_from_this()); + } else if (scalar.value) { + return std::make_shared(shared_from_this()); + } else { + return std::make_shared(shared_from_this()); + } + } + if (datum.is_array()) { + return MakeBodyMaskFromBitmap(datum.array_as(), exec_context); + } + DCHECK(datum.is_chunked_array()); + return MakeBodyMaskFromBitmap(datum.chunked_array(), exec_context); +} + +Result AllPassBranchMask::DoApplyCond(const Expression& expr, + const ExecBatch& input, + ExecContext* exec_context) const { + DCHECK_EQ(input.length, length_); + auto input_with_sel_vec = input; + input_with_sel_vec.selection_vector = nullptr; + return ExecuteScalarExpression(expr, input_with_sel_vec, exec_context); +} + +Result> AllPassBranchMask::GetSelectionVector() const { + return nullptr; +} + +Result> AllPassBranchMask::MakeBodyMaskFromBitmap( + const std::shared_ptr& bitmap, ExecContext* exec_context) const { + DCHECK_EQ(bitmap->length(), length_); + + Int32Builder body_builder(exec_context->memory_pool()); + Int32Builder rest_builder(exec_context->memory_pool()); + RETURN_NOT_OK(body_builder.Reserve(length_)); + RETURN_NOT_OK(rest_builder.Reserve(length_)); + + ArraySpan span(*bitmap->data()); + int32_t i = 0; + VisitArraySpanInline( + span, + [&](bool mask) { + if (mask) { + body_builder.UnsafeAppend(i); + } else { + rest_builder.UnsafeAppend(i); + } + ++i; + }, + [&]() { ++i; }); + + ARROW_ASSIGN_OR_RAISE(auto body_arr, body_builder.Finish()); + ARROW_ASSIGN_OR_RAISE(auto rest_arr, rest_builder.Finish()); + auto body = std::make_shared(body_arr->data()); + auto rest = std::make_shared(rest_arr->data()); + return std::make_shared(std::move(body), std::move(rest), length_); +} + +Result> AllPassBranchMask::MakeBodyMaskFromBitmap( + const std::shared_ptr& bitmap, ExecContext* exec_context) const { + DCHECK_EQ(bitmap->length(), length_); + + Int32Builder body_builder(exec_context->memory_pool()); + Int32Builder rest_builder(exec_context->memory_pool()); + RETURN_NOT_OK(body_builder.Reserve(length_)); + RETURN_NOT_OK(rest_builder.Reserve(length_)); + + int32_t i = 0; + for (const auto& chunk : bitmap->chunks()) { + DCHECK_EQ(chunk->type()->id(), Type::BOOL); + ArraySpan span(*chunk->data()); + VisitArraySpanInline( + span, + [&](bool mask) { + if (mask) { + body_builder.UnsafeAppend(i); + } else { + rest_builder.UnsafeAppend(i); + } + ++i; + }, + [&]() { ++i; }); + } + + ARROW_ASSIGN_OR_RAISE(auto body_arr, body_builder.Finish()); + ARROW_ASSIGN_OR_RAISE(auto rest_arr, rest_builder.Finish()); + auto body = std::make_shared(body_arr->data()); + auto rest = std::make_shared(rest_arr->data()); + return std::make_shared(std::move(body), std::move(rest), length_); +} + +Result ConditionalBranchMask::DoApplyCond(const Expression& expr, + const ExecBatch& input, + ExecContext* exec_context) const { + auto sparse_input = input; + sparse_input.selection_vector = selection_vector_; + return ExecuteScalarExpression(expr, sparse_input, exec_context); +} + +Result> ConditionalBranchMask::MakeBodyMaskFromBitmap( + const std::shared_ptr& bitmap, ExecContext* exec_context) const { + DCHECK_EQ(bitmap->length(), length_); + + Int32Builder body_builder(exec_context->memory_pool()); + Int32Builder rest_builder(exec_context->memory_pool()); + RETURN_NOT_OK(body_builder.Reserve(length_)); + RETURN_NOT_OK(rest_builder.Reserve(length_)); + + for (int64_t i = 0; i < selection_vector_->length(); ++i) { + auto index = selection_vector_->indices()[i]; + if (!bitmap->IsNull(index)) { + if (bitmap->Value(index)) { + body_builder.UnsafeAppend(index); + } else { + rest_builder.UnsafeAppend(index); + } + } + } + + ARROW_ASSIGN_OR_RAISE(auto body_arr, body_builder.Finish()); + ARROW_ASSIGN_OR_RAISE(auto rest_arr, rest_builder.Finish()); + auto body = std::make_shared(body_arr->data()); + auto rest = std::make_shared(rest_arr->data()); + return std::make_shared(std::move(body), std::move(rest), length_); +} + +Result> ConditionalBranchMask::MakeBodyMaskFromBitmap( + const std::shared_ptr& bitmap, ExecContext* exec_context) const { + DCHECK_EQ(bitmap->length(), length_); + + std::vector boolean_arrays(bitmap->num_chunks()); + std::transform(bitmap->chunks().begin(), bitmap->chunks().end(), boolean_arrays.begin(), + [](const auto& chunk) { + DCHECK_EQ(chunk->type()->id(), Type::BOOL); + return checked_cast(chunk.get()); + }); + + Int32Builder body_builder(exec_context->memory_pool()); + Int32Builder rest_builder(exec_context->memory_pool()); + RETURN_NOT_OK(body_builder.Reserve(length_)); + RETURN_NOT_OK(rest_builder.Reserve(length_)); + + ChunkResolver resolver(bitmap->chunks()); + ChunkLocation location; + for (int64_t i = 0; i < selection_vector_->length(); ++i) { + auto index = selection_vector_->indices()[i]; + location = resolver.ResolveWithHint(index, location); + if (boolean_arrays[location.chunk_index]->IsValid(location.index_in_chunk)) { + if (boolean_arrays[location.chunk_index]->Value(location.index_in_chunk)) { + body_builder.UnsafeAppend(index); + } else { + rest_builder.UnsafeAppend(index); + } + } + } + + ARROW_ASSIGN_OR_RAISE(auto body_arr, body_builder.Finish()); + ARROW_ASSIGN_OR_RAISE(auto rest_arr, rest_builder.Finish()); + auto body = std::make_shared(body_arr->data()); + auto rest = std::make_shared(rest_arr->data()); + return std::make_shared(std::move(body), std::move(rest), length_); +} + +Result ConditionalBodyMask::ApplyCond(const Expression& expr, + const ExecBatch& input, + ExecContext* exec_context) const { + auto sparse_input = input; + sparse_input.selection_vector = body_; + return ExecuteScalarExpression(expr, sparse_input, exec_context); +} + +Result ConditionalExec::Execute(const ExecBatch& input, + ExecContext* exec_context) const&& { + DCHECK(!branches.empty()); + + BranchResults results; + results.Reserve(branches.size()); + ARROW_ASSIGN_OR_RAISE(auto branch_mask, InitBranchMask(input, exec_context)); + for (const auto& branch : branches) { + if (branch_mask->empty()) { + break; + } + ARROW_ASSIGN_OR_RAISE(auto body_mask, + ApplyCond(branch_mask, branch.cond, input, exec_context)); + if (body_mask->empty()) { + ARROW_ASSIGN_OR_RAISE(branch_mask, body_mask->NextBranchMask()); + continue; + } + ARROW_ASSIGN_OR_RAISE(auto body_result, + ApplyBody(body_mask, branch.body, input, exec_context)); + DCHECK(body_result.type()->Equals(*result_type)); + ARROW_ASSIGN_OR_RAISE(auto selection_vector, body_mask->GetSelectionVector()); + results.Emplace(std::move(body_result), std::move(selection_vector)); + ARROW_ASSIGN_OR_RAISE(branch_mask, body_mask->NextBranchMask()); + } + return MultiplexResults(input, results, exec_context); +} + +Result ConditionalExec::MultiplexResults(const ExecBatch& input, + const BranchResults& results, + ExecContext* exec_context) const { + if (results.empty()) { + return MakeArrayOfNull(result_type.GetSharedPtr(), input.length, + exec_context->memory_pool()); + } + + if (results.size() == 1) { + if (const auto& result = results.body_results()[0]; + results.selection_vectors()[0] == nullptr || + results.selection_vectors()[0]->length() == + (input.selection_vector ? input.selection_vector->length() : input.length)) { + return result; + } + } + + std::vector choose_args; + choose_args.reserve(results.size() + 1); + ARROW_ASSIGN_OR_RAISE(auto indices, ChooseIndices(results.selection_vectors(), + input.length, exec_context)); + choose_args.emplace_back(std::move(indices)); + choose_args.insert(choose_args.end(), results.body_results().begin(), + results.body_results().end()); + return CallFunction("choose", choose_args, exec_context); +} + +Result ConditionalExec::ChooseIndices( + const std::vector>& selection_vectors, + int64_t length, ExecContext* exec_context) const { + const int64_t validity_bytes = bit_util::BytesForBits(length); + ARROW_ASSIGN_OR_RAISE( + std::shared_ptr validity_buf, + AllocateResizableBuffer(validity_bytes, exec_context->memory_pool())); + auto validity_data = validity_buf->mutable_data_as(); + std::memset(validity_data, 0, validity_bytes); + + ARROW_ASSIGN_OR_RAISE( + std::shared_ptr indices_buf, + AllocateResizableBuffer(length * sizeof(int32_t), exec_context->memory_pool())); + auto indices_data = indices_buf->mutable_data_as(); + for (int32_t index = 0; index < static_cast(selection_vectors.size()); + ++index) { + DCHECK_NE(selection_vectors[index], nullptr); + DCHECK_GT(selection_vectors[index]->length(), 0); + auto row_ids = selection_vectors[index]->indices(); + for (int64_t i = 0; i < selection_vectors[index]->length(); ++i) { + const int32_t row_id = row_ids[i]; + DCHECK_EQ(bit_util::GetBit(validity_data, row_id), false); + bit_util::SetBitTo(validity_data, row_id, true); + indices_data[row_id] = index; + } + } + + return ArrayData::Make(int32(), length, + {std::move(validity_buf), std::move(indices_buf)}); +} + +} // namespace arrow::compute::internal diff --git a/cpp/src/arrow/compute/special/conditional_internal.h b/cpp/src/arrow/compute/special/conditional_internal.h new file mode 100644 index 000000000000..6e01d163ff2d --- /dev/null +++ b/cpp/src/arrow/compute/special/conditional_internal.h @@ -0,0 +1,376 @@ +// 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. + +#pragma once + +#include +#include + +#include "arrow/compute/exec.h" +#include "arrow/compute/special/special_form_internal.h" +#include "arrow/compute/visibility.h" + +namespace arrow::compute::internal { + +struct ARROW_COMPUTE_EXPORT BodyMask; + +struct ARROW_COMPUTE_EXPORT BranchMask : public std::enable_shared_from_this { + virtual ~BranchMask() = default; + + Result> ApplyCond(const Expression& expr, + const ExecBatch& input, + ExecContext* exec_context) const { + ARROW_ASSIGN_OR_RAISE(auto datum, DoApplyCond(expr, input, exec_context)); + return MakeBodyMaskFromDatum(datum, exec_context); + } + + virtual bool empty() const = 0; + + protected: + virtual Result DoApplyCond(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const = 0; + + virtual Result> GetSelectionVector() const = 0; + + virtual Result> MakeBodyMaskFromBitmap( + const std::shared_ptr& bitmap, ExecContext* exec_context) const = 0; + + virtual Result> MakeBodyMaskFromBitmap( + const std::shared_ptr& bitmap, ExecContext* exec_context) const = 0; + + private: + Result> MakeBodyMaskFromDatum( + const Datum& datum, ExecContext* exec_context) const; + + friend struct DelegateBodyMask; +}; + +struct ARROW_COMPUTE_EXPORT AllPassBranchMask : public BranchMask { + explicit AllPassBranchMask(int64_t length) : length_(length) {} + + bool empty() const override { return false; } + + protected: + Result DoApplyCond(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const override; + + Result> GetSelectionVector() const override; + + Result> MakeBodyMaskFromBitmap( + const std::shared_ptr& bitmap, + ExecContext* exec_context) const override; + + Result> MakeBodyMaskFromBitmap( + const std::shared_ptr& bitmap, + ExecContext* exec_context) const override; + + private: + int64_t length_; +}; + +struct ARROW_COMPUTE_EXPORT AllFailBranchMask : public BranchMask { + AllFailBranchMask() = default; + + bool empty() const override { return true; } + + protected: + Result DoApplyCond(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const override { + DCHECK(false); + return Status::Invalid("AllFailBranchMask::DoApplyCond should not be called"); + } + + Result> GetSelectionVector() const override { + DCHECK(false); + return Status::Invalid("AllFailBranchMask::GetSelectionVector should not be called"); + } + + Result> MakeBodyMaskFromBitmap( + const std::shared_ptr& bitmap, + ExecContext* exec_context) const override { + DCHECK(false); + return Status::Invalid( + "AllFailBranchMask::MakeBodyMaskFromBitmap should not be called"); + } + + Result> MakeBodyMaskFromBitmap( + const std::shared_ptr& bitmap, + ExecContext* exec_context) const override { + DCHECK(false); + return Status::Invalid( + "AllFailBranchMask::MakeBodyMaskFromBitmap should not be called"); + } +}; + +struct ARROW_COMPUTE_EXPORT ConditionalBranchMask : public BranchMask { + ConditionalBranchMask(std::shared_ptr selection_vector, int64_t length) + : selection_vector_(std::move(selection_vector)), length_(length) {} + + bool empty() const override { return selection_vector_->length() == 0; } + + protected: + Result DoApplyCond(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const override; + + Result> GetSelectionVector() const override { + return selection_vector_; + } + + Result> MakeBodyMaskFromBitmap( + const std::shared_ptr& bitmap, + ExecContext* exec_context) const override; + + Result> MakeBodyMaskFromBitmap( + const std::shared_ptr& bitmap, + ExecContext* exec_context) const override; + + protected: + std::shared_ptr selection_vector_ = nullptr; + int64_t length_ = 0; +}; + +struct ARROW_COMPUTE_EXPORT BodyMask : public std::enable_shared_from_this { + virtual ~BodyMask() = default; + + virtual bool empty() const = 0; + + virtual Result ApplyCond(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const = 0; + + virtual Result> GetSelectionVector() const = 0; + + virtual Result> NextBranchMask() const = 0; +}; + +struct ARROW_COMPUTE_EXPORT DelegateBodyMask : public BodyMask { + explicit DelegateBodyMask(std::shared_ptr branch_mask) + : branch_mask_(std::move(branch_mask)) {} + + protected: + Result DelegateApplyCond(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const { + return branch_mask_->DoApplyCond(expr, input, exec_context); + } + + Result> DelegateGetSelectionVector() const { + return branch_mask_->GetSelectionVector(); + } + + protected: + std::shared_ptr branch_mask_; +}; + +struct ARROW_COMPUTE_EXPORT AllNullBodyMask : public DelegateBodyMask { + using DelegateBodyMask::DelegateBodyMask; + + bool empty() const override { return true; } + + Result ApplyCond(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const override { + DCHECK(false); + return Status::Invalid("AllNullBodyMask::ApplyCond should not be called"); + } + + Result> GetSelectionVector() const override { + DCHECK(false); + return Status::Invalid("AllNullBodyMask::GetSelectionVector should not be called"); + } + + Result> NextBranchMask() const override { + return std::make_shared(); + } +}; + +struct ARROW_COMPUTE_EXPORT AllPassBodyMask : public DelegateBodyMask { + using DelegateBodyMask::DelegateBodyMask; + + bool empty() const override { return false; } + + Result ApplyCond(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const override { + return DelegateApplyCond(expr, input, exec_context); + } + + Result> GetSelectionVector() const override { + return DelegateGetSelectionVector(); + } + + Result> NextBranchMask() const override { + return std::make_shared(); + } +}; + +struct ARROW_COMPUTE_EXPORT AllFailBodyMask : public DelegateBodyMask { + using DelegateBodyMask::DelegateBodyMask; + + bool empty() const override { return true; } + + Result ApplyCond(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const override { + DCHECK(false); + return Status::Invalid("AllFailBodyMask::ApplyCond should not be called"); + } + + Result> GetSelectionVector() const override { + DCHECK(false); + return Status::Invalid("AllFailBodyMask::GetSelectionVector should not be called"); + } + + Result> NextBranchMask() const override { + return branch_mask_; + } +}; + +struct ARROW_COMPUTE_EXPORT Branch { + Expression cond; + Expression body; +}; + +struct ARROW_COMPUTE_EXPORT ConditionalBodyMask : public BodyMask { + ConditionalBodyMask(std::shared_ptr body, + std::shared_ptr rest, int64_t length) + : body_(std::move(body)), rest_(std::move(rest)), length_(length) {} + + bool empty() const override { return body_->length() == 0; } + + Result ApplyCond(const Expression& expr, const ExecBatch& input, + ExecContext* exec_context) const override; + + Result> GetSelectionVector() const override { + return body_; + } + + Result> NextBranchMask() const override { + if (!rest_) { + return std::make_shared(); + } + return std::make_shared(rest_, length_); + } + + private: + std::shared_ptr body_; + std::shared_ptr rest_; + int64_t length_; +}; + +struct ARROW_COMPUTE_EXPORT ConditionalExec { + ConditionalExec(const std::vector& branches, const TypeHolder& result_type) + : branches(branches), result_type(result_type) {} + + Result Execute(const ExecBatch& input, ExecContext* exec_context) const&&; + + private: + struct BranchResults { + void Reserve(int64_t size) { + body_results_.reserve(size); + selection_vectors_.reserve(size); + } + + void Emplace(Datum body_result, std::shared_ptr selection_vector) { + body_results_.emplace_back(std::move(body_result)); + selection_vectors_.emplace_back(std::move(selection_vector)); + } + + bool empty() const { return body_results_.empty(); } + + size_t size() const { return body_results_.size(); } + + const std::vector& body_results() const { return body_results_; } + + const std::vector>& selection_vectors() const { + return selection_vectors_; + } + + private: + std::vector body_results_; + std::vector> selection_vectors_; + }; + + Result> InitBranchMask( + const ExecBatch& input, ExecContext* exec_context) const { + if (input.selection_vector) { + return std::make_shared(input.selection_vector, + input.length); + } + return std::make_shared(input.length); + } + + Result> ApplyCond( + const std::shared_ptr& branch_mask, const Expression& cond, + const ExecBatch& input, ExecContext* exec_context) const { + return branch_mask->ApplyCond(cond, input, exec_context); + } + + Result ApplyBody(const std::shared_ptr& body_mask, + const Expression& body, const ExecBatch& input, + ExecContext* exec_context) const { + return body_mask->ApplyCond(body, input, exec_context); + } + + Result MultiplexResults(const ExecBatch& input, const BranchResults& results, + ExecContext* exec_context) const; + + Result ChooseIndices( + const std::vector>& selection_vectors, + int64_t length, ExecContext* exec_context) const; + + protected: + const std::vector& branches; + const TypeHolder& result_type; +}; + +class ARROW_COMPUTE_EXPORT ConditionalSpecialExecutor : public SpecialExecutor { + public: + ConditionalSpecialExecutor(std::vector branches, TypeHolder out_type) + : SpecialExecutor(std::move(out_type), /*options=*/nullptr), + branches(std::move(branches)) {} + + Result Execute(const ExecBatch& input, + ExecContext* exec_context) const override { + return ConditionalExec(branches, out_type_).Execute(input, exec_context); + } + + private: + std::vector branches; +}; + +template +class ConditionalSpecialForm + : public FunctionBackedSpecialForm> { + public: + using FunctionBackedSpecialForm< + ConditionalSpecialForm>::FunctionBackedSpecialForm; + + ARROW_DISALLOW_COPY_AND_ASSIGN(ConditionalSpecialForm); + ARROW_DEFAULT_MOVE_AND_ASSIGN(ConditionalSpecialForm); + + protected: + Result> BindWithBoundCall( + Expression::Call call, ExecContext* exec_context) const { + // Shouldn't have options. This is guaranteed by the call binding. + DCHECK_EQ(call.options, nullptr); + + auto branches = + static_cast(this)->GetBranches(std::move(call.arguments)); + return std::make_unique(std::move(branches), + std::move(call.type)); + } + + friend class FunctionBackedSpecialForm>; +}; + +} // namespace arrow::compute::internal diff --git a/cpp/src/arrow/compute/special/conditional_test.cc b/cpp/src/arrow/compute/special/conditional_test.cc new file mode 100644 index 000000000000..f1d1de04aa71 --- /dev/null +++ b/cpp/src/arrow/compute/special/conditional_test.cc @@ -0,0 +1,31 @@ +// 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 + +#include "arrow/compute/special/conditional_internal.h" + +namespace arrow::compute::internal { + +TEST(AllPassBranchMask, Emptiness) { + for (auto length : {0, 42}) { + auto branch_mask = std::make_shared(length); + EXPECT_FALSE(branch_mask->empty()); + } +} + +} // namespace arrow::compute::internal diff --git a/cpp/src/arrow/compute/special/if_else_special.cc b/cpp/src/arrow/compute/special/if_else_special.cc index 09c93d052f2b..5bdf3fe17a18 100644 --- a/cpp/src/arrow/compute/special/if_else_special.cc +++ b/cpp/src/arrow/compute/special/if_else_special.cc @@ -15,631 +15,14 @@ // specific language governing permissions and limitations // under the License. -#include "arrow/array/builder_primitive.h" -#include "arrow/chunk_resolver.h" -#include "arrow/compute/api_vector.h" -#include "arrow/compute/exec.h" -#include "arrow/compute/expression.h" -#include "arrow/compute/expression_internal.h" -#include "arrow/compute/registry.h" -#include "arrow/compute/special_form.h" -#include "arrow/util/logging_internal.h" +#include "arrow/compute/special/conditional_internal.h" namespace arrow::compute { namespace { -struct BodyMask; - -struct BranchMask : public std::enable_shared_from_this { - virtual ~BranchMask() = default; - - Result> ApplyCond(const Expression& expr, - const ExecBatch& input, - ExecContext* exec_context) const { - ARROW_ASSIGN_OR_RAISE(auto datum, DoApplyCond(expr, input, exec_context)); - return MakeBodyMaskFromDatum(datum, exec_context); - } - - virtual bool empty() const = 0; - - protected: - virtual Result DoApplyCond(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const = 0; - - virtual Result> GetSelectionVector() const = 0; - - virtual Result> MakeBodyMaskFromBitmap( - const std::shared_ptr& bitmap, ExecContext* exec_context) const = 0; - - virtual Result> MakeBodyMaskFromBitmap( - const std::shared_ptr& bitmap, ExecContext* exec_context) const = 0; - - private: - Result> MakeBodyMaskFromDatum( - const Datum& datum, ExecContext* exec_context) const; - - friend struct DelegateBodyMask; -}; - -struct BodyMask : public std::enable_shared_from_this { - virtual ~BodyMask() = default; - - virtual bool empty() const = 0; - - virtual Result ApplyCond(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const = 0; - - virtual Result> GetSelectionVector() const = 0; - - virtual Result> NextBranchMask() const = 0; -}; - -struct AllPassBranchMask : public BranchMask { - explicit AllPassBranchMask(int64_t length) : length_(length) {} - - bool empty() const override { return false; } - - protected: - Result DoApplyCond(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const override { - DCHECK_EQ(input.length, length_); - auto input_with_sel_vec = input; - input_with_sel_vec.selection_vector = nullptr; - return ExecuteScalarExpression(expr, input_with_sel_vec, exec_context); - } - - Result> GetSelectionVector() const override { - return nullptr; - } - - Result> MakeBodyMaskFromBitmap( - const std::shared_ptr& bitmap, - ExecContext* exec_context) const override; - - Result> MakeBodyMaskFromBitmap( - const std::shared_ptr& bitmap, - ExecContext* exec_context) const override; - - private: - int64_t length_; -}; - -struct AllFailBranchMask : public BranchMask { - AllFailBranchMask() = default; - - bool empty() const override { return true; } - - protected: - Result DoApplyCond(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const override { - DCHECK(false); - return Status::Invalid("AllFailBranchMask::DoApplyCond should not be called"); - } - - Result> GetSelectionVector() const override { - DCHECK(false); - return Status::Invalid("AllFailBranchMask::GetSelectionVector should not be called"); - } - - Result> MakeBodyMaskFromBitmap( - const std::shared_ptr& bitmap, - ExecContext* exec_context) const override { - DCHECK(false); - return Status::Invalid( - "AllFailBranchMask::MakeBodyMaskFromBitmap should not be called"); - } - - Result> MakeBodyMaskFromBitmap( - const std::shared_ptr& bitmap, - ExecContext* exec_context) const override { - DCHECK(false); - return Status::Invalid( - "AllFailBranchMask::MakeBodyMaskFromBitmap should not be called"); - } -}; - -struct DelegateBodyMask : public BodyMask { - explicit DelegateBodyMask(std::shared_ptr branch_mask) - : branch_mask_(std::move(branch_mask)) {} - - protected: - Result DelegateApplyCond(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const { - return branch_mask_->DoApplyCond(expr, input, exec_context); - } - - Result> DelegateGetSelectionVector() const { - return branch_mask_->GetSelectionVector(); - } - - protected: - std::shared_ptr branch_mask_; -}; - -struct AllNullBodyMask : public DelegateBodyMask { - using DelegateBodyMask::DelegateBodyMask; - - bool empty() const override { return true; } - - Result ApplyCond(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const override { - DCHECK(false); - return Status::Invalid("AllNullBodyMask::ApplyCond should not be called"); - } - - Result> GetSelectionVector() const override { - DCHECK(false); - return Status::Invalid("AllNullBodyMask::GetSelectionVector should not be called"); - } - - Result> NextBranchMask() const override { - return std::make_shared(); - } -}; - -struct AllPassBodyMask : public DelegateBodyMask { - using DelegateBodyMask::DelegateBodyMask; - - bool empty() const override { return false; } - - Result ApplyCond(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const override { - return DelegateApplyCond(expr, input, exec_context); - } - - Result> GetSelectionVector() const override { - return DelegateGetSelectionVector(); - } - - Result> NextBranchMask() const override { - return std::make_shared(); - } -}; - -struct AllFailBodyMask : public DelegateBodyMask { - using DelegateBodyMask::DelegateBodyMask; - - bool empty() const override { return true; } - - Result ApplyCond(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const override { - DCHECK(false); - return Status::Invalid("AllFailBodyMask::ApplyCond should not be called"); - } - - Result> GetSelectionVector() const override { - DCHECK(false); - return Status::Invalid("AllFailBodyMask::GetSelectionVector should not be called"); - } - - Result> NextBranchMask() const override { - return branch_mask_; - } -}; - -Result> BranchMask::MakeBodyMaskFromDatum( - const Datum& datum, ExecContext* exec_context) const { - DCHECK(datum.type()->id() == Type::BOOL); - if (datum.is_scalar()) { - auto scalar = datum.scalar_as(); - if (!scalar.is_valid) { - return std::make_shared(shared_from_this()); - } else if (scalar.value) { - return std::make_shared(shared_from_this()); - } else { - return std::make_shared(shared_from_this()); - } - } - if (datum.is_array()) { - return MakeBodyMaskFromBitmap(datum.array_as(), exec_context); - } - DCHECK(datum.is_chunked_array()); - return MakeBodyMaskFromBitmap(datum.chunked_array(), exec_context); -} - -struct ConditionalBranchMask : public BranchMask { - ConditionalBranchMask(std::shared_ptr selection_vector, int64_t length) - : selection_vector_(std::move(selection_vector)), length_(length) {} - - bool empty() const override { return selection_vector_->length() == 0; } - - protected: - Result DoApplyCond(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const override { - auto sparse_input = input; - sparse_input.selection_vector = selection_vector_; - return ExecuteScalarExpression(expr, sparse_input, exec_context); - } - - Result> GetSelectionVector() const override { - return selection_vector_; - } - - Result> MakeBodyMaskFromBitmap( - const std::shared_ptr& bitmap, - ExecContext* exec_context) const override; - - Result> MakeBodyMaskFromBitmap( - const std::shared_ptr& bitmap, - ExecContext* exec_context) const override; - - protected: - std::shared_ptr selection_vector_ = nullptr; - int64_t length_ = 0; -}; - -struct ConditionalBodyMask : public BodyMask { - ConditionalBodyMask(std::shared_ptr body, - std::shared_ptr rest, int64_t length) - : body_(std::move(body)), rest_(std::move(rest)), length_(length) {} - - bool empty() const override { return body_->length() == 0; } - - Result ApplyCond(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const override { - auto sparse_input = input; - sparse_input.selection_vector = body_; - return ExecuteScalarExpression(expr, sparse_input, exec_context); - } - - Result> GetSelectionVector() const override { - return body_; - } - - Result> NextBranchMask() const override { - if (!rest_) { - return std::make_shared(); - } - return std::make_shared(rest_, length_); - } - - private: - std::shared_ptr body_; - std::shared_ptr rest_; - int64_t length_; -}; - -Result> AllPassBranchMask::MakeBodyMaskFromBitmap( - const std::shared_ptr& bitmap, ExecContext* exec_context) const { - DCHECK_EQ(bitmap->length(), length_); - - Int32Builder body_builder(exec_context->memory_pool()); - Int32Builder rest_builder(exec_context->memory_pool()); - RETURN_NOT_OK(body_builder.Reserve(length_)); - RETURN_NOT_OK(rest_builder.Reserve(length_)); - - ArraySpan span(*bitmap->data()); - int32_t i = 0; - VisitArraySpanInline( - span, - [&](bool mask) { - if (mask) { - body_builder.UnsafeAppend(i); - } else { - rest_builder.UnsafeAppend(i); - } - ++i; - }, - [&]() { ++i; }); - - ARROW_ASSIGN_OR_RAISE(auto body_arr, body_builder.Finish()); - ARROW_ASSIGN_OR_RAISE(auto rest_arr, rest_builder.Finish()); - auto body = std::make_shared(body_arr->data()); - auto rest = std::make_shared(rest_arr->data()); - return std::make_shared(std::move(body), std::move(rest), length_); -} - -Result> AllPassBranchMask::MakeBodyMaskFromBitmap( - const std::shared_ptr& bitmap, ExecContext* exec_context) const { - DCHECK_EQ(bitmap->length(), length_); - - Int32Builder body_builder(exec_context->memory_pool()); - Int32Builder rest_builder(exec_context->memory_pool()); - RETURN_NOT_OK(body_builder.Reserve(length_)); - RETURN_NOT_OK(rest_builder.Reserve(length_)); - - int32_t i = 0; - for (const auto& chunk : bitmap->chunks()) { - DCHECK_EQ(chunk->type()->id(), Type::BOOL); - ArraySpan span(*chunk->data()); - VisitArraySpanInline( - span, - [&](bool mask) { - if (mask) { - body_builder.UnsafeAppend(i); - } else { - rest_builder.UnsafeAppend(i); - } - ++i; - }, - [&]() { ++i; }); - } - - ARROW_ASSIGN_OR_RAISE(auto body_arr, body_builder.Finish()); - ARROW_ASSIGN_OR_RAISE(auto rest_arr, rest_builder.Finish()); - auto body = std::make_shared(body_arr->data()); - auto rest = std::make_shared(rest_arr->data()); - return std::make_shared(std::move(body), std::move(rest), length_); -} - -Result> ConditionalBranchMask::MakeBodyMaskFromBitmap( - const std::shared_ptr& bitmap, ExecContext* exec_context) const { - DCHECK_EQ(bitmap->length(), length_); - - Int32Builder body_builder(exec_context->memory_pool()); - Int32Builder rest_builder(exec_context->memory_pool()); - RETURN_NOT_OK(body_builder.Reserve(length_)); - RETURN_NOT_OK(rest_builder.Reserve(length_)); - - for (int64_t i = 0; i < selection_vector_->length(); ++i) { - auto index = selection_vector_->indices()[i]; - if (!bitmap->IsNull(index)) { - if (bitmap->Value(index)) { - body_builder.UnsafeAppend(index); - } else { - rest_builder.UnsafeAppend(index); - } - } - } - - ARROW_ASSIGN_OR_RAISE(auto body_arr, body_builder.Finish()); - ARROW_ASSIGN_OR_RAISE(auto rest_arr, rest_builder.Finish()); - auto body = std::make_shared(body_arr->data()); - auto rest = std::make_shared(rest_arr->data()); - return std::make_shared(std::move(body), std::move(rest), length_); -} - -Result> ConditionalBranchMask::MakeBodyMaskFromBitmap( - const std::shared_ptr& bitmap, ExecContext* exec_context) const { - DCHECK_EQ(bitmap->length(), length_); - - std::vector boolean_arrays(bitmap->num_chunks()); - std::transform(bitmap->chunks().begin(), bitmap->chunks().end(), boolean_arrays.begin(), - [](const auto& chunk) { - DCHECK_EQ(chunk->type()->id(), Type::BOOL); - return checked_cast(chunk.get()); - }); - - Int32Builder body_builder(exec_context->memory_pool()); - Int32Builder rest_builder(exec_context->memory_pool()); - RETURN_NOT_OK(body_builder.Reserve(length_)); - RETURN_NOT_OK(rest_builder.Reserve(length_)); - - ChunkResolver resolver(bitmap->chunks()); - ChunkLocation location; - for (int64_t i = 0; i < selection_vector_->length(); ++i) { - auto index = selection_vector_->indices()[i]; - location = resolver.ResolveWithHint(index, location); - if (boolean_arrays[location.chunk_index]->IsValid(location.index_in_chunk)) { - if (boolean_arrays[location.chunk_index]->Value(location.index_in_chunk)) { - body_builder.UnsafeAppend(index); - } else { - rest_builder.UnsafeAppend(index); - } - } - } - - ARROW_ASSIGN_OR_RAISE(auto body_arr, body_builder.Finish()); - ARROW_ASSIGN_OR_RAISE(auto rest_arr, rest_builder.Finish()); - auto body = std::make_shared(body_arr->data()); - auto rest = std::make_shared(rest_arr->data()); - return std::make_shared(std::move(body), std::move(rest), length_); -} - -struct Branch { - Expression cond; - Expression body; -}; - -struct ConditionalExec { - ConditionalExec(const std::vector& branches, const TypeHolder& result_type) - : branches(branches), result_type(result_type) {} - - Result Execute(const ExecBatch& input, ExecContext* exec_context) const&& { - DCHECK(!branches.empty()); - - BranchResults results; - results.Reserve(branches.size()); - ARROW_ASSIGN_OR_RAISE(auto branch_mask, InitBranchMask(input, exec_context)); - for (const auto& branch : branches) { - if (branch_mask->empty()) { - break; - } - ARROW_ASSIGN_OR_RAISE(auto body_mask, - ApplyCond(branch_mask, branch.cond, input, exec_context)); - if (body_mask->empty()) { - ARROW_ASSIGN_OR_RAISE(branch_mask, body_mask->NextBranchMask()); - continue; - } - ARROW_ASSIGN_OR_RAISE(auto body_result, - ApplyBody(body_mask, branch.body, input, exec_context)); - DCHECK(body_result.type()->Equals(*result_type)); - ARROW_ASSIGN_OR_RAISE(auto selection_vector, body_mask->GetSelectionVector()); - results.Emplace(std::move(body_result), std::move(selection_vector)); - ARROW_ASSIGN_OR_RAISE(branch_mask, body_mask->NextBranchMask()); - } - return MultiplexResults(input, results, exec_context); - } - - private: - struct BranchResults { - void Reserve(int64_t size) { - body_results_.reserve(size); - selection_vectors_.reserve(size); - } - - void Emplace(Datum body_result, std::shared_ptr selection_vector) { - body_results_.emplace_back(std::move(body_result)); - selection_vectors_.emplace_back(std::move(selection_vector)); - } - - bool empty() const { return body_results_.empty(); } - - size_t size() const { return body_results_.size(); } - - const std::vector& body_results() const { return body_results_; } - - const std::vector>& selection_vectors() const { - return selection_vectors_; - } - - private: - std::vector body_results_; - std::vector> selection_vectors_; - }; - - Result> InitBranchMask( - const ExecBatch& input, ExecContext* exec_context) const { - if (input.selection_vector) { - return std::make_shared(input.selection_vector, - input.length); - } - return std::make_shared(input.length); - } - - Result> ApplyCond( - const std::shared_ptr& branch_mask, const Expression& cond, - const ExecBatch& input, ExecContext* exec_context) const { - return branch_mask->ApplyCond(cond, input, exec_context); - } - - Result ApplyBody(const std::shared_ptr& body_mask, - const Expression& body, const ExecBatch& input, - ExecContext* exec_context) const { - return body_mask->ApplyCond(body, input, exec_context); - } - - Result MultiplexResults(const ExecBatch& input, const BranchResults& results, - ExecContext* exec_context) const { - if (results.empty()) { - return MakeArrayOfNull(result_type.GetSharedPtr(), input.length, - exec_context->memory_pool()); - } - - if (results.size() == 1) { - if (const auto& result = results.body_results()[0]; - results.selection_vectors()[0] == nullptr || - results.selection_vectors()[0]->length() == - (input.selection_vector ? input.selection_vector->length() - : input.length)) { - return result; - } - } - - std::vector choose_args; - choose_args.reserve(results.size() + 1); - ARROW_ASSIGN_OR_RAISE(auto indices, ChooseIndices(results.selection_vectors(), - input.length, exec_context)); - choose_args.emplace_back(std::move(indices)); - choose_args.insert(choose_args.end(), results.body_results().begin(), - results.body_results().end()); - return CallFunction("choose", choose_args, exec_context); - } - - Result ChooseIndices( - const std::vector>& selection_vectors, - int64_t length, ExecContext* exec_context) const { - const int64_t validity_bytes = bit_util::BytesForBits(length); - ARROW_ASSIGN_OR_RAISE( - std::shared_ptr validity_buf, - AllocateResizableBuffer(validity_bytes, exec_context->memory_pool())); - auto validity_data = validity_buf->mutable_data_as(); - std::memset(validity_data, 0, validity_bytes); - - ARROW_ASSIGN_OR_RAISE( - std::shared_ptr indices_buf, - AllocateResizableBuffer(length * sizeof(int32_t), exec_context->memory_pool())); - auto indices_data = indices_buf->mutable_data_as(); - for (int32_t index = 0; index < static_cast(selection_vectors.size()); - ++index) { - DCHECK_NE(selection_vectors[index], nullptr); - DCHECK_GT(selection_vectors[index]->length(), 0); - auto row_ids = selection_vectors[index]->indices(); - for (int64_t i = 0; i < selection_vectors[index]->length(); ++i) { - const int32_t row_id = row_ids[i]; - DCHECK_EQ(bit_util::GetBit(validity_data, row_id), false); - bit_util::SetBitTo(validity_data, row_id, true); - indices_data[row_id] = index; - } - } - - return ArrayData::Make(int32(), length, - {std::move(validity_buf), std::move(indices_buf)}); - } - - protected: - const std::vector& branches; - const TypeHolder& result_type; -}; - -class ConditionalSpecialExecutor : public SpecialExecutor { - public: - ConditionalSpecialExecutor(std::vector branches, TypeHolder out_type) - : SpecialExecutor(std::move(out_type), /*options=*/nullptr), - branches(std::move(branches)) {} - - Result Execute(const ExecBatch& input, - ExecContext* exec_context) const override { - return ConditionalExec(branches, out_type_).Execute(input, exec_context); - } - - private: - std::vector branches; -}; - -template -class FunctionBackedSpecialForm : public SpecialForm { - public: - using SpecialForm::SpecialForm; - - ARROW_DISALLOW_COPY_AND_ASSIGN(FunctionBackedSpecialForm); - ARROW_DEFAULT_MOVE_AND_ASSIGN(FunctionBackedSpecialForm); - - Result> Bind( - std::vector& arguments, std::shared_ptr options, - ExecContext* exec_context) const override { - DCHECK(std::all_of(arguments.begin(), arguments.end(), - [](const Expression& argument) { return argument.IsBound(); })); - Expression::Call call; - call.function_name = name(); - call.arguments = std::move(arguments); - call.options = std::move(options); - ARROW_ASSIGN_OR_RAISE( - auto bound, BindNonRecursive(call, /*insert_implicit_casts=*/true, exec_context)); - auto bound_call = CallNotNull(bound); - auto bound_call_copy = *bound_call; - arguments = std::move(bound_call->arguments); - options = std::move(bound_call->options); - return static_cast(this)->BindWithBoundCall(std::move(bound_call_copy), - exec_context); - } -}; - -template -class ConditionalSpecialForm - : public FunctionBackedSpecialForm> { - public: - using FunctionBackedSpecialForm< - ConditionalSpecialForm>::FunctionBackedSpecialForm; - - ARROW_DISALLOW_COPY_AND_ASSIGN(ConditionalSpecialForm); - ARROW_DEFAULT_MOVE_AND_ASSIGN(ConditionalSpecialForm); - - protected: - Result> BindWithBoundCall( - Expression::Call call, ExecContext* exec_context) const { - // Shouldn't have options. This is guaranteed by the call binding. - DCHECK_EQ(call.options, nullptr); - - auto branches = - static_cast(this)->GetBranches(std::move(call.arguments)); - return std::make_unique(std::move(branches), - std::move(call.type)); - } - - friend class FunctionBackedSpecialForm>; -}; +using internal::Branch; +using internal::ConditionalSpecialForm; class IfElseSpecialForm : public ConditionalSpecialForm { public: diff --git a/cpp/src/arrow/compute/special/if_else_special_test.cc b/cpp/src/arrow/compute/special/if_else_special_test.cc index c9ee018f61ae..572db18b7782 100644 --- a/cpp/src/arrow/compute/special/if_else_special_test.cc +++ b/cpp/src/arrow/compute/special/if_else_special_test.cc @@ -34,14 +34,14 @@ using internal::ExpectBindsTo; using internal::kBoringSchema; using internal::no_change; -TEST(Expression, ToString) { +TEST(IfElseSpecial, ToString) { EXPECT_EQ( if_else_special(field_ref("cond"), field_ref("if_true"), field_ref("if_false")) .ToString(), "if_else_special(cond, if_true, if_false)"); } -TEST(Expression, Equality) { +TEST(IfElseSpecial, Equality) { EXPECT_EQ(if_else_special(literal(true), field_ref("a"), field_ref("b")), if_else_special(literal(true), field_ref("a"), field_ref("b"))); EXPECT_NE(if_else_special(literal(true), field_ref("a"), field_ref("b")), @@ -54,7 +54,7 @@ TEST(Expression, Equality) { call("if_else", {literal(true), field_ref("a"), field_ref("b")})); } -TEST(Expression, Hash) { +TEST(IfElseSpecial, Hash) { std::unordered_set set; EXPECT_TRUE( @@ -70,19 +70,19 @@ TEST(Expression, Hash) { EXPECT_EQ(set.size(), 2); } -TEST(Expression, IsScalarExpression) { +TEST(IfElseSpecial, IsScalarExpression) { EXPECT_TRUE(if_else_special(field_ref("cond"), field_ref("a"), field_ref("b")) .IsScalarExpression()); } -TEST(Expression, IsSatisfiable) { +TEST(IfElseSpecial, IsSatisfiable) { auto Bind = [](Expression expr) { return expr.Bind(*kBoringSchema).ValueOrDie(); }; EXPECT_TRUE(Bind(if_else_special(field_ref("bool"), field_ref("i32"), field_ref("i32"))) .IsSatisfiable()); } -TEST(Expression, FieldsInExpression) { +TEST(IfElseSpecial, FieldsInExpression) { auto ExpectFieldsAre = [](Expression expr, std::vector expected) { EXPECT_THAT(FieldsInExpression(expr), testing::ContainerEq(expected)); }; @@ -103,7 +103,7 @@ TEST(Expression, FieldsInExpression) { {"a", "b", "c"}); } -TEST(Expression, ExpressionHasFieldRefs) { +TEST(IfElseSpecial, ExpressionHasFieldRefs) { EXPECT_FALSE( ExpressionHasFieldRefs(if_else_special(literal(true), literal(1), literal(0)))); EXPECT_TRUE( @@ -114,7 +114,7 @@ TEST(Expression, ExpressionHasFieldRefs) { ExpressionHasFieldRefs(if_else_special(literal(true), literal(0), field_ref("a")))); } -TEST(Expression, BindSpecialForm) { +TEST(IfElseSpecial, BindSpecialForm) { { auto expr = if_else_special(field_ref("bool"), field_ref("i8"), field_ref("i8")); EXPECT_FALSE(expr.IsBound()); diff --git a/cpp/src/arrow/compute/special/special_form_internal.h b/cpp/src/arrow/compute/special/special_form_internal.h new file mode 100644 index 000000000000..9953b94a2ffe --- /dev/null +++ b/cpp/src/arrow/compute/special/special_form_internal.h @@ -0,0 +1,54 @@ +// 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. + +#pragma once + +#include "arrow/compute/expression_internal.h" +#include "arrow/compute/special_form.h" +#include "arrow/util/logging_internal.h" + +namespace arrow::compute::internal { + +template +class FunctionBackedSpecialForm : public SpecialForm { + public: + using SpecialForm::SpecialForm; + + ARROW_DISALLOW_COPY_AND_ASSIGN(FunctionBackedSpecialForm); + ARROW_DEFAULT_MOVE_AND_ASSIGN(FunctionBackedSpecialForm); + + Result> Bind( + std::vector& arguments, std::shared_ptr options, + ExecContext* exec_context) const override { + DCHECK(std::all_of(arguments.begin(), arguments.end(), + [](const Expression& argument) { return argument.IsBound(); })); + Expression::Call call; + call.function_name = name(); + call.arguments = std::move(arguments); + call.options = std::move(options); + ARROW_ASSIGN_OR_RAISE( + auto bound, BindNonRecursive(call, /*insert_implicit_casts=*/true, exec_context)); + auto bound_call = CallNotNull(bound); + auto bound_call_copy = *bound_call; + arguments = std::move(bound_call->arguments); + options = std::move(bound_call->options); + return static_cast(this)->BindWithBoundCall(std::move(bound_call_copy), + exec_context); + } +}; + +} // namespace arrow::compute::internal diff --git a/cpp/src/arrow/compute/special_form.h b/cpp/src/arrow/compute/special_form.h index 5588e6a73f2e..6b74079c8320 100644 --- a/cpp/src/arrow/compute/special_form.h +++ b/cpp/src/arrow/compute/special_form.h @@ -15,6 +15,8 @@ // specific language governing permissions and limitations // under the License. +#pragma once + #include "arrow/compute/expression.h" namespace arrow::compute { From bd3dfa513e64feffa399c10485f91668fde2293a Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Mon, 6 Oct 2025 22:48:21 -0700 Subject: [PATCH 35/71] Merge --- cpp/src/arrow/compute/exec.cc | 48 +++++++- cpp/src/arrow/compute/exec.h | 7 +- cpp/src/arrow/compute/exec_test.cc | 122 +++++++++++--------- cpp/src/arrow/compute/test_util_internal.cc | 12 ++ cpp/src/arrow/compute/test_util_internal.h | 4 + 5 files changed, 134 insertions(+), 59 deletions(-) diff --git a/cpp/src/arrow/compute/exec.cc b/cpp/src/arrow/compute/exec.cc index e45b831ffd49..82c534a58208 100644 --- a/cpp/src/arrow/compute/exec.cc +++ b/cpp/src/arrow/compute/exec.cc @@ -72,6 +72,18 @@ ExecContext* threaded_exec_context() { return &threaded_ctx; } +ExecBatch::ExecBatch(std::vector values, int64_t length, + std::shared_ptr selection_vector) + : values(std::move(values)), + length(length), + selection_vector(std::move(selection_vector)) { +#ifndef NDEBUG + if (this->selection_vector) { + DCHECK_OK(this->selection_vector->Validate(this->length)); + } +#endif +} + ExecBatch::ExecBatch(const RecordBatch& batch) : values(batch.num_columns()), length(batch.num_rows()) { auto columns = batch.column_data(); @@ -1449,8 +1461,8 @@ const CpuInfo* ExecContext::cpu_info() const { return CpuInfo::GetInstance(); } SelectionVector::SelectionVector(std::shared_ptr data) : data_(std::move(data)) { - DCHECK_EQ(Type::INT32, data_->type->id()); - DCHECK_EQ(0, data_->GetNullCount()); + DCHECK_NE(data_, nullptr); + DCHECK_EQ(data_->type->id(), Type::INT32); indices_ = data_->GetValues(1); } @@ -1458,6 +1470,38 @@ SelectionVector::SelectionVector(const Array& arr) : SelectionVector(arr.data()) int64_t SelectionVector::length() const { return data_->length; } +Status SelectionVector::Validate(int64_t max_index) const { + if (data_ == nullptr) { + return Status::Invalid("SelectionVector not initialized"); + } + ARROW_CHECK_NE(indices_, nullptr); + if (data_->type->id() != Type::INT32) { + return Status::Invalid("SelectionVector must be of type int32"); + } + if (data_->GetNullCount() != 0) { + return Status::Invalid("SelectionVector cannot contain nulls"); + } + for (int64_t i = 1; i < length(); ++i) { + if (indices_[i - 1] > indices_[i]) { + return Status::Invalid("SelectionVector indices must be sorted"); + } + } + for (int64_t i = 0; i < length(); ++i) { + if (indices_[i] < 0) { + return Status::Invalid("SelectionVector indices must be non-negative"); + } + } + if (max_index >= 0) { + for (int64_t i = 0; i < length(); ++i) { + if (indices_[i] > max_index) { + return Status::Invalid("SelectionVector index ", indices_[i], " exceeds maximum ", + max_index); + } + } + } + return Status::OK(); +} + void SelectionVectorSpan::SetSlice(int64_t offset, int64_t length, int32_t index_back_shift) { DCHECK_NE(indices_, nullptr); diff --git a/cpp/src/arrow/compute/exec.h b/cpp/src/arrow/compute/exec.h index e96b383ef624..7cf11ab7c141 100644 --- a/cpp/src/arrow/compute/exec.h +++ b/cpp/src/arrow/compute/exec.h @@ -144,6 +144,8 @@ class ARROW_EXPORT SelectionVector { const int32_t* indices() const { return indices_; } int64_t length() const; + Status Validate(int64_t max_index = -1) const; + private: std::shared_ptr data_; const int32_t* indices_; @@ -196,10 +198,7 @@ constexpr int64_t kUnsequencedIndex = -1; struct ARROW_EXPORT ExecBatch { ExecBatch() = default; ExecBatch(std::vector values, int64_t length, - std::shared_ptr selection_vector = NULLPTR) - : values(std::move(values)), - length(length), - selection_vector(std::move(selection_vector)) {} + std::shared_ptr selection_vector = NULLPTR); explicit ExecBatch(const RecordBatch& batch); diff --git a/cpp/src/arrow/compute/exec_test.cc b/cpp/src/arrow/compute/exec_test.cc index 483774cf04c1..c87d44c965d8 100644 --- a/cpp/src/arrow/compute/exec_test.cc +++ b/cpp/src/arrow/compute/exec_test.cc @@ -36,11 +36,11 @@ #include "arrow/compute/kernel.h" #include "arrow/compute/ordering.h" #include "arrow/compute/registry.h" +#include "arrow/compute/test_util_internal.h" #include "arrow/memory_pool.h" #include "arrow/record_batch.h" #include "arrow/scalar.h" #include "arrow/status.h" -#include "arrow/testing/generator.h" #include "arrow/type.h" #include "arrow/util/bit_util.h" #include "arrow/util/bitmap_ops.h" @@ -148,11 +148,38 @@ TEST(ExecContext, BasicWorkings) { } TEST(SelectionVector, Basics) { - auto indices = ArrayFromJSON(int32(), "[0, 3]"); - auto sel_vector = std::make_shared(*indices); + auto sel_vector = SelectionVectorFromJSON("[0, 42]"); - ASSERT_EQ(indices->length(), sel_vector->length()); - ASSERT_EQ(3, sel_vector->indices()[1]); + ASSERT_EQ(sel_vector->length(), 2); + ASSERT_EQ(sel_vector->indices()[0], 0); + ASSERT_EQ(sel_vector->indices()[1], 42); +} + +TEST(SelectionVector, Validate) { + { + auto sel_vector = SelectionVectorFromJSON("[]"); + ASSERT_OK(sel_vector->Validate()); + } + { + auto sel_vector = SelectionVectorFromJSON("[0, null, 42]"); + ASSERT_RAISES(Invalid, sel_vector->Validate()); + } + { + auto sel_vector = SelectionVectorFromJSON("[42, 0]"); + ASSERT_RAISES(Invalid, sel_vector->Validate()); + } + { + auto sel_vector = SelectionVectorFromJSON("[-42, 0]"); + ASSERT_RAISES(Invalid, sel_vector->Validate()); + } + { + auto sel_vector = SelectionVectorFromJSON("[0, 1]"); + ASSERT_OK(sel_vector->Validate(/*max_index=*/1)); + } + { + auto sel_vector = SelectionVectorFromJSON("[0, 42]"); + ASSERT_RAISES(Invalid, sel_vector->Validate(/*max_index=*/1)); + } } TEST(SelectionVectorSpan, Basics) { @@ -921,22 +948,11 @@ TEST_F(TestExecSpanIterator, ZeroLengthInputs) { CheckArgs(input); } -std::shared_ptr MakeSelectionVector(const std::string& json) { - return std::make_shared(*ArrayFromJSON(int32(), json)); -} - -std::shared_ptr MakeSelectionVectorUntil(int64_t length) { - auto res = gen::Step()->Generate(length); - DCHECK_OK(res.status()); - auto arr = res.ValueUnsafe(); - return std::make_shared(*arr); -} - TEST_F(TestExecSpanIterator, SelectionSpanBasic) { ExecBatch batch( {Datum(GetInt32Array(30)), Datum(GetInt32Array(30)), Datum(std::make_shared(5)), Datum(MakeNullScalar(boolean()))}, - 30, MakeSelectionVector("[1, 2, 7, 29]")); + 30, SelectionVectorFromJSON("[1, 2, 7, 29]")); CheckIteration(batch, /*chunksize=*/7, {7, 7, 7, 7, 2}, {2, 1, 0, 0, 1}); CheckIteration(batch, /*chunksize=*/10, {10, 10, 10}, {3, 0, 1}); @@ -948,7 +964,7 @@ TEST_F(TestExecSpanIterator, SelectionSpanChunked) { ExecBatch batch({Datum(GetInt32Chunked({0, 20, 10})), Datum(GetInt32Chunked({15, 15})), Datum(GetInt32Array(30)), Datum(std::make_shared(5)), Datum(MakeNullScalar(boolean()))}, - 30, MakeSelectionVector("[1, 2, 7, 29]")); + 30, SelectionVectorFromJSON("[1, 2, 7, 29]")); CheckIteration(batch, /*chunksize=*/7, {7, 7, 1, 5, 7, 3}, {2, 1, 0, 0, 0, 1}); CheckIteration(batch, /*chunksize=*/10, {10, 5, 5, 10}, {3, 0, 0, 1}); @@ -959,17 +975,17 @@ TEST_F(TestExecSpanIterator, SelectionSpanChunked) { // ---------------------------------------------------------------------- // Scalar function execution -template -void VisitSelectionVector(const SelectionVectorSpan& selection, int64_t length, - OnSelectionFn&& on_selection, - OnNonSelectionFn&& on_non_selection) { +template +void VisitIndicesWithSelection(int64_t length, const SelectionVectorSpan& selection, + OnSelectedFn&& on_selected, + OnNonSelectedFn&& on_non_selected) { int64_t selected = 0; for (int64_t i = 0; i < length; ++i) { if (selected < selection.length() && i == selection[selected]) { - on_selection(i); + on_selected(i); ++selected; } else { - on_non_selection(i); + on_non_selected(i); } } } @@ -990,8 +1006,8 @@ void AssertArraysEqualSparseWithSelection(const Array& src, int64_t src_offset = src.data()->offset; int64_t dst_offset = dst.data()->offset; - VisitSelectionVector( - selection, src.length(), + VisitIndicesWithSelection( + src.length(), selection, [&](int64_t i) { // Selected values should match ASSERT_EQ(bit_util::GetBit(src_validity, src_offset + i), @@ -1025,8 +1041,8 @@ void AssertArraysEqualDenseWithSelection(const Array& src, int64_t src_offset = src.data()->offset; int64_t dst_offset = dst.data()->offset; - VisitSelectionVector( - selection, src.length(), + VisitIndicesWithSelection( + src.length(), selection, [&](int64_t i) { // Selected values should match ASSERT_EQ(bit_util::GetBit(src_validity, src_offset + i), @@ -1107,8 +1123,8 @@ Status SelectiveExecCopyArrayData(KernelContext* ctx, const ExecSpan& batch, int64_t dst_validity_offset = out_arr->offset; uint8_t* dst = out_arr->buffers[1]->mutable_data() + out_arr->offset * value_size; const uint8_t* src = arg0.buffers[1].data + arg0.offset * value_size; - VisitSelectionVector( - selection, batch.length, + VisitIndicesWithSelection( + batch.length, selection, [&](int64_t i) { // Copy the selected value std::memcpy(dst + i * value_size, src + i * value_size, value_size); @@ -1143,8 +1159,8 @@ Status SelectiveExecCopyArraySpan(KernelContext* ctx, const ExecSpan& batch, int64_t dst_validity_offset = out_arr->offset; uint8_t* dst = out_arr->buffers[1].data + out_arr->offset * value_size; const uint8_t* src = arg0.buffers[1].data + arg0.offset * value_size; - VisitSelectionVector( - selection, batch.length, + VisitIndicesWithSelection( + batch.length, selection, [&](int64_t i) { // Copy the selected value std::memcpy(dst + i * value_size, src + i * value_size, value_size); @@ -1312,8 +1328,8 @@ Status SelectiveExecStateful(KernelContext* ctx, const ExecSpan& batch, uint8_t* dst_validity = out_arr->buffers[0].data; int64_t dst_validity_offset = out_arr->offset; int32_t* dst = out_arr->GetValues(1); - VisitSelectionVector( - selection, batch.length, + VisitIndicesWithSelection( + batch.length, selection, [&](int64_t i) { // Copy the selected value dst[i] = arg0_data[i] * multiplier; @@ -1345,8 +1361,8 @@ Status SelectiveExecAddInt32(KernelContext* ctx, const ExecSpan& batch, uint8_t* dst_validity = out_arr->buffers[0].data; int64_t dst_validity_offset = out_arr->offset; int32_t* out_data = out_arr->GetValues(1); - VisitSelectionVector( - selection, batch.length, + VisitIndicesWithSelection( + batch.length, selection, [&](int64_t i) { // Copy the selected value out_data[i] = left_data[i] + right_data[i]; @@ -1707,16 +1723,16 @@ class TestCallScalarFunctionPreallocationCases : public TestCallScalarFunction { std::shared_ptr GetTestArray() { return GetUInt8Array(100, 0.2); } std::vector> GetTestSelectionVectors() { - return {MakeSelectionVector("[]"), - MakeSelectionVector("[0]"), - MakeSelectionVector("[42]"), - MakeSelectionVector("[99]"), - MakeSelectionVector("[0, 1, 2, 3, 4]"), - MakeSelectionVector("[0, 42, 99]"), - MakeSelectionVectorUntil(40), - MakeSelectionVectorUntil(41), - MakeSelectionVectorUntil(99), - MakeSelectionVectorUntil(100)}; + return {SelectionVectorFromJSON("[]"), + SelectionVectorFromJSON("[0]"), + SelectionVectorFromJSON("[42]"), + SelectionVectorFromJSON("[99]"), + SelectionVectorFromJSON("[0, 1, 2, 3, 4]"), + SelectionVectorFromJSON("[0, 42, 99]"), + MakeSelectionVectorTo(40), + MakeSelectionVectorTo(41), + MakeSelectionVectorTo(99), + MakeSelectionVectorTo(100)}; } template @@ -1981,9 +1997,9 @@ class TestCallScalarFunctionBasicNonStandardCases : public TestCallScalarFunctio std::shared_ptr GetTestArray() { return GetUInt8Array(1000, 0.2); } std::vector> GetTestSelectionVectors() { - return {MakeSelectionVector("[]"), MakeSelectionVector("[0]"), - MakeSelectionVector("[999]"), MakeSelectionVectorUntil(400), - MakeSelectionVectorUntil(401), MakeSelectionVectorUntil(1000)}; + return {SelectionVectorFromJSON("[]"), SelectionVectorFromJSON("[0]"), + SelectionVectorFromJSON("[999]"), MakeSelectionVectorTo(400), + MakeSelectionVectorTo(401), MakeSelectionVectorTo(1000)}; } template @@ -2146,9 +2162,9 @@ class TestCallScalarFunctionStatefulKernel : public TestCallScalarFunction { } std::vector> GetTestSelectionVectors() { - return {MakeSelectionVector("[]"), MakeSelectionVector("[0]"), - MakeSelectionVector("[4]"), MakeSelectionVectorUntil(2), - MakeSelectionVectorUntil(5)}; + return {SelectionVectorFromJSON("[]"), SelectionVectorFromJSON("[0]"), + SelectionVectorFromJSON("[4]"), MakeSelectionVectorTo(2), + MakeSelectionVectorTo(5)}; } template @@ -2226,7 +2242,7 @@ class TestCallScalarFunctionScalarFunction : public TestCallScalarFunction { static constexpr int32_t kExpectedResult = 12; std::vector> GetTestSelectionVectors() { - return {MakeSelectionVector("[]"), MakeSelectionVector("[0]")}; + return {SelectionVectorFromJSON("[]"), SelectionVectorFromJSON("[0]")}; } void DoTestBasic(const FunctionCaller* caller, const std::vector& args, diff --git a/cpp/src/arrow/compute/test_util_internal.cc b/cpp/src/arrow/compute/test_util_internal.cc index 219028943cb3..6407799399c2 100644 --- a/cpp/src/arrow/compute/test_util_internal.cc +++ b/cpp/src/arrow/compute/test_util_internal.cc @@ -24,6 +24,7 @@ #include "arrow/record_batch.h" #include "arrow/scalar.h" #include "arrow/table.h" +#include "arrow/testing/generator.h" #include "arrow/testing/gtest_util.h" #include "arrow/type.h" #include "arrow/util/logging_internal.h" @@ -120,4 +121,15 @@ void ValidateOutput(const Datum& output) { } } +std::shared_ptr SelectionVectorFromJSON(const std::string& json) { + return std::make_shared(*ArrayFromJSON(int32(), json)); +} + +std::shared_ptr MakeSelectionVectorTo(int64_t length) { + auto res = gen::Step()->Generate(length); + DCHECK_OK(res.status()); + auto arr = res.ValueUnsafe(); + return std::make_shared(*arr); +} + } // namespace arrow::compute diff --git a/cpp/src/arrow/compute/test_util_internal.h b/cpp/src/arrow/compute/test_util_internal.h index 6a172b07692e..b7f4e1261993 100644 --- a/cpp/src/arrow/compute/test_util_internal.h +++ b/cpp/src/arrow/compute/test_util_internal.h @@ -39,4 +39,8 @@ ExecBatch ExecBatchFromJSON(const std::vector& types, void ValidateOutput(const Datum& output); +std::shared_ptr SelectionVectorFromJSON(const std::string& json); + +std::shared_ptr MakeSelectionVectorTo(int64_t length); + } // namespace arrow::compute From 1a6b82820f59b6ea3c311873d6f9cc20500ec80a Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Mon, 6 Oct 2025 23:58:41 -0700 Subject: [PATCH 36/71] Port --- cpp/src/arrow/compute/exec.cc | 10 +++++----- cpp/src/arrow/compute/exec.h | 2 +- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/cpp/src/arrow/compute/exec.cc b/cpp/src/arrow/compute/exec.cc index 82c534a58208..f8b59df98289 100644 --- a/cpp/src/arrow/compute/exec.cc +++ b/cpp/src/arrow/compute/exec.cc @@ -1470,7 +1470,7 @@ SelectionVector::SelectionVector(const Array& arr) : SelectionVector(arr.data()) int64_t SelectionVector::length() const { return data_->length; } -Status SelectionVector::Validate(int64_t max_index) const { +Status SelectionVector::Validate(int64_t values_length) const { if (data_ == nullptr) { return Status::Invalid("SelectionVector not initialized"); } @@ -1491,11 +1491,11 @@ Status SelectionVector::Validate(int64_t max_index) const { return Status::Invalid("SelectionVector indices must be non-negative"); } } - if (max_index >= 0) { + if (values_length >= 0) { for (int64_t i = 0; i < length(); ++i) { - if (indices_[i] > max_index) { - return Status::Invalid("SelectionVector index ", indices_[i], " exceeds maximum ", - max_index); + if (indices_[i] >= values_length) { + return Status::Invalid("SelectionVector index ", indices_[i], + " >= values length ", values_length); } } } diff --git a/cpp/src/arrow/compute/exec.h b/cpp/src/arrow/compute/exec.h index 7cf11ab7c141..b6116b4f08ed 100644 --- a/cpp/src/arrow/compute/exec.h +++ b/cpp/src/arrow/compute/exec.h @@ -144,7 +144,7 @@ class ARROW_EXPORT SelectionVector { const int32_t* indices() const { return indices_; } int64_t length() const; - Status Validate(int64_t max_index = -1) const; + Status Validate(int64_t values_length = -1) const; private: std::shared_ptr data_; From 21badec310bdc78ebb51e2c6ce64bbb493178689 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Tue, 7 Oct 2025 00:06:51 -0700 Subject: [PATCH 37/71] Fix all scalar execbatch not match selection vector --- cpp/src/arrow/compute/exec.cc | 12 ------------ cpp/src/arrow/compute/exec.h | 5 ++++- cpp/src/arrow/compute/expression.cc | 9 ++++++++- 3 files changed, 12 insertions(+), 14 deletions(-) diff --git a/cpp/src/arrow/compute/exec.cc b/cpp/src/arrow/compute/exec.cc index f8b59df98289..c14921f32ba6 100644 --- a/cpp/src/arrow/compute/exec.cc +++ b/cpp/src/arrow/compute/exec.cc @@ -72,18 +72,6 @@ ExecContext* threaded_exec_context() { return &threaded_ctx; } -ExecBatch::ExecBatch(std::vector values, int64_t length, - std::shared_ptr selection_vector) - : values(std::move(values)), - length(length), - selection_vector(std::move(selection_vector)) { -#ifndef NDEBUG - if (this->selection_vector) { - DCHECK_OK(this->selection_vector->Validate(this->length)); - } -#endif -} - ExecBatch::ExecBatch(const RecordBatch& batch) : values(batch.num_columns()), length(batch.num_rows()) { auto columns = batch.column_data(); diff --git a/cpp/src/arrow/compute/exec.h b/cpp/src/arrow/compute/exec.h index b6116b4f08ed..f254eb7d5bb5 100644 --- a/cpp/src/arrow/compute/exec.h +++ b/cpp/src/arrow/compute/exec.h @@ -198,7 +198,10 @@ constexpr int64_t kUnsequencedIndex = -1; struct ARROW_EXPORT ExecBatch { ExecBatch() = default; ExecBatch(std::vector values, int64_t length, - std::shared_ptr selection_vector = NULLPTR); + std::shared_ptr selection_vector = NULLPTR) + : values(std::move(values)), + length(length), + selection_vector(std::move(selection_vector)) {} explicit ExecBatch(const RecordBatch& batch); diff --git a/cpp/src/arrow/compute/expression.cc b/cpp/src/arrow/compute/expression.cc index c884e6d205f2..2fc25d7f279b 100644 --- a/cpp/src/arrow/compute/expression.cc +++ b/cpp/src/arrow/compute/expression.cc @@ -837,12 +837,19 @@ Result ExecuteScalarExpression(const Expression& expr, const ExecBatch& i } int64_t input_length; + std::shared_ptr input_selection_vector = nullptr; if (!arguments.empty() && all_scalar) { // all inputs are scalar, so use a 1-long batch to avoid // computing input.length equivalent outputs input_length = 1; } else { input_length = input.length; + input_selection_vector = input.selection_vector; +#ifndef NDEBUG + if (input_selection_vector) { + RETURN_NOT_OK(input_selection_vector->Validate(input.length)); + } +#endif } auto executor = compute::detail::KernelExecutor::MakeScalar(); @@ -857,7 +864,7 @@ Result ExecuteScalarExpression(const Expression& expr, const ExecBatch& i compute::detail::DatumAccumulator listener; RETURN_NOT_OK(executor->Execute( - ExecBatch(arguments, input_length, input.selection_vector), &listener)); + ExecBatch(arguments, input_length, std::move(input_selection_vector)), &listener)); const auto out = executor->WrapResults(arguments, listener.values()); #ifndef NDEBUG DCHECK_OK(executor->CheckResultType(out, call->function_name.c_str())); From 780738e90a124b372e28de15a86f485afdf315e4 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Tue, 7 Oct 2025 01:17:56 -0700 Subject: [PATCH 38/71] Some branch mask tests --- .../compute/special/conditional_internal.h | 15 ++++++++++++--- .../arrow/compute/special/conditional_test.cc | 19 +++++++++++++++++-- 2 files changed, 29 insertions(+), 5 deletions(-) diff --git a/cpp/src/arrow/compute/special/conditional_internal.h b/cpp/src/arrow/compute/special/conditional_internal.h index 6e01d163ff2d..290dc8e3bcdb 100644 --- a/cpp/src/arrow/compute/special/conditional_internal.h +++ b/cpp/src/arrow/compute/special/conditional_internal.h @@ -118,7 +118,11 @@ struct ARROW_COMPUTE_EXPORT AllFailBranchMask : public BranchMask { struct ARROW_COMPUTE_EXPORT ConditionalBranchMask : public BranchMask { ConditionalBranchMask(std::shared_ptr selection_vector, int64_t length) - : selection_vector_(std::move(selection_vector)), length_(length) {} + : selection_vector_(std::move(selection_vector)), length_(length) { +#ifndef NDEBUG + DCHECK_OK(selection_vector_->Validate(length_)); +#endif + } bool empty() const override { return selection_vector_->length() == 0; } @@ -243,7 +247,12 @@ struct ARROW_COMPUTE_EXPORT Branch { struct ARROW_COMPUTE_EXPORT ConditionalBodyMask : public BodyMask { ConditionalBodyMask(std::shared_ptr body, std::shared_ptr rest, int64_t length) - : body_(std::move(body)), rest_(std::move(rest)), length_(length) {} + : body_(std::move(body)), rest_(std::move(rest)), length_(length) { +#ifndef NDEBUG + DCHECK_OK(body_->Validate(length_)); + DCHECK_OK(rest_->Validate(length_)); +#endif + } bool empty() const override { return body_->length() == 0; } @@ -255,7 +264,7 @@ struct ARROW_COMPUTE_EXPORT ConditionalBodyMask : public BodyMask { } Result> NextBranchMask() const override { - if (!rest_) { + if (rest_->length() == 0) { return std::make_shared(); } return std::make_shared(rest_, length_); diff --git a/cpp/src/arrow/compute/special/conditional_test.cc b/cpp/src/arrow/compute/special/conditional_test.cc index f1d1de04aa71..17ff230929eb 100644 --- a/cpp/src/arrow/compute/special/conditional_test.cc +++ b/cpp/src/arrow/compute/special/conditional_test.cc @@ -18,13 +18,28 @@ #include #include "arrow/compute/special/conditional_internal.h" +#include "arrow/compute/test_util_internal.h" namespace arrow::compute::internal { TEST(AllPassBranchMask, Emptiness) { for (auto length : {0, 42}) { - auto branch_mask = std::make_shared(length); - EXPECT_FALSE(branch_mask->empty()); + auto all_pass = std::make_shared(length); + EXPECT_FALSE(all_pass->empty()); + } +} + +TEST(AllFailBranchMask, Emptiness) { + auto all_fail = std::make_shared(); + EXPECT_TRUE(all_fail->empty()); +} + +TEST(ConditionalBranchMask, Emptiness) { + for (const auto& selection : + {SelectionVectorFromJSON("[]"), SelectionVectorFromJSON("[0]"), + SelectionVectorFromJSON("[0, 41]")}) { + auto conditional = std::make_shared(selection, /*length=*/42); + EXPECT_EQ(conditional->empty(), selection->length() == 0); } } From ee0ac7aa10aab7dbc7bbecf748986dae1f6f1d74 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Tue, 7 Oct 2025 02:26:47 -0700 Subject: [PATCH 39/71] Some tests --- cpp/src/arrow/compute/special/conditional.cc | 1 + .../arrow/compute/special/conditional_test.cc | 97 +++++++++++++++++-- 2 files changed, 92 insertions(+), 6 deletions(-) diff --git a/cpp/src/arrow/compute/special/conditional.cc b/cpp/src/arrow/compute/special/conditional.cc index bebafcadae40..2a3bad34122f 100644 --- a/cpp/src/arrow/compute/special/conditional.cc +++ b/cpp/src/arrow/compute/special/conditional.cc @@ -122,6 +122,7 @@ Result> AllPassBranchMask::MakeBodyMaskFromBitma Result ConditionalBranchMask::DoApplyCond(const Expression& expr, const ExecBatch& input, ExecContext* exec_context) const { + DCHECK_EQ(input.length, length_); auto sparse_input = input; sparse_input.selection_vector = selection_vector_; return ExecuteScalarExpression(expr, sparse_input, exec_context); diff --git a/cpp/src/arrow/compute/special/conditional_test.cc b/cpp/src/arrow/compute/special/conditional_test.cc index 17ff230929eb..97de8b90089f 100644 --- a/cpp/src/arrow/compute/special/conditional_test.cc +++ b/cpp/src/arrow/compute/special/conditional_test.cc @@ -19,27 +19,112 @@ #include "arrow/compute/special/conditional_internal.h" #include "arrow/compute/test_util_internal.h" +#include "arrow/testing/generator.h" +#include "arrow/testing/gtest_util.h" +#include "arrow/util/checked_cast.h" namespace arrow::compute::internal { TEST(AllPassBranchMask, Emptiness) { for (auto length : {0, 42}) { - auto all_pass = std::make_shared(length); - EXPECT_FALSE(all_pass->empty()); + auto branch_mask = std::make_shared(length); + EXPECT_FALSE(branch_mask->empty()); + } +} + +TEST(AllPassBranchMask, ApplyCond) { + for (auto length : {0, 42}) { + { + auto branch_mask = std::make_shared(length); + ASSERT_OK_AND_ASSIGN(auto body_mask, + branch_mask->ApplyCond(literal(true), ExecBatch({}, length), + default_exec_context())); + auto casted = checked_cast(body_mask.get()); + EXPECT_NE(casted, nullptr); + } + + { + auto branch_mask = std::make_shared(length); + ASSERT_OK_AND_ASSIGN(Expression expr, + field_ref(0).Bind(Schema({field("", boolean())}))); + ASSERT_OK_AND_ASSIGN( + auto body_mask, + branch_mask->ApplyCond(expr, ExecBatch({BooleanScalar(true)}, length), + default_exec_context())); + auto casted = checked_cast(body_mask.get()); + EXPECT_NE(casted, nullptr); + } + + { + auto branch_mask = std::make_shared(length); + ASSERT_OK_AND_ASSIGN(auto body_mask, + branch_mask->ApplyCond(literal(false), ExecBatch({}, length), + default_exec_context())); + auto casted = checked_cast(body_mask.get()); + EXPECT_NE(casted, nullptr); + } + + { + auto branch_mask = std::make_shared(length); + ASSERT_OK_AND_ASSIGN(Expression expr, + field_ref(0).Bind(Schema({field("", boolean())}))); + ASSERT_OK_AND_ASSIGN( + auto body_mask, + branch_mask->ApplyCond(expr, ExecBatch({BooleanScalar(false)}, length), + default_exec_context())); + auto casted = checked_cast(body_mask.get()); + EXPECT_NE(casted, nullptr); + } + + { + auto branch_mask = std::make_shared(length); + ASSERT_OK_AND_ASSIGN( + auto body_mask, + branch_mask->ApplyCond(literal(BooleanScalar()), ExecBatch({}, length), + default_exec_context())); + auto casted = checked_cast(body_mask.get()); + EXPECT_NE(casted, nullptr); + } + + { + auto branch_mask = std::make_shared(length); + ASSERT_OK_AND_ASSIGN(Expression expr, + field_ref(0).Bind(Schema({field("", boolean())}))); + ASSERT_OK_AND_ASSIGN(auto body_mask, branch_mask->ApplyCond( + expr, ExecBatch({BooleanScalar()}, length), + default_exec_context())); + auto casted = checked_cast(body_mask.get()); + EXPECT_NE(casted, nullptr); + } + + for (auto boolean_scalar : + {MakeScalar(true), MakeScalar(false), MakeNullScalar(boolean())}) { + auto branch_mask = std::make_shared(length); + ASSERT_OK_AND_ASSIGN(Expression expr, + field_ref(0).Bind(Schema({field("", boolean())}))); + ASSERT_OK_AND_ASSIGN(auto boolean_array, + gen::Constant(boolean_scalar)->Generate(length)); + ASSERT_OK_AND_ASSIGN( + auto body_mask, + branch_mask->ApplyCond(expr, ExecBatch({std::move(boolean_array)}, length), + default_exec_context())); + auto casted = checked_cast(body_mask.get()); + EXPECT_NE(casted, nullptr); + } } } TEST(AllFailBranchMask, Emptiness) { - auto all_fail = std::make_shared(); - EXPECT_TRUE(all_fail->empty()); + auto branch_mask = std::make_shared(); + EXPECT_TRUE(branch_mask->empty()); } TEST(ConditionalBranchMask, Emptiness) { for (const auto& selection : {SelectionVectorFromJSON("[]"), SelectionVectorFromJSON("[0]"), SelectionVectorFromJSON("[0, 41]")}) { - auto conditional = std::make_shared(selection, /*length=*/42); - EXPECT_EQ(conditional->empty(), selection->length() == 0); + auto branch_mask = std::make_shared(selection, /*length=*/42); + EXPECT_EQ(branch_mask->empty(), selection->length() == 0); } } From d559fb95f9717f249ea886f3d6aaa3f36f52188d Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Wed, 8 Oct 2025 18:34:52 -0700 Subject: [PATCH 40/71] Simplify branch and body masks and complete most tests (#60) * Refine * Simplify branch and body masks and complete most tests --- cpp/src/arrow/compute/special/conditional.cc | 101 ++-- .../compute/special/conditional_internal.h | 149 +++--- .../arrow/compute/special/conditional_test.cc | 442 +++++++++++++++--- 3 files changed, 474 insertions(+), 218 deletions(-) diff --git a/cpp/src/arrow/compute/special/conditional.cc b/cpp/src/arrow/compute/special/conditional.cc index 2a3bad34122f..c726b51bc092 100644 --- a/cpp/src/arrow/compute/special/conditional.cc +++ b/cpp/src/arrow/compute/special/conditional.cc @@ -23,13 +23,13 @@ namespace arrow::compute::internal { -Result> BranchMask::MakeBodyMaskFromDatum( +Result> BranchMask::MakeBodyMask( const Datum& datum, ExecContext* exec_context) const { DCHECK(datum.type()->id() == Type::BOOL); if (datum.is_scalar()) { auto scalar = datum.scalar_as(); if (!scalar.is_valid) { - return std::make_shared(shared_from_this()); + return std::make_shared(); } else if (scalar.value) { return std::make_shared(shared_from_this()); } else { @@ -43,13 +43,23 @@ Result> BranchMask::MakeBodyMaskFromDatum( return MakeBodyMaskFromBitmap(datum.chunked_array(), exec_context); } -Result AllPassBranchMask::DoApplyCond(const Expression& expr, - const ExecBatch& input, - ExecContext* exec_context) const { - DCHECK_EQ(input.length, length_); - auto input_with_sel_vec = input; - input_with_sel_vec.selection_vector = nullptr; - return ExecuteScalarExpression(expr, input_with_sel_vec, exec_context); +Result> BranchMask::FromSelectionVector( + std::shared_ptr selection, int64_t length) { + DCHECK_NE(selection, nullptr); + +#ifndef NDEBUG + RETURN_NOT_OK(selection->Validate(length)); +#endif + + if (selection->length() == 0) { + return std::make_shared(); + } + + if (selection->length() == length) { + return std::make_shared(length); + } + + return std::make_shared(std::move(selection), length); } Result> AllPassBranchMask::GetSelectionVector() const { @@ -61,9 +71,9 @@ Result> AllPassBranchMask::MakeBodyMaskFromBitma DCHECK_EQ(bitmap->length(), length_); Int32Builder body_builder(exec_context->memory_pool()); - Int32Builder rest_builder(exec_context->memory_pool()); + Int32Builder remainder_builder(exec_context->memory_pool()); RETURN_NOT_OK(body_builder.Reserve(length_)); - RETURN_NOT_OK(rest_builder.Reserve(length_)); + RETURN_NOT_OK(remainder_builder.Reserve(length_)); ArraySpan span(*bitmap->data()); int32_t i = 0; @@ -73,17 +83,18 @@ Result> AllPassBranchMask::MakeBodyMaskFromBitma if (mask) { body_builder.UnsafeAppend(i); } else { - rest_builder.UnsafeAppend(i); + remainder_builder.UnsafeAppend(i); } ++i; }, [&]() { ++i; }); ARROW_ASSIGN_OR_RAISE(auto body_arr, body_builder.Finish()); - ARROW_ASSIGN_OR_RAISE(auto rest_arr, rest_builder.Finish()); + ARROW_ASSIGN_OR_RAISE(auto remainder_arr, remainder_builder.Finish()); auto body = std::make_shared(body_arr->data()); - auto rest = std::make_shared(rest_arr->data()); - return std::make_shared(std::move(body), std::move(rest), length_); + auto remainder = std::make_shared(remainder_arr->data()); + return std::make_shared(std::move(body), std::move(remainder), + length_); } Result> AllPassBranchMask::MakeBodyMaskFromBitmap( @@ -91,9 +102,9 @@ Result> AllPassBranchMask::MakeBodyMaskFromBitma DCHECK_EQ(bitmap->length(), length_); Int32Builder body_builder(exec_context->memory_pool()); - Int32Builder rest_builder(exec_context->memory_pool()); + Int32Builder remainder_builder(exec_context->memory_pool()); RETURN_NOT_OK(body_builder.Reserve(length_)); - RETURN_NOT_OK(rest_builder.Reserve(length_)); + RETURN_NOT_OK(remainder_builder.Reserve(length_)); int32_t i = 0; for (const auto& chunk : bitmap->chunks()) { @@ -105,7 +116,7 @@ Result> AllPassBranchMask::MakeBodyMaskFromBitma if (mask) { body_builder.UnsafeAppend(i); } else { - rest_builder.UnsafeAppend(i); + remainder_builder.UnsafeAppend(i); } ++i; }, @@ -113,19 +124,11 @@ Result> AllPassBranchMask::MakeBodyMaskFromBitma } ARROW_ASSIGN_OR_RAISE(auto body_arr, body_builder.Finish()); - ARROW_ASSIGN_OR_RAISE(auto rest_arr, rest_builder.Finish()); + ARROW_ASSIGN_OR_RAISE(auto remainder_arr, remainder_builder.Finish()); auto body = std::make_shared(body_arr->data()); - auto rest = std::make_shared(rest_arr->data()); - return std::make_shared(std::move(body), std::move(rest), length_); -} - -Result ConditionalBranchMask::DoApplyCond(const Expression& expr, - const ExecBatch& input, - ExecContext* exec_context) const { - DCHECK_EQ(input.length, length_); - auto sparse_input = input; - sparse_input.selection_vector = selection_vector_; - return ExecuteScalarExpression(expr, sparse_input, exec_context); + auto remainder = std::make_shared(remainder_arr->data()); + return std::make_shared(std::move(body), std::move(remainder), + length_); } Result> ConditionalBranchMask::MakeBodyMaskFromBitmap( @@ -133,9 +136,9 @@ Result> ConditionalBranchMask::MakeBodyMaskFromB DCHECK_EQ(bitmap->length(), length_); Int32Builder body_builder(exec_context->memory_pool()); - Int32Builder rest_builder(exec_context->memory_pool()); + Int32Builder remainder_builder(exec_context->memory_pool()); RETURN_NOT_OK(body_builder.Reserve(length_)); - RETURN_NOT_OK(rest_builder.Reserve(length_)); + RETURN_NOT_OK(remainder_builder.Reserve(length_)); for (int64_t i = 0; i < selection_vector_->length(); ++i) { auto index = selection_vector_->indices()[i]; @@ -143,16 +146,17 @@ Result> ConditionalBranchMask::MakeBodyMaskFromB if (bitmap->Value(index)) { body_builder.UnsafeAppend(index); } else { - rest_builder.UnsafeAppend(index); + remainder_builder.UnsafeAppend(index); } } } ARROW_ASSIGN_OR_RAISE(auto body_arr, body_builder.Finish()); - ARROW_ASSIGN_OR_RAISE(auto rest_arr, rest_builder.Finish()); + ARROW_ASSIGN_OR_RAISE(auto remainder_arr, remainder_builder.Finish()); auto body = std::make_shared(body_arr->data()); - auto rest = std::make_shared(rest_arr->data()); - return std::make_shared(std::move(body), std::move(rest), length_); + auto remainder = std::make_shared(remainder_arr->data()); + return std::make_shared(std::move(body), std::move(remainder), + length_); } Result> ConditionalBranchMask::MakeBodyMaskFromBitmap( @@ -167,9 +171,9 @@ Result> ConditionalBranchMask::MakeBodyMaskFromB }); Int32Builder body_builder(exec_context->memory_pool()); - Int32Builder rest_builder(exec_context->memory_pool()); + Int32Builder remainder_builder(exec_context->memory_pool()); RETURN_NOT_OK(body_builder.Reserve(length_)); - RETURN_NOT_OK(rest_builder.Reserve(length_)); + RETURN_NOT_OK(remainder_builder.Reserve(length_)); ChunkResolver resolver(bitmap->chunks()); ChunkLocation location; @@ -180,24 +184,17 @@ Result> ConditionalBranchMask::MakeBodyMaskFromB if (boolean_arrays[location.chunk_index]->Value(location.index_in_chunk)) { body_builder.UnsafeAppend(index); } else { - rest_builder.UnsafeAppend(index); + remainder_builder.UnsafeAppend(index); } } } ARROW_ASSIGN_OR_RAISE(auto body_arr, body_builder.Finish()); - ARROW_ASSIGN_OR_RAISE(auto rest_arr, rest_builder.Finish()); + ARROW_ASSIGN_OR_RAISE(auto remainder_arr, remainder_builder.Finish()); auto body = std::make_shared(body_arr->data()); - auto rest = std::make_shared(rest_arr->data()); - return std::make_shared(std::move(body), std::move(rest), length_); -} - -Result ConditionalBodyMask::ApplyCond(const Expression& expr, - const ExecBatch& input, - ExecContext* exec_context) const { - auto sparse_input = input; - sparse_input.selection_vector = body_; - return ExecuteScalarExpression(expr, sparse_input, exec_context); + auto remainder = std::make_shared(remainder_arr->data()); + return std::make_shared(std::move(body), std::move(remainder), + length_); } Result ConditionalExec::Execute(const ExecBatch& input, @@ -212,13 +209,13 @@ Result ConditionalExec::Execute(const ExecBatch& input, break; } ARROW_ASSIGN_OR_RAISE(auto body_mask, - ApplyCond(branch_mask, branch.cond, input, exec_context)); + EvaluateCond(branch_mask, branch.cond, input, exec_context)); if (body_mask->empty()) { ARROW_ASSIGN_OR_RAISE(branch_mask, body_mask->NextBranchMask()); continue; } ARROW_ASSIGN_OR_RAISE(auto body_result, - ApplyBody(body_mask, branch.body, input, exec_context)); + EvaluateBody(body_mask, branch.body, input, exec_context)); DCHECK(body_result.type()->Equals(*result_type)); ARROW_ASSIGN_OR_RAISE(auto selection_vector, body_mask->GetSelectionVector()); results.Emplace(std::move(body_result), std::move(selection_vector)); diff --git a/cpp/src/arrow/compute/special/conditional_internal.h b/cpp/src/arrow/compute/special/conditional_internal.h index 290dc8e3bcdb..6c074d041ff1 100644 --- a/cpp/src/arrow/compute/special/conditional_internal.h +++ b/cpp/src/arrow/compute/special/conditional_internal.h @@ -26,50 +26,49 @@ namespace arrow::compute::internal { +/// Structures to manage masks for branching expressions. +/// Mostly for efficient short circuiting of if-else chains. +/// The whole abstraction is as follows: +/// - A branch represents an compound expression of condition and body; +/// - A branch mask represents the set of rows (in an ExecBatch) to be evaluated for the +/// branch condition; +/// - When a branch mask is applied to an condition expression, it produces a body mask; +/// - A body mask represents the set of rows (in an ExecBatch) to be evaluated for the +/// branch body. +/// - A body mask also produces the next branch mask, representing the set of rows +/// remaining to be evaluated for the next branch in the chain. + struct ARROW_COMPUTE_EXPORT BodyMask; struct ARROW_COMPUTE_EXPORT BranchMask : public std::enable_shared_from_this { virtual ~BranchMask() = default; - Result> ApplyCond(const Expression& expr, - const ExecBatch& input, - ExecContext* exec_context) const { - ARROW_ASSIGN_OR_RAISE(auto datum, DoApplyCond(expr, input, exec_context)); - return MakeBodyMaskFromDatum(datum, exec_context); - } - virtual bool empty() const = 0; - protected: - virtual Result DoApplyCond(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const = 0; - virtual Result> GetSelectionVector() const = 0; + Result> MakeBodyMask(const Datum& datum, + ExecContext* exec_context) const; + + static Result> FromSelectionVector( + std::shared_ptr selection, int64_t length); + + protected: virtual Result> MakeBodyMaskFromBitmap( const std::shared_ptr& bitmap, ExecContext* exec_context) const = 0; virtual Result> MakeBodyMaskFromBitmap( const std::shared_ptr& bitmap, ExecContext* exec_context) const = 0; - - private: - Result> MakeBodyMaskFromDatum( - const Datum& datum, ExecContext* exec_context) const; - - friend struct DelegateBodyMask; }; struct ARROW_COMPUTE_EXPORT AllPassBranchMask : public BranchMask { explicit AllPassBranchMask(int64_t length) : length_(length) {} - bool empty() const override { return false; } - - protected: - Result DoApplyCond(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const override; + bool empty() const override { return length_ == 0; } Result> GetSelectionVector() const override; + protected: Result> MakeBodyMaskFromBitmap( const std::shared_ptr& bitmap, ExecContext* exec_context) const override; @@ -87,18 +86,12 @@ struct ARROW_COMPUTE_EXPORT AllFailBranchMask : public BranchMask { bool empty() const override { return true; } - protected: - Result DoApplyCond(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const override { - DCHECK(false); - return Status::Invalid("AllFailBranchMask::DoApplyCond should not be called"); - } - Result> GetSelectionVector() const override { DCHECK(false); return Status::Invalid("AllFailBranchMask::GetSelectionVector should not be called"); } + protected: Result> MakeBodyMaskFromBitmap( const std::shared_ptr& bitmap, ExecContext* exec_context) const override { @@ -126,14 +119,11 @@ struct ARROW_COMPUTE_EXPORT ConditionalBranchMask : public BranchMask { bool empty() const override { return selection_vector_->length() == 0; } - protected: - Result DoApplyCond(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const override; - Result> GetSelectionVector() const override { return selection_vector_; } + protected: Result> MakeBodyMaskFromBitmap( const std::shared_ptr& bitmap, ExecContext* exec_context) const override; @@ -152,43 +142,16 @@ struct ARROW_COMPUTE_EXPORT BodyMask : public std::enable_shared_from_this ApplyCond(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const = 0; - virtual Result> GetSelectionVector() const = 0; virtual Result> NextBranchMask() const = 0; }; -struct ARROW_COMPUTE_EXPORT DelegateBodyMask : public BodyMask { - explicit DelegateBodyMask(std::shared_ptr branch_mask) - : branch_mask_(std::move(branch_mask)) {} - - protected: - Result DelegateApplyCond(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const { - return branch_mask_->DoApplyCond(expr, input, exec_context); - } - - Result> DelegateGetSelectionVector() const { - return branch_mask_->GetSelectionVector(); - } - - protected: - std::shared_ptr branch_mask_; -}; - -struct ARROW_COMPUTE_EXPORT AllNullBodyMask : public DelegateBodyMask { - using DelegateBodyMask::DelegateBodyMask; +struct ARROW_COMPUTE_EXPORT AllNullBodyMask : public BodyMask { + AllNullBodyMask() = default; bool empty() const override { return true; } - Result ApplyCond(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const override { - DCHECK(false); - return Status::Invalid("AllNullBodyMask::ApplyCond should not be called"); - } - Result> GetSelectionVector() const override { DCHECK(false); return Status::Invalid("AllNullBodyMask::GetSelectionVector should not be called"); @@ -199,18 +162,21 @@ struct ARROW_COMPUTE_EXPORT AllNullBodyMask : public DelegateBodyMask { } }; +struct ARROW_COMPUTE_EXPORT DelegateBodyMask : public BodyMask { + explicit DelegateBodyMask(std::shared_ptr branch_mask) + : branch_mask_(std::move(branch_mask)) {} + + protected: + std::shared_ptr branch_mask_; +}; + struct ARROW_COMPUTE_EXPORT AllPassBodyMask : public DelegateBodyMask { using DelegateBodyMask::DelegateBodyMask; - bool empty() const override { return false; } - - Result ApplyCond(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const override { - return DelegateApplyCond(expr, input, exec_context); - } + bool empty() const override { return branch_mask_->empty(); } Result> GetSelectionVector() const override { - return DelegateGetSelectionVector(); + return branch_mask_->GetSelectionVector(); } Result> NextBranchMask() const override { @@ -223,12 +189,6 @@ struct ARROW_COMPUTE_EXPORT AllFailBodyMask : public DelegateBodyMask { bool empty() const override { return true; } - Result ApplyCond(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const override { - DCHECK(false); - return Status::Invalid("AllFailBodyMask::ApplyCond should not be called"); - } - Result> GetSelectionVector() const override { DCHECK(false); return Status::Invalid("AllFailBodyMask::GetSelectionVector should not be called"); @@ -246,33 +206,27 @@ struct ARROW_COMPUTE_EXPORT Branch { struct ARROW_COMPUTE_EXPORT ConditionalBodyMask : public BodyMask { ConditionalBodyMask(std::shared_ptr body, - std::shared_ptr rest, int64_t length) - : body_(std::move(body)), rest_(std::move(rest)), length_(length) { + std::shared_ptr remainder, int64_t length) + : body_(std::move(body)), remainder_(std::move(remainder)), length_(length) { #ifndef NDEBUG DCHECK_OK(body_->Validate(length_)); - DCHECK_OK(rest_->Validate(length_)); + DCHECK_OK(remainder_->Validate(length_)); #endif } bool empty() const override { return body_->length() == 0; } - Result ApplyCond(const Expression& expr, const ExecBatch& input, - ExecContext* exec_context) const override; - Result> GetSelectionVector() const override { return body_; } Result> NextBranchMask() const override { - if (rest_->length() == 0) { - return std::make_shared(); - } - return std::make_shared(rest_, length_); + return BranchMask::FromSelectionVector(remainder_, length_); } private: std::shared_ptr body_; - std::shared_ptr rest_; + std::shared_ptr remainder_; int64_t length_; }; @@ -312,22 +266,29 @@ struct ARROW_COMPUTE_EXPORT ConditionalExec { Result> InitBranchMask( const ExecBatch& input, ExecContext* exec_context) const { if (input.selection_vector) { - return std::make_shared(input.selection_vector, - input.length); + return BranchMask::FromSelectionVector(input.selection_vector, input.length); } return std::make_shared(input.length); } - Result> ApplyCond( + Result> EvaluateCond( const std::shared_ptr& branch_mask, const Expression& cond, const ExecBatch& input, ExecContext* exec_context) const { - return branch_mask->ApplyCond(cond, input, exec_context); + auto input_with_selection = input; + ARROW_ASSIGN_OR_RAISE(input_with_selection.selection_vector, + branch_mask->GetSelectionVector()); + ARROW_ASSIGN_OR_RAISE( + auto datum, ExecuteScalarExpression(cond, input_with_selection, exec_context)); + return branch_mask->MakeBodyMask(datum, exec_context); } - Result ApplyBody(const std::shared_ptr& body_mask, - const Expression& body, const ExecBatch& input, - ExecContext* exec_context) const { - return body_mask->ApplyCond(body, input, exec_context); + Result EvaluateBody(const std::shared_ptr& body_mask, + const Expression& body, const ExecBatch& input, + ExecContext* exec_context) const { + auto input_with_selection = input; + ARROW_ASSIGN_OR_RAISE(input_with_selection.selection_vector, + body_mask->GetSelectionVector()); + return ExecuteScalarExpression(body, input_with_selection, exec_context); } Result MultiplexResults(const ExecBatch& input, const BranchResults& results, diff --git a/cpp/src/arrow/compute/special/conditional_test.cc b/cpp/src/arrow/compute/special/conditional_test.cc index 97de8b90089f..51587b35728a 100644 --- a/cpp/src/arrow/compute/special/conditional_test.cc +++ b/cpp/src/arrow/compute/special/conditional_test.cc @@ -15,102 +15,187 @@ // specific language governing permissions and limitations // under the License. +#include "arrow/compute/special/conditional_internal.h" + #include -#include "arrow/compute/special/conditional_internal.h" #include "arrow/compute/test_util_internal.h" +#include "arrow/testing/builder.h" #include "arrow/testing/generator.h" #include "arrow/testing/gtest_util.h" #include "arrow/util/checked_cast.h" namespace arrow::compute::internal { +namespace { + +template +void CheckBranchMask(const std::shared_ptr& branch_mask) { + auto casted = checked_cast(branch_mask.get()); + EXPECT_NE(casted, nullptr); +} + +template +void CheckMakeBodyMask(const std::shared_ptr& branch_mask, + const Datum& datum) { + ASSERT_OK_AND_ASSIGN(auto body_mask, + branch_mask->MakeBodyMask(datum, default_exec_context())); + auto casted = checked_cast(body_mask.get()); + EXPECT_NE(casted, nullptr); +} + +template +void CheckNextBranchMask(const std::shared_ptr& body_mask) { + ASSERT_OK_AND_ASSIGN(auto branch_mask, body_mask->NextBranchMask()); + CheckBranchMask(branch_mask); +} + +void AssertSelectionVectorsEqual(const std::shared_ptr& expected, + const std::shared_ptr& actual) { + if (expected == nullptr) { + EXPECT_EQ(actual, nullptr); + return; + } + + if (actual == nullptr) { + EXPECT_EQ(expected, nullptr); + return; + } + + ASSERT_EQ(expected->length(), actual->length()); + AssertArraysEqual(*MakeArray(expected->data()), *MakeArray(actual->data())); +} + +template +void CheckMakeBodyMaskAndSelection( + const std::shared_ptr& branch_mask, const Datum& datum, + const std::shared_ptr& expected_body_selection) { + ASSERT_OK_AND_ASSIGN(auto body_mask, + branch_mask->MakeBodyMask(datum, default_exec_context())); + auto casted = checked_cast(body_mask.get()); + EXPECT_NE(casted, nullptr); + ASSERT_OK_AND_ASSIGN(auto body_selection, body_mask->GetSelectionVector()); + AssertSelectionVectorsEqual(expected_body_selection, body_selection); +} + +template +void CheckNextBranchMaskAndSelection( + const std::shared_ptr& body_mask, + const std::shared_ptr& expected_branch_selection) { + ASSERT_OK_AND_ASSIGN(auto branch_mask, body_mask->NextBranchMask()); + auto casted = checked_cast(branch_mask.get()); + EXPECT_NE(casted, nullptr); + ASSERT_OK_AND_ASSIGN(auto branch_selection, branch_mask->GetSelectionVector()); + AssertSelectionVectorsEqual(expected_branch_selection, branch_selection); +} + +const auto kTrueScalar = MakeScalar(true); +const auto kFalseScalar = MakeScalar(false); +const auto kNullScalar = MakeNullScalar(boolean()); + +const auto kTrueLiteral = literal(true); +const auto kFalseLiteral = literal(false); +const auto kNullLiteral = literal(kNullScalar); + +} // namespace + +TEST(BranchMask, FromSelectionVector) { + { + ASSERT_OK_AND_ASSIGN(auto branch_mask, BranchMask::FromSelectionVector( + SelectionVectorFromJSON("[]"), 0)); + CheckBranchMask(branch_mask); + } + + { + ASSERT_OK_AND_ASSIGN(auto branch_mask, BranchMask::FromSelectionVector( + SelectionVectorFromJSON("[]"), 42)); + CheckBranchMask(branch_mask); + } + + { + ASSERT_OK_AND_ASSIGN(auto branch_mask, BranchMask::FromSelectionVector( + SelectionVectorFromJSON("[0]"), 1)); + CheckBranchMask(branch_mask); + } + + { + ASSERT_OK_AND_ASSIGN(auto branch_mask, + BranchMask::FromSelectionVector(MakeSelectionVectorTo(42), 42)); + CheckBranchMask(branch_mask); + } + + { + ASSERT_OK_AND_ASSIGN(auto branch_mask, BranchMask::FromSelectionVector( + SelectionVectorFromJSON("[0]"), 42)); + CheckBranchMask(branch_mask); + } +} + TEST(AllPassBranchMask, Emptiness) { for (auto length : {0, 42}) { auto branch_mask = std::make_shared(length); - EXPECT_FALSE(branch_mask->empty()); + EXPECT_EQ(branch_mask->empty(), length == 0); } } -TEST(AllPassBranchMask, ApplyCond) { +TEST(AllPassBranchMask, GetSelectionVector) { for (auto length : {0, 42}) { - { - auto branch_mask = std::make_shared(length); - ASSERT_OK_AND_ASSIGN(auto body_mask, - branch_mask->ApplyCond(literal(true), ExecBatch({}, length), - default_exec_context())); - auto casted = checked_cast(body_mask.get()); - EXPECT_NE(casted, nullptr); - } + auto branch_mask = std::make_shared(length); + EXPECT_EQ(branch_mask->GetSelectionVector(), nullptr); + } +} - { - auto branch_mask = std::make_shared(length); - ASSERT_OK_AND_ASSIGN(Expression expr, - field_ref(0).Bind(Schema({field("", boolean())}))); - ASSERT_OK_AND_ASSIGN( - auto body_mask, - branch_mask->ApplyCond(expr, ExecBatch({BooleanScalar(true)}, length), - default_exec_context())); - auto casted = checked_cast(body_mask.get()); - EXPECT_NE(casted, nullptr); - } +TEST(AllPassBranchMask, MakeBodyMask) { + for (auto length : {0, 42}) { + CheckMakeBodyMask(std::make_shared(length), + Datum(kTrueScalar)); + CheckMakeBodyMask(std::make_shared(length), + Datum(kFalseScalar)); + CheckMakeBodyMask(std::make_shared(length), + Datum(kNullScalar)); + // TODO: Consider short-cutting for these masks. { - auto branch_mask = std::make_shared(length); - ASSERT_OK_AND_ASSIGN(auto body_mask, - branch_mask->ApplyCond(literal(false), ExecBatch({}, length), - default_exec_context())); - auto casted = checked_cast(body_mask.get()); - EXPECT_NE(casted, nullptr); - } + ASSERT_OK_AND_ASSIGN(auto boolean_array, + gen::Constant(kTrueScalar)->Generate(length)); + CheckMakeBodyMaskAndSelection( + std::make_shared(length), Datum(boolean_array), + MakeSelectionVectorTo(length)); - { - auto branch_mask = std::make_shared(length); - ASSERT_OK_AND_ASSIGN(Expression expr, - field_ref(0).Bind(Schema({field("", boolean())}))); - ASSERT_OK_AND_ASSIGN( - auto body_mask, - branch_mask->ApplyCond(expr, ExecBatch({BooleanScalar(false)}, length), - default_exec_context())); - auto casted = checked_cast(body_mask.get()); - EXPECT_NE(casted, nullptr); + auto boolean_chunked_array = + std::make_shared(ArrayVector{boolean_array, boolean_array}); + CheckMakeBodyMaskAndSelection( + std::make_shared(length * 2), Datum(boolean_chunked_array), + MakeSelectionVectorTo(length * 2)); } - { - auto branch_mask = std::make_shared(length); - ASSERT_OK_AND_ASSIGN( - auto body_mask, - branch_mask->ApplyCond(literal(BooleanScalar()), ExecBatch({}, length), - default_exec_context())); - auto casted = checked_cast(body_mask.get()); - EXPECT_NE(casted, nullptr); - } + for (const auto& scalar : {kFalseScalar, kNullScalar}) { + ASSERT_OK_AND_ASSIGN(auto boolean_array, gen::Constant(scalar)->Generate(length)); + CheckMakeBodyMaskAndSelection( + std::make_shared(length), Datum(boolean_array), + SelectionVectorFromJSON("[]")); - { - auto branch_mask = std::make_shared(length); - ASSERT_OK_AND_ASSIGN(Expression expr, - field_ref(0).Bind(Schema({field("", boolean())}))); - ASSERT_OK_AND_ASSIGN(auto body_mask, branch_mask->ApplyCond( - expr, ExecBatch({BooleanScalar()}, length), - default_exec_context())); - auto casted = checked_cast(body_mask.get()); - EXPECT_NE(casted, nullptr); + auto boolean_chunked_array = + std::make_shared(ArrayVector{boolean_array, boolean_array}); + CheckMakeBodyMaskAndSelection( + std::make_shared(length * 2), Datum(boolean_chunked_array), + SelectionVectorFromJSON("[]")); } + } - for (auto boolean_scalar : - {MakeScalar(true), MakeScalar(false), MakeNullScalar(boolean())}) { - auto branch_mask = std::make_shared(length); - ASSERT_OK_AND_ASSIGN(Expression expr, - field_ref(0).Bind(Schema({field("", boolean())}))); - ASSERT_OK_AND_ASSIGN(auto boolean_array, - gen::Constant(boolean_scalar)->Generate(length)); - ASSERT_OK_AND_ASSIGN( - auto body_mask, - branch_mask->ApplyCond(expr, ExecBatch({std::move(boolean_array)}, length), - default_exec_context())); - auto casted = checked_cast(body_mask.get()); - EXPECT_NE(casted, nullptr); - } + { + auto boolean_array = ArrayFromJSON(boolean(), "[true, false, null, true]"); + CheckMakeBodyMaskAndSelection( + std::make_shared(boolean_array->length()), + Datum(boolean_array), SelectionVectorFromJSON("[0, 3]")); + } + + { + auto boolean_chunked_array = ChunkedArrayFromJSON( + boolean(), {"[true, false, null, true]", "[true, false, null, true]"}); + CheckMakeBodyMaskAndSelection( + std::make_shared(boolean_chunked_array->length()), + Datum(boolean_chunked_array), SelectionVectorFromJSON("[0, 3, 4, 7]")); } } @@ -122,10 +207,223 @@ TEST(AllFailBranchMask, Emptiness) { TEST(ConditionalBranchMask, Emptiness) { for (const auto& selection : {SelectionVectorFromJSON("[]"), SelectionVectorFromJSON("[0]"), - SelectionVectorFromJSON("[0, 41]")}) { + SelectionVectorFromJSON("[0, 41]"), MakeSelectionVectorTo(42)}) { auto branch_mask = std::make_shared(selection, /*length=*/42); EXPECT_EQ(branch_mask->empty(), selection->length() == 0); } } +TEST(ConditionalBranchMask, GetSelectionVector) { + for (const auto& selection : + {SelectionVectorFromJSON("[]"), SelectionVectorFromJSON("[0]"), + SelectionVectorFromJSON("[0, 41]"), MakeSelectionVectorTo(42)}) { + auto branch_mask = std::make_shared(selection, /*length=*/42); + ASSERT_OK_AND_ASSIGN(auto got, branch_mask->GetSelectionVector()); + AssertSelectionVectorsEqual(selection, got); + } +} + +TEST(ConditionalBranchMask, MakeBodyMask) { + const int64_t length = 42; + for (const auto& selection : + {SelectionVectorFromJSON("[]"), SelectionVectorFromJSON("[0]"), + SelectionVectorFromJSON("[0, 41]"), MakeSelectionVectorTo(length)}) { + CheckMakeBodyMask( + std::make_shared(selection, length), Datum(kTrueScalar)); + CheckMakeBodyMask( + std::make_shared(selection, length), Datum(kFalseScalar)); + CheckMakeBodyMask( + std::make_shared(selection, length), Datum(kNullScalar)); + + { + ASSERT_OK_AND_ASSIGN(auto boolean_array, + gen::Constant(kTrueScalar)->Generate(length)); + CheckMakeBodyMaskAndSelection( + std::make_shared(selection, length), + Datum(boolean_array), selection); + + auto boolean_chunked_array = + std::make_shared(ArrayVector{boolean_array, boolean_array}); + CheckMakeBodyMaskAndSelection( + std::make_shared(selection, length * 2), + Datum(boolean_chunked_array), selection); + } + + for (const auto& scalar : {kFalseScalar, kNullScalar}) { + ASSERT_OK_AND_ASSIGN(auto boolean_array, gen::Constant(scalar)->Generate(length)); + CheckMakeBodyMaskAndSelection( + std::make_shared(selection, length), + Datum(boolean_array), SelectionVectorFromJSON("[]")); + + auto boolean_chunked_array = + std::make_shared(ArrayVector{boolean_array, boolean_array}); + CheckMakeBodyMaskAndSelection( + std::make_shared(selection, length * 2), + Datum(boolean_chunked_array), SelectionVectorFromJSON("[]")); + } + } + + { + auto selection = SelectionVectorFromJSON("[2, 3]"); + auto boolean_array = ArrayFromJSON(boolean(), "[true, false, null, true]"); + CheckMakeBodyMaskAndSelection( + std::make_shared(selection, boolean_array->length()), + Datum(boolean_array), SelectionVectorFromJSON("[3]")); + } + + { + auto selection = SelectionVectorFromJSON("[2, 4]"); + auto boolean_chunked_array = ChunkedArrayFromJSON( + boolean(), {"[true, false, null, true]", "[true, false, null, true]"}); + CheckMakeBodyMaskAndSelection( + std::make_shared(selection, + boolean_chunked_array->length()), + Datum(boolean_chunked_array), SelectionVectorFromJSON("[4]")); + } +} + +TEST(AllNullBodyMask, Emptiness) { + auto body_mask = std::make_shared(); + EXPECT_TRUE(body_mask->empty()); +} + +TEST(AllNullBodyMask, NextBranchMask) { + auto body_mask = std::make_shared(); + ASSERT_OK_AND_ASSIGN(auto next_branch_mask, body_mask->NextBranchMask()); + CheckBranchMask(next_branch_mask); +} + +TEST(AllPassBodyMask, Emptiness) { + for (const auto& branch_masks : std::vector>{ + std::make_shared(0), + std::make_shared(42), + std::make_shared(SelectionVectorFromJSON("[]"), 0), + std::make_shared(SelectionVectorFromJSON("[]"), 42), + std::make_shared(SelectionVectorFromJSON("[0]"), 42), + std::make_shared(SelectionVectorFromJSON("[0, 41]"), + 42), + std::make_shared(MakeSelectionVectorTo(42), 42)}) { + auto body_mask = std::make_shared(branch_masks); + EXPECT_EQ(body_mask->empty(), branch_masks->empty()); + } +} + +TEST(AllPassBodyMask, GetSelectionVector) { + const int64_t length = 42; + for (const auto& branch_mask : std::vector>{ + std::make_shared(length), + std::make_shared(SelectionVectorFromJSON("[]"), length), + std::make_shared(SelectionVectorFromJSON("[0]"), + length), + std::make_shared(SelectionVectorFromJSON("[0, 41]"), + length), + std::make_shared(MakeSelectionVectorTo(length), + length)}) { + auto body_mask = std::make_shared(branch_mask); + ASSERT_OK_AND_ASSIGN(auto body_selection, body_mask->GetSelectionVector()); + ASSERT_OK_AND_ASSIGN(auto branch_selection, branch_mask->GetSelectionVector()); + AssertSelectionVectorsEqual(branch_selection, body_selection); + } +} + +TEST(AllPassBodyMask, NextBranchMask) { + const int64_t length = 42; + for (const auto& branch_mask : std::vector>{ + std::make_shared(length), + std::make_shared(SelectionVectorFromJSON("[]"), length), + std::make_shared(SelectionVectorFromJSON("[0]"), + length), + std::make_shared(SelectionVectorFromJSON("[0, 41]"), + length), + std::make_shared(MakeSelectionVectorTo(length), + length)}) { + auto body_mask = std::make_shared(branch_mask); + CheckNextBranchMask(body_mask); + } +} + +TEST(AllFailBodyMask, Emptiness) { + for (const auto& branch_masks : std::vector>{ + std::make_shared(0), + std::make_shared(42), + std::make_shared(SelectionVectorFromJSON("[]"), 0), + std::make_shared(SelectionVectorFromJSON("[]"), 42), + std::make_shared(SelectionVectorFromJSON("[0]"), 42), + std::make_shared(SelectionVectorFromJSON("[0, 41]"), + 42), + std::make_shared(MakeSelectionVectorTo(42), 42)}) { + auto body_mask = std::make_shared(branch_masks); + EXPECT_TRUE(body_mask->empty()); + } +} + +TEST(AllFailBodyMask, NextBranchMask) { + const int64_t length = 42; + + { + auto branch_mask = std::make_shared(length); + ASSERT_OK_AND_ASSIGN(auto branch_selection, branch_mask->GetSelectionVector()); + auto body_mask = std::make_shared(branch_mask); + CheckNextBranchMaskAndSelection(body_mask, branch_selection); + } + + for (const auto& branch_mask : + {std::make_shared(SelectionVectorFromJSON("[]"), length), + std::make_shared(SelectionVectorFromJSON("[0]"), length), + std::make_shared(SelectionVectorFromJSON("[0, 41]"), + length), + std::make_shared(MakeSelectionVectorTo(length), length)}) { + ASSERT_OK_AND_ASSIGN(auto branch_selection, branch_mask->GetSelectionVector()); + auto body_mask = std::make_shared(branch_mask); + CheckNextBranchMaskAndSelection(body_mask, branch_selection); + } +} + +TEST(ConditionalBodyMask, Emptiness) { + const int64_t length = 42; + for (const auto& selection : + {SelectionVectorFromJSON("[]"), SelectionVectorFromJSON("[0]"), + SelectionVectorFromJSON("[0, 41]"), MakeSelectionVectorTo(length)}) { + auto body_mask = std::make_shared( + selection, SelectionVectorFromJSON("[1]"), length); + EXPECT_EQ(body_mask->empty(), selection->length() == 0); + } +} + +TEST(ConditionalBodyMask, GetSelectionVector) { + const int64_t length = 42; + for (const auto& selection : + {SelectionVectorFromJSON("[]"), SelectionVectorFromJSON("[0]"), + SelectionVectorFromJSON("[0, 41]"), MakeSelectionVectorTo(length)}) { + auto body_mask = std::make_shared( + selection, SelectionVectorFromJSON("[1]"), length); + ASSERT_OK_AND_ASSIGN(auto got, body_mask->GetSelectionVector()); + AssertSelectionVectorsEqual(selection, got); + } +} + +TEST(ConditionalBodyMask, NextBranchMask) { + const int64_t length = 42; + + { + auto body_mask = std::make_shared( + SelectionVectorFromJSON("[]"), SelectionVectorFromJSON("[]"), length); + CheckNextBranchMask(body_mask); + } + + for (const auto& remainder : + {SelectionVectorFromJSON("[0]"), SelectionVectorFromJSON("[0, 41]")}) { + auto body_mask = std::make_shared(SelectionVectorFromJSON("[]"), + remainder, length); + CheckNextBranchMaskAndSelection(body_mask, remainder); + } + + { + auto remainder = MakeSelectionVectorTo(length); + auto body_mask = std::make_shared(SelectionVectorFromJSON("[]"), + remainder, length); + CheckNextBranchMaskAndSelection(body_mask, nullptr); + } +} + } // namespace arrow::compute::internal From 1d9dca2e1a83a1ef42a202b884f15b19c408907c Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Thu, 9 Oct 2025 01:45:16 -0700 Subject: [PATCH 41/71] Conditional exec tests --- .../compute/special/conditional_internal.h | 10 +- .../arrow/compute/special/conditional_test.cc | 153 ++++++++++++++++++ 2 files changed, 158 insertions(+), 5 deletions(-) diff --git a/cpp/src/arrow/compute/special/conditional_internal.h b/cpp/src/arrow/compute/special/conditional_internal.h index 6c074d041ff1..1f482e2e8152 100644 --- a/cpp/src/arrow/compute/special/conditional_internal.h +++ b/cpp/src/arrow/compute/special/conditional_internal.h @@ -199,11 +199,6 @@ struct ARROW_COMPUTE_EXPORT AllFailBodyMask : public DelegateBodyMask { } }; -struct ARROW_COMPUTE_EXPORT Branch { - Expression cond; - Expression body; -}; - struct ARROW_COMPUTE_EXPORT ConditionalBodyMask : public BodyMask { ConditionalBodyMask(std::shared_ptr body, std::shared_ptr remainder, int64_t length) @@ -230,6 +225,11 @@ struct ARROW_COMPUTE_EXPORT ConditionalBodyMask : public BodyMask { int64_t length_; }; +struct ARROW_COMPUTE_EXPORT Branch { + Expression cond; + Expression body; +}; + struct ARROW_COMPUTE_EXPORT ConditionalExec { ConditionalExec(const std::vector& branches, const TypeHolder& result_type) : branches(branches), result_type(result_type) {} diff --git a/cpp/src/arrow/compute/special/conditional_test.cc b/cpp/src/arrow/compute/special/conditional_test.cc index 51587b35728a..e59aaa5f7373 100644 --- a/cpp/src/arrow/compute/special/conditional_test.cc +++ b/cpp/src/arrow/compute/special/conditional_test.cc @@ -426,4 +426,157 @@ TEST(ConditionalBodyMask, NextBranchMask) { } } +TEST(ConditionalExec, Basic) { + auto schm = schema({field("", boolean()), field("", boolean()), field("", utf8()), + field("", utf8()), field("", utf8())}); + + ASSERT_OK_AND_ASSIGN(auto b0, field_ref(0).Bind(*schm)); + ASSERT_OK_AND_ASSIGN(auto b1, field_ref(1).Bind(*schm)); + ASSERT_OK_AND_ASSIGN(auto sa, field_ref(2).Bind(*schm)); + ASSERT_OK_AND_ASSIGN(auto sb, field_ref(3).Bind(*schm)); + ASSERT_OK_AND_ASSIGN(auto sc, field_ref(4).Bind(*schm)); + + auto batch = ExecBatchFromJSON({boolean(), boolean(), utf8(), utf8(), utf8()}, + R"([[true, true, "a0", "b0", "c0"], + [false, false, "a1", "b1", "c1"], + [null, true, "a2", "b2", "c2"], + [true, true, "a3", "b3", "c3"]])"); + + { + ASSERT_OK_AND_ASSIGN( + auto result, ConditionalExec({Branch{kFalseLiteral, sa}, Branch{kTrueLiteral, sb}, + Branch{kTrueLiteral, sc}}, + TypeHolder(utf8())) + .Execute(batch, default_exec_context())); + AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["b0", "b1", "b2", "b3"])")), + result); + } + + { + ASSERT_OK_AND_ASSIGN( + auto result, ConditionalExec({Branch{kFalseLiteral, sa}, Branch{kNullLiteral, sb}, + Branch{kTrueLiteral, sc}}, + TypeHolder(utf8())) + .Execute(batch, default_exec_context())); + AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"([null, null, null, null])")), + result); + } + + { + ASSERT_OK_AND_ASSIGN(auto result, + ConditionalExec({Branch{kFalseLiteral, sa}, Branch{b0, sb}, + Branch{kTrueLiteral, sc}}, + TypeHolder(utf8())) + .Execute(batch, default_exec_context())); + AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["b0", "c1", null, "b3"])")), + result); + } + + { + ASSERT_OK_AND_ASSIGN(auto result, + ConditionalExec({Branch{kFalseLiteral, sa}, Branch{b1, sb}, + Branch{kTrueLiteral, sc}}, + TypeHolder(utf8())) + .Execute(batch, default_exec_context())); + AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["b0", "c1", "b2", "b3"])")), + result); + } + + { + ASSERT_OK_AND_ASSIGN( + auto result, ConditionalExec({Branch{kTrueLiteral, sa}, Branch{kTrueLiteral, sb}, + Branch{kTrueLiteral, sc}}, + TypeHolder(utf8())) + .Execute(batch, default_exec_context())); + AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "a1", "a2", "a3"])")), + result); + } + + { + ASSERT_OK_AND_ASSIGN( + auto result, ConditionalExec({Branch{kNullLiteral, sa}, Branch{kTrueLiteral, sb}, + Branch{kTrueLiteral, sc}}, + TypeHolder(utf8())) + .Execute(batch, default_exec_context())); + AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"([null, null, null, null])")), + result); + } + + { + ASSERT_OK_AND_ASSIGN(auto result, + ConditionalExec({Branch{b0, sa}, Branch{kTrueLiteral, sb}, + Branch{kTrueLiteral, sc}}, + TypeHolder(utf8())) + .Execute(batch, default_exec_context())); + AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "b1", null, "a3"])")), + result); + } + + { + ASSERT_OK_AND_ASSIGN(auto result, + ConditionalExec({Branch{b0, sa}, Branch{kNullLiteral, sb}, + Branch{kTrueLiteral, sc}}, + TypeHolder(utf8())) + .Execute(batch, default_exec_context())); + AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", null, null, "a3"])")), + result); + } + + { + ASSERT_OK_AND_ASSIGN(auto result, ConditionalExec({Branch{b0, sa}, Branch{b0, sb}, + Branch{kTrueLiteral, sc}}, + TypeHolder(utf8())) + .Execute(batch, default_exec_context())); + AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "c1", null, "a3"])")), + result); + } + + { + ASSERT_OK_AND_ASSIGN(auto result, ConditionalExec({Branch{b0, sa}, Branch{b1, sb}, + Branch{kTrueLiteral, sc}}, + TypeHolder(utf8())) + .Execute(batch, default_exec_context())); + AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "c1", null, "a3"])")), + result); + } + + { + ASSERT_OK_AND_ASSIGN(auto result, + ConditionalExec({Branch{b1, sa}, Branch{kTrueLiteral, sb}, + Branch{kTrueLiteral, sc}}, + TypeHolder(utf8())) + .Execute(batch, default_exec_context())); + AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "b1", "a2", "a3"])")), + result); + } + + { + ASSERT_OK_AND_ASSIGN(auto result, + ConditionalExec({Branch{b1, sa}, Branch{kNullLiteral, sb}, + Branch{kTrueLiteral, sc}}, + TypeHolder(utf8())) + .Execute(batch, default_exec_context())); + AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", null, "a2", "a3"])")), + result); + } + + { + ASSERT_OK_AND_ASSIGN(auto result, ConditionalExec({Branch{b1, sa}, Branch{b0, sb}, + Branch{kTrueLiteral, sc}}, + TypeHolder(utf8())) + .Execute(batch, default_exec_context())); + AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "c1", "a2", "a3"])")), + result); + } + + { + ASSERT_OK_AND_ASSIGN(auto result, ConditionalExec({Branch{b1, sa}, Branch{b1, sb}, + Branch{kTrueLiteral, sc}}, + TypeHolder(utf8())) + .Execute(batch, default_exec_context())); + AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "c1", "a2", "a3"])")), + result); + } +} + } // namespace arrow::compute::internal From 4906ef2a89dcf127f46ec694ce1a3ff3b3d141a6 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Thu, 9 Oct 2025 14:59:39 -0700 Subject: [PATCH 42/71] More shortcut body mask generation --- cpp/src/arrow/compute/special/conditional.cc | 99 ++++++++++++--- .../arrow/compute/special/conditional_test.cc | 119 +++++++++++------- 2 files changed, 158 insertions(+), 60 deletions(-) diff --git a/cpp/src/arrow/compute/special/conditional.cc b/cpp/src/arrow/compute/special/conditional.cc index c726b51bc092..3e2616bdd5ec 100644 --- a/cpp/src/arrow/compute/special/conditional.cc +++ b/cpp/src/arrow/compute/special/conditional.cc @@ -23,24 +23,93 @@ namespace arrow::compute::internal { +namespace { + +template +std::shared_ptr MakeTrivialBodyMask( + CheckAllNullFn&& check_all_null, CheckAllTrueFn&& check_all_true, + CheckAllFalseFn&& check_all_false, std::shared_ptr branch_mask) { + if (check_all_null()) { + return std::make_shared(); + } else if (check_all_true()) { + return std::make_shared(std::move(branch_mask)); + } else if (check_all_false()) { + return std::make_shared(std::move(branch_mask)); + } + + return nullptr; +} + +} // namespace + Result> BranchMask::MakeBodyMask( const Datum& datum, ExecContext* exec_context) const { DCHECK(datum.type()->id() == Type::BOOL); if (datum.is_scalar()) { auto scalar = datum.scalar_as(); - if (!scalar.is_valid) { - return std::make_shared(); - } else if (scalar.value) { - return std::make_shared(shared_from_this()); - } else { - return std::make_shared(shared_from_this()); - } + auto body_mask = MakeTrivialBodyMask( + [&]() { return !scalar.is_valid; }, [&]() { return scalar.value; }, + [&]() { return !scalar.value; }, shared_from_this()); + DCHECK_NE(body_mask, nullptr); + return body_mask; } + if (datum.is_array()) { - return MakeBodyMaskFromBitmap(datum.array_as(), exec_context); + auto boolean_array = datum.array_as(); + if (auto body_mask = MakeTrivialBodyMask( + [&]() { return boolean_array->null_count() == boolean_array->length(); }, + [&]() { + return boolean_array->null_count() == 0 && + boolean_array->true_count() == boolean_array->length(); + }, + [&]() { + return boolean_array->null_count() == 0 && + boolean_array->false_count() == boolean_array->length(); + }, + shared_from_this()); + body_mask) { + return body_mask; + } + return MakeBodyMaskFromBitmap(std::move(boolean_array), exec_context); } + DCHECK(datum.is_chunked_array()); - return MakeBodyMaskFromBitmap(datum.chunked_array(), exec_context); + auto chunked_array = datum.chunked_array(); + DCHECK(std::all_of(chunked_array->chunks().begin(), chunked_array->chunks().end(), + [](const std::shared_ptr& chunk) { + return chunk->type()->id() == Type::BOOL; + })); + if (auto body_mask = MakeTrivialBodyMask( + [&]() { + return std::all_of(chunked_array->chunks().begin(), + chunked_array->chunks().end(), + [&](const std::shared_ptr& chunk) { + return chunk->null_count() == chunk->length(); + }); + }, + [&]() { + return std::all_of( + chunked_array->chunks().begin(), chunked_array->chunks().end(), + [&](const std::shared_ptr& chunk) { + auto boolean_array = checked_cast(chunk.get()); + return boolean_array->null_count() == 0 && + boolean_array->true_count() == boolean_array->length(); + }); + }, + [&]() { + return std::all_of( + chunked_array->chunks().begin(), chunked_array->chunks().end(), + [&](const std::shared_ptr& chunk) { + auto boolean_array = checked_cast(chunk.get()); + return boolean_array->null_count() == 0 && + boolean_array->false_count() == boolean_array->length(); + }); + }, + shared_from_this()); + body_mask) { + return body_mask; + } + return MakeBodyMaskFromBitmap(std::move(chunked_array), exec_context); } Result> BranchMask::FromSelectionVector( @@ -91,8 +160,8 @@ Result> AllPassBranchMask::MakeBodyMaskFromBitma ARROW_ASSIGN_OR_RAISE(auto body_arr, body_builder.Finish()); ARROW_ASSIGN_OR_RAISE(auto remainder_arr, remainder_builder.Finish()); - auto body = std::make_shared(body_arr->data()); - auto remainder = std::make_shared(remainder_arr->data()); + auto body = std::make_shared(*body_arr); + auto remainder = std::make_shared(*remainder_arr); return std::make_shared(std::move(body), std::move(remainder), length_); } @@ -125,8 +194,8 @@ Result> AllPassBranchMask::MakeBodyMaskFromBitma ARROW_ASSIGN_OR_RAISE(auto body_arr, body_builder.Finish()); ARROW_ASSIGN_OR_RAISE(auto remainder_arr, remainder_builder.Finish()); - auto body = std::make_shared(body_arr->data()); - auto remainder = std::make_shared(remainder_arr->data()); + auto body = std::make_shared(*body_arr); + auto remainder = std::make_shared(*remainder_arr); return std::make_shared(std::move(body), std::move(remainder), length_); } @@ -153,8 +222,8 @@ Result> ConditionalBranchMask::MakeBodyMaskFromB ARROW_ASSIGN_OR_RAISE(auto body_arr, body_builder.Finish()); ARROW_ASSIGN_OR_RAISE(auto remainder_arr, remainder_builder.Finish()); - auto body = std::make_shared(body_arr->data()); - auto remainder = std::make_shared(remainder_arr->data()); + auto body = std::make_shared(*body_arr); + auto remainder = std::make_shared(*remainder_arr); return std::make_shared(std::move(body), std::move(remainder), length_); } diff --git a/cpp/src/arrow/compute/special/conditional_test.cc b/cpp/src/arrow/compute/special/conditional_test.cc index e59aaa5f7373..a26b758c9586 100644 --- a/cpp/src/arrow/compute/special/conditional_test.cc +++ b/cpp/src/arrow/compute/special/conditional_test.cc @@ -89,13 +89,13 @@ void CheckNextBranchMaskAndSelection( AssertSelectionVectorsEqual(expected_branch_selection, branch_selection); } +const auto kNullScalar = MakeNullScalar(boolean()); const auto kTrueScalar = MakeScalar(true); const auto kFalseScalar = MakeScalar(false); -const auto kNullScalar = MakeNullScalar(boolean()); +const auto kNullLiteral = literal(kNullScalar); const auto kTrueLiteral = literal(true); const auto kFalseLiteral = literal(false); -const auto kNullLiteral = literal(kNullScalar); } // namespace @@ -145,44 +145,55 @@ TEST(AllPassBranchMask, GetSelectionVector) { } } -TEST(AllPassBranchMask, MakeBodyMask) { - for (auto length : {0, 42}) { - CheckMakeBodyMask(std::make_shared(length), - Datum(kTrueScalar)); - CheckMakeBodyMask(std::make_shared(length), - Datum(kFalseScalar)); +TEST(AllPassBranchMask, MakeBodyMaskTrival) { + const int64_t length = 42; + + CheckMakeBodyMask(std::make_shared(length), + Datum(kNullScalar)); + CheckMakeBodyMask(std::make_shared(length), + Datum(kTrueScalar)); + CheckMakeBodyMask(std::make_shared(length), + Datum(kFalseScalar)); + + { + ASSERT_OK_AND_ASSIGN(auto boolean_array, + gen::Constant(kNullScalar)->Generate(length)); CheckMakeBodyMask(std::make_shared(length), - Datum(kNullScalar)); + Datum(boolean_array)); - // TODO: Consider short-cutting for these masks. - { - ASSERT_OK_AND_ASSIGN(auto boolean_array, - gen::Constant(kTrueScalar)->Generate(length)); - CheckMakeBodyMaskAndSelection( - std::make_shared(length), Datum(boolean_array), - MakeSelectionVectorTo(length)); + auto boolean_chunked_array = + std::make_shared(ArrayVector{boolean_array, boolean_array}); + CheckMakeBodyMask(std::make_shared(length * 2), + Datum(boolean_chunked_array)); + } - auto boolean_chunked_array = - std::make_shared(ArrayVector{boolean_array, boolean_array}); - CheckMakeBodyMaskAndSelection( - std::make_shared(length * 2), Datum(boolean_chunked_array), - MakeSelectionVectorTo(length * 2)); - } + { + ASSERT_OK_AND_ASSIGN(auto boolean_array, + gen::Constant(kTrueScalar)->Generate(length)); + CheckMakeBodyMaskAndSelection( + std::make_shared(length), Datum(boolean_array), nullptr); - for (const auto& scalar : {kFalseScalar, kNullScalar}) { - ASSERT_OK_AND_ASSIGN(auto boolean_array, gen::Constant(scalar)->Generate(length)); - CheckMakeBodyMaskAndSelection( - std::make_shared(length), Datum(boolean_array), - SelectionVectorFromJSON("[]")); + auto boolean_chunked_array = + std::make_shared(ArrayVector{boolean_array, boolean_array}); + CheckMakeBodyMaskAndSelection( + std::make_shared(length * 2), Datum(boolean_chunked_array), + nullptr); + } - auto boolean_chunked_array = - std::make_shared(ArrayVector{boolean_array, boolean_array}); - CheckMakeBodyMaskAndSelection( - std::make_shared(length * 2), Datum(boolean_chunked_array), - SelectionVectorFromJSON("[]")); - } + { + ASSERT_OK_AND_ASSIGN(auto boolean_array, + gen::Constant(kFalseScalar)->Generate(length)); + CheckMakeBodyMask(std::make_shared(length), + Datum(boolean_array)); + + auto boolean_chunked_array = + std::make_shared(ArrayVector{boolean_array, boolean_array}); + CheckMakeBodyMask(std::make_shared(length * 2), + Datum(boolean_chunked_array)); } +} +TEST(AllPassBranchMask, MakeBodyMask) { { auto boolean_array = ArrayFromJSON(boolean(), "[true, false, null, true]"); CheckMakeBodyMaskAndSelection( @@ -223,8 +234,9 @@ TEST(ConditionalBranchMask, GetSelectionVector) { } } -TEST(ConditionalBranchMask, MakeBodyMask) { +TEST(ConditionalBranchMask, MakeBodyMaskTrivial) { const int64_t length = 42; + for (const auto& selection : {SelectionVectorFromJSON("[]"), SelectionVectorFromJSON("[0]"), SelectionVectorFromJSON("[0, 41]"), MakeSelectionVectorTo(length)}) { @@ -235,34 +247,51 @@ TEST(ConditionalBranchMask, MakeBodyMask) { CheckMakeBodyMask( std::make_shared(selection, length), Datum(kNullScalar)); + { + ASSERT_OK_AND_ASSIGN(auto boolean_array, + gen::Constant(kNullScalar)->Generate(length)); + CheckMakeBodyMask( + std::make_shared(selection, length), + Datum(boolean_array)); + + auto boolean_chunked_array = + std::make_shared(ArrayVector{boolean_array, boolean_array}); + CheckMakeBodyMask( + std::make_shared(selection, length * 2), + Datum(boolean_chunked_array)); + } + { ASSERT_OK_AND_ASSIGN(auto boolean_array, gen::Constant(kTrueScalar)->Generate(length)); - CheckMakeBodyMaskAndSelection( + CheckMakeBodyMaskAndSelection( std::make_shared(selection, length), Datum(boolean_array), selection); auto boolean_chunked_array = std::make_shared(ArrayVector{boolean_array, boolean_array}); - CheckMakeBodyMaskAndSelection( + CheckMakeBodyMaskAndSelection( std::make_shared(selection, length * 2), Datum(boolean_chunked_array), selection); } - for (const auto& scalar : {kFalseScalar, kNullScalar}) { - ASSERT_OK_AND_ASSIGN(auto boolean_array, gen::Constant(scalar)->Generate(length)); - CheckMakeBodyMaskAndSelection( + { + ASSERT_OK_AND_ASSIGN(auto boolean_array, + gen::Constant(kFalseScalar)->Generate(length)); + CheckMakeBodyMask( std::make_shared(selection, length), - Datum(boolean_array), SelectionVectorFromJSON("[]")); + Datum(boolean_array)); auto boolean_chunked_array = std::make_shared(ArrayVector{boolean_array, boolean_array}); - CheckMakeBodyMaskAndSelection( + CheckMakeBodyMask( std::make_shared(selection, length * 2), - Datum(boolean_chunked_array), SelectionVectorFromJSON("[]")); + Datum(boolean_chunked_array)); } } +} +TEST(ConditionalBranchMask, MakeBodyMask) { { auto selection = SelectionVectorFromJSON("[2, 3]"); auto boolean_array = ArrayFromJSON(boolean(), "[true, false, null, true]"); @@ -437,10 +466,10 @@ TEST(ConditionalExec, Basic) { ASSERT_OK_AND_ASSIGN(auto sc, field_ref(4).Bind(*schm)); auto batch = ExecBatchFromJSON({boolean(), boolean(), utf8(), utf8(), utf8()}, - R"([[true, true, "a0", "b0", "c0"], + R"([[true, true, "a0", "b0", "c0"], [false, false, "a1", "b1", "c1"], - [null, true, "a2", "b2", "c2"], - [true, true, "a3", "b3", "c3"]])"); + [null, true, "a2", "b2", "c2"], + [true, true, "a3", "b3", "c3"]])"); { ASSERT_OK_AND_ASSIGN( From 7ba20137443c34bf62dab97c3fc6d7d00665ca9d Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Thu, 9 Oct 2025 22:45:04 -0700 Subject: [PATCH 43/71] Refine --- cpp/src/arrow/CMakeLists.txt | 2 +- cpp/src/arrow/compute/special/CMakeLists.txt | 2 +- .../special/{conditional.cc => conditional_special.cc} | 2 +- ...{conditional_internal.h => conditional_special_internal.h} | 0 .../{conditional_test.cc => conditional_special_test.cc} | 2 +- cpp/src/arrow/compute/special/if_else_special.cc | 2 +- cpp/src/arrow/compute/special/if_else_special_test.cc | 4 ---- 7 files changed, 5 insertions(+), 9 deletions(-) rename cpp/src/arrow/compute/special/{conditional.cc => conditional_special.cc} (99%) rename cpp/src/arrow/compute/special/{conditional_internal.h => conditional_special_internal.h} (100%) rename cpp/src/arrow/compute/special/{conditional_test.cc => conditional_special_test.cc} (99%) diff --git a/cpp/src/arrow/CMakeLists.txt b/cpp/src/arrow/CMakeLists.txt index 4042f4e9f8ff..befe18fa16dd 100644 --- a/cpp/src/arrow/CMakeLists.txt +++ b/cpp/src/arrow/CMakeLists.txt @@ -802,7 +802,7 @@ if(ARROW_COMPUTE) compute/row/grouper.cc compute/row/row_encoder_internal.cc compute/row/row_internal.cc - compute/special/conditional.cc + compute/special/conditional_special.cc compute/special/if_else_special.cc compute/util.cc compute/util_internal.cc) diff --git a/cpp/src/arrow/compute/special/CMakeLists.txt b/cpp/src/arrow/compute/special/CMakeLists.txt index 1ba6d11819b4..c98b981751ec 100644 --- a/cpp/src/arrow/compute/special/CMakeLists.txt +++ b/cpp/src/arrow/compute/special/CMakeLists.txt @@ -25,7 +25,7 @@ add_arrow_compute_test(if_else_special_test add_arrow_compute_test(special_internal_test SOURCES - conditional_test.cc + conditional_special_test.cc EXTRA_LINK_LIBS arrow_compute_testing) diff --git a/cpp/src/arrow/compute/special/conditional.cc b/cpp/src/arrow/compute/special/conditional_special.cc similarity index 99% rename from cpp/src/arrow/compute/special/conditional.cc rename to cpp/src/arrow/compute/special/conditional_special.cc index 3e2616bdd5ec..41fde622e09a 100644 --- a/cpp/src/arrow/compute/special/conditional.cc +++ b/cpp/src/arrow/compute/special/conditional_special.cc @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -#include "arrow/compute/special/conditional_internal.h" +#include "arrow/compute/special/conditional_special_internal.h" #include "arrow/array/array_primitive.h" #include "arrow/array/builder_primitive.h" diff --git a/cpp/src/arrow/compute/special/conditional_internal.h b/cpp/src/arrow/compute/special/conditional_special_internal.h similarity index 100% rename from cpp/src/arrow/compute/special/conditional_internal.h rename to cpp/src/arrow/compute/special/conditional_special_internal.h diff --git a/cpp/src/arrow/compute/special/conditional_test.cc b/cpp/src/arrow/compute/special/conditional_special_test.cc similarity index 99% rename from cpp/src/arrow/compute/special/conditional_test.cc rename to cpp/src/arrow/compute/special/conditional_special_test.cc index a26b758c9586..c56faae08333 100644 --- a/cpp/src/arrow/compute/special/conditional_test.cc +++ b/cpp/src/arrow/compute/special/conditional_special_test.cc @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -#include "arrow/compute/special/conditional_internal.h" +#include "arrow/compute/special/conditional_special_internal.h" #include diff --git a/cpp/src/arrow/compute/special/if_else_special.cc b/cpp/src/arrow/compute/special/if_else_special.cc index 5bdf3fe17a18..39d6612978cb 100644 --- a/cpp/src/arrow/compute/special/if_else_special.cc +++ b/cpp/src/arrow/compute/special/if_else_special.cc @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -#include "arrow/compute/special/conditional_internal.h" +#include "arrow/compute/special/conditional_special_internal.h" namespace arrow::compute { diff --git a/cpp/src/arrow/compute/special/if_else_special_test.cc b/cpp/src/arrow/compute/special/if_else_special_test.cc index 572db18b7782..2ff9144b854c 100644 --- a/cpp/src/arrow/compute/special/if_else_special_test.cc +++ b/cpp/src/arrow/compute/special/if_else_special_test.cc @@ -26,8 +26,6 @@ namespace arrow::compute { -namespace { - using internal::add; using internal::cast; using internal::ExpectBindsTo; @@ -281,8 +279,6 @@ Expression assert_sv_empty(Expression arg) { return call("assert_sv_empty", {std::move(arg)}); } -} // namespace - TEST(IfElseSpecialForm, Basic) { { ARROW_SCOPED_TRACE("if (b != 0) then a / b else b"); From 046dde558cc1d21a5e1c2226db724ff0233d7e04 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Thu, 9 Oct 2025 23:10:13 -0700 Subject: [PATCH 44/71] Refine --- .../special/conditional_special_test.cc | 126 ++++++++---------- 1 file changed, 59 insertions(+), 67 deletions(-) diff --git a/cpp/src/arrow/compute/special/conditional_special_test.cc b/cpp/src/arrow/compute/special/conditional_special_test.cc index c56faae08333..4c39f9dda3a5 100644 --- a/cpp/src/arrow/compute/special/conditional_special_test.cc +++ b/cpp/src/arrow/compute/special/conditional_special_test.cc @@ -455,7 +455,13 @@ TEST(ConditionalBodyMask, NextBranchMask) { } } -TEST(ConditionalExec, Basic) { +TEST(ConditionalSpecialExecutor, Basic) { + ConditionalSpecialExecutor executor({}, TypeHolder(utf8())); + EXPECT_EQ(executor.out_type().id(), Type::STRING); + EXPECT_EQ(executor.options(), nullptr); +} + +TEST(ConditionalSpecialExecutor, Execute) { auto schm = schema({field("", boolean()), field("", boolean()), field("", utf8()), field("", utf8()), field("", utf8())}); @@ -472,137 +478,123 @@ TEST(ConditionalExec, Basic) { [true, true, "a3", "b3", "c3"]])"); { - ASSERT_OK_AND_ASSIGN( - auto result, ConditionalExec({Branch{kFalseLiteral, sa}, Branch{kTrueLiteral, sb}, - Branch{kTrueLiteral, sc}}, - TypeHolder(utf8())) - .Execute(batch, default_exec_context())); + ConditionalSpecialExecutor executor( + {Branch{kFalseLiteral, sa}, Branch{kTrueLiteral, sb}, Branch{kTrueLiteral, sc}}, + TypeHolder(utf8())); + ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["b0", "b1", "b2", "b3"])")), result); } { - ASSERT_OK_AND_ASSIGN( - auto result, ConditionalExec({Branch{kFalseLiteral, sa}, Branch{kNullLiteral, sb}, - Branch{kTrueLiteral, sc}}, - TypeHolder(utf8())) - .Execute(batch, default_exec_context())); + ConditionalSpecialExecutor executor( + {Branch{kFalseLiteral, sa}, Branch{kNullLiteral, sb}, Branch{kTrueLiteral, sc}}, + TypeHolder(utf8())); + ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"([null, null, null, null])")), result); } { - ASSERT_OK_AND_ASSIGN(auto result, - ConditionalExec({Branch{kFalseLiteral, sa}, Branch{b0, sb}, - Branch{kTrueLiteral, sc}}, - TypeHolder(utf8())) - .Execute(batch, default_exec_context())); + ConditionalSpecialExecutor executor( + {Branch{kFalseLiteral, sa}, Branch{b0, sb}, Branch{kTrueLiteral, sc}}, + TypeHolder(utf8())); + ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["b0", "c1", null, "b3"])")), result); } { - ASSERT_OK_AND_ASSIGN(auto result, - ConditionalExec({Branch{kFalseLiteral, sa}, Branch{b1, sb}, - Branch{kTrueLiteral, sc}}, - TypeHolder(utf8())) - .Execute(batch, default_exec_context())); + ConditionalSpecialExecutor executor( + {Branch{kFalseLiteral, sa}, Branch{b1, sb}, Branch{kTrueLiteral, sc}}, + TypeHolder(utf8())); + ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["b0", "c1", "b2", "b3"])")), result); } { - ASSERT_OK_AND_ASSIGN( - auto result, ConditionalExec({Branch{kTrueLiteral, sa}, Branch{kTrueLiteral, sb}, - Branch{kTrueLiteral, sc}}, - TypeHolder(utf8())) - .Execute(batch, default_exec_context())); + ConditionalSpecialExecutor executor( + {Branch{kTrueLiteral, sa}, Branch{kTrueLiteral, sb}, Branch{kTrueLiteral, sc}}, + TypeHolder(utf8())); + ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "a1", "a2", "a3"])")), result); } { - ASSERT_OK_AND_ASSIGN( - auto result, ConditionalExec({Branch{kNullLiteral, sa}, Branch{kTrueLiteral, sb}, - Branch{kTrueLiteral, sc}}, - TypeHolder(utf8())) - .Execute(batch, default_exec_context())); + ConditionalSpecialExecutor executor( + {Branch{kNullLiteral, sa}, Branch{kTrueLiteral, sb}, Branch{kTrueLiteral, sc}}, + TypeHolder(utf8())); + ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"([null, null, null, null])")), result); } { - ASSERT_OK_AND_ASSIGN(auto result, - ConditionalExec({Branch{b0, sa}, Branch{kTrueLiteral, sb}, - Branch{kTrueLiteral, sc}}, - TypeHolder(utf8())) - .Execute(batch, default_exec_context())); + ConditionalSpecialExecutor executor( + {Branch{b0, sa}, Branch{kTrueLiteral, sb}, Branch{kTrueLiteral, sc}}, + TypeHolder(utf8())); + ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "b1", null, "a3"])")), result); } { - ASSERT_OK_AND_ASSIGN(auto result, - ConditionalExec({Branch{b0, sa}, Branch{kNullLiteral, sb}, - Branch{kTrueLiteral, sc}}, - TypeHolder(utf8())) - .Execute(batch, default_exec_context())); + ConditionalSpecialExecutor executor( + {Branch{b0, sa}, Branch{kNullLiteral, sb}, Branch{kTrueLiteral, sc}}, + TypeHolder(utf8())); + ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", null, null, "a3"])")), result); } { - ASSERT_OK_AND_ASSIGN(auto result, ConditionalExec({Branch{b0, sa}, Branch{b0, sb}, - Branch{kTrueLiteral, sc}}, - TypeHolder(utf8())) - .Execute(batch, default_exec_context())); + ConditionalSpecialExecutor executor( + {Branch{b0, sa}, Branch{b0, sb}, Branch{kTrueLiteral, sc}}, TypeHolder(utf8())); + ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "c1", null, "a3"])")), result); } { - ASSERT_OK_AND_ASSIGN(auto result, ConditionalExec({Branch{b0, sa}, Branch{b1, sb}, - Branch{kTrueLiteral, sc}}, - TypeHolder(utf8())) - .Execute(batch, default_exec_context())); + ConditionalSpecialExecutor executor( + {Branch{b0, sa}, Branch{b1, sb}, Branch{kTrueLiteral, sc}}, TypeHolder(utf8())); + ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "c1", null, "a3"])")), result); } { - ASSERT_OK_AND_ASSIGN(auto result, - ConditionalExec({Branch{b1, sa}, Branch{kTrueLiteral, sb}, - Branch{kTrueLiteral, sc}}, - TypeHolder(utf8())) - .Execute(batch, default_exec_context())); + ConditionalSpecialExecutor executor( + {Branch{b1, sa}, Branch{kTrueLiteral, sb}, Branch{kTrueLiteral, sc}}, + TypeHolder(utf8())); + ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "b1", "a2", "a3"])")), result); } { - ASSERT_OK_AND_ASSIGN(auto result, - ConditionalExec({Branch{b1, sa}, Branch{kNullLiteral, sb}, - Branch{kTrueLiteral, sc}}, - TypeHolder(utf8())) - .Execute(batch, default_exec_context())); + ConditionalSpecialExecutor executor( + {Branch{b1, sa}, Branch{kNullLiteral, sb}, Branch{kTrueLiteral, sc}}, + TypeHolder(utf8())); + ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", null, "a2", "a3"])")), result); } { - ASSERT_OK_AND_ASSIGN(auto result, ConditionalExec({Branch{b1, sa}, Branch{b0, sb}, - Branch{kTrueLiteral, sc}}, - TypeHolder(utf8())) - .Execute(batch, default_exec_context())); + ConditionalSpecialExecutor executor( + {Branch{b1, sa}, Branch{b0, sb}, Branch{kTrueLiteral, sc}}, TypeHolder(utf8())); + ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "c1", "a2", "a3"])")), result); } { - ASSERT_OK_AND_ASSIGN(auto result, ConditionalExec({Branch{b1, sa}, Branch{b1, sb}, - Branch{kTrueLiteral, sc}}, - TypeHolder(utf8())) - .Execute(batch, default_exec_context())); + ConditionalSpecialExecutor executor( + {Branch{b1, sa}, Branch{b1, sb}, Branch{kTrueLiteral, sc}}, TypeHolder(utf8())); + ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "c1", "a2", "a3"])")), result); } From 84f643b231c1fba3db7948fa55f44d3f86e7be60 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Thu, 9 Oct 2025 23:17:53 -0700 Subject: [PATCH 45/71] Refine --- .../special/conditional_special_test.cc | 36 ++++++++----------- 1 file changed, 15 insertions(+), 21 deletions(-) diff --git a/cpp/src/arrow/compute/special/conditional_special_test.cc b/cpp/src/arrow/compute/special/conditional_special_test.cc index 4c39f9dda3a5..5361dc9d3107 100644 --- a/cpp/src/arrow/compute/special/conditional_special_test.cc +++ b/cpp/src/arrow/compute/special/conditional_special_test.cc @@ -456,7 +456,7 @@ TEST(ConditionalBodyMask, NextBranchMask) { } TEST(ConditionalSpecialExecutor, Basic) { - ConditionalSpecialExecutor executor({}, TypeHolder(utf8())); + ConditionalSpecialExecutor executor({}, utf8()); EXPECT_EQ(executor.out_type().id(), Type::STRING); EXPECT_EQ(executor.options(), nullptr); } @@ -480,7 +480,7 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( {Branch{kFalseLiteral, sa}, Branch{kTrueLiteral, sb}, Branch{kTrueLiteral, sc}}, - TypeHolder(utf8())); + utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["b0", "b1", "b2", "b3"])")), result); @@ -489,7 +489,7 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( {Branch{kFalseLiteral, sa}, Branch{kNullLiteral, sb}, Branch{kTrueLiteral, sc}}, - TypeHolder(utf8())); + utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"([null, null, null, null])")), result); @@ -497,8 +497,7 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{kFalseLiteral, sa}, Branch{b0, sb}, Branch{kTrueLiteral, sc}}, - TypeHolder(utf8())); + {Branch{kFalseLiteral, sa}, Branch{b0, sb}, Branch{kTrueLiteral, sc}}, utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["b0", "c1", null, "b3"])")), result); @@ -506,8 +505,7 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{kFalseLiteral, sa}, Branch{b1, sb}, Branch{kTrueLiteral, sc}}, - TypeHolder(utf8())); + {Branch{kFalseLiteral, sa}, Branch{b1, sb}, Branch{kTrueLiteral, sc}}, utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["b0", "c1", "b2", "b3"])")), result); @@ -516,7 +514,7 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( {Branch{kTrueLiteral, sa}, Branch{kTrueLiteral, sb}, Branch{kTrueLiteral, sc}}, - TypeHolder(utf8())); + utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "a1", "a2", "a3"])")), result); @@ -525,7 +523,7 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( {Branch{kNullLiteral, sa}, Branch{kTrueLiteral, sb}, Branch{kTrueLiteral, sc}}, - TypeHolder(utf8())); + utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"([null, null, null, null])")), result); @@ -533,8 +531,7 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{b0, sa}, Branch{kTrueLiteral, sb}, Branch{kTrueLiteral, sc}}, - TypeHolder(utf8())); + {Branch{b0, sa}, Branch{kTrueLiteral, sb}, Branch{kTrueLiteral, sc}}, utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "b1", null, "a3"])")), result); @@ -542,8 +539,7 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{b0, sa}, Branch{kNullLiteral, sb}, Branch{kTrueLiteral, sc}}, - TypeHolder(utf8())); + {Branch{b0, sa}, Branch{kNullLiteral, sb}, Branch{kTrueLiteral, sc}}, utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", null, null, "a3"])")), result); @@ -551,7 +547,7 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{b0, sa}, Branch{b0, sb}, Branch{kTrueLiteral, sc}}, TypeHolder(utf8())); + {Branch{b0, sa}, Branch{b0, sb}, Branch{kTrueLiteral, sc}}, utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "c1", null, "a3"])")), result); @@ -559,7 +555,7 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{b0, sa}, Branch{b1, sb}, Branch{kTrueLiteral, sc}}, TypeHolder(utf8())); + {Branch{b0, sa}, Branch{b1, sb}, Branch{kTrueLiteral, sc}}, utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "c1", null, "a3"])")), result); @@ -567,8 +563,7 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{b1, sa}, Branch{kTrueLiteral, sb}, Branch{kTrueLiteral, sc}}, - TypeHolder(utf8())); + {Branch{b1, sa}, Branch{kTrueLiteral, sb}, Branch{kTrueLiteral, sc}}, utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "b1", "a2", "a3"])")), result); @@ -576,8 +571,7 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{b1, sa}, Branch{kNullLiteral, sb}, Branch{kTrueLiteral, sc}}, - TypeHolder(utf8())); + {Branch{b1, sa}, Branch{kNullLiteral, sb}, Branch{kTrueLiteral, sc}}, utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", null, "a2", "a3"])")), result); @@ -585,7 +579,7 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{b1, sa}, Branch{b0, sb}, Branch{kTrueLiteral, sc}}, TypeHolder(utf8())); + {Branch{b1, sa}, Branch{b0, sb}, Branch{kTrueLiteral, sc}}, utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "c1", "a2", "a3"])")), result); @@ -593,7 +587,7 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{b1, sa}, Branch{b1, sb}, Branch{kTrueLiteral, sc}}, TypeHolder(utf8())); + {Branch{b1, sa}, Branch{b1, sb}, Branch{kTrueLiteral, sc}}, utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "c1", "a2", "a3"])")), result); From f9f30318adb74771ac756c58ac47c87de8661303 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Fri, 10 Oct 2025 02:02:25 -0700 Subject: [PATCH 46/71] Pre-specified selection vector tests --- cpp/src/arrow/compute/exec_internal.h | 20 ++++ .../compute/special/conditional_special.cc | 13 ++- .../special/conditional_special_test.cc | 95 +++++++++++++++++++ 3 files changed, 124 insertions(+), 4 deletions(-) diff --git a/cpp/src/arrow/compute/exec_internal.h b/cpp/src/arrow/compute/exec_internal.h index 1c4a49f8e484..de81f3fb9767 100644 --- a/cpp/src/arrow/compute/exec_internal.h +++ b/cpp/src/arrow/compute/exec_internal.h @@ -29,6 +29,7 @@ #include "arrow/compute/kernel.h" #include "arrow/status.h" #include "arrow/util/visibility.h" +#include "arrow/visit_data_inline.h" namespace arrow { namespace compute { @@ -171,6 +172,25 @@ Status PropagateNulls(KernelContext* ctx, const ExecSpan& batch, ArrayData* out) ARROW_EXPORT void PropagateNullsSpans(const ExecSpan& batch, ArraySpan* out); +template +typename ::arrow::internal::call_traits::enable_if_return::type +VisitSelectionVectorSpanInline(const SelectionVectorSpan& selection, + OnSelectionFn&& on_selection) { + for (int64_t i = 0; i < selection.length(); ++i) { + RETURN_NOT_OK(on_selection(selection[i])); + } + return Status::OK(); +} + +template +typename ::arrow::internal::call_traits::enable_if_return::type +VisitSelectionVectorSpanInline(const SelectionVectorSpan& selection, + OnSelectionFn&& on_selection) { + for (int64_t i = 0; i < selection.length(); ++i) { + on_selection(selection[i]); + } +} + } // namespace detail } // namespace compute } // namespace arrow diff --git a/cpp/src/arrow/compute/special/conditional_special.cc b/cpp/src/arrow/compute/special/conditional_special.cc index 41fde622e09a..db70524fffb0 100644 --- a/cpp/src/arrow/compute/special/conditional_special.cc +++ b/cpp/src/arrow/compute/special/conditional_special.cc @@ -302,12 +302,17 @@ Result ConditionalExec::MultiplexResults(const ExecBatch& input, } if (results.size() == 1) { - if (const auto& result = results.body_results()[0]; - results.selection_vectors()[0] == nullptr || - results.selection_vectors()[0]->length() == - (input.selection_vector ? input.selection_vector->length() : input.length)) { + const auto& result = results.body_results()[0]; + if (results.selection_vectors()[0] == nullptr) { return result; } + if (input.selection_vector == nullptr) { + DCHECK_NE(results.selection_vectors()[0]->length(), input.length); + } else { + if (results.selection_vectors()[0]->length() == input.selection_vector->length()) { + return result; + } + } } std::vector choose_args; diff --git a/cpp/src/arrow/compute/special/conditional_special_test.cc b/cpp/src/arrow/compute/special/conditional_special_test.cc index 5361dc9d3107..cf98471bdbcc 100644 --- a/cpp/src/arrow/compute/special/conditional_special_test.cc +++ b/cpp/src/arrow/compute/special/conditional_special_test.cc @@ -19,6 +19,7 @@ #include +#include "arrow/compute/exec_internal.h" #include "arrow/compute/test_util_internal.h" #include "arrow/testing/builder.h" #include "arrow/testing/generator.h" @@ -97,6 +98,10 @@ const auto kNullLiteral = literal(kNullScalar); const auto kTrueLiteral = literal(true); const auto kFalseLiteral = literal(false); +Expression nullify_selected(Expression expr) { + return call("nullify_selected", {std::move(expr)}); +} + } // namespace TEST(BranchMask, FromSelectionVector) { @@ -594,4 +599,94 @@ TEST(ConditionalSpecialExecutor, Execute) { } } +class ConditionalSpecialExecutorTest : public ::testing::Test { + protected: + static void SetUpTestSuite() { ASSERT_OK(RegisterNullifySelectedFunction()); } + + protected: + static Status NullifySelectedExec(KernelContext* ctx, const ExecSpan& span, + ExecResult* out) { + return Status::NotImplemented("Not implemented"); + } + + static Status NullifySelectedSelectiveExec(KernelContext* ctx, const ExecSpan& batch, + const SelectionVectorSpan& selection, + ExecResult* out) { + DCHECK_EQ(1, batch.num_values()); + int value_size = batch[0].type()->byte_width(); + const ArraySpan& arg0 = batch[0].array; + ArraySpan* out_arr = out->array_span_mutable(); + uint8_t* dst_validity = out_arr->buffers[0].data; + int64_t dst_validity_offset = out_arr->offset; + uint8_t* dst = out_arr->buffers[1].data + out_arr->offset * value_size; + uint8_t* src_validity = arg0.buffers[0].data; + int64_t src_validity_offset = arg0.offset; + const uint8_t* src = arg0.buffers[1].data + arg0.offset * value_size; + bit_util::SetBitmap(dst_validity, dst_validity_offset, batch.length); + std::memcpy(dst, src, batch.length * value_size); + compute::detail::VisitSelectionVectorSpanInline(selection, [&](int64_t i) { + DCHECK(!src_validity || bit_util::GetBit(src_validity, src_validity_offset + i)); + bit_util::ClearBit(dst_validity, dst_validity_offset + i); + }); + return Status::OK(); + } + + static Status RegisterNullifySelectedFunction() { + auto registry = GetFunctionRegistry(); + + auto func = std::make_shared(std::move("nullify_selected"), + Arity::Unary(), FunctionDoc::Empty()); + ScalarKernel kernel({InputType(int32())}, internal::FirstType, NullifySelectedExec, + NullifySelectedSelectiveExec); + kernel.can_write_into_slices = false; + kernel.null_handling = NullHandling::COMPUTED_PREALLOCATE; + kernel.mem_allocation = MemAllocation::PREALLOCATE; + RETURN_NOT_OK(func->AddKernel(kernel)); + RETURN_NOT_OK(registry->AddFunction(std::move(func))); + + return Status::OK(); + } +}; + +TEST_F(ConditionalSpecialExecutorTest, ExecuteWithSelection) { + auto schm = schema({field("a", int32())}); + ASSERT_OK_AND_ASSIGN(auto body, nullify_selected(field_ref("a")).Bind(*schm)); + ConditionalSpecialExecutor executor({Branch{kTrueLiteral, body}}, int32()); + auto batch = ExecBatchFromJSON({int32()}, R"([[10], [11], [12], [13]])"); + + { + auto batch_with_selection = batch; + batch_with_selection.selection_vector = SelectionVectorFromJSON("[]"); + ASSERT_OK_AND_ASSIGN(auto result, + executor.Execute(batch_with_selection, default_exec_context())); + // Empty selection short-circuits to all-null output - the kernel is not invoked. + AssertDatumsEqual(Datum(ArrayFromJSON(int32(), R"([null, null, null, null])")), + result); + } + + { + auto batch_with_selection = batch; + batch_with_selection.selection_vector = SelectionVectorFromJSON("[0]"); + ASSERT_OK_AND_ASSIGN(auto result, + executor.Execute(batch_with_selection, default_exec_context())); + AssertDatumsEqual(Datum(ArrayFromJSON(int32(), R"([null, 11, 12, 13])")), result); + } + + { + auto batch_with_selection = batch; + batch_with_selection.selection_vector = SelectionVectorFromJSON("[0, 1, 2]"); + ASSERT_OK_AND_ASSIGN(auto result, + executor.Execute(batch_with_selection, default_exec_context())); + AssertDatumsEqual(Datum(ArrayFromJSON(int32(), R"([null, null, null, 13])")), result); + } + + { + auto batch_with_selection = batch; + batch_with_selection.selection_vector = SelectionVectorFromJSON("[0, 1, 2, 3]"); + // Full selection short-circuits to non-selective execution. + ASSERT_RAISES(NotImplemented, + executor.Execute(batch_with_selection, default_exec_context())); + } +} + } // namespace arrow::compute::internal From 148c5c879e465d35b2b8525bc89f3a0c632faa03 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Fri, 10 Oct 2025 13:55:33 -0700 Subject: [PATCH 47/71] Fine-grained selection vector check --- .../special/conditional_special_test.cc | 119 +++++++++---- cpp/src/arrow/compute/test_util_internal.cc | 164 ++++++++++++++++++ cpp/src/arrow/compute/test_util_internal.h | 8 + 3 files changed, 261 insertions(+), 30 deletions(-) diff --git a/cpp/src/arrow/compute/special/conditional_special_test.cc b/cpp/src/arrow/compute/special/conditional_special_test.cc index cf98471bdbcc..835e46882e8c 100644 --- a/cpp/src/arrow/compute/special/conditional_special_test.cc +++ b/cpp/src/arrow/compute/special/conditional_special_test.cc @@ -51,22 +51,6 @@ void CheckNextBranchMask(const std::shared_ptr& body_mask) { CheckBranchMask(branch_mask); } -void AssertSelectionVectorsEqual(const std::shared_ptr& expected, - const std::shared_ptr& actual) { - if (expected == nullptr) { - EXPECT_EQ(actual, nullptr); - return; - } - - if (actual == nullptr) { - EXPECT_EQ(expected, nullptr); - return; - } - - ASSERT_EQ(expected->length(), actual->length()); - AssertArraysEqual(*MakeArray(expected->data()), *MakeArray(actual->data())); -} - template void CheckMakeBodyMaskAndSelection( const std::shared_ptr& branch_mask, const Datum& datum, @@ -475,6 +459,23 @@ TEST(ConditionalSpecialExecutor, Execute) { ASSERT_OK_AND_ASSIGN(auto sa, field_ref(2).Bind(*schm)); ASSERT_OK_AND_ASSIGN(auto sb, field_ref(3).Bind(*schm)); ASSERT_OK_AND_ASSIGN(auto sc, field_ref(4).Bind(*schm)); + auto unreachable = [&](Expression argument) -> Expression { + EXPECT_OK_AND_ASSIGN(auto bound, + unreachable_special(std::move(argument)).Bind(*schm)); + return bound; + }; + auto assert_empty_selection = [&](Expression argument) -> Expression { + EXPECT_OK_AND_ASSIGN(auto bound, + assert_empty_selection_special(std::move(argument)).Bind(*schm)); + return bound; + }; + auto assert_selection_eq = + [&](Expression argument, std::shared_ptr selection) -> Expression { + EXPECT_OK_AND_ASSIGN( + auto bound, assert_selection_eq_special(std::move(argument), std::move(selection)) + .Bind(*schm)); + return bound; + }; auto batch = ExecBatchFromJSON({boolean(), boolean(), utf8(), utf8(), utf8()}, R"([[true, true, "a0", "b0", "c0"], @@ -484,7 +485,9 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{kFalseLiteral, sa}, Branch{kTrueLiteral, sb}, Branch{kTrueLiteral, sc}}, + {Branch{assert_empty_selection(kFalseLiteral), unreachable(sa)}, + Branch{assert_empty_selection(kTrueLiteral), assert_empty_selection(sb)}, + Branch{unreachable(kTrueLiteral), unreachable(sc)}}, utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["b0", "b1", "b2", "b3"])")), @@ -493,7 +496,9 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{kFalseLiteral, sa}, Branch{kNullLiteral, sb}, Branch{kTrueLiteral, sc}}, + {Branch{assert_empty_selection(kFalseLiteral), unreachable(sa)}, + Branch{assert_empty_selection(kNullLiteral), unreachable(sb)}, + Branch{unreachable(kTrueLiteral), unreachable(sc)}}, utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"([null, null, null, null])")), @@ -502,7 +507,12 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{kFalseLiteral, sa}, Branch{b0, sb}, Branch{kTrueLiteral, sc}}, utf8()); + {Branch{assert_empty_selection(kFalseLiteral), unreachable(sa)}, + Branch{assert_empty_selection(b0), + assert_selection_eq(sb, SelectionVectorFromJSON("[0, 3]"))}, + Branch{assert_selection_eq(kTrueLiteral, SelectionVectorFromJSON("[1]")), + assert_selection_eq(sc, SelectionVectorFromJSON("[1]"))}}, + utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["b0", "c1", null, "b3"])")), result); @@ -510,7 +520,12 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{kFalseLiteral, sa}, Branch{b1, sb}, Branch{kTrueLiteral, sc}}, utf8()); + {Branch{assert_empty_selection(kFalseLiteral), unreachable(sa)}, + Branch{assert_empty_selection(b1), + assert_selection_eq(sb, SelectionVectorFromJSON("[0, 2, 3]"))}, + Branch{assert_selection_eq(kTrueLiteral, SelectionVectorFromJSON("[1]")), + assert_selection_eq(sc, SelectionVectorFromJSON("[1]"))}}, + utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["b0", "c1", "b2", "b3"])")), result); @@ -518,7 +533,9 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{kTrueLiteral, sa}, Branch{kTrueLiteral, sb}, Branch{kTrueLiteral, sc}}, + {Branch{assert_empty_selection(kTrueLiteral), assert_empty_selection(sa)}, + Branch{unreachable(kTrueLiteral), unreachable(sb)}, + Branch{unreachable(kTrueLiteral), unreachable(sc)}}, utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "a1", "a2", "a3"])")), @@ -527,7 +544,9 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{kNullLiteral, sa}, Branch{kTrueLiteral, sb}, Branch{kTrueLiteral, sc}}, + {Branch{assert_empty_selection(kNullLiteral), unreachable(sa)}, + Branch{unreachable(kTrueLiteral), unreachable(sb)}, + Branch{unreachable(kTrueLiteral), unreachable(sc)}}, utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"([null, null, null, null])")), @@ -536,7 +555,12 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{b0, sa}, Branch{kTrueLiteral, sb}, Branch{kTrueLiteral, sc}}, utf8()); + {Branch{assert_empty_selection(b0), + assert_selection_eq(sa, SelectionVectorFromJSON("[0, 3]"))}, + Branch{assert_selection_eq(kTrueLiteral, SelectionVectorFromJSON("[1]")), + assert_selection_eq(sb, SelectionVectorFromJSON("[1]"))}, + Branch{unreachable(kTrueLiteral), unreachable(sc)}}, + utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "b1", null, "a3"])")), result); @@ -544,7 +568,12 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{b0, sa}, Branch{kNullLiteral, sb}, Branch{kTrueLiteral, sc}}, utf8()); + {Branch{assert_empty_selection(b0), + assert_selection_eq(sa, SelectionVectorFromJSON("[0, 3]"))}, + Branch{assert_selection_eq(kNullLiteral, SelectionVectorFromJSON("[1]")), + unreachable(sb)}, + Branch{unreachable(kTrueLiteral), unreachable(sc)}}, + utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", null, null, "a3"])")), result); @@ -552,7 +581,12 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{b0, sa}, Branch{b0, sb}, Branch{kTrueLiteral, sc}}, utf8()); + {Branch{assert_empty_selection(b0), + assert_selection_eq(sa, SelectionVectorFromJSON("[0, 3]"))}, + Branch{assert_selection_eq(b0, SelectionVectorFromJSON("[1]")), unreachable(sb)}, + Branch{assert_selection_eq(kTrueLiteral, SelectionVectorFromJSON("[1]")), + assert_selection_eq(sc, SelectionVectorFromJSON("[1]"))}}, + utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "c1", null, "a3"])")), result); @@ -560,7 +594,12 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{b0, sa}, Branch{b1, sb}, Branch{kTrueLiteral, sc}}, utf8()); + {Branch{assert_empty_selection(b0), + assert_selection_eq(sa, SelectionVectorFromJSON("[0, 3]"))}, + Branch{assert_selection_eq(b1, SelectionVectorFromJSON("[1]")), unreachable(sb)}, + Branch{assert_selection_eq(kTrueLiteral, SelectionVectorFromJSON("[1]")), + assert_selection_eq(sc, SelectionVectorFromJSON("[1]"))}}, + utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "c1", null, "a3"])")), result); @@ -568,7 +607,12 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{b1, sa}, Branch{kTrueLiteral, sb}, Branch{kTrueLiteral, sc}}, utf8()); + {Branch{assert_empty_selection(b1), + assert_selection_eq(sa, SelectionVectorFromJSON("[0, 2, 3]"))}, + Branch{assert_selection_eq(kTrueLiteral, SelectionVectorFromJSON("[1]")), + assert_selection_eq(sb, SelectionVectorFromJSON("[1]"))}, + Branch{unreachable(kTrueLiteral), unreachable(sc)}}, + utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "b1", "a2", "a3"])")), result); @@ -576,7 +620,12 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{b1, sa}, Branch{kNullLiteral, sb}, Branch{kTrueLiteral, sc}}, utf8()); + {Branch{assert_empty_selection(b1), + assert_selection_eq(sa, SelectionVectorFromJSON("[0, 2, 3]"))}, + Branch{assert_selection_eq(kNullLiteral, SelectionVectorFromJSON("[1]")), + unreachable(sb)}, + Branch{unreachable(kTrueLiteral), unreachable(sc)}}, + utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", null, "a2", "a3"])")), result); @@ -584,7 +633,12 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{b1, sa}, Branch{b0, sb}, Branch{kTrueLiteral, sc}}, utf8()); + {Branch{assert_empty_selection(b1), + assert_selection_eq(sa, SelectionVectorFromJSON("[0, 2, 3]"))}, + Branch{assert_selection_eq(b0, SelectionVectorFromJSON("[1]")), unreachable(sb)}, + Branch{assert_selection_eq(kTrueLiteral, SelectionVectorFromJSON("[1]")), + assert_selection_eq(sc, SelectionVectorFromJSON("[1]"))}}, + utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "c1", "a2", "a3"])")), result); @@ -592,7 +646,12 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{b1, sa}, Branch{b1, sb}, Branch{kTrueLiteral, sc}}, utf8()); + {Branch{assert_empty_selection(b1), + assert_selection_eq(sa, SelectionVectorFromJSON("[0, 2, 3]"))}, + Branch{assert_selection_eq(b1, SelectionVectorFromJSON("[1]")), unreachable(sb)}, + Branch{assert_selection_eq(kTrueLiteral, SelectionVectorFromJSON("[1]")), + assert_selection_eq(sc, SelectionVectorFromJSON("[1]"))}}, + utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "c1", "a2", "a3"])")), result); diff --git a/cpp/src/arrow/compute/test_util_internal.cc b/cpp/src/arrow/compute/test_util_internal.cc index 6407799399c2..af3c389ba6fe 100644 --- a/cpp/src/arrow/compute/test_util_internal.cc +++ b/cpp/src/arrow/compute/test_util_internal.cc @@ -20,6 +20,7 @@ #include "arrow/array/array_base.h" #include "arrow/array/validate.h" #include "arrow/chunked_array.h" +#include "arrow/compute/special_form.h" #include "arrow/datum.h" #include "arrow/record_batch.h" #include "arrow/scalar.h" @@ -132,4 +133,167 @@ std::shared_ptr MakeSelectionVectorTo(int64_t length) { return std::make_shared(*arr); } +void AssertSelectionVectorsEqual(const std::shared_ptr& expected, + const std::shared_ptr& actual) { + if (expected == nullptr) { + EXPECT_EQ(actual, nullptr); + return; + } + + if (actual == nullptr) { + EXPECT_EQ(expected, nullptr); + return; + } + + ASSERT_EQ(expected->length(), actual->length()); + AssertArraysEqual(*MakeArray(expected->data()), *MakeArray(actual->data())); +} + +namespace { + +class TrivialSpecialExecutor : public SpecialExecutor { + public: + explicit TrivialSpecialExecutor(Expression argument) + : SpecialExecutor(argument.type()), argument_(std::move(argument)) {} + + Result Execute(const ExecBatch& input, + ExecContext* exec_context) const override { + RETURN_NOT_OK(PreExecute(input, exec_context)); + return ExecuteScalarExpression(argument_, input, exec_context); + } + + protected: + virtual Status PreExecute(const ExecBatch& input, ExecContext* exec_context) const { + return Status::OK(); + } + + protected: + Expression argument_; +}; + +class UnreachableSpecialExecutor : public TrivialSpecialExecutor { + public: + explicit UnreachableSpecialExecutor(Expression argument) + : TrivialSpecialExecutor(std::move(argument)) {} + + protected: + Status PreExecute(const ExecBatch& input, ExecContext* exec_context) const override { + return Status::Invalid("Unreachable"); + } +}; + +class UnreachableSpecialForm : public SpecialForm { + public: + UnreachableSpecialForm() : SpecialForm("unreachable") {} + + protected: + Result> Bind( + std::vector& arguments, std::shared_ptr options, + ExecContext* exec_context) const override { + DCHECK_EQ(arguments.size(), 1); + return std::make_unique(arguments[0]); + } +}; + +class AssertEmptySelectionSpecialExecutor : public TrivialSpecialExecutor { + public: + explicit AssertEmptySelectionSpecialExecutor(Expression argument) + : TrivialSpecialExecutor(std::move(argument)) {} + + protected: + Status PreExecute(const ExecBatch& input, ExecContext* exec_context) const override { + if (input.selection_vector) { + return Status::Invalid("There shouldn't be a selection vector"); + } + return Status::OK(); + } +}; + +class AssertEmptySelectionSpecialForm : public SpecialForm { + public: + AssertEmptySelectionSpecialForm() : SpecialForm("assert_empty_selection") {} + + protected: + Result> Bind( + std::vector& arguments, std::shared_ptr options, + ExecContext* exec_context) const override { + DCHECK_EQ(arguments.size(), 1); + return std::make_unique(arguments[0]); + } +}; + +class AssertSelectionEqualSpecialExecutor : public TrivialSpecialExecutor { + public: + explicit AssertSelectionEqualSpecialExecutor(Expression argument, + std::shared_ptr expected) + : TrivialSpecialExecutor(std::move(argument)), expected_(std::move(expected)) { + DCHECK_NE(expected_, nullptr); + } + + protected: + Status PreExecute(const ExecBatch& input, ExecContext* exec_context) const override { + if (input.selection_vector == nullptr) { + return Status::Invalid("There should be a selection vector"); + } + if (!SelectionVectorsEqual(expected_, input.selection_vector)) { + return Status::Invalid("Selection vector does not match expected"); + } + return Status::OK(); + } + + private: + bool SelectionVectorsEqual(const std::shared_ptr& left, + const std::shared_ptr& right) const { + DCHECK_NE(left, nullptr); + DCHECK_NE(right, nullptr); + return MakeArray(left->data())->Equals(MakeArray(right->data())); + } + + std::shared_ptr expected_; +}; + +class AssertSelectionEqualSpecialForm : public SpecialForm { + public: + explicit AssertSelectionEqualSpecialForm(std::shared_ptr expected) + : SpecialForm("assert_selection_equal"), expected_(std::move(expected)) { + DCHECK_NE(expected_, nullptr); + } + + protected: + Result> Bind( + std::vector& arguments, std::shared_ptr options, + ExecContext* exec_context) const override { + DCHECK_EQ(arguments.size(), 1); + return std::make_unique(arguments[0], expected_); + } + + private: + std::shared_ptr expected_; +}; + +} // namespace + +Expression unreachable_special(Expression argument) { + Expression::Special special; + special.special_form = std::make_shared(); + special.arguments.push_back(std::move(argument)); + return Expression(std::move(special)); +} + +Expression assert_empty_selection_special(Expression argument) { + Expression::Special special; + special.special_form = std::make_shared(); + special.arguments.push_back(std::move(argument)); + return Expression(std::move(special)); +} + +Expression assert_selection_eq_special(Expression argument, + std::shared_ptr expected) { + Expression::Special special; + special.special_form = + std::make_shared(std::move(expected)); + special.arguments.push_back(std::move(argument)); + return Expression(std::move(special)); +} + } // namespace arrow::compute diff --git a/cpp/src/arrow/compute/test_util_internal.h b/cpp/src/arrow/compute/test_util_internal.h index b7f4e1261993..9d4fbc283026 100644 --- a/cpp/src/arrow/compute/test_util_internal.h +++ b/cpp/src/arrow/compute/test_util_internal.h @@ -43,4 +43,12 @@ std::shared_ptr SelectionVectorFromJSON(const std::string& json std::shared_ptr MakeSelectionVectorTo(int64_t length); +void AssertSelectionVectorsEqual(const std::shared_ptr& expected, + const std::shared_ptr& actual); + +Expression unreachable_special(Expression argument); +Expression assert_empty_selection_special(Expression argument); +Expression assert_selection_eq_special(Expression argument, + std::shared_ptr expected); + } // namespace arrow::compute From 509c9b9808134c451d239d2929c44b0ff7488809 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Fri, 10 Oct 2025 16:03:30 -0700 Subject: [PATCH 48/71] Refine --- .../special/conditional_special_test.cc | 181 ++++++++---------- cpp/src/arrow/compute/test_util_internal.cc | 4 +- cpp/src/arrow/compute/test_util_internal.h | 2 +- 3 files changed, 81 insertions(+), 106 deletions(-) diff --git a/cpp/src/arrow/compute/special/conditional_special_test.cc b/cpp/src/arrow/compute/special/conditional_special_test.cc index 835e46882e8c..ff08a368e0d1 100644 --- a/cpp/src/arrow/compute/special/conditional_special_test.cc +++ b/cpp/src/arrow/compute/special/conditional_special_test.cc @@ -82,8 +82,24 @@ const auto kNullLiteral = literal(kNullScalar); const auto kTrueLiteral = literal(true); const auto kFalseLiteral = literal(false); -Expression nullify_selected(Expression expr) { - return call("nullify_selected", {std::move(expr)}); +Expression bind_unreachable(Expression argument, const Schema& schm) { + EXPECT_OK_AND_ASSIGN(auto bound, unreachable_special(std::move(argument)).Bind(schm)); + return bound; +} + +Expression bind_assert_selection_empty(Expression argument, const Schema& schm) { + EXPECT_OK_AND_ASSIGN(auto bound, + assert_selection_empty_special(std::move(argument)).Bind(schm)); + return bound; +} + +Expression bind_assert_selection_equal(Expression argument, + std::shared_ptr selection, + const Schema& schm) { + EXPECT_OK_AND_ASSIGN( + auto bound, + assert_selection_eq_special(std::move(argument), std::move(selection)).Bind(schm)); + return bound; } } // namespace @@ -454,27 +470,20 @@ TEST(ConditionalSpecialExecutor, Execute) { auto schm = schema({field("", boolean()), field("", boolean()), field("", utf8()), field("", utf8()), field("", utf8())}); - ASSERT_OK_AND_ASSIGN(auto b0, field_ref(0).Bind(*schm)); - ASSERT_OK_AND_ASSIGN(auto b1, field_ref(1).Bind(*schm)); - ASSERT_OK_AND_ASSIGN(auto sa, field_ref(2).Bind(*schm)); - ASSERT_OK_AND_ASSIGN(auto sb, field_ref(3).Bind(*schm)); - ASSERT_OK_AND_ASSIGN(auto sc, field_ref(4).Bind(*schm)); + auto b0 = field_ref(0); + auto b1 = field_ref(1); + auto sa = field_ref(2); + auto sb = field_ref(3); + auto sc = field_ref(4); auto unreachable = [&](Expression argument) -> Expression { - EXPECT_OK_AND_ASSIGN(auto bound, - unreachable_special(std::move(argument)).Bind(*schm)); - return bound; + return bind_unreachable(std::move(argument), *schm); }; - auto assert_empty_selection = [&](Expression argument) -> Expression { - EXPECT_OK_AND_ASSIGN(auto bound, - assert_empty_selection_special(std::move(argument)).Bind(*schm)); - return bound; + auto assert_selection_empty = [&](Expression argument) -> Expression { + return bind_assert_selection_empty(std::move(argument), *schm); }; auto assert_selection_eq = [&](Expression argument, std::shared_ptr selection) -> Expression { - EXPECT_OK_AND_ASSIGN( - auto bound, assert_selection_eq_special(std::move(argument), std::move(selection)) - .Bind(*schm)); - return bound; + return bind_assert_selection_equal(std::move(argument), std::move(selection), *schm); }; auto batch = ExecBatchFromJSON({boolean(), boolean(), utf8(), utf8(), utf8()}, @@ -485,8 +494,8 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{assert_empty_selection(kFalseLiteral), unreachable(sa)}, - Branch{assert_empty_selection(kTrueLiteral), assert_empty_selection(sb)}, + {Branch{assert_selection_empty(kFalseLiteral), unreachable(sa)}, + Branch{assert_selection_empty(kTrueLiteral), assert_selection_empty(sb)}, Branch{unreachable(kTrueLiteral), unreachable(sc)}}, utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); @@ -496,8 +505,8 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{assert_empty_selection(kFalseLiteral), unreachable(sa)}, - Branch{assert_empty_selection(kNullLiteral), unreachable(sb)}, + {Branch{assert_selection_empty(kFalseLiteral), unreachable(sa)}, + Branch{assert_selection_empty(kNullLiteral), unreachable(sb)}, Branch{unreachable(kTrueLiteral), unreachable(sc)}}, utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); @@ -507,8 +516,8 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{assert_empty_selection(kFalseLiteral), unreachable(sa)}, - Branch{assert_empty_selection(b0), + {Branch{assert_selection_empty(kFalseLiteral), unreachable(sa)}, + Branch{assert_selection_empty(b0), assert_selection_eq(sb, SelectionVectorFromJSON("[0, 3]"))}, Branch{assert_selection_eq(kTrueLiteral, SelectionVectorFromJSON("[1]")), assert_selection_eq(sc, SelectionVectorFromJSON("[1]"))}}, @@ -520,8 +529,8 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{assert_empty_selection(kFalseLiteral), unreachable(sa)}, - Branch{assert_empty_selection(b1), + {Branch{assert_selection_empty(kFalseLiteral), unreachable(sa)}, + Branch{assert_selection_empty(b1), assert_selection_eq(sb, SelectionVectorFromJSON("[0, 2, 3]"))}, Branch{assert_selection_eq(kTrueLiteral, SelectionVectorFromJSON("[1]")), assert_selection_eq(sc, SelectionVectorFromJSON("[1]"))}}, @@ -533,7 +542,7 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{assert_empty_selection(kTrueLiteral), assert_empty_selection(sa)}, + {Branch{assert_selection_empty(kTrueLiteral), assert_selection_empty(sa)}, Branch{unreachable(kTrueLiteral), unreachable(sb)}, Branch{unreachable(kTrueLiteral), unreachable(sc)}}, utf8()); @@ -544,7 +553,7 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{assert_empty_selection(kNullLiteral), unreachable(sa)}, + {Branch{assert_selection_empty(kNullLiteral), unreachable(sa)}, Branch{unreachable(kTrueLiteral), unreachable(sb)}, Branch{unreachable(kTrueLiteral), unreachable(sc)}}, utf8()); @@ -555,7 +564,7 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{assert_empty_selection(b0), + {Branch{assert_selection_empty(b0), assert_selection_eq(sa, SelectionVectorFromJSON("[0, 3]"))}, Branch{assert_selection_eq(kTrueLiteral, SelectionVectorFromJSON("[1]")), assert_selection_eq(sb, SelectionVectorFromJSON("[1]"))}, @@ -568,7 +577,7 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{assert_empty_selection(b0), + {Branch{assert_selection_empty(b0), assert_selection_eq(sa, SelectionVectorFromJSON("[0, 3]"))}, Branch{assert_selection_eq(kNullLiteral, SelectionVectorFromJSON("[1]")), unreachable(sb)}, @@ -581,7 +590,7 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{assert_empty_selection(b0), + {Branch{assert_selection_empty(b0), assert_selection_eq(sa, SelectionVectorFromJSON("[0, 3]"))}, Branch{assert_selection_eq(b0, SelectionVectorFromJSON("[1]")), unreachable(sb)}, Branch{assert_selection_eq(kTrueLiteral, SelectionVectorFromJSON("[1]")), @@ -594,7 +603,7 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{assert_empty_selection(b0), + {Branch{assert_selection_empty(b0), assert_selection_eq(sa, SelectionVectorFromJSON("[0, 3]"))}, Branch{assert_selection_eq(b1, SelectionVectorFromJSON("[1]")), unreachable(sb)}, Branch{assert_selection_eq(kTrueLiteral, SelectionVectorFromJSON("[1]")), @@ -607,7 +616,7 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{assert_empty_selection(b1), + {Branch{assert_selection_empty(b1), assert_selection_eq(sa, SelectionVectorFromJSON("[0, 2, 3]"))}, Branch{assert_selection_eq(kTrueLiteral, SelectionVectorFromJSON("[1]")), assert_selection_eq(sb, SelectionVectorFromJSON("[1]"))}, @@ -620,7 +629,7 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{assert_empty_selection(b1), + {Branch{assert_selection_empty(b1), assert_selection_eq(sa, SelectionVectorFromJSON("[0, 2, 3]"))}, Branch{assert_selection_eq(kNullLiteral, SelectionVectorFromJSON("[1]")), unreachable(sb)}, @@ -633,7 +642,7 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{assert_empty_selection(b1), + {Branch{assert_selection_empty(b1), assert_selection_eq(sa, SelectionVectorFromJSON("[0, 2, 3]"))}, Branch{assert_selection_eq(b0, SelectionVectorFromJSON("[1]")), unreachable(sb)}, Branch{assert_selection_eq(kTrueLiteral, SelectionVectorFromJSON("[1]")), @@ -646,7 +655,7 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{assert_empty_selection(b1), + {Branch{assert_selection_empty(b1), assert_selection_eq(sa, SelectionVectorFromJSON("[0, 2, 3]"))}, Branch{assert_selection_eq(b1, SelectionVectorFromJSON("[1]")), unreachable(sb)}, Branch{assert_selection_eq(kTrueLiteral, SelectionVectorFromJSON("[1]")), @@ -658,93 +667,59 @@ TEST(ConditionalSpecialExecutor, Execute) { } } -class ConditionalSpecialExecutorTest : public ::testing::Test { - protected: - static void SetUpTestSuite() { ASSERT_OK(RegisterNullifySelectedFunction()); } - - protected: - static Status NullifySelectedExec(KernelContext* ctx, const ExecSpan& span, - ExecResult* out) { - return Status::NotImplemented("Not implemented"); - } - - static Status NullifySelectedSelectiveExec(KernelContext* ctx, const ExecSpan& batch, - const SelectionVectorSpan& selection, - ExecResult* out) { - DCHECK_EQ(1, batch.num_values()); - int value_size = batch[0].type()->byte_width(); - const ArraySpan& arg0 = batch[0].array; - ArraySpan* out_arr = out->array_span_mutable(); - uint8_t* dst_validity = out_arr->buffers[0].data; - int64_t dst_validity_offset = out_arr->offset; - uint8_t* dst = out_arr->buffers[1].data + out_arr->offset * value_size; - uint8_t* src_validity = arg0.buffers[0].data; - int64_t src_validity_offset = arg0.offset; - const uint8_t* src = arg0.buffers[1].data + arg0.offset * value_size; - bit_util::SetBitmap(dst_validity, dst_validity_offset, batch.length); - std::memcpy(dst, src, batch.length * value_size); - compute::detail::VisitSelectionVectorSpanInline(selection, [&](int64_t i) { - DCHECK(!src_validity || bit_util::GetBit(src_validity, src_validity_offset + i)); - bit_util::ClearBit(dst_validity, dst_validity_offset + i); - }); - return Status::OK(); - } - - static Status RegisterNullifySelectedFunction() { - auto registry = GetFunctionRegistry(); - - auto func = std::make_shared(std::move("nullify_selected"), - Arity::Unary(), FunctionDoc::Empty()); - ScalarKernel kernel({InputType(int32())}, internal::FirstType, NullifySelectedExec, - NullifySelectedSelectiveExec); - kernel.can_write_into_slices = false; - kernel.null_handling = NullHandling::COMPUTED_PREALLOCATE; - kernel.mem_allocation = MemAllocation::PREALLOCATE; - RETURN_NOT_OK(func->AddKernel(kernel)); - RETURN_NOT_OK(registry->AddFunction(std::move(func))); - - return Status::OK(); - } -}; - -TEST_F(ConditionalSpecialExecutorTest, ExecuteWithSelection) { +TEST(ConditionalSpecialExecutor, ExecuteWithSelection) { auto schm = schema({field("a", int32())}); - ASSERT_OK_AND_ASSIGN(auto body, nullify_selected(field_ref("a")).Bind(*schm)); - ConditionalSpecialExecutor executor({Branch{kTrueLiteral, body}}, int32()); + + auto a = field_ref("a"); + auto unreachable = [&](Expression argument) -> Expression { + return bind_unreachable(std::move(argument), *schm); + }; + auto assert_selection_empty = [&](Expression argument) -> Expression { + return bind_assert_selection_empty(std::move(argument), *schm); + }; + auto assert_selection_eq = + [&](Expression argument, std::shared_ptr selection) -> Expression { + return bind_assert_selection_equal(std::move(argument), std::move(selection), *schm); + }; + auto batch = ExecBatchFromJSON({int32()}, R"([[10], [11], [12], [13]])"); { + // Empty selection short-circuits to all-null output - the kernel is not invoked. + const auto selection = SelectionVectorFromJSON("[]"); auto batch_with_selection = batch; - batch_with_selection.selection_vector = SelectionVectorFromJSON("[]"); + batch_with_selection.selection_vector = selection; + ConditionalSpecialExecutor executor( + {Branch{unreachable(kTrueLiteral), unreachable(a)}}, int32()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch_with_selection, default_exec_context())); - // Empty selection short-circuits to all-null output - the kernel is not invoked. AssertDatumsEqual(Datum(ArrayFromJSON(int32(), R"([null, null, null, null])")), result); } - { + for (const auto& selection : + {SelectionVectorFromJSON("[0]"), SelectionVectorFromJSON("[1]"), + SelectionVectorFromJSON("[0, 1, 2]")}) { auto batch_with_selection = batch; - batch_with_selection.selection_vector = SelectionVectorFromJSON("[0]"); + batch_with_selection.selection_vector = selection; + ConditionalSpecialExecutor executor( + {Branch{assert_selection_eq(kTrueLiteral, selection), + assert_selection_eq(a, selection)}}, + int32()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch_with_selection, default_exec_context())); - AssertDatumsEqual(Datum(ArrayFromJSON(int32(), R"([null, 11, 12, 13])")), result); } { + // Full selection short-circuits to non-selective execution. + const auto selection = SelectionVectorFromJSON("[0, 1, 2, 3]"); auto batch_with_selection = batch; - batch_with_selection.selection_vector = SelectionVectorFromJSON("[0, 1, 2]"); + batch_with_selection.selection_vector = selection; + ConditionalSpecialExecutor executor( + {Branch{assert_selection_empty(kTrueLiteral), assert_selection_empty(a)}}, + int32()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch_with_selection, default_exec_context())); - AssertDatumsEqual(Datum(ArrayFromJSON(int32(), R"([null, null, null, 13])")), result); - } - - { - auto batch_with_selection = batch; - batch_with_selection.selection_vector = SelectionVectorFromJSON("[0, 1, 2, 3]"); - // Full selection short-circuits to non-selective execution. - ASSERT_RAISES(NotImplemented, - executor.Execute(batch_with_selection, default_exec_context())); } } diff --git a/cpp/src/arrow/compute/test_util_internal.cc b/cpp/src/arrow/compute/test_util_internal.cc index af3c389ba6fe..52775275dddf 100644 --- a/cpp/src/arrow/compute/test_util_internal.cc +++ b/cpp/src/arrow/compute/test_util_internal.cc @@ -211,7 +211,7 @@ class AssertEmptySelectionSpecialExecutor : public TrivialSpecialExecutor { class AssertEmptySelectionSpecialForm : public SpecialForm { public: - AssertEmptySelectionSpecialForm() : SpecialForm("assert_empty_selection") {} + AssertEmptySelectionSpecialForm() : SpecialForm("assert_selection_empty") {} protected: Result> Bind( @@ -280,7 +280,7 @@ Expression unreachable_special(Expression argument) { return Expression(std::move(special)); } -Expression assert_empty_selection_special(Expression argument) { +Expression assert_selection_empty_special(Expression argument) { Expression::Special special; special.special_form = std::make_shared(); special.arguments.push_back(std::move(argument)); diff --git a/cpp/src/arrow/compute/test_util_internal.h b/cpp/src/arrow/compute/test_util_internal.h index 9d4fbc283026..c90f6a437647 100644 --- a/cpp/src/arrow/compute/test_util_internal.h +++ b/cpp/src/arrow/compute/test_util_internal.h @@ -47,7 +47,7 @@ void AssertSelectionVectorsEqual(const std::shared_ptr& expecte const std::shared_ptr& actual); Expression unreachable_special(Expression argument); -Expression assert_empty_selection_special(Expression argument); +Expression assert_selection_empty_special(Expression argument); Expression assert_selection_eq_special(Expression argument, std::shared_ptr expected); From 7937232fc00721f9e9433e0484393c58bf6a44d6 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Fri, 10 Oct 2025 20:01:04 -0700 Subject: [PATCH 49/71] Reorg --- .../special/conditional_special_test.cc | 371 +++++++++++++----- cpp/src/arrow/compute/test_util_internal.cc | 163 -------- cpp/src/arrow/compute/test_util_internal.h | 8 - 3 files changed, 269 insertions(+), 273 deletions(-) diff --git a/cpp/src/arrow/compute/special/conditional_special_test.cc b/cpp/src/arrow/compute/special/conditional_special_test.cc index ff08a368e0d1..13dca618f298 100644 --- a/cpp/src/arrow/compute/special/conditional_special_test.cc +++ b/cpp/src/arrow/compute/special/conditional_special_test.cc @@ -30,6 +30,22 @@ namespace arrow::compute::internal { namespace { +void AssertSelectionVectorsEqual(const std::shared_ptr& expected, + const std::shared_ptr& actual) { + if (expected == nullptr) { + EXPECT_EQ(actual, nullptr); + return; + } + + if (actual == nullptr) { + EXPECT_EQ(expected, nullptr); + return; + } + + ASSERT_EQ(expected->length(), actual->length()); + AssertArraysEqual(*MakeArray(expected->data()), *MakeArray(actual->data())); +} + template void CheckBranchMask(const std::shared_ptr& branch_mask) { auto casted = checked_cast(branch_mask.get()); @@ -78,30 +94,6 @@ const auto kNullScalar = MakeNullScalar(boolean()); const auto kTrueScalar = MakeScalar(true); const auto kFalseScalar = MakeScalar(false); -const auto kNullLiteral = literal(kNullScalar); -const auto kTrueLiteral = literal(true); -const auto kFalseLiteral = literal(false); - -Expression bind_unreachable(Expression argument, const Schema& schm) { - EXPECT_OK_AND_ASSIGN(auto bound, unreachable_special(std::move(argument)).Bind(schm)); - return bound; -} - -Expression bind_assert_selection_empty(Expression argument, const Schema& schm) { - EXPECT_OK_AND_ASSIGN(auto bound, - assert_selection_empty_special(std::move(argument)).Bind(schm)); - return bound; -} - -Expression bind_assert_selection_equal(Expression argument, - std::shared_ptr selection, - const Schema& schm) { - EXPECT_OK_AND_ASSIGN( - auto bound, - assert_selection_eq_special(std::move(argument), std::move(selection)).Bind(schm)); - return bound; -} - } // namespace TEST(BranchMask, FromSelectionVector) { @@ -460,6 +452,177 @@ TEST(ConditionalBodyMask, NextBranchMask) { } } +namespace { + +class TrivialSpecialExecutor : public SpecialExecutor { + public: + explicit TrivialSpecialExecutor(Expression argument) + : SpecialExecutor(argument.type()), argument_(std::move(argument)) {} + + Result Execute(const ExecBatch& input, + ExecContext* exec_context) const override { + RETURN_NOT_OK(PreExecute(input, exec_context)); + return ExecuteScalarExpression(argument_, input, exec_context); + } + + protected: + virtual Status PreExecute(const ExecBatch& input, ExecContext* exec_context) const { + return Status::OK(); + } + + protected: + Expression argument_; +}; + +class UnreachableSpecialExecutor : public TrivialSpecialExecutor { + public: + explicit UnreachableSpecialExecutor(Expression argument) + : TrivialSpecialExecutor(std::move(argument)) {} + + protected: + Status PreExecute(const ExecBatch& input, ExecContext* exec_context) const override { + return Status::Invalid("Unreachable"); + } +}; + +class UnreachableSpecialForm : public SpecialForm { + public: + UnreachableSpecialForm() : SpecialForm("unreachable") {} + + protected: + Result> Bind( + std::vector& arguments, std::shared_ptr options, + ExecContext* exec_context) const override { + DCHECK_EQ(arguments.size(), 1); + return std::make_unique(arguments[0]); + } +}; + +class AssertEmptySelectionSpecialExecutor : public TrivialSpecialExecutor { + public: + explicit AssertEmptySelectionSpecialExecutor(Expression argument) + : TrivialSpecialExecutor(std::move(argument)) {} + + protected: + Status PreExecute(const ExecBatch& input, ExecContext* exec_context) const override { + if (input.selection_vector) { + return Status::Invalid("There shouldn't be a selection vector"); + } + return Status::OK(); + } +}; + +class AssertEmptySelectionSpecialForm : public SpecialForm { + public: + AssertEmptySelectionSpecialForm() : SpecialForm("assert_selection_empty") {} + + protected: + Result> Bind( + std::vector& arguments, std::shared_ptr options, + ExecContext* exec_context) const override { + DCHECK_EQ(arguments.size(), 1); + return std::make_unique(arguments[0]); + } +}; + +class AssertSelectionEqualSpecialExecutor : public TrivialSpecialExecutor { + public: + explicit AssertSelectionEqualSpecialExecutor(Expression argument, + std::shared_ptr expected) + : TrivialSpecialExecutor(std::move(argument)), expected_(std::move(expected)) { + DCHECK_NE(expected_, nullptr); + } + + protected: + Status PreExecute(const ExecBatch& input, ExecContext* exec_context) const override { + if (input.selection_vector == nullptr) { + return Status::Invalid("There should be a selection vector"); + } + if (!SelectionVectorsEqual(expected_, input.selection_vector)) { + return Status::Invalid("Selection vector does not match expected"); + } + return Status::OK(); + } + + private: + bool SelectionVectorsEqual(const std::shared_ptr& left, + const std::shared_ptr& right) const { + DCHECK_NE(left, nullptr); + DCHECK_NE(right, nullptr); + return MakeArray(left->data())->Equals(MakeArray(right->data())); + } + + std::shared_ptr expected_; +}; + +class AssertSelectionEqualSpecialForm : public SpecialForm { + public: + explicit AssertSelectionEqualSpecialForm(std::shared_ptr expected) + : SpecialForm("assert_selection_equal"), expected_(std::move(expected)) { + DCHECK_NE(expected_, nullptr); + } + + protected: + Result> Bind( + std::vector& arguments, std::shared_ptr options, + ExecContext* exec_context) const override { + DCHECK_EQ(arguments.size(), 1); + return std::make_unique(arguments[0], expected_); + } + + private: + std::shared_ptr expected_; +}; + +Expression unreachable_special(Expression argument) { + Expression::Special special; + special.special_form = std::make_shared(); + special.arguments.push_back(std::move(argument)); + return Expression(std::move(special)); +} + +Expression assert_selection_empty_special(Expression argument) { + Expression::Special special; + special.special_form = std::make_shared(); + special.arguments.push_back(std::move(argument)); + return Expression(std::move(special)); +} + +Expression assert_selection_eq_special(Expression argument, + std::shared_ptr expected) { + Expression::Special special; + special.special_form = + std::make_shared(std::move(expected)); + special.arguments.push_back(std::move(argument)); + return Expression(std::move(special)); +} + +const auto kNullLiteral = literal(kNullScalar); +const auto kTrueLiteral = literal(true); +const auto kFalseLiteral = literal(false); + +Expression unreachable(Expression argument, const Schema& schm) { + EXPECT_OK_AND_ASSIGN(auto bound, unreachable_special(std::move(argument)).Bind(schm)); + return bound; +} + +Expression assert_selection_empty(Expression argument, const Schema& schm) { + EXPECT_OK_AND_ASSIGN(auto bound, + assert_selection_empty_special(std::move(argument)).Bind(schm)); + return bound; +} + +Expression assert_selection_equal(Expression argument, + std::shared_ptr selection, + const Schema& schm) { + EXPECT_OK_AND_ASSIGN( + auto bound, + assert_selection_eq_special(std::move(argument), std::move(selection)).Bind(schm)); + return bound; +} + +} // namespace + TEST(ConditionalSpecialExecutor, Basic) { ConditionalSpecialExecutor executor({}, utf8()); EXPECT_EQ(executor.out_type().id(), Type::STRING); @@ -475,15 +638,15 @@ TEST(ConditionalSpecialExecutor, Execute) { auto sa = field_ref(2); auto sb = field_ref(3); auto sc = field_ref(4); - auto unreachable = [&](Expression argument) -> Expression { - return bind_unreachable(std::move(argument), *schm); + auto unreachable_sp = [&](Expression argument) -> Expression { + return unreachable(std::move(argument), *schm); }; - auto assert_selection_empty = [&](Expression argument) -> Expression { - return bind_assert_selection_empty(std::move(argument), *schm); + auto assert_selection_empty_sp = [&](Expression argument) -> Expression { + return assert_selection_empty(std::move(argument), *schm); }; - auto assert_selection_eq = + auto assert_selection_eq_sp = [&](Expression argument, std::shared_ptr selection) -> Expression { - return bind_assert_selection_equal(std::move(argument), std::move(selection), *schm); + return assert_selection_equal(std::move(argument), std::move(selection), *schm); }; auto batch = ExecBatchFromJSON({boolean(), boolean(), utf8(), utf8(), utf8()}, @@ -494,9 +657,9 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{assert_selection_empty(kFalseLiteral), unreachable(sa)}, - Branch{assert_selection_empty(kTrueLiteral), assert_selection_empty(sb)}, - Branch{unreachable(kTrueLiteral), unreachable(sc)}}, + {Branch{assert_selection_empty_sp(kFalseLiteral), unreachable_sp(sa)}, + Branch{assert_selection_empty_sp(kTrueLiteral), assert_selection_empty_sp(sb)}, + Branch{unreachable_sp(kTrueLiteral), unreachable_sp(sc)}}, utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["b0", "b1", "b2", "b3"])")), @@ -505,9 +668,9 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{assert_selection_empty(kFalseLiteral), unreachable(sa)}, - Branch{assert_selection_empty(kNullLiteral), unreachable(sb)}, - Branch{unreachable(kTrueLiteral), unreachable(sc)}}, + {Branch{assert_selection_empty_sp(kFalseLiteral), unreachable_sp(sa)}, + Branch{assert_selection_empty_sp(kNullLiteral), unreachable_sp(sb)}, + Branch{unreachable_sp(kTrueLiteral), unreachable_sp(sc)}}, utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"([null, null, null, null])")), @@ -516,11 +679,11 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{assert_selection_empty(kFalseLiteral), unreachable(sa)}, - Branch{assert_selection_empty(b0), - assert_selection_eq(sb, SelectionVectorFromJSON("[0, 3]"))}, - Branch{assert_selection_eq(kTrueLiteral, SelectionVectorFromJSON("[1]")), - assert_selection_eq(sc, SelectionVectorFromJSON("[1]"))}}, + {Branch{assert_selection_empty_sp(kFalseLiteral), unreachable_sp(sa)}, + Branch{assert_selection_empty_sp(b0), + assert_selection_eq_sp(sb, SelectionVectorFromJSON("[0, 3]"))}, + Branch{assert_selection_eq_sp(kTrueLiteral, SelectionVectorFromJSON("[1]")), + assert_selection_eq_sp(sc, SelectionVectorFromJSON("[1]"))}}, utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["b0", "c1", null, "b3"])")), @@ -529,11 +692,11 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{assert_selection_empty(kFalseLiteral), unreachable(sa)}, - Branch{assert_selection_empty(b1), - assert_selection_eq(sb, SelectionVectorFromJSON("[0, 2, 3]"))}, - Branch{assert_selection_eq(kTrueLiteral, SelectionVectorFromJSON("[1]")), - assert_selection_eq(sc, SelectionVectorFromJSON("[1]"))}}, + {Branch{assert_selection_empty_sp(kFalseLiteral), unreachable_sp(sa)}, + Branch{assert_selection_empty_sp(b1), + assert_selection_eq_sp(sb, SelectionVectorFromJSON("[0, 2, 3]"))}, + Branch{assert_selection_eq_sp(kTrueLiteral, SelectionVectorFromJSON("[1]")), + assert_selection_eq_sp(sc, SelectionVectorFromJSON("[1]"))}}, utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["b0", "c1", "b2", "b3"])")), @@ -542,9 +705,9 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{assert_selection_empty(kTrueLiteral), assert_selection_empty(sa)}, - Branch{unreachable(kTrueLiteral), unreachable(sb)}, - Branch{unreachable(kTrueLiteral), unreachable(sc)}}, + {Branch{assert_selection_empty_sp(kTrueLiteral), assert_selection_empty_sp(sa)}, + Branch{unreachable_sp(kTrueLiteral), unreachable_sp(sb)}, + Branch{unreachable_sp(kTrueLiteral), unreachable_sp(sc)}}, utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "a1", "a2", "a3"])")), @@ -553,9 +716,9 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{assert_selection_empty(kNullLiteral), unreachable(sa)}, - Branch{unreachable(kTrueLiteral), unreachable(sb)}, - Branch{unreachable(kTrueLiteral), unreachable(sc)}}, + {Branch{assert_selection_empty_sp(kNullLiteral), unreachable_sp(sa)}, + Branch{unreachable_sp(kTrueLiteral), unreachable_sp(sb)}, + Branch{unreachable_sp(kTrueLiteral), unreachable_sp(sc)}}, utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"([null, null, null, null])")), @@ -564,11 +727,11 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{assert_selection_empty(b0), - assert_selection_eq(sa, SelectionVectorFromJSON("[0, 3]"))}, - Branch{assert_selection_eq(kTrueLiteral, SelectionVectorFromJSON("[1]")), - assert_selection_eq(sb, SelectionVectorFromJSON("[1]"))}, - Branch{unreachable(kTrueLiteral), unreachable(sc)}}, + {Branch{assert_selection_empty_sp(b0), + assert_selection_eq_sp(sa, SelectionVectorFromJSON("[0, 3]"))}, + Branch{assert_selection_eq_sp(kTrueLiteral, SelectionVectorFromJSON("[1]")), + assert_selection_eq_sp(sb, SelectionVectorFromJSON("[1]"))}, + Branch{unreachable_sp(kTrueLiteral), unreachable_sp(sc)}}, utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "b1", null, "a3"])")), @@ -577,11 +740,11 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{assert_selection_empty(b0), - assert_selection_eq(sa, SelectionVectorFromJSON("[0, 3]"))}, - Branch{assert_selection_eq(kNullLiteral, SelectionVectorFromJSON("[1]")), - unreachable(sb)}, - Branch{unreachable(kTrueLiteral), unreachable(sc)}}, + {Branch{assert_selection_empty_sp(b0), + assert_selection_eq_sp(sa, SelectionVectorFromJSON("[0, 3]"))}, + Branch{assert_selection_eq_sp(kNullLiteral, SelectionVectorFromJSON("[1]")), + unreachable_sp(sb)}, + Branch{unreachable_sp(kTrueLiteral), unreachable_sp(sc)}}, utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", null, null, "a3"])")), @@ -590,11 +753,12 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{assert_selection_empty(b0), - assert_selection_eq(sa, SelectionVectorFromJSON("[0, 3]"))}, - Branch{assert_selection_eq(b0, SelectionVectorFromJSON("[1]")), unreachable(sb)}, - Branch{assert_selection_eq(kTrueLiteral, SelectionVectorFromJSON("[1]")), - assert_selection_eq(sc, SelectionVectorFromJSON("[1]"))}}, + {Branch{assert_selection_empty_sp(b0), + assert_selection_eq_sp(sa, SelectionVectorFromJSON("[0, 3]"))}, + Branch{assert_selection_eq_sp(b0, SelectionVectorFromJSON("[1]")), + unreachable_sp(sb)}, + Branch{assert_selection_eq_sp(kTrueLiteral, SelectionVectorFromJSON("[1]")), + assert_selection_eq_sp(sc, SelectionVectorFromJSON("[1]"))}}, utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "c1", null, "a3"])")), @@ -603,11 +767,12 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{assert_selection_empty(b0), - assert_selection_eq(sa, SelectionVectorFromJSON("[0, 3]"))}, - Branch{assert_selection_eq(b1, SelectionVectorFromJSON("[1]")), unreachable(sb)}, - Branch{assert_selection_eq(kTrueLiteral, SelectionVectorFromJSON("[1]")), - assert_selection_eq(sc, SelectionVectorFromJSON("[1]"))}}, + {Branch{assert_selection_empty_sp(b0), + assert_selection_eq_sp(sa, SelectionVectorFromJSON("[0, 3]"))}, + Branch{assert_selection_eq_sp(b1, SelectionVectorFromJSON("[1]")), + unreachable_sp(sb)}, + Branch{assert_selection_eq_sp(kTrueLiteral, SelectionVectorFromJSON("[1]")), + assert_selection_eq_sp(sc, SelectionVectorFromJSON("[1]"))}}, utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "c1", null, "a3"])")), @@ -616,11 +781,11 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{assert_selection_empty(b1), - assert_selection_eq(sa, SelectionVectorFromJSON("[0, 2, 3]"))}, - Branch{assert_selection_eq(kTrueLiteral, SelectionVectorFromJSON("[1]")), - assert_selection_eq(sb, SelectionVectorFromJSON("[1]"))}, - Branch{unreachable(kTrueLiteral), unreachable(sc)}}, + {Branch{assert_selection_empty_sp(b1), + assert_selection_eq_sp(sa, SelectionVectorFromJSON("[0, 2, 3]"))}, + Branch{assert_selection_eq_sp(kTrueLiteral, SelectionVectorFromJSON("[1]")), + assert_selection_eq_sp(sb, SelectionVectorFromJSON("[1]"))}, + Branch{unreachable_sp(kTrueLiteral), unreachable_sp(sc)}}, utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "b1", "a2", "a3"])")), @@ -629,11 +794,11 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{assert_selection_empty(b1), - assert_selection_eq(sa, SelectionVectorFromJSON("[0, 2, 3]"))}, - Branch{assert_selection_eq(kNullLiteral, SelectionVectorFromJSON("[1]")), - unreachable(sb)}, - Branch{unreachable(kTrueLiteral), unreachable(sc)}}, + {Branch{assert_selection_empty_sp(b1), + assert_selection_eq_sp(sa, SelectionVectorFromJSON("[0, 2, 3]"))}, + Branch{assert_selection_eq_sp(kNullLiteral, SelectionVectorFromJSON("[1]")), + unreachable_sp(sb)}, + Branch{unreachable_sp(kTrueLiteral), unreachable_sp(sc)}}, utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", null, "a2", "a3"])")), @@ -642,11 +807,12 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{assert_selection_empty(b1), - assert_selection_eq(sa, SelectionVectorFromJSON("[0, 2, 3]"))}, - Branch{assert_selection_eq(b0, SelectionVectorFromJSON("[1]")), unreachable(sb)}, - Branch{assert_selection_eq(kTrueLiteral, SelectionVectorFromJSON("[1]")), - assert_selection_eq(sc, SelectionVectorFromJSON("[1]"))}}, + {Branch{assert_selection_empty_sp(b1), + assert_selection_eq_sp(sa, SelectionVectorFromJSON("[0, 2, 3]"))}, + Branch{assert_selection_eq_sp(b0, SelectionVectorFromJSON("[1]")), + unreachable_sp(sb)}, + Branch{assert_selection_eq_sp(kTrueLiteral, SelectionVectorFromJSON("[1]")), + assert_selection_eq_sp(sc, SelectionVectorFromJSON("[1]"))}}, utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "c1", "a2", "a3"])")), @@ -655,11 +821,12 @@ TEST(ConditionalSpecialExecutor, Execute) { { ConditionalSpecialExecutor executor( - {Branch{assert_selection_empty(b1), - assert_selection_eq(sa, SelectionVectorFromJSON("[0, 2, 3]"))}, - Branch{assert_selection_eq(b1, SelectionVectorFromJSON("[1]")), unreachable(sb)}, - Branch{assert_selection_eq(kTrueLiteral, SelectionVectorFromJSON("[1]")), - assert_selection_eq(sc, SelectionVectorFromJSON("[1]"))}}, + {Branch{assert_selection_empty_sp(b1), + assert_selection_eq_sp(sa, SelectionVectorFromJSON("[0, 2, 3]"))}, + Branch{assert_selection_eq_sp(b1, SelectionVectorFromJSON("[1]")), + unreachable_sp(sb)}, + Branch{assert_selection_eq_sp(kTrueLiteral, SelectionVectorFromJSON("[1]")), + assert_selection_eq_sp(sc, SelectionVectorFromJSON("[1]"))}}, utf8()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(utf8(), R"(["a0", "c1", "a2", "a3"])")), @@ -671,15 +838,15 @@ TEST(ConditionalSpecialExecutor, ExecuteWithSelection) { auto schm = schema({field("a", int32())}); auto a = field_ref("a"); - auto unreachable = [&](Expression argument) -> Expression { - return bind_unreachable(std::move(argument), *schm); + auto unreachable_sp = [&](Expression argument) -> Expression { + return unreachable(std::move(argument), *schm); }; - auto assert_selection_empty = [&](Expression argument) -> Expression { - return bind_assert_selection_empty(std::move(argument), *schm); + auto assert_selection_empty_sp = [&](Expression argument) -> Expression { + return assert_selection_empty(std::move(argument), *schm); }; - auto assert_selection_eq = + auto assert_selection_eq_sp = [&](Expression argument, std::shared_ptr selection) -> Expression { - return bind_assert_selection_equal(std::move(argument), std::move(selection), *schm); + return assert_selection_equal(std::move(argument), std::move(selection), *schm); }; auto batch = ExecBatchFromJSON({int32()}, R"([[10], [11], [12], [13]])"); @@ -690,7 +857,7 @@ TEST(ConditionalSpecialExecutor, ExecuteWithSelection) { auto batch_with_selection = batch; batch_with_selection.selection_vector = selection; ConditionalSpecialExecutor executor( - {Branch{unreachable(kTrueLiteral), unreachable(a)}}, int32()); + {Branch{unreachable_sp(kTrueLiteral), unreachable_sp(a)}}, int32()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch_with_selection, default_exec_context())); AssertDatumsEqual(Datum(ArrayFromJSON(int32(), R"([null, null, null, null])")), @@ -703,8 +870,8 @@ TEST(ConditionalSpecialExecutor, ExecuteWithSelection) { auto batch_with_selection = batch; batch_with_selection.selection_vector = selection; ConditionalSpecialExecutor executor( - {Branch{assert_selection_eq(kTrueLiteral, selection), - assert_selection_eq(a, selection)}}, + {Branch{assert_selection_eq_sp(kTrueLiteral, selection), + assert_selection_eq_sp(a, selection)}}, int32()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch_with_selection, default_exec_context())); @@ -716,7 +883,7 @@ TEST(ConditionalSpecialExecutor, ExecuteWithSelection) { auto batch_with_selection = batch; batch_with_selection.selection_vector = selection; ConditionalSpecialExecutor executor( - {Branch{assert_selection_empty(kTrueLiteral), assert_selection_empty(a)}}, + {Branch{assert_selection_empty_sp(kTrueLiteral), assert_selection_empty_sp(a)}}, int32()); ASSERT_OK_AND_ASSIGN(auto result, executor.Execute(batch_with_selection, default_exec_context())); diff --git a/cpp/src/arrow/compute/test_util_internal.cc b/cpp/src/arrow/compute/test_util_internal.cc index 52775275dddf..95cc4c51c693 100644 --- a/cpp/src/arrow/compute/test_util_internal.cc +++ b/cpp/src/arrow/compute/test_util_internal.cc @@ -133,167 +133,4 @@ std::shared_ptr MakeSelectionVectorTo(int64_t length) { return std::make_shared(*arr); } -void AssertSelectionVectorsEqual(const std::shared_ptr& expected, - const std::shared_ptr& actual) { - if (expected == nullptr) { - EXPECT_EQ(actual, nullptr); - return; - } - - if (actual == nullptr) { - EXPECT_EQ(expected, nullptr); - return; - } - - ASSERT_EQ(expected->length(), actual->length()); - AssertArraysEqual(*MakeArray(expected->data()), *MakeArray(actual->data())); -} - -namespace { - -class TrivialSpecialExecutor : public SpecialExecutor { - public: - explicit TrivialSpecialExecutor(Expression argument) - : SpecialExecutor(argument.type()), argument_(std::move(argument)) {} - - Result Execute(const ExecBatch& input, - ExecContext* exec_context) const override { - RETURN_NOT_OK(PreExecute(input, exec_context)); - return ExecuteScalarExpression(argument_, input, exec_context); - } - - protected: - virtual Status PreExecute(const ExecBatch& input, ExecContext* exec_context) const { - return Status::OK(); - } - - protected: - Expression argument_; -}; - -class UnreachableSpecialExecutor : public TrivialSpecialExecutor { - public: - explicit UnreachableSpecialExecutor(Expression argument) - : TrivialSpecialExecutor(std::move(argument)) {} - - protected: - Status PreExecute(const ExecBatch& input, ExecContext* exec_context) const override { - return Status::Invalid("Unreachable"); - } -}; - -class UnreachableSpecialForm : public SpecialForm { - public: - UnreachableSpecialForm() : SpecialForm("unreachable") {} - - protected: - Result> Bind( - std::vector& arguments, std::shared_ptr options, - ExecContext* exec_context) const override { - DCHECK_EQ(arguments.size(), 1); - return std::make_unique(arguments[0]); - } -}; - -class AssertEmptySelectionSpecialExecutor : public TrivialSpecialExecutor { - public: - explicit AssertEmptySelectionSpecialExecutor(Expression argument) - : TrivialSpecialExecutor(std::move(argument)) {} - - protected: - Status PreExecute(const ExecBatch& input, ExecContext* exec_context) const override { - if (input.selection_vector) { - return Status::Invalid("There shouldn't be a selection vector"); - } - return Status::OK(); - } -}; - -class AssertEmptySelectionSpecialForm : public SpecialForm { - public: - AssertEmptySelectionSpecialForm() : SpecialForm("assert_selection_empty") {} - - protected: - Result> Bind( - std::vector& arguments, std::shared_ptr options, - ExecContext* exec_context) const override { - DCHECK_EQ(arguments.size(), 1); - return std::make_unique(arguments[0]); - } -}; - -class AssertSelectionEqualSpecialExecutor : public TrivialSpecialExecutor { - public: - explicit AssertSelectionEqualSpecialExecutor(Expression argument, - std::shared_ptr expected) - : TrivialSpecialExecutor(std::move(argument)), expected_(std::move(expected)) { - DCHECK_NE(expected_, nullptr); - } - - protected: - Status PreExecute(const ExecBatch& input, ExecContext* exec_context) const override { - if (input.selection_vector == nullptr) { - return Status::Invalid("There should be a selection vector"); - } - if (!SelectionVectorsEqual(expected_, input.selection_vector)) { - return Status::Invalid("Selection vector does not match expected"); - } - return Status::OK(); - } - - private: - bool SelectionVectorsEqual(const std::shared_ptr& left, - const std::shared_ptr& right) const { - DCHECK_NE(left, nullptr); - DCHECK_NE(right, nullptr); - return MakeArray(left->data())->Equals(MakeArray(right->data())); - } - - std::shared_ptr expected_; -}; - -class AssertSelectionEqualSpecialForm : public SpecialForm { - public: - explicit AssertSelectionEqualSpecialForm(std::shared_ptr expected) - : SpecialForm("assert_selection_equal"), expected_(std::move(expected)) { - DCHECK_NE(expected_, nullptr); - } - - protected: - Result> Bind( - std::vector& arguments, std::shared_ptr options, - ExecContext* exec_context) const override { - DCHECK_EQ(arguments.size(), 1); - return std::make_unique(arguments[0], expected_); - } - - private: - std::shared_ptr expected_; -}; - -} // namespace - -Expression unreachable_special(Expression argument) { - Expression::Special special; - special.special_form = std::make_shared(); - special.arguments.push_back(std::move(argument)); - return Expression(std::move(special)); -} - -Expression assert_selection_empty_special(Expression argument) { - Expression::Special special; - special.special_form = std::make_shared(); - special.arguments.push_back(std::move(argument)); - return Expression(std::move(special)); -} - -Expression assert_selection_eq_special(Expression argument, - std::shared_ptr expected) { - Expression::Special special; - special.special_form = - std::make_shared(std::move(expected)); - special.arguments.push_back(std::move(argument)); - return Expression(std::move(special)); -} - } // namespace arrow::compute diff --git a/cpp/src/arrow/compute/test_util_internal.h b/cpp/src/arrow/compute/test_util_internal.h index c90f6a437647..b7f4e1261993 100644 --- a/cpp/src/arrow/compute/test_util_internal.h +++ b/cpp/src/arrow/compute/test_util_internal.h @@ -43,12 +43,4 @@ std::shared_ptr SelectionVectorFromJSON(const std::string& json std::shared_ptr MakeSelectionVectorTo(int64_t length); -void AssertSelectionVectorsEqual(const std::shared_ptr& expected, - const std::shared_ptr& actual); - -Expression unreachable_special(Expression argument); -Expression assert_selection_empty_special(Expression argument); -Expression assert_selection_eq_special(Expression argument, - std::shared_ptr expected); - } // namespace arrow::compute From b4b7d76cf866b03ed918ebc46141479b8b2e4411 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Mon, 13 Oct 2025 17:00:00 -0700 Subject: [PATCH 50/71] Fix uninitialized seletective exec --- cpp/src/arrow/compute/kernel.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/arrow/compute/kernel.h b/cpp/src/arrow/compute/kernel.h index 6c7cccbed713..a84e93c50602 100644 --- a/cpp/src/arrow/compute/kernel.h +++ b/cpp/src/arrow/compute/kernel.h @@ -591,7 +591,7 @@ struct ARROW_EXPORT ScalarKernel : public Kernel { /// through the KernelContext. ArrayKernelExec exec; - ArrayKernelSelectiveExec selective_exec; + ArrayKernelSelectiveExec selective_exec = NULLPTR; /// \brief Writing execution results into larger contiguous allocations /// requires that the kernel be able to write into sliced output ArrayData*, From 69a5dbb09c28f551470006ea7df4a68ad5c766e2 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Mon, 13 Oct 2025 17:00:14 -0700 Subject: [PATCH 51/71] WIP if else special tests --- .../compute/special/if_else_special_test.cc | 1372 ++++++----------- 1 file changed, 451 insertions(+), 921 deletions(-) diff --git a/cpp/src/arrow/compute/special/if_else_special_test.cc b/cpp/src/arrow/compute/special/if_else_special_test.cc index 2ff9144b854c..b5b8e86436a2 100644 --- a/cpp/src/arrow/compute/special/if_else_special_test.cc +++ b/cpp/src/arrow/compute/special/if_else_special_test.cc @@ -19,9 +19,12 @@ #include #include +#include "arrow/array/concatenate.h" +#include "arrow/array/util.h" #include "arrow/compute/api_special.h" #include "arrow/compute/expression_test_internal.h" #include "arrow/compute/special_form.h" +#include "arrow/compute/test_util_internal.h" #include "arrow/util/logging_internal.h" namespace arrow::compute { @@ -212,860 +215,314 @@ TEST(IfElseSpecial, BindSpecialForm) { } } -void AssertEqualIgnoreShape(const Datum& expected, const Datum& result) { - if (expected.kind() == result.kind()) { - AssertDatumsEqual(expected, result); - return; - } - if (expected.is_scalar()) { - ASSERT_OK_AND_ASSIGN(auto expected_array, - MakeArrayFromScalar(*expected.scalar(), result.length())); - AssertDatumsEqual(expected_array, result); - return; - } - if (result.is_scalar()) { - ASSERT_OK_AND_ASSIGN(auto result_array, - MakeArrayFromScalar(*result.scalar(), expected.length())); - AssertDatumsEqual(expected, result_array); - return; - } -} - -Result ExecuteExpr(const Expression& expr, const std::shared_ptr& schema, - const ExecBatch& batch, - ExecContext* exec_context = default_exec_context()) { - ARROW_ASSIGN_OR_RAISE(auto bound, expr.Bind(*schema, exec_context)); - return ExecuteScalarExpression(bound, batch, exec_context); -} +// class IfElseSpecialFormTest : public ::testing::Test { +// protected: +// static void SetUpTestSuite() { ASSERT_OK(RegisterAuxilaryFunctions()); } -#define AssertExprRaisesWithMessage(expr, schema, batch, ENUM, message, ...) \ - { \ - ASSERT_RAISES_WITH_MESSAGE(ENUM, message, \ - ExecuteExpr(expr, schema, batch, ##__VA_ARGS__)); \ - } - -void AssertExprEqualIgnoreShape(const Expression& expr, - const std::shared_ptr& schema, - const ExecBatch& batch, const Datum& expected, - ExecContext* exec_context = default_exec_context()) { - ASSERT_OK_AND_ASSIGN(auto result, ExecuteExpr(expr, schema, batch, exec_context)); - AssertEqualIgnoreShape(expected, result); -} - -void AssertExprEqualExprsIgnoreShape(const Expression& expr, - const std::vector& exprs, - const std::shared_ptr& schema, - const ExecBatch& batch, - ExecContext* exec_context = default_exec_context()) { - ASSERT_OK_AND_ASSIGN(auto expected, ExecuteExpr(expr, schema, batch, exec_context)); - for (const auto& e : exprs) { - ARROW_SCOPED_TRACE(e.ToString()); - AssertExprEqualIgnoreShape(e, schema, batch, expected, exec_context); - } -} - -auto kBooleanNull = literal(MakeNullScalar(boolean())); -auto kIntNull = literal(MakeNullScalar(int32())); - -Expression if_else_regular(Expression cond, Expression if_true, Expression if_false) { - return call("if_else", {std::move(cond), std::move(if_true), std::move(if_false)}); -} -Expression unreachable(Expression arg) { return call("unreachable", {std::move(arg)}); } -Expression sv_suppress(Expression arg) { return call("sv_suppress", {std::move(arg)}); } -Expression assert_sv_exist(Expression arg) { - return call("assert_sv_exist", {std::move(arg)}); -} -Expression assert_sv_empty(Expression arg) { - return call("assert_sv_empty", {std::move(arg)}); -} +// protected: +// static Status UnreachableExec(KernelContext*, const ExecSpan&, ExecResult*) { +// return Status::Invalid("Unreachable"); +// } -TEST(IfElseSpecialForm, Basic) { - { - ARROW_SCOPED_TRACE("if (b != 0) then a / b else b"); - auto cond = call("not_equal", {field_ref("b"), literal(0)}); - auto if_true = call("divide", {field_ref("a"), field_ref("b")}); - auto if_false = field_ref("b"); - auto schema = arrow::schema({field("a", int32()), field("b", int32())}); - auto rb = RecordBatchFromJSON(schema, R"([ - [1, 1], - [2, 1], - [3, 0], - [4, 1], - [5, 1] - ])"); - auto batch = ExecBatch(*rb); - auto if_else_sp = if_else_special(cond, if_true, if_false); - { - auto expected = ArrayFromJSON(int32(), "[1, 2, 0, 4, 5]"); - AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); - } - { - ARROW_SCOPED_TRACE("(if (b != 0) then a / b else b) + 1"); - auto plus_one = call("add", {if_else_sp, literal(1)}); - { - auto expected = ArrayFromJSON(int32(), "[2, 3, 1, 5, 6]"); - AssertExprEqualIgnoreShape(plus_one, schema, batch, expected); - } - { - ARROW_SCOPED_TRACE( - "if ((if (b != 0) then a / b else b) + 1 != 1) then a / b else b"); - auto cond = call("not_equal", {plus_one, literal(1)}); - auto if_true = call("divide", {field_ref("a"), field_ref("b")}); - auto if_false = field_ref("b"); - auto if_else_sp = if_else_special(cond, if_true, if_false); - auto expected = ArrayFromJSON(int32(), "[1, 2, 0, 4, 5]"); - AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); - } - } - } - { - ARROW_SCOPED_TRACE("if (b != 0) then a else a"); - auto cond = call("not_equal", {field_ref("b"), literal(0)}); - auto if_true = field_ref("a"); - auto if_false = field_ref("a"); - auto schema = arrow::schema({field("a", int32()), field("b", int32())}); - auto rb = RecordBatchFromJSON(schema, R"([ - [1, 1], - [2, 1], - [3, 0], - [4, 1], - [5, 1] - ])"); - auto batch = ExecBatch(*rb); - auto if_else_sp = if_else_special(cond, if_true, if_false); - auto expected = ArrayFromJSON(int32(), "[1, 2, 3, 4, 5]"); - AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); - } -} +// static ArrayKernelSelectiveExec AssertSelection() { +// return [](KernelContext*, const ExecSpan&, const SelectionVectorSpan&, ExecResult*) +// { +// return Status::Invalid("Unreachable"); +// }; +// } -TEST(IfElseSpecialForm, ImplicitCast) { - auto schema = arrow::schema({field("i8", int8()), field("i32", int32())}); - for (const auto& if_else_sp : - {if_else_special(literal(true), field_ref("i8"), field_ref("i32")), - if_else_special(literal(true), field_ref("i32"), field_ref("i8")), - // Literal will be downcast. - if_else_special(literal(true), field_ref("i32"), - literal(static_cast(0)))}) { - ARROW_SCOPED_TRACE(if_else_sp.ToString()); - ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema)); - ASSERT_EQ(bound.type()->id(), Type::INT32); - } -} +// static Status IdentityExec(KernelContext*, const ExecSpan& span, ExecResult* out) { +// DCHECK_EQ(span.num_values(), 1); +// const auto& arg = span[0]; +// DCHECK(arg.is_array()); +// *out->array_data_mutable() = *arg.array.ToArrayData(); +// return Status::OK(); +// } -class IfElseSpecialFormTest : public ::testing::Test { - protected: - static void SetUpTestSuite() { ASSERT_OK(RegisterAuxilaryFunctions()); } +// static Status AssertSelectionVectorNotExistExec(KernelContext* kernel_ctx, +// const ExecSpan& span, ExecResult* +// out) { +// return IdentityExec(kernel_ctx, span, out); +// } - protected: - static Status UnreachableExec(KernelContext*, const ExecSpan&, ExecResult*) { - return Status::Invalid("Unreachable"); - } +// static Status AssertSelectionVectorExistExec(KernelContext* kernel_ctx, +// const ExecSpan& span, +// const SelectionVectorSpan& selection, +// ExecResult* out) { +// for (int32_t i = 0; i < selection.length(); ++i) { +// EXPECT_GE(selection[i], 0); +// EXPECT_LT(selection[i], span.length); +// } +// return IdentityExec(kernel_ctx, span, out); +// } - static Status IdentityExec(KernelContext*, const ExecSpan& span, ExecResult* out) { - DCHECK_EQ(span.num_values(), 1); - const auto& arg = span[0]; - DCHECK(arg.is_array()); - *out->array_data_mutable() = *arg.array.ToArrayData(); - return Status::OK(); - } +// static Status RegisterAuxilaryFunctions() { +// auto registry = GetFunctionRegistry(); + +// { +// auto register_unreachable_func = [&](const std::string& name) -> Status { +// auto func = +// std::make_shared(name, Arity::Unary(), +// FunctionDoc::Empty()); + +// ScalarKernel kernel({InputType::Any()}, internal::FirstType, UnreachableExec); +// kernel.can_write_into_slices = false; +// kernel.null_handling = NullHandling::COMPUTED_NO_PREALLOCATE; +// kernel.mem_allocation = MemAllocation::NO_PREALLOCATE; +// RETURN_NOT_OK(func->AddKernel(kernel)); +// RETURN_NOT_OK(registry->AddFunction(std::move(func))); +// return Status::OK(); +// }; + +// RETURN_NOT_OK(register_unreachable_func("unreachable")); +// } - static Status AssertSelectionVectorNotExistExec(KernelContext* kernel_ctx, - const ExecSpan& span, ExecResult* out) { - return IdentityExec(kernel_ctx, span, out); - } +// { +// auto register_sv_awareness_func = [&](const std::string& name, +// bool sv_awareness) -> Status { +// auto func = +// std::make_shared(name, Arity::Unary(), +// FunctionDoc::Empty()); + +// ScalarKernel kernel({InputType::Any()}, internal::FirstType, IdentityExec); +// kernel.can_write_into_slices = false; +// kernel.null_handling = NullHandling::COMPUTED_NO_PREALLOCATE; +// kernel.mem_allocation = MemAllocation::NO_PREALLOCATE; +// RETURN_NOT_OK(func->AddKernel(kernel)); +// RETURN_NOT_OK(registry->AddFunction(std::move(func))); +// return Status::OK(); +// }; + +// RETURN_NOT_OK(register_sv_awareness_func("sv_suppress", false)); +// } - static Status AssertSelectionVectorExistExec(KernelContext* kernel_ctx, - const ExecSpan& span, - const SelectionVectorSpan& selection, - ExecResult* out) { - for (int32_t i = 0; i < selection.length(); ++i) { - EXPECT_GE(selection[i], 0); - EXPECT_LT(selection[i], span.length); - } - return IdentityExec(kernel_ctx, span, out); - } +// { +// auto register_assert_sv_func = [&](const std::string& name, +// bool sv_existence) -> Status { +// auto func = +// std::make_shared(name, Arity::Unary(), +// FunctionDoc::Empty()); + +// ArrayKernelExec exec = AssertSelectionVectorNotExistExec; +// ArrayKernelSelectiveExec selective_exec = nullptr; +// if (sv_existence) { +// selective_exec = AssertSelectionVectorExistExec; +// } +// ScalarKernel kernel({InputType::Any()}, internal::FirstType, std::move(exec), +// std::move(selective_exec), +// /*init=*/nullptr); +// kernel.can_write_into_slices = false; +// kernel.null_handling = NullHandling::COMPUTED_NO_PREALLOCATE; +// kernel.mem_allocation = MemAllocation::NO_PREALLOCATE; +// RETURN_NOT_OK(func->AddKernel(kernel)); +// RETURN_NOT_OK(registry->AddFunction(std::move(func))); +// return Status::OK(); +// }; + +// RETURN_NOT_OK(register_assert_sv_func("assert_sv_exist", /*sv_existence=*/true)); +// RETURN_NOT_OK(register_assert_sv_func("assert_sv_empty", +// /*sv_existence=*/false)); +// } - static Status RegisterAuxilaryFunctions() { - auto registry = GetFunctionRegistry(); +// return Status::OK(); +// } - { - auto register_unreachable_func = [&](const std::string& name) -> Status { - auto func = - std::make_shared(name, Arity::Unary(), FunctionDoc::Empty()); +// static std::vector SuppressSelectionVectorAwareForIfElse( +// const Expression& cond, const Expression& if_true, const Expression& if_false) { +// auto suppress_if_else_recursive = +// [&](const Expression& expr) -> std::vector { +// if (const auto& sp = expr.special(); sp && sp->special_form->name() == "if_else") +// { +// const auto& cond = sp->arguments[0]; +// const auto& if_true = sp->arguments[1]; +// const auto& if_false = sp->arguments[2]; +// return SuppressSelectionVectorAwareForIfElse(cond, if_true, if_false); +// } else { +// return {expr}; +// } +// }; +// auto suppressed_conds = suppress_if_else_recursive(cond); +// auto suppressed_if_trues = suppress_if_else_recursive(if_true); +// auto suppressed_if_falses = suppress_if_else_recursive(if_false); +// std::vector result; +// for (const auto& suppressed_cond : suppressed_conds) { +// for (const auto& suppressed_if_true : suppressed_if_trues) { +// for (const auto& suppressed_if_false : suppressed_if_falses) { +// result.emplace_back( +// if_else_special(suppressed_cond, suppressed_if_true, +// suppressed_if_false)); +// result.emplace_back(if_else_special(sv_suppress(suppressed_cond), +// suppressed_if_true, +// suppressed_if_false)); +// result.emplace_back(if_else_special(suppressed_cond, +// sv_suppress(suppressed_if_true), +// sv_suppress(suppressed_if_false))); +// } +// } +// } +// return result; +// } - ScalarKernel kernel({InputType::Any()}, internal::FirstType, UnreachableExec); - kernel.can_write_into_slices = false; - kernel.null_handling = NullHandling::COMPUTED_NO_PREALLOCATE; - kernel.mem_allocation = MemAllocation::NO_PREALLOCATE; - RETURN_NOT_OK(func->AddKernel(kernel)); - RETURN_NOT_OK(registry->AddFunction(std::move(func))); - return Status::OK(); - }; +// static void CheckIfElseIgnoreShape(const Expression& cond, const Expression& if_true, +// const Expression& if_false, +// const std::shared_ptr& schema, +// const ExecBatch& batch, +// ExecContext* exec_context = +// default_exec_context()) { +// auto if_else = if_else_regular(cond, if_true, if_false); +// auto exprs = SuppressSelectionVectorAwareForIfElse(cond, if_true, if_false); +// AssertExprEqualExprsIgnoreShape(if_else, exprs, schema, batch, exec_context); +// } +// }; - RETURN_NOT_OK(register_unreachable_func("unreachable")); - } +namespace { - { - auto register_sv_awareness_func = [&](const std::string& name, - bool sv_awareness) -> Status { - auto func = - std::make_shared(name, Arity::Unary(), FunctionDoc::Empty()); - - ScalarKernel kernel({InputType::Any()}, internal::FirstType, IdentityExec); - kernel.can_write_into_slices = false; - kernel.null_handling = NullHandling::COMPUTED_NO_PREALLOCATE; - kernel.mem_allocation = MemAllocation::NO_PREALLOCATE; - RETURN_NOT_OK(func->AddKernel(kernel)); - RETURN_NOT_OK(registry->AddFunction(std::move(func))); - return Status::OK(); - }; - - RETURN_NOT_OK(register_sv_awareness_func("sv_suppress", false)); - } +Expression if_else(Expression cond, Expression if_true, Expression if_false) { + return call("if_else", {std::move(cond), std::move(if_true), std::move(if_false)}); +} - { - auto register_assert_sv_func = [&](const std::string& name, - bool sv_existence) -> Status { - auto func = - std::make_shared(name, Arity::Unary(), FunctionDoc::Empty()); +Result ExecuteExpr(Expression expr, const Schema& schema, const ExecBatch& batch, + ExecContext* exec_context = default_exec_context()) { + ARROW_ASSIGN_OR_RAISE(auto bound, expr.Bind(schema, exec_context)); + return ExecuteScalarExpression(bound, batch, exec_context); +} - ArrayKernelExec exec = AssertSelectionVectorNotExistExec; - ArrayKernelSelectiveExec selective_exec = nullptr; - if (sv_existence) { - selective_exec = AssertSelectionVectorExistExec; - } - ScalarKernel kernel({InputType::Any()}, internal::FirstType, std::move(exec), - std::move(selective_exec), - /*init=*/nullptr); - kernel.can_write_into_slices = false; - kernel.null_handling = NullHandling::COMPUTED_NO_PREALLOCATE; - kernel.mem_allocation = MemAllocation::NO_PREALLOCATE; - RETURN_NOT_OK(func->AddKernel(kernel)); - RETURN_NOT_OK(registry->AddFunction(std::move(func))); - return Status::OK(); - }; - - RETURN_NOT_OK(register_assert_sv_func("assert_sv_exist", /*sv_existence=*/true)); - RETURN_NOT_OK(register_assert_sv_func("assert_sv_empty", /*sv_existence=*/false)); - } +Result ExecuteIfElse(Expression cond, Expression if_true, Expression if_false, + const Schema& schema, const ExecBatch& batch, + ExecContext* exec_context = default_exec_context()) { + return ExecuteExpr(if_else(std::move(cond), std::move(if_true), std::move(if_false)), + schema, batch, exec_context); +} - return Status::OK(); - } +Result ExecuteIfElseSpecial(Expression cond, Expression if_true, + Expression if_false, const Schema& schema, + const ExecBatch& batch, + ExecContext* exec_context = default_exec_context()) { + return ExecuteExpr( + if_else_special(std::move(cond), std::move(if_true), std::move(if_false)), schema, + batch, exec_context); +} - static std::vector SuppressSelectionVectorAwareForIfElse( - const Expression& cond, const Expression& if_true, const Expression& if_false) { - auto suppress_if_else_recursive = - [&](const Expression& expr) -> std::vector { - if (const auto& sp = expr.special(); sp && sp->special_form->name() == "if_else") { - const auto& cond = sp->arguments[0]; - const auto& if_true = sp->arguments[1]; - const auto& if_false = sp->arguments[2]; - return SuppressSelectionVectorAwareForIfElse(cond, if_true, if_false); - } else { - return {expr}; - } - }; - auto suppressed_conds = suppress_if_else_recursive(cond); - auto suppressed_if_trues = suppress_if_else_recursive(if_true); - auto suppressed_if_falses = suppress_if_else_recursive(if_false); - std::vector result; - for (const auto& suppressed_cond : suppressed_conds) { - for (const auto& suppressed_if_true : suppressed_if_trues) { - for (const auto& suppressed_if_false : suppressed_if_falses) { - result.emplace_back( - if_else_special(suppressed_cond, suppressed_if_true, suppressed_if_false)); - result.emplace_back(if_else_special(sv_suppress(suppressed_cond), - suppressed_if_true, suppressed_if_false)); - result.emplace_back(if_else_special(suppressed_cond, - sv_suppress(suppressed_if_true), - sv_suppress(suppressed_if_false))); - } - } - } - return result; - } +void AssertDatumsEqualIgnoreShape(const Datum& expected, const Datum& result) { + DCHECK(expected.is_scalar() || expected.is_array() || expected.is_chunked_array()); + DCHECK(result.is_scalar() || result.is_array() || result.is_chunked_array()); - static void CheckIfElseIgnoreShape(const Expression& cond, const Expression& if_true, - const Expression& if_false, - const std::shared_ptr& schema, - const ExecBatch& batch, - ExecContext* exec_context = default_exec_context()) { - auto if_else = if_else_regular(cond, if_true, if_false); - auto exprs = SuppressSelectionVectorAwareForIfElse(cond, if_true, if_false); - AssertExprEqualExprsIgnoreShape(if_else, exprs, schema, batch, exec_context); + if (expected.kind() == result.kind()) { + AssertDatumsEqual(expected, result); + return; } -}; -TEST_F(IfElseSpecialFormTest, AuxilaryFunction) { - auto schema = arrow::schema({field("a", boolean())}); - auto a = field_ref("a"); - auto batch = ExecBatch({*ArrayFromJSON(boolean(), "[null, true, false]")}, 3); - { - ARROW_SCOPED_TRACE("unreachable"); - AssertExprRaisesWithMessage(unreachable(a), schema, batch, Invalid, - "Invalid: Unreachable"); - } - { - ARROW_SCOPED_TRACE("assert selection vector existence"); - { - ARROW_SCOPED_TRACE("if (a) then a else a"); - auto cond = a; - auto if_true = a; - auto if_false = a; - auto expected = ArrayFromJSON(boolean(), "[null, true, false]"); - { - auto if_else_sp = - if_else_special(cond, assert_sv_exist(if_true), assert_sv_exist(if_false)); - AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); - } - { - auto if_else_sp = - if_else_special(cond, sv_suppress(if_true), assert_sv_exist(if_false)); - AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); - } - { - auto if_else_sp = - if_else_special(cond, assert_sv_exist(if_true), sv_suppress(if_false)); - AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); - } - { - auto if_else_sp = - if_else_special(cond, sv_suppress(if_true), assert_sv_empty(if_false)); - AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); - } - { - auto if_else_sp = - if_else_special(cond, assert_sv_empty(if_true), sv_suppress(if_false)); - AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); - } - { - auto if_else_sp = if_else_special(cond, assert_sv_empty(if_true), if_false); - AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); - } - { - auto if_else_sp = if_else_special(cond, if_true, assert_sv_empty(if_false)); - AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); - } + int64_t length = expected.is_scalar() ? result.length() : expected.length(); + auto to_array = [&](const Datum& datum) -> Result> { + if (datum.is_scalar()) { + return MakeArrayFromScalar(*datum.scalar(), length); } - { - ARROW_SCOPED_TRACE("if (true) then a else false"); - auto cond = literal(true); - auto if_true = a; - auto if_false = literal(false); - auto expected = ArrayFromJSON(boolean(), "[null, true, false]"); - { - auto if_else_sp = if_else_special(cond, assert_sv_exist(if_true), if_false); - AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); - } - { - auto if_else_sp = if_else_special(cond, assert_sv_empty(if_true), if_false); - AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); - } + if (datum.is_array()) { + return datum.make_array(); } - } -} + DCHECK(datum.is_chunked_array()); + return Concatenate(datum.chunked_array()->chunks(), default_memory_pool()); + }; -TEST_F(IfElseSpecialFormTest, SelectionVectorExistence) { - auto schema = arrow::schema({field("b_null", boolean()), field("b_true", boolean()), - field("b_false", boolean()), field("b", boolean()), - field("i1", int32()), field("i2", int32())}); - auto b_null = field_ref("b_null"); - auto b_true = field_ref("b_true"); - auto b_false = field_ref("b_false"); - auto b = field_ref("b"); - auto i1 = field_ref("i1"); - auto i2 = field_ref("i2"); - auto batch = ExecBatch( - { - *ArrayFromJSON(boolean(), "[null, null, null]"), - *ArrayFromJSON(boolean(), "[true, true, true]"), - *ArrayFromJSON(boolean(), "[false, false, false]"), - *ArrayFromJSON(boolean(), "[null, true, false]"), - *ArrayFromJSON(int32(), "[0, 1, -1]"), - *ArrayFromJSON(int32(), "[0, -1, 1]"), - }, - 3); + ASSERT_OK_AND_ASSIGN(auto expected_array, to_array(expected)); + ASSERT_OK_AND_ASSIGN(auto result_array, to_array(result)); +} - { - ARROW_SCOPED_TRACE("all null condition"); - auto expected = ArrayFromJSON(int32(), "[null, null, null]"); - for (const auto& cond : - {kBooleanNull, b_null, assert_sv_empty(kBooleanNull), assert_sv_empty(b_null)}) { - ARROW_SCOPED_TRACE("cond: " + cond.ToString()); - auto if_else_sp = if_else_special(cond, unreachable(i1), unreachable(i2)); - AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); - } - } - { - ARROW_SCOPED_TRACE("all true condition"); - auto expected = ArrayFromJSON(int32(), "[0, 1, -1]"); - for (const auto& if_else_sp : - {if_else_special(literal(true), i1, unreachable(i2)), - if_else_special(b_true, i1, unreachable(i2)), - if_else_special(assert_sv_empty(literal(true)), assert_sv_empty(i1), - unreachable(i2)), - if_else_special(assert_sv_empty(b_true), assert_sv_exist(i1), unreachable(i2)), - if_else_special(assert_sv_empty(b_true), assert_sv_empty(i1), - sv_suppress(unreachable(i2)))}) { - ARROW_SCOPED_TRACE(if_else_sp.ToString()); - AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); - } - } - { - ARROW_SCOPED_TRACE("all false condition"); - auto expected = ArrayFromJSON(int32(), "[0, -1, 1]"); - for (const auto& if_else_sp : - {if_else_special(literal(false), unreachable(i1), i2), - if_else_special(b_false, unreachable(i1), i2), - if_else_special(assert_sv_empty(literal(false)), unreachable(i1), - assert_sv_empty(i2)), - if_else_special(assert_sv_empty(b_false), unreachable(i1), assert_sv_exist(i2)), - if_else_special(assert_sv_empty(b_false), sv_suppress(unreachable(i1)), - assert_sv_empty(i2))}) { - ARROW_SCOPED_TRACE(if_else_sp.ToString()); - AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); - } - } - { - ARROW_SCOPED_TRACE("even condition"); - auto expected = ArrayFromJSON(int32(), "[null, 1, 1]"); - for (const auto& if_else_sp : - {if_else_special(b, i1, i2), - if_else_special(assert_sv_empty(b), assert_sv_exist(i1), assert_sv_exist(i2)), - if_else_special(sv_suppress(b), assert_sv_exist(i1), assert_sv_exist(i2)), - if_else_special(b, sv_suppress(i1), assert_sv_empty(i2)), - if_else_special(b, assert_sv_empty(i1), sv_suppress(i2))}) { - ARROW_SCOPED_TRACE(if_else_sp.ToString()); - AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); - } - } - { - ARROW_SCOPED_TRACE("literal bodies"); - auto expected = ArrayFromJSON(int32(), "[null, 1, 1]"); - for (const auto& if_else_sp : - {if_else_special(assert_sv_empty(b), assert_sv_empty(literal(1)), - assert_sv_exist(i2)), - if_else_special(assert_sv_empty(b), assert_sv_exist(i1), - assert_sv_empty(literal(1))), - if_else_special(assert_sv_empty(b), assert_sv_empty(literal(1)), - assert_sv_empty(literal(1)))}) { - ARROW_SCOPED_TRACE(if_else_sp.ToString()); - AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); - } - } - { - ARROW_SCOPED_TRACE("nested"); - auto expected = ArrayFromJSON(int32(), "[null, 1, 1]"); - for (const auto& if_else_sp : - {if_else_special(b, if_else_special(b, i1, i2), if_else_special(b, i1, i2)), - // The nested if_else_special will see a selection vector. - if_else_special(b, assert_sv_exist(if_else_special(b, i1, i2)), - assert_sv_exist(if_else_special(b, i1, i2))), - // The arguments of nested if_else_special will see a selection vector. - if_else_special(b, - if_else_special(assert_sv_empty(literal(true)), - assert_sv_exist(i1), unreachable(i2)), - if_else_special(assert_sv_empty(literal(false)), - unreachable(i1), assert_sv_exist(i2))), - if_else_special(b, - if_else_special(assert_sv_exist(b), assert_sv_exist(i1), - assert_sv_exist(i2)), - if_else_special(assert_sv_exist(b), assert_sv_exist(i1), - assert_sv_exist(i2))), - // Selection vector existences with some argument of the nested if_else_special - // being selection vector unaware. - if_else_special( - b, - assert_sv_exist(if_else_special(sv_suppress(literal(true)), - assert_sv_exist(i1), unreachable(i2))), - assert_sv_exist(if_else_special(assert_sv_empty(literal(false)), - unreachable(i1), assert_sv_exist(i2)))), - if_else_special(b, - assert_sv_exist(if_else_special( - sv_suppress(b), assert_sv_exist(i1), assert_sv_exist(i2))), - assert_sv_exist(if_else_special( - assert_sv_exist(b), unreachable(i1), assert_sv_exist(i2)))), - if_else_special( - b, - assert_sv_exist(if_else_special(assert_sv_empty(literal(true)), - sv_suppress(i1), unreachable(i2))), - assert_sv_exist(if_else_special(assert_sv_empty(literal(false)), - unreachable(i1), assert_sv_exist(i2)))), - if_else_special( - b, - assert_sv_exist(if_else_special(assert_sv_exist(b), sv_suppress(i1), - assert_sv_empty(i2))), - assert_sv_exist(if_else_special(assert_sv_exist(b), assert_sv_exist(i1), - assert_sv_exist(i2)))), - if_else_special( - b, - assert_sv_exist(if_else_special(assert_sv_empty(literal(true)), - assert_sv_empty(i1), - sv_suppress(unreachable(i2)))), - assert_sv_exist(if_else_special(assert_sv_empty(literal(false)), - unreachable(i1), assert_sv_exist(i2)))), - if_else_special( - b, - assert_sv_exist(if_else_special(assert_sv_exist(b), assert_sv_empty(i1), - sv_suppress(i2))), - assert_sv_exist(if_else_special(assert_sv_exist(b), assert_sv_exist(i1), - assert_sv_exist(i2)))), - if_else_special( - b, - assert_sv_exist(if_else_special(assert_sv_empty(literal(true)), - assert_sv_exist(i1), unreachable(i2))), - assert_sv_exist(if_else_special(assert_sv_empty(literal(false)), - unreachable(i1), assert_sv_exist(i2)))), - if_else_special( - b, - assert_sv_exist(if_else_special(assert_sv_exist(b), assert_sv_exist(i1), - assert_sv_exist(i2))), - assert_sv_exist(if_else_special(sv_suppress(b), assert_sv_exist(i1), - assert_sv_exist(i2)))), - if_else_special( - b, - assert_sv_exist(if_else_special(assert_sv_empty(literal(true)), - assert_sv_exist(i1), unreachable(i2))), - assert_sv_exist(if_else_special(assert_sv_empty(literal(false)), - sv_suppress(unreachable(i1)), - assert_sv_empty(i2)))), - if_else_special( - b, - assert_sv_exist(if_else_special(assert_sv_exist(b), assert_sv_exist(i1), - assert_sv_exist(i2))), - assert_sv_exist(if_else_special(assert_sv_exist(b), sv_suppress(i1), - assert_sv_empty(i2)))), - if_else_special( - b, - assert_sv_exist(if_else_special(assert_sv_empty(literal(true)), - assert_sv_exist(i1), unreachable(i2))), - assert_sv_exist(if_else_special(assert_sv_empty(literal(false)), - assert_sv_empty(unreachable(i1)), - sv_suppress(i2)))), - if_else_special( - b, - assert_sv_exist(if_else_special(assert_sv_exist(b), assert_sv_exist(i1), - assert_sv_exist(i2))), - assert_sv_exist(if_else_special(assert_sv_exist(b), assert_sv_empty(i1), - sv_suppress(i2))))}) { - ARROW_SCOPED_TRACE(if_else_sp.ToString()); - AssertExprEqualIgnoreShape(if_else_sp, schema, batch, expected); - } - } +void CheckIfElseSpecial(Expression cond, Expression if_true, Expression if_false, + const Schema& schema, const ExecBatch& batch, + ExecContext* exec_context = default_exec_context()) { + ASSERT_OK_AND_ASSIGN( + auto expected, ExecuteIfElse(cond, if_true, if_false, schema, batch, exec_context)); + ASSERT_OK_AND_ASSIGN(auto result, ExecuteIfElseSpecial(cond, if_true, if_false, schema, + batch, exec_context)); + AssertDatumsEqualIgnoreShape(expected, result); } -TEST_F(IfElseSpecialFormTest, SelectionVectorExistenceExecChunkSize) { - ExecContext exec_context; - constexpr int64_t num_rows = 8; - auto schema = arrow::schema({field("b", boolean()), field("i", int32())}); - auto b = field_ref("b"); - auto i = field_ref("i"); - auto batch = ExecBatch( - { - *ArrayFromJSON(boolean(), "[true, true, true, true, true, true, true, false]"), - *ArrayFromJSON(int32(), "[1, 2, 3, 4, 5, 6, 7, 8]"), - }, - num_rows); - { - ARROW_SCOPED_TRACE("exec_chunksize >= batch_size"); - for (auto chunksize : {num_rows, num_rows + 1}) { - ARROW_SCOPED_TRACE("exec_chunksize: " + std::to_string(chunksize)); - exec_context.set_exec_chunksize(chunksize); - { - ARROW_SCOPED_TRACE("all literal bodies"); - auto if_else_sp = - if_else_special(b, assert_sv_empty(literal(1)), assert_sv_empty(literal(0))); - ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); - ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); - } - { - ARROW_SCOPED_TRACE("array true body"); - auto if_else_sp = - if_else_special(b, assert_sv_exist(i), assert_sv_empty(literal(0))); - ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); - ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); - } - { - ARROW_SCOPED_TRACE("array false body"); - auto if_else_sp = - if_else_special(b, assert_sv_empty(literal(1)), assert_sv_exist(i)); - ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); - ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); - } - { - ARROW_SCOPED_TRACE("all array bodies"); - auto if_else_sp = if_else_special(b, assert_sv_exist(i), assert_sv_exist(i)); - ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); - ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); - } - } - } - { - ARROW_SCOPED_TRACE("exec_chunksize < batch_size"); - for (auto chunksize : {num_rows / 2, num_rows - 1}) { - ARROW_SCOPED_TRACE("exec_chunksize: " + std::to_string(chunksize)); - exec_context.set_exec_chunksize(num_rows - 1); - { - ARROW_SCOPED_TRACE("all literal bodies"); - auto if_else_sp = - if_else_special(b, assert_sv_empty(literal(1)), assert_sv_empty(literal(0))); - ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); - ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); - } - { - ARROW_SCOPED_TRACE("array true body"); - auto if_else_sp = - if_else_special(b, assert_sv_exist(i), assert_sv_empty(literal(0))); - // ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); - // ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); - } - { - ARROW_SCOPED_TRACE("array true body"); - auto if_else_sp = - if_else_special(b, assert_sv_empty(i), assert_sv_empty(literal(0))); - ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); - ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); - } - { - ARROW_SCOPED_TRACE("array false body"); - auto if_else_sp = - if_else_special(b, assert_sv_empty(literal(1)), assert_sv_exist(i)); - ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); - ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); - } - { - ARROW_SCOPED_TRACE("all array bodies"); - auto if_else_sp = if_else_special(b, assert_sv_exist(i), assert_sv_exist(i)); - ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); - ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); - } - } - } - { - ARROW_SCOPED_TRACE("nested"); - { - ARROW_SCOPED_TRACE("exec_chunksize == batch_size"); - exec_context.set_exec_chunksize(num_rows); - auto if_else_sp = if_else_special( - b, - assert_sv_exist(if_else_special(assert_sv_exist(b), assert_sv_exist(i), - assert_sv_exist(i))), - assert_sv_exist(if_else_special(assert_sv_exist(b), assert_sv_exist(i), - assert_sv_exist(i)))); - ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); - ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); - } - { - ARROW_SCOPED_TRACE("exec_chunksize < batch_size"); - exec_context.set_exec_chunksize(num_rows - 2); - auto if_else_sp = if_else_special( - b, - assert_sv_exist(if_else_special(assert_sv_exist(b), assert_sv_exist(i), - assert_sv_exist(i))), - assert_sv_exist(if_else_special( - assert_sv_exist(b), - assert_sv_exist(if_else_special(assert_sv_exist(b), assert_sv_exist(i), - assert_sv_exist(i))), - assert_sv_exist(i)))); - ASSERT_OK_AND_ASSIGN(auto bound, if_else_sp.Bind(*schema, &exec_context)); - ASSERT_OK(ExecuteScalarExpression(bound, batch, &exec_context)); - } - } +} // namespace + +// GH-XXXXX: +TEST(IfElseSpecial, IfNotZeroThenDivide) { + // if (b != 0) then (a / b) else b + auto cond = call("not_equal", {field_ref("b"), literal(0)}); + auto if_true = call("divide", {field_ref("a"), field_ref("b")}); + auto if_false = field_ref("b"); + + auto schm = schema({field("a", int32()), field("b", int32())}); + auto batch = ExecBatchFromJSON({int32(), int32()}, + R"([[1, 1], + [2, 1], + [3, 0], + [4, 1], + [5, 1]])"); + + ASSERT_OK_AND_ASSIGN(auto result, + ExecuteIfElseSpecial(cond, if_true, if_false, *schm, batch)); + auto expected = ArrayFromJSON(int32(), "[1, 2, 0, 4, 5]"); + + AssertDatumsEqual(expected, result); } -// TODO: Selection vector existence of chunked array. - -TEST_F(IfElseSpecialFormTest, ResultShape) { - auto schema = arrow::schema({field("b_null", boolean()), field("b_true", boolean()), - field("b_false", boolean())}); - auto b_null = field_ref("b_null"); - auto b_true = field_ref("b_true"); - auto b_false = field_ref("b_false"); - auto batch = ExecBatch( - { - *ArrayFromJSON(boolean(), "[null, null, null]"), - *ArrayFromJSON(boolean(), "[true, true, true]"), - *ArrayFromJSON(boolean(), "[false, false, false]"), - }, - 3); - { - ARROW_SCOPED_TRACE("if (null) then 1 else 0"); - auto expected = ArrayFromJSON(int32(), "[null, null, null]"); - for (const auto& cond : {kBooleanNull, b_null}) { - ARROW_SCOPED_TRACE("cond: " + cond.ToString()); - for (const auto& if_else_sp : - SuppressSelectionVectorAwareForIfElse(cond, literal(1), literal(0))) { - ARROW_SCOPED_TRACE(if_else_sp.ToString()); - ASSERT_OK_AND_ASSIGN(auto result, ExecuteExpr(if_else_sp, schema, batch)); - AssertDatumsEqual(expected, result); - } - } - } - { - ARROW_SCOPED_TRACE("if (true) then 1 else 0"); - auto expected = MakeScalar(1); - for (const auto& cond : {literal(true), b_true}) { - ARROW_SCOPED_TRACE("cond: " + cond.ToString()); - for (const auto& if_else_sp : - SuppressSelectionVectorAwareForIfElse(cond, literal(1), literal(0))) { - ARROW_SCOPED_TRACE(if_else_sp.ToString()); - ASSERT_OK_AND_ASSIGN(auto result, ExecuteExpr(if_else_sp, schema, batch)); - AssertDatumsEqual(expected, result); - } - } - } - { - ARROW_SCOPED_TRACE("if (false) then 1 else 0"); - auto expected = MakeScalar(0); - for (const auto& cond : {literal(false), b_false}) { - ARROW_SCOPED_TRACE("cond: " + cond.ToString()); - for (const auto& if_else_sp : - SuppressSelectionVectorAwareForIfElse(cond, literal(1), literal(0))) { - ARROW_SCOPED_TRACE(if_else_sp.ToString()); - ASSERT_OK_AND_ASSIGN(auto result, ExecuteExpr(if_else_sp, schema, batch)); - AssertDatumsEqual(expected, result); - } - } - } - { - ARROW_SCOPED_TRACE("if (if (null) true then false) then 1 else 0"); - auto expected = ArrayFromJSON(int32(), "[null, null, null]"); - for (const auto& nested_cond : {kBooleanNull, b_null}) { - auto cond = if_else_special(nested_cond, literal(true), literal(false)); - ARROW_SCOPED_TRACE("nested cond: " + cond.ToString()); - for (const auto& if_else_sp : - SuppressSelectionVectorAwareForIfElse(cond, literal(1), literal(0))) { - ARROW_SCOPED_TRACE(if_else_sp.ToString()); - ASSERT_OK_AND_ASSIGN(auto result, ExecuteExpr(if_else_sp, schema, batch)); - AssertDatumsEqual(expected, result); - } - } - } - { - ARROW_SCOPED_TRACE("if (if (true) then true else false) then 1 else 0"); - auto expected = MakeScalar(1); - for (const auto& nested_cond : {literal(true), b_true}) { - auto cond = if_else_special(nested_cond, nested_cond, literal(false)); - ARROW_SCOPED_TRACE("nested cond: " + cond.ToString()); - for (const auto& if_else_sp : - SuppressSelectionVectorAwareForIfElse(cond, literal(1), literal(0))) { - ARROW_SCOPED_TRACE(if_else_sp.ToString()); - ASSERT_OK_AND_ASSIGN(auto result, ExecuteExpr(if_else_sp, schema, batch)); - AssertDatumsEqual(expected, result); - } - } - } - { - ARROW_SCOPED_TRACE("if (if (false) then true else false) then 1 else 0"); - auto expected = MakeScalar(0); - for (const auto& nested_cond : {literal(false), b_false}) { - auto cond = if_else_special(nested_cond, literal(true), nested_cond); - ARROW_SCOPED_TRACE("nested cond: " + cond.ToString()); - for (const auto& if_else_sp : - SuppressSelectionVectorAwareForIfElse(cond, literal(1), literal(0))) { - ARROW_SCOPED_TRACE(if_else_sp.ToString()); - ASSERT_OK_AND_ASSIGN(auto result, ExecuteExpr(if_else_sp, schema, batch)); - AssertDatumsEqual(expected, result); +TEST(IfElseSpecial, ExecuteScalar) { + for (const auto& cond_expr : + {literal(MakeNullScalar(boolean())), literal(true), literal(false)}) { + std::vector> literals_list{ + {literal(MakeNullScalar(int32())), literal(0), literal(1), + literal(std::numeric_limits::min()), + literal(std::numeric_limits::max())}, + {literal(MakeNullScalar(utf8())), literal(""), literal("foo"), literal("bar")}}; + for (const auto& literals : literals_list) { + for (const auto& if_true_expr : literals) { + for (const auto& if_false_expr : literals) { + for (int64_t length : {1, 2}) { + CheckIfElseSpecial(cond_expr, if_true_expr, if_false_expr, *schema({}), + ExecBatch({}, length)); + } + } } } } - // TODO: Non-scalar branch bodies. } -namespace { +TEST(IfElseSpecial, ExecuteBasic) { + const int64_t length = 7; -auto kCanonicalSchema = arrow::schema({field("a", boolean()), field("b", int32())}); - -auto kCanonicalA = field_ref("a"); -auto kCanonicalB = field_ref("b"); - -const auto kCanonicalBooleanCols = {kBooleanNull, literal(true), literal(false), - kCanonicalA}; -const auto kCanonicalIntCols = {kIntNull, literal(42), kCanonicalB}; - -const std::vector kCanonicalBatches = { - ExecBatch(*RecordBatchFromJSON(kCanonicalSchema, R"([ - [null, 0], - [null, null], - [null, 1] - ])")), - ExecBatch(*RecordBatchFromJSON(kCanonicalSchema, R"([ - [true, 0], - [true, null], - [true, 1] - ])")), - ExecBatch(*RecordBatchFromJSON(kCanonicalSchema, R"([ - [false, 0], - [false, null], - [false, 1] - ])")), - ExecBatch(*RecordBatchFromJSON(kCanonicalSchema, R"([ - [false, 0], - [true, null], - [null, 1] - ])")), -}; + auto schm = schema( + {field("cond", boolean()), field("if_true", int32()), field("if_false", int32())}); -} // namespace + auto cond_arr = + ArrayFromJSON(boolean(), "[null, true, false, true, false, null, true]"); + auto cond_chunked = ChunkedArrayFromJSON( + boolean(), {"[null, true, false]", "[]", "[true, false, null, true]"}); -TEST_F(IfElseSpecialFormTest, Simple) { - const auto& schema = kCanonicalSchema; - const auto& boolean_datums = kCanonicalBooleanCols; - const auto& int_datums = kCanonicalIntCols; - const auto& batches = kCanonicalBatches; - for (const auto& cond : boolean_datums) { - for (const auto& if_true : int_datums) { - for (const auto& if_false : int_datums) { - for (const auto& batch : batches) { - CheckIfElseIgnoreShape(cond, if_true, if_false, schema, batch); - } - } - } - } -} + auto if_true_arr = ArrayFromJSON(int32(), "[1, 2, 3, 4, 5, 6, 7]"); + auto if_true_chunked = ChunkedArrayFromJSON(int32(), {"[1, 2]", "[3, 4, 5]", "[6, 7]"}); -TEST_F(IfElseSpecialFormTest, NestedSimple) { - const auto& schema = kCanonicalSchema; - const auto& a = kCanonicalA; - const auto& b = kCanonicalB; - ExecBatch batch(*RecordBatchFromJSON(kCanonicalSchema, R"([ - [false, 0], - [true, null], - [null, 1] - ])")); - for (const auto& cond : { - if_else_special(a, kBooleanNull, a), - if_else_special(a, a, literal(true)), - }) { - for (const auto& if_true : { - if_else_special(a, kIntNull, b), - if_else_special(a, b, literal(42)), - }) { - for (const auto& if_false : { - if_else_special(a, kIntNull, b), - if_else_special(a, b, literal(42)), - }) { - CheckIfElseIgnoreShape(cond, if_true, if_false, schema, batch); - } - } - } -} + auto if_false_arr = ArrayFromJSON(int32(), "[10, 20, 30, 40, 50, 60, 70]"); + auto if_false_chunked = + ChunkedArrayFromJSON(int32(), {"[10, 20, 30]", "[]", "[40, 50, 60, 70]"}); -// TODO: Deprecate this test due to slowness. -TEST_F(IfElseSpecialFormTest, NestedConditionComplex) { - const auto& batches = kCanonicalBatches; - const auto& schema = kCanonicalSchema; - const auto& boolean_datums = kCanonicalBooleanCols; - const auto& int_datums = kCanonicalIntCols; - for (const auto& nested_cond : boolean_datums) { - for (const auto& nested_if_true : boolean_datums) { - for (const auto& nested_if_false : boolean_datums) { - auto nested_if_else_sp = - if_else_special(nested_cond, nested_if_true, nested_if_false); - for (const auto& if_true : int_datums) { - for (const auto& if_false : int_datums) { - for (const auto& batch : batches) { - CheckIfElseIgnoreShape(nested_if_else_sp, if_true, if_false, schema, batch); + for (const auto& cond_expr : {literal(MakeNullScalar(boolean())), literal(true), + literal(false), field_ref("cond")}) { + for (const auto& if_true_expr : + {literal(MakeNullScalar(int32())), literal(42), + literal(std::numeric_limits::max()), field_ref("if_true")}) { + for (const auto& if_false_expr : + {literal(MakeNullScalar(int32())), literal(24), + literal(std::numeric_limits::min()), field_ref("if_false")}) { + ARROW_SCOPED_TRACE( + "expression: " + + if_else_special(cond_expr, if_true_expr, if_false_expr).ToString()); + + for (const auto& cond_datum : {Datum(cond_arr), Datum(cond_chunked)}) { + ARROW_SCOPED_TRACE("cond: " + cond_datum.ToString()); + for (const auto& if_true_datum : {Datum(if_true_arr), Datum(if_true_chunked)}) { + ARROW_SCOPED_TRACE("if_true: " + if_true_datum.ToString()); + for (const auto& if_false_datum : + {Datum(if_false_arr), Datum(if_false_chunked)}) { + ARROW_SCOPED_TRACE("if_false: " + if_false_datum.ToString()); + CheckIfElseSpecial( + cond_expr, if_true_expr, if_false_expr, *schm, + ExecBatch({cond_datum, if_true_datum, if_false_datum}, length)); } } } @@ -1074,48 +531,40 @@ TEST_F(IfElseSpecialFormTest, NestedConditionComplex) { } } -// TODO: Deprecate this test due to slowness. -TEST_F(IfElseSpecialFormTest, NestedBodyComplex) { - const auto& batches = kCanonicalBatches; - const auto& schema = kCanonicalSchema; - const auto& boolean_datums = kCanonicalBooleanCols; - const auto& int_datums = kCanonicalIntCols; - for (const auto& cond : boolean_datums) { - for (const auto& nested_cond : boolean_datums) { - for (const auto& nested_if_true : int_datums) { - for (const auto& nested_if_false : int_datums) { - auto nested_if_else_sp = - if_else_special(nested_cond, nested_if_true, nested_if_false); - for (const auto& batch : batches) { - CheckIfElseIgnoreShape(cond, nested_if_else_sp, nested_if_else_sp, schema, - batch); - } - } - } - } - } -} +TEST(IfElseSpecial, ExecuteBasicDebug) { + const int64_t length = 7; -// TODO: Deprecate this test due to slowness. -TEST_F(IfElseSpecialFormTest, NestedComplex) { - const auto& batches = kCanonicalBatches; - const auto& schema = kCanonicalSchema; - const auto& boolean_datums = kCanonicalBooleanCols; - const auto& int_datums = kCanonicalIntCols; - for (const auto& cond_nested_cond : boolean_datums) { - for (const auto& cond_nested_if_true : boolean_datums) { - for (const auto& cond_nested_if_false : boolean_datums) { - auto cond = - if_else_special(cond_nested_cond, cond_nested_if_true, cond_nested_if_false); - for (const auto& nested_cond : boolean_datums) { - for (const auto& nested_if_true : int_datums) { - for (const auto& nested_if_false : int_datums) { - auto nested_if_else_sp = - if_else_special(nested_cond, nested_if_true, nested_if_false); - for (const auto& batch : batches) { - CheckIfElseIgnoreShape(cond, nested_if_else_sp, nested_if_else_sp, schema, - batch); - } + auto schm = schema( + {field("cond", boolean()), field("if_true", int32()), field("if_false", int32())}); + + auto cond_arr = + ArrayFromJSON(boolean(), "[null, true, false, true, false, null, true]"); + // auto cond_chunked = ChunkedArrayFromJSON( + // boolean(), {"[null, true, false]", "[]", "[true, false, null, true]"}); + + // auto if_true_arr = ArrayFromJSON(int32(), "[1, 2, 3, 4, 5, 6, 7]"); + auto if_true_chunked = ChunkedArrayFromJSON(int32(), {"[1, 2]", "[3, 4, 5]", "[6, 7]"}); + + auto if_false_arr = ArrayFromJSON(int32(), "[10, 20, 30, 40, 50, 60, 70]"); + // auto if_false_chunked = + // ChunkedArrayFromJSON(int32(), {"[10, 20, 30]", "[]", "[40, 50, 60, 70]"}); + + for (const auto& cond_expr : {field_ref("cond")}) { + for (const auto& if_true_expr : {field_ref("if_true")}) { + for (const auto& if_false_expr : {literal(24)}) { + ARROW_SCOPED_TRACE( + "expression: " + + if_else_special(cond_expr, if_true_expr, if_false_expr).ToString()); + + for (const auto& cond_datum : {Datum(cond_arr)}) { + ARROW_SCOPED_TRACE("cond: " + cond_datum.ToString()); + for (const auto& if_true_datum : {Datum(if_true_chunked)}) { + ARROW_SCOPED_TRACE("if_true: " + if_true_datum.ToString()); + for (const auto& if_false_datum : {Datum(if_false_arr)}) { + ARROW_SCOPED_TRACE("if_false: " + if_false_datum.ToString()); + CheckIfElseSpecial( + cond_expr, if_true_expr, if_false_expr, *schm, + ExecBatch({cond_datum, if_true_datum, if_false_datum}, length)); } } } @@ -1124,81 +573,162 @@ TEST_F(IfElseSpecialFormTest, NestedComplex) { } } -// TODO: ChunkedArray. - // namespace { -// template -// Status ConstantKernelExec(KernelContext*, const ExecSpan& span, ExecResult* out) { -// DCHECK_EQ(span.num_values(), 1); -// DCHECK_EQ(span.length, 1); -// DCHECK(out->is_array_span()); -// DCHECK_EQ(out->length(), 1); -// if constexpr (!selection_vector_aware) { -// if (span.selection_vector->length() > 0) { -// return Status::Invalid("There is a selection vector"); -// } -// } -// int32_t* out_values = out->array_span_mutable()->GetValues(1); -// *out_values = 0; -// return Status::OK(); -// } -// static Status RegisterConstantFunctions() { -// auto registry = GetFunctionRegistry(); +// auto kCanonicalSchema = arrow::schema({field("a", boolean()), field("b", int32())}); + +// auto kCanonicalA = field_ref("a"); +// auto kCanonicalB = field_ref("b"); + +// const auto kCanonicalBooleanCols = {kBooleanNull, literal(true), literal(false), +// kCanonicalA}; +// const auto kCanonicalIntCols = {kIntNull, literal(42), kCanonicalB}; + +// const std::vector kCanonicalBatches = { +// ExecBatch(*RecordBatchFromJSON(kCanonicalSchema, R"([ +// [null, 0], +// [null, null], +// [null, 1] +// ])")), +// ExecBatch(*RecordBatchFromJSON(kCanonicalSchema, R"([ +// [true, 0], +// [true, null], +// [true, 1] +// ])")), +// ExecBatch(*RecordBatchFromJSON(kCanonicalSchema, R"([ +// [false, 0], +// [false, null], +// [false, 1] +// ])")), +// ExecBatch(*RecordBatchFromJSON(kCanonicalSchema, R"([ +// [false, 0], +// [true, null], +// [null, 1] +// ])")), +// }; -// auto register_test_func = [&](const std::string& name, -// bool selection_vector_aware) -> Status { -// auto zero = -// std::make_shared(name, Arity::Unary(), FunctionDoc::Empty()); +// } // namespace -// ArrayKernelExec exec; -// if (selection_vector_aware) { -// exec = ConstantKernelExec; -// } else { -// exec = ConstantKernelExec; +// TEST_F(IfElseSpecialFormTest, Simple) { +// const auto& schema = kCanonicalSchema; +// const auto& boolean_datums = kCanonicalBooleanCols; +// const auto& int_datums = kCanonicalIntCols; +// const auto& batches = kCanonicalBatches; +// for (const auto& cond : boolean_datums) { +// for (const auto& if_true : int_datums) { +// for (const auto& if_false : int_datums) { +// for (const auto& batch : batches) { +// CheckIfElseIgnoreShape(cond, if_true, if_false, schema, batch); +// } +// } // } -// ScalarKernel kernel({InputType::Any()}, OutputType{int32()}, std::move(exec)); -// kernel.selection_vector_aware = selection_vector_aware; -// kernel.can_write_into_slices = true; -// kernel.null_handling = NullHandling::OUTPUT_NOT_NULL; -// kernel.mem_allocation = MemAllocation::PREALLOCATE; -// RETURN_NOT_OK(zero->AddKernel(kernel)); -// RETURN_NOT_OK(registry->AddFunction(std::move(zero))); -// return Status::OK(); -// }; +// } +// } -// RETURN_NOT_OK(register_test_func("zero_panic", false)); -// RETURN_NOT_OK(register_test_func("zero_calm", true)); +// TEST_F(IfElseSpecialFormTest, NestedSimple) { +// const auto& schema = kCanonicalSchema; +// const auto& a = kCanonicalA; +// const auto& b = kCanonicalB; +// ExecBatch batch(*RecordBatchFromJSON(kCanonicalSchema, R"([ +// [false, 0], +// [true, null], +// [null, 1] +// ])")); +// for (const auto& cond : { +// if_else_special(a, kBooleanNull, a), +// if_else_special(a, a, literal(true)), +// }) { +// for (const auto& if_true : { +// if_else_special(a, kIntNull, b), +// if_else_special(a, b, literal(42)), +// }) { +// for (const auto& if_false : { +// if_else_special(a, kIntNull, b), +// if_else_special(a, b, literal(42)), +// }) { +// CheckIfElseIgnoreShape(cond, if_true, if_false, schema, batch); +// } +// } +// } +// } -// return Status::OK(); +// // TODO: Deprecate this test due to slowness. +// TEST_F(IfElseSpecialFormTest, NestedConditionComplex) { +// const auto& batches = kCanonicalBatches; +// const auto& schema = kCanonicalSchema; +// const auto& boolean_datums = kCanonicalBooleanCols; +// const auto& int_datums = kCanonicalIntCols; +// for (const auto& nested_cond : boolean_datums) { +// for (const auto& nested_if_true : boolean_datums) { +// for (const auto& nested_if_false : boolean_datums) { +// auto nested_if_else_sp = +// if_else_special(nested_cond, nested_if_true, nested_if_false); +// for (const auto& if_true : int_datums) { +// for (const auto& if_false : int_datums) { +// for (const auto& batch : batches) { +// CheckIfElseIgnoreShape(nested_if_else_sp, if_true, if_false, schema, +// batch); +// } +// } +// } +// } +// } +// } // } -// } // namespace +// // TODO: Deprecate this test due to slowness. +// TEST_F(IfElseSpecialFormTest, NestedBodyComplex) { +// const auto& batches = kCanonicalBatches; +// const auto& schema = kCanonicalSchema; +// const auto& boolean_datums = kCanonicalBooleanCols; +// const auto& int_datums = kCanonicalIntCols; +// for (const auto& cond : boolean_datums) { +// for (const auto& nested_cond : boolean_datums) { +// for (const auto& nested_if_true : int_datums) { +// for (const auto& nested_if_false : int_datums) { +// auto nested_if_else_sp = +// if_else_special(nested_cond, nested_if_true, nested_if_false); +// for (const auto& batch : batches) { +// CheckIfElseIgnoreShape(cond, nested_if_else_sp, nested_if_else_sp, +// schema, +// batch); +// } +// } +// } +// } +// } +// } -// TEST(IfElseSpecialForm, Reference) { -// ASSERT_OK(RegisterConstantFunctions()); - -// auto schema = arrow::schema({field("a", int32()), field("b", int32())}); -// std::vector batches = { -// ExecBatch(*RecordBatchFromJSON(schema, R"([])")), -// ExecBatch(*RecordBatchFromJSON(schema, R"([ -// [1, 0], -// [1, 0], -// [1, 0] -// ])")), -// ExecBatch(*RecordBatchFromJSON(schema, R"([ -// [1, 0], -// [null, 0], -// [1, null] -// ])")), -// }; -// for (const auto& batch : batches) { -// auto expr = call("zero_panic", {literal(42)}); -// ASSERT_OK_AND_ASSIGN(auto bound, expr.Bind(*schema)); -// ASSERT_OK_AND_ASSIGN(auto result, ExecuteScalarExpression(bound, batch)); -// std::cout << result.ToString() << std::endl; +// // TODO: Deprecate this test due to slowness. +// TEST_F(IfElseSpecialFormTest, NestedComplex) { +// const auto& batches = kCanonicalBatches; +// const auto& schema = kCanonicalSchema; +// const auto& boolean_datums = kCanonicalBooleanCols; +// const auto& int_datums = kCanonicalIntCols; +// for (const auto& cond_nested_cond : boolean_datums) { +// for (const auto& cond_nested_if_true : boolean_datums) { +// for (const auto& cond_nested_if_false : boolean_datums) { +// auto cond = +// if_else_special(cond_nested_cond, cond_nested_if_true, +// cond_nested_if_false); +// for (const auto& nested_cond : boolean_datums) { +// for (const auto& nested_if_true : int_datums) { +// for (const auto& nested_if_false : int_datums) { +// auto nested_if_else_sp = +// if_else_special(nested_cond, nested_if_true, nested_if_false); +// for (const auto& batch : batches) { +// CheckIfElseIgnoreShape(cond, nested_if_else_sp, nested_if_else_sp, +// schema, +// batch); +// } +// } +// } +// } +// } +// } // } -// // TODO: The result shape of exec_chunksize, chunked input, and preallocate. // } +// // TODO: ChunkedArray. + } // namespace arrow::compute From 52627a5d4dacb65746d0da6a5cce01d23da92e90 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Tue, 14 Oct 2025 00:23:15 -0700 Subject: [PATCH 52/71] Fix --- .../compute/special/if_else_special_test.cc | 43 ------------------- 1 file changed, 43 deletions(-) diff --git a/cpp/src/arrow/compute/special/if_else_special_test.cc b/cpp/src/arrow/compute/special/if_else_special_test.cc index b5b8e86436a2..897644b81967 100644 --- a/cpp/src/arrow/compute/special/if_else_special_test.cc +++ b/cpp/src/arrow/compute/special/if_else_special_test.cc @@ -512,7 +512,6 @@ TEST(IfElseSpecial, ExecuteBasic) { ARROW_SCOPED_TRACE( "expression: " + if_else_special(cond_expr, if_true_expr, if_false_expr).ToString()); - for (const auto& cond_datum : {Datum(cond_arr), Datum(cond_chunked)}) { ARROW_SCOPED_TRACE("cond: " + cond_datum.ToString()); for (const auto& if_true_datum : {Datum(if_true_arr), Datum(if_true_chunked)}) { @@ -531,48 +530,6 @@ TEST(IfElseSpecial, ExecuteBasic) { } } -TEST(IfElseSpecial, ExecuteBasicDebug) { - const int64_t length = 7; - - auto schm = schema( - {field("cond", boolean()), field("if_true", int32()), field("if_false", int32())}); - - auto cond_arr = - ArrayFromJSON(boolean(), "[null, true, false, true, false, null, true]"); - // auto cond_chunked = ChunkedArrayFromJSON( - // boolean(), {"[null, true, false]", "[]", "[true, false, null, true]"}); - - // auto if_true_arr = ArrayFromJSON(int32(), "[1, 2, 3, 4, 5, 6, 7]"); - auto if_true_chunked = ChunkedArrayFromJSON(int32(), {"[1, 2]", "[3, 4, 5]", "[6, 7]"}); - - auto if_false_arr = ArrayFromJSON(int32(), "[10, 20, 30, 40, 50, 60, 70]"); - // auto if_false_chunked = - // ChunkedArrayFromJSON(int32(), {"[10, 20, 30]", "[]", "[40, 50, 60, 70]"}); - - for (const auto& cond_expr : {field_ref("cond")}) { - for (const auto& if_true_expr : {field_ref("if_true")}) { - for (const auto& if_false_expr : {literal(24)}) { - ARROW_SCOPED_TRACE( - "expression: " + - if_else_special(cond_expr, if_true_expr, if_false_expr).ToString()); - - for (const auto& cond_datum : {Datum(cond_arr)}) { - ARROW_SCOPED_TRACE("cond: " + cond_datum.ToString()); - for (const auto& if_true_datum : {Datum(if_true_chunked)}) { - ARROW_SCOPED_TRACE("if_true: " + if_true_datum.ToString()); - for (const auto& if_false_datum : {Datum(if_false_arr)}) { - ARROW_SCOPED_TRACE("if_false: " + if_false_datum.ToString()); - CheckIfElseSpecial( - cond_expr, if_true_expr, if_false_expr, *schm, - ExecBatch({cond_datum, if_true_datum, if_false_datum}, length)); - } - } - } - } - } - } -} - // namespace { // auto kCanonicalSchema = arrow::schema({field("a", boolean()), field("b", int32())}); From b50a8743874edbe56a86129337ce8df30436d8b3 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Tue, 14 Oct 2025 11:50:21 -0700 Subject: [PATCH 53/71] Reorder tests --- .../compute/special/if_else_special_test.cc | 66 +++++++------------ 1 file changed, 23 insertions(+), 43 deletions(-) diff --git a/cpp/src/arrow/compute/special/if_else_special_test.cc b/cpp/src/arrow/compute/special/if_else_special_test.cc index 897644b81967..7e144f29c120 100644 --- a/cpp/src/arrow/compute/special/if_else_special_test.cc +++ b/cpp/src/arrow/compute/special/if_else_special_test.cc @@ -440,49 +440,6 @@ void CheckIfElseSpecial(Expression cond, Expression if_true, Expression if_false } // namespace -// GH-XXXXX: -TEST(IfElseSpecial, IfNotZeroThenDivide) { - // if (b != 0) then (a / b) else b - auto cond = call("not_equal", {field_ref("b"), literal(0)}); - auto if_true = call("divide", {field_ref("a"), field_ref("b")}); - auto if_false = field_ref("b"); - - auto schm = schema({field("a", int32()), field("b", int32())}); - auto batch = ExecBatchFromJSON({int32(), int32()}, - R"([[1, 1], - [2, 1], - [3, 0], - [4, 1], - [5, 1]])"); - - ASSERT_OK_AND_ASSIGN(auto result, - ExecuteIfElseSpecial(cond, if_true, if_false, *schm, batch)); - auto expected = ArrayFromJSON(int32(), "[1, 2, 0, 4, 5]"); - - AssertDatumsEqual(expected, result); -} - -TEST(IfElseSpecial, ExecuteScalar) { - for (const auto& cond_expr : - {literal(MakeNullScalar(boolean())), literal(true), literal(false)}) { - std::vector> literals_list{ - {literal(MakeNullScalar(int32())), literal(0), literal(1), - literal(std::numeric_limits::min()), - literal(std::numeric_limits::max())}, - {literal(MakeNullScalar(utf8())), literal(""), literal("foo"), literal("bar")}}; - for (const auto& literals : literals_list) { - for (const auto& if_true_expr : literals) { - for (const auto& if_false_expr : literals) { - for (int64_t length : {1, 2}) { - CheckIfElseSpecial(cond_expr, if_true_expr, if_false_expr, *schema({}), - ExecBatch({}, length)); - } - } - } - } - } -} - TEST(IfElseSpecial, ExecuteBasic) { const int64_t length = 7; @@ -530,6 +487,29 @@ TEST(IfElseSpecial, ExecuteBasic) { } } +// GH-41094: Maskable execution of division that otherwise would error on division by +// zero. +TEST(IfElseSpecial, IfNotZeroThenDivide) { + // if (b != 0) then (a / b) else b + auto cond = call("not_equal", {field_ref("b"), literal(0)}); + auto if_true = call("divide", {field_ref("a"), field_ref("b")}); + auto if_false = field_ref("b"); + + auto schm = schema({field("a", int32()), field("b", int32())}); + auto batch = ExecBatchFromJSON({int32(), int32()}, + R"([[1, 1], + [2, 1], + [3, 0], + [4, 1], + [5, 1]])"); + + ASSERT_OK_AND_ASSIGN(auto result, + ExecuteIfElseSpecial(cond, if_true, if_false, *schm, batch)); + auto expected = ArrayFromJSON(int32(), "[1, 2, 0, 4, 5]"); + + AssertDatumsEqual(expected, result); +} + // namespace { // auto kCanonicalSchema = arrow::schema({field("a", boolean()), field("b", int32())}); From efbd80947e728360b29fd1cbdc344e1b508d6b5f Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Tue, 14 Oct 2025 21:39:07 -0700 Subject: [PATCH 54/71] Refine if else tests --- .../arrow/compute/expression_test_internal.h | 4 + .../compute/special/if_else_special_test.cc | 400 +++++++++++++++--- 2 files changed, 349 insertions(+), 55 deletions(-) diff --git a/cpp/src/arrow/compute/expression_test_internal.h b/cpp/src/arrow/compute/expression_test_internal.h index 25d3d6787e16..ccb3a5768933 100644 --- a/cpp/src/arrow/compute/expression_test_internal.h +++ b/cpp/src/arrow/compute/expression_test_internal.h @@ -57,6 +57,10 @@ inline Expression add(Expression l, Expression r) { return call("add", {std::move(l), std::move(r)}); } +inline Expression sub(Expression l, Expression r) { + return call("subtract", {std::move(l), std::move(r)}); +} + inline std::string make_range_json(int start, int end) { std::string result = "["; for (int i = start; i <= end; ++i) { diff --git a/cpp/src/arrow/compute/special/if_else_special_test.cc b/cpp/src/arrow/compute/special/if_else_special_test.cc index 7e144f29c120..effcb9384e16 100644 --- a/cpp/src/arrow/compute/special/if_else_special_test.cc +++ b/cpp/src/arrow/compute/special/if_else_special_test.cc @@ -34,6 +34,7 @@ using internal::cast; using internal::ExpectBindsTo; using internal::kBoringSchema; using internal::no_change; +using internal::sub; TEST(IfElseSpecial, ToString) { EXPECT_EQ( @@ -387,21 +388,12 @@ Result ExecuteExpr(Expression expr, const Schema& schema, const ExecBatch return ExecuteScalarExpression(bound, batch, exec_context); } -Result ExecuteIfElse(Expression cond, Expression if_true, Expression if_false, - const Schema& schema, const ExecBatch& batch, - ExecContext* exec_context = default_exec_context()) { - return ExecuteExpr(if_else(std::move(cond), std::move(if_true), std::move(if_false)), - schema, batch, exec_context); -} - -Result ExecuteIfElseSpecial(Expression cond, Expression if_true, - Expression if_false, const Schema& schema, - const ExecBatch& batch, - ExecContext* exec_context = default_exec_context()) { - return ExecuteExpr( - if_else_special(std::move(cond), std::move(if_true), std::move(if_false)), schema, - batch, exec_context); -} +// Result ExecuteIfElse(Expression cond, Expression if_true, Expression if_false, +// const Schema& schema, const ExecBatch& batch, +// ExecContext* exec_context = default_exec_context()) { +// return ExecuteExpr(if_else(std::move(cond), std::move(if_true), std::move(if_false)), +// schema, batch, exec_context); +// } void AssertDatumsEqualIgnoreShape(const Datum& expected, const Datum& result) { DCHECK(expected.is_scalar() || expected.is_array() || expected.is_chunked_array()); @@ -428,65 +420,363 @@ void AssertDatumsEqualIgnoreShape(const Datum& expected, const Datum& result) { ASSERT_OK_AND_ASSIGN(auto result_array, to_array(result)); } -void CheckIfElseSpecial(Expression cond, Expression if_true, Expression if_false, +using MakeIfElseFunc = std::function; + +using MakeExprContainingIfElseFunc = + std::function; + +void CheckIfElseSpecial(MakeExprContainingIfElseFunc make_expr_containing_if_else, const Schema& schema, const ExecBatch& batch, ExecContext* exec_context = default_exec_context()) { - ASSERT_OK_AND_ASSIGN( - auto expected, ExecuteIfElse(cond, if_true, if_false, schema, batch, exec_context)); - ASSERT_OK_AND_ASSIGN(auto result, ExecuteIfElseSpecial(cond, if_true, if_false, schema, - batch, exec_context)); + auto if_else_expr = make_expr_containing_if_else(if_else); + ASSERT_OK_AND_ASSIGN(auto expected, + ExecuteExpr(if_else_expr, schema, batch, exec_context)); + auto if_else_sp_expr = make_expr_containing_if_else(if_else_special); + ASSERT_OK_AND_ASSIGN(auto result, + ExecuteExpr(if_else_sp_expr, schema, batch, exec_context)); AssertDatumsEqualIgnoreShape(expected, result); } +void CheckIfElseSpecial(Expression cond, Expression if_true, Expression if_false, + const Schema& schema, const ExecBatch& batch, + ExecContext* exec_context = default_exec_context()) { + CheckIfElseSpecial( + [=](MakeIfElseFunc make_if_else) { return if_else(cond, if_true, if_false); }, + schema, batch, exec_context); +} + } // namespace -TEST(IfElseSpecial, ExecuteBasic) { +class TestIfElseSpecialExecute : public ::testing::Test { + protected: const int64_t length = 7; - auto schm = schema( - {field("cond", boolean()), field("if_true", int32()), field("if_false", int32())}); - - auto cond_arr = + std::shared_ptr schm = + schema({field("boolean1", boolean()), field("boolean2", boolean()), + field("int1", int32()), field("int2", int32())}); + Expression boolean1 = field_ref("boolean1"); + Expression boolean2 = field_ref("boolean2"); + Expression int1 = field_ref("int1"); + Expression int2 = field_ref("int2"); + + std::vector boolean_literals = {literal(MakeNullScalar(boolean())), + literal(true), literal(false)}; + std::vector boolean_fields = {boolean1, boolean2}; + std::vector boolean_complex_exprs = {and_(boolean1, boolean2), + or_(boolean1, boolean2)}; + + std::vector int_literals = {literal(MakeNullScalar(int32())), literal(42)}; + std::vector int_fields = {int1, int2}; + std::vector int_complex_exprs = {add(int1, int2), sub(int1, int2)}; + + std::vector boolean_scalars = {Datum(MakeNullScalar(boolean())), + Datum(MakeScalar(true)), + Datum(MakeScalar(false))}; + std::shared_ptr boolean1_arr = ArrayFromJSON(boolean(), "[null, true, false, true, false, null, true]"); - auto cond_chunked = ChunkedArrayFromJSON( + std::shared_ptr boolean1_chunked = ChunkedArrayFromJSON( boolean(), {"[null, true, false]", "[]", "[true, false, null, true]"}); - - auto if_true_arr = ArrayFromJSON(int32(), "[1, 2, 3, 4, 5, 6, 7]"); - auto if_true_chunked = ChunkedArrayFromJSON(int32(), {"[1, 2]", "[3, 4, 5]", "[6, 7]"}); - - auto if_false_arr = ArrayFromJSON(int32(), "[10, 20, 30, 40, 50, 60, 70]"); - auto if_false_chunked = - ChunkedArrayFromJSON(int32(), {"[10, 20, 30]", "[]", "[40, 50, 60, 70]"}); - - for (const auto& cond_expr : {literal(MakeNullScalar(boolean())), literal(true), - literal(false), field_ref("cond")}) { - for (const auto& if_true_expr : - {literal(MakeNullScalar(int32())), literal(42), - literal(std::numeric_limits::max()), field_ref("if_true")}) { - for (const auto& if_false_expr : - {literal(MakeNullScalar(int32())), literal(24), - literal(std::numeric_limits::min()), field_ref("if_false")}) { - ARROW_SCOPED_TRACE( - "expression: " + - if_else_special(cond_expr, if_true_expr, if_false_expr).ToString()); - for (const auto& cond_datum : {Datum(cond_arr), Datum(cond_chunked)}) { - ARROW_SCOPED_TRACE("cond: " + cond_datum.ToString()); - for (const auto& if_true_datum : {Datum(if_true_arr), Datum(if_true_chunked)}) { - ARROW_SCOPED_TRACE("if_true: " + if_true_datum.ToString()); - for (const auto& if_false_datum : - {Datum(if_false_arr), Datum(if_false_chunked)}) { - ARROW_SCOPED_TRACE("if_false: " + if_false_datum.ToString()); - CheckIfElseSpecial( - cond_expr, if_true_expr, if_false_expr, *schm, - ExecBatch({cond_datum, if_true_datum, if_false_datum}, length)); + std::shared_ptr boolean2_arr = + ArrayFromJSON(boolean(), "[true, false, true, true, null, false, true]"); + std::shared_ptr boolean2_chunked = ChunkedArrayFromJSON( + boolean(), {"[true, false]", "[true]", "[true]", "[null, false, true]"}); + std::vector boolean_arrays = {Datum(boolean1_arr), Datum(boolean1_chunked), + Datum(boolean2_arr), Datum(boolean2_chunked)}; + + std::vector int_scalars = {Datum(MakeNullScalar(int32())), + Datum(MakeScalar(42))}; + std::shared_ptr int1_arr = ArrayFromJSON(int32(), "[0, 1, 2, 3, 4, 5, 6]"); + std::shared_ptr int1_chunked = + ChunkedArrayFromJSON(int32(), {"[0, 1]", "[2, 3, 4]", "[5, 6]"}); + std::shared_ptr int2_arr = ArrayFromJSON(int32(), "[0, 10, 20, 30, 40, 50, 60]"); + std::shared_ptr int2_chunked = + ChunkedArrayFromJSON(int32(), {"[0]", "[10, 20]", "[30, 40, 50]", "[]", "[60]"}); + std::vector int_arrays = {Datum(int1_arr), Datum(int1_chunked), Datum(int2_arr), + Datum(int2_chunked)}; + + protected: + void DoTestBasic(const std::vector& cond_exprs, + const std::vector& if_true_exprs, + const std::vector& if_false_exprs, + const std::vector& boolean_datums, + const std::vector& int_datums) { + for (const auto& cond_expr : cond_exprs) { + for (const auto& if_true_expr : if_true_exprs) { + for (const auto& if_false_expr : if_false_exprs) { + ARROW_SCOPED_TRACE( + "if_else_special: " + + if_else_special(cond_expr, if_true_expr, if_false_expr).ToString()); + for (const auto& b1_datum : boolean_datums) { + for (const auto& b2_datum : boolean_datums) { + for (const auto& i1_datum : int_datums) { + for (const auto& i2_datum : int_datums) { + ExecBatch batch({b1_datum, b2_datum, i1_datum, i2_datum}, length); + ARROW_SCOPED_TRACE("batch: " + batch.ToString()); + CheckIfElseSpecial(cond_expr, if_true_expr, if_false_expr, *schm, + batch); + } + } } } } } } } +}; + +TEST_F(TestIfElseSpecialExecute, AllLiterals) { + DoTestBasic(boolean_literals, int_literals, int_literals, boolean_arrays, int_arrays); +} + +TEST_F(TestIfElseSpecialExecute, AllScalars) { + DoTestBasic(boolean_fields, int_fields, int_fields, boolean_scalars, int_scalars); +} + +TEST_F(TestIfElseSpecialExecute, FieldWithArrays) { + DoTestBasic(boolean_fields, int_fields, int_fields, boolean_arrays, int_arrays); +} + +TEST_F(TestIfElseSpecialExecute, ComplexExprsWithArrays) { + DoTestBasic(boolean_complex_exprs, int_complex_exprs, int_complex_exprs, boolean_arrays, + int_arrays); +} + +// TEST(IfElseSpecial, ExecuteNestedCond) { +// auto boolean_exprs = {boolean1, boolean2, and_(boolean1, boolean2), +// or_(boolean1, boolean2)}; +// for (const auto& cond_inner_expr : boolean_exprs) { +// for (const auto& if_true_inner_expr : boolean_exprs) { +// for (const auto& if_false_inner_expr : boolean_exprs) { +// ARROW_SCOPED_TRACE( +// "expression: if (" + +// if_else_special(cond_inner_expr, if_true_inner_expr, if_false_inner_expr) +// .ToString() + +// ") then int1 else int2"); +// for (const auto& boolean1_datum : +// {Datum(boolean1_arr), Datum(boolean1_chunked)}) { +// ARROW_SCOPED_TRACE("boolean1: " + boolean1_datum.ToString()); +// for (const auto& boolean2_datum : +// {Datum(boolean2_arr), Datum(boolean2_chunked)}) { +// ARROW_SCOPED_TRACE("boolean2: " + boolean2_datum.ToString()); +// for (const auto& int1_datum : {Datum(int1_arr), Datum(int1_chunked)}) { +// ARROW_SCOPED_TRACE("int1: " + int1_datum.ToString()); +// for (const auto& int2_datum : {Datum(int2_arr), Datum(int2_chunked)}) { +// ARROW_SCOPED_TRACE("int2: " + int2_datum.ToString()); +// CheckIfElseSpecial( +// [=](MakeIfElseFunc make_if_else) { +// return make_if_else( +// make_if_else(cond_inner_expr, if_true_inner_expr, +// if_false_inner_expr), +// int1, int2); +// }, +// *schm, +// ExecBatch({Datum(boolean1_arr), Datum(boolean2_arr), +// Datum(int1_arr), +// Datum(int2_arr)}, +// length)); +// } +// } +// } +// } +// } +// } +// } +// } + +// TEST(IfElseSpecial, ExecuteNestedBody) { +// const int64_t length = 7; + +// auto schm = schema({field("boolean1", boolean()), field("boolean2", boolean()), +// field("int1", int32()), field("int2", int32())}); +// auto boolean1 = field_ref("boolean1"); +// auto boolean2 = field_ref("boolean2"); +// auto int1 = field_ref("int1"); +// auto int2 = field_ref("int2"); + +// auto boolean1_arr = +// ArrayFromJSON(boolean(), "[null, true, false, true, false, null, true]"); +// auto boolean1_chunked = ChunkedArrayFromJSON( +// boolean(), {"[null, true, false]", "[]", "[true, false, null, true]"}); + +// auto boolean2_arr = +// ArrayFromJSON(boolean(), "[true, false, true, true, null, false, true]"); +// auto boolean2_chunked = ChunkedArrayFromJSON( +// boolean(), {"[true, false]", "[true]", "[true]", "[null, false, true]"}); + +// auto int1_arr = ArrayFromJSON(int32(), "[0, 1, 2, 3, 4, 5, 6]"); +// auto int1_chunked = ChunkedArrayFromJSON(int32(), {"[0, 1]", "[2, 3, 4]", "[5, +// 6]"}); + +// auto int2_arr = ArrayFromJSON(int32(), "[0, 10, 20, 30, 40, 50, 60]"); +// auto int2_chunked = +// ChunkedArrayFromJSON(int32(), {"[0]", "[10, 20]", "[30, 40, 50]", "[]", +// "[60]"}); + +// auto boolean_exprs = {boolean1, boolean2, and_(boolean1, boolean2), +// or_(boolean1, boolean2)}; +// auto int_exprs = {int1, int2, add(int1, int2)}; +// for (const auto& cond_outer_expr : boolean_exprs) { +// for (const auto& cond_inner_expr : boolean_exprs) { +// for (const auto& if_true_inner_expr : int_exprs) { +// for (const auto& if_false_inner_expr : int_exprs) { +// for (const auto& if_false_outer_expr : int_exprs) { +// ARROW_SCOPED_TRACE( +// "expression: if (" + cond_outer_expr.ToString() + ") then (" + +// if_else_special(cond_inner_expr, if_true_inner_expr, +// if_false_inner_expr) +// .ToString() + +// ") else (" + if_false_outer_expr.ToString() + ")"); +// for (const auto& boolean1_datum : +// {Datum(boolean1_arr), Datum(boolean1_chunked)}) { +// ARROW_SCOPED_TRACE("boolean1: " + boolean1_datum.ToString()); +// for (const auto& boolean2_datum : +// {Datum(boolean2_arr), Datum(boolean2_chunked)}) { +// ARROW_SCOPED_TRACE("boolean2: " + boolean2_datum.ToString()); +// for (const auto& int1_datum : {Datum(int1_arr), Datum(int1_chunked)}) +// { +// ARROW_SCOPED_TRACE("int1: " + int1_datum.ToString()); +// for (const auto& int2_datum : {Datum(int2_arr), +// Datum(int2_chunked)}) +// { +// ARROW_SCOPED_TRACE("int2: " + int2_datum.ToString()); +// CheckIfElseSpecial( +// [=](MakeIfElseFunc make_if_else) { +// return make_if_else( +// cond_outer_expr, +// make_if_else(cond_inner_expr, if_true_inner_expr, +// if_false_inner_expr), +// if_false_outer_expr); +// }, +// *schm, +// ExecBatch({Datum(boolean1_arr), Datum(boolean2_arr), +// Datum(int1_arr), Datum(int2_arr)}, +// length)); +// } +// } +// } +// } +// } + +// for (const auto& if_true_outer_expr : int_exprs) { +// ARROW_SCOPED_TRACE( +// "expression: if (" + cond_outer_expr.ToString() + ") then (" + +// if_true_outer_expr.ToString() + ") else (" + +// if_else_special(cond_inner_expr, if_true_inner_expr, +// if_false_inner_expr) +// .ToString() + +// ")"); +// for (const auto& boolean1_datum : +// {Datum(boolean1_arr), Datum(boolean1_chunked)}) { +// ARROW_SCOPED_TRACE("boolean1: " + boolean1_datum.ToString()); +// for (const auto& boolean2_datum : +// {Datum(boolean2_arr), Datum(boolean2_chunked)}) { +// ARROW_SCOPED_TRACE("boolean2: " + boolean2_datum.ToString()); +// for (const auto& int1_datum : {Datum(int1_arr), Datum(int1_chunked)}) +// { +// ARROW_SCOPED_TRACE("int1: " + int1_datum.ToString()); +// for (const auto& int2_datum : {Datum(int2_arr), +// Datum(int2_chunked)}) +// { +// ARROW_SCOPED_TRACE("int2: " + int2_datum.ToString()); +// CheckIfElseSpecial( +// [=](MakeIfElseFunc make_if_else) { +// return make_if_else( +// cond_outer_expr, if_true_outer_expr, +// make_if_else(cond_inner_expr, if_true_inner_expr, +// if_false_inner_expr)); +// }, +// *schm, +// ExecBatch({Datum(boolean1_arr), Datum(boolean2_arr), +// Datum(int1_arr), Datum(int2_arr)}, +// length)); +// } +// } +// } +// } +// } +// } +// } +// } +// } +// } + +// TEST(IfElseSpecial, ExecuteNested) { +// // if (if (boolean1) then (boolean2) else (int1 == int2))) then if () +// // if (if (boolean1) then (boolean2) else (int1 == int2)) then +// // if (boolean1) then (int1) else (int1 + int2) +// // else +// // if (boolean2) then (int2) else (int1 + int2) +// const int64_t length = 7; + +// auto schm = schema({field("boolean1", boolean()), field("boolean2", boolean()), +// field("int1", int32()), field("int2", int32())}); +// auto boolean1 = field_ref("boolean1"); +// auto boolean2 = field_ref("boolean2"); +// auto int1 = field_ref("int1"); +// auto int2 = field_ref("int2"); + +// auto boolean1_arr = +// ArrayFromJSON(boolean(), "[null, true, false, true, false, null, true]"); +// auto boolean1_chunked = ChunkedArrayFromJSON( +// boolean(), {"[null, true, false]", "[]", "[true, false, null, true]"}); + +// auto boolean2_arr = +// ArrayFromJSON(boolean(), "[true, false, true, true, null, false, true]"); +// auto boolean2_chunked = ChunkedArrayFromJSON( +// boolean(), {"[true, false]", "[true]", "[true]", "[null, false, true]"}); + +// auto int1_arr = ArrayFromJSON(int32(), "[0, 1, 2, 3, 4, 5, 6]"); +// auto int1_chunked = ChunkedArrayFromJSON(int32(), {"[0, 1]", "[2, 3, 4]", "[5, +// 6]"}); + +// auto int2_arr = ArrayFromJSON(int32(), "[0, 10, 20, 30, 40, 50, 60]"); +// auto int2_chunked = +// ChunkedArrayFromJSON(int32(), {"[0]", "[10, 20]", "[30, 40, 50]", "[]", +// "[60]"}); + +// for (const auto& cond_inner_expr : +// {literal(MakeNullScalar(boolean())), literal(true), literal(false), boolean1, +// boolean2, equal(int1, int2)}) { +// for (const auto& if_true_inner_expr : +// {literal(MakeNullScalar(int32())), literal(42), +// literal(std::numeric_limits::max()), int1, int2, add(int1, int2)}) +// { +// for (const auto& if_false_inner_expr : +// {literal(MakeNullScalar(int32())), literal(24), +// literal(std::numeric_limits::min()), int2, add(b, c)}) { +// ARROW_SCOPED_TRACE( +// "expression: " + +// if_else_special(cond_expr, if_true_expr, if_false_expr).ToString()); +// for (const auto& cond_datum : {Datum(a_arr), Datum(a_chunked)}) { +// ARROW_SCOPED_TRACE("cond: " + cond_datum.ToString()); +// for (const auto& if_true_datum : {Datum(b_arr), Datum(b_chunked)}) { +// ARROW_SCOPED_TRACE("if_true: " + if_true_datum.ToString()); +// for (const auto& if_false_datum : {Datum(c_arr), Datum(c_chunked)}) { +// ARROW_SCOPED_TRACE("if_false: " + if_false_datum.ToString()); +// CheckIfElseSpecial( +// cond_expr, if_true_expr, if_false_expr, *schm, +// ExecBatch({cond_datum, if_true_datum, if_false_datum}, length)); +// } +// } +// } +// } +// } +// } +// } + +namespace { + +Result ExecuteIfElseSpecial(Expression cond, Expression if_true, + Expression if_false, const Schema& schema, + const ExecBatch& batch, + ExecContext* exec_context = default_exec_context()) { + return ExecuteExpr( + if_else_special(std::move(cond), std::move(if_true), std::move(if_false)), schema, + batch, exec_context); } +} // namespace + // GH-41094: Maskable execution of division that otherwise would error on division by // zero. TEST(IfElseSpecial, IfNotZeroThenDivide) { From 0c655b7e1200b4f209eb018f05fff41a84f11e18 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Tue, 14 Oct 2025 23:26:36 -0700 Subject: [PATCH 55/71] WIP --- .../compute/special/if_else_special_test.cc | 90 +++++++++++++++++-- 1 file changed, 85 insertions(+), 5 deletions(-) diff --git a/cpp/src/arrow/compute/special/if_else_special_test.cc b/cpp/src/arrow/compute/special/if_else_special_test.cc index effcb9384e16..c74bcc5aaa26 100644 --- a/cpp/src/arrow/compute/special/if_else_special_test.cc +++ b/cpp/src/arrow/compute/special/if_else_special_test.cc @@ -447,7 +447,7 @@ void CheckIfElseSpecial(Expression cond, Expression if_true, Expression if_false } // namespace -class TestIfElseSpecialExecute : public ::testing::Test { +class TestExecuteIfElseSpecial : public ::testing::Test { protected: const int64_t length = 7; @@ -522,25 +522,105 @@ class TestIfElseSpecialExecute : public ::testing::Test { } } } + + void DoTestNestedCond(const std::vector& cond_exprs, + const std::vector& if_true_exprs, + const std::vector& if_false_exprs, + const std::vector& boolean_datums, + const std::vector& int_datums) { + for (const auto& cond_expr : cond_exprs) { + for (const auto& if_true_expr : if_true_exprs) { + for (const auto& if_false_expr : if_false_exprs) { + ARROW_SCOPED_TRACE( + "if_else_special: " + + if_else_special(if_else_special(cond_expr, if_true_expr, if_false_expr), + int1, int2) + .ToString()); + for (const auto& b1_datum : boolean_datums) { + for (const auto& b2_datum : boolean_datums) { + for (const auto& i1_datum : int_datums) { + for (const auto& i2_datum : int_datums) { + ExecBatch batch({b1_datum, b2_datum, i1_datum, i2_datum}, length); + ARROW_SCOPED_TRACE("batch: " + batch.ToString()); + CheckIfElseSpecial( + [=](MakeIfElseFunc make_if_else) { + return make_if_else( + make_if_else(cond_expr, if_true_expr, if_false_expr), int1, + int2); + }, + *schm, batch); + } + } + } + } + } + } + } + } + + void DoTestNestedCond(const std::vector& cond_exprs, + const std::vector& if_true_exprs, + const std::vector& if_false_exprs, + const ExecBatch& batch) { + for (const auto& cond_expr : cond_exprs) { + for (const auto& if_true_expr : if_true_exprs) { + for (const auto& if_false_expr : if_false_exprs) { + ARROW_SCOPED_TRACE( + "if_else_special: " + + if_else_special(if_else_special(cond_expr, if_true_expr, if_false_expr), + int1, int2) + .ToString()); + ARROW_SCOPED_TRACE("batch: " + batch.ToString()); + CheckIfElseSpecial( + [=](MakeIfElseFunc make_if_else) { + return make_if_else(make_if_else(cond_expr, if_true_expr, if_false_expr), + int1, int2); + }, + *schm, batch); + } + } + } + } }; -TEST_F(TestIfElseSpecialExecute, AllLiterals) { +TEST_F(TestExecuteIfElseSpecial, AllLiterals) { DoTestBasic(boolean_literals, int_literals, int_literals, boolean_arrays, int_arrays); } -TEST_F(TestIfElseSpecialExecute, AllScalars) { +TEST_F(TestExecuteIfElseSpecial, AllScalars) { DoTestBasic(boolean_fields, int_fields, int_fields, boolean_scalars, int_scalars); } -TEST_F(TestIfElseSpecialExecute, FieldWithArrays) { +TEST_F(TestExecuteIfElseSpecial, FieldWithArrays) { DoTestBasic(boolean_fields, int_fields, int_fields, boolean_arrays, int_arrays); } -TEST_F(TestIfElseSpecialExecute, ComplexExprsWithArrays) { +TEST_F(TestExecuteIfElseSpecial, ComplexExprsWithArrays) { DoTestBasic(boolean_complex_exprs, int_complex_exprs, int_complex_exprs, boolean_arrays, int_arrays); } +TEST_F(TestExecuteIfElseSpecial, NestedCondAllLiterals) { + DoTestNestedCond(boolean_literals, boolean_literals, boolean_literals, boolean_arrays, + int_arrays); +} + +TEST_F(TestExecuteIfElseSpecial, NestedCondAllScalars) { + DoTestNestedCond(boolean_fields, boolean_fields, boolean_fields, boolean_scalars, + int_scalars); +} + +TEST_F(TestExecuteIfElseSpecial, NestedCondFieldWithArrays) { + DoTestNestedCond(boolean_fields, boolean_fields, boolean_fields, boolean_arrays, + int_arrays); +} + +// TEST_F(TestExecuteIfElseSpecial, NarrowDown) { +// DoTestNestedCond( +// {boolean1}, {boolean2}, {boolean1}, +// ExecBatch({boolean1_arr, boolean1_chunked, int1_arr, int2_chunked}, length)); +// } + // TEST(IfElseSpecial, ExecuteNestedCond) { // auto boolean_exprs = {boolean1, boolean2, and_(boolean1, boolean2), // or_(boolean1, boolean2)}; From 73d9be856d0b8bcd8cfbd0ef2a8c80ba2cff55b0 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Wed, 15 Oct 2025 12:49:50 -0700 Subject: [PATCH 56/71] Almost done --- .../compute/special/if_else_special_test.cc | 105 +++++++++++++++--- 1 file changed, 91 insertions(+), 14 deletions(-) diff --git a/cpp/src/arrow/compute/special/if_else_special_test.cc b/cpp/src/arrow/compute/special/if_else_special_test.cc index c74bcc5aaa26..eaf176675f1e 100644 --- a/cpp/src/arrow/compute/special/if_else_special_test.cc +++ b/cpp/src/arrow/compute/special/if_else_special_test.cc @@ -558,25 +558,79 @@ class TestExecuteIfElseSpecial : public ::testing::Test { } } - void DoTestNestedCond(const std::vector& cond_exprs, + // void DoTestNestedCond(const std::vector& cond_exprs, + // const std::vector& if_true_exprs, + // const std::vector& if_false_exprs, + // const ExecBatch& batch) { + // for (const auto& cond_expr : cond_exprs) { + // for (const auto& if_true_expr : if_true_exprs) { + // for (const auto& if_false_expr : if_false_exprs) { + // ARROW_SCOPED_TRACE( + // "if_else_special: " + + // if_else_special(if_else_special(cond_expr, if_true_expr, if_false_expr), + // int1, int2) + // .ToString()); + // ARROW_SCOPED_TRACE("batch: " + batch.ToString()); + // CheckIfElseSpecial( + // [=](MakeIfElseFunc make_if_else) { + // return make_if_else(make_if_else(cond_expr, if_true_expr, + // if_false_expr), + // int1, int2); + // }, + // *schm, batch); + // } + // } + // } + // } + + void DoTestNestedBody(const std::vector& cond_exprs, const std::vector& if_true_exprs, const std::vector& if_false_exprs, - const ExecBatch& batch) { + const std::vector& boolean_datums, + const std::vector& int_datums) { for (const auto& cond_expr : cond_exprs) { for (const auto& if_true_expr : if_true_exprs) { for (const auto& if_false_expr : if_false_exprs) { - ARROW_SCOPED_TRACE( - "if_else_special: " + - if_else_special(if_else_special(cond_expr, if_true_expr, if_false_expr), - int1, int2) - .ToString()); - ARROW_SCOPED_TRACE("batch: " + batch.ToString()); - CheckIfElseSpecial( - [=](MakeIfElseFunc make_if_else) { - return make_if_else(make_if_else(cond_expr, if_true_expr, if_false_expr), - int1, int2); - }, - *schm, batch); + for (const auto& b1_datum : boolean_datums) { + for (const auto& b2_datum : boolean_datums) { + for (const auto& i1_datum : int_datums) { + for (const auto& i2_datum : int_datums) { + ExecBatch batch({b1_datum, b2_datum, i1_datum, i2_datum}, length); + ARROW_SCOPED_TRACE("batch: " + batch.ToString()); + { + ARROW_SCOPED_TRACE( + "if_else_special: " + + if_else_special( + boolean1, + if_else_special(cond_expr, if_true_expr, if_false_expr), int2) + .ToString()); + CheckIfElseSpecial( + [=](MakeIfElseFunc make_if_else) { + return make_if_else( + boolean1, + make_if_else(cond_expr, if_true_expr, if_false_expr), int2); + }, + *schm, batch); + } + { + ARROW_SCOPED_TRACE( + "if_else_special: " + + if_else_special( + boolean1, int1, + if_else_special(cond_expr, if_true_expr, if_false_expr)) + .ToString()); + CheckIfElseSpecial( + [=](MakeIfElseFunc make_if_else) { + return make_if_else( + boolean1, int1, + make_if_else(cond_expr, if_true_expr, if_false_expr)); + }, + *schm, batch); + } + } + } + } + } } } } @@ -615,6 +669,29 @@ TEST_F(TestExecuteIfElseSpecial, NestedCondFieldWithArrays) { int_arrays); } +TEST_F(TestExecuteIfElseSpecial, NestedCondComplexExprsWithArrays) { + DoTestNestedCond(boolean_complex_exprs, boolean_complex_exprs, boolean_complex_exprs, + boolean_arrays, int_arrays); +} + +TEST_F(TestExecuteIfElseSpecial, NestedBodyAllLiterals) { + DoTestNestedBody(boolean_literals, int_literals, int_literals, boolean_arrays, + int_arrays); +} + +TEST_F(TestExecuteIfElseSpecial, NestedBodyAllScalars) { + DoTestNestedBody(boolean_fields, int_fields, int_fields, boolean_scalars, int_scalars); +} + +TEST_F(TestExecuteIfElseSpecial, NestedBodyFieldWithArrays) { + DoTestNestedBody(boolean_fields, int_fields, int_fields, boolean_arrays, int_arrays); +} + +TEST_F(TestExecuteIfElseSpecial, NestedBodyComplexExprsWithArrays) { + DoTestNestedBody(boolean_complex_exprs, int_complex_exprs, int_complex_exprs, + boolean_arrays, int_arrays); +} + // TEST_F(TestExecuteIfElseSpecial, NarrowDown) { // DoTestNestedCond( // {boolean1}, {boolean2}, {boolean1}, From 35d61185f7b9e672d9dc7db7b3b0c01754e1df38 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Tue, 4 Nov 2025 17:10:26 -0800 Subject: [PATCH 57/71] Complete the test --- .../compute/special/if_else_special_test.cc | 894 ++++-------------- 1 file changed, 185 insertions(+), 709 deletions(-) diff --git a/cpp/src/arrow/compute/special/if_else_special_test.cc b/cpp/src/arrow/compute/special/if_else_special_test.cc index eaf176675f1e..efd973c2d9d9 100644 --- a/cpp/src/arrow/compute/special/if_else_special_test.cc +++ b/cpp/src/arrow/compute/special/if_else_special_test.cc @@ -216,166 +216,6 @@ TEST(IfElseSpecial, BindSpecialForm) { } } -// class IfElseSpecialFormTest : public ::testing::Test { -// protected: -// static void SetUpTestSuite() { ASSERT_OK(RegisterAuxilaryFunctions()); } - -// protected: -// static Status UnreachableExec(KernelContext*, const ExecSpan&, ExecResult*) { -// return Status::Invalid("Unreachable"); -// } - -// static ArrayKernelSelectiveExec AssertSelection() { -// return [](KernelContext*, const ExecSpan&, const SelectionVectorSpan&, ExecResult*) -// { -// return Status::Invalid("Unreachable"); -// }; -// } - -// static Status IdentityExec(KernelContext*, const ExecSpan& span, ExecResult* out) { -// DCHECK_EQ(span.num_values(), 1); -// const auto& arg = span[0]; -// DCHECK(arg.is_array()); -// *out->array_data_mutable() = *arg.array.ToArrayData(); -// return Status::OK(); -// } - -// static Status AssertSelectionVectorNotExistExec(KernelContext* kernel_ctx, -// const ExecSpan& span, ExecResult* -// out) { -// return IdentityExec(kernel_ctx, span, out); -// } - -// static Status AssertSelectionVectorExistExec(KernelContext* kernel_ctx, -// const ExecSpan& span, -// const SelectionVectorSpan& selection, -// ExecResult* out) { -// for (int32_t i = 0; i < selection.length(); ++i) { -// EXPECT_GE(selection[i], 0); -// EXPECT_LT(selection[i], span.length); -// } -// return IdentityExec(kernel_ctx, span, out); -// } - -// static Status RegisterAuxilaryFunctions() { -// auto registry = GetFunctionRegistry(); - -// { -// auto register_unreachable_func = [&](const std::string& name) -> Status { -// auto func = -// std::make_shared(name, Arity::Unary(), -// FunctionDoc::Empty()); - -// ScalarKernel kernel({InputType::Any()}, internal::FirstType, UnreachableExec); -// kernel.can_write_into_slices = false; -// kernel.null_handling = NullHandling::COMPUTED_NO_PREALLOCATE; -// kernel.mem_allocation = MemAllocation::NO_PREALLOCATE; -// RETURN_NOT_OK(func->AddKernel(kernel)); -// RETURN_NOT_OK(registry->AddFunction(std::move(func))); -// return Status::OK(); -// }; - -// RETURN_NOT_OK(register_unreachable_func("unreachable")); -// } - -// { -// auto register_sv_awareness_func = [&](const std::string& name, -// bool sv_awareness) -> Status { -// auto func = -// std::make_shared(name, Arity::Unary(), -// FunctionDoc::Empty()); - -// ScalarKernel kernel({InputType::Any()}, internal::FirstType, IdentityExec); -// kernel.can_write_into_slices = false; -// kernel.null_handling = NullHandling::COMPUTED_NO_PREALLOCATE; -// kernel.mem_allocation = MemAllocation::NO_PREALLOCATE; -// RETURN_NOT_OK(func->AddKernel(kernel)); -// RETURN_NOT_OK(registry->AddFunction(std::move(func))); -// return Status::OK(); -// }; - -// RETURN_NOT_OK(register_sv_awareness_func("sv_suppress", false)); -// } - -// { -// auto register_assert_sv_func = [&](const std::string& name, -// bool sv_existence) -> Status { -// auto func = -// std::make_shared(name, Arity::Unary(), -// FunctionDoc::Empty()); - -// ArrayKernelExec exec = AssertSelectionVectorNotExistExec; -// ArrayKernelSelectiveExec selective_exec = nullptr; -// if (sv_existence) { -// selective_exec = AssertSelectionVectorExistExec; -// } -// ScalarKernel kernel({InputType::Any()}, internal::FirstType, std::move(exec), -// std::move(selective_exec), -// /*init=*/nullptr); -// kernel.can_write_into_slices = false; -// kernel.null_handling = NullHandling::COMPUTED_NO_PREALLOCATE; -// kernel.mem_allocation = MemAllocation::NO_PREALLOCATE; -// RETURN_NOT_OK(func->AddKernel(kernel)); -// RETURN_NOT_OK(registry->AddFunction(std::move(func))); -// return Status::OK(); -// }; - -// RETURN_NOT_OK(register_assert_sv_func("assert_sv_exist", /*sv_existence=*/true)); -// RETURN_NOT_OK(register_assert_sv_func("assert_sv_empty", -// /*sv_existence=*/false)); -// } - -// return Status::OK(); -// } - -// static std::vector SuppressSelectionVectorAwareForIfElse( -// const Expression& cond, const Expression& if_true, const Expression& if_false) { -// auto suppress_if_else_recursive = -// [&](const Expression& expr) -> std::vector { -// if (const auto& sp = expr.special(); sp && sp->special_form->name() == "if_else") -// { -// const auto& cond = sp->arguments[0]; -// const auto& if_true = sp->arguments[1]; -// const auto& if_false = sp->arguments[2]; -// return SuppressSelectionVectorAwareForIfElse(cond, if_true, if_false); -// } else { -// return {expr}; -// } -// }; -// auto suppressed_conds = suppress_if_else_recursive(cond); -// auto suppressed_if_trues = suppress_if_else_recursive(if_true); -// auto suppressed_if_falses = suppress_if_else_recursive(if_false); -// std::vector result; -// for (const auto& suppressed_cond : suppressed_conds) { -// for (const auto& suppressed_if_true : suppressed_if_trues) { -// for (const auto& suppressed_if_false : suppressed_if_falses) { -// result.emplace_back( -// if_else_special(suppressed_cond, suppressed_if_true, -// suppressed_if_false)); -// result.emplace_back(if_else_special(sv_suppress(suppressed_cond), -// suppressed_if_true, -// suppressed_if_false)); -// result.emplace_back(if_else_special(suppressed_cond, -// sv_suppress(suppressed_if_true), -// sv_suppress(suppressed_if_false))); -// } -// } -// } -// return result; -// } - -// static void CheckIfElseIgnoreShape(const Expression& cond, const Expression& if_true, -// const Expression& if_false, -// const std::shared_ptr& schema, -// const ExecBatch& batch, -// ExecContext* exec_context = -// default_exec_context()) { -// auto if_else = if_else_regular(cond, if_true, if_false); -// auto exprs = SuppressSelectionVectorAwareForIfElse(cond, if_true, if_false); -// AssertExprEqualExprsIgnoreShape(if_else, exprs, schema, batch, exec_context); -// } -// }; - namespace { Expression if_else(Expression cond, Expression if_true, Expression if_false) { @@ -388,13 +228,6 @@ Result ExecuteExpr(Expression expr, const Schema& schema, const ExecBatch return ExecuteScalarExpression(bound, batch, exec_context); } -// Result ExecuteIfElse(Expression cond, Expression if_true, Expression if_false, -// const Schema& schema, const ExecBatch& batch, -// ExecContext* exec_context = default_exec_context()) { -// return ExecuteExpr(if_else(std::move(cond), std::move(if_true), std::move(if_false)), -// schema, batch, exec_context); -// } - void AssertDatumsEqualIgnoreShape(const Datum& expected, const Datum& result) { DCHECK(expected.is_scalar() || expected.is_array() || expected.is_chunked_array()); DCHECK(result.is_scalar() || result.is_array() || result.is_chunked_array()); @@ -441,7 +274,9 @@ void CheckIfElseSpecial(Expression cond, Expression if_true, Expression if_false const Schema& schema, const ExecBatch& batch, ExecContext* exec_context = default_exec_context()) { CheckIfElseSpecial( - [=](MakeIfElseFunc make_if_else) { return if_else(cond, if_true, if_false); }, + [=](MakeIfElseFunc make_if_else) { + return make_if_else(std::move(cond), std::move(if_true), std::move(if_false)); + }, schema, batch, exec_context); } @@ -482,6 +317,8 @@ class TestExecuteIfElseSpecial : public ::testing::Test { boolean(), {"[true, false]", "[true]", "[true]", "[null, false, true]"}); std::vector boolean_arrays = {Datum(boolean1_arr), Datum(boolean1_chunked), Datum(boolean2_arr), Datum(boolean2_chunked)}; + std::vector boolean1_arrays = {Datum(boolean1_arr), Datum(boolean1_chunked)}; + std::vector boolean2_arrays = {Datum(boolean2_arr), Datum(boolean2_chunked)}; std::vector int_scalars = {Datum(MakeNullScalar(int32())), Datum(MakeScalar(42))}; @@ -493,434 +330,231 @@ class TestExecuteIfElseSpecial : public ::testing::Test { ChunkedArrayFromJSON(int32(), {"[0]", "[10, 20]", "[30, 40, 50]", "[]", "[60]"}); std::vector int_arrays = {Datum(int1_arr), Datum(int1_chunked), Datum(int2_arr), Datum(int2_chunked)}; + std::vector int1_arrays = {Datum(int1_arr), Datum(int1_chunked)}; + std::vector int2_arrays = {Datum(int2_arr), Datum(int2_chunked)}; protected: void DoTestBasic(const std::vector& cond_exprs, const std::vector& if_true_exprs, - const std::vector& if_false_exprs, - const std::vector& boolean_datums, - const std::vector& int_datums) { - for (const auto& cond_expr : cond_exprs) { - for (const auto& if_true_expr : if_true_exprs) { - for (const auto& if_false_expr : if_false_exprs) { - ARROW_SCOPED_TRACE( - "if_else_special: " + - if_else_special(cond_expr, if_true_expr, if_false_expr).ToString()); - for (const auto& b1_datum : boolean_datums) { - for (const auto& b2_datum : boolean_datums) { - for (const auto& i1_datum : int_datums) { - for (const auto& i2_datum : int_datums) { - ExecBatch batch({b1_datum, b2_datum, i1_datum, i2_datum}, length); - ARROW_SCOPED_TRACE("batch: " + batch.ToString()); - CheckIfElseSpecial(cond_expr, if_true_expr, if_false_expr, *schm, - batch); - } - } - } - } - } - } - } - } + const std::vector& if_false_exprs, const ExecBatch& batch); - void DoTestNestedCond(const std::vector& cond_exprs, - const std::vector& if_true_exprs, - const std::vector& if_false_exprs, - const std::vector& boolean_datums, - const std::vector& int_datums) { - for (const auto& cond_expr : cond_exprs) { - for (const auto& if_true_expr : if_true_exprs) { - for (const auto& if_false_expr : if_false_exprs) { - ARROW_SCOPED_TRACE( - "if_else_special: " + - if_else_special(if_else_special(cond_expr, if_true_expr, if_false_expr), - int1, int2) - .ToString()); - for (const auto& b1_datum : boolean_datums) { - for (const auto& b2_datum : boolean_datums) { - for (const auto& i1_datum : int_datums) { - for (const auto& i2_datum : int_datums) { - ExecBatch batch({b1_datum, b2_datum, i1_datum, i2_datum}, length); - ARROW_SCOPED_TRACE("batch: " + batch.ToString()); - CheckIfElseSpecial( - [=](MakeIfElseFunc make_if_else) { - return make_if_else( - make_if_else(cond_expr, if_true_expr, if_false_expr), int1, - int2); - }, - *schm, batch); - } - } - } + void DoTestNestedCond(const std::vector& nested_cond_exprs, + const std::vector& nested_if_true_exprs, + const std::vector& nested_if_false_exprs, + const std::vector& outer_if_true_exprs, + const std::vector& outer_if_false_exprs, + const ExecBatch& batch); + + void DoTestNestedBody(const std::vector& cond_exprs, + const std::vector& nested_cond_exprs, + const std::vector& nested_if_true_exprs, + const std::vector& nested_if_false_exprs, + const std::vector& other_branch_exprs, + const ExecBatch& batch); + + void WithDatumCombinations(const std::vector& boolean1_datums, + const std::vector& boolean2_datums, + const std::vector& int1_datums, + const std::vector& int2_datums, + std::function test_func) { + for (const auto& b1_datum : boolean1_datums) { + for (const auto& b2_datum : boolean2_datums) { + for (const auto& i1_datum : int1_datums) { + for (const auto& i2_datum : int2_datums) { + ExecBatch batch({b1_datum, b2_datum, i1_datum, i2_datum}, length); + ARROW_SCOPED_TRACE("batch: " + batch.ToString()); + test_func(batch); } } } } } +}; - // void DoTestNestedCond(const std::vector& cond_exprs, - // const std::vector& if_true_exprs, - // const std::vector& if_false_exprs, - // const ExecBatch& batch) { - // for (const auto& cond_expr : cond_exprs) { - // for (const auto& if_true_expr : if_true_exprs) { - // for (const auto& if_false_expr : if_false_exprs) { - // ARROW_SCOPED_TRACE( - // "if_else_special: " + - // if_else_special(if_else_special(cond_expr, if_true_expr, if_false_expr), - // int1, int2) - // .ToString()); - // ARROW_SCOPED_TRACE("batch: " + batch.ToString()); - // CheckIfElseSpecial( - // [=](MakeIfElseFunc make_if_else) { - // return make_if_else(make_if_else(cond_expr, if_true_expr, - // if_false_expr), - // int1, int2); - // }, - // *schm, batch); - // } - // } - // } - // } - - void DoTestNestedBody(const std::vector& cond_exprs, - const std::vector& if_true_exprs, - const std::vector& if_false_exprs, - const std::vector& boolean_datums, - const std::vector& int_datums) { - for (const auto& cond_expr : cond_exprs) { - for (const auto& if_true_expr : if_true_exprs) { - for (const auto& if_false_expr : if_false_exprs) { - for (const auto& b1_datum : boolean_datums) { - for (const auto& b2_datum : boolean_datums) { - for (const auto& i1_datum : int_datums) { - for (const auto& i2_datum : int_datums) { - ExecBatch batch({b1_datum, b2_datum, i1_datum, i2_datum}, length); - ARROW_SCOPED_TRACE("batch: " + batch.ToString()); - { - ARROW_SCOPED_TRACE( - "if_else_special: " + - if_else_special( - boolean1, - if_else_special(cond_expr, if_true_expr, if_false_expr), int2) - .ToString()); - CheckIfElseSpecial( - [=](MakeIfElseFunc make_if_else) { - return make_if_else( - boolean1, - make_if_else(cond_expr, if_true_expr, if_false_expr), int2); - }, - *schm, batch); - } - { - ARROW_SCOPED_TRACE( - "if_else_special: " + - if_else_special( - boolean1, int1, - if_else_special(cond_expr, if_true_expr, if_false_expr)) - .ToString()); - CheckIfElseSpecial( - [=](MakeIfElseFunc make_if_else) { - return make_if_else( - boolean1, int1, - make_if_else(cond_expr, if_true_expr, if_false_expr)); - }, - *schm, batch); - } - } - } - } - } - } +void TestExecuteIfElseSpecial::DoTestBasic(const std::vector& cond_exprs, + const std::vector& if_true_exprs, + const std::vector& if_false_exprs, + const ExecBatch& batch) { + for (const auto& cond_expr : cond_exprs) { + for (const auto& if_true_expr : if_true_exprs) { + for (const auto& if_false_expr : if_false_exprs) { + ARROW_SCOPED_TRACE( + "if_else_special: " + + if_else_special(cond_expr, if_true_expr, if_false_expr).ToString()); + CheckIfElseSpecial(cond_expr, if_true_expr, if_false_expr, *schm, batch); } } } -}; +} -TEST_F(TestExecuteIfElseSpecial, AllLiterals) { - DoTestBasic(boolean_literals, int_literals, int_literals, boolean_arrays, int_arrays); +TEST_F(TestExecuteIfElseSpecial, BasicAllLiterals) { + DoTestBasic(boolean_literals, int_literals, int_literals, ExecBatch({}, length)); } -TEST_F(TestExecuteIfElseSpecial, AllScalars) { - DoTestBasic(boolean_fields, int_fields, int_fields, boolean_scalars, int_scalars); +TEST_F(TestExecuteIfElseSpecial, BasicAllScalars) { + WithDatumCombinations(boolean_scalars, boolean_scalars, int_scalars, int_scalars, + [&](const ExecBatch& batch) { + DoTestBasic(boolean_fields, {int1}, {int2}, batch); + }); } -TEST_F(TestExecuteIfElseSpecial, FieldWithArrays) { - DoTestBasic(boolean_fields, int_fields, int_fields, boolean_arrays, int_arrays); +TEST_F(TestExecuteIfElseSpecial, BasicArrays) { + WithDatumCombinations(boolean_arrays, boolean_arrays, int_arrays, int_arrays, + [&](const ExecBatch& batch) { + DoTestBasic(boolean_fields, {int1}, {int2}, batch); + }); } -TEST_F(TestExecuteIfElseSpecial, ComplexExprsWithArrays) { - DoTestBasic(boolean_complex_exprs, int_complex_exprs, int_complex_exprs, boolean_arrays, - int_arrays); +TEST_F(TestExecuteIfElseSpecial, BasicComplexExprs) { + WithDatumCombinations(boolean1_arrays, boolean2_arrays, int1_arrays, int2_arrays, + [&](const ExecBatch& batch) { + DoTestBasic(boolean_complex_exprs, int_complex_exprs, + int_complex_exprs, batch); + }); +} + +void TestExecuteIfElseSpecial::DoTestNestedCond( + const std::vector& nested_cond_exprs, + const std::vector& nested_if_true_exprs, + const std::vector& nested_if_false_exprs, + const std::vector& outer_if_true_exprs, + const std::vector& outer_if_false_exprs, const ExecBatch& batch) { + for (const auto& nested_cond_expr : nested_cond_exprs) { + for (const auto& nested_if_true_expr : nested_if_true_exprs) { + for (const auto& nested_if_false_expr : nested_if_false_exprs) { + for (const auto& outer_if_true_expr : outer_if_true_exprs) { + for (const auto& outer_if_false_expr : outer_if_false_exprs) { + ARROW_SCOPED_TRACE( + "if_else_special: " + + if_else_special(if_else_special(nested_cond_expr, nested_if_true_expr, + nested_if_false_expr), + outer_if_true_expr, outer_if_false_expr) + .ToString()); + CheckIfElseSpecial( + [=](MakeIfElseFunc make_if_else) { + return make_if_else(make_if_else(nested_cond_expr, nested_if_true_expr, + nested_if_false_expr), + outer_if_true_expr, outer_if_false_expr); + }, + *schm, batch); + } + } + } + } + } } TEST_F(TestExecuteIfElseSpecial, NestedCondAllLiterals) { - DoTestNestedCond(boolean_literals, boolean_literals, boolean_literals, boolean_arrays, - int_arrays); + DoTestNestedCond(boolean_literals, boolean_literals, boolean_literals, int_literals, + int_literals, ExecBatch({}, length)); } TEST_F(TestExecuteIfElseSpecial, NestedCondAllScalars) { - DoTestNestedCond(boolean_fields, boolean_fields, boolean_fields, boolean_scalars, - int_scalars); + WithDatumCombinations(boolean_scalars, boolean_scalars, int_scalars, int_scalars, + [&](const ExecBatch& batch) { + DoTestNestedCond(boolean_fields, boolean_fields, boolean_fields, + {int1}, {int2}, batch); + }); +} + +TEST_F(TestExecuteIfElseSpecial, NestedCondArrays) { + WithDatumCombinations(boolean1_arrays, boolean2_arrays, int1_arrays, int2_arrays, + [&](const ExecBatch& batch) { + DoTestNestedCond(boolean_fields, boolean_fields, boolean_fields, + {int1}, {int2}, batch); + }); } -TEST_F(TestExecuteIfElseSpecial, NestedCondFieldWithArrays) { - DoTestNestedCond(boolean_fields, boolean_fields, boolean_fields, boolean_arrays, - int_arrays); +TEST_F(TestExecuteIfElseSpecial, NestedCondComplexExprs) { + WithDatumCombinations(boolean1_arrays, boolean2_arrays, int1_arrays, int2_arrays, + [&](const ExecBatch& batch) { + DoTestNestedCond(boolean_complex_exprs, boolean_complex_exprs, + boolean_complex_exprs, {int1}, {int2}, batch); + }); } -TEST_F(TestExecuteIfElseSpecial, NestedCondComplexExprsWithArrays) { - DoTestNestedCond(boolean_complex_exprs, boolean_complex_exprs, boolean_complex_exprs, - boolean_arrays, int_arrays); +void TestExecuteIfElseSpecial::DoTestNestedBody( + const std::vector& cond_exprs, + const std::vector& nested_cond_exprs, + const std::vector& nested_if_true_exprs, + const std::vector& nested_if_false_exprs, + const std::vector& other_branch_exprs, const ExecBatch& batch) { + for (const auto& cond_expr : cond_exprs) { + for (const auto& nested_cond_expr : nested_cond_exprs) { + for (const auto& nested_if_true_expr : nested_if_true_exprs) { + for (const auto& nested_if_false_expr : nested_if_false_exprs) { + for (const auto& other_branch_expr : other_branch_exprs) { + { + ARROW_SCOPED_TRACE( + "if_else_special: " + + if_else_special(cond_expr, + if_else_special(nested_cond_expr, nested_if_true_expr, + nested_if_false_expr), + other_branch_expr) + .ToString()); + CheckIfElseSpecial( + [=](MakeIfElseFunc make_if_else) { + return make_if_else( + cond_expr, + make_if_else(nested_cond_expr, nested_if_true_expr, + nested_if_false_expr), + other_branch_expr); + }, + *schm, batch); + } + { + ARROW_SCOPED_TRACE( + "if_else_special: " + + if_else_special(cond_expr, other_branch_expr, + if_else_special(nested_cond_expr, nested_if_true_expr, + nested_if_false_expr)) + .ToString()); + CheckIfElseSpecial( + [=](MakeIfElseFunc make_if_else) { + return make_if_else( + cond_expr, other_branch_expr, + make_if_else(nested_cond_expr, nested_if_true_expr, + nested_if_false_expr)); + }, + *schm, batch); + } + } + } + } + } + } } TEST_F(TestExecuteIfElseSpecial, NestedBodyAllLiterals) { - DoTestNestedBody(boolean_literals, int_literals, int_literals, boolean_arrays, - int_arrays); + DoTestNestedBody(boolean_literals, boolean_literals, int_literals, int_literals, + int_literals, ExecBatch({}, length)); } TEST_F(TestExecuteIfElseSpecial, NestedBodyAllScalars) { - DoTestNestedBody(boolean_fields, int_fields, int_fields, boolean_scalars, int_scalars); + WithDatumCombinations(boolean_scalars, boolean_scalars, int_scalars, int_scalars, + [&](const ExecBatch& batch) { + DoTestNestedBody({boolean1}, boolean_fields, int_fields, + int_fields, {int1}, batch); + }); } TEST_F(TestExecuteIfElseSpecial, NestedBodyFieldWithArrays) { - DoTestNestedBody(boolean_fields, int_fields, int_fields, boolean_arrays, int_arrays); + WithDatumCombinations(boolean1_arrays, boolean2_arrays, int1_arrays, int2_arrays, + [&](const ExecBatch& batch) { + DoTestNestedBody({boolean1}, boolean_fields, int_fields, + int_fields, {int1}, batch); + }); } TEST_F(TestExecuteIfElseSpecial, NestedBodyComplexExprsWithArrays) { - DoTestNestedBody(boolean_complex_exprs, int_complex_exprs, int_complex_exprs, - boolean_arrays, int_arrays); + WithDatumCombinations(boolean1_arrays, boolean2_arrays, int1_arrays, int2_arrays, + [&](const ExecBatch& batch) { + DoTestNestedBody({boolean1}, boolean_complex_exprs, + int_complex_exprs, int_complex_exprs, {int1}, + batch); + }); } -// TEST_F(TestExecuteIfElseSpecial, NarrowDown) { -// DoTestNestedCond( -// {boolean1}, {boolean2}, {boolean1}, -// ExecBatch({boolean1_arr, boolean1_chunked, int1_arr, int2_chunked}, length)); -// } - -// TEST(IfElseSpecial, ExecuteNestedCond) { -// auto boolean_exprs = {boolean1, boolean2, and_(boolean1, boolean2), -// or_(boolean1, boolean2)}; -// for (const auto& cond_inner_expr : boolean_exprs) { -// for (const auto& if_true_inner_expr : boolean_exprs) { -// for (const auto& if_false_inner_expr : boolean_exprs) { -// ARROW_SCOPED_TRACE( -// "expression: if (" + -// if_else_special(cond_inner_expr, if_true_inner_expr, if_false_inner_expr) -// .ToString() + -// ") then int1 else int2"); -// for (const auto& boolean1_datum : -// {Datum(boolean1_arr), Datum(boolean1_chunked)}) { -// ARROW_SCOPED_TRACE("boolean1: " + boolean1_datum.ToString()); -// for (const auto& boolean2_datum : -// {Datum(boolean2_arr), Datum(boolean2_chunked)}) { -// ARROW_SCOPED_TRACE("boolean2: " + boolean2_datum.ToString()); -// for (const auto& int1_datum : {Datum(int1_arr), Datum(int1_chunked)}) { -// ARROW_SCOPED_TRACE("int1: " + int1_datum.ToString()); -// for (const auto& int2_datum : {Datum(int2_arr), Datum(int2_chunked)}) { -// ARROW_SCOPED_TRACE("int2: " + int2_datum.ToString()); -// CheckIfElseSpecial( -// [=](MakeIfElseFunc make_if_else) { -// return make_if_else( -// make_if_else(cond_inner_expr, if_true_inner_expr, -// if_false_inner_expr), -// int1, int2); -// }, -// *schm, -// ExecBatch({Datum(boolean1_arr), Datum(boolean2_arr), -// Datum(int1_arr), -// Datum(int2_arr)}, -// length)); -// } -// } -// } -// } -// } -// } -// } -// } - -// TEST(IfElseSpecial, ExecuteNestedBody) { -// const int64_t length = 7; - -// auto schm = schema({field("boolean1", boolean()), field("boolean2", boolean()), -// field("int1", int32()), field("int2", int32())}); -// auto boolean1 = field_ref("boolean1"); -// auto boolean2 = field_ref("boolean2"); -// auto int1 = field_ref("int1"); -// auto int2 = field_ref("int2"); - -// auto boolean1_arr = -// ArrayFromJSON(boolean(), "[null, true, false, true, false, null, true]"); -// auto boolean1_chunked = ChunkedArrayFromJSON( -// boolean(), {"[null, true, false]", "[]", "[true, false, null, true]"}); - -// auto boolean2_arr = -// ArrayFromJSON(boolean(), "[true, false, true, true, null, false, true]"); -// auto boolean2_chunked = ChunkedArrayFromJSON( -// boolean(), {"[true, false]", "[true]", "[true]", "[null, false, true]"}); - -// auto int1_arr = ArrayFromJSON(int32(), "[0, 1, 2, 3, 4, 5, 6]"); -// auto int1_chunked = ChunkedArrayFromJSON(int32(), {"[0, 1]", "[2, 3, 4]", "[5, -// 6]"}); - -// auto int2_arr = ArrayFromJSON(int32(), "[0, 10, 20, 30, 40, 50, 60]"); -// auto int2_chunked = -// ChunkedArrayFromJSON(int32(), {"[0]", "[10, 20]", "[30, 40, 50]", "[]", -// "[60]"}); - -// auto boolean_exprs = {boolean1, boolean2, and_(boolean1, boolean2), -// or_(boolean1, boolean2)}; -// auto int_exprs = {int1, int2, add(int1, int2)}; -// for (const auto& cond_outer_expr : boolean_exprs) { -// for (const auto& cond_inner_expr : boolean_exprs) { -// for (const auto& if_true_inner_expr : int_exprs) { -// for (const auto& if_false_inner_expr : int_exprs) { -// for (const auto& if_false_outer_expr : int_exprs) { -// ARROW_SCOPED_TRACE( -// "expression: if (" + cond_outer_expr.ToString() + ") then (" + -// if_else_special(cond_inner_expr, if_true_inner_expr, -// if_false_inner_expr) -// .ToString() + -// ") else (" + if_false_outer_expr.ToString() + ")"); -// for (const auto& boolean1_datum : -// {Datum(boolean1_arr), Datum(boolean1_chunked)}) { -// ARROW_SCOPED_TRACE("boolean1: " + boolean1_datum.ToString()); -// for (const auto& boolean2_datum : -// {Datum(boolean2_arr), Datum(boolean2_chunked)}) { -// ARROW_SCOPED_TRACE("boolean2: " + boolean2_datum.ToString()); -// for (const auto& int1_datum : {Datum(int1_arr), Datum(int1_chunked)}) -// { -// ARROW_SCOPED_TRACE("int1: " + int1_datum.ToString()); -// for (const auto& int2_datum : {Datum(int2_arr), -// Datum(int2_chunked)}) -// { -// ARROW_SCOPED_TRACE("int2: " + int2_datum.ToString()); -// CheckIfElseSpecial( -// [=](MakeIfElseFunc make_if_else) { -// return make_if_else( -// cond_outer_expr, -// make_if_else(cond_inner_expr, if_true_inner_expr, -// if_false_inner_expr), -// if_false_outer_expr); -// }, -// *schm, -// ExecBatch({Datum(boolean1_arr), Datum(boolean2_arr), -// Datum(int1_arr), Datum(int2_arr)}, -// length)); -// } -// } -// } -// } -// } - -// for (const auto& if_true_outer_expr : int_exprs) { -// ARROW_SCOPED_TRACE( -// "expression: if (" + cond_outer_expr.ToString() + ") then (" + -// if_true_outer_expr.ToString() + ") else (" + -// if_else_special(cond_inner_expr, if_true_inner_expr, -// if_false_inner_expr) -// .ToString() + -// ")"); -// for (const auto& boolean1_datum : -// {Datum(boolean1_arr), Datum(boolean1_chunked)}) { -// ARROW_SCOPED_TRACE("boolean1: " + boolean1_datum.ToString()); -// for (const auto& boolean2_datum : -// {Datum(boolean2_arr), Datum(boolean2_chunked)}) { -// ARROW_SCOPED_TRACE("boolean2: " + boolean2_datum.ToString()); -// for (const auto& int1_datum : {Datum(int1_arr), Datum(int1_chunked)}) -// { -// ARROW_SCOPED_TRACE("int1: " + int1_datum.ToString()); -// for (const auto& int2_datum : {Datum(int2_arr), -// Datum(int2_chunked)}) -// { -// ARROW_SCOPED_TRACE("int2: " + int2_datum.ToString()); -// CheckIfElseSpecial( -// [=](MakeIfElseFunc make_if_else) { -// return make_if_else( -// cond_outer_expr, if_true_outer_expr, -// make_if_else(cond_inner_expr, if_true_inner_expr, -// if_false_inner_expr)); -// }, -// *schm, -// ExecBatch({Datum(boolean1_arr), Datum(boolean2_arr), -// Datum(int1_arr), Datum(int2_arr)}, -// length)); -// } -// } -// } -// } -// } -// } -// } -// } -// } -// } - -// TEST(IfElseSpecial, ExecuteNested) { -// // if (if (boolean1) then (boolean2) else (int1 == int2))) then if () -// // if (if (boolean1) then (boolean2) else (int1 == int2)) then -// // if (boolean1) then (int1) else (int1 + int2) -// // else -// // if (boolean2) then (int2) else (int1 + int2) -// const int64_t length = 7; - -// auto schm = schema({field("boolean1", boolean()), field("boolean2", boolean()), -// field("int1", int32()), field("int2", int32())}); -// auto boolean1 = field_ref("boolean1"); -// auto boolean2 = field_ref("boolean2"); -// auto int1 = field_ref("int1"); -// auto int2 = field_ref("int2"); - -// auto boolean1_arr = -// ArrayFromJSON(boolean(), "[null, true, false, true, false, null, true]"); -// auto boolean1_chunked = ChunkedArrayFromJSON( -// boolean(), {"[null, true, false]", "[]", "[true, false, null, true]"}); - -// auto boolean2_arr = -// ArrayFromJSON(boolean(), "[true, false, true, true, null, false, true]"); -// auto boolean2_chunked = ChunkedArrayFromJSON( -// boolean(), {"[true, false]", "[true]", "[true]", "[null, false, true]"}); - -// auto int1_arr = ArrayFromJSON(int32(), "[0, 1, 2, 3, 4, 5, 6]"); -// auto int1_chunked = ChunkedArrayFromJSON(int32(), {"[0, 1]", "[2, 3, 4]", "[5, -// 6]"}); - -// auto int2_arr = ArrayFromJSON(int32(), "[0, 10, 20, 30, 40, 50, 60]"); -// auto int2_chunked = -// ChunkedArrayFromJSON(int32(), {"[0]", "[10, 20]", "[30, 40, 50]", "[]", -// "[60]"}); - -// for (const auto& cond_inner_expr : -// {literal(MakeNullScalar(boolean())), literal(true), literal(false), boolean1, -// boolean2, equal(int1, int2)}) { -// for (const auto& if_true_inner_expr : -// {literal(MakeNullScalar(int32())), literal(42), -// literal(std::numeric_limits::max()), int1, int2, add(int1, int2)}) -// { -// for (const auto& if_false_inner_expr : -// {literal(MakeNullScalar(int32())), literal(24), -// literal(std::numeric_limits::min()), int2, add(b, c)}) { -// ARROW_SCOPED_TRACE( -// "expression: " + -// if_else_special(cond_expr, if_true_expr, if_false_expr).ToString()); -// for (const auto& cond_datum : {Datum(a_arr), Datum(a_chunked)}) { -// ARROW_SCOPED_TRACE("cond: " + cond_datum.ToString()); -// for (const auto& if_true_datum : {Datum(b_arr), Datum(b_chunked)}) { -// ARROW_SCOPED_TRACE("if_true: " + if_true_datum.ToString()); -// for (const auto& if_false_datum : {Datum(c_arr), Datum(c_chunked)}) { -// ARROW_SCOPED_TRACE("if_false: " + if_false_datum.ToString()); -// CheckIfElseSpecial( -// cond_expr, if_true_expr, if_false_expr, *schm, -// ExecBatch({cond_datum, if_true_datum, if_false_datum}, length)); -// } -// } -// } -// } -// } -// } -// } - namespace { Result ExecuteIfElseSpecial(Expression cond, Expression if_true, @@ -957,162 +591,4 @@ TEST(IfElseSpecial, IfNotZeroThenDivide) { AssertDatumsEqual(expected, result); } -// namespace { - -// auto kCanonicalSchema = arrow::schema({field("a", boolean()), field("b", int32())}); - -// auto kCanonicalA = field_ref("a"); -// auto kCanonicalB = field_ref("b"); - -// const auto kCanonicalBooleanCols = {kBooleanNull, literal(true), literal(false), -// kCanonicalA}; -// const auto kCanonicalIntCols = {kIntNull, literal(42), kCanonicalB}; - -// const std::vector kCanonicalBatches = { -// ExecBatch(*RecordBatchFromJSON(kCanonicalSchema, R"([ -// [null, 0], -// [null, null], -// [null, 1] -// ])")), -// ExecBatch(*RecordBatchFromJSON(kCanonicalSchema, R"([ -// [true, 0], -// [true, null], -// [true, 1] -// ])")), -// ExecBatch(*RecordBatchFromJSON(kCanonicalSchema, R"([ -// [false, 0], -// [false, null], -// [false, 1] -// ])")), -// ExecBatch(*RecordBatchFromJSON(kCanonicalSchema, R"([ -// [false, 0], -// [true, null], -// [null, 1] -// ])")), -// }; - -// } // namespace - -// TEST_F(IfElseSpecialFormTest, Simple) { -// const auto& schema = kCanonicalSchema; -// const auto& boolean_datums = kCanonicalBooleanCols; -// const auto& int_datums = kCanonicalIntCols; -// const auto& batches = kCanonicalBatches; -// for (const auto& cond : boolean_datums) { -// for (const auto& if_true : int_datums) { -// for (const auto& if_false : int_datums) { -// for (const auto& batch : batches) { -// CheckIfElseIgnoreShape(cond, if_true, if_false, schema, batch); -// } -// } -// } -// } -// } - -// TEST_F(IfElseSpecialFormTest, NestedSimple) { -// const auto& schema = kCanonicalSchema; -// const auto& a = kCanonicalA; -// const auto& b = kCanonicalB; -// ExecBatch batch(*RecordBatchFromJSON(kCanonicalSchema, R"([ -// [false, 0], -// [true, null], -// [null, 1] -// ])")); -// for (const auto& cond : { -// if_else_special(a, kBooleanNull, a), -// if_else_special(a, a, literal(true)), -// }) { -// for (const auto& if_true : { -// if_else_special(a, kIntNull, b), -// if_else_special(a, b, literal(42)), -// }) { -// for (const auto& if_false : { -// if_else_special(a, kIntNull, b), -// if_else_special(a, b, literal(42)), -// }) { -// CheckIfElseIgnoreShape(cond, if_true, if_false, schema, batch); -// } -// } -// } -// } - -// // TODO: Deprecate this test due to slowness. -// TEST_F(IfElseSpecialFormTest, NestedConditionComplex) { -// const auto& batches = kCanonicalBatches; -// const auto& schema = kCanonicalSchema; -// const auto& boolean_datums = kCanonicalBooleanCols; -// const auto& int_datums = kCanonicalIntCols; -// for (const auto& nested_cond : boolean_datums) { -// for (const auto& nested_if_true : boolean_datums) { -// for (const auto& nested_if_false : boolean_datums) { -// auto nested_if_else_sp = -// if_else_special(nested_cond, nested_if_true, nested_if_false); -// for (const auto& if_true : int_datums) { -// for (const auto& if_false : int_datums) { -// for (const auto& batch : batches) { -// CheckIfElseIgnoreShape(nested_if_else_sp, if_true, if_false, schema, -// batch); -// } -// } -// } -// } -// } -// } -// } - -// // TODO: Deprecate this test due to slowness. -// TEST_F(IfElseSpecialFormTest, NestedBodyComplex) { -// const auto& batches = kCanonicalBatches; -// const auto& schema = kCanonicalSchema; -// const auto& boolean_datums = kCanonicalBooleanCols; -// const auto& int_datums = kCanonicalIntCols; -// for (const auto& cond : boolean_datums) { -// for (const auto& nested_cond : boolean_datums) { -// for (const auto& nested_if_true : int_datums) { -// for (const auto& nested_if_false : int_datums) { -// auto nested_if_else_sp = -// if_else_special(nested_cond, nested_if_true, nested_if_false); -// for (const auto& batch : batches) { -// CheckIfElseIgnoreShape(cond, nested_if_else_sp, nested_if_else_sp, -// schema, -// batch); -// } -// } -// } -// } -// } -// } - -// // TODO: Deprecate this test due to slowness. -// TEST_F(IfElseSpecialFormTest, NestedComplex) { -// const auto& batches = kCanonicalBatches; -// const auto& schema = kCanonicalSchema; -// const auto& boolean_datums = kCanonicalBooleanCols; -// const auto& int_datums = kCanonicalIntCols; -// for (const auto& cond_nested_cond : boolean_datums) { -// for (const auto& cond_nested_if_true : boolean_datums) { -// for (const auto& cond_nested_if_false : boolean_datums) { -// auto cond = -// if_else_special(cond_nested_cond, cond_nested_if_true, -// cond_nested_if_false); -// for (const auto& nested_cond : boolean_datums) { -// for (const auto& nested_if_true : int_datums) { -// for (const auto& nested_if_false : int_datums) { -// auto nested_if_else_sp = -// if_else_special(nested_cond, nested_if_true, nested_if_false); -// for (const auto& batch : batches) { -// CheckIfElseIgnoreShape(cond, nested_if_else_sp, nested_if_else_sp, -// schema, -// batch); -// } -// } -// } -// } -// } -// } -// } -// } - -// // TODO: ChunkedArray. - } // namespace arrow::compute From a399f485ae59b53ff05692221398baf28adae9ba Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Fri, 7 Nov 2025 14:49:47 -0800 Subject: [PATCH 58/71] Rewrite benchmark --- .../special/if_else_special_benchmark.cc | 381 +++++++++--------- cpp/src/arrow/util/macros.h | 1 + 2 files changed, 188 insertions(+), 194 deletions(-) diff --git a/cpp/src/arrow/compute/special/if_else_special_benchmark.cc b/cpp/src/arrow/compute/special/if_else_special_benchmark.cc index 11c6380b2725..ac42ab28bdce 100644 --- a/cpp/src/arrow/compute/special/if_else_special_benchmark.cc +++ b/cpp/src/arrow/compute/special/if_else_special_benchmark.cc @@ -17,7 +17,8 @@ #include "benchmark/benchmark.h" -#include "arrow/compute/exec.h" +#include "arrow/compute/api_special.h" +#include "arrow/compute/exec_internal.h" #include "arrow/compute/expression.h" #include "arrow/compute/function.h" #include "arrow/compute/function_internal.h" @@ -27,272 +28,264 @@ #include "arrow/testing/random.h" #include "arrow/util/logging.h" -namespace arrow { - -namespace compute { +namespace arrow::compute { namespace { -struct PayloadOptions : public FunctionOptions { - explicit PayloadOptions(int64_t load = 0); - static constexpr char const kTypeName[] = "PayloadOptions"; - static PayloadOptions Defaults() { return PayloadOptions(); } - int64_t load = 0; +// A trivial kernel that just keeps the CPU busy for a specified number of iterations per +// input row. Has both regular and selective variants. Used to benchmark the overhead of +// the execution framework. + +struct SpinOptions : public FunctionOptions { + explicit SpinOptions(int64_t count = 0); + static constexpr char kTypeName[] = "SpinOptions"; + static SpinOptions Defaults() { return SpinOptions(); } + int64_t count = 0; }; -static auto kPayloadOptionsType = internal::GetFunctionOptionsType( - arrow::internal::DataMember("load", &PayloadOptions::load)); +static auto kSpinOptionsType = internal::GetFunctionOptionsType( + arrow::internal::DataMember("count", &SpinOptions::count)); -PayloadOptions::PayloadOptions(int64_t load) - : FunctionOptions(kPayloadOptionsType), load(load) {} +SpinOptions::SpinOptions(int64_t count) + : FunctionOptions(kSpinOptionsType), count(count) {} -const PayloadOptions* GetDefaultPayloadOptions() { - static const auto kDefaultPayloadOptions = PayloadOptions::Defaults(); - return &kDefaultPayloadOptions; +const SpinOptions* GetDefaultSpinOptions() { + static const auto kDefaultSpinOptions = SpinOptions::Defaults(); + return &kDefaultSpinOptions; } -using PayloadState = internal::OptionsWrapper; +using SpinState = internal::OptionsWrapper; -void DoLoad(int64_t load, int64_t length) { - for (int64_t i = 0; i < length; ++i) { - volatile int64_t j = load; - while (j-- > 0) { - // Do nothing, just burn CPU cycles - } +inline void Spin(volatile int64_t count) { + while (count-- > 0) { + // Do nothing, just burn CPU cycles. } } -Status PayloadExec(KernelContext* ctx, const ExecSpan& span, ExecResult* out) { +Status SpinExec(KernelContext* ctx, const ExecSpan& span, ExecResult* out) { ARROW_CHECK_EQ(span.num_values(), 1); const auto& arg = span[0]; ARROW_CHECK(arg.is_array()); - int64_t load = PayloadState::Get(ctx).load; - DoLoad(load, arg.length()); + int64_t count = SpinState::Get(ctx).count; + for (int64_t i = 0; i < arg.length(); ++i) { + Spin(count); + } *out->array_data_mutable() = *arg.array.ToArrayData(); return Status::OK(); } -Status PayloadSelectiveExec(KernelContext* ctx, const ExecSpan& span, - const SelectionVectorSpan selection_span, ExecResult* out) { +Status SpinSelectiveExec(KernelContext* ctx, const ExecSpan& span, + const SelectionVectorSpan& selection_span, ExecResult* out) { ARROW_CHECK_EQ(span.num_values(), 1); const auto& arg = span[0]; ARROW_CHECK(arg.is_array()); - int64_t load = PayloadState::Get(ctx).load; - DoLoad(load, span.length); + int64_t count = SpinState::Get(ctx).count; + detail::VisitSelectionVectorSpanInline(selection_span, [&](int64_t i) { Spin(count); }); *out->array_data_mutable() = *arg.array.ToArrayData(); return Status::OK(); } -Status RegisterAuxilaryFunctions() { +Status RegisterSpinFunction() { auto registry = GetFunctionRegistry(); - { - if (registry->CanAddFunctionOptionsType(kPayloadOptionsType).ok()) { - RETURN_NOT_OK(registry->AddFunctionOptionsType(kPayloadOptionsType)); - } + if (registry->CanAddFunctionOptionsType(kSpinOptionsType).ok()) { + RETURN_NOT_OK(registry->AddFunctionOptionsType(kSpinOptionsType)); } - { - auto register_payload_func = [&](const std::string& name, - bool sv_awareness) -> Status { - auto func = std::make_shared( - name, Arity::Unary(), FunctionDoc::Empty(), GetDefaultPayloadOptions()); - - ScalarKernel kernel({InputType::Any()}, internal::FirstType, PayloadExec, - PayloadState::Init, PayloadSelectiveExec); - kernel.selection_vector_aware = sv_awareness; - kernel.can_write_into_slices = false; - kernel.null_handling = NullHandling::COMPUTED_NO_PREALLOCATE; - kernel.mem_allocation = MemAllocation::NO_PREALLOCATE; - RETURN_NOT_OK(func->AddKernel(kernel)); - if (registry->CanAddFunction(func, /*allow_overwrite=*/false).ok()) { - RETURN_NOT_OK(registry->AddFunction(std::move(func))); - } - return Status::OK(); - }; - - RETURN_NOT_OK(register_payload_func("payload_sv_aware", true)); - RETURN_NOT_OK(register_payload_func("payload_sv_unaware", false)); - } - - return Status::OK(); -} -Expression if_else_regular(Expression cond, Expression if_true, Expression if_false) { - return call("if_else", {std::move(cond), std::move(if_true), std::move(if_false)}); -} + auto register_spin_function = [&](std::string name, ArrayKernelExec exec, + ArrayKernelSelectiveExec selective_exec) { + auto func = std::make_shared( + std::move(name), Arity::Unary(), FunctionDoc::Empty(), GetDefaultSpinOptions()); + ScalarKernel kernel({InputType::Any()}, internal::FirstType, exec, selective_exec, + SpinState::Init); + kernel.can_write_into_slices = false; + kernel.null_handling = NullHandling::COMPUTED_NO_PREALLOCATE; + kernel.mem_allocation = MemAllocation::NO_PREALLOCATE; + RETURN_NOT_OK(func->AddKernel(kernel)); + if (registry->CanAddFunction(func, /*allow_overwrite=*/false).ok()) { + RETURN_NOT_OK(registry->AddFunction(std::move(func))); + } + return Status::OK(); + }; -Expression sv_suppress(Expression arg) { - return call("payload_sv_unaware", {std::move(arg)}); -} + // Register two variants, one with selective exec and one without. + RETURN_NOT_OK(register_spin_function("spin_selective", SpinExec, SpinSelectiveExec)); + RETURN_NOT_OK(register_spin_function("spin", SpinExec, /*selective_exec=*/nullptr)); -Expression heavy(Expression arg) { - return call("payload_sv_aware", {std::move(arg)}, - std::make_shared(/*load=*/1024)); -} - -Expression sv_unaware_if_else_regular(Expression cond, Expression if_true, - Expression if_false) { - return if_else_regular(std::move(cond), sv_suppress(std::move(if_true)), - sv_suppress(std::move(if_false))); + return Status::OK(); } -Expression sv_unaware_if_else_special(Expression cond, Expression if_true, - Expression if_false) { - return if_else_special(std::move(cond), sv_suppress(std::move(if_true)), - sv_suppress(std::move(if_false))); +Expression if_else(Expression cond, Expression if_true, Expression if_false) { + return call("if_else", {std::move(cond), std::move(if_true), std::move(if_false)}); } -auto kBooleanNull = literal(MakeNullScalar(boolean())); - -void BenchmarkIfElse( - benchmark::State& state, - std::function if_else_func, - Expression cond, Expression if_true, Expression if_false, - const std::shared_ptr& schema, const ExecBatch& batch) { - ARROW_CHECK_OK(RegisterAuxilaryFunctions()); +auto kBooleanNullScalar = MakeNullScalar(boolean()); +auto kTrueScalar = MakeScalar(true); +auto kFalseScalar = MakeScalar(false); + +using MakeIfElseFunc = std::function; + +void BenchmarkIfElse(benchmark::State& state, MakeIfElseFunc make_if_else_func, + const std::string& spin_function, int64_t if_true_kernel_intensity, + int64_t if_false_kernel_intensity, Datum cond_datum, + int64_t length) { + ARROW_CHECK_EQ(cond_datum.type()->id(), Type::BOOL); + + static auto registered = RegisterSpinFunction(); + ARROW_CHECK_OK(registered); + + auto expr = make_if_else_func( + field_ref(0), + call(spin_function, {field_ref(1)}, SpinOptions(if_true_kernel_intensity)), + call(spin_function, {field_ref(2)}, SpinOptions(if_false_kernel_intensity))); + auto bound = expr.Bind(*schema({field("", cond_datum.type()), field("", int32()), + field("", int32())})) + .ValueOrDie(); + if (cond_datum.is_arraylike()) { + ARROW_CHECK_EQ(cond_datum.length(), length); + } + auto if_true_datum = ConstantArrayGenerator::Int32(length, 1); + auto if_false_datum = ConstantArrayGenerator::Int32(length, 0); + auto batch = ExecBatch{ + {std::move(cond_datum), std::move(if_true_datum), std::move(if_false_datum)}, + length}; - auto if_else = if_else_func(std::move(cond), std::move(if_true), std::move(if_false)); - auto bound = if_else.Bind(*schema).ValueOrDie(); for (auto _ : state) { ARROW_CHECK_OK(ExecuteScalarExpression(bound, batch).status()); } - state.SetItemsProcessed(batch.length * state.iterations()); + state.SetItemsProcessed(state.iterations() * length); } } // namespace -#ifdef BM -# error ("BM is defined") -#else -# define BENCHMARK_IF_ELSE_WITH_BASELINE_AND_SV_SUPPRESS(BM, name, ...) \ - BM(name##_regular, if_else_regular, ##__VA_ARGS__); \ - BM(name##_special, if_else_special, ##__VA_ARGS__); \ - BM(name##_regular_sv_unaware, sv_unaware_if_else_regular, ##__VA_ARGS__); \ - BM(name##_special_sv_unaware, sv_unaware_if_else_special, ##__VA_ARGS__); -#endif - -#define BENCHMARK_IF_ELSE(BM, name, if_else, arg_names, args, ...) \ - BENCHMARK_CAPTURE(BM, name, if_else, ##__VA_ARGS__) \ - ->ArgNames(arg_names) \ +// For each benchmark, expand to three variants: +// - Baseline: regular if_else with regular spin kernel. +// - Special: if_else_special with non-selective spin kernel - triggering dense +// execution. +// - SpecialSelective: if_else_special with selective spin kernel - triggering (more +// efficient) sparse execution. +#define BENCHMARK_IF_ELSE_WITH_BASELINE_AND_SPECIAL(BM, name, ...) \ + BM(ARROW_CONCAT(name, Baseline), if_else, "spin", ##__VA_ARGS__); \ + BM(ARROW_CONCAT(name, Special), if_else_special, "spin", ##__VA_ARGS__); \ + BM(ARROW_CONCAT(name, SpecialSelective), if_else_special, "spin_selective", \ + ##__VA_ARGS__); + +#define BENCHMARK_IF_ELSE(BM, name, if_else, spin_func, arg_names, args, ...) \ + BENCHMARK_CAPTURE(BM, name, if_else, spin_func, ##__VA_ARGS__) \ + ->ArgNames(arg_names) \ ->ArgsProduct(args) -template -static void BM_IfElseTrivialCond( - benchmark::State& state, - std::function if_else_func, - Expression cond, Args&&...) { +// Benchmark with scalar condition, see if short-circuiting takes place. +static void BM_IfElseScalarCond(benchmark::State& state, MakeIfElseFunc make_if_else_func, + std::string spin_func, Datum cond_datum) { const int64_t num_rows = state.range(0); - auto schema = arrow::schema({field("i1", int32()), field("i2", int32())}); - - auto i1 = ConstantArrayGenerator::Int32(num_rows, 1); - auto i2 = ConstantArrayGenerator::Int32(num_rows, 0); - ExecBatch batch{std::vector{std::move(i1), std::move(i2)}, num_rows}; - - BenchmarkIfElse(state, std::move(if_else_func), std::move(cond), field_ref("i1"), - field_ref("i2"), schema, batch); + BenchmarkIfElse(state, std::move(make_if_else_func), spin_func, + /*if_true_kernel_intensity=*/0, + /*if_false_kernel_intensity=*/0, std::move(cond_datum), num_rows); } -const std::vector kNumRowsArgNames{"num_rows"}; -const std::vector kNumRowsArg = benchmark::CreateRange(1, 64 * 1024, 32); +const std::string kNumRowsArgName = "num_rows"; +const std::vector kNumRowsArg{1, 4 * 1024, 64 * 1024}; -#define BM(name, if_else, ...) \ - BENCHMARK_IF_ELSE(BM_IfElseTrivialCond, name, if_else, kNumRowsArgNames, \ +#define BM(name, if_else, spin_func, ...) \ + BENCHMARK_IF_ELSE(BM_IfElseScalarCond, name, if_else, spin_func, {kNumRowsArgName}, \ {kNumRowsArg}, ##__VA_ARGS__) -BENCHMARK_IF_ELSE_WITH_BASELINE_AND_SV_SUPPRESS(BM, literal_null, kBooleanNull) -BENCHMARK_IF_ELSE_WITH_BASELINE_AND_SV_SUPPRESS(BM, literal_true, literal(true)) -BENCHMARK_IF_ELSE_WITH_BASELINE_AND_SV_SUPPRESS(BM, literal_false, literal(false)) +BENCHMARK_IF_ELSE_WITH_BASELINE_AND_SPECIAL(BM, Null, kBooleanNullScalar) +BENCHMARK_IF_ELSE_WITH_BASELINE_AND_SPECIAL(BM, True, kTrueScalar) +BENCHMARK_IF_ELSE_WITH_BASELINE_AND_SPECIAL(BM, False, kFalseScalar) #undef BM -namespace { - -void BenchmarkIfElseWithCondArray( - benchmark::State& state, - std::function if_else_func, - int64_t num_rows, double true_probability, double null_probability) { +static void BenchmarkIfElseArrayCond(benchmark::State& state, + MakeIfElseFunc make_if_else_func, + std::string spin_func, int64_t num_rows, + double true_probability, double null_probability, + int64_t if_true_kernel_intensity, + int64_t if_false_kernel_intensity) { random::RandomArrayGenerator rag(42); - auto b = rag.Boolean(num_rows, true_probability, null_probability); - auto schema = arrow::schema({field("b", boolean())}); + auto cond_datum = rag.Boolean(num_rows, true_probability, null_probability); - ExecBatch batch{std::vector{std::move(b)}, num_rows}; - - BenchmarkIfElse(state, std::move(if_else_func), field_ref("b"), literal(1), literal(0), - schema, batch); + BenchmarkIfElse(state, std::move(make_if_else_func), spin_func, + if_true_kernel_intensity, if_false_kernel_intensity, + std::move(cond_datum), num_rows); } -} // namespace - -template -static void BM_IfElseNumRows( - benchmark::State& state, - std::function if_else_func, - Args&&...) { - const int64_t num_rows = state.range(0); - - BenchmarkIfElseWithCondArray(state, std::move(if_else_func), num_rows, - /*true_probability=*/0.5, /*null_probability=*/0.0); +// Benchmark that: +// - Both branches are evenly heavy. +// - Array condition of tunable null probability. +// See if skipping evaluating both true/false branches takes place. +static void BM_IfElseEvenBranchesNullProbability(benchmark::State& state, + MakeIfElseFunc make_if_else_func, + std::string spin_func) { + const double null_probability = state.range(0) / 100.0; + + BenchmarkIfElseArrayCond(state, std::move(make_if_else_func), spin_func, + /*num_rows=*/16 * 1024, /*true_probability=*/0.5, + null_probability, + /*if_true_kernel_intensity=*/0, + /*if_false_kernel_intensity=*/0); } -#define BM(name, if_else, ...) \ - BENCHMARK_IF_ELSE(BM_IfElseNumRows, name, if_else, kNumRowsArgNames, {kNumRowsArg}) -BENCHMARK_IF_ELSE_WITH_BASELINE_AND_SV_SUPPRESS(BM, num_rows) -#undef BM +const std::string kNullProbabilityArgName = "null_probability"; +const std::vector kNullProbabilityArg{0, 25, 50, 100}; -template -static void BM_IfElseNullProbability( - benchmark::State& state, - std::function if_else_func, - Args&&...) { - const int64_t num_rows = state.range(0); - const double null_probability = state.range(1) / 100.0; +#define BM(name, if_else, spin_func, ...) \ + BENCHMARK_IF_ELSE(BM_IfElseEvenBranchesNullProbability, name, if_else, spin_func, \ + {kNullProbabilityArgName}, {kNullProbabilityArg}, ##__VA_ARGS__) +BENCHMARK_IF_ELSE_WITH_BASELINE_AND_SPECIAL(BM, ) +#undef BM - BenchmarkIfElseWithCondArray(state, std::move(if_else_func), num_rows, - /*true_probability=*/0.5, null_probability); +// Benchmark that: +// - Both branches are evenly heavy. +// - Array condition of tunable true probability. +// See the performance of different selectiveties and short-circuiting for extreme cases. +static void BM_IfElseEvenBranchesTrueProbability(benchmark::State& state, + MakeIfElseFunc make_if_else_func, + std::string spin_func) { + const double true_probability = state.range(0) / 100.0; + + BenchmarkIfElseArrayCond(state, std::move(make_if_else_func), spin_func, + /*num_rows=*/16 * 1024, true_probability, + /*null_probability=*/0, + /*if_true_kernel_intensity=*/0, + /*if_false_kernel_intensity=*/0); } -const std::vector kNumRowsAndNullProbabilityArgNames{"num_rows", - "null_probability"}; -const std::vector> kNumRowsAndNullProbabilityArgs{ - {4 * 1024, 64 * 1024}, {0, 50, 90, 100}}; +const std::string kTrueProbabilityArgName = "true_probability"; +const std::vector kTrueProbabilityArg{0, 25, 50, 100}; -#define BM(name, if_else, ...) \ - BENCHMARK_IF_ELSE(BM_IfElseNullProbability, name, if_else, \ - kNumRowsAndNullProbabilityArgNames, kNumRowsAndNullProbabilityArgs) -BENCHMARK_IF_ELSE_WITH_BASELINE_AND_SV_SUPPRESS(BM, null_probability); +#define BM(name, if_else, spin_func, ...) \ + BENCHMARK_IF_ELSE(BM_IfElseEvenBranchesTrueProbability, name, if_else, spin_func, \ + {kTrueProbabilityArgName}, {kTrueProbabilityArg}, ##__VA_ARGS__) +BENCHMARK_IF_ELSE_WITH_BASELINE_AND_SPECIAL(BM, ) #undef BM -template -void BM_IfElseWithOneHeavySide( - benchmark::State& state, - std::function if_else_func, - Args&&...) { - const double heavy_ratio = state.range(0) / 100.0; - constexpr int64_t num_rows = 65536; - - random::RandomArrayGenerator rag(42); - auto b = - rag.Boolean(num_rows, /*true_probability=*/heavy_ratio, /*null_probability=*/0.0); - auto i = ConstantArrayGenerator::Int32(num_rows, 42); - auto schema = arrow::schema({field("b", boolean()), field("i", int32())}); - - ExecBatch batch{std::vector{std::move(b), std::move(i)}, num_rows}; - - BenchmarkIfElse(state, std::move(if_else_func), field_ref("b"), heavy(field_ref("i")), - literal(0), schema, batch); +// Benchmark that: +// - False branch is tunably heavier than true branch. +// - Array condition of tunable true probability. +// See the performance benefit of maskable execution when skipping the heavy false branch. +static void BM_IfElseHeavyFalse(benchmark::State& state, MakeIfElseFunc make_if_else_func, + std::string spin_func) { + const double true_probability = state.range(0) / 100.0; + const int64_t heaviness = state.range(1); + + BenchmarkIfElseArrayCond(state, std::move(make_if_else_func), spin_func, + /*num_rows=*/16 * 1024, /*true_probability=*/true_probability, + /*null_probability=*/0, + /*if_true_kernel_intensity=*/0, heaviness); } -const std::vector kHeavyRatioArgNames{"heavy_ratio"}; -const std::vector kHeavyRatioArgs{{0, 25, 50, 75, 100}}; +const std::string kHeavinessArgName = "heaviness"; +const std::vector kHeavinessArg{0, 10, 100}; -#define BM(name, if_else, ...) \ - BENCHMARK_IF_ELSE(BM_IfElseWithOneHeavySide, name, if_else, kHeavyRatioArgNames, \ - {kHeavyRatioArgs}) -BENCHMARK_IF_ELSE_WITH_BASELINE_AND_SV_SUPPRESS(BM, one_heavy_side); +#define BM(name, if_else, spin_func, ...) \ + BENCHMARK_IF_ELSE(BM_IfElseHeavyFalse, name, if_else, spin_func, \ + ARROW_ALLOW_COMMA({kTrueProbabilityArgName, kHeavinessArgName}), \ + ARROW_ALLOW_COMMA({{10, 50, 90}, kHeavinessArg}), ##__VA_ARGS__) +BENCHMARK_IF_ELSE_WITH_BASELINE_AND_SPECIAL(BM, ) #undef BM -} // namespace compute - -} // namespace arrow +} // namespace arrow::compute diff --git a/cpp/src/arrow/util/macros.h b/cpp/src/arrow/util/macros.h index 55bc1eeb1d2d..f8fb81a94d07 100644 --- a/cpp/src/arrow/util/macros.h +++ b/cpp/src/arrow/util/macros.h @@ -22,6 +22,7 @@ #define ARROW_EXPAND(x) x #define ARROW_STRINGIFY(x) #x #define ARROW_CONCAT(x, y) x##y +#define ARROW_ALLOW_COMMA(...) __VA_ARGS__ // From Google gutil #ifndef ARROW_DISALLOW_COPY_AND_ASSIGN From 8bf870c7bb239d82f1d862dcd8ae00a2a589abcc Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Mon, 10 Nov 2025 13:04:15 -0800 Subject: [PATCH 59/71] Default params --- .../special/if_else_special_benchmark.cc | 19 ++++++------------- 1 file changed, 6 insertions(+), 13 deletions(-) diff --git a/cpp/src/arrow/compute/special/if_else_special_benchmark.cc b/cpp/src/arrow/compute/special/if_else_special_benchmark.cc index ac42ab28bdce..656db7c9b9c6 100644 --- a/cpp/src/arrow/compute/special/if_else_special_benchmark.cc +++ b/cpp/src/arrow/compute/special/if_else_special_benchmark.cc @@ -199,12 +199,10 @@ BENCHMARK_IF_ELSE_WITH_BASELINE_AND_SPECIAL(BM, True, kTrueScalar) BENCHMARK_IF_ELSE_WITH_BASELINE_AND_SPECIAL(BM, False, kFalseScalar) #undef BM -static void BenchmarkIfElseArrayCond(benchmark::State& state, - MakeIfElseFunc make_if_else_func, - std::string spin_func, int64_t num_rows, - double true_probability, double null_probability, - int64_t if_true_kernel_intensity, - int64_t if_false_kernel_intensity) { +static void BenchmarkIfElseArrayCond( + benchmark::State& state, MakeIfElseFunc make_if_else_func, std::string spin_func, + int64_t num_rows, double true_probability = 0.5, double null_probability = 0, + int64_t if_true_kernel_intensity = 0, int64_t if_false_kernel_intensity = 0) { random::RandomArrayGenerator rag(42); auto cond_datum = rag.Boolean(num_rows, true_probability, null_probability); @@ -224,9 +222,7 @@ static void BM_IfElseEvenBranchesNullProbability(benchmark::State& state, BenchmarkIfElseArrayCond(state, std::move(make_if_else_func), spin_func, /*num_rows=*/16 * 1024, /*true_probability=*/0.5, - null_probability, - /*if_true_kernel_intensity=*/0, - /*if_false_kernel_intensity=*/0); + null_probability); } const std::string kNullProbabilityArgName = "null_probability"; @@ -248,10 +244,7 @@ static void BM_IfElseEvenBranchesTrueProbability(benchmark::State& state, const double true_probability = state.range(0) / 100.0; BenchmarkIfElseArrayCond(state, std::move(make_if_else_func), spin_func, - /*num_rows=*/16 * 1024, true_probability, - /*null_probability=*/0, - /*if_true_kernel_intensity=*/0, - /*if_false_kernel_intensity=*/0); + /*num_rows=*/16 * 1024, true_probability); } const std::string kTrueProbabilityArgName = "true_probability"; From e5b99f8b9a5f2d840f4097fb195b1b727afc9f93 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Wed, 12 Nov 2025 18:17:14 -0800 Subject: [PATCH 60/71] Add comment of overall rational for special form --- cpp/src/arrow/compute/api_special.h | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/cpp/src/arrow/compute/api_special.h b/cpp/src/arrow/compute/api_special.h index 77b0f2b20899..5ce6dd5d268e 100644 --- a/cpp/src/arrow/compute/api_special.h +++ b/cpp/src/arrow/compute/api_special.h @@ -19,6 +19,32 @@ namespace arrow::compute { +/// The concept of "special form" is borrowed from Lisp +/// (https://courses.cs.northwestern.edu/325/readings/special-forms.html). A special form +/// is used to implement evaluation strategies +/// (https://en.wikipedia.org/wiki/Evaluation_strategy) other than the default +/// call-by-value strategy used by Arrow expression evaluation. Velox also uses this term. +/// +/// In a call-by-value strategy, all arguments are evaluated before the function +/// invocation. For example, consider a regular function call +/// if_else(a, b, c) +/// Under call-by-value semantics, the expressions `a`, `b`, and `c` are all evaluated +/// before calling `if_else`. This can lead to unintuitive behavior when subexpressions +/// have observable side effects. For instance, +/// if_else(not_equal(a, 0), divide(b, a), 0) +/// should never produce a divide-by-zero error in most programming languages. However, +/// under call-by-value semantics, `divide(b, a)` is evaluated regardless of the +/// condition, so a divide-by-zero error can still occur. To address this, a special form +/// for `if_else` would be needed, namely `if_else_special`, that follows a +/// call-by-name-like evaluation strategy, where, for each row in a batch, only one of `b` +/// or `c` is evaluated based on the corresponding value of `a`. +/// +/// Each API in this file is intended to refer to a concrete special form. In addition to +/// the aforementioned `if_else_special`, the design anticipates variants of conditional +/// constructs such as `case_when_special` and `coalesce_special`, as well as boolean +/// operators with short-circuit semantics, such as `and_special` and `or_special`, some +/// of which may not be implemented yet. + /// TODO: Doc. ARROW_EXPORT Expression if_else_special(Expression cond, Expression if_true, Expression if_false); From b1bd8479d531c7201f4652a979d7956c7a7d1a9e Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Thu, 13 Nov 2025 18:31:52 -0800 Subject: [PATCH 61/71] Refine overall special form comment --- cpp/src/arrow/compute/api_special.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/src/arrow/compute/api_special.h b/cpp/src/arrow/compute/api_special.h index 5ce6dd5d268e..ae6c0b0c0c21 100644 --- a/cpp/src/arrow/compute/api_special.h +++ b/cpp/src/arrow/compute/api_special.h @@ -36,8 +36,8 @@ namespace arrow::compute { /// under call-by-value semantics, `divide(b, a)` is evaluated regardless of the /// condition, so a divide-by-zero error can still occur. To address this, a special form /// for `if_else` would be needed, namely `if_else_special`, that follows a -/// call-by-name-like evaluation strategy, where, for each row in a batch, only one of `b` -/// or `c` is evaluated based on the corresponding value of `a`. +/// call-by-name-like evaluation strategy, where, for each row in a batch, only one of the +/// branches is evaluated based on the corresponding value of condition. /// /// Each API in this file is intended to refer to a concrete special form. In addition to /// the aforementioned `if_else_special`, the design anticipates variants of conditional From c98ace80a0b628d33e98ca879faf822efde2b09b Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Fri, 14 Nov 2025 13:41:45 -0800 Subject: [PATCH 62/71] Comment on special form and executor --- cpp/src/arrow/compute/special_form.h | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/cpp/src/arrow/compute/special_form.h b/cpp/src/arrow/compute/special_form.h index 6b74079c8320..dd262b8d7a0c 100644 --- a/cpp/src/arrow/compute/special_form.h +++ b/cpp/src/arrow/compute/special_form.h @@ -21,6 +21,23 @@ namespace arrow::compute { +/// @brief A bound representation of a special form that can be invoked directly on an +/// ExecBatch to produce a result Datum during expression evaluation. Conceptually, a +/// SpecialExecutor plays a role similar to a function kernel: it encapsulates the +/// concrete execution logic of a special form after binding. All child expressions, +/// as well as the special form itself, have already been bound to the input schema +/// and argument types at this point. +/// +/// Unlike regular expression evaluation under the default call‑by‑value strategy [1], +/// the child expressions of a special form are not evaluated ahead of time. A +/// SpecialExecutor is therefore free to choose its own evaluation strategy, such as +/// call-by-name [2], deciding when and whether to evaluate individual child expressions +/// according to the semantics of the special form, such as conditional branching or +/// boolean short-circuiting. Implementations may leverage mechanisms such as +/// selection-vector-based masked execution to cope with the vectorized execution. +/// +/// [1] https://en.wikipedia.org/wiki/Evaluation_strategy#Call_by_value +/// [2] https://en.wikipedia.org/wiki/Evaluation_strategy#Call_by_name class ARROW_EXPORT SpecialExecutor { public: explicit SpecialExecutor(TypeHolder out_type, @@ -40,6 +57,14 @@ class ARROW_EXPORT SpecialExecutor { const std::shared_ptr options_; }; +/// @brief An unbound representation of a special form, which can be bound to produce +/// a SpecialExecutor for execution during expression evaluation. +/// +/// A SpecialForm is intentionally immutable and independent of any concrete input schema, +/// argument types, options, or data. It defines how a special form should be bound to a +/// given set of bound arguments and options, to produce a concrete SpecialExecutor for +/// one particular invocation. Therefor implementations are naturally stateless and may be +/// modeled as singletons. class ARROW_EXPORT SpecialForm { public: explicit SpecialForm(std::string name) : name_(std::move(name)) {} From 51546dc87c812d5c3332fbf36a1ab48807bcf76f Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Fri, 14 Nov 2025 17:02:04 -0800 Subject: [PATCH 63/71] Update branch abstraction doc --- .../arrow/compute/special/conditional_special_internal.h | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/cpp/src/arrow/compute/special/conditional_special_internal.h b/cpp/src/arrow/compute/special/conditional_special_internal.h index 1f482e2e8152..0bec8a5baa99 100644 --- a/cpp/src/arrow/compute/special/conditional_special_internal.h +++ b/cpp/src/arrow/compute/special/conditional_special_internal.h @@ -26,13 +26,12 @@ namespace arrow::compute::internal { -/// Structures to manage masks for branching expressions. -/// Mostly for efficient short circuiting of if-else chains. -/// The whole abstraction is as follows: -/// - A branch represents an compound expression of condition and body; +/// Structures to model masks for branching expressions, mostly for efficient +/// short-circuiting of branching chains. The whole abstraction is as follows: +/// - A branch represents a compound expression of condition and body; /// - A branch mask represents the set of rows (in an ExecBatch) to be evaluated for the /// branch condition; -/// - When a branch mask is applied to an condition expression, it produces a body mask; +/// - When a branch mask is applied to a condition expression, it produces a body mask; /// - A body mask represents the set of rows (in an ExecBatch) to be evaluated for the /// branch body. /// - A body mask also produces the next branch mask, representing the set of rows From c4e5d5cc64a2fe9600b53228750d76d48b3c5b87 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Fri, 14 Nov 2025 23:23:42 -0800 Subject: [PATCH 64/71] Small refine and more comments --- .../compute/special/conditional_special.cc | 4 -- .../special/conditional_special_internal.h | 60 +++++++++++++++---- 2 files changed, 47 insertions(+), 17 deletions(-) diff --git a/cpp/src/arrow/compute/special/conditional_special.cc b/cpp/src/arrow/compute/special/conditional_special.cc index db70524fffb0..05e9b792c347 100644 --- a/cpp/src/arrow/compute/special/conditional_special.cc +++ b/cpp/src/arrow/compute/special/conditional_special.cc @@ -131,10 +131,6 @@ Result> BranchMask::FromSelectionVector( return std::make_shared(std::move(selection), length); } -Result> AllPassBranchMask::GetSelectionVector() const { - return nullptr; -} - Result> AllPassBranchMask::MakeBodyMaskFromBitmap( const std::shared_ptr& bitmap, ExecContext* exec_context) const { DCHECK_EQ(bitmap->length(), length_); diff --git a/cpp/src/arrow/compute/special/conditional_special_internal.h b/cpp/src/arrow/compute/special/conditional_special_internal.h index 0bec8a5baa99..6405180bf4db 100644 --- a/cpp/src/arrow/compute/special/conditional_special_internal.h +++ b/cpp/src/arrow/compute/special/conditional_special_internal.h @@ -23,6 +23,7 @@ #include "arrow/compute/exec.h" #include "arrow/compute/special/special_form_internal.h" #include "arrow/compute/visibility.h" +#include "arrow/util/unreachable.h" namespace arrow::compute::internal { @@ -39,33 +40,57 @@ namespace arrow::compute::internal { struct ARROW_COMPUTE_EXPORT BodyMask; +/// @brief A mask representing the set of rows to be evaluated for a branch condition. +/// Being empty indicates that the entire branch, in addition to all the subsequent +/// branches, are concluded. Otherwise, a selection vector can be obtained to evaluate the +/// branch condition, whose result will be used to further produce a body mask. struct ARROW_COMPUTE_EXPORT BranchMask : public std::enable_shared_from_this { virtual ~BranchMask() = default; + /// @brief Check if the branch mask is empty, in which case no rows are to be evaluated. virtual bool empty() const = 0; + /// @brief Get the selection vector representing the rows to be evaluated. Null + /// indicates that all rows are to be evaluated. virtual Result> GetSelectionVector() const = 0; + /// @brief Create a body mask for this branch from the given datum, which is the result + /// of evaluating the condition under this branch mask. All possible trivial cases, such + /// as constant true/false/null and all-true/false/null arrays, are handled here. If + /// none of the trivial cases apply, the call is forwarded to the concrete + /// implementations of MakeBodyMaskFromBitmap(). Result> MakeBodyMask(const Datum& datum, ExecContext* exec_context) const; + /// @brief Create a branch mask from the given selection vector. Based on the content of + /// the selection vector, it may return concrete branch mask implementations that can + /// take advantage of short-circuiting. static Result> FromSelectionVector( std::shared_ptr selection, int64_t length); protected: + /// @brief Create a body mask from the given bitmap, which is the result of evaluating + /// the condition under this branch mask. virtual Result> MakeBodyMaskFromBitmap( const std::shared_ptr& bitmap, ExecContext* exec_context) const = 0; + /// @brief Create a body mask from the given chunked bitmap, which is the result of + /// evaluating the condition under this branch mask. virtual Result> MakeBodyMaskFromBitmap( const std::shared_ptr& bitmap, ExecContext* exec_context) const = 0; }; +/// @brief A branch mask that evaluates the condition for all rows. The selection vector +/// obtained is null, indicating that all rows are to be evaluated. And the body mask +/// produced is solely based on the condition result. struct ARROW_COMPUTE_EXPORT AllPassBranchMask : public BranchMask { explicit AllPassBranchMask(int64_t length) : length_(length) {} bool empty() const override { return length_ == 0; } - Result> GetSelectionVector() const override; + Result> GetSelectionVector() const override { + return NULLPTR; + } protected: Result> MakeBodyMaskFromBitmap( @@ -80,34 +105,36 @@ struct ARROW_COMPUTE_EXPORT AllPassBranchMask : public BranchMask { int64_t length_; }; +/// @brief A branch mask that evaluates the condition for no rows and concludes +/// all the subsequent branches. One should never try to obtain a selection vector or +/// produce a body mask from this branch mask. struct ARROW_COMPUTE_EXPORT AllFailBranchMask : public BranchMask { AllFailBranchMask() = default; bool empty() const override { return true; } Result> GetSelectionVector() const override { - DCHECK(false); - return Status::Invalid("AllFailBranchMask::GetSelectionVector should not be called"); + Unreachable("AllFailBranchMask::GetSelectionVector should not be called"); } protected: Result> MakeBodyMaskFromBitmap( const std::shared_ptr& bitmap, ExecContext* exec_context) const override { - DCHECK(false); - return Status::Invalid( - "AllFailBranchMask::MakeBodyMaskFromBitmap should not be called"); + Unreachable("AllFailBranchMask::MakeBodyMaskFromBitmap should not be called"); } Result> MakeBodyMaskFromBitmap( const std::shared_ptr& bitmap, ExecContext* exec_context) const override { - DCHECK(false); - return Status::Invalid( - "AllFailBranchMask::MakeBodyMaskFromBitmap should not be called"); + Unreachable("AllFailBranchMask::MakeBodyMaskFromBitmap should not be called"); } }; +/// @brief A branch mask that evaluates the condition for rows indicated by the given +/// selection vector, which is also the one obtained from it. The body mask produced, if +/// no short-circuiting available, is one that with a selection vector that is +/// conceptually AND-ing the branch mask's selection vector and the condition result. struct ARROW_COMPUTE_EXPORT ConditionalBranchMask : public BranchMask { ConditionalBranchMask(std::shared_ptr selection_vector, int64_t length) : selection_vector_(std::move(selection_vector)), length_(length) { @@ -136,13 +163,22 @@ struct ARROW_COMPUTE_EXPORT ConditionalBranchMask : public BranchMask { int64_t length_ = 0; }; +/// @brief A mask representing the set of rows to be evaluated for a branch body. Being +/// empty indicates that no rows are to be evaluated for the body. Otherwise, a selection +/// vector can be obtained to evaluate the branch body, whose result will be used as part +/// of the final result. In addition, a branch mask is produced for the next branch. struct ARROW_COMPUTE_EXPORT BodyMask : public std::enable_shared_from_this { virtual ~BodyMask() = default; + /// @brief Check if the body mask is empty, in which case no rows are to be evaluated. virtual bool empty() const = 0; + /// @brief Get the selection vector representing the rows to be evaluated. Null + /// indicates that all rows are to be evaluated. virtual Result> GetSelectionVector() const = 0; + /// @brief Create a branch mask for the next branch. May return concrete branch mask + /// implementations for short-circuiting. virtual Result> NextBranchMask() const = 0; }; @@ -152,8 +188,7 @@ struct ARROW_COMPUTE_EXPORT AllNullBodyMask : public BodyMask { bool empty() const override { return true; } Result> GetSelectionVector() const override { - DCHECK(false); - return Status::Invalid("AllNullBodyMask::GetSelectionVector should not be called"); + Unreachable("AllNullBodyMask::GetSelectionVector should not be called"); } Result> NextBranchMask() const override { @@ -189,8 +224,7 @@ struct ARROW_COMPUTE_EXPORT AllFailBodyMask : public DelegateBodyMask { bool empty() const override { return true; } Result> GetSelectionVector() const override { - DCHECK(false); - return Status::Invalid("AllFailBodyMask::GetSelectionVector should not be called"); + Unreachable("AllFailBodyMask::GetSelectionVector should not be called"); } Result> NextBranchMask() const override { From 89b6b20ae27118da651e6c9a60f8a79954fc27b4 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Tue, 18 Nov 2025 00:39:20 -0800 Subject: [PATCH 65/71] DCHECK for the last branch remaining no rows --- cpp/src/arrow/compute/special/conditional_special.cc | 2 ++ 1 file changed, 2 insertions(+) diff --git a/cpp/src/arrow/compute/special/conditional_special.cc b/cpp/src/arrow/compute/special/conditional_special.cc index 05e9b792c347..6b4fe41b4bbe 100644 --- a/cpp/src/arrow/compute/special/conditional_special.cc +++ b/cpp/src/arrow/compute/special/conditional_special.cc @@ -286,6 +286,8 @@ Result ConditionalExec::Execute(const ExecBatch& input, results.Emplace(std::move(body_result), std::move(selection_vector)); ARROW_ASSIGN_OR_RAISE(branch_mask, body_mask->NextBranchMask()); } + // Should have no remaining rows. + DCHECK(branch_mask->empty()); return MultiplexResults(input, results, exec_context); } From bb89a964fed548338122c03fa07ac1d1e3afd064 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Tue, 18 Nov 2025 00:39:37 -0800 Subject: [PATCH 66/71] Comments for body mask impls --- .../special/conditional_special_internal.h | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/cpp/src/arrow/compute/special/conditional_special_internal.h b/cpp/src/arrow/compute/special/conditional_special_internal.h index 6405180bf4db..766aca179dba 100644 --- a/cpp/src/arrow/compute/special/conditional_special_internal.h +++ b/cpp/src/arrow/compute/special/conditional_special_internal.h @@ -182,6 +182,11 @@ struct ARROW_COMPUTE_EXPORT BodyMask : public std::enable_shared_from_this> NextBranchMask() const = 0; }; +/// @brief A body mask that emits nulls for all rows and produces an all-fail branch mask +/// for the next branch. For example, the body mask for branch: +/// [... else] if (null) ... +// XXX Only works for null policy of intersection (any operands null -> null). Other +// variants may needed for different null policies. struct ARROW_COMPUTE_EXPORT AllNullBodyMask : public BodyMask { AllNullBodyMask() = default; @@ -196,6 +201,8 @@ struct ARROW_COMPUTE_EXPORT AllNullBodyMask : public BodyMask { } }; +/// @brief A body mask that delegates certain operations to an underlying branch mask. +/// Subclasses can override behaviors as needed. struct ARROW_COMPUTE_EXPORT DelegateBodyMask : public BodyMask { explicit DelegateBodyMask(std::shared_ptr branch_mask) : branch_mask_(std::move(branch_mask)) {} @@ -204,6 +211,10 @@ struct ARROW_COMPUTE_EXPORT DelegateBodyMask : public BodyMask { std::shared_ptr branch_mask_; }; +/// @brief A body mask that evaluates the body for all rows indicated by the underlying +/// branch mask, and produces an all-fail branch mask for the next branch. For example, +/// the body mask for branch: +/// [... else] if (true) ... struct ARROW_COMPUTE_EXPORT AllPassBodyMask : public DelegateBodyMask { using DelegateBodyMask::DelegateBodyMask; @@ -218,6 +229,9 @@ struct ARROW_COMPUTE_EXPORT AllPassBodyMask : public DelegateBodyMask { } }; +/// @brief A body mask that evaluates the body for no rows, and pass through the +/// underlying branch mask for the next branch. For example, the body mask for branch: +/// [... else] if (false) ... struct ARROW_COMPUTE_EXPORT AllFailBodyMask : public DelegateBodyMask { using DelegateBodyMask::DelegateBodyMask; @@ -232,6 +246,8 @@ struct ARROW_COMPUTE_EXPORT AllFailBodyMask : public DelegateBodyMask { } }; +/// @brief A body mask that evaluates the body for rows indicated by the given selection +/// vector, and produces a branch mask for the next branch from the remainder rows. struct ARROW_COMPUTE_EXPORT ConditionalBodyMask : public BodyMask { ConditionalBodyMask(std::shared_ptr body, std::shared_ptr remainder, int64_t length) From c1dff53d6702d494767da99e6609a34cfef10405 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Tue, 18 Nov 2025 16:59:36 -0800 Subject: [PATCH 67/71] Comment on function backed special form crtp --- cpp/src/arrow/compute/special/special_form_internal.h | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/cpp/src/arrow/compute/special/special_form_internal.h b/cpp/src/arrow/compute/special/special_form_internal.h index 9953b94a2ffe..aaf09453d4c0 100644 --- a/cpp/src/arrow/compute/special/special_form_internal.h +++ b/cpp/src/arrow/compute/special/special_form_internal.h @@ -23,6 +23,14 @@ namespace arrow::compute::internal { +/// @brief A CRTP base class for special forms whose binding are backed by a function +/// call. +/// +/// Many special forms share the same binding logic as its non-special function +/// counterpart, e.g., implicit casts and output type resolution. This class encapsulates +/// the binding logic for such special forms, instantiating a Call instance and binding +/// it, then delegating the actual binding of the special form to the derived class via +/// BindWithBoundCall() with the bound Call instance. template class FunctionBackedSpecialForm : public SpecialForm { public: From cdd682468b25dacb39a4c0dc6f59535586b434b5 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Wed, 19 Nov 2025 18:23:29 -0800 Subject: [PATCH 68/71] More comment --- .../compute/special/conditional_special.cc | 11 ++++ .../special/conditional_special_internal.h | 57 ++++++++++++++++++- 2 files changed, 65 insertions(+), 3 deletions(-) diff --git a/cpp/src/arrow/compute/special/conditional_special.cc b/cpp/src/arrow/compute/special/conditional_special.cc index 6b4fe41b4bbe..d3cc52e5cc84 100644 --- a/cpp/src/arrow/compute/special/conditional_special.cc +++ b/cpp/src/arrow/compute/special/conditional_special.cc @@ -295,19 +295,30 @@ Result ConditionalExec::MultiplexResults(const ExecBatch& input, const BranchResults& results, ExecContext* exec_context) const { if (results.empty()) { + // No branches were taken, return an array of nulls. return MakeArrayOfNull(result_type.GetSharedPtr(), input.length, exec_context->memory_pool()); } if (results.size() == 1) { + // Single branch taken. const auto& result = results.body_results()[0]; if (results.selection_vectors()[0] == nullptr) { + // This branch has no selection vector, then this branch covers all rows, regardless + // of the existence of outer selection vector, return as is. return result; } if (input.selection_vector == nullptr) { + // This branch has a selection vector but there is no outer selection vector, then + // this branch must not be covering all rows - other branches might just have all + // failed. And we need to go through the choose path to fill in nulls for the rows + // not covered. DCHECK_NE(results.selection_vectors()[0]->length(), input.length); } else { if (results.selection_vectors()[0]->length() == input.selection_vector->length()) { + // This branch has a selection vector and there is outer selection vector, and + // their lengths equal, then this branch must be covering all rows under the outer + // selection vector, return as is. return result; } } diff --git a/cpp/src/arrow/compute/special/conditional_special_internal.h b/cpp/src/arrow/compute/special/conditional_special_internal.h index 766aca179dba..636fa625746e 100644 --- a/cpp/src/arrow/compute/special/conditional_special_internal.h +++ b/cpp/src/arrow/compute/special/conditional_special_internal.h @@ -35,8 +35,38 @@ namespace arrow::compute::internal { /// - When a branch mask is applied to a condition expression, it produces a body mask; /// - A body mask represents the set of rows (in an ExecBatch) to be evaluated for the /// branch body. -/// - A body mask also produces the next branch mask, representing the set of rows -/// remaining to be evaluated for the next branch in the chain. +/// - A body mask also preserves information from its originating branch mask and derives +/// the next branch mask, representing the set of rows remaining to be evaluated for the +/// next branch in the chain. +/// +/// For example, consider the following conditional special form: +/// if_else_sp(/*cond=*/eq(a, 'x'), /*if_true=*/foo(b), /*if_false=*/bar(c)) +/// is being evaluated on an ExecBatch: +/// [a: ['x', 'y', 'x', 'z'], b: ['b0', 'b1', 'b2', 'b3'], c: ['c0', 'c1', 'c2', 'c3']] +/// We'll have an initial branch mask that passes all rows: +/// BranchMask0: [0, 1, 2, 3] // all rows in the batch +/// (In practice we can have a specialized branch mask implementation that doesn't +/// necessarily store all the row indices when all rows are to be evaluated.) +/// Then BranchMask0 is applied to the condition eq(a, 'x'), producing the condition +/// result: +/// [true, false, true, false] +/// Which is then used by BranchMask0 to make a body mask: +/// BodyMask0: [0, 2] // rows with true condition out of [0, 1, 2, 3] +/// It is then applied to the first branch body foo(b), producing the result for this +/// branch: +/// [foo('b0'), foo('b2')] // at rows [0, 2] +/// After that, BodyMask0 produces the next branch mask: +/// BranchMask1: [1, 3] // [0, 1, 2, 3] - [0, 2] +/// Which is then applied to the next branch condition, which is an implicit true literal +/// in this case, producing the condition result: +/// [true, true] // at rows [1, 3] +/// Which is then used by BranchMask1 to make a body mask: +/// BodyMask1: [1, 3] // rows with true condition out of [1, 3] +/// It is then applied to the second branch body bar(c), producing the result for this +/// branch: +/// [bar('c1'), bar('c3')] // at rows [1, 3] +/// Finally, the results from all branches are combined to produce the final result: +/// [foo('b0'), bar('c1'), foo('b2'), bar('c3')] struct ARROW_COMPUTE_EXPORT BodyMask; @@ -279,6 +309,15 @@ struct ARROW_COMPUTE_EXPORT Branch { Expression body; }; +/// @brief A simple structure that assembles the process of executing a sequence of +/// branches, including: +/// - Iterating all branches by: +/// - Evaluating each branch condition under the branch mask; +/// - Producing the body mask from the condition result; +/// - Evaluating the branch body under the body mask; +/// - Producing the next branch mask from the body mask. +/// - Collecting all branch body results and selection vectors, and multiplexing them into +/// the final result. struct ARROW_COMPUTE_EXPORT ConditionalExec { ConditionalExec(const std::vector& branches, const TypeHolder& result_type) : branches(branches), result_type(result_type) {} @@ -286,6 +325,8 @@ struct ARROW_COMPUTE_EXPORT ConditionalExec { Result Execute(const ExecBatch& input, ExecContext* exec_context) const&&; private: + /// @brief A simple helper structure to collect branch body results and their + /// corresponding selection vectors. struct BranchResults { void Reserve(int64_t size) { body_results_.reserve(size); @@ -312,6 +353,13 @@ struct ARROW_COMPUTE_EXPORT ConditionalExec { std::vector> selection_vectors_; }; + /// @brief Get the initial branch mask based on the existence of the selection vector in + /// the given ExecBatch. + /// + /// If a selection vector exists in the input batch, it implies that we are under a + /// masked execution, e.g., within another outer conditional special form. In this case, + /// the initial selection vector should be respected by all the branches, and thus + /// treated as the initial branch mask and propagated to the rest. Result> InitBranchMask( const ExecBatch& input, ExecContext* exec_context) const { if (input.selection_vector) { @@ -340,14 +388,17 @@ struct ARROW_COMPUTE_EXPORT ConditionalExec { return ExecuteScalarExpression(body, input_with_selection, exec_context); } + /// @brief Multiplex all branch body results into the final result based on their + /// corresponding selection vectors. Result MultiplexResults(const ExecBatch& input, const BranchResults& results, ExecContext* exec_context) const; + /// @brief A helper function to choose indices from multiple selection vectors. Result ChooseIndices( const std::vector>& selection_vectors, int64_t length, ExecContext* exec_context) const; - protected: + private: const std::vector& branches; const TypeHolder& result_type; }; From a16d448a3b674f173c2717b42dc1a7b9ce3ac946 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Wed, 19 Nov 2025 23:29:13 -0800 Subject: [PATCH 69/71] Refine and comment --- .../compute/special/conditional_special.cc | 71 +++++++++++-------- .../special/conditional_special_internal.h | 14 ++-- 2 files changed, 48 insertions(+), 37 deletions(-) diff --git a/cpp/src/arrow/compute/special/conditional_special.cc b/cpp/src/arrow/compute/special/conditional_special.cc index d3cc52e5cc84..f90da25c97e6 100644 --- a/cpp/src/arrow/compute/special/conditional_special.cc +++ b/cpp/src/arrow/compute/special/conditional_special.cc @@ -291,6 +291,46 @@ Result ConditionalExec::Execute(const ExecBatch& input, return MultiplexResults(input, results, exec_context); } +namespace { + +/// @brief Results multiplexing is done by invoking a "choose" function to choose values +/// from each branch result based on the selection vectors. This function prepares the +/// choose indices from the branch selection vectors. For the example in +/// ConditionalExec::MultiplexResults's doc string, the choose indices will be: +/// [0, 1, 2, 0, 1, 2, 0] +Result ChooseIndices( + const std::vector>& selection_vectors, + int64_t length, ExecContext* exec_context) { + const int64_t validity_bytes = bit_util::BytesForBits(length); + ARROW_ASSIGN_OR_RAISE( + std::shared_ptr validity_buf, + AllocateResizableBuffer(validity_bytes, exec_context->memory_pool())); + auto validity_data = validity_buf->mutable_data_as(); + std::memset(validity_data, 0, validity_bytes); + + ARROW_ASSIGN_OR_RAISE( + std::shared_ptr indices_buf, + AllocateResizableBuffer(length * sizeof(int32_t), exec_context->memory_pool())); + auto indices_data = indices_buf->mutable_data_as(); + for (int32_t index = 0; index < static_cast(selection_vectors.size()); + ++index) { + DCHECK_NE(selection_vectors[index], nullptr); + DCHECK_GT(selection_vectors[index]->length(), 0); + auto row_ids = selection_vectors[index]->indices(); + for (int64_t i = 0; i < selection_vectors[index]->length(); ++i) { + const int32_t row_id = row_ids[i]; + DCHECK_EQ(bit_util::GetBit(validity_data, row_id), false); + bit_util::SetBitTo(validity_data, row_id, true); + indices_data[row_id] = index; + } + } + + return ArrayData::Make(int32(), length, + {std::move(validity_buf), std::move(indices_buf)}); +} + +} // namespace + Result ConditionalExec::MultiplexResults(const ExecBatch& input, const BranchResults& results, ExecContext* exec_context) const { @@ -334,35 +374,4 @@ Result ConditionalExec::MultiplexResults(const ExecBatch& input, return CallFunction("choose", choose_args, exec_context); } -Result ConditionalExec::ChooseIndices( - const std::vector>& selection_vectors, - int64_t length, ExecContext* exec_context) const { - const int64_t validity_bytes = bit_util::BytesForBits(length); - ARROW_ASSIGN_OR_RAISE( - std::shared_ptr validity_buf, - AllocateResizableBuffer(validity_bytes, exec_context->memory_pool())); - auto validity_data = validity_buf->mutable_data_as(); - std::memset(validity_data, 0, validity_bytes); - - ARROW_ASSIGN_OR_RAISE( - std::shared_ptr indices_buf, - AllocateResizableBuffer(length * sizeof(int32_t), exec_context->memory_pool())); - auto indices_data = indices_buf->mutable_data_as(); - for (int32_t index = 0; index < static_cast(selection_vectors.size()); - ++index) { - DCHECK_NE(selection_vectors[index], nullptr); - DCHECK_GT(selection_vectors[index]->length(), 0); - auto row_ids = selection_vectors[index]->indices(); - for (int64_t i = 0; i < selection_vectors[index]->length(); ++i) { - const int32_t row_id = row_ids[i]; - DCHECK_EQ(bit_util::GetBit(validity_data, row_id), false); - bit_util::SetBitTo(validity_data, row_id, true); - indices_data[row_id] = index; - } - } - - return ArrayData::Make(int32(), length, - {std::move(validity_buf), std::move(indices_buf)}); -} - } // namespace arrow::compute::internal diff --git a/cpp/src/arrow/compute/special/conditional_special_internal.h b/cpp/src/arrow/compute/special/conditional_special_internal.h index 636fa625746e..88aa75bab0da 100644 --- a/cpp/src/arrow/compute/special/conditional_special_internal.h +++ b/cpp/src/arrow/compute/special/conditional_special_internal.h @@ -389,15 +389,17 @@ struct ARROW_COMPUTE_EXPORT ConditionalExec { } /// @brief Multiplex all branch body results into the final result based on their - /// corresponding selection vectors. + /// corresponding selection vectors. For example, given three branch body results and + /// selection vectors: + /// [a, -, -, d, -, -, g], [0, 3, 6] + /// [-, b, -, -, e, -, -], [1, 4] + /// [-, -, c, -, -, f, -], [2, 5] + /// Note each branch result has the same length as the input batch, the non-selected + /// rows are indicated by '-'. The multiplexed result will be: + /// [a, b, c, d, e, f, g] Result MultiplexResults(const ExecBatch& input, const BranchResults& results, ExecContext* exec_context) const; - /// @brief A helper function to choose indices from multiple selection vectors. - Result ChooseIndices( - const std::vector>& selection_vectors, - int64_t length, ExecContext* exec_context) const; - private: const std::vector& branches; const TypeHolder& result_type; From 18e56fd73237bda86a9fd3c009677008ccac3e14 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Wed, 19 Nov 2025 23:44:56 -0800 Subject: [PATCH 70/71] More comment --- cpp/src/arrow/compute/special/conditional_special.cc | 2 ++ 1 file changed, 2 insertions(+) diff --git a/cpp/src/arrow/compute/special/conditional_special.cc b/cpp/src/arrow/compute/special/conditional_special.cc index f90da25c97e6..9fc63a09fbbc 100644 --- a/cpp/src/arrow/compute/special/conditional_special.cc +++ b/cpp/src/arrow/compute/special/conditional_special.cc @@ -271,11 +271,13 @@ Result ConditionalExec::Execute(const ExecBatch& input, ARROW_ASSIGN_OR_RAISE(auto branch_mask, InitBranchMask(input, exec_context)); for (const auto& branch : branches) { if (branch_mask->empty()) { + // No more rows to evaluate. break; } ARROW_ASSIGN_OR_RAISE(auto body_mask, EvaluateCond(branch_mask, branch.cond, input, exec_context)); if (body_mask->empty()) { + // No rows taken for this branch. ARROW_ASSIGN_OR_RAISE(branch_mask, body_mask->NextBranchMask()); continue; } From 91c417c22d8fbd441de0fe9fac9f8c62845e7404 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Wed, 19 Nov 2025 23:45:51 -0800 Subject: [PATCH 71/71] More comments --- cpp/src/arrow/compute/api_special.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/arrow/compute/api_special.h b/cpp/src/arrow/compute/api_special.h index ae6c0b0c0c21..8173d7957018 100644 --- a/cpp/src/arrow/compute/api_special.h +++ b/cpp/src/arrow/compute/api_special.h @@ -45,7 +45,7 @@ namespace arrow::compute { /// operators with short-circuit semantics, such as `and_special` and `or_special`, some /// of which may not be implemented yet. -/// TODO: Doc. +/// @brief Construct an Expression representing an if-else special form. ARROW_EXPORT Expression if_else_special(Expression cond, Expression if_true, Expression if_false);