Add binary format support for Faiss IVF #2

junqiu-lei · 2024-07-09T00:52:03Z

Because opensearch-project#1784 is merged, I rebased junqiu-lei:binary-ivf against opensearch-project:feature/binary-format

Description

Raising this PR in Heemin's forked k-NN repository to facilitate the review of the code changes from opensearch-project#1784

It contains code change in Java and JNI layer to support:

train model with binary format index
add data_type field when train model
create IVF index with binary format
query IVF index with binary format

Will add tests in another PR.

Example workflow

1. Create binary format train index

PUT /train-index HTTP/1.1
Host: localhost:9200
Content-Type: application/json
Content-Length: 481

{
  "settings": {
    "number_of_shards": 3,
    "number_of_replicas": 0,
    "index.knn": true
  },
  "mappings": {
    "properties": {
      "train-field": {
        "type": "knn_vector",
        "dimension": 8,
        "data_type": "binary",
        "method": {
          "name": "hnsw",
          "space_type": "hammingbit",
          "engine": "faiss",
          "parameters": {
            "ef_construction": 128,
            "m": 24
          }
        }
      }
    }
  }
}

2. Ingest tran index

POST /_bulk HTTP/1.1
Host: localhost:9200
Content-Type: application/json
Content-Length: 7067

{ "index": { "_index": "train-index", "_id": "1" } }
{ "train-field": [1]}
{ "index": { "_index": "train-index", "_id": "2" } }
{ "train-field": [2]}
{ "index": { "_index": "train-index", "_id": "3" } }
{ "train-field": [4]}
{ "index": { "_index": "train-index", "_id": "4" } }
........

3. Create train model

POST /_plugins/_knn/models/my-model/_train HTTP/1.1
Host: localhost:9200
Content-Type: application/json
Content-Length: 313

{
  "training_index": "train-index",
  "training_field": "train-field",
  "dimension": 8,
  "description": "My model description",
  "data_type": "binary",
  "method": {
    "name": "ivf",
    "engine": "faiss",
    "space_type": "hammingbit",
    "parameters": {
      "nlist": 4,
      "nprobes":2 
    }
  }
}

4. Create IVF binary format target index

PUT /target-index HTTP/1.1
Host: localhost:9200
Content-Type: application/json
Content-Length: 242

{
  "settings": {
    "number_of_shards": 1,
    "number_of_replicas": 1,
    "index.knn": true
  },
  "mappings": {
    "properties": {
      "target-field": {
        "type": "knn_vector",
        "model_id": "my-model"
      }
    }
  }
}

5. Bulk target index

POST /_bulk HTTP/1.1
Host: localhost:9200
Content-Type: application/json
Content-Length: 931

{ "index": { "_index": "target-index", "_id": "1" } }
{ "target-field": [2]}
{ "index": { "_index": "target-index", "_id": "2" } }
{ "target-field": [3]}
.......

6. Query target index

GET /target-index/_search HTTP/1.1
Host: localhost:9200
Content-Type: application/json
Content-Length: 110

{
  "query": {
    "knn": {
      "target-field": {
        "vector": [10],
        "k": 5
      }
    }
  }
}

7. Query result

{
    "took": 2898,
    "timed_out": false,
    "_shards": {
        "total": 1,
        "successful": 1,
        "skipped": 0,
        "failed": 0
    },
    "hits": {
        "total": {
            "value": 5,
            "relation": "eq"
        },
        "max_score": 1.0,
        "hits": [
            {
                "_index": "target-index",
                "_id": "9",
                "_score": 1.0,
                "_source": {
                    "target-field": [
                        10
                    ]
                }
            },
            {
                "_index": "target-index",
                "_id": "1",
                "_score": 0.5,
                "_source": {
                    "target-field": [
                        2
                    ]
                }
            },
            {
                "_index": "target-index",
                "_id": "7",
                "_score": 0.5,
                "_source": {
                    "target-field": [
                        8
                    ]
                }
            },
            {
                "_index": "target-index",
                "_id": "10",
                "_score": 0.5,
                "_source": {
                    "target-field": [
                        11
                    ]
                }
            },
            {
                "_index": "target-index",
                "_id": "2",
                "_score": 0.33333334,
                "_source": {
                    "target-field": [
                        3
                    ]
                }
            }
        ]
    }
}

