Add EXCLUDE to partiql-eval

partiql-eval/src/main/kotlin/org/partiql/eval/internal/Compiler.kt

@@ -2,6 +2,7 @@ package org.partiql.eval.internal
 
 import org.partiql.eval.internal.operator.Operator
 import org.partiql.eval.internal.operator.rel.RelDistinct
+import org.partiql.eval.internal.operator.rel.RelExclude
 import org.partiql.eval.internal.operator.rel.RelFilter
 import org.partiql.eval.internal.operator.rel.RelJoinInner
 import org.partiql.eval.internal.operator.rel.RelJoinLeft
@@ -185,4 +186,9 @@ internal class Compiler(
         val condition = visitRex(node.predicate, ctx)
         return RelFilter(input, condition)
     }
+
+    override fun visitRelOpExclude(node: Rel.Op.Exclude, ctx: Unit): Operator {
+        val input = visitRel(node.input, ctx)
+        return RelExclude(input, node.paths)
+    }
 }

partiql-eval/src/main/kotlin/org/partiql/eval/internal/operator/rel/RelExclude.kt

@@ -0,0 +1,199 @@
+package org.partiql.eval.internal.operator.rel
+
+import org.partiql.eval.internal.Record
+import org.partiql.eval.internal.operator.Operator
+import org.partiql.plan.Rel
+import org.partiql.plan.relOpExcludeTypeCollIndex
+import org.partiql.plan.relOpExcludeTypeCollWildcard
+import org.partiql.plan.relOpExcludeTypeStructKey
+import org.partiql.plan.relOpExcludeTypeStructSymbol
+import org.partiql.plan.relOpExcludeTypeStructWildcard
+import org.partiql.value.BagValue
+import org.partiql.value.CollectionValue
+import org.partiql.value.ListValue
+import org.partiql.value.PartiQLValue
+import org.partiql.value.PartiQLValueExperimental
+import org.partiql.value.PartiQLValueType
+import org.partiql.value.SexpValue
+import org.partiql.value.StructValue
+import org.partiql.value.bagValue
+import org.partiql.value.listValue
+import org.partiql.value.sexpValue
+import org.partiql.value.structValue
+
+internal class RelExclude(
+    private val input: Operator.Relation,
+    private val exclusions: List<Rel.Op.Exclude.Path>
+) : Operator.Relation {
+
+    override fun open() {
+        input.open()
+    }
+
+    override fun next(): Record? {
+        val record = input.next() ?: return null
+        return exclusions.fold(record) { rec, path -> exclude(rec, path) }
+    }
+
+    override fun close() {
+        input.close()
+    }
+
+    @OptIn(PartiQLValueExperimental::class)
+    private fun exclude(
+        record: Record,
+        path: Rel.Op.Exclude.Path
+    ): Record {
+        val values = record.values
+        val value = values.getOrNull(path.root.ref)
+        val newValues = if (value != null) {
+            values[path.root.ref] = exclude(value, path.steps)
+            values
+        } else {
+            values
+        }
+        return Record(newValues)
+    }
+
+    @OptIn(PartiQLValueExperimental::class)
+    private fun exclude(
+        structValue: StructValue<*>,
+        exclusions: List<Rel.Op.Exclude.Step>
+    ): PartiQLValue {
+        val structSymbolsToRemove = mutableSetOf<String>()
+        val structKeysToRemove = mutableSetOf<String>() // keys stored as lowercase strings
+        val branches = mutableMapOf<Rel.Op.Exclude.Type, List<Rel.Op.Exclude.Step>>()
+        exclusions.forEach { exclusion ->
+            when (exclusion.substeps.isEmpty()) {
+                true -> {
+                    when (val leafType = exclusion.type) {
+                        is Rel.Op.Exclude.Type.StructWildcard -> {
+                            // struct wildcard at current level. return empty struct
+                            return structValue<PartiQLValue>()
+                        }
+                        is Rel.Op.Exclude.Type.StructSymbol -> structSymbolsToRemove.add(leafType.symbol)
+                        is Rel.Op.Exclude.Type.StructKey -> structKeysToRemove.add(leafType.key.lowercase())
+                        else -> { /* coll step; do nothing */ }
+                    }
+                }
+                false -> {
+                    when (exclusion.type) {
+                        is Rel.Op.Exclude.Type.StructWildcard, is Rel.Op.Exclude.Type.StructSymbol, is Rel.Op.Exclude.Type.StructKey -> branches[exclusion.type] =
+                            exclusion.substeps
+                        else -> { /* coll step; do nothing */ }
+                    }
+                }
+            }
+        }
+        val finalStruct = structValue.entries.mapNotNull { structField ->
+            if (structSymbolsToRemove.contains(structField.first) || structKeysToRemove.contains(structField.first.lowercase())) {
+                // struct attr is to be removed at current level
+                null
+            } else {
+                // deeper level exclusions
+                val name = structField.first
+                var value = structField.second
+                // apply struct key exclusions at deeper levels
+                val structKey = relOpExcludeTypeStructKey(name)
+                branches[structKey]?.let {
+                    value = exclude(value, it)
+                }
+                // apply struct symbol exclusions at deeper levels
+                val structSymbol = relOpExcludeTypeStructSymbol(name)
+                branches[structSymbol]?.let {
+                    value = exclude(value, it)
+                }
+                // apply struct wildcard exclusions at deeper levels
+                val structWildcard = relOpExcludeTypeStructWildcard()
+                branches[structWildcard]?.let {
+                    value = exclude(value, it)
+                }
+                Pair(name, value)
+            }
+        }
+        return structValue(finalStruct)
+    }
+
+    /**
+     * Returns a [PartiQLValue] created from an iterable of [coll]. Requires [type] to be a collection type
+     * (i.e. [PartiQLValueType.LIST], [PartiQLValueType.BAG], or [PartiQLValueType.SEXP]).
+     */
+    @OptIn(PartiQLValueExperimental::class)
+    private fun newCollValue(type: PartiQLValueType, coll: Iterable<PartiQLValue>): PartiQLValue {
+        return when (type) {
+            PartiQLValueType.LIST -> listValue(coll)
+            PartiQLValueType.BAG -> bagValue(coll)
+            PartiQLValueType.SEXP -> sexpValue(coll)
+            else -> error("Collection type required")
+        }
+    }
+
+    @OptIn(PartiQLValueExperimental::class)
+    private fun exclude(
+        coll: CollectionValue<*>,
+        type: PartiQLValueType,
+        exclusions: List<Rel.Op.Exclude.Step>
+    ): PartiQLValue {
+        val indexesToRemove = mutableSetOf<Int>()
+        val branches = mutableMapOf<Rel.Op.Exclude.Type, List<Rel.Op.Exclude.Step>>()
+        exclusions.forEach { exclusion ->
+            when (exclusion.substeps.isEmpty()) {
+                true -> {
+                    when (val leafType = exclusion.type) {
+                        is Rel.Op.Exclude.Type.CollWildcard -> {
+                            // collection wildcard at current level. return empty collection
+                            return newCollValue(type, emptyList())
+                        }
+                        is Rel.Op.Exclude.Type.CollIndex -> {
+                            indexesToRemove.add(leafType.index)
+                        }
+                        else -> { /* struct step; do nothing */ }
+                    }
+                }
+                false -> {
+                    when (exclusion.type) {
+                        is Rel.Op.Exclude.Type.CollWildcard, is Rel.Op.Exclude.Type.CollIndex -> branches[exclusion.type] =
+                            exclusion.substeps
+                        else -> { /* struct step; do nothing */ }
+                    }
+                }
+            }
+        }
+        val finalColl = coll.mapIndexedNotNull { index, element ->
+            if (indexesToRemove.contains(index)) {
+                // coll index is to be removed at current level
+                null
+            } else {
+                // deeper level exclusions
+                var value = element
+                if (coll is ListValue || coll is SexpValue) {
+                    // apply collection index exclusions at deeper levels for lists and sexps
+                    val collIndex = relOpExcludeTypeCollIndex(index)
+                    branches[collIndex]?.let {
+                        value = exclude(element, it)
+                    }
+                }
+                // apply collection wildcard exclusions at deeper levels for lists, bags, and sexps
+                val collWildcard = relOpExcludeTypeCollWildcard()
+                branches[collWildcard]?.let {
+                    value = exclude(value, it)
+                }
+                value
+            }
+        }
+        return newCollValue(type, finalColl)
+    }
+
+    @OptIn(PartiQLValueExperimental::class)
+    private fun exclude(initialPartiQLValue: PartiQLValue, exclusions: List<Rel.Op.Exclude.Step>): PartiQLValue {
+        return when (initialPartiQLValue) {
+            is StructValue<*> -> exclude(initialPartiQLValue, exclusions)
+            is BagValue<*> -> exclude(initialPartiQLValue, PartiQLValueType.BAG, exclusions)
+            is ListValue<*> -> exclude(initialPartiQLValue, PartiQLValueType.LIST, exclusions)
+            is SexpValue<*> -> exclude(initialPartiQLValue, PartiQLValueType.SEXP, exclusions)
+            else -> {
+                initialPartiQLValue
+            }
+        }
+    }
+}

