vectorizing find one op

This commit is contained in:
Mark Wolters 2024-05-10 17:49:20 -04:00
parent 085b7b969e
commit d5a96e81c4
2 changed files with 19 additions and 1 deletions

View File

@ -55,8 +55,9 @@ public class DataApiFindOneOpDispenser extends DataApiOpDispenser {
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) {

View File

@ -127,6 +127,11 @@ public abstract class DataApiOpDispenser extends BaseOpDispenser<DataApiBaseOp,
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 @@ public abstract class DataApiOpDispenser extends BaseOpDispenser<DataApiBaseOp,
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);