Skip to content

ES|QL: Enable semantic search in FORK #125960

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Apr 1, 2025
Merged

ES|QL: Enable semantic search in FORK #125960

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 20 additions & 0 deletions x-pack/plugin/esql/qa/testFixtures/src/main/resources/fork.csv-spec
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 changes: 18 additions & 0 deletions x-pack/plugin/esql/qa/testFixtures/src/main/resources/rrf.csv-spec
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 changes: 33 additions & 6 deletions .../java/org/elasticsearch/xpack/esql/expression/function/fulltext/QueryBuilderResolver.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,10 @@
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;
Expand All @@ -22,6 +24,8 @@
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}
Expand All @@ -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),
Expand Down Expand Up @@ -69,12 +69,29 @@ private static Set<String> indexNames(LogicalPlan plan) {
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) {
Copy link
Member

@carlosdelest carlosdelest Apr 1, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should Fork implement forEachExpression instead? That would make things easier for the clients that will not need to take it into account, and be in line with the Composite pattern...

Copy link
Contributor Author

@ioanatia ioanatia Apr 1, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No, Fork should not implement forEachExpression.
There is a slippery slope if we decide to do that since in the future we might decide we need to implement anyMatch, forEachExpressionUp, forEachExpressionDown etc and that's not what we want.

We either treat Fork as the outlier it currently is (because we have nothing like it in ES|QL atm) or we go back and make Fork a n-ary plan (it is currently a UnaryPlan).
If Fork sub plans are treated as children then we would most likely not need to make any big changes to QueryBuilderResolver.
While it sound good in theory to just make Fork a n-ary plan, we tried this already when we did the POC and found out it has its own set of challenges.
So for now I am focusing on getting Fork to work as it should even if it means having this special handling in a couple of places.

Copy link
Member

@carlosdelest carlosdelest Apr 1, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If you already went down that route, that's enough to me 😅

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 {
Expand All @@ -91,5 +108,15 @@ public FullTextFunctionsRewritable rewrite(QueryRewriteContext ctx) throws IOExc
}
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 -> {
Copy link
Member

@carlosdelest carlosdelest Apr 1, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same here, we may override transformExpressionsDown for Fork

var subPlans = fork.subPlans()
.stream()
.map(subPlan -> subPlan.transformExpressionsDown(FullTextFunction.class, rule))
.collect(Collectors.toList());
return new Fork(fork.source(), fork.child(), subPlans);
});
}
}
}