partiql-eval/src/test/kotlin/org/partiql/eval/internal/PartiQLEngineDefaultTest.kt

@@ -226,6 +226,72 @@ class PartiQLEngineDefaultTest {
                     "c" to stringValue("z"),
                 )
             ),
+            SuccessTestCase(
+                input = """
+                    SELECT t
+                    EXCLUDE t.a.b
+                    FROM <<
+                        {'a': {'b': 2}, 'foo': 'bar', 'foo2': 'bar2'}
+                    >> AS t
+                """.trimIndent(),
+                expected = bagValue(
+                    structValue(
+                        "t" to structValue(
+                            "a" to structValue<PartiQLValue>(
+                                // field `b` excluded
+                            ),
+                            "foo" to stringValue("bar"),
+                            "foo2" to stringValue("bar2")
+                        )
+                    ),
+                )
+            ),
+            SuccessTestCase(
+                input = """
+                    SELECT *
+                    EXCLUDE
+                        t.a.b.c[*].field_x
+                    FROM [{
+                        'a': {
+                            'b': {
+                                'c': [
+                                    {                    -- c[0]; field_x to be removed
+                                        'field_x': 0, 
+                                        'field_y': 0
+                                    },
+                                    {                    -- c[1]; field_x to be removed
+                                        'field_x': 1,
+                                        'field_y': 1
+                                    },
+                                    {                    -- c[2]; field_x to be removed
+                                        'field_x': 2,
+                                        'field_y': 2
+                                    }
+                                ]
+                            }
+                        }
+                    }] AS t
+                """.trimIndent(),
+                expected = bagValue(
+                    structValue(
+                        "a" to structValue(
+                            "b" to structValue(
+                                "c" to bagValue( // TODO: should be ListValue; currently, Rex.ExprCollection doesn't return lists
+                                    structValue(
+                                        "field_y" to int32Value(0)
+                                    ),
+                                    structValue(
+                                        "field_y" to int32Value(1)
+                                    ),
+                                    structValue(
+                                        "field_y" to int32Value(2)
+                                    )
+                                )
+                            )
+                        )
+                    )
+                )
+            )
         )
     }
     public class SuccessTestCase @OptIn(PartiQLValueExperimental::class) constructor(