Merge pull request #1271 from partiql/modeling-paths

Modifies how we model paths in the plan

partiql-lang/src/test/kotlin/org/partiql/lang/planner/transforms/PartiQLSchemaInferencerTests.kt

@@ -35,6 +35,7 @@ import org.partiql.types.BagType
 import org.partiql.types.ListType
 import org.partiql.types.SexpType
 import org.partiql.types.StaticType
+import org.partiql.types.StaticType.Companion.ANY
 import org.partiql.types.StaticType.Companion.BAG
 import org.partiql.types.StaticType.Companion.BOOL
 import org.partiql.types.StaticType.Companion.DATE
@@ -2386,6 +2387,124 @@ class PartiQLSchemaInferencerTests {
                     )
                 )
             ),
+            SuccessTestCase(
+                name = "Pathing into resolved local variable without qualification and with sensitivity",
+                query = """
+                    SELECT address."street" AS s FROM employer;
+                """,
+                catalog = "pql",
+                catalogPath = listOf("main"),
+                expected = BagType(
+                    StructType(
+                        fields = mapOf(
+                            "s" to STRING,
+                        ),
+                        contentClosed = true,
+                        constraints = setOf(
+                            TupleConstraint.Open(false),
+                            TupleConstraint.UniqueAttrs(true),
+                            TupleConstraint.Ordered
+                        )
+                    )
+                )
+            ),
+            SuccessTestCase(
+                name = "Pathing into resolved local variable without qualification and with indexing syntax",
+                query = """
+                    SELECT address['street'] AS s FROM employer;
+                """,
+                catalog = "pql",
+                catalogPath = listOf("main"),
+                expected = BagType(
+                    StructType(
+                        fields = mapOf(
+                            "s" to STRING,
+                        ),
+                        contentClosed = true,
+                        constraints = setOf(
+                            TupleConstraint.Open(false),
+                            TupleConstraint.UniqueAttrs(true),
+                            TupleConstraint.Ordered
+                        )
+                    )
+                )
+            ),
+            SuccessTestCase(
+                name = "Pathing into resolved local variable without qualification and with indexing syntax and fully-qualified FROM",
+                query = """
+                    SELECT e.address['street'] AS s FROM "pql"."main"."employer" AS e;
+                """,
+                expected = BagType(
+                    StructType(
+                        fields = mapOf(
+                            "s" to STRING,
+                        ),
+                        contentClosed = true,
+                        constraints = setOf(
+                            TupleConstraint.Open(false),
+                            TupleConstraint.UniqueAttrs(true),
+                            TupleConstraint.Ordered
+                        )
+                    )
+                )
+            ),
+            ErrorTestCase(
+                name = "Show that we can't use [<string>] to reference a value in a schema. It can only be used on tuples.",
+                query = """
+                    SELECT VALUE 1 FROM "pql"."main"['employer'] AS e;
+                """,
+                expected = BagType(INT4),
+                problemHandler = assertProblemExists {
+                    Problem(
+                        UNKNOWN_PROBLEM_LOCATION,
+                        PlanningProblemDetails.UndefinedVariable("main", true)
+                    )
+                }
+            ),
+            ErrorTestCase(
+                name = "Show that we can't use [<string>] to reference a schema in a catalog. It can only be used on tuples.",
+                query = """
+                    SELECT VALUE 1 FROM "pql"['main']."employer" AS e;
+                """,
+                expected = BagType(INT4),
+                problemHandler = assertProblemExists {
+                    Problem(
+                        UNKNOWN_PROBLEM_LOCATION,
+                        PlanningProblemDetails.UndefinedVariable("pql", true)
+                    )
+                }
+            ),
+            SuccessTestCase(
+                name = "Tuple indexing syntax on literal tuple with literal string key",
+                query = """
+                    { 'aBc': 1, 'AbC': 2.0 }['AbC'];
+                """,
+                expected = DECIMAL
+            ),
+            // This should fail because the Spec says tuple indexing MUST use a literal string or explicit cast.
+            ErrorTestCase(
+                name = "Array indexing syntax on literal tuple with non-literal and non-cast key",
+                query = """
+                    { 'aBc': 1, 'AbC': 2.0 }['Ab' || 'C'];
+                """,
+                expected = MISSING,
+                problemHandler = assertProblemExists {
+                    Problem(
+                        sourceLocation = UNKNOWN_PROBLEM_LOCATION,
+                        details = PlanningProblemDetails.ExpressionAlwaysReturnsNullOrMissing
+                    )
+                }
+            ),
+            // The reason this is ANY is because we do not have support for constant-folding. We don't know what
+            //  CAST('Ab' || 'C' AS STRING) will evaluate to, and therefore, we don't know what the indexing operation
+            //  will return.
+            SuccessTestCase(
+                name = "Tuple indexing syntax on literal tuple with explicit cast key",
+                query = """
+                    { 'aBc': 1, 'AbC': 2.0 }[CAST('Ab' || 'C' AS STRING)];
+                """,
+                expected = ANY
+            ),
         )
 
         @JvmStatic

