[fix](join): avoid join amplification in scalar subquery window rewrite

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/AggScalarSubQueryToWindowFunction.java

@@ -95,6 +95,7 @@ public class AggScalarSubQueryToWindowFunction extends DefaultPlanRewriter<JobCo
     private final List<LogicalPlan> innerPlans = Lists.newArrayList();
     private final List<AggregateFunction> functions = Lists.newArrayList();
     private final Map<Expression, Expression> innerOuterSlotMap = Maps.newHashMap();
+    private final List<Slot> partitionBySlots = Lists.newArrayList();
 
     /**
      * the entrance of this rule. we only override one visitor: visitLogicalFilter
@@ -130,6 +131,7 @@ private Optional<LogicalApply<Plan, Plan>> findApply(LogicalFilter<? extends Pla
     }
 
     private boolean check(LogicalFilter<? extends Plan> outerFilter, LogicalApply<Plan, Plan> apply) {
+        resetState();
         outerPlans.addAll(apply.child(0).collect(LogicalPlan.class::isInstance));
         innerPlans.addAll(apply.child(1).collect(LogicalPlan.class::isInstance));
 
@@ -138,10 +140,18 @@ && checkApply(apply)
                 && checkAggregate()
                 && checkJoin()
                 && checkProject()
-                && checkRelation(apply.getCorrelationSlot())
+                && checkRelation(apply)
                 && checkFilter(outerFilter);
     }
 
+    private void resetState() {
+        outerPlans.clear();
+        innerPlans.clear();
+        functions.clear();
+        innerOuterSlotMap.clear();
+        partitionBySlots.clear();
+    }
+
     // check children's nodes because query process will be changed
     private boolean checkPlanType() {
         return outerPlans.stream().allMatch(p -> OUTER_SUPPORTED_PLAN.stream().anyMatch(c -> c.isInstance(p)))
@@ -250,7 +260,8 @@ private boolean checkProject() {
      * 2. outer table list - inner table list should only remain 1 table
      * 3. the remaining table in step 2 should be correlated table for inner plan
      */
