Skip to content

Commit e1690d6

Browse files
lidavidmbkietz
lidavidm
authored and
bkietz
committed
ARROW-12751: [C++] Implement minimum/maximum kernels
This is a bit messy, but implements a variadic scalar maximum/minimum kernel. Closes #10390 from lidavidm/arrow-12751 Authored-by: David Li <li.davidm96@gmail.com> Signed-off-by: Benjamin Kietzman <bengilgit@gmail.com>
1 parent 5baebba commit e1690d6

8 files changed

Lines changed: 770 additions & 30 deletions

File tree

‎cpp/src/arrow/compute/api_scalar.cc‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,16 @@ SCALAR_ARITHMETIC_BINARY(Multiply, "multiply", "multiply_checked")
6363
SCALAR_ARITHMETIC_BINARY(Divide, "divide", "divide_checked")
6464
SCALAR_ARITHMETIC_BINARY(Power, "power", "power_checked")
6565

66+
Result<Datum> ElementWiseMax(const std::vector<Datum>& args,
67+
ElementWiseAggregateOptions options, ExecContext* ctx) {
68+
return CallFunction("element_wise_max", args, &options, ctx);
69+
}
70+
71+
Result<Datum> ElementWiseMin(const std::vector<Datum>& args,
72+
ElementWiseAggregateOptions options, ExecContext* ctx) {
73+
return CallFunction("element_wise_min", args, &options, ctx);
74+
}
75+
6676
// ----------------------------------------------------------------------
6777
// Set-related operations
6878

‎cpp/src/arrow/compute/api_scalar.h‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,11 @@ struct ArithmeticOptions : public FunctionOptions {
4242
bool check_overflow;
4343
};
4444

45+
struct ARROW_EXPORT ElementWiseAggregateOptions : public FunctionOptions {
46+
ElementWiseAggregateOptions() : skip_nulls(true) {}
47+
bool skip_nulls;
48+
};
49+
4550
struct ARROW_EXPORT MatchSubstringOptions : public FunctionOptions {
4651
explicit MatchSubstringOptions(std::string pattern, bool ignore_case = false)
4752
: pattern(std::move(pattern)), ignore_case(ignore_case) {}
@@ -253,6 +258,30 @@ Result<Datum> Power(const Datum& left, const Datum& right,
253258
ArithmeticOptions options = ArithmeticOptions(),
254259
ExecContext* ctx = NULLPTR);
255260

261+
/// \brief Find the element-wise maximum of any number of arrays or scalars.
262+
/// Array values must be the same length.
263+
///
264+
/// \param[in] args arrays or scalars to operate on.
265+
/// \param[in] options options for handling nulls, optional
266+
/// \param[in] ctx the function execution context, optional
267+
/// \return the element-wise maximum
268+
ARROW_EXPORT
269+
Result<Datum> ElementWiseMax(const std::vector<Datum>& args,
270+
ElementWiseAggregateOptions options = {},
271+
ExecContext* ctx = NULLPTR);
272+
273+
/// \brief Find the element-wise minimum of any number of arrays or scalars.
274+
/// Array values must be the same length.
275+
///
276+
/// \param[in] args arrays or scalars to operate on.
277+
/// \param[in] options options for handling nulls, optional
278+
/// \param[in] ctx the function execution context, optional
279+
/// \return the element-wise minimum
280+
ARROW_EXPORT
281+
Result<Datum> ElementWiseMin(const std::vector<Datum>& args,
282+
ElementWiseAggregateOptions options = {},
283+
ExecContext* ctx = NULLPTR);
284+
256285
/// \brief Compare a numeric array with a scalar.
257286
///
258287
/// \param[in] left datum to compare, must be an Array

‎cpp/src/arrow/compute/kernels/codegen_internal.h‎

Lines changed: 41 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -303,8 +303,12 @@ struct BoxScalar;
303303
template <typename Type>
304304
struct BoxScalar<Type, enable_if_has_c_type<Type>> {
305305
using T = typename GetOutputType<Type>::T;
306-
using ScalarType = typename TypeTraits<Type>::ScalarType;
307-
static void Box(T val, Scalar* out) { checked_cast<ScalarType*>(out)->value = val; }
306+
static void Box(T val, Scalar* out) {
307+
// Enables BoxScalar<Int64Type> to work on a (for example) Time64Scalar
308+
T* mutable_data = reinterpret_cast<T*>(
309+
checked_cast<::arrow::internal::PrimitiveScalarBase*>(out)->mutable_data());
310+
*mutable_data = val;
311+
}
308312
};
309313

310314
template <typename Type>
@@ -1093,6 +1097,41 @@ ArrayKernelExec GeneratePhysicalInteger(detail::GetTypeId get_id) {
10931097
}
10941098
}
10951099

