DrillWindowQueryTest: use proper way to decide if the query is ordered (#15118)

sql/src/main/java/org/apache/druid/sql/calcite/planner/DruidPlanner.java

@@ -148,6 +148,7 @@ public class DruidPlanner implements Closeable
     catch (SqlParseException e1) {
       throw translateException(e1);
     }
+    hook.captureSqlNode(root);
     handler = createHandler(root);
     handler.validate();
     plannerContext.setResourceActions(handler.resourceActions());

sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerHook.java

@@ -23,6 +23,8 @@ import org.apache.calcite.interpreter.BindableRel;
 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 @@ import org.apache.druid.sql.calcite.rel.DruidRel;
  * 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);

sql/src/test/java/org/apache/druid/sql/calcite/BaseCalciteQueryTest.java

@@ -1513,6 +1513,7 @@ public class BaseCalciteQueryTest extends CalciteTestBase
       }
       catch (AssertionError e) {
         log.info("sql: %s", sql);
+        log.info(resultsToString("Expected", expectedResults));
         log.info(resultsToString("Actual", queryResults.results));
         throw e;
       }

sql/src/test/java/org/apache/druid/sql/calcite/DrillWindowQueryTest.java

@@ -25,6 +25,8 @@ import com.google.common.collect.ImmutableMap;
 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.segment.writeout.OnHeapMemorySegmentWriteOutMediumFactor
 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 class DrillWindowQueryTest extends BaseCalciteQueryTest
       List<Object[]> 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 class DrillWindowQueryTest extends BaseCalciteQueryTest
       }
     }
 
-    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 class DrillWindowQueryTest extends BaseCalciteQueryTest
           .skipVectorize(true)
           .queryContext(ImmutableMap.of(
               PlannerContext.CTX_ENABLE_WINDOW_FNS, true,
+              PlannerCaptureHook.NEED_CAPTURE_HOOK, true,
               QueryContexts.ENABLE_DEBUG, true)
               )
           .sql(testCase.getQueryString())

sql/src/test/java/org/apache/druid/sql/calcite/QueryTestRunner.java

@@ -86,14 +86,14 @@ public class QueryTestRunner
    */
   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 class QueryTestRunner
   public abstract static class BaseExecuteQuery extends QueryRunStep
   {
     protected final List<QueryResults> results = new ArrayList<>();
-    protected final boolean doCapture;
 
     public BaseExecuteQuery(QueryTestBuilder builder)
     {
       super(builder);
-      doCapture = builder.expectedLogicalPlan != null;
     }
 
     public List<QueryResults> results()
@@ -278,7 +276,7 @@ public class QueryTestRunner
     )
     {
       try {
-        final PlannerCaptureHook capture = doCapture ? new PlannerCaptureHook() : null;
+        final PlannerCaptureHook capture = getCaptureHook();
         final DirectStatement stmt = sqlStatementFactory.directStatement(query);
         stmt.setHook(capture);
         final Sequence<Object[]> results = stmt.execute().getResults();
@@ -300,6 +298,14 @@ public class QueryTestRunner
       }
     }
 
+    private PlannerCaptureHook getCaptureHook()
+    {
+      if (builder.getQueryContext().containsKey(PlannerCaptureHook.NEED_CAPTURE_HOOK) || builder.expectedLogicalPlan != null) {
+        return new PlannerCaptureHook();
+      }
+      return null;
+    }
+
     public static Pair<RowSignature, List<Object[]>> getResults(
         final SqlStatementFactory sqlStatementFactory,
         final SqlQueryPlus query

sql/src/test/java/org/apache/druid/sql/calcite/planner/PlannerCaptureHook.java

@@ -23,12 +23,16 @@ import org.apache.calcite.interpreter.BindableRel;
 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 class PlannerCaptureHook implements PlannerHook
     // 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 class PlannerCaptureHook implements PlannerHook
   {
     return insertNode;
   }
+
+  public SqlNode getSqlNode()
+  {
+    return sqlNode;
+  }
 }