github-mirrors/elasticsearch
1
0
Fork 0
mirror of https://github.com/elastic/elasticsearch.git synced 2025-06-28 09:28:55 -04:00
Code Issues Projects Releases Packages Wiki Activity

Enable semantic search in FORK (#125960)

Browse source
This commit is contained in:
Ioana Tagirta 2025-04-01 13:10:50 +02:00 • committed by GitHub
parent 3fbcb3cf9a
commit fd1c008c35
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 71 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

20
x-pack/plugin/esql/qa/testFixtures/src/main/resources/fork.csv-spec
View file

@ -92,3 +92,23 @@ fork1 | 10052
fork2 | 10099
fork2 | 10100
;
forkWithSemanticSearchAndScore
required_capability: fork
required_capability: semantic_text_field_caps
required_capability: metadata_score
FROM semantic_text METADATA _id, _score
| FORK ( WHERE semantic_text_field:"something" | SORT _score DESC | LIMIT 2)
( WHERE semantic_text_field:"something else" | SORT _score DESC | LIMIT 2)
| EVAL _score = round(_score, 4)
| SORT _fork, _score, _id
| KEEP _fork, _score, _id, semantic_text_field
;
_fork:keyword | _score:double | _id:keyword | semantic_text_field:text
fork1 | 2.156063961865257E18 | 3 | be excellent to each other
fork1 | 5.603396578413904E18 | 2 | all we have to decide is what to do with the time that is given to us
fork2 | 2.3447541759648727E18 | 3 | be excellent to each other
fork2 | 6.093784261960139E18 | 2 | all we have to decide is what to do with the time that is given to us
;

18
x-pack/plugin/esql/qa/testFixtures/src/main/resources/rrf.csv-spec
View file

@ -109,3 +109,21 @@ _score:double | author:keyword | title:keyword | _fork
0.0161 | Ursula K. Le Guin | The Word For World i | fork2
0.0159 | Ursula K. Le Guin | The Dispossessed | fork2
;
rrfWithSemanticSearch
required_capability: rrf
required_capability: semantic_text_field_caps
required_capability: metadata_score
FROM semantic_text METADATA _id, _score, _index
| FORK ( WHERE semantic_text_field:"something" | SORT _score DESC | LIMIT 2)
( WHERE semantic_text_field:"something else" | SORT _score DESC | LIMIT 2)
| RRF
| EVAL _score = round(_score, 4)
| KEEP _fork, _score, _id, semantic_text_field
;
_fork:keyword | _score:double | _id:keyword | semantic_text_field:keyword
[fork1, fork2] | 0.0328 | 2 | all we have to decide is what to do with the time that is given to us
[fork1, fork2] | 0.0323 | 3 | be excellent to each other
;

39
x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/QueryBuilderResolver.java
View file

@ -12,8 +12,10 @@ import org.elasticsearch.action.ResolvedIndices;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryRewriteContext;
import org.elasticsearch.index.query.Rewriteable;
import org.elasticsearch.xpack.esql.core.expression.Expression;
import org.elasticsearch.xpack.esql.core.util.Holder;
import org.elasticsearch.xpack.esql.plan.logical.EsRelation;
import org.elasticsearch.xpack.esql.plan.logical.Fork;
import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
import org.elasticsearch.xpack.esql.planner.TranslatorHandler;
import org.elasticsearch.xpack.esql.plugin.TransportActionServices;
@ -22,6 +24,8 @@ import org.elasticsearch.xpack.esql.session.IndexResolver;
import java.io.IOException;
import java.util.HashSet;
import java.util.Set;
import java.util.function.Function;
import java.util.stream.Collectors;
/**
* Some {@link FullTextFunction} implementations such as {@link org.elasticsearch.xpack.esql.expression.function.fulltext.Match}
@ -34,11 +38,7 @@ public final class QueryBuilderResolver {
private QueryBuilderResolver() {}
public static void resolveQueryBuilders(LogicalPlan plan, TransportActionServices services, ActionListener<LogicalPlan> listener) {
var hasFullTextFunctions = plan.anyMatch(p -> {
Holder<Boolean> hasFullTextFunction = new Holder<>(false);
p.forEachExpression(FullTextFunction.class, unused -> hasFullTextFunction.set(true));
return hasFullTextFunction.get();
});
var hasFullTextFunctions = hasFullTextFunctions(plan);
if (hasFullTextFunctions) {
Rewriteable.rewriteAndFetch(
new FullTextFunctionsRewritable(plan),
@ -69,12 +69,29 @@ public final class QueryBuilderResolver {
return indexNames;
}
private static boolean hasFullTextFunctions(LogicalPlan plan) {
return plan.anyMatch(p -> {
Holder<Boolean> hasFullTextFunction = new Holder<>(false);
p.forEachExpression(FullTextFunction.class, unused -> hasFullTextFunction.set(true));
if (p instanceof Fork fork) {
fork.subPlans().forEach(subPlan -> {
if (hasFullTextFunctions(subPlan)) {
hasFullTextFunction.set(true);
}
});
}
return hasFullTextFunction.get();
});
}
private record FullTextFunctionsRewritable(LogicalPlan plan) implements Rewriteable<QueryBuilderResolver.FullTextFunctionsRewritable> {
@Override
public FullTextFunctionsRewritable rewrite(QueryRewriteContext ctx) throws IOException {
Holder<IOException> exceptionHolder = new Holder<>();
Holder<Boolean> updated = new Holder<>(false);
LogicalPlan newPlan = plan.transformExpressionsDown(FullTextFunction.class, f -> {
LogicalPlan newPlan = transformPlan(plan, f -> {
QueryBuilder builder = f.queryBuilder(), initial = builder;
builder = builder == null ? f.asQuery(TranslatorHandler.TRANSLATOR_HANDLER).toQueryBuilder() : builder;
try {
@ -91,5 +108,15 @@ public final class QueryBuilderResolver {
}
return updated.get() ? new FullTextFunctionsRewritable(newPlan) : this;
}
private LogicalPlan transformPlan(LogicalPlan plan, Function<FullTextFunction, ? extends Expression> rule) {
return plan.transformExpressionsDown(FullTextFunction.class, rule).transformDown(Fork.class, fork -> {
var subPlans = fork.subPlans()
.stream()
.map(subPlan -> subPlan.transformExpressionsDown(FullTextFunction.class, rule))
.collect(Collectors.toList());
return new Fork(fork.source(), fork.child(), subPlans);
});
}
}
}