Allow empty values within partitioned fragment

In some cases, a subplan containing a remote exchange (which results in a
partitioned fragment) can be replaced with an empty Values node. In the
process, the remote exchange is removed, which leaves the Values node
within the downstream partitioned fragment.

The plan fragmenter cannot currently deal with that scenario. It fails
with an error indicating the partitioning is already set. However,
it is for an empty Values to appear within a partitioned fragment,
since there are no issues with rows being "replicated" across
multiple fragments.

core/trino-main/src/main/java/io/trino/sql/planner/PlanFragmenter.java

@@ -347,7 +347,11 @@ public PlanNode visitMergeProcessor(MergeProcessorNode node, RewriteContext<Frag
         @Override
         public PlanNode visitValues(ValuesNode node, RewriteContext<FragmentProperties> context)
         {
-            context.get().setSingleNodeDistribution();
+            // An empty values node is compatible with any distribution, so
+            // don't attempt to overwrite one's already been chosen
+            if (node.getRowCount() != 0 || !context.get().hasDistribution()) {
+                context.get().setSingleNodeDistribution();
+            }
             return context.defaultRewrite(node, context.get());
         }
 
@@ -435,6 +439,11 @@ public List<SubPlan> getChildren()
             return children;
         }
 
+        public boolean hasDistribution()
+        {
+            return partitioningHandle.isPresent();
+        }
+
         public FragmentProperties setSingleNodeDistribution()
         {
             if (partitioningHandle.isPresent() && partitioningHandle.get().isSingleNode()) {

plugin/trino-hive/pom.xml

@@ -430,6 +430,24 @@
             <scope>test</scope>
         </dependency>
 
+        <dependency>
+            <groupId>org.junit.jupiter</groupId>
+            <artifactId>junit-jupiter-api</artifactId>
+            <scope>test</scope>
+        </dependency>
+
+        <dependency>
+            <groupId>org.junit.jupiter</groupId>
+            <artifactId>junit-jupiter-engine</artifactId>
+            <scope>test</scope>
+        </dependency>
+
+        <dependency>
+            <groupId>org.junit.jupiter</groupId>
+            <artifactId>junit-jupiter-params</artifactId>
+            <scope>test</scope>
+        </dependency>
+
         <dependency>
             <groupId>org.openjdk.jmh</groupId>
             <artifactId>jmh-core</artifactId>
@@ -455,6 +473,28 @@
         </dependency>
     </dependencies>
 
+    <build>
+        <plugins>
+            <plugin>
+                <groupId>org.apache.maven.plugins</groupId>
+                <artifactId>maven-surefire-plugin</artifactId>
+                <dependencies>
+                    <!-- allow both JUnit and TestNG -->
+                    <dependency>
+                        <groupId>org.apache.maven.surefire</groupId>
+                        <artifactId>surefire-junit-platform</artifactId>
+                        <version>${dep.plugin.surefire.version}</version>
+                    </dependency>
+                    <dependency>
+                        <groupId>org.apache.maven.surefire</groupId>
+                        <artifactId>surefire-testng</artifactId>
+                        <version>${dep.plugin.surefire.version}</version>
+                    </dependency>
+                </dependencies>
+            </plugin>
+        </plugins>
+    </build>
+
     <profiles>
         <profile>
             <id>default</id>

plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestIssue14317.java

@@ -0,0 +1,70 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.plugin.hive;
+
+import io.trino.sql.query.QueryAssertions;
+import io.trino.testing.DistributedQueryRunner;
+import org.testng.annotations.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Test for https://github.com/trinodb/trino/issues/14317
+ * <p>
+ * The only known reproduction involves a connector that exposes partitioned tables
+ * and supports predicate pushdown.
+ */
+public class TestIssue14317
+{
+    @Test
+    public void test()
+            throws Exception
+    {
+        HiveQueryRunner.Builder<?> builder = HiveQueryRunner.builder()
+                .setCreateTpchSchemas(false);
+
+        try (DistributedQueryRunner queryRunner = builder.build();
+                QueryAssertions assertions = new QueryAssertions(queryRunner);) {
+            queryRunner.execute("CREATE SCHEMA s");
+
+            queryRunner.execute("CREATE TABLE s.t (a bigint, b bigint)");
+            queryRunner.execute("CREATE TABLE s.u (c bigint, d bigint) WITH (partitioned_by = array['d'])");
+            queryRunner.execute("CREATE TABLE s.v (e bigint, f bigint)");
+
+            queryRunner.execute("INSERT INTO s.t VALUES (5, 6)");
+            queryRunner.execute("INSERT INTO s.u VALUES (5, 6)");
+            queryRunner.execute("INSERT INTO s.v VALUES (5, 6)");
+
+            assertThat(assertions.query("""
+                    WITH t1 AS (
+                      SELECT a
+                      FROM (
+                        SELECT a, ROW_NUMBER() OVER (PARTITION BY a) AS rn
+                        FROM s.t)
+                      WHERE rn = 1),
+                    t2 AS (SELECT c FROM s.u WHERE d - 5 = 8)
+                    SELECT v.e
+                    FROM s.v
+                      INNER JOIN t1 on v.e = t1.a
+                      INNER JOIN t2 ON v.e = t2.c
+                    """))
+                    .returnsEmptyResult();
+
+            queryRunner.execute("DROP TABLE s.t");
+            queryRunner.execute("DROP TABLE s.u");
+            queryRunner.execute("DROP TABLE s.v");
+            queryRunner.execute("DROP SCHEMA s");
+        }
+    }
+}