From 585230f854a90fb005410787d475fef524e458d3 Mon Sep 17 00:00:00 2001
From: Sarah Gilmore <sgilmore@mathworks.com>
Date: Wed, 12 Jun 2024 17:09:03 -0400
Subject: [PATCH] Add readRecordBatch method

---
 .../io/ipc/proxy/record_batch_file_reader.cc  | 27 ++++++++++++++++++-
 .../io/ipc/proxy/record_batch_file_reader.h   |  2 ++
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/matlab/src/cpp/arrow/matlab/io/ipc/proxy/record_batch_file_reader.cc b/matlab/src/cpp/arrow/matlab/io/ipc/proxy/record_batch_file_reader.cc
index b2e66eb03d483..f435a3c1edf5b 100644
--- a/matlab/src/cpp/arrow/matlab/io/ipc/proxy/record_batch_file_reader.cc
+++ b/matlab/src/cpp/arrow/matlab/io/ipc/proxy/record_batch_file_reader.cc
@@ -16,17 +16,21 @@
 // under the License.
 
 #include "arrow/io/file.h"
-#include "arrow/io/interfaces.h"
 #include "arrow/ipc/reader.h"
 #include "arrow/matlab/error/error.h"
 #include "arrow/matlab/io/ipc/proxy/record_batch_file_reader.h"
+#include "arrow/matlab/tabular/proxy/record_batch.h"
+#include "arrow/matlab/tabular/proxy/record_batch.h"
 #include "arrow/util/utf8.h"
 
+#include "libmexclass/proxy/ProxyManager.h"
 
 namespace arrow::matlab::io::ipc::proxy {
 
 RecordBatchFileReader::RecordBatchFileReader(std::shared_ptr<arrow::ipc::RecordBatchFileReader> reader) : reader{std::move(reader)} {
   REGISTER_METHOD(RecordBatchFileReader, getNumRecordBatches);
+  REGISTER_METHOD(RecordBatchFileReader, readRecordBatch);
+
 }
 
 libmexclass::proxy::MakeResult RecordBatchFileReader::make(const libmexclass::proxy::FunctionArguments& constructor_arguments) {
@@ -59,6 +63,27 @@ void RecordBatchFileReader::getNumRecordBatches(libmexclass::proxy::method::Cont
     context.outputs[0] = factory.createScalar(num_batches);
 }
 
+void RecordBatchFileReader::readRecordBatch(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<int32_t> matlab_index_mda = opts[0]["Index"];
+    const auto matlab_index = matlab_index_mda[0];
+
+    const auto arrow_index = matlab_index - 1;
+
+    MATLAB_ASSIGN_OR_ERROR_WITH_CONTEXT(const auto record_batch, reader->ReadRecordBatch(arrow_index),
+                                       context, "arrow:matlab:ipc:readfailed");
+
+  auto record_batch_proxy = std::make_shared<RecordBatchProxy>(std::move(record_batch));
+  const auto record_batch_proxy_id = libmexclass::proxy::ProxyManager::manageProxy(record_batch_proxy);
+
+  mda::ArrayFactory factory;
+  const auto record_batch_proxyy_id_mda = factory.createScalar(record_batch_proxy_id);
+  context.outputs[0] = record_batch_proxyy_id_mda;
+}
+
 
 
 } // namespace arrow::matlab::io::ipc::proxy
\ No newline at end of file
diff --git a/matlab/src/cpp/arrow/matlab/io/ipc/proxy/record_batch_file_reader.h b/matlab/src/cpp/arrow/matlab/io/ipc/proxy/record_batch_file_reader.h
index c0bf811648937..69e5cfef6c532 100644
--- a/matlab/src/cpp/arrow/matlab/io/ipc/proxy/record_batch_file_reader.h
+++ b/matlab/src/cpp/arrow/matlab/io/ipc/proxy/record_batch_file_reader.h
@@ -33,6 +33,8 @@ namespace arrow::matlab::io::ipc::proxy {
 
     void getNumRecordBatches(libmexclass::proxy::method::Context& context);
 
+    void readRecordBatch(libmexclass::proxy::method::Context& context);
+
     std::shared_ptr<arrow::ipc::RecordBatchFileReader> reader;
   
   };