[nereids] pull up join from union all rule

fe/fe-core/src/main/java/org/apache/doris/catalog/constraint/ForeignKeyConstraint.java

@@ -62,6 +62,10 @@ public String getReferencedColumnName(String column) {
         return foreignToReference.get(column);
     }
 
+    public ImmutableMap<String, String> getForeignToReference() {
+        return foreignToReference;
+    }
+
     public Map<Column, Column> getForeignToPrimary(TableIf curTable) {
         ImmutableMap.Builder<Column, Column> columnBuilder = new ImmutableMap.Builder<>();
         TableIf refTable = referencedTable.toTableIf();

fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java

@@ -87,6 +87,7 @@
 import org.apache.doris.nereids.rules.rewrite.PruneOlapScanPartition;
 import org.apache.doris.nereids.rules.rewrite.PruneOlapScanTablet;
 import org.apache.doris.nereids.rules.rewrite.PullUpCteAnchor;
+import org.apache.doris.nereids.rules.rewrite.PullUpJoinFromUnionAll;
 import org.apache.doris.nereids.rules.rewrite.PullUpProjectUnderApply;
 import org.apache.doris.nereids.rules.rewrite.PullUpProjectUnderLimit;
 import org.apache.doris.nereids.rules.rewrite.PullUpProjectUnderTopN;
@@ -288,6 +289,21 @@ public class Rewriter extends AbstractBatchJobExecutor {
 
             // this rule should invoke after infer predicate and push down distinct, and before push down limit
             custom(RuleType.ELIMINATE_JOIN_BY_FOREIGN_KEY, EliminateJoinByFK::new),
+            // this rule should be after topic "Column pruning and infer predicate"
+            topic("Join pull up",
+                    topDown(
+                        new EliminateFilter(),
+                        new PushDownFilterThroughProject(),
+                        new MergeProjects()
+                    ),
+                    topDown(
+                        new PullUpJoinFromUnionAll()
+                    ),
+                    custom(RuleType.COLUMN_PRUNING, ColumnPruning::new),
+                    bottomUp(RuleSet.PUSH_DOWN_FILTERS),
+                    custom(RuleType.ELIMINATE_UNNECESSARY_PROJECT, EliminateUnnecessaryProject::new)
+            ),
+
             topic("Limit optimization",
                     // TODO: the logical plan should not contains any phase information,
                     //       we should refactor like AggregateStrategies, e.g. LimitStrategies,
@@ -340,6 +356,7 @@ public class Rewriter extends AbstractBatchJobExecutor {
                     custom(RuleType.ELIMINATE_SORT, EliminateSort::new),
                     bottomUp(new EliminateEmptyRelation())
             ),
+
             // this rule batch must keep at the end of rewrite to do some plan check
             topic("Final rewrite and check",
                     custom(RuleType.CHECK_DATA_TYPES, CheckDataTypes::new),

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

@@ -249,6 +249,7 @@ public enum RuleType {
 
     // split limit
     SPLIT_LIMIT(RuleTypeClass.REWRITE),
+    PULL_UP_JOIN_FROM_UNIONALL(RuleTypeClass.REWRITE),
     // limit push down
     PUSH_LIMIT_THROUGH_JOIN(RuleTypeClass.REWRITE),
     PUSH_LIMIT_THROUGH_PROJECT_JOIN(RuleTypeClass.REWRITE),

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

@@ -326,7 +326,11 @@ private boolean isPredicateCompatible(BiMap<Slot, Slot> equalSlots, Map<Column,
                         .map(e -> e.rewriteUp(
                                 s -> s instanceof Slot ? primarySlotToForeign.getOrDefault(s, (Slot) s) : s))
                         .collect(Collectors.toSet());
-                return columnWithPredicates.get(fp.getKey()).containsAll(primaryPredicates);
+                if (columnWithPredicates.get(fp.getKey()) == null && !columnWithPredicates.isEmpty()) {
+                    return false;
+                } else {
+                    return columnWithPredicates.get(fp.getKey()).containsAll(primaryPredicates);
+                }
             });
         }
     }

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalSetOperation.java

@@ -40,7 +40,6 @@
 import org.apache.doris.nereids.util.TypeCoercionUtils;
 import org.apache.doris.qe.SessionVariable;
 
-import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableList.Builder;
 
@@ -125,7 +124,6 @@ public List<NamedExpression> buildNewOutputs() {
 
     // If the right child is nullable, need to ensure that the left child is also nullable
     private List<Slot> resetNullableForLeftOutputs() {
-        Preconditions.checkState(children.size() == 2);
         List<Slot> resetNullableForLeftOutputs = new ArrayList<>();
         for (int i = 0; i < child(1).getOutput().size(); ++i) {
             if (child(1).getOutput().get(i).nullable() && !child(0).getOutput().get(i).nullable()) {

regression-test/data/nereids_tpcds_shape_sf1000_p0/shape/query11.out

@@ -2,64 +2,60 @@
 -- !ds_shape_11 --
 PhysicalCteAnchor ( cteId=CTEId#0 )
 --PhysicalCteProducer ( cteId=CTEId#0 )
-----PhysicalUnion
-------PhysicalProject
---------hashAgg[GLOBAL]
-----------PhysicalDistribute
-------------hashAgg[LOCAL]
+----PhysicalProject
+------hashJoin[INNER_JOIN] hashCondition=((PULL_UP_UNIFIED_OUTPUT_ALIAS = customer.c_customer_sk)) otherCondition=()
+--------PhysicalDistribute
+----------PhysicalProject
+------------PhysicalUnion
 --------------PhysicalProject
-----------------hashJoin[INNER_JOIN] hashCondition=((customer.c_customer_sk = store_sales.ss_customer_sk)) otherCondition=() build RFs:RF1 c_customer_sk->[ss_customer_sk]
+----------------hashAgg[GLOBAL]
 ------------------PhysicalDistribute
---------------------hashJoin[INNER_JOIN] hashCondition=((store_sales.ss_sold_date_sk = date_dim.d_date_sk)) otherCondition=() build RFs:RF0 d_date_sk->[ss_sold_date_sk]
+--------------------hashAgg[LOCAL]
 ----------------------PhysicalProject
-------------------------PhysicalOlapScan[store_sales] apply RFs: RF0 RF1
-----------------------PhysicalDistribute
-------------------------PhysicalProject
---------------------------filter(d_year IN (1998, 1999))
-----------------------------PhysicalOlapScan[date_dim]
-------------------PhysicalDistribute
---------------------PhysicalProject
-----------------------PhysicalOlapScan[customer]
-------PhysicalProject
---------hashAgg[GLOBAL]
-----------PhysicalDistribute
-------------hashAgg[LOCAL]
+------------------------hashJoin[INNER_JOIN] hashCondition=((store_sales.ss_sold_date_sk = date_dim.d_date_sk)) otherCondition=() build RFs:RF0 d_date_sk->[ss_sold_date_sk]
+--------------------------PhysicalProject
+----------------------------PhysicalOlapScan[store_sales] apply RFs: RF0
+--------------------------PhysicalDistribute
+----------------------------PhysicalProject
+------------------------------filter(d_year IN (1998, 1999))
+--------------------------------PhysicalOlapScan[date_dim]
 --------------PhysicalProject
-----------------hashJoin[INNER_JOIN] hashCondition=((customer.c_customer_sk = web_sales.ws_bill_customer_sk)) otherCondition=() build RFs:RF3 c_customer_sk->[ws_bill_customer_sk]
+----------------hashAgg[GLOBAL]
 ------------------PhysicalDistribute
---------------------hashJoin[INNER_JOIN] hashCondition=((web_sales.ws_sold_date_sk = date_dim.d_date_sk)) otherCondition=() build RFs:RF2 d_date_sk->[ws_sold_date_sk]
+--------------------hashAgg[LOCAL]
 ----------------------PhysicalProject
-------------------------PhysicalOlapScan[web_sales] apply RFs: RF2 RF3
-----------------------PhysicalDistribute
-------------------------PhysicalProject
---------------------------filter(d_year IN (1998, 1999))
-----------------------------PhysicalOlapScan[date_dim]
-------------------PhysicalDistribute
---------------------PhysicalProject
-----------------------PhysicalOlapScan[customer]
+------------------------hashJoin[INNER_JOIN] hashCondition=((web_sales.ws_sold_date_sk = date_dim.d_date_sk)) otherCondition=() build RFs:RF1 d_date_sk->[ws_sold_date_sk]
+--------------------------PhysicalProject
+----------------------------PhysicalOlapScan[web_sales] apply RFs: RF1
+--------------------------PhysicalDistribute
+----------------------------PhysicalProject
+------------------------------filter(d_year IN (1998, 1999))
+--------------------------------PhysicalOlapScan[date_dim]
+--------PhysicalDistribute
+----------PhysicalProject
+------------PhysicalOlapScan[customer]
 --PhysicalResultSink
 ----PhysicalTopN[MERGE_SORT]
 ------PhysicalDistribute
 --------PhysicalTopN[LOCAL_SORT]
 ----------PhysicalProject
-------------hashJoin[INNER_JOIN] hashCondition=((t_s_firstyear.customer_id = t_w_secyear.customer_id)) otherCondition=((if((year_total > 0.00), (cast(year_total as DECIMALV3(38, 8)) / year_total), 0.000000) > if((year_total > 0.00), (cast(year_total as DECIMALV3(38, 8)) / year_total), 0.000000)))
---------------PhysicalProject
-----------------hashJoin[INNER_JOIN] hashCondition=((t_s_secyear.customer_id = t_s_firstyear.customer_id)) otherCondition=()
-------------------hashJoin[INNER_JOIN] hashCondition=((t_s_firstyear.customer_id = t_w_firstyear.customer_id)) otherCondition=()
---------------------PhysicalDistribute
-----------------------PhysicalProject
-------------------------filter((t_s_firstyear.dyear = 1998) and (t_s_firstyear.sale_type = 's') and (t_s_firstyear.year_total > 0.00))
---------------------------PhysicalCteConsumer ( cteId=CTEId#0 )
---------------------PhysicalDistribute
-----------------------PhysicalProject
-------------------------filter((t_w_firstyear.dyear = 1998) and (t_w_firstyear.sale_type = 'w') and (t_w_firstyear.year_total > 0.00))
---------------------------PhysicalCteConsumer ( cteId=CTEId#0 )
-------------------PhysicalDistribute
---------------------PhysicalProject
-----------------------filter((t_s_secyear.dyear = 1999) and (t_s_secyear.sale_type = 's'))
-------------------------PhysicalCteConsumer ( cteId=CTEId#0 )
+------------hashJoin[INNER_JOIN] hashCondition=((t_s_firstyear.customer_id = t_w_firstyear.customer_id)) otherCondition=((if((year_total > 0.00), (cast(year_total as DECIMALV3(38, 8)) / year_total), 0.000000) > if((year_total > 0.00), (cast(year_total as DECIMALV3(38, 8)) / year_total), 0.000000)))
 --------------PhysicalDistribute
 ----------------PhysicalProject
-------------------filter((t_w_secyear.dyear = 1999) and (t_w_secyear.sale_type = 'w'))
+------------------filter((t_w_firstyear.dyear = 1998) and (t_w_firstyear.sale_type = 'w') and (t_w_firstyear.year_total > 0.00))
 --------------------PhysicalCteConsumer ( cteId=CTEId#0 )
+--------------hashJoin[INNER_JOIN] hashCondition=((t_s_secyear.customer_id = t_s_firstyear.customer_id)) otherCondition=()
+----------------PhysicalDistribute
+------------------PhysicalProject
+--------------------filter((t_s_secyear.dyear = 1999) and (t_s_secyear.sale_type = 's'))
+----------------------PhysicalCteConsumer ( cteId=CTEId#0 )
+----------------hashJoin[INNER_JOIN] hashCondition=((t_s_firstyear.customer_id = t_w_secyear.customer_id)) otherCondition=()
+------------------PhysicalDistribute
+--------------------PhysicalProject
+----------------------filter((t_w_secyear.dyear = 1999) and (t_w_secyear.sale_type = 'w'))
+------------------------PhysicalCteConsumer ( cteId=CTEId#0 )
+------------------PhysicalDistribute
+--------------------PhysicalProject
+----------------------filter((t_s_firstyear.dyear = 1998) and (t_s_firstyear.sale_type = 's') and (t_s_firstyear.year_total > 0.00))
+------------------------PhysicalCteConsumer ( cteId=CTEId#0 )