Adds support for planning/transpilation of SELECT *

Adds support for planning/transpilation of SELECT a.* Adds support for planning/transpilation of SELECT a.b.* Adds support for planning/transpilation of SELECT a.*, a.b
partiql · Oct 3, 2023 · 3f353b6 · 3f353b6
1 parent c136f02
commit 3f353b6
Show file tree

Hide file tree

Showing 10 changed files with 1,185 additions and 48 deletions.
diff --git a/lib/partiql-transpiler/src/main/kotlin/org/partiql/transpiler/sql/RexToSql.kt b/lib/partiql-transpiler/src/main/kotlin/org/partiql/transpiler/sql/RexToSql.kt
@@ -3,6 +3,8 @@ package org.partiql.transpiler.sql
 import org.partiql.ast.Ast
 import org.partiql.ast.Expr
 import org.partiql.ast.Identifier
+import org.partiql.ast.Select
+import org.partiql.ast.SetQuantifier
 import org.partiql.plan.Fn
 import org.partiql.plan.PlanNode
 import org.partiql.plan.Rel
@@ -11,6 +13,7 @@ import org.partiql.plan.visitor.PlanBaseVisitor
 import org.partiql.transpiler.TranspilerProblem
 import org.partiql.transpiler.sql.SqlTransform.Companion.translate
 import org.partiql.types.StaticType
+import org.partiql.types.StructType
 import org.partiql.value.PartiQLValueExperimental
 import org.partiql.value.StringValue
 
