diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/PartialPlanMatcher.java b/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/PartialPlanMatcher.java new file mode 100644 index 0000000000000..1820007721d34 --- /dev/null +++ b/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/PartialPlanMatcher.java @@ -0,0 +1,95 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.esql; + +import org.elasticsearch.xpack.esql.core.expression.Expression; +import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan; +import org.hamcrest.Description; +import org.hamcrest.TypeSafeMatcher; + +import java.util.List; + +/** + * A matcher that compares a plan with a partial plan. + */ +public class PartialPlanMatcher extends TypeSafeMatcher { + private final LogicalPlan expected; + + public static PartialPlanMatcher matchesPartialPlan(PlanBuilder expected) { + return new PartialPlanMatcher(expected.build()); + } + + public static PartialPlanMatcher matchesPartialPlan(LogicalPlan expected) { + return new PartialPlanMatcher(expected); + } + + PartialPlanMatcher(LogicalPlan expected) { + this.expected = expected; + } + + @Override + protected boolean matchesSafely(LogicalPlan actual) { + return matches(expected, actual) == null; + } + + @Override + public void describeTo(Description description) { + description.appendText(expected.toString()); + } + + @Override + protected void describeMismatchSafely(LogicalPlan item, Description mismatchDescription) { + super.describeMismatchSafely(item, mismatchDescription); + mismatchDescription.appendText(System.lineSeparator()).appendText(matches(expected, item)); + } + + static String matches(LogicalPlan expected, LogicalPlan actual) { + if (expected.getClass() != actual.getClass()) { + return "Mismatch in plan types: Expected [" + expected.getClass() + "], found [" + actual.getClass() + "]"; + } + + List expectedProperties = expected.nodeProperties(); + List actualProperties = actual.nodeProperties(); + + for (int i = 0; i < expectedProperties.size(); i++) { + Object expectedProperty = expectedProperties.get(i); + Object actualProperty = actualProperties.get(i); + + // Only check equality if expected is not null. This way we can be lenient by simply + // leaving out properties. + if (expectedProperty == null || expectedProperty instanceof PlanBuilder.Ignore) { + continue; + } + + if (actualProperty == null) { + return "Expected [" + expectedProperty.getClass() + "], found null."; + } + + if (expectedProperty instanceof LogicalPlan l) { + if (actualProperty instanceof LogicalPlan rightPlan) { + String subMatch = matches(l, rightPlan); + if (subMatch != null) { + return subMatch; + } + } else {} + } else if (expectedProperty instanceof Expression e) { + if (actualProperty instanceof Expression rightExpr) { + if (e.semanticEquals(rightExpr) == false) { + return "Mismatch in expressions: Expected [" + e + "], found [" + rightExpr + "]"; + } + } else { + return "Mismatch in types: Expected [" + expectedProperty.getClass() + "], found [" + actualProperty.getClass() + "]"; + } + } else if (expectedProperty.equals(actualProperty) == false) { + return "Mismatch in properties: Expected [" + expectedProperty + "], found [" + actualProperty + "]"; + } + } + + return null; + } +} diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/PlanBuilder.java b/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/PlanBuilder.java new file mode 100644 index 0000000000000..c222194dc1112 --- /dev/null +++ b/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/PlanBuilder.java @@ -0,0 +1,255 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.esql; + +import org.elasticsearch.index.IndexMode; +import org.elasticsearch.xpack.esql.core.expression.Attribute; +import org.elasticsearch.xpack.esql.core.expression.Expression; +import org.elasticsearch.xpack.esql.core.expression.Literal; +import org.elasticsearch.xpack.esql.core.tree.Source; +import org.elasticsearch.xpack.esql.core.type.DataType; +import org.elasticsearch.xpack.esql.core.util.Holder; +import org.elasticsearch.xpack.esql.index.EsIndex; +import org.elasticsearch.xpack.esql.plan.logical.EsRelation; +import org.elasticsearch.xpack.esql.plan.logical.Filter; +import org.elasticsearch.xpack.esql.plan.logical.Limit; +import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan; +import org.elasticsearch.xpack.esql.plan.logical.OrderBy; +import org.elasticsearch.xpack.esql.plan.logical.TopN; +import org.elasticsearch.xpack.esql.plan.logical.local.LocalRelation; +import org.elasticsearch.xpack.esql.plan.logical.local.LocalSupplier; + +import java.lang.reflect.ParameterizedType; +import java.util.ArrayList; +import java.util.Collection; +import java.util.Iterator; +import java.util.List; +import java.util.ListIterator; +import java.util.Map; + +public class PlanBuilder { + private final LogicalPlan currentPlan; + private final List currentExpressions = new ArrayList<>(); + private final List> currentExpressionLists = new ArrayList<>(); + + private PlanBuilder(LogicalPlan plan) { + this.currentPlan = plan; + } + + public LogicalPlan build() { + return currentPlan; + } + + public PlanBuilder withChild(PlanBuilder child) { + return withChild(child.currentPlan); + } + + public PlanBuilder withChild(LogicalPlan child) { + List children = currentPlan.children(); + if (children.size() == 1 && children.get(0) instanceof MockChild) { + children = new ArrayList<>(); + } + + children.add(child); + return new PlanBuilder(currentPlan.replaceChildren(children)); + } + + public PlanBuilder withExpression(Expression expression) { + currentExpressions.add(expression); + + Holder i = new Holder<>(0); + return new PlanBuilder(currentPlan.transformPropertiesOnly(Expression.class, expr -> { + int idx = i.get(); + if (idx < currentExpressions.size() == false) { + return expr; + } + i.set(idx + 1); + return currentExpressions.get(idx); + })); + } + + public PlanBuilder withExpressionList(List expressions) { + currentExpressionLists.add(expressions); + + Holder i = new Holder<>(0); + return new PlanBuilder(currentPlan.transformPropertiesOnly(List.class, prop -> { + if (prop instanceof ParameterizedType pt + && pt.getRawType() == List.class + && pt.getActualTypeArguments()[0].getClass().isAssignableFrom(Expression.class)) { + int idx = i.get(); + if (idx < currentExpressionLists.size() == false) { + return prop; + } + i.set(idx + 1); + return currentExpressionLists.get(idx); + } + return prop; + })); + } + + public static PlanBuilder limit() { + return new PlanBuilder(new Limit(null, new MockExpression(), new MockChild())); + } + + public static PlanBuilder relation() { + return new PlanBuilder(new EsRelation((Source) null, (EsIndex) null, (IndexMode) null)); + } + + public static PlanBuilder localRelation() { + return new PlanBuilder(new LocalRelation((Source) null, new MockList<>(), null)); + } + + public static PlanBuilder localRelation(List output, LocalSupplier supplier) { + return new PlanBuilder(new LocalRelation((Source) null, output, supplier)); + } + + public static PlanBuilder filter() { + return new PlanBuilder(new Filter((Source) null, new MockChild(), new MockExpression())); + } + + public static PlanBuilder orderBy() { + return new PlanBuilder(new OrderBy((Source) null, new MockChild(), new MockList<>())); + } + + public static PlanBuilder topN() { + return new PlanBuilder(new TopN((Source) null, new MockChild(), new MockList<>(), new MockExpression())); + } + + interface Ignore {} + + private static class MockChild extends EsRelation implements Ignore { + MockChild() { + super(Source.EMPTY, new EsIndex("mock_index", Map.of()), IndexMode.STANDARD); + } + } + + private static class MockExpression extends Literal implements Ignore { + MockExpression() { + super(Source.EMPTY, 0, DataType.NULL); + } + } + + private static class MockList implements List, Ignore { + MockList() {} + + @Override + public int size() { + return 0; + } + + @Override + public boolean isEmpty() { + return false; + } + + @Override + public boolean contains(Object o) { + return false; + } + + @Override + public Iterator iterator() { + return null; + } + + @Override + public Object[] toArray() { + return new Object[0]; + } + + @Override + public T[] toArray(T[] a) { + return null; + } + + @Override + public boolean add(E e) { + return false; + } + + @Override + public boolean remove(Object o) { + return false; + } + + @Override + public boolean containsAll(Collection c) { + return false; + } + + @Override + public boolean addAll(Collection c) { + return false; + } + + @Override + public boolean addAll(int index, Collection c) { + return false; + } + + @Override + public boolean removeAll(Collection c) { + return false; + } + + @Override + public boolean retainAll(Collection c) { + return false; + } + + @Override + public void clear() { + + } + + @Override + public E get(int index) { + return null; + } + + @Override + public E set(int index, E element) { + return null; + } + + @Override + public void add(int index, E element) { + + } + + @Override + public E remove(int index) { + return null; + } + + @Override + public int indexOf(Object o) { + return 0; + } + + @Override + public int lastIndexOf(Object o) { + return 0; + } + + @Override + public ListIterator listIterator() { + return null; + } + + @Override + public ListIterator listIterator(int index) { + return null; + } + + @Override + public List subList(int fromIndex, int toIndex) { + return List.of(); + } + } +} diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/PlanContainsMatcher.java b/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/PlanContainsMatcher.java new file mode 100644 index 0000000000000..5b04b026460d2 --- /dev/null +++ b/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/PlanContainsMatcher.java @@ -0,0 +1,58 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.esql; + +import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan; +import org.hamcrest.Description; +import org.hamcrest.TypeSafeMatcher; + +import java.util.List; + +/** + * This validates if the actual plan contains the expected plan anywhere within the tree + */ +public class PlanContainsMatcher extends TypeSafeMatcher { + private final LogicalPlan expected; + + public static PlanContainsMatcher containsPlan(PlanBuilder expected) { + return new PlanContainsMatcher(expected.build()); + } + + public static PlanContainsMatcher containsPlan(LogicalPlan expected) { + return new PlanContainsMatcher(expected); + } + + PlanContainsMatcher(LogicalPlan expected) { + this.expected = expected; + } + + @Override + protected boolean matchesSafely(LogicalPlan actual) { + return matches(expected, actual); + } + + @Override + public void describeTo(Description description) { + + } + + private static List findMatchingChildren(LogicalPlan expected, LogicalPlan actual) { + return actual.collect(p -> p.getClass() == expected.getClass()); + } + + private static boolean matches(LogicalPlan expected, LogicalPlan initialActual) { + List candidates = findMatchingChildren(expected, initialActual); + if (candidates.isEmpty()) { + return false; + } + + // TODO: Deal with multiple matching children case + LogicalPlan actual = candidates.get(0); + return PartialPlanMatcher.matches(expected, actual) == null; + } +} diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizerTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizerTests.java index dd48cfac29ea0..a4edf06bdae73 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizerTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizerTests.java @@ -19,6 +19,7 @@ import org.elasticsearch.dissect.DissectParser; import org.elasticsearch.index.IndexMode; import org.elasticsearch.xpack.esql.EsqlTestUtils; +import org.elasticsearch.xpack.esql.PlanBuilder; import org.elasticsearch.xpack.esql.VerificationException; import org.elasticsearch.xpack.esql.action.EsqlCapabilities; import org.elasticsearch.xpack.esql.analysis.Analyzer; @@ -167,6 +168,7 @@ import static org.elasticsearch.xpack.esql.EsqlTestUtils.referenceAttribute; import static org.elasticsearch.xpack.esql.EsqlTestUtils.singleValue; import static org.elasticsearch.xpack.esql.EsqlTestUtils.withDefaultLimitWarning; +import static org.elasticsearch.xpack.esql.PartialPlanMatcher.matchesPartialPlan; import static org.elasticsearch.xpack.esql.analysis.Analyzer.NO_FIELDS; import static org.elasticsearch.xpack.esql.analysis.AnalyzerTestUtils.analyze; import static org.elasticsearch.xpack.esql.core.expression.Literal.NULL; @@ -190,7 +192,6 @@ import static org.hamcrest.Matchers.containsInAnyOrder; import static org.hamcrest.Matchers.containsString; import static org.hamcrest.Matchers.empty; -import static org.hamcrest.Matchers.emptyArray; import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.everyItem; import static org.hamcrest.Matchers.hasItem; @@ -212,9 +213,7 @@ public void testEmptyProjections() { | drop salary """); - var relation = as(plan, LocalRelation.class); - assertThat(relation.output(), is(empty())); - assertThat(relation.supplier().get(), emptyArray()); + assertThat(plan, matchesPartialPlan(PlanBuilder.localRelation(List.of(), EmptyLocalSupplier.EMPTY))); } public void testEmptyProjectionInStat() { @@ -224,9 +223,7 @@ public void testEmptyProjectionInStat() { | drop c """); - var relation = as(plan, LocalRelation.class); - assertThat(relation.output(), is(empty())); - assertThat(relation.supplier().get(), emptyArray()); + assertThat(plan, matchesPartialPlan(PlanBuilder.localRelation(List.of(), EmptyLocalSupplier.EMPTY))); } /** @@ -247,7 +244,6 @@ public void testEmptyProjectInStatWithEval() { | eval x = 1, c2 = c*2 | drop c, c2 """); - var project = as(plan, Project.class); var eval = as(project.child(), Eval.class); var limit = as(eval.child(), Limit.class);