From a7cfde39238fbab1eb3ef4f90ca4f4a7b5def9dc Mon Sep 17 00:00:00 2001 From: Nik Everett Date: Tue, 25 Nov 2025 14:50:31 -0500 Subject: [PATCH 1/2] WIP --- .../local/PushExpressionsToFieldLoad.java | 109 +++++++++++++----- 1 file changed, 77 insertions(+), 32 deletions(-) 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..d019be3730047 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 @@ -34,9 +34,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 +87,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 +95,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 EsqlProject 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 +141,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(EsqlProject 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()); From 98ba1bb62c7fea1fcbb34e18ce87e63155f31372 Mon Sep 17 00:00:00 2001 From: Nik Everett Date: Tue, 25 Nov 2025 19:45:10 -0500 Subject: [PATCH 2/2] fix project --- muted-tests.yml | 3 --- .../rules/logical/local/PushExpressionsToFieldLoad.java | 5 +++-- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/muted-tests.yml b/muted-tests.yml index 16d503633bc78..235b1f60de5cd 100644 --- a/muted-tests.yml +++ b/muted-tests.yml @@ -438,9 +438,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 d019be3730047..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; @@ -102,7 +103,7 @@ private LogicalPlan doRule(LogicalPlan plan) { if (addedAttrs.isEmpty()) { return plan; } - if (plan instanceof EsqlProject project) { + if (plan instanceof Project project) { return transformProject(project); } if (plan instanceof EsRelation rel) { @@ -145,7 +146,7 @@ private Expression transformExpression(Expression e, BlockLoaderExpression ble) return replaceFieldsForFieldTransformations(e, fuse); } - private LogicalPlan transformProject(EsqlProject project) { + 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());