diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownAggThroughJoinOnPkFk.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownAggThroughJoinOnPkFk.java index f1d6e0816d112c..160578a5dcfb6e 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownAggThroughJoinOnPkFk.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownAggThroughJoinOnPkFk.java @@ -42,6 +42,7 @@ import java.util.BitSet; import java.util.HashMap; import java.util.HashSet; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Map.Entry; @@ -60,6 +61,15 @@ * | Agg(group by fk) * | | * pk fk + * + * Constraints applied on the pattern: + * - Join is Inner Join (no semi/anti/outer join, no mark join). + * - otherJoinConjuncts is empty: only equi-join conditions are allowed, + * and each condition must reference exactly one slot from each side. + * - All GROUP BY expressions are Slot (complex grouping expressions + * are not supported). + * - If an intermediate LogicalProject exists, it must be isAllSlots + * (projections that introduce computations are not supported). */ public class PushDownAggThroughJoinOnPkFk implements RewriteRuleFactory { @Override @@ -100,7 +110,7 @@ public List buildRules() { LogicalAggregate newAgg = eliminatePrimaryOutput(agg, subJoin, primaryAndForeign.first, primaryAndForeign.second); if (newAgg == null) { - return null; + continue; } LogicalJoin newJoin = innerJoinCluster .constructJoinWithPrimary(e.getKey(), subJoin, primaryAndForeign.first); @@ -279,7 +289,7 @@ private LogicalAggregate eliminatePrimaryOutput(LogicalAggregate agg, Plan * a b */ static class InnerJoinCluster { - private final Map> innerJoins = new HashMap<>(); + private final Map> innerJoins = new LinkedHashMap<>(); private final List leaf = new ArrayList<>(); void collectContiguousInnerJoins(Plan plan) { @@ -345,9 +355,25 @@ boolean isValid() { currentPlan = entry.getValue(); forbiddenJoin.add(entry.getKey()); } else if (currentBitset.intersects(entry.getKey())) { + // The new join shares leaves with the current accumulated plan. + // We need to connect the newChild with current plan + BitSet entryBitset = entry.getKey(); + + BitSet newBits = (BitSet) entryBitset.clone(); + newBits.andNot(currentBitset); + LogicalJoin entryJoin = entry.getValue(); + + // Determine the new child: it must be a single leaf + Plan newChild; + if (newBits.cardinality() == 1) { + newChild = leaf.get(newBits.nextSetBit(0)); + } else { + return null; + } + + currentPlan = entryJoin.withChildren(newChild, currentPlan); addJoin = true; - currentBitset.or(entry.getKey()); - currentPlan = currentPlan.withChildren(currentPlan, entry.getValue()); + currentBitset.or(entryBitset); forbiddenJoin.add(entry.getKey()); } } diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/PushDownAggThroughJoinOnPkFkTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/PushDownAggThroughJoinOnPkFkTest.java index 52dd0c55b6133e..f3d2ec147cd9b0 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/PushDownAggThroughJoinOnPkFkTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/PushDownAggThroughJoinOnPkFkTest.java @@ -49,6 +49,46 @@ protected void runBeforeAll() throws Exception { + ")\n" + "DUPLICATE KEY(id3)\n" + "DISTRIBUTED BY HASH(id3) BUCKETS 10\n" + + "PROPERTIES (\"replication_num\" = \"1\")\n", + // Tables for tree-structured multi-join test (like TPC-H Q10): + // t_primary acts as the "bridge" table (like customer): + // - pk(pk_id) + // - fk(fk_to_other) references t_other_primary(other_pk_id) + // t_foreign1 acts like orders: + // - fk(fk_to_primary) references t_primary(pk_id) + // t_foreign2 acts like lineitem (no pk/fk constraints, just for tree structure) + // t_other_primary acts like nation: + // - pk(other_pk_id) + "CREATE TABLE IF NOT EXISTS t_primary (\n" + + " pk_id int not null,\n" + + " name char,\n" + + " fk_to_other int not null\n" + + ")\n" + + "DUPLICATE KEY(pk_id)\n" + + "DISTRIBUTED BY HASH(pk_id) BUCKETS 10\n" + + "PROPERTIES (\"replication_num\" = \"1\")\n", + "CREATE TABLE IF NOT EXISTS t_foreign1 (\n" + + " fk1_id int not null,\n" + + " name char,\n" + + " fk_to_primary int not null\n" + + ")\n" + + "DUPLICATE KEY(fk1_id)\n" + + "DISTRIBUTED BY HASH(fk1_id) BUCKETS 10\n" + + "PROPERTIES (\"replication_num\" = \"1\")\n", + "CREATE TABLE IF NOT EXISTS t_foreign2 (\n" + + " fk2_id int not null,\n" + + " name char,\n" + + " fk_to_foreign1 int not null\n" + + ")\n" + + "DUPLICATE KEY(fk2_id)\n" + + "DISTRIBUTED BY HASH(fk2_id) BUCKETS 10\n" + + "PROPERTIES (\"replication_num\" = \"1\")\n", + "CREATE TABLE IF NOT EXISTS t_other_primary (\n" + + " other_pk_id int not null,\n" + + " name char\n" + + ")\n" + + "DUPLICATE KEY(other_pk_id)\n" + + "DISTRIBUTED BY HASH(other_pk_id) BUCKETS 10\n" + "PROPERTIES (\"replication_num\" = \"1\")\n" ); addConstraint("Alter table pri add constraint pk primary key (id1)"); @@ -56,7 +96,15 @@ protected void runBeforeAll() throws Exception { + "references pri(id1)"); addConstraint("Alter table foreign_null add constraint f_not_null foreign key (id3)\n" + "references pri(id1)"); - connectContext.getSessionVariable().setDisableNereidsRules("PRUNE_EMPTY_PARTITION"); + // Constraints for tree-structured multi-join test + addConstraint("Alter table t_primary add constraint pk_t_primary primary key (pk_id)"); + addConstraint("Alter table t_foreign1 add constraint fk1_to_primary foreign key (fk_to_primary)\n" + + "references t_primary(pk_id)"); + addConstraint("Alter table t_other_primary add constraint pk_other primary key (other_pk_id)"); + addConstraint("Alter table t_primary add constraint fk_primary_to_other foreign key (fk_to_other)\n" + + "references t_other_primary(other_pk_id)"); + connectContext.getSessionVariable().setDisableNereidsRules( + "PRUNE_EMPTY_PARTITION,ELIMINATE_JOIN_BY_FK"); } @Test @@ -179,4 +227,80 @@ void testMissSlot() { .matches(logicalAggregate(logicalProject(logicalJoin()))) .printlnTree(); } + + @Test + void testPushDownAggThroughTreeJoin() { + // Test pushing agg through a tree-structured join (like TPC-H Q10): + // Join(t_primary.fk_to_other = t_other_primary.other_pk_id) ← root (pk-fk: other is pk) + // / \ + // Join(t_foreign1.fk1_id = t_foreign2.fk_to_foreign1) t_other_primary + // / \ + // Join(t_primary.pk_id = t_foreign1.fk_to_primary) t_foreign2 + // / \ + // t_primary t_foreign1 + // The pk-fk join with t_other_primary as pk is recognized, + // and agg should be pushed down through it. + String sql = "select t_primary.pk_id, t_primary.fk_to_other, sum(t_foreign2.fk2_id) " + + "from t_primary " + + "inner join t_foreign1 on t_primary.pk_id = t_foreign1.fk_to_primary " + + "inner join t_foreign2 on t_foreign1.fk1_id = t_foreign2.fk_to_foreign1 " + + "inner join t_other_primary on t_primary.fk_to_other = t_other_primary.other_pk_id " + + "group by t_primary.pk_id, t_primary.fk_to_other"; + PlanChecker.from(connectContext) + .analyze(sql) + .rewrite() + .matches(logicalJoin(logicalAggregate(), any())) + .printlnTree(); + } + + @Test + void testPushDownAggThroughTreeJoinOtherSide() { + // When the pk-fk join's primary is t_primary (the connector table that + // also joins to other tables), the remaining leaves can still form a + // connected chain t_foreign1 ↔ t_foreign2 ↔ t_other_primary via joins + // {1,2} and {2,3}, allowing constructPlan to assemble them into a + // single plan. The agg should be pushed down through the pk-fk join. + // Note: the join to t_other_primary uses t_foreign2.name = t_other_primary.name + // (a non-PK-FK join) so the remaining leaves form a connected chain. + String sql = "select t_primary.pk_id, t_foreign1.fk_to_primary, sum(t_foreign2.fk2_id) " + + "from t_primary " + + "inner join t_foreign1 on t_primary.pk_id = t_foreign1.fk_to_primary " + + "inner join t_foreign2 on t_foreign1.fk1_id = t_foreign2.fk_to_foreign1 " + + "inner join t_other_primary on t_foreign2.name = t_other_primary.name " + + "group by t_primary.pk_id, t_foreign1.fk_to_primary"; + // In this case, pk-fk is t_primary(pk)↔t_foreign1(fk). t_primary is the + // connector table but the remaining leaves {t_foreign1, t_foreign2, t_other_primary} + // are connected (t_foreign1↔t_foreign2, t_foreign2↔t_other_primary). + // constructPlan successfully assembles them, and the agg is pushed down + // below the pk-fk join. + PlanChecker.from(connectContext) + .analyze(sql) + .rewrite() + .matches(logicalJoin(logicalAggregate(), any())) + .printlnTree(); + } + + @Test + void testSkipIneligibleEdgeAndSucceedOnLaterEdge() { + // Regression: pushAgg iterates all PK/FK edges but used to return null + // on the first ineligible candidate instead of continuing to later edges. + // The lower edge {t_primary, t_foreign1} (t_primary is pk) is visited + // first on Java 17 HashMap iteration, and sum(t_primary.pk_id) + // makes eliminatePrimaryOutput return null because the Sum aggregate + // function references a primary-side column that cannot be rewritten. + // With the fix, the rule skips that edge and succeeds on the root edge + // {t_other_primary, t_primary} (t_other_primary is pk), which pushes + // the aggregate through because the query outputs no t_other_primary cols. + String sql = "select t_primary.pk_id, t_primary.fk_to_other, sum(t_primary.pk_id) " + + "from t_primary " + + "inner join t_foreign1 on t_primary.pk_id = t_foreign1.fk_to_primary " + + "inner join t_foreign2 on t_foreign1.fk1_id = t_foreign2.fk_to_foreign1 " + + "inner join t_other_primary on t_primary.fk_to_other = t_other_primary.other_pk_id " + + "group by t_primary.pk_id, t_primary.fk_to_other"; + PlanChecker.from(connectContext) + .analyze(sql) + .rewrite() + .matches(logicalJoin(logicalAggregate(), any())) + .printlnTree(); + } }