diff --git a/src/main/java/ru/rt/restream/reindexer/vector/params/BaseKnnSearchParam.java b/src/main/java/ru/rt/restream/reindexer/vector/params/BaseKnnSearchParam.java index a644532..73664d6 100644 --- a/src/main/java/ru/rt/restream/reindexer/vector/params/BaseKnnSearchParam.java +++ b/src/main/java/ru/rt/restream/reindexer/vector/params/BaseKnnSearchParam.java @@ -30,7 +30,7 @@ * Common parameters for all types of KNN indices. * *
If all parameters are specified, the filtering will be performed in such a way that all conditions are met.
- * At least one of these parameters must be specified.
+ * At least one of these parameters must be specified, except for HNSW streaming search where both may be omitted.
*/
@Getter
@AllArgsConstructor(access = AccessLevel.PACKAGE)
@@ -58,11 +58,15 @@ public class BaseKnnSearchParam implements KnnSearchParam {
public void serializeBy(ByteBuffer buffer) {
buffer.putVarUInt32(KNN_QUERY_TYPE_BASE)
.putVarUInt32(KNN_QUERY_PARAMS_VERSION);
- serializeKAndRadius(buffer);
+ serializeKAndRadius(buffer, false);
}
void serializeKAndRadius(ByteBuffer buffer) {
- checkValues();
+ serializeKAndRadius(buffer, false);
+ }
+
+ void serializeKAndRadius(ByteBuffer buffer, boolean allowEmpty) {
+ checkValues(allowEmpty);
int mask = 0;
if (k != null) {
mask |= KNN_SERIALIZE_WITH_K;
@@ -79,8 +83,8 @@ void serializeKAndRadius(ByteBuffer buffer) {
}
}
- private void checkValues() {
- if (k == null && radius == null) {
+ private void checkValues(boolean allowEmpty) {
+ if (k == null && radius == null && !allowEmpty) {
throw new IllegalArgumentException("Both params (k and radius) cannot be null");
}
if (k != null && k <= 0) {
diff --git a/src/main/java/ru/rt/restream/reindexer/vector/params/IndexHnswSearchParam.java b/src/main/java/ru/rt/restream/reindexer/vector/params/IndexHnswSearchParam.java
index 5144c71..be32438 100644
--- a/src/main/java/ru/rt/restream/reindexer/vector/params/IndexHnswSearchParam.java
+++ b/src/main/java/ru/rt/restream/reindexer/vector/params/IndexHnswSearchParam.java
@@ -52,7 +52,7 @@ public class IndexHnswSearchParam implements KnnSearchParam {
public void serializeBy(ByteBuffer buffer) {
buffer.putVarUInt32(KNN_QUERY_TYPE_HNSW)
.putVarUInt32(KNN_QUERY_PARAMS_VERSION);
- base.serializeKAndRadius(buffer);
+ base.serializeKAndRadius(buffer, true);
buffer.putVarInt32(ef);
}
diff --git a/src/main/java/ru/rt/restream/reindexer/vector/params/KnnParams.java b/src/main/java/ru/rt/restream/reindexer/vector/params/KnnParams.java
index d399962..c5cbaa4 100644
--- a/src/main/java/ru/rt/restream/reindexer/vector/params/KnnParams.java
+++ b/src/main/java/ru/rt/restream/reindexer/vector/params/KnnParams.java
@@ -21,6 +21,10 @@
* Factories for KnnSearchParams.
*/
public class KnnParams {
+ public static BaseKnnSearchParam base() {
+ return new BaseKnnSearchParam(null, null);
+ }
+
@Deprecated
public static BaseKnnSearchParam base(int k) {
checkK(k);
@@ -49,8 +53,12 @@ public static IndexHnswSearchParam hnsw(int k, int ef) {
}
public static IndexHnswSearchParam hnsw(@NonNull BaseKnnSearchParam base, int ef) {
- if (base.getK() != null && ef < base.getK()) {
- throw new IllegalArgumentException("Minimal value of 'ef' must be greater than or equal to 'k'");
+ if (base.getK() != null) {
+ if (ef < base.getK()) {
+ throw new IllegalArgumentException("Minimal value of 'ef' must be greater than or equal to 'k'");
+ }
+ } else if (ef <= 0) {
+ throw new IllegalArgumentException("Minimal value of 'ef' must be greater than 0");
}
return new IndexHnswSearchParam(base, ef);
}
diff --git a/src/test/java/ru/rt/restream/reindexer/connector/FloatVectorHnswTest.java b/src/test/java/ru/rt/restream/reindexer/connector/FloatVectorHnswTest.java
index 486371c..2ce5e5d 100644
--- a/src/test/java/ru/rt/restream/reindexer/connector/FloatVectorHnswTest.java
+++ b/src/test/java/ru/rt/restream/reindexer/connector/FloatVectorHnswTest.java
@@ -30,14 +30,21 @@
import ru.rt.restream.reindexer.annotations.Reindex;
import ru.rt.restream.reindexer.db.DbBaseTest;
import ru.rt.restream.reindexer.exceptions.ReindexerException;
+import ru.rt.restream.reindexer.util.Pair;
+import ru.rt.restream.reindexer.vector.params.IndexHnswSearchParam;
import ru.rt.restream.reindexer.vector.params.KnnParams;
import java.util.Arrays;
+import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ThreadLocalRandom;
import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.greaterThanOrEqualTo;
import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.lessThanOrEqualTo;
import static org.hamcrest.Matchers.notNullValue;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static ru.rt.restream.reindexer.Query.Condition.EQ;
@@ -242,6 +249,90 @@ public void testSearchWithNotHnswNorBaseParams_isException() {
.toList());
}
+ @Test
+ public void testStreamingKnnHnsw_isOk() {
+ final int maxElements = 500;
+ final int limit = 10;
+ final int offset = 5;
+ final String streamingNamespace = "items_streaming_knn";
+
+ db.openNamespace(streamingNamespace, NamespaceOptions.defaultOptions(), StreamingKnnItem.class);
+ for (int id = 0; id < maxElements; id++) {
+ db.insert(streamingNamespace, new StreamingKnnItem(id, id % 2 == 0, randVect(128)));
+ }
+
+ float[] queryVector = randVect(128);
+ IndexHnswSearchParam streamParams = KnnParams.hnsw(KnnParams.base(), 1);
+ IndexHnswSearchParam refParams = KnnParams.hnsw(KnnParams.k(maxElements), maxElements * 2);
+
+ Pair, float[]> refResult = db.query(streamingNamespace, StreamingKnnItem.class)
+ .where("filter_bool", EQ, true)
+ .whereKnn("vec", queryVector, refParams)
+ .withRank()
+ .executeAllWithRank();
+ List
, float[]> streamResult = db.query(streamingNamespace, StreamingKnnItem.class)
+ .where("filter_bool", EQ, true)
+ .whereKnn("vec", queryVector, streamParams)
+ .limit(limit)
+ .withRank()
+ .executeAllWithRank();
+ List
, float[]> streamOffsetResult = db.query(streamingNamespace, StreamingKnnItem.class)
+ .where("filter_bool", EQ, true)
+ .whereKnn("vec", queryVector, streamParams)
+ .offset(offset)
+ .limit(limit)
+ .withRank()
+ .executeAllWithRank();
+ List