From 7bfa195101f5261e357358e22a51ce04c457cf92 Mon Sep 17 00:00:00 2001
From: Sarah Gilmore <sgilmore@mathworks.com>
Date: Thu, 31 Aug 2023 10:39:34 -0400
Subject: [PATCH] Update getChunk() and getType() proxy methods to return
 struct arrays

---
 .../arrow/matlab/array/proxy/chunked_array.cc | 25 +++++++++++--------
 .../src/matlab/+arrow/+array/ChunkedArray.m   | 18 ++++++-------
 2 files changed, 22 insertions(+), 21 deletions(-)

diff --git a/matlab/src/cpp/arrow/matlab/array/proxy/chunked_array.cc b/matlab/src/cpp/arrow/matlab/array/proxy/chunked_array.cc
index f31a984d0b70e..82f314cf9cad2 100644
--- a/matlab/src/cpp/arrow/matlab/array/proxy/chunked_array.cc
+++ b/matlab/src/cpp/arrow/matlab/array/proxy/chunked_array.cc
@@ -24,6 +24,7 @@
 #include "arrow/matlab/array/proxy/wrap.h"
 
 #include "libmexclass/proxy/ProxyManager.h"
+#include <iostream>
 
 namespace arrow::matlab::array::proxy {
 
@@ -67,6 +68,7 @@ namespace arrow::matlab::array::proxy {
             auto array_proxy = std::static_pointer_cast<proxy::Array>(proxy);
             auto array = array_proxy->unwrap();
             arrays.push_back(array);
+            std::cout << array->type_id() << std::endl;
         }
 
         MATLAB_ASSIGN_OR_ERROR(auto chunked_array, 
@@ -125,14 +127,12 @@ namespace arrow::matlab::array::proxy {
         
         
         const auto array_proxy_id = libmexclass::proxy::ProxyManager::manageProxy(array_proxy);
-        const auto array_proxy_id_mda = factory.createScalar(array_proxy_id);
-        const auto array_type_id_mda = factory.createScalar(static_cast<int32_t>(array->type_id()));
-        
-        context.outputs[0] = array_proxy_id_mda;
-        context.outputs[1] = array_type_id_mda;
+        const auto type_id = static_cast<int64_t>(array->type_id());
 
-        auto length_mda = factory.createScalar(chunked_array->num_chunks());
-        context.outputs[0] = length_mda;
+        mda::StructArray output = factory.createStructArray({1, 1}, {"ProxyID", "TypeID"});
+        output[0]["ProxyID"] = factory.createScalar(array_proxy_id);
+        output[0]["TypeID"] = factory.createScalar(type_id);
+        context.outputs[0] = output;
     }
 
 
@@ -146,10 +146,13 @@ namespace arrow::matlab::array::proxy {
                                             context,
                                             error::ARRAY_FAILED_TO_CREATE_TYPE_PROXY);
 
-        auto type_id = type_proxy->unwrap()->id();
-        auto proxy_id = libmexclass::proxy::ProxyManager::manageProxy(type_proxy);
 
-        context.outputs[0] = factory.createScalar(proxy_id);
-        context.outputs[1] = factory.createScalar(static_cast<int64_t>(type_id));
+        const auto proxy_id = libmexclass::proxy::ProxyManager::manageProxy(type_proxy);
+        const auto type_id = static_cast<int32_t>(type_proxy->unwrap()->id());
+
+        mda::StructArray output = factory.createStructArray({1, 1}, {"ProxyID", "TypeID"});
+        output[0]["ProxyID"] = factory.createScalar(proxy_id);
+        output[0]["TypeID"] = factory.createScalar(type_id);
+        context.outputs[0] = output;
     }
 }
diff --git a/matlab/src/matlab/+arrow/+array/ChunkedArray.m b/matlab/src/matlab/+arrow/+array/ChunkedArray.m
index 5f39946798d38..518026c0ab12a 100644
--- a/matlab/src/matlab/+arrow/+array/ChunkedArray.m
+++ b/matlab/src/matlab/+arrow/+array/ChunkedArray.m
@@ -17,8 +17,6 @@
 
 classdef ChunkedArray < matlab.mixin.CustomDisplay & ...
                         matlab.mixin.Scalar
-    %UNTITLED Summary of this class goes here
-    %   Detailed explanation goes here
 
     properties(Hidden, SetAccess=private, GetAccess=public)
         Proxy
@@ -48,17 +46,17 @@
         end
 
         function type = get.Type(obj)
-            [typeID, proxyID] = obj.Proxy.getType();
-            traits = arrow.type.traits.traits(arrow.type.ID(typeID));
-            proxy = libmexclass.proxy.Proxy(Name=traits.TypeProxyClassName, ID=proxyID);
+            typeStruct = obj.Proxy.getType();
+            traits = arrow.type.traits.traits(arrow.type.ID(typeStruct.TypeID));
+            proxy = libmexclass.proxy.Proxy(Name=traits.TypeProxyClassName, ID=typeStruct.ProxyID);
             type = traits.TypeConstructor(proxy);
         end
 
-        function array = chunk(obj, index)
-            index = arrow.internal.validate.index.numeric(index, "int32");
-            [typeID, proxyID] = obj.Proxy.getChunk(index);
-            traits = arrow.type.traits.traits(arrow.type.ID(typeID));
-            proxy = libmexclass.proxy.Proxy(Name=traits.ArrayProxyClassName, ID=proxyID);
+        function array = chunk(obj, idx)
+            idx = arrow.internal.validate.index.numeric(idx, "int32", AllowNonScalar=false);
+            chunkStruct = obj.Proxy.getChunk(struct(Index=idx));
+            traits = arrow.type.traits.traits(arrow.type.ID(chunkStruct.TypeID));
+            proxy = libmexclass.proxy.Proxy(Name=traits.ArrayProxyClassName, ID=chunkStruct.ProxyID);
             array = traits.ArrayConstructor(proxy);
         end
     end