Add EXCLUDE to the parser, ast, plan, and plan schema inferencer

CHANGELOG.md

@@ -34,13 +34,16 @@ Thank you to all who have contributed!
 - Support parsing, planning, and evaluation of Bitwise AND operator (&).
   - The Bitwise And Operator only works for integer operands.
   - The operator precedence may change based on the pending operator precedence [RFC](https://github.com/partiql/partiql-docs/issues/50).
+- **EXPERIMENTAL** Adds `EXCLUDE` to parser, ast, plan, and plan schema inferencer
+  - This feature is marked as experimental until an RFC is added https://github.com/partiql/partiql-spec/issues/39
+  - NOTE: this feature is not currently implemented in the evaluator
 
 ### Changed
 
 ### Deprecated
 
 ### Fixed
-
+- Fixes schema inferencer behavior for ORDER BY
 ### Removed
 
 ### Security

partiql-ast/src/main/kotlin/org/partiql/ast/helpers/ToLegacyAst.kt

@@ -17,6 +17,7 @@ import com.amazon.ionelement.api.ionSymbol
 import com.amazon.ionelement.api.metaContainerOf
 import org.partiql.ast.AstNode
 import org.partiql.ast.DatetimeField
+import org.partiql.ast.Exclude
 import org.partiql.ast.Expr
 import org.partiql.ast.From
 import org.partiql.ast.GraphMatch
@@ -665,14 +666,15 @@ private class AstTranslator(val metas: Map<String, MetaContainer>) : AstBaseVisi
         }
         val project = visitSelect(node.select, ctx)
         val from = visitFrom(node.from, ctx)
+        val exclude = node.exclude?.let { visitExclude(it, ctx) }
         val fromLet = node.let?.let { visitLet(it, ctx) }
         val where = node.where?.let { visitExpr(it, ctx) }
         val groupBy = node.groupBy?.let { visitGroupBy(it, ctx) }
         val having = node.having?.let { visitExpr(it, ctx) }
         val orderBy = node.orderBy?.let { visitOrderBy(it, ctx) }
         val limit = node.limit?.let { visitExpr(it, ctx) }
         val offset = node.offset?.let { visitExpr(it, ctx) }