-    private boolean checkRelation(List<Slot> correlatedSlots) {
+    private boolean checkRelation(LogicalApply<Plan, Plan> apply) {
+        List<Slot> correlatedSlots = apply.getCorrelationSlot();
         List<CatalogRelation> outerTables = outerPlans.stream().filter(CatalogRelation.class::isInstance)
                 .map(CatalogRelation.class::cast)
                 .collect(Collectors.toList());
@@ -278,6 +289,19 @@ private boolean checkRelation(List<Slot> correlatedSlots) {
                 .filter(node -> outerIds.contains(node.getTable().getId()))
                 .map(LogicalRelation.class::cast)
                 .map(LogicalRelation::getOutputExprIdSet).flatMap(Collection::stream).collect(Collectors.toSet());
+        partitionBySlots.addAll(apply.left().getOutput().stream()
+                .filter(slot -> correlatedRelationOutput.contains(slot.getExprId()))
+                .collect(Collectors.toList()));
+        if (partitionBySlots.isEmpty()) {
+            return false;
+        }
+        if (!partitionBySlots.stream().map(Slot::getExprId).collect(Collectors.toSet())
+                .equals(correlatedRelationOutput)) {
+            // The rewrite is only safe when the window can partition by the full visible row image of the
+            // outer-only relation. If some of those slots have been pruned from apply.left(), duplicate
+            // outer rows can collapse into one partition and multiply the aggregate result.
+            return false;
+        }
         return correlatedSlots.stream().allMatch(e -> correlatedRelationOutput.contains(e.getExprId()));
     }
 
@@ -346,7 +370,7 @@ private Plan rewrite(LogicalFilter<? extends Plan> filter, LogicalApply<Plan, Pl
             function = ((NullableAggregateFunction) function).withAlwaysNullable(false);
         }
 
-        WindowExpression windowFunction = createWindowFunction(apply.getCorrelationSlot(),
+        WindowExpression windowFunction = createWindowFunction(
                 (AggregateFunction) ExpressionUtils.replace(function, innerOuterSlotMap));
         NamedExpression windowFunctionAlias = new Alias(windowFunction);
 
@@ -368,9 +392,10 @@ private Plan rewrite(LogicalFilter<? extends Plan> filter, LogicalApply<Plan, Pl
         return windowFilter;
     }
 
-    private WindowExpression createWindowFunction(List<Slot> correlatedSlots, AggregateFunction function) {
-        // partition by clause is set by all the correlated slots.
-        return new WindowExpression(function, ImmutableList.copyOf(correlatedSlots), Collections.emptyList());
+    private WindowExpression createWindowFunction(AggregateFunction function) {
+        // Partition by all visible slots from the outer-only relation so join-expanded rows from
+        // that relation do not change the scalar subquery result.
+        return new WindowExpression(function, ImmutableList.copyOf(partitionBySlots), Collections.emptyList());
     }
 
     private static class ExpressionIdenticalChecker extends DefaultExpressionVisitor<Boolean, Expression> {

fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/AggScalarSubQueryToWindowFunctionTest.java

@@ -19,15 +19,24 @@
 
 import org.apache.doris.nereids.datasets.tpch.TPCHTestBase;
 import org.apache.doris.nereids.datasets.tpch.TPCHUtils;
+import org.apache.doris.nereids.trees.expressions.Expression;
+import org.apache.doris.nereids.trees.expressions.NamedExpression;
+import org.apache.doris.nereids.trees.expressions.Slot;
+import org.apache.doris.nereids.trees.expressions.WindowExpression;
 import org.apache.doris.nereids.trees.plans.Plan;
 import org.apache.doris.nereids.trees.plans.logical.LogicalWindow;
 import org.apache.doris.nereids.util.MemoPatternMatchSupported;
 import org.apache.doris.nereids.util.PlanChecker;
 
+import com.google.common.collect.ImmutableSet;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+
 public class AggScalarSubQueryToWindowFunctionTest extends TPCHTestBase implements MemoPatternMatchSupported {
     private static final String SQL_TEMPLATE = "    select\n"
             + "        sum(l_extendedprice) / 7.0 as avg_yearly\n"
@@ -338,6 +347,103 @@ public void testNotMatchTheRule() {
         }
     }
 
+    @Test
+    public void testWindowPartitionsByOuterOnlyRelationSlots() throws Exception {
+        createTable("CREATE TABLE tpch.fact_window_safe (\n"
+                + "  id INT,\n"
+                + "  k INT,\n"
+                + "  v INT\n"
+                + ") ENGINE=OLAP\n"
+                + "DUPLICATE KEY(id)\n"
+                + "DISTRIBUTED BY HASH(id) BUCKETS 1\n"
+                + "PROPERTIES ('replication_num' = '1')");
+        createTable("CREATE TABLE tpch.dim_window_safe (\n"
+                + "  did INT,\n"
+                + "  k INT,\n"
+                + "  tag INT\n"
+                + ") ENGINE=OLAP\n"
+                + "DUPLICATE KEY(did)\n"
+                + "DISTRIBUTED BY HASH(did) BUCKETS 1\n"
+                + "PROPERTIES ('replication_num' = '1')");
+
+        String sql = "SELECT d.did, d.k, d.tag, f.id, f.v "
+                + "FROM fact_window_safe f, dim_window_safe d "
+                + "WHERE f.k = d.k "
+                + "  AND f.v * 2 > ("
+                + "    SELECT SUM(f2.v) "
+                + "    FROM fact_window_safe f2 "
+                + "    WHERE f2.k = d.k"
+                + "  )";
+
+        Plan plan = PlanChecker.from(createCascadesContext(sql))
+                .analyze(sql)
+                .applyBottomUp(new PullUpProjectUnderApply())
+                .applyTopDown(new PushDownFilterThroughProject())
+                .customRewrite(new EliminateUnnecessaryProject())
+                .customRewrite(new AggScalarSubQueryToWindowFunction())
+                .rewrite()
+                .getPlan();
+
+        List<LogicalWindow<Plan>> windows = plan.collectToList(LogicalWindow.class::isInstance);
+        Assertions.assertEquals(1, windows.size());
+
+        LogicalWindow<Plan> window = windows.get(0);
+        List<NamedExpression> windowExpressions = window.getWindowExpressions();
+        Assertions.assertEquals(1, windowExpressions.size());
+
+        WindowExpression windowExpression = (WindowExpression) windowExpressions.get(0).child(0);
+        Set<String> partitionKeys = windowExpression.getPartitionKeys().stream()
+                .map(Expression.class::cast)
+                .filter(Slot.class::isInstance)
+                .map(Slot.class::cast)
+                .map(Slot::getName)
+                .collect(Collectors.toSet());
+        Assertions.assertEquals(ImmutableSet.of("did", "k", "tag"), partitionKeys);
+    }
+
+    @Test
+    public void testNotMatchWhenOuterOnlyRelationOutputIsPruned() throws Exception {
+        createTable("CREATE TABLE tpch.fact_window_pruned (\n"
+                + "  id INT,\n"
+                + "  k INT,\n"
+                + "  v INT\n"
+                + ") ENGINE=OLAP\n"
+                + "DUPLICATE KEY(id)\n"
+                + "DISTRIBUTED BY HASH(id) BUCKETS 1\n"
+                + "PROPERTIES ('replication_num' = '1')");
+        createTable("CREATE TABLE tpch.dim_window_pruned (\n"
+                + "  did INT,\n"
+                + "  k INT,\n"
+                + "  tag INT\n"
+                + ") ENGINE=OLAP\n"
+                + "DUPLICATE KEY(did)\n"
+                + "DISTRIBUTED BY HASH(did) BUCKETS 1\n"
+                + "PROPERTIES ('replication_num' = '1')");
+
+        String sql = "SELECT t.id, t.k, t.v "
+                + "FROM ("
+                + "    SELECT f.id, d.k, f.v "
+                + "    FROM fact_window_pruned f, dim_window_pruned d "
+                + "    WHERE f.k = d.k"
+                + ") t "
+                + "WHERE t.v * 2 > ("
+                + "    SELECT SUM(f2.v) "
+                + "    FROM fact_window_pruned f2 "
+                + "    WHERE f2.k = t.k"
+                + "  )";
+
+        Plan plan = PlanChecker.from(createCascadesContext(sql))
+                .analyze(sql)
+                .applyBottomUp(new PullUpProjectUnderApply())
+                .applyTopDown(new PushDownFilterThroughProject())
+                .customRewrite(new EliminateUnnecessaryProject())
+                .customRewrite(new AggScalarSubQueryToWindowFunction())
+                .rewrite()
+                .getPlan();
+
+        Assertions.assertFalse(plan.anyMatch(LogicalWindow.class::isInstance));
+    }
+
     private void check(String sql) {
         System.out.printf("Test:\n%s\n\n", sql);
         Plan plan = PlanChecker.from(createCascadesContext(sql))

regression-test/data/nereids_rules_p0/subquery_to_window_function/correlated_scalar_subquery_to_window_function.out

@@ -0,0 +1,13 @@
+-- This file is automatically generated. You should know what you did if you want to edit this
+-- !d26072 --
+10	2	1	7
+11	2	1	7
+20	4	2	10
+30	5	3	8
+
+-- !d26072_no_change --
+2	1	7
+2	1	7
+4	2	10
+5	3	8
+

regression-test/suites/nereids_rules_p0/subquery_to_window_function/correlated_scalar_subquery_to_window_function.groovy

@@ -0,0 +1,88 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you 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.
+
+suite("correlated_scalar_subquery_to_window_function") {
+    sql "DROP TABLE IF EXISTS fact FORCE"
+    sql "DROP TABLE IF EXISTS dim FORCE"
+
+    sql """
+        CREATE TABLE fact (
+          id INT,
+          k INT,
+          v INT
+        ) ENGINE=OLAP
+        DUPLICATE KEY(id)
+        DISTRIBUTED BY HASH(id) BUCKETS 1
+        PROPERTIES ('replication_num' = '1')
+    """
+
+    sql """
+        CREATE TABLE dim (
+          did INT,
+          k INT,
+          tag INT
+        ) ENGINE=OLAP
+        DUPLICATE KEY(did)
+        DISTRIBUTED BY HASH(did) BUCKETS 1
+        PROPERTIES ('replication_num' = '1')
+    """
+
+    sql """
+        insert into fact values
+            (1, 1, 5),
+            (2, 1, 7),
+            (3, 2, 4),
+            (4, 2, 10),
+            (5, 3, 8)
+    """
+
+    sql """
+        insert into dim values
+            (10, 1, 1),
+            (11, 1, 1),
+            (20, 2, 1),
+            (30, 3, 0),
+            (31, null, 1)
+    """
+
+    order_qt_d26072 """
+        SELECT d.did, f.id, f.k, f.v
+        FROM fact f, dim d
+        WHERE f.k = d.k
+          AND f.v * 2 > (
+            SELECT SUM(f2.v)
+            FROM fact f2
+            WHERE f2.k = d.k
+          )
+        ORDER BY d.did, f.id
+    """
+
+    order_qt_d26072_no_change """
+        SELECT t.id, t.k, t.v
+        FROM (
+            SELECT f.id, d.k, f.v
+            FROM fact f, dim d
+            WHERE f.k = d.k
+        ) t
+        WHERE t.v * 2 > (
+            SELECT SUM(f2.v)
+            FROM fact f2
+            WHERE f2.k = t.k
+          )
+        ORDER BY t.id, t.k, t.v
+    """
+}