diff --git a/cpp/src/arrow/CMakeLists.txt b/cpp/src/arrow/CMakeLists.txt index 2e5c67e07b6e..befe18fa16dd 100644 --- a/cpp/src/arrow/CMakeLists.txt +++ b/cpp/src/arrow/CMakeLists.txt @@ -802,6 +802,8 @@ if(ARROW_COMPUTE) compute/row/grouper.cc compute/row/row_encoder_internal.cc compute/row/row_internal.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/CMakeLists.txt b/cpp/src/arrow/compute/CMakeLists.txt index 6c530a76e18f..6434f9edf75f 100644 --- a/cpp/src/arrow/compute/CMakeLists.txt +++ b/cpp/src/arrow/compute/CMakeLists.txt @@ -179,3 +179,5 @@ add_arrow_compute_benchmark(function_benchmark) 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..8173d7957018 --- /dev/null +++ b/cpp/src/arrow/compute/api_special.h @@ -0,0 +1,52 @@ +// 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 { + +/// 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 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 +/// 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. + +/// @brief Construct an Expression representing an if-else special form. +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/exec.cc b/cpp/src/arrow/compute/exec.cc index 1be398fdae9e..c14921f32ba6 100644 --- a/cpp/src/arrow/compute/exec.cc +++ b/cpp/src/arrow/compute/exec.cc @@ -50,6 +50,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 { @@ -359,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); @@ -367,6 +369,12 @@ 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(); + if (selection_vector_) { + selection_length_ = selection_vector_->length(); + } else { + selection_length_ = 0; + } return Status::OK(); } @@ -403,7 +411,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; @@ -442,6 +450,13 @@ bool ExecSpanIterator::Next(ExecSpan* span) { PromoteExecSpanScalars(span); } + if (!have_all_scalars_ || promote_if_all_scalars_) { + if (selection_vector_) { + DCHECK_NE(selection_span, nullptr); + *selection_span = SelectionVectorSpan(selection_vector_->indices()); + } + } + initialized_ = true; } else if (position_ == length_) { // We've emitted at least one span and we're at the end so we are done @@ -465,6 +480,23 @@ bool ExecSpanIterator::Next(ExecSpan* 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; @@ -694,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()); } @@ -781,17 +820,41 @@ 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 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); } + if (batch.selection_vector && !kernel_->selective_exec) { + return ExecuteSelectiveDense(batch, listener); + } + + return ExecuteBatch(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 ExecuteBatch(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 @@ -811,19 +874,40 @@ 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_); - } else { - // Outputs have just one element - return outputs[0]; + 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 ExecuteBatch(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())); + } + ARROW_ASSIGN_OR_RAISE( + ExecBatch input, + ExecBatch::Make(std::move(values), batch.selection_vector->length())); + + DatumAccumulator dense_listener; + RETURN_NOT_OK(ExecuteBatch(input, &dense_listener)); + Datum dense_result = WrapResults(input.values, dense_listener.values()); + + ARROW_ASSIGN_OR_RAISE(auto result, + Scatter(dense_result, *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 @@ -842,6 +926,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(); @@ -853,10 +942,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(); } @@ -866,10 +955,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)); } @@ -877,7 +966,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; @@ -888,7 +978,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, selection, out)); // Output type didn't change DCHECK(out->is_array_span()); return Status::OK(); @@ -903,8 +993,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()); @@ -917,7 +1012,7 @@ class ScalarExecutor : public KernelExecutorImpl { out_arr->null_count = 0; } - RETURN_NOT_OK(kernel_->exec(kernel_ctx_, input, &output)); + RETURN_NOT_OK(InvokeKernel(input, selection_ptr, &output)); // Output type didn't change DCHECK(output.is_array_data()); @@ -983,6 +1078,15 @@ class ScalarExecutor : public KernelExecutorImpl { return Status::OK(); } + 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, *selection, 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 @@ -1345,18 +1449,53 @@ 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); } SelectionVector::SelectionVector(const Array& arr) : SelectionVector(arr.data()) {} -int32_t SelectionVector::length() const { return static_cast(data_->length); } +int64_t SelectionVector::length() const { return data_->length; } + +Status SelectionVector::Validate(int64_t values_length) 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 (values_length >= 0) { + for (int64_t i = 0; i < length(); ++i) { + if (indices_[i] >= values_length) { + return Status::Invalid("SelectionVector index ", indices_[i], + " >= values length ", values_length); + } + } + } + return Status::OK(); +} -Result> SelectionVector::FromMask( - const BooleanArray& arr) { - return Status::NotImplemented("FromMask"); +void SelectionVectorSpan::SetSlice(int64_t offset, int64_t length, + int32_t index_back_shift) { + DCHECK_NE(indices_, nullptr); + offset_ = offset; + length_ = length; + 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 dae7e1ea6868..f254eb7d5bb5 100644 --- a/cpp/src/arrow/compute/exec.h +++ b/cpp/src/arrow/compute/exec.h @@ -140,17 +140,41 @@ class ARROW_EXPORT SelectionVector { explicit SelectionVector(const Array& arr); - /// \brief Create SelectionVector from boolean mask - static Result> FromMask(const BooleanArray& arr); - + std::shared_ptr data() const { return data_; } const int32_t* indices() const { return indices_; } - int32_t length() const; + int64_t length() const; + + Status Validate(int64_t values_length = -1) 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, 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 index_back_shift = 0); + + int32_t operator[](int64_t i) const { + return indices_[i + offset_] - index_back_shift_; + } + + int64_t length() const { return length_; } + + private: + const int32_t* indices_; + int64_t length_; + int64_t offset_; + int32_t index_back_shift_; +}; + /// An index to represent that a batch does not belong to an ordered stream constexpr int64_t kUnsequencedIndex = -1; @@ -173,8 +197,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 +223,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 +238,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 diff --git a/cpp/src/arrow/compute/exec_internal.h b/cpp/src/arrow/compute/exec_internal.h index 7e4f364a9288..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 { @@ -66,12 +67,15 @@ 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 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; } private: ExecSpanIterator(const std::vector& args, int64_t length, int64_t max_chunksize); @@ -83,6 +87,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_; @@ -93,6 +98,8 @@ 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_; }; @@ -165,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/exec_test.cc b/cpp/src/arrow/compute/exec_test.cc index 8314ad1d5c3f..c87d44c965d8 100644 --- a/cpp/src/arrow/compute/exec_test.cc +++ b/cpp/src/arrow/compute/exec_test.cc @@ -30,11 +30,13 @@ #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" #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" @@ -146,11 +148,56 @@ 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) { + 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) { @@ -732,13 +779,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()) { @@ -764,12 +821,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: @@ -881,9 +948,158 @@ 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, 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}); + 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, 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}); + CheckIteration(batch, /*chunksize=*/20, {15, 5, 10}, {3, 0, 1}); + CheckIteration(batch, /*chunksize=*/30, {15, 5, 10}, {3, 0, 1}); +} + // ---------------------------------------------------------------------- // Scalar function execution +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_selected(i); + ++selected; + } else { + on_non_selected(i); + } + } +} + +constexpr uint8_t kNonSelectedByte = 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; + + VisitIndicesWithSelection( + src.length(), selection, + [&](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 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], kNonSelectedByte); + } + }); +} + +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; + + VisitIndicesWithSelection( + src.length(), selection, + [&](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) { DCHECK_EQ(1, batch.num_values()); int value_size = batch[0].type()->byte_width(); @@ -896,6 +1112,32 @@ Status ExecCopyArrayData(KernelContext*, const ExecSpan& batch, ExecResult* out) return Status::OK(); } +Status SelectiveExecCopyArrayData(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; + 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; + VisitIndicesWithSelection( + batch.length, selection, + [&](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, kNonSelectedByte, value_size); + }); + return Status::OK(); +} + Status ExecCopyArraySpan(KernelContext*, const ExecSpan& batch, ExecResult* out) { DCHECK_EQ(1, batch.num_values()); int value_size = batch[0].type()->byte_width(); @@ -907,6 +1149,31 @@ Status ExecCopyArraySpan(KernelContext*, const ExecSpan& batch, ExecResult* out) return Status::OK(); } +Status SelectiveExecCopyArraySpan(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; + const uint8_t* src = arg0.buffers[1].data + arg0.offset * value_size; + VisitIndicesWithSelection( + batch.length, selection, + [&](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, kNonSelectedByte, value_size); + }); + return Status::OK(); +} + 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 @@ -923,6 +1190,24 @@ 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) { + // 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, ExecResult* out) { // Validity preallocated, but not the data @@ -934,6 +1219,18 @@ Status ExecNoPreallocatedData(KernelContext* ctx, const ExecSpan& batch, return ExecCopyArrayData(ctx, batch, out); } +Status SelectiveExecNoPreallocatedData(KernelContext* ctx, const ExecSpan& batch, + const SelectionVectorSpan& selection, + ExecResult* 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, ExecResult* out) { // Neither validity nor data preallocated @@ -949,6 +1246,22 @@ Status ExecNoPreallocatedAnything(KernelContext* ctx, const ExecSpan& batch, return ExecNoPreallocatedData(ctx, batch, out); } +Status SelectiveExecNoPreallocatedAnything(KernelContext* ctx, const ExecSpan& batch, + const SelectionVectorSpan& selection, + ExecResult* 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 { public: explicit ExampleOptions(std::shared_ptr value); @@ -1003,6 +1316,33 @@ Status ExecStateful(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) return Status::OK(); } +Status SelectiveExecStateful(KernelContext* ctx, const ExecSpan& batch, + const SelectionVectorSpan& selection, ExecResult* 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); + VisitIndicesWithSelection( + batch.length, selection, + [&](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); + memset(dst + i, kNonSelectedByte, sizeof(int32_t)); + }); + return Status::OK(); +} + 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); @@ -1013,6 +1353,29 @@ Status ExecAddInt32(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) return Status::OK(); } +Status SelectiveExecAddInt32(KernelContext* ctx, const ExecSpan& batch, + const SelectionVectorSpan& selection, ExecResult* 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); + VisitIndicesWithSelection( + batch.length, selection, + [&](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); + memset(out_data + i, kNonSelectedByte, sizeof(int32_t)); + }); + return Status::OK(); +} + class TestCallScalarFunction : public TestComputeInternals { protected: static bool initialized_; @@ -1023,9 +1386,13 @@ class TestCallScalarFunction : public TestComputeInternals { if (!initialized_) { initialized_ = true; AddCopyFunctions(); + AddSelectiveCopyFunctions(); AddNoPreallocateFunctions(); + AddSelectiveNoPreallocateFunctions(); AddStatefulFunction(); + AddSelectiveStatefulFunction(); AddScalarFunction(); + AddSelectiveScalarFunction(); } } @@ -1052,6 +1419,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(); @@ -1074,6 +1469,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(); @@ -1087,6 +1508,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(); @@ -1095,6 +1531,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; @@ -1103,11 +1550,23 @@ class FunctionCaller { public: virtual ~FunctionCaller() = default; + virtual std::string name() const = 0; + + virtual Result Call(const std::vector& args, + std::shared_ptr selection, + const FunctionOptions* options = NULLPTR, + ExecContext* ctx = NULLPTR) const = 0; + virtual Result Call(const std::vector& args, const FunctionOptions* options, - ExecContext* ctx = NULLPTR) = 0; + ExecContext* ctx = NULLPTR) const { + return Call(args, nullptr, options, ctx); + } + virtual Result Call(const std::vector& args, - ExecContext* ctx = NULLPTR) = 0; + ExecContext* ctx = NULLPTR) const { + return Call(args, /*options=*/nullptr, ctx); + } }; using FunctionCallerMaker = std::function>( @@ -1117,6 +1576,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); } @@ -1126,13 +1587,13 @@ class SimpleFunctionCaller : public FunctionCaller { return Make(func_name); } - Result Call(const std::vector& args, const FunctionOptions* options, - ExecContext* ctx) 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); } - Result Call(const std::vector& args, ExecContext* ctx) override { - return CallFunction(func_name, args, ctx); - } std::string func_name; }; @@ -1142,6 +1603,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, @@ -1165,235 +1628,670 @@ class ExecFunctionCaller : public FunctionCaller { return Make(func_name, std::move(in_types)); } - Result Call(const std::vector& args, const FunctionOptions* options, - ExecContext* ctx) 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); } - Result Call(const std::vector& args, ExecContext* ctx) override { - return Call(args, nullptr, ctx); - } std::shared_ptr func_exec; }; -class TestCallScalarFunctionArgumentValidation : public TestCallScalarFunction { - protected: - void DoTest(FunctionCallerMaker caller_maker); -}; +// Call the function via expression with an optional selection vector. +class ExpressionFunctionCaller : public FunctionCaller { + public: + 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)); + } -void TestCallScalarFunctionArgumentValidation::DoTest(FunctionCallerMaker caller_maker) { - ASSERT_OK_AND_ASSIGN(auto test_copy, caller_maker("test_copy", {int32()})); + std::string name() const override { return "expression_caller"; } - // Copy accepts only a single array argument - Datum d1(GetInt32Array(10)); + static Result> Make(std::string func_name, + std::vector in_types) { + return std::make_shared(std::move(func_name), + std::move(in_types)); + } - // Too many args - std::vector args = {d1, d1}; - ASSERT_RAISES(Invalid, test_copy->Call(args)); + 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)); + 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); + } - // Too few - args = {}; - ASSERT_RAISES(Invalid, test_copy->Call(args)); + Result Call(const std::vector& args, const FunctionOptions* options, + ExecContext* ctx) const override { + return Call(args, /*selection=*/nullptr, options, ctx); + } - // 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})); -} + static Result> Maker(const std::string& func_name, + std::vector in_types) { + return Make(func_name, std::move(in_types)); + } -TEST_F(TestCallScalarFunctionArgumentValidation, SimpleCall) { - TestCallScalarFunctionArgumentValidation::DoTest(SimpleFunctionCaller::Maker); -} + private: + std::string func_name_; + std::shared_ptr schema_; +}; -TEST_F(TestCallScalarFunctionArgumentValidation, ExecCall) { - TestCallScalarFunctionArgumentValidation::DoTest(ExecFunctionCaller::Maker); -} +class TestCallScalarFunctionArgumentValidation : public TestCallScalarFunction {}; -class TestCallScalarFunctionPreallocationCases : public TestCallScalarFunction { - protected: - void DoTest(FunctionCallerMaker caller_maker); -}; +TEST_F(TestCallScalarFunctionArgumentValidation, Basic) { + for (const auto& caller_maker : {SimpleFunctionCaller::Maker, ExecFunctionCaller::Maker, + ExpressionFunctionCaller::Maker}) { + ASSERT_OK_AND_ASSIGN(auto test_copy, caller_maker("test_copy", {int32()})); + ARROW_SCOPED_TRACE(test_copy->name()); + ResetContexts(); -void TestCallScalarFunctionPreallocationCases::DoTest(FunctionCallerMaker caller_maker) { - double null_prob = 0.2; + // Copy accepts only a single array argument + Datum d1(GetInt32Array(10)); - auto arr = GetUInt8Array(100, null_prob); + // Too many args + std::vector args = {d1, d1}; + ASSERT_RAISES(Invalid, test_copy->Call(args)); - auto CheckFunction = [&](std::shared_ptr test_copy) { - ResetContexts(); + // 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})); + } +} + +class TestCallScalarFunctionPreallocationCases : public TestCallScalarFunction { + protected: + std::shared_ptr GetTestArray() { return GetUInt8Array(100, 0.2); } + + std::vector> GetTestSelectionVectors() { + 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 + 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(arr)}; - ASSERT_OK_AND_ASSIGN(Datum result, test_copy->Call(args)); - ASSERT_EQ(Datum::ARRAY, result.kind()); - AssertArraysEqual(*arr, *result.make_array()); + 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(arr)}; + std::vector args = {Datum(input)}; exec_ctx_->set_exec_chunksize(80); - ASSERT_OK_AND_ASSIGN(Datum result, test_copy->Call(args, exec_ctx_.get())); - AssertArraysEqual(*arr, *result.make_array()); + 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(arr)}; + std::vector args = {Datum(input)}; exec_ctx_->set_exec_chunksize(11); - ASSERT_OK_AND_ASSIGN(Datum result, test_copy->Call(args, exec_ctx_.get())); - AssertArraysEqual(*arr, *result.make_array()); + 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 - { - 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); - } + 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(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)); + 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) { + 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}) { + ASSERT_OK_AND_ASSIGN(auto test_copy, caller_maker(name, {uint8()})); + ARROW_SCOPED_TRACE(test_copy->name()); + ResetContexts(); + + DoTestBasic(test_copy.get(), *arr, /*selection=*/nullptr, [&](const Datum& result) { + ASSERT_EQ(Datum::ARRAY, result.kind()); + AssertArraysEqual(*arr, *result.make_array()); + }); } - }; + } +} - 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, BasicSelectiveSparse) { + auto arr = GetTestArray(); + 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, + 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()); + AssertArraysEqualSparseWithSelection(*arr, selection_span, *result.make_array()); + }); + } + } } -TEST_F(TestCallScalarFunctionPreallocationCases, SimpleCaller) { - TestCallScalarFunctionPreallocationCases::DoTest(SimpleFunctionCaller::Maker); +TEST_F(TestCallScalarFunctionPreallocationCases, BasicSelectiveDense) { + auto arr = GetTestArray(); + auto selections = GetTestSelectionVectors(); + 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, ExecCaller) { - TestCallScalarFunctionPreallocationCases::DoTest(ExecFunctionCaller::Maker); +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, + ExpressionFunctionCaller::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 = GetTestSelectionVectors(); + 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()); + 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 = GetTestSelectionVectors(); + 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(); + 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}) { + ASSERT_OK_AND_ASSIGN(auto test_copy, caller_maker(name, {uint8()})); + ARROW_SCOPED_TRACE(test_copy->name()); + ResetContexts(); + + 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, IndependentPreallocateSelectiveSparse) { + auto arr = GetTestArray(); + 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, + ExpressionFunctionCaller::Maker(name, {uint8()})); + for (const auto& selection : selections) { + const int64_t exec_chunksize = 40; + ResetContexts(); + + 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 = GetTestSelectionVectors(); + 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) { + const int64_t exec_chunksize = 40; + ResetContexts(); + + DoTestIndependentPreallocate(test_copy.get(), exec_chunksize, *arr, selection, + [&](const Datum& result) { + AssertChunkedExecResultsEqualDenseWithSelection( + exec_chunksize, *arr, selection.get(), result); + }); + } + } +} + +// 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: - void DoTest(FunctionCallerMaker caller_maker); -}; + std::shared_ptr GetTestArray() { return GetUInt8Array(1000, 0.2); } -void TestCallScalarFunctionBasicNonStandardCases::DoTest( - 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 + std::vector> GetTestSelectionVectors() { + return {SelectionVectorFromJSON("[]"), SelectionVectorFromJSON("[0]"), + SelectionVectorFromJSON("[999]"), MakeSelectionVectorTo(400), + MakeSelectionVectorTo(401), MakeSelectionVectorTo(1000)}; + } - double null_prob = 0.2; + 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); + } - auto arr = GetUInt8Array(1000, null_prob); - std::vector args = {Datum(arr)}; + 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); + } +}; - auto CheckFunction = [&](std::shared_ptr test_nopre) { - ResetContexts(); +TEST_F(TestCallScalarFunctionBasicNonStandardCases, Basic) { + auto arr = GetTestArray(); + 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}) { + ASSERT_OK_AND_ASSIGN(auto test_nopre, caller_maker(name, {uint8()})); + ARROW_SCOPED_TRACE(test_nopre->name()); + ResetContexts(); + + DoTestBasic(test_nopre.get(), *arr, /*selection=*/nullptr, + [&](const Datum& result) { + ASSERT_EQ(Datum::ARRAY, result.kind()); + AssertArraysEqual(*arr, *result.make_array(), /*verbose=*/true); + }); + } + } +} - // The default should be a single array output - { - ASSERT_OK_AND_ASSIGN(Datum result, test_nopre->Call(args)); - AssertArraysEqual(*arr, *result.make_array(), 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()); + }); } + } +} - // 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(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()); + }); } - }; + } +} - 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, SplitExecution) { + auto arr = GetTestArray(); + 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}) { + ASSERT_OK_AND_ASSIGN(auto test_nopre, caller_maker(name, {uint8()})); + ARROW_SCOPED_TRACE(test_nopre->name()); + ResetContexts(); + + 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, SimpleCall) { - TestCallScalarFunctionBasicNonStandardCases::DoTest(SimpleFunctionCaller::Maker); +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, ExecCall) { - TestCallScalarFunctionBasicNonStandardCases::DoTest(ExecFunctionCaller::Maker); +TEST_F(TestCallScalarFunctionBasicNonStandardCases, SplitExecutionSelectiveDense) { + 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) { + const int64_t exec_chunksize = 400; + ResetContexts(); + + DoTestSplitExecution(test_nopre.get(), exec_chunksize, *arr, selection, + [&](const Datum& result) { + AssertChunkedExecResultsEqualDenseWithSelection( + exec_chunksize, *arr, selection.get(), result); + }); + } + } } class TestCallScalarFunctionStatefulKernel : public TestCallScalarFunction { protected: - void DoTest(FunctionCallerMaker caller_maker); -}; + 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]"); + } -void TestCallScalarFunctionStatefulKernel::DoTest(FunctionCallerMaker caller_maker) { - ASSERT_OK_AND_ASSIGN(auto test_stateful, caller_maker("test_stateful", {int32()})); + std::vector> GetTestSelectionVectors() { + return {SelectionVectorFromJSON("[]"), SelectionVectorFromJSON("[0]"), + SelectionVectorFromJSON("[4]"), MakeSelectionVectorTo(2), + MakeSelectionVectorTo(5)}; + } - auto input = ArrayFromJSON(int32(), "[1, 2, 3, null, 5]"); - auto multiplier = std::make_shared(2); - auto expected = ArrayFromJSON(int32(), "[2, 4, 6, null, 10]"); + 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); + } +}; - 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, Basic) { + 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(); + + DoTestBasic( + test_stateful.get(), *input, multiplier, /*selection=*/nullptr, + [&](const Datum& result) { AssertArraysEqual(*expected, *result.make_array()); }); + } } -TEST_F(TestCallScalarFunctionStatefulKernel, Simplecall) { - TestCallScalarFunctionStatefulKernel::DoTest(SimpleFunctionCaller::Maker); +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, ExecCall) { - TestCallScalarFunctionStatefulKernel::DoTest(ExecFunctionCaller::Maker); +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 { protected: - void DoTest(FunctionCallerMaker caller_maker); -}; + std::vector GetTestArgs() { + return {Datum(std::make_shared(5)), + Datum(std::make_shared(7))}; + } -void TestCallScalarFunctionScalarFunction::DoTest(FunctionCallerMaker caller_maker) { - ASSERT_OK_AND_ASSIGN(auto test_scalar_add_int32, - caller_maker("test_scalar_add_int32", {int32(), int32()})); + static constexpr int32_t kExpectedResult = 12; + + std::vector> GetTestSelectionVectors() { + return {SelectionVectorFromJSON("[]"), SelectionVectorFromJSON("[0]")}; + } - 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()); + 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(12); - ASSERT_TRUE(expected->Equals(*result.scalar())); + auto expected = std::make_shared(kExpectedResult); + ASSERT_TRUE(expected->Equals(*result.scalar())); + } +}; + +TEST_F(TestCallScalarFunctionScalarFunction, Basic) { + 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(); + + DoTestBasic(test_scalar_add_int32.get(), args, /*selection=*/nullptr); + } } -TEST_F(TestCallScalarFunctionScalarFunction, SimpleCall) { - TestCallScalarFunctionScalarFunction::DoTest(SimpleFunctionCaller::Maker); +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, ExecCall) { - TestCallScalarFunctionScalarFunction::DoTest(ExecFunctionCaller::Maker); +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); + } } TEST(Ordering, IsSuborderOf) { diff --git a/cpp/src/arrow/compute/expression.cc b/cpp/src/arrow/compute/expression.cc index 3c2ec1004022..2fc25d7f279b 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) { @@ -110,6 +121,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,6 +138,10 @@ const DataType* Expression::type() const { return parameter->type.type; } + if (const Special* special = this->special()) { + return special->type.type; + } + return CallNotNull(*this)->type.type; } @@ -169,6 +190,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,25 +269,39 @@ bool Expression::Equals(const Expression& other) const { return ref->Equals(*other.field_ref()); } - auto call = CallNotNull(*this); - auto other_call = CallNotNull(other); + 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 (call->function_name != other_call->function_name || - call->kernel != other_call->kernel) { + 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); - for (size_t i = 0; i < call->arguments.size(); ++i) { - if (!call->arguments[i].Equals(other_call->arguments[i])) { + if (special->special_form->name() != other_special->special_form->name()) { return false; } + + return args_and_options_eq(special, other_special); } - if (call->options == other_call->options) return true; - if (call->options && other_call->options) { - return call->options->Equals(*other_call->options); + auto call = CallNotNull(*this); + auto other_call = CallNotNull(other); + + if (call->function_name != other_call->function_name || + call->kernel != other_call->kernel) { + return false; } - return false; + + return args_and_options_eq(call, other_call); } bool Expression::Identical(const Expression& l, const Expression& r) { @@ -276,6 +320,10 @@ size_t Expression::hash() const { return ref->hash(); } + if (auto special = this->special()) { + return special->hash; + } + return CallNotNull(*this)->hash; } @@ -300,6 +348,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 +410,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,7 +588,49 @@ inline std::vector GetTypesWithSmallestLiteralRepresentation( return types; } -// Produce a bound Expression from unbound Call and bound arguments. +template +Result BindImpl(Expression expr, const TypeOrSchema& in, + compute::ExecContext* exec_context) { + if (exec_context == nullptr) { + compute::ExecContext exec_context; + return BindImpl(std::move(expr), in, &exec_context); + } + + if (expr.literal()) return expr; + + if (const FieldRef* ref = expr.field_ref()) { + ARROW_ASSIGN_OR_RAISE(FieldPath path, ref->FindOne(in)); + + Expression::Parameter param = *expr.parameter(); + param.indices.resize(path.indices().size()); + std::copy(path.indices().begin(), path.indices().end(), param.indices.begin()); + ARROW_ASSIGN_OR_RAISE(auto field, path.Get(in)); + param.type = field->type(); + 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)); + } + 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)); + } + + auto call = *CallNotNull(expr); + for (auto& argument : call.arguments) { + ARROW_ASSIGN_OR_RAISE(argument, BindImpl(std::move(argument), in, exec_context)); + } + return BindNonRecursive(call, /*insert_implicit_casts=*/true, exec_context); +} + +} // namespace + Result BindNonRecursive(Expression::Call call, bool insert_implicit_casts, compute::ExecContext* exec_context) { DCHECK(std::all_of(call.arguments.begin(), call.arguments.end(), @@ -603,37 +697,6 @@ Result BindNonRecursive(Expression::Call call, bool insert_implicit_ return Expression(std::move(call)); } -template -Result BindImpl(Expression expr, const TypeOrSchema& in, - compute::ExecContext* exec_context) { - if (exec_context == nullptr) { - compute::ExecContext exec_context; - return BindImpl(std::move(expr), in, &exec_context); - } - - if (expr.literal()) return expr; - - if (const FieldRef* ref = expr.field_ref()) { - ARROW_ASSIGN_OR_RAISE(FieldPath path, ref->FindOne(in)); - - Expression::Parameter param = *expr.parameter(); - param.indices.resize(path.indices().size()); - std::copy(path.indices().begin(), path.indices().end(), param.indices.begin()); - ARROW_ASSIGN_OR_RAISE(auto field, path.Get(in)); - param.type = field->type(); - return Expression{std::move(param)}; - } - - 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); -} - -} // namespace - Result Expression::Bind(const TypeHolder& in, compute::ExecContext* exec_context) const { return BindImpl(*this, *in.type, exec_context); @@ -758,6 +821,10 @@ Result ExecuteScalarExpression(const Expression& expr, const ExecBatch& i return field; } + if (auto special = expr.special()) { + return special->special_executor->Execute(input, exec_context); + } + auto call = CallNotNull(expr); std::vector arguments(call->arguments.size()); @@ -770,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(); @@ -789,7 +863,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, 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())); @@ -817,12 +892,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 +914,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..9d2c3eb81e1a 100644 --- a/cpp/src/arrow/compute/expression.h +++ b/cpp/src/arrow/compute/expression.h @@ -60,6 +60,20 @@ class ARROW_EXPORT Expression { void ComputeHash(); }; + struct Special { + std::shared_ptr special_form; + std::vector arguments; + std::shared_ptr options; + // Cached hash value + size_t hash; + + // post-Bind properties: + std::shared_ptr special_executor; + TypeHolder type; + + void ComputeHash(); + }; + std::string ToString() const; bool Equals(const Expression& other) const; size_t hash() const; @@ -112,6 +126,8 @@ 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; @@ -131,11 +147,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_; }; diff --git a/cpp/src/arrow/compute/expression_internal.h b/cpp/src/arrow/compute/expression_internal.h index 50b08a3d7a0f..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 @@ -44,6 +46,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) { @@ -290,5 +298,9 @@ inline Result> GetFunction( return GetCastFunction(*to_type); } +// 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/expression_test.cc b/cpp/src/arrow/compute/expression_test.cc index bbab57feebb7..c680708c02c9 100644 --- a/cpp/src/arrow/compute/expression_test.cc +++ b/cpp/src/arrow/compute/expression_test.cc @@ -28,11 +28,12 @@ #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/testing/gtest_util.h" +#include "arrow/compute/special_form.h" #include "arrow/testing/matchers.h" +#include "arrow/util/logging_internal.h" using testing::Eq; using testing::HasSubstr; @@ -47,50 +48,47 @@ using internal::checked_pointer_cast; namespace compute { -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))); -} +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; -Expression true_unless_null(Expression argument) { - return call("true_unless_null", {std::move(argument)}); -} +class EchoSpecialExecutor : public SpecialExecutor { + public: + EchoSpecialExecutor(Expression argument) + : SpecialExecutor(argument.type()), argument_(std::move(argument)) {} -Expression add(Expression l, Expression r) { - return call("add", {std::move(l), std::move(r)}); -} + Result Execute(const ExecBatch& input, + ExecContext* exec_context) const override { + return ExecuteScalarExpression(argument_, input, exec_context); + } + + private: + Expression argument_; +}; -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); +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]); } - result += "]"; - return result; -} +}; -const auto no_change = std::nullopt; +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)); +} TEST(ExpressionUtils, Comparison) { auto cmp_name = [](Datum l, Datum r) { @@ -323,6 +321,8 @@ 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)"); } TEST(Expression, Equality) { @@ -375,6 +375,19 @@ 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(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)))); } Expression null_literal(const std::shared_ptr& type) { @@ -402,7 +415,11 @@ 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(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_EQ(set.size(), 10); } TEST(Expression, IsScalarExpression) { @@ -422,6 +439,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()); } TEST(Expression, IsSatisfiable) { @@ -491,6 +510,8 @@ TEST(Expression, IsSatisfiable) { // fill_na) EXPECT_TRUE(Bind(call("is_null", {never_true})).IsSatisfiable()); } + + EXPECT_TRUE(Bind(echo_special(field_ref("i32"))).IsSatisfiable()); } TEST(Expression, FieldsInExpression) { @@ -518,6 +539,9 @@ TEST(Expression, FieldsInExpression) { equal(field_ref("b"), literal(2))), not_(less(field_ref("c"), literal(3)))), {"a", "b", "c"}); + + ExpectFieldsAre(echo_special(literal(42)), {}); + ExpectFieldsAre(echo_special(field_ref("a")), {"a"}); } TEST(Expression, ExpressionHasFieldRefs) { @@ -540,6 +564,9 @@ 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(echo_special(literal(42)))); + EXPECT_TRUE(ExpressionHasFieldRefs(echo_special(field_ref("a")))); } TEST(Expression, BindLiteral) { @@ -555,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"); @@ -948,6 +957,40 @@ TEST(Expression, BindNestedCall) { EXPECT_TRUE(expr.IsBound()); } +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())); + } +} + TEST(Expression, ExecuteFieldRef) { auto ExpectRefIs = [](FieldRef ref, Datum in, Datum expected) { auto expr = field_ref(ref); @@ -1042,15 +1085,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; + }; + + compute::ExecContext exec_context; - 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)); + 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); @@ -1215,6 +1275,58 @@ TEST(Expression, ExecuteDictionaryTransparent) { ])")); } +TEST(Expression, NaiveExecuteSpecialForm) { + ExpectExecute(echo_special(field_ref("i32")), + ArrayFromJSON(struct_({field("i32", int32())}), R"([ + {"i32": 0}, + {"i32": 1}, + {"i32": 2} + ])")); +} + +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/expression_test_internal.h b/cpp/src/arrow/compute/expression_test_internal.h new file mode 100644 index 000000000000..ccb3a5768933 --- /dev/null +++ b/cpp/src/arrow/compute/expression_test_internal.h @@ -0,0 +1,94 @@ + +// 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 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) { + 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/function.cc b/cpp/src/arrow/compute/function.cc index b0b12a690f86..b9c557aed274 100644 --- a/cpp/src/arrow/compute/function.cc +++ b/cpp/src/arrow/compute/function.cc @@ -412,6 +412,14 @@ Status Function::Validate() const { Status ScalarFunction::AddKernel(std::vector in_types, OutputType out_type, 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), std::move(constraint)); +} + +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())); if (arity_.is_varargs && in_types.size() != 1) { @@ -419,7 +427,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..d86c8ab3b5ea 100644 --- a/cpp/src/arrow/compute/function.h +++ b/cpp/src/arrow/compute/function.h @@ -311,6 +311,11 @@ 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, + 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. Status AddKernel(ScalarKernel kernel); diff --git a/cpp/src/arrow/compute/kernel.h b/cpp/src/arrow/compute/kernel.h index 0d4f9d6ff436..a84e93c50602 100644 --- a/cpp/src/arrow/compute/kernel.h +++ b/cpp/src/arrow/compute/kernel.h @@ -555,19 +555,35 @@ struct ARROW_EXPORT Kernel { /// employed this may not be possible. using ArrayKernelExec = 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 /// and memory pre-allocation preferences. struct ARROW_EXPORT ScalarKernel : public Kernel { ScalarKernel() = default; + ScalarKernel(std::shared_ptr sig, ArrayKernelExec 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, + ArrayKernelSelectiveExec selective_exec, KernelInit init = NULLPTR) + : Kernel(std::move(in_types), std::move(out_type), std::move(init)), + exec(std::move(exec)), + selective_exec(std::move(selective_exec)) {} + ScalarKernel(std::shared_ptr sig, ArrayKernelExec exec, KernelInit init = NULLPTR) - : Kernel(std::move(sig), init), exec(exec) {} + : ScalarKernel(std::move(sig), std::move(exec), NULLPTR, std::move(init)) {} 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) {} + : 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 @@ -575,6 +591,8 @@ struct ARROW_EXPORT ScalarKernel : public Kernel { /// through the KernelContext. ArrayKernelExec 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*, /// including sliced output validity bitmaps. Some kernel implementations may diff --git a/cpp/src/arrow/compute/special/CMakeLists.txt b/cpp/src/arrow/compute/special/CMakeLists.txt new file mode 100644 index 000000000000..c98b981751ec --- /dev/null +++ b/cpp/src/arrow/compute/special/CMakeLists.txt @@ -0,0 +1,32 @@ +# 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_test(special_internal_test + SOURCES + conditional_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/conditional_special.cc b/cpp/src/arrow/compute/special/conditional_special.cc new file mode 100644 index 000000000000..9fc63a09fbbc --- /dev/null +++ b/cpp/src/arrow/compute/special/conditional_special.cc @@ -0,0 +1,379 @@ +// 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_special_internal.h" + +#include "arrow/array/array_primitive.h" +#include "arrow/array/builder_primitive.h" +#include "arrow/visit_data_inline.h" + +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(); + 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()) { + 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()); + 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( + 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::MakeBodyMaskFromBitmap( + const std::shared_ptr& bitmap, ExecContext* exec_context) const { + DCHECK_EQ(bitmap->length(), length_); + + Int32Builder body_builder(exec_context->memory_pool()); + Int32Builder remainder_builder(exec_context->memory_pool()); + RETURN_NOT_OK(body_builder.Reserve(length_)); + RETURN_NOT_OK(remainder_builder.Reserve(length_)); + + ArraySpan span(*bitmap->data()); + int32_t i = 0; + VisitArraySpanInline( + span, + [&](bool mask) { + if (mask) { + body_builder.UnsafeAppend(i); + } else { + remainder_builder.UnsafeAppend(i); + } + ++i; + }, + [&]() { ++i; }); + + 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); + auto remainder = std::make_shared(*remainder_arr); + return std::make_shared(std::move(body), std::move(remainder), + 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 remainder_builder(exec_context->memory_pool()); + RETURN_NOT_OK(body_builder.Reserve(length_)); + RETURN_NOT_OK(remainder_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 { + remainder_builder.UnsafeAppend(i); + } + ++i; + }, + [&]() { ++i; }); + } + + 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); + auto remainder = std::make_shared(*remainder_arr); + return std::make_shared(std::move(body), std::move(remainder), + 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 remainder_builder(exec_context->memory_pool()); + RETURN_NOT_OK(body_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]; + if (!bitmap->IsNull(index)) { + if (bitmap->Value(index)) { + body_builder.UnsafeAppend(index); + } else { + remainder_builder.UnsafeAppend(index); + } + } + } + + 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); + auto remainder = std::make_shared(*remainder_arr); + return std::make_shared(std::move(body), std::move(remainder), + 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 remainder_builder(exec_context->memory_pool()); + RETURN_NOT_OK(body_builder.Reserve(length_)); + RETURN_NOT_OK(remainder_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 { + remainder_builder.UnsafeAppend(index); + } + } + } + + 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()); + return std::make_shared(std::move(body), std::move(remainder), + length_); +} + +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()) { + // 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; + } + ARROW_ASSIGN_OR_RAISE(auto body_result, + 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)); + ARROW_ASSIGN_OR_RAISE(branch_mask, body_mask->NextBranchMask()); + } + // Should have no remaining rows. + DCHECK(branch_mask->empty()); + 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 { + 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; + } + } + } + + 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); +} + +} // 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 new file mode 100644 index 000000000000..88aa75bab0da --- /dev/null +++ b/cpp/src/arrow/compute/special/conditional_special_internal.h @@ -0,0 +1,448 @@ +// 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" +#include "arrow/util/unreachable.h" + +namespace arrow::compute::internal { + +/// 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 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 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; + +/// @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 { + return NULLPTR; + } + + protected: + 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_; +}; + +/// @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 { + Unreachable("AllFailBranchMask::GetSelectionVector should not be called"); + } + + protected: + Result> MakeBodyMaskFromBitmap( + const std::shared_ptr& bitmap, + ExecContext* exec_context) const override { + Unreachable("AllFailBranchMask::MakeBodyMaskFromBitmap should not be called"); + } + + Result> MakeBodyMaskFromBitmap( + const std::shared_ptr& bitmap, + ExecContext* exec_context) const override { + 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) { +#ifndef NDEBUG + DCHECK_OK(selection_vector_->Validate(length_)); +#endif + } + + bool empty() const override { return selection_vector_->length() == 0; } + + Result> GetSelectionVector() const override { + return selection_vector_; + } + + protected: + 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; +}; + +/// @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; +}; + +/// @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; + + bool empty() const override { return true; } + + Result> GetSelectionVector() const override { + Unreachable("AllNullBodyMask::GetSelectionVector should not be called"); + } + + Result> NextBranchMask() const override { + return std::make_shared(); + } +}; + +/// @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)) {} + + protected: + 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; + + bool empty() const override { return branch_mask_->empty(); } + + Result> GetSelectionVector() const override { + return branch_mask_->GetSelectionVector(); + } + + Result> NextBranchMask() const override { + return std::make_shared(); + } +}; + +/// @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; + + bool empty() const override { return true; } + + Result> GetSelectionVector() const override { + Unreachable("AllFailBodyMask::GetSelectionVector should not be called"); + } + + Result> NextBranchMask() const override { + return branch_mask_; + } +}; + +/// @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) + : body_(std::move(body)), remainder_(std::move(remainder)), length_(length) { +#ifndef NDEBUG + DCHECK_OK(body_->Validate(length_)); + DCHECK_OK(remainder_->Validate(length_)); +#endif + } + + bool empty() const override { return body_->length() == 0; } + + Result> GetSelectionVector() const override { + return body_; + } + + Result> NextBranchMask() const override { + return BranchMask::FromSelectionVector(remainder_, length_); + } + + private: + std::shared_ptr body_; + std::shared_ptr remainder_; + int64_t length_; +}; + +struct ARROW_COMPUTE_EXPORT Branch { + Expression cond; + 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) {} + + 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); + 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_; + }; + + /// @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) { + return BranchMask::FromSelectionVector(input.selection_vector, input.length); + } + return std::make_shared(input.length); + } + + Result> EvaluateCond( + const std::shared_ptr& branch_mask, const Expression& cond, + const ExecBatch& input, ExecContext* exec_context) const { + 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 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); + } + + /// @brief Multiplex all branch body results into the final result based on their + /// 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; + + private: + 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_special_test.cc b/cpp/src/arrow/compute/special/conditional_special_test.cc new file mode 100644 index 000000000000..13dca618f298 --- /dev/null +++ b/cpp/src/arrow/compute/special/conditional_special_test.cc @@ -0,0 +1,893 @@ +// 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_special_internal.h" + +#include + +#include "arrow/compute/exec_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 { + +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()); + 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); +} + +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 kNullScalar = MakeNullScalar(boolean()); +const auto kTrueScalar = MakeScalar(true); +const auto kFalseScalar = MakeScalar(false); + +} // 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_EQ(branch_mask->empty(), length == 0); + } +} + +TEST(AllPassBranchMask, GetSelectionVector) { + for (auto length : {0, 42}) { + auto branch_mask = std::make_shared(length); + EXPECT_EQ(branch_mask->GetSelectionVector(), nullptr); + } +} + +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(boolean_array)); + + auto boolean_chunked_array = + std::make_shared(ArrayVector{boolean_array, boolean_array}); + CheckMakeBodyMask(std::make_shared(length * 2), + Datum(boolean_chunked_array)); + } + + { + ASSERT_OK_AND_ASSIGN(auto boolean_array, + gen::Constant(kTrueScalar)->Generate(length)); + CheckMakeBodyMaskAndSelection( + std::make_shared(length), Datum(boolean_array), nullptr); + + auto boolean_chunked_array = + std::make_shared(ArrayVector{boolean_array, boolean_array}); + CheckMakeBodyMaskAndSelection( + std::make_shared(length * 2), Datum(boolean_chunked_array), + nullptr); + } + + { + 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( + 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]")); + } +} + +TEST(AllFailBranchMask, Emptiness) { + auto branch_mask = std::make_shared(); + EXPECT_TRUE(branch_mask->empty()); +} + +TEST(ConditionalBranchMask, Emptiness) { + for (const auto& selection : + {SelectionVectorFromJSON("[]"), SelectionVectorFromJSON("[0]"), + 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, MakeBodyMaskTrivial) { + 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(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( + 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); + } + + { + ASSERT_OK_AND_ASSIGN(auto boolean_array, + gen::Constant(kFalseScalar)->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)); + } + } +} + +TEST(ConditionalBranchMask, MakeBodyMask) { + { + 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 { + +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); + EXPECT_EQ(executor.options(), nullptr); +} + +TEST(ConditionalSpecialExecutor, Execute) { + auto schm = schema({field("", boolean()), field("", boolean()), field("", utf8()), + field("", utf8()), field("", utf8())}); + + 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_sp = [&](Expression argument) -> Expression { + return unreachable(std::move(argument), *schm); + }; + auto assert_selection_empty_sp = [&](Expression argument) -> Expression { + return assert_selection_empty(std::move(argument), *schm); + }; + auto assert_selection_eq_sp = + [&](Expression argument, std::shared_ptr selection) -> Expression { + return assert_selection_equal(std::move(argument), std::move(selection), *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"]])"); + + { + ConditionalSpecialExecutor executor( + {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"])")), + result); + } + + { + ConditionalSpecialExecutor executor( + {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])")), + result); + } + + { + ConditionalSpecialExecutor executor( + {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"])")), + result); + } + + { + ConditionalSpecialExecutor executor( + {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"])")), + result); + } + + { + ConditionalSpecialExecutor executor( + {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"])")), + result); + } + + { + ConditionalSpecialExecutor executor( + {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])")), + result); + } + + { + ConditionalSpecialExecutor executor( + {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"])")), + result); + } + + { + ConditionalSpecialExecutor executor( + {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"])")), + result); + } + + { + ConditionalSpecialExecutor executor( + {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"])")), + result); + } + + { + ConditionalSpecialExecutor executor( + {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"])")), + result); + } + + { + ConditionalSpecialExecutor executor( + {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"])")), + result); + } + + { + ConditionalSpecialExecutor executor( + {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"])")), + result); + } + + { + ConditionalSpecialExecutor executor( + {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"])")), + result); + } + + { + ConditionalSpecialExecutor executor( + {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"])")), + result); + } +} + +TEST(ConditionalSpecialExecutor, ExecuteWithSelection) { + auto schm = schema({field("a", int32())}); + + auto a = field_ref("a"); + auto unreachable_sp = [&](Expression argument) -> Expression { + return unreachable(std::move(argument), *schm); + }; + auto assert_selection_empty_sp = [&](Expression argument) -> Expression { + return assert_selection_empty(std::move(argument), *schm); + }; + auto assert_selection_eq_sp = + [&](Expression argument, std::shared_ptr selection) -> Expression { + return 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 = selection; + ConditionalSpecialExecutor executor( + {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])")), + result); + } + + for (const auto& selection : + {SelectionVectorFromJSON("[0]"), SelectionVectorFromJSON("[1]"), + SelectionVectorFromJSON("[0, 1, 2]")}) { + auto batch_with_selection = batch; + batch_with_selection.selection_vector = selection; + ConditionalSpecialExecutor executor( + {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())); + } + + { + // 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 = selection; + ConditionalSpecialExecutor executor( + {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())); + } +} + +} // 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 new file mode 100644 index 000000000000..39d6612978cb --- /dev/null +++ b/cpp/src/arrow/compute/special/if_else_special.cc @@ -0,0 +1,64 @@ +// 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_special_internal.h" + +namespace arrow::compute { + +namespace { + +using internal::Branch; +using internal::ConditionalSpecialForm; + +class IfElseSpecialForm : public ConditionalSpecialForm { + public: + IfElseSpecialForm() : ConditionalSpecialForm("if_else") {} + + ARROW_DISALLOW_COPY_AND_ASSIGN(IfElseSpecialForm); + ARROW_DEFAULT_MOVE_AND_ASSIGN(IfElseSpecialForm); + + 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]); + + return std::vector{{std::move(cond), std::move(if_true)}, + {literal(true), std::move(if_false)}}; + } + + friend class ConditionalSpecialForm; +}; + +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/if_else_special_benchmark.cc b/cpp/src/arrow/compute/special/if_else_special_benchmark.cc new file mode 100644 index 000000000000..656db7c9b9c6 --- /dev/null +++ b/cpp/src/arrow/compute/special/if_else_special_benchmark.cc @@ -0,0 +1,284 @@ +// 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/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" +#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::compute { + +namespace { + +// 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 kSpinOptionsType = internal::GetFunctionOptionsType( + arrow::internal::DataMember("count", &SpinOptions::count)); + +SpinOptions::SpinOptions(int64_t count) + : FunctionOptions(kSpinOptionsType), count(count) {} + +const SpinOptions* GetDefaultSpinOptions() { + static const auto kDefaultSpinOptions = SpinOptions::Defaults(); + return &kDefaultSpinOptions; +} + +using SpinState = internal::OptionsWrapper; + +inline void Spin(volatile int64_t count) { + while (count-- > 0) { + // Do nothing, just burn CPU cycles. + } +} + +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 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 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 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 RegisterSpinFunction() { + auto registry = GetFunctionRegistry(); + + if (registry->CanAddFunctionOptionsType(kSpinOptionsType).ok()) { + RETURN_NOT_OK(registry->AddFunctionOptionsType(kSpinOptionsType)); + } + + 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(); + }; + + // 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)); + + return Status::OK(); +} + +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 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}; + + for (auto _ : state) { + ARROW_CHECK_OK(ExecuteScalarExpression(bound, batch).status()); + } + + state.SetItemsProcessed(state.iterations() * length); +} + +} // namespace + +// 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) + +// 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); + + 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::string kNumRowsArgName = "num_rows"; +const std::vector kNumRowsArg{1, 4 * 1024, 64 * 1024}; + +#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_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 + +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); + + BenchmarkIfElse(state, std::move(make_if_else_func), spin_func, + if_true_kernel_intensity, if_false_kernel_intensity, + std::move(cond_datum), num_rows); +} + +// 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); +} + +const std::string kNullProbabilityArgName = "null_probability"; +const std::vector kNullProbabilityArg{0, 25, 50, 100}; + +#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 + +// 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); +} + +const std::string kTrueProbabilityArgName = "true_probability"; +const std::vector kTrueProbabilityArg{0, 25, 50, 100}; + +#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 + +// 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::string kHeavinessArgName = "heaviness"; +const std::vector kHeavinessArg{0, 10, 100}; + +#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 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 new file mode 100644 index 000000000000..efd973c2d9d9 --- /dev/null +++ b/cpp/src/arrow/compute/special/if_else_special_test.cc @@ -0,0 +1,594 @@ +// 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 + +#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 { + +using internal::add; +using internal::cast; +using internal::ExpectBindsTo; +using internal::kBoringSchema; +using internal::no_change; +using internal::sub; + +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(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")), + 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(IfElseSpecial, 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(IfElseSpecial, IsScalarExpression) { + EXPECT_TRUE(if_else_special(field_ref("cond"), field_ref("a"), field_ref("b")) + .IsScalarExpression()); +} + +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(IfElseSpecial, 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(IfElseSpecial, 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(IfElseSpecial, 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())); + } +} + +namespace { + +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)}); +} + +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); +} + +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()); + + if (expected.kind() == result.kind()) { + AssertDatumsEqual(expected, result); + return; + } + + 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); + } + if (datum.is_array()) { + return datum.make_array(); + } + DCHECK(datum.is_chunked_array()); + return Concatenate(datum.chunked_array()->chunks(), default_memory_pool()); + }; + + ASSERT_OK_AND_ASSIGN(auto expected_array, to_array(expected)); + ASSERT_OK_AND_ASSIGN(auto result_array, to_array(result)); +} + +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()) { + 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 make_if_else(std::move(cond), std::move(if_true), std::move(if_false)); + }, + schema, batch, exec_context); +} + +} // namespace + +class TestExecuteIfElseSpecial : public ::testing::Test { + protected: + const int64_t length = 7; + + 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]"); + std::shared_ptr boolean1_chunked = ChunkedArrayFromJSON( + boolean(), {"[null, true, false]", "[]", "[true, false, null, true]"}); + 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 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))}; + 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)}; + 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 ExecBatch& 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 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, BasicAllLiterals) { + DoTestBasic(boolean_literals, int_literals, int_literals, ExecBatch({}, length)); +} + +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, BasicArrays) { + WithDatumCombinations(boolean_arrays, boolean_arrays, int_arrays, int_arrays, + [&](const ExecBatch& batch) { + DoTestBasic(boolean_fields, {int1}, {int2}, batch); + }); +} + +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, int_literals, + int_literals, ExecBatch({}, length)); +} + +TEST_F(TestExecuteIfElseSpecial, NestedCondAllScalars) { + 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, 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); + }); +} + +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, boolean_literals, int_literals, int_literals, + int_literals, ExecBatch({}, length)); +} + +TEST_F(TestExecuteIfElseSpecial, NestedBodyAllScalars) { + 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) { + 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) { + 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); + }); +} + +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) { + // 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 arrow::compute 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/special_form_internal.h b/cpp/src/arrow/compute/special/special_form_internal.h new file mode 100644 index 000000000000..aaf09453d4c0 --- /dev/null +++ b/cpp/src/arrow/compute/special/special_form_internal.h @@ -0,0 +1,62 @@ +// 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 { + +/// @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: + 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 new file mode 100644 index 000000000000..dd262b8d7a0c --- /dev/null +++ b/cpp/src/arrow/compute/special_form.h @@ -0,0 +1,84 @@ +// 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.h" + +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, + 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; + + protected: + const TypeHolder out_type_; + 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)) {} + + virtual ~SpecialForm() = default; + + const std::string& name() const { return name_; } + + virtual Result> Bind( + std::vector& arguments, std::shared_ptr options, + ExecContext* exec_context) const = 0; + + private: + std::string name_; +}; + +} // namespace arrow::compute diff --git a/cpp/src/arrow/compute/test_util_internal.cc b/cpp/src/arrow/compute/test_util_internal.cc index 219028943cb3..95cc4c51c693 100644 --- a/cpp/src/arrow/compute/test_util_internal.cc +++ b/cpp/src/arrow/compute/test_util_internal.cc @@ -20,10 +20,12 @@ #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" #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 +122,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 diff --git a/cpp/src/arrow/compute/type_fwd.h b/cpp/src/arrow/compute/type_fwd.h index 016d97a0dbc2..f88f12a7c235 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 SpecialForm; +class SpecialExecutor; + ARROW_EXPORT ExecContext* default_exec_context(); ARROW_EXPORT ExecContext* threaded_exec_context(); 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