Add binary format support with IVF method in Faiss Engine #1784

junqiu-lei · 2024-07-02T16:45:14Z

Resolved comments in PR heemin32#2, which was friendly to check the file diffs when #1781 wasn't merged. Because #1781 now is merged, I rebased junqiu-lei:binary-ivf against opensearch-project:feature/binary-format

Description

This PR will support using binary format with Faiss IVF method, it mainly have changes:

introduce data_type field when train model
support train model with binary format index
create binary format IVF target index with with binary format model
query IVF index with binary format

JNI layer related refactor works will be complete in another PR tracked by #1846

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
                    ]
                }
            }
        ]
    }
}

Issues Resolved

part of #1767

Check List

New functionality includes testing.
- All tests pass
New functionality has been documented.
- New functionality has javadoc added
Commits are signed as per the DCO using --signoff

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.

jni/include/faiss_wrapper.h

heemin32 · 2024-07-17T16:33:29Z

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

-                floats.stream().map(ArrayUtils::toPrimitive).toArray(float[][]::new)
-            )
-        );
+        long memoryAddress = trainingDataAllocation.getMemoryAddress();


Please add a unit test to validate the code.

src/main/java/org/opensearch/knn/plugin/transport/TrainingModelRequest.java

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

src/main/java/org/opensearch/knn/index/memory/NativeMemoryLoadStrategy.java

heemin32 · 2024-07-17T16:47:53Z

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

@@ -248,6 +263,28 @@ public static ValidationException validateKnnField(
            return exception;
        }

