From 18536d9012a46604e3a2a1ee520c93ed61918a78 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Mon, 23 Oct 2023 16:54:28 +0200 Subject: [PATCH] DrillWindowQueryTest: use proper way to decide if the query is ordered (#15118) --- .../druid/sql/calcite/planner/DruidPlanner.java | 1 + .../druid/sql/calcite/planner/PlannerHook.java | 6 ++++++ .../druid/sql/calcite/BaseCalciteQueryTest.java | 1 + .../druid/sql/calcite/DrillWindowQueryTest.java | 15 +++++++++------ .../druid/sql/calcite/QueryTestRunner.java | 16 +++++++++++----- .../sql/calcite/planner/PlannerCaptureHook.java | 15 +++++++++++++++ 6 files changed, 43 insertions(+), 11 deletions(-) rename sql/src/{main => test}/java/org/apache/druid/sql/calcite/planner/PlannerCaptureHook.java (88%) diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/planner/DruidPlanner.java b/sql/src/main/java/org/apache/druid/sql/calcite/planner/DruidPlanner.java index e47411361b2e8..4b697a0d5dfaf 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/planner/DruidPlanner.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/planner/DruidPlanner.java @@ -148,6 +148,7 @@ public void validate() catch (SqlParseException e1) { throw translateException(e1); } + hook.captureSqlNode(root); handler = createHandler(root); handler.validate(); plannerContext.setResourceActions(handler.resourceActions()); diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerHook.java b/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerHook.java index a65f59d4d1c7b..fd245b096d82e 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerHook.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerHook.java @@ -23,6 +23,8 @@ import org.apache.calcite.rel.RelRoot; import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.sql.SqlInsert; +import org.apache.calcite.sql.SqlNode; +import org.apache.druid.guice.annotations.UnstableApi; import org.apache.druid.sql.calcite.rel.DruidRel; /** @@ -32,9 +34,13 @@ * none at the points where tests want to verify, except for the opportunity to * capture the native query. */ +@UnstableApi public interface PlannerHook { void captureSql(String sql); + default void captureSqlNode(SqlNode node) + { + } void captureQueryRel(RelRoot rootQueryRel); void captureDruidRel(DruidRel druidRel); void captureBindableRel(BindableRel bindableRel); diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/BaseCalciteQueryTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/BaseCalciteQueryTest.java index 6b2d12c6cce0c..8384307fa750b 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/BaseCalciteQueryTest.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/BaseCalciteQueryTest.java @@ -1513,6 +1513,7 @@ public void verify(String sql, QueryResults queryResults) } catch (AssertionError e) { log.info("sql: %s", sql); + log.info(resultsToString("Expected", expectedResults)); log.info(resultsToString("Actual", queryResults.results)); throw e; } diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/DrillWindowQueryTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/DrillWindowQueryTest.java index 7a2b9b70f1a67..bf13b88e29176 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/DrillWindowQueryTest.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/DrillWindowQueryTest.java @@ -25,6 +25,8 @@ import com.google.common.collect.Iterators; import com.google.common.io.ByteStreams; import com.google.inject.Injector; +import org.apache.calcite.sql.SqlNode; +import org.apache.calcite.sql2rel.SqlToRelConverter; import org.apache.commons.io.FileUtils; import org.apache.druid.common.config.NullHandling; import org.apache.druid.data.input.InputRow; @@ -52,6 +54,7 @@ import org.apache.druid.sql.calcite.NotYetSupported.Modes; import org.apache.druid.sql.calcite.NotYetSupported.NotYetSupportedProcessor; import org.apache.druid.sql.calcite.QueryTestRunner.QueryResults; +import org.apache.druid.sql.calcite.planner.PlannerCaptureHook; import org.apache.druid.sql.calcite.planner.PlannerContext; import org.apache.druid.sql.calcite.util.SpecificSegmentsQuerySegmentWalker; import org.apache.druid.timeline.DataSegment; @@ -363,7 +366,8 @@ public void verify(String sql, QueryResults queryResults) List expectedResults = parseResults(currentRowSignature, expectedResultsText); try { Assert.assertEquals(StringUtils.format("result count: %s", sql), expectedResultsText.size(), results.size()); - if (!isOrdered(sql)) { + if (!isOrdered(queryResults)) { + // in case the resultset is not ordered; order via the same comparator before comparision results.sort(new ArrayRowCmp()); expectedResults.sort(new ArrayRowCmp()); } @@ -377,12 +381,10 @@ public void verify(String sql, QueryResults queryResults) } } - private boolean isOrdered(String sql) + private boolean isOrdered(QueryResults queryResults) { - // FIXME: SqlToRelConverter.isOrdered(null) would be better - sql = sql.toLowerCase(Locale.ENGLISH).replace('\n', ' '); - sql = sql.substring(sql.lastIndexOf(')')); - return sql.contains("order"); + SqlNode sqlNode = ((PlannerCaptureHook) queryResults.capture).getSqlNode(); + return SqlToRelConverter.isOrdered(sqlNode); } } @@ -478,6 +480,7 @@ public void windowQueryTest() .skipVectorize(true) .queryContext(ImmutableMap.of( PlannerContext.CTX_ENABLE_WINDOW_FNS, true, + PlannerCaptureHook.NEED_CAPTURE_HOOK, true, QueryContexts.ENABLE_DEBUG, true) ) .sql(testCase.getQueryString()) diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/QueryTestRunner.java b/sql/src/test/java/org/apache/druid/sql/calcite/QueryTestRunner.java index 0bd1f1305eb9a..1dd1df4eea8ce 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/QueryTestRunner.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/QueryTestRunner.java @@ -86,14 +86,14 @@ public interface QueryRunStepFactory */ public abstract static class QueryRunStep { - private final QueryTestBuilder builder; + protected final QueryTestBuilder builder; public QueryRunStep(final QueryTestBuilder builder) { this.builder = builder; } - public QueryTestBuilder builder() + public final QueryTestBuilder builder() { return builder; } @@ -207,12 +207,10 @@ public void run() public abstract static class BaseExecuteQuery extends QueryRunStep { protected final List results = new ArrayList<>(); - protected final boolean doCapture; public BaseExecuteQuery(QueryTestBuilder builder) { super(builder); - doCapture = builder.expectedLogicalPlan != null; } public List results() @@ -278,7 +276,7 @@ public QueryResults runQuery( ) { try { - final PlannerCaptureHook capture = doCapture ? new PlannerCaptureHook() : null; + final PlannerCaptureHook capture = getCaptureHook(); final DirectStatement stmt = sqlStatementFactory.directStatement(query); stmt.setHook(capture); final Sequence results = stmt.execute().getResults(); @@ -300,6 +298,14 @@ public QueryResults runQuery( } } + private PlannerCaptureHook getCaptureHook() + { + if (builder.getQueryContext().containsKey(PlannerCaptureHook.NEED_CAPTURE_HOOK) || builder.expectedLogicalPlan != null) { + return new PlannerCaptureHook(); + } + return null; + } + public static Pair> getResults( final SqlStatementFactory sqlStatementFactory, final SqlQueryPlus query diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerCaptureHook.java b/sql/src/test/java/org/apache/druid/sql/calcite/planner/PlannerCaptureHook.java similarity index 88% rename from sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerCaptureHook.java rename to sql/src/test/java/org/apache/druid/sql/calcite/planner/PlannerCaptureHook.java index bdf50a8a0c588..bf69e1bf162b6 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerCaptureHook.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/planner/PlannerCaptureHook.java @@ -23,12 +23,16 @@ import org.apache.calcite.rel.RelRoot; import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.sql.SqlInsert; +import org.apache.calcite.sql.SqlNode; import org.apache.druid.sql.calcite.rel.DruidRel; public class PlannerCaptureHook implements PlannerHook { + public static final String NEED_CAPTURE_HOOK = "need_capture_hook"; + private RelRoot relRoot; private SqlInsert insertNode; + private SqlNode sqlNode; @Override public void captureSql(String sql) @@ -36,6 +40,12 @@ public void captureSql(String sql) // Not used at present. Add a field to capture this if you need it. } + @Override + public void captureSqlNode(SqlNode node) + { + this.sqlNode = node; + } + @Override public void captureQueryRel(RelRoot rootQueryRel) { @@ -75,4 +85,9 @@ public SqlInsert insertNode() { return insertNode; } + + public SqlNode getSqlNode() + { + return sqlNode; + } }