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/RAPIDS_BRANCH b/RAPIDS_BRANCH index ba2906d0..5595d1b3 100644 --- a/RAPIDS_BRANCH +++ b/RAPIDS_BRANCH @@ -1 +1 @@ -main +pull-request/2035 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/pom.xml b/bench/pom.xml index d4bc7fbc..f5955873 100644 --- a/bench/pom.xml +++ b/bench/pom.xml @@ -11,7 +11,7 @@ com.nvidia.cuvs.lucene.benchmarks cuvs-lucene-benchmarks - 26.10.0 + 26.10.0-SNAPSHOT jar cuvs-lucene-benchmarks @@ -85,7 +85,7 @@ com.nvidia.cuvs.lucene cuvs-lucene - 26.10.0 + 26.10.0-SNAPSHOT commons-io 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/build.sh b/build.sh index 0ae9381a..e7bfe225 100755 --- a/build.sh +++ b/build.sh @@ -8,7 +8,7 @@ set -e -u -o pipefail ARGS="$*" NUMARGS=$# -VERSION="26.10.0" # 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 9b8fcaaf..50e06c3a 100644 --- a/examples/pom.xml +++ b/examples/pom.xml @@ -11,7 +11,7 @@ com.nvidia.cuvs.lucene.examples examples - 26.10.0 + 26.10.0-SNAPSHOT examples @@ -100,7 +100,7 @@ com.nvidia.cuvs.lucene cuvs-lucene - 26.10.0 + 26.10.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..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 @@ -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; @@ -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; @@ -30,9 +31,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; @@ -46,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); @@ -56,6 +63,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 +78,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 +114,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 +127,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 70f246d2..1676478f 100644 --- a/pom.xml +++ b/pom.xml @@ -12,7 +12,7 @@ com.nvidia.cuvs.lucene cuvs-lucene - 26.10.0 + 26.10.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.10.0 + 26.10.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/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 ccc61eae..3e62bd71 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; @@ -35,7 +35,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 { @@ -53,7 +54,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); } /** @@ -63,8 +64,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); } /** @@ -83,7 +97,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 c783660d..4aa8a82a 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; @@ -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 = @@ -430,7 +446,10 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits } } topK = Math.min(knnCollector.k() + 10, mask[0].cardinality()); - maskLength = mask[0].length(); + // numDocs must be the total vector count so cuVS sizes the prefilter to cover every ordinal. + // BitSet.length() is (highest set bit + 1), which under a selective filter is smaller than the + // vector count, leaving the trailing ordinals outside the filter and thus default-accepted. + maskLength = acceptedOrds.length(); } try { @@ -443,6 +462,9 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits new CagraSearchParams.Builder() .withItopkSize(Math.max(collector.getiTopK(), topK)) .withSearchWidth(collector.getSearchWidth()) + .withThreadBlockSize(collector.getThreadBlockSize()) + .withMaxIterations(collector.getMaxIterations()) + .withAlgo(collector.getSearchAlgo()) .build(); } else { // Setting itopK as topK because in any case iTopK should be ATLEAST equal to topK @@ -509,7 +531,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 || ord < 0) { continue; } float correctedScore = scoreCorrectionFunction.apply(score); @@ -600,6 +626,27 @@ private static void checkVersion(int versionMeta, int versionVectorData, IndexIn } } + /** + * Returns the {@link CagraIndex} for the given field, or {@code null} if unavailable + * (e.g., during a merge or when the field is missing). + * + * @param field the vector field name + * @return the CAGRA index, or {@code null} + */ + public CagraIndex getCagraIndexForField(String field) { + if (cuvsIndices == null) return null; + FieldInfo info = fieldInfos.fieldInfo(field); + if (info == null) return null; + GPUIndex gpuIndex = cuvsIndices.get(info.number); + if (gpuIndex == null) return null; + return gpuIndex.getCagraIndex(); + } + + /** 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 new file mode 100644 index 00000000..af30ebef --- /dev/null +++ b/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCache.java @@ -0,0 +1,318 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.FilterBitsetHandle; +import java.io.IOException; +import java.util.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 java.util.logging.Logger; +import org.apache.lucene.index.IndexReader; +import org.apache.lucene.search.Query; + +/** + * 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 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. 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

