GH-43541: [C++] Check accepted device allocation types before executing kernel

cpp/src/arrow/array/data.cc

@@ -68,6 +68,15 @@ static inline void AdjustNonNullable(Type::type type_id, int64_t length,
 
 namespace internal {
 
+void AdjustNonNullable(ArrayData* array_data) {
+ int64_t null_count = kUnknownNullCount;
+ AdjustNonNullable(array_data->type->id(), array_data->length, &array_data->buffers,
+ &null_count);
+ if (null_count != kUnknownNullCount) {
+ array_data->null_count.store(null_count, std::memory_order_release);
+ }
+}
+
 bool IsNullSparseUnion(const ArrayData& data, int64_t i) {
  auto* union_type = checked_cast<const SparseUnionType*>(data.type.get());
  const auto* types = reinterpret_cast<const int8_t*>(data.buffers[1]->data());

cpp/src/arrow/array/data.h

@@ -36,6 +36,9 @@
 namespace arrow {
 
 namespace internal {
+
+ARROW_EXPORT void AdjustNonNullable(ArrayData* array_data);
+
 // ----------------------------------------------------------------------
 // Null handling for types without a validity bitmap and the dictionary type
 
@@ -342,6 +345,17 @@ struct ARROW_EXPORT ArrayData {
  /// ...
  /// }
  bool MayHaveLogicalNulls() const {
+ if (ARROW_PREDICT_FALSE(device_type() != DeviceAllocationType::kCPU)) {
+ // Some arrays are malformed in that they have kUnknownNullCount, but no
+ // validity bitmap and that makes MayHaveLogicalNulls() return true when
+ // it should return false.
+ //
+ // Not a problem if we're on the CPU because ArrayData::GetNullCount()
+ // will eventually return 0. But it's a problem if we're trying to make
+ // cheap decisions without fully calculating the null count and want to
+ // defer the processing of the bitmap to later.
+ internal::AdjustNonNullable(const_cast<ArrayData*>(this));
+ }
  if (buffers[0] != NULLPTR) {
  return null_count.load() != 0;
  }

cpp/src/arrow/compute/function.cc

@@ -102,13 +102,38 @@ Status Function::CheckArity(size_t num_args) const {
 namespace {
 
 Status CheckOptions(const Function& function, const FunctionOptions* options) {
- if (options == nullptr && function.doc().options_required) {
+ if (ARROW_PREDICT_FALSE(options == nullptr && function.doc().options_required)) {
  return Status::Invalid("Function '", function.name(),
  "' cannot be called without options");
  }
  return Status::OK();
 }
 
+Status BadDeviceTypeStatus(const std::string& func_name,
+ DeviceAllocationTypeSet expected_device_type_set,
+ int offending_arg_index, const Datum& offending_arg) {
+ std::string ordinal;
+ switch (offending_arg_index) {
+ case 0:
+ ordinal = "1st";
+ break;
+ case 1:
+ ordinal = "2nd";
+ break;
+ case 2:
+ ordinal = "3rd";
+ break;
+ default:
+ ordinal = std::to_string(offending_arg_index + 1);
+ ordinal += "th";
+ break;
+ }
+ return Status::NotImplemented("'", func_name, "' expects ", ordinal,
+ " argument's device allocation types(s) to be in ",
+ expected_device_type_set.ToString(), " but got ",
+ offending_arg.device_types().ToString(), ".");
+}
+
 } // namespace
 
 namespace detail {
@@ -248,6 +273,15 @@ struct FunctionExecutorImpl : public FunctionExecutor {
  }
  args_with_cast[i] = std::move(arg);
  }
+ {
+ DeviceAllocationTypeSet expected_device_type_set;
+ int offending_arg_index = 0;
+ if (ARROW_PREDICT_FALSE(!kernel->signature->MatchesDeviceAllocationTypes(
+ args_with_cast, &expected_device_type_set, &offending_arg_index))) {
+ return BadDeviceTypeStatus(func_name, expected_device_type_set,
+ offending_arg_index, args[offending_arg_index]);
+ }
+ }
 
  detail::DatumAccumulator listener;
 

cpp/src/arrow/compute/kernel.cc

@@ -23,6 +23,7 @@
 #include <string>
 
 #include "arrow/buffer.h"
+#include "arrow/chunked_array.h"
 #include "arrow/compute/exec.h"
 #include "arrow/device_allocation_type_set.h"
 #include "arrow/result.h"
@@ -431,6 +432,25 @@ bool InputType::Matches(const Datum& value) const {
  return Matches(*value.type());
 }
 
+bool InputType::MatchesDeviceAllocationType(const Datum& value) const {
+ DCHECK(Matches(value));
+ switch (value.kind()) {
+ case Datum::NONE:
+ case Datum::RECORD_BATCH:
+ case Datum::TABLE:
+ break;
+ case Datum::ARRAY:
+ return accepted_device_types_.contains(value.array()->device_type());
+ case Datum::CHUNKED_ARRAY:
+ return accepted_device_types_.Contains(value.chunked_array()->device_types());
+ case Datum::SCALAR:
+ // Scalars are asssumed as always residing in CPU memory for now.
+ return accepted_device_types_.contains(DeviceAllocationType::kCPU);
+ }
+ DCHECK(false) << "MatchesDeviceAllocationType expects ARRAY, CHUNKED_ARRAY or SCALAR";
+ return false;
+}
+
 const std::shared_ptr<DataType>& InputType::type() const {
  DCHECK_EQ(InputType::EXACT_TYPE, kind_);
  return type_;
@@ -529,6 +549,50 @@ bool KernelSignature::MatchesInputs(const std::vector<TypeHolder>& types) const
  return true;
 }
 
+bool KernelSignature::MatchesDeviceAllocationTypes(
+ const std::vector<Datum>& args, DeviceAllocationTypeSet* out_expected_device_types,
+ int* out_offending_arg_index) const {
+ DeviceAllocationTypeSet expected_device_types;
+ int offending_arg_index = 0;
+ bool matches = true;
+ if (is_varargs_) {
+ for (size_t i = 0; i < args.size(); ++i) {
+ auto& param_type = in_types_[std::min(i, in_types_.size() - 1)];
+ DCHECK(param_type.Matches(*args[i].type()));
+ if (!param_type.MatchesDeviceAllocationType(args[i])) {
+ matches = false;
+ expected_device_types = param_type.accepted_device_types();
+ offending_arg_index = static_cast<int>(i);
+ break;
+ }
+ }
+ } else {
+ DCHECK(args.size() == in_types_.size());
+ if (args.size() != in_types_.size()) {
+ matches = false;
+ } else {
+ for (size_t i = 0; i < in_types_.size(); ++i) {
+ auto& param_type = in_types_[i];
+ DCHECK(param_type.Matches(*args[i].type()));
+ if (!param_type.MatchesDeviceAllocationType(args[i])) {
+ matches = false;
+ offending_arg_index = static_cast<int>(i);
+ expected_device_types = param_type.accepted_device_types();
+ break;
+ }
+ }
+ }
+ }
+
+ if (out_expected_device_types) {
+ *out_expected_device_types = expected_device_types;
+ }
+ if (out_offending_arg_index) {
+ *out_offending_arg_index = offending_arg_index;
+ }
+ return matches;
+}
+
 size_t KernelSignature::Hash() const {
  if (hash_code_ != 0) {
  return hash_code_;

cpp/src/arrow/compute/kernel.h

@@ -29,6 +29,7 @@
 #include <vector>
 
 #include "arrow/buffer.h"
+#include "arrow/chunked_array.h" // for DeviceAllocationTypeSet
 #include "arrow/compute/exec.h"
 #include "arrow/datum.h"
 #include "arrow/device_allocation_type_set.h"
@@ -194,16 +195,37 @@ class ARROW_EXPORT InputType {
  USE_TYPE_MATCHER
  };
 
- /// \brief Accept any value type
- InputType() : kind_(ANY_TYPE) {}
-
- /// \brief Accept an exact value type.
+ /// \brief Accept any value type of values allocated on the given set of device
+ /// types.
+ explicit InputType(DeviceAllocationTypeSet accepted_devices)
+ : kind_(ANY_TYPE), accepted_device_types_(std::move(accepted_devices)) {}
+
+ /// \brief Accept an exact value type of values allocated on the given set of device
+ /// types.
+ InputType(std::shared_ptr<DataType> type, DeviceAllocationTypeSet accepted_devices)
+ : kind_(EXACT_TYPE),
+ type_(std::move(type)),
+ accepted_device_types_(accepted_devices) {}
+
+ /// \brief Use the passed TypeMatcher to check the type of arguments allocated on the
+ /// given set of device types.
+ InputType(std::shared_ptr<TypeMatcher> type_matcher,
+ DeviceAllocationTypeSet accepted_devices)
+ : kind_(USE_TYPE_MATCHER),
+ type_matcher_(std::move(type_matcher)),
+ accepted_device_types_(std::move(accepted_devices)) {}
+
+ /// \brief Accept any value type of values allocated on the CPU.
+ InputType() : InputType(DeviceAllocationTypeSet::CpuOnly()) {}
+
+ /// \brief Accept an exact value type of values allocated on the CPU.
  InputType(std::shared_ptr<DataType> type) // NOLINT implicit construction
- : kind_(EXACT_TYPE), type_(std::move(type)) {}
+ : InputType(std::move(type), DeviceAllocationTypeSet::CpuOnly()) {}
 
- /// \brief Use the passed TypeMatcher to type check.
+ /// \brief Use the passed TypeMatcher to type check the type of arguments allocated on
+ /// the CPU.
  InputType(std::shared_ptr<TypeMatcher> type_matcher) // NOLINT implicit construction
- : kind_(USE_TYPE_MATCHER), type_matcher_(std::move(type_matcher)) {}
+ : InputType(std::move(type_matcher), DeviceAllocationTypeSet::CpuOnly()) {}
 
  /// \brief Match any type with the given Type::type. Uses a TypeMatcher for
  /// its implementation.
@@ -242,6 +264,13 @@ class ARROW_EXPORT InputType {
  /// \brief Return true if the type matches this InputType
  bool Matches(const DataType& type) const;
 
+ /// \brief Return true if the Datum's device allocation type matches this
+ /// argument's accepted device allocation type set (and only allows scalar or
+ /// array-like Datums).
+ ///
+ /// \pre Matches(value) == true
+ bool MatchesDeviceAllocationType(const Datum& value) const;
+
  /// \brief The type matching rule that this InputType uses.
  Kind kind() const { return kind_; }
 
@@ -255,26 +284,32 @@ class ARROW_EXPORT InputType {
  /// and will assert in debug builds.
  const TypeMatcher& type_matcher() const;
 
+ /// \brief The device allocation types that are accepted for this input type.
+ DeviceAllocationTypeSet accepted_device_types() const { return accepted_device_types_; }
+
  private:
  void CopyInto(const InputType& other) {
  this->kind_ = other.kind_;
  this->type_ = other.type_;
  this->type_matcher_ = other.type_matcher_;
+ this->accepted_device_types_ = other.accepted_device_types_;
  }
 
  void MoveInto(InputType&& other) {
  this->kind_ = other.kind_;
  this->type_ = std::move(other.type_);
  this->type_matcher_ = std::move(other.type_matcher_);
+ this->accepted_device_types_ = std::move(other.accepted_device_types_);
  }
 
  Kind kind_;
-
  // For EXACT_TYPE Kind
  std::shared_ptr<DataType> type_;
-
  // For USE_TYPE_MATCHER Kind
  std::shared_ptr<TypeMatcher> type_matcher_;
+
+ // For matching device types. CPU-only by default.
+ DeviceAllocationTypeSet accepted_device_types_;
 };
 
 /// \brief Container to capture both exact and input-dependent output types.
@@ -368,6 +403,15 @@ class ARROW_EXPORT KernelSignature {
  /// value descriptors.
  bool MatchesInputs(const std::vector<TypeHolder>& types) const;
 
+ /// \brief Return true if the signature is compatible with the list of
+ /// execution arguments regarding device allocation types.
+ ///
+ /// \pre MatchesInputs(GetTypes(args)) == true
+ bool MatchesDeviceAllocationTypes(
+ const std::vector<Datum>& args,
+ DeviceAllocationTypeSet* out_expected_device_types = NULLPTR,
+ int* offending_arg_index = NULLPTR) const;
+
  /// \brief Returns true if the input types of each signature are
  /// equal. Well-formed functions should have a deterministic output type
  /// given input types, but currently it is the responsibility of the

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

@@ -54,6 +54,7 @@ using internal::BitBlockCounter;
 using internal::CheckIndexBounds;
 using internal::CopyBitmap;
 using internal::CountSetBits;
+using internal::may_have_validity_bitmap;
 using internal::OptionalBitBlockCounter;
 using internal::OptionalBitIndexer;
 
@@ -68,32 +69,50 @@ using TakeState = OptionsWrapper<TakeOptions>;
 // ----------------------------------------------------------------------
 // DropNull Implementation
 
+/// \pre values.MayHaveNulls()
+/// \pre may_have_validity_bitmap(values_type_id)
 std::shared_ptr<arrow::BooleanArray> MakeDropNullFilter(const Array& values) {
+ DCHECK(may_have_validity_bitmap(values.type()->id()));
  auto& bitmap_buffer = values.null_bitmap();
+ DCHECK(bitmap_buffer);
  return std::make_shared<BooleanArray>(values.length(), bitmap_buffer, nullptr, 0,
  values.offset());
 }
 
 Result<Datum> DropNullArray(const std::shared_ptr<Array>& values, ExecContext* ctx) {
- if (values->null_count() == 0) {
+ if (!values->data()->MayHaveLogicalNulls()) {
  return values;
  }
- if (values->null_count() == values->length()) {
- return MakeEmptyArray(values->type(), ctx->memory_pool());
- }
- if (values->type()->id() == Type::type::NA) {
+ if (ARROW_PREDICT_FALSE(values->type()->id() == Type::type::NA)) {
  return std::make_shared<NullArray>(0);
  }
+ if (values->device_type() == DeviceAllocationType::kCPU) {
+ const auto null_count = values->null_count();
+ if (null_count == 0) {
+ return values;
+ }
+ if (null_count == values->length()) {
+ return MakeEmptyArray(values->type(), ctx->memory_pool());
+ }
+ }
+ // TODO(GH-43851): Handle logical nulls in REE and Unions arrays correctly
+ if (!may_have_validity_bitmap(values->type()->id())) {
+ return values;
+ }
  auto drop_null_filter = Datum{MakeDropNullFilter(*values)};
  return Filter(values, drop_null_filter, FilterOptions::Defaults(), ctx);
 }
 
 Result<Datum> DropNullChunkedArray(const std::shared_ptr<ChunkedArray>& values,
  ExecContext* ctx) {
- if (values->null_count() == 0) {
+ // The null count in a chunked array is calculated by summing the null counts
+ // of all chunks, so by this point it's safe to access it without triggering
+ // access to potentially non-CPU data.
+ auto null_count = values->null_count();
+ if (null_count == 0) {
  return values;
  }
- if (values->null_count() == values->length()) {
+ if (null_count == values->length()) {
  return ChunkedArray::MakeEmpty(values->type(), ctx->memory_pool());
  }
  std::vector<std::shared_ptr<Array>> new_chunks;