Enable the IN clause to also compare the value of single-pair structs

cli/test/org/partiql/cli/CliTest.kt

@@ -155,6 +155,30 @@ class CliTest {
         assertEquals("{a:1}\n{b:1}\n", actual)
     }
 
+    @Test
+    fun withSelectIn() {
+        makeCli(
+            query = "SELECT age FROM << { 'name': 'John', 'age': 22 } >> WHERE name IN (SELECT name FROM <<{ 'name': 'John' }>>)",
+            output = FileOutputStream(testFile),
+            outputFormat = OutputFormat.ION_TEXT
+        ).run()
+        val actual = testFile.bufferedReader().use { it.readText() }
+
+        assertEquals("{age:22}\n", actual)
+    }
+
+    @Test
+    fun withSelectInButEmpty() {
+        makeCli(
+            query = "SELECT age FROM << { 'name': 'John', 'age': 22 } >> WHERE name IN (SELECT name FROM <<{ 'name': 'john' }>>)",
+            output = FileOutputStream(testFile),
+            outputFormat = OutputFormat.ION_TEXT
+        ).run()
+        val actual = testFile.bufferedReader().use { it.readText() }
+
+        assertEquals("\n", actual)
+    }
+
     @Test
     fun withoutInput() {
         val subject = makeCli("1")

lang/src/org/partiql/lang/eval/EvaluatingCompiler.kt

@@ -61,6 +61,7 @@ import org.partiql.lang.types.TypedOpParameter
 import org.partiql.lang.types.UnknownArguments
 import org.partiql.lang.types.UnsupportedTypeCheckException
 import org.partiql.lang.types.toTypedOpParameter
+import org.partiql.lang.util.asIonStruct
 import org.partiql.lang.util.bigDecimalOf
 import org.partiql.lang.util.checkThreadInterrupted
 import org.partiql.lang.util.codePointSequence
@@ -732,8 +733,11 @@ internal class EvaluatingCompiler(
         val args = expr.operands
         val leftThunk = compileAstExpr(args[0])
         val rightOp = args[1]
+        fun PartiqlAst.Expr.Struct.isSingleValueOptimizedStruct() = this.fields.size == 1 && this.fields[0].second is PartiqlAst.Expr.Lit
 
-        fun isOptimizedCase(values: List<PartiqlAst.Expr>): Boolean = values.all { it is PartiqlAst.Expr.Lit && !it.value.isNull }
+        fun isOptimizedCase(values: List<PartiqlAst.Expr>): Boolean = values.all {
+            (it is PartiqlAst.Expr.Lit && !it.value.isNull) || (it is PartiqlAst.Expr.Struct && it.isSingleValueOptimizedStruct())
+        }
 
         fun optimizedCase(values: List<PartiqlAst.Expr>): ThunkEnv {
             // Put all the literals in the sequence into a pre-computed map to be checked later by the thunk.
@@ -743,6 +747,13 @@ internal class EvaluatingCompiler(
             // NOTE: we cannot use a [HashSet<>] here because [ExprValue] does not implement [Object.hashCode] or
             // [Object.equals].
             val precomputedLiteralsMap = values
+                .asSequence()
+                .map { value ->
+                    when (value) {
+                        is PartiqlAst.Expr.Struct -> if (value.isSingleValueOptimizedStruct()) value.fields[0].second else value
+                        else -> value
+                    }
+                }
                 .filterIsInstance<PartiqlAst.Expr.Lit>()
                 .mapTo(TreeSet<ExprValue>(DEFAULT_COMPARATOR)) {
                     valueFactory.newFromIonValue(
@@ -798,6 +809,12 @@ internal class EvaluatingCompiler(
                                 when (it.type) {
                                     ExprValueType.NULL -> nullSeen = true
                                     ExprValueType.MISSING -> missingSeen = true
+                                    // Allow comparison with 1-pair structs to remain compatible SQL-92
+                                    ExprValueType.STRUCT -> {
+                                        if (it.ionValue.asIonStruct().size() != 1) return@forEach
+                                        if (it.iterator().next().exprEquals(leftValue))
+                                            return@thunkEnvOperands valueFactory.newBoolean(true)
+                                    }
                                     // short-circuit to TRUE on the first matching value
                                     else -> if (it.exprEquals(leftValue)) {
                                         return@thunkEnvOperands valueFactory.newBoolean(true)

lang/test/org/partiql/lang/eval/EvaluatingCompilerInTests.kt

@@ -50,4 +50,73 @@ class EvaluatingCompilerInTests : EvaluatorTestBase() {
             ),
         )
     }
+
+    /**
+     * Note: if we start doing compile-time reduction of literal values, we'll need to change `1+1` in the unoptimized bags
+     * to something else
+     */
+    @ParameterizedTest
+    @ArgumentsSource(InStructOperatorTestCases::class)
+    fun inStructOperatorTest(tc: EvaluatorTestCase) =
+        runEvaluatorTestCase(tc, EvaluationSession.standard())
+    class InStructOperatorTestCases : ArgumentsProviderBase() {
+        private val optimizedBag = "<< {'a': 1}, {'b':2}, {'c':3}>>"
+        private val multiKeyBag = "<< {'a': 1, 'b': 1}, {'a':2, 'b': 2}>>"
+        private val unoptimizedBag = "<< {'a': 1}, {'b': (1+1) }, {'c':3}>>"
+        override fun getParameters(): List<Any> = listOf(
+            // These cases get the optimized thunk since the right-operand consists solely of known literal values
+            EvaluatorTestCase("0 IN $optimizedBag", "FALSE"),
+            EvaluatorTestCase("1 IN $optimizedBag", "TRUE"),
+            EvaluatorTestCase("2 IN $optimizedBag", "TRUE"),
+            EvaluatorTestCase("3 IN $optimizedBag", "TRUE"),
+            EvaluatorTestCase("4 IN $optimizedBag", "FALSE"),
+
+            // These cases should all fail due to multi-key structs. We only compare single-key structs.
+            EvaluatorTestCase("0 IN $multiKeyBag", "FALSE"),
+            EvaluatorTestCase("1 IN $multiKeyBag", "FALSE"),
+            EvaluatorTestCase("2 IN $multiKeyBag", "FALSE"),
+
+            // These cases get the un-optimized thunk since the right-operand does not consist solely of known literal values
+            EvaluatorTestCase("0 IN $unoptimizedBag", "FALSE"),
+            EvaluatorTestCase("1 IN $unoptimizedBag", "TRUE"),
+            EvaluatorTestCase("2 IN $unoptimizedBag", "TRUE"),
+            EvaluatorTestCase("3 IN $unoptimizedBag", "TRUE"),
+            EvaluatorTestCase("4 IN $unoptimizedBag", "FALSE"),
+
+        )
+    }
+
+    /**
+     * Note: if we start doing compile-time reduction of literal values, we'll need to change `1+1` in the unoptimized bags
+     * to something else
+     */
+    @ParameterizedTest
+    @ArgumentsSource(InSelectOperatorTestCases::class)
+    fun inSelectOperatorTest(tc: EvaluatorTestCase) =
+        runEvaluatorTestCase(tc, EvaluationSession.standard())
+    class InSelectOperatorTestCases : ArgumentsProviderBase() {
+        private val optimizedQuery: String = "SELECT a FROM << {'a': 1}, {'a':2}, {'a':3} >>"
+        private val unoptimizedQuery: String = "SELECT a FROM << {'a': 1}, {'a':(1+1)}, {'a':3} >>"
+        private val multiPairStructQuery: String = "SELECT a, b FROM << {'a': 1, 'b': 1}, {'a':2, 'b': 2 } >>"
+        override fun getParameters(): List<Any> = listOf(
+            // These cases get the optimized thunk since the right-operand consists solely of known literal values
+            EvaluatorTestCase("0 IN ($optimizedQuery)", "FALSE"),
+            EvaluatorTestCase("1 IN ($optimizedQuery)", "TRUE"),
+            EvaluatorTestCase("2 IN ($optimizedQuery)", "TRUE"),
+            EvaluatorTestCase("3 IN ($optimizedQuery)", "TRUE"),
+            EvaluatorTestCase("4 IN ($optimizedQuery)", "FALSE"),
+
+            // These cases should all fail due to multi-key structs. We only compare single-key structs.
+            EvaluatorTestCase("0 IN ($multiPairStructQuery)", "FALSE"),
+            EvaluatorTestCase("1 IN ($multiPairStructQuery)", "FALSE"),
+            EvaluatorTestCase("2 IN ($multiPairStructQuery)", "FALSE"),
+
+            // These cases get the un-optimized thunk since the right-operand does not consist solely of known literal values
+            EvaluatorTestCase("0 IN ($unoptimizedQuery)", "FALSE"),
+            EvaluatorTestCase("1 IN ($unoptimizedQuery)", "TRUE"),
+            EvaluatorTestCase("2 IN ($unoptimizedQuery)", "TRUE"),
+            EvaluatorTestCase("3 IN ($unoptimizedQuery)", "TRUE"),
+            EvaluatorTestCase("4 IN ($unoptimizedQuery)", "FALSE"),
+        )
+    }
 }