66package org .opensearch .sql .api ;
77
88import org .antlr .v4 .runtime .tree .ParseTree ;
9+ import org .apache .calcite .plan .hep .HepPlanner ;
10+ import org .apache .calcite .plan .hep .HepProgramBuilder ;
911import org .apache .calcite .rel .RelCollation ;
1012import org .apache .calcite .rel .RelCollations ;
1113import org .apache .calcite .rel .RelNode ;
14+ import org .apache .calcite .rel .RelRoot ;
1215import org .apache .calcite .rel .core .Sort ;
1316import org .apache .calcite .rel .logical .LogicalSort ;
17+ import org .apache .calcite .rel .rules .CoreRules ;
18+ import org .apache .calcite .sql .SqlNode ;
19+ import org .apache .calcite .tools .Frameworks ;
20+ import org .apache .calcite .tools .Planner ;
1421import org .opensearch .sql .ast .statement .Query ;
1522import org .opensearch .sql .ast .statement .Statement ;
1623import org .opensearch .sql .ast .tree .UnresolvedPlan ;
2633 * {@code UnifiedQueryPlanner} provides a high-level API for parsing and analyzing queries using the
2734 * Calcite-based query engine. It serves as the primary integration point for external consumers
2835 * such as Spark or command-line tools, abstracting away Calcite internals.
36+ *
37+ * <p>Two planning paths are supported:
38+ *
39+ * <ul>
40+ * <li><b>PPL</b>: ANTLR parser → AST → CalciteRelNodeVisitor → RelNode
41+ * <li><b>SQL</b>: Calcite native SqlParser → SqlValidator → SqlToRelConverter → RelNode
42+ * </ul>
2943 */
3044public class UnifiedQueryPlanner {
31- /** The parser instance responsible for converting query text into a parse tree. */
45+ /** The parser instance responsible for converting query text into a parse tree (PPL only) . */
3246 private final Parser parser ;
3347
3448 /** Unified query context containing CalcitePlanContext with all configuration. */
3549 private final UnifiedQueryContext context ;
3650
37- /** AST-to-RelNode visitor that builds logical plans from the parsed AST. */
51+ /** AST-to-RelNode visitor that builds logical plans from the parsed AST (PPL only) . */
3852 private final CalciteRelNodeVisitor relNodeVisitor =
3953 new CalciteRelNodeVisitor (new EmptyDataSourceService ());
4054
@@ -44,19 +58,23 @@ public class UnifiedQueryPlanner {
4458 * @param context the unified query context containing CalcitePlanContext
4559 */
4660 public UnifiedQueryPlanner (UnifiedQueryContext context ) {
47- this .parser = buildQueryParser (context .getPlanContext ().queryType );
4861 this .context = context ;
62+ this .parser =
63+ context .getPlanContext ().queryType == QueryType .PPL ? new PPLSyntaxParser () : null ;
4964 }
5065
5166 /**
5267 * Parses and analyzes a query string into a Calcite logical plan (RelNode). TODO: Generate
5368 * optimal physical plan to fully unify query execution and leverage Calcite's optimizer.
5469 *
55- * @param query the raw query string in PPL or other supported syntax
70+ * @param query the raw query string in PPL or SQL syntax
5671 * @return a logical plan representing the query
5772 */
5873 public RelNode plan (String query ) {
5974 try {
75+ if (context .getPlanContext ().queryType == QueryType .SQL ) {
76+ return planWithCalcite (query );
77+ }
6078 return preserveCollation (analyze (parse (query )));
6179 } catch (SyntaxCheckException e ) {
6280 // Re-throw syntax error without wrapping
@@ -66,11 +84,29 @@ public RelNode plan(String query) {
6684 }
6785 }
6886
69- private Parser buildQueryParser (QueryType queryType ) {
70- if (queryType == QueryType .PPL ) {
71- return new PPLSyntaxParser ();
87+ /**
88+ * Plans a SQL query using Calcite's native parser and validator. The pipeline is: SqlParser →
89+ * SqlValidator → SqlToRelConverter → HepPlanner optimization → RelNode.
90+ */
91+ private RelNode planWithCalcite (String query ) throws Exception {
92+ Planner planner = Frameworks .getPlanner (context .getPlanContext ().config );
93+ try {
94+ SqlNode parsed = planner .parse (query );
95+ SqlNode validated = planner .validate (parsed );
96+ RelRoot relRoot = planner .rel (validated );
97+ return optimize (relRoot .rel );
98+ } finally {
99+ planner .close ();
72100 }
73- throw new IllegalArgumentException ("Unsupported query type: " + queryType );
101+ }
102+
103+ /** Applies heuristic optimization rules to the logical plan. */
104+ private RelNode optimize (RelNode rel ) {
105+ HepPlanner hepPlanner =
106+ new HepPlanner (
107+ new HepProgramBuilder ().addRuleInstance (CoreRules .AGGREGATE_CASE_TO_FILTER ).build ());
108+ hepPlanner .setRoot (rel );
109+ return hepPlanner .findBestExp ();
74110 }
75111
76112 private UnresolvedPlan parse (String query ) {
0 commit comments