Skip to content

HNSW streaming KNN search support #138

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
maximbogatyrev wants to merge 1 commit into
base: master
Choose a base branch
from
Open

HNSW streaming KNN search support #138

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 9 additions & 5 deletions src/main/java/ru/rt/restream/reindexer/vector/params/BaseKnnSearchParam.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@
* Common parameters for all types of KNN indices.
*
* <p>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)
Expand Down Expand Up @@ -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;
Expand All @@ -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) {
Expand Down
2 changes: 1 addition & 1 deletion src/main/java/ru/rt/restream/reindexer/vector/params/IndexHnswSearchParam.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}

Expand Down
12 changes: 10 additions & 2 deletions src/main/java/ru/rt/restream/reindexer/vector/params/KnnParams.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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");

@MadSchemas MadSchemas Jul 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

А откуда появилось вот это ограничение? Насколько я помню, в ядро для стриминг-режима можно передать ef=0 и там пройдёт выборка

@maximbogatyrev maximbogatyrev Jul 10, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ну, вообще вроде не должны проходить валидацию на этапе разбора knn-параметров (ни через десериализацию, ни через c++ query-builder):

void KnnSearchParams::Validate() const {
	// Empty 'k' and empty 'radius' both are allowed for the streaming KNN search over an HNSW index
	std::visit(
		overloaded{[](const KnnSearchParamsBase& p) { checkKNotZero(p); }, [](const BruteForceSearchParams& p) { checkKandRadius(p); },
				   [](const HnswSearchParams& p) {
					   checkKNotZero(p);
					   if (p.K() && p.Ef() < *p.K()) {
						   throw Error{errQueryExec, "Ef should not be less than k in hnsw query"};
					   }
					   if (p.Ef() < 1) {
						   throw Error{errQueryExec, "Value of 'ef' - {} is out of bounds: [1,{}]", p.Ef(),
									   std::numeric_limits<size_t>::max()};
					   }
				   },
				   [](const IvfSearchParams& p) {
					   checkKandRadius(p);
					   if (p.NProbe() == 0) {
						   throw Error{errQueryExec, "Nprobe should not be 0"};
					   }
				   }},
		toVariant());
}

}
return new IndexHnswSearchParam(base, ef);
}
Expand Down
107 changes: 107 additions & 0 deletions src/test/java/ru/rt/restream/reindexer/connector/FloatVectorHnswTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<List<StreamingKnnItem>, float[]> refResult = db.query(streamingNamespace, StreamingKnnItem.class)
.where("filter_bool", EQ, true)
.whereKnn("vec", queryVector, refParams)
.withRank()
.executeAllWithRank();
List<StreamingKnnItem> refItems = refResult.getFirst();
float[] refRanks = refResult.getSecond();

assertThat(refItems.isEmpty(), is(false));

Map<Integer, Float> 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<List<StreamingKnnItem>, float[]> streamResult = db.query(streamingNamespace, StreamingKnnItem.class)
.where("filter_bool", EQ, true)
.whereKnn("vec", queryVector, streamParams)
.limit(limit)
.withRank()
.executeAllWithRank();
List<StreamingKnnItem> 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<List<StreamingKnnItem>, float[]> streamOffsetResult = db.query(streamingNamespace, StreamingKnnItem.class)
.where("filter_bool", EQ, true)
.whereKnn("vec", queryVector, streamParams)
.offset(offset)
.limit(limit)
.withRank()
.executeAllWithRank();
List<StreamingKnnItem> 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<VectorItem> 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}),
Expand All @@ -268,6 +359,22 @@ private static List<VectorItem> 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
Expand Down
Loading