Merge pull request #7 from RelationalAI/hnr-enhance-arrow-result

Enhance arrow parsing results type

Author: NRHelmi

README.md

@@ -86,7 +86,7 @@ In order to use the `rai-sdk-java`, you need add this dependency to your project
     <dependency>
         <groupId>com.relationalai</groupId>
         <artifactId>rai-sdk</artifactId>
-        <version>0.0.2</version>
+        <version>0.1.0-alpha</version>
     </dependency>
 
 You need also to point maven to the SDK GitHub packages repository in the project's POM:

pom.xml

@@ -21,7 +21,7 @@
     <description>The RelationalAI Software Development Kit (SDK) for Java</description>
     <groupId>com.relationalai</groupId>
     <artifactId>rai-sdk-pom</artifactId>
-    <version>0.0.2</version>
+    <version>0.1.0-alpha</version>
     <packaging>pom</packaging>
     <url></url>
 

rai-sdk-examples/pom.xml

@@ -20,7 +20,7 @@
     <parent>
         <groupId>com.relationalai</groupId>
         <artifactId>rai-sdk-pom</artifactId>
-        <version>0.0.2</version>
+        <version>0.1.0-alpha</version>
     </parent>
 
     <name>RelationalAI SDK for Java Examples</name>

rai-sdk/pom.xml

@@ -20,7 +20,7 @@
     <parent>
         <groupId>com.relationalai</groupId>
         <artifactId>rai-sdk-pom</artifactId>
-        <version>0.0.2</version>
+        <version>0.1.0-alpha</version>
     </parent>
 
     <name>RelationalAI SDK for Java Package</name>

rai-sdk/src/main/java/com/relationalai/Client.java

@@ -19,6 +19,7 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.jsoniter.any.Any;
 import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.FieldVector;
 import org.apache.arrow.vector.VectorSchemaRoot;
 import org.apache.arrow.vector.ipc.ArrowStreamReader;
 import org.apache.arrow.vector.types.pojo.Field;
@@ -306,6 +307,7 @@ private String parseMultipartResponse(HttpResponse<byte[]> response) throws IOEx
                 result.add(out.toString(StandardCharsets.UTF_8));
             } else if (partContentType.toLowerCase().equals("application/vnd.apache.arrow.stream")) {
                 result.add(parseArrowResponse(out));
+                out.close();
             } else {
                 throw new HttpError(statusCode, String.format("unknown part content type: %s", partContentType));
             }
@@ -317,27 +319,43 @@ private String parseMultipartResponse(HttpResponse<byte[]> response) throws IOEx
     }
 
     private String parseArrowResponse(ByteArrayOutputStream out) throws IOException {
-        Map<String, String> result = new HashMap<>();
+        List<Object> output = new ArrayList<>();
 
-        RootAllocator allocator = new RootAllocator(Integer.MAX_VALUE);
         ByteArrayInputStream in = new ByteArrayInputStream(out.toByteArray());
-        ArrowStreamReader reader = new ArrowStreamReader(in, allocator);
-
-        VectorSchemaRoot readBatch = reader.getVectorSchemaRoot();
-
-        for (Field f : readBatch.getSchema().getFields()) {
-            result.put(f.getName(), String.valueOf(readBatch.getVector(f)));
-        }
-
-        while(reader.loadNextBatch()) {
-            readBatch = reader.getVectorSchemaRoot();
-
-            for (Field f : readBatch.getSchema().getFields()) {
-                result.put(f.getName(), String.valueOf(readBatch.getVector(f)));
+        List<FieldVector> fieldVectors = null;
+        RootAllocator allocator = new RootAllocator(Long.MAX_VALUE);
+        try(ArrowStreamReader arrowStreamReader = new ArrowStreamReader(in, allocator)){
+
+            VectorSchemaRoot root = arrowStreamReader.getVectorSchemaRoot();
+            fieldVectors = root.getFieldVectors();
+
+            while(arrowStreamReader.loadNextBatch()) {
+                 for(FieldVector fieldVector : fieldVectors) {
+                     Map<String, List<Object>> col = new HashMap<>();
+
+                     col.put(fieldVector.getField().getName(), new ArrayList<>());
+                     for (int i = 0; i < fieldVector.getValueCount(); i ++) {
+                         col.get(fieldVector.getField().getName()).add(fieldVector.getObject(i));
+                     }
+                     output.add(col);
+                 }
+            }
+        } finally {
+            if (in != null ) {
+                in.close();
             }
+            if (fieldVectors != null) {
+                for (FieldVector fieldVector : fieldVectors) {
+                    if (fieldVector != null) {
+                        fieldVector.close();
+                    }
+                }
+            }
+            allocator.close();
         }
+
         // serialize map to json string
-        return new ObjectMapper().writeValueAsString(result);
+        return new ObjectMapper().writeValueAsString(output);
     }
 
     static void printRequest(HttpRequest request) {

rai-sdk/src/test/java/com/relationalai/ExecuteAsyncTest.java

@@ -16,6 +16,7 @@
 
 package com.relationalai;
 
+import com.jsoniter.JsonIterator;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.TestInstance;
@@ -35,17 +36,17 @@ void testExecuteAsync() throws HttpError, InterruptedException, IOException {
         var query = "x, x^2, x^3, x^4 from x in {1; 2; 3; 4; 5}";
 
         var rsp = client.executeAsyncWait(databaseName, engineName, query, true);
+
         assertEquals(
-                rsp.get("results").toString(),
-                "[{\"v1\":\"[1, 2, 3, 4, 5]\",\"v2\":\"[1, 4, 9, 16, 25]\",\"v3\":\"[1, 8, 27, 64, 125]\",\"v4\":\"[1, 16, 81, 256, 625]\"}]"
-        );
+                rsp.get("results"),
+                JsonIterator.deserialize("[[{\"v1\":[1,2,3,4,5]},{\"v2\":[1,4,9,16,25]},{\"v3\":[1,8,27,64,125]},{\"v4\":[1,16,81,256,625]}]]"));
         assertEquals(
-                rsp.get("metadata").toString(),
-                "[{\"relationId\":\"/:output/Int64/Int64/Int64/Int64\",\"types\":[\":output\",\"Int64\",\"Int64\",\"Int64\",\"Int64\"]}]"
+                rsp.get("metadata"),
+                JsonIterator.deserialize("[{\"relationId\":\"/:output/Int64/Int64/Int64/Int64\",\"types\":[\":output\",\"Int64\",\"Int64\",\"Int64\",\"Int64\"]}]")
         );
         assertEquals(
-                rsp.get("problems").toString(),
-                "[]"
+                rsp.get("problems"),
+                JsonIterator.deserialize("[]")
         );
     }