fix(calcite): handle Limit, Alias projection, RelationSubquery, and Values dual-table in visitor

dai-chen · dai-chen · commit b1c59701b59b · 2026-05-18T11:41:56.000-07:00
Add visitLimit to support LIMIT/OFFSET clauses. Handle Alias nodes in
project list by referencing already-computed aggregate fields instead of
re-analyzing. Add visitRelationSubquery for derived tables in FROM
clause. Fix visitValues to treat a single empty row as a dual-table
(SELECT without FROM). Make bucketNullable lookup null-safe with
getOrDefault in visitAggregation.

Add integration tests covering LIMIT OFFSET, aggregate with alias,
GROUP BY without bucket nullable, SELECT with alias, derived table
subquery, and SELECT without FROM clause.

Signed-off-by: Chen Dai &lt;daichen@amazon.com&gt;
diff --git a/api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerSqlTest.java b/api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerSqlTest.java
@@ -259,4 +259,96 @@ SELECT department, count(*)
             """)
         .assertErrorMessage("Encountered");
   }
+
+  @Test
+  public void testSqlLimitOffset() {
+    givenQuery(
+            """
+            SELECT name
+            FROM catalog.employees
+            LIMIT 10 OFFSET 5\
+            """)
+        .assertPlan(
+            """
+            LogicalProject(name=[$1])
+              LogicalSort(offset=[5], fetch=[10])
+                LogicalTableScan(table=[[catalog, employees]])
+            """);
+  }
+
+  @Test
+  public void testSqlAggregateWithAlias() {
+    givenQuery(
+            """
+            SELECT department, COUNT(*) AS cnt
+            FROM catalog.employees
+            GROUP BY department\
+            """)
+        .assertPlan(
+            """
+            LogicalAggregate(group=[{0}], COUNT(*)=[COUNT()])
+              LogicalProject(department=[$3])
+                LogicalTableScan(table=[[catalog, employees]])
+            """);
+  }
+
+  @Test
+  public void testSqlGroupByWithoutBucketNullable() {
+    givenQuery(
+            """
+            SELECT age, COUNT(*) AS cnt
+            FROM catalog.employees
+            GROUP BY age\
+            """)
+        .assertPlan(
+            """
+            LogicalAggregate(group=[{0}], COUNT(*)=[COUNT()])
+              LogicalProject(age=[$2])
+                LogicalTableScan(table=[[catalog, employees]])
+            """);
+  }
+
+  @Test
+  public void testSqlSelectWithAlias() {
+    givenQuery(
+            """
+            SELECT age AS employee_age, name AS employee_name
+            FROM catalog.employees\
+            """)
+        .assertPlan(
+            """
+            LogicalProject(employee_age=[$2], employee_name=[$1])
+              LogicalTableScan(table=[[catalog, employees]])
+            """);
+  }
+
+  @Test
+  public void testSqlDerivedTableInFromClause() {
+    // SELECT ... FROM (SELECT ...) AS t — exercises visitRelationSubquery override.
+    givenQuery(
+            """
+            SELECT t.id
+            FROM (SELECT id, name FROM catalog.employees WHERE age > 30) AS t\
+            """)
+        .assertPlan(
+            """
+            LogicalProject(t.id=[$0])
+              LogicalFilter(condition=[>($2, 30)])
+                LogicalTableScan(table=[[catalog, employees]])
+            """);
+  }
+
+  @Test
+  public void testSqlSelectWithoutFromClause() {
+    // SELECT 1 — exercises visitValues dual-table case (single empty row).
+    givenQuery(
+            """
+            SELECT 1\
+            """)
+        .assertPlan(
+            """
+            LogicalSort(sort0=[$0], dir0=[ASC])
+              LogicalValues(tuples=[[]])
+            """);
+  }
 }
diff --git a/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java b/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java
@@ -132,6 +132,7 @@
 import org.opensearch.sql.ast.tree.Head;
 import org.opensearch.sql.ast.tree.Join;
 import org.opensearch.sql.ast.tree.Kmeans;
+import org.opensearch.sql.ast.tree.Limit;
 import org.opensearch.sql.ast.tree.Lookup;
 import org.opensearch.sql.ast.tree.Lookup.OutputStrategy;
 import org.opensearch.sql.ast.tree.ML;
@@ -146,6 +147,7 @@
 import org.opensearch.sql.ast.tree.RareTopN;
 import org.opensearch.sql.ast.tree.Regex;
 import org.opensearch.sql.ast.tree.Relation;
+import org.opensearch.sql.ast.tree.RelationSubquery;
 import org.opensearch.sql.ast.tree.Rename;
 import org.opensearch.sql.ast.tree.Replace;
 import org.opensearch.sql.ast.tree.ReplacePair;
@@ -541,6 +543,25 @@ private List<RexNode> expandProjectFields(
               .filter(addedFields::add)
               .forEach(field -> expandedFields.add(context.relBuilder.field(field)));
         }
+        case Alias alias -> {
+          String aliasName =
+              Strings.isNullOrEmpty(alias.getAlias()) ? alias.getName() : alias.getAlias();
+          // When an aggregate was already computed (its output name matches a current field),
+          // reference the existing field instead of re-analyzing (which would return null).
+          if (alias.getDelegated() instanceof AggregateFunction) {
+            String aggFieldName = alias.getName();
+            if (currentFields.contains(aliasName)) {
+              expandedFields.add(context.relBuilder.field(aliasName));
+            } else if (aggFieldName != null && currentFields.contains(aggFieldName)) {
+              expandedFields.add(
+                  context.relBuilder.alias(context.relBuilder.field(aggFieldName), aliasName));
+            } else {
+              expandedFields.add(rexVisitor.analyze(alias, context));
+            }
+          } else {
+            expandedFields.add(rexVisitor.analyze(alias, context));
+          }
+        }
         default ->
             throw new IllegalStateException(
                 "Unexpected expression type in project list: " + expr.getClass().getSimpleName());
@@ -763,6 +784,13 @@ public RelNode visitHead(Head node, CalcitePlanContext context) {
     return context.relBuilder.peek();
   }
 
+  @Override
+  public RelNode visitLimit(Limit node, CalcitePlanContext context) {
+    visitChildren(node, context);
+    context.relBuilder.limit(node.getOffset(), node.getLimit());
+    return context.relBuilder.peek();
+  }
+
   /**
    * Insert a reversed sort node after finding the original sort in the tree. This rebuilds the tree
    * with the reversed sort inserted right after the original sort.
@@ -1621,7 +1649,8 @@ private Pair<List<RexNode>, List<AggCall>> resolveAttributesForAggregation(
   @Override
   public RelNode visitAggregation(Aggregation node, CalcitePlanContext context) {
     Argument.ArgumentMap statsArgs = Argument.ArgumentMap.of(node.getArgExprList());
-    Boolean bucketNullable = (Boolean) statsArgs.get(Argument.BUCKET_NULLABLE).getValue();
+    boolean bucketNullable =
+        (Boolean) statsArgs.getOrDefault(Argument.BUCKET_NULLABLE, Literal.TRUE).getValue();
     int nGroup = node.getGroupExprList().size() + (Objects.nonNull(node.getSpan()) ? 1 : 0);
     BitSet nonNullGroupMask = new BitSet(nGroup);
     if (!bucketNullable) {
@@ -1928,6 +1957,15 @@ public RelNode visitSubqueryAlias(SubqueryAlias node, CalcitePlanContext context
     return context.relBuilder.peek();
   }
 
+  @Override
+  public RelNode visitRelationSubquery(RelationSubquery node, CalcitePlanContext context) {
+    // Derived table in FROM clause: SELECT ... FROM (SELECT ...) AS t
+    // Visit the inner query, then apply the alias as the relation name.
+    visitChildren(node, context);
+    context.relBuilder.as(node.getAliasAsTableName());
+    return context.relBuilder.peek();
+  }
+
   @Override
   public RelNode visitLookup(Lookup node, CalcitePlanContext context) {
     // 1. resolve source side
@@ -4122,12 +4160,14 @@ public RelNode visitMvExpand(MvExpand mvExpand, CalcitePlanContext context) {
 
   @Override
   public RelNode visitValues(Values values, CalcitePlanContext context) {
-    if (values.getValues() == null || values.getValues().isEmpty()) {
+    List<List<Literal>> rows = values.getValues();
+    if (rows == null || rows.isEmpty() || (rows.size() == 1 && rows.get(0).isEmpty())) {
+      // Single empty row (dual table) = SELECT without FROM. Provide a 0-column relation
+      // so subsequent Project can evaluate constant expressions.
       context.relBuilder.values(context.relBuilder.getTypeFactory().builder().build());
       return context.relBuilder.peek();
-    } else {
-      throw new CalciteUnsupportedException("Explicit values node is unsupported in Calcite");
     }
+    throw new CalciteUnsupportedException("Inline VALUES with literal rows is unsupported");
   }
 
   @Override
