Merge pull request #1316 from partiql/eval-tests

Updates pivot on rebase and join tests

partiql-eval/src/main/kotlin/org/partiql/eval/internal/Compiler.kt

@@ -8,9 +8,12 @@ import org.partiql.eval.internal.operator.rel.RelJoinOuterFull
 import org.partiql.eval.internal.operator.rel.RelJoinRight
 import org.partiql.eval.internal.operator.rel.RelProject
 import org.partiql.eval.internal.operator.rel.RelScan
+import org.partiql.eval.internal.operator.rel.RelScanIndexed
 import org.partiql.eval.internal.operator.rex.ExprCase
 import org.partiql.eval.internal.operator.rex.ExprCollection
 import org.partiql.eval.internal.operator.rex.ExprLiteral
+import org.partiql.eval.internal.operator.rex.ExprPathKey
+import org.partiql.eval.internal.operator.rex.ExprPivot
 import org.partiql.eval.internal.operator.rex.ExprSelect
 import org.partiql.eval.internal.operator.rex.ExprStruct
 import org.partiql.eval.internal.operator.rex.ExprTupleUnion
@@ -30,10 +33,13 @@ internal object Compiler {
     }
 
     private object PlanToCodeTransformer : PlanBaseVisitor<Operator, Unit>() {
+
         override fun defaultReturn(node: PlanNode, ctx: Unit): Operator {
             TODO("Not yet implemented")
         }
 
+        // REX
+
         override fun visitRexOpStruct(node: Rex.Op.Struct, ctx: Unit): Operator {
             val fields = node.fields.map {
                 ExprStruct.Field(visitRex(it.k, ctx), visitRex(it.v, ctx))
@@ -45,33 +51,53 @@ internal object Compiler {
             return ExprVar(node.ref)
         }
 
+        override fun visitRexOpPathKey(node: Rex.Op.Path.Key, ctx: Unit): Operator {
+            val root = visitRex(node.root, ctx)
+            val key = visitRex(node.key, ctx)
+            return ExprPathKey(root, key)
+        }
+
         override fun visitRexOpCollection(node: Rex.Op.Collection, ctx: Unit): Operator {
             val values = node.values.map { visitRex(it, ctx) }
             return ExprCollection(values)
         }
 
-        override fun visitRelOpProject(node: Rel.Op.Project, ctx: Unit): Operator {
-            val input = visitRel(node.input, ctx)
-            val projections = node.projections.map { visitRex(it, ctx) }
-            return RelProject(input, projections)
-        }
-
         override fun visitRexOpSelect(node: Rex.Op.Select, ctx: Unit): Operator {
             val rel = visitRel(node.rel, ctx)
             val constructor = visitRex(node.constructor, ctx)
             return ExprSelect(rel, constructor)
         }
 
-        override fun visitRelOpScan(node: Rel.Op.Scan, ctx: Unit): Operator {
-            val rex = visitRex(node.rex, ctx)
-            return RelScan(rex)
+        override fun visitRexOpPivot(node: Rex.Op.Pivot, ctx: Unit): Operator {
+            val rel = visitRel(node.rel, ctx)
+            val key = visitRex(node.key, ctx)
+            val value = visitRex(node.value, ctx)
+            return ExprPivot(rel, key, value)
         }
 
         @OptIn(PartiQLValueExperimental::class)
         override fun visitRexOpLit(node: Rex.Op.Lit, ctx: Unit): Operator {
             return ExprLiteral(node.value)
         }
 
+        // REL
+
+        override fun visitRelOpScan(node: Rel.Op.Scan, ctx: Unit): Operator {
+            val rex = visitRex(node.rex, ctx)
+            return RelScan(rex)
+        }
+
+        override fun visitRelOpProject(node: Rel.Op.Project, ctx: Unit): Operator {
+            val input = visitRel(node.input, ctx)
+            val projections = node.projections.map { visitRex(it, ctx) }
+            return RelProject(input, projections)
+        }
+
+        override fun visitRelOpScanIndexed(node: Rel.Op.ScanIndexed, ctx: Unit): Operator {
+            val rex = visitRex(node.rex, ctx)
+            return RelScanIndexed(rex)
+        }
+
         override fun visitRel(node: Rel, ctx: Unit): Operator.Relation {
             return super.visitRelOp(node.op, ctx) as Operator.Relation
         }

partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rel/RelScanIndexed.kt

@@ -0,0 +1,38 @@
+package org.partiql.eval.internal.operator.rel
+
+import org.partiql.eval.internal.Record
+import org.partiql.eval.internal.operator.Operator
+import org.partiql.value.CollectionValue
+import org.partiql.value.PartiQLValue
+import org.partiql.value.PartiQLValueExperimental
+import org.partiql.value.int64Value
+
+@OptIn(PartiQLValueExperimental::class)
+internal class RelScanIndexed(
+    private val expr: Operator.Expr
+) : Operator.Relation {
+
+    private lateinit var iterator: Iterator<PartiQLValue>
+    private var index: Long = 0
+
+    override fun open() {
+        val r = expr.eval(Record.empty)
+        index = 0
+        iterator = when (r) {
+            is CollectionValue<*> -> r.iterator()
+            else -> iterator { yield(r) }
+        }
+    }
+
+    override fun next(): Record? {
+        if (!iterator.hasNext()) {
+            return null
+        }
+        val i = index
+        val v = iterator.next()
+        index += 1
+        return Record.of(v, int64Value(i))
+    }
+
+    override fun close() {}
+}

partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rex/ExprPathKey.kt

@@ -10,8 +10,8 @@ import org.partiql.value.check
 import org.partiql.value.missingValue
 
 internal class ExprPathKey(
-    val root: Operator.Expr,
-    val key: Operator.Expr
+    @JvmField val root: Operator.Expr,
+    @JvmField val key: Operator.Expr
 ) : Operator.Expr {
 
     @OptIn(PartiQLValueExperimental::class)

partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rex/ExprPivot.kt

@@ -0,0 +1,30 @@
+package org.partiql.eval.internal.operator.rex
+
+import org.partiql.eval.internal.Record
+import org.partiql.eval.internal.operator.Operator
+import org.partiql.value.PartiQLValue
+import org.partiql.value.PartiQLValueExperimental
+import org.partiql.value.StringValue
+import org.partiql.value.check
+import org.partiql.value.structValue
+
+@OptIn(PartiQLValueExperimental::class)
+internal class ExprPivot(
+    private val input: Operator.Relation,
+    private val key: Operator.Expr,
+    private val value: Operator.Expr,
+) : Operator.Expr {
+
+    override fun eval(record: Record): PartiQLValue {
+        input.open()
+        val fields = mutableListOf<Pair<String, PartiQLValue>>()
+        while (true) {
+            val row = input.next() ?: break
+            val k = key.eval(row).check<StringValue>()
+            val v = value.eval(row)
+            fields.add(k.value!! to v)
+        }
+        input.close()
+        return structValue(fields)
+    }
+}

partiql-eval/src/test/kotlin/org/partiql/eval/internal/PartiQLEngineDefaultTest.kt

@@ -10,9 +10,11 @@ import org.partiql.planner.PartiQLPlannerBuilder
 import org.partiql.value.BagValue
 import org.partiql.value.PartiQLValue
 import org.partiql.value.PartiQLValueExperimental
+import org.partiql.value.StructValue
 import org.partiql.value.bagValue
 import org.partiql.value.boolValue
 import org.partiql.value.int32Value
+import org.partiql.value.int64Value
 import org.partiql.value.io.PartiQLValueIonWriterBuilder
 import org.partiql.value.missingValue
 import org.partiql.value.nullValue
@@ -23,10 +25,11 @@ import kotlin.test.assertEquals
 
 /**
  * This holds sanity tests during the development of the [PartiQLEngine.default] implementation.
+<<<<<<< HEAD
+=======
  *
- * TODO need to update implementations
+>>>>>>> 1772f0ed (Updates JOIN tests)
  */
-@Disabled
 class PartiQLEngineDefaultTest {
 
     private val engine = PartiQLEngine.default()
@@ -82,7 +85,7 @@ class PartiQLEngineDefaultTest {
     @OptIn(PartiQLValueExperimental::class)
     @Test
     fun testJoinInner() {
-        val statement = parser.parse("SELECT a, b FROM << { 'a': 1 } >> t, << { 'b': 2 } >> s;").root
+        val statement = parser.parse("SELECT t.a, s.b FROM << { 'a': 1 } >> t, << { 'b': 2 } >> s;").root
         val session = PartiQLPlanner.Session("q", "u")
         val plan = planner.plan(statement, session)
 
@@ -97,7 +100,7 @@ class PartiQLEngineDefaultTest {
     @OptIn(PartiQLValueExperimental::class)
     @Test
     fun testJoinLeft() {
-        val statement = parser.parse("SELECT a, b FROM << { 'a': 1 } >> t LEFT JOIN << { 'b': 2 } >> s ON false;").root
+        val statement = parser.parse("SELECT t.a, s.b FROM << { 'a': 1 } >> t LEFT JOIN << { 'b': 2 } >> s ON false;").root
         val session = PartiQLPlanner.Session("q", "u")
         val plan = planner.plan(statement, session)
 
@@ -113,8 +116,7 @@ class PartiQLEngineDefaultTest {
     @Test
     fun testJoinOuterFull() {
         val statement =
-            parser.parse("SELECT a, b FROM << { 'a': 1 } >> t FULL OUTER JOIN << { 'b': 2 } >> s ON false;").root
-
+            parser.parse("SELECT t.a, s.b FROM << { 'a': 1 } >> t FULL OUTER JOIN << { 'b': 2 } >> s ON false;").root
         val session = PartiQLPlanner.Session("q", "u")
         val plan = planner.plan(statement, session)
 
@@ -128,14 +130,14 @@ class PartiQLEngineDefaultTest {
         val output = result.value as BagValue<*>
 
         val expected = bagValue(
-            structValue(
-                "a" to int32Value(1),
-                "b" to nullValue()
-            ),
             structValue(
                 "a" to nullValue(),
                 "b" to int32Value(2)
             ),
+            structValue(
+                "a" to int32Value(1),
+                "b" to nullValue()
+            ),
         )
         assertEquals(expected, output, comparisonString(expected, output))
     }
@@ -194,7 +196,7 @@ class PartiQLEngineDefaultTest {
     @Test
     fun testJoinOuterFullOnTrue() {
         val statement =
-            parser.parse("SELECT a, b FROM << { 'a': 1 } >> t FULL OUTER JOIN << { 'b': 2 } >> s ON TRUE;").root
+            parser.parse("SELECT t.a, s.b FROM << { 'a': 1 } >> t FULL OUTER JOIN << { 'b': 2 } >> s ON TRUE;").root
 
         val session = PartiQLPlanner.Session("q", "u")
         val plan = planner.plan(statement, session)
@@ -412,19 +414,22 @@ class PartiQLEngineDefaultTest {
             SELECT VALUE {
                 'a': 1,
                 'b': NULL,
-                c : d
+                t.c : t.d
             }
             FROM <<
                 { 'c': 'hello', 'd': 'world' }
-            >>
+            >> AS t
         """.trimIndent()
         val statement = parser.parse(source).root
         val session = PartiQLPlanner.Session("q", "u")
         val plan = planner.plan(statement, session)
 
         val prepared = engine.prepare(plan.plan)
-        val result = engine.execute(prepared) as PartiQLResult.Value
-        val output = result.value
+        val result = engine.execute(prepared)
+        if (result is PartiQLResult.Error) {
+            throw result.cause
+        }
+        val output = (result as PartiQLResult.Value).value
 
         val expected: PartiQLValue = bagValue(
             structValue(
@@ -435,4 +440,63 @@ class PartiQLEngineDefaultTest {
         )
         assertEquals(expected, output)
     }
+
+    @OptIn(PartiQLValueExperimental::class)
+    @Test
+    fun testScanIndexed() {
+        val statement = parser.parse("SELECT v, i FROM << 'a', 'b', 'c' >> AS v AT i").root
+        val session = PartiQLPlanner.Session("q", "u")
+        val plan = planner.plan(statement, session)
+
+        val prepared = engine.prepare(plan.plan)
+        val result = engine.execute(prepared) as PartiQLResult.Value
+        val output = result.value
+
+        val expected = bagValue(
+            structValue(
+                "v" to stringValue("a"),
+                "i" to int64Value(0),
+            ),
+            structValue(
+                "v" to stringValue("b"),
+                "i" to int64Value(1),
+            ),
+            structValue(
+                "v" to stringValue("c"),
+                "i" to int64Value(2),
+            ),
+        )
+        assertEquals(expected, output, comparisonString(expected, output))
+    }
+
+    @OptIn(PartiQLValueExperimental::class)
+    @Test
+    fun testPivot() {
+        val statement = parser.parse(
+            """
+           PIVOT x.v AT x.k FROM << 
+                { 'k': 'a', 'v': 'x' },
+                { 'k': 'b', 'v': 'y' },
+                { 'k': 'c', 'v': 'z' }
+           >> AS x
+            """.trimIndent()
+        ).root
+        val session = PartiQLPlanner.Session("q", "u")
+        val plan = planner.plan(statement, session)
+
+        val prepared = engine.prepare(plan.plan)
+        val result = engine.execute(prepared)
+        if (result is PartiQLResult.Error) {
+            throw result.cause
+        }
+        result as PartiQLResult.Value
+        val output = result.value as StructValue<*>
+
+        val expected = structValue(
+            "a" to stringValue("x"),
+            "b" to stringValue("y"),
+            "c" to stringValue("z"),
+        )
+        assertEquals(expected, output, comparisonString(expected, output))
+    }
 }

partiql-planner/src/main/kotlin/org/partiql/planner/internal/transforms/PlanTransform.kt

@@ -115,13 +115,13 @@ internal object PlanTransform : PlanBaseVisitor<PlanNode, ProblemCallback>() {
 
     override fun visitRexOpPathIndex(node: Rex.Op.Path.Index, ctx: ProblemCallback): PlanNode {
         val root = visitRex(node.root, ctx)
-        val key = visitRex(node.root, ctx)
+        val key = visitRex(node.key, ctx)
         return org.partiql.plan.Rex.Op.Path.Index(root, key)
     }
 
     override fun visitRexOpPathKey(node: Rex.Op.Path.Key, ctx: ProblemCallback): PlanNode {
         val root = visitRex(node.root, ctx)
-        val key = visitRex(node.root, ctx)
+        val key = visitRex(node.key, ctx)
         return org.partiql.plan.Rex.Op.Path.Key(root, key)
     }