@@ -158,6 +183,9 @@
true
false
+
+ ${project.basedir}/license-header.txt
+
@@ -219,6 +247,7 @@
jar
+ com.nvidia.cuvs.lucene
all,-missing
${project.build.directory}/javadocs
diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java
index ccc61eae..1abb9e58 100644
--- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java
+++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java
@@ -1,5 +1,5 @@
/*
- * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*/
package com.nvidia.cuvs.lucene;
@@ -7,6 +7,7 @@
import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.assertIsSupported;
import com.nvidia.cuvs.LibraryException;
+import com.nvidia.cuvs.spi.CuVSProvider;
import java.io.IOException;
import org.apache.lucene.codecs.KnnVectorsFormat;
import org.apache.lucene.codecs.KnnVectorsReader;
@@ -39,6 +40,7 @@ public class CuVS2510GPUVectorsFormat extends KnnVectorsFormat {
static {
try {
+ CuVSProvider.provider().enableRMMAsyncMemory();
LUCENE_PROVIDER = LuceneProvider.getInstance("99");
FLAT_VECTORS_FORMAT =
LUCENE_PROVIDER.getLuceneFlatVectorsFormatInstance(DefaultFlatVectorScorer.INSTANCE);
diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java
index c783660d..8c29fff3 100644
--- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java
+++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java
@@ -1,5 +1,5 @@
/*
- * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*/
package com.nvidia.cuvs.lucene;
@@ -430,7 +430,10 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits
}
}
topK = Math.min(knnCollector.k() + 10, mask[0].cardinality());
- maskLength = mask[0].length();
+ // numDocs must be the total vector count so cuVS sizes the prefilter to cover every ordinal.
+ // BitSet.length() is (highest set bit + 1), which under a selective filter is smaller than the
+ // vector count, leaving the trailing ordinals outside the filter and thus default-accepted.
+ maskLength = acceptedOrds.length();
}
try {
@@ -443,6 +446,9 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits
new CagraSearchParams.Builder()
.withItopkSize(Math.max(collector.getiTopK(), topK))
.withSearchWidth(collector.getSearchWidth())
+ .withThreadBlockSize(collector.getThreadBlockSize())
+ .withMaxIterations(collector.getMaxIterations())
+ .withAlgo(collector.getSearchAlgo())
.build();
} else {
// Setting itopK as topK because in any case iTopK should be ATLEAST equal to topK
@@ -509,7 +515,11 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits
if (knnCollector.earlyTerminated()) {
break;
}
- if (ord < 0) {
+ // Empty top-k slots (fewer than k passing candidates) carry a sentinel distance of FLT_MAX.
+ // Prefer this over the neighbor-index sentinel: the index sentinel is not uniform across
+ // CAGRA search algorithms (single-CTA emits 0x7FFFFFFF, multi-CTA 0xFFFFFFFF), so the
+ // distance is the reliable, algorithm-independent signal for an empty slot.
+ if (score == Float.MAX_VALUE) {
continue;
}
float correctedScore = scoreCorrectionFunction.apply(score);
@@ -600,6 +610,22 @@ private static void checkVersion(int versionMeta, int versionVectorData, IndexIn
}
}
+ /**
+ * Returns the {@link CagraIndex} for the given field, or {@code null} if unavailable
+ * (e.g., during a merge or when the field is missing).
+ *
+ * @param field the vector field name
+ * @return the CAGRA index, or {@code null}
+ */
+ public CagraIndex getCagraIndexForField(String field) {
+ if (cuvsIndices == null) return null;
+ FieldInfo info = fieldInfos.fieldInfo(field);
+ if (info == null) return null;
+ GPUIndex gpuIndex = cuvsIndices.get(info.number);
+ if (gpuIndex == null) return null;
+ return gpuIndex.getCagraIndex();
+ }
+
/**
* Gets the instance of FieldInfos.
*
diff --git a/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java b/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java
new file mode 100644
index 00000000..7e2a3528
--- /dev/null
+++ b/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java
@@ -0,0 +1,141 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package com.nvidia.cuvs.lucene;
+
+import com.nvidia.cuvs.FilterBitsetHandle;
+import java.io.IOException;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import org.apache.lucene.search.Query;
+
+/**
+ * Shared LRU cache mapping (filter Query, per-segment reader keys, field) → {@link
+ * FilterBitsetHandle}.
+ *
+ * Host-side cache holding packed bitset arrays; the device-side upload is managed inside {@link
+ * FilterBitsetHandle} itself. Entries are evicted in LRU order.
+ *
+ *
Concurrency
+ *
+ * Values are stored as {@link CompletableFuture}s so that concurrent misses on the same key
+ * build the handle exactly once: the first thread inserts an incomplete future and builds; others
+ * find the future and await it (without holding the cache lock, so builds for different keys still
+ * run in parallel).
+ *
+ *
Lifetime is reference-counted (see {@link FilterBitsetHandle}). A cached handle carries one
+ * cache-owned reference, released via {@link FilterBitsetHandle#close()} when the entry is evicted.
+ * {@link #acquire} returns a handle with an additional reference already taken; the caller
+ * must {@link FilterBitsetHandle#decRef()} it after use. Because the free happens only when the last
+ * reference is dropped, an eviction concurrent with an in-flight search cannot free a handle still
+ * in use, and a replaced entry cannot leak.
+ */
+final class FilterBitsetCache {
+
+ static final int MAX_HOST_ENTRIES = 16;
+
+ /** Builds the handle for a key on a cache miss. */
+ @FunctionalInterface
+ interface FilterBuilder {
+ FilterBitsetHandle build() throws IOException;
+ }
+
+ private record FilterCacheKey(Query filter, List segReaderKeys, String field) {
+ @Override
+ public boolean equals(Object o) {
+ if (!(o instanceof FilterCacheKey other)) return false;
+ return Objects.equals(filter, other.filter)
+ && Objects.equals(segReaderKeys, other.segReaderKeys)
+ && Objects.equals(field, other.field);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(filter, segReaderKeys, field);
+ }
+ }
+
+ private static final LinkedHashMap> CACHE =
+ new LinkedHashMap<>(MAX_HOST_ENTRIES + 2, 0.75f, /* access-order= */ true) {
+ @Override
+ protected boolean removeEldestEntry(
+ Map.Entry> eldest) {
+ if (size() > MAX_HOST_ENTRIES) {
+ releaseCacheRef(eldest.getValue());
+ return true;
+ }
+ return false;
+ }
+ };
+
+ private FilterBitsetCache() {}
+
+ /**
+ * Returns the handle for the given key, building it once if absent. The returned handle has a
+ * reference taken on the caller's behalf that must be released with {@link
+ * FilterBitsetHandle#decRef()} once the search completes.
+ */
+ static FilterBitsetHandle acquire(
+ Query filter, List segReaderKeys, String field, FilterBuilder builder)
+ throws IOException {
+ FilterCacheKey key = new FilterCacheKey(filter, segReaderKeys, field);
+ while (true) {
+ CompletableFuture future;
+ boolean owner = false;
+ synchronized (FilterBitsetCache.class) {
+ future = CACHE.get(key);
+ if (future == null) {
+ future = new CompletableFuture<>();
+ CACHE.put(key, future);
+ owner = true;
+ }
+ }
+
+ if (owner) {
+ // Build outside the cache lock so misses on other keys are not serialized behind this one.
+ try {
+ future.complete(builder.build());
+ } catch (Throwable t) {
+ // Drop the failed entry so a later call rebuilds, then propagate to this thread and any
+ // waiters (which observe the exception via join() and retry).
+ synchronized (FilterBitsetCache.class) {
+ CACHE.remove(key, future);
+ }
+ future.completeExceptionally(t);
+ if (t instanceof IOException io) throw io;
+ if (t instanceof RuntimeException re) throw re;
+ if (t instanceof Error err) throw err;
+ throw new RuntimeException(t);
+ }
+ }
+
+ FilterBitsetHandle handle;
+ try {
+ handle = future.join();
+ } catch (CompletionException e) {
+ // The owning thread's build failed and removed the entry; retry so exactly one thread
+ // rebuilds.
+ continue;
+ }
+
+ if (handle.tryIncRef()) {
+ return handle;
+ }
+ // Rare: the entry was evicted and its cache reference released between completion and our
+ // acquisition. It is already gone from the map, so retry and rebuild.
+ }
+ }
+
+ /** Releases the cache's reference once the (possibly in-flight) build completes. */
+ private static void releaseCacheRef(CompletableFuture future) {
+ future.whenComplete(
+ (handle, err) -> {
+ if (handle != null) handle.close();
+ });
+ }
+}
diff --git a/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java b/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java
index 07e64fb1..b237d4bc 100644
--- a/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java
+++ b/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java
@@ -1,5 +1,5 @@
/*
- * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*/
package com.nvidia.cuvs.lucene;
@@ -9,19 +9,23 @@
import com.nvidia.cuvs.CagraIndexParams;
import com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType;
import com.nvidia.cuvs.CagraIndexParams.HnswHeuristicType;
+import com.nvidia.cuvs.CagraQuery;
import com.nvidia.cuvs.CuVSDeviceMatrix;
import com.nvidia.cuvs.CuVSHostMatrix;
import com.nvidia.cuvs.CuVSMatrix;
import com.nvidia.cuvs.CuVSMatrix.Builder;
import com.nvidia.cuvs.CuVSMatrix.DataType;
import com.nvidia.cuvs.CuVSResources;
+import com.nvidia.cuvs.FilterBitsetHandle;
import com.nvidia.cuvs.GPUInfoProvider;
import com.nvidia.cuvs.HnswIndex;
import com.nvidia.cuvs.HnswIndexParams;
+import com.nvidia.cuvs.MultiPartitionSearchResults;
import com.nvidia.cuvs.TieredIndex;
import com.nvidia.cuvs.spi.CuVSProvider;
import java.lang.invoke.MethodHandle;
import java.nio.file.Path;
+import java.util.List;
import java.util.logging.Level;
class FilterCuVSProvider implements CuVSProvider {
@@ -54,6 +58,22 @@ public CagraIndex.Builder newCagraIndexBuilder(CuVSResources cuVSResources)
return delegate.newCagraIndexBuilder(cuVSResources);
}
+ @Override
+ public FilterBitsetHandle newFilterBitsetHandle(long[] combinedLongs) {
+ return delegate.newFilterBitsetHandle(combinedLongs);
+ }
+
+ @Override
+ public MultiPartitionSearchResults searchCagraMultiPartition(
+ CuVSResources resources,
+ List indices,
+ CagraQuery query,
+ int k,
+ FilterBitsetHandle filter)
+ throws Throwable {
+ return delegate.searchCagraMultiPartition(resources, indices, query, k, filter);
+ }
+
@Override
public HnswIndex.Builder newHnswIndexBuilder(CuVSResources cuVSResources)
throws UnsupportedOperationException {
@@ -145,6 +165,11 @@ public HnswIndex hnswIndexFromCagra(HnswIndexParams arg0, CagraIndex arg1) throw
return delegate.hnswIndexFromCagra(arg0, arg1);
}
+ @Override
+ public void enableRMMAsyncMemory() {
+ delegate.enableRMMAsyncMemory();
+ }
+
@Override
public void enableRMMManagedPooledMemory(int arg0, int arg1) {
delegate.enableRMMManagedPooledMemory(arg0, arg1);
diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java
index 6a9daae7..0cd1802d 100644
--- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java
+++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java
@@ -1,20 +1,69 @@
/*
- * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*/
package com.nvidia.cuvs.lucene;
+import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.getCuVSResourcesInstance;
+
+import com.nvidia.cuvs.CagraIndex;
+import com.nvidia.cuvs.CagraQuery;
+import com.nvidia.cuvs.CagraSearchParams;
+import com.nvidia.cuvs.CuVSMatrix;
+import com.nvidia.cuvs.CuVSResources;
+import com.nvidia.cuvs.FilterBitsetHandle;
+import com.nvidia.cuvs.MultiPartitionCagraSearch;
+import com.nvidia.cuvs.MultiPartitionSearchResults;
import java.io.IOException;
-import org.apache.lucene.index.LeafReader;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Comparator;
+import java.util.List;
+import org.apache.lucene.codecs.KnnVectorsReader;
+import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat;
+import org.apache.lucene.index.CodecReader;
+import org.apache.lucene.index.FilterLeafReader;
+import org.apache.lucene.index.FloatVectorValues;
+import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.search.Explanation;
+import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.KnnFloatVectorQuery;
+import org.apache.lucene.search.MatchNoDocsQuery;
import org.apache.lucene.search.Query;
+import org.apache.lucene.search.QueryVisitor;
+import org.apache.lucene.search.ScoreDoc;
+import org.apache.lucene.search.ScoreMode;
+import org.apache.lucene.search.Scorer;
+import org.apache.lucene.search.ScorerSupplier;
import org.apache.lucene.search.TopDocs;
+import org.apache.lucene.search.Weight;
import org.apache.lucene.search.knn.KnnCollectorManager;
import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.FixedBitSet;
/**
- * Extends upon KnnFloatVectorQuery for GPU-only search.
+ * Extends {@link KnnFloatVectorQuery} for GPU-only search.
+ *
+ * When all index segments use {@link CuVS2510GPUVectorsReader}, {@link #rewrite} delegates a
+ * single multi-partition search to cuVS, passing one Lucene segment per cuVS partition. cuVS
+ * runs the per-partition CAGRA searches, applies distance post-processing, and performs the
+ * cross-partition top-k merge internally; the returned arrays are mapped to Lucene doc IDs on
+ * the host. The effective CAGRA algorithm (SINGLE_CTA or MULTI_KERNEL) is selected by cuVS
+ * based on {@code searchAlgo} and {@code itopk_size}, with MULTI_KERNEL handling k beyond
+ * SINGLE_CTA's per-partition cap.
+ *
+ *
If the query has an explicit {@code filter}, or if any segment carries live-document deletes,
+ * the combined acceptance mask (filter ∩ liveDocs) is packed across all segments into a single
+ * {@link FilterBitsetHandle}. The host-side packed arrays are cached per unique
+ * (filter, reader-state, field) triple via {@link FilterBitsetCache}; the device upload is cached
+ * inside the handle itself across threads.
+ *
+ *
Falls back to the standard per-segment Lucene path when the optimized path cannot be
+ * applied: mixed segment types, a missing CAGRA index for the field on any segment, or segments
+ * whose built CAGRA graphs differ in degree (a single multi-partition request requires a uniform
+ * graph degree, and a small segment can have its degree truncated at build time).
*
* @since 25.10
*/
@@ -22,24 +71,192 @@ public class GPUKnnFloatVectorQuery extends KnnFloatVectorQuery {
private final int iTopK;
private final int searchWidth;
+ private final int threadBlockSize;
+ private final int maxIterations;
+ private final CagraSearchParams.SearchAlgo searchAlgo;
/**
- * Initializes {@link GPUKnnFloatVectorQuery}
+ * Initializes {@link GPUKnnFloatVectorQuery} with {@link CagraSearchParams.SearchAlgo#AUTO},
+ * and max_iterations auto-selected (0).
*
- * @param field the vector field name
- * @param target the vector target query
- * @param k the topK value
- * @param filter instance of the Query
- * @param iTopK the iTopK value
- * @param searchWidth the search width
+ * @param field the vector field name
+ * @param target the query vector
+ * @param k the number of nearest neighbors to return
+ * @param filter optional pre-filter query
+ * @param iTopK CAGRA itopk_size parameter
+ * @param searchWidth CAGRA search_width parameter
*/
public GPUKnnFloatVectorQuery(
String field, float[] target, int k, Query filter, int iTopK, int searchWidth) {
+ this(field, target, k, filter, iTopK, searchWidth, 0, 0, CagraSearchParams.SearchAlgo.AUTO);
+ }
+
+ /**
+ * Initializes {@link GPUKnnFloatVectorQuery}.
+ *
+ * @param field the vector field name
+ * @param target the query vector
+ * @param k the number of nearest neighbors to return
+ * @param filter optional pre-filter query
+ * @param iTopK CAGRA itopk_size parameter
+ * @param searchWidth CAGRA search_width parameter
+ * @param threadBlockSize CAGRA thread_block_size (0 = auto)
+ * @param maxIterations CAGRA max_iterations (0 = auto)
+ * @param searchAlgo CAGRA search algorithm
+ */
+ public GPUKnnFloatVectorQuery(
+ String field,
+ float[] target,
+ int k,
+ Query filter,
+ int iTopK,
+ int searchWidth,
+ int threadBlockSize,
+ int maxIterations,
+ CagraSearchParams.SearchAlgo searchAlgo) {
super(field, target, k, filter);
this.iTopK = iTopK;
this.searchWidth = searchWidth;
+ this.threadBlockSize = threadBlockSize;
+ this.maxIterations = maxIterations;
+ this.searchAlgo = searchAlgo;
+ }
+
+ // -------------------------------------------------------------------------
+ // Optimized multi-segment path
+ // -------------------------------------------------------------------------
+
+ @Override
+ public Query rewrite(IndexSearcher indexSearcher) throws IOException {
+ IndexReader reader = indexSearcher.getIndexReader();
+ List leaves = reader.leaves();
+ if (leaves.isEmpty()) {
+ return new MatchNoDocsQuery();
+ }
+
+ // Collect a CuVS2510GPUVectorsReader for every segment; fall back if any segment lacks one,
+ // has no CAGRA index for this field, or has a CAGRA graph whose degree differs from the other
+ // segments'. The multi-partition search requires every partition to share one graph degree
+ // (a small segment can have its degree truncated at build time), so a non-uniform set is
+ // routed to the per-segment path instead.
+ List gpuReaders = new ArrayList<>(leaves.size());
+ long commonGraphDegree = -1;
+ for (LeafReaderContext ctx : leaves) {
+ CuVS2510GPUVectorsReader gpuReader = unwrapGpuReader(ctx, field);
+ if (gpuReader == null) {
+ return super.rewrite(indexSearcher);
+ }
+ CagraIndex cagraIndex = gpuReader.getCagraIndexForField(field);
+ if (cagraIndex == null) {
+ return super.rewrite(indexSearcher);
+ }
+ long graphDegree = cagraIndex.getGraphDegree();
+ if (commonGraphDegree == -1) {
+ commonGraphDegree = graphDegree;
+ } else if (graphDegree != commonGraphDegree) {
+ return super.rewrite(indexSearcher);
+ }
+ gpuReaders.add(gpuReader);
+ }
+
+ // Build a single filter handle encoding (filter ∩ liveDocs) across every segment whenever
+ // any filtering is required — either an explicit Lucene filter, or live-document deletes in
+ // at least one segment. With the single-query multi-partition API, there is no other channel
+ // for per-segment liveDocs, so they must be folded into the FilterBitsetHandle.
+ boolean hasExplicitFilter = (filter != null);
+ boolean hasDeletes = false;
+ for (LeafReaderContext ctx : leaves) {
+ if (ctx.reader().getLiveDocs() != null) {
+ hasDeletes = true;
+ break;
+ }
+ }
+ // Build the per-segment CagraIndex list (one entry per Lucene segment / cuVS partition).
+ CuVSResources resources = getCuVSResourcesInstance();
+ List cagraIndices = new ArrayList<>(leaves.size());
+
+ FilterBitsetHandle filterHandle = null;
+ try {
+ if (hasExplicitFilter || hasDeletes) {
+ filterHandle = buildOrGetCachedFilterHandle(indexSearcher, leaves, gpuReaders);
+ }
+
+ float[] target = getTargetCopy();
+ CagraSearchParams searchParams =
+ new CagraSearchParams.Builder()
+ .withItopkSize(Math.max(iTopK, k))
+ .withSearchWidth(searchWidth)
+ .withThreadBlockSize(threadBlockSize)
+ .withMaxIterations(maxIterations)
+ .withAlgo(searchAlgo)
+ .build();
+
+ // Upload the query vector to device once; the same matrix view is searched against every
+ // partition by the cuVS multi-partition API.
+ CuVSMatrix.Builder vectorBuilder =
+ CuVSMatrix.deviceBuilder(resources, 1, target.length, CuVSMatrix.DataType.FLOAT);
+ vectorBuilder.addVector(target);
+
+ ScoreDoc[] scoreDocs;
+ try (CuVSMatrix queryVector = vectorBuilder.build()) {
+ for (int i = 0; i < leaves.size(); i++) {
+ cagraIndices.add(gpuReaders.get(i).getCagraIndexForField(field));
+ }
+
+ CagraQuery cagraQuery =
+ new CagraQuery.Builder(resources)
+ .withTopK(k)
+ .withSearchParams(searchParams)
+ .withQueryVectors(queryVector)
+ .build();
+
+ MultiPartitionSearchResults results =
+ MultiPartitionCagraSearch.search(resources, cagraIndices, cagraQuery, k, filterHandle);
+
+ if (results.count() == 0) {
+ return new MatchNoDocsQuery();
+ }
+
+ // Map (segmentIdx, ordinal) → global Lucene doc ID; compute normalized score.
+ // Stash FloatVectorValues per segment: results.count() can be far larger than the
+ // number of segments, and getFloatVectorValues() allocates on each call.
+ FloatVectorValues[] segVectorValues = new FloatVectorValues[leaves.size()];
+ scoreDocs = new ScoreDoc[results.count()];
+ for (int j = 0; j < results.count(); j++) {
+ int segIdx = results.getPartitionIndex(j);
+ int ordinal = results.getOrdinal(j);
+ float dist = results.getDistance(j);
+
+ FloatVectorValues fvv = segVectorValues[segIdx];
+ if (fvv == null) {
+ fvv = gpuReaders.get(segIdx).getFloatVectorValues(field);
+ segVectorValues[segIdx] = fvv;
+ }
+ LeafReaderContext ctx = leaves.get(segIdx);
+ int globalDoc = ctx.docBase + fvv.ordToDoc(ordinal);
+ float score = 1.0f / (1.0f + dist);
+ scoreDocs[j] = new ScoreDoc(globalDoc, score);
+ }
+
+ Arrays.sort(scoreDocs, Comparator.comparingDouble((ScoreDoc sd) -> sd.score).reversed());
+
+ return docAndScoreQuery(scoreDocs);
+ }
+
+ } catch (Throwable t) {
+ Utils.handleThrowable(t);
+ throw new AssertionError("handleThrowable always throws"); // unreachable
+ } finally {
+ // Release this query's reference. A cached handle is kept alive by the cache's own reference
+ // until eviction; an uncacheable handle holds only this reference and is freed here.
+ if (filterHandle != null) filterHandle.decRef();
+ }
}
+ // -------------------------------------------------------------------------
+ // Per-segment fallback path (used when k > 1024 or not all GPU segments)
+ // -------------------------------------------------------------------------
+
@Override
protected TopDocs approximateSearch(
LeafReaderContext context,
@@ -47,12 +264,321 @@ protected TopDocs approximateSearch(
int visitedLimit,
KnnCollectorManager knnCollectorManager)
throws IOException {
-
GPUPerLeafCuVSKnnCollector results =
- new GPUPerLeafCuVSKnnCollector(k, visitedLimit, iTopK, searchWidth);
-
- LeafReader reader = context.reader();
- reader.searchNearestVectors(field, this.getTargetCopy(), results, acceptDocs);
+ new GPUPerLeafCuVSKnnCollector(
+ k, visitedLimit, iTopK, searchWidth, threadBlockSize, maxIterations, searchAlgo);
+ context.reader().searchNearestVectors(field, getTargetCopy(), results, acceptDocs);
return results.topDocs();
}
+
+ // -------------------------------------------------------------------------
+ // Filter handle construction
+ // -------------------------------------------------------------------------
+
+ /**
+ * Returns a {@link FilterBitsetHandle} encoding ({@link #filter} ∩ liveDocs) for every segment,
+ * pulling from {@link FilterBitsetCache} when the reader state is unchanged.
+ *
+ * Cache key uses per-reader keys (not just core keys) so that liveDocs changes — which happen
+ * when a reader is reopened after deletes — automatically invalidate the cached bitset.
+ *
+ *
The returned handle carries a reference the caller must release with {@link
+ * FilterBitsetHandle#decRef()} once the search completes.
+ */
+ private FilterBitsetHandle buildOrGetCachedFilterHandle(
+ IndexSearcher indexSearcher,
+ List leaves,
+ List gpuReaders)
+ throws IOException {
+
+ List segReaderKeys = new ArrayList<>(leaves.size());
+ for (LeafReaderContext ctx : leaves) {
+ var helper = ctx.reader().getReaderCacheHelper();
+ if (helper == null) {
+ // This reader can't be cached; build an uncached handle owned outright by the caller.
+ return buildFilterHandle(indexSearcher, leaves, gpuReaders);
+ }
+ segReaderKeys.add(helper.getKey());
+ }
+
+ return FilterBitsetCache.acquire(
+ filter,
+ segReaderKeys,
+ field,
+ () -> buildFilterHandle(indexSearcher, leaves, gpuReaders));
+ }
+
+ /**
+ * Evaluates {@link #filter} per segment (when set), intersects with liveDocs, and packs the
+ * result into a new {@link FilterBitsetHandle}. When {@link #filter} is {@code null}, the
+ * handle encodes liveDocs alone — this path is taken when one or more segments have deletes
+ * but no explicit Lucene filter was supplied.
+ */
+ private FilterBitsetHandle buildFilterHandle(
+ IndexSearcher indexSearcher,
+ List leaves,
+ List gpuReaders)
+ throws IOException {
+
+ Weight filterWeight = null;
+ if (filter != null) {
+ filterWeight =
+ indexSearcher.createWeight(
+ indexSearcher.rewrite(filter), ScoreMode.COMPLETE_NO_SCORES, 1.0f);
+ }
+
+ int numSegments = leaves.size();
+ long[] segBitOffsets = new long[numSegments];
+ long totalBits = 0;
+ for (int i = 0; i < numSegments; i++) {
+ segBitOffsets[i] = totalBits;
+ int numOrds = gpuReaders.get(i).getFloatVectorValues(field).size();
+ totalBits += (((long) numOrds + 63) / 64) * 64;
+ }
+ long[] combinedLongs = new long[(int) (totalBits / 64)];
+
+ for (int i = 0; i < numSegments; i++) {
+ LeafReaderContext ctx = leaves.get(i);
+ Bits liveDocs = ctx.reader().getLiveDocs();
+ // When filterWeight is null, accept all live documents (acceptDocs == liveDocs, which may
+ // itself be null to mean "all docs accepted" in this segment).
+ Bits acceptDocs = (filterWeight != null) ? evalFilter(filterWeight, ctx, liveDocs) : liveDocs;
+ FloatVectorValues fvv = gpuReaders.get(i).getFloatVectorValues(field);
+ Bits acceptedOrds = fvv.getAcceptOrds(acceptDocs);
+ int numOrds = fvv.size();
+ int longOffset = (int) (segBitOffsets[i] / 64);
+ packOrdsToLongs(acceptedOrds, numOrds, combinedLongs, longOffset);
+ }
+
+ // Only the combined bitset is transported; cuVS recomputes the per-partition bit offsets from
+ // the index sizes. segBitOffsets is used above solely to pack each segment's slice (64-bit
+ // word-aligned), matching that recomputation.
+ return FilterBitsetHandle.create(combinedLongs);
+ }
+
+ /**
+ * Evaluates {@code filterWeight} in {@code ctx} and intersects with {@code liveDocs}.
+ * Returns a {@link Bits} over local doc IDs where {@code get(doc)} is true for accepted docs.
+ */
+ private static Bits evalFilter(Weight filterWeight, LeafReaderContext ctx, Bits liveDocs)
+ throws IOException {
+ ScorerSupplier scorerSupplier = filterWeight.scorerSupplier(ctx);
+ if (scorerSupplier == null) {
+ // No scorer: the filter matches no documents in this segment.
+ return new Bits.MatchNoBits(ctx.reader().maxDoc());
+ }
+
+ int maxDoc = ctx.reader().maxDoc();
+ FixedBitSet filterBits = new FixedBitSet(maxDoc);
+ Scorer scorer = scorerSupplier.get(Long.MAX_VALUE);
+ DocIdSetIterator it = scorer.iterator();
+ int doc;
+ while ((doc = it.nextDoc()) != DocIdSetIterator.NO_MORE_DOCS) {
+ filterBits.set(doc);
+ }
+
+ if (liveDocs == null) return filterBits;
+
+ // Intersect: accept only docs that pass both filter and liveDocs.
+ return new Bits() {
+ @Override
+ public boolean get(int i) {
+ return filterBits.get(i) && liveDocs.get(i);
+ }
+
+ @Override
+ public int length() {
+ return maxDoc;
+ }
+ };
+ }
+
+ /**
+ * Packs {@code numOrds} ordinal bits from {@code bits} into {@code dest} starting at long index
+ * {@code destLongOffset}. {@code bits == null} means all ordinals accepted (Lucene convention).
+ */
+ private static void packOrdsToLongs(Bits bits, int numOrds, long[] dest, int destLongOffset) {
+ if (bits == null) {
+ // All ordinals accepted: fill with all-ones, masking the last partial word.
+ int numLongs = (numOrds + 63) / 64;
+ Arrays.fill(dest, destLongOffset, destLongOffset + numLongs, -1L);
+ int tail = numOrds % 64;
+ if (tail != 0) {
+ dest[destLongOffset + numLongs - 1] = (1L << tail) - 1L;
+ }
+ return;
+ }
+ for (int i = 0; i < numOrds; i++) {
+ if (bits.get(i)) {
+ dest[destLongOffset + i / 64] |= (1L << (i % 64));
+ }
+ }
+ }
+
+ // -------------------------------------------------------------------------
+ // Helpers
+ // -------------------------------------------------------------------------
+
+ /**
+ * Unwraps the {@link LeafReaderContext}'s reader to a {@link CuVS2510GPUVectorsReader}, or
+ * returns {@code null} if the reader is not of that type.
+ */
+ private static CuVS2510GPUVectorsReader unwrapGpuReader(LeafReaderContext ctx, String field) {
+ var unwrapped = FilterLeafReader.unwrap(ctx.reader());
+ if (!(unwrapped instanceof CodecReader)) return null;
+ KnnVectorsReader vr = ((CodecReader) unwrapped).getVectorReader();
+ // A per-field codec wraps the vectors reader; unwrap to this field's delegate so the
+ // multi-partition GPU path is not silently bypassed when the CuVS format is used via
+ // PerFieldKnnVectorsFormat.
+ if (vr instanceof PerFieldKnnVectorsFormat.FieldsReader perField) {
+ vr = perField.getFieldReader(field);
+ }
+ return (vr instanceof CuVS2510GPUVectorsReader gpuReader) ? gpuReader : null;
+ }
+
+ /**
+ * Builds a {@link Query} that matches exactly the given pre-scored documents.
+ *
+ * Partitions {@code scoreDocs} by segment (using {@link ScoreDoc#shardIndex} as the segment
+ * offset relative to {@link LeafReaderContext#docBase}), then returns a {@link Scorer} per
+ * segment that iterates those docs in ascending doc-ID order and replays their pre-computed
+ * scores.
+ */
+ private static Query docAndScoreQuery(ScoreDoc[] scoreDocs) {
+ return new Query() {
+ @Override
+ public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost)
+ throws IOException {
+ return new Weight(this) {
+ @Override
+ public ScorerSupplier scorerSupplier(LeafReaderContext ctx) {
+ int base = ctx.docBase;
+ int maxDoc = base + ctx.reader().maxDoc();
+ // Collect docs belonging to this segment; re-sort by local doc ID ascending.
+ int[] localDocs = new int[scoreDocs.length];
+ float[] localScores = new float[scoreDocs.length];
+ int count = 0;
+ float max = Float.NEGATIVE_INFINITY;
+ for (ScoreDoc sd : scoreDocs) {
+ if (sd.doc >= base && sd.doc < maxDoc) {
+ localDocs[count] = sd.doc - base;
+ localScores[count] = sd.score * boost;
+ if (localScores[count] > max) max = localScores[count];
+ count++;
+ }
+ }
+ if (count == 0) return null;
+ final int n = count;
+ final float maxScore = max;
+ Integer[] idx = new Integer[n];
+ for (int i = 0; i < n; i++) idx[i] = i;
+ Arrays.sort(idx, Comparator.comparingInt(i -> localDocs[i]));
+ final int[] sortedDocs = new int[n];
+ final float[] sortedScores = new float[n];
+ for (int i = 0; i < n; i++) {
+ sortedDocs[i] = localDocs[idx[i]];
+ sortedScores[i] = localScores[idx[i]];
+ }
+
+ return new ScorerSupplier() {
+ @Override
+ public Scorer get(long leadCost) {
+ return new Scorer() {
+ private int pos = -1;
+
+ @Override
+ public DocIdSetIterator iterator() {
+ return new DocIdSetIterator() {
+ @Override
+ public int docID() {
+ return pos < 0 ? -1 : (pos >= n ? NO_MORE_DOCS : sortedDocs[pos]);
+ }
+
+ @Override
+ public int nextDoc() {
+ pos++;
+ return docID();
+ }
+
+ @Override
+ public int advance(int target) {
+ if (pos < 0) pos = 0;
+ while (pos < n && sortedDocs[pos] < target) pos++;
+ return docID();
+ }
+
+ @Override
+ public long cost() {
+ return n;
+ }
+ };
+ }
+
+ @Override
+ public float getMaxScore(int upTo) {
+ // The precomputed segment max is a valid upper bound for any upTo.
+ return maxScore;
+ }
+
+ @Override
+ public float score() {
+ return sortedScores[pos];
+ }
+
+ @Override
+ public int docID() {
+ return pos < 0
+ ? -1
+ : (pos >= n ? DocIdSetIterator.NO_MORE_DOCS : sortedDocs[pos]);
+ }
+ };
+ }
+
+ @Override
+ public long cost() {
+ return n;
+ }
+ };
+ }
+
+ @Override
+ public boolean isCacheable(LeafReaderContext ctx) {
+ return false;
+ }
+
+ @Override
+ public Explanation explain(LeafReaderContext ctx, int doc) {
+ for (ScoreDoc sd : scoreDocs) {
+ if (sd.doc == ctx.docBase + doc) {
+ return Explanation.match(
+ sd.score * boost,
+ "GPU multi-segment CAGRA search, product of:",
+ Explanation.match(sd.score, "raw CAGRA score, 1 / (1 + distance)"),
+ Explanation.match(boost, "boost"));
+ }
+ }
+ return Explanation.noMatch("not a GPU search result");
+ }
+ };
+ }
+
+ @Override
+ public String toString(String field) {
+ return "GPUDocAndScoreQuery";
+ }
+
+ @Override
+ public void visit(QueryVisitor visitor) {}
+
+ @Override
+ public boolean equals(Object o) {
+ return this == o;
+ }
+
+ @Override
+ public int hashCode() {
+ return System.identityHashCode(this);
+ }
+ };
+ }
}
diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java b/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java
index 421b2baa..6413c221 100644
--- a/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java
+++ b/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java
@@ -1,9 +1,10 @@
/*
- * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*/
package com.nvidia.cuvs.lucene;
+import com.nvidia.cuvs.CagraSearchParams;
import org.apache.lucene.search.TopKnnCollector;
/**
@@ -15,6 +16,9 @@ class GPUPerLeafCuVSKnnCollector extends TopKnnCollector {
private int iTopK;
private int searchWidth;
+ private int threadBlockSize;
+ private int maxIterations;
+ private CagraSearchParams.SearchAlgo searchAlgo;
/**
* Initializes {@link GPUPerLeafCuVSKnnCollector}
@@ -22,11 +26,24 @@ class GPUPerLeafCuVSKnnCollector extends TopKnnCollector {
* @param topK the topK value
* @param iTopK the iTopK value
* @param searchWidth the search width
+ * @param threadBlockSize CAGRA thread_block_size (0 = auto; controls worker_queue_size)
+ * @param maxIterations CAGRA max_iterations (0 = auto)
+ * @param searchAlgo the CAGRA search algorithm
*/
- public GPUPerLeafCuVSKnnCollector(int topK, int visitLimit, int iTopK, int searchWidth) {
+ public GPUPerLeafCuVSKnnCollector(
+ int topK,
+ int visitLimit,
+ int iTopK,
+ int searchWidth,
+ int threadBlockSize,
+ int maxIterations,
+ CagraSearchParams.SearchAlgo searchAlgo) {
super(topK, visitLimit);
this.iTopK = iTopK > topK ? iTopK : topK;
this.searchWidth = searchWidth;
+ this.threadBlockSize = threadBlockSize;
+ this.maxIterations = maxIterations;
+ this.searchAlgo = searchAlgo;
}
public int getiTopK() {
@@ -36,4 +53,16 @@ public int getiTopK() {
public int getSearchWidth() {
return searchWidth;
}
+
+ public int getThreadBlockSize() {
+ return threadBlockSize;
+ }
+
+ public int getMaxIterations() {
+ return maxIterations;
+ }
+
+ public CagraSearchParams.SearchAlgo getSearchAlgo() {
+ return searchAlgo;
+ }
}
diff --git a/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java b/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java
index 9e259e27..e1bb46f6 100644
--- a/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java
+++ b/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java
@@ -1,8 +1,7 @@
/*
- * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*/
-
package com.nvidia.cuvs.lucene;
import com.nvidia.cuvs.CuVSResources;
@@ -42,9 +41,17 @@ public static void setCuVSResourcesInstance(CuVSResources resources) {
cuVSResources.set(resources);
}
+ /** System property controlling the workspace pool size per resources handle (in bytes). */
+ public static final String WORKSPACE_POOL_SIZE_PROPERTY = "com.nvidia.cuvs.workspacePoolSize";
+
private static CuVSResources cuVSResourcesOrNull() {
try {
- return CuVSResources.create();
+ CuVSResources resources = CuVSResources.create();
+ String poolSizeProp = System.getProperty(WORKSPACE_POOL_SIZE_PROPERTY);
+ if (poolSizeProp != null) {
+ resources.setWorkspacePool(Long.parseLong(poolSizeProp));
+ }
+ return resources;
} catch (UnsupportedOperationException uoe) {
log.log(
Level.WARNING,
diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestFilterBitsetCache.java b/src/test/java/com/nvidia/cuvs/lucene/TestFilterBitsetCache.java
new file mode 100644
index 00000000..01e8b050
--- /dev/null
+++ b/src/test/java/com/nvidia/cuvs/lucene/TestFilterBitsetCache.java
@@ -0,0 +1,236 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package com.nvidia.cuvs.lucene;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
+import com.nvidia.cuvs.FilterBitsetHandle;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import org.junit.Test;
+
+/**
+ * Concurrency tests for {@link FilterBitsetCache}. These exercise the compute-once and
+ * reference-counting behavior in isolation by injecting a fake {@link FilterBitsetHandle}, so no GPU
+ * or native library is required.
+ */
+public class TestFilterBitsetCache {
+
+ /**
+ * Fake handle faithfully modeling the {@link FilterBitsetHandle} reference-counting contract, so
+ * the cache's use of it can be observed. {@link #decRef()} throws if released more times than
+ * referenced, turning any over-release by the cache into a test failure.
+ */
+ private static final class CountingHandle implements FilterBitsetHandle {
+ final AtomicInteger refCount = new AtomicInteger(1); // initial reference, as in the real handle
+ final AtomicInteger closeCalls = new AtomicInteger();
+ final AtomicBoolean freed = new AtomicBoolean();
+ private final AtomicBoolean initialReleased = new AtomicBoolean();
+
+ @Override
+ public boolean tryIncRef() {
+ int c;
+ do {
+ c = refCount.get();
+ if (c == 0) return false;
+ } while (!refCount.compareAndSet(c, c + 1));
+ return true;
+ }
+
+ @Override
+ public void decRef() {
+ int c = refCount.decrementAndGet();
+ if (c == 0) {
+ freed.set(true);
+ } else if (c < 0) {
+ throw new IllegalStateException("decRef() called more times than references were taken");
+ }
+ }
+
+ @Override
+ public void close() {
+ closeCalls.incrementAndGet();
+ if (initialReleased.compareAndSet(false, true)) {
+ decRef();
+ }
+ }
+ }
+
+ private static List keys(String s) {
+ return List.of(s);
+ }
+
+ /** Concurrent misses on the same key build the handle exactly once and each get their own ref. */
+ @Test
+ public void computeOnceUnderConcurrentAcquire() throws Exception {
+ final int threads = 32;
+ AtomicInteger buildCount = new AtomicInteger();
+ CountingHandle handle = new CountingHandle();
+ final String field = "computeOnce";
+
+ ExecutorService pool = Executors.newFixedThreadPool(threads);
+ CountDownLatch start = new CountDownLatch(1);
+ List> results = new ArrayList<>();
+ try {
+ for (int i = 0; i < threads; i++) {
+ results.add(
+ pool.submit(
+ () -> {
+ start.await();
+ return FilterBitsetCache.acquire(
+ null,
+ keys(field),
+ field,
+ () -> {
+ buildCount.incrementAndGet();
+ // Widen the race window so waiters reach the future before it completes.
+ try {
+ Thread.sleep(20);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ return handle;
+ });
+ }));
+ }
+ start.countDown();
+ for (Future r : results) {
+ assertSame(handle, r.get(10, TimeUnit.SECONDS));
+ }
+ } finally {
+ pool.shutdownNow();
+ pool.awaitTermination(5, TimeUnit.SECONDS);
+ }
+
+ assertEquals("builder must run exactly once for a key", 1, buildCount.get());
+ // One cache reference plus one caller reference per acquiring thread.
+ assertEquals(1 + threads, handle.refCount.get());
+ assertFalse(handle.freed.get());
+
+ // Releasing every caller reference leaves the cache's own reference intact.
+ for (int i = 0; i < threads; i++) {
+ handle.decRef();
+ }
+ assertEquals(1, handle.refCount.get());
+ assertFalse(handle.freed.get());
+ }
+
+ /** Eviction releases exactly the cache's reference, freeing evicted handles and keeping the rest. */
+ @Test
+ public void evictionReleasesCacheReferenceExactlyOnce() throws Exception {
+ final int max = FilterBitsetCache.MAX_HOST_ENTRIES;
+ final int total = 3 * max;
+ final int firstRetained = total - max; // oldest `total - max` entries are evicted
+
+ List handles = new ArrayList<>();
+ for (int i = 0; i < total; i++) {
+ CountingHandle h = new CountingHandle();
+ handles.add(h);
+ final String field = "evict-" + i;
+ FilterBitsetHandle got = FilterBitsetCache.acquire(null, keys(field), field, () -> h);
+ assertSame(h, got);
+ got.decRef(); // caller finished; the cache keeps its reference
+ }
+
+ for (int i = 0; i < firstRetained; i++) {
+ CountingHandle h = handles.get(i);
+ assertTrue("handle " + i + " should have been evicted and freed", h.freed.get());
+ assertEquals("evicted handle closed exactly once", 1, h.closeCalls.get());
+ }
+ for (int i = firstRetained; i < total; i++) {
+ CountingHandle h = handles.get(i);
+ assertFalse("handle " + i + " should still be cached", h.freed.get());
+ assertEquals("retained handle must not be closed", 0, h.closeCalls.get());
+ }
+ }
+
+ /** A failed build is not cached, so a later acquire rebuilds. */
+ @Test
+ public void failedBuildIsNotCachedAndCanRetry() throws Exception {
+ final String field = "failing";
+ AtomicInteger buildCount = new AtomicInteger();
+
+ try {
+ FilterBitsetCache.acquire(
+ null,
+ keys(field),
+ field,
+ () -> {
+ buildCount.incrementAndGet();
+ throw new IOException("boom");
+ });
+ fail("expected IOException to propagate");
+ } catch (IOException expected) {
+ assertEquals("boom", expected.getMessage());
+ }
+ assertEquals(1, buildCount.get());
+
+ CountingHandle handle = new CountingHandle();
+ FilterBitsetHandle got =
+ FilterBitsetCache.acquire(
+ null,
+ keys(field),
+ field,
+ () -> {
+ buildCount.incrementAndGet();
+ return handle;
+ });
+ assertSame(handle, got);
+ assertEquals("failed entry must not be cached; retry rebuilds", 2, buildCount.get());
+ got.decRef();
+ }
+
+ /**
+ * Hammer acquire/decRef concurrently across several keys. Any unbalanced reference handling by the
+ * cache trips {@link CountingHandle#decRef()}'s below-zero guard and fails the test.
+ */
+ @Test
+ public void concurrentAcquireReleaseIsBalanced() throws Exception {
+ final int keySpace = 8; // < MAX_HOST_ENTRIES, so these keys are never evicted mid-run
+ final int threads = 16;
+ final int iterationsPerThread = 500;
+
+ ExecutorService pool = Executors.newFixedThreadPool(threads);
+ CountDownLatch start = new CountDownLatch(1);
+ List> tasks = new ArrayList<>();
+ try {
+ for (int t = 0; t < threads; t++) {
+ final int seed = t;
+ tasks.add(
+ pool.submit(
+ () -> {
+ start.await();
+ for (int i = 0; i < iterationsPerThread; i++) {
+ final String field = "stress-" + ((seed + i) % keySpace);
+ // A fresh handle per build; a cached key reuses whatever was built first.
+ FilterBitsetHandle h =
+ FilterBitsetCache.acquire(null, keys(field), field, CountingHandle::new);
+ h.decRef();
+ }
+ return null;
+ }));
+ }
+ start.countDown();
+ for (Future f : tasks) {
+ f.get(30, TimeUnit.SECONDS); // surfaces any exception thrown inside a task
+ }
+ } finally {
+ pool.shutdownNow();
+ pool.awaitTermination(5, TimeUnit.SECONDS);
+ }
+ }
+}
diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestMultiSegmentGPUFilterConcurrency.java b/src/test/java/com/nvidia/cuvs/lucene/TestMultiSegmentGPUFilterConcurrency.java
new file mode 100644
index 00000000..fb8e92c6
--- /dev/null
+++ b/src/test/java/com/nvidia/cuvs/lucene/TestMultiSegmentGPUFilterConcurrency.java
@@ -0,0 +1,194 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package com.nvidia.cuvs.lucene;
+
+import static com.nvidia.cuvs.lucene.TestUtils.generateDataset;
+import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported;
+import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.ThreadLocalRandom;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import org.apache.lucene.codecs.Codec;
+import org.apache.lucene.document.Document;
+import org.apache.lucene.document.Field;
+import org.apache.lucene.document.KnnFloatVectorField;
+import org.apache.lucene.document.StringField;
+import org.apache.lucene.index.DirectoryReader;
+import org.apache.lucene.index.IndexWriter;
+import org.apache.lucene.index.IndexWriterConfig;
+import org.apache.lucene.index.NoMergePolicy;
+import org.apache.lucene.index.Term;
+import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.Query;
+import org.apache.lucene.search.ScoreDoc;
+import org.apache.lucene.search.TermQuery;
+import org.apache.lucene.store.ByteBuffersDirectory;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.tests.util.LuceneTestCase;
+import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.junit.Test;
+
+/**
+ * End-to-end concurrency test for the filter-bitset path of multi-segment GPU search. Many threads
+ * issue filtered {@link GPUKnnFloatVectorQuery} searches across a rotating set of distinct filters
+ * (more than {@link FilterBitsetCache#MAX_HOST_ENTRIES}) so the shared filter cache continuously
+ * builds, caches, evicts, and reuses handles under contention. Every returned hit must satisfy its
+ * filter, which validates that the reference-counted eviction never frees a handle still in use.
+ */
+@SuppressSysoutChecks(bugUrl = "")
+public class TestMultiSegmentGPUFilterConcurrency extends LuceneTestCase {
+
+ // Use the real GPU-search codec (returns the CuVS format directly, not wrapped) so the reader is
+ // a CuVS2510GPUVectorsReader and GPUKnnFloatVectorQuery takes the multi-partition GPU path rather
+ // than falling back to per-segment search.
+ private static Codec codec;
+ private static final String VECTOR_FIELD = "vectors";
+ private static final String CATEGORY_FIELD = "cat";
+ // More categories than MAX_HOST_ENTRIES so distinct-filter churn forces cache eviction.
+ private static final int NUM_CATEGORIES = FilterBitsetCache.MAX_HOST_ENTRIES + 8;
+
+ private static Directory directory;
+ private static DirectoryReader reader;
+ private static IndexSearcher searcher;
+ private static float[][] queryVectors;
+ private static int topK;
+ private static List filters;
+ private static List> acceptedDocsPerFilter;
+
+ @BeforeClass
+ public static void beforeClass() throws Exception {
+ assumeTrue("cuVS not supported", isSupported());
+ codec = new CuVS2510GPUSearchCodec();
+
+ int datasetSize = 2000;
+ int dimensions = 128;
+ int numQueries = 64;
+ topK = 10;
+
+ directory = newDirectory(new ByteBuffersDirectory());
+ // Disable merges so several commits yield several GPU segments -> multi-partition search.
+ IndexWriterConfig config =
+ new IndexWriterConfig().setCodec(codec).setMergePolicy(NoMergePolicy.INSTANCE);
+ float[][] dataset = generateDataset(random(), datasetSize, dimensions);
+
+ try (IndexWriter writer = new IndexWriter(directory, config)) {
+ final int commitEvery = datasetSize / 4;
+ for (int i = 0; i < datasetSize; i++) {
+ Document doc = new Document();
+ doc.add(new StringField("id", String.valueOf(i), Field.Store.YES));
+ doc.add(new StringField(CATEGORY_FIELD, "c" + (i % NUM_CATEGORIES), Field.Store.NO));
+ doc.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN));
+ writer.addDocument(doc);
+ if ((i + 1) % commitEvery == 0) {
+ writer.commit();
+ }
+ }
+ writer.commit();
+ }
+
+ reader = DirectoryReader.open(directory);
+ searcher = new IndexSearcher(reader);
+ queryVectors = generateDataset(random(), numQueries, dimensions);
+
+ // One distinct filter per category, plus the set of global doc IDs each filter accepts.
+ filters = new ArrayList<>(NUM_CATEGORIES);
+ acceptedDocsPerFilter = new ArrayList<>(NUM_CATEGORIES);
+ for (int c = 0; c < NUM_CATEGORIES; c++) {
+ Query filter = new TermQuery(new Term(CATEGORY_FIELD, "c" + c));
+ filters.add(filter);
+ Set accepted = new HashSet<>();
+ for (ScoreDoc sd : searcher.search(filter, datasetSize).scoreDocs) {
+ accepted.add(sd.doc);
+ }
+ acceptedDocsPerFilter.add(accepted);
+ }
+ }
+
+ @Test
+ public void testConcurrentFilteredSearchAcrossManyFilters() throws Exception {
+ final int numThreads = 6;
+ final int iterationsPerThread = 150;
+
+ List errors = new CopyOnWriteArrayList<>();
+ AtomicInteger nonEmptyResults = new AtomicInteger();
+
+ ExecutorService pool = Executors.newFixedThreadPool(numThreads);
+ CountDownLatch go = new CountDownLatch(1);
+ List> tasks = new ArrayList<>();
+ try {
+ for (int t = 0; t < numThreads; t++) {
+ tasks.add(
+ pool.submit(
+ () -> {
+ ThreadLocalRandom rnd = ThreadLocalRandom.current();
+ go.await();
+ for (int i = 0; i < iterationsPerThread; i++) {
+ int c = rnd.nextInt(NUM_CATEGORIES);
+ Query filter = filters.get(c);
+ Set accepted = acceptedDocsPerFilter.get(c);
+ float[] q = queryVectors[rnd.nextInt(queryVectors.length)];
+
+ GPUKnnFloatVectorQuery query =
+ new GPUKnnFloatVectorQuery(VECTOR_FIELD, q, topK, filter, topK, 1);
+ ScoreDoc[] hits = searcher.search(query, topK).scoreDocs;
+ if (hits.length > 0) {
+ nonEmptyResults.incrementAndGet();
+ }
+ for (ScoreDoc hit : hits) {
+ if (!accepted.contains(hit.doc)) {
+ throw new AssertionError(
+ "hit " + hit.doc + " not accepted by filter c" + c);
+ }
+ }
+ }
+ return null;
+ }));
+ }
+ go.countDown();
+ for (Future f : tasks) {
+ try {
+ f.get(120, TimeUnit.SECONDS);
+ } catch (Throwable t) {
+ errors.add(t);
+ }
+ }
+ } finally {
+ pool.shutdownNow();
+ pool.awaitTermination(10, TimeUnit.SECONDS);
+ }
+
+ if (!errors.isEmpty()) {
+ throw new AssertionError(
+ errors.size() + " search thread(s) failed; first: " + errors.get(0), errors.get(0));
+ }
+ assertTrue("expected at least some non-empty filtered results", nonEmptyResults.get() > 0);
+ }
+
+ @AfterClass
+ public static void afterClass() throws IOException {
+ if (reader != null) reader.close();
+ if (directory != null) directory.close();
+ reader = null;
+ directory = null;
+ searcher = null;
+ filters = null;
+ acceptedDocsPerFilter = null;
+ queryVectors = null;
+ }
+}
diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestPerSegmentGPUFilterSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestPerSegmentGPUFilterSearch.java
new file mode 100644
index 00000000..df2ba211
--- /dev/null
+++ b/src/test/java/com/nvidia/cuvs/lucene/TestPerSegmentGPUFilterSearch.java
@@ -0,0 +1,91 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package com.nvidia.cuvs.lucene;
+
+import static com.nvidia.cuvs.lucene.TestUtils.generateDataset;
+import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported;
+import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN;
+
+import org.apache.lucene.document.Document;
+import org.apache.lucene.document.KnnFloatVectorField;
+import org.apache.lucene.index.DirectoryReader;
+import org.apache.lucene.index.IndexWriter;
+import org.apache.lucene.index.IndexWriterConfig;
+import org.apache.lucene.index.LeafReader;
+import org.apache.lucene.search.ScoreDoc;
+import org.apache.lucene.search.TopKnnCollector;
+import org.apache.lucene.store.ByteBuffersDirectory;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.tests.util.LuceneTestCase;
+import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks;
+import org.apache.lucene.util.FixedBitSet;
+import org.junit.Test;
+
+/**
+ * Exercises the per-segment (single-index) GPU search with a selective prefilter. It drives {@link
+ * CuVS2510GPUVectorsReader#search} directly via {@code LeafReader.searchNearestVectors} with a
+ * synthesized {@code acceptDocs}, so the path is exercised regardless of any higher-level query
+ * routing (which now sends the multi-partition case elsewhere).
+ *
+ * A highly selective {@code acceptDocs} rejects the trailing ordinals of the segment — the regime
+ * where the per-segment prefilter previously mis-sized itself (it passed {@code BitSet.length()},
+ * the highest set bit + 1, instead of the vector count) and default-accepted the trailing ordinals,
+ * returning filtered-out vectors. Every returned hit must be accepted.
+ */
+@SuppressSysoutChecks(bugUrl = "")
+public class TestPerSegmentGPUFilterSearch extends LuceneTestCase {
+
+ private static final String VECTOR_FIELD = "vectors";
+ private static final int NUM_CATEGORIES = 24;
+
+ @Test
+ public void testSelectiveFilterDoesNotLeakTrailingOrdinals() throws Exception {
+ assumeTrue("cuVS not supported", isSupported());
+
+ // A segment size that is not a multiple of 32 exercises the last partial uint32 word of the
+ // prefilter bitset, which is where the trailing ordinals used to leak.
+ final int datasetSize = 500;
+ final int dimensions = 128;
+ final int topK = 10;
+
+ try (Directory directory = newDirectory(new ByteBuffersDirectory())) {
+ float[][] dataset = generateDataset(random(), datasetSize, dimensions);
+ IndexWriterConfig config = new IndexWriterConfig().setCodec(new CuVS2510GPUSearchCodec());
+ try (IndexWriter writer = new IndexWriter(directory, config)) {
+ for (int i = 0; i < datasetSize; i++) {
+ Document doc = new Document();
+ doc.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN));
+ writer.addDocument(doc);
+ }
+ writer.commit();
+ }
+
+ try (DirectoryReader reader = DirectoryReader.open(directory)) {
+ assertEquals("expected a single segment", 1, reader.leaves().size());
+ LeafReader leaf = reader.leaves().get(0).reader();
+ float[][] queries = generateDataset(random(), 32, dimensions);
+
+ // Category c = ordinal % NUM_CATEGORIES; acceptDocs keeps ~1/24 of the ordinals, so the
+ // segment's trailing ordinals are rejected. (docId == vector ordinal for a dense segment.)
+ for (int c = 0; c < NUM_CATEGORIES; c++) {
+ FixedBitSet acceptDocs = new FixedBitSet(datasetSize);
+ for (int i = c; i < datasetSize; i += NUM_CATEGORIES) {
+ acceptDocs.set(i);
+ }
+ for (float[] q : queries) {
+ TopKnnCollector collector = new TopKnnCollector(topK, Integer.MAX_VALUE);
+ leaf.searchNearestVectors(VECTOR_FIELD, q, collector, acceptDocs);
+ for (ScoreDoc hit : collector.topDocs().scoreDocs) {
+ assertTrue(
+ "per-segment search returned doc " + hit.doc + " outside filter category " + c,
+ acceptDocs.get(hit.doc));
+ }
+ }
+ }
+ }
+ }
+ }
+}