New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add on disk 4x compression with Faiss #2425

Open

naveentatikonda wants to merge 9 commits into opensearch-project:2.x from naveentatikonda:faiss_ondisk_4x

+1,408 −190

Member

naveentatikonda commented Jan 23, 2025

Description

Add on disk 4x compression with Faiss which accepts fp32 vectors as input and dynamically quantizes them into byte sized vectors using the Faiss SQ8 quantizer.

Related Issues

Resolves #1723

Check List

New functionality includes testing.
New functionality has been documented.
API changes companion pull request created.
Commits are signed per the DCO using --signoff.
Public documentation issue/PR created.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

naveentatikonda added Features backport 2.x labels

naveentatikonda force-pushed the faiss_ondisk_4x branch from 101c94f to 58a66e4 Compare

January 23, 2025 05:18

naveentatikonda added the v2.19.0 label

naveentatikonda force-pushed the faiss_ondisk_4x branch 2 times, most recently from 6721604 to 7739606 Compare

January 23, 2025 23:03

naveentatikonda changed the base branch from main to 2.x

January 23, 2025 23:03

naveentatikonda added backport main and removed backport 2.x labels

naveentatikonda force-pushed the faiss_ondisk_4x branch from 7739606 to e09cdfd Compare

January 23, 2025 23:05

shatejas reviewed

View reviewed changes

Collaborator

shatejas left a comment

Reviewed partially

jni/src/faiss_index_service.cpp Outdated Show resolved Hide resolved

jni/src/faiss_index_service.cpp

@@ @@ -155,6 +155,47 @@ void IndexService::writeIndex( @@
                   }
               }
