From 4986f1dfbde140f445c60cdb600fc96ee1ec332e Mon Sep 17 00:00:00 2001 From: Sarah Gilmore Date: Tue, 3 Dec 2024 15:17:54 -0500 Subject: [PATCH] Add RecordBatchWriter proxy class --- .../io/ipc/proxy/record_batch_writer.cc | 76 +++++++++++++++++++ .../matlab/io/ipc/proxy/record_batch_writer.h | 42 ++++++++++ 2 files changed, 118 insertions(+) create mode 100644 matlab/src/cpp/arrow/matlab/io/ipc/proxy/record_batch_writer.cc create mode 100644 matlab/src/cpp/arrow/matlab/io/ipc/proxy/record_batch_writer.h 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 new file mode 100644 index 0000000000000..a29cf72ba2e52 --- /dev/null +++ b/matlab/src/cpp/arrow/matlab/io/ipc/proxy/record_batch_writer.cc @@ -0,0 +1,76 @@ +// 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/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/matlab/io/ipc/proxy/record_batch_writer.h" + +#include "libmexclass/proxy/ProxyManager.h" + +namespace arrow::matlab::io::ipc::proxy { + +RecordBatchWriter::RecordBatchWriter( + const std::shared_ptr 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 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(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 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(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 new file mode 100644 index 0000000000000..4ecaad41b05d6 --- /dev/null +++ b/matlab/src/cpp/arrow/matlab/io/ipc/proxy/record_batch_writer.h @@ -0,0 +1,42 @@ +// 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 "libmexclass/proxy/Proxy.h" + +namespace arrow::matlab::io::ipc::proxy { + +class RecordBatchWriter : public libmexclass::proxy::Proxy { + public: + RecordBatchWriter(std::shared_ptr writer); + + virtual ~RecordBatchWriter() = default; + + static libmexclass::proxy::MakeResult make( + const libmexclass::proxy::FunctionArguments& constructor_arguments); + + protected: + std::shared_ptr 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