Adds ORDER BY typing

partiql · Oct 9, 2023 · b0839c7 · b0839c7
1 parent 19837a2
commit b0839c7
Show file tree

Hide file tree

Showing 3 changed files with 69 additions and 28 deletions.
diff --git a/...-lang/src/test/kotlin/org/partiql/lang/planner/transforms/PartiQLSchemaInferencerTests.kt b/...-lang/src/test/kotlin/org/partiql/lang/planner/transforms/PartiQLSchemaInferencerTests.kt
@@ -83,10 +83,15 @@ class PartiQLSchemaInferencerTests {
     @Execution(ExecutionMode.CONCURRENT)
     fun testJoins(tc: TestCase) = runTest(tc)
 
+    // @ParameterizedTest
+    // @MethodSource("excludeCases")
+    // @Execution(ExecutionMode.CONCURRENT)
+    // fun testExclude(tc: TestCase) = runTest(tc)
+
     @ParameterizedTest
-    @MethodSource("excludeCases")
+    @MethodSource("orderByCases")
     @Execution(ExecutionMode.CONCURRENT)
-    fun testExclude(tc: TestCase) = runTest(tc)
+    fun testOrderBy(tc: TestCase) = runTest(tc)
 
     companion object {
 
@@ -222,6 +227,26 @@ class PartiQLSchemaInferencerTests {
                 query = "SEXP ( 1, 2, 3 )",
                 expected = SexpType(INT),
             ),
+            SuccessTestCase(
+                name = "SELECT from array",
+                query = "SELECT VALUE x FROM [ 1, 2, 3 ] as x",
+                expected = BagType(INT),
+            ),
+            SuccessTestCase(
+                name = "SELECT from array",
+                query = "SELECT x FROM [ 1, 2, 3 ] as x",
+                expected = BagType(
+                    StructType(
+                        fields = listOf(StructType.Field("x", INT)),
+                        contentClosed = true,
+                        constraints = setOf(
+                            TupleConstraint.Open(false),
+                            TupleConstraint.UniqueAttrs(true),
+                            TupleConstraint.Ordered
+                        )
+                    )
+                )
+            ),
         )
 
         @JvmStatic
@@ -260,8 +285,6 @@ class PartiQLSchemaInferencerTests {
                 query = "SELECT VALUE a FROM [ 0 ] AS a WHERE CURRENT_USER = 'hello'",
                 expected = BagType(INT)
             ),
-            // TODO discuss how this is not an ERROR case. It's nonsense, but is not an error. The PartiQL `=` always
-            // returns a boolean so this is valid query with no typing errors, just a bit silly.
             SuccessTestCase(
                 name = "Current User in WHERE",
                 query = "SELECT VALUE a FROM [ 0 ] AS a WHERE CURRENT_USER = 5",
@@ -2079,6 +2102,31 @@ class PartiQLSchemaInferencerTests {
                 )
             ),
         )
+
+        @JvmStatic
+        fun orderByCases() = listOf(
+            SuccessTestCase(
+                name = "ORDER BY int",
+                catalog = CATALOG_AWS,
+                catalogPath = listOf("ddb"),
+                query = "SELECT * FROM pets ORDER BY id",
+                expected = TABLE_AWS_DDB_PETS_LIST
+            ),
+            SuccessTestCase(
+                name = "ORDER BY str",
+                catalog = CATALOG_AWS,
+                catalogPath = listOf("ddb"),
+                query = "SELECT * FROM pets ORDER BY breed",
+                expected = TABLE_AWS_DDB_PETS_LIST
+            ),
+            SuccessTestCase(
+                name = "ORDER BY str",
+                catalog = CATALOG_AWS,
+                catalogPath = listOf("ddb"),
+                query = "SELECT * FROM pets ORDER BY unknown_col",
+                expected = TABLE_AWS_DDB_PETS_LIST
+            ),
+        )
     }
 
     sealed class TestCase {
@@ -2545,27 +2593,6 @@ class PartiQLSchemaInferencerTests {
                     )
                 }
             ),
