Adds projection alias support to ORDER BY clause

CHANGELOG.md

@@ -41,6 +41,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
       [#707](https://github.com/partiql/partiql-lang-kotlin/issues/707), [#683](https://github.com/partiql/partiql-lang-kotlin/issues/683),
       and [#730](https://github.com/partiql/partiql-lang-kotlin/issues/730)
 - Parsing of `INSERT` DML with `ON CONFLICT DO REPLACE EXCLUDED` based on [RFC-0011](https://github.com/partiql/partiql-docs/blob/main/RFCs/0011-partiql-insert.md)
+- Enabled projection alias support for ORDER BY clause
 
 #### Experimental Planner Additions
 

lang/src/org/partiql/lang/ast/IsTransformedOrderByAliasMeta.kt

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2019 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://aws.amazon.com/apache2.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.lang.ast
+
+import org.partiql.lang.eval.visitors.OrderBySortSpecVisitorTransform
+
+/**
+ * A [Meta] to help the [OrderBySortSpecVisitorTransform] to know when the OrderBy SortSpec has already been transformed. It
+ * essentially helps to turn
+ *
+ * ```SELECT a + 1 AS b FROM c ORDER BY b```
+ *
+ * into
+ *
+ * ```SELECT a + 1 AS b FROM c ORDER BY a + 1```
+ *
+ * even when there are multiple transforms over the AST.
+ */
+class IsTransformedOrderByAliasMeta private constructor() : Meta {
+    override val tag = TAG
+    companion object {
+        const val TAG = "\$is_transformed_order_by_alias"
+
+        val instance = IsTransformedOrderByAliasMeta()
+        val deserializer = MemoizedMetaDeserializer(TAG, instance)
+    }
+}

lang/src/org/partiql/lang/eval/EvaluatingCompiler.kt

@@ -1754,9 +1754,11 @@ internal class EvaluatingCompiler(
                         // Grouping is not needed -- simply project the results from the FROM clause directly.
                         thunkFactory.thunkEnv(metas) { env ->
 
+                            val sourcedRows = sourceThunks(env)
+
                             val orderedRows = when (orderByThunk) {
-                                null -> sourceThunks(env)
-                                else -> evalOrderBy(sourceThunks(env), orderByThunk, orderByLocationMeta)
+                                null -> sourcedRows
+                                else -> evalOrderBy(sourcedRows, orderByThunk, orderByLocationMeta)
                             }
 
                             val projectedRows = orderedRows.map { (joinedValues, projectEnv) ->

lang/src/org/partiql/lang/eval/visitors/OrderBySortSpecVisitorTransform.kt

@@ -0,0 +1,79 @@
+/*
+ * Copyright 2022 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://aws.amazon.com/apache2.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.lang.eval.visitors
+
+import org.partiql.lang.ast.IsTransformedOrderByAliasMeta
+import org.partiql.lang.domains.PartiqlAst
+import org.partiql.lang.domains.metaContainerOf
+import org.partiql.pig.runtime.SymbolPrimitive
+
+/**
+ * A [PartiqlAst.VisitorTransform] to replace the [PartiqlAst.SortSpec] of a [PartiqlAst.OrderBy] with a reference to
+ * a [PartiqlAst.ProjectItem]'s [PartiqlAst.Expr] if an alias is provided. Also utilizes [IsTransformedOrderByAliasMeta]
+ * to enforce idempotency (delivering the same result through multiple passes).
+ *
+ * Turns:
+ *
+ * ```SELECT a + 1 AS b FROM c ORDER BY b```
+ *
+ * Into:
+ *
+ * ```SELECT a + 1 AS b FROM c ORDER BY a + 1```
+ */
+internal class OrderBySortSpecVisitorTransform : VisitorTransformBase() {
+
+    private val projectionAliases: MutableMap<String, PartiqlAst.Expr> = mutableMapOf()
+
+    /**
+     * Nests itself to ensure ORDER BYs don't have access to the same [projectionAliases]
+     */
+    override fun transformExprSelect(node: PartiqlAst.Expr.Select): PartiqlAst.Expr {
+        return OrderBySortSpecVisitorTransform().transformExprSelectEvaluationOrder(node)
+    }
+
+    /**
+     * Uses default transform and adds the alias to the [projectionAliases] map
+     */
+    override fun transformProjectItemProjectExpr_asAlias(node: PartiqlAst.ProjectItem.ProjectExpr): SymbolPrimitive? {
+        val transformedAlias = super.transformProjectItemProjectExpr_asAlias(node)
+        if (node.asAlias != null) { projectionAliases[node.asAlias.text] = node.expr }
+        return transformedAlias
+    }
+
+    /**
+     * Uses the [OrderByAliasSupport] class to transform any encountered IDs in ORDER BY <sortSpec> into the appropriate
+     * expression using the [projectionAliases] while ensuring idempotency via [IsTransformedOrderByAliasMeta]
+     */
+    override fun transformSortSpec_expr(node: PartiqlAst.SortSpec): PartiqlAst.Expr {
+        val newExpr = when (node.expr.metas.containsKey(IsTransformedOrderByAliasMeta.TAG)) {
+            true -> super.transformSortSpec_expr(node)
+            false -> OrderByAliasSupport(projectionAliases).transformSortSpec_expr(node)
+        }
+        return newExpr.copy(metas = newExpr.metas + metaContainerOf(IsTransformedOrderByAliasMeta.instance))
+    }
+
+    /**
+     * A [PartiqlAst.VisitorTransform] that converts any found Expr.Id's into what it is mapped to in [aliases].
+     */
+    private class OrderByAliasSupport(val aliases: Map<String, PartiqlAst.Expr>) : VisitorTransformBase() {
+        override fun transformExprId(node: PartiqlAst.Expr.Id): PartiqlAst.Expr {
+            val transformedExpr = super.transformExprId(node)
+            return when (node.case) {
+                is PartiqlAst.CaseSensitivity.CaseSensitive -> aliases[node.name.text] ?: transformedExpr
+                else -> aliases[node.name.text.toLowerCase()] ?: aliases[node.name.text.toUpperCase()] ?: transformedExpr
+            }
+        }
+    }
+}

lang/src/org/partiql/lang/eval/visitors/VisitorTransforms.kt

@@ -14,6 +14,7 @@ fun basicVisitorTransforms() = PipelinedVisitorTransform(
     FromSourceAliasVisitorTransform(),
     GroupByItemAliasVisitorTransform(),
     AggregateSupportVisitorTransform(),
+    OrderBySortSpecVisitorTransform(),
 
     // [GroupByPathExpressionVisitorTransform] requires:
     //   - the synthetic from source aliases added by [FromSourceAliasVisitorTransform]

lang/test/org/partiql/lang/eval/EvaluatingCompilerOrderByTests.kt

@@ -271,6 +271,17 @@ class EvaluatingCompilerOrderByTests : EvaluatorTestBase() {
                 "SELECT * FROM [{'a': <<5>>}, {'a': <<1>>}, {'a': <<10>>}] ORDER BY a DESC",
                 "[{'a': <<10>>}, {'a': <<5>>}, {'a': <<1>>}]"
             ),
+            // testing alias support
+            EvaluatorTestCase(
+                "SELECT a AS b FROM [{'a': <<5>>}, {'a': <<1>>}, {'a': <<10>>}] ORDER BY b DESC",
+                "[{'b': <<10>>}, {'b': <<5>>}, {'b': <<1>>}]"
+            ),
+
+            // testing nested alias support
+            EvaluatorTestCase(
+                "SELECT b AS \"C\" FROM (SELECT a AS b FROM [{'a': <<5>>}, {'a': <<1>>}, {'a': <<10>>}] ORDER BY b DESC) ORDER BY \"C\"",
+                "[{'C': <<1>>}, {'C': <<5>>}, {'C': <<10>>}]"
+            ),
         )
     }
 

lang/test/org/partiql/lang/eval/visitors/OrderBySortSpecVisitorTransformTests.kt

@@ -0,0 +1,224 @@
+/*
+ * Copyright 2022 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://aws.amazon.com/apache2.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.lang.eval.visitors
+
+import org.junit.jupiter.params.ParameterizedTest
+import org.junit.jupiter.params.provider.ArgumentsSource
+import org.partiql.lang.util.ArgumentsProviderBase
+
+class OrderBySortSpecVisitorTransformTests : VisitorTransformTestBase() {
+
+    class ArgsProvider : ArgumentsProviderBase() {
+        override fun getParameters(): List<Any> = listOf(
+            // Simplest Case
+            TransformTestCase(
+                """
+                SELECT a AS b 
+                FROM foo
+                ORDER BY b
+            """,
+                """
+                SELECT a AS b 
+                FROM foo
+                ORDER BY a
+            """
+            ),
+            // Different Projection Aliases
+            TransformTestCase(
+                """
+                SELECT a AS b 
+                FROM (
+                  SELECT c AS d
+                  FROM e
+                  ORDER BY d
+                )
+                ORDER BY b
+            """,
+                """
+                SELECT a AS b 
+                FROM (
+                  SELECT c AS d
+                  FROM e
+                  ORDER BY c
+                )
+                ORDER BY a
+            """
+            ),
+            // Same projection alias
+            TransformTestCase(
+                """
+                SELECT a AS b 
+                FROM (
+                  SELECT c AS b
+                  FROM e
+                  ORDER BY b
+                )
+                ORDER BY b
+            """,
+                """
+                SELECT a AS b 
+                FROM (
+                  SELECT c AS b
+                  FROM e
+                  ORDER BY c
+                )
+                ORDER BY a
+            """
+            ),
+            // Complex projection expressions with same alias
+            TransformTestCase(
+                """
+                SELECT a + b AS b 
+                FROM (
+                  SELECT b + a AS b
+                  FROM e
+                  ORDER BY b
+                )
+                ORDER BY b
+            """,
+                """
+                SELECT a + b AS b 
+                FROM (
+                  SELECT b + a AS b
+                  FROM e
+                  ORDER BY b + a
+                )
+                ORDER BY a + b
+            """
+            ),
+            // Projection Aliases are lower-case while ORDER BY sort spec is case sensitive
+            TransformTestCase(
+                """
+                SELECT a + b AS b 
+                FROM (
+                  SELECT b + a AS b
+                  FROM e
+                  ORDER BY b
+                )
+                ORDER BY "b"
+            """,
+                """
+                SELECT a + b AS b 
+                FROM (
+                  SELECT b + a AS b
+                  FROM e
+                  ORDER BY b + a
+                )
+                ORDER BY a + b
+            """
+            ),
+            // Projection Aliases are Case Sensitive while ORDER BY sort spec is NOT
+            TransformTestCase(
+                """
+                SELECT a + b AS "B"
+                FROM (
+                  SELECT b + a AS "B"
+                  FROM e
+                  ORDER BY b
+                )
+                ORDER BY b
+            """,
+                """
+                SELECT a + b AS "B"
+                FROM (
+                  SELECT b + a AS "B"
+                  FROM e
+                  ORDER BY b + a
+                )
+                ORDER BY a + b
+            """
+            ),
+            // Projection Aliases are Case Insensitive while ORDER BY sort spec is Sensitive
+            TransformTestCase(
+                """
+                SELECT a + b AS b 
+                FROM (
+                  SELECT b + a AS b
+                  FROM e
+                  ORDER BY b
+                )
+                ORDER BY "B"
+            """,
+                """
+                SELECT a + b AS b 
+                FROM (
+                  SELECT b + a AS b
+                  FROM e
+                  ORDER BY b + a
+                )
+                ORDER BY "B"
+            """
+            ),
+            // Projection Aliases and ORDER BY sort specs are Case Insensitive, but sort specs have different cases
+            TransformTestCase(
+                """
+                SELECT a + b AS b 
+                FROM (
+                  SELECT b + a AS b
+                  FROM e
+                  ORDER BY b
+                )
+                ORDER BY B
+            """,
+                """
+                SELECT a + b AS b 
+                FROM (
+                  SELECT b + a AS b
+                  FROM e
+                  ORDER BY b + a
+                )
+                ORDER BY a + b
+            """
+            ),
+            // Multiple Sort Specs
+            TransformTestCase(
+                """
+                SELECT a + b AS b 
+                FROM (
+                  SELECT b + a AS b
+                  FROM e
+                  ORDER BY b
+                )
+                ORDER BY B, a, b
+            """,
+                """
+                SELECT a + b AS b 
+                FROM (
+                  SELECT b + a AS b
+                  FROM e
+                  ORDER BY b + a
+                )
+                ORDER BY a + b, a, a + b
+            """
+            ),
+            TransformTestCase(
+                """
+                SELECT (a * -1) AS a
+                FROM << { 'a': 1 }, { 'a': 2 } >>
+                ORDER BY a
+            """,
+                """
+                SELECT (a * -1) AS a
+                FROM << { 'a': 1 }, { 'a': 2 } >>
+                ORDER BY (a * -1)
+            """
+            ),
+        )
+    }
+
+    @ParameterizedTest
+    @ArgumentsSource(ArgsProvider::class)
+    fun test(tc: TransformTestCase) = runTestForIdempotentTransform(tc, OrderBySortSpecVisitorTransform())
+}