Skip to content

Commit 52690c6

Browse files
zhuqi-lucasalamb
andauthored
fix: Partial AggregateMode will generate duplicate field names which will fail DFSchema construct (#17706)
* fix: Partial AggregateMode will generate duplicate field names which will fail DFSchema construct * Update datafusion/common/src/dfschema.rs Co-authored-by: Andrew Lamb <[email protected]> * fmt --------- Co-authored-by: Andrew Lamb <[email protected]>
1 parent 79fc94a commit 52690c6

2 files changed

Lines changed: 117 additions & 4 deletions

File tree

‎datafusion/common/src/dfschema.rs‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1129,7 +1129,11 @@ impl TryFrom<SchemaRef> for DFSchema {
11291129
field_qualifiers: vec![None; field_count],
11301130
functional_dependencies: FunctionalDependencies::empty(),
11311131
};
1132-
dfschema.check_names()?;
1132+
// Without checking names, because schema here may have duplicate field names.
1133+
// For example, Partial AggregateMode will generate duplicate field names from
1134+
// state_fields.
1135+
// See <https://github.com/apache/datafusion/issues/17715>
1136+
// dfschema.check_names()?;
11331137
Ok(dfschema)
11341138
}
11351139
}

‎datafusion/core/tests/dataframe/mod.rs‎

Lines changed: 112 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ use arrow::datatypes::{
3232
};
3333
use arrow::error::ArrowError;
3434
use arrow::util::pretty::pretty_format_batches;
35+
use arrow_schema::{SortOptions, TimeUnit};
3536
use datafusion::{assert_batches_eq, dataframe};
3637
use datafusion_functions_aggregate::count::{count_all, count_all_window};
3738
use datafusion_functions_aggregate::expr_fn::{
@@ -63,8 +64,8 @@ use datafusion::test_util::{
6364
use datafusion_catalog::TableProvider;
6465
use datafusion_common::test_util::{batches_to_sort_string, batches_to_string};
6566
use 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
};
6970
use datafusion_common_runtime::SpawnedTask;
7071
use 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;
8385
use datafusion_physical_expr::expressions::Column;
8486
use datafusion_physical_expr::Partitioning;
8587
use 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
8996
async 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

Comments
 (0)