query/optimizer/
parallelize_scan.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::BinaryHeap;
16use std::sync::Arc;
17
18use common_telemetry::debug;
19use datafusion::config::ConfigOptions;
20use datafusion::physical_optimizer::PhysicalOptimizerRule;
21use datafusion::physical_plan::ExecutionPlan;
22use datafusion::physical_plan::sorts::sort::SortExec;
23use datafusion_common::tree_node::{Transformed, TreeNode};
24use datafusion_common::{DataFusionError, Result};
25use store_api::region_engine::PartitionRange;
26use table::table::scan::RegionScanExec;
27
28#[derive(Debug)]
29pub struct ParallelizeScan;
30
31impl PhysicalOptimizerRule for ParallelizeScan {
32    fn optimize(
33        &self,
34        plan: Arc<dyn ExecutionPlan>,
35        config: &ConfigOptions,
36    ) -> Result<Arc<dyn ExecutionPlan>> {
37        Self::do_optimize(plan, config)
38    }
39
40    fn name(&self) -> &str {
41        "parallelize_scan"
42    }
43
44    fn schema_check(&self) -> bool {
45        true
46    }
47}
48
49impl ParallelizeScan {
50    fn do_optimize(
51        plan: Arc<dyn ExecutionPlan>,
52        config: &ConfigOptions,
53    ) -> Result<Arc<dyn ExecutionPlan>> {
54        let mut first_order_expr = None;
55
56        let result = plan
57            .transform_down(|plan| {
58                if let Some(sort_exec) = plan.downcast_ref::<SortExec>() {
59                    // save the first order expr
60                    first_order_expr = Some(sort_exec.expr().first()).cloned();
61                } else if let Some(region_scan_exec) = plan.downcast_ref::<RegionScanExec>() {
62                    let expected_partition_num = config.execution.target_partitions;
63                    if region_scan_exec.is_partition_set()
64                        || region_scan_exec.scanner_type().as_str() == "SinglePartition"
65                    {
66                        return Ok(Transformed::no(plan));
67                    }
68
69                    let ranges = region_scan_exec.get_partition_ranges();
70                    let total_range_num = ranges.len();
71
72                    // assign ranges to each partition
73                    let mut partition_ranges =
74                        Self::assign_partition_range(ranges, expected_partition_num);
75                    debug!(
76                        "Assign {total_range_num} ranges to {expected_partition_num} partitions"
77                    );
78
79                    // Sort the ranges in each partition based on the order expr
80                    //
81                    // This optimistically assumes that the first order expr is on the time index column
82                    // to skip the validation of the order expr. As it's not harmful if this condition
83                    // is not met.
84                    if let Some(order_expr) = &first_order_expr
85                        && order_expr.options.descending
86                    {
87                        for ranges in partition_ranges.iter_mut() {
88                            // Primary: end descending (larger end first)
89                            // Secondary: start descending (shorter range first when ends are equal)
90                            ranges.sort_by(|a, b| {
91                                b.end.cmp(&a.end).then_with(|| b.start.cmp(&a.start))
92                            });
93                        }
94                    } else {
95                        for ranges in partition_ranges.iter_mut() {
96                            // Primary: start ascending (smaller start first)
97                            // Secondary: end ascending (shorter range first when starts are equal)
98                            ranges.sort_by(|a, b| {
99                                a.start.cmp(&b.start).then_with(|| a.end.cmp(&b.end))
100                            });
101                        }
102                    }
103
104                    // update the partition ranges
105                    let new_exec = region_scan_exec
106                        .with_new_partitions(partition_ranges, expected_partition_num)
107                        .map_err(|e| DataFusionError::External(e.into_inner()))?;
108                    return Ok(Transformed::yes(Arc::new(new_exec)));
109                }
110
111                // The plan might be modified, but it's modified in-place so we always return
112                // Transformed::no(plan) to indicate there is no "new child"
113                Ok(Transformed::no(plan))
114            })?
115            .data;
116
117        Ok(result)
118    }
119
120    /// Distribute [`PartitionRange`]s to each partition.
121    ///
122    /// Currently we assign ranges to partitions according to their rows so each partition
123    /// has similar number of rows. This method always return `expected_partition_num` partitions.
124    pub fn assign_partition_range(
125        mut ranges: Vec<PartitionRange>,
126        expected_partition_num: usize,
127    ) -> Vec<Vec<PartitionRange>> {
128        if ranges.is_empty() {
129            // Returns a single partition with no range.
130            return vec![vec![]; expected_partition_num];
131        }
132
133        if ranges.len() == 1 {
134            let mut vec = vec![vec![]; expected_partition_num];
135            vec[0] = ranges;
136            return vec;
137        }
138
139        // Sort ranges by number of rows in descending order.
140        ranges.sort_by_key(|b| std::cmp::Reverse(b.num_rows));
141        let mut partition_ranges = vec![vec![]; expected_partition_num];
142
143        #[derive(Eq, PartialEq)]
144        struct HeapNode {
145            num_rows: usize,
146            partition_idx: usize,
147        }
148
149        impl Ord for HeapNode {
150            fn cmp(&self, other: &Self) -> std::cmp::Ordering {
151                // Reverse for min-heap.
152                self.num_rows.cmp(&other.num_rows).reverse()
153            }
154        }
155
156        impl PartialOrd for HeapNode {
157            fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
158                Some(self.cmp(other))
159            }
160        }
161
162        let mut part_heap =
163            BinaryHeap::from_iter((0..expected_partition_num).map(|partition_idx| HeapNode {
164                num_rows: 0,
165                partition_idx,
166            }));
167
168        // Assigns the range to the partition with the smallest number of rows.
169        for range in ranges {
170            let mut node = part_heap.pop().unwrap();
171            let partition_idx = node.partition_idx;
172            node.num_rows += range.num_rows;
173            partition_ranges[partition_idx].push(range);
174            part_heap.push(node);
175        }
176
177        partition_ranges
178    }
179}
180
181#[cfg(test)]
182mod test {
183    use common_time::Timestamp;
184    use common_time::timestamp::TimeUnit;
185
186    use super::*;
187
188    #[test]
189    fn test_assign_partition_range() {
190        let ranges = vec![
191            PartitionRange {
192                start: Timestamp::new(0, TimeUnit::Second),
193                end: Timestamp::new(10, TimeUnit::Second),
194                num_rows: 100,
195                identifier: 1,
196            },
197            PartitionRange {
198                start: Timestamp::new(10, TimeUnit::Second),
199                end: Timestamp::new(20, TimeUnit::Second),
200                num_rows: 200,
201                identifier: 2,
202            },
203            PartitionRange {
204                start: Timestamp::new(20, TimeUnit::Second),
205                end: Timestamp::new(30, TimeUnit::Second),
206                num_rows: 150,
207                identifier: 3,
208            },
209            PartitionRange {
210                start: Timestamp::new(30, TimeUnit::Second),
211                end: Timestamp::new(40, TimeUnit::Second),
212                num_rows: 250,
213                identifier: 4,
214            },
215        ];
216
217        // assign to 2 partitions
218        let expected_partition_num = 2;
219        let result =
220            ParallelizeScan::assign_partition_range(ranges.clone(), expected_partition_num);
221        let expected = vec![
222            vec![
223                PartitionRange {
224                    start: Timestamp::new(30, TimeUnit::Second),
225                    end: Timestamp::new(40, TimeUnit::Second),
226                    num_rows: 250,
227                    identifier: 4,
228                },
229                PartitionRange {
230                    start: Timestamp::new(0, TimeUnit::Second),
231                    end: Timestamp::new(10, TimeUnit::Second),
232                    num_rows: 100,
233                    identifier: 1,
234                },
235            ],
236            vec![
237                PartitionRange {
238                    start: Timestamp::new(10, TimeUnit::Second),
239                    end: Timestamp::new(20, TimeUnit::Second),
240                    num_rows: 200,
241                    identifier: 2,
242                },
243                PartitionRange {
244                    start: Timestamp::new(20, TimeUnit::Second),
245                    end: Timestamp::new(30, TimeUnit::Second),
246                    num_rows: 150,
247                    identifier: 3,
248                },
249            ],
250        ];
251        assert_eq!(result, expected);
252
253        // assign 4 ranges to 5 partitions.
254        let expected_partition_num = 5;
255        let result = ParallelizeScan::assign_partition_range(ranges, expected_partition_num);
256        let expected = vec![
257            vec![PartitionRange {
258                start: Timestamp::new(30, TimeUnit::Second),
259                end: Timestamp::new(40, TimeUnit::Second),
260                num_rows: 250,
261                identifier: 4,
262            }],
263            vec![PartitionRange {
264                start: Timestamp::new(0, TimeUnit::Second),
265                end: Timestamp::new(10, TimeUnit::Second),
266                num_rows: 100,
267                identifier: 1,
268            }],
269            vec![PartitionRange {
270                start: Timestamp::new(10, TimeUnit::Second),
271                end: Timestamp::new(20, TimeUnit::Second),
272                num_rows: 200,
273                identifier: 2,
274            }],
275            vec![],
276            vec![PartitionRange {
277                start: Timestamp::new(20, TimeUnit::Second),
278                end: Timestamp::new(30, TimeUnit::Second),
279                num_rows: 150,
280                identifier: 3,
281            }],
282        ];
283        assert_eq!(result, expected);
284
285        // assign 0 ranges to 5 partitions. Should return 5 empty ranges.
286        let result = ParallelizeScan::assign_partition_range(vec![], 5);
287        assert_eq!(result.len(), 5);
288    }
289
290    #[test]
291    fn test_assign_unbalance_partition_range() {
292        let ranges = vec![
293            PartitionRange {
294                start: Timestamp::new(0, TimeUnit::Second),
295                end: Timestamp::new(10, TimeUnit::Second),
296                num_rows: 100,
297                identifier: 1,
298            },
299            PartitionRange {
300                start: Timestamp::new(10, TimeUnit::Second),
301                end: Timestamp::new(20, TimeUnit::Second),
302                num_rows: 200,
303                identifier: 2,
304            },
305            PartitionRange {
306                start: Timestamp::new(20, TimeUnit::Second),
307                end: Timestamp::new(30, TimeUnit::Second),
308                num_rows: 150,
309                identifier: 3,
310            },
311            PartitionRange {
312                start: Timestamp::new(30, TimeUnit::Second),
313                end: Timestamp::new(40, TimeUnit::Second),
314                num_rows: 2500,
315                identifier: 4,
316            },
317        ];
318
319        // assign to 2 partitions
320        let expected_partition_num = 2;
321        let result =
322            ParallelizeScan::assign_partition_range(ranges.clone(), expected_partition_num);
323        let expected = vec![
324            vec![PartitionRange {
325                start: Timestamp::new(30, TimeUnit::Second),
326                end: Timestamp::new(40, TimeUnit::Second),
327                num_rows: 2500,
328                identifier: 4,
329            }],
330            vec![
331                PartitionRange {
332                    start: Timestamp::new(10, TimeUnit::Second),
333                    end: Timestamp::new(20, TimeUnit::Second),
334                    num_rows: 200,
335                    identifier: 2,
336                },
337                PartitionRange {
338                    start: Timestamp::new(20, TimeUnit::Second),
339                    end: Timestamp::new(30, TimeUnit::Second),
340                    num_rows: 150,
341                    identifier: 3,
342                },
343                PartitionRange {
344                    start: Timestamp::new(0, TimeUnit::Second),
345                    end: Timestamp::new(10, TimeUnit::Second),
346                    num_rows: 100,
347                    identifier: 1,
348                },
349            ],
350        ];
351        assert_eq!(result, expected);
352    }
353}
query/optimizer/parallelize_scan.rs

query/optimizer/
parallelize_scan.rs