jmazanec15 · 2024-07-10T19:19:25Z

jni/include/faiss_wrapper.h

@@ -29,6 +29,12 @@ namespace knn_jni {
                                     jlong vectorsAddressJ, jint dimJ, jstring indexPathJ, jbyteArray templateIndexJ,
                                     jobject parametersJ);

+        // Create an index with ids and vectors. Instead of creating a new index, this function creates the index
+        // based off of the template index passed in. The index is serialized to indexPathJ.
+        void CreateBinaryIndexFromTemplate(knn_jni::JNIUtilInterface * jniUtil, JNIEnv * env, jintArray idsJ,


Why introduce Binary<...> methods, but for Free take parameter?

It is to reflect underlying cpp method.

Faiss has two different method between non binary and binary for create index method. However, there is only one method, delete` to free the memory.

jmazanec15 · 2024-07-10T19:20:05Z

jni/src/faiss_wrapper.cpp

+        throw std::runtime_error("Template index cannot be null");
+    }
+
+    // Set thread count if it is passed in as a parameter. Setting this variable will only impact the current thread


Is it possible to avoid code duplication?

I tried to add it into faiss_index_service, but failed with some c++ null pointer error, will try again with workable solution

I'll raise another PR to try to add it into faiss_index_service

src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java

jmazanec15 · 2024-07-10T19:23:39Z

src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapper.java

@@ -596,8 +597,13 @@ protected List<Field> getFieldsForByteVector(final byte[] array, final FieldType
        return fields;
    }

-    protected void parseCreateField(ParseContext context, int dimension, SpaceType spaceType, MethodComponentContext methodComponentContext)
-        throws IOException {
+    protected void parseCreateField(


Why this need to be changed?

when we create model based index, we don't specify the data_type, which makes vectorDataType to be default float in KNNVectorFieldMapper, so I update this function to be able to pass in the correct vector datatype for model based index from modelMetaData

src/main/java/org/opensearch/knn/index/query/KNNWeight.java

src/main/java/org/opensearch/knn/index/util/ModelInfoExtractor.java

src/main/java/org/opensearch/knn/indices/ModelMetadata.java

navneet1v · 2024-07-11T04:37:19Z

@junqiu-lei in this PR I can see heemin's change and your change both. Is that intensional?

junqiu-lei · 2024-07-11T16:15:47Z

@junqiu-lei in this PR I can see heemin's change and your change both. Is that intensional?

@navneet1v It might because Heemin updated his branch which I targeted on this PR, let me rebase and update it.

junqiu-lei · 2024-07-11T16:22:13Z

@junqiu-lei in this PR I can see heemin's change and your change both. Is that intensional?

@navneet1v It might because Heemin updated his branch which I targeted on this PR, let me rebase and update it.

Rebased and only my commit on this PR.

src/main/java/org/opensearch/knn/index/util/Faiss.java

src/test/java/org/opensearch/knn/index/codec/KNNCodecTestCase.java

naveentatikonda · 2024-07-11T20:20:16Z

src/main/java/org/opensearch/knn/indices/ModelMetadata.java

@@ -128,6 +137,7 @@ public ModelMetadata(
        this.error = Objects.requireNonNull(error, "error must not be null");
        this.trainingNodeAssignment = Objects.requireNonNull(trainingNodeAssignment, "node assignment must not be null");
        this.methodComponentContext = Objects.requireNonNull(methodComponentContext, "method context must not be null");
+        this.vectorDataType = Objects.requireNonNull(vectorDataType, "vector data type must not be null");


Can you update Java doc of the constructor?

src/main/java/org/opensearch/knn/indices/ModelMetadata.java

src/test/java/org/opensearch/knn/KNNSingleNodeTestCase.java

src/test/java/org/opensearch/knn/plugin/action/RestKNNStatsHandlerIT.java

naveentatikonda · 2024-07-11T21:23:51Z

@junqiu-lei @heemin32 Are you adding validation check to not use an encoder with binary datatype because I believe as of now all the encoders are not supported with Binary datatype ?

heemin32 · 2024-07-11T21:26:38Z

@junqiu-lei @heemin32 Are you adding validation check to not use an encoder with binary datatype because I believe as of now all the encoders are not supported with Binary datatype ?

We need. https://github.com/opensearch-project/k-NN/pull/1781/files#diff-c6d56dded611534c5be8f8afdba6bf320f7892d47465deee44bd6390f2bd1f12R347

src/main/java/org/opensearch/knn/training/TrainingJob.java

src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java

src/main/java/org/opensearch/knn/index/codec/transfer/VectorTransfer.java

jni/src/faiss_wrapper.cpp

src/main/java/org/opensearch/knn/index/IndexUtil.java

src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java

src/main/java/org/opensearch/knn/indices/ModelMetadata.java

src/main/java/org/opensearch/knn/training/TrainingJob.java

Signed-off-by: Junqiu Lei <junqiu@amazon.com>

junqiu-lei · 2024-07-16T23:35:48Z

Closing this PR, the comments on this PR are resolved and continue being reviewing on opensearch-project#1784

junqiu-lei requested a review from heemin32 as a code owner July 9, 2024 00:52

heemin32 force-pushed the binary-java branch from e62fb76 to b2a1332 Compare July 9, 2024 20:27

junqiu-lei force-pushed the binary-ivf branch 2 times, most recently from 404056a to 735ccda Compare July 9, 2024 23:51

junqiu-lei mentioned this pull request Jul 10, 2024

Add binary format support with IVF method in Faiss Engine opensearch-project/k-NN#1784

Merged

5 tasks

junqiu-lei force-pushed the binary-ivf branch from 735ccda to 49e6725 Compare July 10, 2024 18:28

jmazanec15 reviewed Jul 10, 2024

View reviewed changes

junqiu-lei force-pushed the binary-ivf branch from 49e6725 to 7493dd9 Compare July 10, 2024 21:15

heemin32 force-pushed the binary-java branch 4 times, most recently from 612a612 to 29815af Compare July 10, 2024 22:50

junqiu-lei force-pushed the binary-ivf branch from 7493dd9 to fb53162 Compare July 11, 2024 00:03

heemin32 force-pushed the binary-java branch from 29815af to 930a27e Compare July 11, 2024 04:21

junqiu-lei force-pushed the binary-ivf branch from fb53162 to 15044a2 Compare July 11, 2024 16:18

ryanbogan reviewed Jul 11, 2024

View reviewed changes

src/main/java/org/opensearch/knn/index/util/Faiss.java Show resolved Hide resolved

ryanbogan reviewed Jul 11, 2024

View reviewed changes

src/test/java/org/opensearch/knn/index/codec/KNNCodecTestCase.java Outdated Show resolved Hide resolved

junqiu-lei force-pushed the binary-ivf branch 3 times, most recently from b3fc615 to b8b127d Compare July 11, 2024 19:39

naveentatikonda reviewed Jul 11, 2024

View reviewed changes

src/main/java/org/opensearch/knn/indices/ModelMetadata.java Show resolved Hide resolved

naveentatikonda reviewed Jul 11, 2024

View reviewed changes

src/test/java/org/opensearch/knn/KNNSingleNodeTestCase.java Show resolved Hide resolved

naveentatikonda reviewed Jul 11, 2024

View reviewed changes

src/test/java/org/opensearch/knn/plugin/action/RestKNNStatsHandlerIT.java Show resolved Hide resolved

junqiu-lei force-pushed the binary-ivf branch 2 times, most recently from f5529f2 to 90479ed Compare July 11, 2024 21:21

naveentatikonda reviewed Jul 11, 2024

View reviewed changes

src/main/java/org/opensearch/knn/training/TrainingJob.java Outdated Show resolved Hide resolved

heemin32 force-pushed the binary-java branch from 930a27e to eacc478 Compare July 11, 2024 23:20

junqiu-lei force-pushed the binary-ivf branch 9 times, most recently from c22e654 to 2fce648 Compare July 13, 2024 03:11

Vikasht34 reviewed Jul 13, 2024

View reviewed changes

src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java Outdated Show resolved Hide resolved

src/main/java/org/opensearch/knn/index/codec/transfer/VectorTransfer.java Outdated Show resolved Hide resolved

heemin32 force-pushed the binary-java branch from eacc478 to 5139b16 Compare July 15, 2024 17:35

junqiu-lei changed the base branch from binary-java to feature/binary-format July 15, 2024 17:49