From f49653f4b944c9fa489eccdc084121c6bbaacbce Mon Sep 17 00:00:00 2001 From: punAhuja Date: Mon, 24 Feb 2025 16:57:35 +0530 Subject: [PATCH 1/3] Added tunable parameter useHNSW to enable/disable HNSW --- .../sandbox/vectorsearch/CuVSCodec.java | 3 ++- .../vectorsearch/CuVSVectorsFormat.java | 19 +++++++++++++++---- .../vectorsearch/CuVSVectorsReader.java | 6 ++++-- .../vectorsearch/CuVSVectorsWriter.java | 12 ++++++++++-- 4 files changed, 31 insertions(+), 9 deletions(-) diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/vectorsearch/CuVSCodec.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/vectorsearch/CuVSCodec.java index ac94fffaf504..eaf4c6127d54 100644 --- a/lucene/sandbox/src/java/org/apache/lucene/sandbox/vectorsearch/CuVSCodec.java +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/vectorsearch/CuVSCodec.java @@ -36,9 +36,10 @@ public CuVSCodec(String name, Codec delegate) { super(name, delegate); KnnVectorsFormat format; try { + boolean useHNSW = Boolean.parseBoolean(System.getProperty("lucene.cuvs.hnsw", "true")); format = new CuVSVectorsFormat( - 1, 128, 64, MergeStrategy.NON_TRIVIAL_MERGE, IndexType.CAGRA); + 1, 128, 64, MergeStrategy.NON_TRIVIAL_MERGE, IndexType.CAGRA, useHNSW); setKnnFormat(format); } catch (LibraryException ex) { Logger log = Logger.getLogger(CuVSCodec.class.getName()); diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/vectorsearch/CuVSVectorsFormat.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/vectorsearch/CuVSVectorsFormat.java index 705929fd86fe..ec759613f4f7 100644 --- a/lucene/sandbox/src/java/org/apache/lucene/sandbox/vectorsearch/CuVSVectorsFormat.java +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/vectorsearch/CuVSVectorsFormat.java @@ -48,6 +48,7 @@ public class CuVSVectorsFormat extends KnnVectorsFormat { public static final int DEFAULT_GRAPH_DEGREE = 64; public static final MergeStrategy DEFAULT_MERGE_STRATEGY = MergeStrategy.NON_TRIVIAL_MERGE; public static final IndexType DEFAULT_INDEX_TYPE = IndexType.CAGRA; + public static final boolean DEFAULT_USE_HNSW = true; static CuVSResources resources = cuVSResourcesOrNull(); @@ -61,6 +62,7 @@ public class CuVSVectorsFormat extends KnnVectorsFormat { final int graphDegree; final MergeStrategy mergeStrategy; final CuVSVectorsWriter.IndexType indexType; // the index type to build, when writing + final boolean useHNSW; public CuVSVectorsFormat() { this( @@ -68,7 +70,8 @@ public CuVSVectorsFormat() { DEFAULT_INTERMEDIATE_GRAPH_DEGREE, DEFAULT_GRAPH_DEGREE, DEFAULT_MERGE_STRATEGY, - DEFAULT_INDEX_TYPE); + DEFAULT_INDEX_TYPE, + DEFAULT_USE_HNSW); } public CuVSVectorsFormat( @@ -76,7 +79,8 @@ public CuVSVectorsFormat( int intGraphDegree, int graphDegree, MergeStrategy mergeStrategy, - IndexType indexType) + IndexType indexType, + boolean useHNSW) throws LibraryException { super("CuVSVectorsFormat"); this.mergeStrategy = mergeStrategy; @@ -84,8 +88,13 @@ public CuVSVectorsFormat( this.intGraphDegree = intGraphDegree; this.graphDegree = graphDegree; this.indexType = indexType; + this.useHNSW = useHNSW; } + public boolean isHNSWEnabled() { + return useHNSW; +} + private static CuVSResources cuVSResourcesOrNull() { try { resources = CuVSResources.create(); @@ -124,14 +133,15 @@ public CuVSVectorsWriter fieldsWriter(SegmentWriteState state) throws IOExceptio mergeStrategy, indexType, resources, - flatWriter); + flatWriter, + useHNSW); } @Override public CuVSVectorsReader fieldsReader(SegmentReadState state) throws IOException { checkSupported(); var flatReader = flatVectorsFormat.fieldsReader(state); - return new CuVSVectorsReader(state, resources, flatReader); + return new CuVSVectorsReader(state, resources, flatReader, useHNSW); } @Override @@ -147,6 +157,7 @@ public String toString() { sb.append("graphDegree=").append(graphDegree); sb.append("mergeStrategy=").append(mergeStrategy); sb.append("resources=").append(resources); + sb.append("useHNSW=").append(useHNSW); sb.append(")"); return sb.toString(); } diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/vectorsearch/CuVSVectorsReader.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/vectorsearch/CuVSVectorsReader.java index 97c12798e6fb..241ce72edbac 100644 --- a/lucene/sandbox/src/java/org/apache/lucene/sandbox/vectorsearch/CuVSVectorsReader.java +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/vectorsearch/CuVSVectorsReader.java @@ -66,6 +66,7 @@ public class CuVSVectorsReader extends KnnVectorsReader { private static final Logger log = Logger.getLogger(CuVSVectorsReader.class.getName()); private final CuVSResources resources; + private final boolean useHNSW; private final FlatVectorsReader flatVectorsReader; // for reading the raw vectors private final FieldInfos fieldInfos; private final IntObjectHashMap fields; @@ -73,11 +74,12 @@ public class CuVSVectorsReader extends KnnVectorsReader { private final IndexInput cuvsIndexInput; public CuVSVectorsReader( - SegmentReadState state, CuVSResources resources, FlatVectorsReader flatReader) + SegmentReadState state, CuVSResources resources, FlatVectorsReader flatReader, boolean useHNSW) throws IOException { this.resources = resources; this.flatVectorsReader = flatReader; this.fieldInfos = state.fieldInfos; + this.useHNSW = useHNSW; this.fields = new IntObjectHashMap<>(); String metaFileName = @@ -260,7 +262,7 @@ private CuVSIndex loadCuVSIndex(FieldEntry fieldEntry) throws IOException { } len = fieldEntry.hnswIndexLength(); - if (len > 0) { + if (useHNSW && len > 0) { long off = fieldEntry.hnswIndexOffset(); try (var slice = cuvsIndexInput.slice("hnsw index", off, len); var in = new IndexInputInputStream(slice)) { diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/vectorsearch/CuVSVectorsWriter.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/vectorsearch/CuVSVectorsWriter.java index e7670484ed14..7be9e55ff0cc 100644 --- a/lucene/sandbox/src/java/org/apache/lucene/sandbox/vectorsearch/CuVSVectorsWriter.java +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/vectorsearch/CuVSVectorsWriter.java @@ -127,6 +127,8 @@ public boolean hnsw() { return hnsw; } } + + private final boolean useHNSW; public CuVSVectorsWriter( SegmentWriteState state, @@ -136,7 +138,8 @@ public CuVSVectorsWriter( MergeStrategy mergeStrategy, IndexType indexType, CuVSResources resources, - FlatVectorsWriter flatVectorsWriter) + FlatVectorsWriter flatVectorsWriter, + boolean useHNSW) throws IOException { super(); this.mergeStrategy = mergeStrategy; @@ -146,6 +149,7 @@ public CuVSVectorsWriter( this.graphDegree = graphDegree; this.resources = resources; this.flatVectorsWriter = flatVectorsWriter; + this.useHNSW = useHNSW; this.infoStream = state.infoStream; String metaFileName = @@ -257,6 +261,10 @@ private void writeBruteForceIndex(OutputStream os, float[][] vectors) throws Thr } private void writeHNSWIndex(OutputStream os, float[][] vectors) throws Throwable { + if (!useHNSW) { // Skip HNSW writing if disabled + return; + } + if (vectors.length < 2) { throw new IllegalArgumentException(vectors.length + " vectors, less than min [2] required"); } @@ -335,7 +343,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, float[][] vectors) throws I } hnswIndexOffset = cuvsIndex.getFilePointer(); - if (indexType.hnsw()) { + if (useHNSW && indexType.hnsw()) { var hnswIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); if (vectors.length > MIN_CAGRA_INDEX_SIZE) { try { From 40072077750bdc307d87992016d5a29895f7db89 Mon Sep 17 00:00:00 2001 From: punAhuja Date: Mon, 3 Mar 2025 15:21:31 +0530 Subject: [PATCH 2/3] Initializing useHnsw properly --- .../lucene/sandbox/vectorsearch/CuVSVectorsReader.java | 6 +++++- .../lucene/sandbox/vectorsearch/CuVSVectorsWriter.java | 7 ++++++- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/vectorsearch/CuVSVectorsReader.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/vectorsearch/CuVSVectorsReader.java index 241ce72edbac..fe0c61d34575 100644 --- a/lucene/sandbox/src/java/org/apache/lucene/sandbox/vectorsearch/CuVSVectorsReader.java +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/vectorsearch/CuVSVectorsReader.java @@ -79,7 +79,8 @@ public CuVSVectorsReader( this.resources = resources; this.flatVectorsReader = flatReader; this.fieldInfos = state.fieldInfos; - this.useHNSW = useHNSW; + this.useHNSW = Boolean.getBoolean("lucene.cuvs.hnsw"); + log.info("CuVSVectorsReader initialized. useHNSW=" + this.useHNSW); this.fields = new IntObjectHashMap<>(); String metaFileName = @@ -238,6 +239,8 @@ private IntObjectHashMap loadCuVSIndices() throws IOException { } private CuVSIndex loadCuVSIndex(FieldEntry fieldEntry) throws IOException { + log.info("Loading CuVS index for field: "); + CagraIndex cagraIndex = null; BruteForceIndex bruteForceIndex = null; HnswIndex hnswIndex = null; @@ -263,6 +266,7 @@ private CuVSIndex loadCuVSIndex(FieldEntry fieldEntry) throws IOException { len = fieldEntry.hnswIndexLength(); if (useHNSW && len > 0) { + log.info("Attempting to load HNSW index."); long off = fieldEntry.hnswIndexOffset(); try (var slice = cuvsIndexInput.slice("hnsw index", off, len); var in = new IndexInputInputStream(slice)) { diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/vectorsearch/CuVSVectorsWriter.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/vectorsearch/CuVSVectorsWriter.java index 7be9e55ff0cc..9dc22da77a19 100644 --- a/lucene/sandbox/src/java/org/apache/lucene/sandbox/vectorsearch/CuVSVectorsWriter.java +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/vectorsearch/CuVSVectorsWriter.java @@ -149,7 +149,8 @@ public CuVSVectorsWriter( this.graphDegree = graphDegree; this.resources = resources; this.flatVectorsWriter = flatVectorsWriter; - this.useHNSW = useHNSW; + this.useHNSW = Boolean.getBoolean("lucene.cuvs.hnsw"); + log.info("CuVSVectorsWriter initialized. useHNSW=" + this.useHNSW); this.infoStream = state.infoStream; String metaFileName = @@ -262,9 +263,13 @@ private void writeBruteForceIndex(OutputStream os, float[][] vectors) throws Thr private void writeHNSWIndex(OutputStream os, float[][] vectors) throws Throwable { if (!useHNSW) { // Skip HNSW writing if disabled + log.warning("Skipping HNSW indexing because useHNSW is false."); return; } + if (vectors.length == 0) { + log.warning("HNSW indexing failed because no vectors were provided."); + } if (vectors.length < 2) { throw new IllegalArgumentException(vectors.length + " vectors, less than min [2] required"); } From f99c724ee40615c0b41bc602de6378f9253d54ec Mon Sep 17 00:00:00 2001 From: punAhuja Date: Fri, 25 Apr 2025 17:54:54 +0530 Subject: [PATCH 3/3] Creating hnswQuery if useHnsw is true --- .../lucene/sandbox/vectorsearch/CuVSCodec.java | 1 - .../vectorsearch/CuVSVectorsReader.java | 18 +++++++++++++++++- 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/vectorsearch/CuVSCodec.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/vectorsearch/CuVSCodec.java index 415f5bdcc3a6..eaf4c6127d54 100644 --- a/lucene/sandbox/src/java/org/apache/lucene/sandbox/vectorsearch/CuVSCodec.java +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/vectorsearch/CuVSCodec.java @@ -40,7 +40,6 @@ public CuVSCodec(String name, Codec delegate) { format = new CuVSVectorsFormat( 1, 128, 64, MergeStrategy.NON_TRIVIAL_MERGE, IndexType.CAGRA, useHNSW); - format = new CuVSVectorsFormat(1, 128, 64, MergeStrategy.NON_TRIVIAL_MERGE, IndexType.CAGRA); setKnnFormat(format); } catch (LibraryException ex) { Logger log = Logger.getLogger(CuVSCodec.class.getName()); diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/vectorsearch/CuVSVectorsReader.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/vectorsearch/CuVSVectorsReader.java index 577ce35fc019..4505a88eb1a9 100644 --- a/lucene/sandbox/src/java/org/apache/lucene/sandbox/vectorsearch/CuVSVectorsReader.java +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/vectorsearch/CuVSVectorsReader.java @@ -356,7 +356,23 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits assert topK > 0 : "Expected topK > 0, got:" + topK; Map result; - if (knnCollector.k() <= 1024 && cuvsIndex.getCagraIndex() != null) { + if(useHNSW && cuvsIndex.getHNSWIndex() != null) { + log.info("Searching with HNSW index"); + var hnswQuery = new com.nvidia.cuvs.HnswQuery.Builder() + .withQueryVectors(new float[][] { target }) + .withTopK(knnCollector.k()) + .build(); + List> searchResult = null; + try { + searchResult = cuvsIndex.getHNSWIndex().search(hnswQuery).getResults(); + }catch (Throwable t) { + handleThrowable(t); + } + + assert searchResult.size() == 1; + result = searchResult.getFirst(); + } + else if (knnCollector.k() <= 1024 && cuvsIndex.getCagraIndex() != null) { // log.info("searching cagra index"); CagraSearchParams searchParams = new CagraSearchParams.Builder(resources)