[fix](union) Local shuffle for union operator (apache#56048)

Author: Gabriel39

be/src/pipeline/exec/union_sink_operator.cpp

@@ -57,7 +57,9 @@ UnionSinkOperatorX::UnionSinkOperatorX(int child_id, int sink_id, ObjectPool* po
           _first_materialized_child_idx(tnode.union_node.first_materialized_child_idx),
           _row_descriptor(descs, tnode.row_tuples, tnode.nullable_tuples),
           _cur_child_id(child_id),
-          _child_size(tnode.num_children) {}
+          _child_size(tnode.num_children),
+          _distribute_exprs(tnode.__isset.distribute_expr_lists ? tnode.distribute_expr_lists[0]
+                                                                : std::vector<TExpr> {}) {}
 
 Status UnionSinkOperatorX::init(const TPlanNode& tnode, RuntimeState* state) {
     RETURN_IF_ERROR(DataSinkOperatorX::init(tnode, state));

be/src/pipeline/exec/union_sink_operator.h

@@ -94,6 +94,16 @@ class UnionSinkOperatorX final : public DataSinkOperatorX<UnionSinkLocalState> {
         return _followed_by_shuffled_operator;
     }
 
+    DataDistribution required_data_distribution() const override {
+        if (_child->is_serial_operator() && _followed_by_shuffled_operator) {
+            return DataDistribution(ExchangeType::HASH_SHUFFLE, _distribute_exprs);
+        }
+        if (_child->is_serial_operator()) {
+            return DataDistribution(ExchangeType::PASSTHROUGH);
+        }
+        return DataDistribution(ExchangeType::NOOP);
+    }
+
     bool is_shuffled_operator() const override { return _followed_by_shuffled_operator; }
 
 private:
@@ -113,6 +123,7 @@ class UnionSinkOperatorX final : public DataSinkOperatorX<UnionSinkLocalState> {
     const RowDescriptor _row_descriptor;
     const int _cur_child_id;
     const int _child_size;
+    const std::vector<TExpr> _distribute_exprs;
     int children_count() const { return _child_size; }
     bool is_child_passthrough(int child_idx) const {
         DCHECK_LT(child_idx, _child_size);
@@ -152,4 +163,4 @@ class UnionSinkOperatorX final : public DataSinkOperatorX<UnionSinkLocalState> {
 };
 
 } // namespace pipeline
-} // namespace doris
+} // namespace doris