From f505d2964c875f88fd6e24b7818117699321e69c Mon Sep 17 00:00:00 2001
From: John Ed Quinn <johqunn@amazon.com>
Date: Fri, 8 Dec 2023 14:55:56 -0800
Subject: [PATCH] Adds support for DISTINCT in the Planner

---
 .../internal/transforms/RelConverter.kt       | 25 ++++++++++++++++---
 .../planner/internal/typer/PlanTyper.kt       |  4 ++-
 2 files changed, 25 insertions(+), 4 deletions(-)

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
index 15afd9c3b..d7b79db8d 100644
--- 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
@@ -24,6 +24,7 @@ import org.partiql.ast.GroupBy
 import org.partiql.ast.OrderBy
 import org.partiql.ast.Select
 import org.partiql.ast.SetOp
+import org.partiql.ast.SetQuantifier
 import org.partiql.ast.Sort
 import org.partiql.ast.builder.ast
 import org.partiql.ast.helpers.toBinder
@@ -37,6 +38,7 @@ import org.partiql.planner.internal.ir.rel
 import org.partiql.planner.internal.ir.relBinding
 import org.partiql.planner.internal.ir.relOpAggregate
 import org.partiql.planner.internal.ir.relOpAggregateCall
+import org.partiql.planner.internal.ir.relOpDistinct
 import org.partiql.planner.internal.ir.relOpErr
 import org.partiql.planner.internal.ir.relOpExcept
 import org.partiql.planner.internal.ir.relOpExclude
@@ -146,14 +148,31 @@ internal object RelConverter {
             rel = convertExclude(rel, sel.exclude)
             // append SQL projection if present
             rel = when (val projection = sel.select) {
-                is Select.Project -> visitSelectProject(projection, rel)
-                is Select.Value -> visitSelectValue(projection, rel)
+                is Select.Project -> {
+                    val project = visitSelectProject(projection, rel)
+                    visitSetQuantifier(projection.setq, project)
+                }
+                is Select.Value -> {
+                    val project = visitSelectValue(projection, rel)
+                    visitSetQuantifier(projection.setq, project)
+                }
                 is Select.Star -> error("AST not normalized, found project star")
-                else -> rel // skip PIVOT and SELECT VALUE
+                is Select.Pivot -> rel // Skip PIVOT
             }
             return rel
         }
 
+        /**
+         * Given a non-null [setQuantifier], this will return a [Rel] of [Rel.Op.Distinct] wrapping the [input].
+         * If [setQuantifier] is null or ALL, this will return the [input].
+         */
+        private fun visitSetQuantifier(setQuantifier: SetQuantifier?, input: Rel): Rel {
+            return when (setQuantifier) {
+                SetQuantifier.DISTINCT -> rel(input.type, relOpDistinct(input))
+                SetQuantifier.ALL, null -> input
+            }
+        }
+
         override fun visitSelectProject(node: Select.Project, input: Rel): Rel {
             // this ignores aggregations
             val schema = mutableListOf<Rel.Binding>()
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
index d81e42b6e..9497d5e4f 100644
--- 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
@@ -38,6 +38,7 @@ import org.partiql.planner.internal.ir.rel
 import org.partiql.planner.internal.ir.relBinding
 import org.partiql.planner.internal.ir.relOpAggregate
 import org.partiql.planner.internal.ir.relOpAggregateCall
+import org.partiql.planner.internal.ir.relOpDistinct
 import org.partiql.planner.internal.ir.relOpErr
 import org.partiql.planner.internal.ir.relOpFilter
 import org.partiql.planner.internal.ir.relOpJoin
@@ -197,7 +198,8 @@ internal class PlanTyper(
         }
 
         override fun visitRelOpDistinct(node: Rel.Op.Distinct, ctx: Rel.Type?): Rel {
-            TODO("Type RelOp Distinct")
+            val input = visitRel(node.input, ctx)
+            return rel(input.type, relOpDistinct(input))
         }
 
         override fun visitRelOpFilter(node: Rel.Op.Filter, ctx: Rel.Type?): Rel {