forked from opensearch-project/k-NN
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Integration With Qunatization Config
Signed-off-by: VIKASH TIWARI <viktari@amazon.com>
- Loading branch information
Showing
7 changed files
with
149 additions
and
70 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
78 changes: 78 additions & 0 deletions
78
src/main/java/org/opensearch/knn/index/codec/nativeindex/IndexBuildHelper.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,78 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.knn.index.codec.nativeindex; | ||
|
|
||
| import lombok.experimental.UtilityClass; | ||
| import org.opensearch.knn.index.codec.nativeindex.model.BuildIndexParams; | ||
| import org.opensearch.knn.index.codec.transfer.OffHeapVectorTransfer; | ||
| import org.opensearch.knn.index.quantizationService.QuantizationService; | ||
| import org.opensearch.knn.index.vectorvalues.KNNVectorValues; | ||
| import org.opensearch.knn.quantization.models.quantizationOutput.QuantizationOutput; | ||
| import org.opensearch.knn.quantization.models.quantizationState.QuantizationState; | ||
|
|
||
| import java.io.IOException; | ||
|
|
||
| @UtilityClass | ||
| class IndexBuildHelper { | ||
|
|
||
| // private static final QuantizationService quantizationService = QuantizationService.getInstance(); | ||
|
|
||
| /** | ||
| * Processes and transfers the vector based on whether quantization is applied or not. | ||
| * | ||
| * @param knnVectorValues the KNN vector values to be processed. | ||
| * @param indexBuildSetup the setup containing quantization state and output, along with other parameters. | ||
| * @param vectorTransfer the off-heap vector transfer utility. | ||
| * @param append flag indicating whether to append or overwrite the transfer buffer. | ||
| * @return boolean indicating whether the transfer was successful. | ||
| * @throws IOException if an I/O error occurs during vector transfer. | ||
| */ | ||
| static boolean processAndTransferVector( | ||
| KNNVectorValues<?> knnVectorValues, | ||
| IndexBuildSetup indexBuildSetup, | ||
| OffHeapVectorTransfer vectorTransfer, | ||
| boolean append | ||
| ) throws IOException { | ||
| QuantizationService quantizationService = QuantizationService.getInstance(); | ||
| if (indexBuildSetup.getQuantizationState() != null && indexBuildSetup.getQuantizationOutput() != null) { | ||
| quantizationService.quantize( | ||
| indexBuildSetup.getQuantizationState(), | ||
| knnVectorValues.getVector(), | ||
| indexBuildSetup.getQuantizationOutput() | ||
| ); | ||
| return vectorTransfer.transfer(indexBuildSetup.getQuantizationOutput().getQuantizedVector(), append); | ||
| } else { | ||
| return vectorTransfer.transfer(knnVectorValues.conditionalCloneVector(), append); | ||
| } | ||
| } | ||
|
|
||
| /** | ||
| * Prepares the quantization setup including bytes per vector and dimensions. | ||
| * | ||
| * @param knnVectorValues the KNN vector values. | ||
| * @param indexInfo the index build parameters. | ||
| * @return an instance of QuantizationSetup containing relevant information. | ||
| */ | ||
| static IndexBuildSetup prepareIndexBuild(KNNVectorValues<?> knnVectorValues, BuildIndexParams indexInfo) { | ||
| QuantizationState quantizationState = indexInfo.getQuantizationState(); | ||
| QuantizationOutput quantizationOutput = null; | ||
| QuantizationService quantizationService = QuantizationService.getInstance(); | ||
|
|
||
| int bytesPerVector; | ||
| int dimensions; | ||
|
|
||
| if (quantizationState != null) { | ||
| bytesPerVector = quantizationState.getBytesPerVector(); | ||
| dimensions = quantizationState.getDimensions(); | ||
| quantizationOutput = quantizationService.createQuantizationOutput(quantizationState.getQuantizationParams()); | ||
| } else { | ||
| bytesPerVector = knnVectorValues.bytesPerVector(); | ||
| dimensions = knnVectorValues.dimension(); | ||
| } | ||
|
|
||
| return new IndexBuildSetup(bytesPerVector, dimensions, quantizationOutput, quantizationState); | ||
| } | ||
| } |
40 changes: 40 additions & 0 deletions
40
src/main/java/org/opensearch/knn/index/codec/nativeindex/IndexBuildSetup.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,40 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.knn.index.codec.nativeindex; | ||
|
|
||
| import lombok.AllArgsConstructor; | ||
| import lombok.Getter; | ||
| import org.opensearch.knn.quantization.models.quantizationOutput.QuantizationOutput; | ||
| import org.opensearch.knn.quantization.models.quantizationState.QuantizationState; | ||
|
|
||
| /** | ||
| * IndexBuildSetup encapsulates the configuration and parameters required for building an index. | ||
| * This includes the size of each vector, the dimensions of the vectors, and any quantization-related | ||
| * settings such as the output and state of quantization. | ||
| */ | ||
| @Getter | ||
| @AllArgsConstructor | ||
| class IndexBuildSetup { | ||
| /** | ||
| * The number of bytes per vector. | ||
| */ | ||
| private final int bytesPerVector; | ||
|
|
||
| /** | ||
| * The number of dimensions in the vector. | ||
| */ | ||
| private final int dimensions; | ||
|
|
||
| /** | ||
| * The quantization output that will hold the quantized vector. | ||
| */ | ||
| private final QuantizationOutput quantizationOutput; | ||
|
|
||
| /** | ||
| * The state of quantization, which may include parameters and trained models. | ||
| */ | ||
| private final QuantizationState quantizationState; | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters