From d9d728eadd87f72e747cbe9ebf2e5c49245ee69f Mon Sep 17 00:00:00 2001 From: Mark Wolters Date: Thu, 9 May 2024 17:42:29 -0400 Subject: [PATCH] implementing options behaviors --- .../DataApiDeleteOneOpDispenser.java | 8 ++++ .../DataApiFindOneAndDeleteOpDispenser.java | 5 +++ .../DataApiFindOneOpDispenser.java | 6 +++ .../opdispensers/DataApiFindOpDispenser.java | 6 +++ .../DataApiInsertManyOpDispenser.java | 17 ++++++++- .../opdispensers/DataApiOpDispenser.java | 38 +++++++++++++++++++ 6 files changed, 79 insertions(+), 1 deletion(-) diff --git a/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiDeleteOneOpDispenser.java b/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiDeleteOneOpDispenser.java index 2cb71a3c2..c5044a007 100644 --- a/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiDeleteOneOpDispenser.java +++ b/nb-adapters/adapter-dataapi/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/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiFindOneAndDeleteOpDispenser.java b/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiFindOneAndDeleteOpDispenser.java index b63a6386c..7a85f9995 100644 --- a/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiFindOneAndDeleteOpDispenser.java +++ b/nb-adapters/adapter-dataapi/src/main/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/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiFindOneOpDispenser.java b/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiFindOneOpDispenser.java index a07c0be97..7731672bf 100644 --- a/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiFindOneOpDispenser.java +++ b/nb-adapters/adapter-dataapi/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/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiFindOpDispenser.java b/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiFindOpDispenser.java index 3c9af8fa3..4c7bfd3c7 100644 --- a/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiFindOpDispenser.java +++ b/nb-adapters/adapter-dataapi/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/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiInsertManyOpDispenser.java b/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiInsertManyOpDispenser.java index e2aea7d51..10788cfea 100644 --- a/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiInsertManyOpDispenser.java +++ b/nb-adapters/adapter-dataapi/src/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 createOpFunction(ParsedOp op) { private InsertManyOptions getInsertManyOptions(ParsedOp op, long l) { InsertManyOptions options = new InsertManyOptions(); - + Optional> optionsFunction = op.getAsOptionalFunction("options", Map.class); + if (optionsFunction.isPresent()) { + Map optionFields = optionsFunction.get().apply(l); + for(Map.Entry 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/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiOpDispenser.java b/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiOpDispenser.java index f249824d1..f13559a56 100644 --- a/nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiOpDispenser.java +++ b/nb-adapters/adapter-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> projectionFunction = op.getAsOptionalFunction("projection", Map.class); + if (projectionFunction.isPresent()) { + Map> projectionFields = projectionFunction.get().apply(l); + for (Map.Entry> field : projectionFields.entrySet()) { + List 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; + } + }