+ * + *

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 Logger LOG = Logger.getLogger(FilterBitsetCache.class.getName()); + + /** + * 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 + interface FilterBuilder { + FilterBitsetHandle build() throws IOException; + } + + 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(segReaderKey, other.segReaderKey) + && Objects.equals(field, other.field); + } + + @Override + public int hashCode() { + return Objects.hash(filter, segReaderKey, field); + } + } + + /** A cached future together with the byte size charged to the budget for its entry. */ + private record CacheEntry(CompletableFuture future, long bytes) {} + + 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 final Set registered = new HashSet<>(); + + // ---- Configuration and accounting (guarded by this) ---- + + private final boolean enabled; + + /** Fixed byte budget for this cache instance. */ + private final long maxBytes; + + private boolean workingSetWarned; + private boolean oversizedWarned; + + /** Total bytes charged to the budget across all live cache entries. */ + private long totalBytes; + + FilterBitsetCache(FilterBitsetCacheConfig config) { + Objects.requireNonNull(config, "config"); + enabled = config.enabled(); + maxBytes = config.maxBytes(); + } + + /** Whether caching is enabled. When disabled, callers build uncached, caller-owned handles. */ + synchronized boolean isEnabled() { + return enabled; + } + + /** + * 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}. + * + *

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 (maxBytes < workingSetBytes && !workingSetWarned) { + workingSetWarned = 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) + */ + 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 (this) { + 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 (this) { + CacheEntry entry = cache.get(key); + if (entry == null) { + future = new CompletableFuture<>(); + 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); + } + } + + 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 (this) { + 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; + 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. + } + } + + /** + * 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 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 — + * 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. + */ + void ensureCloseListener(IndexReader.CacheHelper helper) { + Object segReaderKey = helper.getKey(); + synchronized (this) { + if (!registered.add(segReaderKey)) return; // a listener is already registered for this reader + } + helper.addClosedListener(this::invalidateReader); + } + + /** Evicts and releases every cache entry belonging to the given segment reader key. */ + void invalidateReader(Object segReaderKey) { + List> toRelease = new ArrayList<>(); + 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)) { + CacheEntry ce = e.getValue(); + totalBytes -= ce.bytes(); + toRelease.add(ce.future()); + it.remove(); + } + } + } + for (CompletableFuture future : toRelease) { + releaseCacheRef(future); + } + } + + /** Evicts and releases every cache entry, resetting the byte budget accounting. */ + void clear() { + List> toRelease = new ArrayList<>(); + synchronized (this) { + 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. */ + 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( + (handle, err) -> { + if (handle != null) handle.close(); + }); + } +} 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..576678c2 --- /dev/null +++ b/src/main/java/com/nvidia/cuvs/lucene/FilterBitsetCacheConfig.java @@ -0,0 +1,39 @@ +/* + * 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. + * + *

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; 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, 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/FilterCuVSProvider.java b/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java index 07e64fb1..c4449b79 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, + List filters) + throws Throwable { + return delegate.searchCagraMultiPartition(resources, indices, query, k, filters); + } + @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..9e2ca291 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUKnnFloatVectorQuery.java @@ -1,20 +1,71 @@ /* - * 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.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; +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 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 + * whose built CAGRA graphs differ in degree (a single multi-partition request requires a uniform + * graph degree, and a small segment can have its degree truncated at build time). * * @since 25.10 */ @@ -22,24 +73,197 @@ public class GPUKnnFloatVectorQuery extends KnnFloatVectorQuery { private final int iTopK; private final int searchWidth; + private final int threadBlockSize; + private final int maxIterations; + private final CagraSearchParams.SearchAlgo searchAlgo; /** - * Initializes {@link GPUKnnFloatVectorQuery} + * Initializes {@link GPUKnnFloatVectorQuery} with {@link CagraSearchParams.SearchAlgo#AUTO}, + * and max_iterations auto-selected (0). * - * @param field the vector field name - * @param target the vector target query - * @param k the topK value - * @param filter instance of the Query - * @param iTopK the iTopK value - * @param searchWidth the search width + * @param field the vector field name + * @param target the query vector + * @param k the number of nearest neighbors to return + * @param filter optional pre-filter query + * @param iTopK CAGRA itopk_size parameter + * @param searchWidth CAGRA search_width parameter */ public GPUKnnFloatVectorQuery( String field, float[] target, int k, Query filter, int iTopK, int searchWidth) { + this(field, target, k, filter, iTopK, searchWidth, 0, 0, CagraSearchParams.SearchAlgo.AUTO); + } + + /** + * Initializes {@link GPUKnnFloatVectorQuery}. + * + * @param field the vector field name + * @param target the query vector + * @param k the number of nearest neighbors to return + * @param filter optional pre-filter query + * @param iTopK CAGRA itopk_size parameter + * @param searchWidth CAGRA search_width parameter + * @param threadBlockSize CAGRA thread_block_size (0 = auto) + * @param maxIterations CAGRA max_iterations (0 = auto) + * @param searchAlgo CAGRA search algorithm + */ + public GPUKnnFloatVectorQuery( + String field, + float[] target, + int k, + Query filter, + int iTopK, + int searchWidth, + int threadBlockSize, + int maxIterations, + CagraSearchParams.SearchAlgo searchAlgo) { super(field, target, k, filter); this.iTopK = iTopK; this.searchWidth = searchWidth; + this.threadBlockSize = threadBlockSize; + this.maxIterations = maxIterations; + this.searchAlgo = searchAlgo; } + // ------------------------------------------------------------------------- + // Optimized multi-segment path + // ------------------------------------------------------------------------- + + @Override + public Query rewrite(IndexSearcher indexSearcher) throws IOException { + IndexReader reader = indexSearcher.getIndexReader(); + List leaves = reader.leaves(); + if (leaves.isEmpty()) { + return new MatchNoDocsQuery(); + } + + // Collect a CuVS2510GPUVectorsReader for every segment; fall back if any segment lacks one, + // has no CAGRA index for this field, or has a CAGRA graph whose degree differs from the other + // segments'. The multi-partition search requires every partition to share one graph degree + // (a small segment can have its degree truncated at build time), so a non-uniform set is + // routed to the per-segment path instead. + List gpuReaders = new ArrayList<>(leaves.size()); + long commonGraphDegree = -1; + for (LeafReaderContext ctx : leaves) { + CuVS2510GPUVectorsReader gpuReader = unwrapGpuReader(ctx, field); + if (gpuReader == null) { + return super.rewrite(indexSearcher); + } + CagraIndex cagraIndex = gpuReader.getCagraIndexForField(field); + if (cagraIndex == null) { + return super.rewrite(indexSearcher); + } + long graphDegree = cagraIndex.getGraphDegree(); + if (commonGraphDegree == -1) { + commonGraphDegree = graphDegree; + } else if (graphDegree != commonGraphDegree) { + return super.rewrite(indexSearcher); + } + gpuReaders.add(gpuReader); + } + + // Build 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) { + if (ctx.reader().getLiveDocs() != null) { + hasDeletes = true; + break; + } + } + // Build the per-segment CagraIndex list (one entry per Lucene segment / cuVS partition). + CuVSResources resources = getCuVSResourcesInstance(); + List cagraIndices = new ArrayList<>(leaves.size()); + + List filterHandles = null; + try { + if (hasExplicitFilter || hasDeletes) { + filterHandles = buildPerSegmentFilterHandles(indexSearcher, leaves, gpuReaders); + } + + float[] target = getTargetCopy(); + CagraSearchParams searchParams = + new CagraSearchParams.Builder() + .withItopkSize(Math.max(iTopK, k)) + .withSearchWidth(searchWidth) + .withThreadBlockSize(threadBlockSize) + .withMaxIterations(maxIterations) + .withAlgo(searchAlgo) + .build(); + + // Upload the query vector to device once; the same matrix view is searched against every + // partition by the cuVS multi-partition API. + CuVSMatrix.Builder vectorBuilder = + CuVSMatrix.deviceBuilder(resources, 1, target.length, CuVSMatrix.DataType.FLOAT); + vectorBuilder.addVector(target); + + ScoreDoc[] scoreDocs; + try (CuVSMatrix queryVector = vectorBuilder.build()) { + for (int i = 0; i < leaves.size(); i++) { + cagraIndices.add(gpuReaders.get(i).getCagraIndexForField(field)); + } + + CagraQuery cagraQuery = + new CagraQuery.Builder(resources) + .withTopK(k) + .withSearchParams(searchParams) + .withQueryVectors(queryVector) + .build(); + + MultiPartitionSearchResults results = + MultiPartitionCagraSearch.search(resources, cagraIndices, cagraQuery, k, filterHandles); + + if (results.count() == 0) { + return new MatchNoDocsQuery(); + } + + // Map (segmentIdx, ordinal) → global Lucene doc ID; compute normalized score. + // Stash FloatVectorValues per segment: results.count() can be far larger than the + // number of segments, and getFloatVectorValues() allocates on each call. + FloatVectorValues[] segVectorValues = new FloatVectorValues[leaves.size()]; + scoreDocs = new ScoreDoc[results.count()]; + for (int j = 0; j < results.count(); j++) { + int segIdx = results.getPartitionIndex(j); + int ordinal = results.getOrdinal(j); + float dist = results.getDistance(j); + + FloatVectorValues fvv = segVectorValues[segIdx]; + if (fvv == null) { + fvv = gpuReaders.get(segIdx).getFloatVectorValues(field); + segVectorValues[segIdx] = fvv; + } + LeafReaderContext ctx = leaves.get(segIdx); + int globalDoc = ctx.docBase + fvv.ordToDoc(ordinal); + float score = 1.0f / (1.0f + dist); + scoreDocs[j] = new ScoreDoc(globalDoc, score); + } + + Arrays.sort(scoreDocs, Comparator.comparingDouble((ScoreDoc sd) -> sd.score).reversed()); + + return docAndScoreQuery(scoreDocs); + } + + } catch (Throwable t) { + Utils.handleThrowable(t); + throw new AssertionError("handleThrowable always throws"); // unreachable + } finally { + // Release this query's reference 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(); + } + } + } + } + + // ------------------------------------------------------------------------- + // Per-segment fallback path (used when k > 1024 or not all GPU segments) + // ------------------------------------------------------------------------- + @Override protected TopDocs approximateSearch( LeafReaderContext context, @@ -47,12 +271,354 @@ 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 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). + * + *

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. + * + *

Each non-null handle carries a reference the caller must release with {@link + * FilterBitsetHandle#decRef()} once the search completes. + */ + private List buildPerSegmentFilterHandles( + 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); + } + final Weight sharedFilterWeight = filterWeight; + + final int n = leaves.size(); + boolean[] needsFilter = new boolean[n]; + FloatVectorValues[] fvvs = new FloatVectorValues[n]; + long[] entryBytes = new long[n]; + 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. + if (filter == null && ctx.reader().getLiveDocs() == null) { + 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; + if (cache.isEnabled()) { + workingSetBytesByCache.merge(cache, segBytes, Long::sum); + } + } + + // 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()); + } + + List handles = new ArrayList<>(n); + 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))); + } + } + 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; + } + } + + /** + * 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, 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; + 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); + } + + /** + * Evaluates {@code filterWeight} in {@code ctx} and intersects with {@code liveDocs}. + * Returns a {@link Bits} over local doc IDs where {@code get(doc)} is true for accepted docs. + */ + private static Bits evalFilter(Weight filterWeight, LeafReaderContext ctx, Bits liveDocs) + throws IOException { + ScorerSupplier scorerSupplier = filterWeight.scorerSupplier(ctx); + if (scorerSupplier == null) { + // No scorer: the filter matches no documents in this segment. + return new Bits.MatchNoBits(ctx.reader().maxDoc()); + } + + int maxDoc = ctx.reader().maxDoc(); + FixedBitSet filterBits = new FixedBitSet(maxDoc); + Scorer scorer = scorerSupplier.get(Long.MAX_VALUE); + DocIdSetIterator it = scorer.iterator(); + int doc; + while ((doc = it.nextDoc()) != DocIdSetIterator.NO_MORE_DOCS) { + filterBits.set(doc); + } + + if (liveDocs == null) return filterBits; + + // Intersect: accept only docs that pass both filter and liveDocs. + return new Bits() { + @Override + public boolean get(int i) { + return filterBits.get(i) && liveDocs.get(i); + } + + @Override + public int length() { + return maxDoc; + } + }; + } + + /** + * Packs {@code numOrds} ordinal bits from {@code bits} into {@code dest} starting at long index + * {@code destLongOffset}. {@code bits == null} means all ordinals accepted (Lucene convention). + */ + private static void packOrdsToLongs(Bits bits, int numOrds, long[] dest, int destLongOffset) { + if (bits == null) { + // All ordinals accepted: fill with all-ones, masking the last partial word. + int numLongs = (numOrds + 63) / 64; + Arrays.fill(dest, destLongOffset, destLongOffset + numLongs, -1L); + int tail = numOrds % 64; + if (tail != 0) { + dest[destLongOffset + numLongs - 1] = (1L << tail) - 1L; + } + return; + } + for (int i = 0; i < numOrds; i++) { + if (bits.get(i)) { + dest[destLongOffset + i / 64] |= (1L << (i % 64)); + } + } + } + + // ------------------------------------------------------------------------- + // Helpers + // ------------------------------------------------------------------------- + + /** + * Unwraps the {@link LeafReaderContext}'s reader to a {@link CuVS2510GPUVectorsReader}, or + * returns {@code null} if the reader is not of that type. + */ + private static CuVS2510GPUVectorsReader unwrapGpuReader(LeafReaderContext ctx, String field) { + var unwrapped = FilterLeafReader.unwrap(ctx.reader()); + if (!(unwrapped instanceof CodecReader)) return null; + KnnVectorsReader vr = ((CodecReader) unwrapped).getVectorReader(); + // A per-field codec wraps the vectors reader; unwrap to this field's delegate so the + // multi-partition GPU path is not silently bypassed when the CuVS format is used via + // PerFieldKnnVectorsFormat. + if (vr instanceof PerFieldKnnVectorsFormat.FieldsReader perField) { + vr = perField.getFieldReader(field); + } + return (vr instanceof CuVS2510GPUVectorsReader gpuReader) ? gpuReader : null; + } + + /** + * Builds a {@link Query} that matches exactly the given pre-scored documents. + * + *

Partitions {@code scoreDocs} by segment (using {@link ScoreDoc#shardIndex} as the segment + * offset relative to {@link LeafReaderContext#docBase}), then returns a {@link Scorer} per + * segment that iterates those docs in ascending doc-ID order and replays their pre-computed + * scores. + */ + private static Query docAndScoreQuery(ScoreDoc[] scoreDocs) { + return new Query() { + @Override + public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) + throws IOException { + return new Weight(this) { + @Override + public ScorerSupplier scorerSupplier(LeafReaderContext ctx) { + int base = ctx.docBase; + int maxDoc = base + ctx.reader().maxDoc(); + // Collect docs belonging to this segment; re-sort by local doc ID ascending. + int[] localDocs = new int[scoreDocs.length]; + float[] localScores = new float[scoreDocs.length]; + int count = 0; + float max = Float.NEGATIVE_INFINITY; + for (ScoreDoc sd : scoreDocs) { + if (sd.doc >= base && sd.doc < maxDoc) { + localDocs[count] = sd.doc - base; + localScores[count] = sd.score * boost; + if (localScores[count] > max) max = localScores[count]; + count++; + } + } + if (count == 0) return null; + final int n = count; + final float maxScore = max; + Integer[] idx = new Integer[n]; + for (int i = 0; i < n; i++) idx[i] = i; + Arrays.sort(idx, Comparator.comparingInt(i -> localDocs[i])); + final int[] sortedDocs = new int[n]; + final float[] sortedScores = new float[n]; + for (int i = 0; i < n; i++) { + sortedDocs[i] = localDocs[idx[i]]; + sortedScores[i] = localScores[idx[i]]; + } + + return new ScorerSupplier() { + @Override + public Scorer get(long leadCost) { + return new Scorer() { + private int pos = -1; + + @Override + public DocIdSetIterator iterator() { + return new DocIdSetIterator() { + @Override + public int docID() { + return pos < 0 ? -1 : (pos >= n ? NO_MORE_DOCS : sortedDocs[pos]); + } + + @Override + public int nextDoc() { + pos++; + return docID(); + } + + @Override + public int advance(int target) { + if (pos < 0) pos = 0; + while (pos < n && sortedDocs[pos] < target) pos++; + return docID(); + } + + @Override + public long cost() { + return n; + } + }; + } + + @Override + public float getMaxScore(int upTo) { + // The precomputed segment max is a valid upper bound for any upTo. + return maxScore; + } + + @Override + public float score() { + return sortedScores[pos]; + } + + @Override + public int docID() { + return pos < 0 + ? -1 + : (pos >= n ? DocIdSetIterator.NO_MORE_DOCS : sortedDocs[pos]); + } + }; + } + + @Override + public long cost() { + return n; + } + }; + } + + @Override + public boolean isCacheable(LeafReaderContext ctx) { + return false; + } + + @Override + public Explanation explain(LeafReaderContext ctx, int doc) { + for (ScoreDoc sd : scoreDocs) { + if (sd.doc == ctx.docBase + doc) { + return Explanation.match( + sd.score * boost, + "GPU multi-segment CAGRA search, product of:", + Explanation.match(sd.score, "raw CAGRA score, 1 / (1 + distance)"), + Explanation.match(boost, "boost")); + } + } + return Explanation.noMatch("not a GPU search result"); + } + }; + } + + @Override + public String toString(String field) { + return "GPUDocAndScoreQuery"; + } + + @Override + public void visit(QueryVisitor visitor) {} + + @Override + public boolean equals(Object o) { + return this == o; + } + + @Override + public int hashCode() { + return System.identityHashCode(this); + } + }; + } } diff --git a/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java b/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java index 421b2baa..6413c221 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java +++ b/src/main/java/com/nvidia/cuvs/lucene/GPUPerLeafCuVSKnnCollector.java @@ -1,9 +1,10 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.lucene; +import com.nvidia.cuvs.CagraSearchParams; import org.apache.lucene.search.TopKnnCollector; /** @@ -15,6 +16,9 @@ class GPUPerLeafCuVSKnnCollector extends TopKnnCollector { private int iTopK; private int searchWidth; + private int threadBlockSize; + private int maxIterations; + private CagraSearchParams.SearchAlgo searchAlgo; /** * Initializes {@link GPUPerLeafCuVSKnnCollector} @@ -22,11 +26,24 @@ class GPUPerLeafCuVSKnnCollector extends TopKnnCollector { * @param topK the topK value * @param iTopK the iTopK value * @param searchWidth the search width + * @param threadBlockSize CAGRA thread_block_size (0 = auto; controls worker_queue_size) + * @param maxIterations CAGRA max_iterations (0 = auto) + * @param searchAlgo the CAGRA search algorithm */ - public GPUPerLeafCuVSKnnCollector(int topK, int visitLimit, int iTopK, int searchWidth) { + public GPUPerLeafCuVSKnnCollector( + int topK, + int visitLimit, + int iTopK, + int searchWidth, + int threadBlockSize, + int maxIterations, + CagraSearchParams.SearchAlgo searchAlgo) { super(topK, visitLimit); this.iTopK = iTopK > topK ? iTopK : topK; this.searchWidth = searchWidth; + this.threadBlockSize = threadBlockSize; + this.maxIterations = maxIterations; + this.searchAlgo = searchAlgo; } public int getiTopK() { @@ -36,4 +53,16 @@ public int getiTopK() { public int getSearchWidth() { return searchWidth; } + + public int getThreadBlockSize() { + return threadBlockSize; + } + + public int getMaxIterations() { + return maxIterations; + } + + public CagraSearchParams.SearchAlgo getSearchAlgo() { + return searchAlgo; + } } diff --git a/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java b/src/main/java/com/nvidia/cuvs/lucene/ThreadLocalCuVSResourcesProvider.java index 9e259e27..f2a14c36 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,22 +41,84 @@ 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 final long RMM_ALIGNMENT_BYTES = 256; + private static CuVSResources cuVSResourcesOrNull() { + CuVSResources resources = null; try { - return CuVSResources.create(); + // 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/TestFilterBitsetCache.java b/src/test/java/com/nvidia/cuvs/lucene/TestFilterBitsetCache.java new file mode 100644 index 00000000..686aacc4 --- /dev/null +++ b/src/test/java/com/nvidia/cuvs/lucene/TestFilterBitsetCache.java @@ -0,0 +1,425 @@ +/* + * 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.After; +import org.junit.Before; +import org.junit.Test; + +/** + * 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. + */ +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; + + private FilterBitsetCache cache; + + @Before + public void resetCache() { + cache = new FilterBitsetCache(new FilterBitsetCacheConfig(true, HUGE_BUDGET)); + } + + @After + public void restoreDefaults() { + cache.clear(); + cache = null; + } + + /** + * 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 Object segKey(String s) { + return 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 cache.acquire( + null, + segKey(field), + field, + ENTRY_BYTES, + () -> { + 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()); + } + + /** Once the byte budget is exceeded, the eldest entry is evicted and freed; newer ones stay. */ + @Test + public void byteCapEvictsEldestWhenBudgetExceeded() throws Exception { + // Budget holds two 64-byte entries; inserting a third must evict exactly the oldest. + cache = new FilterBitsetCache(new FilterBitsetCacheConfig(true, 2 * ENTRY_BYTES)); + + List handles = new ArrayList<>(); + for (int i = 0; i < 3; i++) { + CountingHandle h = new CountingHandle(); + handles.add(h); + final String field = "bcap-" + i; + FilterBitsetHandle got = + cache.acquire(null, segKey(field), field, ENTRY_BYTES, () -> h); + assertSame(h, got); + got.decRef(); // caller finished; the cache keeps its reference + } + + 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", + cache.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; + cache.acquire(null, segKey(field), field, ENTRY_BYTES, () -> h).decRef(); + } + assertTrue("entries charged to the budget", cache.currentBytesForTests() > 0); + + 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, cache.currentBytesForTests()); + } + + /** An entry larger than the whole budget is never cached; each acquire rebuilds it uncached. */ + @Test + public void oversizedEntryIsNotCached() throws Exception { + 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 = + cache.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, cache.currentBytesForTests()); + + CountingHandle h2 = new CountingHandle(); + FilterBitsetHandle got2 = + cache.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 public default configuration has a deterministic 128 MiB budget. */ + @Test + public void defaultConfiguration() { + assertTrue(FilterBitsetCacheConfig.DEFAULT.enabled()); + 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. */ + @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 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. */ + @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. + cache.acquire( + null, + keyA, + "f", + ENTRY_BYTES, + () -> { + buildA.incrementAndGet(); + return hA; + }) + .decRef(); + cache.acquire(null, keyB, "f", ENTRY_BYTES, () -> hB).decRef(); + + // Invalidate segment A only, as its reader's close listener would. + cache.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. + cache.acquire( + null, + keyA, + "f", + ENTRY_BYTES, + () -> { + 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 { + final String field = "failing"; + AtomicInteger buildCount = new AtomicInteger(); + + try { + cache.acquire( + null, + segKey(field), + field, + ENTRY_BYTES, + () -> { + buildCount.incrementAndGet(); + throw new IOException("boom"); + }); + fail("expected IOException to propagate"); + } catch (IOException expected) { + assertEquals("boom", expected.getMessage()); + } + assertEquals(1, buildCount.get()); + assertEquals("failed build must not charge the budget", 0, cache.currentBytesForTests()); + + CountingHandle handle = new CountingHandle(); + FilterBitsetHandle got = + cache.acquire( + null, + segKey(field), + field, + ENTRY_BYTES, + () -> { + 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. The budget is + * left huge (see {@link #resetCache}) so no eviction happens mid-run. + */ + @Test + public void concurrentAcquireReleaseIsBalanced() throws Exception { + final int keySpace = 8; + 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 = + cache.acquire( + null, segKey(field), field, ENTRY_BYTES, 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..0c76956c --- /dev/null +++ b/src/test/java/com/nvidia/cuvs/lucene/TestMultiSegmentGPUFilterConcurrency.java @@ -0,0 +1,217 @@ +/* + * 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 com.nvidia.cuvs.spi.CuVSProvider; +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.codecs.KnnVectorsFormat; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.NoMergePolicy; +import org.apache.lucene.index.Term; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TermQuery; +import org.apache.lucene.store.ByteBuffersDirectory; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; + +/** + * End-to-end concurrency test for the filter-bitset path of multi-segment GPU search. Many threads + * issue filtered {@link GPUKnnFloatVectorQuery} searches across a rotating set of distinct filters + * (more than 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 { + + // 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"; + 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; + private static IndexSearcher searcher; + private static float[][] queryVectors; + private static int topK; + private static List filters; + private static List> acceptedDocsPerFilter; + + @BeforeClass + public static void beforeClass() throws Exception { + try { + CuVSProvider.provider().enableRMMAsyncMemory(); + } catch (UnsupportedOperationException unsupported) { + assumeTrue("cuVS not supported: " + unsupported.getMessage(), false); + } + assumeTrue("cuVS not supported", isSupported()); + codec = + new CuVS2510GPUSearchCodec( + new GPUSearchParams.Builder().build(), TEST_CACHE_CONFIG); + + 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(); + } + + // 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); + + // One distinct filter per category, plus the set of global doc IDs each filter accepts. + filters = new ArrayList<>(NUM_CATEGORIES); + acceptedDocsPerFilter = new ArrayList<>(NUM_CATEGORIES); + for (int c = 0; c < NUM_CATEGORIES; c++) { + Query filter = new TermQuery(new Term(CATEGORY_FIELD, "c" + c)); + filters.add(filter); + Set accepted = new HashSet<>(); + for (ScoreDoc sd : searcher.search(filter, datasetSize).scoreDocs) { + accepted.add(sd.doc); + } + acceptedDocsPerFilter.add(accepted); + } + } + + @Test + public void testConcurrentFilteredSearchAcrossManyFilters() throws Exception { + final int numThreads = 6; + final int iterationsPerThread = 150; + + List errors = new CopyOnWriteArrayList<>(); + AtomicInteger nonEmptyResults = new AtomicInteger(); + + ExecutorService pool = Executors.newFixedThreadPool(numThreads); + CountDownLatch go = new CountDownLatch(1); + List> tasks = new ArrayList<>(); + try { + for (int t = 0; t < numThreads; t++) { + tasks.add( + pool.submit( + () -> { + ThreadLocalRandom rnd = ThreadLocalRandom.current(); + go.await(); + for (int i = 0; i < iterationsPerThread; i++) { + int c = rnd.nextInt(NUM_CATEGORIES); + Query filter = filters.get(c); + Set accepted = acceptedDocsPerFilter.get(c); + float[] q = queryVectors[rnd.nextInt(queryVectors.length)]; + + GPUKnnFloatVectorQuery query = + new GPUKnnFloatVectorQuery(VECTOR_FIELD, q, topK, filter, topK, 1); + ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; + if (hits.length > 0) { + nonEmptyResults.incrementAndGet(); + } + for (ScoreDoc hit : hits) { + if (!accepted.contains(hit.doc)) { + throw new AssertionError( + "hit " + hit.doc + " not accepted by filter c" + c); + } + } + } + return null; + })); + } + go.countDown(); + for (Future f : tasks) { + try { + f.get(120, TimeUnit.SECONDS); + } catch (Throwable t) { + errors.add(t); + } + } + } finally { + pool.shutdownNow(); + pool.awaitTermination(10, TimeUnit.SECONDS); + } + + if (!errors.isEmpty()) { + throw new AssertionError( + errors.size() + " search thread(s) failed; first: " + errors.get(0), errors.get(0)); + } + assertTrue("expected at least some non-empty filtered results", nonEmptyResults.get() > 0); + } + + @AfterClass + public static void afterClass() throws IOException { + if (reader != null) reader.close(); + if (directory != null) directory.close(); + reader = null; + directory = null; + searcher = null; + filters = null; + acceptedDocsPerFilter = null; + queryVectors = null; + } +} diff --git a/src/test/java/com/nvidia/cuvs/lucene/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()); diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestPerSegmentGPUFilterSearch.java b/src/test/java/com/nvidia/cuvs/lucene/TestPerSegmentGPUFilterSearch.java new file mode 100644 index 00000000..df2ba211 --- /dev/null +++ b/src/test/java/com/nvidia/cuvs/lucene/TestPerSegmentGPUFilterSearch.java @@ -0,0 +1,91 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; + +import org.apache.lucene.document.Document; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TopKnnCollector; +import org.apache.lucene.store.ByteBuffersDirectory; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.apache.lucene.util.FixedBitSet; +import org.junit.Test; + +/** + * Exercises the per-segment (single-index) GPU search with a selective prefilter. It drives {@link + * CuVS2510GPUVectorsReader#search} directly via {@code LeafReader.searchNearestVectors} with a + * synthesized {@code acceptDocs}, so the path is exercised regardless of any higher-level query + * routing (which now sends the multi-partition case elsewhere). + * + *

A highly selective {@code acceptDocs} rejects the trailing ordinals of the segment — the regime + * where the per-segment prefilter previously mis-sized itself (it passed {@code BitSet.length()}, + * the highest set bit + 1, instead of the vector count) and default-accepted the trailing ordinals, + * returning filtered-out vectors. Every returned hit must be accepted. + */ +@SuppressSysoutChecks(bugUrl = "") +public class TestPerSegmentGPUFilterSearch extends LuceneTestCase { + + private static final String VECTOR_FIELD = "vectors"; + private static final int NUM_CATEGORIES = 24; + + @Test + public void testSelectiveFilterDoesNotLeakTrailingOrdinals() throws Exception { + assumeTrue("cuVS not supported", isSupported()); + + // A segment size that is not a multiple of 32 exercises the last partial uint32 word of the + // prefilter bitset, which is where the trailing ordinals used to leak. + final int datasetSize = 500; + final int dimensions = 128; + final int topK = 10; + + try (Directory directory = newDirectory(new ByteBuffersDirectory())) { + float[][] dataset = generateDataset(random(), datasetSize, dimensions); + IndexWriterConfig config = new IndexWriterConfig().setCodec(new CuVS2510GPUSearchCodec()); + try (IndexWriter writer = new IndexWriter(directory, config)) { + for (int i = 0; i < datasetSize; i++) { + Document doc = new Document(); + doc.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); + writer.addDocument(doc); + } + writer.commit(); + } + + try (DirectoryReader reader = DirectoryReader.open(directory)) { + assertEquals("expected a single segment", 1, reader.leaves().size()); + LeafReader leaf = reader.leaves().get(0).reader(); + float[][] queries = generateDataset(random(), 32, dimensions); + + // Category c = ordinal % NUM_CATEGORIES; acceptDocs keeps ~1/24 of the ordinals, so the + // segment's trailing ordinals are rejected. (docId == vector ordinal for a dense segment.) + for (int c = 0; c < NUM_CATEGORIES; c++) { + FixedBitSet acceptDocs = new FixedBitSet(datasetSize); + for (int i = c; i < datasetSize; i += NUM_CATEGORIES) { + acceptDocs.set(i); + } + for (float[] q : queries) { + TopKnnCollector collector = new TopKnnCollector(topK, Integer.MAX_VALUE); + leaf.searchNearestVectors(VECTOR_FIELD, q, collector, acceptDocs); + for (ScoreDoc hit : collector.topDocs().scoreDocs) { + assertTrue( + "per-segment search returned doc " + hit.doc + " outside filter category " + c, + acceptDocs.get(hit.doc)); + } + } + } + } + } + } +} 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))); + } +}