[GitHub] [pinot] xiangfu0 commented on a diff in pull request #10873: Use pinot-query-planner utils to extract table names in pinot-controller

["xiangfu0 (via GitHub)" <gi...@apache.org>]

xiangfu0 commented on code in PR #10873:
URL: https://github.com/apache/pinot/pull/10873#discussion_r1235942012


##########
pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java:
##########
@@ -253,6 +254,104 @@ public void testQueryWithHint()
     }
   }
 
+  @Test
+  public void testGetTableNamesForQuery() {
+    // A simple filter query with one table
+    String query = "Select * from a where col1 = 'a'";
+    List<String> tableNames = _queryEnvironment.getTableNamesForQuery(query);
+    Assert.assertEquals(tableNames.size(), 1);
+    Assert.assertEquals(tableNames.get(0), "a");
+
+    // query with IN / NOT IN clause
+    query = "SELECT COUNT(*) FROM a WHERE col1 IN (SELECT col1 FROM b) "
+        + "and col1 NOT IN (SELECT col1 from c)";
+    tableNames = _queryEnvironment.getTableNamesForQuery(query);
+    Assert.assertEquals(tableNames.size(), 3);
+    Collections.sort(tableNames);
+    Assert.assertEquals(tableNames.get(0), "a");
+    Assert.assertEquals(tableNames.get(1), "b");
+    Assert.assertEquals(tableNames.get(2), "c");
+
+    // query with JOIN clause
+    query = "SELECT a.col1, b.col2 FROM a JOIN b ON a.col3 = b.col3 WHERE a.col1 = 'a'";
+    tableNames = _queryEnvironment.getTableNamesForQuery(query);
+    Assert.assertEquals(tableNames.size(), 2);
+    Collections.sort(tableNames);
+    Assert.assertEquals(tableNames.get(0), "a");
+    Assert.assertEquals(tableNames.get(1), "b");
+
+    // query with WHERE clause JOIN
+    query = "SELECT a.col1, b.col2 FROM a, b WHERE a.col3 = b.col3 AND a.col1 = 'a'";
+    tableNames = _queryEnvironment.getTableNamesForQuery(query);
+    Assert.assertEquals(tableNames.size(), 2);
+    Collections.sort(tableNames);
+    Assert.assertEquals(tableNames.get(0), "a");
+    Assert.assertEquals(tableNames.get(1), "b");
+
+    // query with JOIN clause and table alias
+    query = "SELECT A.col1, B.col2 FROM a AS A JOIN b AS B ON A.col3 = B.col3 WHERE A.col1 = 'a'";
+    tableNames = _queryEnvironment.getTableNamesForQuery(query);
+    Assert.assertEquals(tableNames.size(), 2);
+    Collections.sort(tableNames);
+    Assert.assertEquals(tableNames.get(0), "a");
+    Assert.assertEquals(tableNames.get(1), "b");
+
+    // query with UNION clause
+    query = "SELECT * FROM a UNION ALL SELECT * FROM b UNION ALL SELECT * FROM c";
+    tableNames = _queryEnvironment.getTableNamesForQuery(query);
+    Assert.assertEquals(tableNames.size(), 3);
+    Collections.sort(tableNames);
+    Assert.assertEquals(tableNames.get(0), "a");
+    Assert.assertEquals(tableNames.get(1), "b");
+    Assert.assertEquals(tableNames.get(2), "c");
+
+    // query with UNION clause and table alias
+    query = "SELECT * FROM (SELECT * FROM a) AS t1 UNION SELECT * FROM ( SELECT * FROM b) AS t2";
+    tableNames = _queryEnvironment.getTableNamesForQuery(query);
+    Assert.assertEquals(tableNames.size(), 2);
+    Collections.sort(tableNames);
+    Assert.assertEquals(tableNames.get(0), "a");
+    Assert.assertEquals(tableNames.get(1), "b");
+
+    // query with UNION clause and table alias using WITH clause
+    query = "WITH tmp1 AS (SELECT * FROM a), \n"
+        + "tmp2 AS (SELECT * FROM b) \n"
+        + "SELECT * FROM tmp1 UNION ALL SELECT * FROM tmp2";
+    tableNames = _queryEnvironment.getTableNamesForQuery(query);
+    Assert.assertEquals(tableNames.size(), 2);
+    Collections.sort(tableNames);
+    Assert.assertEquals(tableNames.get(0), "a");
+    Assert.assertEquals(tableNames.get(1), "b");
+
+    // query with aliases, JOIN, IN/NOT-IN, group-by
+    query = "with tmp as (select col1, sum(col3) as col3, count(*) from a where col1 = 'a' group by col1), "
+        + "tmp2 as (select A.col1, B.col3 from b as A JOIN c AS B on A.col1 = B.col1) "
+        + "select sum(col3) from tmp where col1 in (select col1 from tmp2) and col1 not in (select col1 from d)";
+    tableNames = _queryEnvironment.getTableNamesForQuery(query);
+    Assert.assertEquals(tableNames.size(), 4);
+    Assert.assertEquals(tableNames.get(0), "a");
+    Assert.assertEquals(tableNames.get(1), "b");
+    Assert.assertEquals(tableNames.get(2), "c");
+    Assert.assertEquals(tableNames.get(3), "d");
+
+    // query with aliases, JOIN, IN/NOT-IN, group-by and explain
+    query = "explain plan for with tmp as (select col1, sum(col3) as col3, count(*) from a where col1 = 'a' "
+        + "group by col1), tmp2 as (select A.col1, B.col3 from b as A JOIN c AS B on A.col1 = B.col1) "
+        + "select sum(col3) from tmp where col1 in (select col1 from tmp2) and col1 not in (select col1 from d)";
+    tableNames = _queryEnvironment.getTableNamesForQuery(query);
+    Assert.assertEquals(tableNames.size(), 4);
+    Assert.assertEquals(tableNames.get(0), "a");
+    Assert.assertEquals(tableNames.get(1), "b");
+    Assert.assertEquals(tableNames.get(2), "c");
+    Assert.assertEquals(tableNames.get(3), "d");
+
+    // test for self join queries
+    query = "SELECT a.col1 FROM a JOIN(SELECT col2 FROM a) as self ON a.col1=self.col2 ";
+    tableNames = _queryEnvironment.getTableNamesForQuery(query);
+    Assert.assertEquals(tableNames.size(), 1);
+    Assert.assertEquals(tableNames.get(0), "a");
+  }

Review Comment:
   Please also add lateral join example introduced from: https://github.com/apache/pinot/pull/10933 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org