From ff9949f27ac355ed78055e81381ad94d316f2760 Mon Sep 17 00:00:00 2001 From: James Xia Date: Tue, 31 Mar 2026 16:13:38 -0700 Subject: [PATCH 01/41] Add GPU-optimized multi-segment CAGRA search in GPUKnnFloatVectorQuery Overrides rewrite() to run all segment searches into a shared device buffer and merge with cuvsSelectK entirely on GPU, eliminating per-segment D2H copies and CPU-side TopDocs.merge(). Falls back to the standard Lucene per-segment path when any segment lacks a CAGRA index, an explicit filter is set, or k > 1024. Also adds ordToDoc() and getCagraIndexForField() helpers to CuVS2510GPUVectorsReader to support result decoding. Fixes for Lucene 10.2 API changes: CodecReader moved to org.apache.lucene.index; createRewrittenQuery() removed and replaced with an inline docAndScoreQuery() implementation using the public Weight/ScorerSupplier API. --- .../cuvs/lucene/CuVS2510GPUVectorsReader.java | 31 ++ .../cuvs/lucene/GPUKnnFloatVectorQuery.java | 347 +++++++++++++++++- 2 files changed, 365 insertions(+), 13 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java index c783660d..71169af0 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java @@ -600,6 +600,37 @@ private static void checkVersion(int versionMeta, int versionVectorData, IndexIn } } + /** + * Maps a local segment ordinal to a Lucene doc ID within this segment. + * + *

Used by {@link GPUKnnFloatVectorQuery} after a multi-segment GPU search to convert + * select_k result ordinals to doc IDs before adding {@code docBase}. + * + * @param field the vector field name + * @param ordinal the local ordinal returned by CAGRA + * @return the Lucene doc ID within this segment + * @throws IOException if the vector values cannot be read + */ + public int ordToDoc(String field, int ordinal) throws IOException { + return flatVectorsReader.getFloatVectorValues(field).ordToDoc(ordinal); + } + + /** + * 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/GPUKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index 6a9daae7..a6ea5690 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -4,17 +4,58 @@ */ 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.MultiSegmentCagraSearch; +import com.nvidia.cuvs.MultiSegmentSearchResults; import java.io.IOException; -import org.apache.lucene.index.LeafReader; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.BitSet; +import java.util.Comparator; +import java.util.List; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.index.CodecReader; +import org.apache.lucene.index.FilterLeafReader; +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; /** - * Extends upon KnnFloatVectorQuery for GPU-only search. + * Extends {@link KnnFloatVectorQuery} for GPU-only search. + * + *

When all index segments use {@link CuVS2510GPUVectorsReader} and the query uses CAGRA + * (k ≤ 1024, no explicit filter), {@link #rewrite} runs a globally-optimized + * multi-segment search: + *

    + *
  1. All per-segment CAGRA searches write into a single shared device buffer without any + * per-segment device-to-host copy or stream synchronization.
    2. + *
  2. A single {@code cuvsSelectK} call finds the global top-k entirely on GPU.
    3. + *
  3. Results are copied to host in one pass and mapped to Lucene doc IDs.
    4. + *
+ * + *

Falls back to the standard per-segment Lucene path when the optimized path cannot be + * applied (mixed segment types, explicit query filter, k > 1024, brute-force + * fallback needed, or missing CAGRA index). * * @since 25.10 */ @@ -24,14 +65,14 @@ public class GPUKnnFloatVectorQuery extends KnnFloatVectorQuery { private final int searchWidth; /** - * Initializes {@link GPUKnnFloatVectorQuery} + * Initializes {@link GPUKnnFloatVectorQuery}. * - * @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) { @@ -40,6 +81,98 @@ public GPUKnnFloatVectorQuery( this.searchWidth = searchWidth; } + // ------------------------------------------------------------------------- + // Optimized multi-segment path + // ------------------------------------------------------------------------- + + @Override + public Query rewrite(IndexSearcher indexSearcher) throws IOException { + // Only apply the optimized path when there is no explicit filter. + // Live-document filtering (deletions) is handled via acceptDocs below. + if (filter != null) { + return super.rewrite(indexSearcher); + } + // CAGRA search is capped at k=1024. + if (k > 1024) { + return super.rewrite(indexSearcher); + } + + 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 or has no CAGRA index for this field. + List gpuReaders = new ArrayList<>(leaves.size()); + for (LeafReaderContext ctx : leaves) { + CuVS2510GPUVectorsReader gpuReader = unwrapGpuReader(ctx); + if (gpuReader == null || gpuReader.getCagraIndexForField(field) == null) { + return super.rewrite(indexSearcher); + } + gpuReaders.add(gpuReader); + } + + // Build one CagraIndex + CagraQuery per segment. + CuVSResources resources = getCuVSResourcesInstance(); + List cagraIndices = new ArrayList<>(leaves.size()); + List cagraQueries = new ArrayList<>(leaves.size()); + + try { + float[] target = getTargetCopy(); + CagraSearchParams searchParams = + new CagraSearchParams.Builder() + .withItopkSize(Math.max(iTopK, k)) + .withSearchWidth(searchWidth) + .build(); + + for (int i = 0; i < leaves.size(); i++) { + LeafReaderContext ctx = leaves.get(i); + cagraIndices.add(gpuReaders.get(i).getCagraIndexForField(field)); + + // Pass live-document bits as the prefilter so deleted docs are excluded. + Bits liveDocs = ctx.reader().getLiveDocs(); + CagraQuery query = + buildCagraQuery(resources, target, k, searchParams, liveDocs, gpuReaders.get(i), ctx); + cagraQueries.add(query); + } + + MultiSegmentSearchResults results = + MultiSegmentCagraSearch.search(resources, cagraIndices, cagraQueries, k); + + if (results.count() == 0) { + return new MatchNoDocsQuery(); + } + + // Map (segmentIdx, ordinal) → global Lucene doc ID; compute normalized score. + ScoreDoc[] scoreDocs = new ScoreDoc[results.count()]; + for (int j = 0; j < results.count(); j++) { + int segIdx = results.getSegmentIndex(j); + int ordinal = results.getOrdinal(j); + float dist = results.getDistance(j); + + LeafReaderContext ctx = leaves.get(segIdx); + int localDoc = gpuReaders.get(segIdx).ordToDoc(field, ordinal); + int globalDoc = ctx.docBase + localDoc; + 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) { + if (t instanceof IOException) throw (IOException) t; + if (t instanceof RuntimeException) throw (RuntimeException) t; + throw new RuntimeException("Multi-segment GPU search failed", t); + } + } + + // ------------------------------------------------------------------------- + // Per-segment fallback path (used when filter != null or k > 1024) + // ------------------------------------------------------------------------- + @Override protected TopDocs approximateSearch( LeafReaderContext context, @@ -47,12 +180,200 @@ 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); + context.reader().searchNearestVectors(field, getTargetCopy(), results, acceptDocs); return results.topDocs(); } + + // ------------------------------------------------------------------------- + // 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) { + var unwrapped = FilterLeafReader.unwrap(ctx.reader()); + if (!(unwrapped instanceof CodecReader)) return null; + KnnVectorsReader vr = ((CodecReader) unwrapped).getVectorReader(); + return (vr instanceof CuVS2510GPUVectorsReader gpuReader) ? gpuReader : null; + } + + /** + * Builds a {@link CagraQuery} for a single segment, incorporating live-document filtering + * via a prefilter bitset when deletions are present. + */ + private CagraQuery buildCagraQuery( + CuVSResources resources, + float[] target, + int topK, + CagraSearchParams searchParams, + Bits liveDocs, + CuVS2510GPUVectorsReader gpuReader, + LeafReaderContext ctx) + throws IOException { + CuVSMatrix.Builder vectorBuilder = + CuVSMatrix.deviceBuilder(resources, 1, target.length, CuVSMatrix.DataType.FLOAT); + vectorBuilder.addVector(target); + CuVSMatrix queryVector = vectorBuilder.build(); + + CagraQuery.Builder queryBuilder = + new CagraQuery.Builder(resources) + .withTopK(topK) + .withSearchParams(searchParams) + .withQueryVectors(queryVector); + + if (liveDocs != null) { + // Convert liveDocs to a BitSet over vector ordinals so CAGRA can filter on GPU. + var rawValues = gpuReader.getFloatVectorValues(field); + Bits acceptedOrds = rawValues.getAcceptOrds(liveDocs); + int length = acceptedOrds.length(); + BitSet mask = new BitSet(length); + for (int i = 0; i < length; i++) { + if (acceptedOrds.get(i)) mask.set(i); + } + queryBuilder.withPrefilter(mask, length); + } + + return queryBuilder.build(); + } + + /** + * 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; + for (ScoreDoc sd : scoreDocs) { + if (sd.doc >= base && sd.doc < maxDoc) { + localDocs[count] = sd.doc - base; + localScores[count] = sd.score * boost; + count++; + } + } + if (count == 0) return null; + final int n = count; + 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) { + while (pos < n && sortedDocs[pos] < target) pos++; + return docID(); + } + + @Override + public long cost() { + return n; + } + }; + } + + @Override + public float getMaxScore(int upTo) { + return Float.MAX_VALUE; + } + + @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, "GPU multi-segment CAGRA search"); + } + } + 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); + } + }; + } } From bdb8e661d11aff35419ca4fc22ad01abb94e3977 Mon Sep 17 00:00:00 2001 From: James Xia Date: Wed, 1 Apr 2026 17:10:36 -0700 Subject: [PATCH 02/41] Enable async memory resource and fix query vector upload in GPU search - CuVS2510GPUVectorsFormat: call CuVSProvider.provider().enableRMMAsyncMemory() in the static initializer so that cuda_async_memory_resource is active for the lifetime of the codec. This makes CAGRA workspace deallocations stream-ordered and non-blocking, which is required for the CudaStreamPool to provide any parallelism benefit. - GPUKnnFloatVectorQuery: upload the query vector to device once before the per-segment loop and share the resulting CuVSMatrix across all CagraQuery instances, reducing host-to-device copies from O(numSegments) to 1 per query. Wrap the shared device matrix in try-with-resources to close the RMM allocation promptly after MultiSegmentCagraSearch.search() returns. - FilterCuVSProvider: delegate enableRMMAsyncMemory() to the wrapped provider. --- .../cuvs/lucene/CuVS2510GPUVectorsFormat.java | 2 + .../cuvs/lucene/FilterCuVSProvider.java | 5 ++ .../cuvs/lucene/GPUKnnFloatVectorQuery.java | 76 ++++++++++--------- 3 files changed, 47 insertions(+), 36 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java index ccc61eae..f16ea026 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java @@ -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/FilterCuVSProvider.java b/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java index 07e64fb1..ac7dbc4e 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java +++ b/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java @@ -145,6 +145,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 a6ea5690..562ade99 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -127,36 +127,45 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException { .withSearchWidth(searchWidth) .build(); - for (int i = 0; i < leaves.size(); i++) { - LeafReaderContext ctx = leaves.get(i); - cagraIndices.add(gpuReaders.get(i).getCagraIndexForField(field)); - - // Pass live-document bits as the prefilter so deleted docs are excluded. - Bits liveDocs = ctx.reader().getLiveDocs(); - CagraQuery query = - buildCagraQuery(resources, target, k, searchParams, liveDocs, gpuReaders.get(i), ctx); - cagraQueries.add(query); - } - - MultiSegmentSearchResults results = - MultiSegmentCagraSearch.search(resources, cagraIndices, cagraQueries, k); - - if (results.count() == 0) { - return new MatchNoDocsQuery(); - } - - // Map (segmentIdx, ordinal) → global Lucene doc ID; compute normalized score. - ScoreDoc[] scoreDocs = new ScoreDoc[results.count()]; - for (int j = 0; j < results.count(); j++) { - int segIdx = results.getSegmentIndex(j); - int ordinal = results.getOrdinal(j); - float dist = results.getDistance(j); - - LeafReaderContext ctx = leaves.get(segIdx); - int localDoc = gpuReaders.get(segIdx).ordToDoc(field, ordinal); - int globalDoc = ctx.docBase + localDoc; - float score = 1.0f / (1.0f + dist); - scoreDocs[j] = new ScoreDoc(globalDoc, score); + // Upload the query vector to device once and share it across all per-segment CagraQueries. + 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++) { + LeafReaderContext ctx = leaves.get(i); + cagraIndices.add(gpuReaders.get(i).getCagraIndexForField(field)); + + // Pass live-document bits as the prefilter so deleted docs are excluded. + Bits liveDocs = ctx.reader().getLiveDocs(); + CagraQuery query = + buildCagraQuery( + resources, queryVector, k, searchParams, liveDocs, gpuReaders.get(i), ctx); + cagraQueries.add(query); + } + + MultiSegmentSearchResults results = + MultiSegmentCagraSearch.search(resources, cagraIndices, cagraQueries, k); + + if (results.count() == 0) { + return new MatchNoDocsQuery(); + } + + // Map (segmentIdx, ordinal) → global Lucene doc ID; compute normalized score. + scoreDocs = new ScoreDoc[results.count()]; + for (int j = 0; j < results.count(); j++) { + int segIdx = results.getSegmentIndex(j); + int ordinal = results.getOrdinal(j); + float dist = results.getDistance(j); + + LeafReaderContext ctx = leaves.get(segIdx); + int localDoc = gpuReaders.get(segIdx).ordToDoc(field, ordinal); + int globalDoc = ctx.docBase + localDoc; + float score = 1.0f / (1.0f + dist); + scoreDocs[j] = new ScoreDoc(globalDoc, score); + } } Arrays.sort(scoreDocs, Comparator.comparingDouble((ScoreDoc sd) -> sd.score).reversed()); @@ -207,18 +216,13 @@ private static CuVS2510GPUVectorsReader unwrapGpuReader(LeafReaderContext ctx) { */ private CagraQuery buildCagraQuery( CuVSResources resources, - float[] target, + CuVSMatrix queryVector, int topK, CagraSearchParams searchParams, Bits liveDocs, CuVS2510GPUVectorsReader gpuReader, LeafReaderContext ctx) throws IOException { - CuVSMatrix.Builder vectorBuilder = - CuVSMatrix.deviceBuilder(resources, 1, target.length, CuVSMatrix.DataType.FLOAT); - vectorBuilder.addVector(target); - CuVSMatrix queryVector = vectorBuilder.build(); - CagraQuery.Builder queryBuilder = new CagraQuery.Builder(resources) .withTopK(topK) From 7a8b26b8e1a1d866a091c2dbc445d2d72bef88be Mon Sep 17 00:00:00 2001 From: James Xia Date: Mon, 13 Apr 2026 07:16:16 -0700 Subject: [PATCH 03/41] Expose CAGRA SearchAlgo parameter --- .../cuvs/lucene/CuVS2510GPUVectorsReader.java | 1 + .../cuvs/lucene/GPUKnnFloatVectorQuery.java | 29 +++++++++++++++++-- .../lucene/GPUPerLeafCuVSKnnCollector.java | 15 +++++++++- 3 files changed, 42 insertions(+), 3 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java index 71169af0..c99d43a9 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java @@ -443,6 +443,7 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits new CagraSearchParams.Builder() .withItopkSize(Math.max(collector.getiTopK(), topK)) .withSearchWidth(collector.getSearchWidth()) + .withAlgo(collector.getSearchAlgo()) .build(); } else { // Setting itopK as topK because in any case iTopK should be ATLEAST equal to topK diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index 562ade99..6974c9d1 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -63,9 +63,10 @@ public class GPUKnnFloatVectorQuery extends KnnFloatVectorQuery { private final int iTopK; private final int searchWidth; + private final CagraSearchParams.SearchAlgo searchAlgo; /** - * Initializes {@link GPUKnnFloatVectorQuery}. + * Initializes {@link GPUKnnFloatVectorQuery} with {@link CagraSearchParams.SearchAlgo#AUTO}. * * @param field the vector field name * @param target the query vector @@ -76,9 +77,32 @@ public class GPUKnnFloatVectorQuery extends KnnFloatVectorQuery { */ public GPUKnnFloatVectorQuery( String field, float[] target, int k, Query filter, int iTopK, int searchWidth) { + this(field, target, k, filter, iTopK, searchWidth, 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 searchAlgo CAGRA search algorithm + */ + public GPUKnnFloatVectorQuery( + String field, + float[] target, + int k, + Query filter, + int iTopK, + int searchWidth, + CagraSearchParams.SearchAlgo searchAlgo) { super(field, target, k, filter); this.iTopK = iTopK; this.searchWidth = searchWidth; + this.searchAlgo = searchAlgo; } // ------------------------------------------------------------------------- @@ -125,6 +149,7 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException { new CagraSearchParams.Builder() .withItopkSize(Math.max(iTopK, k)) .withSearchWidth(searchWidth) + .withAlgo(searchAlgo) .build(); // Upload the query vector to device once and share it across all per-segment CagraQueries. @@ -190,7 +215,7 @@ protected TopDocs approximateSearch( KnnCollectorManager knnCollectorManager) throws IOException { GPUPerLeafCuVSKnnCollector results = - new GPUPerLeafCuVSKnnCollector(k, visitedLimit, iTopK, searchWidth); + new GPUPerLeafCuVSKnnCollector(k, visitedLimit, iTopK, searchWidth, searchAlgo); context.reader().searchNearestVectors(field, getTargetCopy(), results, acceptDocs); return results.topDocs(); } diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java b/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java index 421b2baa..6686415c 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java @@ -4,6 +4,7 @@ */ package com.nvidia.cuvs.lucene; +import com.nvidia.cuvs.CagraSearchParams; import org.apache.lucene.search.TopKnnCollector; /** @@ -15,6 +16,7 @@ class GPUPerLeafCuVSKnnCollector extends TopKnnCollector { private int iTopK; private int searchWidth; + private CagraSearchParams.SearchAlgo searchAlgo; /** * Initializes {@link GPUPerLeafCuVSKnnCollector} @@ -22,11 +24,18 @@ class GPUPerLeafCuVSKnnCollector extends TopKnnCollector { * @param topK the topK value * @param iTopK the iTopK value * @param searchWidth the search width + * @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, + CagraSearchParams.SearchAlgo searchAlgo) { super(topK, visitLimit); this.iTopK = iTopK > topK ? iTopK : topK; this.searchWidth = searchWidth; + this.searchAlgo = searchAlgo; } public int getiTopK() { @@ -36,4 +45,8 @@ public int getiTopK() { public int getSearchWidth() { return searchWidth; } + + public CagraSearchParams.SearchAlgo getSearchAlgo() { + return searchAlgo; + } } From aeda8bf2a50cfa45f28a1ff640b843c2d8d9461f Mon Sep 17 00:00:00 2001 From: James Xia Date: Mon, 13 Apr 2026 16:50:53 -0700 Subject: [PATCH 04/41] Add persistent kernel support and fix runner hash stability GPUKnnFloatVectorQuery / GPUPerLeafCuVSKnnCollector: - Add persistent, persistentLifetime, and persistentDeviceUsage parameters, threaded through all constructor overloads and forwarded to CagraSearchParams.Builder in both the multi-segment rewrite() path and the per-segment approximateSearch() fallback path. - Add threadBlockSize parameter (0 = auto) to allow tuning of the persistent kernel's worker_queue_size, which determines how many concurrent query threads can run without latency increase. Fix persistent-runner hash instability across segments (rewrite() path): - When max_iterations is 0 (auto), CAGRA computes it from each segment's dataset size. Different-sized segments produce different values, causing a distinct runner hash per segment and a destroy/recreate cycle on every search call. - Add computeMaxIterations(), which mirrors adjust_search_params() from search_plan.cuh, and call it once using the largest segment's graph size and degree. All segments then share the same max_iterations, producing a stable runner hash across the full multi-segment query. CuVS2510GPUVectorsReader: - Forward threadBlockSize, persistent, persistentLifetime, and persistentDeviceUsage from GPUPerLeafCuVSKnnCollector to CagraSearchParams.Builder in the per-segment fallback path. --- .../cuvs/lucene/CuVS2510GPUVectorsReader.java | 4 + .../cuvs/lucene/GPUKnnFloatVectorQuery.java | 122 ++++++++++++++++-- .../lucene/GPUPerLeafCuVSKnnCollector.java | 34 ++++- 3 files changed, 148 insertions(+), 12 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java index c99d43a9..4cb4b652 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java @@ -443,7 +443,11 @@ 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()) .withAlgo(collector.getSearchAlgo()) + .withPersistent(collector.isPersistent()) + .withPersistentLifetime(collector.getPersistentLifetime()) + .withPersistentDeviceUsage(collector.getPersistentDeviceUsage()) .build(); } else { // Setting itopK as topK because in any case iTopK should be ATLEAST equal to topK diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index 6974c9d1..be6b508c 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -63,10 +63,15 @@ public class GPUKnnFloatVectorQuery extends KnnFloatVectorQuery { private final int iTopK; private final int searchWidth; + private final int threadBlockSize; private final CagraSearchParams.SearchAlgo searchAlgo; + private final boolean persistent; + private final float persistentLifetime; + private final float persistentDeviceUsage; /** - * Initializes {@link GPUKnnFloatVectorQuery} with {@link CagraSearchParams.SearchAlgo#AUTO}. + * Initializes {@link GPUKnnFloatVectorQuery} with {@link CagraSearchParams.SearchAlgo#AUTO}, + * persistent kernel disabled, and max_iterations auto-selected (0). * * @param field the vector field name * @param target the query vector @@ -77,19 +82,20 @@ public class GPUKnnFloatVectorQuery extends KnnFloatVectorQuery { */ public GPUKnnFloatVectorQuery( String field, float[] target, int k, Query filter, int iTopK, int searchWidth) { - this(field, target, k, filter, iTopK, searchWidth, CagraSearchParams.SearchAlgo.AUTO); + this(field, target, k, filter, iTopK, searchWidth, 0, CagraSearchParams.SearchAlgo.AUTO); } /** - * Initializes {@link GPUKnnFloatVectorQuery}. + * Initializes {@link GPUKnnFloatVectorQuery} with persistent kernel disabled. * - * @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 searchAlgo CAGRA search algorithm + * @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 searchAlgo CAGRA search algorithm */ public GPUKnnFloatVectorQuery( String field, @@ -98,11 +104,57 @@ public GPUKnnFloatVectorQuery( Query filter, int iTopK, int searchWidth, + int threadBlockSize, CagraSearchParams.SearchAlgo searchAlgo) { + this( + field, + target, + k, + filter, + iTopK, + searchWidth, + threadBlockSize, + searchAlgo, + false, + 2.0f, + 1.0f); + } + + /** + * 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; controls worker_queue_size) + * @param searchAlgo CAGRA search algorithm + * @param persistent whether to use the persistent kernel + * @param persistentLifetime persistent kernel lifetime in seconds + * @param persistentDeviceUsage fraction of GPU SMs for the persistent kernel + */ + public GPUKnnFloatVectorQuery( + String field, + float[] target, + int k, + Query filter, + int iTopK, + int searchWidth, + int threadBlockSize, + CagraSearchParams.SearchAlgo searchAlgo, + boolean persistent, + float persistentLifetime, + float persistentDeviceUsage) { super(field, target, k, filter); this.iTopK = iTopK; this.searchWidth = searchWidth; + this.threadBlockSize = threadBlockSize; this.searchAlgo = searchAlgo; + this.persistent = persistent; + this.persistentLifetime = persistentLifetime; + this.persistentDeviceUsage = persistentDeviceUsage; } // ------------------------------------------------------------------------- @@ -138,6 +190,21 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException { gpuReaders.add(gpuReader); } + // Compute max_iterations from the largest segment so all segments hash identically, + // preventing persistent-runner destroy/recreate churn when segment sizes differ. + int maxDatasetSize = 0; + int graphDegree = 0; + for (int i = 0; i < gpuReaders.size(); i++) { + CuVSMatrix graph = gpuReaders.get(i).getCagraIndexForField(field).getGraph(); + int segSize = (int) graph.size(); + if (segSize > maxDatasetSize) { + maxDatasetSize = segSize; + graphDegree = (int) graph.columns(); + } + } + int maxIterations = + computeMaxIterations(Math.max(iTopK, k), searchWidth, maxDatasetSize, graphDegree); + // Build one CagraIndex + CagraQuery per segment. CuVSResources resources = getCuVSResourcesInstance(); List cagraIndices = new ArrayList<>(leaves.size()); @@ -149,7 +216,12 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException { new CagraSearchParams.Builder() .withItopkSize(Math.max(iTopK, k)) .withSearchWidth(searchWidth) + .withMaxIterations(maxIterations) + .withThreadBlockSize(threadBlockSize) .withAlgo(searchAlgo) + .withPersistent(persistent) + .withPersistentLifetime(persistentLifetime) + .withPersistentDeviceUsage(persistentDeviceUsage) .build(); // Upload the query vector to device once and share it across all per-segment CagraQueries. @@ -215,7 +287,16 @@ protected TopDocs approximateSearch( KnnCollectorManager knnCollectorManager) throws IOException { GPUPerLeafCuVSKnnCollector results = - new GPUPerLeafCuVSKnnCollector(k, visitedLimit, iTopK, searchWidth, searchAlgo); + new GPUPerLeafCuVSKnnCollector( + k, + visitedLimit, + iTopK, + searchWidth, + threadBlockSize, + searchAlgo, + persistent, + persistentLifetime, + persistentDeviceUsage); context.reader().searchNearestVectors(field, getTargetCopy(), results, acceptDocs); return results.topDocs(); } @@ -224,6 +305,25 @@ protected TopDocs approximateSearch( // Helpers // ------------------------------------------------------------------------- + /** + * Mirrors {@code adjust_search_params()} from {@code search_plan.cuh}: computes the + * {@code max_iterations} value that CAGRA would auto-select for the given parameters. + * + *

Called with the largest segment's size so that all segments produce the same + * value and therefore the same persistent-runner hash, preventing destroy/recreate churn. + */ + private static int computeMaxIterations( + int itopkSize, int searchWidth, int datasetSize, int graphDegree) { + int maxIter = itopkSize / searchWidth; + long numReachableNodes = 1; + long factor = Math.max(2L, graphDegree / 2); + while (numReachableNodes < datasetSize) { + numReachableNodes *= factor; + maxIter++; + } + return maxIter; + } + /** * Unwraps the {@link LeafReaderContext}'s reader to a {@link CuVS2510GPUVectorsReader}, or * returns {@code null} if the reader is not of that type. diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java b/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java index 6686415c..1761bfa5 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java @@ -16,7 +16,11 @@ class GPUPerLeafCuVSKnnCollector extends TopKnnCollector { private int iTopK; private int searchWidth; + private int threadBlockSize; private CagraSearchParams.SearchAlgo searchAlgo; + private boolean persistent; + private float persistentLifetime; + private float persistentDeviceUsage; /** * Initializes {@link GPUPerLeafCuVSKnnCollector} @@ -24,18 +28,30 @@ 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 searchAlgo the CAGRA search algorithm + * @param persistent whether to use the persistent kernel + * @param persistentLifetime persistent kernel lifetime in seconds + * @param persistentDeviceUsage fraction of GPU SMs for the persistent kernel */ public GPUPerLeafCuVSKnnCollector( int topK, int visitLimit, int iTopK, int searchWidth, - CagraSearchParams.SearchAlgo searchAlgo) { + int threadBlockSize, + CagraSearchParams.SearchAlgo searchAlgo, + boolean persistent, + float persistentLifetime, + float persistentDeviceUsage) { super(topK, visitLimit); this.iTopK = iTopK > topK ? iTopK : topK; this.searchWidth = searchWidth; + this.threadBlockSize = threadBlockSize; this.searchAlgo = searchAlgo; + this.persistent = persistent; + this.persistentLifetime = persistentLifetime; + this.persistentDeviceUsage = persistentDeviceUsage; } public int getiTopK() { @@ -46,7 +62,23 @@ public int getSearchWidth() { return searchWidth; } + public int getThreadBlockSize() { + return threadBlockSize; + } + public CagraSearchParams.SearchAlgo getSearchAlgo() { return searchAlgo; } + + public boolean isPersistent() { + return persistent; + } + + public float getPersistentLifetime() { + return persistentLifetime; + } + + public float getPersistentDeviceUsage() { + return persistentDeviceUsage; + } } From 252126176cd7069d725b5c1559a2665631f65296 Mon Sep 17 00:00:00 2001 From: James Xia Date: Thu, 16 Apr 2026 16:34:35 -0700 Subject: [PATCH 05/41] Remove persistent kernel support; add workspace pool configuration Remove persistent kernel mode: - Drop persistent, persistentLifetime, and persistentDeviceUsage fields and parameters from GPUKnnFloatVectorQuery, GPUPerLeafCuVSKnnCollector, and CuVS2510GPUVectorsReader. The persistent kernel is superseded by the native multi-segment kernel (cuvsCagraSearchMultiSegment) which achieves better concurrency without the per-runner lifecycle overhead. - Collapse the 11-argument GPUKnnFloatVectorQuery constructor (which only existed to accept persistent parameters) into the standard 8-argument form. - Remove stale comments that described max_iterations uniformity in terms of persistent-runner hash stability; replace with accurate explanation (consistent search quality across segments of different sizes). Add workspace pool configuration: - Add WORKSPACE_POOL_SIZE_PROPERTY constant to ThreadLocalCuVSResourcesProvider. - On resources creation, read com.nvidia.cuvs.workspacePoolSize system property and call setWorkspacePool() if set, so callers can pre-warm the per-thread RMM pool without modifying cuvs-lucene source. --- .../cuvs/lucene/CuVS2510GPUVectorsReader.java | 3 - .../cuvs/lucene/GPUKnnFloatVectorQuery.java | 73 +++---------------- .../lucene/GPUPerLeafCuVSKnnCollector.java | 25 +------ .../ThreadLocalCuVSResourcesProvider.java | 10 ++- 4 files changed, 20 insertions(+), 91 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java index 4cb4b652..caa7fabc 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java @@ -445,9 +445,6 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits .withSearchWidth(collector.getSearchWidth()) .withThreadBlockSize(collector.getThreadBlockSize()) .withAlgo(collector.getSearchAlgo()) - .withPersistent(collector.isPersistent()) - .withPersistentLifetime(collector.getPersistentLifetime()) - .withPersistentDeviceUsage(collector.getPersistentDeviceUsage()) .build(); } else { // Setting itopK as topK because in any case iTopK should be ATLEAST equal to topK diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index be6b508c..8de69f4e 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -65,13 +65,9 @@ public class GPUKnnFloatVectorQuery extends KnnFloatVectorQuery { private final int searchWidth; private final int threadBlockSize; private final CagraSearchParams.SearchAlgo searchAlgo; - private final boolean persistent; - private final float persistentLifetime; - private final float persistentDeviceUsage; - /** * Initializes {@link GPUKnnFloatVectorQuery} with {@link CagraSearchParams.SearchAlgo#AUTO}, - * persistent kernel disabled, and max_iterations auto-selected (0). + * and max_iterations auto-selected (0). * * @param field the vector field name * @param target the query vector @@ -86,7 +82,7 @@ public GPUKnnFloatVectorQuery( } /** - * Initializes {@link GPUKnnFloatVectorQuery} with persistent kernel disabled. + * Initializes {@link GPUKnnFloatVectorQuery}. * * @param field the vector field name * @param target the query vector @@ -106,55 +102,11 @@ public GPUKnnFloatVectorQuery( int searchWidth, int threadBlockSize, CagraSearchParams.SearchAlgo searchAlgo) { - this( - field, - target, - k, - filter, - iTopK, - searchWidth, - threadBlockSize, - searchAlgo, - false, - 2.0f, - 1.0f); - } - - /** - * 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; controls worker_queue_size) - * @param searchAlgo CAGRA search algorithm - * @param persistent whether to use the persistent kernel - * @param persistentLifetime persistent kernel lifetime in seconds - * @param persistentDeviceUsage fraction of GPU SMs for the persistent kernel - */ - public GPUKnnFloatVectorQuery( - String field, - float[] target, - int k, - Query filter, - int iTopK, - int searchWidth, - int threadBlockSize, - CagraSearchParams.SearchAlgo searchAlgo, - boolean persistent, - float persistentLifetime, - float persistentDeviceUsage) { super(field, target, k, filter); this.iTopK = iTopK; this.searchWidth = searchWidth; this.threadBlockSize = threadBlockSize; this.searchAlgo = searchAlgo; - this.persistent = persistent; - this.persistentLifetime = persistentLifetime; - this.persistentDeviceUsage = persistentDeviceUsage; } // ------------------------------------------------------------------------- @@ -190,8 +142,8 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException { gpuReaders.add(gpuReader); } - // Compute max_iterations from the largest segment so all segments hash identically, - // preventing persistent-runner destroy/recreate churn when segment sizes differ. + // Compute max_iterations from the largest segment so all segments use the same value, + // ensuring consistent search quality across segments of different sizes. int maxDatasetSize = 0; int graphDegree = 0; for (int i = 0; i < gpuReaders.size(); i++) { @@ -219,9 +171,6 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException { .withMaxIterations(maxIterations) .withThreadBlockSize(threadBlockSize) .withAlgo(searchAlgo) - .withPersistent(persistent) - .withPersistentLifetime(persistentLifetime) - .withPersistentDeviceUsage(persistentDeviceUsage) .build(); // Upload the query vector to device once and share it across all per-segment CagraQueries. @@ -263,10 +212,11 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException { 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); + Arrays.sort(scoreDocs, Comparator.comparingDouble((ScoreDoc sd) -> sd.score).reversed()); + + return docAndScoreQuery(scoreDocs); + } } catch (Throwable t) { if (t instanceof IOException) throw (IOException) t; @@ -293,10 +243,7 @@ protected TopDocs approximateSearch( iTopK, searchWidth, threadBlockSize, - searchAlgo, - persistent, - persistentLifetime, - persistentDeviceUsage); + searchAlgo); context.reader().searchNearestVectors(field, getTargetCopy(), results, acceptDocs); return results.topDocs(); } @@ -310,7 +257,7 @@ protected TopDocs approximateSearch( * {@code max_iterations} value that CAGRA would auto-select for the given parameters. * *

Called with the largest segment's size so that all segments produce the same - * value and therefore the same persistent-runner hash, preventing destroy/recreate churn. + * value, ensuring consistent search quality regardless of segment size. */ private static int computeMaxIterations( int itopkSize, int searchWidth, int datasetSize, int graphDegree) { diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java b/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java index 1761bfa5..1c6d5ec9 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java @@ -18,9 +18,6 @@ class GPUPerLeafCuVSKnnCollector extends TopKnnCollector { private int searchWidth; private int threadBlockSize; private CagraSearchParams.SearchAlgo searchAlgo; - private boolean persistent; - private float persistentLifetime; - private float persistentDeviceUsage; /** * Initializes {@link GPUPerLeafCuVSKnnCollector} @@ -30,9 +27,6 @@ class GPUPerLeafCuVSKnnCollector extends TopKnnCollector { * @param searchWidth the search width * @param threadBlockSize CAGRA thread_block_size (0 = auto; controls worker_queue_size) * @param searchAlgo the CAGRA search algorithm - * @param persistent whether to use the persistent kernel - * @param persistentLifetime persistent kernel lifetime in seconds - * @param persistentDeviceUsage fraction of GPU SMs for the persistent kernel */ public GPUPerLeafCuVSKnnCollector( int topK, @@ -40,18 +34,12 @@ public GPUPerLeafCuVSKnnCollector( int iTopK, int searchWidth, int threadBlockSize, - CagraSearchParams.SearchAlgo searchAlgo, - boolean persistent, - float persistentLifetime, - float persistentDeviceUsage) { + CagraSearchParams.SearchAlgo searchAlgo) { super(topK, visitLimit); this.iTopK = iTopK > topK ? iTopK : topK; this.searchWidth = searchWidth; this.threadBlockSize = threadBlockSize; this.searchAlgo = searchAlgo; - this.persistent = persistent; - this.persistentLifetime = persistentLifetime; - this.persistentDeviceUsage = persistentDeviceUsage; } public int getiTopK() { @@ -70,15 +58,4 @@ public CagraSearchParams.SearchAlgo getSearchAlgo() { return searchAlgo; } - public boolean isPersistent() { - return persistent; - } - - public float getPersistentLifetime() { - return persistentLifetime; - } - - public float getPersistentDeviceUsage() { - return persistentDeviceUsage; - } } diff --git a/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java b/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java index 9e259e27..69d37a36 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java +++ b/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java @@ -42,9 +42,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, From 9abe58aa1da1d9e82fea9aaa9171bc135e761b1b Mon Sep 17 00:00:00 2001 From: Vivek Narang Date: Wed, 6 May 2026 18:39:21 -0400 Subject: [PATCH 06/41] Modify example to demo prefilter. --- .../examples/IndexAndSearchonGPUExample.java | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/examples/src/main/java/com/nvidia/cuvs/lucene/examples/IndexAndSearchonGPUExample.java b/examples/src/main/java/com/nvidia/cuvs/lucene/examples/IndexAndSearchonGPUExample.java index 834544fe..fbf31fe2 100644 --- a/examples/src/main/java/com/nvidia/cuvs/lucene/examples/IndexAndSearchonGPUExample.java +++ b/examples/src/main/java/com/nvidia/cuvs/lucene/examples/IndexAndSearchonGPUExample.java @@ -30,9 +30,12 @@ import org.apache.lucene.index.IndexWriterConfig; import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.Term; import org.apache.lucene.search.IndexSearcher; import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.Query; import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TermQuery; import org.apache.lucene.search.TopDocs; import org.apache.lucene.store.Directory; import org.apache.lucene.store.FSDirectory; @@ -56,6 +59,7 @@ public static void main(String[] args) throws Exception { final int COMMIT_FREQ = 2000; final String ID_FIELD = "id"; final String VECTOR_FIELD = "vector_field"; + final String STATUS_FIELD = "status"; int numDocs = 2000; int dimension = 32; @@ -70,6 +74,9 @@ public static void main(String[] args) throws Exception { Document document = new Document(); document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); + // Alternatively flip the status's for documents as an example. + document.add( + new StringField(STATUS_FIELD, i % 2 == 0 ? "active" : "in-active", Field.Store.YES)); indexWriter.addDocument(document); count -= 1; if (count == 0) { @@ -103,8 +110,11 @@ public static void main(String[] args) throws Exception { float[] queryVector = generateDataset(random, 1, dimension)[0]; log.log(Level.FINE, "Query vector: " + Arrays.toString(queryVector)); + // Make a filter query + Query filter = new TermQuery(new Term(STATUS_FIELD, "active")); + KnnFloatVectorQuery query = - new GPUKnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK, null, topK, 1); + new GPUKnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK, filter, topK, 1); TopDocs results = searcher.search(query, topK); log.log(Level.FINE, "Search results (" + results.totalHits + " total hits):"); @@ -113,12 +123,15 @@ public static void main(String[] args) throws Exception { ScoreDoc scoreDoc = results.scoreDocs[i]; Document doc = searcher.storedFields().document(scoreDoc.doc); String id = doc.get(ID_FIELD); + String status = doc.get(STATUS_FIELD); log.log( Level.FINE, " Rank " + (i + 1) + ": doc " + scoreDoc.doc + + " status " + + status + " (id=" + id + "), score=" From 92f4f3f29920161c3c8b4c66db17a0b94e2ac7e7 Mon Sep 17 00:00:00 2001 From: James Xia Date: Sat, 16 May 2026 10:10:04 -0700 Subject: [PATCH 07/41] Enable GPU multi-segment search with explicit Lucene filter queries GPUKnnFloatVectorQuery previously fell back to the per-segment CPU path whenever a filter query was present. This change extends the GPU fast path to handle filters by packing the accepted-ordinal bitset for every segment into a single FilterBitsetHandle and passing it to MultiSegmentCagraSearch. FilterBitsetCache (new): 16-entry shared LRU cache keyed by (filter Query, per-segment reader-cache keys, field). Per-reader keys are used rather than core keys so that liveDocs changes caused by a reader reopen automatically invalidate the cached bitset. GPUKnnFloatVectorQuery changes: - Remove the early bail-out for filter != null - Add buildOrGetCachedFilterHandle: checks FilterBitsetCache, falls back to buildFilterHandle on miss - Add buildFilterHandle: evaluates the filter Weight per segment via FloatVectorValues.getAcceptOrds (intersecting with liveDocs), then packs the result into combinedLongs / segBitOffsets for FilterBitsetHandle - When a handle is present, CagraQuery is built without a per-query prefilter (the bitset already encodes the intersection) --- .../nvidia/cuvs/lucene/FilterBitsetCache.java | 64 ++++++ .../cuvs/lucene/GPUKnnFloatVectorQuery.java | 206 ++++++++++++++++-- .../lucene/GPUPerLeafCuVSKnnCollector.java | 1 - 3 files changed, 247 insertions(+), 24 deletions(-) create mode 100644 src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java 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..5d422c7a --- /dev/null +++ b/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java @@ -0,0 +1,64 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.FilterBitsetHandle; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +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 (per-thread LRU). Entries are evicted in LRU order and closed + * on eviction to free the host arrays and signal the device cache. + */ +final class FilterBitsetCache { + + private static final int MAX_HOST_ENTRIES = 16; + + 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) { + eldest.getValue().close(); + return true; + } + return false; + } + }; + + private FilterBitsetCache() {} + + static synchronized FilterBitsetHandle get( + Query filter, List segReaderKeys, String field) { + return CACHE.get(new FilterCacheKey(filter, segReaderKeys, field)); + } + + static synchronized void put( + Query filter, List segReaderKeys, String field, FilterBitsetHandle handle) { + CACHE.put(new FilterCacheKey(filter, segReaderKeys, field), handle); + } +} diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index 8de69f4e..e5677b50 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -11,6 +11,7 @@ import com.nvidia.cuvs.CagraSearchParams; import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.CuVSResources; +import com.nvidia.cuvs.FilterBitsetHandle; import com.nvidia.cuvs.MultiSegmentCagraSearch; import com.nvidia.cuvs.MultiSegmentSearchResults; import java.io.IOException; @@ -22,6 +23,7 @@ import org.apache.lucene.codecs.KnnVectorsReader; 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; @@ -39,13 +41,13 @@ 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 {@link KnnFloatVectorQuery} for GPU-only search. * *

When all index segments use {@link CuVS2510GPUVectorsReader} and the query uses CAGRA - * (k ≤ 1024, no explicit filter), {@link #rewrite} runs a globally-optimized - * multi-segment search: + * (k ≤ 1024), {@link #rewrite} runs a globally-optimized multi-segment search: *

    *
  1. All per-segment CAGRA searches write into a single shared device buffer without any * per-segment device-to-host copy or stream synchronization.
    2. @@ -53,9 +55,13 @@ *
  2. Results are copied to host in one pass and mapped to Lucene doc IDs.
    3. *
* - *

Falls back to the standard per-segment Lucene path when the optimized path cannot be - * applied (mixed segment types, explicit query filter, k > 1024, brute-force - * fallback needed, or missing CAGRA index). + *

When an explicit {@code filter} query is present, each segment's matching doc IDs are + * intersected with {@code liveDocs} and packed into a combined bitset transferred to the GPU once + * per unique (filter, reader-state, field) triple via {@link FilterBitsetCache}. + * + *

Falls back to the standard per-segment Lucene path when the optimized path cannot be applied + * (mixed segment types, k > 1024, brute-force fallback needed, or missing CAGRA + * index). * * @since 25.10 */ @@ -65,6 +71,7 @@ public class GPUKnnFloatVectorQuery extends KnnFloatVectorQuery { private final int searchWidth; private final int threadBlockSize; private final CagraSearchParams.SearchAlgo searchAlgo; + /** * Initializes {@link GPUKnnFloatVectorQuery} with {@link CagraSearchParams.SearchAlgo#AUTO}, * and max_iterations auto-selected (0). @@ -115,11 +122,6 @@ public GPUKnnFloatVectorQuery( @Override public Query rewrite(IndexSearcher indexSearcher) throws IOException { - // Only apply the optimized path when there is no explicit filter. - // Live-document filtering (deletions) is handled via acceptDocs below. - if (filter != null) { - return super.rewrite(indexSearcher); - } // CAGRA search is capped at k=1024. if (k > 1024) { return super.rewrite(indexSearcher); @@ -157,6 +159,13 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException { int maxIterations = computeMaxIterations(Math.max(iTopK, k), searchWidth, maxDatasetSize, graphDegree); + // Build filter handle when an explicit filter query is present. + // The handle encodes (filter ∩ liveDocs) per segment and is cached by reader state. + FilterBitsetHandle filterHandle = null; + if (filter != null) { + filterHandle = buildOrGetCachedFilterHandle(indexSearcher, leaves, gpuReaders); + } + // Build one CagraIndex + CagraQuery per segment. CuVSResources resources = getCuVSResourcesInstance(); List cagraIndices = new ArrayList<>(leaves.size()); @@ -184,16 +193,30 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException { LeafReaderContext ctx = leaves.get(i); cagraIndices.add(gpuReaders.get(i).getCagraIndexForField(field)); - // Pass live-document bits as the prefilter so deleted docs are excluded. - Bits liveDocs = ctx.reader().getLiveDocs(); - CagraQuery query = - buildCagraQuery( - resources, queryVector, k, searchParams, liveDocs, gpuReaders.get(i), ctx); + CagraQuery query; + if (filterHandle != null) { + // Filter (and liveDocs) are encoded in the handle; no per-query prefilter. + query = + new CagraQuery.Builder(resources) + .withTopK(k) + .withSearchParams(searchParams) + .withQueryVectors(queryVector) + .build(); + } else { + // No explicit filter; pass live-document bits as the prefilter. + Bits liveDocs = ctx.reader().getLiveDocs(); + query = + buildCagraQuery( + resources, queryVector, k, searchParams, liveDocs, gpuReaders.get(i), ctx); + } cagraQueries.add(query); } MultiSegmentSearchResults results = - MultiSegmentCagraSearch.search(resources, cagraIndices, cagraQueries, k); + filterHandle != null + ? MultiSegmentCagraSearch.search( + resources, cagraIndices, cagraQueries, k, filterHandle) + : MultiSegmentCagraSearch.search(resources, cagraIndices, cagraQueries, k); if (results.count() == 0) { return new MatchNoDocsQuery(); @@ -226,7 +249,7 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException { } // ------------------------------------------------------------------------- - // Per-segment fallback path (used when filter != null or k > 1024) + // Per-segment fallback path (used when k > 1024 or not all GPU segments) // ------------------------------------------------------------------------- @Override @@ -238,16 +261,153 @@ protected TopDocs approximateSearch( throws IOException { GPUPerLeafCuVSKnnCollector results = new GPUPerLeafCuVSKnnCollector( - k, - visitedLimit, - iTopK, - searchWidth, - threadBlockSize, - searchAlgo); + k, visitedLimit, iTopK, searchWidth, threadBlockSize, 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. + */ + 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 without caching. + return buildFilterHandle(indexSearcher, leaves, gpuReaders); + } + segReaderKeys.add(helper.getKey()); + } + + FilterBitsetHandle cached = FilterBitsetCache.get(filter, segReaderKeys, field); + if (cached != null) return cached; + + FilterBitsetHandle handle = buildFilterHandle(indexSearcher, leaves, gpuReaders); + FilterBitsetCache.put(filter, segReaderKeys, field, handle); + return handle; + } + + /** + * Evaluates {@link #filter} per segment, intersects with liveDocs, and packs the result into a + * new {@link FilterBitsetHandle}. + */ + private FilterBitsetHandle buildFilterHandle( + IndexSearcher indexSearcher, + List leaves, + List gpuReaders) + throws IOException { + + Weight 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(); + Bits acceptDocs = evalFilter(filterWeight, ctx, 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); + } + + return new FilterBitsetHandle(combinedLongs, segBitOffsets, totalBits); + } + + /** + * 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) { + int maxDoc = ctx.reader().maxDoc(); + return new Bits() { + @Override + public boolean get(int i) { + return false; + } + + @Override + public int length() { + return 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 // ------------------------------------------------------------------------- diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java b/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java index 1c6d5ec9..6e6387cd 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java @@ -57,5 +57,4 @@ public int getThreadBlockSize() { public CagraSearchParams.SearchAlgo getSearchAlgo() { return searchAlgo; } - } From c49a1a0dc6a9ff3178bd908d91036837a3854ba4 Mon Sep 17 00:00:00 2001 From: James Xia Date: Thu, 21 May 2026 07:04:13 -0700 Subject: [PATCH 08/41] Adapt GPU search path to merged multi-partition cuVS API cuVS's search_multi_partition was refactored to take a single queries matrix and return globally merged top-k outputs (partition_ids, neighbors, distances) instead of per-partition vectors of matrix views. This commit updates GPUKnnFloatVectorQuery to match. Changes in rewrite(): - Build one shared CagraQuery (queryVector + searchParams), not one per Lucene segment; drop the per-segment construction loop. - The remaining loop only collects the CagraIndex per segment for the cuVS partition list. - Single call to MultiPartitionCagraSearch.search(...), passing the optional FilterBitsetHandle as the last argument (no more filterHandle != null branching between overloads). Filter handling: - Since a single shared CagraQuery cannot carry per-segment liveDocs via withPrefilter anymore, fold liveDocs into the FilterBitsetHandle alongside any explicit Lucene filter. The handle is now built whenever filter != null OR any segment has deletes. - buildFilterHandle extended to accept filter == null, in which case acceptDocs is just liveDocs per segment. - Removed buildCagraQuery helper that constructed per-segment liveDocs bitsets; the unused java.util.BitSet import is removed with it. --- .../cuvs/lucene/GPUKnnFloatVectorQuery.java | 134 +++++++----------- 1 file changed, 50 insertions(+), 84 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index e5677b50..a4f5d6e9 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -12,12 +12,11 @@ import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.CuVSResources; import com.nvidia.cuvs.FilterBitsetHandle; -import com.nvidia.cuvs.MultiSegmentCagraSearch; -import com.nvidia.cuvs.MultiSegmentSearchResults; +import com.nvidia.cuvs.MultiPartitionCagraSearch; +import com.nvidia.cuvs.MultiPartitionSearchResults; import java.io.IOException; import java.util.ArrayList; import java.util.Arrays; -import java.util.BitSet; import java.util.Comparator; import java.util.List; import org.apache.lucene.codecs.KnnVectorsReader; @@ -47,17 +46,16 @@ * Extends {@link KnnFloatVectorQuery} for GPU-only search. * *

When all index segments use {@link CuVS2510GPUVectorsReader} and the query uses CAGRA - * (k ≤ 1024), {@link #rewrite} runs a globally-optimized multi-segment search: - *

    - *
  1. All per-segment CAGRA searches write into a single shared device buffer without any - * per-segment device-to-host copy or stream synchronization.
    2. - *
  2. A single {@code cuvsSelectK} call finds the global top-k entirely on GPU.
    3. - *
  3. Results are copied to host in one pass and mapped to Lucene doc IDs.
    4. - *
+ * (k ≤ 1024), {@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. * - *

When an explicit {@code filter} query is present, each segment's matching doc IDs are - * intersected with {@code liveDocs} and packed into a combined bitset transferred to the GPU once - * per unique (filter, reader-state, field) triple via {@link FilterBitsetCache}. + *

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, k > 1024, brute-force fallback needed, or missing CAGRA @@ -159,17 +157,26 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException { int maxIterations = computeMaxIterations(Math.max(iTopK, k), searchWidth, maxDatasetSize, graphDegree); - // Build filter handle when an explicit filter query is present. - // The handle encodes (filter ∩ liveDocs) per segment and is cached by reader state. + // 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; + } + } FilterBitsetHandle filterHandle = null; - if (filter != null) { + if (hasExplicitFilter || hasDeletes) { filterHandle = buildOrGetCachedFilterHandle(indexSearcher, leaves, gpuReaders); } - // Build one CagraIndex + CagraQuery per segment. + // Build the per-segment CagraIndex list (one entry per Lucene segment / cuVS partition). CuVSResources resources = getCuVSResourcesInstance(); List cagraIndices = new ArrayList<>(leaves.size()); - List cagraQueries = new ArrayList<>(leaves.size()); try { float[] target = getTargetCopy(); @@ -182,7 +189,8 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException { .withAlgo(searchAlgo) .build(); - // Upload the query vector to device once and share it across all per-segment CagraQueries. + // 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); @@ -190,33 +198,18 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException { ScoreDoc[] scoreDocs; try (CuVSMatrix queryVector = vectorBuilder.build()) { for (int i = 0; i < leaves.size(); i++) { - LeafReaderContext ctx = leaves.get(i); cagraIndices.add(gpuReaders.get(i).getCagraIndexForField(field)); - - CagraQuery query; - if (filterHandle != null) { - // Filter (and liveDocs) are encoded in the handle; no per-query prefilter. - query = - new CagraQuery.Builder(resources) - .withTopK(k) - .withSearchParams(searchParams) - .withQueryVectors(queryVector) - .build(); - } else { - // No explicit filter; pass live-document bits as the prefilter. - Bits liveDocs = ctx.reader().getLiveDocs(); - query = - buildCagraQuery( - resources, queryVector, k, searchParams, liveDocs, gpuReaders.get(i), ctx); - } - cagraQueries.add(query); } - MultiSegmentSearchResults results = - filterHandle != null - ? MultiSegmentCagraSearch.search( - resources, cagraIndices, cagraQueries, k, filterHandle) - : MultiSegmentCagraSearch.search(resources, cagraIndices, cagraQueries, k); + 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(); @@ -225,7 +218,7 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException { // Map (segmentIdx, ordinal) → global Lucene doc ID; compute normalized score. scoreDocs = new ScoreDoc[results.count()]; for (int j = 0; j < results.count(); j++) { - int segIdx = results.getSegmentIndex(j); + int segIdx = results.getPartitionIndex(j); int ordinal = results.getOrdinal(j); float dist = results.getDistance(j); @@ -302,8 +295,10 @@ private FilterBitsetHandle buildOrGetCachedFilterHandle( } /** - * Evaluates {@link #filter} per segment, intersects with liveDocs, and packs the result into a - * new {@link FilterBitsetHandle}. + * 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, @@ -311,9 +306,12 @@ private FilterBitsetHandle buildFilterHandle( List gpuReaders) throws IOException { - Weight filterWeight = - indexSearcher.createWeight( - indexSearcher.rewrite(filter), ScoreMode.COMPLETE_NO_SCORES, 1.0f); + 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]; @@ -328,7 +326,9 @@ private FilterBitsetHandle buildFilterHandle( for (int i = 0; i < numSegments; i++) { LeafReaderContext ctx = leaves.get(i); Bits liveDocs = ctx.reader().getLiveDocs(); - Bits acceptDocs = evalFilter(filterWeight, ctx, liveDocs); + // 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(); @@ -442,40 +442,6 @@ private static CuVS2510GPUVectorsReader unwrapGpuReader(LeafReaderContext ctx) { return (vr instanceof CuVS2510GPUVectorsReader gpuReader) ? gpuReader : null; } - /** - * Builds a {@link CagraQuery} for a single segment, incorporating live-document filtering - * via a prefilter bitset when deletions are present. - */ - private CagraQuery buildCagraQuery( - CuVSResources resources, - CuVSMatrix queryVector, - int topK, - CagraSearchParams searchParams, - Bits liveDocs, - CuVS2510GPUVectorsReader gpuReader, - LeafReaderContext ctx) - throws IOException { - CagraQuery.Builder queryBuilder = - new CagraQuery.Builder(resources) - .withTopK(topK) - .withSearchParams(searchParams) - .withQueryVectors(queryVector); - - if (liveDocs != null) { - // Convert liveDocs to a BitSet over vector ordinals so CAGRA can filter on GPU. - var rawValues = gpuReader.getFloatVectorValues(field); - Bits acceptedOrds = rawValues.getAcceptOrds(liveDocs); - int length = acceptedOrds.length(); - BitSet mask = new BitSet(length); - for (int i = 0; i < length; i++) { - if (acceptedOrds.get(i)) mask.set(i); - } - queryBuilder.withPrefilter(mask, length); - } - - return queryBuilder.build(); - } - /** * Builds a {@link Query} that matches exactly the given pre-scored documents. * From 7f394de4a247958c86c58a6b4a4019fa101058f0 Mon Sep 17 00:00:00 2001 From: James Xia Date: Fri, 29 May 2026 17:39:13 -0700 Subject: [PATCH 09/41] Allow rewrite() to use multi-partition path for k beyond 1024 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The k>1024 short-circuit at the top of rewrite() forced large-k queries through Lucene's per-leaf approximateSearch fallback, where they fail either at the codec reader's k<=1024 gate or — when that gate doesn't fire — by trying to use a brute-force index that callers typically do not build. The multi-partition path is the right place to serve these queries: with N segments at SINGLE_CTA's itopk_size cap of 512, the per-partition × num_segments candidate pool is more than enough to assemble a global top-k well above 1024. Companion to the cuVS change that decouples the multi-partition kernel's per-partition output count from the global topk, allowing each partition to emit up to itopk_size candidates and merging across partitions to produce the final top-k. - Drop the k>1024 early-return in rewrite(). Feasibility is now enforced by cuVS's itopk_size * num_partitions >= topk RAFT_EXPECTS in search_multi_partition. - Clamp itopk_size to 512 when building CagraSearchParams, since SINGLE_CTA rejects larger values. The clamp is a no-op when the caller-requested itopk_size is already <= 512. - Update the class Javadoc to reflect that the multi-partition path is no longer k-capped. The per-leaf approximateSearch fallback's k<=1024 gate and brute-force- index assumption are unchanged; they fire only when some segment lacks a GPU reader, which is independent of this change. --- .../cuvs/lucene/GPUKnnFloatVectorQuery.java | 23 ++++++++++--------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index a4f5d6e9..58a185dd 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -45,11 +45,13 @@ /** * Extends {@link KnnFloatVectorQuery} for GPU-only search. * - *

When all index segments use {@link CuVS2510GPUVectorsReader} and the query uses CAGRA - * (k ≤ 1024), {@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. + *

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 multi-partition path supports k beyond the per-partition SINGLE_CTA cap by + * having each partition emit up to {@code itopk_size} candidates and letting the cross-partition + * select_k assemble the global top-k. * *

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 @@ -120,11 +122,6 @@ public GPUKnnFloatVectorQuery( @Override public Query rewrite(IndexSearcher indexSearcher) throws IOException { - // CAGRA search is capped at k=1024. - if (k > 1024) { - return super.rewrite(indexSearcher); - } - IndexReader reader = indexSearcher.getIndexReader(); List leaves = reader.leaves(); if (leaves.isEmpty()) { @@ -180,9 +177,13 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException { try { float[] target = getTargetCopy(); + // SINGLE_CTA caps itopk_size at 512. When k exceeds that, the multi-partition path emits up + // to itopk_size candidates per segment and merges across segments to produce the global k. + // Feasibility (itopk_size * num_segments >= k) is enforced by cuVS in search_multi_partition. + int effectiveItopk = Math.min(Math.max(iTopK, k), 512); CagraSearchParams searchParams = new CagraSearchParams.Builder() - .withItopkSize(Math.max(iTopK, k)) + .withItopkSize(effectiveItopk) .withSearchWidth(searchWidth) .withMaxIterations(maxIterations) .withThreadBlockSize(threadBlockSize) From 1b0355a9f6d699c11f169e022404ba6b2221c043 Mon Sep 17 00:00:00 2001 From: James Xia Date: Mon, 1 Jun 2026 14:21:00 -0700 Subject: [PATCH 10/41] Revert "Allow rewrite() to use multi-partition path for k beyond 1024" This reverts commit 7f394de4a247958c86c58a6b4a4019fa101058f0. --- .../cuvs/lucene/GPUKnnFloatVectorQuery.java | 23 +++++++++---------- 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index 58a185dd..a4f5d6e9 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -45,13 +45,11 @@ /** * 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 multi-partition path supports k beyond the per-partition SINGLE_CTA cap by - * having each partition emit up to {@code itopk_size} candidates and letting the cross-partition - * select_k assemble the global top-k. + *

When all index segments use {@link CuVS2510GPUVectorsReader} and the query uses CAGRA + * (k ≤ 1024), {@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. * *

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 @@ -122,6 +120,11 @@ public GPUKnnFloatVectorQuery( @Override public Query rewrite(IndexSearcher indexSearcher) throws IOException { + // CAGRA search is capped at k=1024. + if (k > 1024) { + return super.rewrite(indexSearcher); + } + IndexReader reader = indexSearcher.getIndexReader(); List leaves = reader.leaves(); if (leaves.isEmpty()) { @@ -177,13 +180,9 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException { try { float[] target = getTargetCopy(); - // SINGLE_CTA caps itopk_size at 512. When k exceeds that, the multi-partition path emits up - // to itopk_size candidates per segment and merges across segments to produce the global k. - // Feasibility (itopk_size * num_segments >= k) is enforced by cuVS in search_multi_partition. - int effectiveItopk = Math.min(Math.max(iTopK, k), 512); CagraSearchParams searchParams = new CagraSearchParams.Builder() - .withItopkSize(effectiveItopk) + .withItopkSize(Math.max(iTopK, k)) .withSearchWidth(searchWidth) .withMaxIterations(maxIterations) .withThreadBlockSize(threadBlockSize) From 05d61b48ab81aec6c0a0608ce4982b042749ad00 Mon Sep 17 00:00:00 2001 From: James Xia Date: Mon, 1 Jun 2026 16:22:14 -0700 Subject: [PATCH 11/41] Remove computeMaxIterations from GPUKnnFloatVectorQuery computeMaxIterations was added to stabilize the persistent-kernel runner hash across segments by pre-computing max_iterations on the Java side from the largest segment, so every segment received the same value. Persistent-kernel support was subsequently removed, eliminating that rationale. The cuVS multi-partition C++ entry (search_multi_partition in cagra_search.cuh) already iterates partitions to find max_dataset_size and max_graph_degree, then calls adjust_search_params(), which produces the same max_iterations value the Java side was computing. Passing max_iterations = 0 (the CagraSearchParams default) lets that path run naturally. - Delete the per-segment loop that computed max_dataset_size and graph_degree. - Drop .withMaxIterations() from the multi-partition search builder. - Delete the computeMaxIterations helper and its Javadoc. No behavior change; the C++ side computes the identical value. --- .../cuvs/lucene/GPUKnnFloatVectorQuery.java | 35 ------------------- 1 file changed, 35 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index a4f5d6e9..a01c995e 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -142,21 +142,6 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException { gpuReaders.add(gpuReader); } - // Compute max_iterations from the largest segment so all segments use the same value, - // ensuring consistent search quality across segments of different sizes. - int maxDatasetSize = 0; - int graphDegree = 0; - for (int i = 0; i < gpuReaders.size(); i++) { - CuVSMatrix graph = gpuReaders.get(i).getCagraIndexForField(field).getGraph(); - int segSize = (int) graph.size(); - if (segSize > maxDatasetSize) { - maxDatasetSize = segSize; - graphDegree = (int) graph.columns(); - } - } - int maxIterations = - computeMaxIterations(Math.max(iTopK, k), searchWidth, maxDatasetSize, graphDegree); - // 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 @@ -184,7 +169,6 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException { new CagraSearchParams.Builder() .withItopkSize(Math.max(iTopK, k)) .withSearchWidth(searchWidth) - .withMaxIterations(maxIterations) .withThreadBlockSize(threadBlockSize) .withAlgo(searchAlgo) .build(); @@ -412,25 +396,6 @@ private static void packOrdsToLongs(Bits bits, int numOrds, long[] dest, int des // Helpers // ------------------------------------------------------------------------- - /** - * Mirrors {@code adjust_search_params()} from {@code search_plan.cuh}: computes the - * {@code max_iterations} value that CAGRA would auto-select for the given parameters. - * - *

Called with the largest segment's size so that all segments produce the same - * value, ensuring consistent search quality regardless of segment size. - */ - private static int computeMaxIterations( - int itopkSize, int searchWidth, int datasetSize, int graphDegree) { - int maxIter = itopkSize / searchWidth; - long numReachableNodes = 1; - long factor = Math.max(2L, graphDegree / 2); - while (numReachableNodes < datasetSize) { - numReachableNodes *= factor; - maxIter++; - } - return maxIter; - } - /** * Unwraps the {@link LeafReaderContext}'s reader to a {@link CuVS2510GPUVectorsReader}, or * returns {@code null} if the reader is not of that type. From 2160b6311e8a1ee8bbbf58210382cd49ba7950d7 Mon Sep 17 00:00:00 2001 From: James Xia Date: Tue, 2 Jun 2026 14:40:21 -0700 Subject: [PATCH 12/41] Drop k>1024 gate from GPUKnnFloatVectorQuery.rewrite() MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The gate forced k>1024 queries through Lucene's per-leaf fallback, where they hit the codec reader's k<=1024 check and tried to invoke a brute-force index that GPU-indexed segments typically do not write — producing a confusing NPE. cuVS now handles k>1024 in the multi-partition path via MULTI_KERNEL, routed automatically by the AUTO heuristic on itopk_size or selected explicitly via cagraSearchAlgo. Feasibility is validated at the C++ side per algo: SINGLE_CTA still requires itopk_size * num_partitions >= topk, MULTI_KERNEL requires topk <= itopk_size. Failures surface as RAFT_EXPECTS messages instead of NPEs. --- .../cuvs/lucene/GPUKnnFloatVectorQuery.java | 22 ++++++++----------- 1 file changed, 9 insertions(+), 13 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index a01c995e..42985384 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -45,11 +45,13 @@ /** * Extends {@link KnnFloatVectorQuery} for GPU-only search. * - *

When all index segments use {@link CuVS2510GPUVectorsReader} and the query uses CAGRA - * (k ≤ 1024), {@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. + *

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 @@ -57,9 +59,8 @@ * (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, k > 1024, brute-force fallback needed, or missing CAGRA - * index). + *

Falls back to the standard per-segment Lucene path when the optimized path cannot be + * applied (mixed segment types or missing CAGRA index for the field on any segment). * * @since 25.10 */ @@ -120,11 +121,6 @@ public GPUKnnFloatVectorQuery( @Override public Query rewrite(IndexSearcher indexSearcher) throws IOException { - // CAGRA search is capped at k=1024. - if (k > 1024) { - return super.rewrite(indexSearcher); - } - IndexReader reader = indexSearcher.getIndexReader(); List leaves = reader.leaves(); if (leaves.isEmpty()) { From 51e03b109ab4247483c73c10617e31ebe67d9fdd Mon Sep 17 00:00:00 2001 From: jolorunyomi Date: Thu, 25 Jun 2026 10:08:33 -0500 Subject: [PATCH 13/41] Updates to pom to buidl snapshot --- .gitignore | 11 ++++-- license-header.txt | 4 +++ pom.xml | 36 +++++++++++++++++-- .../cuvs/lucene/AcceleratedHNSWParams.java | 3 +- .../cuvs/lucene/AcceleratedHNSWUtils.java | 3 +- .../cuvs/lucene/CagraIndexParamsFactory.java | 1 - .../cuvs/lucene/CuVS2510GPUSearchCodec.java | 2 +- .../cuvs/lucene/CuVS2510GPUVectorsFormat.java | 2 +- .../cuvs/lucene/CuVS2510GPUVectorsReader.java | 2 +- .../cuvs/lucene/CuVS2510GPUVectorsWriter.java | 2 +- .../com/nvidia/cuvs/lucene/FieldWriter.java | 3 +- .../nvidia/cuvs/lucene/FilterBitsetCache.java | 2 +- .../cuvs/lucene/FilterCuVSProvider.java | 2 +- .../lucene/FilterCuVSServiceProvider.java | 2 +- .../nvidia/cuvs/lucene/GPUBuiltHnswGraph.java | 2 +- .../nvidia/cuvs/lucene/GPUFieldWriter.java | 2 +- .../java/com/nvidia/cuvs/lucene/GPUIndex.java | 2 +- .../cuvs/lucene/GPUKnnFloatVectorQuery.java | 2 +- .../lucene/GPUPerLeafCuVSKnnCollector.java | 2 +- .../nvidia/cuvs/lucene/GPUSearchParams.java | 3 +- .../cuvs/lucene/IndexInputInputStream.java | 2 +- .../cuvs/lucene/IndexOutputOutputStream.java | 2 +- .../lucene/Lucene101AcceleratedHNSWCodec.java | 2 +- .../Lucene99AcceleratedHNSWVectorsFormat.java | 2 +- .../Lucene99AcceleratedHNSWVectorsWriter.java | 2 +- ...neAcceleratedHNSWBinaryQuantizedCodec.java | 2 +- ...ratedHNSWBinaryQuantizedVectorsFormat.java | 2 +- ...ratedHNSWBinaryQuantizedVectorsWriter.java | 2 +- ...neAcceleratedHNSWScalarQuantizedCodec.java | 2 +- ...ratedHNSWScalarQuantizedVectorsFormat.java | 2 +- ...ratedHNSWScalarQuantizedVectorsWriter.java | 2 +- .../nvidia/cuvs/lucene/LuceneProvider.java | 2 +- .../ThreadLocalCuVSResourcesProvider.java | 3 +- .../java/com/nvidia/cuvs/lucene/Utils.java | 2 +- .../TestAcceleratedHNSWDeletedDocuments.java | 2 +- .../lucene/TestAcceleratedHNSWParams.java | 3 +- .../nvidia/cuvs/lucene/TestBackCompat.java | 2 +- ...TestCagraToHnswSerializationAndSearch.java | 2 +- ...ializationAndSearchWithFallbackWriter.java | 2 +- ...stCuVSAcceleratedHNSWDeletedDocuments.java | 2 +- .../lucene/TestCuVSAcceleratedHNSWGaps.java | 2 +- .../cuvs/lucene/TestCuVSDeletedDocuments.java | 2 +- .../com/nvidia/cuvs/lucene/TestCuVSGaps.java | 2 +- .../TestCuVSRandomizedHNSWVectorSearch.java | 2 +- .../TestCuVSRandomizedVectorSearch.java | 2 +- .../cuvs/lucene/TestCuVSVectorsFormat.java | 2 +- .../cuvs/lucene/TestGPUSearchParams.java | 3 +- .../lucene/TestIndexOutputOutputStream.java | 2 +- ...tLucene99AcceleratedHNSWVectorsFormat.java | 2 +- .../com/nvidia/cuvs/lucene/TestMerge.java | 2 +- .../TestMultithreadedCuVSGPUSearch.java | 3 +- .../lucene/TestQuantizedVectorsFormats.java | 2 +- .../com/nvidia/cuvs/lucene/TestUtils.java | 2 +- 53 files changed, 95 insertions(+), 63 deletions(-) create mode 100644 license-header.txt diff --git a/.gitignore b/.gitignore index f7ee48d3..d288869d 100644 --- a/.gitignore +++ b/.gitignore @@ -1,9 +1,14 @@ # SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 +# Maven *.jar -target -**/.DS_Store -bin +target/ + +# IDE - IntelliJ +.idea/ +*.iml + +# IDE - Eclipse .project cuvs-workdir diff --git a/license-header.txt b/license-header.txt new file mode 100644 index 00000000..f9b093f8 --- /dev/null +++ b/license-header.txt @@ -0,0 +1,4 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) $YEAR, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ diff --git a/pom.xml b/pom.xml index e13e5d2b..602d40c8 100644 --- a/pom.xml +++ b/pom.xml @@ -12,7 +12,7 @@ com.nvidia.cuvs.lucene cuvs-lucene - 26.08.0 + 26.08.0-SNAPSHOT cuvs-lucene jar @@ -36,6 +36,18 @@ https://github.com/rapidsai/cuvs-lucene + + + rapidsai + rapidsai + cuvs@rapids.ai + + Maintainer + Committer + + + + 22 22 @@ -75,10 +87,26 @@ com.nvidia.cuvs cuvs-java - 26.08.0 + 26.08.0-SNAPSHOT + + + central-snapshots + https://central.sonatype.com/repository/maven-snapshots/ + true + false + + + + + + central + https://central.sonatype.com/repository/maven-snapshots/ + + + @@ -104,6 +132,9 @@ true false + + ${project.basedir}/license-header.txt + @@ -165,6 +196,7 @@ jar + com.nvidia.cuvs.lucene all,-missing ${project.build.directory}/javadocs diff --git a/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java b/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java index dffb5004..8a69da71 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java +++ b/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.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.CagraIndexParams.CagraGraphBuildAlgo; diff --git a/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index 2a8ab02f..c0964d93 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.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 static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.getCuVSResourcesInstance; diff --git a/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java b/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java index 20fed04b..3a881eb7 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java @@ -2,7 +2,6 @@ * 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.CagraIndexParams; diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java index 2de761a7..9eb5317b 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.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; diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java index f16ea026..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; diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java index caa7fabc..99d9cceb 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; diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java index e8fb304e..757ca057 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.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; diff --git a/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java b/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java index 627d931d..1bdbf998 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.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 static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.quantizeFloatVectorsToBinary; diff --git a/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java b/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java index 5d422c7a..3b269798 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java +++ b/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.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; diff --git a/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java b/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java index ac7dbc4e..6337a67a 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; diff --git a/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSServiceProvider.java b/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSServiceProvider.java index 83cebafa..69f5cac6 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSServiceProvider.java +++ b/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSServiceProvider.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java b/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java index b0b6286d..62203ad9 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUFieldWriter.java b/src/main/java/com/nvidia/cuvs/lucene/GPUFieldWriter.java index 45255da4..d159c5c1 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUFieldWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUFieldWriter.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; diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUIndex.java b/src/main/java/com/nvidia/cuvs/lucene/GPUIndex.java index 76cd916b..5cf047bc 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUIndex.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUIndex.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index 42985384..eb23f996 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.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; diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java b/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java index 6e6387cd..28443bde 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.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; diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUSearchParams.java b/src/main/java/com/nvidia/cuvs/lucene/GPUSearchParams.java index cc445268..f8458cf2 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUSearchParams.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUSearchParams.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.CagraIndexParams.CagraGraphBuildAlgo; diff --git a/src/main/java/com/nvidia/cuvs/lucene/IndexInputInputStream.java b/src/main/java/com/nvidia/cuvs/lucene/IndexInputInputStream.java index 3ab14122..1effe159 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/IndexInputInputStream.java +++ b/src/main/java/com/nvidia/cuvs/lucene/IndexInputInputStream.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/IndexOutputOutputStream.java b/src/main/java/com/nvidia/cuvs/lucene/IndexOutputOutputStream.java index b5fc120b..82e06737 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/IndexOutputOutputStream.java +++ b/src/main/java/com/nvidia/cuvs/lucene/IndexOutputOutputStream.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java index b4c5a33d..e9f4d6fe 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.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; diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java index 3c42707c..61c9b41c 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.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; diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index 35209ce5..0e558669 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.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; diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java index f2c1aa37..0b1653bc 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java index 0f8d9602..d1810bc5 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java index d8a98cc2..e00b2d07 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedCodec.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedCodec.java index 0c7736a0..0705ed0a 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedCodec.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsFormat.java index 534390b8..8d599a54 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsFormat.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java index 21b4be3f..a542cf89 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java index 7635e323..9caa2a6a 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.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; diff --git a/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java b/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java index 69d37a36..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; diff --git a/src/main/java/com/nvidia/cuvs/lucene/Utils.java b/src/main/java/com/nvidia/cuvs/lucene/Utils.java index c55aa7c3..e4a20d2b 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Utils.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Utils.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; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java index d164af55..cd4aaa6e 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.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; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java index 2273af8c..9dc110f8 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.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 static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_BEAM_WIDTH; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java b/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java index 2de6e660..0de813ff 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java index 96be6c01..67e6b0ad 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.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; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java index 04ad20a0..8999d8e8 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.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; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWDeletedDocuments.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWDeletedDocuments.java index d3dc866e..10d3ceee 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWDeletedDocuments.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWDeletedDocuments.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; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWGaps.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWGaps.java index b0dae38e..fa0936e6 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWGaps.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWGaps.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; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java index de868c7c..7303d04d 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.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; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java index a3e2f59a..31958ef0 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.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; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedHNSWVectorSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedHNSWVectorSearch.java index c2e2f2f6..e630a0c3 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedHNSWVectorSearch.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedHNSWVectorSearch.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; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java index 3e77e2ed..85720c30 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.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; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java index 1c75cb23..2b960bf3 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.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; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestGPUSearchParams.java b/src/test/java/com/nvidia/cuvs/lucene/TestGPUSearchParams.java index 483b689c..41f48005 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestGPUSearchParams.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestGPUSearchParams.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 static com.nvidia.cuvs.lucene.GPUSearchParams.DEFAULT_CAGRA_GRAPH_BUILD_ALGO; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestIndexOutputOutputStream.java b/src/test/java/com/nvidia/cuvs/lucene/TestIndexOutputOutputStream.java index 04b22b3d..b1eba881 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestIndexOutputOutputStream.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestIndexOutputOutputStream.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java b/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java index 2428fc50..9bb14022 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.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; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java b/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java index bcc29178..a42f55b3 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestMerge.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; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestMultithreadedCuVSGPUSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestMultithreadedCuVSGPUSearch.java index 25dc13b8..ba03861f 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestMultithreadedCuVSGPUSearch.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestMultithreadedCuVSGPUSearch.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 static com.nvidia.cuvs.lucene.TestUtils.generateDataset; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java b/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java index 4c3c2809..5140782f 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.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; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java b/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java index 05e37f84..c1c5066b 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; From 7f99c978cc93f5ef18f8b8480c679f5936b5b312 Mon Sep 17 00:00:00 2001 From: James Xia Date: Thu, 25 Jun 2026 18:34:00 -0700 Subject: [PATCH 14/41] Bump cuvs-lucene to 26.08.0-SNAPSHOT --- bench/pom.xml | 4 ++-- build.sh | 2 +- examples/pom.xml | 4 ++-- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/bench/pom.xml b/bench/pom.xml index de461074..9e4e6893 100644 --- a/bench/pom.xml +++ b/bench/pom.xml @@ -11,7 +11,7 @@ com.nvidia.cuvs.lucene.benchmarks cuvs-lucene-benchmarks - 26.02.0 + 26.08.0-SNAPSHOT jar cuvs-lucene-benchmarks @@ -31,7 +31,7 @@ com.nvidia.cuvs.lucene cuvs-lucene - 26.08.0 + 26.08.0-SNAPSHOT commons-io diff --git a/build.sh b/build.sh index 4ec8c087..bf4f8955 100755 --- a/build.sh +++ b/build.sh @@ -8,7 +8,7 @@ set -e -u -o pipefail ARGS="$*" NUMARGS=$# -VERSION="26.08.0" # Note: The version is updated automatically when ci/release/update-version.sh is invoked +VERSION="26.08.0-SNAPSHOT" # Note: The version is updated automatically when ci/release/update-version.sh is invoked GROUP_ID="com.nvidia.cuvs.lucene" function hasArg { diff --git a/examples/pom.xml b/examples/pom.xml index 1bc5cdd4..195d8bdf 100644 --- a/examples/pom.xml +++ b/examples/pom.xml @@ -11,7 +11,7 @@ com.nvidia.cuvs.lucene.examples examples - 26.08.0 + 26.08.0-SNAPSHOT examples @@ -51,7 +51,7 @@ com.nvidia.cuvs.lucene cuvs-lucene - 26.08.0 + 26.08.0-SNAPSHOT From f6a55423fd05b4ecf75237ec03e4c7973a690756 Mon Sep 17 00:00:00 2001 From: James Xia Date: Tue, 30 Jun 2026 15:00:56 -0700 Subject: [PATCH 15/41] Adapt to multi-partition CAGRA Java API refactor cuvs-java moved FilterBitsetHandle/MultiPartitionCagraSearch to the base source set behind CuVSProvider: - GPUKnnFloatVectorQuery: new FilterBitsetHandle(...) -> FilterBitsetHandle.create(...). - FilterCuVSProvider: forward the two new CuVSProvider methods (newFilterBitsetHandle, searchCagraMultiPartition) to the delegate. --- .../cuvs/lucene/FilterCuVSProvider.java | 21 +++++++++++++++++++ .../cuvs/lucene/GPUKnnFloatVectorQuery.java | 2 +- 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java b/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java index 6337a67a..e044b37d 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java +++ b/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java @@ -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,23 @@ public CagraIndex.Builder newCagraIndexBuilder(CuVSResources cuVSResources) return delegate.newCagraIndexBuilder(cuVSResources); } + @Override + public FilterBitsetHandle newFilterBitsetHandle( + long[] combinedLongs, long[] partBitOffsets, long totalBits) { + return delegate.newFilterBitsetHandle(combinedLongs, partBitOffsets, totalBits); + } + + @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 { diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index eb23f996..bbfb5062 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -316,7 +316,7 @@ private FilterBitsetHandle buildFilterHandle( packOrdsToLongs(acceptedOrds, numOrds, combinedLongs, longOffset); } - return new FilterBitsetHandle(combinedLongs, segBitOffsets, totalBits); + return FilterBitsetHandle.create(combinedLongs, segBitOffsets, totalBits); } /** From 91684e84e850f88249ad691f30495fa49863cdab Mon Sep 17 00:00:00 2001 From: James Xia Date: Tue, 30 Jun 2026 17:17:11 -0700 Subject: [PATCH 16/41] Fix duplicate tag in pom.xml The merge of origin/main (c840848) combined two independently-added blocks: the GCS-mirror + Maven Central block from #166 (releases only, snapshots disabled) and the central-snapshots block from the snapshot-build change (snapshots enabled). Two tags is invalid XML, so the POM failed to parse. Consolidate both into a single block, preserving all three repositories: the GCS mirror and Maven Central for releases, and central-snapshots for snapshot resolution. --- pom.xml | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/pom.xml b/pom.xml index c50099b0..0433ecc4 100644 --- a/pom.xml +++ b/pom.xml @@ -80,6 +80,12 @@ false + + central-snapshots + https://central.sonatype.com/repository/maven-snapshots/ + true + false + @@ -145,15 +151,6 @@ - - - central-snapshots - https://central.sonatype.com/repository/maven-snapshots/ - true - false - - - central From bdf69dd2be0a6a3e86add123737c629378adcbac Mon Sep 17 00:00:00 2001 From: James Xia Date: Tue, 30 Jun 2026 17:21:17 -0700 Subject: [PATCH 17/41] Expose CAGRA max_iterations as a configurable search setting max_iterations was already plumbed through cuvs-java (CagraSearchParams.Builder.withMaxIterations) and the C-API, but cuvs-lucene always left it at the default (0 = auto). Thread it through so callers can bound CAGRA search iterations, following the existing threadBlockSize pattern: - GPUKnnFloatVectorQuery: new maxIterations field; the full constructor takes it (between threadBlockSize and searchAlgo). The 6-arg constructor is unchanged and delegates with 0. - GPUPerLeafCuVSKnnCollector: carries maxIterations with a getter. - Both search paths (multi-segment rewrite and per-leaf fallback) pass it into CagraSearchParams.Builder.withMaxIterations. 0 = auto preserves the previous behavior, so existing callers are unaffected. --- .../com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java | 1 + .../com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java | 9 +++++++-- .../nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java | 8 ++++++++ 3 files changed, 16 insertions(+), 2 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java index 99d9cceb..86955beb 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java @@ -444,6 +444,7 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits .withItopkSize(Math.max(collector.getiTopK(), topK)) .withSearchWidth(collector.getSearchWidth()) .withThreadBlockSize(collector.getThreadBlockSize()) + .withMaxIterations(collector.getMaxIterations()) .withAlgo(collector.getSearchAlgo()) .build(); } else { diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index bbfb5062..6d9ef33c 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -69,6 +69,7 @@ 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; /** @@ -84,7 +85,7 @@ public class GPUKnnFloatVectorQuery extends KnnFloatVectorQuery { */ public GPUKnnFloatVectorQuery( String field, float[] target, int k, Query filter, int iTopK, int searchWidth) { - this(field, target, k, filter, iTopK, searchWidth, 0, CagraSearchParams.SearchAlgo.AUTO); + this(field, target, k, filter, iTopK, searchWidth, 0, 0, CagraSearchParams.SearchAlgo.AUTO); } /** @@ -97,6 +98,7 @@ public GPUKnnFloatVectorQuery( * @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( @@ -107,11 +109,13 @@ public GPUKnnFloatVectorQuery( 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; } @@ -166,6 +170,7 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException { .withItopkSize(Math.max(iTopK, k)) .withSearchWidth(searchWidth) .withThreadBlockSize(threadBlockSize) + .withMaxIterations(maxIterations) .withAlgo(searchAlgo) .build(); @@ -234,7 +239,7 @@ protected TopDocs approximateSearch( throws IOException { GPUPerLeafCuVSKnnCollector results = new GPUPerLeafCuVSKnnCollector( - k, visitedLimit, iTopK, searchWidth, threadBlockSize, searchAlgo); + k, visitedLimit, iTopK, searchWidth, threadBlockSize, maxIterations, searchAlgo); context.reader().searchNearestVectors(field, getTargetCopy(), results, acceptDocs); return results.topDocs(); } diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java b/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java index 28443bde..6413c221 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java @@ -17,6 +17,7 @@ class GPUPerLeafCuVSKnnCollector extends TopKnnCollector { private int iTopK; private int searchWidth; private int threadBlockSize; + private int maxIterations; private CagraSearchParams.SearchAlgo searchAlgo; /** @@ -26,6 +27,7 @@ class GPUPerLeafCuVSKnnCollector extends TopKnnCollector { * @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( @@ -34,11 +36,13 @@ public GPUPerLeafCuVSKnnCollector( 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; } @@ -54,6 +58,10 @@ public int getThreadBlockSize() { return threadBlockSize; } + public int getMaxIterations() { + return maxIterations; + } + public CagraSearchParams.SearchAlgo getSearchAlgo() { return searchAlgo; } From 62f6302d7576b85cc0ba165548e3ac5aaed5a903 Mon Sep 17 00:00:00 2001 From: James Xia Date: Tue, 30 Jun 2026 18:30:43 -0700 Subject: [PATCH 18/41] Remove unnecessary formatting updates --- build.sh | 2 +- .../cuvs/lucene/examples/IndexAndSearchonGPUExample.java | 2 +- .../java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java | 3 ++- src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java | 3 ++- .../java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java | 1 + .../java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java | 2 +- .../java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java | 2 +- src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java | 3 ++- .../java/com/nvidia/cuvs/lucene/FilterCuVSServiceProvider.java | 2 +- src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java | 2 +- src/main/java/com/nvidia/cuvs/lucene/GPUFieldWriter.java | 2 +- src/main/java/com/nvidia/cuvs/lucene/GPUIndex.java | 2 +- src/main/java/com/nvidia/cuvs/lucene/GPUSearchParams.java | 3 ++- .../java/com/nvidia/cuvs/lucene/IndexInputInputStream.java | 2 +- .../java/com/nvidia/cuvs/lucene/IndexOutputOutputStream.java | 2 +- .../com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java | 2 +- .../cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java | 2 +- .../cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java | 2 +- .../cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java | 2 +- .../LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java | 2 +- .../LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java | 2 +- .../cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedCodec.java | 2 +- .../LuceneAcceleratedHNSWScalarQuantizedVectorsFormat.java | 2 +- .../LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java | 2 +- src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java | 2 +- src/main/java/com/nvidia/cuvs/lucene/Utils.java | 2 +- .../cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java | 2 +- .../java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java | 3 ++- src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java | 2 +- .../cuvs/lucene/TestCagraToHnswSerializationAndSearch.java | 2 +- ...estCagraToHnswSerializationAndSearchWithFallbackWriter.java | 2 +- .../cuvs/lucene/TestCuVSAcceleratedHNSWDeletedDocuments.java | 2 +- .../com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWGaps.java | 2 +- .../java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java | 2 +- src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java | 2 +- .../nvidia/cuvs/lucene/TestCuVSRandomizedHNSWVectorSearch.java | 2 +- .../com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java | 2 +- .../java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java | 2 +- src/test/java/com/nvidia/cuvs/lucene/TestGPUSearchParams.java | 3 ++- .../com/nvidia/cuvs/lucene/TestIndexOutputOutputStream.java | 2 +- .../cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java | 2 +- src/test/java/com/nvidia/cuvs/lucene/TestMerge.java | 2 +- .../com/nvidia/cuvs/lucene/TestMultithreadedCuVSGPUSearch.java | 3 ++- .../com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java | 2 +- src/test/java/com/nvidia/cuvs/lucene/TestUtils.java | 2 +- 45 files changed, 52 insertions(+), 44 deletions(-) diff --git a/build.sh b/build.sh index bf4f8955..3c86351c 100755 --- a/build.sh +++ b/build.sh @@ -1,6 +1,6 @@ #!/bin/bash -# 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 set -e -u -o pipefail diff --git a/examples/src/main/java/com/nvidia/cuvs/lucene/examples/IndexAndSearchonGPUExample.java b/examples/src/main/java/com/nvidia/cuvs/lucene/examples/IndexAndSearchonGPUExample.java index fbf31fe2..dd920001 100644 --- a/examples/src/main/java/com/nvidia/cuvs/lucene/examples/IndexAndSearchonGPUExample.java +++ b/examples/src/main/java/com/nvidia/cuvs/lucene/examples/IndexAndSearchonGPUExample.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.examples; diff --git a/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java b/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java index 8a69da71..dffb5004 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java +++ b/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java @@ -1,7 +1,8 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ + package com.nvidia.cuvs.lucene; import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; diff --git a/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index c0964d93..2a8ab02f 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -1,7 +1,8 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ + package com.nvidia.cuvs.lucene; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.getCuVSResourcesInstance; diff --git a/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java b/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java index 3a881eb7..20fed04b 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java @@ -2,6 +2,7 @@ * 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.CagraIndexParams; diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java index 9eb5317b..2de761a7 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java index 757ca057..e8fb304e 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java b/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java index 1bdbf998..627d931d 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java @@ -1,7 +1,8 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ + package com.nvidia.cuvs.lucene; import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.quantizeFloatVectorsToBinary; diff --git a/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSServiceProvider.java b/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSServiceProvider.java index 69f5cac6..83cebafa 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSServiceProvider.java +++ b/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSServiceProvider.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java b/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java index 62203ad9..b0b6286d 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUFieldWriter.java b/src/main/java/com/nvidia/cuvs/lucene/GPUFieldWriter.java index d159c5c1..45255da4 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUFieldWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUFieldWriter.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUIndex.java b/src/main/java/com/nvidia/cuvs/lucene/GPUIndex.java index 5cf047bc..76cd916b 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUIndex.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUIndex.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUSearchParams.java b/src/main/java/com/nvidia/cuvs/lucene/GPUSearchParams.java index f8458cf2..cc445268 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUSearchParams.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUSearchParams.java @@ -1,7 +1,8 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ + package com.nvidia.cuvs.lucene; import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; diff --git a/src/main/java/com/nvidia/cuvs/lucene/IndexInputInputStream.java b/src/main/java/com/nvidia/cuvs/lucene/IndexInputInputStream.java index 1effe159..3ab14122 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/IndexInputInputStream.java +++ b/src/main/java/com/nvidia/cuvs/lucene/IndexInputInputStream.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/IndexOutputOutputStream.java b/src/main/java/com/nvidia/cuvs/lucene/IndexOutputOutputStream.java index 82e06737..b5fc120b 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/IndexOutputOutputStream.java +++ b/src/main/java/com/nvidia/cuvs/lucene/IndexOutputOutputStream.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java index e9f4d6fe..b4c5a33d 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java index 61c9b41c..3c42707c 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index 0e558669..35209ce5 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java index 0b1653bc..f2c1aa37 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java index d1810bc5..0f8d9602 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java index e00b2d07..d8a98cc2 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedCodec.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedCodec.java index 0705ed0a..0c7736a0 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedCodec.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsFormat.java index 8d599a54..534390b8 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsFormat.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java index a542cf89..21b4be3f 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java index 9caa2a6a..7635e323 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/main/java/com/nvidia/cuvs/lucene/Utils.java b/src/main/java/com/nvidia/cuvs/lucene/Utils.java index e4a20d2b..c55aa7c3 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Utils.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Utils.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java index cd4aaa6e..d164af55 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java index 9dc110f8..2273af8c 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java @@ -1,7 +1,8 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ + package com.nvidia.cuvs.lucene; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_BEAM_WIDTH; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java b/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java index 0de813ff..2de6e660 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java index 67e6b0ad..96be6c01 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java index 8999d8e8..04ad20a0 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearchWithFallbackWriter.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWDeletedDocuments.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWDeletedDocuments.java index 10d3ceee..d3dc866e 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWDeletedDocuments.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWDeletedDocuments.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWGaps.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWGaps.java index fa0936e6..b0dae38e 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWGaps.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSAcceleratedHNSWGaps.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java index 7303d04d..de868c7c 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java index 31958ef0..a3e2f59a 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedHNSWVectorSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedHNSWVectorSearch.java index e630a0c3..c2e2f2f6 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedHNSWVectorSearch.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedHNSWVectorSearch.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java index 85720c30..3e77e2ed 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java index 2b960bf3..1c75cb23 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestGPUSearchParams.java b/src/test/java/com/nvidia/cuvs/lucene/TestGPUSearchParams.java index 41f48005..483b689c 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestGPUSearchParams.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestGPUSearchParams.java @@ -1,7 +1,8 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ + package com.nvidia.cuvs.lucene; import static com.nvidia.cuvs.lucene.GPUSearchParams.DEFAULT_CAGRA_GRAPH_BUILD_ALGO; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestIndexOutputOutputStream.java b/src/test/java/com/nvidia/cuvs/lucene/TestIndexOutputOutputStream.java index b1eba881..04b22b3d 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestIndexOutputOutputStream.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestIndexOutputOutputStream.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java b/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java index 9bb14022..2428fc50 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java b/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java index a42f55b3..bcc29178 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestMerge.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestMultithreadedCuVSGPUSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestMultithreadedCuVSGPUSearch.java index ba03861f..25dc13b8 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestMultithreadedCuVSGPUSearch.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestMultithreadedCuVSGPUSearch.java @@ -1,7 +1,8 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ + package com.nvidia.cuvs.lucene; import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java b/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java index 5140782f..4c3c2809 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java b/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java index c1c5066b..05e37f84 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; From 480ae2898c1c684b13c1bba01efc17c898d733b0 Mon Sep 17 00:00:00 2001 From: James Xia Date: Fri, 3 Jul 2026 14:27:12 -0700 Subject: [PATCH 19/41] Adapt multi-partition filter to the offset-free cuVS API MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit cuVS now recomputes per-partition bit offsets from the index sizes, so the combined bitset is passed as a plain BITSET filter with no offsets or total- bit count. Update the consumers of `FilterBitsetHandle` accordingly: - `GPUKnnFloatVectorQuery.buildFilterHandle` calls the 1-arg `FilterBitsetHandle.create(combinedLongs)`. The per-segment offsets are still computed locally to pack each segment's slice into the combined bitset (64-bit word-aligned, matching cuVS's recomputation) — they're just no longer transported. - `FilterCuVSProvider.newFilterBitsetHandle` delegate override updated to the 1-arg signature. --- src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java | 5 ++--- .../java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java | 5 ++++- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java b/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java index e044b37d..b237d4bc 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java +++ b/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java @@ -59,9 +59,8 @@ public CagraIndex.Builder newCagraIndexBuilder(CuVSResources cuVSResources) } @Override - public FilterBitsetHandle newFilterBitsetHandle( - long[] combinedLongs, long[] partBitOffsets, long totalBits) { - return delegate.newFilterBitsetHandle(combinedLongs, partBitOffsets, totalBits); + public FilterBitsetHandle newFilterBitsetHandle(long[] combinedLongs) { + return delegate.newFilterBitsetHandle(combinedLongs); } @Override diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index 6d9ef33c..ee604449 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -321,7 +321,10 @@ private FilterBitsetHandle buildFilterHandle( packOrdsToLongs(acceptedOrds, numOrds, combinedLongs, longOffset); } - return FilterBitsetHandle.create(combinedLongs, segBitOffsets, totalBits); + // 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); } /** From 672e03fe72df37fd18db410d991cb4add2a63b5b Mon Sep 17 00:00:00 2001 From: James Xia Date: Thu, 16 Jul 2026 15:08:19 -0700 Subject: [PATCH 20/41] Stash FloatVectorValues per segment in multi-segment result mapping The result-mapping loop in GPUKnnFloatVectorQuery called ordToDoc() once per result, each call allocating a fresh FloatVectorValues via the reader. Since the number of results can be far larger than the number of segments, cache one FloatVectorValues per segment and reuse it across all results landing in that segment, reducing allocations from O(results) to O(distinct segments touched). Removes the now-unused CuVS2510GPUVectorsReader.ordToDoc() helper. --- .../cuvs/lucene/CuVS2510GPUVectorsReader.java | 15 --------------- .../cuvs/lucene/GPUKnnFloatVectorQuery.java | 11 +++++++++-- 2 files changed, 9 insertions(+), 17 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java index 86955beb..d0854ed5 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java @@ -603,21 +603,6 @@ private static void checkVersion(int versionMeta, int versionVectorData, IndexIn } } - /** - * Maps a local segment ordinal to a Lucene doc ID within this segment. - * - *

Used by {@link GPUKnnFloatVectorQuery} after a multi-segment GPU search to convert - * select_k result ordinals to doc IDs before adding {@code docBase}. - * - * @param field the vector field name - * @param ordinal the local ordinal returned by CAGRA - * @return the Lucene doc ID within this segment - * @throws IOException if the vector values cannot be read - */ - public int ordToDoc(String field, int ordinal) throws IOException { - return flatVectorsReader.getFloatVectorValues(field).ordToDoc(ordinal); - } - /** * Returns the {@link CagraIndex} for the given field, or {@code null} if unavailable * (e.g., during a merge or when the field is missing). diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index ee604449..e6ef76c4 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -201,15 +201,22 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException { } // 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 localDoc = gpuReaders.get(segIdx).ordToDoc(field, ordinal); - int globalDoc = ctx.docBase + localDoc; + int globalDoc = ctx.docBase + fvv.ordToDoc(ordinal); float score = 1.0f / (1.0f + dist); scoreDocs[j] = new ScoreDoc(globalDoc, score); } From ce0237e2ec0229976fd7c39786f842081d64fde8 Mon Sep 17 00:00:00 2001 From: James Xia Date: Thu, 16 Jul 2026 21:38:30 -0700 Subject: [PATCH 21/41] Fix NPE when advance() precedes nextDoc() in docAndScoreQuery iterator --- src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java | 1 + 1 file changed, 1 insertion(+) diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index e6ef76c4..3934e603 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -481,6 +481,7 @@ public int nextDoc() { @Override public int advance(int target) { + if (pos < 0) pos = 0; while (pos < n && sortedDocs[pos] < target) pos++; return docID(); } From a8fa2d93f510aef559e2ff0a867b98c73f44ac7d Mon Sep 17 00:00:00 2001 From: James Xia Date: Fri, 17 Jul 2026 16:16:58 -0700 Subject: [PATCH 22/41] Reference-count and compute-once the filter bitset cache FilterBitsetCache had two concurrency problems. get() and put() were synchronized independently, so concurrent misses on the same key both built a FilterBitsetHandle and the second put() overwrote the first without closing it, leaking device memory. And eviction closed a handle that another thread's in-flight search could still be using, freeing its device allocation underneath the search. Rework the cache around FilterBitsetHandle's new reference counting. Values are now per-key CompletableFutures, so a key is built exactly once while builds for different keys still run in parallel. A cached handle carries one cache-owned reference, released via close() on eviction; acquire() takes an additional reference for the caller, which GPUKnnFloatVectorQuery releases in a finally after the search. Because the allocation is freed only when the last reference drops, eviction can no longer free a handle still in use, and a replaced entry cannot leak. Add TestFilterBitsetCache (compute-once, eviction refcount, failed-build retry, concurrent balance) and TestMultiSegmentGPUFilterConcurrency (end-to-end filtered multi-segment GPU search under cache eviction). --- .../nvidia/cuvs/lucene/FilterBitsetCache.java | 103 +++++++- .../cuvs/lucene/GPUKnnFloatVectorQuery.java | 30 ++- .../cuvs/lucene/TestFilterBitsetCache.java | 236 ++++++++++++++++++ .../TestMultiSegmentGPUFilterConcurrency.java | 192 ++++++++++++++ 4 files changed, 536 insertions(+), 25 deletions(-) create mode 100644 src/test/java/com/nvidia/cuvs/lucene/TestFilterBitsetCache.java create mode 100644 src/test/java/com/nvidia/cuvs/lucene/TestMultiSegmentGPUFilterConcurrency.java diff --git a/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java b/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java index 3b269798..7e2a3528 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java +++ b/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java @@ -5,23 +5,45 @@ 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 (per-thread LRU). Entries are evicted in LRU order and closed - * on eviction to free the host arrays and signal the device cache. + *

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 { - private static final int MAX_HOST_ENTRIES = 16; + 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 @@ -38,12 +60,13 @@ public int hashCode() { } } - private static final LinkedHashMap CACHE = + private static final LinkedHashMap> CACHE = new LinkedHashMap<>(MAX_HOST_ENTRIES + 2, 0.75f, /* access-order= */ true) { @Override - protected boolean removeEldestEntry(Map.Entry eldest) { + protected boolean removeEldestEntry( + Map.Entry> eldest) { if (size() > MAX_HOST_ENTRIES) { - eldest.getValue().close(); + releaseCacheRef(eldest.getValue()); return true; } return false; @@ -52,13 +75,67 @@ protected boolean removeEldestEntry(Map.Entry segReaderKeys, String field) { - return CACHE.get(new FilterCacheKey(filter, segReaderKeys, field)); + /** + * 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. + } } - static synchronized void put( - Query filter, List segReaderKeys, String field, FilterBitsetHandle handle) { - CACHE.put(new FilterCacheKey(filter, segReaderKeys, field), handle); + /** 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/GPUKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index 3934e603..1c576319 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -154,16 +154,16 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException { break; } } - FilterBitsetHandle filterHandle = null; - if (hasExplicitFilter || hasDeletes) { - filterHandle = buildOrGetCachedFilterHandle(indexSearcher, leaves, gpuReaders); - } - // 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() @@ -230,6 +230,10 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException { if (t instanceof IOException) throw (IOException) t; if (t instanceof RuntimeException) throw (RuntimeException) t; throw new RuntimeException("Multi-segment GPU search failed", t); + } 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(); } } @@ -261,6 +265,9 @@ protected TopDocs approximateSearch( * *

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, @@ -272,18 +279,17 @@ private FilterBitsetHandle buildOrGetCachedFilterHandle( for (LeafReaderContext ctx : leaves) { var helper = ctx.reader().getReaderCacheHelper(); if (helper == null) { - // This reader can't be cached; build without caching. + // This reader can't be cached; build an uncached handle owned outright by the caller. return buildFilterHandle(indexSearcher, leaves, gpuReaders); } segReaderKeys.add(helper.getKey()); } - FilterBitsetHandle cached = FilterBitsetCache.get(filter, segReaderKeys, field); - if (cached != null) return cached; - - FilterBitsetHandle handle = buildFilterHandle(indexSearcher, leaves, gpuReaders); - FilterBitsetCache.put(filter, segReaderKeys, field, handle); - return handle; + return FilterBitsetCache.acquire( + filter, + segReaderKeys, + field, + () -> buildFilterHandle(indexSearcher, leaves, gpuReaders)); } /** 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..6b4f0784 --- /dev/null +++ b/src/test/java/com/nvidia/cuvs/lucene/TestMultiSegmentGPUFilterConcurrency.java @@ -0,0 +1,192 @@ +/* + * 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.apache.lucene.tests.util.TestUtil; +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 { + + private static final Codec codec = + TestUtil.alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); + 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 IOException { + assumeTrue("cuVS not supported", isSupported()); + + 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; + } +} From 863d38d5ed67223cff13aae95797c95c2937e1e6 Mon Sep 17 00:00:00 2001 From: James Xia Date: Sat, 18 Jul 2026 08:48:15 -0700 Subject: [PATCH 23/41] Gate per-segment CAGRA decode on the distance sentinel The per-segment search decoder skipped empty result slots by checking ord < 0, which only catches the 0xFFFFFFFF sentinel index. Single-CTA emits 0x7FFFFFFF for empty slots, which is positive and slipped through as a bogus ordinal under high filter selectivity. Gate on the FLT_MAX sentinel distance instead, which every search algorithm honors and which is robust regardless of how the index is mapped downstream. --- .../com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java index d0854ed5..837a57c8 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java @@ -512,7 +512,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); From c2b4fc0e91a29ce1753933470e47d09350be1b8a Mon Sep 17 00:00:00 2001 From: James Xia Date: Sat, 18 Jul 2026 08:49:08 -0700 Subject: [PATCH 24/41] Exercise the multi-partition GPU path in the filter concurrency test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The test built its index with TestUtil.alwaysKnnVectorsFormat, whose wrapped reader is not a CuVS2510GPUVectorsReader, so unwrapGpuReader returned null and every query fell back to per-segment search — the multi-partition path was never tested. Use CuVS2510GPUSearchCodec, which exposes the CuVS reader directly, so GPUKnnFloatVectorQuery takes the multi-partition path the test is meant to cover. --- .../lucene/TestMultiSegmentGPUFilterConcurrency.java | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestMultiSegmentGPUFilterConcurrency.java b/src/test/java/com/nvidia/cuvs/lucene/TestMultiSegmentGPUFilterConcurrency.java index 6b4f0784..fb8e92c6 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestMultiSegmentGPUFilterConcurrency.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestMultiSegmentGPUFilterConcurrency.java @@ -40,7 +40,6 @@ 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.tests.util.TestUtil; import org.junit.AfterClass; import org.junit.BeforeClass; import org.junit.Test; @@ -55,8 +54,10 @@ @SuppressSysoutChecks(bugUrl = "") public class TestMultiSegmentGPUFilterConcurrency extends LuceneTestCase { - private static final Codec codec = - TestUtil.alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); + // 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. @@ -71,8 +72,9 @@ public class TestMultiSegmentGPUFilterConcurrency extends LuceneTestCase { private static List> acceptedDocsPerFilter; @BeforeClass - public static void beforeClass() throws IOException { + public static void beforeClass() throws Exception { assumeTrue("cuVS not supported", isSupported()); + codec = new CuVS2510GPUSearchCodec(); int datasetSize = 2000; int dimensions = 128; From 4427afb552d452565fa950001d4ec717593dd59e Mon Sep 17 00:00:00 2001 From: James Xia Date: Sat, 18 Jul 2026 09:35:14 -0700 Subject: [PATCH 25/41] Fix per-segment prefilter dropping trailing ordinals under selective filters MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The per-segment search passed BitSet.length() as the prefilter's numDocs (the "total dataset vectors" cuVS uses to size the filter). BitSet.length() is the highest set bit plus one, so under a selective filter — where the trailing ordinals are rejected — it is smaller than the vector count. cuVS then treats ordinals beyond that length as outside the filter and default-accepts them, returning filtered-out vectors. Under a permissive filter the last ordinal is usually accepted, so length() matched the vector count and the bug was masked. Pass the vector count (acceptedOrds.length()) instead, so the prefilter covers every ordinal. Add TestPerSegmentGPUFilterSearch, which forces the per-segment path (alwaysKnnVectorsFormat wraps the reader) with selective filters and asserts every hit satisfies its filter. --- .../cuvs/lucene/CuVS2510GPUVectorsReader.java | 5 +- .../lucene/TestPerSegmentGPUFilterSearch.java | 100 ++++++++++++++++++ 2 files changed, 104 insertions(+), 1 deletion(-) create mode 100644 src/test/java/com/nvidia/cuvs/lucene/TestPerSegmentGPUFilterSearch.java diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java index 837a57c8..8c29fff3 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java @@ -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 { 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..221ca2d9 --- /dev/null +++ b/src/test/java/com/nvidia/cuvs/lucene/TestPerSegmentGPUFilterSearch.java @@ -0,0 +1,100 @@ +/* + * 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.util.HashSet; +import java.util.Set; +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.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.apache.lucene.tests.util.TestUtil; +import org.junit.Test; + +/** + * Exercises the per-segment (single-index) GPU search fallback with a selective prefilter. + * + *

{@code TestUtil.alwaysKnnVectorsFormat} wraps the segment reader, so {@code + * GPUKnnFloatVectorQuery} cannot unwrap it to a {@link CuVS2510GPUVectorsReader} and falls back to + * per-segment search ({@link CuVS2510GPUVectorsReader#search}) rather than the multi-partition path. + * A highly selective filter rejects the trailing ordinals of the segment — the regime where the + * per-segment prefilter previously mis-sized itself and default-accepted those ordinals, returning + * filtered-out vectors. Every returned hit must satisfy its filter. + */ +@SuppressSysoutChecks(bugUrl = "") +public class TestPerSegmentGPUFilterSearch extends LuceneTestCase { + + private static final Codec codec = + TestUtil.alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); + private static final String VECTOR_FIELD = "vectors"; + private static final String CATEGORY_FIELD = "cat"; + 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(codec); + try (IndexWriter writer = new IndexWriter(directory, config)) { + for (int i = 0; i < datasetSize; i++) { + Document doc = new Document(); + 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); + } + writer.commit(); + } + + try (DirectoryReader reader = DirectoryReader.open(directory)) { + IndexSearcher searcher = new IndexSearcher(reader); + float[][] queries = generateDataset(random(), 32, dimensions); + + // Each category accepts ~1/24 of the docs, so the trailing ordinals are rejected. + for (int c = 0; c < NUM_CATEGORIES; c++) { + Query filter = new TermQuery(new Term(CATEGORY_FIELD, "c" + c)); + Set accepted = new HashSet<>(); + for (ScoreDoc sd : searcher.search(filter, datasetSize).scoreDocs) { + accepted.add(sd.doc); + } + for (float[] q : queries) { + GPUKnnFloatVectorQuery query = + new GPUKnnFloatVectorQuery(VECTOR_FIELD, q, topK, filter, topK, 1); + for (ScoreDoc hit : searcher.search(query, topK).scoreDocs) { + assertTrue( + "per-segment search returned doc " + hit.doc + " not accepted by filter c" + c, + accepted.contains(hit.doc)); + } + } + } + } + } + } +} From e32040f1e6b0edce06e77ea6114c53c00409d648 Mon Sep 17 00:00:00 2001 From: James Xia Date: Sat, 18 Jul 2026 14:47:08 -0700 Subject: [PATCH 26/41] Take the multi-partition path for PerField GPU segments, gated on graph degree unwrapGpuReader only recognized a directly-wrapped CuVS reader, so a real index -- where PerFieldKnnVectorsFormat.FieldsReader wraps the per-field reader -- never matched and every query silently fell back to per-segment search, leaving the multi-partition path (and its filter cache) unused. Unwrap the PerField reader so those segments take the multi-partition path. That path requires every partition to share one built CAGRA graph degree; cuvsCagraSearchMultiPartition rejects a mismatch, and a small segment can have its degree truncated at build time. Gate the multi-partition path on all segments reporting the same CagraIndex.getGraphDegree(), falling back to the per-segment path (which handles heterogeneous segments, liveDocs, and gaps natively) otherwise. Rework TestPerSegmentGPUFilterSearch to drive leaf.searchNearestVectors directly, since the unwrap change now routes its former setup through the multi-partition path; the reworked test exercises the per-segment prefilter regardless of higher-level routing. --- .../cuvs/lucene/GPUKnnFloatVectorQuery.java | 35 ++++++++++-- .../lucene/TestPerSegmentGPUFilterSearch.java | 57 ++++++++----------- 2 files changed, 53 insertions(+), 39 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index 1c576319..b1b2c83d 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -20,6 +20,7 @@ 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; @@ -60,7 +61,9 @@ * 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 or missing CAGRA index for the field on any segment). + * 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 */ @@ -131,12 +134,26 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException { return new MatchNoDocsQuery(); } - // Collect a CuVS2510GPUVectorsReader for every segment; fall back if any segment - // lacks one or has no CAGRA index for this field. + // 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); - if (gpuReader == null || gpuReader.getCagraIndexForField(field) == null) { + 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); @@ -417,10 +434,16 @@ private static void packOrdsToLongs(Bits bits, int numOrds, long[] dest, int des * 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) { + 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; } diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestPerSegmentGPUFilterSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestPerSegmentGPUFilterSearch.java index 221ca2d9..df2ba211 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestPerSegmentGPUFilterSearch.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestPerSegmentGPUFilterSearch.java @@ -9,45 +9,36 @@ import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; -import java.util.HashSet; -import java.util.Set; -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.Term; -import org.apache.lucene.search.IndexSearcher; -import org.apache.lucene.search.Query; +import org.apache.lucene.index.LeafReader; import org.apache.lucene.search.ScoreDoc; -import org.apache.lucene.search.TermQuery; +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.tests.util.TestUtil; +import org.apache.lucene.util.FixedBitSet; import org.junit.Test; /** - * Exercises the per-segment (single-index) GPU search fallback with a selective prefilter. + * 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). * - *

{@code TestUtil.alwaysKnnVectorsFormat} wraps the segment reader, so {@code - * GPUKnnFloatVectorQuery} cannot unwrap it to a {@link CuVS2510GPUVectorsReader} and falls back to - * per-segment search ({@link CuVS2510GPUVectorsReader#search}) rather than the multi-partition path. - * A highly selective filter rejects the trailing ordinals of the segment — the regime where the - * per-segment prefilter previously mis-sized itself and default-accepted those ordinals, returning - * filtered-out vectors. Every returned hit must satisfy its filter. + *

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 Codec codec = - TestUtil.alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat()); private static final String VECTOR_FIELD = "vectors"; - private static final String CATEGORY_FIELD = "cat"; private static final int NUM_CATEGORIES = 24; @Test @@ -62,11 +53,10 @@ public void testSelectiveFilterDoesNotLeakTrailingOrdinals() throws Exception { try (Directory directory = newDirectory(new ByteBuffersDirectory())) { float[][] dataset = generateDataset(random(), datasetSize, dimensions); - IndexWriterConfig config = new IndexWriterConfig().setCodec(codec); + 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 StringField(CATEGORY_FIELD, "c" + (i % NUM_CATEGORIES), Field.Store.NO)); doc.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); writer.addDocument(doc); } @@ -74,23 +64,24 @@ public void testSelectiveFilterDoesNotLeakTrailingOrdinals() throws Exception { } try (DirectoryReader reader = DirectoryReader.open(directory)) { - IndexSearcher searcher = new IndexSearcher(reader); + assertEquals("expected a single segment", 1, reader.leaves().size()); + LeafReader leaf = reader.leaves().get(0).reader(); float[][] queries = generateDataset(random(), 32, dimensions); - // Each category accepts ~1/24 of the docs, so the trailing ordinals are rejected. + // 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++) { - Query filter = new TermQuery(new Term(CATEGORY_FIELD, "c" + c)); - Set accepted = new HashSet<>(); - for (ScoreDoc sd : searcher.search(filter, datasetSize).scoreDocs) { - accepted.add(sd.doc); + FixedBitSet acceptDocs = new FixedBitSet(datasetSize); + for (int i = c; i < datasetSize; i += NUM_CATEGORIES) { + acceptDocs.set(i); } for (float[] q : queries) { - GPUKnnFloatVectorQuery query = - new GPUKnnFloatVectorQuery(VECTOR_FIELD, q, topK, filter, topK, 1); - for (ScoreDoc hit : searcher.search(query, topK).scoreDocs) { + 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 + " not accepted by filter c" + c, - accepted.contains(hit.doc)); + "per-segment search returned doc " + hit.doc + " outside filter category " + c, + acceptDocs.get(hit.doc)); } } } From 8751542b6052ac8391232a915502b14c86ac1035 Mon Sep 17 00:00:00 2001 From: James Xia Date: Sat, 18 Jul 2026 15:51:29 -0700 Subject: [PATCH 27/41] Include boost in GPU search score explanation Weight.explain returned the raw CAGRA score while the scorer multiplies it by the query boost, so the explained value disagreed with the actual score whenever boost != 1.0 (and would fail Lucene's explain-consistency checks). Report sd.score * boost, with the raw score and boost as detail factors. --- .../java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index b1b2c83d..8e1dc2f3 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -557,7 +557,11 @@ public boolean isCacheable(LeafReaderContext ctx) { public Explanation explain(LeafReaderContext ctx, int doc) { for (ScoreDoc sd : scoreDocs) { if (sd.doc == ctx.docBase + doc) { - return Explanation.match(sd.score, "GPU multi-segment CAGRA search"); + 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"); From 21e28f3b6e2176b10824bca7577281d8e41439dd Mon Sep 17 00:00:00 2001 From: James Xia Date: Sat, 18 Jul 2026 16:02:56 -0700 Subject: [PATCH 28/41] Return a real max score to enable score-based pruning getMaxScore returned Float.MAX_VALUE, which disables TOP_SCORES/MaxScore pruning. Precompute the segment's maximum (boosted) score once during scorer setup and return it, ignoring upTo. The segment-wide max is a valid upper bound for any upTo window, so pruning stays correct while becoming effective. --- .../java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index 8e1dc2f3..44031534 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -469,15 +469,18 @@ public ScorerSupplier scorerSupplier(LeafReaderContext ctx) { 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])); @@ -524,7 +527,8 @@ public long cost() { @Override public float getMaxScore(int upTo) { - return Float.MAX_VALUE; + // The precomputed segment max is a valid upper bound for any upTo. + return maxScore; } @Override From d389dfcea27215287a4c64a6282a40774b571a7c Mon Sep 17 00:00:00 2001 From: James Xia Date: Sat, 18 Jul 2026 16:17:37 -0700 Subject: [PATCH 29/41] Address review comments - Propagate Errors (and IOException/RuntimeException) unchanged from the rewrite catch via Utils.handleThrowable instead of wrapping everything in RuntimeException, matching the codebase idiom. - Widen numOrds to long before adding 63 in the bit-offset calc to avoid a potential int overflow for very large segments. - Replace the hand-rolled all-false Bits with Bits.MatchNoBits for the no-scorer (empty filter) case. --- .../cuvs/lucene/GPUKnnFloatVectorQuery.java | 21 +++++-------------- 1 file changed, 5 insertions(+), 16 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index 44031534..0cd1802d 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -244,9 +244,8 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException { } } catch (Throwable t) { - if (t instanceof IOException) throw (IOException) t; - if (t instanceof RuntimeException) throw (RuntimeException) t; - throw new RuntimeException("Multi-segment GPU search failed", 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. @@ -334,7 +333,7 @@ private FilterBitsetHandle buildFilterHandle( for (int i = 0; i < numSegments; i++) { segBitOffsets[i] = totalBits; int numOrds = gpuReaders.get(i).getFloatVectorValues(field).size(); - totalBits += ((long) (numOrds + 63) / 64) * 64; + totalBits += (((long) numOrds + 63) / 64) * 64; } long[] combinedLongs = new long[(int) (totalBits / 64)]; @@ -365,18 +364,8 @@ private static Bits evalFilter(Weight filterWeight, LeafReaderContext ctx, Bits throws IOException { ScorerSupplier scorerSupplier = filterWeight.scorerSupplier(ctx); if (scorerSupplier == null) { - int maxDoc = ctx.reader().maxDoc(); - return new Bits() { - @Override - public boolean get(int i) { - return false; - } - - @Override - public int length() { - return maxDoc; - } - }; + // No scorer: the filter matches no documents in this segment. + return new Bits.MatchNoBits(ctx.reader().maxDoc()); } int maxDoc = ctx.reader().maxDoc(); From 7880820765884861ae31ec192a833e9f9b782e6c Mon Sep 17 00:00:00 2001 From: James Xia Date: Tue, 21 Jul 2026 18:09:49 -0700 Subject: [PATCH 30/41] cuvs-lucene: per-segment filter bitsets and per-segment cache keying MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adapt to the new cuvs-java per-partition multi-partition search API and fix the cache-staleness concern by keying per segment instead of per segment-set: - GPUKnnFloatVectorQuery builds one FilterBitsetHandle per segment (filter ∩ that segment's liveDocs, packed into its own long[]), or null for a segment with no explicit filter and no deletes, and passes the List to MultiPartitionCagraSearch.search. Every non-null handle is reference-released once per query. - FilterBitsetCache keys on (filter, single-segment reader key, field) instead of the whole segment-key list, so an index update only invalidates the changed segments' entries rather than the entire cache. - FilterCuVSProvider passthrough updated to the new SPI signature. - TestFilterBitsetCache updated for the per-segment acquire key. --- .../nvidia/cuvs/lucene/FilterBitsetCache.java | 17 +-- .../cuvs/lucene/FilterCuVSProvider.java | 4 +- .../cuvs/lucene/GPUKnnFloatVectorQuery.java | 143 +++++++++--------- .../cuvs/lucene/TestFilterBitsetCache.java | 14 +- 4 files changed, 87 insertions(+), 91 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java b/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java index 7e2a3528..2882d7cf 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java +++ b/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java @@ -7,7 +7,6 @@ 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; @@ -15,8 +14,9 @@ import org.apache.lucene.search.Query; /** - * Shared LRU cache mapping (filter Query, per-segment reader keys, field) → {@link - * FilterBitsetHandle}. + * Shared LRU cache mapping (filter Query, single-segment reader key, field) → {@link + * FilterBitsetHandle}. One entry per segment, so an index update only invalidates the changed + * segments' entries rather than the whole cache. * *

Host-side cache holding packed bitset arrays; the device-side upload is managed inside {@link * FilterBitsetHandle} itself. Entries are evicted in LRU order. @@ -45,18 +45,18 @@ interface FilterBuilder { FilterBitsetHandle build() throws IOException; } - private record FilterCacheKey(Query filter, List segReaderKeys, String field) { + private record FilterCacheKey(Query filter, Object segReaderKey, 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(segReaderKey, other.segReaderKey) && Objects.equals(field, other.field); } @Override public int hashCode() { - return Objects.hash(filter, segReaderKeys, field); + return Objects.hash(filter, segReaderKey, field); } } @@ -81,9 +81,8 @@ private FilterBitsetCache() {} * 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); + Query filter, Object segReaderKey, String field, FilterBuilder builder) throws IOException { + FilterCacheKey key = new FilterCacheKey(filter, segReaderKey, field); while (true) { CompletableFuture future; boolean owner = false; diff --git a/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java b/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java index b237d4bc..c4449b79 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java +++ b/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java @@ -69,9 +69,9 @@ public MultiPartitionSearchResults searchCagraMultiPartition( List indices, CagraQuery query, int k, - FilterBitsetHandle filter) + List filters) throws Throwable { - return delegate.searchCagraMultiPartition(resources, indices, query, k, filter); + return delegate.searchCagraMultiPartition(resources, indices, query, k, filters); } @Override diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index 0cd1802d..5099ba7b 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -55,10 +55,10 @@ * 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. + * the acceptance mask (filter ∩ liveDocs) is packed into one {@link FilterBitsetHandle} per segment + * and passed as that partition's filter. The host-side packed arrays are cached per unique + * (filter, single-segment reader key, 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 @@ -159,10 +159,10 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException { 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. + // Build one filter handle per segment encoding (filter ∩ that segment's liveDocs) whenever any + // filtering is required — either an explicit Lucene filter, or live-document deletes in at least + // one segment. Each segment's handle becomes that partition's filter; a segment with neither an + // explicit filter nor deletes gets a null entry (unfiltered for that partition). boolean hasExplicitFilter = (filter != null); boolean hasDeletes = false; for (LeafReaderContext ctx : leaves) { @@ -175,10 +175,10 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException { CuVSResources resources = getCuVSResourcesInstance(); List cagraIndices = new ArrayList<>(leaves.size()); - FilterBitsetHandle filterHandle = null; + List filterHandles = null; try { if (hasExplicitFilter || hasDeletes) { - filterHandle = buildOrGetCachedFilterHandle(indexSearcher, leaves, gpuReaders); + filterHandles = buildPerSegmentFilterHandles(indexSearcher, leaves, gpuReaders); } float[] target = getTargetCopy(); @@ -211,7 +211,7 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException { .build(); MultiPartitionSearchResults results = - MultiPartitionCagraSearch.search(resources, cagraIndices, cagraQuery, k, filterHandle); + MultiPartitionCagraSearch.search(resources, cagraIndices, cagraQuery, k, filterHandles); if (results.count() == 0) { return new MatchNoDocsQuery(); @@ -247,9 +247,14 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException { 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(); + // Release this query's reference on each per-segment handle. 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 (filterHandles != null) { + for (FilterBitsetHandle handle : filterHandles) { + if (handle != null) handle.decRef(); + } + } } } @@ -276,45 +281,19 @@ protected TopDocs approximateSearch( // ------------------------------------------------------------------------- /** - * Returns a {@link FilterBitsetHandle} encoding ({@link #filter} ∩ liveDocs) for every segment, - * pulling from {@link FilterBitsetCache} when the reader state is unchanged. + * Returns one {@link FilterBitsetHandle} per segment (in {@code leaves} order) encoding ({@link + * #filter} ∩ that segment's liveDocs), pulling each from {@link FilterBitsetCache} when the + * reader state is unchanged. A segment with neither an explicit filter nor deletes gets a {@code + * null} entry (unfiltered for that partition). * - *

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 cache key uses the single segment's per-reader key (not just the core key), so liveDocs + * changes — which happen when a reader is reopened after deletes — invalidate that segment's + * cached bitset, and an index update only affects the changed segments' entries. * - *

The returned handle carries a reference the caller must release with {@link + *

Each non-null 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( + private List buildPerSegmentFilterHandles( IndexSearcher indexSearcher, List leaves, List gpuReaders) @@ -326,34 +305,52 @@ private FilterBitsetHandle buildFilterHandle( indexSearcher.createWeight( indexSearcher.rewrite(filter), ScoreMode.COMPLETE_NO_SCORES, 1.0f); } + final Weight sharedFilterWeight = filterWeight; - 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++) { + List handles = new ArrayList<>(leaves.size()); + for (int i = 0; i < leaves.size(); 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); + CuVS2510GPUVectorsReader gpuReader = gpuReaders.get(i); + // No explicit filter and no deletes in this segment: no filter for this partition. + if (filter == null && ctx.reader().getLiveDocs() == null) { + handles.add(null); + continue; + } + var helper = ctx.reader().getReaderCacheHelper(); + if (helper == null) { + // This reader can't be cached; build an uncached handle owned outright by the caller. + handles.add(buildSegmentFilterHandle(sharedFilterWeight, ctx, gpuReader)); + } else { + handles.add( + FilterBitsetCache.acquire( + filter, + helper.getKey(), + field, + () -> buildSegmentFilterHandle(sharedFilterWeight, ctx, gpuReader))); + } } + return handles; + } - // 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} (when non-null) in {@code ctx}, intersects with liveDocs, and + * packs the accepted ordinals of this one segment into a new {@link FilterBitsetHandle}. When + * {@code filterWeight} is {@code null}, the handle encodes liveDocs alone — the path taken for a + * segment with deletes but no explicit Lucene filter. + */ + private FilterBitsetHandle buildSegmentFilterHandle( + Weight filterWeight, LeafReaderContext ctx, CuVS2510GPUVectorsReader gpuReader) + throws IOException { + 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 = gpuReader.getFloatVectorValues(field); + Bits acceptedOrds = fvv.getAcceptOrds(acceptDocs); + int numOrds = fvv.size(); + long[] segLongs = new long[(int) (((long) numOrds + 63) / 64)]; + packOrdsToLongs(acceptedOrds, numOrds, segLongs, 0); + return FilterBitsetHandle.create(segLongs); } /** diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestFilterBitsetCache.java b/src/test/java/com/nvidia/cuvs/lucene/TestFilterBitsetCache.java index 01e8b050..9df434a7 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestFilterBitsetCache.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestFilterBitsetCache.java @@ -70,8 +70,8 @@ public void close() { } } - private static List keys(String s) { - return List.of(s); + private static Object segKey(String s) { + return s; } /** Concurrent misses on the same key build the handle exactly once and each get their own ref. */ @@ -93,7 +93,7 @@ public void computeOnceUnderConcurrentAcquire() throws Exception { start.await(); return FilterBitsetCache.acquire( null, - keys(field), + segKey(field), field, () -> { buildCount.incrementAndGet(); @@ -141,7 +141,7 @@ public void evictionReleasesCacheReferenceExactlyOnce() throws Exception { CountingHandle h = new CountingHandle(); handles.add(h); final String field = "evict-" + i; - FilterBitsetHandle got = FilterBitsetCache.acquire(null, keys(field), field, () -> h); + FilterBitsetHandle got = FilterBitsetCache.acquire(null, segKey(field), field, () -> h); assertSame(h, got); got.decRef(); // caller finished; the cache keeps its reference } @@ -167,7 +167,7 @@ public void failedBuildIsNotCachedAndCanRetry() throws Exception { try { FilterBitsetCache.acquire( null, - keys(field), + segKey(field), field, () -> { buildCount.incrementAndGet(); @@ -183,7 +183,7 @@ public void failedBuildIsNotCachedAndCanRetry() throws Exception { FilterBitsetHandle got = FilterBitsetCache.acquire( null, - keys(field), + segKey(field), field, () -> { buildCount.incrementAndGet(); @@ -218,7 +218,7 @@ public void concurrentAcquireReleaseIsBalanced() throws Exception { 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); + FilterBitsetCache.acquire(null, segKey(field), field, CountingHandle::new); h.decRef(); } return null; From 29aa10d4f5ce484123c583f9e8ad4481917fd2f4 Mon Sep 17 00:00:00 2001 From: James Xia Date: Tue, 21 Jul 2026 18:49:55 -0700 Subject: [PATCH 31/41] cuvs-lucene: evict filter cache entries on segment close Integrate the per-segment filter bitset cache with the segment lifecycle so obsolete entries don't linger until LRU eviction: - FilterBitsetCache.ensureCloseListener registers a one-time close listener per segment reader (tracked so each reader registers once); when the reader's core closes (merge, reopen with new deletes, or index close) its cache entries are evicted and their device allocations released. - invalidateReader(key) evicts and releases every entry for a segment reader key; the close listener calls it. - GPUKnnFloatVectorQuery registers the listener before caching each segment. This mirrors Lucene's LRUQueryCache (core key + closed listener). Adds TestFilterBitsetCache.invalidateReaderEvictsOnlyThatSegment. --- .../nvidia/cuvs/lucene/FilterBitsetCache.java | 42 +++++++++++++++++++ .../cuvs/lucene/GPUKnnFloatVectorQuery.java | 2 + .../cuvs/lucene/TestFilterBitsetCache.java | 42 +++++++++++++++++++ 3 files changed, 86 insertions(+) diff --git a/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java b/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java index 2882d7cf..0a7be8e1 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java +++ b/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java @@ -6,11 +6,16 @@ import com.nvidia.cuvs.FilterBitsetHandle; import java.io.IOException; +import java.util.ArrayList; +import java.util.HashSet; import java.util.LinkedHashMap; +import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.Set; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; +import org.apache.lucene.index.IndexReader; import org.apache.lucene.search.Query; /** @@ -73,6 +78,9 @@ protected boolean removeEldestEntry( } }; + // Segment reader keys with a registered close listener, so each reader registers exactly once. + private static final Set REGISTERED = new HashSet<>(); + private FilterBitsetCache() {} /** @@ -130,6 +138,40 @@ static FilterBitsetHandle acquire( } } + /** + * Registers a one-time close listener on {@code helper} so every cache entry for this segment + * reader is evicted (and its device allocation released) as soon as the reader's core closes — + * e.g. when the segment is merged away or the reader is reopened with new deletes. This ties the + * cache to the segment lifecycle (mirroring Lucene's own query cache) so obsolete entries do not + * linger until LRU eviction. + */ + static void ensureCloseListener(IndexReader.CacheHelper helper) { + Object segReaderKey = helper.getKey(); + synchronized (FilterBitsetCache.class) { + if (!REGISTERED.add(segReaderKey)) return; // a listener is already registered for this reader + } + helper.addClosedListener(FilterBitsetCache::invalidateReader); + } + + /** Evicts and releases every cache entry belonging to the given segment reader key. */ + static void invalidateReader(Object segReaderKey) { + List> toRelease = new ArrayList<>(); + synchronized (FilterBitsetCache.class) { + REGISTERED.remove(segReaderKey); + var it = CACHE.entrySet().iterator(); + while (it.hasNext()) { + Map.Entry> e = it.next(); + if (Objects.equals(e.getKey().segReaderKey(), segReaderKey)) { + toRelease.add(e.getValue()); + it.remove(); + } + } + } + for (CompletableFuture future : toRelease) { + releaseCacheRef(future); + } + } + /** Releases the cache's reference once the (possibly in-flight) build completes. */ private static void releaseCacheRef(CompletableFuture future) { future.whenComplete( diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index 5099ba7b..19bd265a 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -321,6 +321,8 @@ private List buildPerSegmentFilterHandles( // This reader can't be cached; build an uncached handle owned outright by the caller. handles.add(buildSegmentFilterHandle(sharedFilterWeight, ctx, gpuReader)); } else { + // Evict this segment's cache entries when its reader closes (merge/reopen), not just on LRU. + FilterBitsetCache.ensureCloseListener(helper); handles.add( FilterBitsetCache.acquire( filter, diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestFilterBitsetCache.java b/src/test/java/com/nvidia/cuvs/lucene/TestFilterBitsetCache.java index 9df434a7..c4c3c9cb 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestFilterBitsetCache.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestFilterBitsetCache.java @@ -158,6 +158,48 @@ public void evictionReleasesCacheReferenceExactlyOnce() throws Exception { } } + /** invalidateReader evicts and frees exactly the given segment's entries, leaving others cached. */ + @Test + public void invalidateReaderEvictsOnlyThatSegment() throws Exception { + CountingHandle hA = new CountingHandle(); + CountingHandle hB = new CountingHandle(); + Object keyA = segKey("segA"); + Object keyB = segKey("segB"); + AtomicInteger buildA = new AtomicInteger(); + + // Cache one entry per segment, releasing the caller ref so the cache keeps its own. + FilterBitsetCache.acquire( + null, + keyA, + "f", + () -> { + buildA.incrementAndGet(); + return hA; + }) + .decRef(); + FilterBitsetCache.acquire(null, keyB, "f", () -> hB).decRef(); + + // Invalidate segment A only, as its reader's close listener would. + FilterBitsetCache.invalidateReader(keyA); + + assertTrue("segment A handle freed on invalidation", hA.freed.get()); + assertEquals("segment A handle closed exactly once", 1, hA.closeCalls.get()); + assertFalse("segment B handle must remain cached", hB.freed.get()); + assertEquals("segment B handle must not be closed", 0, hB.closeCalls.get()); + + // Segment A's entry is gone, so a later acquire rebuilds it. + FilterBitsetCache.acquire( + null, + keyA, + "f", + () -> { + buildA.incrementAndGet(); + return new CountingHandle(); + }) + .decRef(); + assertEquals("invalidated segment rebuilds on next acquire", 2, buildA.get()); + } + /** A failed build is not cached, so a later acquire rebuilds. */ @Test public void failedBuildIsNotCachedAndCanRetry() throws Exception { From d3e87b8e16b8652dee6ff5e002a9fb64dfd462f4 Mon Sep 17 00:00:00 2001 From: James Xia Date: Tue, 21 Jul 2026 19:32:18 -0700 Subject: [PATCH 32/41] byte-budget filter cache with disable/clear controls Replace the fixed 16-entry filter bitset cache bound with a byte budget and add the client controls the reviewer asked for: - Byte-cap eviction: track total cached bytes and evict LRU entries until within budget (never the entry the in-flight query needs), plus a generous MAX_ENTRIES_GUARD against pathological tiny-entry streams. Host and device footprints match one-for-one, so one budget bounds both. - Controls: on by default; cuvs.lucene.filterBitsetCache.enabled and .maxBytes system properties plus setEnabled/setMaxBytes/clear(). Disabling routes every segment through the uncached, caller-owned build path. - Sizing: the budget defaults to min(16 x working set, 512 MiB), derived lazily from the first search's working set. An explicit budget smaller than that minimum working set fails loudly once; later index growth past the budget only warns and leans on eviction. A single segment larger than the whole budget is served uncached with a one-time warning. GPUKnnFloatVectorQuery computes each filtered segment's byte size in a pre-pass and reports the per-query working set before acquiring. Cache tests are reworked for the byte budget and reset the process-static cache in setup. --- .../nvidia/cuvs/lucene/FilterBitsetCache.java | 224 ++++++++++++++++-- .../cuvs/lucene/GPUKnnFloatVectorQuery.java | 43 +++- .../cuvs/lucene/TestFilterBitsetCache.java | 145 ++++++++++-- .../TestMultiSegmentGPUFilterConcurrency.java | 13 +- 4 files changed, 370 insertions(+), 55 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java b/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java index 0a7be8e1..4ce73e5c 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java +++ b/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java @@ -15,6 +15,7 @@ import java.util.Set; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; +import java.util.logging.Logger; import org.apache.lucene.index.IndexReader; import org.apache.lucene.search.Query; @@ -24,7 +25,17 @@ * segments' entries rather than the whole cache. * *

Host-side cache holding packed bitset arrays; the device-side upload is managed inside {@link - * FilterBitsetHandle} itself. Entries are evicted in LRU order. + * FilterBitsetHandle} itself. Entries are evicted in LRU order once the total cached size exceeds + * the configured byte budget. + * + *

Sizing

+ * + *

The cache is bounded by a byte budget rather than an entry count, since per-segment bitsets + * vary widely in size. The budget covers the host-side packed arrays; the device-side allocation + * mirrors it one-for-one, so a single budget bounds both. By default the budget is derived lazily + * from the first search's working set (16× the bytes needed to hold one bit per vector across the + * searched segments, capped at {@link #CEILING_BYTES}); it can also be set explicitly via {@link + * #setMaxBytes} or the {@code cuvs.lucene.filterBitsetCache.maxBytes} system property. * *

Concurrency

* @@ -42,7 +53,22 @@ */ final class FilterBitsetCache { - static final int MAX_HOST_ENTRIES = 16; + private static final Logger LOG = Logger.getLogger(FilterBitsetCache.class.getName()); + + static final String PROP_ENABLED = "cuvs.lucene.filterBitsetCache.enabled"; + static final String PROP_MAX_BYTES = "cuvs.lucene.filterBitsetCache.maxBytes"; + + /** Upper bound on the lazily-derived default budget. */ + static final long CEILING_BYTES = 512L * 1024 * 1024; + + /** Multiplier applied to the first search's working set when deriving the default budget. */ + static final long DEFAULT_BUDGET_MULTIPLIER = 16; + + /** + * Secondary guard on entry count so a pathological stream of tiny entries cannot grow the map + * unbounded even when the byte budget is generous. The byte budget is the primary bound. + */ + static final int MAX_ENTRIES_GUARD = 4096; /** Builds the handle for a key on a cache miss. */ @FunctionalInterface @@ -65,41 +91,150 @@ public int hashCode() { } } - 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; - } - }; + /** A cached future together with the byte size charged to the budget for its entry. */ + private record CacheEntry(CompletableFuture future, long bytes) {} + + private static final LinkedHashMap CACHE = + new LinkedHashMap<>(16, 0.75f, /* access-order= */ true); // Segment reader keys with a registered close listener, so each reader registers exactly once. private static final Set REGISTERED = new HashSet<>(); + // ---- Configuration (guarded by FilterBitsetCache.class) ---- + + private static boolean enabled = + Boolean.parseBoolean(System.getProperty(PROP_ENABLED, "true")); + + /** Byte budget; {@code <= 0} means "derive lazily from the first search's working set". */ + private static long maxBytes = Long.parseLong(System.getProperty(PROP_MAX_BYTES, "0")); + + /** Whether {@link #maxBytes} was configured explicitly (vs. left to lazy derivation). */ + private static boolean explicitBudget = maxBytes > 0; + + /** Whether the one-time hard min-size check against an explicit budget has run. */ + private static boolean minSizeChecked; + + private static boolean growthWarned; + private static boolean oversizedWarned; + + /** Total bytes charged to the budget across all live cache entries. */ + private static long totalBytes; + private FilterBitsetCache() {} + /** Whether caching is enabled. When disabled, callers build uncached, caller-owned handles. */ + static synchronized boolean isEnabled() { + return enabled; + } + + /** Enables or disables caching. Disabling does not evict existing entries; call {@link #clear}. */ + static synchronized void setEnabled(boolean value) { + enabled = value; + } + + /** + * Sets the byte budget. A positive value is honored as an explicit budget; a non-positive value + * restores lazy derivation from the next search's working set. + */ + static synchronized void setMaxBytes(long bytes) { + maxBytes = bytes; + explicitBudget = bytes > 0; + minSizeChecked = false; + growthWarned = false; + oversizedWarned = false; + } + + /** + * Informs the cache of the total byte working set of the search about to run — the sum, over the + * segments needing a filter, of the bytes to hold one bit per vector. Called once per search + * before any {@link #acquire}. + * + *

On the first informative call this resolves a lazily-derived budget. When the budget was set + * explicitly and is smaller than even this minimum working set, it fails loudly once (the caller + * should raise the budget or disable the cache). Later index growth that pushes the working set + * past the budget does not fail — the byte-cap eviction and the oversized-entry skip absorb it, + * with a one-time warning. + */ + static synchronized void onSearchWorkingSet(long workingSetBytes) { + if (workingSetBytes <= 0) return; + if (!explicitBudget && maxBytes <= 0) { + maxBytes = Math.min(DEFAULT_BUDGET_MULTIPLIER * workingSetBytes, CEILING_BYTES); + return; + } + if (explicitBudget && !minSizeChecked) { + minSizeChecked = true; + if (maxBytes < workingSetBytes) { + throw new IllegalStateException( + "filter bitset cache maxBytes=" + + maxBytes + + " is smaller than the minimum working set " + + workingSetBytes + + " bytes for this query; raise " + + PROP_MAX_BYTES + + " or disable the cache"); + } + } + if (maxBytes < workingSetBytes && !growthWarned) { + growthWarned = true; + LOG.warning( + "filter bitset cache maxBytes=" + + maxBytes + + " is smaller than the current working set " + + workingSetBytes + + " bytes; entries will be evicted or skipped, reducing cache effectiveness"); + } + } + /** * 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. + * + * @param entryBytes byte size charged to the budget for this entry (the packed bitset length) */ static FilterBitsetHandle acquire( - Query filter, Object segReaderKey, String field, FilterBuilder builder) throws IOException { + Query filter, Object segReaderKey, String field, long entryBytes, FilterBuilder builder) + throws IOException { + // A single entry larger than the whole budget is never cached; hand back a caller-owned handle. + boolean oversized; + synchronized (FilterBitsetCache.class) { + oversized = maxBytes > 0 && entryBytes > maxBytes; + if (oversized && !oversizedWarned) { + oversizedWarned = true; + LOG.warning( + "filter bitset entry of " + + entryBytes + + " bytes exceeds the cache budget " + + maxBytes + + " bytes; serving it uncached"); + } + } + if (oversized) { + return builder.build(); + } + FilterCacheKey key = new FilterCacheKey(filter, segReaderKey, field); while (true) { CompletableFuture future; boolean owner = false; + List> evicted = null; synchronized (FilterBitsetCache.class) { - future = CACHE.get(key); - if (future == null) { + CacheEntry entry = CACHE.get(key); + if (entry == null) { future = new CompletableFuture<>(); - CACHE.put(key, future); + CACHE.put(key, new CacheEntry(future, entryBytes)); + totalBytes += entryBytes; owner = true; + evicted = new ArrayList<>(); + evictToBudget(key, evicted); + } else { + future = entry.future(); + } + } + + if (evicted != null) { + for (CompletableFuture f : evicted) { + releaseCacheRef(f); } } @@ -111,7 +246,11 @@ static FilterBitsetHandle acquire( // 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); + CacheEntry cur = CACHE.get(key); + if (cur != null && cur.future() == future) { + CACHE.remove(key); + totalBytes -= cur.bytes(); + } } future.completeExceptionally(t); if (t instanceof IOException io) throw io; @@ -138,6 +277,26 @@ static FilterBitsetHandle acquire( } } + /** + * Evicts eldest (least-recently-used) entries until the total size is within the byte budget and + * the entry count is within the guard, never evicting {@code keep} (the entry needed by the + * in-flight acquire). Runs under the cache lock; the evicted futures are collected into {@code + * out} so their cache references are released after the lock is dropped. + */ + private static void evictToBudget( + FilterCacheKey keep, List> out) { + long budget = maxBytes > 0 ? maxBytes : Long.MAX_VALUE; + var it = CACHE.entrySet().iterator(); + while ((totalBytes > budget || CACHE.size() > MAX_ENTRIES_GUARD) && it.hasNext()) { + Map.Entry e = it.next(); + if (e.getKey().equals(keep)) continue; // access-order puts `keep` last; never evict it + CacheEntry ce = e.getValue(); + totalBytes -= ce.bytes(); + out.add(ce.future()); + it.remove(); + } + } + /** * Registers a one-time close listener on {@code helper} so every cache entry for this segment * reader is evicted (and its device allocation released) as soon as the reader's core closes — @@ -160,9 +319,11 @@ static void invalidateReader(Object segReaderKey) { REGISTERED.remove(segReaderKey); var it = CACHE.entrySet().iterator(); while (it.hasNext()) { - Map.Entry> e = it.next(); + Map.Entry e = it.next(); if (Objects.equals(e.getKey().segReaderKey(), segReaderKey)) { - toRelease.add(e.getValue()); + CacheEntry ce = e.getValue(); + totalBytes -= ce.bytes(); + toRelease.add(ce.future()); it.remove(); } } @@ -172,6 +333,27 @@ static void invalidateReader(Object segReaderKey) { } } + /** Evicts and releases every cache entry, resetting the byte budget accounting. */ + static void clear() { + List> toRelease = new ArrayList<>(); + synchronized (FilterBitsetCache.class) { + for (CacheEntry ce : CACHE.values()) { + toRelease.add(ce.future()); + } + CACHE.clear(); + REGISTERED.clear(); + totalBytes = 0; + } + for (CompletableFuture future : toRelease) { + releaseCacheRef(future); + } + } + + /** Total bytes currently charged to the budget. Test hook. */ + static synchronized long currentBytesForTests() { + return totalBytes; + } + /** Releases the cache's reference once the (possibly in-flight) build completes. */ private static void releaseCacheRef(CompletableFuture future) { future.whenComplete( diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index 19bd265a..16a552ee 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -307,19 +307,45 @@ private List buildPerSegmentFilterHandles( } final Weight sharedFilterWeight = filterWeight; - List handles = new ArrayList<>(leaves.size()); - for (int i = 0; i < leaves.size(); i++) { + final int n = leaves.size(); + boolean[] needsFilter = new boolean[n]; + FloatVectorValues[] fvvs = new FloatVectorValues[n]; + long[] entryBytes = new long[n]; + long workingSetBytes = 0; + for (int i = 0; i < n; i++) { LeafReaderContext ctx = leaves.get(i); - CuVS2510GPUVectorsReader gpuReader = gpuReaders.get(i); // No explicit filter and no deletes in this segment: no filter for this partition. if (filter == null && ctx.reader().getLiveDocs() == null) { + continue; + } + needsFilter[i] = true; + FloatVectorValues fvv = gpuReaders.get(i).getFloatVectorValues(field); + fvvs[i] = fvv; + // Bytes to hold one bit per vector, word-aligned: this segment's cache-entry size. + long segBytes = (((long) fvv.size() + 63) / 64) * 8; + entryBytes[i] = segBytes; + workingSetBytes += segBytes; + } + + boolean cacheEnabled = FilterBitsetCache.isEnabled(); + if (cacheEnabled) { + // Resolves the lazy default budget (and validates an explicit one) before any acquire. + FilterBitsetCache.onSearchWorkingSet(workingSetBytes); + } + + List handles = new ArrayList<>(n); + for (int i = 0; i < n; i++) { + if (!needsFilter[i]) { handles.add(null); continue; } - var helper = ctx.reader().getReaderCacheHelper(); + LeafReaderContext ctx = leaves.get(i); + FloatVectorValues fvv = fvvs[i]; + // When caching is disabled, route every segment through the uncached, caller-owned path. + var helper = cacheEnabled ? ctx.reader().getReaderCacheHelper() : null; if (helper == null) { // This reader can't be cached; build an uncached handle owned outright by the caller. - handles.add(buildSegmentFilterHandle(sharedFilterWeight, ctx, gpuReader)); + handles.add(buildSegmentFilterHandle(sharedFilterWeight, ctx, fvv)); } else { // Evict this segment's cache entries when its reader closes (merge/reopen), not just on LRU. FilterBitsetCache.ensureCloseListener(helper); @@ -328,7 +354,8 @@ private List buildPerSegmentFilterHandles( filter, helper.getKey(), field, - () -> buildSegmentFilterHandle(sharedFilterWeight, ctx, gpuReader))); + entryBytes[i], + () -> buildSegmentFilterHandle(sharedFilterWeight, ctx, fvv))); } } return handles; @@ -341,13 +368,11 @@ private List buildPerSegmentFilterHandles( * segment with deletes but no explicit Lucene filter. */ private FilterBitsetHandle buildSegmentFilterHandle( - Weight filterWeight, LeafReaderContext ctx, CuVS2510GPUVectorsReader gpuReader) - throws IOException { + Weight filterWeight, LeafReaderContext ctx, FloatVectorValues fvv) throws IOException { 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 = gpuReader.getFloatVectorValues(field); Bits acceptedOrds = fvv.getAcceptOrds(acceptDocs); int numOrds = fvv.size(); long[] segLongs = new long[(int) (((long) numOrds + 63) / 64)]; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestFilterBitsetCache.java b/src/test/java/com/nvidia/cuvs/lucene/TestFilterBitsetCache.java index c4c3c9cb..63bf6320 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestFilterBitsetCache.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestFilterBitsetCache.java @@ -21,15 +21,39 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; +import org.junit.After; +import org.junit.Before; 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. + * Concurrency and sizing tests for {@link FilterBitsetCache}. These exercise the compute-once, + * reference-counting, and byte-budget behavior in isolation by injecting a fake {@link + * FilterBitsetHandle}, so no GPU or native library is required. + * + *

The cache is static, so each test resets it via {@link Before}/{@link After}. */ public class TestFilterBitsetCache { + /** A byte size comfortably larger than any single test's total entries, so nothing evicts. */ + private static final long HUGE_BUDGET = 1L << 40; + + /** Representative per-entry byte size used by tests that don't care about sizing. */ + private static final long ENTRY_BYTES = 64; + + @Before + public void resetCache() { + FilterBitsetCache.clear(); + FilterBitsetCache.setEnabled(true); + FilterBitsetCache.setMaxBytes(HUGE_BUDGET); + } + + @After + public void restoreDefaults() { + FilterBitsetCache.clear(); + FilterBitsetCache.setEnabled(true); + FilterBitsetCache.setMaxBytes(0); // back to lazy-default derivation + } + /** * 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 @@ -95,6 +119,7 @@ public void computeOnceUnderConcurrentAcquire() throws Exception { null, segKey(field), field, + ENTRY_BYTES, () -> { buildCount.incrementAndGet(); // Widen the race window so waiters reach the future before it completes. @@ -129,33 +154,102 @@ public void computeOnceUnderConcurrentAcquire() throws Exception { assertFalse(handle.freed.get()); } - /** Eviction releases exactly the cache's reference, freeing evicted handles and keeping the rest. */ + /** Once the byte budget is exceeded, the eldest entry is evicted and freed; newer ones stay. */ @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 + public void byteCapEvictsEldestWhenBudgetExceeded() throws Exception { + // Budget holds two 64-byte entries; inserting a third must evict exactly the oldest. + FilterBitsetCache.setMaxBytes(2 * ENTRY_BYTES); List handles = new ArrayList<>(); - for (int i = 0; i < total; i++) { + for (int i = 0; i < 3; i++) { CountingHandle h = new CountingHandle(); handles.add(h); - final String field = "evict-" + i; - FilterBitsetHandle got = FilterBitsetCache.acquire(null, segKey(field), field, () -> h); + final String field = "bcap-" + i; + FilterBitsetHandle got = + FilterBitsetCache.acquire(null, segKey(field), field, ENTRY_BYTES, () -> 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()); + assertTrue("oldest handle should have been evicted and freed", handles.get(0).freed.get()); + assertEquals("evicted handle closed exactly once", 1, handles.get(0).closeCalls.get()); + assertFalse("second handle should still be cached", handles.get(1).freed.get()); + assertFalse("third handle should still be cached", handles.get(2).freed.get()); + assertTrue( + "total bytes must stay within budget", + FilterBitsetCache.currentBytesForTests() <= 2 * ENTRY_BYTES); + } + + /** clear() releases every cache reference (freeing all handles) and resets byte accounting. */ + @Test + public void clearFreesAllAndResets() throws Exception { + List handles = new ArrayList<>(); + for (int i = 0; i < 3; i++) { + CountingHandle h = new CountingHandle(); + handles.add(h); + final String field = "clr-" + i; + FilterBitsetCache.acquire(null, segKey(field), field, ENTRY_BYTES, () -> h).decRef(); } - for (int i = firstRetained; i < total; i++) { + assertTrue("entries charged to the budget", FilterBitsetCache.currentBytesForTests() > 0); + + FilterBitsetCache.clear(); + + for (int i = 0; i < handles.size(); 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()); + assertTrue("handle " + i + " freed by clear()", h.freed.get()); + assertEquals("handle " + i + " closed exactly once", 1, h.closeCalls.get()); } + assertEquals("byte accounting reset", 0, FilterBitsetCache.currentBytesForTests()); + } + + /** An entry larger than the whole budget is never cached; each acquire rebuilds it uncached. */ + @Test + public void oversizedEntryIsNotCached() throws Exception { + FilterBitsetCache.setMaxBytes(2 * ENTRY_BYTES); + final long oversized = 4 * ENTRY_BYTES; // larger than the whole budget + final String field = "oversized"; + AtomicInteger buildCount = new AtomicInteger(); + + CountingHandle h1 = new CountingHandle(); + FilterBitsetHandle got1 = + FilterBitsetCache.acquire( + null, + segKey(field), + field, + oversized, + () -> { + buildCount.incrementAndGet(); + return h1; + }); + assertSame(h1, got1); + got1.decRef(); // caller-owned; nothing else holds it + assertTrue("uncached oversized handle freed once caller releases it", h1.freed.get()); + assertEquals("nothing charged to the budget", 0, FilterBitsetCache.currentBytesForTests()); + + CountingHandle h2 = new CountingHandle(); + FilterBitsetHandle got2 = + FilterBitsetCache.acquire( + null, + segKey(field), + field, + oversized, + () -> { + buildCount.incrementAndGet(); + return h2; + }); + assertSame(h2, got2); + got2.decRef(); + assertEquals("oversized entry is rebuilt every time, never cached", 2, buildCount.get()); + } + + /** The enabled flag round-trips; the query path uses it to route around the cache. */ + @Test + public void enabledFlagToggles() { + assertTrue(FilterBitsetCache.isEnabled()); + FilterBitsetCache.setEnabled(false); + assertFalse(FilterBitsetCache.isEnabled()); + FilterBitsetCache.setEnabled(true); + assertTrue(FilterBitsetCache.isEnabled()); } /** invalidateReader evicts and frees exactly the given segment's entries, leaving others cached. */ @@ -172,12 +266,13 @@ public void invalidateReaderEvictsOnlyThatSegment() throws Exception { null, keyA, "f", + ENTRY_BYTES, () -> { buildA.incrementAndGet(); return hA; }) .decRef(); - FilterBitsetCache.acquire(null, keyB, "f", () -> hB).decRef(); + FilterBitsetCache.acquire(null, keyB, "f", ENTRY_BYTES, () -> hB).decRef(); // Invalidate segment A only, as its reader's close listener would. FilterBitsetCache.invalidateReader(keyA); @@ -192,6 +287,7 @@ public void invalidateReaderEvictsOnlyThatSegment() throws Exception { null, keyA, "f", + ENTRY_BYTES, () -> { buildA.incrementAndGet(); return new CountingHandle(); @@ -211,6 +307,7 @@ public void failedBuildIsNotCachedAndCanRetry() throws Exception { null, segKey(field), field, + ENTRY_BYTES, () -> { buildCount.incrementAndGet(); throw new IOException("boom"); @@ -220,6 +317,7 @@ public void failedBuildIsNotCachedAndCanRetry() throws Exception { assertEquals("boom", expected.getMessage()); } assertEquals(1, buildCount.get()); + assertEquals("failed build must not charge the budget", 0, FilterBitsetCache.currentBytesForTests()); CountingHandle handle = new CountingHandle(); FilterBitsetHandle got = @@ -227,6 +325,7 @@ public void failedBuildIsNotCachedAndCanRetry() throws Exception { null, segKey(field), field, + ENTRY_BYTES, () -> { buildCount.incrementAndGet(); return handle; @@ -238,11 +337,12 @@ public void failedBuildIsNotCachedAndCanRetry() throws Exception { /** * 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. + * cache trips {@link CountingHandle#decRef()}'s below-zero guard and fails the test. The budget is + * left huge (see {@link #resetCache}) so no eviction happens mid-run. */ @Test public void concurrentAcquireReleaseIsBalanced() throws Exception { - final int keySpace = 8; // < MAX_HOST_ENTRIES, so these keys are never evicted mid-run + final int keySpace = 8; final int threads = 16; final int iterationsPerThread = 500; @@ -260,7 +360,8 @@ public void concurrentAcquireReleaseIsBalanced() throws Exception { 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, segKey(field), field, CountingHandle::new); + FilterBitsetCache.acquire( + null, segKey(field), field, ENTRY_BYTES, CountingHandle::new); h.decRef(); } return null; diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestMultiSegmentGPUFilterConcurrency.java b/src/test/java/com/nvidia/cuvs/lucene/TestMultiSegmentGPUFilterConcurrency.java index fb8e92c6..ca182b7f 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestMultiSegmentGPUFilterConcurrency.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestMultiSegmentGPUFilterConcurrency.java @@ -47,7 +47,7 @@ /** * 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 + * (more than the cache's default budget spans) 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. */ @@ -60,8 +60,9 @@ public class TestMultiSegmentGPUFilterConcurrency extends LuceneTestCase { 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; + // More distinct filters than the default budget spans (16x one query's working set), so their + // churn forces byte-cap cache eviction. + private static final int NUM_CATEGORIES = (int) FilterBitsetCache.DEFAULT_BUDGET_MULTIPLIER + 8; private static Directory directory; private static DirectoryReader reader; @@ -74,6 +75,10 @@ public class TestMultiSegmentGPUFilterConcurrency extends LuceneTestCase { @BeforeClass public static void beforeClass() throws Exception { assumeTrue("cuVS not supported", isSupported()); + // Start from a clean, lazily-derived budget so churn is driven only by this test's filters. + FilterBitsetCache.clear(); + FilterBitsetCache.setEnabled(true); + FilterBitsetCache.setMaxBytes(0); codec = new CuVS2510GPUSearchCodec(); int datasetSize = 2000; @@ -184,6 +189,8 @@ public void testConcurrentFilteredSearchAcrossManyFilters() throws Exception { public static void afterClass() throws IOException { if (reader != null) reader.close(); if (directory != null) directory.close(); + FilterBitsetCache.clear(); + FilterBitsetCache.setMaxBytes(0); reader = null; directory = null; searcher = null; From 0fd1a9e8860d388131938113b57d31c286b5290f Mon Sep 17 00:00:00 2001 From: James Xia Date: Tue, 21 Jul 2026 19:49:46 -0700 Subject: [PATCH 33/41] back filter bitset cache with a sized device pool Default the cuVS filter device pool (com.nvidia.cuvs.filterBitsetPoolSize) to the resolved cache byte budget once it is known, before the first bitset upload, so the RMM pool backing filter device allocations tracks the cache size instead of churning cudaMalloc/cudaFree on every eviction/rebuild. An operator who sets the property explicitly overrides this default; the value is published only when unset. Adds filterPoolSizeDefaultsToBudget. --- .../nvidia/cuvs/lucene/FilterBitsetCache.java | 21 +++++++++++++++++++ .../cuvs/lucene/TestFilterBitsetCache.java | 21 +++++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java b/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java index 4ce73e5c..3f5350bb 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java +++ b/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java @@ -58,6 +58,13 @@ final class FilterBitsetCache { static final String PROP_ENABLED = "cuvs.lucene.filterBitsetCache.enabled"; static final String PROP_MAX_BYTES = "cuvs.lucene.filterBitsetCache.maxBytes"; + /** + * cuVS property (read by cuvs-java's filter resources) sizing the pre-warmed RMM pool that backs + * filter bitset device allocations. Defaulted to the resolved cache byte budget so the device pool + * tracks the cache; an operator can set it explicitly to override. + */ + static final String PROP_FILTER_POOL_BYTES = "com.nvidia.cuvs.filterBitsetPoolSize"; + /** Upper bound on the lazily-derived default budget. */ static final long CEILING_BYTES = 512L * 1024 * 1024; @@ -159,6 +166,7 @@ static synchronized void onSearchWorkingSet(long workingSetBytes) { if (workingSetBytes <= 0) return; if (!explicitBudget && maxBytes <= 0) { maxBytes = Math.min(DEFAULT_BUDGET_MULTIPLIER * workingSetBytes, CEILING_BYTES); + publishFilterPoolSize(); return; } if (explicitBudget && !minSizeChecked) { @@ -183,6 +191,19 @@ static synchronized void onSearchWorkingSet(long workingSetBytes) { + workingSetBytes + " bytes; entries will be evicted or skipped, reducing cache effectiveness"); } + publishFilterPoolSize(); + } + + /** + * Defaults the cuVS filter device pool size ({@link #PROP_FILTER_POOL_BYTES}) to the resolved + * cache budget, so cuvs-java pre-warms a growable RMM pool sized to the cache. Set once, before + * the first bitset upload, and only when the operator has not set the property explicitly. + */ + private static void publishFilterPoolSize() { + if (maxBytes <= 0) return; + if (System.getProperty(PROP_FILTER_POOL_BYTES) == null) { + System.setProperty(PROP_FILTER_POOL_BYTES, Long.toString(maxBytes)); + } } /** diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestFilterBitsetCache.java b/src/test/java/com/nvidia/cuvs/lucene/TestFilterBitsetCache.java index 63bf6320..e361375c 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestFilterBitsetCache.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestFilterBitsetCache.java @@ -54,6 +54,27 @@ public void restoreDefaults() { FilterBitsetCache.setMaxBytes(0); // back to lazy-default derivation } + /** The device pool property defaults to the resolved budget and never overrides an explicit one. */ + @Test + public void filterPoolSizeDefaultsToBudget() { + String prop = FilterBitsetCache.PROP_FILTER_POOL_BYTES; + String saved = System.getProperty(prop); + System.clearProperty(prop); + try { + FilterBitsetCache.setMaxBytes(0); // lazy: budget derived from the first working set + FilterBitsetCache.onSearchWorkingSet(1_000); + long expected = FilterBitsetCache.DEFAULT_BUDGET_MULTIPLIER * 1_000; // below the ceiling + assertEquals(Long.toString(expected), System.getProperty(prop)); + + // A later, larger working set must not overwrite the already-published pool size. + FilterBitsetCache.onSearchWorkingSet(4_000); + assertEquals(Long.toString(expected), System.getProperty(prop)); + } finally { + if (saved == null) System.clearProperty(prop); + else System.setProperty(prop, saved); + } + } + /** * 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 From 17fceeb7b1ba98e9bed2b774a00f71e193905e70 Mon Sep 17 00:00:00 2001 From: James Xia Date: Wed, 22 Jul 2026 08:21:39 -0700 Subject: [PATCH 34/41] align workspace pool size to RMM's 256-byte requirement RMM's pool_memory_resource requires the initial pool size to be a multiple of 256 bytes. Round com.nvidia.cuvs.workspacePoolSize up before passing it to setWorkspacePool, and skip non-positive values so a misconfigured size falls back to the default workspace resource instead of throwing and disabling cuVS for the thread. --- .../cuvs/lucene/ThreadLocalCuVSResourcesProvider.java | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java b/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java index e1bb46f6..4d98b61c 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java +++ b/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java @@ -49,7 +49,11 @@ private static CuVSResources cuVSResourcesOrNull() { CuVSResources resources = CuVSResources.create(); String poolSizeProp = System.getProperty(WORKSPACE_POOL_SIZE_PROPERTY); if (poolSizeProp != null) { - resources.setWorkspacePool(Long.parseLong(poolSizeProp)); + long poolBytes = Long.parseLong(poolSizeProp); + if (poolBytes > 0) { + // RMM's pool_memory_resource requires the initial size to be a multiple of 256 bytes. + resources.setWorkspacePool((poolBytes + 255L) & ~255L); + } } return resources; } catch (UnsupportedOperationException uoe) { From 361daf39efae5abac68c1dd6445fb56b69999cfe Mon Sep 17 00:00:00 2001 From: James Xia Date: Wed, 22 Jul 2026 08:58:17 -0700 Subject: [PATCH 35/41] Update cuvs version --- bench/pom.xml | 4 ++-- build.sh | 2 +- examples/pom.xml | 4 ++-- pom.xml | 4 ++-- 4 files changed, 7 insertions(+), 7 deletions(-) diff --git a/bench/pom.xml b/bench/pom.xml index 4271d802..f5955873 100644 --- a/bench/pom.xml +++ b/bench/pom.xml @@ -11,7 +11,7 @@ com.nvidia.cuvs.lucene.benchmarks cuvs-lucene-benchmarks - 26.08.0-SNAPSHOT + 26.10.0-SNAPSHOT jar cuvs-lucene-benchmarks @@ -85,7 +85,7 @@ com.nvidia.cuvs.lucene cuvs-lucene - 26.08.0-SNAPSHOT + 26.10.0-SNAPSHOT commons-io diff --git a/build.sh b/build.sh index 3c86351c..e7bfe225 100755 --- a/build.sh +++ b/build.sh @@ -8,7 +8,7 @@ set -e -u -o pipefail ARGS="$*" NUMARGS=$# -VERSION="26.08.0-SNAPSHOT" # Note: The version is updated automatically when ci/release/update-version.sh is invoked +VERSION="26.10.0-SNAPSHOT" # Note: The version is updated automatically when ci/release/update-version.sh is invoked GROUP_ID="com.nvidia.cuvs.lucene" function hasArg { diff --git a/examples/pom.xml b/examples/pom.xml index fd277c3e..50e06c3a 100644 --- a/examples/pom.xml +++ b/examples/pom.xml @@ -11,7 +11,7 @@ com.nvidia.cuvs.lucene.examples examples - 26.08.0-SNAPSHOT + 26.10.0-SNAPSHOT examples @@ -100,7 +100,7 @@ com.nvidia.cuvs.lucene cuvs-lucene - 26.08.0-SNAPSHOT + 26.10.0-SNAPSHOT diff --git a/pom.xml b/pom.xml index 0433ecc4..1676478f 100644 --- a/pom.xml +++ b/pom.xml @@ -12,7 +12,7 @@ com.nvidia.cuvs.lucene cuvs-lucene - 26.08.0-SNAPSHOT + 26.10.0-SNAPSHOT cuvs-lucene jar @@ -147,7 +147,7 @@ com.nvidia.cuvs cuvs-java - 26.08.0-SNAPSHOT + 26.10.0-SNAPSHOT From 11fce3f61664b049374238a0b530242253f93e41 Mon Sep 17 00:00:00 2001 From: James Xia Date: Thu, 23 Jul 2026 07:57:07 -0700 Subject: [PATCH 36/41] Add double-check for hardening --- .../java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java index 8c29fff3..eee7a530 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java @@ -519,7 +519,7 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits // 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) { + if (score == Float.MAX_VALUE || ord < 0) { continue; } float correctedScore = scoreCorrectionFunction.apply(score); From 094fa781a776c5618772381e17ad0b3e5857a3b1 Mon Sep 17 00:00:00 2001 From: James Xia Date: Thu, 23 Jul 2026 11:17:26 -0700 Subject: [PATCH 37/41] Make filter bitset cache codec-local Replace the static, JVM-configured filter cache with per-format cache instances configured through immutable codec parameters. Remove global property handling and the automatic cuvs-java device-pool bridge. Account for working sets per distinct cache, preserve handle lifecycle during failures, and update cache and concurrency tests. --- .../cuvs/lucene/CuVS2510GPUSearchCodec.java | 54 ++++++-- .../cuvs/lucene/CuVS2510GPUVectorsFormat.java | 21 ++- .../cuvs/lucene/CuVS2510GPUVectorsReader.java | 21 +++ .../nvidia/cuvs/lucene/FilterBitsetCache.java | 131 ++++++------------ .../cuvs/lucene/FilterBitsetCacheConfig.java | 28 ++++ .../cuvs/lucene/GPUKnnFloatVectorQuery.java | 72 ++++++---- .../cuvs/lucene/TestFilterBitsetCache.java | 122 ++++++++-------- .../TestMultiSegmentGPUFilterConcurrency.java | 13 +- 8 files changed, 274 insertions(+), 188 deletions(-) create mode 100644 src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCacheConfig.java diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java index 2de761a7..45a16740 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.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; @@ -29,8 +29,11 @@ public class CuVS2510GPUSearchCodec extends FilterCodec { * @throws Exception */ public CuVS2510GPUSearchCodec() throws Exception { - this(NAME, LuceneProvider.getCodec("101")); - initializeFormat(new GPUSearchParams.Builder().build()); + this( + NAME, + LuceneProvider.getCodec("101"), + new GPUSearchParams.Builder().build(), + FilterBitsetCacheConfig.DEFAULT); } /** @@ -41,8 +44,11 @@ public CuVS2510GPUSearchCodec() throws Exception { * @param delegate the delegate codec */ public CuVS2510GPUSearchCodec(String name, Codec delegate) { - super(name, delegate); - initializeFormat(new GPUSearchParams.Builder().build()); + this( + name, + delegate, + new GPUSearchParams.Builder().build(), + FilterBitsetCacheConfig.DEFAULT); } /** @@ -53,18 +59,48 @@ public CuVS2510GPUSearchCodec(String name, Codec delegate) { * @throws Exception Exception raised when initializing the codec */ public CuVS2510GPUSearchCodec(GPUSearchParams params) throws Exception { - this(NAME, LuceneProvider.getCodec("101")); - initializeFormat(params); + this(NAME, LuceneProvider.getCodec("101"), params, FilterBitsetCacheConfig.DEFAULT); + } + + /** + * Initialize the codec with GPU search and filter-bitset-cache parameters. + * + * @param params GPU index and search parameters + * @param filterCacheConfig filter-bitset-cache configuration + * @throws Exception Exception raised when initializing the codec + */ + public CuVS2510GPUSearchCodec( + GPUSearchParams params, FilterBitsetCacheConfig filterCacheConfig) throws Exception { + this(NAME, LuceneProvider.getCodec("101"), params, filterCacheConfig); + } + + /** + * Initialize a named codec with explicit delegate, GPU search, and filter-cache parameters. + * + * @param name the name of the codec + * @param delegate the delegate codec + * @param params GPU index and search parameters + * @param filterCacheConfig filter-bitset-cache configuration + */ + public CuVS2510GPUSearchCodec( + String name, + Codec delegate, + GPUSearchParams params, + FilterBitsetCacheConfig filterCacheConfig) { + super(name, delegate); + initializeFormat(params, filterCacheConfig); } /** * Initialize the {@link CuVS2510GPUVectorsFormat} instance using {@link GPUSearchParams}. * * @param params an instance of {@link GPUSearchParams} + * @param filterCacheConfig filter-bitset-cache configuration */ - private void initializeFormat(GPUSearchParams params) { + private void initializeFormat( + GPUSearchParams params, FilterBitsetCacheConfig filterCacheConfig) { try { - format = new CuVS2510GPUVectorsFormat(params); + format = new CuVS2510GPUVectorsFormat(params, filterCacheConfig); setKnnFormat(format); } catch (LibraryException ex) { log.log( diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java index 1abb9e58..e3967ff1 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java @@ -36,7 +36,8 @@ public class CuVS2510GPUVectorsFormat extends KnnVectorsFormat { public static final int VERSION_START = 0; public static final int VERSION_CURRENT = VERSION_START; - private GPUSearchParams gpuSearchParams; + private final GPUSearchParams gpuSearchParams; + private final FilterBitsetCache filterBitsetCache; static { try { @@ -55,7 +56,7 @@ public class CuVS2510GPUVectorsFormat extends KnnVectorsFormat { * @throws LibraryException if the native library fails to load */ public CuVS2510GPUVectorsFormat() { - this(new GPUSearchParams.Builder().build()); + this(new GPUSearchParams.Builder().build(), FilterBitsetCacheConfig.DEFAULT); } /** @@ -65,8 +66,21 @@ public CuVS2510GPUVectorsFormat() { * @throws LibraryException if the native library fails to load */ public CuVS2510GPUVectorsFormat(GPUSearchParams gpuSearchParams) { + this(gpuSearchParams, FilterBitsetCacheConfig.DEFAULT); + } + + /** + * Initializes the format with GPU search and filter-bitset-cache parameters. + * + * @param gpuSearchParams GPU index and search parameters + * @param filterCacheConfig filter-bitset-cache configuration + * @throws LibraryException if the native library fails to load + */ + public CuVS2510GPUVectorsFormat( + GPUSearchParams gpuSearchParams, FilterBitsetCacheConfig filterCacheConfig) { super("CuVS2510GPUVectorsFormat"); this.gpuSearchParams = gpuSearchParams; + this.filterBitsetCache = new FilterBitsetCache(filterCacheConfig); } /** @@ -85,7 +99,8 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException @Override public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { assertIsSupported(); - return new CuVS2510GPUVectorsReader(state, FLAT_VECTORS_FORMAT.fieldsReader(state)); + return new CuVS2510GPUVectorsReader( + state, FLAT_VECTORS_FORMAT.fieldsReader(state), filterBitsetCache); } /** diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java index eee7a530..4aa8a82a 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java @@ -66,6 +66,7 @@ public class CuVS2510GPUVectorsReader extends KnnVectorsReader { private final IntObjectHashMap fields; private final IntObjectHashMap cuvsIndices; private final IndexInput cuvsIndexInput; + private final FilterBitsetCache filterBitsetCache; static { try { @@ -86,7 +87,22 @@ public class CuVS2510GPUVectorsReader extends KnnVectorsReader { */ public CuVS2510GPUVectorsReader(SegmentReadState state, FlatVectorsReader flatReader) throws IOException { + this(state, flatReader, new FilterBitsetCache(FilterBitsetCacheConfig.DEFAULT)); + } + + /** + * Initializes the reader with the cache owned by its vectors format. + * + * @param state instance of the SegmentReadState + * @param flatReader instance of the FlatVectorsReader + * @param filterBitsetCache filter cache shared by readers from the same vectors format + * @throws IOException I/O exception + */ + CuVS2510GPUVectorsReader( + SegmentReadState state, FlatVectorsReader flatReader, FilterBitsetCache filterBitsetCache) + throws IOException { this.flatVectorsReader = flatReader; + this.filterBitsetCache = filterBitsetCache; this.fieldInfos = state.fieldInfos; this.fields = new IntObjectHashMap<>(); String metaFileName = @@ -626,6 +642,11 @@ public CagraIndex getCagraIndexForField(String field) { return gpuIndex.getCagraIndex(); } + /** Returns the filter cache owned by the vectors format that created this reader. */ + FilterBitsetCache getFilterBitsetCache() { + return filterBitsetCache; + } + /** * 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 index 3f5350bb..82037968 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java +++ b/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java @@ -34,8 +34,8 @@ * vary widely in size. The budget covers the host-side packed arrays; the device-side allocation * mirrors it one-for-one, so a single budget bounds both. By default the budget is derived lazily * from the first search's working set (16× the bytes needed to hold one bit per vector across the - * searched segments, capped at {@link #CEILING_BYTES}); it can also be set explicitly via {@link - * #setMaxBytes} or the {@code cuvs.lucene.filterBitsetCache.maxBytes} system property. + * searched segments, capped at {@link #CEILING_BYTES}); it can also be set explicitly through + * {@link FilterBitsetCacheConfig}. * *

Concurrency

* @@ -55,16 +55,6 @@ final class FilterBitsetCache { private static final Logger LOG = Logger.getLogger(FilterBitsetCache.class.getName()); - static final String PROP_ENABLED = "cuvs.lucene.filterBitsetCache.enabled"; - static final String PROP_MAX_BYTES = "cuvs.lucene.filterBitsetCache.maxBytes"; - - /** - * cuVS property (read by cuvs-java's filter resources) sizing the pre-warmed RMM pool that backs - * filter bitset device allocations. Defaulted to the resolved cache byte budget so the device pool - * tracks the cache; an operator can set it explicitly to override. - */ - static final String PROP_FILTER_POOL_BYTES = "com.nvidia.cuvs.filterBitsetPoolSize"; - /** Upper bound on the lazily-derived default budget. */ static final long CEILING_BYTES = 512L * 1024 * 1024; @@ -101,56 +91,43 @@ public int hashCode() { /** A cached future together with the byte size charged to the budget for its entry. */ private record CacheEntry(CompletableFuture future, long bytes) {} - private static final LinkedHashMap CACHE = + private final LinkedHashMap cache = new LinkedHashMap<>(16, 0.75f, /* access-order= */ true); // Segment reader keys with a registered close listener, so each reader registers exactly once. - private static final Set REGISTERED = new HashSet<>(); + private final Set registered = new HashSet<>(); - // ---- Configuration (guarded by FilterBitsetCache.class) ---- + // ---- Configuration and accounting (guarded by this) ---- - private static boolean enabled = - Boolean.parseBoolean(System.getProperty(PROP_ENABLED, "true")); + private final boolean enabled; /** Byte budget; {@code <= 0} means "derive lazily from the first search's working set". */ - private static long maxBytes = Long.parseLong(System.getProperty(PROP_MAX_BYTES, "0")); + private long maxBytes; /** Whether {@link #maxBytes} was configured explicitly (vs. left to lazy derivation). */ - private static boolean explicitBudget = maxBytes > 0; + private final boolean explicitBudget; /** Whether the one-time hard min-size check against an explicit budget has run. */ - private static boolean minSizeChecked; + private boolean minSizeChecked; - private static boolean growthWarned; - private static boolean oversizedWarned; + private boolean growthWarned; + private boolean oversizedWarned; /** Total bytes charged to the budget across all live cache entries. */ - private static long totalBytes; + private long totalBytes; - private FilterBitsetCache() {} + FilterBitsetCache(FilterBitsetCacheConfig config) { + Objects.requireNonNull(config, "config"); + enabled = config.enabled(); + maxBytes = config.maxBytes(); + explicitBudget = maxBytes > 0; + } /** Whether caching is enabled. When disabled, callers build uncached, caller-owned handles. */ - static synchronized boolean isEnabled() { + synchronized boolean isEnabled() { return enabled; } - /** Enables or disables caching. Disabling does not evict existing entries; call {@link #clear}. */ - static synchronized void setEnabled(boolean value) { - enabled = value; - } - - /** - * Sets the byte budget. A positive value is honored as an explicit budget; a non-positive value - * restores lazy derivation from the next search's working set. - */ - static synchronized void setMaxBytes(long bytes) { - maxBytes = bytes; - explicitBudget = bytes > 0; - minSizeChecked = false; - growthWarned = false; - oversizedWarned = false; - } - /** * Informs the cache of the total byte working set of the search about to run — the sum, over the * segments needing a filter, of the bytes to hold one bit per vector. Called once per search @@ -162,11 +139,10 @@ static synchronized void setMaxBytes(long bytes) { * past the budget does not fail — the byte-cap eviction and the oversized-entry skip absorb it, * with a one-time warning. */ - static synchronized void onSearchWorkingSet(long workingSetBytes) { + synchronized void onSearchWorkingSet(long workingSetBytes) { if (workingSetBytes <= 0) return; if (!explicitBudget && maxBytes <= 0) { maxBytes = Math.min(DEFAULT_BUDGET_MULTIPLIER * workingSetBytes, CEILING_BYTES); - publishFilterPoolSize(); return; } if (explicitBudget && !minSizeChecked) { @@ -177,9 +153,7 @@ static synchronized void onSearchWorkingSet(long workingSetBytes) { + maxBytes + " is smaller than the minimum working set " + workingSetBytes - + " bytes for this query; raise " - + PROP_MAX_BYTES - + " or disable the cache"); + + " bytes for this query; raise the configured maxBytes or disable the cache"); } } if (maxBytes < workingSetBytes && !growthWarned) { @@ -191,19 +165,6 @@ static synchronized void onSearchWorkingSet(long workingSetBytes) { + workingSetBytes + " bytes; entries will be evicted or skipped, reducing cache effectiveness"); } - publishFilterPoolSize(); - } - - /** - * Defaults the cuVS filter device pool size ({@link #PROP_FILTER_POOL_BYTES}) to the resolved - * cache budget, so cuvs-java pre-warms a growable RMM pool sized to the cache. Set once, before - * the first bitset upload, and only when the operator has not set the property explicitly. - */ - private static void publishFilterPoolSize() { - if (maxBytes <= 0) return; - if (System.getProperty(PROP_FILTER_POOL_BYTES) == null) { - System.setProperty(PROP_FILTER_POOL_BYTES, Long.toString(maxBytes)); - } } /** @@ -213,12 +174,12 @@ private static void publishFilterPoolSize() { * * @param entryBytes byte size charged to the budget for this entry (the packed bitset length) */ - static FilterBitsetHandle acquire( + FilterBitsetHandle acquire( Query filter, Object segReaderKey, String field, long entryBytes, FilterBuilder builder) throws IOException { // A single entry larger than the whole budget is never cached; hand back a caller-owned handle. boolean oversized; - synchronized (FilterBitsetCache.class) { + synchronized (this) { oversized = maxBytes > 0 && entryBytes > maxBytes; if (oversized && !oversizedWarned) { oversizedWarned = true; @@ -239,11 +200,11 @@ static FilterBitsetHandle acquire( CompletableFuture future; boolean owner = false; List> evicted = null; - synchronized (FilterBitsetCache.class) { - CacheEntry entry = CACHE.get(key); + synchronized (this) { + CacheEntry entry = cache.get(key); if (entry == null) { future = new CompletableFuture<>(); - CACHE.put(key, new CacheEntry(future, entryBytes)); + cache.put(key, new CacheEntry(future, entryBytes)); totalBytes += entryBytes; owner = true; evicted = new ArrayList<>(); @@ -266,10 +227,10 @@ static FilterBitsetHandle acquire( } 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) { - CacheEntry cur = CACHE.get(key); + synchronized (this) { + CacheEntry cur = cache.get(key); if (cur != null && cur.future() == future) { - CACHE.remove(key); + cache.remove(key); totalBytes -= cur.bytes(); } } @@ -304,11 +265,11 @@ static FilterBitsetHandle acquire( * in-flight acquire). Runs under the cache lock; the evicted futures are collected into {@code * out} so their cache references are released after the lock is dropped. */ - private static void evictToBudget( + private void evictToBudget( FilterCacheKey keep, List> out) { long budget = maxBytes > 0 ? maxBytes : Long.MAX_VALUE; - var it = CACHE.entrySet().iterator(); - while ((totalBytes > budget || CACHE.size() > MAX_ENTRIES_GUARD) && it.hasNext()) { + var it = cache.entrySet().iterator(); + while ((totalBytes > budget || cache.size() > MAX_ENTRIES_GUARD) && it.hasNext()) { Map.Entry e = it.next(); if (e.getKey().equals(keep)) continue; // access-order puts `keep` last; never evict it CacheEntry ce = e.getValue(); @@ -325,20 +286,20 @@ private static void evictToBudget( * cache to the segment lifecycle (mirroring Lucene's own query cache) so obsolete entries do not * linger until LRU eviction. */ - static void ensureCloseListener(IndexReader.CacheHelper helper) { + void ensureCloseListener(IndexReader.CacheHelper helper) { Object segReaderKey = helper.getKey(); - synchronized (FilterBitsetCache.class) { - if (!REGISTERED.add(segReaderKey)) return; // a listener is already registered for this reader + synchronized (this) { + if (!registered.add(segReaderKey)) return; // a listener is already registered for this reader } - helper.addClosedListener(FilterBitsetCache::invalidateReader); + helper.addClosedListener(this::invalidateReader); } /** Evicts and releases every cache entry belonging to the given segment reader key. */ - static void invalidateReader(Object segReaderKey) { + void invalidateReader(Object segReaderKey) { List> toRelease = new ArrayList<>(); - synchronized (FilterBitsetCache.class) { - REGISTERED.remove(segReaderKey); - var it = CACHE.entrySet().iterator(); + synchronized (this) { + registered.remove(segReaderKey); + var it = cache.entrySet().iterator(); while (it.hasNext()) { Map.Entry e = it.next(); if (Objects.equals(e.getKey().segReaderKey(), segReaderKey)) { @@ -355,14 +316,14 @@ static void invalidateReader(Object segReaderKey) { } /** Evicts and releases every cache entry, resetting the byte budget accounting. */ - static void clear() { + void clear() { List> toRelease = new ArrayList<>(); - synchronized (FilterBitsetCache.class) { - for (CacheEntry ce : CACHE.values()) { + synchronized (this) { + for (CacheEntry ce : cache.values()) { toRelease.add(ce.future()); } - CACHE.clear(); - REGISTERED.clear(); + cache.clear(); + registered.clear(); totalBytes = 0; } for (CompletableFuture future : toRelease) { @@ -371,7 +332,7 @@ static void clear() { } /** Total bytes currently charged to the budget. Test hook. */ - static synchronized long currentBytesForTests() { + synchronized long currentBytesForTests() { return totalBytes; } diff --git a/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCacheConfig.java b/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCacheConfig.java new file mode 100644 index 00000000..5cb99408 --- /dev/null +++ b/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCacheConfig.java @@ -0,0 +1,28 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.nvidia.cuvs.lucene; + +/** + * Configuration for the filter-bitset cache used by multi-segment GPU searches. + * + *

This configuration is local to one vectors-format instance. It does not configure cuvs-java's + * process-wide filter-bitset device pool. + * + * @param enabled whether filter bitsets are cached between searches + * @param maxBytes maximum bytes cached by this format; {@code 0} derives a budget from the first + * filtered search + */ +public record FilterBitsetCacheConfig(boolean enabled, long maxBytes) { + + /** Default configuration used by SPI-created codecs and vector formats. */ + public static final FilterBitsetCacheConfig DEFAULT = new FilterBitsetCacheConfig(true, 0); + + public FilterBitsetCacheConfig { + if (maxBytes < 0) { + throw new IllegalArgumentException("maxBytes must be non-negative"); + } + } +} diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index 16a552ee..644f5a64 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -18,7 +18,9 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.Comparator; +import java.util.IdentityHashMap; import java.util.List; +import java.util.Map; import org.apache.lucene.codecs.KnnVectorsReader; import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat; import org.apache.lucene.index.CodecReader; @@ -311,7 +313,8 @@ private List buildPerSegmentFilterHandles( boolean[] needsFilter = new boolean[n]; FloatVectorValues[] fvvs = new FloatVectorValues[n]; long[] entryBytes = new long[n]; - long workingSetBytes = 0; + FilterBitsetCache[] caches = new FilterBitsetCache[n]; + Map workingSetBytesByCache = new IdentityHashMap<>(); for (int i = 0; i < n; i++) { LeafReaderContext ctx = leaves.get(i); // No explicit filter and no deletes in this segment: no filter for this partition. @@ -319,46 +322,59 @@ private List buildPerSegmentFilterHandles( continue; } needsFilter[i] = true; + FilterBitsetCache cache = gpuReaders.get(i).getFilterBitsetCache(); + caches[i] = cache; FloatVectorValues fvv = gpuReaders.get(i).getFloatVectorValues(field); fvvs[i] = fvv; // Bytes to hold one bit per vector, word-aligned: this segment's cache-entry size. long segBytes = (((long) fvv.size() + 63) / 64) * 8; entryBytes[i] = segBytes; - workingSetBytes += segBytes; + if (cache.isEnabled()) { + workingSetBytesByCache.merge(cache, segBytes, Long::sum); + } } - boolean cacheEnabled = FilterBitsetCache.isEnabled(); - if (cacheEnabled) { - // Resolves the lazy default budget (and validates an explicit one) before any acquire. - FilterBitsetCache.onSearchWorkingSet(workingSetBytes); + // Resolve each cache's lazy default budget (and validate an explicit one) before any acquire. + for (Map.Entry entry : workingSetBytesByCache.entrySet()) { + entry.getKey().onSearchWorkingSet(entry.getValue()); } List handles = new ArrayList<>(n); - for (int i = 0; i < n; i++) { - if (!needsFilter[i]) { - handles.add(null); - continue; + try { + for (int i = 0; i < n; i++) { + if (!needsFilter[i]) { + handles.add(null); + continue; + } + LeafReaderContext ctx = leaves.get(i); + FloatVectorValues fvv = fvvs[i]; + FilterBitsetCache cache = caches[i]; + // When caching is disabled, route every segment through the uncached, caller-owned path. + var helper = cache.isEnabled() ? ctx.reader().getReaderCacheHelper() : null; + if (helper == null) { + // This reader can't be cached; build an uncached handle owned outright by the caller. + handles.add(buildSegmentFilterHandle(sharedFilterWeight, ctx, fvv)); + } else { + // Evict entries when this segment reader closes (merge/reopen), not just on LRU. + cache.ensureCloseListener(helper); + handles.add( + cache.acquire( + filter, + helper.getKey(), + field, + entryBytes[i], + () -> buildSegmentFilterHandle(sharedFilterWeight, ctx, fvv))); + } } - LeafReaderContext ctx = leaves.get(i); - FloatVectorValues fvv = fvvs[i]; - // When caching is disabled, route every segment through the uncached, caller-owned path. - var helper = cacheEnabled ? ctx.reader().getReaderCacheHelper() : null; - if (helper == null) { - // This reader can't be cached; build an uncached handle owned outright by the caller. - handles.add(buildSegmentFilterHandle(sharedFilterWeight, ctx, fvv)); - } else { - // Evict this segment's cache entries when its reader closes (merge/reopen), not just on LRU. - FilterBitsetCache.ensureCloseListener(helper); - handles.add( - FilterBitsetCache.acquire( - filter, - helper.getKey(), - field, - entryBytes[i], - () -> buildSegmentFilterHandle(sharedFilterWeight, ctx, fvv))); + return handles; + } catch (IOException | RuntimeException | Error e) { + // The caller cannot release a partially constructed list, so release every acquired handle + // here before propagating the failure. + for (FilterBitsetHandle handle : handles) { + if (handle != null) handle.decRef(); } + throw e; } - return handles; } /** diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestFilterBitsetCache.java b/src/test/java/com/nvidia/cuvs/lucene/TestFilterBitsetCache.java index e361375c..4bdff245 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestFilterBitsetCache.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestFilterBitsetCache.java @@ -29,8 +29,6 @@ * Concurrency and sizing tests for {@link FilterBitsetCache}. These exercise the compute-once, * reference-counting, and byte-budget behavior in isolation by injecting a fake {@link * FilterBitsetHandle}, so no GPU or native library is required. - * - *

The cache is static, so each test resets it via {@link Before}/{@link After}. */ public class TestFilterBitsetCache { @@ -40,39 +38,17 @@ public class TestFilterBitsetCache { /** Representative per-entry byte size used by tests that don't care about sizing. */ private static final long ENTRY_BYTES = 64; + private FilterBitsetCache cache; + @Before public void resetCache() { - FilterBitsetCache.clear(); - FilterBitsetCache.setEnabled(true); - FilterBitsetCache.setMaxBytes(HUGE_BUDGET); + cache = new FilterBitsetCache(new FilterBitsetCacheConfig(true, HUGE_BUDGET)); } @After public void restoreDefaults() { - FilterBitsetCache.clear(); - FilterBitsetCache.setEnabled(true); - FilterBitsetCache.setMaxBytes(0); // back to lazy-default derivation - } - - /** The device pool property defaults to the resolved budget and never overrides an explicit one. */ - @Test - public void filterPoolSizeDefaultsToBudget() { - String prop = FilterBitsetCache.PROP_FILTER_POOL_BYTES; - String saved = System.getProperty(prop); - System.clearProperty(prop); - try { - FilterBitsetCache.setMaxBytes(0); // lazy: budget derived from the first working set - FilterBitsetCache.onSearchWorkingSet(1_000); - long expected = FilterBitsetCache.DEFAULT_BUDGET_MULTIPLIER * 1_000; // below the ceiling - assertEquals(Long.toString(expected), System.getProperty(prop)); - - // A later, larger working set must not overwrite the already-published pool size. - FilterBitsetCache.onSearchWorkingSet(4_000); - assertEquals(Long.toString(expected), System.getProperty(prop)); - } finally { - if (saved == null) System.clearProperty(prop); - else System.setProperty(prop, saved); - } + cache.clear(); + cache = null; } /** @@ -136,7 +112,7 @@ public void computeOnceUnderConcurrentAcquire() throws Exception { pool.submit( () -> { start.await(); - return FilterBitsetCache.acquire( + return cache.acquire( null, segKey(field), field, @@ -179,7 +155,7 @@ public void computeOnceUnderConcurrentAcquire() throws Exception { @Test public void byteCapEvictsEldestWhenBudgetExceeded() throws Exception { // Budget holds two 64-byte entries; inserting a third must evict exactly the oldest. - FilterBitsetCache.setMaxBytes(2 * ENTRY_BYTES); + cache = new FilterBitsetCache(new FilterBitsetCacheConfig(true, 2 * ENTRY_BYTES)); List handles = new ArrayList<>(); for (int i = 0; i < 3; i++) { @@ -187,7 +163,7 @@ public void byteCapEvictsEldestWhenBudgetExceeded() throws Exception { handles.add(h); final String field = "bcap-" + i; FilterBitsetHandle got = - FilterBitsetCache.acquire(null, segKey(field), field, ENTRY_BYTES, () -> h); + cache.acquire(null, segKey(field), field, ENTRY_BYTES, () -> h); assertSame(h, got); got.decRef(); // caller finished; the cache keeps its reference } @@ -198,7 +174,7 @@ public void byteCapEvictsEldestWhenBudgetExceeded() throws Exception { assertFalse("third handle should still be cached", handles.get(2).freed.get()); assertTrue( "total bytes must stay within budget", - FilterBitsetCache.currentBytesForTests() <= 2 * ENTRY_BYTES); + cache.currentBytesForTests() <= 2 * ENTRY_BYTES); } /** clear() releases every cache reference (freeing all handles) and resets byte accounting. */ @@ -209,31 +185,31 @@ public void clearFreesAllAndResets() throws Exception { CountingHandle h = new CountingHandle(); handles.add(h); final String field = "clr-" + i; - FilterBitsetCache.acquire(null, segKey(field), field, ENTRY_BYTES, () -> h).decRef(); + cache.acquire(null, segKey(field), field, ENTRY_BYTES, () -> h).decRef(); } - assertTrue("entries charged to the budget", FilterBitsetCache.currentBytesForTests() > 0); + assertTrue("entries charged to the budget", cache.currentBytesForTests() > 0); - FilterBitsetCache.clear(); + cache.clear(); for (int i = 0; i < handles.size(); i++) { CountingHandle h = handles.get(i); assertTrue("handle " + i + " freed by clear()", h.freed.get()); assertEquals("handle " + i + " closed exactly once", 1, h.closeCalls.get()); } - assertEquals("byte accounting reset", 0, FilterBitsetCache.currentBytesForTests()); + assertEquals("byte accounting reset", 0, cache.currentBytesForTests()); } /** An entry larger than the whole budget is never cached; each acquire rebuilds it uncached. */ @Test public void oversizedEntryIsNotCached() throws Exception { - FilterBitsetCache.setMaxBytes(2 * ENTRY_BYTES); + cache = new FilterBitsetCache(new FilterBitsetCacheConfig(true, 2 * ENTRY_BYTES)); final long oversized = 4 * ENTRY_BYTES; // larger than the whole budget final String field = "oversized"; AtomicInteger buildCount = new AtomicInteger(); CountingHandle h1 = new CountingHandle(); FilterBitsetHandle got1 = - FilterBitsetCache.acquire( + cache.acquire( null, segKey(field), field, @@ -245,11 +221,11 @@ public void oversizedEntryIsNotCached() throws Exception { assertSame(h1, got1); got1.decRef(); // caller-owned; nothing else holds it assertTrue("uncached oversized handle freed once caller releases it", h1.freed.get()); - assertEquals("nothing charged to the budget", 0, FilterBitsetCache.currentBytesForTests()); + assertEquals("nothing charged to the budget", 0, cache.currentBytesForTests()); CountingHandle h2 = new CountingHandle(); FilterBitsetHandle got2 = - FilterBitsetCache.acquire( + cache.acquire( null, segKey(field), field, @@ -263,14 +239,48 @@ public void oversizedEntryIsNotCached() throws Exception { assertEquals("oversized entry is rebuilt every time, never cached", 2, buildCount.get()); } - /** The enabled flag round-trips; the query path uses it to route around the cache. */ + /** The public default configuration preserves the previous enabled, lazily-sized behavior. */ + @Test + public void defaultConfiguration() { + assertTrue(FilterBitsetCacheConfig.DEFAULT.enabled()); + assertEquals(0, FilterBitsetCacheConfig.DEFAULT.maxBytes()); + } + + /** Entries and lifecycle operations are isolated between cache instances. */ + @Test + public void cacheInstancesAreIndependent() throws Exception { + FilterBitsetCache other = + new FilterBitsetCache(new FilterBitsetCacheConfig(true, HUGE_BUDGET)); + CountingHandle first = new CountingHandle(); + CountingHandle second = new CountingHandle(); + Object key = segKey("shared-reader-key"); + + cache.acquire(null, key, "f", ENTRY_BYTES, () -> first).decRef(); + other.acquire(null, key, "f", ENTRY_BYTES, () -> second).decRef(); + + assertFalse(first.freed.get()); + assertFalse(second.freed.get()); + assertEquals(ENTRY_BYTES, cache.currentBytesForTests()); + assertEquals(ENTRY_BYTES, other.currentBytesForTests()); + + cache.clear(); + + assertTrue("clearing the first cache releases its handle", first.freed.get()); + assertFalse("clearing the first cache must not affect the second", second.freed.get()); + assertEquals(0, cache.currentBytesForTests()); + assertEquals(ENTRY_BYTES, other.currentBytesForTests()); + + other.clear(); + assertTrue("the second cache releases its own handle", second.freed.get()); + } + + /** The enabled flag reflects the instance's immutable configuration. */ @Test - public void enabledFlagToggles() { - assertTrue(FilterBitsetCache.isEnabled()); - FilterBitsetCache.setEnabled(false); - assertFalse(FilterBitsetCache.isEnabled()); - FilterBitsetCache.setEnabled(true); - assertTrue(FilterBitsetCache.isEnabled()); + public void enabledFlagReflectsConfiguration() { + assertTrue(cache.isEnabled()); + FilterBitsetCache disabled = + new FilterBitsetCache(new FilterBitsetCacheConfig(false, 0)); + assertFalse(disabled.isEnabled()); } /** invalidateReader evicts and frees exactly the given segment's entries, leaving others cached. */ @@ -283,7 +293,7 @@ public void invalidateReaderEvictsOnlyThatSegment() throws Exception { AtomicInteger buildA = new AtomicInteger(); // Cache one entry per segment, releasing the caller ref so the cache keeps its own. - FilterBitsetCache.acquire( + cache.acquire( null, keyA, "f", @@ -293,10 +303,10 @@ public void invalidateReaderEvictsOnlyThatSegment() throws Exception { return hA; }) .decRef(); - FilterBitsetCache.acquire(null, keyB, "f", ENTRY_BYTES, () -> hB).decRef(); + cache.acquire(null, keyB, "f", ENTRY_BYTES, () -> hB).decRef(); // Invalidate segment A only, as its reader's close listener would. - FilterBitsetCache.invalidateReader(keyA); + cache.invalidateReader(keyA); assertTrue("segment A handle freed on invalidation", hA.freed.get()); assertEquals("segment A handle closed exactly once", 1, hA.closeCalls.get()); @@ -304,7 +314,7 @@ public void invalidateReaderEvictsOnlyThatSegment() throws Exception { assertEquals("segment B handle must not be closed", 0, hB.closeCalls.get()); // Segment A's entry is gone, so a later acquire rebuilds it. - FilterBitsetCache.acquire( + cache.acquire( null, keyA, "f", @@ -324,7 +334,7 @@ public void failedBuildIsNotCachedAndCanRetry() throws Exception { AtomicInteger buildCount = new AtomicInteger(); try { - FilterBitsetCache.acquire( + cache.acquire( null, segKey(field), field, @@ -338,11 +348,11 @@ public void failedBuildIsNotCachedAndCanRetry() throws Exception { assertEquals("boom", expected.getMessage()); } assertEquals(1, buildCount.get()); - assertEquals("failed build must not charge the budget", 0, FilterBitsetCache.currentBytesForTests()); + assertEquals("failed build must not charge the budget", 0, cache.currentBytesForTests()); CountingHandle handle = new CountingHandle(); FilterBitsetHandle got = - FilterBitsetCache.acquire( + cache.acquire( null, segKey(field), field, @@ -381,7 +391,7 @@ public void concurrentAcquireReleaseIsBalanced() throws Exception { final String field = "stress-" + ((seed + i) % keySpace); // A fresh handle per build; a cached key reuses whatever was built first. FilterBitsetHandle h = - FilterBitsetCache.acquire( + cache.acquire( null, segKey(field), field, ENTRY_BYTES, CountingHandle::new); h.decRef(); } diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestMultiSegmentGPUFilterConcurrency.java b/src/test/java/com/nvidia/cuvs/lucene/TestMultiSegmentGPUFilterConcurrency.java index ca182b7f..e9ae71fb 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestMultiSegmentGPUFilterConcurrency.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestMultiSegmentGPUFilterConcurrency.java @@ -75,11 +75,12 @@ public class TestMultiSegmentGPUFilterConcurrency extends LuceneTestCase { @BeforeClass public static void beforeClass() throws Exception { assumeTrue("cuVS not supported", isSupported()); - // Start from a clean, lazily-derived budget so churn is driven only by this test's filters. - FilterBitsetCache.clear(); - FilterBitsetCache.setEnabled(true); - FilterBitsetCache.setMaxBytes(0); - codec = new CuVS2510GPUSearchCodec(); + // The explicit codec config applies while writing. DirectoryReader.open resolves the codec + // through SPI for reading, whose no-argument instance uses the same default, lazily-derived + // cache budget. + codec = + new CuVS2510GPUSearchCodec( + new GPUSearchParams.Builder().build(), FilterBitsetCacheConfig.DEFAULT); int datasetSize = 2000; int dimensions = 128; @@ -189,8 +190,6 @@ public void testConcurrentFilteredSearchAcrossManyFilters() throws Exception { public static void afterClass() throws IOException { if (reader != null) reader.close(); if (directory != null) directory.close(); - FilterBitsetCache.clear(); - FilterBitsetCache.setMaxBytes(0); reader = null; directory = null; searcher = null; From e7e50b5509cb914558c80139d6e1a4b3b57a6722 Mon Sep 17 00:00:00 2001 From: James Xia Date: Thu, 23 Jul 2026 14:27:57 -0700 Subject: [PATCH 38/41] Make async RMM initialization application-controlled Remove the process-wide allocator switch from vectors format class initialization. Document explicit startup configuration and update examples, benchmarks, and concurrency tests to enable async RMM before creating cuVS resources. --- README.md | 15 +++++++++++++++ .../benchmarks/CagraIndexingBenchmarks.java | 4 +++- .../lucene/benchmarks/CagraSearchBenchmarks.java | 4 +++- .../examples/IndexAndSearchonGPUExample.java | 4 ++++ .../cuvs/lucene/CuVS2510GPUVectorsFormat.java | 2 -- .../TestMultiSegmentGPUFilterConcurrency.java | 6 ++++++ .../lucene/TestMultithreadedCuVSGPUSearch.java | 8 +++++++- 7 files changed, 38 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 4162e19d..34ca8efa 100644 --- a/README.md +++ b/README.md @@ -62,6 +62,21 @@ The example below plugs the GPU-accelerated HNSW codec into a standard Lucene `I Before running it, make sure cuVS is installed and available on your system library load path. The cuVS [tarball install instructions](https://docs.rapids.ai/api/cuvs/stable/build/#download-extract) show how to set this up. +### RMM async allocation for GPU search + +Applications using `CuVS2510GPUSearchCodec` can opt into RMM's stream-ordered asynchronous device +allocator during startup: + +```java +CuVSProvider.provider().enableRMMAsyncMemory(); +``` + +Call this before creating any cuVS resources, codecs, writers, or readers. The setting affects the +entire process on the current CUDA device, so allocator policy belongs to the application rather +than an individual Lucene codec. Async allocation is optional for correctness and recommended for +GPU workloads with repeated device allocations, especially concurrent or multi-stream searches. +Applications that do not opt in use the default RMM device-memory resource. + In a Maven project that includes the `cuvs-lucene` dependency shown above, create `src/main/java/com/nvidia/cuvs/lucene/examples/HelloCuvsLucene.java`: ```java diff --git a/bench/src/main/java/com/nvidia/cuvs/lucene/benchmarks/CagraIndexingBenchmarks.java b/bench/src/main/java/com/nvidia/cuvs/lucene/benchmarks/CagraIndexingBenchmarks.java index 119f95a1..01cf540c 100644 --- a/bench/src/main/java/com/nvidia/cuvs/lucene/benchmarks/CagraIndexingBenchmarks.java +++ b/bench/src/main/java/com/nvidia/cuvs/lucene/benchmarks/CagraIndexingBenchmarks.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.benchmarks; @@ -9,6 +9,7 @@ import static com.nvidia.cuvs.lucene.benchmarks.Utils.index; import com.nvidia.cuvs.lucene.CuVS2510GPUSearchCodec; +import com.nvidia.cuvs.spi.CuVSProvider; import java.io.IOException; import java.nio.file.Path; import java.nio.file.Paths; @@ -57,6 +58,7 @@ public class CagraIndexingBenchmarks { @Setup(Level.Trial) public void setup() throws Exception { + CuVSProvider.provider().enableRMMAsyncMemory(); random = new Random(222); indexDirPath = Paths.get(UUID.randomUUID().toString()); codec = new CuVS2510GPUSearchCodec(); diff --git a/bench/src/main/java/com/nvidia/cuvs/lucene/benchmarks/CagraSearchBenchmarks.java b/bench/src/main/java/com/nvidia/cuvs/lucene/benchmarks/CagraSearchBenchmarks.java index 45f35caf..111d8c14 100644 --- a/bench/src/main/java/com/nvidia/cuvs/lucene/benchmarks/CagraSearchBenchmarks.java +++ b/bench/src/main/java/com/nvidia/cuvs/lucene/benchmarks/CagraSearchBenchmarks.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.benchmarks; @@ -11,6 +11,7 @@ import com.nvidia.cuvs.lucene.CuVS2510GPUSearchCodec; import com.nvidia.cuvs.lucene.GPUKnnFloatVectorQuery; +import com.nvidia.cuvs.spi.CuVSProvider; import java.io.IOException; import java.nio.file.Path; import java.nio.file.Paths; @@ -62,6 +63,7 @@ public class CagraSearchBenchmarks { @Setup(Level.Trial) public void setup() throws Exception { + CuVSProvider.provider().enableRMMAsyncMemory(); random = new Random(222); indexDirPath = Paths.get(UUID.randomUUID().toString()); codec = new CuVS2510GPUSearchCodec(); diff --git a/examples/src/main/java/com/nvidia/cuvs/lucene/examples/IndexAndSearchonGPUExample.java b/examples/src/main/java/com/nvidia/cuvs/lucene/examples/IndexAndSearchonGPUExample.java index dd920001..1fd0a39f 100644 --- a/examples/src/main/java/com/nvidia/cuvs/lucene/examples/IndexAndSearchonGPUExample.java +++ b/examples/src/main/java/com/nvidia/cuvs/lucene/examples/IndexAndSearchonGPUExample.java @@ -10,6 +10,7 @@ import com.nvidia.cuvs.lucene.CuVS2510GPUSearchCodec; import com.nvidia.cuvs.lucene.GPUKnnFloatVectorQuery; import com.nvidia.cuvs.lucene.GPUSearchParams; +import com.nvidia.cuvs.spi.CuVSProvider; import java.io.File; import java.nio.file.Path; import java.nio.file.Paths; @@ -49,6 +50,9 @@ public class IndexAndSearchonGPUExample { public static void main(String[] args) throws Exception { + // Select the process-wide RMM allocator before creating any cuVS resources or codecs. + CuVSProvider.provider().enableRMMAsyncMemory(); + GPUSearchParams params = new GPUSearchParams.Builder().build(); Codec codec = new CuVS2510GPUSearchCodec(params); IndexWriterConfig config = new IndexWriterConfig().setCodec(codec).setUseCompoundFile(false); diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java index e3967ff1..3e62bd71 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java @@ -7,7 +7,6 @@ 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; @@ -41,7 +40,6 @@ 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/test/java/com/nvidia/cuvs/lucene/TestMultiSegmentGPUFilterConcurrency.java b/src/test/java/com/nvidia/cuvs/lucene/TestMultiSegmentGPUFilterConcurrency.java index e9ae71fb..a2957d2b 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestMultiSegmentGPUFilterConcurrency.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestMultiSegmentGPUFilterConcurrency.java @@ -9,6 +9,7 @@ import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; +import com.nvidia.cuvs.spi.CuVSProvider; import java.io.IOException; import java.util.ArrayList; import java.util.HashSet; @@ -74,6 +75,11 @@ public class TestMultiSegmentGPUFilterConcurrency extends LuceneTestCase { @BeforeClass public static void beforeClass() throws Exception { + try { + CuVSProvider.provider().enableRMMAsyncMemory(); + } catch (UnsupportedOperationException unsupported) { + assumeTrue("cuVS not supported: " + unsupported.getMessage(), false); + } assumeTrue("cuVS not supported", isSupported()); // The explicit codec config applies while writing. DirectoryReader.open resolves the codec // through SPI for reading, whose no-argument instance uses the same default, lazily-derived diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestMultithreadedCuVSGPUSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestMultithreadedCuVSGPUSearch.java index 25dc13b8..a7215ca8 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestMultithreadedCuVSGPUSearch.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestMultithreadedCuVSGPUSearch.java @@ -1,5 +1,5 @@ /* - * 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 */ @@ -9,6 +9,7 @@ import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; +import com.nvidia.cuvs.spi.CuVSProvider; import java.io.IOException; import java.util.Arrays; import java.util.Random; @@ -58,6 +59,11 @@ public class TestMultithreadedCuVSGPUSearch extends LuceneTestCase { @BeforeClass public static void beforeClass() throws IOException { + try { + CuVSProvider.provider().enableRMMAsyncMemory(); + } catch (UnsupportedOperationException unsupported) { + assumeTrue("cuVS not supported: " + unsupported.getMessage(), false); + } assumeTrue("cuVS not supported", isSupported()); random = random(); directory = newDirectory(new ByteBuffersDirectory()); From b77bfbf12bbba97e87656bdd9de6c1875edec256 Mon Sep 17 00:00:00 2001 From: James Xia Date: Fri, 24 Jul 2026 08:32:56 -0700 Subject: [PATCH 39/41] Use fixed and explicitly configurable filter cache budgets Replace first-query-derived sizing with a deterministic 128 MiB default while retaining FilterBitsetCacheConfig for explicit budgets. Validate configured limits and preserve eviction coverage with a test-specific cache size. --- .../nvidia/cuvs/lucene/FilterBitsetCache.java | 56 +++++-------------- .../cuvs/lucene/FilterBitsetCacheConfig.java | 17 +++++- .../cuvs/lucene/GPUKnnFloatVectorQuery.java | 2 +- .../cuvs/lucene/TestFilterBitsetCache.java | 19 ++++++- .../TestMultiSegmentGPUFilterConcurrency.java | 33 +++++++---- 5 files changed, 68 insertions(+), 59 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java b/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java index 82037968..af30ebef 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java +++ b/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java @@ -32,10 +32,12 @@ * *

The cache is bounded by a byte budget rather than an entry count, since per-segment bitsets * vary widely in size. The budget covers the host-side packed arrays; the device-side allocation - * mirrors it one-for-one, so a single budget bounds both. By default the budget is derived lazily - * from the first search's working set (16× the bytes needed to hold one bit per vector across the - * searched segments, capped at {@link #CEILING_BYTES}); it can also be set explicitly through - * {@link FilterBitsetCacheConfig}. + * mirrors it one-for-one, so a single budget bounds both. The default is the fixed {@link + * FilterBitsetCacheConfig#DEFAULT_MAX_BYTES} budget and can be overridden through {@link + * FilterBitsetCacheConfig}. The budget is not preallocated: packed arrays are retained as entries + * are inserted. Actual host usage can exceed the accounted payload bytes because of object + * overhead, temporary filter construction, and handles that remain referenced by in-flight + * searches after eviction. * *

Concurrency

* @@ -55,12 +57,6 @@ final class FilterBitsetCache { private static final Logger LOG = Logger.getLogger(FilterBitsetCache.class.getName()); - /** Upper bound on the lazily-derived default budget. */ - static final long CEILING_BYTES = 512L * 1024 * 1024; - - /** Multiplier applied to the first search's working set when deriving the default budget. */ - static final long DEFAULT_BUDGET_MULTIPLIER = 16; - /** * Secondary guard on entry count so a pathological stream of tiny entries cannot grow the map * unbounded even when the byte budget is generous. The byte budget is the primary bound. @@ -101,16 +97,10 @@ private record CacheEntry(CompletableFuture future, long byt private final boolean enabled; - /** Byte budget; {@code <= 0} means "derive lazily from the first search's working set". */ - private long maxBytes; - - /** Whether {@link #maxBytes} was configured explicitly (vs. left to lazy derivation). */ - private final boolean explicitBudget; - - /** Whether the one-time hard min-size check against an explicit budget has run. */ - private boolean minSizeChecked; + /** Fixed byte budget for this cache instance. */ + private final long maxBytes; - private boolean growthWarned; + private boolean workingSetWarned; private boolean oversizedWarned; /** Total bytes charged to the budget across all live cache entries. */ @@ -120,7 +110,6 @@ private record CacheEntry(CompletableFuture future, long byt Objects.requireNonNull(config, "config"); enabled = config.enabled(); maxBytes = config.maxBytes(); - explicitBudget = maxBytes > 0; } /** Whether caching is enabled. When disabled, callers build uncached, caller-owned handles. */ @@ -133,31 +122,14 @@ synchronized boolean isEnabled() { * segments needing a filter, of the bytes to hold one bit per vector. Called once per search * before any {@link #acquire}. * - *

On the first informative call this resolves a lazily-derived budget. When the budget was set - * explicitly and is smaller than even this minimum working set, it fails loudly once (the caller - * should raise the budget or disable the cache). Later index growth that pushes the working set - * past the budget does not fail — the byte-cap eviction and the oversized-entry skip absorb it, - * with a one-time warning. + *

If the working set exceeds the fixed budget, the byte-cap eviction and oversized-entry skip + * keep memory bounded, but repeated searches may rebuild entries. A warning is emitted once so + * operators can raise the budget or disable the cache. */ synchronized void onSearchWorkingSet(long workingSetBytes) { if (workingSetBytes <= 0) return; - if (!explicitBudget && maxBytes <= 0) { - maxBytes = Math.min(DEFAULT_BUDGET_MULTIPLIER * workingSetBytes, CEILING_BYTES); - return; - } - if (explicitBudget && !minSizeChecked) { - minSizeChecked = true; - if (maxBytes < workingSetBytes) { - throw new IllegalStateException( - "filter bitset cache maxBytes=" - + maxBytes - + " is smaller than the minimum working set " - + workingSetBytes - + " bytes for this query; raise the configured maxBytes or disable the cache"); - } - } - if (maxBytes < workingSetBytes && !growthWarned) { - growthWarned = true; + if (maxBytes < workingSetBytes && !workingSetWarned) { + workingSetWarned = true; LOG.warning( "filter bitset cache maxBytes=" + maxBytes diff --git a/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCacheConfig.java b/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCacheConfig.java index 5cb99408..576678c2 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCacheConfig.java +++ b/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCacheConfig.java @@ -11,18 +11,29 @@ *

This configuration is local to one vectors-format instance. It does not configure cuvs-java's * process-wide filter-bitset device pool. * + *

The byte budget is a retention cap, not a preallocation. Packed host arrays and their + * device-side mirrors are allocated as cache entries are populated. Actual host usage can exceed + * the accounted payload bytes because of object overhead, temporary filter construction, and + * in-flight handles. + * * @param enabled whether filter bitsets are cached between searches - * @param maxBytes maximum bytes cached by this format; {@code 0} derives a budget from the first - * filtered search + * @param maxBytes maximum bytes cached by this format; must be positive when caching is enabled */ public record FilterBitsetCacheConfig(boolean enabled, long maxBytes) { + /** Fixed 128 MiB cache budget used by SPI-created codecs and vector formats. */ + public static final long DEFAULT_MAX_BYTES = 128L * 1024 * 1024; + /** Default configuration used by SPI-created codecs and vector formats. */ - public static final FilterBitsetCacheConfig DEFAULT = new FilterBitsetCacheConfig(true, 0); + public static final FilterBitsetCacheConfig DEFAULT = + new FilterBitsetCacheConfig(true, DEFAULT_MAX_BYTES); public FilterBitsetCacheConfig { if (maxBytes < 0) { throw new IllegalArgumentException("maxBytes must be non-negative"); } + if (enabled && maxBytes == 0) { + throw new IllegalArgumentException("maxBytes must be positive when caching is enabled"); + } } } diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index 644f5a64..9e2ca291 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -334,7 +334,7 @@ private List buildPerSegmentFilterHandles( } } - // Resolve each cache's lazy default budget (and validate an explicit one) before any acquire. + // Report each cache's complete working set before any acquire so undersizing is diagnosed once. for (Map.Entry entry : workingSetBytesByCache.entrySet()) { entry.getKey().onSearchWorkingSet(entry.getValue()); } diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestFilterBitsetCache.java b/src/test/java/com/nvidia/cuvs/lucene/TestFilterBitsetCache.java index 4bdff245..686aacc4 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestFilterBitsetCache.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestFilterBitsetCache.java @@ -239,11 +239,26 @@ public void oversizedEntryIsNotCached() throws Exception { assertEquals("oversized entry is rebuilt every time, never cached", 2, buildCount.get()); } - /** The public default configuration preserves the previous enabled, lazily-sized behavior. */ + /** The public default configuration has a deterministic 128 MiB budget. */ @Test public void defaultConfiguration() { assertTrue(FilterBitsetCacheConfig.DEFAULT.enabled()); - assertEquals(0, FilterBitsetCacheConfig.DEFAULT.maxBytes()); + assertEquals(128L * 1024 * 1024, FilterBitsetCacheConfig.DEFAULT.maxBytes()); + } + + /** Enabled caches require a real budget; zero remains valid when the cache is disabled. */ + @Test + public void enabledConfigurationRequiresPositiveBudget() { + try { + new FilterBitsetCacheConfig(true, 0); + fail("enabled cache with a zero budget should be rejected"); + } catch (IllegalArgumentException expected) { + assertEquals("maxBytes must be positive when caching is enabled", expected.getMessage()); + } + + FilterBitsetCacheConfig disabled = new FilterBitsetCacheConfig(false, 0); + assertFalse(disabled.enabled()); + assertEquals(0, disabled.maxBytes()); } /** Entries and lifecycle operations are isolated between cache instances. */ diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestMultiSegmentGPUFilterConcurrency.java b/src/test/java/com/nvidia/cuvs/lucene/TestMultiSegmentGPUFilterConcurrency.java index a2957d2b..0c76956c 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestMultiSegmentGPUFilterConcurrency.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestMultiSegmentGPUFilterConcurrency.java @@ -24,6 +24,7 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.KnnVectorsFormat; import org.apache.lucene.document.Document; import org.apache.lucene.document.Field; import org.apache.lucene.document.KnnFloatVectorField; @@ -48,9 +49,9 @@ /** * 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 the cache's default budget spans) 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. + * (more than the test's deliberately small cache budget spans) so the shared filter cache builds, + * caches, evicts, and reuses handles under contention. Every returned hit must satisfy its filter, + * which validates that reference-counted eviction never frees a handle still in use. */ @SuppressSysoutChecks(bugUrl = "") public class TestMultiSegmentGPUFilterConcurrency extends LuceneTestCase { @@ -61,9 +62,10 @@ public class TestMultiSegmentGPUFilterConcurrency extends LuceneTestCase { private static Codec codec; private static final String VECTOR_FIELD = "vectors"; private static final String CATEGORY_FIELD = "cat"; - // More distinct filters than the default budget spans (16x one query's working set), so their - // churn forces byte-cap cache eviction. - private static final int NUM_CATEGORIES = (int) FilterBitsetCache.DEFAULT_BUDGET_MULTIPLIER + 8; + private static final long TEST_CACHE_MAX_BYTES = 4 * 1024; + private static final FilterBitsetCacheConfig TEST_CACHE_CONFIG = + new FilterBitsetCacheConfig(true, TEST_CACHE_MAX_BYTES); + private static final int NUM_CATEGORIES = 24; private static Directory directory; private static DirectoryReader reader; @@ -81,12 +83,9 @@ public static void beforeClass() throws Exception { assumeTrue("cuVS not supported: " + unsupported.getMessage(), false); } assumeTrue("cuVS not supported", isSupported()); - // The explicit codec config applies while writing. DirectoryReader.open resolves the codec - // through SPI for reading, whose no-argument instance uses the same default, lazily-derived - // cache budget. codec = new CuVS2510GPUSearchCodec( - new GPUSearchParams.Builder().build(), FilterBitsetCacheConfig.DEFAULT); + new GPUSearchParams.Builder().build(), TEST_CACHE_CONFIG); int datasetSize = 2000; int dimensions = 128; @@ -114,7 +113,19 @@ public static void beforeClass() throws Exception { writer.commit(); } - reader = DirectoryReader.open(directory); + // DirectoryReader resolves the stored codec name through SPI. Give that singleton the same + // deliberately small cache budget while it constructs this test's segment readers, then restore + // its previous format so the test does not alter later reader construction in this JVM. + CuVS2510GPUSearchCodec spiCodec = + (CuVS2510GPUSearchCodec) Codec.forName("CuVS2510GPUSearchCodec"); + KnnVectorsFormat previousFormat = spiCodec.knnVectorsFormat(); + spiCodec.setKnnFormat( + new CuVS2510GPUVectorsFormat(new GPUSearchParams.Builder().build(), TEST_CACHE_CONFIG)); + try { + reader = DirectoryReader.open(directory); + } finally { + spiCodec.setKnnFormat(previousFormat); + } searcher = new IndexSearcher(reader); queryVectors = generateDataset(random(), numQueries, dimensions); From 6d05bb8cd3de467c1056ac2f219bce12085f7214 Mon Sep 17 00:00:00 2001 From: James Xia Date: Fri, 24 Jul 2026 08:51:47 -0700 Subject: [PATCH 40/41] Prevent resource leaks from invalid workspace pool sizes Parse and validate workspace pool configuration before creating cuVS resources, close resources after initialization failures, and add GPU-free configuration tests. --- .../ThreadLocalCuVSResourcesProvider.java | 72 ++++++++++++++++--- .../TestThreadLocalCuVSResourcesProvider.java | 46 ++++++++++++ 2 files changed, 107 insertions(+), 11 deletions(-) create mode 100644 src/test/java/com/nvidia/cuvs/lucene/TestThreadLocalCuVSResourcesProvider.java diff --git a/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java b/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java index 4d98b61c..f2a14c36 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java +++ b/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java @@ -44,31 +44,81 @@ public static void setCuVSResourcesInstance(CuVSResources 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 final long RMM_ALIGNMENT_BYTES = 256; + private static CuVSResources cuVSResourcesOrNull() { + CuVSResources resources = null; try { - CuVSResources resources = CuVSResources.create(); - String poolSizeProp = System.getProperty(WORKSPACE_POOL_SIZE_PROPERTY); - if (poolSizeProp != null) { - long poolBytes = Long.parseLong(poolSizeProp); - if (poolBytes > 0) { - // RMM's pool_memory_resource requires the initial size to be a multiple of 256 bytes. - resources.setWorkspacePool((poolBytes + 255L) & ~255L); - } + // Resolve configuration before allocating resources so malformed input cannot leak a newly + // created native handle and pinned host buffer. + long poolBytes = + resolveWorkspacePoolBytes(System.getProperty(WORKSPACE_POOL_SIZE_PROPERTY)); + resources = CuVSResources.create(); + if (poolBytes > 0) { + resources.setWorkspacePool(poolBytes); } return resources; } catch (UnsupportedOperationException uoe) { + closeAfterFailedInitialization(resources, uoe); log.log( Level.WARNING, "cuVS is not supported on this platform or java version: " + uoe.getMessage()); } catch (Throwable t) { - if (t instanceof ExceptionInInitializerError ex) { - t = ex.getCause(); + Throwable failure = t; + if (t instanceof ExceptionInInitializerError ex && ex.getCause() != null) { + failure = ex.getCause(); } - log.log(Level.WARNING, "Exception occurred during creation of cuVS resources. " + t); + closeAfterFailedInitialization(resources, failure); + log.log(Level.WARNING, "Exception occurred during creation of cuVS resources. " + failure); } return null; } + /** + * Resolves a raw workspace-pool property value to a 256-byte-aligned size. Zero or an absent + * value disables the per-resources pool. Invalid, negative, or unalignable values warn and also + * disable it. + */ + static long resolveWorkspacePoolBytes(String raw) { + if (raw == null) return 0; + + final long requestedBytes; + try { + requestedBytes = Long.parseLong(raw.trim()); + } catch (NumberFormatException invalid) { + warnInvalidWorkspacePoolSize(raw); + return 0; + } + + if (requestedBytes == 0) return 0; + if (requestedBytes < 0 + || requestedBytes > Long.MAX_VALUE - (RMM_ALIGNMENT_BYTES - 1)) { + warnInvalidWorkspacePoolSize(raw); + return 0; + } + + return (requestedBytes + (RMM_ALIGNMENT_BYTES - 1)) & ~(RMM_ALIGNMENT_BYTES - 1); + } + + private static void warnInvalidWorkspacePoolSize(String raw) { + log.warning( + "Invalid " + + WORKSPACE_POOL_SIZE_PROPERTY + + " value \"" + + raw + + "\"; expected a non-negative byte count that can be aligned to 256 bytes. " + + "Continuing without a workspace pool."); + } + + private static void closeAfterFailedInitialization(CuVSResources resources, Throwable failure) { + if (resources == null) return; + try { + resources.close(); + } catch (Throwable closeFailure) { + failure.addSuppressed(closeFailure); + } + } + /** * Attempts to close the thread's {@link CuVSResources} instance. */ diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestThreadLocalCuVSResourcesProvider.java b/src/test/java/com/nvidia/cuvs/lucene/TestThreadLocalCuVSResourcesProvider.java new file mode 100644 index 00000000..ccf6b8f4 --- /dev/null +++ b/src/test/java/com/nvidia/cuvs/lucene/TestThreadLocalCuVSResourcesProvider.java @@ -0,0 +1,46 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 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 org.junit.Test; + +/** GPU-free tests for thread-local cuVS resource configuration. */ +public class TestThreadLocalCuVSResourcesProvider { + + @Test + public void absentOrZeroValueDisablesPool() { + assertEquals(0, ThreadLocalCuVSResourcesProvider.resolveWorkspacePoolBytes(null)); + assertEquals(0, ThreadLocalCuVSResourcesProvider.resolveWorkspacePoolBytes("0")); + assertEquals(0, ThreadLocalCuVSResourcesProvider.resolveWorkspacePoolBytes(" 0 ")); + } + + @Test + public void positiveValueIsAlignedUp() { + assertEquals(256, ThreadLocalCuVSResourcesProvider.resolveWorkspacePoolBytes("1")); + assertEquals(256, ThreadLocalCuVSResourcesProvider.resolveWorkspacePoolBytes("256")); + assertEquals(512, ThreadLocalCuVSResourcesProvider.resolveWorkspacePoolBytes("257")); + } + + @Test + public void malformedValueDisablesPool() { + assertEquals(0, ThreadLocalCuVSResourcesProvider.resolveWorkspacePoolBytes("")); + assertEquals(0, ThreadLocalCuVSResourcesProvider.resolveWorkspacePoolBytes("1024M")); + } + + @Test + public void negativeValueDisablesPool() { + assertEquals(0, ThreadLocalCuVSResourcesProvider.resolveWorkspacePoolBytes("-1")); + } + + @Test + public void roundingOverflowDisablesPool() { + assertEquals( + 0, + ThreadLocalCuVSResourcesProvider.resolveWorkspacePoolBytes( + Long.toString(Long.MAX_VALUE))); + } +} From 9c55185cb386016052989fe03647371fa139ae69 Mon Sep 17 00:00:00 2001 From: James Xia Date: Fri, 24 Jul 2026 11:18:03 -0700 Subject: [PATCH 41/41] Temporary change for CI --- RAPIDS_BRANCH | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/RAPIDS_BRANCH b/RAPIDS_BRANCH index ba2906d0..5595d1b3 100644 --- a/RAPIDS_BRANCH +++ b/RAPIDS_BRANCH @@ -1 +1 @@ -main +pull-request/2035