[GitHub] [arrow] alexbaden commented on a change in pull request #7263: ARROW-8927: [C++] Support dictionary memo in CUDA IPC ReadRecordBatch functions

[GitBox <gi...@apache.org>]

alexbaden commented on a change in pull request #7263:
URL: https://github.com/apache/arrow/pull/7263#discussion_r443164077



##########
File path: cpp/src/arrow/gpu/cuda_test.cc
##########
@@ -582,13 +584,41 @@ TEST_F(TestCudaArrowIpc, BasicWriteRead) {
 
   std::shared_ptr<RecordBatch> cpu_batch;
   io::BufferReader cpu_reader(std::move(host_buffer));
-  ipc::DictionaryMemo unused_memo;
   ASSERT_OK_AND_ASSIGN(
       cpu_batch, ipc::ReadRecordBatch(batch->schema(), &unused_memo,
                                       ipc::IpcReadOptions::Defaults(), &cpu_reader));
 
   CompareBatch(*batch, *cpu_batch);
 }
 
+TEST_F(TestCudaArrowIpc, DictionaryWriteRead) {
+  std::shared_ptr<RecordBatch> batch;
+  ASSERT_OK(ipc::test::MakeDictionary(&batch));
+
+  ipc::DictionaryMemo dictionary_memo;
+  ASSERT_OK(ipc::CollectDictionaries(*batch, &dictionary_memo));
+
+  std::shared_ptr<CudaBuffer> device_serialized;
+  ASSERT_OK_AND_ASSIGN(device_serialized, SerializeRecordBatch(*batch, context_.get()));
+
+  // Test that ReadRecordBatch works properly
+  std::shared_ptr<RecordBatch> device_batch;
+  ASSERT_OK_AND_ASSIGN(device_batch, ReadRecordBatch(batch->schema(), &dictionary_memo,
+                                                     device_serialized));
+
+  // Copy data from device, read batch, and compare
+  int64_t size = device_serialized->size();
+  ASSERT_OK_AND_ASSIGN(auto host_buffer, AllocateBuffer(size, pool_));
+  ASSERT_OK(device_serialized->CopyToHost(0, size, host_buffer->mutable_data()));
+
+  std::shared_ptr<RecordBatch> cpu_batch;
+  io::BufferReader cpu_reader(std::move(host_buffer));
+  ASSERT_OK_AND_ASSIGN(
+      cpu_batch, ipc::ReadRecordBatch(batch->schema(), &dictionary_memo,
+                                      ipc::IpcReadOptions::Defaults(), &cpu_reader));
+
+  CompareBatch(*batch, *cpu_batch);

Review comment:
       No, the goal of the test is to ensure the dictionary information can be retrieved and re-associated with the record batch if the batch is copied back to CPU, much like `TestCudaArrowIpc.BasicWriteRead`. A better test would be to actually run through the IPC routines, but the `TestCudaBuffer.ExportForIpc` test appears to be disabled. 




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org