diff --git a/README.md b/README.md index d88d340..9b32724 100644 --- a/README.md +++ b/README.md @@ -18,6 +18,40 @@ mvn clean compile package ``` The artifacts would be built and available in the target / folder. +### Using with PyLucene + +PyLucene embeds a JVM and starts it with the classpath passed to `lucene.initVM(...)`. +Because PyLucene's generated Python module only exposes the Java classes it was built +to wrap, use Lucene's service provider lookup to load `cuvs-lucene` codecs from +Python instead of importing `com.nvidia.cuvs.lucene` classes directly. + +Build the PyLucene sidecar jar: + +```sh +mvn clean package -DskipTests +``` + +Then add the generated PyLucene jar to `lucene.CLASSPATH` before starting the VM: + +```python +import os +import lucene + +cuvs_lucene_jar = "target/cuvs-lucene-26.06.0-jar-with-pylucene-dependencies.jar" +lucene.initVM(classpath=os.pathsep.join([lucene.CLASSPATH, cuvs_lucene_jar])) + +from org.apache.lucene.codecs import Codec + +codec = Codec.forName("Lucene101AcceleratedHNSWCodec") +``` + +Use the returned `codec` with `IndexWriterConfig.setCodec(codec)`. The +`jar-with-pylucene-dependencies` artifact includes `cuvs-lucene` and its non-Lucene +runtime dependencies while leaving Lucene itself to PyLucene's own classpath. This +avoids loading a second copy of Lucene classes into the embedded JVM. The regular +`jar-with-dependencies` artifact also merges `META-INF/services` entries and is +available for non-PyLucene Java applications that want a standalone jar. + ### Running Tests ```sh export LD_LIBRARY_PATH={ PATH TO YOUR LOCAL libcuvs_c.so }:$LD_LIBRARY_PATH && mvn clean test diff --git a/bench/pom.xml b/bench/pom.xml index 8f7262a..d9a1344 100644 --- a/bench/pom.xml +++ b/bench/pom.xml @@ -11,7 +11,7 @@ com.nvidia.cuvs.lucene.benchmarks cuvs-lucene-benchmarks - 26.02.0 + 26.06.0 jar cuvs-lucene-benchmarks diff --git a/examples/pylucene_smoke.py b/examples/pylucene_smoke.py new file mode 100644 index 0000000..68bf614 --- /dev/null +++ b/examples/pylucene_smoke.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 + +import os +from pathlib import Path + +import lucene +from lucene import JArray + + +REPO_ROOT = Path(__file__).resolve().parents[1] +CUVS_LUCENE_JAR = ( + REPO_ROOT / "target" / "cuvs-lucene-26.06.0-jar-with-pylucene-dependencies.jar" +) +INDEX_PATH = Path("/tmp/cuvs-lucene-pylucene-index") + + +def fvec(values): + return JArray("float")([float(value) for value in values]) + + +def main(): + lucene.initVM(classpath=os.pathsep.join([lucene.CLASSPATH, str(CUVS_LUCENE_JAR)])) + + from java.nio.file import Paths + from org.apache.lucene.codecs import Codec + from org.apache.lucene.document import Document, Field, KnnFloatVectorField, StringField + from org.apache.lucene.index import ( + DirectoryReader, + IndexWriter, + IndexWriterConfig, + VectorSimilarityFunction, + ) + from org.apache.lucene.search import IndexSearcher, KnnFloatVectorQuery + from org.apache.lucene.store import MMapDirectory + + codec = Codec.forName("Lucene101AcceleratedHNSWCodec") + + directory = MMapDirectory(Paths.get(str(INDEX_PATH))) + config = IndexWriterConfig() + config.setCodec(codec) + config.setUseCompoundFile(False) + + writer = IndexWriter(directory, config) + for doc_id, vector in [("a", [0.1, 0.2, 0.3]), ("b", [0.2, 0.1, 0.4])]: + doc = Document() + doc.add(StringField("id", doc_id, Field.Store.YES)) + doc.add( + KnnFloatVectorField( + "vector", fvec(vector), VectorSimilarityFunction.EUCLIDEAN + ) + ) + writer.addDocument(doc) + writer.close() + + reader = DirectoryReader.open(directory) + searcher = IndexSearcher(reader) + query = KnnFloatVectorQuery("vector", fvec([0.1, 0.2, 0.3]), 2) + + for hit in searcher.search(query, 2).scoreDocs: + print(searcher.doc(hit.doc).get("id"), hit.score) + + reader.close() + directory.close() + + +if __name__ == "__main__": + main() diff --git a/pom.xml b/pom.xml index 19479f1..8b18da2 100644 --- a/pom.xml +++ b/pom.xml @@ -112,9 +112,10 @@ maven-assembly-plugin 3.6.0 - - jar-with-dependencies - + + src/main/assembly/jar-with-dependencies.xml + src/main/assembly/jar-with-pylucene-dependencies.xml + diff --git a/src/main/assembly/jar-with-dependencies.xml b/src/main/assembly/jar-with-dependencies.xml new file mode 100644 index 0000000..3722a87 --- /dev/null +++ b/src/main/assembly/jar-with-dependencies.xml @@ -0,0 +1,22 @@ + + jar-with-dependencies + + jar + + false + + + / + true + true + runtime + + + + + metaInf-services + + + diff --git a/src/main/assembly/jar-with-pylucene-dependencies.xml b/src/main/assembly/jar-with-pylucene-dependencies.xml new file mode 100644 index 0000000..d243257 --- /dev/null +++ b/src/main/assembly/jar-with-pylucene-dependencies.xml @@ -0,0 +1,25 @@ + + jar-with-pylucene-dependencies + + jar + + false + + + / + true + true + runtime + + org.apache.lucene:* + + + + + + metaInf-services + + + diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java index 2de761a..2567fc4 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUSearchCodec.java @@ -29,7 +29,7 @@ public class CuVS2510GPUSearchCodec extends FilterCodec { * @throws Exception */ public CuVS2510GPUSearchCodec() throws Exception { - this(NAME, LuceneProvider.getCodec("101")); + this(NAME, LuceneProvider.getDefaultDelegateCodec()); initializeFormat(new GPUSearchParams.Builder().build()); } @@ -53,7 +53,7 @@ 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")); + this(NAME, LuceneProvider.getDefaultDelegateCodec()); initializeFormat(params); } diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java index b4c5a33..4944b76 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java @@ -30,7 +30,7 @@ public class Lucene101AcceleratedHNSWCodec extends FilterCodec { * @throws Exception */ public Lucene101AcceleratedHNSWCodec() throws Exception { - this(NAME, LuceneProvider.getCodec("101")); + this(NAME, LuceneProvider.getDefaultDelegateCodec()); } /** @@ -52,7 +52,7 @@ public Lucene101AcceleratedHNSWCodec(String name, Codec delegate) { */ public Lucene101AcceleratedHNSWCodec(AcceleratedHNSWParams acceleratedHNSWParams) throws Exception { - this(NAME, LuceneProvider.getCodec("101")); + this(NAME, LuceneProvider.getDefaultDelegateCodec()); initializeFormat(acceleratedHNSWParams); } diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java index f2c1aa3..6af2d3d 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedCodec.java @@ -25,7 +25,7 @@ public class LuceneAcceleratedHNSWBinaryQuantizedCodec extends FilterCodec { private KnnVectorsFormat format; public LuceneAcceleratedHNSWBinaryQuantizedCodec() throws Exception { - this(NAME, LuceneProvider.getCodec("101")); + this(NAME, LuceneProvider.getDefaultDelegateCodec()); } public LuceneAcceleratedHNSWBinaryQuantizedCodec(String name, Codec delegate) { @@ -35,7 +35,7 @@ public LuceneAcceleratedHNSWBinaryQuantizedCodec(String name, Codec delegate) { public LuceneAcceleratedHNSWBinaryQuantizedCodec(AcceleratedHNSWParams acceleratedHNSWParams) throws Exception { - this(NAME, LuceneProvider.getCodec("101")); + this(NAME, LuceneProvider.getDefaultDelegateCodec()); initializeFormat(acceleratedHNSWParams); } diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java index 0f8d960..d70efc5 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java @@ -27,21 +27,26 @@ public class LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat extends KnnVector private static final Logger log = Logger.getLogger(LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.class.getName()); - private static final LuceneProvider LUCENE102_PROVIDER; - private static final LuceneProvider LUCENE99_PROVIDER; - private static final FlatVectorsFormat FLAT_VECTORS_FORMAT; private static final int MAX_DIMENSIONS = 4096; private final AcceleratedHNSWParams acceleratedHNSWParams; - static { + private static LuceneProvider luceneProvider(String version) { try { - LUCENE99_PROVIDER = LuceneProvider.getInstance("99"); - LUCENE102_PROVIDER = LuceneProvider.getInstance("102"); - FLAT_VECTORS_FORMAT = - LUCENE102_PROVIDER.getLuceneFlatVectorsFormatInstance(DefaultFlatVectorScorer.INSTANCE); + return LuceneProvider.getInstance(version); } catch (Exception e) { - throw new ExceptionInInitializerError(e.getMessage()); + throw new UnsupportedOperationException( + "Lucene" + version + " vector formats are not available in this runtime", e); + } + } + + private static FlatVectorsFormat flatVectorsFormat() { + try { + return luceneProvider("102") + .getLuceneFlatVectorsFormatInstance(DefaultFlatVectorScorer.INSTANCE); + } catch (Exception e) { + throw new UnsupportedOperationException( + "Binary quantized vectors require Lucene102 vector formats", e); } } @@ -70,7 +75,7 @@ public LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat( */ @Override public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { - var flatWriter = FLAT_VECTORS_FORMAT.fieldsWriter(state); + var flatWriter = flatVectorsFormat().fieldsWriter(state); if (isSupported()) { log.log( Level.FINE, @@ -85,8 +90,9 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException "GPU based indexing not supported, falling back to using the" + " Lucene102HnswBinaryQuantizedVectorsFormat"); KnnVectorsFormat fallbackFormat = - LUCENE102_PROVIDER.getLuceneHnswBinaryQuantizedVectorsFormatInstance( - acceleratedHNSWParams.getMaxConn(), acceleratedHNSWParams.getBeamWidth()); + luceneProvider("102") + .getLuceneHnswBinaryQuantizedVectorsFormatInstance( + acceleratedHNSWParams.getMaxConn(), acceleratedHNSWParams.getBeamWidth()); return fallbackFormat.fieldsWriter(state); } catch (Exception e) { throw new RuntimeException(e.getMessage()); @@ -100,8 +106,8 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException @Override public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { try { - return LUCENE99_PROVIDER.getLuceneHnswVectorsReaderInstance( - state, FLAT_VECTORS_FORMAT.fieldsReader(state)); + return luceneProvider("99") + .getLuceneHnswVectorsReaderInstance(state, flatVectorsFormat().fieldsReader(state)); } catch (Exception e) { throw new RuntimeException(e.getMessage()); } diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedCodec.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedCodec.java index 0c7736a..4e49027 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedCodec.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedCodec.java @@ -25,7 +25,7 @@ public class LuceneAcceleratedHNSWScalarQuantizedCodec extends FilterCodec { private KnnVectorsFormat format; public LuceneAcceleratedHNSWScalarQuantizedCodec() throws Exception { - this(NAME, LuceneProvider.getCodec("101")); + this(NAME, LuceneProvider.getDefaultDelegateCodec()); } public LuceneAcceleratedHNSWScalarQuantizedCodec(String name, Codec delegate) { @@ -35,7 +35,7 @@ public LuceneAcceleratedHNSWScalarQuantizedCodec(String name, Codec delegate) { public LuceneAcceleratedHNSWScalarQuantizedCodec(AcceleratedHNSWParams acceleratedHNSWParams) throws Exception { - this(NAME, LuceneProvider.getCodec("101")); + this(NAME, LuceneProvider.getDefaultDelegateCodec()); initializeFormat(acceleratedHNSWParams); } diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java index 7635e32..df758cc 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneProvider.java @@ -8,7 +8,9 @@ import java.lang.invoke.VarHandle; import java.lang.reflect.Constructor; import java.lang.reflect.InvocationTargetException; +import java.util.HashMap; import java.util.List; +import java.util.Map; import java.util.logging.Level; import java.util.logging.Logger; import org.apache.lucene.codecs.Codec; @@ -79,7 +81,7 @@ public class LuceneProvider { private static String luceneCodec = BASE + codecs + "LuceneCodec"; private static String luceneCodecFallback = BASE + fallbackCodecs + "LuceneCodec"; - private static LuceneProvider instance; + private static final Map INSTANCES = new HashMap<>(); private static MethodHandles.Lookup lookup = MethodHandles.lookup(); @@ -92,9 +94,12 @@ public class LuceneProvider { private Class scalarQuantizedVectorsFormat; private Class hnswScalarQuantizedVectorsFormat; - public static LuceneProvider getInstance(String version) throws ClassNotFoundException { + public static synchronized LuceneProvider getInstance(String version) + throws ClassNotFoundException { + LuceneProvider instance = INSTANCES.get(version); if (instance == null) { instance = new LuceneProvider(version); + INSTANCES.put(version, instance); } return instance; } @@ -173,6 +178,21 @@ public static Codec getCodec(String version) return (Codec) codecClassConstructor.newInstance(); } + public static Codec getDefaultDelegateCodec() { + for (String version : List.of("101", "99")) { + try { + return getCodec(version); + } catch (ReflectiveOperationException + | SecurityException + | IllegalArgumentException + | LinkageError e) { + log.log(Level.FINE, "Unable to load Lucene" + version + "Codec", e); + } + } + log.log(Level.FINE, "Falling back to the runtime default codec"); + return Codec.getDefault(); + } + public FlatVectorsFormat getLuceneFlatVectorsFormatInstance(FlatVectorsScorer scorer) throws Exception { try { diff --git a/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java b/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java index 2de6e66..1407e0a 100644 --- a/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java +++ b/src/test/java/com/nvidia/cuvs/lucene/TestBackCompat.java @@ -6,6 +6,8 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNotSame; import static org.junit.Assert.assertTrue; import org.apache.lucene.codecs.Codec; @@ -38,4 +40,30 @@ public void testExistingComponents() throws Exception { assertEquals(provider.getStaticIntParam("VERSION_CURRENT"), 0); assertNotEquals(provider.getSimilarityFunctions().size(), 0); } + + @Test + public void testProviderCachesByVersion() throws Exception { + LuceneProvider provider99 = LuceneProvider.getInstance("99"); + LuceneProvider provider102 = LuceneProvider.getInstance("102"); + assertNotSame(provider99, provider102); + } + + @Test + public void testDefaultDelegateCodec() { + assertNotNull(LuceneProvider.getDefaultDelegateCodec()); + } + + @Test + public void testServiceLoadedCodecsCanBeInstantiated() { + String[] codecNames = { + "Lucene101AcceleratedHNSWCodec", + "CuVS2510GPUSearchCodec", + "Lucene101AcceleratedHNSWBinaryQuantizedCodec", + "Lucene101AcceleratedHNSWScalarQuantizedCodec" + }; + for (String codecName : codecNames) { + assertTrue(Codec.availableCodecs().contains(codecName)); + assertEquals(codecName, Codec.forName(codecName).getName()); + } + } }