partiql-plan/src/main/resources/partiql_plan.ion

@@ -78,9 +78,21 @@ rex::{
       steps:  list::[step],
       _: [
         step::[
+          // The key MUST be an integer expression. Ex: a[0], a[1 + 1]
           index::{ key: rex },
-          symbol::{ identifier: '.identifier.symbol' },
+
+          // Case-sensitive lookup. The key MUST be a string expression. Ex: a["b"], a."b", a[CAST(b AS STRING)]
+          key::{ key: rex },
+
+          // Case-insensitive lookup. The key MUST be a literal string. Ex: a.b
+          symbol::{ key: string },
+
+          // For arrays. Ex: a[*]
+          // TODO: Do we need this? According to specification: [1,2,3][*] ⇔ SELECT VALUE v FROM [1, 2, 3] AS v
           wildcard::{},
+
+          // For tuples. Ex: a.*
+          // TODO: Do we need this? According to specification: {'a':1, 'b':2}.* ⇔ SELECT VALUE v FROM UNPIVOT {'a':1, 'b':2} AS v
           unpivot::{},
         ],
       ],

partiql-planner/src/main/kotlin/org/partiql/planner/internal/ir/Nodes.kt

@@ -415,6 +415,7 @@ internal data class Rex(
                     is Symbol -> visitor.visitRexOpPathStepSymbol(this, ctx)
                     is Wildcard -> visitor.visitRexOpPathStepWildcard(this, ctx)
                     is Unpivot -> visitor.visitRexOpPathStepUnpivot(this, ctx)
+                    is Key -> visitor.visitRexOpPathStepKey(this, ctx)
                 }
 
                 internal data class Index(
@@ -436,6 +437,44 @@ internal data class Rex(
                     }
                 }
 
+                /**
+                 * This represents a case-sensitive lookup on a tuple. Ex: a['b'] or a[CAST('a' || 'b' AS STRING)].
+                 * This would normally contain the dot notation for case-sensitive lookup, however, due to
+                 * limitations -- we cannot consolidate these. See [Symbol] for more information.
+                 *
+                 * The main difference is that this does NOT include `a."b"`
+                 */
+                internal data class Key(
+                    @JvmField
+                    internal val key: Rex,
+                ) : Step() {
+                    internal override val children: List<PlanNode> by lazy {
+                        val kids = mutableListOf<PlanNode?>()
+                        kids.add(key)
+                        kids.filterNotNull()
+                    }
+
+                    internal override fun <R, C> accept(visitor: PlanVisitor<R, C>, ctx: C): R =
+                        visitor.visitRexOpPathStepKey(this, ctx)
+
+                    internal companion object {
+                        @JvmStatic
+                        internal fun builder(): RexOpPathStepIndexBuilder = RexOpPathStepIndexBuilder()
+                    }
+                }
+
+                /**
+                 * This represents a lookup on a tuple. We differentiate a [Key] and a [Symbol] at this point in the
+                 * pipeline because we NEED to retain some syntactic knowledge for the following reason: we cannot
+                 * use the syntactic index operation on a schema -- as it is not synonymous with a tuple. In other words,
+                 * `<schema-name>."<value-name>"` is not interchangeable with `<schema-name>['<value-name>']`.
+                 *
+                 * So, in order to temporarily differentiate the `a."b"` from `a['b']` (see [Key]), we need to maintain
+                 * the syntactic difference here. Note that this would potentially be mitigated by typing during the AST to Plan
+                 * transformation.
+                 *
+                 * That being said, this represents a lookup on a tuple such as `a.b` or `a."b"`.
+                 */
                 internal data class Symbol(
                     @JvmField
                     internal val identifier: Identifier.Symbol,

partiql-planner/src/main/kotlin/org/partiql/planner/internal/ir/Plan.kt

@@ -53,6 +53,8 @@ internal fun rexOpPath(root: Rex, steps: List<Rex.Op.Path.Step>): Rex.Op.Path =
 
 internal fun rexOpPathStepIndex(key: Rex): Rex.Op.Path.Step.Index = Rex.Op.Path.Step.Index(key)
 
+internal fun rexOpPathStepKey(key: Rex): Rex.Op.Path.Step.Key = Rex.Op.Path.Step.Key(key)
+
 internal fun rexOpPathStepSymbol(identifier: Identifier.Symbol): Rex.Op.Path.Step.Symbol =
     Rex.Op.Path.Step.Symbol(identifier)
 

partiql-planner/src/main/kotlin/org/partiql/planner/internal/ir/builder/PlanBuilder.kt

@@ -165,6 +165,15 @@ internal class PlanBuilder {
         return builder.build()
     }
 
+    internal fun rexOpPathStepKey(
+        key: Rex? = null,
+        block: RexOpPathStepKeyBuilder.() -> Unit = {},
+    ): Rex.Op.Path.Step.Key {
+        val builder = RexOpPathStepKeyBuilder(key)
+        builder.block()
+        return builder.build()
+    }
+
     internal fun rexOpPathStepSymbol(
         identifier: Identifier.Symbol? = null,
         block: RexOpPathStepSymbolBuilder.() -> Unit = {},

partiql-planner/src/main/kotlin/org/partiql/planner/internal/ir/builder/PlanBuilders.kt

@@ -239,6 +239,16 @@ internal class RexOpPathStepIndexBuilder(
     internal fun build(): Rex.Op.Path.Step.Index = Rex.Op.Path.Step.Index(key = key!!)
 }
 
+internal class RexOpPathStepKeyBuilder(
+    internal var key: Rex? = null,
+) {
+    internal fun key(key: Rex?): RexOpPathStepKeyBuilder = this.apply {
+        this.key = key
+    }
+
+    internal fun build(): Rex.Op.Path.Step.Key = Rex.Op.Path.Step.Key(key = key!!)
+}
+
 internal class RexOpPathStepSymbolBuilder(
     internal var identifier: Identifier.Symbol? = null,
 ) {

partiql-planner/src/main/kotlin/org/partiql/planner/internal/ir/visitor/PlanBaseVisitor.kt

@@ -97,6 +97,7 @@ internal abstract class PlanBaseVisitor<R, C> : PlanVisitor<R, C> {
 
     override fun visitRexOpPathStep(node: Rex.Op.Path.Step, ctx: C): R = when (node) {
         is Rex.Op.Path.Step.Index -> visitRexOpPathStepIndex(node, ctx)
+        is Rex.Op.Path.Step.Key -> visitRexOpPathStepKey(node, ctx)
         is Rex.Op.Path.Step.Symbol -> visitRexOpPathStepSymbol(node, ctx)
         is Rex.Op.Path.Step.Wildcard -> visitRexOpPathStepWildcard(node, ctx)
         is Rex.Op.Path.Step.Unpivot -> visitRexOpPathStepUnpivot(node, ctx)
@@ -105,6 +106,9 @@ internal abstract class PlanBaseVisitor<R, C> : PlanVisitor<R, C> {
     override fun visitRexOpPathStepIndex(node: Rex.Op.Path.Step.Index, ctx: C): R =
         defaultVisit(node, ctx)
 
+    override fun visitRexOpPathStepKey(node: Rex.Op.Path.Step.Key, ctx: C): R =
+        defaultVisit(node, ctx)
+
     override fun visitRexOpPathStepSymbol(node: Rex.Op.Path.Step.Symbol, ctx: C): R =
         defaultVisit(node, ctx)
 

partiql-planner/src/main/kotlin/org/partiql/planner/internal/ir/visitor/PlanVisitor.kt

@@ -62,6 +62,8 @@ internal interface PlanVisitor<R, C> {
 
     fun visitRexOpPathStepIndex(node: Rex.Op.Path.Step.Index, ctx: C): R
 
+    fun visitRexOpPathStepKey(node: Rex.Op.Path.Step.Key, ctx: C): R
+
     fun visitRexOpPathStepSymbol(node: Rex.Op.Path.Step.Symbol, ctx: C): R
 
     fun visitRexOpPathStepWildcard(node: Rex.Op.Path.Step.Wildcard, ctx: C): R

partiql-planner/src/main/kotlin/org/partiql/planner/internal/transforms/PlanTransform.kt

@@ -3,6 +3,10 @@ package org.partiql.planner.internal.transforms
 import org.partiql.errors.ProblemCallback
 import org.partiql.plan.PlanNode
 import org.partiql.plan.partiQLPlan
+import org.partiql.plan.rex
+import org.partiql.plan.rexOpLit
+import org.partiql.plan.rexOpPathStepKey
+import org.partiql.plan.rexOpPathStepSymbol
 import org.partiql.planner.internal.ir.Agg
 import org.partiql.planner.internal.ir.Fn
 import org.partiql.planner.internal.ir.Global
@@ -12,7 +16,9 @@ import org.partiql.planner.internal.ir.Rel
 import org.partiql.planner.internal.ir.Rex
 import org.partiql.planner.internal.ir.Statement
 import org.partiql.planner.internal.ir.visitor.PlanBaseVisitor
+import org.partiql.types.StaticType
 import org.partiql.value.PartiQLValueExperimental
+import org.partiql.value.stringValue
 
 /**
  * This is an internal utility to translate from the internal unresolved plan used for typing to the public plan IR.
@@ -118,10 +124,15 @@ internal object PlanTransform : PlanBaseVisitor<PlanNode, ProblemCallback>() {
             key = visitRex(node.key, ctx),
         )
 
-    override fun visitRexOpPathStepSymbol(node: Rex.Op.Path.Step.Symbol, ctx: ProblemCallback) =
-        org.partiql.plan.Rex.Op.Path.Step.Symbol(
-            identifier = visitIdentifierSymbol(node.identifier, ctx),
-        )
+    @OptIn(PartiQLValueExperimental::class)
+    override fun visitRexOpPathStepSymbol(node: Rex.Op.Path.Step.Symbol, ctx: ProblemCallback) = when (node.identifier.caseSensitivity) {
+        Identifier.CaseSensitivity.SENSITIVE -> rexOpPathStepKey(rex(StaticType.STRING, rexOpLit(stringValue(node.identifier.symbol))))
+        Identifier.CaseSensitivity.INSENSITIVE -> rexOpPathStepSymbol(node.identifier.symbol)
+    }
+
+    override fun visitRexOpPathStepKey(node: Rex.Op.Path.Step.Key, ctx: ProblemCallback): PlanNode = rexOpPathStepKey(
+        key = visitRex(node.key, ctx)
+    )
 
     override fun visitRexOpPathStepWildcard(node: Rex.Op.Path.Step.Wildcard, ctx: ProblemCallback) =
         org.partiql.plan.Rex.Op.Path.Step.Wildcard()
@@ -130,7 +141,7 @@ internal object PlanTransform : PlanBaseVisitor<PlanNode, ProblemCallback>() {
         org.partiql.plan.Rex.Op.Path.Step.Unpivot()
 
     override fun visitRexOpCall(node: Rex.Op.Call, ctx: ProblemCallback) =
-        super.visitRexOpCall(node, ctx) as org.partiql.plan.Rex.Op.Call
+        super.visitRexOpCall(node, ctx) as org.partiql.plan.Rex.Op
 
     override fun visitRexOpCallStatic(node: Rex.Op.Call.Static, ctx: ProblemCallback): org.partiql.plan.Rex.Op {
         val fn = visitFn(node.fn, ctx)

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

@@ -33,6 +33,7 @@ import org.partiql.planner.internal.ir.rexOpCollection
 import org.partiql.planner.internal.ir.rexOpLit
 import org.partiql.planner.internal.ir.rexOpPath
 import org.partiql.planner.internal.ir.rexOpPathStepIndex
+import org.partiql.planner.internal.ir.rexOpPathStepKey
 import org.partiql.planner.internal.ir.rexOpPathStepSymbol
 import org.partiql.planner.internal.ir.rexOpPathStepUnpivot
 import org.partiql.planner.internal.ir.rexOpPathStepWildcard
@@ -46,6 +47,7 @@ import org.partiql.planner.internal.typer.toStaticType
 import org.partiql.types.StaticType
 import org.partiql.types.TimeType
 import org.partiql.value.PartiQLValueExperimental
+import org.partiql.value.StringValue
 import org.partiql.value.boolValue
 import org.partiql.value.int32Value
 import org.partiql.value.int64Value
@@ -145,7 +147,17 @@ internal object RexConverter {
                 when (it) {
                     is Expr.Path.Step.Index -> {
                         val key = visitExprCoerce(it.key, context)
-                        rexOpPathStepIndex(key)
+                        when (val astKey = it.key) {
+                            is Expr.Lit -> when (astKey.value) {
+                                is StringValue -> rexOpPathStepKey(key)
+                                else -> rexOpPathStepIndex(key)
+                            }
+                            is Expr.Cast -> when (astKey.asType is Type.String) {
+                                true -> rexOpPathStepKey(key)
+                                false -> rexOpPathStepIndex(key)
+                            }
+                            else -> rexOpPathStepIndex(key)
+                        }
                     }
                     is Expr.Path.Step.Symbol -> {
                         val identifier = AstToPlan.convert(it.symbol)