From 7d90d22beba7a0de38091a0c6f9411d71d455beb Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Tue, 28 Nov 2023 07:51:48 -0500 Subject: [PATCH 1/3] Minor: restore DataFrame test --- datafusion/core/tests/dataframe/mod.rs | 27 +++++++++++++++++++++++++- 1 file changed, 26 insertions(+), 1 deletion(-) diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs index 10f4574020bfa..6109b4c8194c8 100644 --- a/datafusion/core/tests/dataframe/mod.rs +++ b/datafusion/core/tests/dataframe/mod.rs @@ -37,7 +37,7 @@ use datafusion::error::Result; use datafusion::execution::context::SessionContext; use datafusion::prelude::JoinType; use datafusion::prelude::{CsvReadOptions, ParquetReadOptions}; -use datafusion::test_util::parquet_test_data; +use datafusion::test_util::{parquet_test_data, scan_empty}; use datafusion::{assert_batches_eq, assert_batches_sorted_eq}; use datafusion_common::{DataFusionError, ScalarValue, UnnestOptions}; use datafusion_execution::config::SessionConfig; @@ -1597,3 +1597,28 @@ async fn test_array_agg() -> Result<()> { Ok(()) } + +#[tokio::test] +async fn aggregate_with_alias() -> Result<()> { + let ctx = SessionContext::new(); + let state = ctx.state(); + + let schema = Arc::new(Schema::new(vec![ + Field::new("c1", DataType::Utf8, false), + Field::new("c2", DataType::UInt32, false), + ])); + + let plan = scan_empty(None, schema.as_ref(), None)? + .aggregate(vec![col("c1")], vec![sum(col("c2"))])? + .project(vec![col("c1"), sum(col("c2")).alias("total_salary")])? + .build()?; + + let plan = state.optimize(&plan)?; + let physical_plan = state.create_physical_plan(&Arc::new(plan)).await?; + assert_eq!("c1", physical_plan.schema().field(0).name().as_str()); + assert_eq!( + "total_salary", + physical_plan.schema().field(1).name().as_str() + ); + Ok(()) +} From 2ea11a064b5ae7bb22b6a0e757382dee06e864e4 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Tue, 28 Nov 2023 15:01:50 -0500 Subject: [PATCH 2/3] Move test to a better location --- datafusion/core/src/physical_planner.rs | 25 +++++++++++++++++++++++ datafusion/core/tests/dataframe/mod.rs | 27 +------------------------ 2 files changed, 26 insertions(+), 26 deletions(-) diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 09f0e11dc2b55..7d9dc08703e7d 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -2540,6 +2540,31 @@ mod tests { Ok(()) } + #[tokio::test] + async fn aggregate_with_alias() -> Result<()> { + let ctx = SessionContext::new(); + let state = ctx.state(); + + let schema = Arc::new(Schema::new(vec![ + Field::new("c1", DataType::Utf8, false), + Field::new("c2", DataType::UInt32, false), + ])); + + let plan = scan_empty(None, schema.as_ref(), None)? + .aggregate(vec![col("c1")], vec![sum(col("c2"))])? + .project(vec![col("c1"), sum(col("c2")).alias("total_salary")])? + .build()?; + + let plan = state.optimize(&plan)?; + let physical_plan = state.create_physical_plan(&Arc::new(plan)).await?; + assert_eq!("c1", physical_plan.schema().field(0).name().as_str()); + assert_eq!( + "total_salary", + physical_plan.schema().field(1).name().as_str() + ); + Ok(()) + } + #[tokio::test] async fn test_explain() { let schema = Schema::new(vec![Field::new("id", DataType::Int32, false)]); diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs index 6109b4c8194c8..10f4574020bfa 100644 --- a/datafusion/core/tests/dataframe/mod.rs +++ b/datafusion/core/tests/dataframe/mod.rs @@ -37,7 +37,7 @@ use datafusion::error::Result; use datafusion::execution::context::SessionContext; use datafusion::prelude::JoinType; use datafusion::prelude::{CsvReadOptions, ParquetReadOptions}; -use datafusion::test_util::{parquet_test_data, scan_empty}; +use datafusion::test_util::parquet_test_data; use datafusion::{assert_batches_eq, assert_batches_sorted_eq}; use datafusion_common::{DataFusionError, ScalarValue, UnnestOptions}; use datafusion_execution::config::SessionConfig; @@ -1597,28 +1597,3 @@ async fn test_array_agg() -> Result<()> { Ok(()) } - -#[tokio::test] -async fn aggregate_with_alias() -> Result<()> { - let ctx = SessionContext::new(); - let state = ctx.state(); - - let schema = Arc::new(Schema::new(vec![ - Field::new("c1", DataType::Utf8, false), - Field::new("c2", DataType::UInt32, false), - ])); - - let plan = scan_empty(None, schema.as_ref(), None)? - .aggregate(vec![col("c1")], vec![sum(col("c2"))])? - .project(vec![col("c1"), sum(col("c2")).alias("total_salary")])? - .build()?; - - let plan = state.optimize(&plan)?; - let physical_plan = state.create_physical_plan(&Arc::new(plan)).await?; - assert_eq!("c1", physical_plan.schema().field(0).name().as_str()); - assert_eq!( - "total_salary", - physical_plan.schema().field(1).name().as_str() - ); - Ok(()) -} From 8b77b652ed5956a2561488b8e38ab33596a402ae Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Tue, 28 Nov 2023 15:03:16 -0500 Subject: [PATCH 3/3] simplify test --- datafusion/core/src/physical_planner.rs | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 7d9dc08703e7d..e0f1201aea012 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -2542,21 +2542,17 @@ mod tests { #[tokio::test] async fn aggregate_with_alias() -> Result<()> { - let ctx = SessionContext::new(); - let state = ctx.state(); - let schema = Arc::new(Schema::new(vec![ Field::new("c1", DataType::Utf8, false), Field::new("c2", DataType::UInt32, false), ])); - let plan = scan_empty(None, schema.as_ref(), None)? + let logical_plan = scan_empty(None, schema.as_ref(), None)? .aggregate(vec![col("c1")], vec![sum(col("c2"))])? .project(vec![col("c1"), sum(col("c2")).alias("total_salary")])? .build()?; - let plan = state.optimize(&plan)?; - let physical_plan = state.create_physical_plan(&Arc::new(plan)).await?; + let physical_plan = plan(&logical_plan).await?; assert_eq!("c1", physical_plan.schema().field(0).name().as_str()); assert_eq!( "total_salary",