Fixes variable resolution (#1322)

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

@@ -77,16 +77,18 @@ internal sealed interface ResolvedVar {
      *
      * @property type              Resolved StaticType
      * @property ordinal           Index offset in [TypeEnv]
-     * @property replacementSteps  Path steps to replace.
-     * @property depth             The depth/level of the path match.
+     * @property resolvedSteps     The fully resolved path steps.s
      */
     class Local(
         override val type: StaticType,
         override val ordinal: Int,
         val rootType: StaticType,
-        val replacementSteps: List<BindingName>,
-        override val depth: Int
-    ) : ResolvedVar
+        val resolvedSteps: List<BindingName>,
+    ) : ResolvedVar {
+        // the depth are always going to be 1 because this is local variable.
+        // the global path, however the path length maybe, going to be replaced by a binding name.
+        override val depth: Int = 1
+    }
 
     /**
      * Metadata for a resolved global variable
@@ -233,7 +235,7 @@ internal class Env(
                     catalogs[catalogIndex] = catalogs[catalogIndex].copy(
                         symbols = symbols + listOf(Catalog.Symbol(valuePath, valueType))
                     )
-                    catalogIndex to 0
+                    catalogIndex to catalogs[catalogIndex].symbols.lastIndex
                 }
                 else -> {
                     catalogIndex to index
@@ -325,19 +327,25 @@ internal class Env(
         locals.forEachIndexed { ordinal, binding ->
             val root = path.steps[0]
             if (root.isEquivalentTo(binding.name)) {
-                return ResolvedVar.Local(binding.type, ordinal, binding.type, emptyList(), 1)
+                return ResolvedVar.Local(binding.type, ordinal, binding.type, path.steps)
             }
         }
 
         // 2. Check if this variable is referencing a struct field, carrying ordinals
         val matches = mutableListOf<ResolvedVar.Local>()
         for (ordinal in locals.indices) {
             val rootType = locals[ordinal].type
+            val pathPrefix = BindingName(locals[ordinal].name, BindingCase.SENSITIVE)
             if (rootType is StructType) {
                 val varType = inferStructLookup(rootType, path)
                 if (varType != null) {
                     // we found this path within a struct!
-                    val match = ResolvedVar.Local(varType.resolvedType, ordinal, rootType, varType.replacementPath.steps, varType.replacementPath.steps.size)
+                    val match = ResolvedVar.Local(
+                        varType.resolvedType,
+                        ordinal,
+                        rootType,
+                        listOf(pathPrefix) + varType.replacementPath.steps,
+                    )
                     matches.add(match)
                 }
             }
@@ -413,7 +421,7 @@ internal class Env(
                 }
             }
             // 3. Struct is open
-            else -> null
+            else -> key to StaticType.ANY
         }
     }
 

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

@@ -451,25 +451,47 @@ internal class PlanTyper(
                 return rex(ANY, rexOpErr("Undefined variable ${node.identifier}"))
             }
             val type = resolvedVar.type
-            val op = when (resolvedVar) {
-                is ResolvedVar.Global -> rexOpGlobal(catalogSymbolRef(resolvedVar.ordinal, resolvedVar.position))
-                is ResolvedVar.Local -> rexOpVarResolved(resolvedVar.ordinal) // resolvedLocalPath(resolvedVar)
-            }
-            val variable = rex(type, op)
-            return when (resolvedVar.depth) {
-                path.steps.size -> variable
-                else -> {
-                    val foldedPath = path.steps.subList(resolvedVar.depth, path.steps.size).fold(variable) { current, step ->
-                        when (step.bindingCase) {
-                            BindingCase.SENSITIVE -> rex(ANY, rexOpPathKey(current, rex(STRING, rexOpLit(stringValue(step.name)))))
-                            BindingCase.INSENSITIVE -> rex(ANY, rexOpPathSymbol(current, step.name))
+            return when (resolvedVar) {
+                is ResolvedVar.Global -> {
+                    val variable = rex(type, rexOpGlobal(catalogSymbolRef(resolvedVar.ordinal, resolvedVar.position)))
+                    when (resolvedVar.depth) {
+                        path.steps.size -> variable
+                        else -> {
+                            val foldedPath = foldPath(path.steps, resolvedVar.depth, path.steps.size, variable)
+                            visitRex(foldedPath, ctx)
+                        }
+                    }
+                }
+                is ResolvedVar.Local -> {
+                    val variable = rex(type, rexOpVarResolved(resolvedVar.ordinal))
+                    when {
+                        path.isEquivalentTo(resolvedVar.resolvedSteps) && path.steps.size == resolvedVar.depth -> variable
+                        else -> {
+                            val foldedPath = foldPath(resolvedVar.resolvedSteps, resolvedVar.depth, resolvedVar.resolvedSteps.size, variable)
+                            visitRex(foldedPath, ctx)
                         }
                     }
-                    visitRex(foldedPath, ctx)
                 }
             }
         }
 
+        private fun foldPath(path: List<BindingName>, start: Int, end: Int, global: Rex) =
+            path.subList(start, end).fold(global) { current, step ->
+                when (step.bindingCase) {
+                    BindingCase.SENSITIVE -> rex(ANY, rexOpPathKey(current, rex(STRING, rexOpLit(stringValue(step.name)))))
+                    BindingCase.INSENSITIVE -> rex(ANY, rexOpPathSymbol(current, step.name))
+                }
+            }
+
+        private fun BindingPath.isEquivalentTo(other: List<BindingName>): Boolean {
+            this.steps.forEachIndexed { index, bindingName ->
+                if (bindingName != other[index]) {
+                    return false
+                }
+            }
+            return true
+        }
+
         override fun visitRexOpGlobal(node: Rex.Op.Global, ctx: StaticType?): Rex {
             val catalog = env.catalogs[node.ref.catalog]
             val type = catalog.symbols[node.ref.symbol].type

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

@@ -0,0 +1,164 @@
+package org.partiql.planner
+
+import org.junit.jupiter.api.DynamicContainer
+import org.junit.jupiter.api.DynamicContainer.dynamicContainer
+import org.junit.jupiter.api.DynamicNode
+import org.junit.jupiter.api.DynamicTest
+import org.junit.jupiter.api.TestFactory
+import org.partiql.parser.PartiQLParser
+import org.partiql.plan.PlanNode
+import org.partiql.plan.debug.PlanPrinter
+import org.partiql.planner.test.PartiQLTest
+import org.partiql.planner.test.PartiQLTestProvider
+import org.partiql.planner.util.PlanNodeEquivalentVisitor
+import org.partiql.planner.util.ProblemCollector
+import org.partiql.plugins.memory.MemoryConnector
+import org.partiql.types.BagType
+import org.partiql.types.StaticType
+import org.partiql.types.StructType
+import org.partiql.types.TupleConstraint
+import java.io.File
+import java.nio.file.Path
+import java.time.Instant
+import java.util.stream.Stream
+import kotlin.io.path.toPath
+
+// Prevent Unintentional break of the plan
+// We currently don't have a good way to assert on the result plan
+// so we assert on having the partiql text.
+// The input text and the normalized partiql text should produce identical plan.
+// I.e.,
+// if the input text is `SELECT a,b,c FROM T`
+// the produced plan will be identical as the normalized query:
+// `SELECT "T"['a'] AS "a", "T"['b'] AS "b", "T"['c'] AS "c" FROM "default"."T" AS "T";`
+class PlanTest {
+    val root: Path = this::class.java.getResource("/outputs")!!.toURI().toPath()
+
+    val input = PartiQLTestProvider().apply { load() }
+
+    val session: (PartiQLTest.Key) -> PartiQLPlanner.Session = { key ->
+        PartiQLPlanner.Session(
+            queryId = key.toString(),
+            userId = "user_id",
+            currentCatalog = "default",
+            currentDirectory = listOf(),
+            instant = Instant.now(),
+        )
+    }
+
+    val metadata = MemoryConnector.Metadata.of(
+        "default.t" to BagType(
+            StructType(
+                listOf(
+                    StructType.Field("a", StaticType.BOOL),
+                    StructType.Field("b", StaticType.INT4),
+                    StructType.Field("c", StaticType.STRING),
+                    StructType.Field(
+                        "d",
+                        StructType(
+                            listOf(StructType.Field("e", StaticType.STRING)),
+                            contentClosed = true,
+                            emptyList(),
+                            setOf(TupleConstraint.Open(false)),
+                            emptyMap()
+                        )
+                    ),
+                    StructType.Field("x", StaticType.ANY),
+                    StructType.Field("z", StaticType.STRING),
+                    StructType.Field("v", StaticType.STRING),
+                ),
+                contentClosed = true,
+                emptyList(),
+                setOf(TupleConstraint.Open(false)),
+                emptyMap()
+            )
+        )
+    )
+
+    val pipeline: (PartiQLTest) -> PartiQLPlanner.Result = { test ->
+        val problemCollector = ProblemCollector()
+        val ast = PartiQLParser.default().parse(test.statement).root
+        val planner = PartiQLPlannerBuilder()
+            .addCatalog("default", metadata)
+            .build()
+        planner.plan(ast, session(test.key), problemCollector)
+    }
+
+    @TestFactory
+    fun factory(): Stream<DynamicNode> {
+        val r = root.toFile()
+        return r
+            .listFiles { f -> f.isDirectory }!!
+            .mapNotNull { load(r, it) }
+            .stream()
+    }
+
+    private fun load(parent: File, file: File): DynamicNode? = when {
+        file.isDirectory -> loadD(parent, file)
+        file.extension == "sql" -> loadF(parent, file)
+        else -> null
+    }
+
+    private fun loadD(parent: File, file: File): DynamicContainer {
+        val name = file.name
+        val children = file.listFiles()!!.map { load(file, it) }
+        return dynamicContainer(name, children)
+    }
+
+    private fun loadF(parent: File, file: File): DynamicContainer {
+        val group = parent.name
+        val tests = parse(group, file)
+
+        val children = tests.map {
+            // Prepare
+            val displayName = it.key.toString()
+
+            // Assert
+            DynamicTest.dynamicTest(displayName) {
+                val input = input[it.key] ?: error("no test cases")
+
+                val inputPlan = pipeline.invoke(input).plan
+                val outputPlan = pipeline.invoke(it).plan
+                assert(inputPlan.isEquaivalentTo(outputPlan)) {
+                    buildString {
+                        this.appendLine("expect plan equivalence")
+                        PlanPrinter.append(this, inputPlan)
+                        PlanPrinter.append(this, outputPlan)
+                    }
+                }
+            }
+        }
+        return dynamicContainer(file.nameWithoutExtension, children)
+    }
+
+    private fun parse(group: String, file: File): List<PartiQLTest> {
+        val tests = mutableListOf<PartiQLTest>()
+        var name = ""
+        val statement = StringBuilder()
+        for (line in file.readLines()) {
+            // start of test
+            if (line.startsWith("--#[") and line.endsWith("]")) {
+                name = line.substring(4, line.length - 1)
+                statement.clear()
+            }
+            if (name.isNotEmpty() && line.isNotBlank()) {
+                // accumulating test statement
+                statement.appendLine(line)
+            } else {
+                // skip these lines
+                continue
+            }
+            // Finish & Reset
+            if (line.endsWith(";")) {
+                val key = PartiQLTest.Key(group, name)
+                tests.add(PartiQLTest(key, statement.toString()))
+                name = ""
+                statement.clear()
+            }
+        }
+        return tests
+    }
+
+    private fun PlanNode.isEquaivalentTo(other: PlanNode): Boolean =
+        PlanNodeEquivalentVisitor().visit(this, other)
+}