From aeb1d1b024173d305827b2ac68d74c1d17a18834 Mon Sep 17 00:00:00 2001 From: Nga Tran Date: Wed, 3 Nov 2021 15:50:49 -0400 Subject: [PATCH 1/4] fix: some debug info --- datafusion/src/physical_plan/hash_aggregate.rs | 14 +++++++++++++- datafusion/src/physical_plan/planner.rs | 4 ++++ datafusion/tests/sql.rs | 7 +++++++ 3 files changed, 24 insertions(+), 1 deletion(-) diff --git a/datafusion/src/physical_plan/hash_aggregate.rs b/datafusion/src/physical_plan/hash_aggregate.rs index 33c68273077a6..7a718c71ef3ee 100644 --- a/datafusion/src/physical_plan/hash_aggregate.rs +++ b/datafusion/src/physical_plan/hash_aggregate.rs @@ -137,6 +137,8 @@ impl HashAggregateExec { input_schema: SchemaRef, ) -> Result { let schema = create_schema(&input.schema(), &group_expr, &aggr_expr, mode)?; + println!(" +++++ Schema created for the aggregate: \n {:#?}", schema.clone() ); + println!(" +++++ Input Schema of the aggregate: \n {:#?}", input_schema.clone()); let schema = Arc::new(schema); @@ -208,6 +210,7 @@ impl ExecutionPlan for HashAggregateExec { } async fn execute(&self, partition: usize) -> Result { + println!(" ============== RUNNING HashAggregateExec"); let input = self.input.execute(partition).await?; let group_expr = self.group_expr.iter().map(|x| x.0.clone()).collect(); @@ -777,7 +780,12 @@ async fn compute_hash_aggregate( // 2. convert values to a record batch let timer = elapsed_compute.timer(); let batch = finalize_aggregation(&accumulators, &mode) - .map(|columns| RecordBatch::try_new(schema.clone(), columns)) + .map(|columns| { + println!(" =================== RecordBatch::try_new in HashAggregare's compute_hash_aggregate"); + println!(" ================ schema: {:#?}", schema.clone()); + println!(" ================ columns: {:#?}", columns); + RecordBatch::try_new(schema.clone(), columns) + }) .map_err(DataFusionError::into_arrow_external_error)?; timer.done(); batch @@ -967,6 +975,10 @@ fn create_batch_from_map( .map(|(col, desired_field)| cast(col, desired_field.data_type())) .collect::>>()?; + println!(" =================== RecordBatch::try_new in HashAggregare's create_batch_from_map"); + println!(" ================ schema: {:#?}", output_schema.clone()); + println!(" ================ columns: {:#?}", columns); + RecordBatch::try_new(Arc::new(output_schema.to_owned()), columns) } diff --git a/datafusion/src/physical_plan/planner.rs b/datafusion/src/physical_plan/planner.rs index fd0421b5e0bad..90fc5b31a14a2 100644 --- a/datafusion/src/physical_plan/planner.rs +++ b/datafusion/src/physical_plan/planner.rs @@ -503,6 +503,7 @@ impl DefaultPhysicalPlanner { }) .collect::>>()?; + println!(" ++++++++++++++++++++++++++ schema to HashAggregateExec in create_initial_plan for AggregateMode::Partial: \n {:#?}", physical_input_schema.clone()); let initial_aggr = Arc::new(HashAggregateExec::try_new( AggregateMode::Partial, groups.clone(), @@ -547,6 +548,7 @@ impl DefaultPhysicalPlanner { (initial_aggr, AggregateMode::Final) }; + println!(" ++++++++++++++++++++++++++ schema to HashAggregateExec in create_initial_plan: \n {:#?}", physical_input_schema.clone()); Ok(Arc::new(HashAggregateExec::try_new( next_partition_mode, final_group @@ -1416,6 +1418,7 @@ impl DefaultPhysicalPlanner { { let optimizers = &ctx_state.config.physical_optimizers; debug!("Physical plan:\n{:?}", plan); + println!("NGA NGA NGA ============== Physical plan:\n{:?}", plan); let mut new_plan = plan; for optimizer in optimizers { @@ -1423,6 +1426,7 @@ impl DefaultPhysicalPlanner { observer(new_plan.as_ref(), optimizer.as_ref()) } debug!("Optimized physical plan:\n{:?}", new_plan); + println!("NGA NGA NGA ============== Optimized physical plan:\n{:?}", new_plan); Ok(new_plan) } } diff --git a/datafusion/tests/sql.rs b/datafusion/tests/sql.rs index 6cd1d3822bce9..d8144df6fa0fc 100644 --- a/datafusion/tests/sql.rs +++ b/datafusion/tests/sql.rs @@ -3946,6 +3946,13 @@ async fn query_on_string_dictionary() -> Result<()> { let mut ctx = ExecutionContext::new(); ctx.register_table("test", Arc::new(table))?; + // aggregation min + let sql = "SELECT MIN(d1) FROM test"; + let actual = execute(&mut ctx, sql).await; + let expected = vec![vec!["one"]]; + assert_eq!(expected, actual); + + // Basic SELECT let sql = "SELECT * FROM test"; let actual = execute(&mut ctx, sql).await; From 577210e26d79af39ae01183139614744c83137b3 Mon Sep 17 00:00:00 2001 From: Nga Tran Date: Wed, 3 Nov 2021 17:27:18 -0400 Subject: [PATCH 2/4] fix: aggregation data type should not be Dictionary --- .../src/physical_plan/expressions/min_max.rs | 17 +++++++++++++---- datafusion/src/physical_plan/hash_aggregate.rs | 10 ---------- datafusion/src/physical_plan/planner.rs | 3 +-- datafusion/tests/sql.rs | 13 ++++++------- 4 files changed, 20 insertions(+), 23 deletions(-) diff --git a/datafusion/src/physical_plan/expressions/min_max.rs b/datafusion/src/physical_plan/expressions/min_max.rs index 97486680f2e09..04d3110f076d5 100644 --- a/datafusion/src/physical_plan/expressions/min_max.rs +++ b/datafusion/src/physical_plan/expressions/min_max.rs @@ -38,6 +38,15 @@ use arrow::{ use super::format_state_name; +// min/max aggregation only returns a value for each group +// and should not be a Dictionary data type anymore +fn min_max_aggregate_data_type(input_type: DataType) -> DataType { + match input_type { + DataType::Dictionary(_, value_type) => (*value_type).clone(), + _ => input_type.clone(), + } +} + /// MAX aggregate expression #[derive(Debug)] pub struct Max { @@ -72,7 +81,7 @@ impl AggregateExpr for Max { fn field(&self) -> Result { Ok(Field::new( &self.name, - self.data_type.clone(), + min_max_aggregate_data_type(self.data_type.clone()), self.nullable, )) } @@ -80,7 +89,7 @@ impl AggregateExpr for Max { fn state_fields(&self) -> Result> { Ok(vec![Field::new( &format_state_name(&self.name, "max"), - self.data_type.clone(), + min_max_aggregate_data_type(self.data_type.clone()), true, )]) } @@ -394,7 +403,7 @@ impl AggregateExpr for Min { fn field(&self) -> Result { Ok(Field::new( &self.name, - self.data_type.clone(), + min_max_aggregate_data_type(self.data_type.clone()), self.nullable, )) } @@ -402,7 +411,7 @@ impl AggregateExpr for Min { fn state_fields(&self) -> Result> { Ok(vec![Field::new( &format_state_name(&self.name, "min"), - self.data_type.clone(), + min_max_aggregate_data_type(self.data_type.clone()), true, )]) } diff --git a/datafusion/src/physical_plan/hash_aggregate.rs b/datafusion/src/physical_plan/hash_aggregate.rs index 7a718c71ef3ee..1721333d88cab 100644 --- a/datafusion/src/physical_plan/hash_aggregate.rs +++ b/datafusion/src/physical_plan/hash_aggregate.rs @@ -137,8 +137,6 @@ impl HashAggregateExec { input_schema: SchemaRef, ) -> Result { let schema = create_schema(&input.schema(), &group_expr, &aggr_expr, mode)?; - println!(" +++++ Schema created for the aggregate: \n {:#?}", schema.clone() ); - println!(" +++++ Input Schema of the aggregate: \n {:#?}", input_schema.clone()); let schema = Arc::new(schema); @@ -210,7 +208,6 @@ impl ExecutionPlan for HashAggregateExec { } async fn execute(&self, partition: usize) -> Result { - println!(" ============== RUNNING HashAggregateExec"); let input = self.input.execute(partition).await?; let group_expr = self.group_expr.iter().map(|x| x.0.clone()).collect(); @@ -781,9 +778,6 @@ async fn compute_hash_aggregate( let timer = elapsed_compute.timer(); let batch = finalize_aggregation(&accumulators, &mode) .map(|columns| { - println!(" =================== RecordBatch::try_new in HashAggregare's compute_hash_aggregate"); - println!(" ================ schema: {:#?}", schema.clone()); - println!(" ================ columns: {:#?}", columns); RecordBatch::try_new(schema.clone(), columns) }) .map_err(DataFusionError::into_arrow_external_error)?; @@ -975,10 +969,6 @@ fn create_batch_from_map( .map(|(col, desired_field)| cast(col, desired_field.data_type())) .collect::>>()?; - println!(" =================== RecordBatch::try_new in HashAggregare's create_batch_from_map"); - println!(" ================ schema: {:#?}", output_schema.clone()); - println!(" ================ columns: {:#?}", columns); - RecordBatch::try_new(Arc::new(output_schema.to_owned()), columns) } diff --git a/datafusion/src/physical_plan/planner.rs b/datafusion/src/physical_plan/planner.rs index 90fc5b31a14a2..3bf43fc71842d 100644 --- a/datafusion/src/physical_plan/planner.rs +++ b/datafusion/src/physical_plan/planner.rs @@ -1418,15 +1418,14 @@ impl DefaultPhysicalPlanner { { let optimizers = &ctx_state.config.physical_optimizers; debug!("Physical plan:\n{:?}", plan); - println!("NGA NGA NGA ============== Physical plan:\n{:?}", plan); let mut new_plan = plan; for optimizer in optimizers { new_plan = optimizer.optimize(new_plan, &ctx_state.config)?; observer(new_plan.as_ref(), optimizer.as_ref()) } + debug!("Optimized physical plan short version:\n{}\n", displayable(new_plan.as_ref()).indent().to_string()); debug!("Optimized physical plan:\n{:?}", new_plan); - println!("NGA NGA NGA ============== Optimized physical plan:\n{:?}", new_plan); Ok(new_plan) } } diff --git a/datafusion/tests/sql.rs b/datafusion/tests/sql.rs index d8144df6fa0fc..0ab376de18f08 100644 --- a/datafusion/tests/sql.rs +++ b/datafusion/tests/sql.rs @@ -3946,13 +3946,6 @@ async fn query_on_string_dictionary() -> Result<()> { let mut ctx = ExecutionContext::new(); ctx.register_table("test", Arc::new(table))?; - // aggregation min - let sql = "SELECT MIN(d1) FROM test"; - let actual = execute(&mut ctx, sql).await; - let expected = vec![vec!["one"]]; - assert_eq!(expected, actual); - - // Basic SELECT let sql = "SELECT * FROM test"; let actual = execute(&mut ctx, sql).await; @@ -3983,6 +3976,12 @@ async fn query_on_string_dictionary() -> Result<()> { let expected = vec![vec!["2"]]; assert_eq!(expected, actual); + // aggregation min + let sql = "SELECT MIN(d1) FROM test"; + let actual = execute(&mut ctx, sql).await; + let expected = vec![vec!["one"]]; + assert_eq!(expected, actual); + // grouping let sql = "SELECT d1, COUNT(*) FROM test group by d1"; let mut actual = execute(&mut ctx, sql).await; From ed33eb5d9d6cdd13e8f273c7bb37da1646ccb0e3 Mon Sep 17 00:00:00 2001 From: Nga Tran Date: Wed, 3 Nov 2021 17:36:40 -0400 Subject: [PATCH 3/4] chore: cleanup --- datafusion/src/physical_plan/hash_aggregate.rs | 4 +--- datafusion/src/physical_plan/planner.rs | 7 ++++--- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/datafusion/src/physical_plan/hash_aggregate.rs b/datafusion/src/physical_plan/hash_aggregate.rs index 1721333d88cab..33c68273077a6 100644 --- a/datafusion/src/physical_plan/hash_aggregate.rs +++ b/datafusion/src/physical_plan/hash_aggregate.rs @@ -777,9 +777,7 @@ async fn compute_hash_aggregate( // 2. convert values to a record batch let timer = elapsed_compute.timer(); let batch = finalize_aggregation(&accumulators, &mode) - .map(|columns| { - RecordBatch::try_new(schema.clone(), columns) - }) + .map(|columns| RecordBatch::try_new(schema.clone(), columns)) .map_err(DataFusionError::into_arrow_external_error)?; timer.done(); batch diff --git a/datafusion/src/physical_plan/planner.rs b/datafusion/src/physical_plan/planner.rs index 3bf43fc71842d..eef6349dd6f95 100644 --- a/datafusion/src/physical_plan/planner.rs +++ b/datafusion/src/physical_plan/planner.rs @@ -503,7 +503,6 @@ impl DefaultPhysicalPlanner { }) .collect::>>()?; - println!(" ++++++++++++++++++++++++++ schema to HashAggregateExec in create_initial_plan for AggregateMode::Partial: \n {:#?}", physical_input_schema.clone()); let initial_aggr = Arc::new(HashAggregateExec::try_new( AggregateMode::Partial, groups.clone(), @@ -548,7 +547,6 @@ impl DefaultPhysicalPlanner { (initial_aggr, AggregateMode::Final) }; - println!(" ++++++++++++++++++++++++++ schema to HashAggregateExec in create_initial_plan: \n {:#?}", physical_input_schema.clone()); Ok(Arc::new(HashAggregateExec::try_new( next_partition_mode, final_group @@ -1424,7 +1422,10 @@ impl DefaultPhysicalPlanner { new_plan = optimizer.optimize(new_plan, &ctx_state.config)?; observer(new_plan.as_ref(), optimizer.as_ref()) } - debug!("Optimized physical plan short version:\n{}\n", displayable(new_plan.as_ref()).indent().to_string()); + debug!( + "Optimized physical plan short version:\n{}\n", + displayable(new_plan.as_ref()).indent().to_string() + ); debug!("Optimized physical plan:\n{:?}", new_plan); Ok(new_plan) } From 11b737c397ec840a742eed26c3595fecccd2410b Mon Sep 17 00:00:00 2001 From: Nga Tran Date: Thu, 4 Nov 2021 10:02:37 -0400 Subject: [PATCH 4/4] refactor: address review comments --- .../src/physical_plan/expressions/min_max.rs | 25 +++++++++++-------- datafusion/src/physical_plan/planner.rs | 2 +- datafusion/tests/sql.rs | 6 +++++ 3 files changed, 21 insertions(+), 12 deletions(-) diff --git a/datafusion/src/physical_plan/expressions/min_max.rs b/datafusion/src/physical_plan/expressions/min_max.rs index 04d3110f076d5..9e5b1e095cd6f 100644 --- a/datafusion/src/physical_plan/expressions/min_max.rs +++ b/datafusion/src/physical_plan/expressions/min_max.rs @@ -38,12 +38,15 @@ use arrow::{ use super::format_state_name; -// min/max aggregation only returns a value for each group -// and should not be a Dictionary data type anymore +// Min/max aggregation can take Dictionary encode input but always produces unpacked +// (aka non Dictionary) output. We need to adjust the output data type to reflect this. +// The reason min/max aggregate produces unpacked output because there is only one +// min/max value per group; there is no needs to keep them Dictionary encode fn min_max_aggregate_data_type(input_type: DataType) -> DataType { - match input_type { - DataType::Dictionary(_, value_type) => (*value_type).clone(), - _ => input_type.clone(), + if let DataType::Dictionary(_, value_type) = input_type { + *value_type + } else { + input_type } } @@ -66,7 +69,7 @@ impl Max { Self { name: name.into(), expr, - data_type, + data_type: min_max_aggregate_data_type(data_type), nullable: true, } } @@ -81,7 +84,7 @@ impl AggregateExpr for Max { fn field(&self) -> Result { Ok(Field::new( &self.name, - min_max_aggregate_data_type(self.data_type.clone()), + self.data_type.clone(), self.nullable, )) } @@ -89,7 +92,7 @@ impl AggregateExpr for Max { fn state_fields(&self) -> Result> { Ok(vec![Field::new( &format_state_name(&self.name, "max"), - min_max_aggregate_data_type(self.data_type.clone()), + self.data_type.clone(), true, )]) } @@ -388,7 +391,7 @@ impl Min { Self { name: name.into(), expr, - data_type, + data_type: min_max_aggregate_data_type(data_type), nullable: true, } } @@ -403,7 +406,7 @@ impl AggregateExpr for Min { fn field(&self) -> Result { Ok(Field::new( &self.name, - min_max_aggregate_data_type(self.data_type.clone()), + self.data_type.clone(), self.nullable, )) } @@ -411,7 +414,7 @@ impl AggregateExpr for Min { fn state_fields(&self) -> Result> { Ok(vec![Field::new( &format_state_name(&self.name, "min"), - min_max_aggregate_data_type(self.data_type.clone()), + self.data_type.clone(), true, )]) } diff --git a/datafusion/src/physical_plan/planner.rs b/datafusion/src/physical_plan/planner.rs index eef6349dd6f95..9c6aeb8f65108 100644 --- a/datafusion/src/physical_plan/planner.rs +++ b/datafusion/src/physical_plan/planner.rs @@ -1424,7 +1424,7 @@ impl DefaultPhysicalPlanner { } debug!( "Optimized physical plan short version:\n{}\n", - displayable(new_plan.as_ref()).indent().to_string() + displayable(new_plan.as_ref()).indent() ); debug!("Optimized physical plan:\n{:?}", new_plan); Ok(new_plan) diff --git a/datafusion/tests/sql.rs b/datafusion/tests/sql.rs index 0ab376de18f08..c1424ce8bd1d7 100644 --- a/datafusion/tests/sql.rs +++ b/datafusion/tests/sql.rs @@ -3982,6 +3982,12 @@ async fn query_on_string_dictionary() -> Result<()> { let expected = vec![vec!["one"]]; assert_eq!(expected, actual); + // aggregation max + let sql = "SELECT MAX(d1) FROM test"; + let actual = execute(&mut ctx, sql).await; + let expected = vec![vec!["three"]]; + assert_eq!(expected, actual); + // grouping let sql = "SELECT d1, COUNT(*) FROM test group by d1"; let mut actual = execute(&mut ctx, sql).await;