[2/2] Parameter match on SqlTypeFamily (#1603)

partiql · Oct 2, 2024 · bb2947f · bb2947f
1 parent da22cf8
commit bb2947f
Show file tree

Hide file tree

Showing 20 changed files with 401 additions and 617 deletions.
diff --git a/partiql-planner/src/main/kotlin/org/partiql/planner/internal/Env.kt b/partiql-planner/src/main/kotlin/org/partiql/planner/internal/Env.kt
@@ -1,7 +1,6 @@
 package org.partiql.planner.internal
 
 import org.partiql.planner.internal.casts.CastTable
-import org.partiql.planner.internal.casts.Coercions
 import org.partiql.planner.internal.ir.Ref
 import org.partiql.planner.internal.ir.Rel
 import org.partiql.planner.internal.ir.Rex
@@ -15,6 +14,7 @@ import org.partiql.planner.internal.ir.rexOpCallDynamicCandidate
 import org.partiql.planner.internal.ir.rexOpCastResolved
 import org.partiql.planner.internal.ir.rexOpVarGlobal
 import org.partiql.planner.internal.typer.CompilerType
+import org.partiql.planner.internal.typer.PlanTyper.Companion.toCType
 import org.partiql.planner.internal.typer.Scope.Companion.toPath
 import org.partiql.spi.catalog.Catalog
 import org.partiql.spi.catalog.Catalogs
@@ -23,7 +23,6 @@ import org.partiql.spi.catalog.Name
 import org.partiql.spi.catalog.Session
 import org.partiql.spi.function.Aggregation
 import org.partiql.types.PType
-import org.partiql.types.PType.Kind
 
 /**
  * [Env] is similar to the database type environment from the PartiQL Specification. This includes resolution of
@@ -265,13 +264,12 @@ internal class Env(private val session: Session) {
     /**
      * Check if this function accepts the exact input argument types. Assume same arity.
      */
-
     private fun Aggregation.matches(args: List<PType>): Boolean {
         val parameters = getParameters()
         for (i in args.indices) {
             val a = args[i]
             val p = parameters[i]
-            if (p.getType().kind != Kind.DYNAMIC && a != p.getType()) return false
+            if (p.getMatch(a) != a) return false
         }
         return true
     }
@@ -286,20 +284,19 @@ internal class Env(private val session: Session) {
         val parameters = getParameters()
         val mapping = arrayOfNulls<Ref.Cast?>(args.size)
         for (i in args.indices) {
-            val arg = args[i]
+            val a = args[i]
             val p = parameters[i]
+            val m = p.getMatch(a)
             when {
-                // 1. Exact match
-                arg == p.getType() -> continue
-                // 2. Match ANY, no coercion needed
-                p.getType().kind == Kind.DYNAMIC -> continue
-                // 3. Check for a coercion
-                else -> when (val coercion = Coercions.get(arg, p.getType())) {
-                    null -> return null // short-circuit
-                    else -> mapping[i] = coercion
-                }
+                m == null -> return null
+                m == a -> continue
+                else -> mapping[i] = coercion(a, m)
             }
         }
         return this to mapping
     }
+
+    private fun coercion(arg: PType, target: PType): Ref.Cast {
+        return Ref.Cast(arg.toCType(), target.toCType(), Ref.Cast.Safety.COERCION, true)
+    }
 }
diff --git a/partiql-planner/src/main/kotlin/org/partiql/planner/internal/FnResolver.kt b/partiql-planner/src/main/kotlin/org/partiql/planner/internal/FnResolver.kt
@@ -1,10 +1,10 @@
 package org.partiql.planner.internal
 
-import org.partiql.planner.internal.casts.Coercions
 import org.partiql.planner.internal.ir.Ref
 import org.partiql.planner.internal.typer.CompilerType
 import org.partiql.planner.internal.typer.PlanTyper.Companion.toCType
 import org.partiql.spi.function.Function
