Plan queries with and without EXPLAIN ANALYZE the same way

core/trino-main/src/main/java/io/trino/sql/planner/optimizations/AddExchanges.java

@@ -866,12 +866,8 @@ public PlanWithProperties visitTableUpdate(TableUpdateNode node, PreferredProper
         @Override
         public PlanWithProperties visitExplainAnalyze(ExplainAnalyzeNode node, PreferredProperties preferredProperties)
         {
-            PlanWithProperties child = planChild(node, PreferredProperties.any());
-
-            // if the child is already a gathering exchange, don't add another
-            if ((child.getNode() instanceof ExchangeNode) && ((ExchangeNode) child.getNode()).getType() == ExchangeNode.Type.GATHER) {
-                return rebaseAndDeriveProperties(node, child);
-            }
+            // Same PreferredProperties as OutputNode
+            PlanWithProperties child = planChild(node, PreferredProperties.undistributed());
 
             // Always add an exchange because ExplainAnalyze should be in its own stage
             child = withDerivedProperties(

core/trino-main/src/test/java/io/trino/sql/planner/TestLogicalPlanner.java

@@ -2237,6 +2237,25 @@ public void testExplainAnalyze()
                         node(ExplainAnalyzeNode.class,
                                 exchange(LOCAL, GATHER,
                                         strictTableScan("nation", ImmutableMap.of("regionkey", "regionkey"))))));
+
+        assertDistributedPlan("""
+                              EXPLAIN ANALYZE
+                              SELECT * FROM
+                                  (SELECT * from nation, region)
+                              UNION ALL
+                                  (SELECT * from nation, region)
+                              """,
+                output(
+                        node(ExplainAnalyzeNode.class,
+                                exchange(LOCAL, GATHER,
+                                        exchange(REMOTE, GATHER,
+                                                exchange(REMOTE, GATHER,
+                                                        join(INNER, builder -> builder
+                                                                .left(tableScan("nation", ImmutableMap.of("regionkey_0", "regionkey")))
+                                                                .right(anyTree(tableScan("region", ImmutableMap.of("regionkey_1", "regionkey"))))),
+                                                        join(INNER, builder -> builder
+                                                                .left(tableScan("nation", ImmutableMap.of("regionkey_2", "regionkey")))
+                                                                .right(anyTree(tableScan("region", ImmutableMap.of("regionkey_3", "regionkey")))))))))));
     }
 
     @Test

testing/trino-tests/src/test/java/io/trino/sql/planner/BaseCostBasedPlanTest.java

@@ -44,7 +44,9 @@
 import java.io.UncheckedIOException;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.util.Arrays;
 import java.util.List;
+import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
 import static com.google.common.base.Preconditions.checkState;
@@ -148,6 +150,29 @@ public void test(String queryResourcePath)
         assertThat(generateQueryPlan(readQuery(queryResourcePath))).isEqualTo(read(getQueryPlanResourcePath(queryResourcePath)));
     }
 
+    protected void assertExplainAnalyzePlan(String queryResourcePath)
+    {
+        String query = readQuery(queryResourcePath);
+
+        String queryPlan = generateQueryPlan(query);
+        String explainAnalyzeQueryPlan = generateQueryPlan("EXPLAIN ANALYZE " + query);
+        String[] explainAnalyzeLines = explainAnalyzeQueryPlan.split("\n");
+
+        // for EXPLAIN ANALYZE, the first two lines reflect the additional root fragment containing the ExplainAnalyze operator
+        assertThat(String.join("\n", Arrays.copyOfRange(explainAnalyzeLines, 0, 2)) + "\n")
+                .isEqualTo("""
+                           local exchange (GATHER, SINGLE, [])
+                               remote exchange (GATHER, SINGLE, [])
+                           """);
+
+        // the remaining lines should match the original query plan, except for the indentation
+        explainAnalyzeQueryPlan = Arrays.stream(Arrays.copyOfRange(explainAnalyzeLines, 2, explainAnalyzeLines.length))
+                .map(line -> line.replaceFirst("^ {8}", ""))
+                .collect(Collectors.joining("\n")) + "\n";
+
+        assertThat(queryPlan).isEqualTo(explainAnalyzeQueryPlan);
+    }
+
     private String getQueryPlanResourcePath(String queryResourcePath)
     {
         Path queryPath = Paths.get(queryResourcePath);

testing/trino-tests/src/test/java/io/trino/sql/planner/TestExplainAnalyzePartitionedTpcdsPlan.java

@@ -0,0 +1,66 @@
+/*
+ * 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.sql.planner;
+
+import io.trino.tpcds.Table;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.MethodSource;
+
+import java.util.List;
+
+import static com.google.common.collect.ImmutableList.toImmutableList;
+
+public class TestExplainAnalyzePartitionedTpcdsPlan
+        extends BaseCostBasedPlanTest
+{
+    protected TestExplainAnalyzePartitionedTpcdsPlan()
+    {
+        super("tpcds_sf1000_parquet_part", true);
+    }
+
+    @Override
+    @ParameterizedTest
+    @MethodSource("getQueryResourcePaths")
+    public void test(String queryResourcePath)
+    {
+        assertExplainAnalyzePlan(queryResourcePath);
+    }
+
+    @Override
+    protected List<String> getTableNames()
+    {
+        return io.trino.tpcds.Table.getBaseTables().stream()
+                .filter(table -> table != io.trino.tpcds.Table.DBGEN_VERSION)
+                .map(Table::getName)
+                .collect(toImmutableList());
+    }
+
+    @Override
+    protected String getTableResourceDirectory()
+    {
+        return "iceberg/tpcds/sf1000/partitioned/";
+    }
+
+    @Override
+    protected String getTableTargetDirectory()
+    {
+        return "iceberg-tpcds-sf1000-parquet-part/";
+    }
+
+    @Override
+    protected List<String> getQueryResourcePaths()
+    {
+        return TPCDS_SQL_FILES;
+    }
+}

testing/trino-tests/src/test/java/io/trino/sql/planner/TestExplainAnalyzeTpchPlan.java

@@ -0,0 +1,65 @@
+/*
+ * 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.sql.planner;
+
+import io.trino.tpch.TpchTable;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.MethodSource;
+
+import java.util.List;
+
+import static com.google.common.collect.ImmutableList.toImmutableList;
+
+public class TestExplainAnalyzeTpchPlan
+        extends BaseCostBasedPlanTest
+{
+    public TestExplainAnalyzeTpchPlan()
+    {
+        super("tpch_sf1000_parquet", false);
+    }
+
+    @Override
+    @ParameterizedTest
+    @MethodSource("getQueryResourcePaths")
+    public void test(String queryResourcePath)
+    {
+        assertExplainAnalyzePlan(queryResourcePath);
+    }
+
+    @Override
+    protected List<String> getTableNames()
+    {
+        return TpchTable.getTables().stream()
+                .map(TpchTable::getTableName)
+                .collect(toImmutableList());
+    }
+
+    @Override
+    protected String getTableResourceDirectory()
+    {
+        return "iceberg/tpch/sf1000/unpartitioned/";
+    }
+
+    @Override
+    protected String getTableTargetDirectory()
+    {
+        return "iceberg-tpch-sf1000-parquet/";
+    }
+
+    @Override
+    protected List<String> getQueryResourcePaths()
+    {
+        return TPCH_SQL_FILES;
+    }
+}