Fixes scan_indexed and adds PIVOT (#1297)

partiql · Dec 12, 2023 · ebfc93f · ebfc93f · github-actions · Dec 12, 2023
1 parent 54990be
commit ebfc93f
Show file tree

Hide file tree

Showing 5 changed files with 70 additions and 5 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
@@ -78,6 +78,16 @@ class PartiQLSchemaInferencerTests {
     @Execution(ExecutionMode.CONCURRENT)
     fun testSelectStar(tc: TestCase) = runTest(tc)
 
+    @ParameterizedTest
+    @MethodSource("scanCases")
+    @Execution(ExecutionMode.CONCURRENT)
+    fun testScan(tc: TestCase) = runTest(tc)
+
+    @ParameterizedTest
+    @MethodSource("pivotCases")
+    @Execution(ExecutionMode.CONCURRENT)
+    fun testPivot(tc: TestCase) = runTest(tc)
+
     @ParameterizedTest
     @MethodSource("sessionVariables")
     @Execution(ExecutionMode.CONCURRENT)
@@ -420,6 +430,44 @@ class PartiQLSchemaInferencerTests {
             ),
         )
 
+        @JvmStatic
+        fun scanCases() = listOf(
+            SuccessTestCase(
+                name = "Basic Scan Indexed",
+                key = key("sanity-07"),
+                catalog = "pql",
+                expected = BagType(
+                    StructType(
+                        fields = listOf(
+                            StructType.Field("first", STRING),
+                            StructType.Field("i", INT8),
+                        ),
+                        contentClosed = true,
+                        constraints = setOf(
+                            TupleConstraint.Open(false),
+                            TupleConstraint.UniqueAttrs(true),
+                            TupleConstraint.Ordered
+                        )
+                    )
+                )
+            ),
+        )
+
+        @JvmStatic
+        fun pivotCases() = listOf(
+            SuccessTestCase(
+                name = "Basic PIVOT",
+                key = key("pivot-00"),
+                catalog = "pql",
+                expected = StructType(
+                    contentClosed = false,
+                    constraints = setOf(
+                        TupleConstraint.Open(true),
+                    )
+                )
+            ),
+        )
+
         @JvmStatic
         fun sessionVariables() = listOf(
             SuccessTestCase(

diff --git a/partiql-planner/src/main/kotlin/org/partiql/planner/internal/transforms/RelConverter.kt b/partiql-planner/src/main/kotlin/org/partiql/planner/internal/transforms/RelConverter.kt
@@ -54,6 +54,7 @@ import org.partiql.planner.internal.ir.relOpLimit
 import org.partiql.planner.internal.ir.relOpOffset
 import org.partiql.planner.internal.ir.relOpProject
 import org.partiql.planner.internal.ir.relOpScan
+import org.partiql.planner.internal.ir.relOpScanIndexed
 import org.partiql.planner.internal.ir.relOpSort
 import org.partiql.planner.internal.ir.relOpSortSpec
 import org.partiql.planner.internal.ir.relOpUnion
@@ -271,9 +272,9 @@ internal object RelConverter {
 
         private fun convertScanIndexed(rex: Rex, binding: Rel.Binding, index: Rel.Binding): Rel {
             val schema = listOf(binding, index)
-            val props = setOf(Rel.Prop.ORDERED)
+            val props = emptySet<Rel.Prop>()
             val type = relType(schema, props)
-            val op = relOpScan(rex)
+            val op = relOpScanIndexed(rex)
             return rel(type, op)
         }
 

diff --git a/partiql-planner/src/main/kotlin/org/partiql/planner/internal/typer/PlanTyper.kt b/partiql-planner/src/main/kotlin/org/partiql/planner/internal/typer/PlanTyper.kt
@@ -46,6 +46,7 @@ import org.partiql.planner.internal.ir.relOpLimit
 import org.partiql.planner.internal.ir.relOpOffset
 import org.partiql.planner.internal.ir.relOpProject
 import org.partiql.planner.internal.ir.relOpScan
+import org.partiql.planner.internal.ir.relOpScanIndexed
 import org.partiql.planner.internal.ir.relOpSort
 import org.partiql.planner.internal.ir.relOpUnpivot
 import org.partiql.planner.internal.ir.relType
@@ -60,6 +61,7 @@ import org.partiql.planner.internal.ir.rexOpGlobal
 import org.partiql.planner.internal.ir.rexOpLit
 import org.partiql.planner.internal.ir.rexOpPath
 import org.partiql.planner.internal.ir.rexOpPathStepSymbol
+import org.partiql.planner.internal.ir.rexOpPivot
 import org.partiql.planner.internal.ir.rexOpSelect
 import org.partiql.planner.internal.ir.rexOpStruct
 import org.partiql.planner.internal.ir.rexOpStructField
@@ -159,10 +161,10 @@ internal class PlanTyper(
             val rex = node.rex.type(outer.global())
             // compute rel type
             val valueT = getElementTypeForFromSource(rex.type)
-            val indexT = StaticType.INT
+            val indexT = StaticType.INT8
             val type = ctx!!.copyWithSchema(listOf(valueT, indexT))
             // rewrite
-            val op = relOpScan(rex)
+            val op = relOpScanIndexed(rex)
             return rel(type, op)
         }
 
@@ -817,7 +819,16 @@ internal class PlanTyper(
         }
 
         override fun visitRexOpPivot(node: Rex.Op.Pivot, ctx: StaticType?): Rex {
-            TODO("Type RexOpPivot")
+            val rel = node.rel.type(locals)
+            val typeEnv = TypeEnv(rel.type.schema, ResolutionStrategy.LOCAL)
+            val key = node.key.type(typeEnv)
+            val value = node.value.type(typeEnv)
+            val type = StructType(
+                contentClosed = false,
+                constraints = setOf(TupleConstraint.Open(true))
+            )
+            val op = rexOpPivot(key, value, rel)
+            return rex(type, op)
         }
 
         override fun visitRexOpSubquery(node: Rex.Op.Subquery, ctx: StaticType?): Rex {

diff --git a/partiql-planner/src/testFixtures/resources/inputs/schema_inferencer/pivot.sql b/partiql-planner/src/testFixtures/resources/inputs/schema_inferencer/pivot.sql
@@ -0,0 +1,2 @@
+--#[pivot-00]
+PIVOT t.a AT t.c FROM main.T AS t;
diff --git a/partiql-planner/src/testFixtures/resources/inputs/schema_inferencer/sanity.sql b/partiql-planner/src/testFixtures/resources/inputs/schema_inferencer/sanity.sql
@@ -45,3 +45,6 @@ SELECT
        p.name.*,
        (p.name."first" || ' ' || p.name."last") AS full_name
 FROM main.person AS p;
+
+--#[sanity-07]
+SELECT p.name."first", i FROM main.person AS p AT i;
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,2 @@
		--#[pivot-00]
		PIVOT t.a AT t.c FROM main.T AS t;