This is an automated email from the ASF dual-hosted git repository.
kou pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/master by this push:
new fee9f1d ARROW-9234: [GLib][CUDA] Add support for dictionary memo on
reading record batch from buffer
fee9f1d is described below
commit fee9f1df7b21f372bd30d59ec276ffd289645170
Author: Sutou Kouhei <[email protected]>
AuthorDate: Sat Jun 27 11:12:04 2020 +0900
ARROW-9234: [GLib][CUDA] Add support for dictionary memo on reading record
batch from buffer
This is a follow up task for https://github.com/apache/arrow/pull/7263 .
Closes #7553 from kou/glib-cuda-read-record-batch-dictionary-memo
Authored-by: Sutou Kouhei <[email protected]>
Signed-off-by: Sutou Kouhei <[email protected]>
---
c_glib/arrow-cuda-glib/cuda.cpp | 39 +++++++++++++++++++++++++++++++--------
c_glib/arrow-cuda-glib/cuda.h | 1 +
c_glib/test/test-cuda.rb | 37 +++++++++++++++++++++++++------------
3 files changed, 57 insertions(+), 20 deletions(-)
diff --git a/c_glib/arrow-cuda-glib/cuda.cpp b/c_glib/arrow-cuda-glib/cuda.cpp
index 37f12e8..e3a5644 100644
--- a/c_glib/arrow-cuda-glib/cuda.cpp
+++ b/c_glib/arrow-cuda-glib/cuda.cpp
@@ -24,6 +24,7 @@
#include <arrow-glib/buffer.hpp>
#include <arrow-glib/error.hpp>
#include <arrow-glib/input-stream.hpp>
+#include <arrow-glib/ipc-options.hpp>
#include <arrow-glib/output-stream.hpp>
#include <arrow-glib/readable.hpp>
#include <arrow-glib/record-batch.hpp>
@@ -448,6 +449,7 @@ garrow_cuda_buffer_get_context(GArrowCUDABuffer *buffer)
* garrow_cuda_buffer_read_record_batch:
* @buffer: A #GArrowCUDABuffer.
* @schema: A #GArrowSchema for record batch.
+ * @options: (nullable): A #GArrowReadOptions.
* @error: (nullable): Return location for a #GError or %NULL.
*
* Returns: (transfer full): A newly created #GArrowRecordBatch on
@@ -458,19 +460,40 @@ garrow_cuda_buffer_get_context(GArrowCUDABuffer *buffer)
GArrowRecordBatch *
garrow_cuda_buffer_read_record_batch(GArrowCUDABuffer *buffer,
GArrowSchema *schema,
+ GArrowReadOptions *options,
GError **error)
{
auto arrow_buffer = garrow_cuda_buffer_get_raw(buffer);
auto arrow_schema = garrow_schema_get_raw(schema);
- auto pool = arrow::default_memory_pool();
- auto arrow_record_batch = arrow::cuda::ReadRecordBatch(arrow_schema,
- arrow_buffer,
- pool);
- if (garrow::check(error, arrow_record_batch,
- "[cuda][buffer][read-record-batch]")) {
- return garrow_record_batch_new_raw(&(*arrow_record_batch));
+
+ if (options) {
+ auto arrow_options = garrow_read_options_get_raw(options);
+ auto arrow_dictionary_memo =
+ garrow_read_options_get_dictionary_memo_raw(options);
+ auto arrow_record_batch =
+ arrow::cuda::ReadRecordBatch(arrow_schema,
+ arrow_dictionary_memo,
+ arrow_buffer,
+ arrow_options->memory_pool);
+ if (garrow::check(error, arrow_record_batch,
+ "[cuda][buffer][read-record-batch]")) {
+ return garrow_record_batch_new_raw(&(*arrow_record_batch));
+ } else {
+ return NULL;
+ }
} else {
- return NULL;
+ auto arrow_pool = arrow::default_memory_pool();
+ auto arrow_record_batch =
+ arrow::cuda::ReadRecordBatch(arrow_schema,
+ nullptr,
+ arrow_buffer,
+ arrow_pool);
+ if (garrow::check(error, arrow_record_batch,
+ "[cuda][buffer][read-record-batch]")) {
+ return garrow_record_batch_new_raw(&(*arrow_record_batch));
+ } else {
+ return NULL;
+ }
}
}
diff --git a/c_glib/arrow-cuda-glib/cuda.h b/c_glib/arrow-cuda-glib/cuda.h
index 6cdef99..3c98dae 100644
--- a/c_glib/arrow-cuda-glib/cuda.h
+++ b/c_glib/arrow-cuda-glib/cuda.h
@@ -147,6 +147,7 @@ garrow_cuda_buffer_get_context(GArrowCUDABuffer *buffer);
GArrowRecordBatch *
garrow_cuda_buffer_read_record_batch(GArrowCUDABuffer *buffer,
GArrowSchema *schema,
+ GArrowReadOptions *options,
GError **error);
diff --git a/c_glib/test/test-cuda.rb b/c_glib/test/test-cuda.rb
index d24bab6..c9b3492 100644
--- a/c_glib/test/test-cuda.rb
+++ b/c_glib/test/test-cuda.rb
@@ -79,18 +79,31 @@ end
@buffer.context.allocated_size)
end
- def test_record_batch
- field = Arrow::Field.new("enabled", Arrow::BooleanDataType.new)
- schema = Arrow::Schema.new([field])
- columns = [
- build_boolean_array([true]),
- ]
- cpu_record_batch = Arrow::RecordBatch.new(schema, 1, columns)
-
- buffer = ArrowCUDA::Buffer.new(@context, cpu_record_batch)
- gpu_record_batch = buffer.read_record_batch(schema)
- assert_equal(cpu_record_batch.n_rows,
- gpu_record_batch.n_rows)
+ sub_test_case("#read_record_batch") do
+ def setup
+ super
+ @field = Arrow::Field.new("enabled", Arrow::BooleanDataType.new)
+ @schema = Arrow::Schema.new([@field])
+ @columns = [
+ build_boolean_array([true]),
+ ]
+ @cpu_record_batch = Arrow::RecordBatch.new(@schema, 1, @columns)
+
+ @buffer = ArrowCUDA::Buffer.new(@context, @cpu_record_batch)
+ end
+
+ def test_default
+ gpu_record_batch = @buffer.read_record_batch(@schema)
+ assert_equal(@cpu_record_batch.n_rows,
+ gpu_record_batch.n_rows)
+ end
+
+ def test_options
+ options = Arrow::ReadOptions.new
+ gpu_record_batch = @buffer.read_record_batch(@schema, options)
+ assert_equal(@cpu_record_batch.n_rows,
+ gpu_record_batch.n_rows)
+ end
end
end
