2828
2929import io .tiledb .libtiledb .*;
3030import java .math .BigInteger ;
31+ import java .nio .ByteBuffer ;
32+ import java .nio .ByteOrder ;
3133import java .util .Collections ;
3234import java .util .HashMap ;
3335import java .util .Map ;
@@ -58,6 +60,7 @@ public class Query implements AutoCloseable {
5860 private NativeArray subarray ;
5961
6062 private Map <String , NativeArray > buffers_ ;
63+ private Map <String , ByteBuffer > byteBuffers_ ;
6164 private Map <String , Pair <NativeArray , NativeArray >> var_buffers_ ;
6265 private Map <String , Pair <uint64_tArray , uint64_tArray >> buffer_sizes_ ;
6366
@@ -78,6 +81,7 @@ public Query(Array array, QueryType type) throws TileDBError {
7881 this .querypp = _querypp ;
7982 this .queryp = tiledb .tiledb_query_tpp_value (_querypp );
8083 this .buffers_ = Collections .synchronizedMap (new HashMap <>());
84+ this .byteBuffers_ = Collections .synchronizedMap (new HashMap <>());
8185 this .var_buffers_ = Collections .synchronizedMap (new HashMap <>());
8286 this .buffer_sizes_ = Collections .synchronizedMap (new HashMap <>());
8387 }
@@ -467,6 +471,82 @@ public synchronized Query setBuffer(String attr, NativeArray buffer, long buffer
467471 return this ;
468472 }
469473
474+ /**
475+ * * Sets a NIO ByteBuffer
476+ *
477+ * @param attr The attribute
478+ * @param bufferElements
479+ * @return The NIO ByteBuffer
480+ * @throws TileDBError
481+ */
482+ public synchronized ByteBuffer setBuffer (String attr , long bufferElements ) throws TileDBError {
483+ if (bufferElements <= 0 ) {
484+ throw new TileDBError ("Number of buffer elements must be >= 1" );
485+ }
486+
487+ Datatype dt ;
488+
489+ try (ArraySchema schema = array .getSchema ()) {
490+ try (Domain domain = schema .getDomain ()) {
491+ if (domain .hasDimension (attr )) {
492+ dt = domain .getDimension (attr ).getType ();
493+ } else {
494+ try (Attribute attribute = schema .getAttribute (attr )) {
495+ dt = attribute .getType ();
496+ }
497+ }
498+ }
499+ }
500+
501+ int size = Util .castLongToInt (bufferElements * dt .getNativeSize ());
502+
503+ ByteBuffer buffer = ByteBuffer .allocateDirect (size );
504+
505+ // Set the byte order to the native system's native order
506+ buffer .order (ByteOrder .nativeOrder ());
507+
508+ this .byteBuffers_ .put (attr , buffer );
509+
510+ this .setBuffer (attr , buffer );
511+
512+ return buffer ;
513+ }
514+
515+ /**
516+ * * Sets a NIO ByteBuffer
517+ *
518+ * @param attr The attribute
519+ * @param buffer The input NIO ByteBuffer
520+ * @return The NIO ByteBuffer
521+ * @throws TileDBError
522+ */
523+ public synchronized ByteBuffer setBuffer (String attr , ByteBuffer buffer ) throws TileDBError {
524+ if (buffer .capacity () <= 0 ) {
525+ throw new TileDBError ("Number of buffer elements must be >= 1" );
526+ }
527+
528+ if (!buffer .isDirect ()) {
529+ throw new TileDBError (
530+ "The ByteBuffer provided is not direct. Please provide a direct buffer (ByteBuffer.allocateDirect(...))" );
531+ }
532+
533+ if (!buffer .order ().equals (ByteOrder .nativeOrder ())) {
534+ // TODO: Add a logger component to Query class and a WARN here
535+ buffer .order (ByteOrder .nativeOrder ());
536+ }
537+
538+ this .byteBuffers_ .put (attr , buffer );
539+
540+ uint64_tArray values_array_size = new uint64_tArray (1 );
541+ values_array_size .setitem (0 , BigInteger .valueOf (buffer .capacity ()));
542+
543+ ctx .handleError (
544+ tiledb .tiledb_query_set_buffer_nio (
545+ ctx .getCtxp (), queryp , attr , buffer , values_array_size .cast ()));
546+
547+ return buffer ;
548+ }
549+
470550 /**
471551 * Sets a buffer for a variable-sized getAttribute.
472552 *
@@ -885,6 +965,18 @@ public Object getBuffer(String attr) throws TileDBError {
885965 }
886966 }
887967
968+ /**
969+ * Retrieves the ByteBuffer of attribute attr
970+ *
971+ * @param attr The attribute name
972+ * @return The ByteBuffer
973+ * @throws TileDBError A TileDB exception
974+ */
975+ public ByteBuffer getByteBuffer (String attr ) throws TileDBError {
976+ if (byteBuffers_ .containsKey (attr )) return byteBuffers_ .get (attr );
977+ else throw new TileDBError ("ByteBuffer does not exist for attribute: " + attr );
978+ }
979+
888980 /**
889981 * Return an array containing offsets for a variable attribute buffer
890982 *
0 commit comments