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 refItems = refResult.getFirst(); + float[] refRanks = refResult.getSecond(); + + assertThat(refItems.isEmpty(), is(false)); + + Map refRankById = new HashMap<>(refItems.size()); + for (int i = 0; i < refItems.size(); i++) { + refRankById.put(refItems.get(i).getId(), refRanks[i]); + if (i > 0) { + assertThat(refRanks[i], greaterThanOrEqualTo(refRanks[i - 1])); + } + } + + Pair, float[]> streamResult = db.query(streamingNamespace, StreamingKnnItem.class) + .where("filter_bool", EQ, true) + .whereKnn("vec", queryVector, streamParams) + .limit(limit) + .withRank() + .executeAllWithRank(); + List streamItems = streamResult.getFirst(); + float[] streamRanks = streamResult.getSecond(); + + assertThat(streamItems.size(), lessThanOrEqualTo(limit)); + assertThat(streamItems.size(), is(Math.min(refItems.size(), limit))); + + for (int i = 0; i < streamItems.size(); i++) { + StreamingKnnItem item = streamItems.get(i); + assertThat(refRankById.containsKey(item.getId()), is(true)); + assertThat(streamRanks[i], is(refRankById.get(item.getId()))); + if (i > 0) { + assertThat(streamRanks[i], greaterThanOrEqualTo(streamRanks[i - 1])); + } + } + + Pair, float[]> streamOffsetResult = db.query(streamingNamespace, StreamingKnnItem.class) + .where("filter_bool", EQ, true) + .whereKnn("vec", queryVector, streamParams) + .offset(offset) + .limit(limit) + .withRank() + .executeAllWithRank(); + List streamOffsetItems = streamOffsetResult.getFirst(); + float[] streamOffsetRanks = streamOffsetResult.getSecond(); + + assertThat(streamOffsetItems.size(), lessThanOrEqualTo(limit)); + for (int i = 0; i < streamOffsetItems.size(); i++) { + StreamingKnnItem item = streamOffsetItems.get(i); + assertThat(refRankById.containsKey(item.getId()), is(true)); + assertThat(streamOffsetRanks[i], is(refRankById.get(item.getId()))); + if (i > 0) { + assertThat(streamOffsetRanks[i], greaterThanOrEqualTo(streamOffsetRanks[i - 1])); + } + } + } + + private static float[] randVect(int dimension) { + float[] vector = new float[dimension]; + for (int i = 0; i < dimension; i++) { + vector[i] = ThreadLocalRandom.current().nextFloat(); + } + return vector; + } + private static List getTestVectorItems() { return Arrays.asList( new VectorItem(0, new float[]{0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}), @@ -268,6 +359,22 @@ private static List getTestVectorItems() { ); } + @Getter + @Setter + @NoArgsConstructor + @AllArgsConstructor + public static class StreamingKnnItem { + @Reindex(name = "id", isPrimaryKey = true) + private Integer id; + + @Reindex(name = "filter_bool") + private boolean filterBool; + + @Reindex(name = "vec") + @Hnsw(metric = Metric.L2, dimension = 128, startSize = 500, m = 16, efConstruction = 200) + private float[] vec; + } + @Getter @Setter @NoArgsConstructor