implementing options behaviors

nosqlbench · May 9, 2024 · d9d728e · d9d728e
1 parent 464b6a8
commit d9d728e
Show file tree

Hide file tree

Showing 6 changed files with 79 additions and 1 deletion.
diff --git a/...src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiDeleteOneOpDispenser.java b/...src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiDeleteOneOpDispenser.java
@@ -59,9 +59,17 @@ private DeleteOneOptions getDeleteOneOptions(ParsedOp op, long l) {
         if (sort != null) {
             options = options.sort(sort);
         }
+        float[] vector = getVectorFromOp(op, l);
+        if (vector != null) {
+            options = options.vector(vector);
+        }
         return options;
     }
 
+    private float[] getVectorFromOp(ParsedOp op, long l) {
+        return getVectorValues(op.get("vector", l));
+    }
+
     @Override
     public DataApiBaseOp getOp(long value) {
         return opFunction.apply(value);

diff --git a/...n/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiFindOneAndDeleteOpDispenser.java b/...n/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiFindOneAndDeleteOpDispenser.java
@@ -19,6 +19,7 @@
 import com.datastax.astra.client.Database;
 import com.datastax.astra.client.model.Filter;
 import com.datastax.astra.client.model.FindOneAndDeleteOptions;
+import com.datastax.astra.client.model.Projection;
 import com.datastax.astra.client.model.Sort;
 import io.nosqlbench.adapter.dataapi.DataApiDriverAdapter;
 import io.nosqlbench.adapter.dataapi.ops.DataApiBaseOp;
@@ -59,6 +60,10 @@ private FindOneAndDeleteOptions getFindOneAndDeleteOptions(ParsedOp op, long l)
         if (sort != null) {
             options = options.sort(sort);
         }
+        Projection[] projection = getProjectionFromOp(op, l);
+        if (projection != null) {
+            options = options.projection(projection);
+        }
         return options;
     }
 

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
@@ -19,6 +19,7 @@
 import com.datastax.astra.client.Database;
 import com.datastax.astra.client.model.Filter;
 import com.datastax.astra.client.model.FindOneOptions;
+import com.datastax.astra.client.model.Projection;
 import com.datastax.astra.client.model.Sort;
 import io.nosqlbench.adapter.dataapi.DataApiDriverAdapter;
 import io.nosqlbench.adapter.dataapi.ops.DataApiBaseOp;
@@ -57,6 +58,11 @@ private FindOneOptions getFindOneOptions(ParsedOp op, long l) {
         if (sort != null) {
             options = options.sort(sort);
         }
+        Projection[] projection = getProjectionFromOp(op, l);
+        if (projection != null) {
+            options = options.projection(projection);
+        }
+        options.setIncludeSimilarity(true);
         return options;
     }
 

diff --git a/...aapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiFindOpDispenser.java b/...aapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiFindOpDispenser.java
@@ -19,6 +19,7 @@
 import com.datastax.astra.client.Database;
 import com.datastax.astra.client.model.Filter;
 import com.datastax.astra.client.model.FindOptions;
+import com.datastax.astra.client.model.Projection;
 import com.datastax.astra.client.model.Sort;
 import io.nosqlbench.adapter.dataapi.DataApiDriverAdapter;
 import io.nosqlbench.adapter.dataapi.ops.DataApiBaseOp;
@@ -57,6 +58,11 @@ private FindOptions getFindOptions(ParsedOp op, long l) {
         if (sort != null) {
             options = options.sort(sort);
         }
+        Projection[] projection = getProjectionFromOp(op, l);
+        if (projection != null) {
+            options = options.projection(projection);
+        }
+        options.setIncludeSimilarity(true);
         return options;
     }
 

diff --git a/...rc/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiInsertManyOpDispenser.java b/...rc/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiInsertManyOpDispenser.java
@@ -27,6 +27,8 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
+import java.util.Optional;
 import java.util.function.LongFunction;
 
 public class DataApiInsertManyOpDispenser extends DataApiOpDispenser {
@@ -53,7 +55,20 @@ private LongFunction<DataApiInsertManyOp> createOpFunction(ParsedOp op) {
 
     private InsertManyOptions getInsertManyOptions(ParsedOp op, long l) {
         InsertManyOptions options = new InsertManyOptions();
-
+        Optional<LongFunction<Map>> optionsFunction = op.getAsOptionalFunction("options", Map.class);
+        if (optionsFunction.isPresent()) {
+            Map<String, String> optionFields = optionsFunction.get().apply(l);
+            for(Map.Entry<String,String> entry: optionFields.entrySet()) {
+                switch(entry.getKey()) {
+                    case "chunkSize"->
+                        options = options.chunkSize(Integer.parseInt(entry.getValue()));
+                    case "concurrency" ->
+                        options = options.concurrency(Integer.parseInt(entry.getValue()));
+                    case "ordered" ->
+                        options = options.ordered(Boolean.parseBoolean(entry.getValue()));
+                }
+            }
+        }
         return options;
     }
 

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
@@ -120,4 +120,42 @@ protected Update getUpdates(ParsedOp op, long l) {
         return update;
     }
 
+    protected float[] getVectorValues(Object rawVectorValues) {
+        float[] floatValues;
+        if (rawVectorValues instanceof String) {
+            String[] rawValues = (((String) rawVectorValues).split(","));
+            floatValues = new float[rawValues.length];
+            for (int i = 0; i < rawValues.length; i++) {
+                floatValues[i] = Float.parseFloat(rawValues[i]);
+            }
+        } else {
+            throw new RuntimeException("Invalid type specified for values");
+        }
+        return floatValues;
+    }
+
+    protected Projection[] getProjectionFromOp(ParsedOp op, long l) {
+        Projection[] projection = null;
+        Optional<LongFunction<Map>> projectionFunction = op.getAsOptionalFunction("projection", Map.class);
+        if (projectionFunction.isPresent()) {
+            Map<String,List<String>> projectionFields = projectionFunction.get().apply(l);
+            for (Map.Entry<String,List<String>> field : projectionFields.entrySet()) {
+                List<String> includeFields = field.getValue();
+                StringBuffer sb = new StringBuffer();
+                for (String includeField : includeFields) {
+                    sb.append(includeField).append(",");
+                }
+                sb.deleteCharAt(sb.length() - 1);
+                if (field.getKey().equalsIgnoreCase("include")) {
+                    projection = Projections.include(sb.toString());
+                } else if (field.getKey().equalsIgnoreCase("exclude")) {
+                    projection = Projections.exclude(sb.toString());
+                } else {
+                    logger.error("Projection " + field + " not supported");
+                }
+            }
+        }
+        return projection;
+    }
+
 }