Changed setBuffer to return the Query instance instead of the ByteBuffer

Author: gsvic

src/main/java/io/tiledb/java/api/Query.java

@@ -496,7 +496,7 @@ public synchronized Query setBuffer(String attr, NativeArray buffer, long buffer
    * @return The NIO ByteBuffer
    * @throws TileDBError
    */
-  public synchronized ByteBuffer setBuffer(String attr, long bufferElements) throws TileDBError {
+  public synchronized Query setBuffer(String attr, long bufferElements) throws TileDBError {
     if (bufferElements <= 0) {
       throw new TileDBError("Number of buffer elements must be >= 1");
     }
@@ -521,7 +521,7 @@ public synchronized ByteBuffer setBuffer(String attr, long bufferElements) throw
 
     this.setBuffer(attr, buffer);
 
-    return buffer;
+    return this;
   }
 
   /**
@@ -532,7 +532,7 @@ public synchronized ByteBuffer setBuffer(String attr, long bufferElements) throw
    * @return The NIO ByteBuffer
    * @throws TileDBError
    */
-  public synchronized ByteBuffer setBuffer(String attr, ByteBuffer buffer) throws TileDBError {
+  public synchronized Query setBuffer(String attr, ByteBuffer buffer) throws TileDBError {
     if (buffer.capacity() <= 0) {
       throw new TileDBError("Number of buffer elements must be >= 1");
     }
@@ -566,7 +566,7 @@ public synchronized ByteBuffer setBuffer(String attr, ByteBuffer buffer) throws
         tiledb.tiledb_query_set_buffer_nio(
             ctx.getCtxp(), queryp, attr, buffer, values_array_size.cast()));
 
-    return buffer;
+    return this;
   }
 
   /**
@@ -1054,8 +1054,10 @@ public Object getBuffer(String attr) throws TileDBError {
    * @return The ByteBuffer
    * @throws TileDBError A TileDB exception
    */
-  public ByteBuffer getByteBuffer(String attr) throws TileDBError {
-    if (byteBuffers_.containsKey(attr)) return byteBuffers_.get(attr);
+  public Pair<ByteBuffer, ByteBuffer> getByteBuffer(String attr) throws TileDBError {
+    if (byteBuffers_.containsKey(attr)) return new Pair(null, byteBuffers_.get(attr));
+    else if (varByteBuffers_.containsKey(attr))
+      return new Pair(varByteBuffers_.get(attr).getFirst(), varByteBuffers_.get(attr).getSecond());
     else throw new TileDBError("ByteBuffer does not exist for attribute: " + attr);
   }
 

src/test/java/io/tiledb/java/api/QueryTest.java

@@ -314,8 +314,10 @@ public void queryTestNIOReadArray() throws Exception {
 
       int bufferSize = 4;
 
-      ByteBuffer d1 = query.setBuffer("rows", bufferSize);
-      ByteBuffer d2 = query.setBuffer("cols", bufferSize);
+      query.setBuffer("rows", bufferSize);
+      query.setBuffer("cols", bufferSize);
+      ByteBuffer d1 = query.getByteBuffer("rows").getSecond();
+      ByteBuffer d2 = query.getByteBuffer("cols").getSecond();
 
       query.addRange(0, 1, 4);
       query.addRange(1, 1, 4);
@@ -354,9 +356,10 @@ public void queryTestNIOReadArrayArbitrarySize() throws Exception {
       Query query = new Query(array, TILEDB_READ);
 
       int bufferSize = 4;
-
-      ByteBuffer d1 = query.setBuffer("rows", ByteBuffer.allocateDirect(10));
-      ByteBuffer d2 = query.setBuffer("cols", ByteBuffer.allocateDirect(10));
+      query.setBuffer("rows", ByteBuffer.allocateDirect(10));
+      query.setBuffer("cols", ByteBuffer.allocateDirect(10));
+      ByteBuffer d1 = query.getByteBuffer("rows").getSecond();
+      ByteBuffer d2 = query.getByteBuffer("cols").getSecond();
 
       query.addRange(0, 1, 4);
       query.addRange(1, 1, 4);
@@ -399,10 +402,12 @@ public void arrayReadTest() throws Exception {
         query.addRange(1, 2, 4);
         query.setLayout(TILEDB_ROW_MAJOR);
 
-        ByteBuffer dim1Buffer = query.setBuffer("rows", 3);
-        ByteBuffer dim2Buffer = query.setBuffer("cols", 3);
-        ByteBuffer a1Buffer = query.setBuffer("a1", 3);
-        ByteBuffer a2Buffer = query.setBuffer("a2", 6);
+        query.setBuffer("rows", 3).setBuffer("cols", 3).setBuffer("a1", 3).setBuffer("a2", 6);
+
+        ByteBuffer dim1Buffer = query.getByteBuffer("rows").getSecond();
+        ByteBuffer dim2Buffer = query.getByteBuffer("cols").getSecond();
+        ByteBuffer a1Buffer = query.getByteBuffer("a1").getSecond();
+        ByteBuffer a2Buffer = query.getByteBuffer("a2").getSecond();
 
         // Submit query
         query.submit();
@@ -459,12 +464,16 @@ public void arrayReadTestCustomBufferWithDifferentOrder() throws Exception {
                 ? ByteOrder.LITTLE_ENDIAN
                 : ByteOrder.BIG_ENDIAN;
 
-        ByteBuffer dim1Buffer =
-            query.setBuffer("rows", ByteBuffer.allocateDirect(3 * 4).order(order));
-        ByteBuffer dim2Buffer =
-            query.setBuffer("cols", ByteBuffer.allocateDirect(3 * 4).order(order));
-        ByteBuffer a1Buffer = query.setBuffer("a1", ByteBuffer.allocateDirect(3).order(order));
-        ByteBuffer a2Buffer = query.setBuffer("a2", ByteBuffer.allocateDirect(6 * 4).order(order));
+        query
+            .setBuffer("rows", ByteBuffer.allocateDirect(3 * 4).order(order))
+            .setBuffer("cols", ByteBuffer.allocateDirect(3 * 4).order(order))
+            .setBuffer("a1", ByteBuffer.allocateDirect(3).order(order))
+            .setBuffer("a2", ByteBuffer.allocateDirect(6 * 4).order(order));
+
+        ByteBuffer dim1Buffer = query.getByteBuffer("rows").getSecond();
+        ByteBuffer dim2Buffer = query.getByteBuffer("cols").getSecond();
+        ByteBuffer a1Buffer = query.getByteBuffer("a1").getSecond();
+        ByteBuffer a2Buffer = query.getByteBuffer("a2").getSecond();
 
         // Submit query
         query.submit();
@@ -514,9 +523,9 @@ public void queryTestNIOGetByteBuffer() throws Exception {
 
       query.setBuffer("rows", bufferSize);
 
-      Assert.assertEquals(query.getByteBuffer("rows").capacity(), bufferSize * 4);
-      Assert.assertTrue(query.getByteBuffer("rows").isDirect());
-      Assert.assertEquals(query.getByteBuffer("rows").order(), ByteOrder.nativeOrder());
+      Assert.assertEquals(query.getByteBuffer("rows").getSecond().capacity(), bufferSize * 4);
+      Assert.assertTrue(query.getByteBuffer("rows").getSecond().isDirect());
+      Assert.assertEquals(query.getByteBuffer("rows").getSecond().order(), ByteOrder.nativeOrder());
     }
 
     @Test()
@@ -544,7 +553,8 @@ public void queryTestNIOGetByteBuffeErrors() throws Exception {
               : ByteOrder.BIG_ENDIAN;
 
       // The Byte Order should be automatically changed to the native order
-      ByteBuffer b = query.setBuffer("rows", ByteBuffer.allocateDirect(bufferSize).order(order));
+      query.setBuffer("rows", ByteBuffer.allocateDirect(bufferSize).order(order));
+      ByteBuffer b = query.getByteBuffer("rows").getSecond();
       Assert.assertEquals(b.order(), ByteOrder.nativeOrder());
     }
 
@@ -566,10 +576,12 @@ public void queryTestNIOSetBufferVarChar() throws Exception {
 
       int idx = 0;
 
-      while (offsetsBuffer.hasRemaining()) offsets[idx++] = offsetsBuffer.getLong();
+      while (q.getByteBuffer("a1").getFirst().hasRemaining())
+        offsets[idx++] = offsetsBuffer.getLong();
 
       idx = 0;
-      while (dataBuffer.hasRemaining()) data[idx++] = (char) dataBuffer.get();
+      while (q.getByteBuffer("a1").getSecond().hasRemaining())
+        data[idx++] = (char) dataBuffer.get();
 
       Assert.assertArrayEquals(new long[] {0, 2, 4, 6, 8, 10, 12, 14}, offsets);
       Assert.assertEquals("aabbccddeeffgghh", new String(data));