+        // Return if vector data type does not need to be checked
+        if (expectedVectorDataType == null) {


Try to write a unit test. You may find an issue with current implementation.
It is okay to return single exception but the issue here is you are returning when there is no exception. Then the rest of the validation code won't work.
Also, please check if your code works when expectedDimension is < 0 or null. In such case, your validation won't get executed.

Vikasht34

Can we have Unit test for all c++ functions written.

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

Vikasht34 · 2024-07-17T17:52:13Z

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

+        long memoryAddress = trainingDataAllocation.getMemoryAddress();
+
+        if (IndexUtil.isBinaryIndex(this.vectorDataType)) {
+            byte[][] byteArray = floats.stream()


Why we need to do this conversion , if it's Binary Index then training would be also on Byte array right , why customer will provide Float array to train Binary Index?

Refactored to have FloatTrainingDataConsumer and ByteTrainingDataConsumer to handle train data consuming separately.

src/main/java/org/opensearch/knn/jni/JNIService.java

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

junqiu-lei · 2024-07-18T08:32:36Z

Can we have Unit test for all c++ functions written.

will include jni tests along with refactor work

jmazanec15 · 2024-07-18T20:19:59Z

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

+
+    private boolean isBinaryField(FieldInfo field) {
+        if (field.attributes().containsKey(MODEL_ID)) {
+            Model model = ModelCache.getInstance().get(field.attributes().get(MODEL_ID));


Model can be quite big because it contains the binary blob. Can we read from model metadata instead via ModelDao?

Updated code to use model only once in the addKNNBinaryField function

heemin32 · 2024-07-18T21:23:48Z

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

+    }
+
+    @Override
+    public List<byte[]> getTotalAddedVectors() {


I see this method is called below. Where the implementation is?

this method is used in TestFloatTrainingDataConsumer

// create test float training data consumer class extending FloatTrainingDataConsumer private static class TestFloatTrainingDataConsumer extends FloatTrainingDataConsumer { @Getter private List<Float[]> totalAddedVectors = new ArrayList<>(); public TestFloatTrainingDataConsumer(NativeMemoryAllocation.TrainingDataAllocation trainingDataAllocation) { super(trainingDataAllocation); } @Override public void processTrainingVectors(SearchResponse searchResponse, int vectorsToAdd, String fieldName) { List<Float[]> vectors = extractVectorsFromHits(searchResponse, vectorsToAdd, fieldName); totalAddedVectors.addAll(vectors); setTotalVectorsCountAdded(getTotalVectorsCountAdded() + vectors.size()); accept(vectors); } }

It is wired to have a method only to be used in testing... Can we do differently?

I tried but couldn't find other way to read the whole vectors

You don't need to. For unit test, all you need to verify is that JNILayer method is called properly with expected parameter.

heemin32 · 2024-07-18T21:26:02Z

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

+    }
+
+    @Override
+    public void accept(List<byte[]> byteVectors) {


Why public?

previously implemented from Consumer, which are public, now removed parent implement from Consumer and changed it to protected

It can be even private. I don't see it being used outside of this class.

accept method is used in some tests, but I removed other methods.

heemin32 · 2024-07-18T21:26:09Z

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

+    }
+
+    @Override
+    public List<byte[]> extractVectorsFromHits(SearchResponse searchResponse, int vectorsToAdd, String fieldName) {


Why public?

heemin32 · 2024-07-18T21:26:28Z

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

+    }
+
+    @Override
+    public void processTrainingVectors(SearchResponse searchResponse, int vectorsToAdd, String fieldName) {


I think only this method is needs to be public?

heemin32 · 2024-07-18T21:32:09Z

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

 */
-public class TrainingDataConsumer implements Consumer<List<Float[]>> {
+public abstract class TrainingDataConsumer<T> implements Consumer<List<T>> {


Why does it have to implement Consumer?
I think accept method is called only inside processTrainingVectors

Updated without implement from Consumer

heemin32 · 2024-07-18T23:34:46Z

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

 */
-public class TrainingDataConsumer implements Consumer<List<Float[]>> {
+public abstract class TrainingDataConsumer<T> {


I feel we don't need T here.
I see only public abstract void processTrainingVectors(SearchResponse searchResponse, int vectorsToAdd, String fieldName); is actually used.

heemin32 · 2024-07-18T23:35:22Z

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

-                floats.stream().map(ArrayUtils::toPrimitive).toArray(float[][]::new)
-            )
-        );
+    protected abstract void accept(List<T> vectors);


We can remove accept, extractVectorsFromHits, getTotalAddedVectors which is not being called outside of this class.

Removed extractVectorsFromHits, getTotalAddedVectors and kept accept to expose for some tests

Signed-off-by: Junqiu Lei <[email protected]>

junqiu-lei added Features Introduces a new unit of functionality that satisfies a requirement feature branch labels Jul 2, 2024

heemin32 force-pushed the feature/binary-format branch 8 times, most recently from dafd79b to a913082 Compare July 3, 2024 16:40

junqiu-lei force-pushed the binary-ivf branch 2 times, most recently from 22d044a to 629757b Compare July 9, 2024 00:07

junqiu-lei changed the title ~~Add binary format support for Faiss IVF train model and create index api~~ Add binary format support for Faiss IVF Jul 9, 2024

junqiu-lei force-pushed the binary-ivf branch 2 times, most recently from f726d81 to 648f342 Compare July 9, 2024 00:49

junqiu-lei mentioned this pull request Jul 9, 2024

Add binary format support for Faiss IVF heemin32/k-NN#2

Closed

junqiu-lei force-pushed the binary-ivf branch from 6c96b1a to 404056a Compare July 9, 2024 23:11

junqiu-lei changed the title ~~Add binary format support for Faiss IVF~~ Add binary format support with IVF method in Faiss Engine Jul 9, 2024

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

junqiu-lei marked this pull request as ready for review July 10, 2024 17:31

junqiu-lei requested review from heemin32, navneet1v, VijayanB, vamshin, jmazanec15, naveentatikonda, martin-gaievski, ryanbogan and luyuncheng as code owners July 10, 2024 17:31

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

junqiu-lei self-assigned this Jul 17, 2024

junqiu-lei added the v2.16.0 label Jul 17, 2024

junqiu-lei requested a review from heemin32 July 17, 2024 15:59

jmazanec15 reviewed Jul 17, 2024

View reviewed changes

jni/include/faiss_wrapper.h Outdated Show resolved Hide resolved

heemin32 reviewed Jul 17, 2024

View reviewed changes

Vikasht34 reviewed Jul 17, 2024

View reviewed changes

junqiu-lei force-pushed the binary-ivf branch from a920c96 to cfbf209 Compare July 17, 2024 20:45

junqiu-lei mentioned this pull request Jul 17, 2024

Refactor train model method in JNI #1846

Open

junqiu-lei force-pushed the binary-ivf branch 3 times, most recently from a1fe81b to 93fef1e Compare July 18, 2024 08:30

junqiu-lei force-pushed the binary-ivf branch from 93fef1e to 8cb1692 Compare July 18, 2024 15:52

junqiu-lei mentioned this pull request Jul 18, 2024

[FEATURE] Support of dynamic efsearch and nprobs value during search for binary index #1769

Open

junqiu-lei force-pushed the binary-ivf branch from 8cb1692 to c6832f3 Compare July 18, 2024 18:49

junqiu-lei requested review from heemin32, Vikasht34 and jmazanec15 July 18, 2024 18:54

junqiu-lei force-pushed the binary-ivf branch from c6832f3 to fc41d4c Compare July 18, 2024 20:20

jmazanec15 reviewed Jul 18, 2024

View reviewed changes

junqiu-lei force-pushed the binary-ivf branch from fc41d4c to 4d1000f Compare July 18, 2024 21:05

junqiu-lei requested a review from jmazanec15 July 18, 2024 21:19

heemin32 reviewed Jul 18, 2024

View reviewed changes

junqiu-lei force-pushed the binary-ivf branch from 4d1000f to d227c8f Compare July 18, 2024 22:02

junqiu-lei requested a review from heemin32 July 18, 2024 22:33

heemin32 reviewed Jul 18, 2024

View reviewed changes

Add binary format support with IVF method in Faiss Engine

5685992

Signed-off-by: Junqiu Lei <[email protected]>

junqiu-lei force-pushed the binary-ivf branch from d227c8f to 5685992 Compare July 19, 2024 00:18

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add binary format support with IVF method in Faiss Engine #1784

Add binary format support with IVF method in Faiss Engine #1784

junqiu-lei commented Jul 2, 2024 •

edited

Loading

heemin32 Jul 17, 2024

junqiu-lei Jul 17, 2024

heemin32 Jul 17, 2024

Vikasht34 left a comment

Vikasht34 Jul 17, 2024

junqiu-lei Jul 18, 2024

junqiu-lei commented Jul 18, 2024 •

edited

Loading

jmazanec15 Jul 18, 2024

junqiu-lei Jul 18, 2024

heemin32 Jul 18, 2024

junqiu-lei Jul 18, 2024

heemin32 Jul 18, 2024

junqiu-lei Jul 18, 2024

heemin32 Jul 18, 2024

junqiu-lei Jul 19, 2024

heemin32 Jul 18, 2024

junqiu-lei Jul 18, 2024

heemin32 Jul 18, 2024

junqiu-lei Jul 19, 2024

heemin32 Jul 18, 2024

junqiu-lei Jul 18, 2024

heemin32 Jul 18, 2024

junqiu-lei Jul 18, 2024

heemin32 Jul 18, 2024

junqiu-lei Jul 18, 2024

heemin32 Jul 18, 2024

junqiu-lei Jul 19, 2024

heemin32 Jul 18, 2024

junqiu-lei Jul 19, 2024

Add binary format support with IVF method in Faiss Engine #1784

Are you sure you want to change the base?

Add binary format support with IVF method in Faiss Engine #1784

Conversation

junqiu-lei commented Jul 2, 2024 • edited Loading

Resolved comments in PR heemin32#2, which was friendly to check the file diffs when #1781 wasn't merged. Because #1781 now is merged, I rebased junqiu-lei:binary-ivf against opensearch-project:feature/binary-format

Description

Example workflow

1. Create binary format train index

2. Ingest tran index

3. Create train model

4. Create IVF binary format target index

5. Bulk target index

6. Query target index

7. Query result

Issues Resolved

Check List

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Vikasht34 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

junqiu-lei commented Jul 18, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

junqiu-lei commented Jul 2, 2024 •

edited

Loading

junqiu-lei commented Jul 18, 2024 •

edited

Loading