adding order by back

walterddr · walterddr · commit 137aa1c56fe9 · 2023-04-04T17:04:49.000-07:00
diff --git a/pinot-common/src/main/java/org/apache/pinot/sql/parsers/CalciteSqlParser.java b/pinot-common/src/main/java/org/apache/pinot/sql/parsers/CalciteSqlParser.java
@@ -162,6 +162,10 @@ public static List<String> extractTableNamesFromNode(SqlNode sqlNode) {
         tableNames.addAll(((SqlIdentifier) fromNode).names);
         tableNames.addAll(extractTableNamesFromNode(((SqlSelect) sqlNode).getWhere()));
       }
+    } else if (sqlNode instanceof SqlOrderBy) {
+      for (SqlNode node : ((SqlOrderBy) sqlNode).getOperandList()) {
+        tableNames.addAll(extractTableNamesFromNode(node));
+      }
     } else if (sqlNode instanceof SqlBasicCall) {
       if (((SqlBasicCall) sqlNode).getOperator() instanceof SqlAsOperator) {
         SqlNode firstOperand = ((SqlBasicCall) sqlNode).getOperandList().get(0);
diff --git a/pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlCompilerTest.java b/pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlCompilerTest.java
@@ -3084,6 +3084,19 @@ public void testExtractTableNamesFromNode() {
     Assert.assertEquals(tableNames.get(2), "tbl3");
     Assert.assertEquals(tableNames.get(3), "tbl4");
 
+    // query with aliases, JOIN, IN/NOT-IN, group-by
+    query = "with tmp as (select col1, count(*) from tbl1 where condition1 = filter1 group by col1 order by col2), "
+        + "tmp2 as (select A.col1, B.col2 from tbl2 as A JOIN tbl3 AS B on A.key = B.key) "
+        + "select sum(col1) from tmp where col1 in (select col1 from tmp2) and col1 not in (select col1 from tbl4) "
+        + "order by A.col1";
+    sqlNodeAndOptions = RequestUtils.parseQuery(query);
+    tableNames = CalciteSqlParser.extractTableNamesFromNode(sqlNodeAndOptions.getSqlNode());
+    Assert.assertEquals(tableNames.size(), 4);
+    Assert.assertEquals(tableNames.get(0), "tbl1");
+    Assert.assertEquals(tableNames.get(1), "tbl2");
+    Assert.assertEquals(tableNames.get(2), "tbl3");
+    Assert.assertEquals(tableNames.get(3), "tbl4");
+
     // query with aliases, JOIN, IN/NOT-IN, group-by and explain
     query = "explain plan for with tmp as (select col1, count(*) from tbl1 where condition1 = filter1 group by col1), "
         + "tmp2 as (select A.col1, B.col2 from tbl2 as A JOIN tbl3 AS B on A.key = B.key) "
