Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Don't preserve functional dependency when generating UNION logical plan #12979

Merged
merged 1 commit into from
Oct 20, 2024
Merged

Don't preserve functional dependency when generating UNION logical plan #12979

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 48 additions & 0 deletions datafusion/core/src/dataframe/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2623,6 +2623,54 @@ mod tests {
Ok(())
}

#[tokio::test]
async fn test_aggregate_with_union() -> Result<()> {
let df = test_table().await?;

let df1 = df
.clone()
// GROUP BY `c1`
.aggregate(vec![col("c1")], vec![min(col("c2"))])?
// SELECT `c1` , min(c2) as `result`
.select(vec![col("c1"), min(col("c2")).alias("result")])?;
let df2 = df
.clone()
// GROUP BY `c1`
.aggregate(vec![col("c1")], vec![max(col("c3"))])?
// SELECT `c1` , max(c3) as `result`
.select(vec![col("c1"), max(col("c3")).alias("result")])?;

let df_union = df1.union(df2)?;
let df = df_union
// GROUP BY `c1`
.aggregate(
vec![col("c1")],
vec![sum(col("result")).alias("sum_result")],
)?
// SELECT `c1`, sum(result) as `sum_result`
.select(vec![(col("c1")), col("sum_result")])?;

let df_results = df.collect().await?;

#[rustfmt::skip]
assert_batches_sorted_eq!(
[
"+----+------------+",
"| c1 | sum_result |",
"+----+------------+",
"| a | 84 |",
"| b | 69 |",
"| c | 124 |",
"| d | 126 |",
"| e | 121 |",
"+----+------------+"
],
&df_results
);

Ok(())
}

#[tokio::test]
async fn test_aggregate_subexpr() -> Result<()> {
let df = test_table().await?;
Expand Down
11 changes: 8 additions & 3 deletions datafusion/expr/src/logical_plan/builder.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,8 +49,8 @@ use datafusion_common::display::ToStringifiedPlan;
use datafusion_common::file_options::file_type::FileType;
use datafusion_common::{
get_target_functional_dependencies, internal_err, not_impl_err, plan_datafusion_err,
plan_err, Column, DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue,
TableReference, ToDFSchema, UnnestOptions,
plan_err, Column, DFSchema, DFSchemaRef, DataFusionError, FunctionalDependencies,
Result, ScalarValue, TableReference, ToDFSchema, UnnestOptions,
};
use datafusion_expr_common::type_coercion::binary::type_union_resolution;

Expand Down Expand Up @@ -1402,7 +1402,12 @@ pub fn validate_unique_names<'a>(
pub fn union(left_plan: LogicalPlan, right_plan: LogicalPlan) -> Result<LogicalPlan> {
// Temporarily use the schema from the left input and later rely on the analyzer to
// coerce the two schemas into a common one.
let schema = Arc::clone(left_plan.schema());

// Functional Dependencies doesn't preserve after UNION operation
let schema = (**left_plan.schema()).clone();
let schema =
Arc::new(schema.with_functional_dependencies(FunctionalDependencies::empty())?);

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is clearing out all dependencies the right fix? Could we retain some if they do not harm?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'll wait to merge this PR until tomorrow to give @Sevenannn a chance to respond

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @berkaysynnada, thanks for the review! I don’t think any FD persists when performing UNION on 2 tables.

A simple example would be UNION table t1 with another table t2 which only has 1 row, there always exists such data in t2 which could break the FDs in t1 / t2 after the UNION.

In this case, clearing FDs would be the right fix since we don’t want FDs to get wrongly retained and affect later plans, e.g. aggregation.

Please let me know if you have any further questions regarding this PR, thanks!

Ok(LogicalPlan::Union(Union {
inputs: vec![Arc::new(left_plan), Arc::new(right_plan)],
schema,
Expand Down