Skip to content

fix: correlated exists/not exists subqueries hit the count bug for groupless aggregates #25391

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

Merged
jayzhan211 merged 10 commits into from
Oct 21, 2026
Merged

fix: correlated exists/not exists subqueries hit the count bug for groupless aggregates #25391

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
43e8ba9
fix: count-bug compensation for groupless aggregates
mohammadnaqvi04 Sep 15, 2026
fa4e3fb
fix: dedup by grouping on the join comparison
mohammadnaqvi04 Sep 15, 2026
d64fc1e
test: align dedup test data with checked-in shapes
mohammadnaqvi04 Sep 15, 2026
afdb738
fix: correlated exists/not exists over zero-row subquery limit
mohammadnaqvi04 Sep 16, 2026
ba10f14
test: update stale query plan expectations
mohammadnaqvi04 Sep 16, 2026
589d646
chore: some cleanup
mohammadnaqvi04 Sep 16, 2026
9a70592
chore: cleanup error messages
mohammadnaqvi04 Sep 17, 2026
38d8380
chore: update test post-rebase
mohammadnaqvi04 Sep 19, 2026
d5d7fd6
chore: address feedback
mohammadnaqvi04 Oct 20, 2026
d96dfb9
fix: gate count-bug compensation on bare keys
mohammadnaqvi04 Oct 20, 2026
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
32 changes: 32 additions & 0 deletions datafusion/expr/src/utils.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1574,6 +1574,38 @@ pub fn collect_subquery_cols(
})
}

/// Determine the subquery-side expressions a join filter compares,
/// stopping at the first subquery-referencing sub-expression rather than
/// descending to the bare columns inside of it.
///
/// `t1.a = t2.b` yields `t2.b`, matching [`collect_subquery_cols`]. But
/// `t1.a = CAST(t2.b AS INT)` yields `CAST(t2.b AS INT)` whole, not `t2.b`.
/// The two differ whenever the correlated predicate wraps the subquery's
/// column in an expression, in which case we need to evaluate the full
/// expression before grouping, so that two column values which compare
/// equal after the wrapping expression end up in the same group.
pub fn collect_subquery_join_exprs(
exprs: &[Expr],
subquery_schema: &DFSchema,
) -> Result<Vec<Expr>> {
let mut result: Vec<Expr> = vec![];
for expr in exprs {
expr.apply(|e| {
if e.column_refs()
.into_iter()
.all(|col| subquery_schema.has_column(col))
&& !e.column_refs().is_empty()
&& !result.contains(e)
{
result.push(e.clone());
return Ok(TreeNodeRecursion::Jump);
}
Ok(TreeNodeRecursion::Continue)
})?;
}
Ok(result)
}