@@ -77,9 +80,7 @@ public open class RexToSql(
 
     override fun visitRexOpGlobal(node: Rex.Op.Global, ctx: StaticType): Expr {
         val global = transform.getGlobal(node.ref)
-        if (global == null) {
-            error("Malformed plan, resolved global (\$global ${node.ref}) does not exist")
-        }
+            ?: error("Malformed plan, resolved global (\$global ${node.ref}) does not exist")
         val identifier = global
         val scope = Expr.Var.Scope.DEFAULT
         return Ast.exprVar(identifier, scope)
@@ -91,6 +92,23 @@ public open class RexToSql(
         return Ast.exprPath(root, steps)
     }
 
+    // TODO: Fix this.
+    @OptIn(PartiQLValueExperimental::class)
+    override fun visitRexOpTupleUnion(node: Rex.Op.TupleUnion, ctx: StaticType): Expr {
+        return Ast.create {
+            val args = node.args.map { arg ->
+                when (arg) {
+                    is Rex.Op.TupleUnion.Arg.Struct -> visitRex(arg.v, ctx)
+                    is Rex.Op.TupleUnion.Arg.Spread -> visitRex(arg.v, ctx)
+                }
+            }
+            exprCall(
+                identifierSymbol("TUPLEUNION", Identifier.CaseSensitivity.INSENSITIVE),
+                args = args
+            )
+        }
+    }
+
     private fun visitRexOpPathStep(node: Rex.Op.Path.Step): Expr.Path.Step = when (node) {
         is Rex.Op.Path.Step.Index -> visitRexOpPathStepIndex(node)
         is Rex.Op.Path.Step.Unpivot -> Ast.exprPathStepUnpivot()
@@ -151,18 +169,86 @@ public open class RexToSql(
     }
 
     override fun visitRexOpSelect(node: Rex.Op.Select, ctx: StaticType): Expr {
-        // val typeEnv = node.rel.type.schema
+        val typeEnv = node.rel.type.schema
         val relToSql = RelToSql(transform)
-        // val rexToSql = RexToSql(transform, typeEnv)
+        val rexToSql = RexToSql(transform, typeEnv)
         val sfw = relToSql.apply(node.rel)
         assert(sfw.select != null) { "SELECT from RelToSql should never be null" }
-        if (node.constructor.isDefault(node.rel.type.schema)) {
+        return if (node.constructor.isDefault(node.rel.type.schema)) {
             // SELECT
-            return sfw.build()
+            sfw.build()
         } else {
-            // SELECT VALUE
-            // TODO rewrite the constructor replacing variable references with the projected expressions
-            throw UnsupportedOperationException("SELECT VALUE is not supported")
+            // Attempt to convert TUPLEUNION into SELECT. If unable, go ahead with SELECT VALUE.
+            val setq = getSetQuantifier(sfw.select!!)
+            val select = convertTupleUnionToSqlSelect(node.constructor.op, node.rel, setq)
+                ?: Ast.create { selectValue(rexToSql.apply(node.constructor), setq) }
+            sfw.select = select
+            sfw.build()
+        }
+    }
+
+    /**
+     * Grabs the [SetQuantifier] of a [Select].
+     */
+    private fun getSetQuantifier(select: Select): SetQuantifier? = when (select) {
+        is Select.Project -> select.setq
+        is Select.Value -> select.setq
+        is Select.Star -> select.setq
+        is Select.Pivot -> null
+    }
+
+    /**
+     * Attempts to convert the [op] (TUPLEUNION) into the projections that it holds. For this to occur, all tuples in TUPLEUNION
+     * **must** be closed content [StructType]s. We assume that the [input] is always a [Rel.Op.Project].
+     *
+     * Example:
+     * ```
+     * SELECT VALUE TUPLEUNION({ 'a': <INT>, 'b': <DECIMAL> }, { 'c': <INT> }) FROM t
+     * -- Gets converted into:
+     * SELECT a, b, c FROM t
+     * ```
+     *
+     * If unable to convert into SQL-style projections (due to open content structs, non-struct arguments, etc), we
+     * return null.
+     */
+    private fun convertTupleUnionToSqlSelect(op: Rex.Op, input: Rel, setq: SetQuantifier?): Select.Project? {
+        val relProject = input.op as? Rel.Op.Project ?: error("Malformed plan, the top rel should be a project.")
+        if (op !is Rex.Op.TupleUnion) { return null }
+        if (op.args.size != input.type.schema.size) { return null }
+        if (op.args.size != relProject.projections.size) { return null }
+        val projections = relProject.projections.flatMapIndexed { index, project ->
+            val newRexToSql = RexToSql(transform, relProject.input.type.schema)
+            val newExpr = newRexToSql.apply(project)
+            when (val arg = op.args[index]) {
+                is Rex.Op.TupleUnion.Arg.Spread -> {
+                    val structType = project.type as? StructType ?: return null
+                    if (structType.contentClosed.not()) { return null }
+                    structType.fields.map { field ->
+                        Ast.create {
+                            selectProjectItemExpression(
+                                expr = exprPath(newExpr, listOf(exprPathStepSymbol(identifierSymbol(field.key, caseSensitivity = Identifier.CaseSensitivity.INSENSITIVE)))), // TODO: Sensitivity
+                                asAlias = identifierSymbol(field.key, Identifier.CaseSensitivity.INSENSITIVE) // TODO: Sensitivity
+                            )
+                        }
+                    }
+                }
+                is Rex.Op.TupleUnion.Arg.Struct -> {
+                    listOf(
+                        Ast.create {
+                            selectProjectItemExpression(
+                                expr = newExpr,
+                                asAlias = identifierSymbol(arg.k, Identifier.CaseSensitivity.INSENSITIVE) // TODO: Sensitivity
+                            )
+                        }
+                    )
+                }
+            }
+        }
+        return Ast.create {
+            selectProject(
+                items = projections,
+                setq = setq
+            )
         }
     }
 
@@ -179,7 +265,7 @@ public open class RexToSql(
     /**
      * Returns true iff this [Rex] is the default constructor for [schema] derived from an SQL SELECT.
      *
-     * See [RelConverter.defaultConstructor] to see how the default constructor is created.
+     * See [org.partiql.planner.transforms.RelConverter.defaultConstructor] to see how the default constructor is created.
      */
     @OptIn(PartiQLValueExperimental::class)
     private fun Rex.isDefault(schema: List<Rel.Binding>): Boolean {