[ES\QL] Text embedding function constant folding#135710
[ES\QL] Text embedding function constant folding#135710afoucret merged 19 commits intoelastic:mainfrom
Conversation
|
Pinging @elastic/es-search-relevance (Team:Search Relevance) |
…esql_text_embedding_function_evaluator
dc7ac08 to
aaec4e1
Compare
ℹ️ Important: Docs version tagging👋 Thanks for updating the docs! Just a friendly reminder that our docs are now cumulative. This means all 9.x versions are documented on the same page and published off of the main branch, instead of creating separate pages for each minor version. We use applies_to tags to mark version-specific features and changes. Expand for a quick overviewWhen to use applies_to tags:✅ At the page level to indicate which products/deployments the content applies to (mandatory) What NOT to do:❌ Don't remove or replace information that applies to an older version 🤔 Need help?
|
| import static org.hamcrest.Matchers.equalTo; | ||
|
|
||
| @FunctionName("text_embedding") | ||
| public class TextEmbeddingTests extends AbstractFunctionTestCase { |
There was a problem hiding this comment.
ℹ️ This tests were not added before because they were an issue while the dense vector type was under construction.
| @@ -0,0 +1,6 @@ | |||
| % This is generated by ESQL's AbstractFunctionTestCase. Do not edit it. See ../README.md for how to regenerate it. | |||
There was a problem hiding this comment.
ℹ️ Doc update are caused by the dense vector type being out of snapshot.
| ; | ||
| // end::embedding-eval[] | ||
|
|
||
| input:keyword | embedding:dense_vector |
There was a problem hiding this comment.
can we get a test with multiple text_embedding calls with different query strings?
There was a problem hiding this comment.
I added a CSV test with FORK that should be covering your ask:
FROM semantic_text METADATA _score
| FORK (EVAL query_embedding = TEXT_EMBEDDING("be excellent to each other", "test_dense_inference") | WHERE KNN(semantic_text_dense_field, query_embedding))
(EVAL query_embedding = TEXT_EMBEDDING("live long and prosper", "test_dense_inference") | WHERE KNN(semantic_text_dense_field, query_embedding))
| KEEP semantic_text_field, query_embedding, _score, _fork
| EVAL _score = ROUND(_score, 4)
| SORT _score DESC
| LIMIT 10
;
semantic_text_field:text | query_embedding:dense_vector | _fork:keyword | _score:double
be excellent to each other | [45.0, 55.0, 54.0] | fork1 | 1.0
live long and prosper | [50.0, 57.0, 56.0] | fork2 | 1.0
be excellent to each other | [50.0, 57.0, 56.0] | fork2 | 0.0295
live long and prosper | [45.0, 55.0, 54.0] | fork1 | 0.0295
all we have to decide is what to do with the time that is given to us | [45.0, 55.0, 54.0] | fork1 | 0.0214
all we have to decide is what to do with the time that is given to us | [50.0, 57.0, 56.0] | fork2 | 0.0109
There was a problem hiding this comment.
plus it tests inference function in the context of fork
x-pack/plugin/esql/qa/testFixtures/src/main/resources/text-embedding.csv-spec
Outdated
Show resolved
Hide resolved
...in/esql/src/main/java/org/elasticsearch/xpack/esql/inference/InferenceFunctionEvaluator.java
Show resolved
Hide resolved
| * Example transformation: | ||
| * {@code TEXT_EMBEDDING("hello world", "model1")} → {@code [0.1, 0.2, 0.3, ...]} | ||
| */ | ||
| public class FoldInferenceFunctions implements LogicalPlanPreOptimizerRule { |
There was a problem hiding this comment.
I guess something like this would also be used for folding COMPLETION when it's used with a foldable prompt?
There was a problem hiding this comment.
You are right. There will be something similar for inference plans.
| import static org.mockito.Mockito.mock; | ||
| import static org.mockito.Mockito.when; | ||
|
|
||
| public class InferenceFunctionEvaluatorTests extends ComputeTestCase { |
There was a problem hiding this comment.
should these be skipped for release tests?
There was a problem hiding this comment.
I double-checked and it is not necessary.
The test is passing in release build since it does not use features that change in release build (function registry, writeable, ...)
This PR introduces constant folding optimization for the
TEXT_EMBEDDINGfunction in ESQL.Key components:
InferenceFunctionEvaluator: a new piece of infrastructure that uses anInferenceOperatorto evaluate an inference function. Currently, only folding of constants is supported, but it could be extended to more use cases in the future.Added a rules mechanism to the
LogicalPlanPreOptimizerand introduced a rule that folds inference functions (FoldInferenceFunctions)Adding CSV tests for the TEXT_EMBEDDING function, including usage with other vector functions such as KNN)
Part of #131022