fix(semantic highlighter): add vector similarity queries and bbq_disk support by mromaios · Pull Request #138140 · elastic/elasticsearch

mromaios · 2025-11-17T00:44:57Z

This PR improves the handling of vector based queries in the SemanticTextHighlighter.

The highlighter will now correctly extract the underlying vector data from queries wrapped in VectorSimilarityQuery (e.g. when a similarity threshold is set in a knn query) as well as vector queries on a bbq_disk dense vector type.

# Create an index
PUT test-index
{
  "mappings": {
    "properties": {
      "inference_field": {
        "type": "semantic_text",
        "inference_id": ".multilingual-e5-small-elasticsearch"
      }
    }
  }
}

# Add a document
POST test-index/_doc/1
{
  "inference_field": "This is a test sentence for semantic search"
}

# Issue a highlighting query  
GET /test-index/_search
{
  "from" : 0,
  "size" : 10,
  "sort" : [
    {
      "_score" : {
        "order" : "asc"
      }
    }
  ],
  "query" : {
    "knn" : {
      "field" : "inference_field",
      "query_vector_builder" : {
        "text_embedding" : {
          "model_text" : "sentence"
        }
      },
      "similarity": 0.8,
      "k" : 12,
      "num_candidates" : 90
    }
  },
  "highlight" : {
    "fields" : {
      "inference_field" : {
        "order" : "score",
        "number_of_fragments" : 1
      }
    }
  }
}

---------------------
{
  "took" : 1811,
  "timed_out" : false,
  "_shards" : {
    "total" : 1,
    "successful" : 1,
    "skipped" : 0,
    "failed" : 0
  },
  "hits" : {
    "total" : {
      "value" : 1,
      "relation" : "eq"
    },
    "max_score" : null,
    "hits" : [
      {
        "_index" : "test-index",
        "_id" : "1",
        "_score" : 0.92290497,
        "_source" : {
          "inference_field" : "This is a test sentence for semantic search"
        },
        "highlight" : {
          "inference_field" : [
            "This is a test sentence for semantic search"
          ]
        },
        "sort" : [
          0.92290497
        ]
      }
    ]
  }
}

elasticsearchmachine · 2025-11-17T11:33:31Z

Hi @mromaios, I've created a changelog YAML for you.

...rence/src/main/java/org/elasticsearch/xpack/inference/highlight/SemanticTextHighlighter.java

.../src/test/java/org/elasticsearch/xpack/inference/highlight/SemanticTextHighlighterTests.java

docs/changelog/138140.yaml

elasticsearchmachine · 2025-11-19T01:01:04Z

Pinging @elastic/search-relevance (Team:Search - Relevance)

…om:mromaios/elasticsearch into fix_min_similarity_in_semantic_highlighter

Mikep86

Looks good overall, I like the approach! I think this is good to go with a little more work on the tests 🚀

docs/changelog/138140.yaml

Mikep86 · 2025-11-19T16:48:27Z

...rence/src/main/java/org/elasticsearch/xpack/inference/highlight/SemanticTextHighlighter.java

