Add test to check writing a sliced record doesn't crash

apache · Jun 21, 2024 · d98aa7e · d98aa7e
1 parent 777e755
commit d98aa7e
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/cpp/src/arrow/gpu/cuda_test.cc b/cpp/src/arrow/gpu/cuda_test.cc
@@ -710,6 +710,18 @@ TEST_F(TestCudaArrowIpc, WriteIpcList) {
             payload.body_buffers[0]->address());
 }
 
+TEST_F(TestCudaArrowIpc, WriteIpcSlicedRecord) {
+  std::shared_ptr<RecordBatch> batch;
+  ASSERT_OK(ipc::test::MakeListRecordBatch(&batch));
+
+  ASSERT_OK_AND_ASSIGN(auto batch_device, batch->CopyTo(mm_));
+  auto sliced_batch_device = batch_device->Slice(10);
+
+  ipc::IpcPayload payload;
+  ASSERT_NOT_OK(ipc::GetRecordBatchPayload(*sliced_batch_device,
+                                           ipc::IpcWriteOptions::Defaults(), &payload));
+}
+
 TEST_F(TestCudaArrowIpc, DictionaryWriteRead) {
   std::shared_ptr<RecordBatch> batch;
   ASSERT_OK(ipc::test::MakeDictionary(&batch));

diff --git a/cpp/src/arrow/ipc/writer.cc b/cpp/src/arrow/ipc/writer.cc
@@ -154,6 +154,10 @@ class RecordBatchSerializer {
       return Status::CapacityError("Cannot write arrays larger than 2^31 - 1 in length");
     }
 
+    if (arr.offset() != 0 && arr.device_type() != DeviceAllocationType::kCPU) {
+      return Status::NotImplemented("Cannot compute null count for non-cpu sliced array");
+    }
+
     // push back all common elements
     field_nodes_.push_back({arr.length(), arr.null_count(), 0});