1100+
template <template <typename... Args> class Generator, typename... Args>
1101+
ArrayKernelExec GeneratePhysicalNumeric(detail::GetTypeId get_id) {
1102+
switch (get_id.id) {
1103+
case Type::INT8:
1104+
return Generator<Int8Type, Args...>::Exec;
1105+
case Type::INT16:
1106+
return Generator<Int16Type, Args...>::Exec;
1107+
case Type::INT32:
1108+
case Type::DATE32:
1109+
case Type::TIME32:
1110+
return Generator<Int32Type, Args...>::Exec;
1111+
case Type::INT64:
1112+
case Type::DATE64:
1113+
case Type::TIMESTAMP:
1114+
case Type::TIME64:
1115+
case Type::DURATION:
1116+
return Generator<Int64Type, Args...>::Exec;
1117+
case Type::UINT8:
1118+
return Generator<UInt8Type, Args...>::Exec;
1119+
case Type::UINT16:
1120+
return Generator<UInt16Type, Args...>::Exec;
1121+
case Type::UINT32:
1122+
return Generator<UInt32Type, Args...>::Exec;
1123+
case Type::UINT64:
1124+
return Generator<UInt64Type, Args...>::Exec;
1125+
case Type::FLOAT:
1126+
return Generator<FloatType, Args...>::Exec;
1127+
case Type::DOUBLE:
1128+
return Generator<DoubleType, Args...>::Exec;
1129+
default:
1130+
DCHECK(false);
1131+
return ExecFail;
1132+
}
1133+
}
1134+
10961135
// Generate a kernel given a templated functor for integer types
10971136
//
10981137
// See "Numeric" above for description of the generator functor

‎cpp/src/arrow/compute/kernels/scalar_boolean.cc‎

Lines changed: 28 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,7 @@ inline Bitmap GetBitmap(const ArrayData& arr, int index) {
9595
return Bitmap{arr.buffers[index], arr.offset, arr.length};
9696
}
9797

