diff --git a/src/service/search/datafusion/optimizer/add_sort_and_limit.rs b/src/service/search/datafusion/optimizer/add_sort_and_limit.rs index bc5c7ce786..6609038f96 100644 --- a/src/service/search/datafusion/optimizer/add_sort_and_limit.rs +++ b/src/service/search/datafusion/optimizer/add_sort_and_limit.rs @@ -54,6 +54,9 @@ impl OptimizerRule for AddSortAndLimitRule { plan: LogicalPlan, _config: &dyn OptimizerConfig, ) -> Result> { + if self.limit == 0 { + return Ok(Transformed::new(plan, false, TreeNodeRecursion::Stop)); + } let mut plan = plan.rewrite(&mut AddSortAndLimit::new(self.limit, self.offset))?; plan.tnr = TreeNodeRecursion::Stop; Ok(plan) diff --git a/src/service/search/datafusion/optimizer/limit_join_right_side.rs b/src/service/search/datafusion/optimizer/limit_join_right_side.rs index 351371f74a..7b1ff135ed 100644 --- a/src/service/search/datafusion/optimizer/limit_join_right_side.rs +++ b/src/service/search/datafusion/optimizer/limit_join_right_side.rs @@ -15,18 +15,20 @@ use std::sync::Arc; +use config::TIMESTAMP_COL_NAME; use datafusion::{ common::{ - Result, - tree_node::{Transformed, TreeNode}, + Column, Result, + tree_node::{Transformed, TreeNode, TreeNodeRecursion, TreeNodeRewriter}, }, - logical_expr::LogicalPlan, + logical_expr::{Extension, LogicalPlan, Sort, SortExpr}, optimizer::{OptimizerConfig, OptimizerRule, optimizer::ApplyOrder}, - prelude::Expr, + prelude::{Expr, col}, }; use itertools::Itertools; -use super::utils::AddSortAndLimit; +use super::utils::{AddSortAndLimit, is_contain_deduplication_plan}; +use crate::service::search::datafusion::plan::deduplication::DeduplicationLogicalNode; #[derive(Default, Debug)] pub struct LimitJoinRightSide { @@ -57,6 +59,9 @@ impl OptimizerRule for LimitJoinRightSide { plan: LogicalPlan, _config: &dyn OptimizerConfig, ) -> Result> { + if self.limit == 0 { + return Ok(Transformed::new(plan, false, TreeNodeRecursion::Stop)); + } match plan { LogicalPlan::Join(mut join) => { let right_column = join @@ -70,31 +75,104 @@ impl OptimizerRule for LimitJoinRightSide { } }) .collect_vec(); - if right_column.is_empty() { - let plan = (*join.right) - .clone() - .rewrite(&mut AddSortAndLimit::new(self.limit, 0))? + // limit the right side output size + let mut plan = (*join.right) + .clone() + .rewrite(&mut AddSortAndLimit::new(self.limit, 0))? + .data; + if !right_column.is_empty() { + // deduplication on join key + plan = plan + .rewrite(&mut DeduplicationRewriter::new(right_column))? .data; - join.right = Arc::new(plan); - Ok(Transformed::yes(LogicalPlan::Join(join))) - } else { - let plan = (*join.right) - .clone() - .rewrite(&mut AddSortAndLimit::new_with_deduplication( - self.limit, - 0, - right_column, - ))? - .data; - join.right = Arc::new(plan); - Ok(Transformed::yes(LogicalPlan::Join(join))) } + join.right = Arc::new(plan); + Ok(Transformed::yes(LogicalPlan::Join(join))) } _ => Ok(Transformed::no(plan)), } } } +struct DeduplicationRewriter { + pub deduplication_columns: Vec, +} + +impl DeduplicationRewriter { + pub fn new(deduplication_columns: Vec) -> Self { + Self { + deduplication_columns, + } + } +} + +impl TreeNodeRewriter for DeduplicationRewriter { + type Node = LogicalPlan; + + fn f_down(&mut self, node: LogicalPlan) -> Result> { + if is_contain_deduplication_plan(&node) { + return Ok(Transformed::new(node, false, TreeNodeRecursion::Stop)); + } + + // insert deduplication to first plan that contains deduplication columns + let plan = match node { + LogicalPlan::Projection(_) | LogicalPlan::SubqueryAlias(_) => { + let schema = node.inputs().first().unwrap().schema(); + for column in self.deduplication_columns.iter() { + if schema.field_with_name(None, column.name()).is_err() { + let plan = generate_deduplication_plan( + Arc::new(node), + self.deduplication_columns.clone(), + ); + return Ok(Transformed::new(plan, true, TreeNodeRecursion::Stop)); + } + } + Transformed::no(node) + } + _ => { + let plan = + generate_deduplication_plan(Arc::new(node), self.deduplication_columns.clone()); + Transformed::new(plan, true, TreeNodeRecursion::Stop) + } + }; + + Ok(plan) + } +} + +fn generate_deduplication_plan( + node: Arc, + deduplication_columns: Vec, +) -> LogicalPlan { + let mut sort_columns = Vec::with_capacity(deduplication_columns.len() + 1); + let schema = node.schema().clone(); + + for column in deduplication_columns.iter() { + sort_columns.push(SortExpr { + expr: col(column.name()), + asc: false, + nulls_first: false, + }); + } + + if schema.field_with_name(None, TIMESTAMP_COL_NAME).is_ok() { + sort_columns.push(SortExpr { + expr: col(TIMESTAMP_COL_NAME.to_string()), + asc: false, + nulls_first: false, + }); + } + + let sort = LogicalPlan::Sort(Sort { + expr: sort_columns, + input: node, + fetch: None, + }); + LogicalPlan::Extension(Extension { + node: Arc::new(DeduplicationLogicalNode::new(sort, deduplication_columns)), + }) +} + #[cfg(test)] mod tests { use std::sync::Arc; diff --git a/src/service/search/datafusion/optimizer/utils.rs b/src/service/search/datafusion/optimizer/utils.rs index 84bc345953..56ed1046b8 100644 --- a/src/service/search/datafusion/optimizer/utils.rs +++ b/src/service/search/datafusion/optimizer/utils.rs @@ -18,22 +18,18 @@ use std::sync::Arc; use config::TIMESTAMP_COL_NAME; use datafusion::{ common::{ - Column, DFSchema, Result, + DFSchema, Result, tree_node::{ Transformed, TransformedResult, TreeNode, TreeNodeRecursion, TreeNodeRewriter, }, }, datasource::DefaultTableSource, - logical_expr::{ - Extension, Limit, LogicalPlan, Projection, Sort, SortExpr, TableScan, TableSource, col, - }, + logical_expr::{Limit, LogicalPlan, Projection, Sort, SortExpr, TableScan, TableSource, col}, prelude::Expr, scalar::ScalarValue, }; -use crate::service::search::datafusion::{ - plan::deduplication::DeduplicationLogicalNode, table_provider::empty_table::NewEmptyTable, -}; +use crate::service::search::datafusion::table_provider::empty_table::NewEmptyTable; // check if the plan is a complex query that we can't add sort _timestamp pub fn is_complex_query(plan: &LogicalPlan) -> bool { @@ -53,28 +49,11 @@ pub fn is_complex_query(plan: &LogicalPlan) -> bool { pub struct AddSortAndLimit { pub limit: usize, pub offset: usize, - pub deduplication_columns: Vec, } impl AddSortAndLimit { pub fn new(limit: usize, offset: usize) -> Self { - Self { - limit, - offset, - deduplication_columns: vec![], - } - } - - pub fn new_with_deduplication( - limit: usize, - offset: usize, - deduplication_columns: Vec, - ) -> Self { - Self { - limit, - offset, - deduplication_columns, - } + Self { limit, offset } } } @@ -82,9 +61,6 @@ impl TreeNodeRewriter for AddSortAndLimit { type Node = LogicalPlan; fn f_down(&mut self, node: LogicalPlan) -> Result> { - if self.limit == 0 { - return Ok(Transformed::new(node, false, TreeNodeRecursion::Stop)); - } if is_contain_deduplication_plan(&node) { return Ok(Transformed::new(node, false, TreeNodeRecursion::Stop)); } @@ -145,43 +121,6 @@ impl TreeNodeRewriter for AddSortAndLimit { transformed.tnr = TreeNodeRecursion::Stop; } - // support deduplication on join key - // sort -> deduplication - // only add when is_stop == true - if !self.deduplication_columns.is_empty() && is_stop { - let mut sort_columns = Vec::with_capacity(self.deduplication_columns.len() + 1); - let schema = transformed.data.schema().clone(); - - for column in self.deduplication_columns.iter() { - sort_columns.push(SortExpr { - expr: col(column.name()), - asc: false, - nulls_first: false, - }); - } - - if schema.field_with_name(None, TIMESTAMP_COL_NAME).is_ok() { - sort_columns.push(SortExpr { - expr: col(TIMESTAMP_COL_NAME.to_string()), - asc: false, - nulls_first: false, - }); - } - - let sort = LogicalPlan::Sort(Sort { - expr: sort_columns, - input: Arc::new(transformed.data), - fetch: None, - }); - let dedup = LogicalPlan::Extension(Extension { - node: Arc::new(DeduplicationLogicalNode::new( - sort, - self.deduplication_columns.clone(), - )), - }); - transformed.data = dedup; - } - if let Some(schema) = schema { let plan = transformed.data; let proj = LogicalPlan::Projection(Projection::new_from_schema(Arc::new(plan), schema)); @@ -358,7 +297,7 @@ fn generate_table_source_with_sorted_by_time( } } -fn is_contain_deduplication_plan(plan: &LogicalPlan) -> bool { +pub fn is_contain_deduplication_plan(plan: &LogicalPlan) -> bool { plan.exists(|plan| Ok(matches!(plan, LogicalPlan::Extension(_)))) .unwrap() } diff --git a/src/service/search/datafusion/plan/deduplication_exec.rs b/src/service/search/datafusion/plan/deduplication_exec.rs index 729141012f..6f6d8614d5 100644 --- a/src/service/search/datafusion/plan/deduplication_exec.rs +++ b/src/service/search/datafusion/plan/deduplication_exec.rs @@ -19,8 +19,11 @@ use std::{ task::{Context, Poll}, }; -use arrow::array::{BooleanArray, Float64Array, Int64Array, RecordBatch, StringArray, UInt64Array}; -use arrow_schema::{DataType, SortOptions}; +use arrow::array::{ + BooleanArray, Float64Array, Int64Array, RecordBatch, StringArray, TimestampMicrosecondArray, + UInt64Array, +}; +use arrow_schema::{DataType, SortOptions, TimeUnit}; use config::TIMESTAMP_COL_NAME; use datafusion::{ arrow::datatypes::SchemaRef, @@ -255,6 +258,7 @@ enum Array { UInt64(UInt64Array), Boolean(BooleanArray), Float64(Float64Array), + TimestampMicrosecond(TimestampMicrosecondArray), } impl Array { @@ -265,6 +269,7 @@ impl Array { Array::UInt64(array) => Value::UInt64(array.value(i)), Array::Boolean(array) => Value::Boolean(array.value(i)), Array::Float64(array) => Value::Float64(array.value(i)), + Array::TimestampMicrosecond(array) => Value::Int64(array.value(i)), } } } @@ -318,7 +323,16 @@ fn generate_deduplication_arrays( .unwrap() .clone(), ), - _ => panic!("Unsupported data type"), + DataType::Timestamp(TimeUnit::Microsecond, None) => Array::TimestampMicrosecond( + array + .as_any() + .downcast_ref::() + .unwrap() + .clone(), + ), + _ => { + panic!("Unsupported data type: {}", array.data_type()); + } } }) .collect_vec();