From a305fe8cbd76d1e58d3046444517c1e2cbf5ee0b Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Thu, 1 Oct 2026 19:34:49 +0000 Subject: [PATCH 1/2] Fix scalar CAGRA graph byte encoding --- ...nvidia-cuvs-lucene-acceleratedhnswutils.md | 2 +- ...leratedhnswscalarquantizedvectorswriter.md | 10 +- .../cuvs/lucene/AcceleratedHNSWUtils.java | 10 +- ...ratedHNSWScalarQuantizedVectorsWriter.java | 31 +-- .../cuvs/lucene/TestScalarQuantization.java | 41 ++++ .../lucene/TestScalarQuantizedCagraGraph.java | 190 ++++++++++++++++++ 6 files changed, 248 insertions(+), 36 deletions(-) create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestScalarQuantization.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestScalarQuantizedCagraGraph.java diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md index 219149bc6d..22cb8786e9 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md @@ -160,7 +160,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWU public static List quantizeFloatVectorsToScalar(List floatVectors) ``` -Scalar quantization. +Scalar quantization to unsigned 7-bit values stored in Java bytes. **Parameters** diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md index 9c51cb26e9..e3e66b1e45 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md @@ -57,7 +57,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Build the indexes and writes it to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:241`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:222`_ ### mergeOneField @@ -67,7 +67,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:326`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:307`_ ### finish @@ -77,7 +77,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:357`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:338`_ ### close @@ -87,7 +87,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:377`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:358`_ ### ramBytesUsed @@ -97,6 +97,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:386`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:367`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:56`_ diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index 9c49c07fe0..27ff26e460 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -444,7 +444,7 @@ public static List quantizeFloatVectorsToBinary(List floatVecto } /** - * Scalar quantization. + * Scalar quantization to unsigned 7-bit values stored in Java bytes. * * @param floatVectors A list of float vectors * @return A list of byte scalar representation for the input vectors @@ -459,8 +459,8 @@ public static List quantizeFloatVectorsToScalar(List floatVecto float[] minPerDim = new float[dimensions]; float[] maxPerDim = new float[dimensions]; - Arrays.fill(minPerDim, Float.MAX_VALUE); - Arrays.fill(maxPerDim, Float.MIN_VALUE); + Arrays.fill(minPerDim, Float.POSITIVE_INFINITY); + Arrays.fill(maxPerDim, Float.NEGATIVE_INFINITY); for (float[] vector : floatVectors) { for (int d = 0; d < dimensions; d++) { @@ -476,8 +476,8 @@ public static List quantizeFloatVectorsToScalar(List floatVecto float range = maxPerDim[d] - minPerDim[d]; if (range > 0) { float normalized = (vector[d] - minPerDim[d]) / range; - int quantizedValue = Math.round(normalized * 127.0f) - 64; - quantized[d] = (byte) Math.max(-64, Math.min(63, quantizedValue)); + int quantizedValue = Math.round(normalized * 127.0f); + quantized[d] = (byte) Math.max(0, Math.min(127, quantizedValue)); } else { quantized[d] = 0; } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java index 7141af56ee..75b97b5d97 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java @@ -146,18 +146,6 @@ public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException return writer; } - private static byte signedToUnsignedByte(byte signedByte) { - return (byte) (signedByte & 0xFF); - } - - private static byte[] convertSignedToUnsigned(byte[] signedVector) { - byte[] unsignedVector = new byte[signedVector.length]; - for (int i = 0; i < signedVector.length; i++) { - unsignedVector[i] = signedToUnsignedByte(signedVector[i]); - } - return unsignedVector; - } - /** * Builds the intermediate CAGRA index and builds and writes the HNSW index. * @@ -165,7 +153,7 @@ private static byte[] convertSignedToUnsigned(byte[] signedVector) { * @param vectors quantized vectors * @throws IOException */ - private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOException { + private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOException { if (vectors.size() == 0) { writeEmpty(fieldInfo, hnswMeta); return; @@ -174,18 +162,11 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOE try { int dimensions = fieldInfo.getVectorDimension(); - // Convert 7-bit signed bytes to 8-bit unsigned bytes for cuVS compatibility - List unsignedVectors = new ArrayList<>(vectors.size()); - for (Object signedVector : vectors) { - unsignedVectors.add(convertSignedToUnsigned((byte[]) signedVector)); - } - - // Create CuVSMatrix with BYTE data type (unsigned bytes) - CuVSMatrix dataset = - Utils.createByteMatrix(unsignedVectors, dimensions, getCuVSResourcesInstance()); + // The scalar quantizer already emits nonnegative bytes for cuVS's unsigned BYTE type. + CuVSMatrix dataset = Utils.createByteMatrix(vectors, dimensions, getCuVSResourcesInstance()); if (dataset.size() < 2) { - writeSingleVectorGraph(fieldInfo, unsignedVectors); + writeSingleVectorGraph(fieldInfo, vectors); return; } @@ -207,7 +188,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOE size, dimensions, adjacencyListMatrix, - unsignedVectors, + vectors, acceleratedHNSWParams.getHnswLayers(), params, QuantizationType.SCALAR); @@ -287,7 +268,7 @@ private void writeSortingField(FieldWriter fieldData, Sorter.DocMap sortMap) thr * Builds and writes a single vector graph. * * @param fieldInfo instance of FieldInfo - * @param vectors the list of scalar quantized vectors (already converted to unsigned) + * @param vectors scalar-quantized nonnegative byte vectors * @throws IOException I/O Exceptions */ private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestScalarQuantization.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestScalarQuantization.java new file mode 100644 index 0000000000..48378e865b --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestScalarQuantization.java @@ -0,0 +1,41 @@ +/* + * 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.assertArrayEquals; + +import java.util.List; +import org.junit.Test; + +public class TestScalarQuantization { + + @Test + public void quantizesNegativeOnlyDimensionsAcrossTheFullRange() { + List quantized = + AcceleratedHNSWUtils.quantizeFloatVectorsToScalar( + List.of(new float[] {-10.0f, -4.0f}, new float[] {-5.0f, -2.0f})); + + assertArrayEquals(new byte[] {0, 0}, quantized.get(0)); + assertArrayEquals(new byte[] {127, 127}, quantized.get(1)); + } + + @Test + public void quantizesIntermediateValuesMonotonically() { + List quantized = + AcceleratedHNSWUtils.quantizeFloatVectorsToScalar( + List.of( + new float[] {-1.0f}, + new float[] {-0.5f}, + new float[] {0.0f}, + new float[] {0.5f}, + new float[] {1.0f})); + + assertArrayEquals(new byte[] {0}, quantized.get(0)); + assertArrayEquals(new byte[] {32}, quantized.get(1)); + assertArrayEquals(new byte[] {64}, quantized.get(2)); + assertArrayEquals(new byte[] {95}, quantized.get(3)); + assertArrayEquals(new byte[] {127}, quantized.get(4)); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestScalarQuantizedCagraGraph.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestScalarQuantizedCagraGraph.java new file mode 100644 index 0000000000..8e4a8357f6 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestScalarQuantizedCagraGraph.java @@ -0,0 +1,190 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.assertIsSupported; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; +import static org.junit.Assume.assumeTrue; + +import java.util.ArrayList; +import java.util.Comparator; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.stream.IntStream; +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.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.TestUtil; +import org.apache.lucene.util.InfoStream; +import org.junit.Test; + +/** Checks the GPU-built scalar graph with CPU HNSW search against exact Euclidean neighbors. */ +public class TestScalarQuantizedCagraGraph extends LuceneTestCase { + private static final String VECTOR_FIELD = "vector"; + private static final String ID_FIELD = "id"; + private static final int VECTOR_COUNT = 512; + private static final int DIMENSIONS = 128; + private static final int TOP_K = 10; + private static final int MIN_EXACT_NEIGHBORS = 8; + + @Test + public void testGpuBuiltScalarHnswRetainsRecallAcrossMerge() throws Exception { + requireGpuWhenSelected(); + + float[][] vectors = vectorsWithNegativeAndMixedSignDimensions(); + RecordingInfoStream buildLog = new RecordingInfoStream(); + AcceleratedHNSWParams params = + new AcceleratedHNSWParams.Builder() + .withStrategy(AcceleratedHNSWParams.Strategy.CUSTOM) + .withGraphDegree(32) + .withIntermediateGraphDegree(64) + .withHNSWLayer(1) + .build(); + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec( + TestUtil.alwaysKnnVectorsFormat( + new LuceneAcceleratedHNSWScalarQuantizedVectorsFormat(params))) + .setMaxBufferedDocs(VECTOR_COUNT + 1) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) + .setInfoStream(buildLog); + + try (Directory directory = newDirectory()) { + try (IndexWriter writer = new IndexWriter(directory, config)) { + for (int id = 0; id < vectors.length; id++) { + Document document = new Document(); + document.add(new StringField(ID_FIELD, Integer.toString(id), Field.Store.YES)); + document.add(new KnnFloatVectorField(VECTOR_FIELD, vectors[id], EUCLIDEAN)); + writer.addDocument(document); + if (id == VECTOR_COUNT / 2 - 1) { + writer.commit(); + } + } + writer.commit(); + + try (DirectoryReader reader = DirectoryReader.open(directory)) { + assertEquals(2, reader.leaves().size()); + assertHnswRecallAgainstExactNeighbors(reader, vectors); + } + + long gpuBuildsBeforeMerge = buildLog.gpuWriterOpenCount(); + assertTrue("No scalar GPU writer opened: " + buildLog.messages, gpuBuildsBeforeMerge >= 2); + + writer.forceMerge(1); + writer.commit(); + assertTrue( + "The merge did not open a scalar GPU writer: " + buildLog.messages, + buildLog.gpuWriterOpenCount() > gpuBuildsBeforeMerge); + } + + try (DirectoryReader reader = DirectoryReader.open(directory)) { + assertEquals(1, reader.leaves().size()); + assertHnswRecallAgainstExactNeighbors(reader, vectors); + } + } + } + + private static void requireGpuWhenSelected() { + if (Boolean.getBoolean("cuvs.lucene.tests.requireGpu")) { + assertIsSupported(); + } else { + assumeTrue("cuVS is not supported", isSupported()); + } + } + + private static void assertHnswRecallAgainstExactNeighbors( + DirectoryReader reader, float[][] vectors) throws Exception { + IndexSearcher searcher = new IndexSearcher(reader); + for (int queryId : new int[] {0, 240, 272, VECTOR_COUNT - 1}) { + var results = + searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, vectors[queryId], TOP_K), TOP_K); + List actual = new ArrayList<>(); + for (var hit : results.scoreDocs) { + actual.add(Integer.parseInt(searcher.storedFields().document(hit.doc).get(ID_FIELD))); + } + + assertEquals("Query " + queryId, TOP_K, actual.size()); + assertEquals("Query " + queryId, queryId, actual.get(0).intValue()); + assertEquals( + "Query " + queryId + " returned duplicates", TOP_K, new HashSet<>(actual).size()); + + List exact = + IntStream.range(0, vectors.length) + .boxed() + .sorted( + Comparator.comparingDouble( + (Integer id) -> squaredDistance(vectors[queryId], vectors[id])) + .thenComparingInt(Integer::intValue)) + .limit(TOP_K) + .toList(); + Set expected = new HashSet<>(exact); + long overlap = actual.stream().filter(expected::contains).count(); + assertTrue( + "Query " + queryId + " found " + overlap + "/" + TOP_K + " exact neighbors: " + actual, + overlap >= MIN_EXACT_NEIGHBORS); + } + } + + private static double squaredDistance(float[] left, float[] right) { + double distance = 0; + for (int dimension = 0; dimension < left.length; dimension++) { + double difference = left[dimension] - right[dimension]; + distance += difference * difference; + } + return distance; + } + + private static float[][] vectorsWithNegativeAndMixedSignDimensions() { + float[][] vectors = new float[VECTOR_COUNT][DIMENSIONS]; + for (int id = 0; id < VECTOR_COUNT; id++) { + int x = id % 32; + int y = id / 32; + for (int dimension = 0; dimension < DIMENSIONS; dimension++) { + vectors[id][dimension] = + switch (dimension % 4) { + case 0 -> -20.0f + 0.4f * x + 0.006f * y; // all negative + case 1 -> -7.5f + 0.9f * y + 0.002f * x; // crosses zero + case 2 -> 4.0f + 0.25f * x + 0.003f * y; // all positive + default -> -6.0f + 0.21f * (x + y); // crosses zero + }; + } + } + return vectors; + } + + private static final class RecordingInfoStream extends InfoStream { + private static final String GPU_WRITER_OPENED = + "Lucene99AcceleratedHNSWQuantizedVectorsWriter opened"; + private final List messages = new CopyOnWriteArrayList<>(); + + @Override + public void message(String component, String message) { + messages.add(component + ": " + message); + } + + @Override + public boolean isEnabled(String component) { + return true; + } + + @Override + public void close() {} + + private long gpuWriterOpenCount() { + return messages.stream().filter(message -> message.contains(GPU_WRITER_OPENED)).count(); + } + } +} From 7cc08f13b613907cb871d8e58e343fda86090f3c Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Thu, 1 Oct 2026 21:55:09 +0000 Subject: [PATCH 2/2] Clarify scalar quantization constants --- .../cuvs/lucene/AcceleratedHNSWUtils.java | 5 +- .../cuvs/lucene/TestScalarQuantization.java | 1 + .../lucene/TestScalarQuantizedCagraGraph.java | 52 +++++++++++++------ 3 files changed, 39 insertions(+), 19 deletions(-) diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index 27ff26e460..0ec59676c6 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -454,6 +454,7 @@ public static List quantizeFloatVectorsToScalar(List floatVecto return new ArrayList<>(); } + final int maximumUnsignedSevenBitValue = 127; int dimensions = floatVectors.get(0).length; int numVectors = floatVectors.size(); @@ -476,8 +477,8 @@ public static List quantizeFloatVectorsToScalar(List floatVecto float range = maxPerDim[d] - minPerDim[d]; if (range > 0) { float normalized = (vector[d] - minPerDim[d]) / range; - int quantizedValue = Math.round(normalized * 127.0f); - quantized[d] = (byte) Math.max(0, Math.min(127, quantizedValue)); + int quantizedValue = Math.round(normalized * maximumUnsignedSevenBitValue); + quantized[d] = (byte) Math.max(0, Math.min(maximumUnsignedSevenBitValue, quantizedValue)); } else { quantized[d] = 0; } diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestScalarQuantization.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestScalarQuantization.java index 48378e865b..4fe1a7305f 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestScalarQuantization.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestScalarQuantization.java @@ -32,6 +32,7 @@ public void quantizesIntermediateValuesMonotonically() { new float[] {0.5f}, new float[] {1.0f})); + // Quarter-step inputs have exact expected codes after rounding across the 0..127 range. assertArrayEquals(new byte[] {0}, quantized.get(0)); assertArrayEquals(new byte[] {32}, quantized.get(1)); assertArrayEquals(new byte[] {64}, quantized.get(2)); diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestScalarQuantizedCagraGraph.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestScalarQuantizedCagraGraph.java index 8e4a8357f6..3671ae052e 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestScalarQuantizedCagraGraph.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestScalarQuantizedCagraGraph.java @@ -36,9 +36,17 @@ public class TestScalarQuantizedCagraGraph extends LuceneTestCase { private static final String VECTOR_FIELD = "vector"; private static final String ID_FIELD = "id"; private static final int VECTOR_COUNT = 512; + private static final int GRID_COLUMNS = 32; + private static final int INITIAL_SEGMENT_COUNT = 2; + private static final int DOCUMENTS_PER_SEGMENT = VECTOR_COUNT / INITIAL_SEGMENT_COUNT; + private static final int MERGED_SEGMENT_COUNT = 1; private static final int DIMENSIONS = 128; + private static final int DIMENSION_PATTERN_COUNT = 4; private static final int TOP_K = 10; private static final int MIN_EXACT_NEIGHBORS = 8; + private static final int CAGRA_GRAPH_DEGREE = 32; + private static final int CAGRA_INTERMEDIATE_GRAPH_DEGREE = 64; + private static final int HNSW_LAYER_COUNT = 1; @Test public void testGpuBuiltScalarHnswRetainsRecallAcrossMerge() throws Exception { @@ -49,16 +57,17 @@ public void testGpuBuiltScalarHnswRetainsRecallAcrossMerge() throws Exception { AcceleratedHNSWParams params = new AcceleratedHNSWParams.Builder() .withStrategy(AcceleratedHNSWParams.Strategy.CUSTOM) - .withGraphDegree(32) - .withIntermediateGraphDegree(64) - .withHNSWLayer(1) + .withGraphDegree(CAGRA_GRAPH_DEGREE) + .withIntermediateGraphDegree(CAGRA_INTERMEDIATE_GRAPH_DEGREE) + .withHNSWLayer(HNSW_LAYER_COUNT) .build(); + int maxBufferedDocsWithoutAutomaticFlush = VECTOR_COUNT + 1; IndexWriterConfig config = new IndexWriterConfig() .setCodec( TestUtil.alwaysKnnVectorsFormat( new LuceneAcceleratedHNSWScalarQuantizedVectorsFormat(params))) - .setMaxBufferedDocs(VECTOR_COUNT + 1) + .setMaxBufferedDocs(maxBufferedDocsWithoutAutomaticFlush) .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) .setInfoStream(buildLog); @@ -69,21 +78,23 @@ public void testGpuBuiltScalarHnswRetainsRecallAcrossMerge() throws Exception { document.add(new StringField(ID_FIELD, Integer.toString(id), Field.Store.YES)); document.add(new KnnFloatVectorField(VECTOR_FIELD, vectors[id], EUCLIDEAN)); writer.addDocument(document); - if (id == VECTOR_COUNT / 2 - 1) { + if (id == DOCUMENTS_PER_SEGMENT - 1) { writer.commit(); } } writer.commit(); try (DirectoryReader reader = DirectoryReader.open(directory)) { - assertEquals(2, reader.leaves().size()); + assertEquals(INITIAL_SEGMENT_COUNT, reader.leaves().size()); assertHnswRecallAgainstExactNeighbors(reader, vectors); } long gpuBuildsBeforeMerge = buildLog.gpuWriterOpenCount(); - assertTrue("No scalar GPU writer opened: " + buildLog.messages, gpuBuildsBeforeMerge >= 2); + assertTrue( + "Insufficient scalar GPU writer openings for initial segments: " + buildLog.messages, + gpuBuildsBeforeMerge >= INITIAL_SEGMENT_COUNT); - writer.forceMerge(1); + writer.forceMerge(MERGED_SEGMENT_COUNT); writer.commit(); assertTrue( "The merge did not open a scalar GPU writer: " + buildLog.messages, @@ -91,7 +102,7 @@ public void testGpuBuiltScalarHnswRetainsRecallAcrossMerge() throws Exception { } try (DirectoryReader reader = DirectoryReader.open(directory)) { - assertEquals(1, reader.leaves().size()); + assertEquals(MERGED_SEGMENT_COUNT, reader.leaves().size()); assertHnswRecallAgainstExactNeighbors(reader, vectors); } } @@ -108,7 +119,14 @@ private static void requireGpuWhenSelected() { private static void assertHnswRecallAgainstExactNeighbors( DirectoryReader reader, float[][] vectors) throws Exception { IndexSearcher searcher = new IndexSearcher(reader); - for (int queryId : new int[] {0, 240, 272, VECTOR_COUNT - 1}) { + int centerColumn = GRID_COLUMNS / 2; + int lastRowOfFirstSegmentCenter = DOCUMENTS_PER_SEGMENT - GRID_COLUMNS + centerColumn; + int firstRowOfSecondSegmentCenter = DOCUMENTS_PER_SEGMENT + centerColumn; + // Probe opposite grid corners and adjacent rows across the segment split. + for (int queryId : + new int[] { + 0, lastRowOfFirstSegmentCenter, firstRowOfSecondSegmentCenter, VECTOR_COUNT - 1 + }) { var results = searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, vectors[queryId], TOP_K), TOP_K); List actual = new ArrayList<>(); @@ -150,15 +168,15 @@ private static double squaredDistance(float[] left, float[] right) { private static float[][] vectorsWithNegativeAndMixedSignDimensions() { float[][] vectors = new float[VECTOR_COUNT][DIMENSIONS]; for (int id = 0; id < VECTOR_COUNT; id++) { - int x = id % 32; - int y = id / 32; + int column = id % GRID_COLUMNS; + int row = id / GRID_COLUMNS; for (int dimension = 0; dimension < DIMENSIONS; dimension++) { vectors[id][dimension] = - switch (dimension % 4) { - case 0 -> -20.0f + 0.4f * x + 0.006f * y; // all negative - case 1 -> -7.5f + 0.9f * y + 0.002f * x; // crosses zero - case 2 -> 4.0f + 0.25f * x + 0.003f * y; // all positive - default -> -6.0f + 0.21f * (x + y); // crosses zero + switch (dimension % DIMENSION_PATTERN_COUNT) { + case 0 -> -20.0f + 0.4f * column + 0.006f * row; // all negative + case 1 -> -7.5f + 0.9f * row + 0.002f * column; // crosses zero + case 2 -> 4.0f + 0.25f * column + 0.003f * row; // all positive + default -> -6.0f + 0.21f * (column + row); // crosses zero }; } }