IntenseFinance/nosqlbench
3
0
Fork 0
mirror of https://github.com/nosqlbench/nosqlbench.git synced 2025-02-25 18:55:28 -06:00
Code Issues Packages Projects Releases Wiki Activity

updated FindOptions to include all possible options

Browse Source
This commit is contained in:
Mark Wolters 2024-06-11 13:18:02 -04:00
parent 93dfaec743
commit f37796b8a0
3 changed files with 17 additions and 38 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiFindOpDispenser.java
View File

@ -28,6 +28,7 @@ import io.nosqlbench.adapters.api.templating.ParsedOp;
import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger; import org.apache.logging.log4j.Logger;
import java.util.Optional;
import java.util.function.LongFunction; import java.util.function.LongFunction;
public class DataApiFindOpDispenser extends DataApiOpDispenser { public class DataApiFindOpDispenser extends DataApiOpDispenser {
@ -67,7 +68,22 @@ public class DataApiFindOpDispenser extends DataApiOpDispenser {
if (projection != null) { if (projection != null) {
options = options.projection(projection); options = options.projection(projection);
} }
options.setIncludeSimilarity(true); Optional<LongFunction<Integer>> limitFunction = op.getAsOptionalFunction("limit", Integer.class);
if (limitFunction.isPresent()) {
options = options.limit(limitFunction.get().apply(l));
}
Optional<LongFunction<Integer>> skipFunction = op.getAsOptionalFunction("skip", Integer.class);
if (skipFunction.isPresent()) {
options = options.skip(skipFunction.get().apply(l));
}
Optional<LongFunction<Boolean>> includeSimilarityFunction = op.getAsOptionalFunction("includeSimilarity", Boolean.class);
if (includeSimilarityFunction.isPresent()) {
options.setIncludeSimilarity(includeSimilarityFunction.get().apply(l));
}
Optional<LongFunction<String>> pageStateFunction = op.getAsOptionalFunction("pageState", String.class);
if (pageStateFunction.isPresent()) {
options.setPageState(pageStateFunction.get().apply(l));
}
return options; return options;
} }

20
nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiFindVectorFilterOpDispenser.java
View File

@ -18,9 +18,6 @@ package io.nosqlbench.adapter.dataapi.opdispensers;
import com.datastax.astra.client.Database; import com.datastax.astra.client.Database;
import com.datastax.astra.client.model.Filter; 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.DataApiDriverAdapter;
import io.nosqlbench.adapter.dataapi.ops.DataApiBaseOp; import io.nosqlbench.adapter.dataapi.ops.DataApiBaseOp;
import io.nosqlbench.adapter.dataapi.ops.DataApiFindVectorFilterOp; import io.nosqlbench.adapter.dataapi.ops.DataApiFindVectorFilterOp;
@ -58,23 +55,6 @@ public class DataApiFindVectorFilterOpDispenser extends DataApiOpDispenser {
return op.getConfigOr("limit", 100, l); return op.getConfigOr("limit", 100, l);
} }
private FindOptions getFindOptions(ParsedOp op, long l) {
FindOptions options = new FindOptions();
Sort sort = getSortFromOp(op, l);
float[] vector = getVectorValues(op, l);
if (sort != null) {
options = vector != null ? options.sort(vector, sort) : options.sort(sort);
} else if (vector != null) {
options = options.sort(vector);
}
Projection[] projection = getProjectionFromOp(op, l);
if (projection != null) {
options = options.projection(projection);
}
options.setIncludeSimilarity(true);
return options;
}
@Override @Override
public DataApiBaseOp getOp(long value) { public DataApiBaseOp getOp(long value) {
return opFunction.apply(value); return opFunction.apply(value);

17
nb-adapters/adapter-dataapi/src/main/java/io/nosqlbench/adapter/dataapi/opdispensers/DataApiFindVectorOpDispenser.java
View File

@ -55,23 +55,6 @@ public class DataApiFindVectorOpDispenser extends DataApiOpDispenser {
return op.getConfigOr("limit", 100, l); return op.getConfigOr("limit", 100, l);
} }
private FindOptions getFindOptions(ParsedOp op, long l) {
FindOptions options = new FindOptions();
Sort sort = getSortFromOp(op, l);
float[] vector = getVectorValues(op, l);
if (sort != null) {
options = vector != null ? options.sort(vector, sort) : options.sort(sort);
} else if (vector != null) {
options = options.sort(vector);
}
Projection[] projection = getProjectionFromOp(op, l);
if (projection != null) {
options = options.projection(projection);
}
options.setIncludeSimilarity(true);
return options;
}
@Override @Override
public DataApiBaseOp getOp(long value) { public DataApiBaseOp getOp(long value) {
return opFunction.apply(value); return opFunction.apply(value);