Skip to content

Commit

Permalink
apacheGH-15290: [C++][Compute] Optimize IfElse kernel AAS/ASA case wh…
Browse files Browse the repository at this point in the history
…en the scalar is null
  • Loading branch information
js8544 committed Jan 10, 2023
1 parent 26a426f commit e965e1a
Showing 1 changed file with 38 additions and 3 deletions.
41 changes: 38 additions & 3 deletions cpp/src/arrow/compute/kernels/scalar_if_else.cc
Original file line number Diff line number Diff line change
Expand Up @@ -15,13 +15,16 @@
// specific language governing permissions and limitations
// under the License.

#include <cstring>
#include "arrow/array/builder_nested.h"
#include "arrow/array/builder_primitive.h"
#include "arrow/array/builder_time.h"
#include "arrow/array/builder_union.h"
#include "arrow/compute/api.h"
#include "arrow/compute/kernels/codegen_internal.h"
#include "arrow/compute/kernels/copy_data_internal.h"
#include "arrow/result.h"
#include "arrow/status.h"
#include "arrow/util/bit_block_counter.h"
#include "arrow/util/bit_run_reader.h"
#include "arrow/util/bitmap.h"
Expand Down Expand Up @@ -470,6 +473,10 @@ struct IfElseFunctor<Type,
// copy right data to out_buff
std::memcpy(out_values, right.GetValues<T>(1), right.length * sizeof(T));

if (!left.is_valid) { // left is null scalar, only need to copy right data to out
return Status::OK();
}

// selectively copy values from left data
T left_data = internal::UnboxScalar<Type>::Unbox(left);

Expand All @@ -490,6 +497,10 @@ struct IfElseFunctor<Type,
const T* left_data = left.GetValues<T>(1);
std::memcpy(out_values, left_data, left.length * sizeof(T));

if (!right.is_valid) { // right is null scalar, only need to copy left data to out
return Status::OK();
}

T right_data = internal::UnboxScalar<Type>::Unbox(right);

RunIfElseLoopInverted(cond, [&](int64_t data_offset, int64_t num_elems) {
Expand Down Expand Up @@ -723,12 +734,24 @@ struct IfElseFunctor<Type, enable_if_base_binary<Type>> {
// ASA
static Status Call(KernelContext* ctx, const ArraySpan& cond, const Scalar& left,
const ArraySpan& right, ExecResult* out) {
std::string_view left_data = internal::UnboxScalar<Type>::Unbox(left);
auto left_size = static_cast<OffsetType>(left_data.size());

const auto* right_offsets = right.GetValues<OffsetType>(1);
const uint8_t* right_data = right.buffers[2].data;

if (!left.is_valid) { // left is null scalar, only need to copy left data to out
auto& out_data = *out->array_data();
auto offset_length = (cond.length + 1) * sizeof(OffsetType);
ARROW_ASSIGN_OR_RAISE(out_data.buffers[1], ctx->Allocate(offset_length));
memcpy(out_data.buffers[1]->mutable_data(), right_offsets, offset_length);

auto right_data_length = right_offsets[right.length] - right_offsets[0];
ARROW_ASSIGN_OR_RAISE(out_data.buffers[2], ctx->Allocate(right_data_length));
memcpy(out_data.buffers[2]->mutable_data(), right_data, right_data_length);
return Status::OK();
}

std::string_view left_data = internal::UnboxScalar<Type>::Unbox(left);
auto left_size = static_cast<OffsetType>(left_data.size());

// allocate data buffer conservatively
int64_t data_buff_alloc =
left_size * cond.length + right_offsets[right.length] - right_offsets[0];
Expand All @@ -754,6 +777,18 @@ struct IfElseFunctor<Type, enable_if_base_binary<Type>> {
const auto* left_offsets = left.GetValues<OffsetType>(1);
const uint8_t* left_data = left.buffers[2].data;

if (!right.is_valid) { // right is null scalar, only need to copy left data to out
auto& out_data = *out->array_data();
auto offset_length = (cond.length + 1) * sizeof(OffsetType);
ARROW_ASSIGN_OR_RAISE(out_data.buffers[1], ctx->Allocate(offset_length));
memcpy(out_data.buffers[1]->mutable_data(), left_offsets, offset_length);

auto left_data_length = left_offsets[left.length] - left_offsets[0];
ARROW_ASSIGN_OR_RAISE(out_data.buffers[2], ctx->Allocate(left_data_length));
memcpy(out_data.buffers[2]->mutable_data(), left_data, left_data_length);
return Status::OK();
}

std::string_view right_data = internal::UnboxScalar<Type>::Unbox(right);
auto right_size = static_cast<OffsetType>(right_data.size());

Expand Down

0 comments on commit e965e1a

Please sign in to comment.