Ensure filters are pushed down before completion when it is possible.

Author: afoucret

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PushDownAndCombineFilters.java

@@ -24,6 +24,7 @@
 import org.elasticsearch.xpack.esql.plan.logical.Project;
 import org.elasticsearch.xpack.esql.plan.logical.RegexExtract;
 import org.elasticsearch.xpack.esql.plan.logical.UnaryPlan;
+import org.elasticsearch.xpack.esql.plan.logical.inference.Completion;
 import org.elasticsearch.xpack.esql.plan.logical.join.Join;
 import org.elasticsearch.xpack.esql.plan.logical.join.JoinTypes;
 
@@ -70,6 +71,10 @@ protected LogicalPlan rule(Filter filter) {
             // Push down filters that do not rely on attributes created by RegexExtract
             var attributes = AttributeSet.of(Expressions.asAttributes(re.extractedFields()));
             plan = maybePushDownPastUnary(filter, re, attributes::contains, NO_OP);
+        } else if (child instanceof Completion completion) {
+            // Push down filters that do not rely on attributes created by Cpmpletion
+            var attributes = AttributeSet.of(completion.generatedAttributes());
+            plan = maybePushDownPastUnary(filter, completion, attributes::contains, NO_OP);
         } else if (child instanceof Enrich enrich) {
             // Push down filters that do not rely on attributes created by Enrich
             var attributes = AttributeSet.of(Expressions.asAttributes(enrich.enrichFields()));

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PushDownAndCombineFiltersTests.java

@@ -8,12 +8,15 @@
 package org.elasticsearch.xpack.esql.optimizer.rules.logical;
 
 import org.elasticsearch.index.IndexMode;
+import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.xpack.esql.core.expression.Alias;
 import org.elasticsearch.xpack.esql.core.expression.Attribute;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.FieldAttribute;
+import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.expression.function.aggregate.Count;
+import org.elasticsearch.xpack.esql.expression.function.fulltext.Match;
 import org.elasticsearch.xpack.esql.expression.function.scalar.math.Pow;
 import org.elasticsearch.xpack.esql.expression.function.scalar.string.RLike;
 import org.elasticsearch.xpack.esql.expression.function.scalar.string.WildcardLike;
@@ -28,6 +31,7 @@
 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.inference.Completion;
 import org.elasticsearch.xpack.esql.plan.logical.local.EsqlProject;
 
 import java.util.ArrayList;
@@ -45,9 +49,12 @@
 import static org.elasticsearch.xpack.esql.EsqlTestUtils.greaterThanOf;
 import static org.elasticsearch.xpack.esql.EsqlTestUtils.greaterThanOrEqualOf;
 import static org.elasticsearch.xpack.esql.EsqlTestUtils.lessThanOf;
+import static org.elasticsearch.xpack.esql.EsqlTestUtils.randomLiteral;
+import static org.elasticsearch.xpack.esql.EsqlTestUtils.referenceAttribute;
 import static org.elasticsearch.xpack.esql.EsqlTestUtils.rlike;
 import static org.elasticsearch.xpack.esql.EsqlTestUtils.wildcardLike;
 import static org.elasticsearch.xpack.esql.core.tree.Source.EMPTY;
+import static org.mockito.Mockito.mock;
 
 public class PushDownAndCombineFiltersTests extends ESTestCase {
 
@@ -238,6 +245,39 @@ public void testSelectivelyPushDownFilterPastFunctionAgg() {
         assertEquals(expected, new PushDownAndCombineFilters().apply(fb));
     }
 
+    public void testPushDownFilterPastCompletion() {
+        FieldAttribute a = getFieldAttribute("a");
+        FieldAttribute b = getFieldAttribute("b");
+        EsRelation relation = relation(List.of(a, b));
+
+        GreaterThan conditionA =  greaterThanOf(getFieldAttribute("a"), ONE);
+        Filter filterA = new Filter(EMPTY, relation, conditionA);
+
+        Completion completion = completion(filterA);
+
+        LessThan conditionB = lessThanOf(getFieldAttribute("b"), TWO);
+        Match conditionCompletion =  new Match(EMPTY, completion.targetField(), randomLiteral(DataType.TEXT), mock(Expression.class), mock(QueryBuilder.class));
+        Filter filterB = new Filter(EMPTY, completion, new And(EMPTY, conditionB, conditionCompletion));
+
+        LogicalPlan expectedOptimizedPlan = new Filter(EMPTY,
+            new Completion(
+                EMPTY,
+                new Filter(EMPTY, relation, new And(EMPTY, conditionA, conditionB)),
+                completion.inferenceId(),
+                completion.prompt(),
+                completion.targetField()
+            ),
+            conditionCompletion
+        );
+
+        assertEquals(expectedOptimizedPlan, new PushDownAndCombineFilters().apply(filterB));
+    }
+
+    private static Completion completion(LogicalPlan child) {
+        return new Completion(EMPTY, child, randomLiteral(DataType.TEXT), randomLiteral(DataType.TEXT), referenceAttribute(randomIdentifier(), DataType.TEXT));
+    }
+
+
     private static EsRelation relation() {
         return relation(List.of());
     }