#[cfg(test)]
mod tests {
use super::*;
Expand Down
100 changes: 84 additions & 16 deletions datafusion/optimizer/src/decorrelate.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -426,6 +426,50 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr {
None
};

if self.exists_sub_query
&& self.pull_up_having_expr.is_some()
&& let Some(expr) = conjunction(subquery_filters.clone())
{
let unqualified_expr = expr

@jayzhan211 jayzhan211 Oct 20, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Dropping the qualifier here means the merged filter is resolved against the joined schema, so it becomes ambiguous when the outer table has a column with the same name. Main returned rows for this query; it now fails to plan:

CREATE TABLE o9(k INT, c INT) AS VALUES (1, 100), (3, 100);
CREATE TABLE i9(k INT) AS VALUES (1);
-- DuckDB: 1, 3; main: 1; this PR: Ambiguous reference to unqualified field c
SELECT k FROM o9 WHERE EXISTS (
  SELECT * FROM (
    SELECT count(*) AS c FROM i9 WHERE i9.k = o9.k HAVING count(*) < 10
  ) s WHERE s.c < 5
);

Every column in pull_up_having_expr is a subquery output, so requalifying it with the alias at the top of build_join_with_count_bug fixes it (returns 1, 3):

let pull_up_having_expr = pull_up_having_expr
    .map(|f| {
        f.clone()
            .transform_up(|e| match e {
                Expr::Column(c) => Ok(Transformed::yes(Expr::Column(Column::new(
                    Some(alias),
                    c.name,
                )))),
                _ => Ok(Transformed::no(e)),
            })
            .data()
    })
    .transpose()?;
let pull_up_having_expr = pull_up_having_expr.as_ref();

@mohammadnaqvi04 mohammadnaqvi04 Oct 20, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Thanks for catching this, fixed. This now requalifies every column in pull_up_having_expr with the alias at the top of build_join_with_count_bug. Verified against your repro, now returns 1, 3.

.transform_up(|e| {
if let Expr::Column(Column { name, .. }) = &e {
Ok(Transformed::yes(Expr::Column(
Column::new_unqualified(name),
)))
} else {
Ok(Transformed::no(e))
}
})
.data()?;
let combined = match self.pull_up_having_expr.take() {
Some(existing) => existing.and(unqualified_expr),
None => unqualified_expr,
};
self.pull_up_having_expr = Some(combined);
let new_plan =
LogicalPlanBuilder::from((*plan_filter.input).clone()).build()?;
let mut carried_correlated_cols = correlated_subquery_cols;
if let Some(existing) =
self.correlated_subquery_cols_map.get(&*plan_filter.input)
{
carried_correlated_cols.extend(existing.iter().cloned());
}
self.correlated_subquery_cols_map
.insert(new_plan.clone(), carried_correlated_cols);
if !expr_result_map_for_count_bug.is_empty() {
self.collected_count_expr_map
.insert(new_plan.clone(), expr_result_map_for_count_bug);
} else if let Some(input_map) = self
.collected_count_expr_map
.get(&*plan_filter.input)
.cloned()
{
self.collected_count_expr_map
.insert(new_plan.clone(), input_map);
}
return Ok(Transformed::yes(new_plan));
}

match (&pull_up_expr_opt, &self.pull_up_having_expr) {
(Some(_), Some(_)) => {
// Error path
Expand Down Expand Up @@ -464,15 +508,18 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr {
&mut local_correlated_cols,
);
// add missing columns to Projection
let mut missing_exprs =
self.collect_missing_exprs(&projection.expr, &local_correlated_cols)?;
let mut missing_exprs = self.collect_missing_exprs(
&projection.expr,
&local_correlated_cols,
projection.input.schema(),
)?;

let mut expr_result_map_for_count_bug = HashMap::new();
if let Some(expr_result_map) =
self.collected_count_expr_map.get(&*projection.input)
{
proj_exprs_evaluation_result_on_empty_batch(
&projection.expr,
&missing_exprs,
projection.input.schema(),
expr_result_map,
&mut expr_result_map_for_count_bug,
Expand Down Expand Up @@ -511,7 +558,6 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr {
&self.correlated_subquery_cols_map,
&mut local_correlated_cols,
);

// A grouping set that omits a column the pull up adds either
// keeps its sets as they are (every set already groups by the
// column), or gains the column in every set, `ROLLUP(i.k)`
Expand Down Expand Up @@ -544,6 +590,7 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr {
self.collect_missing_exprs(
&aggregate.group_expr,
&local_correlated_cols,
aggregate.input.schema(),
)?
}
GroupingSetsPullUp::Unsafe => {
Expand All @@ -558,6 +605,7 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr {
self.collect_missing_exprs(
&aggregate.group_expr,
&local_correlated_cols,
aggregate.input.schema(),
)?
}
}
Expand All @@ -566,6 +614,7 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr {
self.collect_missing_exprs(
&aggregate.group_expr,
&local_correlated_cols,
aggregate.input.schema(),
)?
};

Expand Down Expand Up @@ -800,9 +849,10 @@ impl PullUpCorrelatedExpr {
}

fn collect_missing_exprs(
&self,
&mut self,
exprs: &[Expr],
correlated_subquery_cols: &BTreeSet<Column>,
_subquery_schema: &DFSchemaRef,
) -> Result<Vec<Expr>> {
let mut missing_exprs =
Vec::with_capacity(exprs.len() + correlated_subquery_cols.len());
Expand All @@ -813,19 +863,22 @@ impl PullUpCorrelatedExpr {
}
for col in correlated_subquery_cols.iter() {
let col_expr = Expr::Column(col.clone());
if !missing_exprs.contains(&col_expr) {
if !collides_with_existing(&missing_exprs, &col_expr) {
missing_exprs.push(col_expr)
}
}
if let Some(pull_up_having) = &self.pull_up_having_expr {
let filter_apply_columns = pull_up_having.column_refs();
for col in filter_apply_columns {
// add to missing_exprs if not already there
let contains = missing_exprs
for col in pull_up_having.column_refs() {
let col_expr = Expr::Column(col.clone());
// `col` is unqualified but the projection may already have
// it as `agg.c`. Check by bare name so that doesn't get
// added twice.
let already_present = missing_exprs
.iter()
.any(|expr| matches!(expr, Expr::Column(c) if c == col));
if !contains {
missing_exprs.push(Expr::Column(col.clone()))
.any(|expr| matches!(expr, Expr::Column(c) if c.name == col.name))
|| collides_with_existing(&missing_exprs, &col_expr);
if !already_present {
missing_exprs.push(col_expr)
}
}
}
Expand Down Expand Up @@ -925,6 +978,16 @@ fn partitions_by_all(expr: &Expr, cols: &BTreeSet<Column>) -> bool {
})
}

/// True if `candidate` collides with an existing entry's schema name,
/// e.g. a bare column already wrapped in an equally-named `CAST`.
/// Colliding means `LogicalPlanBuilder::project` would reject it as a duplicate.
fn collides_with_existing(exprs: &[Expr], candidate: &Expr) -> bool {
let candidate_name = candidate.schema_name().to_string();
exprs
.iter()
.any(|expr| expr.schema_name().to_string() == candidate_name)
}

fn can_pullup_over_aggregation(expr: &Expr) -> bool {
if let Expr::BinaryExpr(BinaryExpr {
left,
Expand Down Expand Up @@ -979,7 +1042,7 @@ fn remove_duplicated_filter(
Expr::BinaryExpr(b) => b.op.swap() == Some(b.op),
_ => true,
},
"remove_duplicated_filter: in_predicate must use a commutative operator"
"in_predicate must use a commutative operator"
);

Ok(filters
Expand Down Expand Up @@ -1105,9 +1168,14 @@ fn filter_exprs_evaluation_result_on_empty_batch(
let simplifier = ExprSimplifier::new(info);
let result_expr = simplifier.simplify(result_expr)?;
match &result_expr {
// evaluate to false or null on empty batch, no need to pull up
Expr::Literal(ScalarValue::Null, _)
| Expr::Literal(ScalarValue::Boolean(Some(false)), _) => None,
| Expr::Literal(ScalarValue::Boolean(None), _)
| Expr::Literal(ScalarValue::Boolean(Some(false)), _) => {
for (name, exprs) in input_expr_result_map_for_count_bug {
expr_result_map_for_count_bug.insert(name.clone(), exprs.clone());
}
None
}
// evaluate to true on empty batch, need to pull up the expr
Expr::Literal(ScalarValue::Boolean(Some(true)), _) => {
for (name, exprs) in input_expr_result_map_for_count_bug {
Expand Down
Loading
Loading