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/bench/pom.xml b/bench/pom.xml index 04ea66c5..4271d802 100644 --- a/bench/pom.xml +++ b/bench/pom.xml @@ -11,7 +11,7 @@ com.nvidia.cuvs.lucene.benchmarks cuvs-lucene-benchmarks - 26.08.0 + 26.08.0-SNAPSHOT jar cuvs-lucene-benchmarks @@ -85,7 +85,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..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 @@ -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 c7191234..fd277c3e 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 @@ -100,7 +100,7 @@ com.nvidia.cuvs.lucene cuvs-lucene - 26.08.0 + 26.08.0-SNAPSHOT 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..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; @@ -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=" 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 14390391..0433ecc4 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 @@ -68,6 +80,12 @@ false + + central-snapshots + https://central.sonatype.com/repository/maven-snapshots/ + true + false + @@ -129,10 +147,17 @@ com.nvidia.cuvs cuvs-java - 26.08.0 + 26.08.0-SNAPSHOT + + + central + https://central.sonatype.com/repository/maven-snapshots/ + + + @@ -158,6 +183,9 @@ true false + + ${project.basedir}/license-header.txt + @@ -219,6 +247,7 @@ jar + com.nvidia.cuvs.lucene all,-missing ${project.build.directory}/javadocs diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java index ccc61eae..1abb9e58 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; @@ -7,6 +7,7 @@ import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.assertIsSupported; import com.nvidia.cuvs.LibraryException; +import com.nvidia.cuvs.spi.CuVSProvider; import java.io.IOException; import org.apache.lucene.codecs.KnnVectorsFormat; import org.apache.lucene.codecs.KnnVectorsReader; @@ -39,6 +40,7 @@ public class CuVS2510GPUVectorsFormat extends KnnVectorsFormat { static { try { + CuVSProvider.provider().enableRMMAsyncMemory(); LUCENE_PROVIDER = LuceneProvider.getInstance("99"); FLAT_VECTORS_FORMAT = LUCENE_PROVIDER.getLuceneFlatVectorsFormatInstance(DefaultFlatVectorScorer.INSTANCE); diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java index c783660d..86955beb 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; @@ -443,6 +443,9 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits new CagraSearchParams.Builder() .withItopkSize(Math.max(collector.getiTopK(), topK)) .withSearchWidth(collector.getSearchWidth()) + .withThreadBlockSize(collector.getThreadBlockSize()) + .withMaxIterations(collector.getMaxIterations()) + .withAlgo(collector.getSearchAlgo()) .build(); } else { // Setting itopK as topK because in any case iTopK should be ATLEAST equal to topK @@ -600,6 +603,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/FilterBitsetCache.java b/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java new file mode 100644 index 00000000..3b269798 --- /dev/null +++ b/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java @@ -0,0 +1,64 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.FilterBitsetHandle; +import java.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/FilterCuVSProvider.java b/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java index 07e64fb1..b237d4bc 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java +++ b/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; @@ -9,19 +9,23 @@ import com.nvidia.cuvs.CagraIndexParams; import com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType; import com.nvidia.cuvs.CagraIndexParams.HnswHeuristicType; +import com.nvidia.cuvs.CagraQuery; import com.nvidia.cuvs.CuVSDeviceMatrix; import com.nvidia.cuvs.CuVSHostMatrix; import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.CuVSMatrix.Builder; import com.nvidia.cuvs.CuVSMatrix.DataType; import com.nvidia.cuvs.CuVSResources; +import com.nvidia.cuvs.FilterBitsetHandle; import com.nvidia.cuvs.GPUInfoProvider; import com.nvidia.cuvs.HnswIndex; import com.nvidia.cuvs.HnswIndexParams; +import com.nvidia.cuvs.MultiPartitionSearchResults; import com.nvidia.cuvs.TieredIndex; import com.nvidia.cuvs.spi.CuVSProvider; import java.lang.invoke.MethodHandle; import java.nio.file.Path; +import java.util.List; import java.util.logging.Level; class FilterCuVSProvider implements CuVSProvider { @@ -54,6 +58,22 @@ public CagraIndex.Builder newCagraIndexBuilder(CuVSResources cuVSResources) return delegate.newCagraIndexBuilder(cuVSResources); } + @Override + public FilterBitsetHandle newFilterBitsetHandle(long[] combinedLongs) { + return delegate.newFilterBitsetHandle(combinedLongs); + } + + @Override + public MultiPartitionSearchResults searchCagraMultiPartition( + CuVSResources resources, + List indices, + CagraQuery query, + int k, + FilterBitsetHandle filter) + throws Throwable { + return delegate.searchCagraMultiPartition(resources, indices, query, k, filter); + } + @Override public HnswIndex.Builder newHnswIndexBuilder(CuVSResources cuVSResources) throws UnsupportedOperationException { @@ -145,6 +165,11 @@ public HnswIndex hnswIndexFromCagra(HnswIndexParams arg0, CagraIndex arg1) throw return delegate.hnswIndexFromCagra(arg0, arg1); } + @Override + public void enableRMMAsyncMemory() { + delegate.enableRMMAsyncMemory(); + } + @Override public void enableRMMManagedPooledMemory(int arg0, int arg1) { delegate.enableRMMManagedPooledMemory(arg0, arg1); diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java index 6a9daae7..ee604449 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -1,20 +1,66 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.getCuVSResourcesInstance; + +import com.nvidia.cuvs.CagraIndex; +import com.nvidia.cuvs.CagraQuery; +import com.nvidia.cuvs.CagraSearchParams; +import com.nvidia.cuvs.CuVSMatrix; +import com.nvidia.cuvs.CuVSResources; +import com.nvidia.cuvs.FilterBitsetHandle; +import com.nvidia.cuvs.MultiPartitionCagraSearch; +import com.nvidia.cuvs.MultiPartitionSearchResults; import java.io.IOException; -import org.apache.lucene.index.LeafReader; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Comparator; +import java.util.List; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.index.CodecReader; +import org.apache.lucene.index.FilterLeafReader; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexReader; import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.search.DocIdSetIterator; +import org.apache.lucene.search.Explanation; +import org.apache.lucene.search.IndexSearcher; import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.MatchNoDocsQuery; import org.apache.lucene.search.Query; +import org.apache.lucene.search.QueryVisitor; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.ScoreMode; +import org.apache.lucene.search.Scorer; +import org.apache.lucene.search.ScorerSupplier; import org.apache.lucene.search.TopDocs; +import org.apache.lucene.search.Weight; import org.apache.lucene.search.knn.KnnCollectorManager; import org.apache.lucene.util.Bits; +import org.apache.lucene.util.FixedBitSet; /** - * Extends upon KnnFloatVectorQuery for GPU-only search. + * Extends {@link KnnFloatVectorQuery} for GPU-only search. + * + *

