datafusion_physical_optimizer/
hash_join_buffering.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use crate::PhysicalOptimizerRule;
19use datafusion_common::JoinSide;
20use datafusion_common::config::ConfigOptions;
21use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
22use datafusion_physical_plan::ExecutionPlan;
23use datafusion_physical_plan::buffer::BufferExec;
24use datafusion_physical_plan::execution_plan::replace_children_if_necessary;
25use datafusion_physical_plan::joins::HashJoinExec;
26use std::sync::Arc;
27
28/// Looks for all the [HashJoinExec]s in the plan and places a [BufferExec] node with the
29/// configured capacity in the probe side:
30///
31/// ```text
32///            ┌───────────────────┐
33///            │   HashJoinExec    │
34///            └─────▲────────▲────┘
35///          ┌───────┘        └─────────┐
36///          │                          │
37/// ┌────────────────┐         ┌─────────────────┐
38/// │   Build side   │       + │   BufferExec    │
39/// └────────────────┘         └────────▲────────┘
40///                                     │
41///                            ┌────────┴────────┐
42///                            │   Probe side    │
43///                            └─────────────────┘
44/// ```
45///
46/// Which allows eagerly pulling it even before the build side has completely finished.
47#[derive(Debug, Default)]
48pub struct HashJoinBuffering {}
49
50impl HashJoinBuffering {
51    pub fn new() -> Self {
52        Self::default()
53    }
54}
55
56impl PhysicalOptimizerRule for HashJoinBuffering {
57    fn optimize(
58        &self,
59        plan: Arc<dyn ExecutionPlan>,
60        config: &ConfigOptions,
61    ) -> datafusion_common::Result<Arc<dyn ExecutionPlan>> {
62        let capacity = config.execution.hash_join_buffering_capacity;
63        if capacity == 0 {
64            return Ok(plan);
65        }
66
67        plan.transform_down(|plan| {
68            let Some(node) = plan.downcast_ref::<HashJoinExec>() else {
69                return Ok(Transformed::no(plan));
70            };
71            let plan = Arc::clone(&plan);
72            Ok(Transformed::yes(
73                if HashJoinExec::probe_side() == JoinSide::Left {
74                    // Do not stack BufferExec nodes together.
75                    if node.left.is::<BufferExec>() {
76                        return Ok(Transformed::no(plan));
77                    }
78                    replace_children_if_necessary(
79                        plan,
80                        vec![
81                            Arc::new(BufferExec::new(Arc::clone(&node.left), capacity)),
82                            Arc::clone(&node.right),
83                        ],
84                    )?
85                } else {
86                    // Do not stack BufferExec nodes together.
87                    if node.right.is::<BufferExec>() {
88                        return Ok(Transformed::no(plan));
89                    }
90                    replace_children_if_necessary(
91                        plan,
92                        vec![
93                            Arc::clone(&node.left),
94                            Arc::new(BufferExec::new(Arc::clone(&node.right), capacity)),
95                        ],
96                    )?
97                },
98            ))
99        })
100        .data()
101    }
102
103    fn name(&self) -> &str {
104        "HashJoinBuffering"
105    }
106
107    fn schema_check(&self) -> bool {
108        true
109    }
110}
datafusion_physical_optimizer/hash_join_buffering.rs

datafusion_physical_optimizer/
hash_join_buffering.rs
