diff --git a/c_glib/arrow-glib/input-stream.cpp b/c_glib/arrow-glib/input-stream.cpp index b931cf82506..d81e4a32d5d 100644 --- a/c_glib/arrow-glib/input-stream.cpp +++ b/c_glib/arrow-glib/input-stream.cpp @@ -23,12 +23,14 @@ #include #include +#include #include #include #include #include #include +#include G_BEGIN_DECLS @@ -253,6 +255,36 @@ garrow_seekable_input_stream_read_at(GArrowSeekableInputStream *input_stream, } } +/** + * garrow_seekable_input_stream_read_tensor: + * @input_stream: A #GArrowSeekableInputStream. + * @position: The read start position. + * @error: (nullable): Return location for a #GError or %NULL. + * + * Returns: (transfer full) (nullable): + * #GArrowTensor on success, %NULL on error. + * + * Since: 0.4.0 + */ +GArrowTensor * +garrow_seekable_input_stream_read_tensor(GArrowSeekableInputStream *input_stream, + gint64 position, + GError **error) +{ + auto arrow_random_access_file = + garrow_seekable_input_stream_get_raw(input_stream); + + std::shared_ptr arrow_tensor; + auto status = arrow::ipc::ReadTensor(position, + arrow_random_access_file.get(), + &arrow_tensor); + if (garrow_error_check(error, status, "[seekable-input-stream][read-tensor]")) { + return garrow_tensor_new_raw(&arrow_tensor); + } else { + return NULL; + } +} + G_DEFINE_TYPE(GArrowBufferInputStream, \ garrow_buffer_input_stream, \ diff --git a/c_glib/arrow-glib/input-stream.h b/c_glib/arrow-glib/input-stream.h index 51188286376..8a4d362189f 100644 --- a/c_glib/arrow-glib/input-stream.h +++ b/c_glib/arrow-glib/input-stream.h @@ -20,6 +20,7 @@ #pragma once #include +#include G_BEGIN_DECLS @@ -123,6 +124,9 @@ GArrowBuffer *garrow_seekable_input_stream_read_at(GArrowSeekableInputStream *in gint64 position, gint64 n_bytes, GError **error); +GArrowTensor *garrow_seekable_input_stream_read_tensor(GArrowSeekableInputStream *input_stream, + gint64 position, + GError **error); #define GARROW_TYPE_BUFFER_INPUT_STREAM \ diff --git a/c_glib/arrow-glib/output-stream.cpp b/c_glib/arrow-glib/output-stream.cpp index 48c48b8fdc3..ffb6fecb5c1 100644 --- a/c_glib/arrow-glib/output-stream.cpp +++ b/c_glib/arrow-glib/output-stream.cpp @@ -21,13 +21,14 @@ # include #endif -#include #include +#include #include #include #include #include +#include #include G_BEGIN_DECLS @@ -168,6 +169,36 @@ garrow_output_stream_class_init(GArrowOutputStreamClass *klass) g_object_class_install_property(gobject_class, PROP_OUTPUT_STREAM, spec); } +/** + * garrow_output_stream_write_tensor: + * @stream: A #GArrowWriteable. + * @tensor: A #GArrowTensor to be written. + * @error: (nullable): Return location for a #GError or %NULL. + * + * Returns: The number of written bytes on success, -1 on error. + * + * Since: 0.4.0 + */ +gint64 +garrow_output_stream_write_tensor(GArrowOutputStream *stream, + GArrowTensor *tensor, + GError **error) +{ + auto arrow_tensor = garrow_tensor_get_raw(tensor); + auto arrow_stream = garrow_output_stream_get_raw(stream); + int32_t metadata_length; + int64_t body_length; + auto status = arrow::ipc::WriteTensor(*arrow_tensor, + arrow_stream.get(), + &metadata_length, + &body_length); + if (garrow_error_check(error, status, "[output-stream][write-tensor]")) { + return metadata_length + body_length; + } else { + return -1; + } +} + G_DEFINE_TYPE(GArrowFileOutputStream, garrow_file_output_stream, diff --git a/c_glib/arrow-glib/output-stream.h b/c_glib/arrow-glib/output-stream.h index 48b891c1973..c86597bdd26 100644 --- a/c_glib/arrow-glib/output-stream.h +++ b/c_glib/arrow-glib/output-stream.h @@ -22,6 +22,7 @@ #include #include +#include G_BEGIN_DECLS @@ -71,6 +72,10 @@ struct _GArrowOutputStreamClass GType garrow_output_stream_get_type(void) G_GNUC_CONST; +gint64 garrow_output_stream_write_tensor(GArrowOutputStream *stream, + GArrowTensor *tensor, + GError **error); + #define GARROW_TYPE_FILE_OUTPUT_STREAM \ (garrow_file_output_stream_get_type()) diff --git a/c_glib/test/test-tensor.rb b/c_glib/test/test-tensor.rb index 780c9f179e1..bdfc6576a3f 100644 --- a/c_glib/test/test-tensor.rb +++ b/c_glib/test/test-tensor.rb @@ -114,4 +114,13 @@ def test_column_major? not @tensor.column_major? end end + + def test_io + buffer = Arrow::PoolBuffer.new + output = Arrow::BufferOutputStream.new(buffer) + output.write_tensor(@tensor) + input = Arrow::BufferInputStream.new(buffer) + assert_equal(@tensor, + input.read_tensor(0)) + end end