From aef127c3a4f7c8bd8c095636423905acfdf02577 Mon Sep 17 00:00:00 2001 From: jackwener Date: Mon, 21 Aug 2023 15:25:06 +0800 Subject: [PATCH] fix: union_distinct shouldn't remove child distinct --- datafusion/expr/src/logical_plan/builder.rs | 22 +++++++++------------ 1 file changed, 9 insertions(+), 13 deletions(-) diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index 29d7571d36fc..238fbb70ba34 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -593,15 +593,8 @@ impl LogicalPlanBuilder { /// Apply a union, removing duplicate rows pub fn union_distinct(self, plan: LogicalPlan) -> Result { - // unwrap top-level Distincts, to avoid duplication - let left_plan: LogicalPlan = match self.plan { - LogicalPlan::Distinct(Distinct { input }) => (*input).clone(), - _ => self.plan, - }; - let right_plan: LogicalPlan = match plan { - LogicalPlan::Distinct(Distinct { input }) => (*input).clone(), - _ => plan, - }; + let left_plan: LogicalPlan = self.plan; + let right_plan: LogicalPlan = plan; Ok(Self::from(LogicalPlan::Distinct(Distinct { input: Arc::new(union(left_plan, right_plan)?), @@ -1629,13 +1622,16 @@ mod tests { .union_distinct(plan.build()?)? .build()?; - // output has only one union let expected = "\ Distinct:\ \n Union\ - \n TableScan: employee_csv projection=[state, salary]\ - \n TableScan: employee_csv projection=[state, salary]\ - \n TableScan: employee_csv projection=[state, salary]\ + \n Distinct:\ + \n Union\ + \n Distinct:\ + \n Union\ + \n TableScan: employee_csv projection=[state, salary]\ + \n TableScan: employee_csv projection=[state, salary]\ + \n TableScan: employee_csv projection=[state, salary]\ \n TableScan: employee_csv projection=[state, salary]"; assert_eq!(expected, format!("{plan:?}"));