Skip to content

Commit 3acba66

Browse files
dai-chendai-chen
committed
feat(calcite): fix visitor issues for V2 SQL AST integration
Fix CalciteRelNodeVisitor/RexNodeVisitor gaps surfaced when routing the V2 SQL AST through the unified query path: - expandProjectFields: add Alias case so already-computed aggregates are referenced by schema name (was throwing "Unexpected expression type: Alias"). - visitLimit: new override for SQL LIMIT/OFFSET. PPL is unaffected (uses Head node). - visitAggregation: fix BUCKET_NULLABLE NPE when SQL GROUP BY produces an Aggregation without that arg. Use getOrDefault(..., Literal.TRUE). - visitRelationSubquery: new override for derived tables in FROM (SELECT ... FROM (SELECT ...) AS t). - visitValues: handle (a) dual-table [[]] for SELECT without FROM and (b) non-empty inline VALUES via relBuilder.values(fieldNames, flat). - visitWindowFunction: handle AggregateFunction (V2 SQL) + DISTINCT + OVER ORDER BY (was hardcoded to List.of()). Register ROW_NUMBER/RANK/DENSE_RANK in WINDOW_FUNC_MAPPING and bypass aggregate validation for these pure window functions. Add a makeOver overload with a distinct parameter (old signature delegates with false to preserve PPL behavior). Add nine tests in UnifiedQueryPlannerSqlTest, one per fix. Verified zero PPL regressions across 2300+ existing tests. Signed-off-by: Chen Dai <daichen@amazon.com>
1 parent 6421658 commit 3acba66

5 files changed

Lines changed: 261 additions & 14 deletions

File tree

