GH-44922: [MATLAB] Add IPC RecordBatchStreamFileWriter MATLAB class (…

…#44925) ### Rationale for this change To enable support for the IPC Streaming format in the MATLAB interface, we should add a `RecordBatchStreamWriter` class. ### What changes are included in this PR? Added `arrow.io.ipc.RecordBatchStreamWriter` class. **Example Usage:** ```matlab >> city = ["Boston" "Seattle" "Denver" "Juno" "Anchorage" "Chicago"]'; >> daylength = duration(["15:17:01" "15:59:16" "14:59:14" "19:21:23" "14:18:24" "15:13:39"])'; >> matlabTable = table(city, daylength, VariableNames=["City", "DayLength"]); >> recordBatch1 = arrow.recordBatch(matlabTable(1:4, :)) >> recordBatch2 = arrow.recordBatch(matlabTable(5:end, :)); >> writer = arrow.io.ipc.RecordBatchStreamWriter("daylight.arrow", recordBatch1.Schema); >> writer.writeRecordBatch(recordBatch1); >> writer.writeRecordBatch(recordBatch2); >> writer.close(); ``` ### Are these changes tested? Yes. I Parameterized the test cases in `test/arrow/io/ipc/tRecordBatchWriter.m` to test the behavior of both `arrow.io.ipc.RecordBatchFileWriter` AND `arrow.io.ipc.RecordBatchStreamWriter`. ### Are there any user-facing changes? Yes. Users can now use `arrow.io.ipc.RecordBatchStreamWriter` to serialize `RecordBatch`es/`Table`s to the Arrow IPC Streaming format. ### Future Directions 1. #44923 * GitHub Issue: #44922 Lead-authored-by: Sarah Gilmore <[email protected]> Co-authored-by: Sarah Gilmore <[email protected]> Co-authored-by: Kevin Gurney <[email protected]> Co-authored-by: Sutou Kouhei <[email protected]> Signed-off-by: Sarah Gilmore <[email protected]>
apache · Dec 4, 2024 · 2b5a802 · 2b5a802
1 parent ded148c
commit 2b5a802
Show file tree

Hide file tree

Showing 12 changed files with 387 additions and 145 deletions.
diff --git a/matlab/src/cpp/arrow/matlab/io/ipc/proxy/record_batch_file_writer.cc b/matlab/src/cpp/arrow/matlab/io/ipc/proxy/record_batch_file_writer.cc
@@ -18,9 +18,7 @@
 #include "arrow/matlab/io/ipc/proxy/record_batch_file_writer.h"
 #include "arrow/io/file.h"
 #include "arrow/matlab/error/error.h"
-#include "arrow/matlab/tabular/proxy/record_batch.h"
 #include "arrow/matlab/tabular/proxy/schema.h"
-#include "arrow/matlab/tabular/proxy/table.h"
 #include "arrow/util/utf8.h"
 
 #include "libmexclass/proxy/ProxyManager.h"
@@ -29,11 +27,7 @@ namespace arrow::matlab::io::ipc::proxy {
 
 RecordBatchFileWriter::RecordBatchFileWriter(
     const std::shared_ptr<arrow::ipc::RecordBatchWriter> writer)
-    : writer{std::move(writer)} {
-  REGISTER_METHOD(RecordBatchFileWriter, close);
-  REGISTER_METHOD(RecordBatchFileWriter, writeRecordBatch);
-  REGISTER_METHOD(RecordBatchFileWriter, writeTable);
-}
+    : RecordBatchWriter(std::move(writer)) {}
 
 libmexclass::proxy::MakeResult RecordBatchFileWriter::make(
     const libmexclass::proxy::FunctionArguments& constructor_arguments) {
@@ -65,43 +59,4 @@ libmexclass::proxy::MakeResult RecordBatchFileWriter::make(
   return std::make_shared<RecordBatchFileWriterProxy>(std::move(writer));
 }
 
-void RecordBatchFileWriter::writeRecordBatch(
-    libmexclass::proxy::method::Context& context) {
-  namespace mda = ::matlab::data;
-  using RecordBatchProxy = ::arrow::matlab::tabular::proxy::RecordBatch;
-
-  mda::StructArray opts = context.inputs[0];
-  const mda::TypedArray<uint64_t> record_batch_proxy_id_mda =
-      opts[0]["RecordBatchProxyID"];
-  const uint64_t record_batch_proxy_id = record_batch_proxy_id_mda[0];
-
-  auto proxy = libmexclass::proxy::ProxyManager::getProxy(record_batch_proxy_id);
-  auto record_batch_proxy = std::static_pointer_cast<RecordBatchProxy>(proxy);
-  auto record_batch = record_batch_proxy->unwrap();
-
-  MATLAB_ERROR_IF_NOT_OK_WITH_CONTEXT(writer->WriteRecordBatch(*record_batch), context,
-                                      error::IPC_RECORD_BATCH_WRITE_FAILED);
-}
-
-void RecordBatchFileWriter::writeTable(libmexclass::proxy::method::Context& context) {
-  namespace mda = ::matlab::data;
-  using TableProxy = ::arrow::matlab::tabular::proxy::Table;
-
-  mda::StructArray opts = context.inputs[0];
-  const mda::TypedArray<uint64_t> table_proxy_id_mda = opts[0]["TableProxyID"];
-  const uint64_t table_proxy_id = table_proxy_id_mda[0];
-
-  auto proxy = libmexclass::proxy::ProxyManager::getProxy(table_proxy_id);
-  auto table_proxy = std::static_pointer_cast<TableProxy>(proxy);
-  auto table = table_proxy->unwrap();
-
-  MATLAB_ERROR_IF_NOT_OK_WITH_CONTEXT(writer->WriteTable(*table), context,
-                                      error::IPC_RECORD_BATCH_WRITE_FAILED);
-}
-
-void RecordBatchFileWriter::close(libmexclass::proxy::method::Context& context) {
-  MATLAB_ERROR_IF_NOT_OK_WITH_CONTEXT(writer->Close(), context,
-                                      error::IPC_RECORD_BATCH_WRITE_CLOSE_FAILED);
-}
-
-}  // namespace arrow::matlab::io::ipc::proxy
+}  // namespace arrow::matlab::io::ipc::proxy
diff --git a/matlab/src/cpp/arrow/matlab/io/ipc/proxy/record_batch_file_writer.h b/matlab/src/cpp/arrow/matlab/io/ipc/proxy/record_batch_file_writer.h
@@ -16,27 +16,20 @@
 // under the License.
 
 #include "arrow/ipc/writer.h"
+#include "arrow/matlab/io/ipc/proxy/record_batch_writer.h"
+
 #include "libmexclass/proxy/Proxy.h"
 
 namespace arrow::matlab::io::ipc::proxy {
 
-class RecordBatchFileWriter : public libmexclass::proxy::Proxy {
+class RecordBatchFileWriter : public RecordBatchWriter {
  public:
   RecordBatchFileWriter(std::shared_ptr<arrow::ipc::RecordBatchWriter> writer);
 
-  ~RecordBatchFileWriter() = default;
+  virtual ~RecordBatchFileWriter() = default;
 
   static libmexclass::proxy::MakeResult make(
       const libmexclass::proxy::FunctionArguments& constructor_arguments);
-
- protected:
-  std::shared_ptr<arrow::ipc::RecordBatchWriter> writer;
-
-  void writeRecordBatch(libmexclass::proxy::method::Context& context);
-
-  void writeTable(libmexclass::proxy::method::Context& context);
-
-  void close(libmexclass::proxy::method::Context& context);
 };
 
-}  // namespace arrow::matlab::io::ipc::proxy
+}  // namespace arrow::matlab::io::ipc::proxy
diff --git a/matlab/src/cpp/arrow/matlab/io/ipc/proxy/record_batch_stream_writer.cc b/matlab/src/cpp/arrow/matlab/io/ipc/proxy/record_batch_stream_writer.cc
@@ -0,0 +1,64 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/matlab/io/ipc/proxy/record_batch_stream_writer.h"
+#include "arrow/io/file.h"
+#include "arrow/ipc/writer.h"
+#include "arrow/matlab/error/error.h"
+#include "arrow/matlab/tabular/proxy/schema.h"
+#include "arrow/util/utf8.h"
+
+#include "libmexclass/proxy/ProxyManager.h"
+
+namespace arrow::matlab::io::ipc::proxy {
+
+RecordBatchStreamWriter::RecordBatchStreamWriter(
+    const std::shared_ptr<arrow::ipc::RecordBatchWriter> writer)
+    : RecordBatchWriter(std::move(writer)) {}
+
+libmexclass::proxy::MakeResult RecordBatchStreamWriter::make(
+    const libmexclass::proxy::FunctionArguments& constructor_arguments) {
+  namespace mda = ::matlab::data;
+  using RecordBatchStreamWriterProxy =
+      arrow::matlab::io::ipc::proxy::RecordBatchStreamWriter;
+  using SchemaProxy = arrow::matlab::tabular::proxy::Schema;
+
+  const mda::StructArray opts = constructor_arguments[0];
+
+  const mda::StringArray filename_mda = opts[0]["Filename"];
+  const auto filename_utf16 = std::u16string(filename_mda[0]);
+  MATLAB_ASSIGN_OR_ERROR(const auto filename_utf8,
+                         arrow::util::UTF16StringToUTF8(filename_utf16),
+                         error::UNICODE_CONVERSION_ERROR_ID);
+
+  const mda::TypedArray<uint64_t> arrow_schema_proxy_id_mda = opts[0]["SchemaProxyID"];
+  auto proxy = libmexclass::proxy::ProxyManager::getProxy(arrow_schema_proxy_id_mda[0]);
+  auto arrow_schema_proxy = std::static_pointer_cast<SchemaProxy>(proxy);
+  auto arrow_schema = arrow_schema_proxy->unwrap();
+
+  MATLAB_ASSIGN_OR_ERROR(auto output_stream,
+                         arrow::io::FileOutputStream::Open(filename_utf8),
+                         error::FAILED_TO_OPEN_FILE_FOR_WRITE);
+
+  MATLAB_ASSIGN_OR_ERROR(auto writer,
+                         arrow::ipc::MakeStreamWriter(output_stream, arrow_schema),
+                         "arrow:matlab:MakeFailed");
+
+  return std::make_shared<RecordBatchStreamWriterProxy>(std::move(writer));
+}
+
+}  // namespace arrow::matlab::io::ipc::proxy
diff --git a/matlab/src/cpp/arrow/matlab/io/ipc/proxy/record_batch_stream_writer.h b/matlab/src/cpp/arrow/matlab/io/ipc/proxy/record_batch_stream_writer.h
@@ -0,0 +1,35 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/ipc/writer.h"
+#include "arrow/matlab/io/ipc/proxy/record_batch_writer.h"
+
+#include "libmexclass/proxy/Proxy.h"
+
+namespace arrow::matlab::io::ipc::proxy {
+
+class RecordBatchStreamWriter : public RecordBatchWriter {
+ public:
+  RecordBatchStreamWriter(std::shared_ptr<arrow::ipc::RecordBatchWriter> writer);
+
+  virtual ~RecordBatchStreamWriter() = default;
+
+  static libmexclass::proxy::MakeResult make(
+      const libmexclass::proxy::FunctionArguments& constructor_arguments);
+};
+
+}  // namespace arrow::matlab::io::ipc::proxy
diff --git a/matlab/src/cpp/arrow/matlab/io/ipc/proxy/record_batch_writer.cc b/matlab/src/cpp/arrow/matlab/io/ipc/proxy/record_batch_writer.cc
@@ -0,0 +1,75 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/matlab/io/ipc/proxy/record_batch_writer.h"
+#include "arrow/io/file.h"
+#include "arrow/matlab/error/error.h"
+#include "arrow/matlab/tabular/proxy/record_batch.h"
+#include "arrow/matlab/tabular/proxy/schema.h"
+#include "arrow/matlab/tabular/proxy/table.h"
+
+#include "libmexclass/proxy/ProxyManager.h"
+
+namespace arrow::matlab::io::ipc::proxy {
+
+RecordBatchWriter::RecordBatchWriter(
+    const std::shared_ptr<arrow::ipc::RecordBatchWriter> writer)
+    : writer{std::move(writer)} {
+  REGISTER_METHOD(RecordBatchWriter, close);
+  REGISTER_METHOD(RecordBatchWriter, writeRecordBatch);
+  REGISTER_METHOD(RecordBatchWriter, writeTable);
+}
+
+void RecordBatchWriter::writeRecordBatch(libmexclass::proxy::method::Context& context) {
+  namespace mda = ::matlab::data;
+  using RecordBatchProxy = ::arrow::matlab::tabular::proxy::RecordBatch;
+
+  mda::StructArray opts = context.inputs[0];
+  const mda::TypedArray<uint64_t> record_batch_proxy_id_mda =
+      opts[0]["RecordBatchProxyID"];
+  const uint64_t record_batch_proxy_id = record_batch_proxy_id_mda[0];
+
+  auto proxy = libmexclass::proxy::ProxyManager::getProxy(record_batch_proxy_id);
+  auto record_batch_proxy = std::static_pointer_cast<RecordBatchProxy>(proxy);
+  auto record_batch = record_batch_proxy->unwrap();
+
+  MATLAB_ERROR_IF_NOT_OK_WITH_CONTEXT(writer->WriteRecordBatch(*record_batch), context,
+                                      error::IPC_RECORD_BATCH_WRITE_FAILED);
+}
+
+void RecordBatchWriter::writeTable(libmexclass::proxy::method::Context& context) {
+  namespace mda = ::matlab::data;
+  using TableProxy = ::arrow::matlab::tabular::proxy::Table;
+
+  mda::StructArray opts = context.inputs[0];
+  const mda::TypedArray<uint64_t> table_proxy_id_mda = opts[0]["TableProxyID"];
+  const uint64_t table_proxy_id = table_proxy_id_mda[0];
+
+  auto proxy = libmexclass::proxy::ProxyManager::getProxy(table_proxy_id);
+  auto table_proxy = std::static_pointer_cast<TableProxy>(proxy);
+  auto table = table_proxy->unwrap();
+
+  MATLAB_ERROR_IF_NOT_OK_WITH_CONTEXT(writer->WriteTable(*table), context,
+                                      error::IPC_RECORD_BATCH_WRITE_FAILED);
+}
+
+void RecordBatchWriter::close(libmexclass::proxy::method::Context& context) {
+  MATLAB_ERROR_IF_NOT_OK_WITH_CONTEXT(writer->Close(), context,
+                                      error::IPC_RECORD_BATCH_WRITE_CLOSE_FAILED);
+}
+
+}  // namespace arrow::matlab::io::ipc::proxy
diff --git a/matlab/src/cpp/arrow/matlab/io/ipc/proxy/record_batch_writer.h b/matlab/src/cpp/arrow/matlab/io/ipc/proxy/record_batch_writer.h
@@ -0,0 +1,41 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include "arrow/ipc/writer.h"
+#include "libmexclass/proxy/Proxy.h"
+
+namespace arrow::matlab::io::ipc::proxy {
+
+class RecordBatchWriter : public libmexclass::proxy::Proxy {
+ public:
+  RecordBatchWriter(std::shared_ptr<arrow::ipc::RecordBatchWriter> writer);
+
+  virtual ~RecordBatchWriter() = default;
+
+ protected:
+  std::shared_ptr<arrow::ipc::RecordBatchWriter> writer;
+
+  void writeRecordBatch(libmexclass::proxy::method::Context& context);
+
+  void writeTable(libmexclass::proxy::method::Context& context);
+
+  void close(libmexclass::proxy::method::Context& context);
+};
+
+}  // namespace arrow::matlab::io::ipc::proxy
diff --git a/matlab/src/cpp/arrow/matlab/proxy/factory.cc b/matlab/src/cpp/arrow/matlab/proxy/factory.cc
@@ -36,6 +36,7 @@
 #include "arrow/matlab/io/feather/proxy/writer.h"
 #include "arrow/matlab/io/ipc/proxy/record_batch_file_reader.h"
 #include "arrow/matlab/io/ipc/proxy/record_batch_file_writer.h"
+#include "arrow/matlab/io/ipc/proxy/record_batch_stream_writer.h"
 #include "arrow/matlab/tabular/proxy/record_batch.h"
 #include "arrow/matlab/tabular/proxy/schema.h"
 #include "arrow/matlab/tabular/proxy/table.h"
@@ -111,6 +112,8 @@ libmexclass::proxy::MakeResult Factory::make_proxy(
   REGISTER_PROXY(arrow.c.proxy.RecordBatchImporter , arrow::matlab::c::proxy::RecordBatchImporter);
   REGISTER_PROXY(arrow.io.ipc.proxy.RecordBatchFileReader , arrow::matlab::io::ipc::proxy::RecordBatchFileReader);
   REGISTER_PROXY(arrow.io.ipc.proxy.RecordBatchFileWriter , arrow::matlab::io::ipc::proxy::RecordBatchFileWriter);
+  REGISTER_PROXY(arrow.io.ipc.proxy.RecordBatchStreamWriter , arrow::matlab::io::ipc::proxy::RecordBatchStreamWriter);
+
   // clang-format on
 
   return libmexclass::error::Error{error::UNKNOWN_PROXY_ERROR_ID,