Moves ConnectorMetadata to planner session for PartiQLPlanner re-use

CHANGELOG.md

@@ -28,7 +28,8 @@ Thank you to all who have contributed!
 ## [Unreleased]
 
 ### Added
-- Adds the ability to pass a user-defined-function signature to the planner.
+- Adds the ability to define a user-defined-function in ConnectorMetadata
+- Move ConnectorMetadata map from PartiQLPlanner to PartiQLPlanner.Session for planner re-use.
 
 ### Changed
 
@@ -42,7 +43,7 @@ Thank you to all who have contributed!
 
 ### Contributors
 Thank you to all who have contributed!
-- @<your-username>
+- @rchowell
 
 ## [0.14.0-alpha] - 2023-12-15
 

partiql-cli/src/main/kotlin/org/partiql/cli/Main.kt

@@ -52,12 +52,8 @@ object Debug {
 
     private val root = Paths.get(System.getProperty("user.home")).resolve(".partiql/local")
 
-    private val planner = PartiQLPlanner.builder()
-        .catalogs(
-            "local" to LocalConnector.Metadata(root)
-        )
-        .build()
     private val parser = PartiQLParser.default()
+    private val planner = PartiQLPlanner.default()
 
     // !!
     // IMPLEMENT DEBUG BEHAVIOR HERE
@@ -76,6 +72,9 @@ object Debug {
         val sess = PartiQLPlanner.Session(
             queryId = UUID.randomUUID().toString(),
             userId = "debug",
+            catalogs = mapOf(
+                "local" to LocalConnector.Metadata(root)
+            )
         )
         val result = planner.plan(statement, sess).plan
         out.info("-- Plan ----------")

partiql-planner/src/main/kotlin/org/partiql/planner/PartiQLPlanner.kt

@@ -4,6 +4,7 @@ import org.partiql.ast.Statement
 import org.partiql.errors.Problem
 import org.partiql.errors.ProblemCallback
 import org.partiql.plan.PartiQLPlan
+import org.partiql.spi.connector.ConnectorMetadata
 import java.time.Instant
 
 /**
@@ -45,6 +46,7 @@ public interface PartiQLPlanner {
         public val userId: String,
         public val currentCatalog: String? = null,
         public val currentDirectory: List<String> = emptyList(),
+        public val catalogs: Map<String, ConnectorMetadata> = emptyMap(),
         public val instant: Instant = Instant.now(),
     )
 

partiql-planner/src/main/kotlin/org/partiql/planner/PartiQLPlannerBuilder.kt

@@ -1,82 +1,41 @@
 package org.partiql.planner
 
-import org.partiql.spi.connector.ConnectorMetadata
-import org.partiql.types.function.FunctionSignature
-
 /**
  * PartiQLPlannerBuilder is used to programmatically construct a [PartiQLPlanner] implementation.
  *
  * Usage:
  *      PartiQLPlanner.builder()
- *                    .addCatalog("foo", FooConnector())
- *                    .addCatalog("bar", BarConnector())
- *                    .builder()
+ *                    .addPass(myPass)
+ *                    .build()
  */
 public class PartiQLPlannerBuilder {
 
-    private var fns: MutableList<FunctionSignature.Scalar> = mutableListOf()
-    private var catalogs: MutableMap<String, ConnectorMetadata> = mutableMapOf()
-    private var passes: List<PartiQLPlannerPass> = emptyList()
+    private val passes: MutableList<PartiQLPlannerPass> = mutableListOf()
 
     /**
      * Build the builder, return an implementation of a [PartiQLPlanner].
      *
      * @return
      */
-    public fun build(): PartiQLPlanner {
-        val headers = mutableListOf<Header>(PartiQLHeader)
-        if (fns.isNotEmpty()) {
-            val h = object : Header() {
-                override val namespace: String = "UDF"
-                override val functions = fns
-            }
-            headers.add(h)
-        }
-        return PartiQLPlannerDefault(headers, catalogs, passes)
-    }
-
-    /**
-     * Java style method for assigning a Catalog name to [ConnectorMetadata].
-     *
-     * @param catalog
-     * @param metadata
-     * @return
-     */
-    public fun addCatalog(catalog: String, metadata: ConnectorMetadata): PartiQLPlannerBuilder = this.apply {
-        this.catalogs[catalog] = metadata
-    }
-
-    /**
-     * Kotlin style method for assigning Catalog names to [ConnectorMetadata].
-     *
-     * @param catalogs
-     * @return
-     */
-    public fun catalogs(vararg catalogs: Pair<String, ConnectorMetadata>): PartiQLPlannerBuilder = this.apply {
-        this.catalogs = mutableMapOf(*catalogs)
-    }
+    public fun build(): PartiQLPlanner = PartiQLPlannerDefault(passes)
 
     /**
-     * Java style method for adding a user-defined-function.
+     * Java style method for adding a planner pass to this planner builder.
      *
-     * @param function
+     * @param pass
      * @return
      */
-    public fun addFunction(function: FunctionSignature.Scalar): PartiQLPlannerBuilder = this.apply {
-        this.fns.add(function)
+    public fun addPass(pass: PartiQLPlannerPass): PartiQLPlannerBuilder = this.apply {
+        this.passes.add(pass)
     }
 
     /**
-     * Kotlin style method for setting the user-defined functions. This replaces all existing user-defined functions previously passed to the builder.
+     * Kotlin style method for adding a list of planner passes to this planner builder.
      *
-     * @param function
+     * @param passes
      * @return
      */
-    public fun functions(vararg functions: FunctionSignature.Scalar): PartiQLPlannerBuilder = this.apply {
-        this.fns = mutableListOf(*functions)
-    }
-
-    public fun passes(passes: List<PartiQLPlannerPass>): PartiQLPlannerBuilder = this.apply {
-        this.passes = passes
+    public fun addPasses(vararg passes: PartiQLPlannerPass): PartiQLPlannerBuilder = this.apply {
+        this.passes.addAll(passes)
     }
 }

partiql-planner/src/main/kotlin/org/partiql/planner/PartiQLPlannerDefault.kt

@@ -8,14 +8,11 @@ import org.partiql.planner.internal.ir.PartiQLVersion
 import org.partiql.planner.internal.transforms.AstToPlan
 import org.partiql.planner.internal.transforms.PlanTransform
 import org.partiql.planner.internal.typer.PlanTyper
-import org.partiql.spi.connector.ConnectorMetadata
 
 /**
  * Default PartiQL logical query planner.
  */
 internal class PartiQLPlannerDefault(
-    private val headers: List<Header>,
-    private val catalogs: Map<String, ConnectorMetadata>,
     private val passes: List<PartiQLPlannerPass>,
 ) : PartiQLPlanner {
 
@@ -24,8 +21,9 @@ internal class PartiQLPlannerDefault(
         session: PartiQLPlanner.Session,
         onProblem: ProblemCallback,
     ): PartiQLPlanner.Result {
+
         // 0. Initialize the planning environment
-        val env = Env(headers, catalogs, session)
+        val env = Env(session)
 
         // 1. Normalize
         val ast = statement.normalize()

partiql-planner/src/main/kotlin/org/partiql/planner/internal/Env.kt

@@ -1,6 +1,5 @@
 package org.partiql.planner.internal
 
-import org.partiql.planner.Header
 import org.partiql.planner.PartiQLPlanner
 import org.partiql.planner.internal.ir.Agg
 import org.partiql.planner.internal.ir.Catalog
@@ -20,6 +19,7 @@ import org.partiql.spi.connector.ConnectorSession
 import org.partiql.types.StaticType
 import org.partiql.types.StructType
 import org.partiql.types.TupleConstraint
+import org.partiql.types.function.FunctionSignature
 
 /**
  * Handle for associating a catalog name with catalog related metadata objects.
@@ -102,7 +102,7 @@ internal sealed interface ResolvedVar {
         override val type: StaticType,
         override val ordinal: Int,
         override val depth: Int,
-        val position: Int
+        val position: Int,
     ) : ResolvedVar
 }
 
@@ -122,13 +122,9 @@ internal enum class ResolutionStrategy {
 /**
  * PartiQL Planner Global Environment of Catalogs backed by given plugins.
  *
- * @property headers        List of namespaced definitions
- * @property catalogs       List of plugins for global resolution
  * @property session        Session details
  */
 internal class Env(
-    private val headers: List<Header>,
-    private val connectors: Map<String, ConnectorMetadata>,
     private val session: PartiQLPlanner.Session,
 ) {
 
@@ -137,24 +133,44 @@ internal class Env(
      */
     public val catalogs = mutableListOf<Catalog>()
 
+    /**
+     * Catalog Metadata for this query session.
+     */
+    private val connectors = session.catalogs
+
     /**
      * Encapsulate all function resolving logic within [FnResolver].
+     *
+     * TODO we should be using a search_path for resolving functions. This is not possible at the moment, so we flatten
+     *      all builtin functions to live at the top-level. At the moment, we could technically use this to have
+     *      single-level `catalog`.`function`() syntax but that is out-of-scope for this commit.
      */
-    public val fnResolver = FnResolver(headers)
+    public val fnResolver = FnResolver(object : Header() {
+
+        override val namespace: String = "builtins"
+
+        override val functions: List<FunctionSignature.Scalar> =
+            PartiQLHeader.functions + connectors.values.flatMap { it.functions }
+
+        override val operators: List<FunctionSignature.Scalar> =
+            PartiQLHeader.operators + connectors.values.flatMap { it.operators }
+
+        override val aggregations: List<FunctionSignature.Aggregation> =
+            PartiQLHeader.aggregations + connectors.values.flatMap { it.aggregations }
+    })
 
     private val connectorSession = object : ConnectorSession {
         override fun getQueryId(): String = session.queryId
         override fun getUserId(): String = session.userId
     }
 
     /**
-
-     * Leverages a [FunctionResolver] to find a matching function defined in the [Header] scalar function catalog.
+     * Leverages a [FnResolver] to find a matching function defined in the [Header] scalar function catalog.
      */
     internal fun resolveFn(fn: Fn.Unresolved, args: List<Rex>) = fnResolver.resolveFn(fn, args)
 
     /**
-     * Leverages a [FunctionResolver] to find a matching function defined in the [Header] aggregation function catalog.
+     * Leverages a [FnResolver] to find a matching function defined in the [Header] aggregation function catalog.
      */
     internal fun resolveAgg(agg: Agg.Unresolved, args: List<Rex>) = fnResolver.resolveAgg(agg, args)
 
@@ -213,7 +229,11 @@ internal class Env(
             getObjectHandle(cat, catalogPath)?.let { handle ->
                 getObjectDescriptor(handle).let { type ->
                     val depth = calculateMatched(originalPath, catalogPath, handle.second.absolutePath)
-                    val (catalogIndex, valueIndex) = getOrAddCatalogValue(handle.first, handle.second.absolutePath.steps, type)
+                    val (catalogIndex, valueIndex) = getOrAddCatalogValue(
+                        handle.first,
+                        handle.second.absolutePath.steps,
+                        type
+                    )
                     // Return resolution metadata
                     ResolvedVar.Global(type, catalogIndex, depth, valueIndex)
                 }
@@ -224,7 +244,11 @@ internal class Env(
     /**
      * @return a [Pair] where [Pair.first] is the catalog index and [Pair.second] is the value index within that catalog
      */
-    private fun getOrAddCatalogValue(catalogName: String, valuePath: List<String>, valueType: StaticType): Pair<Int, Int> {
+    private fun getOrAddCatalogValue(
+        catalogName: String,
+        valuePath: List<String>,
+        valueType: StaticType,
+    ): Pair<Int, Int> {
         val catalogIndex = getOrAddCatalog(catalogName)
         val symbols = catalogs[catalogIndex].symbols
         return symbols.indexOfFirst { value ->
@@ -388,7 +412,7 @@ internal class Env(
      */
     private class ResolvedPath(
         val replacementPath: BindingPath,
-        val resolvedType: StaticType
+        val resolvedType: StaticType,
     )
 
     /**

partiql-planner/src/main/kotlin/org/partiql/planner/Header.kt → partiql-planner/src/main/kotlin/org/partiql/planner/internal/Header.kt

@@ -1,4 +1,4 @@
-package org.partiql.planner
+package org.partiql.planner.internal
 
 import org.partiql.planner.internal.typer.TypeLattice
 import org.partiql.types.function.FunctionParameter

partiql-planner/src/main/kotlin/org/partiql/planner/PartiQLHeader.kt → partiql-planner/src/main/kotlin/org/partiql/planner/internal/PartiQLHeader.kt

@@ -1,4 +1,4 @@
-package org.partiql.planner
+package org.partiql.planner.internal
 
 import org.partiql.ast.DatetimeField
 import org.partiql.types.function.FunctionParameter

partiql-planner/src/main/kotlin/org/partiql/planner/internal/typer/FnResolver.kt

@@ -1,6 +1,6 @@
 package org.partiql.planner.internal.typer
 
-import org.partiql.planner.Header
+import org.partiql.planner.internal.Header
 import org.partiql.planner.internal.ir.Agg
 import org.partiql.planner.internal.ir.Fn
 import org.partiql.planner.internal.ir.Identifier
@@ -109,7 +109,7 @@ internal sealed class FnMatch<T : FunctionSignature> {
  * at the moment to keep that information (derived from the current TypeLattice) with the [FnResolver].
  */
 @OptIn(PartiQLValueExperimental::class)
-internal class FnResolver(private val headers: List<Header>) {
+internal class FnResolver(private val header: Header) {
 
     /**
      * All headers use the same type lattice (we don't have a design for plugging type systems at the moment).
@@ -140,10 +140,10 @@ internal class FnResolver(private val headers: List<Header>) {
         val (casts, unsafeCasts) = casts()
         unsafeCastSet = unsafeCasts
         // combine all header definitions
-        val fns = headers.flatMap { it.functions }
+        val fns = header.functions
         functions = fns.toFnMap()
-        operators = (headers.flatMap { it.operators } + casts).toFnMap()
-        aggregations = headers.flatMap { it.aggregations }.toFnMap()
+        operators = (header.operators + casts).toFnMap()
+        aggregations = header.aggregations.toFnMap()
     }
 
     /**

partiql-planner/src/test/kotlin/org/partiql/planner/HeaderTest.kt

@@ -2,6 +2,7 @@ package org.partiql.planner
 
 import org.junit.jupiter.api.Disabled
 import org.junit.jupiter.api.Test
+import org.partiql.planner.internal.PartiQLHeader
 
 class HeaderTest {
 

partiql-planner/src/test/kotlin/org/partiql/planner/internal/EnvTest.kt

@@ -3,7 +3,6 @@ package org.partiql.planner.internal
 import org.junit.jupiter.api.Assertions.assertNull
 import org.junit.jupiter.api.BeforeEach
 import org.junit.jupiter.api.Test
-import org.partiql.planner.PartiQLHeader
 import org.partiql.planner.PartiQLPlanner
 import org.partiql.planner.internal.ir.Catalog
 import org.partiql.planner.internal.ir.Rex
@@ -38,15 +37,14 @@ class EnvTest {
     @BeforeEach
     fun init() {
         env = Env(
-            listOf(PartiQLHeader),
-            mapOf(
-                "pql" to LocalConnector.Metadata(root)
-            ),
             PartiQLPlanner.Session(
                 queryId = Random().nextInt().toString(),
                 userId = "test-user",
                 currentCatalog = "pql",
                 currentDirectory = listOf("main"),
+                catalogs = mapOf(
+                    "pql" to LocalConnector.Metadata(root)
+                ),
             )
         )
     }

partiql-planner/src/test/kotlin/org/partiql/planner/internal/typer/FunctionResolverTest.kt

@@ -2,8 +2,7 @@ package org.partiql.planner.internal.typer
 
 import org.junit.jupiter.api.Test
 import org.junit.jupiter.api.fail
-import org.partiql.planner.Header
-import org.partiql.planner.PartiQLHeader
+import org.partiql.planner.internal.Header
 import org.partiql.types.function.FunctionParameter
 import org.partiql.types.function.FunctionSignature
 import org.partiql.value.PartiQLValueExperimental
@@ -67,7 +66,7 @@ class FunctionResolverTest {
             )
         }
 
-        private val resolver = FnResolver(listOf(PartiQLHeader, myHeader))
+        private val resolver = FnResolver(myHeader)
     }
 
     private sealed class Case {