‎api/src/test/java/org/opensearch/sql/api/UnifiedQueryPlannerSqlTest.java‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -259,4 +259,105 @@ SELECT department, count(*)
259259
""")
260260
.assertErrorMessage("Encountered");
261261
}
262+
263+
@Test
264+
public void testSqlLimitOffset() {
265+
givenQuery(
266+
"""
267+
SELECT name
268+
FROM catalog.employees
269+
LIMIT 10 OFFSET 5\
270+
""")
271+
.assertPlanContains("LogicalSort(offset=[5], fetch=[10])");
272+
}
273+
274+
@Test
275+
public void testSqlAggregateWithAlias() {
276+
givenQuery(
277+
"""
278+
SELECT department, COUNT(*) AS cnt
279+
FROM catalog.employees
280+
GROUP BY department\
281+
""")
282+
.assertPlanContains("LogicalAggregate(group=[{0}]")
283+
.assertPlanContains("COUNT()");
284+
}
285+
286+
@Test
287+
public void testSqlWindowFunctionWithOrderBy() {
288+
givenQuery(
289+
"""
290+
SELECT name, SUM(age) OVER (PARTITION BY department ORDER BY id) AS running_sum
291+
FROM catalog.employees\
292+
""")
293+
.assertPlanContains("OVER")
294+
.assertPlanContains("ORDER BY");
295+
}
296+
297+
@Test
298+
public void testSqlWindowRankFunction() {
299+
givenQuery(
300+
"""
301+
SELECT name, RANK() OVER (ORDER BY age DESC) AS rnk
302+
FROM catalog.employees\
303+
""")
304+
.assertPlanContains("RANK() OVER");
305+
}
306+
307+
@Test
308+
public void testSqlWindowDistinctAggregate() {
309+
givenQuery(
310+
"""
311+
SELECT name, COUNT(DISTINCT department) OVER (PARTITION BY department) AS dist_cnt
312+
FROM catalog.employees\
313+
""")
314+
.assertPlanContains("OVER")
315+
.assertPlanContains("COUNT");
316+
}
317+
318+
@Test
319+
public void testSqlGroupByWithoutBucketNullable() {
320+
givenQuery(
321+
"""
322+
SELECT age, COUNT(*) AS cnt
323+
FROM catalog.employees
324+
GROUP BY age\
325+
""")
326+
.assertPlanContains("LogicalAggregate(group=[{0}]")
327+
.assertPlanContains("COUNT()");
328+
}
329+
330+
@Test
331+
public void testSqlSelectWithAlias() {
332+
givenQuery(
333+
"""
334+
SELECT age AS employee_age, name AS employee_name
335+
FROM catalog.employees\
336+
""")
337+
.assertPlanContains("LogicalProject")
338+
.assertPlanContains("LogicalTableScan");
339+
}
340+
341+
@Test
342+
public void testSqlDerivedTableInFromClause() {
343+
// SELECT ... FROM (SELECT ...) AS t — exercises visitRelationSubquery override.
344+
givenQuery(
345+
"""
346+
SELECT t.id
347+
FROM (SELECT id, name FROM catalog.employees WHERE age > 30) AS t\
348+
""")
349+
.assertPlanContains("LogicalProject")
350+
.assertPlanContains("LogicalFilter")
351+
.assertPlanContains("LogicalTableScan");
352+
}
353+
354+
@Test
355+
public void testSqlSelectWithoutFromClause() {
356+
// SELECT 1 — exercises visitValues dual-table case (single empty row).
357+
givenQuery(
358+
"""
359+
SELECT 1\
360+
""")
361+
.assertPlanContains("LogicalValues");
362+
}
262363
}

‎core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java‎

Lines changed: 66 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,7 @@
132132
import org.opensearch.sql.ast.tree.Head;
133133
import org.opensearch.sql.ast.tree.Join;
134134
import org.opensearch.sql.ast.tree.Kmeans;
135+
import org.opensearch.sql.ast.tree.Limit;
135136
import org.opensearch.sql.ast.tree.Lookup;
136137
import org.opensearch.sql.ast.tree.Lookup.OutputStrategy;
137138
import org.opensearch.sql.ast.tree.ML;
@@ -146,6 +147,7 @@
146147
import org.opensearch.sql.ast.tree.RareTopN;
147148
import org.opensearch.sql.ast.tree.Regex;
148149
import org.opensearch.sql.ast.tree.Relation;
150+
import org.opensearch.sql.ast.tree.RelationSubquery;
149151
import org.opensearch.sql.ast.tree.Rename;
150152
import org.opensearch.sql.ast.tree.Replace;
151153
import org.opensearch.sql.ast.tree.ReplacePair;
@@ -541,6 +543,25 @@ private List<RexNode> expandProjectFields(
541543
.filter(addedFields::add)
542544
.forEach(field -> expandedFields.add(context.relBuilder.field(field)));
543545
}
546+
case Alias alias -> {
547+
String aliasName =
548+
Strings.isNullOrEmpty(alias.getAlias()) ? alias.getName() : alias.getAlias();
549+
// When an aggregate was already computed (its output name matches a current field),
550+
// reference the existing field instead of re-analyzing (which would return null).
551+
if (alias.getDelegated() instanceof AggregateFunction) {
552+
String aggFieldName = alias.getName();
553+
if (currentFields.contains(aliasName)) {
554+
expandedFields.add(context.relBuilder.field(aliasName));
555+
} else if (aggFieldName != null && currentFields.contains(aggFieldName)) {
556+
expandedFields.add(
557+
context.relBuilder.alias(context.relBuilder.field(aggFieldName), aliasName));
558+
} else {
559+
expandedFields.add(rexVisitor.analyze(alias, context));
560+
}
561+
} else {
562+
expandedFields.add(rexVisitor.analyze(alias, context));
563+
}
564+
}
544565
default ->
545566
throw new IllegalStateException(
546567
"Unexpected expression type in project list: " + expr.getClass().getSimpleName());
@@ -763,6 +784,13 @@ public RelNode visitHead(Head node, CalcitePlanContext context) {
763784
return context.relBuilder.peek();
764785
}
765786

787+
@Override
788+
public RelNode visitLimit(Limit node, CalcitePlanContext context) {
789+
visitChildren(node, context);
790+
context.relBuilder.limit(node.getOffset(), node.getLimit());
791+
return context.relBuilder.peek();
792+
}
793+
766794
/**
767795
* Insert a reversed sort node after finding the original sort in the tree. This rebuilds the tree
768796
* with the reversed sort inserted right after the original sort.
@@ -1621,7 +1649,8 @@ private Pair<List<RexNode>, List<AggCall>> resolveAttributesForAggregation(
16211649
@Override
16221650
public RelNode visitAggregation(Aggregation node, CalcitePlanContext context) {
16231651
Argument.ArgumentMap statsArgs = Argument.ArgumentMap.of(node.getArgExprList());
1624-
Boolean bucketNullable = (Boolean) statsArgs.get(Argument.BUCKET_NULLABLE).getValue();
1652+
boolean bucketNullable =
1653+
(Boolean) statsArgs.getOrDefault(Argument.BUCKET_NULLABLE, Literal.TRUE).getValue();
16251654
int nGroup = node.getGroupExprList().size() + (Objects.nonNull(node.getSpan()) ? 1 : 0);
16261655
BitSet nonNullGroupMask = new BitSet(nGroup);
16271656
if (!bucketNullable) {
@@ -1928,6 +1957,15 @@ public RelNode visitSubqueryAlias(SubqueryAlias node, CalcitePlanContext context
19281957
return context.relBuilder.peek();
19291958
}
19301959

1960+
@Override
1961+
public RelNode visitRelationSubquery(RelationSubquery node, CalcitePlanContext context) {
1962+
// Derived table in FROM clause: SELECT ... FROM (SELECT ...) AS t
1963+
// Visit the inner query, then apply the alias as the relation name.
1964+
visitChildren(node, context);
1965+
context.relBuilder.as(node.getAliasAsTableName());
1966+
return context.relBuilder.peek();
1967+
}
1968+
19311969
@Override
19321970
public RelNode visitLookup(Lookup node, CalcitePlanContext context) {
19331971
// 1. resolve source side
@@ -4122,12 +4160,36 @@ public RelNode visitMvExpand(MvExpand mvExpand, CalcitePlanContext context) {
41224160

41234161
@Override
41244162
public RelNode visitValues(Values values, CalcitePlanContext context) {
4125-
if (values.getValues() == null || values.getValues().isEmpty()) {
4163+
List<List<Literal>> rows = values.getValues();
4164+
boolean isDualTable =
4165+
rows == null || rows.isEmpty() || (rows.size() == 1 && rows.get(0).isEmpty());
4166+
if (isDualTable) {
4167+
// Single empty row (dual table) = SELECT without FROM. Provide a 0-column relation
4168+
// so subsequent Project can evaluate constant expressions.
41264169
context.relBuilder.values(context.relBuilder.getTypeFactory().builder().build());
41274170
return context.relBuilder.peek();
4128-
} else {
4129-
throw new CalciteUnsupportedException("Explicit values node is unsupported in Calcite");
41304171
}
4172+
// Non-empty inline VALUES (e.g., VALUES (1,'a'),(2,'b')). Build a Calcite values
4173+
// relation by flattening the literal values; RelBuilder.values infers types from
4174+
// the Java objects and assigns default column names.
4175+
int numCols = rows.get(0).size();
4176+
String[] fieldNames = new String[numCols];
4177+
for (int i = 0; i < numCols; i++) {
4178+
fieldNames[i] = "EXPR$" + i;
4179+
}
4180+
Object[] flat = new Object[rows.size() * numCols];
4181+
int idx = 0;
4182+
for (List<Literal> row : rows) {
4183+
if (row.size() != numCols) {
4184+
throw new CalciteUnsupportedException(
4185+
"Inconsistent VALUES row width: expected " + numCols + " got " + row.size());
4186+
}
4187+
for (Literal lit : row) {
4188+
flat[idx++] = lit.getValue();
4189+
}
4190+
}
4191+
context.relBuilder.values(fieldNames, flat);
4192+
return context.relBuilder.peek();
41314193
}
41324194

41334195
@Override

‎core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java‎

Lines changed: 61 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import java.util.List;
1919
import java.util.Locale;
2020
import java.util.Map;
21+
import java.util.Set;
2122
import java.util.stream.Collectors;
2223
import java.util.stream.IntStream;
2324
import javax.annotation.Nullable;
@@ -42,6 +43,7 @@
4243
import org.apache.calcite.util.TimestampString;
4344
import org.apache.logging.log4j.util.Strings;
4445
import org.opensearch.sql.ast.AbstractNodeVisitor;
46+
import org.opensearch.sql.ast.expression.AggregateFunction;
4547
import org.opensearch.sql.ast.expression.Alias;
4648
import org.opensearch.sql.ast.expression.And;
4749
import org.opensearch.sql.ast.expression.Between;
@@ -90,6 +92,10 @@
9092
public class CalciteRexNodeVisitor extends AbstractNodeVisitor<RexNode, CalcitePlanContext> {
9193
private final CalciteRelNodeVisitor planVisitor;
9294

95+
private static final Set<BuiltinFunctionName> PURE_WINDOW_FUNCTIONS =
96+
Set.of(
97+
BuiltinFunctionName.ROW_NUMBER, BuiltinFunctionName.RANK, BuiltinFunctionName.DENSE_RANK);
98+
9399
public RexNode analyze(UnresolvedExpression unresolved, CalcitePlanContext context) {
94100
return unresolved.accept(this, context);
95101
}
@@ -563,47 +569,93 @@ public RexNode visitFunction(Function node, CalcitePlanContext context) {
563569

564570
@Override
565571
public RexNode visitWindowFunction(WindowFunction node, CalcitePlanContext context) {
566-
Function windowFunction = (Function) node.getFunction();
567-
List<RexNode> arguments =
568-
windowFunction.getFuncArgs().stream().map(arg -> analyze(arg, context)).toList();
572+
String funcName;
573+
List<RexNode> arguments;
574+
final boolean isDistinct;
575+
if (node.getFunction() instanceof AggregateFunction aggFunc) {
576+
funcName = aggFunc.getFuncName();
577+
isDistinct = Boolean.TRUE.equals(aggFunc.getDistinct());
578+
List<UnresolvedExpression> argExprs = new java.util.ArrayList<>();
579+
if (aggFunc.getField() != null) {
580+
argExprs.add(aggFunc.getField());
581+
}
582+
argExprs.addAll(aggFunc.getArgList());
583+
arguments = argExprs.stream().map(arg -> analyze(arg, context)).toList();
584+
} else {
585+
Function windowFunction = (Function) node.getFunction();
586+
funcName = windowFunction.getFuncName();
587+
isDistinct = false;
588+
arguments = windowFunction.getFuncArgs().stream().map(arg -> analyze(arg, context)).toList();
589+
}
569590
List<RexNode> partitions =
570591
node.getPartitionByList().stream()
571592
.map(arg -> analyze(arg, context))
572593
.map(this::extractRexNodeFromAlias)
573594
.toList();
574-
return BuiltinFunctionName.ofWindowFunction(windowFunction.getFuncName())
595+
List<RexNode> orderKeys =
596+
node.getSortList().stream()
597+
.map(
598+
pair -> {
599+
RexNode sortField = analyze(pair.getRight(), context);
600+
if (pair.getLeft().getSortOrder()
601+
== org.opensearch.sql.ast.tree.Sort.SortOrder.DESC) {
602+
sortField = context.relBuilder.desc(sortField);
603+
}
604+
if (pair.getLeft().getNullOrder()
605+
== org.opensearch.sql.ast.tree.Sort.NullOrder.NULL_LAST) {
606+
sortField = context.relBuilder.nullsLast(sortField);
607+
} else if (pair.getLeft().getNullOrder()
608+
== org.opensearch.sql.ast.tree.Sort.NullOrder.NULL_FIRST) {
609+
sortField = context.relBuilder.nullsFirst(sortField);
610+
}
611+
return sortField;
612+
})
613+
.toList();
614+
return BuiltinFunctionName.ofWindowFunction(funcName)
575615
.map(
576616
functionName -> {
577617
RexNode field = arguments.isEmpty() ? null : arguments.getFirst();
578618
List<RexNode> args =
579619
(arguments.isEmpty() || arguments.size() == 1)
580620
? Collections.emptyList()
581621
: arguments.subList(1, arguments.size());
622+
// Pure window functions (ROW_NUMBER, RANK, DENSE_RANK) are not registered
623+
// in aggFunctionRegistry, so skip validation for them.
624+
if (PURE_WINDOW_FUNCTIONS.contains(functionName)) {
625+
return PlanUtils.makeOver(
626+
context,
627+
functionName,
628+
field,
629+
args,
630+
partitions,
631+
orderKeys,
632+
node.getWindowFrame());
633+
}
582634
List<RexNode> nodes =
583635
PPLFuncImpTable.INSTANCE.validateAggFunctionSignature(
584636
functionName, field, args, context.rexBuilder);
585637
return nodes != null
586638
? PlanUtils.makeOver(
587639
context,
588640
functionName,
641+
isDistinct,
589642
nodes.getFirst(),
590643
nodes.size() <= 1 ? Collections.emptyList() : nodes.subList(1, nodes.size()),
591644
partitions,
592-
List.of(),
645+
orderKeys,
593646
node.getWindowFrame())
594647
: PlanUtils.makeOver(
595648
context,
596649
functionName,
650+
isDistinct,
597651
field,
598652
args,
599653
partitions,
600-
List.of(),
654+
orderKeys,
601655
node.getWindowFrame());
602656
})
603657
.orElseThrow(
604-
() ->
605-
new UnsupportedOperationException(
606-
"Unexpected window function: " + windowFunction.getFuncName()));
658+
() -> new UnsupportedOperationException("Unexpected window function: " + funcName));
607659
}
608660

609661
/** extract the expression of Alias from a node */

0 commit comments

Comments
 (0)