+              jlong IndexService::initIndexFromTemplate(

Collaborator

shatejas Jan 24, 2025

The only difference between this and initIndex is the index creation call to faiss, can we abstract out the logic and reuse the rest please. you can pass in the pointer returned by faiss to reuse the logic and set the index uniq pointer maybe.

Member Author

naveentatikonda Jan 25, 2025

refactored as discussed and also validated that the index is getting deleted

jni/src/faiss_index_service.cpp Show resolved Hide resolved

src/main/java/org/opensearch/knn/index/KNNIndexShard.java Outdated Show resolved Hide resolved

src/main/java/org/opensearch/knn/index/codec/KNN990Codec/NativeEngines990KnnVectorsWriter.java Outdated Show resolved Hide resolved

src/main/java/org/opensearch/knn/index/codec/KNN990Codec/NativeEngines990KnnVectorsWriter.java Outdated Show resolved Hide resolved

...in/java/org/opensearch/knn/index/codec/nativeindex/MemOptimizedNativeIndexBuildStrategy.java Outdated Show resolved Hide resolved

naveentatikonda marked this pull request as ready for review

January 24, 2025 17:28

naveentatikonda requested review from heemin32, navneet1v, VijayanB, vamshin, jmazanec15, junqiu-lei, martin-gaievski, ryanbogan, luyuncheng and 0ctopus13prime as code owners

January 24, 2025 17:28

naveentatikonda added 8 commits

January 24, 2025 12:30


          Add support for Faiss OnDisk 4x compression

bc0c79f

Signed-off-by: Naveen Tatikonda <[email protected]>


          Predefined configuration changes

e9b1bba

Signed-off-by: Naveen Tatikonda <[email protected]>


          Ingestion and Querying

f8505f2

Signed-off-by: Naveen Tatikonda <[email protected]>


          Optimize create index from template for 4x compression

4bba135

Signed-off-by: Naveen Tatikonda <[email protected]>


          Add Backwards compatibility for Lucene 4x

4cff131

Signed-off-by: Naveen Tatikonda <[email protected]>


          Add ivf validation

1997d18

Signed-off-by: Naveen Tatikonda <[email protected]>


          Fix vector datatype

a879b44

Signed-off-by: Naveen Tatikonda <[email protected]>


          Fix failing tests

0aa10b4

Signed-off-by: Naveen Tatikonda <[email protected]>

naveentatikonda force-pushed the faiss_ondisk_4x branch 2 times, most recently from 18c4a8a to 4449fde Compare

January 25, 2025 17:05

naveentatikonda requested a review from shatejas

January 25, 2025 17:08

vibrantvarun reviewed

View reviewed changes

src/main/java/org/opensearch/knn/index/codec/KNN990Codec/NativeEngines990KnnVectorsWriter.java Outdated Show resolved Hide resolved

src/main/java/org/opensearch/knn/quantization/quantizer/ByteScalarQuantizer.java Show resolved Hide resolved

src/main/java/org/opensearch/knn/quantization/quantizer/ByteScalarQuantizer.java Show resolved Hide resolved


          Address review comments and add tests

2bb9e05

Signed-off-by: Naveen Tatikonda <[email protected]>

naveentatikonda force-pushed the faiss_ondisk_4x branch from 0ea7bc6 to 2bb9e05 Compare

January 26, 2025 05:21

naveentatikonda removed the v2.19.0 label

Member Author

naveentatikonda commented Jan 27, 2025 •

edited

Loading

Update - After looking at the benchmarks and having a chat with other maintainers of k-NN we want to put this feature on hold from releasing in 2.19. Will run more benchmarking tests and compare them with lucene (with rescoring POC) and then decide if we need to switch the default engine for on_disk 4x compression.

navneet1v reviewed

View reviewed changes

src/main/java/org/opensearch/knn/index/codec/KNN990Codec/NativeEngines990KnnVectorsWriter.java Outdated

Comment on lines 263 to 265

+                          if ((quantizationParams.getTypeIdentifier()).equals(
+                              ScalarQuantizationParams.generateTypeIdentifier(ScalarQuantizationType.EIGHT_BIT)
+                          )) {

Collaborator

navneet1v Jan 24, 2025

to ensure that NPE doesn't come in case quantizationParams.getTypeIdentifier() == null, we should reverse the check.

ScalarQuantizationParams.generateTypeIdentifier(ScalarQuantizationType.EIGHT_BIT).equal(quantizationParams.getTypeIdentifier())

src/main/java/org/opensearch/knn/index/codec/KNN990Codec/NativeEngines990KnnVectorsWriter.java Outdated Show resolved Hide resolved

src/main/java/org/opensearch/knn/index/codec/KNN990Codec/NativeEngines990KnnVectorsWriter.java Outdated

Comment on lines 266 to 271

+                              quantizationState = quantizationService.train(quantizationParams, knnVectorValues, totalLiveDocs, fieldInfo);
+                          } else {
+                              initQuantizationStateWriterIfNecessary();
+                              quantizationState = quantizationService.train(quantizationParams, knnVectorValues, totalLiveDocs, fieldInfo);
+                              quantizationStateWriter.writeState(fieldInfo.getFieldNumber(), quantizationState);
+                          }

Collaborator

navneet1v Jan 24, 2025

I think this whole logic can be simplified where we create the QS first and then we just see if writer needs to init and it needs to write the state,.

...in/java/org/opensearch/knn/index/codec/nativeindex/MemOptimizedNativeIndexBuildStrategy.java

+                      return (indexInfo.getQuantizationState() instanceof ByteScalarQuantizationState);
+                  }
+                  private byte[] getIndexTemplate(BuildIndexParams indexInfo) {

Collaborator

navneet1v Jan 24, 2025

nit pick: make the param final for all the functions

src/main/java/org/opensearch/knn/index/codec/nativeindex/QuantizationIndexUtils.java

                       int bytesPerVector;
                       int dimensions;
-                      if (quantizationState != null) {
+                      if (quantizationState != null && !(quantizationState instanceof ByteScalarQuantizationState)) {

Collaborator

navneet1v Jan 24, 2025

need java doc on this. This kind of instanceOf check is making me nervous. can we think of something here.

src/main/java/org/opensearch/knn/quantization/quantizer/MultiBitScalarQuantizer.java

Comment on lines +124 to +126

+                  public QuantizationState train(final TrainingRequest<float[]> trainingRequest, final FieldInfo fieldInfo) throws IOException {
+                      return null;
+                  }

Collaborator

navneet1v Jan 24, 2025

same as above

src/main/java/org/opensearch/knn/quantization/quantizer/OneBitScalarQuantizer.java

Comment on lines +69 to +71

+                  public QuantizationState train(final TrainingRequest<float[]> trainingRequest, final FieldInfo fieldInfo) throws IOException {
+                      return null;
+                  }

Collaborator

navneet1v Jan 24, 2025

same as above

src/main/java/org/opensearch/knn/quantization/quantizer/Quantizer.java

@@ @@ -31,6 +32,8 @@ public interface Quantizer<T, R> { @@
                    */
                   QuantizationState train(TrainingRequest<T> trainingRequest) throws IOException;
+                  QuantizationState train(TrainingRequest<T> trainingRequest, FieldInfo fieldInfo) throws IOException;

Collaborator

navneet1v Jan 24, 2025

please add java doc and also why we need this function? I thought we are keeping QF free from fieldInfo and other things.

src/main/java/org/opensearch/knn/quantization/quantizer/ByteScalarQuantizer.java

+              import static org.opensearch.knn.common.FieldInfoExtractor.extractVectorDataType;
+              import static org.opensearch.knn.index.codec.transfer.OffHeapVectorTransferFactory.getVectorTransfer;
+              public class ByteScalarQuantizer implements Quantizer<float[], byte[]> {

Collaborator

navneet1v Jan 24, 2025

Please add java doc on all your new classes.

src/main/java/org/opensearch/knn/quantization/quantizer/ByteScalarQuantizer.java

Comment on lines +60 to +71

+                      if (sampledIndices.length == 0) {
+                          return null;
+                      }

Collaborator

navneet1v Jan 24, 2025

should we have some logs here.

navneet1v mentioned this pull request

Add a Faiss codec for KNN searches apache/lucene#14178

Open

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

vibrantvarun vibrantvarun left review comments

navneet1v navneet1v left review comments

heemin32 Awaiting requested review from heemin32 heemin32 is a code owner

VijayanB Awaiting requested review from VijayanB VijayanB is a code owner

vamshin Awaiting requested review from vamshin vamshin is a code owner

jmazanec15 Awaiting requested review from jmazanec15 jmazanec15 is a code owner

junqiu-lei Awaiting requested review from junqiu-lei junqiu-lei is a code owner

martin-gaievski Awaiting requested review from martin-gaievski martin-gaievski is a code owner

ryanbogan Awaiting requested review from ryanbogan ryanbogan is a code owner

luyuncheng Awaiting requested review from luyuncheng luyuncheng is a code owner

0ctopus13prime Awaiting requested review from 0ctopus13prime 0ctopus13prime is a code owner

shatejas Awaiting requested review from shatejas shatejas is a code owner

At least 2 approving reviews are required to merge this pull request.

Labels

backport main Features