rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-10-29 23:16:57 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: enhance windowed-sort optimizer rule (#4910)

Browse Source 
* add RegionScanner::metadata

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* skip PartSort when there is no tag column

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* add more sqlness test

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* handle desc

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* fix: should keep part sort on DESC

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* fix clippy

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

---------

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>
This commit is contained in:
Ruihang Xia
2024-10-31 06:15:45 +00:00
committed by GitHub
parent ea6df9ba49
commit 8b60c27c2e
11 changed files with 320 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/query/src/optimizer/parallelize_scan.rs
+24 -5
View File
@@ -17,6 +17,7 @@ use std::sync::Arc;
use common_telemetry::debug;
use datafusion::config::ConfigOptions;
use datafusion::physical_optimizer::PhysicalOptimizerRule;
use datafusion::physical_plan::sorts::sort::SortExec;
use datafusion::physical_plan::ExecutionPlan;
use datafusion_common::tree_node::{Transformed, TreeNode};
use datafusion_common::{DataFusionError, Result};
@@ -48,9 +49,16 @@ impl ParallelizeScan {
plan: Arc<dyn ExecutionPlan>,
config: &ConfigOptions,
) -> Result<Arc<dyn ExecutionPlan>> {
let mut first_order_expr = None;
let result = plan
.transform_down(|plan| {
if let Some(region_scan_exec) = plan.as_any().downcast_ref::<RegionScanExec>() {
if let Some(sort_exec) = plan.as_any().downcast_ref::<SortExec>() {
// save the first order expr
first_order_expr = sort_exec.expr().first().cloned();
} else if let Some(region_scan_exec) =
plan.as_any().downcast_ref::<RegionScanExec>()
{
if region_scan_exec.is_partition_set() {
return Ok(Transformed::no(plan));
}
@@ -66,10 +74,21 @@ impl ParallelizeScan {
"Assign {total_range_num} ranges to {expected_partition_num} partitions"
);
// sort the ranges in each partition
// TODO(ruihang): smart sort!
for ranges in partition_ranges.iter_mut() {
ranges.sort_by(|a, b| a.start.cmp(&b.start));
// Sort the ranges in each partition based on the order expr
//
// This optimistically assumes that the first order expr is on the time index column
// to skip the validation of the order expr. As it's not harmful if this condition
// is not met.
if let Some(order_expr) = &first_order_expr
&& order_expr.options.descending
{
for ranges in partition_ranges.iter_mut() {
ranges.sort_by(|a, b| b.end.cmp(&a.end));
}
} else {
for ranges in partition_ranges.iter_mut() {
ranges.sort_by(|a, b| a.start.cmp(&b.start));
}
}
// update the partition ranges
src/query/src/optimizer/windowed_sort.rs
+21 -8
View File
@@ -77,7 +77,6 @@ impl WindowedSortPhysicalRule {
};
if let Some(first_sort_expr) = sort_exec.expr().first()
&& !first_sort_expr.options.descending
&& let Some(column_expr) = first_sort_expr
.expr
.as_any()
@@ -87,18 +86,28 @@ impl WindowedSortPhysicalRule {
} else {
return Ok(Transformed::no(plan));
}
let first_sort_expr = sort_exec.expr().first().unwrap().clone();
let part_sort_exec = Arc::new(PartSortExec::new(
first_sort_expr.clone(),
scanner_info.partition_ranges.clone(),
sort_exec.input().clone(),
));
// PartSortExec is unnecessary if:
// - there is no tag column, and
// - the sort is ascending on the time index column
let new_input = if scanner_info.tag_columns.is_empty()
&& !first_sort_expr.options.descending
{
sort_exec.input().clone()
} else {
Arc::new(PartSortExec::new(
first_sort_expr.clone(),
scanner_info.partition_ranges.clone(),
sort_exec.input().clone(),
))
};
let windowed_sort_exec = WindowedSortExec::try_new(
first_sort_expr,
sort_exec.fetch(),
scanner_info.partition_ranges,
part_sort_exec,
new_input,
)?;
return Ok(Transformed {
@@ -119,11 +128,13 @@ impl WindowedSortPhysicalRule {
struct ScannerInfo {
partition_ranges: Vec<Vec<PartitionRange>>,
time_index: String,
tag_columns: Vec<String>,
}
fn fetch_partition_range(input: Arc<dyn ExecutionPlan>) -> DataFusionResult<Option<ScannerInfo>> {
let mut partition_ranges = None;
let mut time_index = None;
let mut tag_columns = None;
input.transform_up(|plan| {
// Unappliable case, reset the state.
@@ -139,6 +150,7 @@ fn fetch_partition_range(input: Arc<dyn ExecutionPlan>) -> DataFusionResult<Opti
if let Some(region_scan_exec) = plan.as_any().downcast_ref::<RegionScanExec>() {
partition_ranges = Some(region_scan_exec.get_uncollapsed_partition_ranges());
time_index = region_scan_exec.time_index();
tag_columns = Some(region_scan_exec.tag_columns());
// set distinguish_partition_ranges to true, this is an incorrect workaround
region_scan_exec.with_distinguish_partition_range(true);
@@ -151,6 +163,7 @@ fn fetch_partition_range(input: Arc<dyn ExecutionPlan>) -> DataFusionResult<Opti
ScannerInfo {
partition_ranges: partition_ranges?,
time_index: time_index?,
tag_columns: tag_columns?,
}
};
src/query/src/part_sort.rs
+1
View File
@@ -175,6 +175,7 @@ struct PartSortStream {
input_complete: bool,
schema: SchemaRef,
partition_ranges: Vec<PartitionRange>,
#[allow(dead_code)] // this is used under #[debug_assertions]
partition: usize,
cur_part_idx: usize,
metrics: BaselineMetrics,
src/query/src/window_sort.rs
+1
View File
@@ -270,6 +270,7 @@ pub struct WindowedSortStream {
/// working ranges promise once input stream get a value out of current range, future values will never be in this range
all_avail_working_range: Vec<(TimeRange, BTreeSet<usize>)>,
/// The input partition ranges
#[allow(dead_code)] // this is used under #[debug_assertions]
ranges: Vec<PartitionRange>,
/// Execution metrics
metrics: BaselineMetrics,