Skip to content

Commit cc2913e

Browse files
narangvivek10narangvivek10
committed
Code refactoring and cleanup
1 parent 49d690f commit cc2913e

8 files changed

Lines changed: 64 additions & 75 deletions

‎src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsFormat.java‎

Lines changed: 15 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@
1515
*/
1616
package com.nvidia.cuvs.lucene;
1717

18+
import static com.nvidia.cuvs.lucene.Utils.cuVSResourcesOrNull;
19+
1820
import com.nvidia.cuvs.CuVSResources;
1921
import com.nvidia.cuvs.LibraryException;
2022
import com.nvidia.cuvs.lucene.CuVS2510GPUVectorsWriter.IndexType;
@@ -35,7 +37,7 @@ public class CuVS2510GPUVectorsFormat extends KnnVectorsFormat {
3537

3638
// TODO: fix Lucene version in name, to the final targeted release, if any
3739
static final String CUVS_META_CODEC_NAME = "Lucene102CuVSVectorsFormatMeta";
38-
static final String CUVS_META_CODEC_EXT = "vemc"; // ""cagmf";
40+
static final String CUVS_META_CODEC_EXT = "vemc";
3941
static final String CUVS_INDEX_CODEC_NAME = "Lucene102CuVSVectorsFormatIndex";
4042
static final String CUVS_INDEX_EXT = "vcag";
4143

@@ -48,7 +50,7 @@ public class CuVS2510GPUVectorsFormat extends KnnVectorsFormat {
4850
static final IndexType DEFAULT_INDEX_TYPE = IndexType.CAGRA;
4951
static final int HNSW_GRAPH_LAYERS = 1;
5052

51-
static CuVSResources resources = Utils.cuVSResourcesOrNull();
53+
static CuVSResources resources = cuVSResourcesOrNull();
5254

5355
/** The format for storing, reading, and merging raw vectors on disk. */
5456
private static final FlatVectorsFormat flatVectorsFormat =
@@ -94,17 +96,6 @@ public CuVS2510GPUVectorsFormat(
9496
this.indexType = indexType;
9597
}
9698

97-
/** Tells whether the platform supports cuvs. */
98-
public static boolean supported() {
99-
return resources != null;
100-
}
101-
102-
private static void checkSupported() {
103-
if (!supported()) {
104-
throw new UnsupportedOperationException();
105-
}
106-
}
107-
10899
@Override
109100
public CuVS2510GPUVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException {
110101
checkSupported();
@@ -142,4 +133,15 @@ public String toString() {
142133
sb.append(")");
143134
return sb.toString();
144135
}
136+
137+
/** Tells whether the platform supports cuVS. */
138+
public static boolean supported() {
139+
return resources != null;
140+
}
141+
142+
public static void checkSupported() {
143+
if (!supported()) {
144+
throw new UnsupportedOperationException();
145+
}
146+
}
145147
}

‎src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java‎

Lines changed: 22 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
*/
1616
package com.nvidia.cuvs.lucene;
1717

18+
import static com.nvidia.cuvs.lucene.Utils.cuVSResourcesOrNull;
1819
import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH;
1920
import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN;
2021
import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_NUM_MERGE_WORKER;
@@ -50,9 +51,7 @@ public class Lucene99AcceleratedHNSWVectorsFormat extends KnnVectorsFormat {
5051
static final String HNSW_INDEX_CODEC_NAME = "Lucene99HnswVectorsFormatIndex";
5152
static final String HNSW_INDEX_EXT = "vex";
5253

53-
// Needed to make this public for a test as mocking/env variable manuplation is complicated.
54-
// TODO: Maybe explore a better solution later.
55-
public static CuVSResources resources = cuVSResourcesOrNull();
54+
private static CuVSResources resources = cuVSResourcesOrNull();
5655

5756
/** The format for storing, reading, and merging raw vectors on disk. */
5857
private static final FlatVectorsFormat flatVectorsFormat =
@@ -103,36 +102,11 @@ public Lucene99AcceleratedHNSWVectorsFormat(
103102
this.beamWidth = beamWidth;
104103
}
105104

106-
public static CuVSResources cuVSResourcesOrNull() {
107-
try {
108-
resources = CuVSResources.create();
109-
return resources;
110-
} catch (UnsupportedOperationException uoe) {
111-
log.warning("cuvs is not supported on this platform or java version: " + uoe.getMessage());
112-
} catch (Throwable t) {
113-
if (t instanceof ExceptionInInitializerError ex) {
114-
t = ex.getCause();
115-
}
116-
log.warning("Exception occurred during creation of cuvs resources. " + t);
117-
}
118-
return null;
119-
}
120-
121-
/** Tells whether the platform supports cuvs. */
122-
public static boolean supported() {
123-
return resources != null;
124-
}
125-
126-
private static void checkSupported() {
127-
if (!supported()) {
128-
throw new UnsupportedOperationException();
129-
}
130-
}
131-
132105
@Override
133106
public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException {
134107
var flatWriter = flatVectorsFormat.fieldsWriter(state);
135108
if (supported()) {
109+
log.info("cuVS is supported so using the Lucene99AcceleratedHNSWVectorsWriter");
136110
return new Lucene99AcceleratedHNSWVectorsWriter(
137111
state, cuvsWriterThreads, intGraphDegree, graphDegree, hnswLayers, resources, flatWriter);
138112
} else {
@@ -165,4 +139,23 @@ public String toString() {
165139
sb.append(")");
166140
return sb.toString();
167141
}
142+
143+
public static CuVSResources getResources() {
144+
return resources;
145+
}
146+
147+
public static void setResources(CuVSResources resources) {
148+
Lucene99AcceleratedHNSWVectorsFormat.resources = resources;
149+
}
150+
151+
/** Tells whether the platform supports cuVS. */
152+
public static boolean supported() {
153+
return resources != null;
154+
}
155+
156+
public static void checkSupported() {
157+
if (!supported()) {
158+
throw new UnsupportedOperationException();
159+
}
160+
}
168161
}

‎src/main/java/com/nvidia/cuvs/lucene/Utils.java‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,9 +20,12 @@
2020
import java.io.IOException;
2121
import java.time.Duration;
2222
import java.util.List;
23+
import java.util.logging.Logger;
2324

2425
public class Utils {
2526

27+
static final Logger log = Logger.getLogger(Utils.class.getName());
28+
2629
static void handleThrowable(Throwable t) throws IOException {
2730
switch (t) {
2831
case IOException ioe -> throw ioe;
@@ -54,17 +57,14 @@ static long nanosToMillis(long nanos) {
5457

5558
static CuVSResources cuVSResourcesOrNull() {
5659
try {
57-
CuVS2510GPUVectorsFormat.resources = CuVSResources.create();
58-
return CuVS2510GPUVectorsFormat.resources;
60+
return CuVSResources.create();
5961
} catch (UnsupportedOperationException uoe) {
60-
CuVS2510GPUVectorsFormat.log.warning(
61-
"cuvs is not supported on this platform or java version: " + uoe.getMessage());
62+
log.warning("cuVS is not supported on this platform or java version: " + uoe.getMessage());
6263
} catch (Throwable t) {
6364
if (t instanceof ExceptionInInitializerError ex) {
6465
t = ex.getCause();
6566
}
66-
CuVS2510GPUVectorsFormat.log.warning(
67-
"Exception occurred during creation of cuvs resources. " + t);
67+
log.warning("Exception occurred during creation of cuVS resources. " + t);
6868
}
6969
return null;
7070
}

‎src/test/java/com/nvidia/cuvs/lucene/TestCagraToHnswSerializationAndSearch.java‎

Lines changed: 11 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,8 @@
1515
*/
1616
package com.nvidia.cuvs.lucene;
1717

18-
import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.cuVSResourcesOrNull;
1918
import static com.nvidia.cuvs.lucene.TestUtils.generateDataset;
19+
import static com.nvidia.cuvs.lucene.Utils.cuVSResourcesOrNull;
2020
import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN;
2121

2222
import java.io.File;
@@ -75,26 +75,20 @@ public static void beforeClass() throws Exception {
7575

7676
@Test
7777
public void testCagraToHnswSerializationAndSearch() throws IOException {
78-
log.info("Test Scenario 1 - cuvs supported");
79-
test(
80-
false,
81-
new Integer[] {1869, 1803, 1302, 59, 1497, 108, 1411, 351, 1982},
82-
indexDirPath.get(0));
83-
84-
log.info("Test Scenario 2 - cuvs NOT supported");
85-
test(true, new Integer[] {885, 612, 1795, 1806, 1665}, indexDirPath.get(1));
78+
log.info("Test Scenario 1 - cuVS supported");
79+
test(new Integer[] {1869, 1803, 1302, 59, 1497, 108, 1411, 351, 1982}, indexDirPath.get(0));
80+
81+
log.info("Test Scenario 2 - cuVS NOT supported");
82+
// Set resources to null to simulate that cuVS is not supported.
83+
Lucene99AcceleratedHNSWVectorsFormat.setResources(null);
84+
test(new Integer[] {885, 612, 1795, 1806, 1665}, indexDirPath.get(1));
8685
}
8786

88-
private void test(boolean disableResources, Integer[] expected, Path indexDirPath)
89-
throws IOException {
87+
private void test(Integer[] expected, Path indexDirPath) throws IOException {
9088

9189
Lucene101AcceleratedHNSWCodec codec =
9290
new Lucene101AcceleratedHNSWCodec(32, 128, 64, 3, 16, 100);
9391

94-
if (disableResources) {
95-
Lucene99AcceleratedHNSWVectorsFormat.resources = null;
96-
}
97-
9892
IndexWriterConfig config = new IndexWriterConfig().setCodec(codec).setUseCompoundFile(false);
9993
// TODO: handle random related issues.
10094
int numDocs = 2000; // random.nextInt(100, 1000);
@@ -103,7 +97,7 @@ private void test(boolean disableResources, Integer[] expected, Path indexDirPat
10397
final int COMMIT_FREQ = 2000; // Math.min(numDocs, random.nextInt(100, 1000));
10498
int count = COMMIT_FREQ;
10599
final String ID_FIELD = "id";
106-
final String VECTOR_FIELD = "knn1";
100+
final String VECTOR_FIELD = "vector_field";
107101
float[][] dataset = generateDataset(random, numDocs, dimension);
108102

109103
// Indexing
@@ -183,7 +177,7 @@ private void test(boolean disableResources, Integer[] expected, Path indexDirPat
183177
@AfterClass
184178
public static void afterClass() throws Exception {
185179
// Reset resources.
186-
Lucene99AcceleratedHNSWVectorsFormat.resources = cuVSResourcesOrNull();
180+
Lucene99AcceleratedHNSWVectorsFormat.setResources(cuVSResourcesOrNull());
187181
// Cleanup.
188182
for (Path p : indexDirPath) {
189183
File indexDirPathFile = p.toFile();

‎src/test/java/com/nvidia/cuvs/lucene/TestCuVSDeletedDocuments.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@
5151
import org.junit.BeforeClass;
5252
import org.junit.Test;
5353

54-
@SuppressSysoutChecks(bugUrl = "prints info from within cuvs")
54+
@SuppressSysoutChecks(bugUrl = "prints info from within cuVS")
5555
public class TestCuVSDeletedDocuments extends LuceneTestCase {
5656

5757
protected static Logger log = Logger.getLogger(TestCuVSDeletedDocuments.class.getName());
@@ -61,7 +61,7 @@ public class TestCuVSDeletedDocuments extends LuceneTestCase {
6161

6262
@BeforeClass
6363
public static void beforeClass() throws Exception {
64-
assumeTrue("cuvs not supported", CuVS2510GPUVectorsFormat.supported());
64+
assumeTrue("cuVS not supported", Lucene99AcceleratedHNSWVectorsFormat.supported());
6565
random = random();
6666
}
6767

‎src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@
4848
import org.junit.BeforeClass;
4949
import org.junit.Test;
5050

51-
@SuppressSysoutChecks(bugUrl = "prints info from within cuvs")
51+
@SuppressSysoutChecks(bugUrl = "prints info from within cuVS")
5252
public class TestCuVSGaps extends LuceneTestCase {
5353

5454
protected static Logger log = Logger.getLogger(TestCuVSGaps.class.getName());
@@ -70,7 +70,7 @@ public class TestCuVSGaps extends LuceneTestCase {
7070

7171
@BeforeClass
7272
public static void beforeClass() throws Exception {
73-
assumeTrue("cuvs not supported", CuVS2510GPUVectorsFormat.supported());
73+
assumeTrue("cuVS not supported", CuVS2510GPUVectorsFormat.supported());
7474
directory = newDirectory();
7575
random = random();
7676

@@ -120,7 +120,7 @@ public static void afterClass() throws Exception {
120120

121121
@Test
122122
public void testVectorSearchWithAlternatingDocuments() throws IOException {
123-
assumeTrue("cuvs not supported", CuVS2510GPUVectorsFormat.supported());
123+
assumeTrue("cuVS not supported", CuVS2510GPUVectorsFormat.supported());
124124

125125
// Use the first vector (from document 0) as query
126126
float[] queryVector = dataset[0];
@@ -153,7 +153,7 @@ public void testVectorSearchWithAlternatingDocuments() throws IOException {
153153

154154
@Test
155155
public void testVectorSearchWithFilterAndAlternatingDocuments() throws IOException {
156-
assumeTrue("cuvs not supported", CuVS2510GPUVectorsFormat.supported());
156+
assumeTrue("cuVS not supported", CuVS2510GPUVectorsFormat.supported());
157157

158158
// Use the first vector (from document 0) as query
159159
float[] queryVector = dataset[0];

‎src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedVectorSearch.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@
5151
import org.junit.BeforeClass;
5252
import org.junit.Test;
5353

54-
@SuppressSysoutChecks(bugUrl = "prints info from within cuvs")
54+
@SuppressSysoutChecks(bugUrl = "prints info from within cuVS")
5555
public class TestCuVSRandomizedVectorSearch extends LuceneTestCase {
5656

5757
protected static Logger log = Logger.getLogger(TestCuVSRandomizedVectorSearch.class.getName());
@@ -69,7 +69,7 @@ public class TestCuVSRandomizedVectorSearch extends LuceneTestCase {
6969

7070
@BeforeClass
7171
public static void beforeClass() throws Exception {
72-
assumeTrue("cuvs not supported", CuVS2510GPUVectorsFormat.supported());
72+
assumeTrue("cuVS not supported", CuVS2510GPUVectorsFormat.supported());
7373
directory = newDirectory();
7474

7575
RandomIndexWriter writer =
@@ -184,7 +184,7 @@ private static List<List<Integer>> generateExpectedResults(
184184

185185
@Test
186186
public void testVectorSearchWithFilter() throws IOException {
187-
assumeTrue("cuvs not supported", CuVS2510GPUVectorsFormat.supported());
187+
assumeTrue("cuVS not supported", CuVS2510GPUVectorsFormat.supported());
188188

189189
Random random = random();
190190
int topK = Math.min(random.nextInt(TOP_K_LIMIT) + 1, dataset.length);

‎src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ public class TestCuVSVectorsFormat extends BaseKnnVectorsFormatTestCase {
4141

4242
@BeforeClass
4343
public static void beforeClass() {
44-
assumeTrue("cuvs is not supported", CuVS2510GPUVectorsFormat.supported());
44+
assumeTrue("cuVS is not supported", CuVS2510GPUVectorsFormat.supported());
4545
}
4646

4747
@Override

0 commit comments

Comments
 (0)