Implement ChunkedArray::IsCpu() instead

cpp/src/arrow/chunked_array.cc

@@ -27,6 +27,7 @@
 #include "arrow/array/array_nested.h"
 #include "arrow/array/util.h"
 #include "arrow/array/validate.h"
+#include "arrow/device.h"
 #include "arrow/pretty_print.h"
 #include "arrow/status.h"
 #include "arrow/type.h"
@@ -286,6 +287,15 @@ Status ChunkedArray::ValidateFull() const {
   return ValidateChunks(chunks_, /*full_validation=*/true);
 }
 
+bool ChunkedArray::IsCpu() const {
+  for (const auto& chunk : chunks_) {
+    if (chunk->device_type() != DeviceAllocationType::kCPU) {
+      return false;
+    }
+  }
+  return true;
+}
+
 namespace internal {
 
 bool MultipleChunkIterator::Next(std::shared_ptr<Array>* next_left,

cpp/src/arrow/chunked_array.h

@@ -182,6 +182,9 @@ class ARROW_EXPORT ChunkedArray {
   /// \return Status
   Status ValidateFull() const;
 
+  /// \brief Determine if all chunks are located on the CPU
+  bool IsCpu() const;
+
  protected:
   ArrayVector chunks_;
   std::shared_ptr<DataType> type_;

python/pyarrow/includes/libarrow.pxd

@@ -983,6 +983,8 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
 
         CResult[vector[shared_ptr[CChunkedArray]]] Flatten(CMemoryPool* pool)
 
+        c_bool IsCpu() const
+
         CStatus Validate() const
         CStatus ValidateFull() const
 

python/pyarrow/table.pxi

@@ -1407,6 +1407,15 @@ cdef class ChunkedArray(_PandasConvertible):
         self.init(c_chunked_array)
         return self
 
+    def is_cpu(self):
+        """
+        Whether all chunks in the ChunkedArray are CPU-accessible.
+        """
+        cdef c_bool result
+        with nogil:
+            result = self.chunked_array.IsCpu()
+        return result
+
 
 def chunked_array(arrays, type=None):
     """

python/pyarrow/tests/test_table.py

@@ -3357,3 +3357,41 @@ def test_invalid_non_join_column():
     with pytest.raises(pa.lib.ArrowInvalid) as excinfo:
         t2.join(t1, 'id', join_type='inner')
     assert exp_error_msg in str(excinfo.value)
+
+
+
+@pytest.fixture
+def cuda_context():
+    cuda = pytest.importorskip("pyarrow.cuda")
+    return cuda.Context(0)
+
+
+@pytest.fixture
+def schema():
+    return pa.schema([pa.field('c0', pa.int16()), pa.field('c1', pa.int32())])
+
+
+@pytest.fixture
+def cpu_arrays():
+    return [pa.array([1, 2, 3, 4, 5], pa.int32()),
+            pa.array([-10, -5, 0, 1, 10], pa.int32())]
+
+
+@pytest.fixture
+def cuda_arrays(cuda_context, cpu_arrays):
+    return [arr.copy_to(cuda_context.memory_manager) for arr in cpu_arrays]
+
+
+@pytest.fixture
+def cuda_chunked_arrays(cuda_arrays):
+    return pa.chunked_array(cuda_arrays)
+
+
+@pytest.fixture
+def cpu_chunked_arrays(cpu_arrays):
+    return pa.chunked_array(cpu_arrays)
+
+
+def test_chunked_array_non_cpu(cuda_context, cpu_chunked_arrays, cuda_chunked_arrays):
+    assert cpu_chunked_arrays.is_cpu() is True
+    assert cuda_chunked_arrays.is_cpu() is False