From 77c624949cc9851d95f2d57147d25f14f23c47ac Mon Sep 17 00:00:00 2001 From: punAhuja Date: Tue, 17 Jun 2025 15:32:52 +0530 Subject: [PATCH 01/11] [Java] Support for tiered index exposed --- .../java/com/nvidia/cuvs/TieredIndex.java | 195 ++++++ .../com/nvidia/cuvs/TieredIndexParams.java | 175 +++++ .../com/nvidia/cuvs/TieredIndexQuery.java | 224 +++++++ .../com/nvidia/cuvs/spi/CuVSProvider.java | 5 + .../nvidia/cuvs/spi/UnsupportedProvider.java | 6 + .../nvidia/cuvs/internal/TieredIndexImpl.java | 634 ++++++++++++++++++ .../internal/TieredSearchResultsImpl.java | 64 ++ .../com/nvidia/cuvs/spi/JDKProvider.java | 7 + .../java/com/nvidia/cuvs/TieredIndexIT.java | 267 ++++++++ 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..a08e01b968 --- /dev/null +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndex.java @@ -0,0 +1,195 @@ +/* + * 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.io.InputStream; +import java.util.Objects; + +import com.nvidia.cuvs.spi.CuVSProvider; + +/** + * {@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 configuration parameters used to build this TieredIndex. + * + * @return An instance of {@link TieredIndexParams} containing the index + * configuration + */ + TieredIndexParams getIndexParameters(); + + /** + * 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..be271c6758 --- /dev/null +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndexParams.java @@ -0,0 +1,175 @@ +/* + * 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..2dc69aaf66 --- /dev/null +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndexQuery.java @@ -0,0 +1,224 @@ +/* + * 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.Arrays; +import java.util.List; + +import com.nvidia.cuvs.TieredIndex.TieredIndexType; + +import java.util.BitSet; + +/** + * 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 a0d0f9a55c..37532829b9 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 com.nvidia.cuvs.CagraMergeParams; import java.nio.file.Path; @@ -67,6 +68,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 5fe372dce0..dc63e330cd 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 @@ -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.nio.file.Path; @@ -49,6 +50,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..dd3e860507 --- /dev/null +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java @@ -0,0 +1,634 @@ +/* + * 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 java.lang.foreign.ValueLayout.ADDRESS; +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; +import static com.nvidia.cuvs.internal.common.LinkerHelper.C_INT_BYTE_SIZE; +import static com.nvidia.cuvs.internal.common.LinkerHelper.C_POINTER; +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.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.common.Util.prepareTensor; +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 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.cuvsResources_t; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsStreamGet; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsStreamSync; +import static com.nvidia.cuvs.internal.panama.headers_h.cudaMemcpy; +import static com.nvidia.cuvs.internal.panama.headers_h.cudaStream_t; + +import java.io.InputStream; +import java.io.OutputStream; +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; + +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.cuvsTieredIndex; +import com.nvidia.cuvs.internal.panama.cuvsTieredIndexParams; +import com.nvidia.cuvs.internal.panama.cuvsFilter; + +/** + * {@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; + } + if (dataset != null) dataset.close(); + } + + /** + * 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(resources, tieredIndexParameters) + : MemorySegment.NULL; + + // Get host data + MemorySegment hostDataSeg = dataset != null ? ((DatasetImpl) dataset).seg : + Util.buildMemorySegment(resources.getArena(), vectors); + + Arena arena = resources.getArena(); + long cuvsRes = resources.getMemorySegment().get(cuvsResources_t, 0); + MemorySegment stream = arena.allocate(cudaStream_t); + var returnValue = cuvsStreamGet(cuvsRes, stream); + checkCuVSError(returnValue, "cuvsStreamGet"); + + // TieredIndex REQUIRES device memory - allocate it + MemorySegment datasetD = arena.allocate(C_POINTER); + long datasetSize = C_FLOAT_BYTE_SIZE * rows * cols; + returnValue = cuvsRMMAlloc(cuvsRes, datasetD, datasetSize); + checkCuVSError(returnValue, "cuvsRMMAlloc"); + + MemorySegment datasetDP = datasetD.get(C_POINTER, 0); + + // Copy host to device + returnValue = cudaMemcpy(datasetDP, hostDataSeg, datasetSize, 1); // cudaMemcpyHostToDevice + checkCudaError(returnValue, "cudaMemcpy"); + + // Create tensor from device memory + long datasetShape[] = { rows, cols }; + MemorySegment datasetTensor = prepareTensor(arena, datasetDP, datasetShape, 2, 32, 2, 2, 1); + + MemorySegment index = arena.allocate(cuvsTieredIndex_t); + returnValue = cuvsTieredIndexCreate(index); + checkCuVSError(returnValue, "cuvsTieredIndexCreate"); + + returnValue = cuvsStreamSync(cuvsRes); + checkCuVSError(returnValue, "cuvsStreamSync"); + + returnValue = cuvsTieredIndexBuild(cuvsRes, indexParamsMemorySegment, datasetTensor, index); + checkCuVSError(returnValue, "cuvsTieredIndexBuild"); + + // Clean up device memory after build + returnValue = cuvsRMMFree(cuvsRes, datasetDP, datasetSize); + checkCuVSError(returnValue, "cuvsRMMFree"); + + return new IndexReference(index); + } + } + + /** + * 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; + Arena arena = resources.getArena(); + + // Allocate HOST memory for results + SequenceLayout neighborsLayout = MemoryLayout.sequenceLayout(numBlocks, C_LONG); + SequenceLayout distancesLayout = MemoryLayout.sequenceLayout(numBlocks, C_FLOAT); + MemorySegment neighborsSeg = arena.allocate(neighborsLayout); + MemorySegment distancesSeg = arena.allocate(distancesLayout); + + // Get host query data + MemorySegment hostQueriesSeg = Util.buildMemorySegment(arena, query.getQueryVectors()); + + long cuvsRes = resources.getMemorySegment().get(cuvsResources_t, 0); + MemorySegment stream = arena.allocate(cudaStream_t); + int returnValue = cuvsStreamGet(cuvsRes, stream); + checkCuVSError(returnValue, "cuvsStreamGet"); + + // Allocate DEVICE memory for all data + MemorySegment queriesD = arena.allocate(C_POINTER); + MemorySegment neighborsD = arena.allocate(C_POINTER); + MemorySegment distancesD = arena.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; + + 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(arena, queriesDP, queriesShape, 2, 32, 2, 2, 1); + long neighborsShape[] = { numQueries, topK }; + MemorySegment neighborsTensor = prepareTensor(arena, neighborsDP, neighborsShape, 0, 64, 2, 2, 1); // 64-bit int + long distancesShape[] = { numQueries, topK }; + MemorySegment distancesTensor = prepareTensor(arena, distancesDP, distancesShape, 2, 32, 2, 2, 1); + + // Sync before prefilter setup + returnValue = cuvsStreamSync(cuvsRes); + checkCuVSError(returnValue, "cuvsStreamSync"); + + // Handle prefilter + MemorySegment prefilter = cuvsFilter.allocate(arena); + MemorySegment prefilterD = arena.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(arena, 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(arena, 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()), + 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 new TieredSearchResultsImpl( + 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 { + 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).seg + : Util.buildMemorySegment(resources.getArena(), extendVectors); + + Arena arena = resources.getArena(); + long cuvsRes = resources.getMemorySegment().get(cuvsResources_t, 0); + MemorySegment stream = arena.allocate(cudaStream_t); + int returnValue = cuvsStreamGet(cuvsRes, stream); + checkCuVSError(returnValue, "cuvsStreamGet"); + + // Allocate device memory for extend data + MemorySegment datasetD = arena.allocate(C_POINTER); + long dataSize = C_FLOAT_BYTE_SIZE * rows * cols; + 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(arena, 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(CuVSResourcesImpl resources, TieredIndexParams params) { + MemorySegment seg = cuvsTieredIndexParams.allocate(resources.getArena()); + + // 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(resources.getArena()); + + 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) { + MemorySegment seg = cuvsCagraSearchParams.allocate(resources.getArena()); + 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 TieredIndexParams} + * + * @return an instance of {@link TieredIndexParams} + */ + @Override + public TieredIndexParams getIndexParameters() { + return tieredIndexParameters; + } + + /** + * 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 = tieredIndexParameters != null && tieredIndexParameters.getCagraParams() != null + ? TieredIndexType.CAGRA + : 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; + } + } +} \ No newline at end of file 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..3e0d5d62f0 --- /dev/null +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredSearchResultsImpl.java @@ -0,0 +1,64 @@ +/* + * 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 java.lang.foreign.MemorySegment; +import java.lang.foreign.SequenceLayout; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +import com.nvidia.cuvs.internal.common.SearchResultsImpl; + +/** + * Implementation of SearchResults for TieredIndex + */ +class TieredSearchResultsImpl extends SearchResultsImpl { + + protected TieredSearchResultsImpl(SequenceLayout neighboursSequenceLayout, SequenceLayout distancesSequenceLayout, + MemorySegment neighboursMemorySegment, MemorySegment distancesMemorySegment, int topK, + List mapping, + long numberOfQueries) { + super(neighboursSequenceLayout, distancesSequenceLayout, neighboursMemorySegment, distancesMemorySegment, topK, + mapping, + numberOfQueries); + readResultMemorySegments(); + } + + /** + * Reads neighbors and distances {@link MemorySegment} and loads the values + * internally into the results structure. + */ + @Override + protected void readResultMemorySegments() { + Map intermediateResultMap = new LinkedHashMap(); + int count = 0; + for (long i = 0; i < topK * numberOfQueries; i++) { + long neighborIdLong = (long) neighboursVarHandle.get(neighboursMemorySegment, 0L, i); + float dst = (float) distancesVarHandle.get(distancesMemorySegment, 0L, i); + if (neighborIdLong != -1L && neighborIdLong != Long.MAX_VALUE) { + int id = (int) neighborIdLong; + intermediateResultMap.put(mapping != null ? mapping.get(id) : id, dst); + } + count++; + if (count == topK) { + results.add(intermediateResultMap); + intermediateResultMap = new LinkedHashMap(); + count = 0; + } + } + } +} 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 e7a2a73e27..81d662c1c5 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 @@ -21,12 +21,14 @@ import com.nvidia.cuvs.CuVSResources; import com.nvidia.cuvs.Dataset; import com.nvidia.cuvs.HnswIndex; +import com.nvidia.cuvs.TieredIndex; import com.nvidia.cuvs.CagraMergeParams; 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 java.nio.file.Files; import java.nio.file.Path; @@ -61,6 +63,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..cb27a4811d --- /dev/null +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/TieredIndexIT.java @@ -0,0 +1,267 @@ +/* + * 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 java.util.Arrays; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; +import java.util.BitSet; +import java.lang.invoke.MethodHandles; + +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import com.carrotsearch.randomizedtesting.RandomizedRunner; + +@RunWith(RandomizedRunner.class) +public class TieredIndexIT extends CuVSTestCase { + + private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass()); + + @Before + public void setup() { + initializeRandom(); + log.info("Random context initialized for test."); + } + + /** + * Tests basic operations of TieredIndex - build, search, and extend. + */ + @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.info("Building initial index..."); + TieredIndex index = TieredIndex.newBuilder(resources) + .withDataset(initialDataset) + .withIndexParams(indexParams) + .build(); + + CagraSearchParams searchParams = new CagraSearchParams.Builder(resources) + .withMaxIterations(20) + .build(); + + TieredIndexQuery query = new TieredIndexQuery.Builder() + .withTopK(3) + .withQueryVectors(queries) + .withSearchParams(searchParams) + .build(); + + log.info("Searching initial index..."); + SearchResults initialResults = index.search(query); + log.info("Initial search results: {}", initialResults.getResults()); + assertEquals(expectedInitialResults, roundResults(initialResults.getResults())); + + log.info("Extending index..."); + index.extend() + .withDataset(extensionVectors) + .execute(); + + log.info("Searching extended index..."); + SearchResults extendedResults = index.search(query); + log.info("Extended search results: {}", extendedResults.getResults()); + assertEquals(expectedExtendedResults, roundResults(extendedResults.getResults())); + } + } + + /** + * Tests error handling and parameter validation. + */ + @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(); + + TieredIndex.newBuilder(resources) + .withIndexParams(indexParams) + .withDataset((float[][]) null) + .build(); + } + } + + /** + * Tests search with different K values. + */ + @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(); + + TieredIndex index = TieredIndex.newBuilder(resources) + .withDataset(dataset) + .withIndexParams(indexParams) + .build(); + + TieredIndexQuery query1 = new TieredIndexQuery.Builder() + .withTopK(1) + .withQueryVectors(queries) + .withSearchParams(new CagraSearchParams.Builder(resources) + .withMaxIterations(20) + .build()) + .build(); + + SearchResults results1 = index.search(query1); + assertEquals(1, results1.getResults().get(0).size()); + + TieredIndexQuery query3 = new TieredIndexQuery.Builder() + .withTopK(3) + .withQueryVectors(queries) + .withSearchParams(new CagraSearchParams.Builder(resources) + .withMaxIterations(20) + .build()) + .build(); + + SearchResults results3 = index.search(query3); + assertEquals(3, results3.getResults().get(0).size()); + } + } + + /** + * Test prefilter functionality with debug logging + */ + @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(); + + TieredIndex index = TieredIndex.newBuilder(resources) + .withDataset(dataset) + .withIndexParams(indexParams) + .build(); + + CagraSearchParams searchParams = new CagraSearchParams.Builder(resources).build(); + + TieredIndexQuery queryWithoutFilter = new TieredIndexQuery.Builder() + .withTopK(3) + .withQueryVectors(queryVectors) + .withSearchParams(searchParams) + .build(); + + SearchResults resultsWithoutFilter = index.search(queryWithoutFilter); + log.info("Results WITHOUT prefilter: {}", resultsWithoutFilter.getResults()); + + BitSet prefilter = new BitSet(4); + prefilter.set(1, true); + prefilter.set(2, true); + // Index 0 and 3 are NOT set, so they should be excluded + + TieredIndexQuery queryWithFilter = new TieredIndexQuery.Builder() + .withTopK(3) + .withQueryVectors(queryVectors) + .withSearchParams(searchParams) + .withPrefilter(prefilter, 4) + .build(); + + SearchResults resultsWithFilter = index.search(queryWithFilter); + log.info("Results WITH prefilter: {}", resultsWithFilter.getResults()); + + Map result = resultsWithFilter.getResults().get(0); + + assertFalse("Index 0 should be filtered out", result.containsKey(0)); + 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 /** From 0916665ecbb28582eff4f0925b135669e3c63a5b Mon Sep 17 00:00:00 2001 From: punAhuja Date: Tue, 17 Jun 2025 15:48:23 +0530 Subject: [PATCH 02/11] Whitespace fixes --- .../java/com/nvidia/cuvs/TieredIndexQuery.java | 2 +- .../com/nvidia/cuvs/internal/TieredIndexImpl.java | 14 +++++++------- .../test/java/com/nvidia/cuvs/TieredIndexIT.java | 4 ++-- 3 files changed, 10 insertions(+), 10 deletions(-) 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 index 2dc69aaf66..42ba133038 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndexQuery.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndexQuery.java @@ -122,7 +122,7 @@ public String toString() { /** * Creates a new Builder instance. - * + * * @return a new Builder instance */ public static Builder newBuilder() { 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 index dd3e860507..f782eb810b 100644 --- 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 @@ -85,7 +85,7 @@ public class TieredIndexImpl implements TieredIndex { private final IndexReference tieredIndexReference; private boolean destroyed; - + /** * Constructor for building the index using specified dataset @@ -225,7 +225,7 @@ public SearchResults search(TieredIndexQuery query) throws Throwable { SequenceLayout distancesLayout = MemoryLayout.sequenceLayout(numBlocks, C_FLOAT); MemorySegment neighborsSeg = arena.allocate(neighborsLayout); MemorySegment distancesSeg = arena.allocate(distancesLayout); - + // Get host query data MemorySegment hostQueriesSeg = Util.buildMemorySegment(arena, query.getQueryVectors()); @@ -282,7 +282,7 @@ public SearchResults search(TieredIndexQuery query) throws Throwable { BitSet concatenatedFilters = concatenate(prefilters, (int)query.getNumDocs()); long filters[] = concatenatedFilters.toLongArray(); MemorySegment hostPrefilterSeg = buildMemorySegment(arena, filters); - + long prefilterDataLength = query.getNumDocs() * prefilters.length; long prefilterShape[] = { (prefilterDataLength + 31) / 32 }; long prefilterLen = prefilterShape[0]; @@ -443,7 +443,7 @@ private static MemorySegment segmentFromIndexParams(CuVSResourcesImpl resources, default -> throw new IllegalArgumentException("Unsupported metric: " + params.getMetric()); }; } - + cuvsTieredIndexParams.metric(seg, metric); int algo = 0; // CUVS_TIERED_INDEX_ALGO_CAGRA @@ -522,8 +522,8 @@ public CuVSResources getCuVSResources() { */ @Override public TieredIndexType getIndexType() { - TieredIndexType indexType = tieredIndexParameters != null && tieredIndexParameters.getCagraParams() != null - ? TieredIndexType.CAGRA + TieredIndexType indexType = tieredIndexParameters != null && tieredIndexParameters.getCagraParams() != null + ? TieredIndexType.CAGRA : TieredIndexType.CAGRA; // Default to CAGRA for now return indexType; } @@ -631,4 +631,4 @@ protected MemorySegment getMemorySegment() { return memorySegment; } } -} \ No newline at end of file +} 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 index cb27a4811d..48442a0c7a 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/TieredIndexIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/TieredIndexIT.java @@ -248,9 +248,9 @@ public void testPrefilter() throws Throwable { SearchResults resultsWithFilter = index.search(queryWithFilter); log.info("Results WITH prefilter: {}", resultsWithFilter.getResults()); - + Map result = resultsWithFilter.getResults().get(0); - + assertFalse("Index 0 should be filtered out", result.containsKey(0)); assertTrue("Index 1 or 2 should be present", result.containsKey(1) || result.containsKey(2)); } From 09215a64e7fcf8f7df545cc61c063fb8e7f8a408 Mon Sep 17 00:00:00 2001 From: punAhuja Date: Tue, 1 Jul 2025 20:16:11 +0530 Subject: [PATCH 03/11] Code format fix --- .../java/com/nvidia/cuvs/TieredIndex.java | 4 +- .../com/nvidia/cuvs/TieredIndexParams.java | 244 ++-- .../com/nvidia/cuvs/TieredIndexQuery.java | 327 ++--- .../com/nvidia/cuvs/spi/CuVSProvider.java | 2 - .../nvidia/cuvs/spi/UnsupportedProvider.java | 1 - .../nvidia/cuvs/internal/TieredIndexImpl.java | 1083 +++++++++-------- .../internal/TieredSearchResultsImpl.java | 68 +- .../com/nvidia/cuvs/spi/JDKProvider.java | 2 - .../java/com/nvidia/cuvs/TieredIndexIT.java | 443 ++++--- 9 files changed, 1102 insertions(+), 1072 deletions(-) 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 index a08e01b968..3cb0e28281 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndex.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndex.java @@ -13,14 +13,12 @@ * 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; -import com.nvidia.cuvs.spi.CuVSProvider; - /** * {@link TieredIndex} encapsulates a Tiered index, along with methods to * interact with it. 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 index be271c6758..f62892b2c1 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndexParams.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndexParams.java @@ -13,7 +13,6 @@ * See the License for the specific language governing permissions and * limitations under the License. */ - package com.nvidia.cuvs; import java.util.Objects; @@ -25,151 +24,150 @@ */ 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; - } + /** + * 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; /** - * Returns the distance metric used for similarity computation. + * Sets the distance metric for similarity computation. * - * @return The {@link Metric} (L2 or Inner Product) + * @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 Metric getMetric() { - return metric; + public Builder metric(Metric metric) { + this.metric = Objects.requireNonNull(metric); + return this; } /** - * Returns the minimum number of rows required to use the ANN algorithm. + * Sets the minimum number of rows required to use the ANN algorithm. * - * @return The minimum row count threshold for ANN algorithm usage + * @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 int getMinAnnRows() { - return minAnnRows; + public Builder minAnnRows(int minAnnRows) { + if (minAnnRows <= 0) { + throw new IllegalArgumentException("minAnnRows must be positive, got: " + minAnnRows); + } + this.minAnnRows = minAnnRows; + return this; } /** - * Returns whether to create an ANN index when extending the dataset. + * Sets whether to create an ANN index when extending the dataset. * - * @return true if ANN index should be created on extend, false otherwise + * @param val true to create ANN index on extend, false otherwise + * @return This Builder instance for method chaining */ - public boolean isCreateAnnIndexOnExtend() { - return createAnnIndexOnExtend; + public Builder createAnnIndexOnExtend(boolean val) { + this.createAnnIndexOnExtend = val; + return this; } /** - * Returns the CAGRA-specific parameters for the ANN algorithm. + * Sets the CAGRA-specific parameters for the ANN algorithm. * - * @return The {@link CagraIndexParams} configuration, or null if not using - * CAGRA + * @param params The {@link CagraIndexParams} configuration for CAGRA + * algorithm + * @return This Builder instance for method chaining + * @throws NullPointerException if params is null */ - public CagraIndexParams getCagraParams() { - return cagraParams; + public Builder withCagraParams(CagraIndexParams params) { + this.cagraParams = Objects.requireNonNull(params); + return this; } /** - * Creates a new Builder for constructing TieredIndexParams. + * Builds and returns a {@link TieredIndexParams} instance with the + * configured parameters. * - * @return A new Builder instance - */ - public static Builder newBuilder() { - return new Builder(); - } - - /** - * Builder class for constructing {@link TieredIndexParams} instances. + * @return A new TieredIndexParams instance + * @throws IllegalStateException if CAGRA params are required but not + * provided */ - 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); - } + 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 index 42ba133038..03bfae070e 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndexQuery.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndexQuery.java @@ -13,15 +13,12 @@ * See the License for the specific language governing permissions and * limitations under the License. */ - package com.nvidia.cuvs; -import java.util.Arrays; -import java.util.List; - 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 @@ -31,194 +28,210 @@ * @since 25.02 */ public class TieredIndexQuery { - private TieredIndexType indexType; - private CagraSearchParams cagraSearchParameters; - private List mapping; + 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 int topK; + private List mapping; + private int topK = 2; 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. + * Sets the index type for this query. * - * @return the TieredIndexType + * @param indexType the index type + * @return an instance of this Builder */ - public TieredIndexType getIndexType() { - return indexType; + public Builder withIndexType(TieredIndexType indexType) { + this.indexType = indexType; + return this; } /** - * Gets the instance of CagraSearchParams initially set. + * Sets the instance of configured CagraSearchParams to be passed for search. * - * @return an instance CagraSearchParams + * @param cagraSearchParams an instance of the configured CagraSearchParams to + * be used for this query + * @return an instance of this Builder */ - public CagraSearchParams getCagraSearchParameters() { - return cagraSearchParameters; + public Builder withSearchParams(CagraSearchParams cagraSearchParams) { + this.cagraSearchParams = cagraSearchParams; + return this; } /** - * Gets the query vector 2D float array. + * Registers the query vectors to be passed in the search call. * - * @return 2D float array + * @param queryVectors 2D float query vector array + * @return an instance of this Builder */ - public float[][] getQueryVectors() { - return queryVectors; + public Builder withQueryVectors(float[][] queryVectors) { + this.queryVectors = queryVectors; + return this; } /** - * Gets the passed map instance. + * Sets the instance of mapping to be used for ID mapping. * - * @return a map of ID mappings + * @param mapping the ID mapping instance + * @return an instance of this Builder */ - public List getMapping() { - return mapping; + public Builder withMapping(List mapping) { + this.mapping = mapping; + return this; } /** - * Gets the topK value. + * Registers the topK value. * - * @return the topK value + * @param topK the topK value used to retrieve the topK results + * @return an instance of this Builder */ - public int getTopK() { - return topK; + public Builder withTopK(int topK) { + this.topK = topK; + return this; } /** - * Gets the prefilter BitSet. + * Sets a BitSet to use as prefilter while searching. * - * @return a BitSet object representing the prefilter + * @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 BitSet getPrefilter() { - return prefilter; + public Builder withPrefilter(BitSet prefilter, int numDocs) { + this.prefilter = prefilter; + this.numDocs = numDocs; + return this; } /** - * Gets the number of documents in this index, as used for prefilter. + * Builds an instance of TieredIndexQuery. * - * @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. + * @return an instance of TieredIndexQuery + * @throws IllegalStateException if required parameters are missing */ - 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); - } + 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 e7fb75c155..cbaf580a32 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 @@ -22,8 +22,6 @@ import com.nvidia.cuvs.Dataset; import com.nvidia.cuvs.HnswIndex; import com.nvidia.cuvs.TieredIndex; -import com.nvidia.cuvs.CagraMergeParams; - import java.nio.file.Path; /** 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 178f790b3e..e9cf0185e7 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 @@ -21,7 +21,6 @@ import com.nvidia.cuvs.Dataset; import com.nvidia.cuvs.HnswIndex; import com.nvidia.cuvs.TieredIndex; - import java.nio.file.Path; /** 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 index f782eb810b..eae1960d7b 100644 --- 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 @@ -15,42 +15,30 @@ */ package com.nvidia.cuvs.internal; -import static java.lang.foreign.ValueLayout.ADDRESS; 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; import static com.nvidia.cuvs.internal.common.LinkerHelper.C_INT_BYTE_SIZE; -import static com.nvidia.cuvs.internal.common.LinkerHelper.C_POINTER; 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.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.common.Util.prepareTensor; +import static com.nvidia.cuvs.internal.panama.headers_h.cudaMemcpy; +import static com.nvidia.cuvs.internal.panama.headers_h.cudaStream_t; +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.cuvsResources_t; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsStreamGet; +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 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.cuvsResources_t; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsStreamGet; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsStreamSync; -import static com.nvidia.cuvs.internal.panama.headers_h.cudaMemcpy; -import static com.nvidia.cuvs.internal.panama.headers_h.cudaStream_t; - -import java.io.InputStream; -import java.io.OutputStream; -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; import com.nvidia.cuvs.CagraIndexParams; import com.nvidia.cuvs.CagraSearchParams; @@ -63,9 +51,15 @@ 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.cuvsTieredIndex; -import com.nvidia.cuvs.internal.panama.cuvsTieredIndexParams; 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 @@ -78,557 +72,594 @@ * @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; + 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; + } + if (dataset != null) dataset.close(); + } + + /** + * 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(resources, tieredIndexParameters) + : MemorySegment.NULL; + + // Get host data + MemorySegment hostDataSeg = + dataset != null + ? ((DatasetImpl) dataset).seg + : Util.buildMemorySegment(resources.getArena(), vectors); + + Arena arena = resources.getArena(); + long cuvsRes = resources.getMemorySegment().get(cuvsResources_t, 0); + MemorySegment stream = arena.allocate(cudaStream_t); + var returnValue = cuvsStreamGet(cuvsRes, stream); + checkCuVSError(returnValue, "cuvsStreamGet"); + + // TieredIndex REQUIRES device memory - allocate it + MemorySegment datasetD = arena.allocate(C_POINTER); + long datasetSize = C_FLOAT_BYTE_SIZE * rows * cols; + returnValue = cuvsRMMAlloc(cuvsRes, datasetD, datasetSize); + checkCuVSError(returnValue, "cuvsRMMAlloc"); + + MemorySegment datasetDP = datasetD.get(C_POINTER, 0); + + // Copy host to device + returnValue = cudaMemcpy(datasetDP, hostDataSeg, datasetSize, 1); // cudaMemcpyHostToDevice + checkCudaError(returnValue, "cudaMemcpy"); + + // Create tensor from device memory + long datasetShape[] = {rows, cols}; + MemorySegment datasetTensor = prepareTensor(arena, datasetDP, datasetShape, 2, 32, 2, 2, 1); + + MemorySegment index = arena.allocate(cuvsTieredIndex_t); + returnValue = cuvsTieredIndexCreate(index); + checkCuVSError(returnValue, "cuvsTieredIndexCreate"); + + returnValue = cuvsStreamSync(cuvsRes); + checkCuVSError(returnValue, "cuvsStreamSync"); + + returnValue = cuvsTieredIndexBuild(cuvsRes, indexParamsMemorySegment, datasetTensor, index); + checkCuVSError(returnValue, "cuvsTieredIndexBuild"); + + // Clean up device memory after build + returnValue = cuvsRMMFree(cuvsRes, datasetDP, datasetSize); + checkCuVSError(returnValue, "cuvsRMMFree"); + + return new IndexReference(index); + } + } + + /** + * 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; + Arena arena = resources.getArena(); + + // Allocate HOST memory for results + SequenceLayout neighborsLayout = MemoryLayout.sequenceLayout(numBlocks, C_LONG); + SequenceLayout distancesLayout = MemoryLayout.sequenceLayout(numBlocks, C_FLOAT); + MemorySegment neighborsSeg = arena.allocate(neighborsLayout); + MemorySegment distancesSeg = arena.allocate(distancesLayout); + + // Get host query data + MemorySegment hostQueriesSeg = Util.buildMemorySegment(arena, query.getQueryVectors()); + + long cuvsRes = resources.getMemorySegment().get(cuvsResources_t, 0); + MemorySegment stream = arena.allocate(cudaStream_t); + int returnValue = cuvsStreamGet(cuvsRes, stream); + checkCuVSError(returnValue, "cuvsStreamGet"); + + // Allocate DEVICE memory for all data + MemorySegment queriesD = arena.allocate(C_POINTER); + MemorySegment neighborsD = arena.allocate(C_POINTER); + MemorySegment distancesD = arena.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; + + 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(arena, queriesDP, queriesShape, 2, 32, 2, 2, 1); + long neighborsShape[] = {numQueries, topK}; + MemorySegment neighborsTensor = + prepareTensor(arena, neighborsDP, neighborsShape, 0, 64, 2, 2, 1); // 64-bit int + long distancesShape[] = {numQueries, topK}; + MemorySegment distancesTensor = + prepareTensor(arena, distancesDP, distancesShape, 2, 32, 2, 2, 1); + + // Sync before prefilter setup + returnValue = cuvsStreamSync(cuvsRes); + checkCuVSError(returnValue, "cuvsStreamSync"); + + // Handle prefilter + MemorySegment prefilter = cuvsFilter.allocate(arena); + MemorySegment prefilterD = arena.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(arena, 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"); - /** - * 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; + MemorySegment prefilterTensor = + prepareTensor(arena, 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()), + 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 new TieredSearchResultsImpl( + neighborsLayout, + distancesLayout, + neighborsSeg, + distancesSeg, + topK, + query.getMapping(), + numQueries); } - - /** - * 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"); + } + + @Override + public ExtendBuilder extend() { + checkNotDestroyed(); + return new ExtendBuilder(this); + } + + /** + * Performs the actual extend operation + */ + private void performExtend(float[][] extendVectors, Dataset extendDataset) throws Throwable { + 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).seg + : Util.buildMemorySegment(resources.getArena(), extendVectors); + + Arena arena = resources.getArena(); + long cuvsRes = resources.getMemorySegment().get(cuvsResources_t, 0); + MemorySegment stream = arena.allocate(cudaStream_t); + int returnValue = cuvsStreamGet(cuvsRes, stream); + checkCuVSError(returnValue, "cuvsStreamGet"); + + // Allocate device memory for extend data + MemorySegment datasetD = arena.allocate(C_POINTER); + long dataSize = C_FLOAT_BYTE_SIZE * rows * cols; + 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(arena, 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; } - /** - * 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; + @Override + public ExtendBuilder withDataset(float[][] vectors) { + this.vectors = vectors; + return this; } - private void checkNotDestroyed() { - if (destroyed) { - throw new IllegalStateException("destroyed"); - } + @Override + public ExtendBuilder withDataset(Dataset dataset) { + this.dataset = dataset; + return this; } - /** - * 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; - } - if (dataset != null) dataset.close(); + 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( + CuVSResourcesImpl resources, TieredIndexParams params) { + MemorySegment seg = cuvsTieredIndexParams.allocate(resources.getArena()); + + // 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()); + }; } - /** - * 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(resources, tieredIndexParameters) - : MemorySegment.NULL; - - // Get host data - MemorySegment hostDataSeg = dataset != null ? ((DatasetImpl) dataset).seg : - Util.buildMemorySegment(resources.getArena(), vectors); - - Arena arena = resources.getArena(); - long cuvsRes = resources.getMemorySegment().get(cuvsResources_t, 0); - MemorySegment stream = arena.allocate(cudaStream_t); - var returnValue = cuvsStreamGet(cuvsRes, stream); - checkCuVSError(returnValue, "cuvsStreamGet"); - - // TieredIndex REQUIRES device memory - allocate it - MemorySegment datasetD = arena.allocate(C_POINTER); - long datasetSize = C_FLOAT_BYTE_SIZE * rows * cols; - returnValue = cuvsRMMAlloc(cuvsRes, datasetD, datasetSize); - checkCuVSError(returnValue, "cuvsRMMAlloc"); - - MemorySegment datasetDP = datasetD.get(C_POINTER, 0); - - // Copy host to device - returnValue = cudaMemcpy(datasetDP, hostDataSeg, datasetSize, 1); // cudaMemcpyHostToDevice - checkCudaError(returnValue, "cudaMemcpy"); - - // Create tensor from device memory - long datasetShape[] = { rows, cols }; - MemorySegment datasetTensor = prepareTensor(arena, datasetDP, datasetShape, 2, 32, 2, 2, 1); + cuvsTieredIndexParams.metric(seg, metric); - MemorySegment index = arena.allocate(cuvsTieredIndex_t); - returnValue = cuvsTieredIndexCreate(index); - checkCuVSError(returnValue, "cuvsTieredIndexCreate"); + int algo = 0; // CUVS_TIERED_INDEX_ALGO_CAGRA + cuvsTieredIndexParams.algo(seg, algo); - returnValue = cuvsStreamSync(cuvsRes); - checkCuVSError(returnValue, "cuvsStreamSync"); + cuvsTieredIndexParams.min_ann_rows(seg, params.getMinAnnRows()); + cuvsTieredIndexParams.create_ann_index_on_extend(seg, params.isCreateAnnIndexOnExtend()); - returnValue = cuvsTieredIndexBuild(cuvsRes, indexParamsMemorySegment, datasetTensor, index); - checkCuVSError(returnValue, "cuvsTieredIndexBuild"); + CagraIndexParams cagraParams = params.getCagraParams(); + if (cagraParams != null) { + MemorySegment cagraParamsSeg = cuvsCagraIndexParams.allocate(resources.getArena()); - // Clean up device memory after build - returnValue = cuvsRMMFree(cuvsRes, datasetDP, datasetSize); - checkCuVSError(returnValue, "cuvsRMMFree"); + 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); - return new IndexReference(index); - } + cuvsTieredIndexParams.cagra_params(seg, cagraParamsSeg); } - /** - * 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; - Arena arena = resources.getArena(); - - // Allocate HOST memory for results - SequenceLayout neighborsLayout = MemoryLayout.sequenceLayout(numBlocks, C_LONG); - SequenceLayout distancesLayout = MemoryLayout.sequenceLayout(numBlocks, C_FLOAT); - MemorySegment neighborsSeg = arena.allocate(neighborsLayout); - MemorySegment distancesSeg = arena.allocate(distancesLayout); - - // Get host query data - MemorySegment hostQueriesSeg = Util.buildMemorySegment(arena, query.getQueryVectors()); - - long cuvsRes = resources.getMemorySegment().get(cuvsResources_t, 0); - MemorySegment stream = arena.allocate(cudaStream_t); - int returnValue = cuvsStreamGet(cuvsRes, stream); - checkCuVSError(returnValue, "cuvsStreamGet"); - - // Allocate DEVICE memory for all data - MemorySegment queriesD = arena.allocate(C_POINTER); - MemorySegment neighborsD = arena.allocate(C_POINTER); - MemorySegment distancesD = arena.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; - - 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(arena, queriesDP, queriesShape, 2, 32, 2, 2, 1); - long neighborsShape[] = { numQueries, topK }; - MemorySegment neighborsTensor = prepareTensor(arena, neighborsDP, neighborsShape, 0, 64, 2, 2, 1); // 64-bit int - long distancesShape[] = { numQueries, topK }; - MemorySegment distancesTensor = prepareTensor(arena, distancesDP, distancesShape, 2, 32, 2, 2, 1); - - // Sync before prefilter setup - returnValue = cuvsStreamSync(cuvsRes); - checkCuVSError(returnValue, "cuvsStreamSync"); - - // Handle prefilter - MemorySegment prefilter = cuvsFilter.allocate(arena); - MemorySegment prefilterD = arena.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(arena, 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(arena, 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()), - 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 new TieredSearchResultsImpl( - neighborsLayout, distancesLayout, neighborsSeg, distancesSeg, topK, query.getMapping(), numQueries); - } + 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) { + MemorySegment seg = cuvsCagraSearchParams.allocate(resources.getArena()); + 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); } - - @Override - public ExtendBuilder extend() { - checkNotDestroyed(); - return new ExtendBuilder(this); + 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); } - - /** - * Performs the actual extend operation - */ - private void performExtend(float[][] extendVectors, Dataset extendDataset) throws Throwable { - 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).seg - : Util.buildMemorySegment(resources.getArena(), extendVectors); - - Arena arena = resources.getArena(); - long cuvsRes = resources.getMemorySegment().get(cuvsResources_t, 0); - MemorySegment stream = arena.allocate(cudaStream_t); - int returnValue = cuvsStreamGet(cuvsRes, stream); - checkCuVSError(returnValue, "cuvsStreamGet"); - - // Allocate device memory for extend data - MemorySegment datasetD = arena.allocate(C_POINTER); - long dataSize = C_FLOAT_BYTE_SIZE * rows * cols; - 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(arena, 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"); + 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 TieredIndexParams} + * + * @return an instance of {@link TieredIndexParams} + */ + @Override + public TieredIndexParams getIndexParameters() { + return tieredIndexParameters; + } + + /** + * 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 = + tieredIndexParameters != null && tieredIndexParameters.getCagraParams() != null + ? TieredIndexType.CAGRA + : 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); } - - /** - * 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); - } + 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; } - /** - * Allocates the configured index parameters in the MemorySegment. - */ - private static MemorySegment segmentFromIndexParams(CuVSResourcesImpl resources, TieredIndexParams params) { - MemorySegment seg = cuvsTieredIndexParams.allocate(resources.getArena()); - - // 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(resources.getArena()); - - 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); + @Override + public Builder from(InputStream inputStream) { + this.inputStream = inputStream; + return this; + } - return seg; + @Override + public Builder withDataset(float[][] vectors) { + this.vectors = vectors; + return this; } - /** - * Allocates the configured search parameters in the MemorySegment. - */ - private MemorySegment segmentFromSearchParams(CagraSearchParams params) { - MemorySegment seg = cuvsCagraSearchParams.allocate(resources.getArena()); - 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; + @Override + public Builder withDataset(Dataset dataset) { + this.dataset = dataset; + return this; } - /** - * Gets an instance of {@link TieredIndexParams} - * - * @return an instance of {@link TieredIndexParams} - */ @Override - public TieredIndexParams getIndexParameters() { - return tieredIndexParameters; + public Builder withIndexParams(TieredIndexParams params) { + this.params = params; + return this; } - /** - * Gets an instance of {@link CuVSResources} - * - * @return an instance of {@link CuVSResources} - */ @Override - public CuVSResources getCuVSResources() { - return resources; + public Builder withIndexType(TieredIndexType indexType) { + this.indexType = indexType; + return this; } - /** - * Gets the index type - * - * @return the index type - */ @Override - public TieredIndexType getIndexType() { - TieredIndexType indexType = tieredIndexParameters != null && tieredIndexParameters.getCagraParams() != null - ? TieredIndexType.CAGRA - : TieredIndexType.CAGRA; // Default to CAGRA for now - return indexType; + 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; /** - * Static method to create a new builder + * Constructs TieredIndexReference and allocate the MemorySegment. */ - 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); + protected IndexReference(CuVSResourcesImpl resources) { + // Don't allocate here - the C function will allocate + memorySegment = MemorySegment.NULL; } /** - * Builder helps configure and create an instance of {@link TieredIndex}. + * Constructs TieredIndexReference with an instance of MemorySegment passed as a + * parameter. */ - 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); - } - } + protected IndexReference(MemorySegment indexMemorySegment) { + this.memorySegment = indexMemorySegment; } - /** - * Holds the memory reference to a Tiered index. + /** + * Gets the instance of index MemorySegment. */ - 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; - } + 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 index 3e0d5d62f0..7500416730 100644 --- 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 @@ -15,50 +15,58 @@ */ package com.nvidia.cuvs.internal; +import com.nvidia.cuvs.internal.common.SearchResultsImpl; import java.lang.foreign.MemorySegment; import java.lang.foreign.SequenceLayout; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; -import com.nvidia.cuvs.internal.common.SearchResultsImpl; - /** * Implementation of SearchResults for TieredIndex */ class TieredSearchResultsImpl extends SearchResultsImpl { - protected TieredSearchResultsImpl(SequenceLayout neighboursSequenceLayout, SequenceLayout distancesSequenceLayout, - MemorySegment neighboursMemorySegment, MemorySegment distancesMemorySegment, int topK, - List mapping, - long numberOfQueries) { - super(neighboursSequenceLayout, distancesSequenceLayout, neighboursMemorySegment, distancesMemorySegment, topK, - mapping, - numberOfQueries); - readResultMemorySegments(); - } + protected TieredSearchResultsImpl( + SequenceLayout neighboursSequenceLayout, + SequenceLayout distancesSequenceLayout, + MemorySegment neighboursMemorySegment, + MemorySegment distancesMemorySegment, + int topK, + List mapping, + long numberOfQueries) { + super( + neighboursSequenceLayout, + distancesSequenceLayout, + neighboursMemorySegment, + distancesMemorySegment, + topK, + mapping, + numberOfQueries); + readResultMemorySegments(); + } - /** + /** * Reads neighbors and distances {@link MemorySegment} and loads the values * internally into the results structure. */ - @Override - protected void readResultMemorySegments() { - Map intermediateResultMap = new LinkedHashMap(); - int count = 0; - for (long i = 0; i < topK * numberOfQueries; i++) { - long neighborIdLong = (long) neighboursVarHandle.get(neighboursMemorySegment, 0L, i); - float dst = (float) distancesVarHandle.get(distancesMemorySegment, 0L, i); - if (neighborIdLong != -1L && neighborIdLong != Long.MAX_VALUE) { - int id = (int) neighborIdLong; - intermediateResultMap.put(mapping != null ? mapping.get(id) : id, dst); - } - count++; - if (count == topK) { - results.add(intermediateResultMap); - intermediateResultMap = new LinkedHashMap(); - count = 0; - } - } + @Override + protected void readResultMemorySegments() { + Map intermediateResultMap = new LinkedHashMap(); + int count = 0; + for (long i = 0; i < topK * numberOfQueries; i++) { + long neighborIdLong = (long) neighboursVarHandle.get(neighboursMemorySegment, 0L, i); + float dst = (float) distancesVarHandle.get(distancesMemorySegment, 0L, i); + if (neighborIdLong != -1L && neighborIdLong != Long.MAX_VALUE) { + int id = (int) neighborIdLong; + intermediateResultMap.put(mapping != null ? mapping.get(id) : id, dst); + } + count++; + if (count == topK) { + results.add(intermediateResultMap); + intermediateResultMap = new LinkedHashMap(); + count = 0; + } } + } } 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 7aa35abfff..b1c188d8af 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 @@ -22,14 +22,12 @@ import com.nvidia.cuvs.Dataset; import com.nvidia.cuvs.HnswIndex; import com.nvidia.cuvs.TieredIndex; -import com.nvidia.cuvs.CagraMergeParams; 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 java.nio.file.Files; import java.nio.file.Path; import java.util.Objects; 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 index 48442a0c7a..7bbdfe7fa9 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/TieredIndexIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/TieredIndexIT.java @@ -19,249 +19,236 @@ 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 java.util.BitSet; -import java.lang.invoke.MethodHandles; - import org.junit.Before; import org.junit.Test; import org.junit.runner.RunWith; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import com.carrotsearch.randomizedtesting.RandomizedRunner; - @RunWith(RandomizedRunner.class) public class TieredIndexIT extends CuVSTestCase { - private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass()); - - @Before - public void setup() { - initializeRandom(); - log.info("Random context initialized for test."); - } - - /** - * Tests basic operations of TieredIndex - build, search, and extend. - */ - @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.info("Building initial index..."); - TieredIndex index = TieredIndex.newBuilder(resources) - .withDataset(initialDataset) - .withIndexParams(indexParams) - .build(); - - CagraSearchParams searchParams = new CagraSearchParams.Builder(resources) - .withMaxIterations(20) - .build(); - - TieredIndexQuery query = new TieredIndexQuery.Builder() - .withTopK(3) - .withQueryVectors(queries) - .withSearchParams(searchParams) - .build(); - - log.info("Searching initial index..."); - SearchResults initialResults = index.search(query); - log.info("Initial search results: {}", initialResults.getResults()); - assertEquals(expectedInitialResults, roundResults(initialResults.getResults())); - - log.info("Extending index..."); - index.extend() - .withDataset(extensionVectors) - .execute(); - - log.info("Searching extended index..."); - SearchResults extendedResults = index.search(query); - log.info("Extended search results: {}", extendedResults.getResults()); - assertEquals(expectedExtendedResults, roundResults(extendedResults.getResults())); - } - } - - /** - * Tests error handling and parameter validation. - */ - @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(); - - TieredIndex.newBuilder(resources) - .withIndexParams(indexParams) - .withDataset((float[][]) null) - .build(); - } + private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass()); + + @Before + public void setup() { + initializeRandom(); + log.info("Random context initialized for test."); + } + + /** + * Tests basic operations of TieredIndex - build, search, and extend. + */ + @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.info("Building initial index..."); + TieredIndex index = + TieredIndex.newBuilder(resources) + .withDataset(initialDataset) + .withIndexParams(indexParams) + .build(); + + CagraSearchParams searchParams = + new CagraSearchParams.Builder(resources).withMaxIterations(20).build(); + + TieredIndexQuery query = + new TieredIndexQuery.Builder() + .withTopK(3) + .withQueryVectors(queries) + .withSearchParams(searchParams) + .build(); + + log.info("Searching initial index..."); + SearchResults initialResults = index.search(query); + log.info("Initial search results: {}", initialResults.getResults()); + assertEquals(expectedInitialResults, roundResults(initialResults.getResults())); + + log.info("Extending index..."); + index.extend().withDataset(extensionVectors).execute(); + + log.info("Searching extended index..."); + SearchResults extendedResults = index.search(query); + log.info("Extended search results: {}", extendedResults.getResults()); + assertEquals(expectedExtendedResults, roundResults(extendedResults.getResults())); } - - /** - * Tests search with different K values. - */ - @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(); - - TieredIndex index = TieredIndex.newBuilder(resources) - .withDataset(dataset) - .withIndexParams(indexParams) - .build(); - - TieredIndexQuery query1 = new TieredIndexQuery.Builder() - .withTopK(1) - .withQueryVectors(queries) - .withSearchParams(new CagraSearchParams.Builder(resources) - .withMaxIterations(20) - .build()) - .build(); - - SearchResults results1 = index.search(query1); - assertEquals(1, results1.getResults().get(0).size()); - - TieredIndexQuery query3 = new TieredIndexQuery.Builder() - .withTopK(3) - .withQueryVectors(queries) - .withSearchParams(new CagraSearchParams.Builder(resources) - .withMaxIterations(20) - .build()) - .build(); - - SearchResults results3 = index.search(query3); - assertEquals(3, results3.getResults().get(0).size()); - } + } + + /** + * Tests error handling and parameter validation. + */ + @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(); + + TieredIndex.newBuilder(resources) + .withIndexParams(indexParams) + .withDataset((float[][]) null) + .build(); } - - /** - * Test prefilter functionality with debug logging - */ - @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(); - - TieredIndex index = TieredIndex.newBuilder(resources) - .withDataset(dataset) - .withIndexParams(indexParams) - .build(); - - CagraSearchParams searchParams = new CagraSearchParams.Builder(resources).build(); - - TieredIndexQuery queryWithoutFilter = new TieredIndexQuery.Builder() - .withTopK(3) - .withQueryVectors(queryVectors) - .withSearchParams(searchParams) - .build(); - - SearchResults resultsWithoutFilter = index.search(queryWithoutFilter); - log.info("Results WITHOUT prefilter: {}", resultsWithoutFilter.getResults()); - - BitSet prefilter = new BitSet(4); - prefilter.set(1, true); - prefilter.set(2, true); - // Index 0 and 3 are NOT set, so they should be excluded - - TieredIndexQuery queryWithFilter = new TieredIndexQuery.Builder() - .withTopK(3) - .withQueryVectors(queryVectors) - .withSearchParams(searchParams) - .withPrefilter(prefilter, 4) - .build(); - - SearchResults resultsWithFilter = index.search(queryWithFilter); - log.info("Results WITH prefilter: {}", resultsWithFilter.getResults()); - - Map result = resultsWithFilter.getResults().get(0); - - assertFalse("Index 0 should be filtered out", result.containsKey(0)); - assertTrue("Index 1 or 2 should be present", result.containsKey(1) || result.containsKey(2)); - } + } + + /** + * Tests search with different K values. + */ + @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(); + + TieredIndex index = + TieredIndex.newBuilder(resources) + .withDataset(dataset) + .withIndexParams(indexParams) + .build(); + + TieredIndexQuery query1 = + new TieredIndexQuery.Builder() + .withTopK(1) + .withQueryVectors(queries) + .withSearchParams( + new CagraSearchParams.Builder(resources).withMaxIterations(20).build()) + .build(); + + SearchResults results1 = index.search(query1); + assertEquals(1, results1.getResults().get(0).size()); + + TieredIndexQuery query3 = + new TieredIndexQuery.Builder() + .withTopK(3) + .withQueryVectors(queries) + .withSearchParams( + new CagraSearchParams.Builder(resources).withMaxIterations(20).build()) + .build(); + + SearchResults results3 = index.search(query3); + assertEquals(3, results3.getResults().get(0).size()); } - - 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()); + } + + /** + * Test prefilter functionality with debug logging + */ + @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(); + + TieredIndex index = + TieredIndex.newBuilder(resources) + .withDataset(dataset) + .withIndexParams(indexParams) + .build(); + + CagraSearchParams searchParams = new CagraSearchParams.Builder(resources).build(); + + TieredIndexQuery queryWithoutFilter = + new TieredIndexQuery.Builder() + .withTopK(3) + .withQueryVectors(queryVectors) + .withSearchParams(searchParams) + .build(); + + SearchResults resultsWithoutFilter = index.search(queryWithoutFilter); + log.info("Results WITHOUT prefilter: {}", resultsWithoutFilter.getResults()); + + BitSet prefilter = new BitSet(4); + prefilter.set(1, true); + prefilter.set(2, true); + // Index 0 and 3 are NOT set, so they should be excluded + + TieredIndexQuery queryWithFilter = + new TieredIndexQuery.Builder() + .withTopK(3) + .withQueryVectors(queryVectors) + .withSearchParams(searchParams) + .withPrefilter(prefilter, 4) + .build(); + + SearchResults resultsWithFilter = index.search(queryWithFilter); + log.info("Results WITH prefilter: {}", resultsWithFilter.getResults()); + + Map result = resultsWithFilter.getResults().get(0); + + assertFalse("Index 0 should be filtered out", result.containsKey(0)); + 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()); + } } From f453d90811cdd640ade539274b6a6f6ed2f55026 Mon Sep 17 00:00:00 2001 From: punAhuja Date: Thu, 3 Jul 2025 15:09:03 +0530 Subject: [PATCH 04/11] Addressed some review comments --- .../java/com/nvidia/cuvs/TieredIndex.java | 1 - .../nvidia/cuvs/internal/TieredIndexImpl.java | 25 ++++++------------- 2 files changed, 7 insertions(+), 19 deletions(-) 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 index 3cb0e28281..b7f9803430 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndex.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndex.java @@ -22,7 +22,6 @@ /** * {@link TieredIndex} encapsulates a Tiered index, along with methods to * interact with it. - * */ public interface TieredIndex { 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 index eae1960d7b..90ddfa496d 100644 --- 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 @@ -21,17 +21,16 @@ 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.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.cudaStream_t; 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.cuvsResources_t; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsStreamGet; 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; @@ -162,24 +161,20 @@ private IndexReference build() throws Throwable { Arena arena = resources.getArena(); long cuvsRes = resources.getMemorySegment().get(cuvsResources_t, 0); - MemorySegment stream = arena.allocate(cudaStream_t); - var returnValue = cuvsStreamGet(cuvsRes, stream); - checkCuVSError(returnValue, "cuvsStreamGet"); // TieredIndex REQUIRES device memory - allocate it - MemorySegment datasetD = arena.allocate(C_POINTER); + MemorySegment datasetD = localArena.allocate(C_POINTER); long datasetSize = C_FLOAT_BYTE_SIZE * rows * cols; - returnValue = cuvsRMMAlloc(cuvsRes, datasetD, datasetSize); + int returnValue = cuvsRMMAlloc(cuvsRes, datasetD, datasetSize); checkCuVSError(returnValue, "cuvsRMMAlloc"); MemorySegment datasetDP = datasetD.get(C_POINTER, 0); // Copy host to device - returnValue = cudaMemcpy(datasetDP, hostDataSeg, datasetSize, 1); // cudaMemcpyHostToDevice - checkCudaError(returnValue, "cudaMemcpy"); + Util.cudaMemcpy(datasetDP, hostDataSeg, datasetSize, HOST_TO_DEVICE); // Create tensor from device memory - long datasetShape[] = {rows, cols}; + long[] datasetShape = {rows, cols}; MemorySegment datasetTensor = prepareTensor(arena, datasetDP, datasetShape, 2, 32, 2, 2, 1); MemorySegment index = arena.allocate(cuvsTieredIndex_t); @@ -232,9 +227,6 @@ public SearchResults search(TieredIndexQuery query) throws Throwable { MemorySegment hostQueriesSeg = Util.buildMemorySegment(arena, query.getQueryVectors()); long cuvsRes = resources.getMemorySegment().get(cuvsResources_t, 0); - MemorySegment stream = arena.allocate(cudaStream_t); - int returnValue = cuvsStreamGet(cuvsRes, stream); - checkCuVSError(returnValue, "cuvsStreamGet"); // Allocate DEVICE memory for all data MemorySegment queriesD = arena.allocate(C_POINTER); @@ -245,7 +237,7 @@ public SearchResults search(TieredIndexQuery query) throws Throwable { long neighborsBytes = C_LONG_BYTE_SIZE * numQueries * topK; // 64-bit for tiered index long distancesBytes = C_FLOAT_BYTE_SIZE * numQueries * topK; - returnValue = cuvsRMMAlloc(cuvsRes, queriesD, queriesBytes); + int returnValue = cuvsRMMAlloc(cuvsRes, queriesD, queriesBytes); checkCuVSError(returnValue, "cuvsRMMAlloc"); returnValue = cuvsRMMAlloc(cuvsRes, neighborsD, neighborsBytes); checkCuVSError(returnValue, "cuvsRMMAlloc"); @@ -377,14 +369,11 @@ private void performExtend(float[][] extendVectors, Dataset extendDataset) throw Arena arena = resources.getArena(); long cuvsRes = resources.getMemorySegment().get(cuvsResources_t, 0); - MemorySegment stream = arena.allocate(cudaStream_t); - int returnValue = cuvsStreamGet(cuvsRes, stream); - checkCuVSError(returnValue, "cuvsStreamGet"); // Allocate device memory for extend data MemorySegment datasetD = arena.allocate(C_POINTER); long dataSize = C_FLOAT_BYTE_SIZE * rows * cols; - returnValue = cuvsRMMAlloc(cuvsRes, datasetD, dataSize); + int returnValue = cuvsRMMAlloc(cuvsRes, datasetD, dataSize); checkCuVSError(returnValue, "cuvsRMMAlloc"); MemorySegment datasetDP = datasetD.get(C_POINTER, 0); From 993c7755c2714d51eda16d46b4fa85a90107f1c6 Mon Sep 17 00:00:00 2001 From: punAhuja Date: Fri, 4 Jul 2025 16:39:21 +0530 Subject: [PATCH 05/11] Addressed review comment -Removed getIndexParameters method --- .../src/main/java/com/nvidia/cuvs/TieredIndex.java | 8 -------- .../com/nvidia/cuvs/internal/TieredIndexImpl.java | 10 ---------- 2 files changed, 18 deletions(-) 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 index b7f9803430..a6bd7e1f6e 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndex.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndex.java @@ -52,14 +52,6 @@ public interface TieredIndex { */ TieredIndexType getIndexType(); - /** - * Returns the configuration parameters used to build this TieredIndex. - * - * @return An instance of {@link TieredIndexParams} containing the index - * configuration - */ - TieredIndexParams getIndexParameters(); - /** * Returns the resources handle associated with this TieredIndex. * 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 index 90ddfa496d..30cd545b4b 100644 --- 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 @@ -512,16 +512,6 @@ private MemorySegment segmentFromSearchParams(CagraSearchParams params) { return seg; } - /** - * Gets an instance of {@link TieredIndexParams} - * - * @return an instance of {@link TieredIndexParams} - */ - @Override - public TieredIndexParams getIndexParameters() { - return tieredIndexParameters; - } - /** * Gets an instance of {@link CuVSResources} * From 4b887016de79f415131c70be6b3d455975d03cb3 Mon Sep 17 00:00:00 2001 From: punAhuja Date: Mon, 7 Jul 2025 18:36:53 +0530 Subject: [PATCH 06/11] Using local arena instead of global arena for better memory management --- .../nvidia/cuvs/internal/TieredIndexImpl.java | 97 ++++++++++--------- 1 file changed, 51 insertions(+), 46 deletions(-) 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 index 30cd545b4b..9c00f9ddbf 100644 --- 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 @@ -175,7 +175,8 @@ private IndexReference build() throws Throwable { // Create tensor from device memory long[] datasetShape = {rows, cols}; - MemorySegment datasetTensor = prepareTensor(arena, datasetDP, datasetShape, 2, 32, 2, 2, 1); + MemorySegment datasetTensor = + prepareTensor(localArena, datasetDP, datasetShape, 2, 32, 2, 2, 1); MemorySegment index = arena.allocate(cuvsTieredIndex_t); returnValue = cuvsTieredIndexCreate(index); @@ -220,18 +221,18 @@ public SearchResults search(TieredIndexQuery query) throws Throwable { // Allocate HOST memory for results SequenceLayout neighborsLayout = MemoryLayout.sequenceLayout(numBlocks, C_LONG); SequenceLayout distancesLayout = MemoryLayout.sequenceLayout(numBlocks, C_FLOAT); - MemorySegment neighborsSeg = arena.allocate(neighborsLayout); - MemorySegment distancesSeg = arena.allocate(distancesLayout); + MemorySegment neighborsSeg = localArena.allocate(neighborsLayout); + MemorySegment distancesSeg = localArena.allocate(distancesLayout); // Get host query data - MemorySegment hostQueriesSeg = Util.buildMemorySegment(arena, query.getQueryVectors()); + MemorySegment hostQueriesSeg = Util.buildMemorySegment(localArena, query.getQueryVectors()); long cuvsRes = resources.getMemorySegment().get(cuvsResources_t, 0); // Allocate DEVICE memory for all data - MemorySegment queriesD = arena.allocate(C_POINTER); - MemorySegment neighborsD = arena.allocate(C_POINTER); - MemorySegment distancesD = arena.allocate(C_POINTER); + 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 @@ -256,21 +257,22 @@ public SearchResults search(TieredIndexQuery query) throws Throwable { // Create tensors from device memory long queriesShape[] = {numQueries, vectorDimension}; - MemorySegment queriesTensor = prepareTensor(arena, queriesDP, queriesShape, 2, 32, 2, 2, 1); + MemorySegment queriesTensor = + prepareTensor(localArena, queriesDP, queriesShape, 2, 32, 2, 2, 1); long neighborsShape[] = {numQueries, topK}; MemorySegment neighborsTensor = - prepareTensor(arena, neighborsDP, neighborsShape, 0, 64, 2, 2, 1); // 64-bit int + prepareTensor(localArena, neighborsDP, neighborsShape, 0, 64, 2, 2, 1); // 64-bit int long distancesShape[] = {numQueries, topK}; MemorySegment distancesTensor = - prepareTensor(arena, distancesDP, distancesShape, 2, 32, 2, 2, 1); + 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(arena); - MemorySegment prefilterD = arena.allocate(C_POINTER); + MemorySegment prefilter = cuvsFilter.allocate(localArena); + MemorySegment prefilterD = localArena.allocate(C_POINTER); MemorySegment prefilterDP = MemorySegment.NULL; long prefilterBytes = 0; @@ -278,7 +280,7 @@ public SearchResults search(TieredIndexQuery query) throws Throwable { BitSet[] prefilters = new BitSet[] {query.getPrefilter()}; BitSet concatenatedFilters = concatenate(prefilters, (int) query.getNumDocs()); long filters[] = concatenatedFilters.toLongArray(); - MemorySegment hostPrefilterSeg = buildMemorySegment(arena, filters); + MemorySegment hostPrefilterSeg = buildMemorySegment(localArena, filters); long prefilterDataLength = query.getNumDocs() * prefilters.length; long prefilterShape[] = {(prefilterDataLength + 31) / 32}; @@ -296,7 +298,7 @@ public SearchResults search(TieredIndexQuery query) throws Throwable { checkCudaError(returnValue, "cudaMemcpy"); MemorySegment prefilterTensor = - prepareTensor(arena, prefilterDP, prefilterShape, 1, 32, 1, 2, 1); + prepareTensor(localArena, prefilterDP, prefilterShape, 1, 32, 1, 2, 1); cuvsFilter.type(prefilter, 1); // BITSET cuvsFilter.addr(prefilter, prefilterTensor.address()); @@ -309,7 +311,7 @@ public SearchResults search(TieredIndexQuery query) throws Throwable { returnValue = cuvsTieredIndexSearch( cuvsRes, - segmentFromSearchParams(query.getCagraSearchParameters()), + segmentFromSearchParams(query.getCagraSearchParameters(), localArena), tieredIndexReference.getMemorySegment(), queriesTensor, neighborsTensor, @@ -358,44 +360,47 @@ public ExtendBuilder extend() { * Performs the actual extend operation */ private void performExtend(float[][] extendVectors, Dataset extendDataset) throws Throwable { - long rows = extendDataset != null ? extendDataset.size() : extendVectors.length; - long cols = extendDataset != null ? extendDataset.dimensions() : extendVectors[0].length; + 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).seg - : Util.buildMemorySegment(resources.getArena(), extendVectors); + // Get host data + MemorySegment hostDataSeg = + extendDataset != null + ? ((DatasetImpl) extendDataset).seg + : Util.buildMemorySegment(localArena, extendVectors); - Arena arena = resources.getArena(); - long cuvsRes = resources.getMemorySegment().get(cuvsResources_t, 0); + Arena arena = resources.getArena(); + long cuvsRes = resources.getMemorySegment().get(cuvsResources_t, 0); - // Allocate device memory for extend data - MemorySegment datasetD = arena.allocate(C_POINTER); - long dataSize = C_FLOAT_BYTE_SIZE * rows * cols; - int returnValue = cuvsRMMAlloc(cuvsRes, datasetD, dataSize); - checkCuVSError(returnValue, "cuvsRMMAlloc"); + // 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); + MemorySegment datasetDP = datasetD.get(C_POINTER, 0); - // Copy host to device - returnValue = cudaMemcpy(datasetDP, hostDataSeg, dataSize, 1); // cudaMemcpyHostToDevice - checkCudaError(returnValue, "cudaMemcpy"); + // 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(arena, datasetDP, datasetShape, 2, 32, 2, 2, 1); + // 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 = cuvsStreamSync(cuvsRes); + checkCuVSError(returnValue, "cuvsStreamSync"); - returnValue = - cuvsTieredIndexExtend(cuvsRes, datasetTensor, tieredIndexReference.getMemorySegment()); - checkCuVSError(returnValue, "cuvsTieredIndexExtend"); + returnValue = + cuvsTieredIndexExtend(cuvsRes, datasetTensor, tieredIndexReference.getMemorySegment()); + checkCuVSError(returnValue, "cuvsTieredIndexExtend"); - // Clean up device memory - returnValue = cuvsRMMFree(cuvsRes, datasetDP, dataSize); - checkCuVSError(returnValue, "cuvsRMMFree"); + // Clean up device memory + returnValue = cuvsRMMFree(cuvsRes, datasetDP, dataSize); + checkCuVSError(returnValue, "cuvsRMMFree"); + } } /** @@ -491,8 +496,8 @@ private static MemorySegment segmentFromIndexParams( /** * Allocates the configured search parameters in the MemorySegment. */ - private MemorySegment segmentFromSearchParams(CagraSearchParams params) { - MemorySegment seg = cuvsCagraSearchParams.allocate(resources.getArena()); + 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()); From 0f7379274e4a02257c2392ca6ffc008253552e0a Mon Sep 17 00:00:00 2001 From: punAhuja Date: Tue, 15 Jul 2025 11:54:16 +0530 Subject: [PATCH 07/11] Addressed review comment -extracting index pointer and using it --- .../java22/com/nvidia/cuvs/internal/TieredIndexImpl.java | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) 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 index 9c00f9ddbf..7d383a1a57 100644 --- 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 @@ -21,6 +21,7 @@ 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; @@ -185,14 +186,18 @@ private IndexReference build() throws Throwable { returnValue = cuvsStreamSync(cuvsRes); checkCuVSError(returnValue, "cuvsStreamSync"); - returnValue = cuvsTieredIndexBuild(cuvsRes, indexParamsMemorySegment, datasetTensor, index); + // 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(index); + return new IndexReference(actualIndexPtr); } } From d8c5d86b973a323eb1f76dafabe89dc335e7f3b8 Mon Sep 17 00:00:00 2001 From: punAhuja Date: Tue, 15 Jul 2025 13:07:45 +0530 Subject: [PATCH 08/11] Fixed some issues caused by upstream changes --- .../nvidia/cuvs/internal/TieredIndexImpl.java | 18 +++-- .../internal/TieredSearchResultsImpl.java | 68 +++++++++---------- 2 files changed, 39 insertions(+), 47 deletions(-) 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 index 7d383a1a57..ee405813c6 100644 --- 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 @@ -31,7 +31,6 @@ 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.cuvsResources_t; 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; @@ -133,7 +132,7 @@ public void destroyIndex() throws Throwable { } finally { destroyed = true; } - if (dataset != null) dataset.close(); + // if (dataset != null) dataset.close(); } /** @@ -157,11 +156,11 @@ private IndexReference build() throws Throwable { // Get host data MemorySegment hostDataSeg = dataset != null - ? ((DatasetImpl) dataset).seg + ? ((DatasetImpl) dataset).asMemorySegment() : Util.buildMemorySegment(resources.getArena(), vectors); Arena arena = resources.getArena(); - long cuvsRes = resources.getMemorySegment().get(cuvsResources_t, 0); + long cuvsRes = resources.getHandle(); // TieredIndex REQUIRES device memory - allocate it MemorySegment datasetD = localArena.allocate(C_POINTER); @@ -223,16 +222,15 @@ public SearchResults search(TieredIndexQuery query) throws Throwable { int vectorDimension = numQueries > 0 ? query.getQueryVectors()[0].length : 0; Arena arena = resources.getArena(); - // Allocate HOST memory for results SequenceLayout neighborsLayout = MemoryLayout.sequenceLayout(numBlocks, C_LONG); SequenceLayout distancesLayout = MemoryLayout.sequenceLayout(numBlocks, C_FLOAT); - MemorySegment neighborsSeg = localArena.allocate(neighborsLayout); - MemorySegment distancesSeg = localArena.allocate(distancesLayout); + MemorySegment neighborsSeg = arena.allocate(neighborsLayout); + MemorySegment distancesSeg = arena.allocate(distancesLayout); // Get host query data MemorySegment hostQueriesSeg = Util.buildMemorySegment(localArena, query.getQueryVectors()); - long cuvsRes = resources.getMemorySegment().get(cuvsResources_t, 0); + long cuvsRes = resources.getHandle(); // Allocate DEVICE memory for all data MemorySegment queriesD = localArena.allocate(C_POINTER); @@ -372,11 +370,11 @@ private void performExtend(float[][] extendVectors, Dataset extendDataset) throw // Get host data MemorySegment hostDataSeg = extendDataset != null - ? ((DatasetImpl) extendDataset).seg + ? ((DatasetImpl) extendDataset).asMemorySegment() : Util.buildMemorySegment(localArena, extendVectors); Arena arena = resources.getArena(); - long cuvsRes = resources.getMemorySegment().get(cuvsResources_t, 0); + long cuvsRes = resources.getHandle(); // Allocate device memory for extend data MemorySegment datasetD = localArena.allocate(C_POINTER); 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 index 7500416730..b82ef6a71a 100644 --- 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 @@ -15,19 +15,23 @@ */ package com.nvidia.cuvs.internal; -import com.nvidia.cuvs.internal.common.SearchResultsImpl; +import com.nvidia.cuvs.SearchResults; import java.lang.foreign.MemorySegment; import java.lang.foreign.SequenceLayout; -import java.util.LinkedHashMap; import java.util.List; import java.util.Map; +import java.util.function.LongToIntFunction; -/** - * Implementation of SearchResults for TieredIndex - */ -class TieredSearchResultsImpl extends SearchResultsImpl { +class TieredSearchResultsImpl implements SearchResults { + private final SequenceLayout neighboursSequenceLayout; + private final SequenceLayout distancesSequenceLayout; + private final MemorySegment neighboursMemorySegment; + private final MemorySegment distancesMemorySegment; + private final int topK; + private final List mapping; + private final long numberOfQueries; - protected TieredSearchResultsImpl( + TieredSearchResultsImpl( SequenceLayout neighboursSequenceLayout, SequenceLayout distancesSequenceLayout, MemorySegment neighboursMemorySegment, @@ -35,38 +39,28 @@ protected TieredSearchResultsImpl( int topK, List mapping, long numberOfQueries) { - super( - neighboursSequenceLayout, - distancesSequenceLayout, - neighboursMemorySegment, - distancesMemorySegment, - topK, - mapping, - numberOfQueries); - readResultMemorySegments(); + this.neighboursSequenceLayout = neighboursSequenceLayout; + this.distancesSequenceLayout = distancesSequenceLayout; + this.neighboursMemorySegment = neighboursMemorySegment; + this.distancesMemorySegment = distancesMemorySegment; + this.topK = topK; + this.mapping = mapping; + this.numberOfQueries = numberOfQueries; } - /** - * Reads neighbors and distances {@link MemorySegment} and loads the values - * internally into the results structure. - */ @Override - protected void readResultMemorySegments() { - Map intermediateResultMap = new LinkedHashMap(); - int count = 0; - for (long i = 0; i < topK * numberOfQueries; i++) { - long neighborIdLong = (long) neighboursVarHandle.get(neighboursMemorySegment, 0L, i); - float dst = (float) distancesVarHandle.get(distancesMemorySegment, 0L, i); - if (neighborIdLong != -1L && neighborIdLong != Long.MAX_VALUE) { - int id = (int) neighborIdLong; - intermediateResultMap.put(mapping != null ? mapping.get(id) : id, dst); - } - count++; - if (count == topK) { - results.add(intermediateResultMap); - intermediateResultMap = new LinkedHashMap(); - count = 0; - } - } + public List> getResults() { + // Use SearchResultsImpl.create to convert the memory segments to results + LongToIntFunction mappingFunction = mapping != null ? (long id) -> mapping.get((int) id) : null; + + return SearchResultsImpl.create( + neighboursSequenceLayout, + distancesSequenceLayout, + neighboursMemorySegment, + distancesMemorySegment, + topK, + mappingFunction, + numberOfQueries) + .getResults(); } } From cb2eba5d17c7173b3eb93d250c52725708477594 Mon Sep 17 00:00:00 2001 From: punAhuja Date: Tue, 15 Jul 2025 22:48:25 +0530 Subject: [PATCH 09/11] Removed commented line --- .../main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java | 1 - 1 file changed, 1 deletion(-) 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 index ee405813c6..d45d90c9aa 100644 --- 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 @@ -132,7 +132,6 @@ public void destroyIndex() throws Throwable { } finally { destroyed = true; } - // if (dataset != null) dataset.close(); } /** From e1a829810907f960e655119bcb212bdb588002d8 Mon Sep 17 00:00:00 2001 From: punAhuja Date: Wed, 16 Jul 2025 14:30:30 +0530 Subject: [PATCH 10/11] Addressed review comments -Cleaned up the test and added some asserts --- .../nvidia/cuvs/internal/TieredIndexImpl.java | 5 +- .../java/com/nvidia/cuvs/TieredIndexIT.java | 52 +++++++------------ 2 files changed, 19 insertions(+), 38 deletions(-) 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 index d45d90c9aa..7e975ae32a 100644 --- 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 @@ -536,10 +536,7 @@ public CuVSResources getCuVSResources() { */ @Override public TieredIndexType getIndexType() { - TieredIndexType indexType = - tieredIndexParameters != null && tieredIndexParameters.getCagraParams() != null - ? TieredIndexType.CAGRA - : TieredIndexType.CAGRA; // Default to CAGRA for now + TieredIndexType indexType = TieredIndexType.CAGRA; // Default to CAGRA for now return indexType; } 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 index 7bbdfe7fa9..10e4290160 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/TieredIndexIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/TieredIndexIT.java @@ -40,12 +40,8 @@ public class TieredIndexIT extends CuVSTestCase { @Before public void setup() { initializeRandom(); - log.info("Random context initialized for test."); } - /** - * Tests basic operations of TieredIndex - build, search, and extend. - */ @Test public void testBasicOperations() throws Throwable { float[][] initialDataset = { @@ -81,7 +77,6 @@ public void testBasicOperations() throws Throwable { .withCagraParams(cagraParams) .build(); - log.info("Building initial index..."); TieredIndex index = TieredIndex.newBuilder(resources) .withDataset(initialDataset) @@ -98,24 +93,16 @@ public void testBasicOperations() throws Throwable { .withSearchParams(searchParams) .build(); - log.info("Searching initial index..."); SearchResults initialResults = index.search(query); - log.info("Initial search results: {}", initialResults.getResults()); assertEquals(expectedInitialResults, roundResults(initialResults.getResults())); - log.info("Extending index..."); index.extend().withDataset(extensionVectors).execute(); - log.info("Searching extended index..."); SearchResults extendedResults = index.search(query); - log.info("Extended search results: {}", extendedResults.getResults()); assertEquals(expectedExtendedResults, roundResults(extendedResults.getResults())); } } - /** - * Tests error handling and parameter validation. - */ @Test(expected = IllegalArgumentException.class) public void testErrorHandling() throws Throwable { try (CuVSResources resources = CuVSResources.create()) { @@ -132,9 +119,6 @@ public void testErrorHandling() throws Throwable { } } - /** - * Tests search with different K values. - */ @Test public void testDifferentKValues() throws Throwable { float[][] dataset = { @@ -169,7 +153,11 @@ public void testDifferentKValues() throws Throwable { .build(); SearchResults results1 = index.search(query1); - assertEquals(1, results1.getResults().get(0).size()); + Map firstResult = results1.getResults().get(0); + + 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() @@ -180,13 +168,21 @@ public void testDifferentKValues() throws Throwable { .build(); SearchResults results3 = index.search(query3); - assertEquals(3, results3.getResults().get(0).size()); + Map thirdResult = results3.getResults().get(0); + + 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 prefilter functionality with debug logging - */ @Test public void testPrefilter() throws Throwable { float[][] dataset = {{0.0f, 0.0f}, {1.0f, 1.0f}, {2.0f, 2.0f}, {3.0f, 3.0f}}; @@ -207,20 +203,9 @@ public void testPrefilter() throws Throwable { CagraSearchParams searchParams = new CagraSearchParams.Builder(resources).build(); - TieredIndexQuery queryWithoutFilter = - new TieredIndexQuery.Builder() - .withTopK(3) - .withQueryVectors(queryVectors) - .withSearchParams(searchParams) - .build(); - - SearchResults resultsWithoutFilter = index.search(queryWithoutFilter); - log.info("Results WITHOUT prefilter: {}", resultsWithoutFilter.getResults()); - BitSet prefilter = new BitSet(4); prefilter.set(1, true); prefilter.set(2, true); - // Index 0 and 3 are NOT set, so they should be excluded TieredIndexQuery queryWithFilter = new TieredIndexQuery.Builder() @@ -231,11 +216,10 @@ public void testPrefilter() throws Throwable { .build(); SearchResults resultsWithFilter = index.search(queryWithFilter); - log.info("Results WITH prefilter: {}", resultsWithFilter.getResults()); - Map result = resultsWithFilter.getResults().get(0); 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)); } } From 9918bc3a50b0f8d564b54088629ae5c2ecd72f78 Mon Sep 17 00:00:00 2001 From: punAhuja Date: Wed, 16 Jul 2025 15:23:16 +0530 Subject: [PATCH 11/11] getArena was removed upstream so made changes accordingly --- .../nvidia/cuvs/internal/TieredIndexImpl.java | 22 ++++---- .../internal/TieredSearchResultsImpl.java | 50 +++++++++---------- .../java/com/nvidia/cuvs/TieredIndexIT.java | 21 ++++++++ 3 files changed, 53 insertions(+), 40 deletions(-) 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 index 7e975ae32a..035bbe2eef 100644 --- 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 @@ -149,16 +149,15 @@ private IndexReference build() throws Throwable { MemorySegment indexParamsMemorySegment = tieredIndexParameters != null - ? segmentFromIndexParams(resources, tieredIndexParameters) + ? segmentFromIndexParams(localArena, tieredIndexParameters) : MemorySegment.NULL; // Get host data MemorySegment hostDataSeg = dataset != null ? ((DatasetImpl) dataset).asMemorySegment() - : Util.buildMemorySegment(resources.getArena(), vectors); + : Util.buildMemorySegment(localArena, vectors); - Arena arena = resources.getArena(); long cuvsRes = resources.getHandle(); // TieredIndex REQUIRES device memory - allocate it @@ -177,7 +176,7 @@ private IndexReference build() throws Throwable { MemorySegment datasetTensor = prepareTensor(localArena, datasetDP, datasetShape, 2, 32, 2, 2, 1); - MemorySegment index = arena.allocate(cuvsTieredIndex_t); + MemorySegment index = localArena.allocate(cuvsTieredIndex_t); returnValue = cuvsTieredIndexCreate(index); checkCuVSError(returnValue, "cuvsTieredIndexCreate"); @@ -219,12 +218,11 @@ public SearchResults search(TieredIndexQuery query) throws Throwable { long numQueries = query.getQueryVectors().length; long numBlocks = (long) topK * numQueries; int vectorDimension = numQueries > 0 ? query.getQueryVectors()[0].length : 0; - Arena arena = resources.getArena(); SequenceLayout neighborsLayout = MemoryLayout.sequenceLayout(numBlocks, C_LONG); SequenceLayout distancesLayout = MemoryLayout.sequenceLayout(numBlocks, C_FLOAT); - MemorySegment neighborsSeg = arena.allocate(neighborsLayout); - MemorySegment distancesSeg = arena.allocate(distancesLayout); + MemorySegment neighborsSeg = localArena.allocate(neighborsLayout); + MemorySegment distancesSeg = localArena.allocate(distancesLayout); // Get host query data MemorySegment hostQueriesSeg = Util.buildMemorySegment(localArena, query.getQueryVectors()); @@ -341,7 +339,7 @@ public SearchResults search(TieredIndexQuery query) throws Throwable { checkCuVSError(returnValue, "cuvsRMMFree"); } - return new TieredSearchResultsImpl( + return TieredSearchResultsImpl.create( neighborsLayout, distancesLayout, neighborsSeg, @@ -372,7 +370,6 @@ private void performExtend(float[][] extendVectors, Dataset extendDataset) throw ? ((DatasetImpl) extendDataset).asMemorySegment() : Util.buildMemorySegment(localArena, extendVectors); - Arena arena = resources.getArena(); long cuvsRes = resources.getHandle(); // Allocate device memory for extend data @@ -446,9 +443,8 @@ public void execute() throws Throwable { /** * Allocates the configured index parameters in the MemorySegment. */ - private static MemorySegment segmentFromIndexParams( - CuVSResourcesImpl resources, TieredIndexParams params) { - MemorySegment seg = cuvsTieredIndexParams.allocate(resources.getArena()); + 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; @@ -476,7 +472,7 @@ private static MemorySegment segmentFromIndexParams( CagraIndexParams cagraParams = params.getCagraParams(); if (cagraParams != null) { - MemorySegment cagraParamsSeg = cuvsCagraIndexParams.allocate(resources.getArena()); + MemorySegment cagraParamsSeg = cuvsCagraIndexParams.allocate(arena); cuvsCagraIndexParams.intermediate_graph_degree( cagraParamsSeg, cagraParams.getIntermediateGraphDegree()); 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 index b82ef6a71a..13c1ee93fa 100644 --- 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 @@ -23,15 +23,13 @@ import java.util.function.LongToIntFunction; class TieredSearchResultsImpl implements SearchResults { - private final SequenceLayout neighboursSequenceLayout; - private final SequenceLayout distancesSequenceLayout; - private final MemorySegment neighboursMemorySegment; - private final MemorySegment distancesMemorySegment; - private final int topK; - private final List mapping; - private final long numberOfQueries; + private final List> results; - TieredSearchResultsImpl( + private TieredSearchResultsImpl(List> results) { + this.results = results; + } + + public static TieredSearchResultsImpl create( SequenceLayout neighboursSequenceLayout, SequenceLayout distancesSequenceLayout, MemorySegment neighboursMemorySegment, @@ -39,28 +37,26 @@ class TieredSearchResultsImpl implements SearchResults { int topK, List mapping, long numberOfQueries) { - this.neighboursSequenceLayout = neighboursSequenceLayout; - this.distancesSequenceLayout = distancesSequenceLayout; - this.neighboursMemorySegment = neighboursMemorySegment; - this.distancesMemorySegment = distancesMemorySegment; - this.topK = topK; - this.mapping = mapping; - this.numberOfQueries = 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() { - // Use SearchResultsImpl.create to convert the memory segments to results - LongToIntFunction mappingFunction = mapping != null ? (long id) -> mapping.get((int) id) : null; - - return SearchResultsImpl.create( - neighboursSequenceLayout, - distancesSequenceLayout, - neighboursMemorySegment, - distancesMemorySegment, - topK, - mappingFunction, - numberOfQueries) - .getResults(); + return results; } } 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 index 10e4290160..ffe1dddd08 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/TieredIndexIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/TieredIndexIT.java @@ -40,6 +40,7 @@ public class TieredIndexIT extends CuVSTestCase { @Before public void setup() { initializeRandom(); + log.debug("Random context initialized for test"); } @Test @@ -77,11 +78,13 @@ public void testBasicOperations() throws Throwable { .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(); @@ -93,12 +96,18 @@ public void testBasicOperations() throws Throwable { .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())); } } @@ -112,6 +121,7 @@ public void testErrorHandling() throws Throwable { 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) @@ -138,11 +148,13 @@ public void testDifferentKValues() throws Throwable { 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() @@ -152,8 +164,10 @@ public void testDifferentKValues() throws Throwable { 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)); @@ -167,8 +181,10 @@ public void testDifferentKValues() throws Throwable { 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)); @@ -195,17 +211,20 @@ public void testPrefilter() throws Throwable { 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() @@ -215,8 +234,10 @@ public void testPrefilter() throws Throwable { .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));