Adds support for DISTINCT in Eval

partiql · Dec 19, 2023 · c5c3d3c · c5c3d3c
1 parent b496a69
commit c5c3d3c
Show file tree

Hide file tree

Showing 4 changed files with 144 additions and 0 deletions.
diff --git a/partiql-eval/src/main/kotlin/org/partiql/eval/internal/Compiler.kt b/partiql-eval/src/main/kotlin/org/partiql/eval/internal/Compiler.kt
@@ -1,6 +1,7 @@
 package org.partiql.eval.internal
 
 import org.partiql.eval.internal.operator.Operator
+import org.partiql.eval.internal.operator.rel.RelDistinct
 import org.partiql.eval.internal.operator.rel.RelFilter
 import org.partiql.eval.internal.operator.rel.RelJoinInner
 import org.partiql.eval.internal.operator.rel.RelJoinLeft
@@ -81,6 +82,11 @@ internal object Compiler {
             return RelFilter(input, condition)
         }
 
+        override fun visitRelOpDistinct(node: Rel.Op.Distinct, ctx: Unit): Operator {
+            val input = visitRel(node.input, ctx)
+            return RelDistinct(input)
+        }
+
         override fun visitRex(node: Rex, ctx: Unit): Operator.Expr {
             return super.visitRexOp(node.op, ctx) as Operator.Expr
         }

diff --git a/partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rel/RelDistinct.kt b/partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rel/RelDistinct.kt
@@ -0,0 +1,31 @@
+package org.partiql.eval.internal.operator.rel
+
+import org.partiql.eval.internal.Record
+import org.partiql.eval.internal.operator.Operator
+
+internal class RelDistinct(
+    val input: Operator.Relation
+) : Operator.Relation {
+
+    private val seen = mutableSetOf<Record>()
+
+    override fun open() {
+        input.open()
+    }
+
+    override fun next(): Record? {
+        var next = input.next()
+        while (next != null) {
+            if (seen.contains(next).not()) {
+                seen.add(next)
+                return next
+            }
+            next = input.next()
+        }
+        return null
+    }
+
+    override fun close() {
+        input.close()
+    }
+}
diff --git a/partiql-eval/src/test/kotlin/org/partiql/eval/internal/PartiQLEngineDefaultTest.kt b/partiql-eval/src/test/kotlin/org/partiql/eval/internal/PartiQLEngineDefaultTest.kt
@@ -307,4 +307,19 @@ class PartiQLEngineDefaultTest {
         )
         assertEquals(expected, output)
     }
+
+    @OptIn(PartiQLValueExperimental::class)
+    @Test
+    fun testDistinct() {
+        val statement = parser.parse("SELECT DISTINCT VALUE t FROM <<true, false, true, false, false, false>> AS t;").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 as BagValue<*>
+
+        val expected = bagValue(sequenceOf(boolValue(true), boolValue(false)))
+        assertEquals(expected, output)
+    }
 }
diff --git a/partiql-planner/src/test/kotlin/org/partiql/planner/internal/typer/PlanTyperTestsPorted.kt b/partiql-planner/src/test/kotlin/org/partiql/planner/internal/typer/PlanTyperTestsPorted.kt
@@ -371,6 +371,98 @@ class PlanTyperTestsPorted {
             ),
         )
 
+        @JvmStatic
+        fun distinctClauseCases() = listOf<TestCase>(
+            SuccessTestCase(
+                name = "Distinct SQL Select",
+                catalog = CATALOG_AWS,
+                query = "SELECT DISTINCT a, b FROM << { 'a': 1, 'b': 'Hello, world!' } >>;",
+                expected = BagType(
+                    StructType(
+                        fields = listOf(
+                            StructType.Field("a", StaticType.INT4),
+                            StructType.Field("b", StaticType.STRING),
+                        ),
+                        contentClosed = true,
+                        constraints = setOf(
+                            TupleConstraint.Open(false),
+                            TupleConstraint.UniqueAttrs(true),
+                            TupleConstraint.Ordered
+                        )
+                    )
+                )
+            ),
+            SuccessTestCase(
+                name = "Distinct SQL Select with Ordering",
+                catalog = CATALOG_AWS,
+                query = "SELECT DISTINCT a, b FROM << { 'a': 1, 'b': 'Hello, world!' } >> ORDER BY a;",
+                expected = ListType(
+                    StructType(
+                        fields = listOf(
+                            StructType.Field("a", StaticType.INT4),
+                            StructType.Field("b", StaticType.STRING),
+                        ),
+                        contentClosed = true,
+                        constraints = setOf(
+                            TupleConstraint.Open(false),
+                            TupleConstraint.UniqueAttrs(true),
+                            TupleConstraint.Ordered
+                        )
+                    )
+                )
+            ),
+            SuccessTestCase(
+                name = "Distinct SQL Select *",
+                catalog = CATALOG_AWS,
+                query = "SELECT DISTINCT * FROM << { 'a': 1, 'b': 'Hello, world!' } >>;",
+                expected = BagType(
+                    StructType(
+                        fields = listOf(
+                            StructType.Field("a", StaticType.INT4),
+                            StructType.Field("b", StaticType.STRING),
+                        ),
+                        contentClosed = true,
+                        constraints = setOf(
+                            TupleConstraint.Open(false),
+                            TupleConstraint.UniqueAttrs(true),
+                            TupleConstraint.Ordered
+                        )
+                    )
+                )
+            ),
+            SuccessTestCase(
+                name = "Distinct SQL Select * with Ordering",
+                catalog = CATALOG_AWS,
+                query = "SELECT DISTINCT * FROM << { 'a': 1, 'b': 'Hello, world!' } >> ORDER BY a;",
+                expected = ListType(
+                    StructType(
+                        fields = listOf(
+                            StructType.Field("a", StaticType.INT4),
+                            StructType.Field("b", StaticType.STRING),
+                        ),
+                        contentClosed = true,
+                        constraints = setOf(
+                            TupleConstraint.Open(false),
+                            TupleConstraint.UniqueAttrs(true),
+                            TupleConstraint.Ordered
+                        )
+                    )
+                )
+            ),
+            SuccessTestCase(
+                name = "Distinct PartiQL Select Value *",
+                catalog = CATALOG_AWS,
+                query = "SELECT DISTINCT VALUE a FROM << { 'a': 1, 'b': 'Hello, world!' } >>;",
+                expected = BagType(StaticType.INT4)
+            ),
+            SuccessTestCase(
+                name = "Distinct PartiQL Select Value * with Ordering",
+                catalog = CATALOG_AWS,
+                query = "SELECT DISTINCT VALUE a FROM << { 'a': 1, 'b': 'Hello, world!' } >> ORDER BY a;",
+                expected = ListType(StaticType.INT4)
+            ),
+        )
+
         @JvmStatic
         fun pivotCases() = listOf(
             SuccessTestCase(