-            SuccessTestCase(
-                name = "ORDER BY int",
-                catalog = CATALOG_AWS,
-                catalogPath = listOf("ddb"),
-                query = "SELECT * FROM pets ORDER BY id",
-                expected = TABLE_AWS_DDB_PETS_LIST
-            ),
-            SuccessTestCase(
-                name = "ORDER BY str",
-                catalog = CATALOG_AWS,
-                catalogPath = listOf("ddb"),
-                query = "SELECT * FROM pets ORDER BY breed",
-                expected = TABLE_AWS_DDB_PETS_LIST
-            ),
-            SuccessTestCase(
-                name = "ORDER BY str",
-                catalog = CATALOG_AWS,
-                catalogPath = listOf("ddb"),
-                query = "SELECT * FROM pets ORDER BY unknown_col",
-                expected = TABLE_AWS_DDB_PETS_LIST
-            ),
             SuccessTestCase(
                 name = "LIMIT INT",
                 catalog = CATALOG_AWS,

diff --git a/partiql-planner/src/main/kotlin/org/partiql/planner/transforms/RelConverter.kt b/partiql-planner/src/main/kotlin/org/partiql/planner/transforms/RelConverter.kt
@@ -60,6 +60,7 @@ import org.partiql.plan.rexOpTupleUnionArgSpread
 import org.partiql.plan.rexOpTupleUnionArgStruct
 import org.partiql.plan.rexOpVarResolved
 import org.partiql.planner.Env
+import org.partiql.types.ListType
 import org.partiql.types.StaticType
 import org.partiql.value.PartiQLValueExperimental
 import org.partiql.value.boolValue
@@ -236,7 +237,7 @@ internal object RelConverter {
         override fun visitSelectProject(node: Select.Project, input: Rel): Rel {
             // this ignores aggregations
             val schema = mutableListOf<Rel.Binding>()
-            val props = emptySet<Rel.Prop>()
+            val props = input.type.props
             val projections = mutableListOf<Rex>()
             node.items.forEach {
                 val (binding, projection) = convertProjectionItem(it)

diff --git a/partiql-planner/src/main/kotlin/org/partiql/planner/typer/PlanTyper.kt b/partiql-planner/src/main/kotlin/org/partiql/planner/typer/PlanTyper.kt
@@ -35,6 +35,7 @@ import org.partiql.plan.relOpLimit
 import org.partiql.plan.relOpOffset
 import org.partiql.plan.relOpProject
 import org.partiql.plan.relOpScan
+import org.partiql.plan.relOpSort
 import org.partiql.plan.relOpUnpivot
 import org.partiql.plan.relType
 import org.partiql.plan.rex
@@ -192,7 +193,19 @@ internal class PlanTyper(
         }
 
         override fun visitRelOpSort(node: Rel.Op.Sort, ctx: Rel.Type?): Rel {
-            TODO("Type RelOp Sort")
+            // compute input schema
+            val input = visitRel(node.input, ctx)
+            // type sub-nodes
+            val typeEnv = TypeEnv(input.type.schema, ResolutionStrategy.LOCAL)
+            val specs = node.specs.map {
+                val rex = it.rex.type(typeEnv)
+                it.copy(rex)
+            }
+            // output schema of a sort is the same as the input
+            val type = input.type.copy(props = setOf(Rel.Prop.ORDERED))
+            // rewrite
+            val op = relOpSort(input, specs)
+            return rel(type, op)
         }
 
         override fun visitRelOpSortSpec(node: Rel.Op.Sort.Spec, ctx: Rel.Type?): Rel {
@@ -249,7 +262,7 @@ internal class PlanTyper(
             val projections = node.projections.map { it.type(typeEnv) }
             // compute output schema
             val schema = projections.map { it.type }
-            val type = ctx!!.copyWithSchema(schema)
+            val type = input.type.copyWithSchema(schema)
             // rewrite
             val op = relOpProject(input, projections)
             return rel(type, op)