98-
struct Invert {
98+
struct InvertOp {
9999
static Status Call(KernelContext* ctx, const Scalar& in, Scalar* out) {
100100
*checked_cast<BooleanScalar*>(out) = InvertScalar(in);
101101
return Status::OK();
@@ -115,8 +115,8 @@ struct Commutative {
115115
}
116116
};
117117

118-
struct And : Commutative<And> {
119-
using Commutative<And>::Call;
118+
struct AndOp : Commutative<AndOp> {
119+
using Commutative<AndOp>::Call;
120120

121121
static Status Call(KernelContext* ctx, const Scalar& left, const Scalar& right,
122122
Scalar* out) {
@@ -147,8 +147,8 @@ struct And : Commutative<And> {
147147
}
148148
};
149149

150-
struct KleeneAnd : Commutative<KleeneAnd> {
151-
using Commutative<KleeneAnd>::Call;
150+
struct KleeneAndOp : Commutative<KleeneAndOp> {
151+
using Commutative<KleeneAndOp>::Call;
152152

153153
static Status Call(KernelContext* ctx, const Scalar& left, const Scalar& right,
154154
Scalar* out) {
@@ -205,7 +205,7 @@ struct KleeneAnd : Commutative<KleeneAnd> {
205205
if (left.GetNullCount() == 0 && right.GetNullCount() == 0) {
206206
out->null_count = 0;
207207
out->buffers[0] = nullptr;
208-
return And::Call(ctx, left, right, out);
208+
return AndOp::Call(ctx, left, right, out);
209209
}
210210
auto compute_word = [](uint64_t left_true, uint64_t left_false, uint64_t right_true,
211211
uint64_t right_false, uint64_t* out_valid,
@@ -218,8 +218,8 @@ struct KleeneAnd : Commutative<KleeneAnd> {
218218
}
219219
};
220220

221-
struct Or : Commutative<Or> {
222-
using Commutative<Or>::Call;
221+
struct OrOp : Commutative<OrOp> {
222+
using Commutative<OrOp>::Call;
223223

224224
static Status Call(KernelContext* ctx, const Scalar& left, const Scalar& right,
225225
Scalar* out) {
@@ -250,8 +250,8 @@ struct Or : Commutative<Or> {
250250
}
251251
};
252252

253-
struct KleeneOr : Commutative<KleeneOr> {
254-
using Commutative<KleeneOr>::Call;
253+
struct KleeneOrOp : Commutative<KleeneOrOp> {
254+
using Commutative<KleeneOrOp>::Call;
255255

256256
static Status Call(KernelContext* ctx, const Scalar& left, const Scalar& right,
257257
Scalar* out) {
@@ -308,7 +308,7 @@ struct KleeneOr : Commutative<KleeneOr> {
308308
if (left.GetNullCount() == 0 && right.GetNullCount() == 0) {
309309
out->null_count = 0;
310310
out->buffers[0] = nullptr;
311-
return Or::Call(ctx, left, right, out);
311+
return OrOp::Call(ctx, left, right, out);
312312
}
313313

314314
static auto compute_word = [](uint64_t left_true, uint64_t left_false,
@@ -323,8 +323,8 @@ struct KleeneOr : Commutative<KleeneOr> {
323323
}
324324
};
325325

326-
struct Xor : Commutative<Xor> {
327-
using Commutative<Xor>::Call;
326+
struct XorOp : Commutative<XorOp> {
327+
using Commutative<XorOp>::Call;
328328

329329
static Status Call(KernelContext* ctx, const Scalar& left, const Scalar& right,
330330
Scalar* out) {
@@ -355,10 +355,10 @@ struct Xor : Commutative<Xor> {
355355
}
356356
};
357357

358-
struct AndNot {
358+
struct AndNotOp {
359359
static Status Call(KernelContext* ctx, const Scalar& left, const Scalar& right,
360360
Scalar* out) {
361-
return And::Call(ctx, left, InvertScalar(right), out);
361+
return AndOp::Call(ctx, left, InvertScalar(right), out);
362362
}
363363

364364
static Status Call(KernelContext* ctx, const Scalar& left, const ArrayData& right,
@@ -373,7 +373,7 @@ struct AndNot {
373373

374374
static Status Call(KernelContext* ctx, const ArrayData& left, const Scalar& right,
375375
ArrayData* out) {
376-
return And::Call(ctx, left, InvertScalar(right), out);
376+
return AndOp::Call(ctx, left, InvertScalar(right), out);
377377
}
378378

379379
static Status Call(KernelContext* ctx, const ArrayData& left, const ArrayData& right,
@@ -385,10 +385,10 @@ struct AndNot {
385385
}
386386
};
387387

388-
struct KleeneAndNot {
388+
struct KleeneAndNotOp {
389389
static Status Call(KernelContext* ctx, const Scalar& left, const Scalar& right,
390390
Scalar* out) {
391-
return KleeneAnd::Call(ctx, left, InvertScalar(right), out);
391+
return KleeneAndOp::Call(ctx, left, InvertScalar(right), out);
392392
}
393393

394394
static Status Call(KernelContext* ctx, const Scalar& left, const ArrayData& right,
@@ -430,15 +430,15 @@ struct KleeneAndNot {
430430

431431
static Status Call(KernelContext* ctx, const ArrayData& left, const Scalar& right,
432432
ArrayData* out) {
433-
return KleeneAnd::Call(ctx, left, InvertScalar(right), out);
433+
return KleeneAndOp::Call(ctx, left, InvertScalar(right), out);
434434
}
435435

436436
static Status Call(KernelContext* ctx, const ArrayData& left, const ArrayData& right,
437437
ArrayData* out) {
438438
if (left.GetNullCount() == 0 && right.GetNullCount() == 0) {
439439
out->null_count = 0;
440440
out->buffers[0] = nullptr;
441-
return AndNot::Call(ctx, left, right, out);
441+
return AndNotOp::Call(ctx, left, right, out);
442442
}
443443

444444
static auto compute_word = [](uint64_t left_true, uint64_t left_false,
@@ -543,20 +543,20 @@ namespace internal {
543543

544544
void RegisterScalarBoolean(FunctionRegistry* registry) {
545545
// These functions can write into sliced output bitmaps
546-
MakeFunction("invert", 1, applicator::SimpleUnary<Invert>, &invert_doc, registry);
547-
MakeFunction("and", 2, applicator::SimpleBinary<And>, &and_doc, registry);
548-
MakeFunction("and_not", 2, applicator::SimpleBinary<AndNot>, &and_not_doc, registry);
549-
MakeFunction("or", 2, applicator::SimpleBinary<Or>, &or_doc, registry);
550-
MakeFunction("xor", 2, applicator::SimpleBinary<Xor>, &xor_doc, registry);
546+
MakeFunction("invert", 1, applicator::SimpleUnary<InvertOp>, &invert_doc, registry);
547+
MakeFunction("and", 2, applicator::SimpleBinary<AndOp>, &and_doc, registry);
548+
MakeFunction("and_not", 2, applicator::SimpleBinary<AndNotOp>, &and_not_doc, registry);
549+
MakeFunction("or", 2, applicator::SimpleBinary<OrOp>, &or_doc, registry);
550+
MakeFunction("xor", 2, applicator::SimpleBinary<XorOp>, &xor_doc, registry);
551551

552552
// The Kleene logic kernels cannot write into sliced output bitmaps
553-
MakeFunction("and_kleene", 2, applicator::SimpleBinary<KleeneAnd>, &and_kleene_doc,
553+
MakeFunction("and_kleene", 2, applicator::SimpleBinary<KleeneAndOp>, &and_kleene_doc,
554554
registry,
555555
/*can_write_into_slices=*/false, NullHandling::COMPUTED_PREALLOCATE);
556-
MakeFunction("and_not_kleene", 2, applicator::SimpleBinary<KleeneAndNot>,
556+
MakeFunction("and_not_kleene", 2, applicator::SimpleBinary<KleeneAndNotOp>,
557557
&and_not_kleene_doc, registry,
558558
/*can_write_into_slices=*/false, NullHandling::COMPUTED_PREALLOCATE);
559-
MakeFunction("or_kleene", 2, applicator::SimpleBinary<KleeneOr>, &or_kleene_doc,
559+
MakeFunction("or_kleene", 2, applicator::SimpleBinary<KleeneOrOp>, &or_kleene_doc,
560560
registry,
561561
/*can_write_into_slices=*/false, NullHandling::COMPUTED_PREALLOCATE);
562562
}

0 commit comments

Comments
 (0)