-        select(setq, project, from, fromLet, where, groupBy, having, orderBy, limit, offset, metas)
+        select(setq, project, exclude, from, fromLet, where, groupBy, having, orderBy, limit, offset, metas)
     }
 
     /**
@@ -750,6 +752,48 @@ private class AstTranslator(val metas: Map<String, MetaContainer>) : AstBaseVisi
         join(type, lhs, rhs, condition, metas)
     }
 
+    override fun visitExclude(node: Exclude, ctx: Ctx): PartiqlAst.ExcludeOp = translate(node) { metas ->
+        val excludeExprs = node.exprs.translate<PartiqlAst.ExcludeExpr>(ctx)
+        excludeOp(excludeExprs, metas)
+    }
+
+    override fun visitExcludeExcludeExpr(node: Exclude.ExcludeExpr, ctx: Ctx) = translate(node) { metas ->
+        val root = visitIdentifierSymbol(node.root, ctx)
+        val steps = node.steps.translate<PartiqlAst.ExcludeStep>(ctx)
+        excludeExpr(root = root, steps = steps, metas)
+    }
+
+    override fun visitExcludeStep(node: Exclude.Step, ctx: Ctx) =
+        super.visitExcludeStep(node, ctx) as PartiqlAst.ExcludeStep
+
+    override fun visitExcludeStepExcludeTupleAttr(node: Exclude.Step.ExcludeTupleAttr, ctx: Ctx) = translate(node) { metas ->
+        val attr = node.symbol.symbol
+        val case = node.symbol.caseSensitivity.toLegacyCaseSensitivity()
+        excludeTupleAttr(identifier(attr, case), metas)
+    }
+
+    override fun visitExcludeStepExcludeCollectionIndex(
+        node: Exclude.Step.ExcludeCollectionIndex,
+        ctx: Ctx
+    ) = translate(node) { metas ->
+        val index = node.index.toLong()
+        excludeCollectionIndex(index, metas)
+    }
+
+    override fun visitExcludeStepExcludeTupleWildcard(
+        node: Exclude.Step.ExcludeTupleWildcard,
+        ctx: Ctx
+    ) = translate(node) { metas ->
+        excludeTupleWildcard(metas)
+    }
+
+    override fun visitExcludeStepExcludeCollectionWildcard(
+        node: Exclude.Step.ExcludeCollectionWildcard,
+        ctx: Ctx
+    ) = translate(node) { metas ->
+        excludeCollectionWildcard(metas)
+    }
+
     override fun visitLet(node: Let, ctx: Ctx) = translate(node) { metas ->
         val bindings = node.bindings.translate<PartiqlAst.LetBinding>(ctx)
         let(bindings, metas)

partiql-ast/src/main/pig/partiql.ion

@@ -167,6 +167,7 @@ may then be further optimized by selecting better implementations of each operat
             (select
                 (setq (? set_quantifier))
                 (project projection)
+                (exclude_clause (? exclude_op))
                 (from from_source)
                 (from_let (? let))
                 (where (? expr))
@@ -220,6 +221,22 @@ may then be further optimized by selecting better implementations of each operat
             // For `<expr> [AS <id>]`
             (project_expr expr::expr as_alias::(? symbol)))
 
+        // EXCLUDE clause
+        (product exclude_op exprs::(* exclude_expr 1))
+
+        (product exclude_expr root::identifier steps::(* exclude_step 1))
+
+        (sum exclude_step
+            // `someRoot.someField` case sensitivity depends on if `someField` is quoted
+            // `someRoot[<string literal>] is equivalent to `someRoot."<string literal>"` (case-sensitive)
+            (exclude_tuple_attr attr::identifier)
+            // `someRoot[<int literal>]`
+            (exclude_collection_index index::int)
+            // `someRoot[*]`]
+            (exclude_tuple_wildcard)
+            // `someRoot.*`
+            (exclude_collection_wildcard))
+
         // A list of LET bindings
         (product let let_bindings::(* let_binding 1))
 

partiql-ast/src/main/resources/partiql_ast.ion

@@ -494,6 +494,7 @@ expr::[
   // The PartiQL `<sfw>` query expression, think SQL `<query specification>`
   s_f_w::{
     select:   select,              // oneof SELECT / SELECT VALUE / PIVOT
+    exclude:  optional::exclude,
     from:     from,
     let:      optional::let,
     where:    optional::expr,
@@ -561,6 +562,22 @@ select::[
   },
 ]
 
+exclude::{
+  exprs: list::[exclude_expr],
+  _: [
+    exclude_expr::{
+      root: '.identifier.symbol',
+      steps: list::[step],
+    },
+    step::[
+      exclude_tuple_attr::{ symbol: '.identifier.symbol' },
+      exclude_collection_index::{ index: int },
+      exclude_tuple_wildcard::{},
+      exclude_collection_wildcard::{},
+    ]
+  ]
+}
+
 // PartiQL FROM Clause Variants — https://partiql.org/dql/from.html
 from::[
 

partiql-lang/src/main/kotlin/org/partiql/lang/eval/visitors/GroupByPathExpressionVisitorTransform.kt

@@ -102,6 +102,7 @@ class GroupByPathExpressionVisitorTransform(
         val projection = currentAndUnshadowedTransformer.transformExprSelect_project(node)
 
         // The scope of the expressions in the FROM clause is the same as that of the parent scope.
+        val exclude = unshadowedTransformer.transformExprSelect_excludeClause(node)
         val from = this.transformExprSelect_from(node)
         val fromLet = unshadowedTransformer.transformExprSelect_fromLet(node)
         val where = unshadowedTransformer.transformExprSelect_where(node)
@@ -116,6 +117,7 @@ class GroupByPathExpressionVisitorTransform(
             PartiqlAst.Expr.Select(
                 setq = node.setq,
                 project = projection,
+                excludeClause = exclude,
                 from = from,
                 fromLet = fromLet,
                 where = where,

partiql-lang/src/main/kotlin/org/partiql/lang/eval/visitors/VisitorTransformBase.kt

@@ -36,6 +36,7 @@ abstract class VisitorTransformBase : PartiqlAst.VisitorTransform() {
      * infinite recursion.
      */
     fun transformExprSelectEvaluationOrder(node: PartiqlAst.Expr.Select): PartiqlAst.Expr {
+        val exclude = transformExprSelect_excludeClause(node)
         val from = transformExprSelect_from(node)
         val fromLet = transformExprSelect_fromLet(node)
         val where = transformExprSelect_where(node)
@@ -51,6 +52,7 @@ abstract class VisitorTransformBase : PartiqlAst.VisitorTransform() {
             PartiqlAst.Expr.Select(
                 setq = setq,
                 project = project,
+                excludeClause = exclude,
                 from = from,
                 fromLet = fromLet,
                 where = where,

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

@@ -37,8 +37,11 @@ import org.partiql.lang.types.TypedOpParameter
 import org.partiql.lang.types.UnknownArguments
 import org.partiql.lang.util.cartesianProduct
 import org.partiql.plan.Arg
+import org.partiql.plan.Attribute
 import org.partiql.plan.Binding
 import org.partiql.plan.Case
+import org.partiql.plan.ExcludeExpr
+import org.partiql.plan.ExcludeStep
 import org.partiql.plan.Plan
 import org.partiql.plan.PlanNode
 import org.partiql.plan.Property
@@ -155,6 +158,150 @@ internal object PlanTyper : PlanRewriter<PlanTyper.Context>() {
         )
     }
 
+    /**
+     * Initial implementation of `EXCLUDE` schema inference. Until an RFC is finalized for `EXCLUDE`
+     * (https://github.com/partiql/partiql-spec/issues/39), this behavior is considered experimental and subject to
+     * change.
+     *
+     * There are still discussion points regarding the following edge cases
+     * - EXCLUDE on a tuple attribute that doesn't exist -- give an error/warning?
+     *   - currently no error
+     * - EXCLUDE on a tuple attribute that has duplicates -- give an error/warning? exclude one? exclude both?
+     *   - currently excludes both w/ no error
+     * - EXCLUDE on a collection index as the last step -- mark element type as optional?
+     *   - currently element type as-is
+     * - EXCLUDE on a collection index w/ remaining path steps -- mark last step's type as optional?
+     *   - currently marks last step's type as optional
+     * - EXCLUDE on a binding tuple variable (e.g. SELECT ... EXCLUDE t FROM t) -- error?
+     *   - currently a parser error
+     * - EXCLUDE on a union type -- give an error/warning? no-op? exclude on each type in union?
+     *   - currently exclude on each union type
+     * - If SELECT list includes an attribute that is excluded, we could consider giving an error in PlanTyper or
+     * some other semantic pass
+     *   - currently does not give an error
+     */
+    override fun visitRelExclude(node: Rel.Exclude, ctx: Context): Rel.Exclude {
+        val input = visitRel(node.input, ctx)
+        val exprs = node.exprs
+        val typeEnv = input.getTypeEnv()
+        val newTypeEnv = exprs.fold(typeEnv) { tEnv, expr ->
+            excludeExpr(tEnv, expr, ctx)
+        }
+        return node.copy(
+            input = input,
+            common = node.common.copy(
+                typeEnv = newTypeEnv,
+                properties = input.getProperties()
+            )
+        )
+    }
+
+    private fun excludeExpr(attrs: List<Attribute>, expr: ExcludeExpr, ctx: Context): List<Attribute> {
+        val resultAttrs = mutableListOf<Attribute>()
+        attrs.forEach { attr ->
+            val rootId = expr.root
+            if (attr.name == rootId || (expr.rootCase == Case.INSENSITIVE && attr.name.equals(expr.root, ignoreCase = true))) {
+                if (expr.steps.isEmpty()) {
+                    throw IllegalStateException("Empty `ExcludeExpr.steps` encountered. This should have been caught by the parser.")
+                } else {
+                    val newType = excludeExprSteps(attr.type, expr.steps, lastStepAsOptional = false, ctx)
+                    resultAttrs.add(
+                        attr.copy(
+                            type = newType
+                        )
+                    )
+                }
+            } else {
+                resultAttrs.add(
+                    attr
+                )
+            }
+        }
+        return resultAttrs
+    }
+
+    private fun excludeExprSteps(type: StaticType, steps: List<ExcludeStep>, lastStepAsOptional: Boolean, ctx: Context): StaticType {
+        fun excludeExprStepsStruct(s: StructType, steps: List<ExcludeStep>, lastStepAsOptional: Boolean): StaticType {
+            val outputFields = mutableListOf<StructType.Field>()
+            val first = steps.first()
+            s.fields.forEach { field ->
+                when (first) {
+                    is ExcludeStep.TupleAttr -> {
+                        if (field.key == first.attr || (first.case == Case.INSENSITIVE && field.key.equals(first.attr, ignoreCase = true))) {
+                            if (steps.size == 1) {
+                                if (lastStepAsOptional) {
+                                    val newField = StructType.Field(field.key, field.value.asOptional())
+                                    outputFields.add(newField)
+                                }
+                            } else {
+                                outputFields.add(StructType.Field(field.key, excludeExprSteps(field.value, steps.drop(1), lastStepAsOptional, ctx)))
+                            }
+                        } else {
+                            outputFields.add(field)
+                        }
+                    }
+                    is ExcludeStep.TupleWildcard -> {
+                        if (steps.size == 1) {
+                            if (lastStepAsOptional) {
+                                val newField = StructType.Field(field.key, field.value.asOptional())
+                                outputFields.add(newField)
+                            }
+                        } else {
+                            outputFields.add(StructType.Field(field.key, excludeExprSteps(field.value, steps.drop(1), lastStepAsOptional, ctx)))
+                        }
+                    }
+                    else -> {
+                        // currently no change to field.value and no error thrown; could consider an error/warning in
+                        // the future
+                        outputFields.add(StructType.Field(field.key, field.value))
+                    }
+                }
+            }
+            return s.copy(fields = outputFields)
+        }
+
+        fun excludeExprStepsCollection(c: CollectionType, steps: List<ExcludeStep>, lastStepAsOptional: Boolean): StaticType {
+            var elementType = c.elementType
+            when (steps.first()) {
+                is ExcludeStep.CollectionIndex -> {
+                    if (steps.size > 1) {
+                        elementType = excludeExprSteps(elementType, steps.drop(1), lastStepAsOptional = true, ctx)
+                    }
+                }
+                is ExcludeStep.CollectionWildcard -> {
+                    if (steps.size > 1) {
+                        elementType =
+                            excludeExprSteps(elementType, steps.drop(1), lastStepAsOptional = lastStepAsOptional, ctx)
+                    }
+                    // currently no change to elementType if collection wildcard is last element; this behavior could
+                    // change based on RFC definition
+                }
+                else -> {
+                    // currently no change to elementType and no error thrown; could consider an error/warning in
+                    // the future
+                }
+            }
+            return when (c) {
+                is BagType -> c.copy(elementType)
+                is ListType -> c.copy(elementType)
+                is SexpType -> c.copy(elementType)
+            }
+        }
+
+        return when (type) {
+            is StructType -> excludeExprStepsStruct(type, steps, lastStepAsOptional)
+            is CollectionType -> excludeExprStepsCollection(type, steps, lastStepAsOptional)
+            is AnyOfType -> {
+                StaticType.unionOf(
+                    type.types.map {
+                        excludeExprSteps(it, steps, lastStepAsOptional, ctx)
+                    }.toSet()
+                )
+            }
+            else -> type
+        }.flatten()
+    }
+
     override fun visitRelUnpivot(node: Rel.Unpivot, ctx: Context): Rel.Unpivot {
         val from = node
 
@@ -227,7 +374,8 @@ internal object PlanTyper : PlanRewriter<PlanTyper.Context>() {
         return node.copy(
             input = input,
             common = node.common.copy(
-                typeEnv = typeEnv
+                typeEnv = typeEnv,
+                properties = input.getProperties()
             )
         )
     }
@@ -1006,6 +1154,7 @@ internal object PlanTyper : PlanRewriter<PlanTyper.Context>() {
         is Rel.Scan -> this.common
         is Rel.Sort -> this.common
         is Rel.Unpivot -> this.common
+        is Rel.Exclude -> this.common
     }
 
     private fun inferPathComponentExprType(

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

@@ -34,6 +34,7 @@ internal object PlanUtils {
         is Rel.Scan -> input.common.typeEnv
         is Rel.Sort -> input.common.typeEnv
         is Rel.Unpivot -> input.common.typeEnv
+        is Rel.Exclude -> input.common.typeEnv
     }
 
     internal fun Rex.addType(type: StaticType): Rex = when (this) {