diff --git a/CHANGELOG.md b/CHANGELOG.md index 009bc369d..4cccec21c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -65,6 +65,7 @@ rewritten to their query representation - fix the parsing of DATE_ADD/DATE_DIFF with an uncapitalized datetime field argument. - fix that wrong offset is set when parsing Ion timestamp with time zone into Datum. - Reimplemented DATE_ADD with interval plus arithmetic. +- Fix ORDER BY statement does not recognize alias from SELECT statement ### Changed diff --git a/partiql-planner/src/main/kotlin/org/partiql/planner/internal/PErrors.kt b/partiql-planner/src/main/kotlin/org/partiql/planner/internal/PErrors.kt index b8139f6c7..654541bd1 100644 --- a/partiql-planner/src/main/kotlin/org/partiql/planner/internal/PErrors.kt +++ b/partiql-planner/src/main/kotlin/org/partiql/planner/internal/PErrors.kt @@ -186,6 +186,26 @@ internal object PErrors { ) } + /** + * @param location see [PError.location] + * @param id see [PError.VAR_REF_AMBIGUOUS] + * @param candidates see [PError.VAR_REF_AMBIGUOUS] + * @return an error representing [PError.VAR_REF_AMBIGUOUS] + */ + internal fun varRefAmbiguous( + location: SourceLocation?, + id: org.partiql.ast.Identifier?, + candidates: List? + ): PError { + return PError( + PError.VAR_REF_AMBIGUOUS, + Severity.ERROR(), + PErrorKind.SEMANTIC(), + location, + mapOf("ID" to id, "CANDIDATES" to candidates) + ) + } + /** * @param path see [PError.INVALID_EXCLUDE_PATH] * @return an error representing [PError.INVALID_EXCLUDE_PATH] diff --git a/partiql-planner/src/main/kotlin/org/partiql/planner/internal/SqlPlanner.kt b/partiql-planner/src/main/kotlin/org/partiql/planner/internal/SqlPlanner.kt index 47e2c6787..40e1a8692 100644 --- a/partiql-planner/src/main/kotlin/org/partiql/planner/internal/SqlPlanner.kt +++ b/partiql-planner/src/main/kotlin/org/partiql/planner/internal/SqlPlanner.kt @@ -10,12 +10,14 @@ import org.partiql.planner.PartiQLPlannerPass import org.partiql.planner.internal.transforms.AstToPlan import org.partiql.planner.internal.transforms.NormalizeFromSource import org.partiql.planner.internal.transforms.NormalizeGroupBy +import org.partiql.planner.internal.transforms.OrderByAliasSupport import org.partiql.planner.internal.transforms.PlanTransform import org.partiql.planner.internal.typer.PlanTyper import org.partiql.spi.Context import org.partiql.spi.catalog.Session import org.partiql.spi.errors.PError import org.partiql.spi.errors.PErrorKind +import org.partiql.spi.errors.PErrorListener import org.partiql.spi.errors.PRuntimeException import org.partiql.spi.types.PType @@ -36,7 +38,7 @@ internal class SqlPlanner( val env = Env(session, ctx.errorListener) // 1. Normalize - val ast = statement.normalize() + val ast = statement.normalize(ctx.errorListener) // 2. AST to Rel/Rex val root = AstToPlan.apply(ast, env) @@ -64,11 +66,12 @@ internal class SqlPlanner( /** * AST normalization */ - private fun Statement.normalize(): Statement { + private fun Statement.normalize(listener: PErrorListener): Statement { // could be a fold, but this is nice for setting breakpoints var ast = this ast = NormalizeFromSource.apply(ast) ast = NormalizeGroupBy.apply(ast) + ast = OrderByAliasSupport(listener).apply(ast) return ast } diff --git a/partiql-planner/src/main/kotlin/org/partiql/planner/internal/transforms/OrderByAliasSupport.kt b/partiql-planner/src/main/kotlin/org/partiql/planner/internal/transforms/OrderByAliasSupport.kt new file mode 100644 index 000000000..5aace0396 --- /dev/null +++ b/partiql-planner/src/main/kotlin/org/partiql/planner/internal/transforms/OrderByAliasSupport.kt @@ -0,0 +1,178 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package org.partiql.planner.internal.transforms + +import org.partiql.ast.Ast.exprQuerySet +import org.partiql.ast.Ast.orderBy +import org.partiql.ast.Ast.sort +import org.partiql.ast.AstNode +import org.partiql.ast.AstRewriter +import org.partiql.ast.OrderBy +import org.partiql.ast.QueryBody +import org.partiql.ast.SelectItem +import org.partiql.ast.Statement +import org.partiql.ast.With +import org.partiql.ast.expr.Expr +import org.partiql.ast.expr.ExprQuerySet +import org.partiql.ast.expr.ExprVarRef +import org.partiql.planner.internal.PErrors +import org.partiql.spi.errors.PErrorListener +import kotlin.collections.MutableMap + +/** + * Replaces ORDER BY aliases with their corresponding SELECT expressions using stack-based scope tracking. + */ +internal class OrderByAliasSupport(val listener: PErrorListener) : AstPass { + override fun apply(statement: Statement): Statement { + return Visitor(listener).visitStatement(statement, Context()) as Statement + } + + /** + * Context for tracking alias scopes across nested queries using a stack-based approach. + * + * Structure breakdown: + * - **Stack (ArrayDeque)**: Each query level gets its own scope pushed/popped from stack + * - **Map (String -> MutableList)**: Maps alias names to their expressions within a scope + * - **List (MutableList)**: Handles duplicate aliases (same name, multiple expressions) + * + * Example for nested query with duplicates: + * ```sql + * SELECT pid AS p FROM ( + * SELECT name AS x, age AS x FROM products ORDER BY x + * ) ORDER BY p + * ``` + * + * Stack operations: + * 1. Push outer scope: {"p" -> [pid]} + * 2. Push inner scope: {"x" -> [name, age]} // List handles duplicate "x" aliases + * 3. Resolve ORDER BY x: detects ambiguity from list size > 1 + * 4. Pop inner scope + * 5. Resolve ORDER BY p: finds single expression from list + * 6. Pop outer scope + */ + data class Context(val aliasMapStack: ArrayDeque>> = ArrayDeque()) + + /** + * Maintains alias scope stack for nested queries and resolves ORDER BY aliases. + */ + private class Visitor(val listener: PErrorListener) : AstRewriter() { + /** + * Pushes new alias scope, processes query, then pops scope. + */ + override fun visitExprQuerySet(node: ExprQuerySet, ctx: Context): AstNode { + // Push new scope + ctx.aliasMapStack.addLast(mutableMapOf()) + + // Visit all statements that may have SELECT or ORDER BY + val body = node.body.let { visitQueryBody(it, ctx) as QueryBody } + val orderBy = node.orderBy?.let { + if (body is QueryBody.SetOp) { + // Skip alias replacement if the query body is set operations + node.orderBy + } else { + visitOrderBy(it, ctx) as OrderBy? + } + } + val with = node.with?.let { visitWith(it, ctx) as With? } + val transformed = if (body !== node.body || orderBy !== node.orderBy || with !== node.with + ) { + exprQuerySet(body, orderBy, node.limit, node.offset, with) + } else { + node + } + + // Pop scope + ctx.aliasMapStack.removeLast() + return transformed + } + + /** + * Collects SELECT aliases into current scope map. + */ + override fun visitSelectItem(node: SelectItem, ctx: Context): AstNode { + if (node is SelectItem.Expr) { + node.asAlias?.let { alias -> + ctx.aliasMapStack.last().getOrPut(alias.text) { mutableListOf() }.add(node.expr) + } + } + + return node + } + + /** + * Replaces ORDER BY aliases with their SELECT expressions. + */ + override fun visitOrderBy(node: OrderBy, ctx: Context): AstNode { + val aliasMap = ctx.aliasMapStack.last() + if (aliasMap.isEmpty()) return node + + val transformedSorts = node.sorts.map { sort -> + val transformedExpr = resolveExpr(sort.expr, aliasMap) + if (transformedExpr != sort.expr) { + sort( + expr = transformedExpr, + order = sort.order, + nulls = sort.nulls + ) + } else { + sort + } + } + return orderBy(transformedSorts) + } + + /** + * Resolves variable references to their aliased expressions. + * Regular identifiers use case-insensitive matching, delimited use case-sensitive. + */ + private fun resolveExpr(expr: Expr, aliasMap: Map>): Expr { + return when (expr) { + is ExprVarRef -> { + val identifier = expr.identifier.identifier + val orderByName = identifier.text + + val candidates = if (identifier.isRegular) { + // This is O(N) look up time for each alias used in order by. In the case of performance concern, + // we need to consider implementing a normalized case-insensitive map for O(1) look up. + aliasMap.filterKeys { it.equals(orderByName, ignoreCase = true) }.values.flatten() + } else { + aliasMap[orderByName] + } + + if (candidates == null) { + expr + } else if (candidates.size == 1) { + candidates[0] + } else { + if (candidates.size > 1) { + val candidateNames = candidates.mapNotNull { + val ref = it + if (ref is ExprVarRef) { + ref.identifier.identifier.text + } else { + null + } + } + listener.report(PErrors.varRefAmbiguous(null, expr.identifier, candidateNames)) + } + expr + } + } + else -> expr + } + } + } +} diff --git a/partiql-planner/src/test/kotlin/org/partiql/planner/internal/transforms/OrderByAliasSupportTest.kt b/partiql-planner/src/test/kotlin/org/partiql/planner/internal/transforms/OrderByAliasSupportTest.kt new file mode 100644 index 000000000..8ba445b9b --- /dev/null +++ b/partiql-planner/src/test/kotlin/org/partiql/planner/internal/transforms/OrderByAliasSupportTest.kt @@ -0,0 +1,475 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package org.partiql.planner.internal.transforms + +import org.junit.jupiter.api.assertThrows +import org.junit.jupiter.params.ParameterizedTest +import org.junit.jupiter.params.provider.MethodSource +import org.partiql.parser.PartiQLParser +import org.partiql.spi.Context +import org.partiql.spi.errors.PError +import org.partiql.spi.errors.PErrorKind +import org.partiql.spi.errors.PRuntimeException +import org.partiql.spi.errors.Severity +import kotlin.test.assertEquals + +class OrderByAliasSupportTest { + + data class SuccessTestCase( + val name: String, + val sql: String, + val transformedSql: String, + val description: String, + ) + + data class FailureTestCase( + val name: String, + val sql: String, + val description: String, + val exception: PRuntimeException + ) + + companion object { + @JvmStatic + fun successTestCases() = listOf( + SuccessTestCase( + name = "regular_alias", + sql = """ + SELECT col AS alias + FROM t + ORDER BY alias + """, + transformedSql = """ + SELECT col AS alias + FROM t + ORDER BY col + """, + description = "Regular column alias should resolve to original column" + ), + SuccessTestCase( + name = "aggregation_alias", + sql = """ + SELECT MAX(price) AS max_price + FROM products + GROUP BY category + ORDER BY max_price + """, + transformedSql = """ + SELECT MAX(price) AS max_price + FROM products + GROUP BY category + ORDER BY MAX(price) + """, + description = "Aggregation alias should resolve to original aggregation function" + ), + SuccessTestCase( + name = "nested_query", + sql = """ + SELECT pid AS p + FROM ( + SELECT productId AS pid + FROM products + ORDER BY pid + ) + ORDER BY p + """, + transformedSql = """ + SELECT pid AS p + FROM ( + SELECT productId AS pid + FROM products + ORDER BY productId + ) + ORDER BY pid + """, + description = "Nested query alias should resolve correctly at each scope level" + ), + SuccessTestCase( + name = "multiple_aliases", + sql = """ + SELECT col1 AS a, col2 AS b + FROM t + ORDER BY b, a + """, + transformedSql = """ + SELECT col1 AS a, col2 AS b + FROM t + ORDER BY col2, col1 + """, + description = "Multiple aliases should resolve in correct order" + ), + SuccessTestCase( + name = "no_alias", + sql = """ + SELECT col + FROM t + ORDER BY col + """, + transformedSql = """ + SELECT col + FROM t + ORDER BY col + """, + description = "No alias case should remain unchanged" + ), + // Test cases from https://github.com/partiql/partiql-lang-kotlin/blob/v0.14.9/partiql-lang/src/test/kotlin/org/partiql/lang/eval/visitors/OrderBySortSpecVisitorTransformTests.kt + SuccessTestCase( + name = "simplest_case", + sql = """ + SELECT a AS b + FROM foo + ORDER BY b + """, + transformedSql = """ + SELECT a AS b + FROM foo + ORDER BY a + """, + description = "Simplest case of alias resolution" + ), + SuccessTestCase( + name = "different_projection_aliases", + sql = """ + SELECT a AS b + FROM ( + SELECT c AS d + FROM e + ORDER BY d + ) + ORDER BY b + """, + transformedSql = """ + SELECT a AS b + FROM ( + SELECT c AS d + FROM e + ORDER BY c + ) + ORDER BY a + """, + description = "Different projection aliases in nested queries" + ), + SuccessTestCase( + name = "same_projection_alias", + sql = """ + SELECT a AS b + FROM ( + SELECT c AS b + FROM e + ORDER BY b + ) + ORDER BY b + """, + transformedSql = """ + SELECT a AS b + FROM ( + SELECT c AS b + FROM e + ORDER BY c + ) + ORDER BY a + """, + description = "Same projection alias in nested queries" + ), + SuccessTestCase( + name = "complex_projection_same_alias", + sql = """ + SELECT a + b AS b + FROM ( + SELECT b + a AS b + FROM e + ORDER BY b + ) + ORDER BY b + """, + transformedSql = """ + SELECT a + b AS b + FROM ( + SELECT b + a AS b + FROM e + ORDER BY b + a + ) + ORDER BY a + b + """, + description = "Complex projection expressions with same alias" + ), + SuccessTestCase( + name = "case_sensitive_alias", + sql = """ + SELECT a + b AS b + FROM ( + SELECT b + a AS b + FROM e + ORDER BY b + ) + ORDER BY "b" + """, + transformedSql = """ + SELECT a + b AS b + FROM ( + SELECT b + a AS b + FROM e + ORDER BY b + a + ) + ORDER BY a + b + """, + description = "Case sensitive ORDER BY with lowercase projection alias" + ), + SuccessTestCase( + name = "case_sensitive_projection", + sql = """ + SELECT a + b AS "B" + FROM ( + SELECT b + a AS "B" + FROM e + ORDER BY b + ) + ORDER BY b + """, + transformedSql = """ + SELECT a + b AS "B" + FROM ( + SELECT b + a AS "B" + FROM e + ORDER BY b + a + ) + ORDER BY a + b + """, + description = "Case sensitive projection alias with case insensitive ORDER BY" + ), + SuccessTestCase( + name = "case_sensitive_mismatch", + sql = """ + SELECT a + b AS b + FROM ( + SELECT b + a AS b + FROM e + ORDER BY b + ) + ORDER BY "B" + """, + transformedSql = """ + SELECT a + b AS b + FROM ( + SELECT b + a AS b + FROM e + ORDER BY b + a + ) + ORDER BY "B" + """, + description = "Case insensitive projection with case sensitive ORDER BY - no match" + ), + SuccessTestCase( + name = "case_insensitive_different_cases", + sql = """ + SELECT a + b AS b + FROM ( + SELECT b + a AS b + FROM e + ORDER BY b + ) + ORDER BY B + """, + transformedSql = """ + SELECT a + b AS b + FROM ( + SELECT b + a AS b + FROM e + ORDER BY b + a + ) + ORDER BY a + b + """, + description = "Case insensitive aliases with different cases in ORDER BY" + ), + SuccessTestCase( + name = "multiple_sort_specs", + sql = """ + SELECT a + b AS b + FROM ( + SELECT b + a AS b + FROM e + ORDER BY b + ) + ORDER BY B, a, b + """, + transformedSql = """ + SELECT a + b AS b + FROM ( + SELECT b + a AS b + FROM e + ORDER BY b + a + ) + ORDER BY a + b, a, a + b + """, + description = "Multiple sort specifications with aliases" + ), + SuccessTestCase( + name = "negative_expression_alias", + sql = """ + SELECT (a * -1) AS a + FROM << { 'a': 1 }, { 'a': 2 } >> + ORDER BY a + """, + transformedSql = """ + SELECT (a * -1) AS a + FROM << { 'a': 1 }, { 'a': 2 } >> + ORDER BY (a * -1) + """, + description = "Negative expression with alias matching original column" + ), + // Set operator test cases + SuccessTestCase( + name = "union_with_order_by_alias", + sql = """ + SELECT col AS alias FROM t1 + UNION + SELECT col AS alias FROM t2 + ORDER BY alias + """, + transformedSql = """ + SELECT col AS alias FROM t1 + UNION + SELECT col AS alias FROM t2 + ORDER BY alias + """, + description = "UNION with ORDER BY alias resolution" + ), + SuccessTestCase( + name = "union_with_order_by_alias_with_inner_order_by", + sql = """ + (SELECT a AS c FROM tA ORDER BY c LIMIT 4 OFFSET 1) + UNION + (SELECT b AS c FROM tB ORDER BY c LIMIT 4 OFFSET 3) + ORDER BY c DESC LIMIT 100 OFFSET 1 + """, + transformedSql = """ + (SELECT a AS c FROM tA ORDER BY a LIMIT 4 OFFSET 1) + UNION + (SELECT b AS c FROM tB ORDER BY b LIMIT 4 OFFSET 3) + ORDER BY c DESC LIMIT 100 OFFSET 1 + """, + description = "UNION with ORDER BY alias resolution" + ), + SuccessTestCase( + name = "intersect_with_order_by_alias", + sql = """ + SELECT price AS p FROM products + INTERSECT + SELECT cost AS p FROM orders + ORDER BY p + """, + transformedSql = """ + SELECT price AS p FROM products + INTERSECT + SELECT cost AS p FROM orders + ORDER BY p + """, + description = "INTERSECT with ORDER BY alias resolution" + ), + SuccessTestCase( + name = "except_with_order_by_alias", + sql = """ + SELECT id AS identifier FROM table1 + EXCEPT + SELECT id AS identifier FROM table2 + ORDER BY identifier + """, + transformedSql = """ + SELECT id AS identifier FROM table1 + EXCEPT + SELECT id AS identifier FROM table2 + ORDER BY identifier + """, + description = "EXCEPT with ORDER BY alias resolution" + ), + SuccessTestCase( + name = "union_all_with_complex_alias", + sql = """ + SELECT a + b AS total FROM t1 + UNION ALL + SELECT x + y AS total FROM t2 + ORDER BY total DESC + """, + transformedSql = """ + SELECT a + b AS total FROM t1 + UNION ALL + SELECT x + y AS total FROM t2 + ORDER BY total DESC + """, + description = "UNION ALL with complex expression alias resolution" + ) + ) + + @JvmStatic + fun failureTestCases() = listOf( + FailureTestCase( + name = "case_insensitive_ambiguous_alias", + sql = """ + SELECT col1 as a, col2 as A + FROM t + ORDER BY a + """, + description = "ambiguous_alias case should remain unchanged", + exception = PRuntimeException(PError(PError.VAR_REF_AMBIGUOUS, Severity.ERROR(), PErrorKind.SEMANTIC(), null, null)) + ), + FailureTestCase( + name = "case_sensitive_ambiguous_alias", + sql = """ + SELECT col1 as a, col2 as a + FROM t + ORDER BY "a" + """, + description = "ambiguous_alias case should remain unchanged", + exception = PRuntimeException(PError(PError.VAR_REF_AMBIGUOUS, Severity.ERROR(), PErrorKind.SEMANTIC(), null, null)) + ), + ) + } + + @ParameterizedTest + @MethodSource("successTestCases") + fun testOrderByAliasResolutionSuccess(testCase: SuccessTestCase) { + // Parse original SQL to AST + val parser = PartiQLParser.standard() + val originalStatement = parser.parse(testCase.sql.trimIndent()) + val ctx = Context.standard() + + // Apply OrderByAliasSupport transform + val transformedStatement = OrderByAliasSupport(ctx.errorListener).apply(originalStatement.statements[0]) + + // Parse expected SQL to AST for comparison + val expectedStatement = parser.parse(testCase.transformedSql.trimIndent()).statements[0] + + assertEquals(expectedStatement, transformedStatement, testCase.description) + } + + @ParameterizedTest + @MethodSource("failureTestCases") + fun testOrderByAliasResolutionFailed(testCase: FailureTestCase) { + // Parse original SQL to AST + val parser = PartiQLParser.standard() + val originalStatement = parser.parse(testCase.sql.trimIndent()) + val ctx = Context.standard() + + val exception = assertThrows { + // Apply OrderByAliasSupport transform + val transformedStatement = OrderByAliasSupport(ctx.errorListener).apply(originalStatement.statements[0]) + } + + assertEquals(testCase.exception.error.code(), exception.error.code()) + assertEquals(testCase.exception.error.kind, exception.error.kind) + assertEquals(testCase.exception.error.severity, exception.error.severity) + } +} diff --git a/test/partiql-tests-runner/src/test/resources/config/eval/skip-eval-core.txt b/test/partiql-tests-runner/src/test/resources/config/eval/skip-eval-core.txt index e0e2ef2f6..f1dbd2992 100644 --- a/test/partiql-tests-runner/src/test/resources/config/eval/skip-eval-core.txt +++ b/test/partiql-tests-runner/src/test/resources/config/eval/skip-eval-core.txt @@ -1,8 +1,4 @@ // TODO: NOT NEEDED FOR V1 RELEASE. Alias tests. Fixing this would not result in an API change. -PERMISSIVE:::testing alias support -STRICT:::testing alias support -PERMISSIVE:::testing nested alias support -STRICT:::testing nested alias support PERMISSIVE:::group and order by count STRICT:::group and order by count