updated FindOptions to include all possible options

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

View File

@ -28,6 +28,7 @@ import io.nosqlbench.adapters.api.templating.ParsedOp;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import java.util.Optional;
import java.util.function.LongFunction;
public class DataApiFindOpDispenser extends DataApiOpDispenser {
@ -67,7 +68,22 @@ public class DataApiFindOpDispenser extends DataApiOpDispenser {
if (projection != null) {
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;
}

View File

@ -18,9 +18,6 @@ package io.nosqlbench.adapter.dataapi.opdispensers;
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;
import io.nosqlbench.adapter.dataapi.ops.DataApiFindVectorFilterOp;
@ -58,23 +55,6 @@ public class DataApiFindVectorFilterOpDispenser extends DataApiOpDispenser {
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
public DataApiBaseOp getOp(long value) {
return opFunction.apply(value);

View File

@ -55,23 +55,6 @@ public class DataApiFindVectorOpDispenser extends DataApiOpDispenser {
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
public DataApiBaseOp getOp(long value) {
return opFunction.apply(value);