pitrou commented on a change in pull request #7263:
URL: https://github.com/apache/arrow/pull/7263#discussion_r433813700
##########
File path: python/pyarrow/_cuda.pyx
##########
@@ -922,7 +922,8 @@ def read_message(object source, pool=None):
return result
-def read_record_batch(object buffer, object schema, pool=None):
+def read_record_batch(object buffer, object schema,
Review comment:
I would make optional arguments keyword-only, i.e.:
```python
def read_record_batch(object buffer, object schema, *,
DictionaryMemo dictionary_memo=None, pool=None):
```
##########
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:
Ok, but `device_batch` isn't compared anywhere... is something missing?
----------------------------------------------------------------
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:
[email protected]
