Merge pull request #1233 from partiql/type-subquery

Adds support for subquery typing in inferencer

CHANGELOG.md

@@ -40,14 +40,16 @@ Thank you to all who have contributed!
 ### Deprecated
 
 ### Fixed
+- Fixes typing of scalar subqueries in the PartiQLSchemaInferencer. Note that usage of `SELECT *` in subqueries
+  is not fully supported. Please make sure to handle InferenceExceptions.
 
 ### Removed
 
 ### Security
 
 ### Contributors
 Thank you to all who have contributed!
-- @<your-username>
+- @johnedquinn
 
 ## [0.13.1] - 2023-09-19
 

partiql-lang/src/main/kotlin/org/partiql/lang/ast/passes/SemanticProblemDetails.kt

@@ -22,6 +22,13 @@ sealed class SemanticProblemDetails(override val severity: ProblemSeverity, val
             }
         )
 
+    data class CoercionError(val actualType: StaticType) : SemanticProblemDetails(
+        severity = ProblemSeverity.ERROR,
+        messageFormatter = {
+            "Unable to coerce $actualType into a single value."
+        }
+    )
+
     object DuplicateAliasesInSelectListItem :
         SemanticProblemDetails(
             severity = ProblemSeverity.ERROR,

partiql-lang/src/main/kotlin/org/partiql/lang/planner/transforms/plan/PlanTyper.kt

@@ -334,12 +334,32 @@ internal object PlanTyper : PlanRewriter<PlanTyper.Context>() {
 
     override fun visitRexQueryScalarSubquery(node: Rex.Query.Scalar.Subquery, ctx: Context): PlanNode {
         val query = visitRex(node.query, ctx) as Rex.Query.Collection
-        when (val queryType = query.grabType() ?: handleMissingType(ctx)) {
+        // If it is SELECT VALUE, do not coerce.
+        if (query.constructor != null) {
+            val type = query.type as? CollectionType
+            return node.copy(query = query, type = type?.elementType?.flatten())
+        }
+        val type = when (val queryType = query.grabType() ?: handleMissingType(ctx)) {
             is CollectionType -> queryType.elementType
             else -> error("Query collection subqueries should always return a CollectionType.")
         }
+        val resultType = when (type) {
+            is StructType -> {
+                if (StaticTypeUtils.isClosedSafe(type) == true && type.fields.size == 1) {
+                    type.fields[0].value
+                } else {
+                    handleCoercionError(ctx, type)
+                    StaticType.ANY
+                }
+            }
+            else -> {
+                handleCoercionError(ctx, type)
+                StaticType.ANY
+            }
+        }
         return node.copy(
-            query = query
+            query = query,
+            type = resultType.flatten()
         )
     }
 
@@ -1313,7 +1333,7 @@ internal object PlanTyper : PlanRewriter<PlanTyper.Context>() {
         return requiredArgumentsMatch
     }
 
-    private fun Rex.isProjectAll(): Boolean {
+    internal fun Rex.isProjectAll(): Boolean {
         return when (this) {
             is Rex.Path -> {
                 val step = this.steps.lastOrNull() ?: return false
@@ -1678,4 +1698,13 @@ internal object PlanTyper : PlanRewriter<PlanTyper.Context>() {
             )
         )
     }
+
+    private fun handleCoercionError(ctx: Context, actualType: StaticType) {
+        ctx.problemHandler.handleProblem(
+            Problem(
+                sourceLocation = UNKNOWN_PROBLEM_LOCATION,
+                details = SemanticProblemDetails.CoercionError(actualType)
+            )
+        )
+    }
 }

partiql-lang/src/main/kotlin/org/partiql/lang/planner/transforms/plan/RexConverter.kt

@@ -9,10 +9,13 @@ import org.partiql.lang.eval.builtins.ExprFunctionCurrentUser
 import org.partiql.lang.eval.err
 import org.partiql.lang.eval.errorContextFrom
 import org.partiql.lang.planner.transforms.AstToPlan
+import org.partiql.lang.planner.transforms.plan.PlanTyper.isProjectAll
 import org.partiql.plan.Case
 import org.partiql.plan.Plan
+import org.partiql.plan.Rel
 import org.partiql.plan.Rex
 import org.partiql.types.StaticType
+import java.util.Locale
 
 /**
  * Some workarounds for transforming a PIG tree without having to create another visitor:
@@ -109,7 +112,7 @@ internal object RexConverter : PartiqlAst.VisitorFold<RexConverter.Ctx>() {
     }
 
     override fun walkExprSessionAttribute(node: PartiqlAst.Expr.SessionAttribute, accumulator: Ctx) = visit(node) {
-        val functionName = when (node.value.text.toUpperCase()) {
+        val functionName = when (node.value.text.uppercase(Locale.getDefault())) {
             EvaluationSession.Constants.CURRENT_USER_KEY -> ExprFunctionCurrentUser.FUNCTION_NAME
             else -> err(
                 "Unsupported session attribute: ${node.value.text}",
@@ -262,59 +265,118 @@ internal object RexConverter : PartiqlAst.VisitorFold<RexConverter.Ctx>() {
     }
 
     override fun walkExprEq(node: PartiqlAst.Expr.Eq, ctx: Ctx) = visit(node) {
+        val (lhs, rhs) = walkComparisonOperands(node.operands)
         Plan.rexBinary(
-            lhs = convert(node.operands[0]),
-            rhs = convert(node.operands[1]),
+            lhs = lhs,
+            rhs = rhs,
             op = Rex.Binary.Op.EQ,
             type = StaticType.BOOL,
         )
     }
 
     override fun walkExprNe(node: PartiqlAst.Expr.Ne, ctx: Ctx) = visit(node) {
+        val (lhs, rhs) = walkComparisonOperands(node.operands)
         Plan.rexBinary(
-            lhs = convert(node.operands[0]),
-            rhs = convert(node.operands[1]),
+            lhs = lhs,
+            rhs = rhs,
             op = Rex.Binary.Op.NEQ,
             type = StaticType.BOOL,
         )
     }
 
     override fun walkExprGt(node: PartiqlAst.Expr.Gt, ctx: Ctx) = visit(node) {
+        val (lhs, rhs) = walkComparisonOperands(node.operands)
         Plan.rexBinary(
-            lhs = convert(node.operands[0]),
-            rhs = convert(node.operands[1]),
+            lhs = lhs,
+            rhs = rhs,
             op = Rex.Binary.Op.GT,
             type = StaticType.BOOL,
         )
     }
 
     override fun walkExprGte(node: PartiqlAst.Expr.Gte, ctx: Ctx) = visit(node) {
+        val (lhs, rhs) = walkComparisonOperands(node.operands)
         Plan.rexBinary(
-            lhs = convert(node.operands[0]),
-            rhs = convert(node.operands[1]),
+            lhs = lhs,
+            rhs = rhs,
             op = Rex.Binary.Op.GTE,
             type = StaticType.BOOL,
         )
     }
 
     override fun walkExprLt(node: PartiqlAst.Expr.Lt, ctx: Ctx) = visit(node) {
+        val (lhs, rhs) = walkComparisonOperands(node.operands)
         Plan.rexBinary(
-            lhs = convert(node.operands[0]),
-            rhs = convert(node.operands[1]),
+            lhs = lhs,
+            rhs = rhs,
             op = Rex.Binary.Op.LT,
             type = StaticType.BOOL,
         )
     }
 
     override fun walkExprLte(node: PartiqlAst.Expr.Lte, ctx: Ctx) = visit(node) {
+        val (lhs, rhs) = walkComparisonOperands(node.operands)
         Plan.rexBinary(
-            lhs = convert(node.operands[0]),
-            rhs = convert(node.operands[1]),
+            lhs = lhs,
+            rhs = rhs,
             op = Rex.Binary.Op.LTE,
             type = StaticType.BOOL,
         )
     }
 
+    /**
+     * Converts Comparison Operands. Also coerces them if one of them is an array.
+     */
+    private fun walkComparisonOperands(operands: List<PartiqlAst.Expr>): Pair<Rex, Rex> {
+        var lhs = convert(operands[0])
+        var rhs = convert(operands[1])
+        if (lhs is Rex.Collection.Array) { rhs = coercePotentialSubquery(rhs) }
+        if (rhs is Rex.Collection.Array) { lhs = coercePotentialSubquery(lhs) }
+        return lhs to rhs
+    }
+
+    /**
+     * We convert the scalar subquery of a SFW into a scalar subquery of a SELECT VALUE.
+     */
+    private fun coercePotentialSubquery(rex: Rex): Rex {
+        var rhs = rex
+        if (rhs is Rex.Query.Scalar.Subquery) {
+            val sfw = rhs.query as? Rex.Query.Collection ?: error("Malformed plan, all scalar subqueries should hold collection queries")
+            val constructor = sfw.constructor ?: run {
+                val relProject = sfw.rel as? Rel.Project ?: error("Malformed plan, the top of a plan should be a projection")
+                getConstructorFromRelProject(relProject)
+            }
+            rhs = rhs.copy(
+                query = rhs.query.copy(
+                    constructor = constructor
+                )
+            )
+        }
+        return rhs
+    }
+
+    private fun getConstructorFromRelProject(relProject: Rel.Project): Rex {
+        return when (relProject.bindings.size) {
+            0 -> error("The Projection should not have held empty bindings.")
+            1 -> {
+                val binding = relProject.bindings.first()
+                if (binding.value.isProjectAll()) {
+                    error("Unimplemented feature: coercion of SELECT *.")
+                }
+                relProject.bindings.first().value
+            }
+            else -> {
+                if (relProject.bindings.any { it.value.isProjectAll() }) {
+                    error("Unimplemented feature: coercion of SELECT *.")
+                }
+                Plan.rexCollectionArray(
+                    relProject.bindings.map { it.value },
+                    type = StaticType.LIST
+                )
+            }
+        }
+    }
+
     override fun walkExprLike(node: PartiqlAst.Expr.Like, ctx: Ctx) = visit(node) {
         when (val escape = node.escape) {
             null -> Plan.rexCall(
@@ -345,12 +407,14 @@ internal object RexConverter : PartiqlAst.VisitorFold<RexConverter.Ctx>() {
         )
     }
 
+    /**
+     * Here, we must visit the RHS. If it is a scalar subquery, we need to grab the underlying collection.
+     */
     override fun walkExprInCollection(node: PartiqlAst.Expr.InCollection, ctx: Ctx) = visit(node) {
         val lhs = convert(node.operands[0])
-        var rhs = convert(node.operands[1])
-        if (rhs is Rex.Query.Scalar.Subquery) {
-            rhs = rhs.query // unpack a scalar subquery coercion
-        }
+        val potentialSubqueryRex = convert(node.operands[1])
+        val potentialSubquery = coercePotentialSubquery(potentialSubqueryRex)
+        val rhs = (potentialSubquery as? Rex.Query.Scalar.Subquery)?.query ?: potentialSubquery
         Plan.rexCall(
             id = Constants.inCollection,
             args = listOf(

partiql-lang/src/main/kotlin/org/partiql/lang/types/StaticTypeUtils.kt

@@ -227,6 +227,22 @@ public object StaticTypeUtils {
     //
     //
 
+    internal fun isClosedSafe(type: StructType): Boolean? = try {
+        isClosed(type)
+    } catch (_: StructClosedMismatch) {
+        null
+    }
+
+    @kotlin.jvm.Throws(StructClosedMismatch::class)
+    internal fun isClosed(type: StructType): Boolean {
+        val contentClosed = type.contentClosed
+        val constraintClosed = type.constraints.contains(TupleConstraint.Open(false))
+        if (contentClosed != constraintClosed) { throw StructClosedMismatch }
+        return contentClosed
+    }
+
+    internal object StructClosedMismatch : Exception()
+
     /**
      * By far, structs have the most complicated logic behind their instance check.
      *