branch-4.0: [fix](local exchange) Do global shuffle for shuffled join/set operator #60758

be/src/pipeline/exec/hashjoin_build_sink.h

@@ -148,6 +148,10 @@ class HashJoinBuildSinkOperatorX MOCK_REMOVE(final)
         return _join_distribution == TJoinDistributionType::BUCKET_SHUFFLE ||
                _join_distribution == TJoinDistributionType::COLOCATE;
     }
+    bool followed_by_shuffled_operator() const override {
+        return (is_shuffled_operator() && !is_colocated_operator()) ||
+               _followed_by_shuffled_operator;
+    }
     std::vector<bool>& is_null_safe_eq_join() { return _is_null_safe_eq_join; }
 
     bool allow_left_semi_direct_return(RuntimeState* state) const {

be/src/pipeline/exec/hashjoin_probe_operator.h

@@ -153,6 +153,10 @@ class HashJoinProbeOperatorX MOCK_REMOVE(final)
         return _join_distribution == TJoinDistributionType::BUCKET_SHUFFLE ||
                _join_distribution == TJoinDistributionType::COLOCATE;
     }
+    bool followed_by_shuffled_operator() const override {
+        return (is_shuffled_operator() && !is_colocated_operator()) ||
+               _followed_by_shuffled_operator;
+    }
 
     bool need_finalize_variant_column() const { return _need_finalize_variant_column; }
 

be/src/pipeline/exec/operator.h

@@ -163,12 +163,14 @@ class OperatorBase {
      */
     [[nodiscard]] virtual bool is_shuffled_operator() const { return false; }
     /**
-     * Return True if this operator is followed by a shuffled operator.
+     * For multiple children's operators, return true if this is a shuffled operator or this is followed by a shuffled operator (HASH JOIN and SET OPERATION).
+     *
+     * For single child's operators, return true if this operator is followed by a shuffled operator.
      * For example, in the plan fragment:
      *   `UNION` -> `SHUFFLED HASH JOIN`
      * The `SHUFFLED HASH JOIN` is a shuffled operator so the UNION operator is followed by a shuffled operator.
      */
-    [[nodiscard]] bool followed_by_shuffled_operator() const {
+    [[nodiscard]] virtual bool followed_by_shuffled_operator() const {
         return _followed_by_shuffled_operator;
     }
     /**

be/src/pipeline/exec/partitioned_hash_join_probe_operator.h

@@ -164,6 +164,9 @@ class PartitionedHashJoinProbeOperatorX final
     bool is_colocated_operator() const override {
         return _inner_probe_operator->is_colocated_operator();
     }
+    bool followed_by_shuffled_operator() const override {
+        return _inner_probe_operator->followed_by_shuffled_operator();
+    }
 
     void update_operator(const TPlanNode& tnode, bool followed_by_shuffled_operator,
                          bool require_bucket_distribution) override {

be/src/pipeline/exec/partitioned_hash_join_sink_operator.h

@@ -150,6 +150,9 @@ class PartitionedHashJoinSinkOperatorX
     bool is_shuffled_operator() const override {
         return _inner_sink_operator->is_shuffled_operator();
     }
+    bool followed_by_shuffled_operator() const override {
+        return _inner_sink_operator->followed_by_shuffled_operator();
+    }
 
     void update_operator(const TPlanNode& tnode, bool followed_by_shuffled_operator,
                          bool require_bucket_distribution) override {

be/src/pipeline/exec/set_probe_sink_operator.h

@@ -115,6 +115,10 @@ class SetProbeSinkOperatorX final : public DataSinkOperatorX<SetProbeSinkLocalSt
 
     bool is_shuffled_operator() const override { return true; }
     bool is_colocated_operator() const override { return _is_colocate; }
+    bool followed_by_shuffled_operator() const override {
+        return (is_shuffled_operator() && !is_colocated_operator()) ||
+               Base::_followed_by_shuffled_operator;
+    }
 
 private:
     void _finalize_probe(SetProbeSinkLocalState<is_intersect>& local_state);

be/src/pipeline/exec/set_sink_operator.h

@@ -119,6 +119,10 @@ class SetSinkOperatorX final : public DataSinkOperatorX<SetSinkLocalState<is_int
 
     bool is_shuffled_operator() const override { return true; }
     bool is_colocated_operator() const override { return _is_colocate; }
+    bool followed_by_shuffled_operator() const override {
+        return (is_shuffled_operator() && !is_colocated_operator()) ||
+               Base::_followed_by_shuffled_operator;
+    }
 
 private:
     template <class HashTableContext, bool is_intersected>