Improve: Zero-copy USearch APIs

Author: ashvardanian

java/cloud/unum/usearch/Index.java

@@ -242,6 +242,23 @@ public void add(long key, float vector[]) {
     c_add_f32(c_ptr, key, vector);
   }
 
+  /**
+   * Adds a vector with a specified key to the index using ByteBuffer (zero-copy).
+   *
+   * @param key    the key associated with the vector
+   * @param vector the vector data as FloatBuffer
+   */
+  public void add(long key, java.nio.FloatBuffer vector) {
+    if (c_ptr == 0) {
+      throw new IllegalStateException("Index already closed");
+    }
+    if (vector.remaining() != dimensions()) {
+      throw new IllegalArgumentException(String.format(
+          "Vector dimensions mismatch: expected %d but got %d", dimensions(), vector.remaining()));
+    }
+    c_add_f32_buffer(c_ptr, key, vector);
+  }
+
   /**
    * Searches for closest vectors to the specified query vector.
    *
@@ -256,6 +273,55 @@ public long[] search(float vector[], long count) {
     return c_search_f32(c_ptr, vector, count);
   }
 
+  /**
+   * Searches for closest vectors to the specified query vector using ByteBuffer
+   * (zero-copy).
+   *
+   * @param vector the query vector data as FloatBuffer
+   * @param count  the number of nearest neighbors to search
+   * @return an array of keys of the nearest neighbors
+   */
+  public long[] search(java.nio.FloatBuffer vector, long count) {
+    if (c_ptr == 0) {
+      throw new IllegalStateException("Index already closed");
+    }
+    if (vector.remaining() != dimensions()) {
+      throw new IllegalArgumentException(String.format(
+          "Vector dimensions mismatch: expected %d but got %d", dimensions(), vector.remaining()));
+    }
+    return c_search_f32_buffer(c_ptr, vector, count);
+  }
+
+  /**
+   * Searches for closest vectors using zero-copy input and zero-allocation
+   * output.
+   *
+   * @param query    the query vector data as FloatBuffer
+   * @param results  the output buffer for result keys
+   * @param maxCount maximum number of results to find
+   * @return actual number of results found
+   */
+  public int searchInto(java.nio.FloatBuffer query, java.nio.LongBuffer results, long maxCount) {
+    if (c_ptr == 0) {
+      throw new IllegalStateException("Index already closed");
+    }
+    if (query.remaining() != dimensions()) {
+      throw new IllegalArgumentException(String.format(
+          "Query vector dimensions mismatch: expected %d but got %d", dimensions(), query.remaining()));
+    }
+    if (results.remaining() < maxCount) {
+      throw new IllegalArgumentException(String.format(
+          "Results buffer too small: need %d but only %d remaining", maxCount, results.remaining()));
+    }
+    int found = c_search_into_f32_buffer(c_ptr, query, results, maxCount);
+    // Advance position by the actual number of results, but don't exceed the
+    // buffer's remaining capacity
+    int currentPosition = results.position();
+    int newPosition = Math.min(currentPosition + found, results.limit());
+    results.position(newPosition);
+    return found;
+  }
+
   /**
    * Return the contents of the vector at key.
    *
@@ -283,6 +349,24 @@ public void add(long key, double vector[]) {
     c_add_f64(c_ptr, key, vector);
   }
 
+  /**
+   * Adds a double precision vector with a specified key to the index using
+   * ByteBuffer (zero-copy).
+   *
+   * @param key    the key associated with the vector
+   * @param vector the double precision vector data as DoubleBuffer
+   */
+  public void add(long key, java.nio.DoubleBuffer vector) {
+    if (c_ptr == 0) {
+      throw new IllegalStateException("Index already closed");
+    }
+    if (vector.remaining() != dimensions()) {
+      throw new IllegalArgumentException(String.format(
+          "Vector dimensions mismatch: expected %d but got %d", dimensions(), vector.remaining()));
+    }
+    c_add_f64_buffer(c_ptr, key, vector);
+  }
+
   /**
    * Searches for closest vectors to the specified double precision query vector.
    *
@@ -297,6 +381,55 @@ public long[] search(double vector[], long count) {
     return c_search_f64(c_ptr, vector, count);
   }
 
+  /**
+   * Searches for closest vectors to the specified double precision query vector
+   * using ByteBuffer (zero-copy).
+   *
+   * @param vector the double precision query vector data as DoubleBuffer
+   * @param count  the number of nearest neighbors to search
+   * @return an array of keys of the nearest neighbors
+   */
+  public long[] search(java.nio.DoubleBuffer vector, long count) {
+    if (c_ptr == 0) {
+      throw new IllegalStateException("Index already closed");
+    }
+    if (vector.remaining() != dimensions()) {
+      throw new IllegalArgumentException(String.format(
+          "Vector dimensions mismatch: expected %d but got %d", dimensions(), vector.remaining()));
+    }
+    return c_search_f64_buffer(c_ptr, vector, count);
+  }
+
+  /**
+   * Searches for closest vectors using zero-copy input and zero-allocation
+   * output.
+   *
+   * @param query    the query vector data as DoubleBuffer
+   * @param results  the output buffer for result keys
+   * @param maxCount maximum number of results to find
+   * @return actual number of results found
+   */
+  public int searchInto(java.nio.DoubleBuffer query, java.nio.LongBuffer results, long maxCount) {
+    if (c_ptr == 0) {
+      throw new IllegalStateException("Index already closed");
+    }
+    if (query.remaining() != dimensions()) {
+      throw new IllegalArgumentException(String.format(
+          "Query vector dimensions mismatch: expected %d but got %d", dimensions(), query.remaining()));
+    }
+    if (results.remaining() < maxCount) {
+      throw new IllegalArgumentException(String.format(
+          "Results buffer too small: need %d but only %d remaining", maxCount, results.remaining()));
+    }
+    int found = c_search_into_f64_buffer(c_ptr, query, results, maxCount);
+    // Advance position by the actual number of results, but don't exceed the
+    // buffer's remaining capacity
+    int currentPosition = results.position();
+    int newPosition = Math.min(currentPosition + found, results.limit());
+    results.position(newPosition);
+    return found;
+  }
+
   /**
    * Adds an int8 quantized vector with a specified key to the index.
    *
@@ -310,6 +443,24 @@ public void add(long key, byte vector[]) {
     c_add_i8(c_ptr, key, vector);
   }
 
+  /**
+   * Adds an int8 quantized vector with a specified key to the index using
+   * ByteBuffer (zero-copy).
+   *
+   * @param key    the key associated with the vector
+   * @param vector the int8 quantized vector data as ByteBuffer
+   */
+  public void add(long key, java.nio.ByteBuffer vector) {
+    if (c_ptr == 0) {
+      throw new IllegalStateException("Index already closed");
+    }
+    if (vector.remaining() != dimensions()) {
+      throw new IllegalArgumentException(String.format(
+          "Vector dimensions mismatch: expected %d but got %d", dimensions(), vector.remaining()));
+    }
+    c_add_i8_buffer(c_ptr, key, vector);
+  }
+
   /**
    * Searches for closest vectors to the specified int8 quantized query vector.
    *
@@ -324,6 +475,55 @@ public long[] search(byte vector[], long count) {
     return c_search_i8(c_ptr, vector, count);
   }
 
+  /**
+   * Searches for closest vectors to the specified int8 quantized query vector
+   * using ByteBuffer (zero-copy).
+   *
+   * @param vector the int8 quantized query vector data as ByteBuffer
+   * @param count  the number of nearest neighbors to search
+   * @return an array of keys of the nearest neighbors
+   */
+  public long[] search(java.nio.ByteBuffer vector, long count) {
+    if (c_ptr == 0) {
+      throw new IllegalStateException("Index already closed");
+    }
+    if (vector.remaining() != dimensions()) {
+      throw new IllegalArgumentException(String.format(
+          "Vector dimensions mismatch: expected %d but got %d", dimensions(), vector.remaining()));
+    }
+    return c_search_i8_buffer(c_ptr, vector, count);
+  }
+
+  /**
+   * Searches for closest vectors using zero-copy input and zero-allocation
+   * output.
+   *
+   * @param query    the query vector data as ByteBuffer
+   * @param results  the output buffer for result keys
+   * @param maxCount maximum number of results to find
+   * @return actual number of results found
+   */
+  public int searchInto(java.nio.ByteBuffer query, java.nio.LongBuffer results, long maxCount) {
+    if (c_ptr == 0) {
+      throw new IllegalStateException("Index already closed");
+    }
+    if (query.remaining() != dimensions()) {
+      throw new IllegalArgumentException(String.format(
+          "Query vector dimensions mismatch: expected %d but got %d", dimensions(), query.remaining()));
+    }
+    if (results.remaining() < maxCount) {
+      throw new IllegalArgumentException(String.format(
+          "Results buffer too small: need %d but only %d remaining", maxCount, results.remaining()));
+    }
+    int found = c_search_into_i8_buffer(c_ptr, query, results, maxCount);
+    // Advance position by the actual number of results, but don't exceed the
+    // buffer's remaining capacity
+    int currentPosition = results.position();
+    int newPosition = Math.min(currentPosition + found, results.limit());
+    results.position(newPosition);
+    return found;
+  }
+
   /**
    * Retrieves the vector at the specified key and populates the provided float
    * buffer.
@@ -696,6 +896,8 @@ private static native long c_create(
 
   private static native boolean c_rename(long ptr, long from, long to);
 
+  private static native long c_memory_usage(long ptr);
+
   private static native float[] c_get(long ptr, long key);
 
   // Overloaded methods:
@@ -717,4 +919,27 @@ private static native long c_create(
 
   private static native void c_get_into_i8(long ptr, long key, byte buffer[]);
 
+  // ByteBuffer overloads for zero-copy operations:
+  private static native void c_add_f32_buffer(long ptr, long key, java.nio.FloatBuffer vector);
+
+  private static native void c_add_f64_buffer(long ptr, long key, java.nio.DoubleBuffer vector);
+
+  private static native void c_add_i8_buffer(long ptr, long key, java.nio.ByteBuffer vector);
+
+  private static native long[] c_search_f32_buffer(long ptr, java.nio.FloatBuffer vector, long count);
+
+  private static native long[] c_search_f64_buffer(long ptr, java.nio.DoubleBuffer vector, long count);
+
+  private static native long[] c_search_i8_buffer(long ptr, java.nio.ByteBuffer vector, long count);
+
+  // Zero-allocation searchInto methods:
+  private static native int c_search_into_f32_buffer(long ptr, java.nio.FloatBuffer query, java.nio.LongBuffer results,
+      long maxCount);
+
+  private static native int c_search_into_f64_buffer(long ptr, java.nio.DoubleBuffer query, java.nio.LongBuffer results,
+      long maxCount);
+
+  private static native int c_search_into_i8_buffer(long ptr, java.nio.ByteBuffer query, java.nio.LongBuffer results,
+      long maxCount);
+
 }