-            @Override
-            public void visitLeaf(Query query) {
+            private void visitLeaf(Query query, Float similarity) {
                if (query instanceof KnnFloatVectorQuery knnQuery) {
-                    queries.add(fieldType.createExactKnnQuery(VectorData.fromFloats(knnQuery.getTargetCopy()), null));
+                    queries.add(fieldType.createExactKnnQuery(VectorData.fromFloats(knnQuery.getTargetCopy()), similarity));
                } else if (query instanceof KnnByteVectorQuery knnQuery) {
-                    queries.add(fieldType.createExactKnnQuery(VectorData.fromBytes(knnQuery.getTargetCopy()), null));
+                    queries.add(fieldType.createExactKnnQuery(VectorData.fromBytes(knnQuery.getTargetCopy()), similarity));
                } else if (query instanceof MatchAllDocsQuery) {
                    queries.add(new MatchAllDocsQuery());
                } else if (query instanceof DenseVectorQuery.Floats floatsQuery) {
-                    queries.add(fieldType.createExactKnnQuery(VectorData.fromFloats(floatsQuery.getQuery()), null));
+                    queries.add(fieldType.createExactKnnQuery(VectorData.fromFloats(floatsQuery.getQuery()), similarity));
+                } else if (query instanceof IVFKnnFloatVectorQuery ivfQuery) {
+                    queries.add(fieldType.createExactKnnQuery(VectorData.fromFloats(ivfQuery.getQuery()), similarity));
+                } else if (query instanceof RescoreKnnVectorQuery rescoreQuery) {
+                    visitLeaf(rescoreQuery.innerQuery(), similarity);
+                } else if (query instanceof VectorSimilarityQuery similarityQuery) {
+                    visitLeaf(similarityQuery.getInnerKnnQuery(), similarityQuery.getSimilarity());
                }
            }
+
+            @Override
+            public void visitLeaf(Query query) {
+                visitLeaf(query, null);
+
+            }


Nice use of recursion to solve the problem here ❤️

.../src/test/java/org/elasticsearch/xpack/inference/highlight/SemanticTextHighlighterTests.java

Mikep86 · 2025-11-19T17:32:31Z

.../src/test/java/org/elasticsearch/xpack/inference/highlight/SemanticTextHighlighterTests.java

+    // TODO: figure out why this fails
+    // @SuppressWarnings("unchecked")
+    // public void testDenseVectorWithDiscBBQ() throws Exception {
+    // var mapperService = createDefaultMapperService(useLegacyFormat);
+    // Map<String, Object> queryMap = (Map<String, Object>) queries.get("dense_vector_1");
+    // float[] vector = readDenseVector(queryMap.get("embeddings"));
+    // var fieldType = (SemanticTextFieldMapper.SemanticTextFieldType)
+    // mapperService.mappingLookup().getFieldType(SEMANTIC_FIELD_E5_DISK_BBQ);
+    //
+    // KnnVectorQueryBuilder knnQuery = new KnnVectorQueryBuilder(
+    // fieldType.getEmbeddingsField().fullPath(),
+    // vector,
+    // 10,
+    // 10,
+    // 10f,
+    // null,
+    // null
+    // );
+    // NestedQueryBuilder nestedQueryBuilder = new NestedQueryBuilder(fieldType.getChunksField().fullPath(), knnQuery, ScoreMode.Max);
+    // var shardRequest = createShardSearchRequest(nestedQueryBuilder);
+    // var sourceToParse = new SourceToParse("0", readSampleDoc(useLegacyFormat), XContentType.JSON);
+    //
+    // String[] expectedPassages = ((List<String>) queryMap.get("expected_with_similarity_threshold")).toArray(String[]::new);
+    // assertHighlightOneDoc(
+    // mapperService,
+    // shardRequest,
+    // sourceToParse,
+    // SEMANTIC_FIELD_E5_DISK_BBQ,
+    // expectedPassages.length,
+    // HighlightBuilder.Order.SCORE,
+    // expectedPassages
+    // );
+    // }


Any insights on the problems with this case?

Yep, me 😓🤦

Sooooo, of course we don't create embeddings in the UT, but read them from the gzipped file.

var sourceToParse = new SourceToParse("0", readSampleDoc(useLegacyFormat), XContentType.JSON);

which I hadn't updated to include the new field.

kderusso

Looks good! On top of Mike's comments, I would love to see a yaml test added under 90_semantic_text_highlighter.yml as well as the associated BWC. Remember that for yaml tests we'll also need to make sure to create a feature/capability for Green CI 👍

…_semantic_highlighter

…ec/test/inference/90_semantic_text_highlighter_bwc.yml Co-authored-by: Kathleen DeRusso <kathleen.derusso@elastic.co>

elasticsearchmachine · 2025-11-24T17:14:57Z

💔 Backport failed

You can use sqren/backport to manually backport by running backport --upstream elastic/elasticsearch --pr 138140

mromaios · 2025-11-24T17:19:49Z

💚 All backports created successfully

Status	Branch	Result
✅	9.2

Questions ?

Please refer to the Backport tool documentation

… support (elastic#138140) (cherry picked from commit 8603933) # Conflicts: # x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferenceFeatures.java

… support (elastic#138140) (cherry picked from commit 8603933) # Conflicts: # x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferenceFeatures.java # x-pack/plugin/inference/src/yamlRestTest/resources/rest-api-spec/test/inference/90_semantic_text_highlighter.yml

… support (elastic#138140) (cherry picked from commit 8603933) # Conflicts: # server/src/main/java/org/elasticsearch/search/vectors/IVFKnnFloatVectorQuery.java # x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferenceFeatures.java # x-pack/plugin/inference/src/yamlRestTest/resources/rest-api-spec/test/inference/90_semantic_text_highlighter.yml

mromaios · 2025-11-25T09:22:54Z

💚 All backports created successfully

Status	Branch	Result
✅	9.1
✅	8.19

Questions ?

Please refer to the Backport tool documentation

… support (#138140) (#138521)

…q_disk support (#138140) (#138555)

…bq_disk support (#138140) (#138558)

… support (elastic#138140)

fix: handle vector similarity queries in semantic text highlighter

390d192

mromaios changed the title ~~fix(semantic highlighter): handle vector similarity queries in semantic text highlighter~~ Nov 17, 2025

elasticsearchmachine added the v9.3.0 label Nov 17, 2025

[CI] Auto commit changes from spotless

a7e9daa

mromaios self-assigned this Nov 17, 2025

mromaios added >bug :SearchOrg/Relevance Label for the Search (solution/org) Relevance team labels Nov 17, 2025

Update docs/changelog/138140.yaml

aff96c7

mromaios commented Nov 17, 2025

View reviewed changes

...rence/src/main/java/org/elasticsearch/xpack/inference/highlight/SemanticTextHighlighter.java Outdated Show resolved Hide resolved

mromaios commented Nov 19, 2025

View reviewed changes

...rence/src/main/java/org/elasticsearch/xpack/inference/highlight/SemanticTextHighlighter.java Outdated Show resolved Hide resolved

mromaios commented Nov 19, 2025

View reviewed changes

.../src/test/java/org/elasticsearch/xpack/inference/highlight/SemanticTextHighlighterTests.java Show resolved Hide resolved

mromaios commented Nov 19, 2025

View reviewed changes

docs/changelog/138140.yaml Outdated Show resolved Hide resolved

mromaios marked this pull request as ready for review November 19, 2025 01:00

elasticsearchmachine added the Team:Search - Relevance The Search organization Search Relevance team label Nov 19, 2025

mromaios marked this pull request as draft November 19, 2025 10:06

mromaios added 2 commits November 19, 2025 14:17

add coverage for bbq_disk

cd5548e

Merge branch 'fix_min_similarity_in_semantic_highlighter' of github.c…

4e58e97

…om:mromaios/elasticsearch into fix_min_similarity_in_semantic_highlighter

mromaios changed the title ~~fix(semantic highlighter): handle vector similarity queries~~ Nov 19, 2025

mromaios requested a review from a team November 19, 2025 13:25

elasticsearchmachine and others added 3 commits November 19, 2025 13:30

[CI] Auto commit changes from spotless

00310ed

Merge branch 'main' into fix_min_similarity_in_semantic_highlighter

721187b

Merge branch 'main' into fix_min_similarity_in_semantic_highlighter

1bf6743

Mikep86 reviewed Nov 19, 2025

View reviewed changes

kderusso reviewed Nov 19, 2025

View reviewed changes

mromaios and others added 4 commits November 20, 2025 11:37

addressing PR comments, improving UT, changelog, yaml

6e9b140

[CI] Auto commit changes from spotless

234bf9f

add yaml tests

709e2ce

Merge remote-tracking branch 'origin/main' into fix_min_similarity_in…

75984f8

…_semantic_highlighter

mromaios and others added 2 commits November 21, 2025 23:18

Update x-pack/plugin/inference/src/yamlRestTest/resources/rest-api-sp…

735c4c3

…ec/test/inference/90_semantic_text_highlighter_bwc.yml Co-authored-by: Kathleen DeRusso <kathleen.derusso@elastic.co>

Merge branch 'main' into fix_min_similarity_in_semantic_highlighter

9b8f8a3

mromaios added branch:9.2 branch:9.1 branch:8.19 labels Nov 24, 2025

elasticsearchmachine added v9.2.2 v8.19.8 v9.1.8 and removed branch:9.2 branch:9.1 branch:8.19 labels Nov 24, 2025

mromaios added the auto-backport Automatically create backport pull requests when merged label Nov 24, 2025

mromaios merged commit 8603933 into elastic:main Nov 24, 2025
34 checks passed

elasticsearchmachine added the backport pending label Nov 24, 2025

mromaios mentioned this pull request Nov 24, 2025

[9.2] fix(semantic highlighter): add vector similarity queries and bbq_disk support (#138140) #138521

Merged

mromaios mentioned this pull request Nov 25, 2025

[9.1] fix(semantic highlighter): add vector similarity queries and bbq_disk support (#138140) #138555

Merged

mromaios mentioned this pull request Nov 25, 2025

[8.19] fix(semantic highlighter): add vector similarity queries and bbq_disk support (#138140) #138558

Merged

mromaios added a commit that referenced this pull request Nov 25, 2025

fix(semantic highlighter): add vector similarity queries and bbq_disk…

ec3c1f7

… support (#138140) (#138521)

mromaios added a commit that referenced this pull request Nov 25, 2025

[9.1] fix(semantic highlighter): add vector similarity queries and bb…

d9dc935

…q_disk support (#138140) (#138555)

mromaios added a commit that referenced this pull request Nov 25, 2025

[8.19] fix(semantic highlighter): add vector similarity queries and b…

0f3f4e9

…bq_disk support (#138140) (#138558)

afoucret pushed a commit to afoucret/elasticsearch that referenced this pull request Nov 26, 2025

fix(semantic highlighter): add vector similarity queries and bbq_disk…

3d5fd95

… support (elastic#138140)

mromaios mentioned this pull request Nov 26, 2025

fix: add missing vector_similarity_support in InferenceFeatures #138644

Merged

ncordon pushed a commit to ncordon/elasticsearch that referenced this pull request Nov 26, 2025

fix(semantic highlighter): add vector similarity queries and bbq_disk…

b8bfd42

… support (elastic#138140)

mromaios mentioned this pull request Nov 26, 2025

chore(SemanticTextHighlighter): improve bwc tests #138662

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

fix(semantic highlighter): add vector similarity queries and bbq_disk support #138140

fix(semantic highlighter): add vector similarity queries and bbq_disk support #138140
mromaios merged 19 commits intoelastic:mainfrom
mromaios:fix_min_similarity_in_semantic_highlighter

mromaios commented Nov 17, 2025 •

edited

Loading

elasticsearchmachine commented Nov 17, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

elasticsearchmachine commented Nov 19, 2025

Mikep86 left a comment

Uh oh!

Mikep86 Nov 19, 2025

Uh oh!

Mikep86 Nov 19, 2025

mromaios Nov 20, 2025

kderusso left a comment

Uh oh!

elasticsearchmachine commented Nov 24, 2025

mromaios commented Nov 24, 2025

mromaios commented Nov 25, 2025

Labels

4 participants

Conversation

mromaios commented Nov 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

elasticsearchmachine commented Nov 17, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

elasticsearchmachine commented Nov 19, 2025

Mikep86 left a comment

Choose a reason for hiding this comment

Uh oh!

Mikep86 Nov 19, 2025

Choose a reason for hiding this comment

Uh oh!

Mikep86 Nov 19, 2025

Choose a reason for hiding this comment

mromaios Nov 20, 2025

Choose a reason for hiding this comment

kderusso left a comment

Choose a reason for hiding this comment

Uh oh!

elasticsearchmachine commented Nov 24, 2025

💔 Backport failed

mromaios commented Nov 24, 2025

💚 All backports created successfully

Questions ?

mromaios commented Nov 25, 2025

💚 All backports created successfully

Questions ?

Labels

4 participants

mromaios commented Nov 17, 2025 •

edited

Loading