ARROW-12959: [C++][R] Option for is_null(NaN) to evaluate to true

cpp/src/arrow/compute/api_scalar.cc

@@ -161,6 +161,8 @@ static auto kMakeStructOptionsType = GetFunctionOptionsType<MakeStructOptions>(
 static auto kDayOfWeekOptionsType = GetFunctionOptionsType<DayOfWeekOptions>(
     DataMember("one_based_numbering", &DayOfWeekOptions::one_based_numbering),
     DataMember("week_start", &DayOfWeekOptions::week_start));
+static auto kNullOptionsType = GetFunctionOptionsType<NullOptions>(
+    DataMember("nan_is_null", &NullOptions::nan_is_null));
 }  // namespace
 }  // namespace internal
 
@@ -291,6 +293,10 @@ DayOfWeekOptions::DayOfWeekOptions(bool one_based_numbering, uint32_t week_start
       week_start(week_start) {}
 constexpr char DayOfWeekOptions::kTypeName[];
 
+NullOptions::NullOptions(bool nan_is_null)
+    : FunctionOptions(internal::kNullOptionsType), nan_is_null(nan_is_null) {}
+constexpr char NullOptions::kTypeName[];
+
 namespace internal {
 void RegisterScalarOptions(FunctionRegistry* registry) {
   DCHECK_OK(registry->AddFunctionOptionsType(kArithmeticOptionsType));
@@ -310,6 +316,7 @@ void RegisterScalarOptions(FunctionRegistry* registry) {
   DCHECK_OK(registry->AddFunctionOptionsType(kSliceOptionsType));
   DCHECK_OK(registry->AddFunctionOptionsType(kMakeStructOptionsType));
   DCHECK_OK(registry->AddFunctionOptionsType(kDayOfWeekOptionsType));
+  DCHECK_OK(registry->AddFunctionOptionsType(kNullOptionsType));
 }
 }  // namespace internal
 
@@ -463,7 +470,6 @@ Result<Datum> Compare(const Datum& left, const Datum& right, CompareOptions opti
 // Validity functions
 
 SCALAR_EAGER_UNARY(IsValid, "is_valid")
-SCALAR_EAGER_UNARY(IsNull, "is_null")
 SCALAR_EAGER_UNARY(IsNan, "is_nan")
 
 Result<Datum> FillNull(const Datum& values, const Datum& fill_value, ExecContext* ctx) {
@@ -483,6 +489,10 @@ Result<Datum> CaseWhen(const Datum& cond, const std::vector<Datum>& cases,
   return CallFunction("case_when", args, ctx);
 }
 
+Result<Datum> IsNull(const Datum& arg, NullOptions options, ExecContext* ctx) {
+  return CallFunction("is_null", {arg}, &options, ctx);
+}
+
 // ----------------------------------------------------------------------
 // Temporal functions
 

cpp/src/arrow/compute/api_scalar.h

@@ -224,6 +224,15 @@ class ARROW_EXPORT SliceOptions : public FunctionOptions {
   int64_t start, stop, step;
 };
 
+class ARROW_EXPORT NullOptions : public FunctionOptions {
+ public:
+  explicit NullOptions(bool nan_is_null = false);
+  constexpr static char const kTypeName[] = "NullOptions";
+  static NullOptions Defaults() { return NullOptions{}; }
+
+  bool nan_is_null = false;
+};
+
 enum CompareOperator : int8_t {
   EQUAL,
   NOT_EQUAL,
@@ -756,13 +765,15 @@ Result<Datum> IsValid(const Datum& values, ExecContext* ctx = NULLPTR);
 /// false otherwise
 ///
 /// \param[in] values input to examine for nullity
+/// \param[in] options NullOptions
 /// \param[in] ctx the function execution context, optional
 /// \return the resulting datum
 ///
 /// \since 1.0.0
 /// \note API not yet finalized
 ARROW_EXPORT
-Result<Datum> IsNull(const Datum& values, ExecContext* ctx = NULLPTR);
+Result<Datum> IsNull(const Datum& values, NullOptions options = NullOptions::Defaults(),
+                     ExecContext* ctx = NULLPTR);
 
 /// \brief IsNan returns true for each element of `values` that is NaN,
 /// false otherwise

cpp/src/arrow/compute/exec/expression.cc

@@ -1154,7 +1154,9 @@ Expression greater_equal(Expression lhs, Expression rhs) {
   return call("greater_equal", {std::move(lhs), std::move(rhs)});
 }
 
-Expression is_null(Expression lhs) { return call("is_null", {std::move(lhs)}); }
+Expression is_null(Expression lhs, bool nan_is_null) {
+  return call("is_null", {std::move(lhs)}, compute::NullOptions(std::move(nan_is_null)));
+}
 
 Expression is_valid(Expression lhs) { return call("is_valid", {std::move(lhs)}); }
 

cpp/src/arrow/compute/exec/expression.h

@@ -255,7 +255,7 @@ ARROW_EXPORT Expression greater(Expression lhs, Expression rhs);
 
 ARROW_EXPORT Expression greater_equal(Expression lhs, Expression rhs);
 
-ARROW_EXPORT Expression is_null(Expression lhs);
+ARROW_EXPORT Expression is_null(Expression lhs, bool nan_is_null = false);
 
 ARROW_EXPORT Expression is_valid(Expression lhs);
 

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

@@ -17,6 +17,7 @@
 
 #include <cmath>
 
+#include "arrow/compute/api_scalar.h"
 #include "arrow/compute/kernels/common.h"
 
 #include "arrow/util/bit_util.h"
@@ -74,21 +75,71 @@ struct IsInfOperator {
   }
 };
 
+using NanOptionsState = OptionsWrapper<NullOptions>;
+
 struct IsNullOperator {
   static Status Call(KernelContext* ctx, const Scalar& in, Scalar* out) {
-    checked_cast<BooleanScalar*>(out)->value = !in.is_valid;
+    const auto& options = NanOptionsState::Get(ctx);
+    bool* out_value = &checked_cast<BooleanScalar*>(out)->value;
+
+    if (in.is_valid) {
+      if (options.nan_is_null && is_floating(in.type->id())) {
+        switch (in.type->id()) {
+          case Type::FLOAT:
+            *out_value = std::isnan(internal::UnboxScalar<FloatType>::Unbox(in));
+            break;
+          case Type::DOUBLE:
+            *out_value = std::isnan(internal::UnboxScalar<DoubleType>::Unbox(in));
+            break;
+          default:
+            return Status::NotImplemented("NaN detection not implemented for type ",
+                                          in.type->ToString());
+        }
+      } else {
+        *out_value = false;
+      }
+    } else {
+      *out_value = true;
+    }
+
     return Status::OK();
   }
 
+  template <typename T>
+  static void SetNanBits(const ArrayData& arr, uint8_t* out_bitmap, int64_t out_offset) {
+    const T* data = arr.GetValues<T>(1);
+    for (int64_t i = 0; i < arr.length; ++i) {
+      if (std::isnan(data[i])) {
+        BitUtil::SetBit(out_bitmap, i + out_offset);
+      }
+    }
+  }
+
   static Status Call(KernelContext* ctx, const ArrayData& arr, ArrayData* out) {
-    if (arr.MayHaveNulls()) {
+    const auto& options = NanOptionsState::Get(ctx);
+
+    uint8_t* out_bitmap = out->buffers[1]->mutable_data();
+    if (arr.GetNullCount() > 0) {
       // Input has nulls => output is the inverted null (validity) bitmap.
-      InvertBitmap(arr.buffers[0]->data(), arr.offset, arr.length,
-                   out->buffers[1]->mutable_data(), out->offset);
+      InvertBitmap(arr.buffers[0]->data(), arr.offset, arr.length, out_bitmap,
+                   out->offset);
     } else {
       // Input has no nulls => output is entirely false.
-      BitUtil::SetBitsTo(out->buffers[1]->mutable_data(), out->offset, out->length,
-                         false);
+      BitUtil::SetBitsTo(out_bitmap, out->offset, out->length, false);
+    }
+
+    if (is_floating(arr.type->id()) && options.nan_is_null) {
+      switch (arr.type->id()) {
+        case Type::FLOAT:
+          SetNanBits<float>(arr, out_bitmap, out->offset);
+          break;
+        case Type::DOUBLE:
+          SetNanBits<double>(arr, out_bitmap, out->offset);
+          break;
+        default:
+          return Status::NotImplemented("NaN detection not implemented for type ",
+                                        arr.type->ToString());
+      }
     }
     return Status::OK();
   }
@@ -104,11 +155,13 @@ struct IsNanOperator {
 void MakeFunction(std::string name, const FunctionDoc* doc,
                   std::vector<InputType> in_types, OutputType out_type,
                   ArrayKernelExec exec, FunctionRegistry* registry,
-                  MemAllocation::type mem_allocation, bool can_write_into_slices) {
+                  MemAllocation::type mem_allocation, bool can_write_into_slices,
+                  const FunctionOptions* default_options = NULLPTR,
+                  KernelInit init = NULLPTR) {
   Arity arity{static_cast<int>(in_types.size())};
-  auto func = std::make_shared<ScalarFunction>(name, arity, doc);
+  auto func = std::make_shared<ScalarFunction>(name, arity, doc, default_options);
 
-  ScalarKernel kernel(std::move(in_types), out_type, exec);
+  ScalarKernel kernel(std::move(in_types), out_type, exec, init);
   kernel.null_handling = NullHandling::OUTPUT_NOT_NULL;
   kernel.can_write_into_slices = can_write_into_slices;
   kernel.mem_allocation = mem_allocation;
@@ -202,9 +255,11 @@ const FunctionDoc is_inf_doc(
     ("For each input value, emit true iff the value is infinite (inf or -inf)."),
     {"values"});
 
-const FunctionDoc is_null_doc("Return true if null",
-                              ("For each input value, emit true iff the value is null."),
-                              {"values"});
+const FunctionDoc is_null_doc(
+    "Return true if null (and optionally NaN)",
+    ("For each input value, emit true iff the value is null.\n"
+     "True may also be emitted for NaN values by setting the `nan_is_null` flag."),
+    {"values"}, "NullOptions");
 
 const FunctionDoc is_nan_doc("Return true if NaN",
                              ("For each input value, emit true iff the value is NaN."),
@@ -213,12 +268,13 @@ const FunctionDoc is_nan_doc("Return true if NaN",
 }  // namespace
 
 void RegisterScalarValidity(FunctionRegistry* registry) {
+  static auto kNullOptions = NullOptions::Defaults();
   MakeFunction("is_valid", &is_valid_doc, {ValueDescr::ANY}, boolean(), IsValidExec,
                registry, MemAllocation::NO_PREALLOCATE, /*can_write_into_slices=*/false);
 
   MakeFunction("is_null", &is_null_doc, {ValueDescr::ANY}, boolean(), IsNullExec,
                registry, MemAllocation::PREALLOCATE,
-               /*can_write_into_slices=*/true);
+               /*can_write_into_slices=*/true, &kNullOptions, NanOptionsState::Init);
 
   DCHECK_OK(registry->AddFunction(MakeIsFiniteFunction("is_finite", &is_finite_doc)));
   DCHECK_OK(registry->AddFunction(MakeIsInfFunction("is_inf", &is_inf_doc)));