@@ -32,6 +32,7 @@ use arrow::datatypes::{
3232} ;
3333use arrow:: error:: ArrowError ;
3434use arrow:: util:: pretty:: pretty_format_batches;
35+ use arrow_schema:: { SortOptions , TimeUnit } ;
3536use datafusion:: { assert_batches_eq, dataframe} ;
3637use datafusion_functions_aggregate:: count:: { count_all, count_all_window} ;
3738use datafusion_functions_aggregate:: expr_fn:: {
@@ -63,8 +64,8 @@ use datafusion::test_util::{
6364use datafusion_catalog:: TableProvider ;
6465use datafusion_common:: test_util:: { batches_to_sort_string, batches_to_string} ;
6566use datafusion_common:: {
66- assert_contains, Constraint , Constraints , DataFusionError , ParamValues , ScalarValue ,
67- TableReference , UnnestOptions ,
67+ assert_contains, Constraint , Constraints , DFSchema , DataFusionError , ParamValues ,
68+ ScalarValue , TableReference , UnnestOptions ,
6869} ;
6970use datafusion_common_runtime:: SpawnedTask ;
7071use datafusion_datasource:: file_format:: format_as_file_type;
@@ -80,10 +81,16 @@ use datafusion_expr::{
8081 LogicalPlanBuilder , ScalarFunctionImplementation , SortExpr , WindowFrame ,
8182 WindowFrameBound , WindowFrameUnits , WindowFunctionDefinition ,
8283} ;
84+ use datafusion_physical_expr:: aggregate:: AggregateExprBuilder ;
8385use datafusion_physical_expr:: expressions:: Column ;
8486use datafusion_physical_expr:: Partitioning ;
8587use datafusion_physical_expr_common:: physical_expr:: PhysicalExpr ;
86- use datafusion_physical_plan:: { displayable, ExecutionPlanProperties } ;
88+ use datafusion_physical_expr_common:: sort_expr:: PhysicalSortExpr ;
89+ use datafusion_physical_plan:: aggregates:: {
90+ AggregateExec , AggregateMode , PhysicalGroupBy ,
91+ } ;
92+ use datafusion_physical_plan:: empty:: EmptyExec ;
93+ use datafusion_physical_plan:: { displayable, ExecutionPlan , ExecutionPlanProperties } ;
8794
8895// Get string representation of the plan
8996async fn physical_plan_to_string ( df : & DataFrame ) -> String {
@@ -6323,3 +6330,105 @@ async fn test_copy_to_preserves_order() -> Result<()> {
63236330 ) ;
63246331 Ok ( ( ) )
63256332}
6333+
6334+ #[ tokio:: test]
6335+ async fn test_duplicate_state_fields_for_dfschema_construct ( ) -> Result < ( ) > {
6336+ let ctx = SessionContext :: new ( ) ;
6337+
6338+ // Simple schema with just the fields we need
6339+ let file_schema = Arc :: new ( Schema :: new ( vec ! [
6340+ Field :: new(
6341+ "timestamp" ,
6342+ DataType :: Timestamp ( TimeUnit :: Nanosecond , Some ( "UTC" . into( ) ) ) ,
6343+ true ,
6344+ ) ,
6345+ Field :: new( "ticker" , DataType :: Utf8 , true ) ,
6346+ Field :: new( "value" , DataType :: Float64 , true ) ,
6347+ Field :: new( "date" , DataType :: Utf8 , false ) ,
6348+ ] ) ) ;
6349+
6350+ let df_schema = DFSchema :: try_from ( file_schema. clone ( ) ) ?;
6351+
6352+ let timestamp = col ( "timestamp" ) ;
6353+ let value = col ( "value" ) ;
6354+ let ticker = col ( "ticker" ) ;
6355+ let date = col ( "date" ) ;
6356+
6357+ let mock_exec = Arc :: new ( EmptyExec :: new ( file_schema. clone ( ) ) ) ;
6358+
6359+ // Build first_value aggregate
6360+ let first_value = Arc :: new (
6361+ AggregateExprBuilder :: new (
6362+ datafusion_functions_aggregate:: first_last:: first_value_udaf ( ) ,
6363+ vec ! [ ctx. create_physical_expr( value. clone( ) , & df_schema) ?] ,
6364+ )
6365+ . alias ( "first_value(value)" )
6366+ . order_by ( vec ! [ PhysicalSortExpr :: new(
6367+ ctx. create_physical_expr( timestamp. clone( ) , & df_schema) ?,
6368+ SortOptions :: new( false , false ) ,
6369+ ) ] )
6370+ . schema ( file_schema. clone ( ) )
6371+ . build ( )
6372+ . expect ( "Failed to build first_value" ) ,
6373+ ) ;
6374+
6375+ // Build last_value aggregate
6376+ let last_value = Arc :: new (
6377+ AggregateExprBuilder :: new (
6378+ datafusion_functions_aggregate:: first_last:: last_value_udaf ( ) ,
6379+ vec ! [ ctx. create_physical_expr( value. clone( ) , & df_schema) ?] ,
6380+ )
6381+ . alias ( "last_value(value)" )
6382+ . order_by ( vec ! [ PhysicalSortExpr :: new(
6383+ ctx. create_physical_expr( timestamp. clone( ) , & df_schema) ?,
6384+ SortOptions :: new( false , false ) ,
6385+ ) ] )
6386+ . schema ( file_schema. clone ( ) )
6387+ . build ( )
6388+ . expect ( "Failed to build last_value" ) ,
6389+ ) ;
6390+
6391+ let partial_agg = AggregateExec :: try_new (
6392+ AggregateMode :: Partial ,
6393+ PhysicalGroupBy :: new_single ( vec ! [
6394+ (
6395+ ctx. create_physical_expr( date. clone( ) , & df_schema) ?,
6396+ "date" . to_string( ) ,
6397+ ) ,
6398+ (
6399+ ctx. create_physical_expr( ticker. clone( ) , & df_schema) ?,
6400+ "ticker" . to_string( ) ,
6401+ ) ,
6402+ ] ) ,
6403+ vec ! [ first_value, last_value] ,
6404+ vec ! [ None , None ] ,
6405+ mock_exec,
6406+ file_schema,
6407+ )
6408+ . expect ( "Failed to build partial agg" ) ;
6409+
6410+ // Assert that the schema field names match the expected names
6411+ let expected_field_names = vec ! [
6412+ "date" ,
6413+ "ticker" ,
6414+ "first_value(value)[first_value]" ,
6415+ "timestamp@0" ,
6416+ "is_set" ,
6417+ "last_value(value)[last_value]" ,
6418+ "timestamp@0" ,
6419+ "is_set" ,
6420+ ] ;
6421+
6422+ let binding = partial_agg. schema ( ) ;
6423+ let actual_field_names: Vec < _ > = binding. fields ( ) . iter ( ) . map ( |f| f. name ( ) ) . collect ( ) ;
6424+ assert_eq ! ( actual_field_names, expected_field_names) ;
6425+
6426+ // Ensure that DFSchema::try_from does not fail
6427+ let partial_agg_exec_schema = DFSchema :: try_from ( partial_agg. schema ( ) ) ;
6428+ assert ! (
6429+ partial_agg_exec_schema. is_ok( ) ,
6430+ "Expected get AggregateExec schema to succeed with duplicate state fields"
6431+ ) ;
6432+
6433+ Ok ( ( ) )
6434+ }
0 commit comments