improvements in the UDF API

Author: DimitrisStaratzis

src/main/java/io/tiledb/cloud/TileDBUDF.java

@@ -3,9 +3,7 @@
 import com.google.gson.Gson;
 import io.tiledb.cloud.rest_api.ApiException;
 import io.tiledb.cloud.rest_api.api.UdfApi;
-import io.tiledb.cloud.rest_api.model.GenericUDF;
-import io.tiledb.cloud.rest_api.model.MultiArrayUDF;
-import io.tiledb.cloud.rest_api.model.ResultFormat;
+import io.tiledb.cloud.rest_api.model.*;
 import io.tiledb.java.api.Pair;
 import org.apache.arrow.compression.CommonsCompressionFactory;
 import org.apache.arrow.memory.RootAllocator;
@@ -21,7 +19,11 @@
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.HashMap;
+import java.util.List;
+
+import static io.tiledb.cloud.TileDBUtils.serializeArgs;
 
 public class TileDBUDF {
     private TileDBClient tileDBClient;
@@ -32,7 +34,6 @@ public TileDBUDF(TileDBClient tileDBClient, String namespace) {
         this.tileDBClient = tileDBClient;
         this.namespace = namespace;
         this.apiInstance = new UdfApi(this.tileDBClient.getApiClient());
-
     }
 
     /**
@@ -57,6 +58,25 @@ public String executeGeneric(GenericUDF genericUDF, HashMap<String, Object> argu
         return null;
     }
 
+    /**
+     * Executes a generic-UDF. A generic-UDF is a UDF that is not using a TIleDB array.
+     *
+     * @param genericUDF The generic UDF definition
+     * @return The result in String format
+     */
+    public String executeGeneric(GenericUDF genericUDF){
+        try {
+            return apiInstance.submitGenericUDFString(namespace, genericUDF, "none");
+        } catch (ApiException e) {
+            System.err.println("Exception when calling UdfApi#submitGenericUDF");
+            System.err.println("Status code: " + e.getCode());
+            System.err.println("Reason: " + e.getResponseBody());
+            System.err.println("Response headers: " + e.getResponseHeaders());
+            e.printStackTrace();
+        }
+        return null;
+    }
+
     /**
      * Executes a generic-UDF. A generic-UDF is a UDF that is not using a TIleDB array.
      *
@@ -70,6 +90,18 @@ public JSONObject executeGenericJSON(GenericUDF genericUDF, HashMap<String, Obje
         return new JSONObject(jsonString);
     }
 
+    /**
+     * Executes a generic-UDF. A generic-UDF is a UDF that is not using a TIleDB array.
+     *
+     * @param genericUDF The generic UDF definition
+     * @return The result as a JSON object
+     */
+    public JSONObject executeGenericJSON(GenericUDF genericUDF){
+        genericUDF.setResultFormat(ResultFormat.JSON);
+        String jsonString =  this.executeGeneric(genericUDF);
+        return new JSONObject(jsonString);
+    }
+
     /**
      * Executes a generic-UDF. A generic-UDF is a UDF that is not using a TIleDB array.
      *
@@ -79,7 +111,19 @@ public JSONObject executeGenericJSON(GenericUDF genericUDF, HashMap<String, Obje
      */
     public JSONArray executeGenericJSONArray(GenericUDF genericUDF, HashMap<String, Object> arguments){
         genericUDF.setResultFormat(ResultFormat.JSON);
-        String jsonString =  this.executeGeneric(genericUDF, arguments);
+        String jsonString = this.executeGeneric(genericUDF, arguments);
+        return new JSONArray(jsonString);
+    }
+
+    /**
+     * Executes a generic-UDF. A generic-UDF is a UDF that is not using a TIleDB array.
+     *
+     * @param genericUDF The generic UDF definition
+     * @return The result as a JSON array object
+     */
+    public JSONArray executeGenericJSONArray(GenericUDF genericUDF){
+        genericUDF.setResultFormat(ResultFormat.JSON);
+        String jsonString = this.executeGeneric(genericUDF);
         return new JSONArray(jsonString);
     }
 
@@ -90,7 +134,8 @@ public JSONArray executeGenericJSONArray(GenericUDF genericUDF, HashMap<String,
      * @param arguments The UDF arguments
      * @return A pair that consists of an ArrayList of all valueVectors and the number of batches read.
      */
-    public io.tiledb.java.api.Pair<ArrayList<ValueVector>, Integer> executeGenericArrow(GenericUDF genericUDF, HashMap<String, Object> arguments){
+    public io.tiledb.java.api.Pair<ArrayList<ValueVector>, Integer> executeGenericArrow(GenericUDF genericUDF,
+                                                                                        HashMap<String, Object> arguments){
         String serializedArgs = serializeArgs(arguments);
         genericUDF.setArgument(serializedArgs);
         genericUDF.setResultFormat(ResultFormat.ARROW);
@@ -103,6 +148,24 @@ public io.tiledb.java.api.Pair<ArrayList<ValueVector>, Integer> executeGenericAr
         return null;
     }
 
+    /**
+     * Executes a generic-UDF. A generic-UDF is a UDF that is not using a TIleDB array.
+     *
+     * @param genericUDF The generic UDF definition
+     * @return A pair that consists of an ArrayList of all valueVectors and the number of batches read.
+     */
+    public io.tiledb.java.api.Pair<ArrayList<ValueVector>, Integer> executeGenericArrow(GenericUDF genericUDF){
+        genericUDF.setResultFormat(ResultFormat.ARROW);
+        try {
+            byte[] bytes = apiInstance.submitGenericUDFBytes(namespace, genericUDF, "none");
+            return TileDBUtils.createValueVectors(bytes);
+        } catch (IOException | ApiException e) {
+            e.printStackTrace();
+        }
+        return null;
+    }
+
+    @Deprecated
     /**
      * Executes an array-UDF. An array-UDF is a UDF applied to a TileDB array
      *
@@ -130,6 +193,40 @@ public String executeSingleArray(MultiArrayUDF multiArrayUDF, HashMap<String, Ob
         return null;
     }
 
+    /**
+     * Executes an array-UDF. An array-UDF is a UDF applied to a TileDB array
+     *
+     * @param genericUDF The UDF to run
+     * @param array The Array to run the UDF on
+     * @param queryRanges the Query ranges
+     * @param xPayer Name of organization or user who should be charged for this request
+     * @return The results as a String
+     */
+    public String executeSingleArray(GenericUDF genericUDF, UDFArrayDetails array,
+                                     QueryRanges queryRanges, String xPayer){
+        List<UDFArrayDetails> arrays = new ArrayList<>();
+        arrays.add(array);
+
+        MultiArrayUDF multiArrayUDF = new MultiArrayUDF();
+        multiArrayUDF.setArgument(genericUDF.getArgument());
+        multiArrayUDF.setUdfInfoName(genericUDF.getUdfInfoName());
+        multiArrayUDF.setRanges(queryRanges);
+        multiArrayUDF.setArrays(arrays);
+        String[] split = breakdownFullURI(array.getUri());
+
+        try {
+            return apiInstance.submitUDFString(split[0], split[1], multiArrayUDF, xPayer, "none", "");
+        } catch (ApiException e) {
+            System.err.println("Exception when calling UdfApi#submitUDF");
+            System.err.println("Status code: " + e.getCode());
+            System.err.println("Reason: " + e.getResponseBody());
+            System.err.println("Response headers: " + e.getResponseHeaders());
+            e.printStackTrace();
+        }
+        return null;
+    }
+
+    @Deprecated
     /**
      * Executes an array-UDF. An array-UDF is a UDF applied to a TileDB array
      *
@@ -145,6 +242,42 @@ public JSONObject executeSingleArrayJSON(MultiArrayUDF multiArrayUDF, HashMap<St
         return new JSONObject(jsonString);
     }
 
+    /**
+     * Executes an array-UDF. An array-UDF is a UDF applied to a TileDB array
+     *
+     * @param genericUDF The UDF to run
+     * @param array The Array to run the UDF on
+     * @param queryRanges the Query ranges
+     * @param xPayer Name of organization or user who should be charged for this request
+     * @return The results as a JSON Object
+     */
+    public JSONObject executeSingleArrayJSON(GenericUDF genericUDF, UDFArrayDetails array,
+                                             QueryRanges queryRanges,  String xPayer){
+
+        List<UDFArrayDetails> arrays = new ArrayList<>();
+        arrays.add(array);
+
+        MultiArrayUDF multiArrayUDF = new MultiArrayUDF();
+        multiArrayUDF.setArgument(genericUDF.getArgument());
+        multiArrayUDF.setUdfInfoName(genericUDF.getUdfInfoName());
+        multiArrayUDF.setRanges(queryRanges);
+        multiArrayUDF.setArrays(arrays);
+        String[] split = breakdownFullURI(array.getUri());
+        multiArrayUDF.setResultFormat(ResultFormat.JSON);
+        String result = "";
+        try {
+            result = apiInstance.submitUDFString(split[0], split[1], multiArrayUDF, xPayer, "none", "");
+        } catch (ApiException e) {
+            System.err.println("Exception when calling UdfApi#submitUDF");
+            System.err.println("Status code: " + e.getCode());
+            System.err.println("Reason: " + e.getResponseBody());
+            System.err.println("Response headers: " + e.getResponseHeaders());
+            e.printStackTrace();
+        }
+        return new JSONObject(result);
+    }
+
+    @Deprecated
     /**
      * Executes an array-UDF. An array-UDF is a UDF applied to a TileDB array
      *
@@ -154,12 +287,48 @@ public JSONObject executeSingleArrayJSON(MultiArrayUDF multiArrayUDF, HashMap<St
      * @param xPayer Name of organization or user who should be charged for this request
      * @return The results as a JSON array
      */
-    public JSONArray executeSingleArrayJSONArray(MultiArrayUDF multiArrayUDF, HashMap<String, Object> arguments, String arrayURI, String xPayer){
+    public JSONArray executeSingleArrayJSONArray(MultiArrayUDF multiArrayUDF, HashMap<String,
+            Object> arguments, String arrayURI, String xPayer){
         multiArrayUDF.setResultFormat(ResultFormat.JSON);
         String jsonString = this.executeSingleArray(multiArrayUDF, arguments, arrayURI, xPayer);
         return new JSONArray(jsonString);
     }
 
+    /**
+     * Executes an array-UDF. An array-UDF is a UDF applied to a TileDB array
+     *
+     * @param genericUDF The UDF to run
+     * @param array The Array to run the UDF on
+     * @param queryRanges the Query ranges
+     * @param xPayer Name of organization or user who should be charged for this request
+     * @return The results as a JSON Object
+     */
+    public JSONArray executeSingleArrayJSONArray(GenericUDF genericUDF, UDFArrayDetails array,
+                                             QueryRanges queryRanges,  String xPayer){
+        List<UDFArrayDetails> arrays = new ArrayList<>();
+        arrays.add(array);
+
+        MultiArrayUDF multiArrayUDF = new MultiArrayUDF();
+        multiArrayUDF.setArgument(genericUDF.getArgument());
+        multiArrayUDF.setUdfInfoName(genericUDF.getUdfInfoName());
+        multiArrayUDF.setRanges(queryRanges);
+        multiArrayUDF.setArrays(arrays);
+        String[] split = breakdownFullURI(array.getUri());
+        multiArrayUDF.setResultFormat(ResultFormat.JSON);
+        String result = "";
+        try {
+            result = apiInstance.submitUDFString(split[0], split[1], multiArrayUDF, xPayer, "none", "");
+        } catch (ApiException e) {
+            System.err.println("Exception when calling UdfApi#submitUDF");
+            System.err.println("Status code: " + e.getCode());
+            System.err.println("Reason: " + e.getResponseBody());
+            System.err.println("Response headers: " + e.getResponseHeaders());
+            e.printStackTrace();
+        }
+        return new JSONArray(result);
+    }
+
+    @Deprecated
     /**
      * Executes an array-UDF. An array-UDF is a UDF applied to a TileDB array
      *
@@ -185,6 +354,37 @@ public Pair<ArrayList<ValueVector>, Integer> executeSingleArrayArrow(MultiArrayU
         return null;
     }
 
+    /**
+     * Executes an array-UDF. An array-UDF is a UDF applied to a TileDB array
+     *
+     * @param genericUDF The UDF to run
+     * @param array The Array to run the UDF on
+     * @param queryRanges the Query ranges
+     * @param xPayer Name of organization or user who should be charged for this request
+     * @return The results in arrow format
+     */
+    public Pair<ArrayList<ValueVector>, Integer> executeSingleArrayArrow(GenericUDF genericUDF, UDFArrayDetails array,
+                                                                         QueryRanges queryRanges,  String xPayer){
+        List<UDFArrayDetails> arrays = new ArrayList<>();
+        arrays.add(array);
+
+        MultiArrayUDF multiArrayUDF = new MultiArrayUDF();
+        multiArrayUDF.setArgument(genericUDF.getArgument());
+        multiArrayUDF.setResultFormat(ResultFormat.ARROW);
+        multiArrayUDF.setUdfInfoName(genericUDF.getUdfInfoName());
+        multiArrayUDF.setRanges(queryRanges);
+        multiArrayUDF.setArrays(arrays);
+        String[] split = breakdownFullURI(array.getUri());
+
+        try {
+            byte[] bytes = apiInstance.submitUDFBytes(split[0], split[1], multiArrayUDF, xPayer, "none", "");
+            return TileDBUtils.createValueVectors(bytes);
+        } catch (IOException | ApiException e) {
+            e.printStackTrace();
+        }
+        return null;
+    }
+
     /**
      * Executes a multi-array-UDF. A multi- array-UDF is a UDF applied to multiple TileDB arrays
      *
@@ -207,6 +407,25 @@ public String executeMultiArray(MultiArrayUDF multiArrayUDF, HashMap<String, Obj
         return null;
     }
 
+    /**
+     * Executes a multi-array-UDF. A multi- array-UDF is a UDF applied to multiple TileDB arrays
+     *
+     * @param multiArrayUDF The multiArrayUDF input object
+     * @return The results as a String
+     */
+    public String executeMultiArray(MultiArrayUDF multiArrayUDF){
+        try {
+            return apiInstance.submitMultiArrayUDFString(this.namespace, multiArrayUDF, "none");
+        } catch (ApiException e) {
+            System.err.println("Exception when calling UdfApi#submitMultiArrayUDFString");
+            System.err.println("Status code: " + e.getCode());
+            System.err.println("Reason: " + e.getResponseBody());
+            System.err.println("Response headers: " + e.getResponseHeaders());
+            e.printStackTrace();
+        }
+        return null;
+    }
+
     /**
      * Executes a multi-array-UDF. A multi- array-UDF is a UDF applied to multiple TileDB arrays
      *
@@ -220,6 +439,18 @@ public JSONObject executeMultiArrayJSON(MultiArrayUDF multiArrayUDF, HashMap<Str
         return new JSONObject(jsonString);
     }
 
+    /**
+     * Executes a multi-array-UDF. A multi- array-UDF is a UDF applied to multiple TileDB arrays
+     *
+     * @param multiArrayUDF The multiArrayUDF input object
+     * @return The results as JSON object
+     */
+    public JSONObject executeMultiArrayJSON(MultiArrayUDF multiArrayUDF){
+        multiArrayUDF.setResultFormat(ResultFormat.JSON);
+        String jsonString = this.executeMultiArray(multiArrayUDF);
+        return new JSONObject(jsonString);
+    }
+
     /**
      * Executes a multi-array-UDF. A multi- array-UDF is a UDF applied to multiple TileDB arrays
      *
@@ -233,6 +464,18 @@ public JSONArray executeMultiArrayJSONArray(MultiArrayUDF multiArrayUDF, HashMap
         return new JSONArray(jsonString);
     }
 
+    /**
+     * Executes a multi-array-UDF. A multi- array-UDF is a UDF applied to multiple TileDB arrays
+     *
+     * @param multiArrayUDF The multiArrayUDF input object
+     * @return The results as JSON Array
+     */
+    public JSONArray executeMultiArrayJSONArray(MultiArrayUDF multiArrayUDF){
+        multiArrayUDF.setResultFormat(ResultFormat.JSON);
+        String jsonString = this.executeMultiArray(multiArrayUDF);
+        return new JSONArray(jsonString);
+    }
+
     /**
      * Executes a multi-array-UDF. A multi- array-UDF is a UDF applied to multiple TileDB arrays
      *
@@ -255,15 +498,21 @@ public Pair<ArrayList<ValueVector>, Integer> executeMultiArrayArrow(MultiArrayUD
     }
 
     /**
-     * Serializes the arguments to a String
+     * Executes a multi-array-UDF. A multi- array-UDF is a UDF applied to multiple TileDB arrays
      *
-     * @param arguments The input arguments in a HashMap
-     * @return The arguments in a String
+     * @param multiArrayUDF The multiArrayUDF input object
+     * @return The results in arrow format
      */
-    private String serializeArgs(HashMap<String, Object> arguments) {
-        if (arguments == null || arguments.isEmpty()) return "";
-        Gson gson = new Gson();
-        return gson.toJson(arguments);
+    public Pair<ArrayList<ValueVector>, Integer> executeMultiArrayArrow(MultiArrayUDF multiArrayUDF){
+        multiArrayUDF.setResultFormat(ResultFormat.ARROW);
+
+        try {
+            byte[] bytes = apiInstance.submitMultiArrayUDFBytes(this.namespace, multiArrayUDF, "none");
+            return TileDBUtils.createValueVectors(bytes);
+        } catch (IOException | ApiException e) {
+            e.printStackTrace();
+        }
+        return null;
     }
 
     /**