Skip to content

Commit 9eeb66f

Browse files
buraksennalamb
buraksenn
and
alamb
authored
fix: do not drop GROUP BY keys using nullable UNIQUE constraints (#26220)
## Which issue does this PR close? - Closes #23819. ## Rationale for this change A nullable UNIQUE column permits multiple NULL rows, so it cannot safely determine other grouping columns. Dropping those columns can merge distinct groups and return incorrect results. ## What changes are included in this PR? - Only rewrite GROUP BY expressions using dependencies valid across NULLs. - Apply the same check to DISTINCT elimination to preserve correct results. - Add regression tests and an upgrade guide note. ## What is the testing strategy for this PR? SLT and plan-builder tests cover nullable and composite UNIQUE keys, DISTINCT after GROUP BY, and valid PRIMARY KEY / NOT NULL UNIQUE cases. The DISTINCT regressions fail without the fix. ## Are there any user-facing changes? Yes: incorrect results are fixed --------- Co-authored-by: Andrew Lamb <[email protected]>
1 parent 6e69fe9 commit 9eeb66f

6 files changed

Lines changed: 307 additions & 79 deletions

File tree

‎datafusion/common/src/functional_dependencies.rs‎

Lines changed: 41 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -472,9 +472,14 @@ pub fn aggregate_functional_dependencies(
472472
// The loop below only re-expresses input dependencies. Skip it when the
473473
// input has none. The GROUP BY-key dependency below always runs.
474474
if !func_dependencies.is_empty() {
475-
// Compute once: this does not change in the loop.
476-
let existing_target_indices =
477-
get_target_functional_dependencies(aggr_input_schema, group_by_input_indices);
475+
// Compare the columns covered by every dependency, whatever its NULL
476+
// semantics: leaving out the ones that do not hold across NULLs would
477+
// hide GROUP BY expressions this determinant does not cover and
478+
// upgrade its mode to `Single` wrongly.
479+
let existing_target_indices = get_target_indices_from_dependencies(
480+
func_dependencies.iter(),
481+
group_by_input_indices,
482+
);
478483
for FunctionalDependence {
479484
source_indices,
480485
nullable,
@@ -500,8 +505,8 @@ pub fn aggregate_functional_dependencies(
500505
new_source_input_indices.push(Some(*input_idx));
501506
}
502507
}
503-
let new_target_indices = get_target_functional_dependencies(
504-
aggr_input_schema,
508+
let new_target_indices = get_target_indices_from_dependencies(
509+
func_dependencies.iter(),
505510
&new_source_input_indices,
506511
);
507512
let mode = if existing_target_indices == new_target_indices
@@ -572,6 +577,12 @@ pub fn aggregate_functional_dependencies(
572577
/// Returns target indices, for the determinant keys that are inside
573578
/// group by expressions.
574579
///
580+
/// Only dependencies that hold when NULL determinant values are grouped
581+
/// together are considered (see
582+
/// [`FunctionalDependence::is_valid_across_nulls`]): a nullable `UNIQUE`
583+
/// column permits multiple NULL rows that may differ in the other columns,
584+
/// so it does not determine them across the single NULL group.
585+
///
575586
/// `group_by_input_indices` holds, for each GROUP BY expression, the index of
576587
/// the `schema` field it references, or `None` for a computed expression.
577588
pub fn get_target_functional_dependencies(
@@ -582,21 +593,27 @@ pub fn get_target_functional_dependencies(
582593
if dependencies.is_empty() {
583594
return None;
584595
}
596+
let dependencies = dependencies
597+
.iter()
598+
.filter(|dependency| dependency.is_valid_across_nulls(schema));
599+
get_target_indices_from_dependencies(dependencies, group_by_input_indices)
600+
}
601+
602+
fn get_target_indices_from_dependencies<'a>(
603+
dependencies: impl IntoIterator<Item = &'a FunctionalDependence>,
604+
group_by_input_indices: &[Option<usize>],
605+
) -> Option<Vec<usize>> {
585606
let mut combined_target_indices = HashSet::new();
586-
for FunctionalDependence {
587-
source_indices,
588-
target_indices,
589-
..
590-
} in &dependencies.deps
591-
{
607+
for dependency in dependencies {
592608
// If the GROUP BY expression contains a determinant key, we can use
593609
// the associated fields after aggregation even if they are not part
594610
// of the GROUP BY expression.
595-
if source_indices
611+
if dependency
612+
.source_indices
596613
.iter()
597614
.all(|source_idx| group_by_input_indices.contains(&Some(*source_idx)))
598615
{
599-
combined_target_indices.extend(target_indices.iter());
616+
combined_target_indices.extend(dependency.target_indices.iter());
600617
}
601618
}
602619
(!combined_target_indices.is_empty()).then_some({
@@ -623,22 +640,21 @@ pub fn get_required_group_by_exprs_indices(
623640
.collect::<Option<Vec<_>>>()?;
624641

625642
groupby_expr_indices.sort_unstable();
626-
for FunctionalDependence {
627-
source_indices,
628-
target_indices,
629-
..
630-
} in &dependencies.deps
631-
{
632-
if source_indices
643+
for dependency in &dependencies.deps {
644+
// If all source indices are among GROUP BY expression indices, we
645+
// can remove target indices from GROUP BY expression indices and
646+
// use source indices instead. The dependency must also hold across
647+
// NULL keys, which GROUP BY collapses into a single group.
648+
if dependency
649+
.source_indices
633650
.iter()
634651
.all(|source_idx| groupby_expr_indices.contains(source_idx))
652+
&& dependency.is_valid_across_nulls(schema)
635653
{
636-
// If all source indices are among GROUP BY expression indices, we
637-
// can remove target indices from GROUP BY expression indices and
638-
// use source indices instead.
639-
groupby_expr_indices = set_difference(&groupby_expr_indices, target_indices);
640654
groupby_expr_indices =
641-
merge_and_order_indices(groupby_expr_indices, source_indices);
655+
set_difference(&groupby_expr_indices, &dependency.target_indices);
656+
groupby_expr_indices =
657+
merge_and_order_indices(groupby_expr_indices, &dependency.source_indices);
642658
}
643659
}
644660
groupby_expr_indices

‎datafusion/expr/src/logical_plan/builder.rs‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3138,6 +3138,63 @@ mod tests {
31383138
Ok(())
31393139
}
31403140

3141+
#[test]
3142+
fn plan_builder_aggregate_nullable_unique_adds_no_implicit_group_by_exprs()
3143+
-> Result<()> {
3144+
// A nullable UNIQUE column permits multiple NULL rows that may differ
3145+
// in the other columns, so it does not determine them across the
3146+
// single NULL group and they must not be added to the GROUP BY.
3147+
let schema = Schema::new(vec![
3148+
Field::new("id", DataType::Int32, true),
3149+
Field::new("state", DataType::Utf8, true),
3150+
Field::new("salary", DataType::Int32, true),
3151+
]);
3152+
let constraints = Constraints::new_unverified(vec![Constraint::Unique(vec![0])]);
3153+
let table_source = table_source_with_constraints(&schema, constraints);
3154+
3155+
let options =
3156+
LogicalPlanBuilderOptions::new().with_add_implicit_group_by_exprs(true);
3157+
let plan = LogicalPlanBuilder::scan("employee_csv", table_source, None)?
3158+
.with_options(options)
3159+
.aggregate(vec![col("id")], vec![sum(col("salary"))])?
3160+
.build()?;
3161+
3162+
assert_snapshot!(plan, @r"
3163+
Aggregate: groupBy=[[employee_csv.id]], aggr=[[sum(employee_csv.salary)]]
3164+
TableScan: employee_csv
3165+
");
3166+
3167+
Ok(())
3168+
}
3169+
3170+
#[test]
3171+
fn plan_builder_aggregate_not_null_unique_implicit_group_by_exprs() -> Result<()> {
3172+
// A UNIQUE column declared NOT NULL has no NULL rows, so it determines
3173+
// the other columns like a PRIMARY KEY and they are added to the
3174+
// GROUP BY.
3175+
let schema = Schema::new(vec![
3176+
Field::new("id", DataType::Int32, false),
3177+
Field::new("state", DataType::Utf8, true),
3178+
Field::new("salary", DataType::Int32, true),
3179+
]);
3180+
let constraints = Constraints::new_unverified(vec![Constraint::Unique(vec![0])]);
3181+
let table_source = table_source_with_constraints(&schema, constraints);
3182+
3183+
let options =
3184+
LogicalPlanBuilderOptions::new().with_add_implicit_group_by_exprs(true);
3185+
let plan = LogicalPlanBuilder::scan("employee_csv", table_source, None)?
3186+
.with_options(options)
3187+
.aggregate(vec![col("id")], vec![sum(col("salary"))])?
3188+
.build()?;
3189+
3190+
assert_snapshot!(plan, @r"
3191+
Aggregate: groupBy=[[employee_csv.id, employee_csv.state, employee_csv.salary]], aggr=[[sum(employee_csv.salary)]]
3192+
TableScan: employee_csv
3193+
");
3194+
3195+
Ok(())
3196+
}
3197+
31413198
#[test]
31423199
fn plan_builder_aggregate_rejects_nested_aggregates() -> Result<()> {
31433200
// https://github.com/apache/datafusion/issues/23812

‎datafusion/optimizer/src/replace_distinct_aggregate.rs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -102,8 +102,11 @@ impl OptimizerRule for ReplaceDistinctWithAggregate {
102102
// is a no-op and we can simply remove it. The dependency mode
103103
// must be `Single`: a `Multi` dependence (e.g. a former key
104104
// downgraded by a join) means equal rows may occur multiple
105-
// times, so the DISTINCT still has work to do.
105+
// times, so the DISTINCT still has work to do. Uniqueness
106+
// must also hold across NULLs: a nullable UNIQUE key can
107+
// contain multiple NULL rows that DISTINCT must collapse.
106108
if dep.mode == Dependency::Single
109+
&& dep.is_valid_across_nulls(input.schema())
107110
&& dep.source_indices.len() >= field_count
108111
&& dep.source_indices[..field_count]
109112
.iter()

0 commit comments

Comments
 (0)