When all index segments use {@link CuVS2510GPUVectorsReader}, {@link #rewrite} delegates a + * single multi-partition search to cuVS, passing one Lucene segment per cuVS partition. cuVS + * runs the per-partition CAGRA searches, applies distance post-processing, and performs the + * cross-partition top-k merge internally; the returned arrays are mapped to Lucene doc IDs on + * the host. The effective CAGRA algorithm (SINGLE_CTA or MULTI_KERNEL) is selected by cuVS + * based on {@code searchAlgo} and {@code itopk_size}, with MULTI_KERNEL handling k beyond + * SINGLE_CTA's per-partition cap. + * + *

If the query has an explicit {@code filter}, or if any segment carries live-document deletes, + * the combined acceptance mask (filter ∩ liveDocs) is packed across all segments into a single + * {@link FilterBitsetHandle}. The host-side packed arrays are cached per unique + * (filter, reader-state, field) triple via {@link FilterBitsetCache}; the device upload is cached + * inside the handle itself across threads. + * + *

Falls back to the standard per-segment Lucene path when the optimized path cannot be + * applied (mixed segment types or missing CAGRA index for the field on any segment). * * @since 25.10 */ @@ -22,24 +68,168 @@ public class GPUKnnFloatVectorQuery extends KnnFloatVectorQuery { private final int iTopK; private final int searchWidth; + private final int threadBlockSize; + private final int maxIterations; + private final CagraSearchParams.SearchAlgo searchAlgo; /** - * Initializes {@link GPUKnnFloatVectorQuery} + * Initializes {@link GPUKnnFloatVectorQuery} with {@link CagraSearchParams.SearchAlgo#AUTO}, + * and max_iterations auto-selected (0). * - * @param field the vector field name - * @param target the vector target query - * @param k the topK value - * @param filter instance of the Query - * @param iTopK the iTopK value - * @param searchWidth the search width + * @param field the vector field name + * @param target the query vector + * @param k the number of nearest neighbors to return + * @param filter optional pre-filter query + * @param iTopK CAGRA itopk_size parameter + * @param searchWidth CAGRA search_width parameter */ public GPUKnnFloatVectorQuery( String field, float[] target, int k, Query filter, int iTopK, int searchWidth) { + this(field, target, k, filter, iTopK, searchWidth, 0, 0, CagraSearchParams.SearchAlgo.AUTO); + } + + /** + * Initializes {@link GPUKnnFloatVectorQuery}. + * + * @param field the vector field name + * @param target the query vector + * @param k the number of nearest neighbors to return + * @param filter optional pre-filter query + * @param iTopK CAGRA itopk_size parameter + * @param searchWidth CAGRA search_width parameter + * @param threadBlockSize CAGRA thread_block_size (0 = auto) + * @param maxIterations CAGRA max_iterations (0 = auto) + * @param searchAlgo CAGRA search algorithm + */ + public GPUKnnFloatVectorQuery( + String field, + float[] target, + int k, + Query filter, + int iTopK, + int searchWidth, + int threadBlockSize, + int maxIterations, + CagraSearchParams.SearchAlgo searchAlgo) { super(field, target, k, filter); this.iTopK = iTopK; this.searchWidth = searchWidth; + this.threadBlockSize = threadBlockSize; + this.maxIterations = maxIterations; + this.searchAlgo = searchAlgo; } + // ------------------------------------------------------------------------- + // Optimized multi-segment path + // ------------------------------------------------------------------------- + + @Override + public Query rewrite(IndexSearcher indexSearcher) throws IOException { + IndexReader reader = indexSearcher.getIndexReader(); + List leaves = reader.leaves(); + if (leaves.isEmpty()) { + return new MatchNoDocsQuery(); + } + + // Collect a CuVS2510GPUVectorsReader for every segment; fall back if any segment + // lacks one 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 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 (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()); + + try { + float[] target = getTargetCopy(); + CagraSearchParams searchParams = + new CagraSearchParams.Builder() + .withItopkSize(Math.max(iTopK, k)) + .withSearchWidth(searchWidth) + .withThreadBlockSize(threadBlockSize) + .withMaxIterations(maxIterations) + .withAlgo(searchAlgo) + .build(); + + // Upload the query vector to device once; the same matrix view is searched against every + // partition by the cuVS multi-partition API. + CuVSMatrix.Builder vectorBuilder = + CuVSMatrix.deviceBuilder(resources, 1, target.length, CuVSMatrix.DataType.FLOAT); + vectorBuilder.addVector(target); + + ScoreDoc[] scoreDocs; + try (CuVSMatrix queryVector = vectorBuilder.build()) { + for (int i = 0; i < leaves.size(); i++) { + cagraIndices.add(gpuReaders.get(i).getCagraIndexForField(field)); + } + + CagraQuery cagraQuery = + new CagraQuery.Builder(resources) + .withTopK(k) + .withSearchParams(searchParams) + .withQueryVectors(queryVector) + .build(); + + MultiPartitionSearchResults results = + MultiPartitionCagraSearch.search(resources, cagraIndices, cagraQuery, k, filterHandle); + + if (results.count() == 0) { + return new MatchNoDocsQuery(); + } + + // Map (segmentIdx, ordinal) → global Lucene doc ID; compute normalized score. + 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); + + 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 k > 1024 or not all GPU segments) + // ------------------------------------------------------------------------- + @Override protected TopDocs approximateSearch( LeafReaderContext context, @@ -47,12 +237,314 @@ protected TopDocs approximateSearch( int visitedLimit, KnnCollectorManager knnCollectorManager) throws IOException { - GPUPerLeafCuVSKnnCollector results = - new GPUPerLeafCuVSKnnCollector(k, visitedLimit, iTopK, searchWidth); - - LeafReader reader = context.reader(); - reader.searchNearestVectors(field, this.getTargetCopy(), results, acceptDocs); + new GPUPerLeafCuVSKnnCollector( + k, visitedLimit, iTopK, searchWidth, threadBlockSize, maxIterations, searchAlgo); + context.reader().searchNearestVectors(field, getTargetCopy(), results, acceptDocs); return results.topDocs(); } + + // ------------------------------------------------------------------------- + // Filter handle construction + // ------------------------------------------------------------------------- + + /** + * Returns a {@link FilterBitsetHandle} encoding ({@link #filter} ∩ liveDocs) for every segment, + * pulling from {@link FilterBitsetCache} when the reader state is unchanged. + * + *

Cache key uses per-reader keys (not just core keys) so that liveDocs changes — which happen + * when a reader is reopened after deletes — automatically invalidate the cached bitset. + */ + 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 (when set), intersects with liveDocs, and packs the + * result into a new {@link FilterBitsetHandle}. When {@link #filter} is {@code null}, the + * handle encodes liveDocs alone — this path is taken when one or more segments have deletes + * but no explicit Lucene filter was supplied. + */ + private FilterBitsetHandle buildFilterHandle( + IndexSearcher indexSearcher, + List leaves, + List gpuReaders) + throws IOException { + + Weight filterWeight = null; + if (filter != null) { + filterWeight = + indexSearcher.createWeight( + indexSearcher.rewrite(filter), ScoreMode.COMPLETE_NO_SCORES, 1.0f); + } + + int numSegments = leaves.size(); + long[] segBitOffsets = new long[numSegments]; + long totalBits = 0; + for (int i = 0; i < numSegments; i++) { + segBitOffsets[i] = totalBits; + int numOrds = gpuReaders.get(i).getFloatVectorValues(field).size(); + totalBits += ((long) (numOrds + 63) / 64) * 64; + } + long[] combinedLongs = new long[(int) (totalBits / 64)]; + + for (int i = 0; i < numSegments; i++) { + LeafReaderContext ctx = leaves.get(i); + Bits liveDocs = ctx.reader().getLiveDocs(); + // When filterWeight is null, accept all live documents (acceptDocs == liveDocs, which may + // itself be null to mean "all docs accepted" in this segment). + Bits acceptDocs = (filterWeight != null) ? evalFilter(filterWeight, ctx, liveDocs) : liveDocs; + FloatVectorValues fvv = gpuReaders.get(i).getFloatVectorValues(field); + Bits acceptedOrds = fvv.getAcceptOrds(acceptDocs); + int numOrds = fvv.size(); + int longOffset = (int) (segBitOffsets[i] / 64); + packOrdsToLongs(acceptedOrds, numOrds, combinedLongs, longOffset); + } + + // Only the combined bitset is transported; cuVS recomputes the per-partition bit offsets from + // the index sizes. segBitOffsets is used above solely to pack each segment's slice (64-bit + // word-aligned), matching that recomputation. + return FilterBitsetHandle.create(combinedLongs); + } + + /** + * Evaluates {@code filterWeight} in {@code ctx} and intersects with {@code liveDocs}. + * Returns a {@link Bits} over local doc IDs where {@code get(doc)} is true for accepted docs. + */ + private static Bits evalFilter(Weight filterWeight, LeafReaderContext ctx, Bits liveDocs) + throws IOException { + ScorerSupplier scorerSupplier = filterWeight.scorerSupplier(ctx); + if (scorerSupplier == null) { + 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 + // ------------------------------------------------------------------------- + + /** + * 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 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); + } + }; + } } diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java b/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java index 421b2baa..6413c221 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java @@ -1,9 +1,10 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; +import com.nvidia.cuvs.CagraSearchParams; import org.apache.lucene.search.TopKnnCollector; /** @@ -15,6 +16,9 @@ class GPUPerLeafCuVSKnnCollector extends TopKnnCollector { private int iTopK; private int searchWidth; + private int threadBlockSize; + private int maxIterations; + private CagraSearchParams.SearchAlgo searchAlgo; /** * Initializes {@link GPUPerLeafCuVSKnnCollector} @@ -22,11 +26,24 @@ class GPUPerLeafCuVSKnnCollector extends TopKnnCollector { * @param topK the topK value * @param iTopK the iTopK value * @param searchWidth the search width + * @param threadBlockSize CAGRA thread_block_size (0 = auto; controls worker_queue_size) + * @param maxIterations CAGRA max_iterations (0 = auto) + * @param searchAlgo the CAGRA search algorithm */ - public GPUPerLeafCuVSKnnCollector(int topK, int visitLimit, int iTopK, int searchWidth) { + public GPUPerLeafCuVSKnnCollector( + int topK, + int visitLimit, + int iTopK, + int searchWidth, + int threadBlockSize, + int maxIterations, + CagraSearchParams.SearchAlgo searchAlgo) { super(topK, visitLimit); this.iTopK = iTopK > topK ? iTopK : topK; this.searchWidth = searchWidth; + this.threadBlockSize = threadBlockSize; + this.maxIterations = maxIterations; + this.searchAlgo = searchAlgo; } public int getiTopK() { @@ -36,4 +53,16 @@ public int getiTopK() { public int getSearchWidth() { return searchWidth; } + + public int getThreadBlockSize() { + return threadBlockSize; + } + + public int getMaxIterations() { + return maxIterations; + } + + public CagraSearchParams.SearchAlgo getSearchAlgo() { + return searchAlgo; + } } diff --git a/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java b/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java index 9e259e27..e1bb46f6 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java +++ b/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java @@ -1,8 +1,7 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ - package com.nvidia.cuvs.lucene; import com.nvidia.cuvs.CuVSResources; @@ -42,9 +41,17 @@ public static void setCuVSResourcesInstance(CuVSResources resources) { cuVSResources.set(resources); } + /** System property controlling the workspace pool size per resources handle (in bytes). */ + public static final String WORKSPACE_POOL_SIZE_PROPERTY = "com.nvidia.cuvs.workspacePoolSize"; + private static CuVSResources cuVSResourcesOrNull() { try { - return CuVSResources.create(); + CuVSResources resources = CuVSResources.create(); + String poolSizeProp = System.getProperty(WORKSPACE_POOL_SIZE_PROPERTY); + if (poolSizeProp != null) { + resources.setWorkspacePool(Long.parseLong(poolSizeProp)); + } + return resources; } catch (UnsupportedOperationException uoe) { log.log( Level.WARNING,