liyafan82 commented on a change in pull request #7748: URL: https://github.com/apache/arrow/pull/7748#discussion_r454957895
########## File path: cpp/src/arrow/compute/kernels/codegen_internal.h ########## @@ -675,11 +704,80 @@ struct ScalarBinary { } }; +// An alternative to ScalarBinary that Applies a scalar operation on only the +// not-null values of arrays. +template <typename OutType, typename Arg0Type, typename Arg1Type, typename Op> +struct ScalarBinaryNotNull { + using OUT = typename GetOutputType<OutType>::T; + + static void ArrayArray(KernelContext* ctx, const ArrayData& arg0, const ArrayData& arg1, + Datum* out) { + ArrayIterator<Arg0Type> arg0_it(arg0); + ArrayIterator<Arg1Type> arg1_it(arg1); + OutputAdapter<OutType>::WriteOnlyValid( + ctx, out, [&]() -> OUT { return Op::template Call(ctx, arg0_it(), arg1_it()); }, + [&]() { + ARROW_UNUSED(arg0_it()); + ARROW_UNUSED(arg1_it()); + }); + } + + static void ArrayScalar(KernelContext* ctx, const ArrayData& arg0, const Scalar& arg1, + Datum* out) { + ArrayIterator<Arg0Type> arg0_it(arg0); + auto arg1_val = UnboxScalar<Arg1Type>::Unbox(arg1); + OutputAdapter<OutType>::WriteOnlyValid( + ctx, out, [&]() -> OUT { return Op::template Call(ctx, arg0_it(), arg1_val); }, + [&]() { ARROW_UNUSED(arg0_it()); }); + } + + static void ScalarArray(KernelContext* ctx, const Scalar& arg0, const ArrayData& arg1, + Datum* out) { + auto arg0_val = UnboxScalar<Arg0Type>::Unbox(arg0); + ArrayIterator<Arg1Type> arg1_it(arg1); + OutputAdapter<OutType>::WriteOnlyValid( + ctx, out, [&]() -> OUT { return Op::template Call(ctx, arg0_val, arg1_it()); }, + [&]() { ARROW_UNUSED(arg1_it()); }); + } + + static void ScalarScalar(KernelContext* ctx, const Scalar& arg0, const Scalar& arg1, + Datum* out) { + if (out->scalar()->is_valid) { + auto arg0_val = UnboxScalar<Arg0Type>::Unbox(arg0); + auto arg1_val = UnboxScalar<Arg1Type>::Unbox(arg1); + BoxScalar<OutType>::Box(Op::template Call(ctx, arg0_val, arg1_val), + out->scalar().get()); + } + } + + static void Exec(KernelContext* ctx, const ExecBatch& batch, Datum* out) { + if (batch[0].kind() == Datum::ARRAY) { + if (batch[1].kind() == Datum::ARRAY) { + return ArrayArray(ctx, *batch[0].array(), *batch[1].array(), out); + } else { + return ArrayScalar(ctx, *batch[0].array(), *batch[1].scalar(), out); + } + } else { + if (batch[1].kind() == Datum::ARRAY) { + // e.g. if we were doing scalar < array, we flip and do array >= scalar Review comment: Nice catch. Both occurences of this comment have been removed. ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: us...@infra.apache.org