Implement RecordBatchFileWriter::writeTable(...)

mathworks · Jun 13, 2024 · 808de49 · 808de49
1 parent bab0f9a
commit 808de49
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 0 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
@@ -20,6 +20,7 @@
 #include "arrow/matlab/io/ipc/proxy/record_batch_file_writer.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"
@@ -30,6 +31,7 @@ namespace arrow::matlab::io::ipc::proxy {
 RecordBatchFileWriter::RecordBatchFileWriter(const std::shared_ptr<arrow::ipc::RecordBatchWriter> writer)
   : writer{std::move(writer)} {
     REGISTER_METHOD(RecordBatchFileWriter, writeRecordBatch);
+    REGISTER_METHOD(RecordBatchFileWriter, writeTable);
   }
 
 libmexclass::proxy::MakeResult RecordBatchFileWriter::make(const libmexclass::proxy::FunctionArguments& constructor_arguments) {
@@ -77,4 +79,20 @@ void RecordBatchFileWriter::writeRecordBatch(libmexclass::proxy::method::Context
                                       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);
+}
+
 } // 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
@@ -33,6 +33,8 @@ class RecordBatchFileWriter : public libmexclass::proxy::Proxy {
 
     void writeRecordBatch(libmexclass::proxy::method::Context& context);
 
+    void writeTable(libmexclass::proxy::method::Context& context);
+
 };
 
 } // namespace arrow::matlab::io::ipc::proxy