vectorizing find one op

nosqlbench · May 10, 2024 · d5a96e8 · d5a96e8
1 parent 085b7b9
commit d5a96e8
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 1 deletion.
diff --git a/...i/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiFindOneOpDispenser.java b/...i/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiFindOneOpDispenser.java
@@ -55,8 +55,9 @@ private LongFunction<DataApiFindOneOp> createOpFunction(ParsedOp op) {
     private FindOneOptions getFindOneOptions(ParsedOp op, long l) {
         FindOneOptions options = new FindOneOptions();
         Sort sort = getSortFromOp(op, l);
+        float[] vector = getVectorValues(op, l);
         if (sort != null) {
-            options = options.sort(sort);
+            options = vector != null ? options.sort(vector, sort) : options.sort(sort);
         }
         Projection[] projection = getProjectionFromOp(op, l);
         if (projection != null) {

diff --git a/...-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiOpDispenser.java b/...-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiOpDispenser.java
@@ -127,6 +127,11 @@ protected Update getUpdates(ParsedOp op, long l) {
         return update;
     }
 
+    protected float[] getVectorValues(ParsedOp op, long l) {
+        Object rawVectorValues = op.get("vector", l);
+        return getVectorValues(rawVectorValues);
+    }
+
     protected float[] getVectorValues(Object rawVectorValues) {
         float[] floatValues;
         if (rawVectorValues instanceof String) {
@@ -135,12 +140,24 @@ protected float[] getVectorValues(Object rawVectorValues) {
             for (int i = 0; i < rawValues.length; i++) {
                 floatValues[i] = Float.parseFloat(rawValues[i]);
             }
+        } else if (rawVectorValues instanceof List) {
+            return getVectorValuesList(rawVectorValues);
         } else {
             throw new RuntimeException("Invalid type specified for values");
         }
         return floatValues;
     }
 
+    protected float[] getVectorValuesList(Object rawVectorValues) {
+        float[] vectorValues = null;
+        List<Object> vectorValuesList = (List<Object>) rawVectorValues;
+        vectorValues = new float[vectorValuesList.size()];
+        for (int i = 0; i < vectorValuesList.size(); i++) {
+            vectorValues[i] = Float.parseFloat(vectorValuesList.get(i).toString());
+        }
+        return vectorValues;
+    }
+
     protected Projection[] getProjectionFromOp(ParsedOp op, long l) {
         Projection[] projection = null;
         Optional<LongFunction<Map>> projectionFunction = op.getAsOptionalFunction("projection", Map.class);