Adding ScalarAggregateOptions to Any and All kernels.

Author: rok

cpp/src/arrow/compute/api_aggregate.cc

@@ -45,12 +45,14 @@ Result<Datum> MinMax(const Datum& value, const ScalarAggregateOptions& options,
   return CallFunction("min_max", {value}, &options, ctx);
 }
 
-Result<Datum> Any(const Datum& value, ExecContext* ctx) {
-  return CallFunction("any", {value}, ctx);
+Result<Datum> Any(const Datum& value, const ScalarAggregateOptions& options,
+                  ExecContext* ctx) {
+  return CallFunction("any", {value}, &options, ctx);
 }
 
-Result<Datum> All(const Datum& value, ExecContext* ctx) {
-  return CallFunction("all", {value}, ctx);
+Result<Datum> All(const Datum& value, const ScalarAggregateOptions& options,
+                  ExecContext* ctx) {
+  return CallFunction("all", {value}, &options, ctx);
 }
 
 Result<Datum> Mode(const Datum& value, const ModeOptions& options, ExecContext* ctx) {

cpp/src/arrow/compute/api_aggregate.h

@@ -201,30 +201,38 @@ Result<Datum> MinMax(
 /// \brief Test whether any element in a boolean array evaluates to true.
 ///
 /// This function returns true if any of the elements in the array evaluates
-/// to true and false otherwise. Null values are skipped.
+/// to true and false otherwise. Null values are ignored by default.
 ///
 /// \param[in] value input datum, expecting a boolean array
+/// \param[in] options see ScalarAggregateOptions for more information
 /// \param[in] ctx the function execution context, optional
 /// \return resulting datum as a BooleanScalar
 ///
 /// \since 3.0.0
 /// \note API not yet finalized
 ARROW_EXPORT
-Result<Datum> Any(const Datum& value, ExecContext* ctx = NULLPTR);
+Result<Datum> Any(
+    const Datum& value,
+    const ScalarAggregateOptions& options = ScalarAggregateOptions::Defaults(),
+    ExecContext* ctx = NULLPTR);
 
 /// \brief Test whether all elements in a boolean array evaluate to true.
 ///
 /// This function returns true if all of the elements in the array evaluate
-/// to true and false otherwise. Null values are skipped.
+/// to true and false otherwise. Null values are ignored by default.
 ///
 /// \param[in] value input datum, expecting a boolean array
+/// \param[in] options see ScalarAggregateOptions for more information
 /// \param[in] ctx the function execution context, optional
 /// \return resulting datum as a BooleanScalar
 
 /// \since 3.0.0
 /// \note API not yet finalized
 ARROW_EXPORT
-Result<Datum> All(const Datum& value, ExecContext* ctx = NULLPTR);
+Result<Datum> All(
+    const Datum& value,
+    const ScalarAggregateOptions& options = ScalarAggregateOptions::Defaults(),
+    ExecContext* ctx = NULLPTR);
 
 /// \brief Calculate the modal (most common) value of a numeric array
 ///

cpp/src/arrow/compute/kernels/aggregate_basic.cc

@@ -142,13 +142,19 @@ Result<std::unique_ptr<KernelState>> MinMaxInit(KernelContext* ctx,
 // Any implementation
 
 struct BooleanAnyImpl : public ScalarAggregator {
+  explicit BooleanAnyImpl(ScalarAggregateOptions options) : options(std::move(options)) {}
+
   Status Consume(KernelContext*, const ExecBatch& batch) override {
     // short-circuit if seen a True already
-    if (this->any == true) {
+    if (options.skip_nulls && this->any == true) {
+      return Status::OK();
+    }
+    // short-circuit if seen a null already
+    if (!options.skip_nulls && this->has_nulls) {
       return Status::OK();
     }
-
     const auto& data = *batch[0].array();
+    this->has_nulls = data.GetNullCount() > 0;
     arrow::internal::OptionalBinaryBitBlockCounter counter(
         data.buffers[0], data.offset, data.buffers[1], data.offset, data.length);
     int64_t position = 0;
@@ -166,32 +172,48 @@ struct BooleanAnyImpl : public ScalarAggregator {
   Status MergeFrom(KernelContext*, KernelState&& src) override {
     const auto& other = checked_cast<const BooleanAnyImpl&>(src);
     this->any |= other.any;
+    this->has_nulls |= other.has_nulls;
     return Status::OK();
   }
 
-  Status Finalize(KernelContext*, Datum* out) override {
-    out->value = std::make_shared<BooleanScalar>(this->any);
+  Status Finalize(KernelContext* ctx, Datum* out) override {
+    if (!options.skip_nulls && this->has_nulls) {
+      out->value = std::make_shared<BooleanScalar>();
+    } else {
+      out->value = std::make_shared<BooleanScalar>(this->any);
+    }
     return Status::OK();
   }
 
   bool any = false;
+  bool has_nulls = false;
+  ScalarAggregateOptions options;
 };
 
 Result<std::unique_ptr<KernelState>> AnyInit(KernelContext*, const KernelInitArgs& args) {
-  return ::arrow::internal::make_unique<BooleanAnyImpl>();
+  const ScalarAggregateOptions options =
+      static_cast<const ScalarAggregateOptions&>(*args.options);
+  return ::arrow::internal::make_unique<BooleanAnyImpl>(
+      static_cast<const ScalarAggregateOptions&>(*args.options));
 }
 
 // ----------------------------------------------------------------------
 // All implementation
 
 struct BooleanAllImpl : public ScalarAggregator {
+  explicit BooleanAllImpl(ScalarAggregateOptions options) : options(std::move(options)) {}
+
   Status Consume(KernelContext*, const ExecBatch& batch) override {
     // short-circuit if seen a false already
-    if (this->all == false) {
+    if (options.skip_nulls && this->all == false) {
+      return Status::OK();
+    }
+    // short-circuit if seen a null already
+    if (!options.skip_nulls && this->has_nulls) {
       return Status::OK();
     }
-
     const auto& data = *batch[0].array();
+    this->has_nulls = data.GetNullCount() > 0;
     arrow::internal::OptionalBinaryBitBlockCounter counter(
         data.buffers[1], data.offset, data.buffers[0], data.offset, data.length);
     int64_t position = 0;
@@ -210,19 +232,27 @@ struct BooleanAllImpl : public ScalarAggregator {
   Status MergeFrom(KernelContext*, KernelState&& src) override {
     const auto& other = checked_cast<const BooleanAllImpl&>(src);
     this->all &= other.all;
+    this->has_nulls |= other.has_nulls;
     return Status::OK();
   }
 
   Status Finalize(KernelContext*, Datum* out) override {
-    out->value = std::make_shared<BooleanScalar>(this->all);
+    if (!options.skip_nulls && this->has_nulls) {
+      out->value = std::make_shared<BooleanScalar>();
+    } else {
+      out->value = std::make_shared<BooleanScalar>(this->all);
+    }
     return Status::OK();
   }
 
   bool all = true;
+  bool has_nulls = false;
+  ScalarAggregateOptions options;
 };
 
 Result<std::unique_ptr<KernelState>> AllInit(KernelContext*, const KernelInitArgs& args) {
-  return ::arrow::internal::make_unique<BooleanAllImpl>();
+  return ::arrow::internal::make_unique<BooleanAllImpl>(
+      static_cast<const ScalarAggregateOptions&>(*args.options));
 }
 
 // ----------------------------------------------------------------------
@@ -408,12 +438,16 @@ const FunctionDoc min_max_doc{"Compute the minimum and maximum values of a numer
                               "ScalarAggregateOptions"};
 
 const FunctionDoc any_doc{"Test whether any element in a boolean array evaluates to true",
-                          ("Null values are ignored."),
-                          {"array"}};
+                          ("Null values are ignored by default.\n"
+                           "This can be changed through ScalarAggregateOptions."),
+                          {"array"},
+                          "ScalarAggregateOptions"};
 
 const FunctionDoc all_doc{"Test whether all elements in a boolean array evaluate to true",
-                          ("Null values are ignored."),
-                          {"array"}};
+                          ("Null values are ignored by default.\n"
+                           "This can be changed through ScalarAggregateOptions."),
+                          {"array"},
+                          "ScalarAggregateOptions"};
 
 const FunctionDoc index_doc{"Find the index of the first occurrence of a given value",
                             ("The result is always computed as an int64_t, regardless\n"
@@ -497,12 +531,14 @@ void RegisterScalarAggregateBasic(FunctionRegistry* registry) {
   DCHECK_OK(registry->AddFunction(std::move(func)));
 
   // any
-  func = std::make_shared<ScalarAggregateFunction>("any", Arity::Unary(), &any_doc);
+  func = std::make_shared<ScalarAggregateFunction>("any", Arity::Unary(), &any_doc,
+                                                   &default_scalar_aggregate_options);
   aggregate::AddBasicAggKernels(aggregate::AnyInit, {boolean()}, boolean(), func.get());
   DCHECK_OK(registry->AddFunction(std::move(func)));
 
   // all
-  func = std::make_shared<ScalarAggregateFunction>("all", Arity::Unary(), &all_doc);
+  func = std::make_shared<ScalarAggregateFunction>("all", Arity::Unary(), &all_doc,
+                                                   &default_scalar_aggregate_options);
   aggregate::AddBasicAggKernels(aggregate::AllInit, {boolean()}, boolean(), func.get());
   DCHECK_OK(registry->AddFunction(std::move(func)));