From 28d2de677e766d345bebb89c3351a38f48e9dcd5 Mon Sep 17 00:00:00 2001 From: Puneet Ahuja <167976912+punAhuja@users.noreply.github.com> Date: Fri, 18 Jul 2025 10:21:06 +0530 Subject: [PATCH] [Java] Support for tiered index (#1028) TieredIndex exposed in the Java layer. Code is written and tested for CAGRA algorithm only. Serialization, Deserialization and merge is still not supported in the C layer. Fixes: #1014 Authors: - Puneet Ahuja (https://github.com/punAhuja) - Vivek Narang (https://github.com/narangvivek10) Approvers: - MithunR (https://github.com/mythrocks) URL: https://github.com/rapidsai/cuvs/pull/1028 --- .../java/com/nvidia/cuvs/TieredIndex.java | 184 +++++ .../com/nvidia/cuvs/TieredIndexParams.java | 173 +++++ .../com/nvidia/cuvs/TieredIndexQuery.java | 237 +++++++ .../com/nvidia/cuvs/spi/CuVSProvider.java | 5 + .../nvidia/cuvs/spi/UnsupportedProvider.java | 6 + .../nvidia/cuvs/internal/TieredIndexImpl.java | 644 ++++++++++++++++++ .../internal/TieredSearchResultsImpl.java | 62 ++ .../com/nvidia/cuvs/spi/JDKProvider.java | 7 + .../java/com/nvidia/cuvs/TieredIndexIT.java | 259 +++++++ java/panama-bindings/headers.h | 1 + 10 files changed, 1578 insertions(+) create mode 100644 java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndex.java create mode 100644 java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndexParams.java create mode 100644 java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndexQuery.java create mode 100644 java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java create mode 100644 java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredSearchResultsImpl.java create mode 100644 java/cuvs-java/src/test/java/com/nvidia/cuvs/TieredIndexIT.java diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndex.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndex.java new file mode 100644 index 0000000000..a6bd7e1f6e --- /dev/null +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndex.java @@ -0,0 +1,184 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs; + +import com.nvidia.cuvs.spi.CuVSProvider; +import java.io.InputStream; +import java.util.Objects; + +/** + * {@link TieredIndex} encapsulates a Tiered index, along with methods to + * interact with it. + */ +public interface TieredIndex { + + /** + * Destroys the underlying native TieredIndex object and releases associated + * resources. + * + * @throws Throwable if an error occurs during index destruction + */ + void destroyIndex() throws Throwable; + + /** + * Searches the index with the specified query and search parameters. + * + * @param query An instance of {@link TieredIndexQuery} describing the queries + * and search parameters + * @return An instance of {@link SearchResults} containing the k-nearest + * neighbors and their distances for each query + * @throws Throwable if an error occurs during the search operation + */ + SearchResults search(TieredIndexQuery query) throws Throwable; + + /** + * Returns the algorithm type backing this TieredIndex. + * + * @return The {@link TieredIndexType} indicating the underlying algorithm + * (e.g., CAGRA) + */ + TieredIndexType getIndexType(); + + /** + * Returns the resources handle associated with this TieredIndex. + * + * @return The {@link CuVSResources} instance used by this index + */ + CuVSResources getCuVSResources(); + + /** + * Creates a new Builder with an instance of {@link CuVSResources}. + * + * @param cuvsResources An instance of {@link CuVSResources} + * @return A new {@link Builder} instance for constructing a TieredIndex + * @throws NullPointerException if cuvsResources is null + */ + static Builder newBuilder(CuVSResources cuvsResources) { + Objects.requireNonNull(cuvsResources); + return CuVSProvider.provider().newTieredIndexBuilder(cuvsResources); + } + + /** + * Returns an ExtendBuilder to add new data to the existing index. + * + * @return An {@link ExtendBuilder} instance for extending the index + */ + ExtendBuilder extend(); + + /** + * Builder interface for constructing {@link TieredIndex} instances. + */ + interface Builder { + + /** + * + * @param inputStream The input stream containing serialized index data + * @return This Builder instance for method chaining + * @throws UnsupportedOperationException as deserialization is not yet + * supported + */ + Builder from(InputStream inputStream); + + /** + * Sets the dataset vectors for building the TieredIndex. + * + * @param vectors A two-dimensional float array containing the dataset + * vectors [n_vectors, dimensions] + * @return This Builder instance for method chaining + */ + Builder withDataset(float[][] vectors); + + /** + * Sets the dataset for building the TieredIndex. + * + * @param dataset A {@link Dataset} instance containing the vectors + * @return This Builder instance for method chaining + */ + Builder withDataset(Dataset dataset); + + /** + * Registers TieredIndex parameters with this Builder. + * + * @param params An instance of {@link TieredIndexParams} containing the + * index configuration + * @return This Builder instance for method chaining + */ + Builder withIndexParams(TieredIndexParams params); + + /** + * Sets the index type for the TieredIndex. + * + * @param indexType The {@link TieredIndexType} to use (currently only CAGRA + * is supported) + * @return This Builder instance for method chaining + */ + Builder withIndexType(TieredIndexType indexType); + + /** + * Builds and returns an instance of TieredIndex with the configured + * parameters. + * + * @return A new {@link TieredIndex} instance + * @throws Throwable if an error occurs during index + * construction + * @throws IllegalArgumentException if both vectors and dataset are provided, + * or if required parameters are missing + */ + TieredIndex build() throws Throwable; + } + + /** + * Enumeration of supported TieredIndex algorithm types. + */ + enum TieredIndexType { + CAGRA + } + + /** + * Builder interface for extending existing {@link TieredIndex} instances with + * new data. + */ + interface ExtendBuilder { + + /** + * Sets the vectors to add to the existing index. + * + * @param vectors A two-dimensional float array containing the new vectors to + * add [n_new_vectors, dimensions] + * @return This ExtendBuilder instance for method chaining + */ + ExtendBuilder withDataset(float[][] vectors); + + /** + * Sets the dataset to add to the existing index. + * + * @param dataset A {@link Dataset} instance containing the new vectors to + * add + * @return This ExtendBuilder instance for method chaining + */ + ExtendBuilder withDataset(Dataset dataset); + + /** + * Executes the extend operation, adding the specified data to the index. + * + * @throws Throwable if an error occurs during the extend + * operation + * @throws IllegalArgumentException if both vectors and dataset are provided, + * or if no data is provided + */ + void execute() throws Throwable; + } +} diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndexParams.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndexParams.java new file mode 100644 index 0000000000..f62892b2c1 --- /dev/null +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndexParams.java @@ -0,0 +1,173 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs; + +import java.util.Objects; + +/** + * Configuration parameters for building a {@link TieredIndex}. + * Only CAGRA is currently supported as the underlying ANN algorithm. + * + */ +public final class TieredIndexParams { + + /** + * Enumeration of supported distance metrics for TieredIndex. + */ + public enum Metric { + /** L2 (Euclidean) distance metric */ + L2, + /** Inner product (cosine similarity) distance metric */ + INNER_PRODUCT + } + + private final Metric metric; + private final int minAnnRows; + private final boolean createAnnIndexOnExtend; + private final CagraIndexParams cagraParams; + + /** + * Private constructor used by the Builder. + * + * @param builder The Builder instance containing the configuration + */ + private TieredIndexParams(Builder builder) { + this.metric = builder.metric; + this.minAnnRows = builder.minAnnRows; + this.createAnnIndexOnExtend = builder.createAnnIndexOnExtend; + this.cagraParams = builder.cagraParams; + } + + /** + * Returns the distance metric used for similarity computation. + * + * @return The {@link Metric} (L2 or Inner Product) + */ + public Metric getMetric() { + return metric; + } + + /** + * Returns the minimum number of rows required to use the ANN algorithm. + * + * @return The minimum row count threshold for ANN algorithm usage + */ + public int getMinAnnRows() { + return minAnnRows; + } + + /** + * Returns whether to create an ANN index when extending the dataset. + * + * @return true if ANN index should be created on extend, false otherwise + */ + public boolean isCreateAnnIndexOnExtend() { + return createAnnIndexOnExtend; + } + + /** + * Returns the CAGRA-specific parameters for the ANN algorithm. + * + * @return The {@link CagraIndexParams} configuration, or null if not using + * CAGRA + */ + public CagraIndexParams getCagraParams() { + return cagraParams; + } + + /** + * Creates a new Builder for constructing TieredIndexParams. + * + * @return A new Builder instance + */ + public static Builder newBuilder() { + return new Builder(); + } + + /** + * Builder class for constructing {@link TieredIndexParams} instances. + */ + public static final class Builder { + private Metric metric = Metric.L2; + private int minAnnRows = 4096; + private boolean createAnnIndexOnExtend = true; + private CagraIndexParams cagraParams = null; + + /** + * Sets the distance metric for similarity computation. + * + * @param metric The {@link Metric} to use (L2 or Inner Product) + * @return This Builder instance for method chaining + * @throws NullPointerException if metric is null + */ + public Builder metric(Metric metric) { + this.metric = Objects.requireNonNull(metric); + return this; + } + + /** + * Sets the minimum number of rows required to use the ANN algorithm. + * + * @param minAnnRows The minimum row count threshold (must be positive) + * @return This Builder instance for method chaining + * @throws IllegalArgumentException if minAnnRows is not positive + */ + public Builder minAnnRows(int minAnnRows) { + if (minAnnRows <= 0) { + throw new IllegalArgumentException("minAnnRows must be positive, got: " + minAnnRows); + } + this.minAnnRows = minAnnRows; + return this; + } + + /** + * Sets whether to create an ANN index when extending the dataset. + * + * @param val true to create ANN index on extend, false otherwise + * @return This Builder instance for method chaining + */ + public Builder createAnnIndexOnExtend(boolean val) { + this.createAnnIndexOnExtend = val; + return this; + } + + /** + * Sets the CAGRA-specific parameters for the ANN algorithm. + * + * @param params The {@link CagraIndexParams} configuration for CAGRA + * algorithm + * @return This Builder instance for method chaining + * @throws NullPointerException if params is null + */ + public Builder withCagraParams(CagraIndexParams params) { + this.cagraParams = Objects.requireNonNull(params); + return this; + } + + /** + * Builds and returns a {@link TieredIndexParams} instance with the + * configured parameters. + * + * @return A new TieredIndexParams instance + * @throws IllegalStateException if CAGRA params are required but not + * provided + */ + public TieredIndexParams build() { + if (cagraParams == null) throw new IllegalStateException("CAGRA params required"); + return new TieredIndexParams(this); + } + } +} diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndexQuery.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndexQuery.java new file mode 100644 index 0000000000..03bfae070e --- /dev/null +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndexQuery.java @@ -0,0 +1,237 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs; + +import com.nvidia.cuvs.TieredIndex.TieredIndexType; +import java.util.Arrays; +import java.util.BitSet; +import java.util.List; + +/** + * TieredIndexQuery holds the search parameters and query vectors to be used + * while + * invoking search. Currently only supports CAGRA index type. + * + * @since 25.02 + */ +public class TieredIndexQuery { + private TieredIndexType indexType; + private CagraSearchParams cagraSearchParameters; + private List mapping; + private float[][] queryVectors; + private int topK; + private BitSet prefilter; + private long numDocs; + + private TieredIndexQuery( + TieredIndexType indexType, + CagraSearchParams cagraSearchParameters, + List mapping, + float[][] queryVectors, + int topK, + BitSet prefilter, + long numDocs) { + super(); + this.indexType = indexType; + this.cagraSearchParameters = cagraSearchParameters; + this.mapping = mapping; + this.queryVectors = queryVectors; + this.topK = topK; + this.prefilter = prefilter; + this.numDocs = numDocs; + } + + /** + * Gets the index type for this query. + * + * @return the TieredIndexType + */ + public TieredIndexType getIndexType() { + return indexType; + } + + /** + * Gets the instance of CagraSearchParams initially set. + * + * @return an instance CagraSearchParams + */ + public CagraSearchParams getCagraSearchParameters() { + return cagraSearchParameters; + } + + /** + * Gets the query vector 2D float array. + * + * @return 2D float array + */ + public float[][] getQueryVectors() { + return queryVectors; + } + + /** + * Gets the passed map instance. + * + * @return a map of ID mappings + */ + public List getMapping() { + return mapping; + } + + /** + * Gets the topK value. + * + * @return the topK value + */ + public int getTopK() { + return topK; + } + + /** + * Gets the prefilter BitSet. + * + * @return a BitSet object representing the prefilter + */ + public BitSet getPrefilter() { + return prefilter; + } + + /** + * Gets the number of documents in this index, as used for prefilter. + * + * @return number of documents as an integer + */ + public long getNumDocs() { + return numDocs; + } + + @Override + public String toString() { + return "TieredIndexQuery [indexType=" + + indexType + + ", cagraSearchParameters=" + + cagraSearchParameters + + ", queryVectors=" + + Arrays.toString(queryVectors) + + ", mapping=" + + mapping + + ", topK=" + + topK + + "]"; + } + + /** + * Creates a new Builder instance. + * + * @return a new Builder instance + */ + public static Builder newBuilder() { + return new Builder(); + } + + /** + * Builder helps configure and create an instance of TieredIndexQuery. + */ + public static class Builder { + private TieredIndexType indexType = TieredIndexType.CAGRA; + private CagraSearchParams cagraSearchParams; + private float[][] queryVectors; + private List mapping; + private int topK = 2; + private BitSet prefilter; + private long numDocs; + + /** + * Sets the index type for this query. + * + * @param indexType the index type + * @return an instance of this Builder + */ + public Builder withIndexType(TieredIndexType indexType) { + this.indexType = indexType; + return this; + } + + /** + * Sets the instance of configured CagraSearchParams to be passed for search. + * + * @param cagraSearchParams an instance of the configured CagraSearchParams to + * be used for this query + * @return an instance of this Builder + */ + public Builder withSearchParams(CagraSearchParams cagraSearchParams) { + this.cagraSearchParams = cagraSearchParams; + return this; + } + + /** + * Registers the query vectors to be passed in the search call. + * + * @param queryVectors 2D float query vector array + * @return an instance of this Builder + */ + public Builder withQueryVectors(float[][] queryVectors) { + this.queryVectors = queryVectors; + return this; + } + + /** + * Sets the instance of mapping to be used for ID mapping. + * + * @param mapping the ID mapping instance + * @return an instance of this Builder + */ + public Builder withMapping(List mapping) { + this.mapping = mapping; + return this; + } + + /** + * Registers the topK value. + * + * @param topK the topK value used to retrieve the topK results + * @return an instance of this Builder + */ + public Builder withTopK(int topK) { + this.topK = topK; + return this; + } + + /** + * Sets a BitSet to use as prefilter while searching. + * + * @param prefilter the BitSet to use as prefilter + * @param numDocs Total number of dataset vectors; used to align the prefilter + * correctly + * @return an instance of this Builder + */ + public Builder withPrefilter(BitSet prefilter, int numDocs) { + this.prefilter = prefilter; + this.numDocs = numDocs; + return this; + } + + /** + * Builds an instance of TieredIndexQuery. + * + * @return an instance of TieredIndexQuery + * @throws IllegalStateException if required parameters are missing + */ + public TieredIndexQuery build() { + return new TieredIndexQuery( + indexType, cagraSearchParams, mapping, queryVectors, topK, prefilter, numDocs); + } + } +} diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java index 8ce8b3c2a2..4dd7fa8eb8 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java @@ -21,6 +21,7 @@ import com.nvidia.cuvs.CuVSResources; import com.nvidia.cuvs.Dataset; import com.nvidia.cuvs.HnswIndex; +import com.nvidia.cuvs.TieredIndex; import java.lang.invoke.MethodHandle; import java.lang.invoke.MethodType; import java.nio.file.Path; @@ -85,6 +86,10 @@ CagraIndex.Builder newCagraIndexBuilder(CuVSResources cuVSResources) HnswIndex.Builder newHnswIndexBuilder(CuVSResources cuVSResources) throws UnsupportedOperationException; + /** Creates a new TieredIndex Builder. */ + TieredIndex.Builder newTieredIndexBuilder(CuVSResources cuVSResources) + throws UnsupportedOperationException; + /** * Merges multiple CAGRA indexes into a single index. * diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java index e7bb0ab663..7e5d1a7e64 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java @@ -20,6 +20,7 @@ import com.nvidia.cuvs.CuVSResources; import com.nvidia.cuvs.Dataset; import com.nvidia.cuvs.HnswIndex; +import com.nvidia.cuvs.TieredIndex; import java.lang.invoke.MethodHandle; import java.nio.file.Path; @@ -48,6 +49,11 @@ public HnswIndex.Builder newHnswIndexBuilder(CuVSResources cuVSResources) { throw new UnsupportedOperationException(); } + @Override + public TieredIndex.Builder newTieredIndexBuilder(CuVSResources cuVSResources) { + throw new UnsupportedOperationException(); + } + @Override public CagraIndex mergeCagraIndexes(CagraIndex[] indexes) throws Throwable { throw new UnsupportedOperationException(); diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java new file mode 100644 index 0000000000..035bbe2eef --- /dev/null +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java @@ -0,0 +1,644 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs.internal; + +import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT; +import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT_BYTE_SIZE; +import static com.nvidia.cuvs.internal.common.LinkerHelper.C_INT_BYTE_SIZE; +import static com.nvidia.cuvs.internal.common.LinkerHelper.C_LONG; +import static com.nvidia.cuvs.internal.common.LinkerHelper.C_LONG_BYTE_SIZE; +import static com.nvidia.cuvs.internal.common.LinkerHelper.C_POINTER; +import static com.nvidia.cuvs.internal.common.Util.*; +import static com.nvidia.cuvs.internal.common.Util.CudaMemcpyKind.*; +import static com.nvidia.cuvs.internal.common.Util.buildMemorySegment; +import static com.nvidia.cuvs.internal.common.Util.checkCuVSError; +import static com.nvidia.cuvs.internal.common.Util.checkCudaError; +import static com.nvidia.cuvs.internal.common.Util.concatenate; +import static com.nvidia.cuvs.internal.common.Util.prepareTensor; +import static com.nvidia.cuvs.internal.panama.headers_h.cudaMemcpy; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsRMMAlloc; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsRMMFree; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsStreamSync; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsTieredIndexBuild; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsTieredIndexCreate; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsTieredIndexDestroy; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsTieredIndexExtend; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsTieredIndexSearch; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsTieredIndex_t; + +import com.nvidia.cuvs.CagraIndexParams; +import com.nvidia.cuvs.CagraSearchParams; +import com.nvidia.cuvs.CuVSResources; +import com.nvidia.cuvs.Dataset; +import com.nvidia.cuvs.SearchResults; +import com.nvidia.cuvs.TieredIndex; +import com.nvidia.cuvs.TieredIndexParams; +import com.nvidia.cuvs.TieredIndexQuery; +import com.nvidia.cuvs.internal.common.Util; +import com.nvidia.cuvs.internal.panama.cuvsCagraIndexParams; +import com.nvidia.cuvs.internal.panama.cuvsCagraSearchParams; +import com.nvidia.cuvs.internal.panama.cuvsFilter; +import com.nvidia.cuvs.internal.panama.cuvsTieredIndexParams; +import java.io.InputStream; +import java.lang.foreign.Arena; +import java.lang.foreign.MemoryLayout; +import java.lang.foreign.MemorySegment; +import java.lang.foreign.SequenceLayout; +import java.util.BitSet; +import java.util.Objects; + +/** + * {@link TieredIndex} encapscaps a Tiered index, along with methods to interact + * with it. + *

+ * TieredIndex is a hybrid index that combines brute force search for small datasets + * with ANN algorithms (like CAGRA) for larger datasets, providing optimal performance + * across different data sizes. + * + * @since 25.02 + */ +public class TieredIndexImpl implements TieredIndex { + private final float[][] vectors; + private final Dataset dataset; + private final CuVSResourcesImpl resources; + private final TieredIndexParams tieredIndexParameters; + private final IndexReference tieredIndexReference; + private boolean destroyed; + + /** + * Constructor for building the index using specified dataset + */ + private TieredIndexImpl( + TieredIndexParams indexParameters, + float[][] vectors, + Dataset dataset, + CuVSResourcesImpl resources) + throws Throwable { + this.tieredIndexParameters = indexParameters; + this.vectors = vectors; + this.dataset = dataset; + this.resources = resources; + this.tieredIndexReference = build(); + this.destroyed = false; + } + + /** + * Constructor for loading the index from an {@link InputStream} + */ + private TieredIndexImpl(InputStream inputStream, CuVSResourcesImpl resources) throws Throwable { + throw new UnsupportedOperationException("Deserialization of TieredIndex is not yet supported"); + } + + /** + * Constructor for creating an index from an existing index reference. + */ + private TieredIndexImpl(IndexReference indexReference, CuVSResourcesImpl resources) { + this.vectors = null; + this.tieredIndexParameters = null; + this.dataset = null; + this.resources = resources; + this.tieredIndexReference = indexReference; + this.destroyed = false; + } + + private void checkNotDestroyed() { + if (destroyed) { + throw new IllegalStateException("destroyed"); + } + } + + /** + * Invokes the native destroy_tiered_index to de-allocate the Tiered index + */ + @Override + public void destroyIndex() throws Throwable { + checkNotDestroyed(); + try { + int returnValue = cuvsTieredIndexDestroy(tieredIndexReference.getMemorySegment()); + checkCuVSError(returnValue, "cuvsTieredIndexDestroy"); + } finally { + destroyed = true; + } + } + + /** + * Translates C build_tiered_index function to Java + * Invokes the native build_tiered_index function via the Panama API to build the + * {@link TieredIndex} + * + * @return an instance of {@link IndexReference} that holds the pointer to the + * index + */ + private IndexReference build() throws Throwable { + try (var localArena = Arena.ofConfined()) { + long rows = dataset != null ? dataset.size() : vectors.length; + long cols = dataset != null ? dataset.dimensions() : (rows > 0 ? vectors[0].length : 0); + + MemorySegment indexParamsMemorySegment = + tieredIndexParameters != null + ? segmentFromIndexParams(localArena, tieredIndexParameters) + : MemorySegment.NULL; + + // Get host data + MemorySegment hostDataSeg = + dataset != null + ? ((DatasetImpl) dataset).asMemorySegment() + : Util.buildMemorySegment(localArena, vectors); + + long cuvsRes = resources.getHandle(); + + // TieredIndex REQUIRES device memory - allocate it + MemorySegment datasetD = localArena.allocate(C_POINTER); + long datasetSize = C_FLOAT_BYTE_SIZE * rows * cols; + int returnValue = cuvsRMMAlloc(cuvsRes, datasetD, datasetSize); + checkCuVSError(returnValue, "cuvsRMMAlloc"); + + MemorySegment datasetDP = datasetD.get(C_POINTER, 0); + + // Copy host to device + Util.cudaMemcpy(datasetDP, hostDataSeg, datasetSize, HOST_TO_DEVICE); + + // Create tensor from device memory + long[] datasetShape = {rows, cols}; + MemorySegment datasetTensor = + prepareTensor(localArena, datasetDP, datasetShape, 2, 32, 2, 2, 1); + + MemorySegment index = localArena.allocate(cuvsTieredIndex_t); + returnValue = cuvsTieredIndexCreate(index); + checkCuVSError(returnValue, "cuvsTieredIndexCreate"); + + returnValue = cuvsStreamSync(cuvsRes); + checkCuVSError(returnValue, "cuvsStreamSync"); + + // Extract the actual index pointer that was written by Create + MemorySegment actualIndexPtr = index.get(C_POINTER, 0); + + returnValue = + cuvsTieredIndexBuild(cuvsRes, indexParamsMemorySegment, datasetTensor, actualIndexPtr); + checkCuVSError(returnValue, "cuvsTieredIndexBuild"); + + // Clean up device memory after build + returnValue = cuvsRMMFree(cuvsRes, datasetDP, datasetSize); + checkCuVSError(returnValue, "cuvsRMMFree"); + + return new IndexReference(actualIndexPtr); + } + } + + /** + * Translates C search_tiered_index function to Java + * Invokes the native search_tiered_index via the Panama API for searching a + * Tiered index. + * + * @param query an instance of {@link TieredIndexQuery} holding the query vectors and + * other parameters + * @return an instance of {@link SearchResults} containing the results + */ + @Override + public SearchResults search(TieredIndexQuery query) throws Throwable { + try (var localArena = Arena.ofConfined()) { + checkNotDestroyed(); + int topK = + query.getMapping() != null + ? Math.min(query.getMapping().size(), query.getTopK()) + : query.getTopK(); + long numQueries = query.getQueryVectors().length; + long numBlocks = (long) topK * numQueries; + int vectorDimension = numQueries > 0 ? query.getQueryVectors()[0].length : 0; + + SequenceLayout neighborsLayout = MemoryLayout.sequenceLayout(numBlocks, C_LONG); + SequenceLayout distancesLayout = MemoryLayout.sequenceLayout(numBlocks, C_FLOAT); + MemorySegment neighborsSeg = localArena.allocate(neighborsLayout); + MemorySegment distancesSeg = localArena.allocate(distancesLayout); + + // Get host query data + MemorySegment hostQueriesSeg = Util.buildMemorySegment(localArena, query.getQueryVectors()); + + long cuvsRes = resources.getHandle(); + + // Allocate DEVICE memory for all data + MemorySegment queriesD = localArena.allocate(C_POINTER); + MemorySegment neighborsD = localArena.allocate(C_POINTER); + MemorySegment distancesD = localArena.allocate(C_POINTER); + + long queriesBytes = C_FLOAT_BYTE_SIZE * numQueries * vectorDimension; + long neighborsBytes = C_LONG_BYTE_SIZE * numQueries * topK; // 64-bit for tiered index + long distancesBytes = C_FLOAT_BYTE_SIZE * numQueries * topK; + + int returnValue = cuvsRMMAlloc(cuvsRes, queriesD, queriesBytes); + checkCuVSError(returnValue, "cuvsRMMAlloc"); + returnValue = cuvsRMMAlloc(cuvsRes, neighborsD, neighborsBytes); + checkCuVSError(returnValue, "cuvsRMMAlloc"); + returnValue = cuvsRMMAlloc(cuvsRes, distancesD, distancesBytes); + checkCuVSError(returnValue, "cuvsRMMAlloc"); + + // Get device pointers + MemorySegment queriesDP = queriesD.get(C_POINTER, 0); + MemorySegment neighborsDP = neighborsD.get(C_POINTER, 0); + MemorySegment distancesDP = distancesD.get(C_POINTER, 0); + + // Copy queries from host to device + returnValue = + cudaMemcpy(queriesDP, hostQueriesSeg, queriesBytes, 1); // cudaMemcpyHostToDevice + checkCudaError(returnValue, "cudaMemcpy"); + + // Create tensors from device memory + long queriesShape[] = {numQueries, vectorDimension}; + MemorySegment queriesTensor = + prepareTensor(localArena, queriesDP, queriesShape, 2, 32, 2, 2, 1); + long neighborsShape[] = {numQueries, topK}; + MemorySegment neighborsTensor = + prepareTensor(localArena, neighborsDP, neighborsShape, 0, 64, 2, 2, 1); // 64-bit int + long distancesShape[] = {numQueries, topK}; + MemorySegment distancesTensor = + prepareTensor(localArena, distancesDP, distancesShape, 2, 32, 2, 2, 1); + + // Sync before prefilter setup + returnValue = cuvsStreamSync(cuvsRes); + checkCuVSError(returnValue, "cuvsStreamSync"); + + // Handle prefilter + MemorySegment prefilter = cuvsFilter.allocate(localArena); + MemorySegment prefilterD = localArena.allocate(C_POINTER); + MemorySegment prefilterDP = MemorySegment.NULL; + long prefilterBytes = 0; + + if (query.getPrefilter() != null) { + BitSet[] prefilters = new BitSet[] {query.getPrefilter()}; + BitSet concatenatedFilters = concatenate(prefilters, (int) query.getNumDocs()); + long filters[] = concatenatedFilters.toLongArray(); + MemorySegment hostPrefilterSeg = buildMemorySegment(localArena, filters); + + long prefilterDataLength = query.getNumDocs() * prefilters.length; + long prefilterShape[] = {(prefilterDataLength + 31) / 32}; + long prefilterLen = prefilterShape[0]; + prefilterBytes = C_INT_BYTE_SIZE * prefilterLen; + + // Allocate device memory for prefilter + returnValue = cuvsRMMAlloc(cuvsRes, prefilterD, prefilterBytes); + checkCuVSError(returnValue, "cuvsRMMAlloc"); + + prefilterDP = prefilterD.get(C_POINTER, 0); + + // Copy prefilter to device + returnValue = cudaMemcpy(prefilterDP, hostPrefilterSeg, prefilterBytes, 1); + checkCudaError(returnValue, "cudaMemcpy"); + + MemorySegment prefilterTensor = + prepareTensor(localArena, prefilterDP, prefilterShape, 1, 32, 1, 2, 1); + + cuvsFilter.type(prefilter, 1); // BITSET + cuvsFilter.addr(prefilter, prefilterTensor.address()); + } else { + cuvsFilter.type(prefilter, 0); // NO_FILTER + cuvsFilter.addr(prefilter, 0); + } + + // Perform search + returnValue = + cuvsTieredIndexSearch( + cuvsRes, + segmentFromSearchParams(query.getCagraSearchParameters(), localArena), + tieredIndexReference.getMemorySegment(), + queriesTensor, + neighborsTensor, + distancesTensor, + prefilter); + checkCuVSError(returnValue, "cuvsTieredIndexSearch"); + + // Copy results from device to host + returnValue = + cudaMemcpy(neighborsSeg, neighborsDP, neighborsBytes, 2); // cudaMemcpyDeviceToHost + checkCudaError(returnValue, "cudaMemcpy"); + returnValue = cudaMemcpy(distancesSeg, distancesDP, distancesBytes, 2); + checkCudaError(returnValue, "cudaMemcpy"); + + // Clean up device memory + returnValue = cuvsRMMFree(cuvsRes, queriesDP, queriesBytes); + checkCuVSError(returnValue, "cuvsRMMFree"); + returnValue = cuvsRMMFree(cuvsRes, neighborsDP, neighborsBytes); + checkCuVSError(returnValue, "cuvsRMMFree"); + returnValue = cuvsRMMFree(cuvsRes, distancesDP, distancesBytes); + checkCuVSError(returnValue, "cuvsRMMFree"); + + if (prefilterDP != MemorySegment.NULL) { + returnValue = cuvsRMMFree(cuvsRes, prefilterDP, prefilterBytes); + checkCuVSError(returnValue, "cuvsRMMFree"); + } + + return TieredSearchResultsImpl.create( + neighborsLayout, + distancesLayout, + neighborsSeg, + distancesSeg, + topK, + query.getMapping(), + numQueries); + } + } + + @Override + public ExtendBuilder extend() { + checkNotDestroyed(); + return new ExtendBuilder(this); + } + + /** + * Performs the actual extend operation + */ + private void performExtend(float[][] extendVectors, Dataset extendDataset) throws Throwable { + try (var localArena = Arena.ofConfined()) { + long rows = extendDataset != null ? extendDataset.size() : extendVectors.length; + long cols = extendDataset != null ? extendDataset.dimensions() : extendVectors[0].length; + + // Get host data + MemorySegment hostDataSeg = + extendDataset != null + ? ((DatasetImpl) extendDataset).asMemorySegment() + : Util.buildMemorySegment(localArena, extendVectors); + + long cuvsRes = resources.getHandle(); + + // Allocate device memory for extend data + MemorySegment datasetD = localArena.allocate(C_POINTER); + long dataSize = C_FLOAT_BYTE_SIZE * rows * cols; + int returnValue = cuvsRMMAlloc(cuvsRes, datasetD, dataSize); + checkCuVSError(returnValue, "cuvsRMMAlloc"); + + MemorySegment datasetDP = datasetD.get(C_POINTER, 0); + + // Copy host to device + returnValue = cudaMemcpy(datasetDP, hostDataSeg, dataSize, 1); // cudaMemcpyHostToDevice + checkCudaError(returnValue, "cudaMemcpy"); + + // Create tensor from device memory + long datasetShape[] = {rows, cols}; + MemorySegment datasetTensor = + prepareTensor(localArena, datasetDP, datasetShape, 2, 32, 2, 2, 1); + + returnValue = cuvsStreamSync(cuvsRes); + checkCuVSError(returnValue, "cuvsStreamSync"); + + returnValue = + cuvsTieredIndexExtend(cuvsRes, datasetTensor, tieredIndexReference.getMemorySegment()); + checkCuVSError(returnValue, "cuvsTieredIndexExtend"); + + // Clean up device memory + returnValue = cuvsRMMFree(cuvsRes, datasetDP, dataSize); + checkCuVSError(returnValue, "cuvsRMMFree"); + } + } + + /** + * ExtendBuilder implementation + */ + public static class ExtendBuilder implements TieredIndex.ExtendBuilder { + private final TieredIndexImpl index; + private float[][] vectors; + private Dataset dataset; + + private ExtendBuilder(TieredIndexImpl index) { + this.index = index; + } + + @Override + public ExtendBuilder withDataset(float[][] vectors) { + this.vectors = vectors; + return this; + } + + @Override + public ExtendBuilder withDataset(Dataset dataset) { + this.dataset = dataset; + return this; + } + + @Override + public void execute() throws Throwable { + if (vectors != null && dataset != null) { + throw new IllegalArgumentException( + "Please specify only one type of dataset (a float[][] or a Dataset instance)"); + } + if (vectors == null && dataset == null) { + throw new IllegalArgumentException("Must provide vectors or dataset"); + } + + index.performExtend(vectors, dataset); + } + } + + /** + * Allocates the configured index parameters in the MemorySegment. + */ + private static MemorySegment segmentFromIndexParams(Arena arena, TieredIndexParams params) { + MemorySegment seg = cuvsTieredIndexParams.allocate(arena); + + // Get the metric from CagraParams if available, otherwise use TieredIndex metric + int metric; + if (params.getCagraParams() != null) { + // Use the metric from CagraParams to ensure consistency + metric = params.getCagraParams().getCuvsDistanceType().value; + } else { + // Fallback to TieredIndex metric + metric = + switch (params.getMetric()) { + case L2 -> 0; + case INNER_PRODUCT -> 1; + default -> + throw new IllegalArgumentException("Unsupported metric: " + params.getMetric()); + }; + } + + cuvsTieredIndexParams.metric(seg, metric); + + int algo = 0; // CUVS_TIERED_INDEX_ALGO_CAGRA + cuvsTieredIndexParams.algo(seg, algo); + + cuvsTieredIndexParams.min_ann_rows(seg, params.getMinAnnRows()); + cuvsTieredIndexParams.create_ann_index_on_extend(seg, params.isCreateAnnIndexOnExtend()); + + CagraIndexParams cagraParams = params.getCagraParams(); + if (cagraParams != null) { + MemorySegment cagraParamsSeg = cuvsCagraIndexParams.allocate(arena); + + cuvsCagraIndexParams.intermediate_graph_degree( + cagraParamsSeg, cagraParams.getIntermediateGraphDegree()); + cuvsCagraIndexParams.graph_degree(cagraParamsSeg, cagraParams.getGraphDegree()); + cuvsCagraIndexParams.build_algo(cagraParamsSeg, cagraParams.getCagraGraphBuildAlgo().value); + cuvsCagraIndexParams.nn_descent_niter( + cagraParamsSeg, cagraParams.getNNDescentNumIterations()); + cuvsCagraIndexParams.metric(cagraParamsSeg, metric); + + cuvsTieredIndexParams.cagra_params(seg, cagraParamsSeg); + } + + cuvsTieredIndexParams.ivf_flat_params(seg, MemorySegment.NULL); + cuvsTieredIndexParams.ivf_pq_params(seg, MemorySegment.NULL); + + return seg; + } + + /** + * Allocates the configured search parameters in the MemorySegment. + */ + private MemorySegment segmentFromSearchParams(CagraSearchParams params, Arena arena) { + MemorySegment seg = cuvsCagraSearchParams.allocate(arena); + cuvsCagraSearchParams.max_queries(seg, params.getMaxQueries()); + cuvsCagraSearchParams.itopk_size(seg, params.getITopKSize()); + cuvsCagraSearchParams.max_iterations(seg, params.getMaxIterations()); + if (params.getCagraSearchAlgo() != null) { + cuvsCagraSearchParams.algo(seg, params.getCagraSearchAlgo().value); + } + cuvsCagraSearchParams.team_size(seg, params.getTeamSize()); + cuvsCagraSearchParams.search_width(seg, params.getSearchWidth()); + cuvsCagraSearchParams.min_iterations(seg, params.getMinIterations()); + cuvsCagraSearchParams.thread_block_size(seg, params.getThreadBlockSize()); + if (params.getHashMapMode() != null) { + cuvsCagraSearchParams.hashmap_mode(seg, params.getHashMapMode().value); + } + cuvsCagraSearchParams.hashmap_max_fill_rate(seg, params.getHashMapMaxFillRate()); + cuvsCagraSearchParams.num_random_samplings(seg, params.getNumRandomSamplings()); + cuvsCagraSearchParams.rand_xor_mask(seg, params.getRandXORMask()); + return seg; + } + + /** + * Gets an instance of {@link CuVSResources} + * + * @return an instance of {@link CuVSResources} + */ + @Override + public CuVSResources getCuVSResources() { + return resources; + } + + /** + * Gets the index type + * + * @return the index type + */ + @Override + public TieredIndexType getIndexType() { + TieredIndexType indexType = TieredIndexType.CAGRA; // Default to CAGRA for now + return indexType; + } + + /** + * Static method to create a new builder + */ + public static TieredIndex.Builder newBuilder(CuVSResources cuvsResources) { + Objects.requireNonNull(cuvsResources); + if (!(cuvsResources instanceof CuVSResourcesImpl)) { + throw new IllegalArgumentException("Unsupported " + cuvsResources); + } + return new TieredIndexImpl.Builder((CuVSResourcesImpl) cuvsResources); + } + + /** + * Builder helps configure and create an instance of {@link TieredIndex}. + */ + public static class Builder implements TieredIndex.Builder { + private CuVSResourcesImpl resources; + private float[][] vectors; + private Dataset dataset; + private TieredIndexParams params; + private TieredIndexType indexType = TieredIndexType.CAGRA; + private InputStream inputStream; + + private Builder(CuVSResourcesImpl resources) { + this.resources = resources; + } + + @Override + public Builder from(InputStream inputStream) { + this.inputStream = inputStream; + return this; + } + + @Override + public Builder withDataset(float[][] vectors) { + this.vectors = vectors; + return this; + } + + @Override + public Builder withDataset(Dataset dataset) { + this.dataset = dataset; + return this; + } + + @Override + public Builder withIndexParams(TieredIndexParams params) { + this.params = params; + return this; + } + + @Override + public Builder withIndexType(TieredIndexType indexType) { + this.indexType = indexType; + return this; + } + + @Override + public TieredIndex build() throws Throwable { + if (inputStream != null) { + return new TieredIndexImpl(inputStream, resources); + } else { + if (vectors != null && dataset != null) { + throw new IllegalArgumentException( + "Please specify only one type of dataset (a float[][] or a Dataset instance)"); + } + if (vectors == null && dataset == null) { + throw new IllegalArgumentException("Must provide vectors or dataset"); + } + if (params == null) { + throw new IllegalStateException("Index parameters must be provided"); + } + return new TieredIndexImpl(params, vectors, dataset, resources); + } + } + } + + /** + * Holds the memory reference to a Tiered index. + */ + public static class IndexReference { + private final MemorySegment memorySegment; + + /** + * Constructs TieredIndexReference and allocate the MemorySegment. + */ + protected IndexReference(CuVSResourcesImpl resources) { + // Don't allocate here - the C function will allocate + memorySegment = MemorySegment.NULL; + } + + /** + * Constructs TieredIndexReference with an instance of MemorySegment passed as a + * parameter. + */ + protected IndexReference(MemorySegment indexMemorySegment) { + this.memorySegment = indexMemorySegment; + } + + /** + * Gets the instance of index MemorySegment. + */ + protected MemorySegment getMemorySegment() { + return memorySegment; + } + } +} diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredSearchResultsImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredSearchResultsImpl.java new file mode 100644 index 0000000000..13c1ee93fa --- /dev/null +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredSearchResultsImpl.java @@ -0,0 +1,62 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs.internal; + +import com.nvidia.cuvs.SearchResults; +import java.lang.foreign.MemorySegment; +import java.lang.foreign.SequenceLayout; +import java.util.List; +import java.util.Map; +import java.util.function.LongToIntFunction; + +class TieredSearchResultsImpl implements SearchResults { + private final List> results; + + private TieredSearchResultsImpl(List> results) { + this.results = results; + } + + public static TieredSearchResultsImpl create( + SequenceLayout neighboursSequenceLayout, + SequenceLayout distancesSequenceLayout, + MemorySegment neighboursMemorySegment, + MemorySegment distancesMemorySegment, + int topK, + List mapping, + long numberOfQueries) { + + // Process the data immediately while the memory segments are still valid + LongToIntFunction mappingFunction = mapping != null ? (long id) -> mapping.get((int) id) : null; + + List> processedResults = + SearchResultsImpl.create( + neighboursSequenceLayout, + distancesSequenceLayout, + neighboursMemorySegment, + distancesMemorySegment, + topK, + mappingFunction, + numberOfQueries) + .getResults(); + + return new TieredSearchResultsImpl(processedResults); + } + + @Override + public List> getResults() { + return results; + } +} diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java index 2b9913f04f..5940e3db48 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java @@ -23,11 +23,13 @@ import com.nvidia.cuvs.CuVSResources; import com.nvidia.cuvs.Dataset; import com.nvidia.cuvs.HnswIndex; +import com.nvidia.cuvs.TieredIndex; import com.nvidia.cuvs.internal.BruteForceIndexImpl; import com.nvidia.cuvs.internal.CagraIndexImpl; import com.nvidia.cuvs.internal.CuVSResourcesImpl; import com.nvidia.cuvs.internal.DatasetImpl; import com.nvidia.cuvs.internal.HnswIndexImpl; +import com.nvidia.cuvs.internal.TieredIndexImpl; import com.nvidia.cuvs.internal.common.Util; import java.lang.foreign.Arena; import java.lang.foreign.MemoryLayout; @@ -85,6 +87,11 @@ public HnswIndex.Builder newHnswIndexBuilder(CuVSResources cuVSResources) { return HnswIndexImpl.newBuilder(Objects.requireNonNull(cuVSResources)); } + @Override + public TieredIndex.Builder newTieredIndexBuilder(CuVSResources cuVSResources) { + return TieredIndexImpl.newBuilder(Objects.requireNonNull(cuVSResources)); + } + @Override public CagraIndex mergeCagraIndexes(CagraIndex[] indexes) throws Throwable { if (indexes == null || indexes.length == 0) { diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/TieredIndexIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/TieredIndexIT.java new file mode 100644 index 0000000000..ffe1dddd08 --- /dev/null +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/TieredIndexIT.java @@ -0,0 +1,259 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +import com.carrotsearch.randomizedtesting.RandomizedRunner; +import java.lang.invoke.MethodHandles; +import java.util.Arrays; +import java.util.BitSet; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +@RunWith(RandomizedRunner.class) +public class TieredIndexIT extends CuVSTestCase { + + private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass()); + + @Before + public void setup() { + initializeRandom(); + log.debug("Random context initialized for test"); + } + + @Test + public void testBasicOperations() throws Throwable { + float[][] initialDataset = { + {0.0f, 0.0f}, + {1.0f, 1.0f}, + {2.0f, 2.0f} + }; + + float[][] queries = { + {0.1f, 0.1f}, + {1.9f, 1.9f} + }; + + float[][] extensionVectors = { + {3.0f, 3.0f}, + {4.0f, 4.0f} + }; + + List> expectedInitialResults = + Arrays.asList(Map.of(0, 0.02f, 1, 1.62f, 2, 7.22f), Map.of(2, 0.02f, 1, 1.62f, 0, 7.22f)); + + List> expectedExtendedResults = + Arrays.asList(Map.of(0, 0.02f, 1, 1.62f, 2, 7.22f), Map.of(2, 0.02f, 3, 2.42f, 1, 1.62f)); + + try (CuVSResources resources = CuVSResources.create()) { + CagraIndexParams cagraParams = + new CagraIndexParams.Builder().withGraphDegree(4).withIntermediateGraphDegree(8).build(); + + TieredIndexParams indexParams = + new TieredIndexParams.Builder() + .minAnnRows(2) + .createAnnIndexOnExtend(true) + .withCagraParams(cagraParams) + .build(); + + log.debug("Building initial TieredIndex with {} vectors", initialDataset.length); + TieredIndex index = + TieredIndex.newBuilder(resources) + .withDataset(initialDataset) + .withIndexParams(indexParams) + .build(); + log.debug("Initial TieredIndex built successfully"); + + CagraSearchParams searchParams = + new CagraSearchParams.Builder(resources).withMaxIterations(20).build(); + + TieredIndexQuery query = + new TieredIndexQuery.Builder() + .withTopK(3) + .withQueryVectors(queries) + .withSearchParams(searchParams) + .build(); + + log.debug("Searching initial index with {} queries", queries.length); + SearchResults initialResults = index.search(query); + log.debug("Initial search completed, validating results"); + assertEquals(expectedInitialResults, roundResults(initialResults.getResults())); + + log.debug("Extending index with {} additional vectors", extensionVectors.length); + index.extend().withDataset(extensionVectors).execute(); + log.debug("Index extension completed"); + + log.debug("Searching extended index"); + SearchResults extendedResults = index.search(query); + log.debug("Extended search completed, validating results"); + assertEquals(expectedExtendedResults, roundResults(extendedResults.getResults())); + } + } + + @Test(expected = IllegalArgumentException.class) + public void testErrorHandling() throws Throwable { + try (CuVSResources resources = CuVSResources.create()) { + CagraIndexParams cagraParams = + new CagraIndexParams.Builder().withGraphDegree(4).withIntermediateGraphDegree(8).build(); + + TieredIndexParams indexParams = + new TieredIndexParams.Builder().minAnnRows(2).withCagraParams(cagraParams).build(); + + log.debug("Testing error handling with null dataset"); + TieredIndex.newBuilder(resources) + .withIndexParams(indexParams) + .withDataset((float[][]) null) + .build(); + } + } + + @Test + public void testDifferentKValues() throws Throwable { + float[][] dataset = { + {0.0f, 0.0f}, + {1.0f, 1.0f}, + {2.0f, 2.0f}, + {3.0f, 3.0f}, + {4.0f, 4.0f} + }; + + float[][] queries = {{0.1f, 0.1f}}; + + try (CuVSResources resources = CuVSResources.create()) { + CagraIndexParams cagraParams = + new CagraIndexParams.Builder().withGraphDegree(4).withIntermediateGraphDegree(8).build(); + + TieredIndexParams indexParams = + new TieredIndexParams.Builder().minAnnRows(2).withCagraParams(cagraParams).build(); + + log.debug("Building TieredIndex for K-value testing with {} vectors", dataset.length); + TieredIndex index = + TieredIndex.newBuilder(resources) + .withDataset(dataset) + .withIndexParams(indexParams) + .build(); + log.debug("TieredIndex built for K-value testing"); + + TieredIndexQuery query1 = + new TieredIndexQuery.Builder() + .withTopK(1) + .withQueryVectors(queries) + .withSearchParams( + new CagraSearchParams.Builder(resources).withMaxIterations(20).build()) + .build(); + + log.debug("Searching with K=1"); + SearchResults results1 = index.search(query1); + Map firstResult = results1.getResults().get(0); + log.debug("K=1 search completed, found {} results", firstResult.size()); + + assertEquals(1, firstResult.size()); + assertTrue("Should contain index 0 (closest vector)", firstResult.containsKey(0)); + assertEquals("Distance to closest vector should be ~0.02", 0.02f, firstResult.get(0), 0.01f); + + TieredIndexQuery query3 = + new TieredIndexQuery.Builder() + .withTopK(3) + .withQueryVectors(queries) + .withSearchParams( + new CagraSearchParams.Builder(resources).withMaxIterations(20).build()) + .build(); + + log.debug("Searching with K=3"); + SearchResults results3 = index.search(query3); + Map thirdResult = results3.getResults().get(0); + log.debug("K=3 search completed, found {} results", thirdResult.size()); + + assertEquals(3, thirdResult.size()); + assertTrue("Should contain index 0", thirdResult.containsKey(0)); + assertTrue("Should contain index 1", thirdResult.containsKey(1)); + assertTrue("Should contain index 2", thirdResult.containsKey(2)); + + float dist0 = thirdResult.get(0); + float dist1 = thirdResult.get(1); + float dist2 = thirdResult.get(2); + + assertTrue("Distance to index 0 should be smallest", dist0 <= dist1 && dist0 <= dist2); + } + } + + @Test + public void testPrefilter() throws Throwable { + float[][] dataset = {{0.0f, 0.0f}, {1.0f, 1.0f}, {2.0f, 2.0f}, {3.0f, 3.0f}}; + float[][] queryVectors = {{0.1f, 0.1f}}; + + try (CuVSResources resources = CuVSResources.create()) { + CagraIndexParams cagraParams = + new CagraIndexParams.Builder().withGraphDegree(4).withIntermediateGraphDegree(8).build(); + + TieredIndexParams indexParams = + new TieredIndexParams.Builder().minAnnRows(2).withCagraParams(cagraParams).build(); + + log.debug("Building TieredIndex for prefilter testing with {} vectors", dataset.length); + TieredIndex index = + TieredIndex.newBuilder(resources) + .withDataset(dataset) + .withIndexParams(indexParams) + .build(); + log.debug("TieredIndex built for prefilter testing"); + + CagraSearchParams searchParams = new CagraSearchParams.Builder(resources).build(); + + BitSet prefilter = new BitSet(4); + prefilter.set(1, true); + prefilter.set(2, true); + log.debug("Created prefilter allowing indices 1 and 2, excluding 0 and 3"); + + TieredIndexQuery queryWithFilter = + new TieredIndexQuery.Builder() + .withTopK(3) + .withQueryVectors(queryVectors) + .withSearchParams(searchParams) + .withPrefilter(prefilter, 4) + .build(); + + log.debug("Searching with prefilter applied"); + SearchResults resultsWithFilter = index.search(queryWithFilter); + Map result = resultsWithFilter.getResults().get(0); + log.debug("Prefilter search completed, validating filtered results"); + + assertFalse("Index 0 should be filtered out", result.containsKey(0)); + assertFalse("Index 3 should be filtered out", result.containsKey(3)); + assertTrue("Index 1 or 2 should be present", result.containsKey(1) || result.containsKey(2)); + } + } + + private List> roundResults(List> results) { + return results.stream() + .map( + queryResult -> + queryResult.entrySet().stream() + .collect( + Collectors.toMap( + Map.Entry::getKey, + entry -> Math.round(entry.getValue() * 100.0f) / 100.0f))) + .collect(Collectors.toList()); + } +} diff --git a/java/panama-bindings/headers.h b/java/panama-bindings/headers.h index 08a381af14..57a13cf1b1 100644 --- a/java/panama-bindings/headers.h +++ b/java/panama-bindings/headers.h @@ -22,6 +22,7 @@ #include #include #include +#include #include /**