diff --git a/muted-tests.yml b/muted-tests.yml
index 07383b10c408b..98b8cf69e385d 100644
--- a/muted-tests.yml
+++ b/muted-tests.yml
@@ -435,9 +435,6 @@ tests:
- class: org.elasticsearch.xpack.esql.heap_attack.HeapAttackLookupJoinIT
method: testLookupExplosionBigString
issue: https://github.com/elastic/elasticsearch/issues/138510
-- class: org.elasticsearch.xpack.esql.qa.single_node.GenerativeForkIT
- method: test {csv-spec:inlinestats.MvMinMvExpand}
- issue: https://github.com/elastic/elasticsearch/issues/137679
- class: org.elasticsearch.xpack.esql.optimizer.rules.physical.local.SubstituteRoundToTests
method: testSubqueryWithCountStarAndDateTrunc {default}
issue: https://github.com/elastic/elasticsearch/issues/138601
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/local/PushExpressionsToFieldLoad.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/local/PushExpressionsToFieldLoad.java
index e700b77cd7163..a6edbd87b3594 100644
--- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/local/PushExpressionsToFieldLoad.java
+++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/local/PushExpressionsToFieldLoad.java
@@ -23,6 +23,7 @@
import org.elasticsearch.xpack.esql.plan.logical.Eval;
import org.elasticsearch.xpack.esql.plan.logical.Filter;
import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
+import org.elasticsearch.xpack.esql.plan.logical.Project;
import org.elasticsearch.xpack.esql.plan.logical.local.EsqlProject;
import org.elasticsearch.xpack.esql.rule.ParameterizedRule;
@@ -34,9 +35,40 @@
import static org.elasticsearch.xpack.esql.core.expression.Attribute.rawTemporaryName;
/**
- * Replaces vector similarity functions with a field attribute that applies
- * the similarity function during value loading, when one side of the function is a literal.
- * It also adds the new field function attribute to the EsRelation output, and adds a projection after it to remove it from the output.
+ * Replaces {@link Expression}s that can be pushed to field loading with a field attribute
+ * that calculates the expression during value loading. See {@link BlockLoaderExpression}
+ * for more about how these loads are implemented and why we do this.
+ *
+ * This rule runs in one downward (aka output-to-read) pass, making four sorts
+ * of transformations:
+ *
+ *
*/
public class PushExpressionsToFieldLoad extends ParameterizedRule {
@@ -56,7 +88,7 @@ private class Rule {
* The primary indices, lazily initialized.
*/
private List primaries;
- private boolean planWasTransformed = false;
+ private boolean addedNewAttribute = false;
private Rule(LocalLogicalOptimizerContext context, LogicalPlan plan) {
this.context = context;
@@ -64,37 +96,38 @@ private Rule(LocalLogicalOptimizerContext context, LogicalPlan plan) {
}
private LogicalPlan doRule(LogicalPlan plan) {
- planWasTransformed = false;
+ addedNewAttribute = false;
if (plan instanceof Eval || plan instanceof Filter || plan instanceof Aggregate) {
- LogicalPlan transformedPlan = plan.transformExpressionsOnly(Expression.class, e -> {
- if (e instanceof BlockLoaderExpression ble) {
- return transformExpression(e, ble);
- }
- return e;
- });
+ return transformPotentialInvocation(plan);
+ }
+ if (addedAttrs.isEmpty()) {
+ return plan;
+ }
+ if (plan instanceof Project project) {
+ return transformProject(project);
+ }
+ if (plan instanceof EsRelation rel) {
+ return transformRelation(rel);
+ }
+ return plan;
+ }
- // TODO rebuild everything one time rather than after each find.
- if (planWasTransformed == false) {
- return plan;
+ private LogicalPlan transformPotentialInvocation(LogicalPlan plan) {
+ LogicalPlan transformedPlan = plan.transformExpressionsOnly(Expression.class, e -> {
+ if (e instanceof BlockLoaderExpression ble) {
+ return transformExpression(e, ble);
}
-
- List previousAttrs = transformedPlan.output();
- // Transforms EsRelation to extract the new attributes
- List addedAttrsList = addedAttrs.values().stream().toList();
- transformedPlan = transformedPlan.transformDown(EsRelation.class, esRelation -> {
- AttributeSet updatedOutput = esRelation.outputSet().combine(AttributeSet.of(addedAttrsList));
- return esRelation.withAttributes(updatedOutput.stream().toList());
- });
- // Transforms Projects so the new attribute is not discarded
- transformedPlan = transformedPlan.transformDown(EsqlProject.class, esProject -> {
- List projections = new ArrayList<>(esProject.projections());
- projections.addAll(addedAttrsList);
- return esProject.withProjections(projections);
- });
-
- return new EsqlProject(Source.EMPTY, transformedPlan, previousAttrs);
+ return e;
+ });
+ if (addedNewAttribute == false) {
+ /*
+ * Either didn't see anything pushable or everything pushable already
+ * has a pushed attribute.
+ */
+ return plan;
}
- return plan;
+ // Found a new pushable attribute, discard it *after* use so we don't modify the output.
+ return new EsqlProject(Source.EMPTY, transformedPlan, transformedPlan.output());
}
private Expression transformExpression(Expression e, BlockLoaderExpression ble) {
@@ -109,10 +142,23 @@ private Expression transformExpression(Expression e, BlockLoaderExpression ble)
if (context.searchStats().supportsLoaderConfig(fuse.field().fieldName(), fuse.config(), preference) == false) {
return e;
}
- planWasTransformed = true;
+ addedNewAttribute = true;
return replaceFieldsForFieldTransformations(e, fuse);
}
+ private LogicalPlan transformProject(Project project) {
+ // Preserve any pushed attributes so we can use them later
+ List projections = new ArrayList<>(project.projections());
+ projections.addAll(addedAttrs.values());
+ return project.withProjections(projections);
+ }
+
+ private LogicalPlan transformRelation(EsRelation rel) {
+ // Add the pushed attribute
+ AttributeSet updatedOutput = rel.outputSet().combine(AttributeSet.of(addedAttrs.values()));
+ return rel.withAttributes(updatedOutput.stream().toList());
+ }
+
private Expression replaceFieldsForFieldTransformations(Expression e, BlockLoaderExpression.PushedBlockLoaderExpression fuse) {
// Change the expression to a reference of the pushed down function on the field
FunctionEsField functionEsField = new FunctionEsField(fuse.field().field(), e.dataType(), fuse.config());