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.

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

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

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.

Thanks for the update but could you try to remove accept method as well? It is not great to have it just because it is used in the test.

The accept is also used in the processTrainingVectors.

Yes. But we can change it to private for that use case.

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

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

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

heemin32 · 2024-07-19T17:03:13Z

src/test/java/org/opensearch/knn/training/VectorReaderTests.java


-        List<Float[]> vectorsConsumed;
+    // create test float training data consumer class extending FloatTrainingDataConsumer
+    private static class TestFloatTrainingDataConsumer extends FloatTrainingDataConsumer {


Why do we need this class? I think we can just use FloatTrainingDataConsumer?

In the TestFloatTrainingDataConsumer processTrainingVectors method, it additionally count the total vectors used for test only.

As I mentioned in another comment, I think we don't need to validate total vectors count. All we need is to verifying that JNIService is called with expected method and parameters.

heemin32 · 2024-07-19T17:04:10Z

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

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


Is there a unit test for this method?

Yes, in VectorReaderTests.java by TestFloatTrainingDataConsumer.

I think, we should name it as FloatTrainingDataConsumerTests

There is a existing test class named FloatTrainingDataConsumerTests, but I think VectorReaderTests can cover the test of processTrainingVectors. Will have another PR for the refactor.

heemin32

Approving. We need to refactor TrainingDataConsumer though.

…-project#1784) Signed-off-by: Junqiu Lei <[email protected]>

) 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 mentioned this pull request Jul 18, 2024

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

Closed

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

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

heemin32 reviewed Jul 18, 2024

View reviewed changes

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

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

heemin32 reviewed Jul 19, 2024

View reviewed changes

heemin32 approved these changes Jul 22, 2024

View reviewed changes

jmazanec15 approved these changes Jul 22, 2024

View reviewed changes

junqiu-lei merged commit abae5d5 into opensearch-project:feature/binary-format Jul 22, 2024
50 of 51 checks passed

junqiu-lei deleted the binary-ivf branch July 22, 2024 17:14

junqiu-lei added a commit to junqiu-lei/k-NN that referenced this pull request Jul 22, 2024

Add binary format support with IVF method in Faiss Engine (opensearch…

77b44bb

…-project#1784) Signed-off-by: Junqiu Lei <[email protected]>

junqiu-lei mentioned this pull request Jul 22, 2024

[Feature to main] Add binary format support with IVF method in Faiss Engine (#1784) #1862

Merged

5 tasks

junqiu-lei added a commit to junqiu-lei/k-NN that referenced this pull request Jul 22, 2024

Add binary format support with IVF method in Faiss Engine (opensearch…

f7da37b

…-project#1784) Signed-off-by: Junqiu Lei <[email protected]>

junqiu-lei added a commit to junqiu-lei/k-NN that referenced this pull request Jul 22, 2024

Add binary format support with IVF method in Faiss Engine (opensearch…

723f761

…-project#1784) Signed-off-by: Junqiu Lei <[email protected]>

junqiu-lei added a commit that referenced this pull request Jul 22, 2024

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

f292124

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

junqiu-lei added a commit that referenced this pull request Jul 22, 2024

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

2091af3

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

junqiu-lei added a commit that referenced this pull request Jul 22, 2024

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

2908920

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

junqiu-lei mentioned this pull request Aug 1, 2024

Refactor train index and create index from template APIs in JNI layer #1918

Open

5 tasks

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

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 19, 2024

heemin32 Jul 19, 2024

junqiu-lei Jul 22, 2024

heemin32 Jul 22, 2024

heemin32 Jul 19, 2024

junqiu-lei Jul 22, 2024

heemin32 Jul 22, 2024

heemin32 Jul 19, 2024

junqiu-lei Jul 22, 2024

heemin32 Jul 22, 2024

junqiu-lei Jul 22, 2024

heemin32 left a comment

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

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

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

heemin32 left a comment

Choose a reason for hiding this comment

junqiu-lei commented Jul 2, 2024 •

edited

Loading