GH-40078: [C++] Import/Export ArrowDeviceArrayStream

cpp/src/arrow/array/array_base.h

@@ -224,6 +224,14 @@ class ARROW_EXPORT Array {
  /// \return Status
  Status ValidateFull() const;
 
+ /// \brief Return the device_type that this array's data is allocated on
+ ///
+ /// This just delegates to calling device_type on the underlying ArrayData
+ /// object which backs this Array.
+ ///
+ /// \return DeviceAllocationType
+ DeviceAllocationType device_type() const { return data_->device_type(); }
+
  protected:
  Array() = default;
  ARROW_DEFAULT_MOVE_AND_ASSIGN(Array);

cpp/src/arrow/array/array_test.cc

@@ -478,6 +478,7 @@ TEST_F(TestArray, TestMakeArrayOfNull) {
  ASSERT_EQ(array->type(), type);
  ASSERT_OK(array->ValidateFull());
  ASSERT_EQ(array->length(), length);
+ ASSERT_EQ(array->device_type(), DeviceAllocationType::kCPU);
  if (is_union(type->id())) {
  ASSERT_EQ(array->null_count(), 0);
  ASSERT_EQ(array->ComputeLogicalNullCount(), length);
@@ -719,6 +720,7 @@ TEST_F(TestArray, TestMakeArrayFromScalar) {
  ASSERT_OK(array->ValidateFull());
  ASSERT_EQ(array->length(), length);
  ASSERT_EQ(array->null_count(), 0);
+ ASSERT_EQ(array->device_type(), DeviceAllocationType::kCPU);
 
  // test case for ARROW-13321
  for (int64_t i : {int64_t{0}, length / 2, length - 1}) {
@@ -744,6 +746,7 @@ TEST_F(TestArray, TestMakeArrayFromScalarSliced) {
  auto sliced = array->Slice(1, 4);
  ASSERT_EQ(sliced->length(), 4);
  ASSERT_EQ(sliced->null_count(), 0);
+ ASSERT_EQ(array->device_type(), DeviceAllocationType::kCPU);
  ARROW_EXPECT_OK(sliced->ValidateFull());
  }
 }
@@ -758,6 +761,7 @@ TEST_F(TestArray, TestMakeArrayFromDictionaryScalar) {
  ASSERT_OK(array->ValidateFull());
  ASSERT_EQ(array->length(), 4);
  ASSERT_EQ(array->null_count(), 0);
+ ASSERT_EQ(array->device_type(), DeviceAllocationType::kCPU);
 
  for (int i = 0; i < 4; i++) {
  ASSERT_OK_AND_ASSIGN(auto item, array->GetScalar(i));
@@ -797,6 +801,7 @@ TEST_F(TestArray, TestMakeEmptyArray) {
  ASSERT_OK_AND_ASSIGN(auto array, MakeEmptyArray(type));
  ASSERT_OK(array->ValidateFull());
  ASSERT_EQ(array->length(), 0);
+
  CheckSpanRoundTrip(*array);
  }
 }

cpp/src/arrow/array/data.cc

@@ -224,6 +224,42 @@ int64_t ArrayData::ComputeLogicalNullCount() const {
  return ArraySpan(*this).ComputeLogicalNullCount();
 }
 
+DeviceAllocationType ArrayData::device_type() const {
+ // we're using 0 as a sentinel value for NOT YET ASSIGNED
+ // there is explicitly no constant DeviceAllocationType to represent
+ // the "UNASSIGNED" case as it is invalid for data to not have an
+ // assigned device type. If it's still 0 at the end, then we return
+ // CPU as the allocation device type
+ int type = 0;
+ for (const auto& buf : buffers) {
+ if (!buf) continue;
+ if (type == 0) {
+ type = static_cast<int>(buf->device_type());
+ } else {
+ DCHECK_EQ(type, static_cast<int>(buf->device_type()));
+ }
+ }
+
+ for (const auto& child : child_data) {
+ if (!child) continue;
+ if (type == 0) {
+ type = static_cast<int>(child->device_type());
+ } else {
+ DCHECK_EQ(type, static_cast<int>(child->device_type()));
+ }
+ }
+
+ if (dictionary) {
+ if (type == 0) {
+ type = static_cast<int>(dictionary->device_type());
+ } else {
+ DCHECK_EQ(type, static_cast<int>(dictionary->device_type()));
+ }
+ }
+
+ return type == 0 ? DeviceAllocationType::kCPU : static_cast<DeviceAllocationType>(type);
+}
+
 // ----------------------------------------------------------------------
 // Methods for ArraySpan
 

cpp/src/arrow/array/data.h

@@ -101,6 +101,11 @@ struct ARROW_EXPORT ArrayData {
  int64_t null_count = kUnknownNullCount, int64_t offset = 0)
  : ArrayData(std::move(type), length, null_count, offset) {
  this->buffers = std::move(buffers);
+#ifndef NDEBUG
+ // in debug mode, call the `device_type` function to trigger
+ // the DCHECKs that validate all the buffers are on the same device
+ ARROW_UNUSED(this->device_type());
+#endif
  }
 
  ArrayData(std::shared_ptr<DataType> type, int64_t length,
@@ -110,6 +115,12 @@ struct ARROW_EXPORT ArrayData {
  : ArrayData(std::move(type), length, null_count, offset) {
  this->buffers = std::move(buffers);
  this->child_data = std::move(child_data);
+#ifndef NDEBUG
+ // in debug mode, call the `device_type` function to trigger
+ // the DCHECKs that validate all the buffers (including children)
+ // are on the same device
+ ARROW_UNUSED(this->device_type());
+#endif
  }
 
  static std::shared_ptr<ArrayData> Make(std::shared_ptr<DataType> type, int64_t length,
@@ -358,6 +369,16 @@ struct ARROW_EXPORT ArrayData {
  /// \see GetNullCount
  int64_t ComputeLogicalNullCount() const;
 
+ /// \brief Returns the device_type of the underlying buffers and children
+ ///
+ /// If there are no buffers in this ArrayData object, it just returns
+ /// DeviceAllocationType::kCPU as a default. We also assume that all buffers
+ /// should be allocated on the same device type and perform DCHECKs to confirm
+ /// this in debug mode.
+ ///
+ /// \return DeviceAllocationType
+ DeviceAllocationType device_type() const;
+
  std::shared_ptr<DataType> type;
  int64_t length = 0;
  mutable std::atomic<int64_t> null_count{0};

cpp/src/arrow/array/util.cc

@@ -548,7 +548,7 @@ class NullArrayFactory {
  }
 
  Status Visit(const StructType& type) {
- for (int i = 0; i < type_->num_fields(); ++i) {
+ for (int i = 0; i < type.num_fields(); ++i) {
  ARROW_ASSIGN_OR_RAISE(out_->child_data[i], CreateChild(type, i, length_));
  }
  return Status::OK();