From c0c98b6a44d4d2278ae93510e2dad3574af119b3 Mon Sep 17 00:00:00 2001
From: Martin Traverso <mtraverso@gmail.com>
Date: Thu, 31 Aug 2023 14:00:31 -0700
Subject: [PATCH] Fix incorrect symbol resolution during predicate pushdown

When the expression evaluated by LayoutConstraintEvaluator contains
a lambda expression, it's incorrect to attempt to resolve a reference
to a lambda expression argument as if it were a column.
---
 .../sql/planner/LookupSymbolResolver.java     |  4 +--
 .../plugin/hive/TestHiveConnectorTest.java    | 28 +++++++++++++++++++
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/core/trino-main/src/main/java/io/trino/sql/planner/LookupSymbolResolver.java b/core/trino-main/src/main/java/io/trino/sql/planner/LookupSymbolResolver.java
index e640385799a5..e2f0b1c65953 100644
--- a/core/trino-main/src/main/java/io/trino/sql/planner/LookupSymbolResolver.java
+++ b/core/trino-main/src/main/java/io/trino/sql/planner/LookupSymbolResolver.java
@@ -19,7 +19,6 @@
 
 import java.util.Map;
 
-import static com.google.common.base.Preconditions.checkArgument;
 import static java.util.Objects.requireNonNull;
 
 public class LookupSymbolResolver
@@ -41,9 +40,8 @@ public LookupSymbolResolver(Map<Symbol, ColumnHandle> assignments, Map<ColumnHan
     public Object getValue(Symbol symbol)
     {
         ColumnHandle column = assignments.get(symbol);
-        checkArgument(column != null, "Missing column assignment for %s", symbol);
 
-        if (!bindings.containsKey(column)) {
+        if (column == null || !bindings.containsKey(column)) {
             return symbol.toSymbolReference();
         }
 
diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveConnectorTest.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveConnectorTest.java
index 3f7d7a2f7819..1d8c94367020 100644
--- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveConnectorTest.java
+++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveConnectorTest.java
@@ -14,6 +14,9 @@
 package io.trino.plugin.hive;
 
 import io.trino.testing.QueryRunner;
+import org.testng.annotations.Test;
+
+import static io.trino.testing.TestingNames.randomNameSuffix;
 
 public class TestHiveConnectorTest
         extends BaseHiveConnectorTest
@@ -24,4 +27,29 @@ protected QueryRunner createQueryRunner()
     {
         return createHiveQueryRunner(HiveQueryRunner.builder());
     }
+
+    @Test
+    public void testPredicatePushdownWithLambdaExpression()
+    {
+        String table = "test_predicate_pushdown_" + randomNameSuffix();
+
+        assertUpdate("""
+            CREATE TABLE %s (v, k)
+            WITH (partitioned_by = ARRAY['k'])
+            AS (VALUES ('value', 'key'))
+            """.formatted(table),
+                1);
+
+        try {
+            assertQuery("""
+                            SELECT *
+                            FROM %s
+                            WHERE k = 'key' AND regexp_replace(v, '(.*)', x -> x[1]) IS NOT NULL
+                            """.formatted(table),
+                    "VALUES ('value', 'key')");
+        }
+        finally {
+            assertUpdate("DROP TABLE " + table);
+        }
+    }
 }