fix: a corner case for join match one (#7016)

fix the below query
```sql
with abc as (select histogram(_timestamp, '1 minute') AS zo_sql_key, count(*) AS zo_sql_num from default GROUP BY zo_sql_key ORDER BY zo_sql_key), edf as (select histogram(_timestamp, '1 minute') AS zo_sql_key, sum(_timestamp) AS zo_sql_num from default GROUP BY zo_sql_key ORDER BY zo_sql_key) select abc.zo_sql_key from abc where zo_sql_key in (select zo_sql_key from edf)
```

---------

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

src/service/search/datafusion/optimizer/add_sort_and_limit.rs

@@ -54,6 +54,9 @@ impl OptimizerRule for AddSortAndLimitRule {
         plan: LogicalPlan,
         _config: &dyn OptimizerConfig,
     ) -> Result<Transformed<LogicalPlan>> {
+        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)

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<Transformed<LogicalPlan>> {
+        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<Column>,
+}
+
+impl DeduplicationRewriter {
+    pub fn new(deduplication_columns: Vec<Column>) -> Self {
+        Self {
+            deduplication_columns,
+        }
+    }
+}
+
+impl TreeNodeRewriter for DeduplicationRewriter {
+    type Node = LogicalPlan;
+
+    fn f_down(&mut self, node: LogicalPlan) -> Result<Transformed<LogicalPlan>> {
+        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<LogicalPlan>,
+    deduplication_columns: Vec<Column>,
+) -> 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;

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<Column>,
 }
 
 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<Column>,
-    ) -> 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<Transformed<LogicalPlan>> {
-        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()
 }

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::<TimestampMicrosecondArray>()
+                        .unwrap()
+                        .clone(),
+                ),
+                _ => {
+                    panic!("Unsupported data type: {}", array.data_type());
+                }
             }
         })
         .collect_vec();