+import org.partiql.types.PType
 import org.partiql.types.PType.Kind
 
 /**
@@ -121,42 +121,33 @@ internal object FnResolver {
         for (i in args.indices) {
             val a = args[i]
             val p = parameters[i]
-            if (a != p.getType()) return false
+            if (p.getMatch(a) != a) return false
         }
         return true
     }
 
     /**
-     * Attempt to match arguments to the parameters; return the implicit casts if necessary.
+     * Attempt to match arguments to the parameters; return the coercions if necessary.
      *
      * @param args
      * @return
      */
     private fun Function.match(args: List<CompilerType>): MatchResult? {
         val parameters = getParameters()
         val mapping = arrayOfNulls<Ref.Cast?>(args.size)
-        var exactInputTypes: Int = 0
+        var exactInputTypes = 0
         for (i in args.indices) {
-            val arg = args[i]
+            val a = args[i]
+            if (a.kind == Kind.UNKNOWN) {
+                continue // skip unknown arguments
+            }
+            // check match
             val p = parameters[i]
+            val m = p.getMatch(a)
             when {
-                // 1. Exact match
-                arg == p.getType() -> {
-                    exactInputTypes++
-                    continue
-                }
-                // 2. Match ANY parameter, no coercion needed
-                p.getType().kind == Kind.DYNAMIC -> continue
-                arg.kind == Kind.UNKNOWN -> continue
-                // 3. Allow for ANY arguments
-                arg.kind == Kind.DYNAMIC -> {
-                    mapping[i] = Ref.Cast(arg, p.getType().toCType(), Ref.Cast.Safety.UNSAFE, true)
-                }
-                // 4. Check for a coercion
-                else -> when (val coercion = Coercions.get(arg, p.getType())) {
-                    null -> return null // short-circuit
-                    else -> mapping[i] = coercion
-                }
+                m == null -> return null // short-circuit
+                m == a -> exactInputTypes++
+                else -> mapping[i] = coercion(a, m)
             }
         }
         return MatchResult(
@@ -165,6 +156,10 @@ internal object FnResolver {
         )
     }
 
+    private fun coercion(arg: PType, target: PType): Ref.Cast {
+        return Ref.Cast(arg.toCType(), target.toCType(), Ref.Cast.Safety.COERCION, true)
+    }
+
     private class MatchResult(
         val match: FnMatch.Static,
         val numberOfExactInputTypes: Int,

diff --git a/partiql-planner/src/main/kotlin/org/partiql/planner/internal/casts/CastTable.kt b/partiql-planner/src/main/kotlin/org/partiql/planner/internal/casts/CastTable.kt
@@ -5,16 +5,13 @@ import org.partiql.planner.internal.ir.Ref.Cast
 import org.partiql.planner.internal.typer.CompilerType
 import org.partiql.types.PType
 import org.partiql.types.PType.Kind
-import org.partiql.value.PartiQLValueExperimental
-import org.partiql.value.PartiQLValueType
 
 /**
  * A place to model type relationships (for now this is to answer CAST inquiries).
  *
  * @property types
  * @property graph      Going with a matrix here (using enum ordinals) as it's simple and avoids walking.
  */
-@OptIn(PartiQLValueExperimental::class)
 internal class CastTable private constructor(
     private val types: Array<Kind>,
     private val graph: Array<Array<Status>>,
@@ -29,8 +26,6 @@ internal class CastTable private constructor(
         }
     }
 
-    private operator fun <T> Array<T>.get(t: PartiQLValueType): T = get(t.ordinal)
-
     /**
      * This represents the Y, M, and N in the table listed in SQL:1999 Section 6.22.
      */
@@ -42,9 +37,7 @@ internal class CastTable private constructor(
 
     companion object {
 
-        private val N = Kind.values().size
-
-        private operator fun <T> Array<T>.set(t: PartiQLValueType, value: T): Unit = this.set(t.ordinal, value)
+        private val N = Kind.entries.size
 
         private fun relationships(block: RelationshipBuilder.() -> Unit): Array<Status> {
             return with(RelationshipBuilder()) {
@@ -60,15 +53,15 @@ internal class CastTable private constructor(
          */
         @JvmStatic
         val partiql: CastTable = run {
-            val types = Kind.values()
+            val types = Kind.entries.toTypedArray()
             val graph = arrayOfNulls<Array<Status>>(N)
             for (type in types) {
                 // initialize all with empty relationships
                 graph[type.ordinal] = Array(N) { Status.NO }
             }
             graph[Kind.DYNAMIC.ordinal] = relationships {
                 cast(Kind.DYNAMIC)
-                Kind.values().filterNot { it == Kind.DYNAMIC }.forEach {
+                Kind.entries.filterNot { it == Kind.DYNAMIC }.forEach {
                     cast(it)
                 }
             }