From 917fb4a2c6fdaf3042649a476e03f4257001c98e Mon Sep 17 00:00:00 2001
From: gsvic <victorasgs@gmail.com>
Date: Wed, 21 Apr 2021 17:36:00 +0300
Subject: [PATCH] Added support for estimating result size of nullable
 attributes

---
 src/main/java/io/tiledb/java/api/Query.java   | 60 +++++++++++++++++++
 .../java/io/tiledb/java/api/QueryTest.java    | 10 ++++
 2 files changed, 70 insertions(+)

diff --git a/src/main/java/io/tiledb/java/api/Query.java b/src/main/java/io/tiledb/java/api/Query.java
index 076dd7d4..89246b4c 100644
--- a/src/main/java/io/tiledb/java/api/Query.java
+++ b/src/main/java/io/tiledb/java/api/Query.java
@@ -345,6 +345,14 @@ public synchronized Pair<String, String> getRangeVar(int dimIdx, BigInteger rang
     }
   }
 
+  /**
+   * Retrieves the estimated result size for a fixed-sized attribute/dimension.
+   *
+   * @param ctx The TileDB Context
+   * @param attribute The attribute name
+   * @return The estimated result size
+   * @throws TileDBError
+   */
   public synchronized int getEstResultSize(Context ctx, String attribute) throws TileDBError {
     SWIGTYPE_p_unsigned_long_long size = tiledb.new_ullp();
 
@@ -354,6 +362,14 @@ public synchronized int getEstResultSize(Context ctx, String attribute) throws T
     return tiledb.ullp_value(size).intValue();
   }
 
+  /**
+   * Retrieves the estimated result size for a var-sized attribute/dimension.
+   *
+   * @param ctx The TileDB Context
+   * @param attribute The attribute name
+   * @return A Pair containing the estimated result size of the offsets and the data buffers
+   * @throws TileDBError
+   */
   public synchronized Pair<Integer, Integer> getEstResultSizeVar(Context ctx, String attribute)
       throws TileDBError {
     SWIGTYPE_p_unsigned_long_long offsetsSize = tiledb.new_ullp();
@@ -367,6 +383,50 @@ public synchronized Pair<Integer, Integer> getEstResultSizeVar(Context ctx, Stri
         tiledb.ullp_value(offsetsSize).intValue(), tiledb.ullp_value(dataSize).intValue());
   }
 
+  /**
+   * Retrieves the estimated result size for a var-sized nullable attribute.
+   *
+   * @param ctx The TileDB Context
+   * @param attribute The attribute name
+   * @return A Pair containing another Pair with the estimated result size of the offsets and the
+   *     data buffers, and the estimated result size of the validity buffer
+   * @throws TileDBError
+   */
+  public synchronized Pair<Pair<Integer, Integer>, Integer> getEstResultSizeVarNullable(
+      Context ctx, String attribute) throws TileDBError {
+    SWIGTYPE_p_unsigned_long_long size = tiledb.new_ullp();
+    SWIGTYPE_p_unsigned_long_long offsets = tiledb.new_ullp();
+    SWIGTYPE_p_unsigned_long_long validity = tiledb.new_ullp();
+
+    ctx.handleError(
+        tiledb.tiledb_query_get_est_result_size_var_nullable(
+            ctx.getCtxp(), queryp, attribute, offsets, size, validity));
+
+    return new Pair(
+        new Pair(tiledb.ullp_value(offsets).intValue(), tiledb.ullp_value(size).intValue()),
+        tiledb.ullp_value(validity).intValue());
+  }
+
+  /**
+   * Retrieves the estimated result size for a fixed-sized nullable attribute.
+   *
+   * @param ctx The TileDB Context
+   * @param attribute The attribute name
+   * @return The estimated result size
+   * @throws TileDBError
+   */
+  public synchronized Pair<Integer, Integer> getEstResultSizeNullable(Context ctx, String attribute)
+      throws TileDBError {
+    SWIGTYPE_p_unsigned_long_long size = tiledb.new_ullp();
+    SWIGTYPE_p_unsigned_long_long validity = tiledb.new_ullp();
+
+    ctx.handleError(
+        tiledb.tiledb_query_get_est_result_size_nullable(
+            ctx.getCtxp(), queryp, attribute, size, validity));
+
+    return new Pair(tiledb.ullp_value(size).intValue(), tiledb.ullp_value(validity).intValue());
+  }
+
   /**
    * Retrieves the number of ranges of the query subarray along a given dimension.
    *
diff --git a/src/test/java/io/tiledb/java/api/QueryTest.java b/src/test/java/io/tiledb/java/api/QueryTest.java
index fd02c114..0b2266af 100644
--- a/src/test/java/io/tiledb/java/api/QueryTest.java
+++ b/src/test/java/io/tiledb/java/api/QueryTest.java
@@ -1124,6 +1124,10 @@ public void denseArrayReadTest() throws Exception {
         query.setBufferNullable("a1", a1Array, a1byteMap);
         query.setBufferNullable("a2", a2Array, a2byteMap);
 
+        Pair<Integer, Integer> estimated = query.getEstResultSizeNullable(ctx, "a1");
+        Assert.assertEquals((int) estimated.getFirst(), 4);
+        Assert.assertEquals((int) estimated.getSecond(), 4);
+
         // Submit query
         query.submit();
 
@@ -1252,6 +1256,12 @@ public void sparseArrayReadTest() throws Exception {
         query.setBufferNullable("a1", a1Array, a1byteMap);
         query.setBufferNullable("a2", a2Offsets, a2Array, a2byteMap);
 
+        Pair<Pair<Integer, Integer>, Integer> estimated =
+            query.getEstResultSizeVarNullable(ctx, "a2");
+        Assert.assertEquals(40, (int) estimated.getFirst().getFirst());
+        Assert.assertEquals(40, (int) estimated.getFirst().getFirst());
+        Assert.assertEquals(10, (int) estimated.getSecond());
+
         // Submit query
         query.submit();