Skip to content

Refactor DecorrelateWhereExists and add back Distinct if needs #5345

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
jackwener merged 9 commits into from
Mar 11, 2023
Merged

Refactor DecorrelateWhereExists and add back Distinct if needs #5345

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
2ce552b
Refactor DecorrelateWhereExists and add back Distinct if needs
ygf11 Feb 20, 2023
45f3212
extract collect_subquery_cols
ygf11 Feb 20, 2023
af25a6e
Add back aggregate
ygf11 Feb 22, 2023
dbaf64c
fix comment
ygf11 Feb 22, 2023
0c28e11
Merge upstream/main into where-exists-distinct
ygf11 Mar 2, 2023
47b1cd5
fix tests
ygf11 Mar 3, 2023
0b4d252
remove unused method
ygf11 Mar 3, 2023
573c48b
fix cargo fmt
ygf11 Mar 3, 2023
8e6c726
Merge upstream/main and fix conflicts
ygf11 Mar 11, 2023
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
119 changes: 119 additions & 0 deletions datafusion/core/tests/sql/joins.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3518,3 +3518,122 @@ async fn not_exists_subquery_to_join_expr_filter() -> Result<()> {

Ok(())
}

#[tokio::test]
async fn exists_distinct_subquery_to_join() -> Result<()> {
let test_repartition_joins = vec![true, false];
for repartition_joins in test_repartition_joins {
let ctx = create_join_context("t1_id", "t2_id", repartition_joins)?;

let sql = "SELECT * FROM t1 WHERE NOT EXISTS(SELECT DISTINCT t2_int FROM t2 WHERE t1.t1_id + 1 > t2.t2_id * 2)";
let msg = format!("Creating logical plan for '{sql}'");
let dataframe = ctx.sql(&("explain ".to_owned() + sql)).await.expect(&msg);
let plan = dataframe.into_optimized_plan()?;

let expected = vec![
"Explain [plan_type:Utf8, plan:Utf8]",
" LeftAnti Join: Filter: CAST(t1.t1_id AS Int64) + Int64(1) > CAST(t2.t2_id AS Int64) * Int64(2) [t1_id:UInt32;N, t1_name:Utf8;N, t1_int:UInt32;N]",
" TableScan: t1 projection=[t1_id, t1_name, t1_int] [t1_id:UInt32;N, t1_name:Utf8;N, t1_int:UInt32;N]",
" Aggregate: groupBy=[[t2.t2_id]], aggr=[[]] [t2_id:UInt32;N]",
" TableScan: t2 projection=[t2_id] [t2_id:UInt32;N]",
];
let formatted = plan.display_indent_schema().to_string();
let actual: Vec<&str> = formatted.trim().lines().collect();
assert_eq!(
expected, actual,
"\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
);
let expected = vec![
"+-------+---------+--------+",
"| t1_id | t1_name | t1_int |",
"+-------+---------+--------+",
"| 11 | a | 1 |",
"+-------+---------+--------+",
];

let results = execute_to_batches(&ctx, sql).await;
assert_batches_sorted_eq!(expected, &results);
}

Ok(())
}

#[tokio::test]
async fn exists_distinct_subquery_to_join_with_expr() -> Result<()> {
let test_repartition_joins = vec![true, false];
for repartition_joins in test_repartition_joins {
let ctx = create_join_context("t1_id", "t2_id", repartition_joins)?;

// `t2_id + t2_int` is in the subquery project.
let sql = "SELECT * FROM t1 WHERE NOT EXISTS(SELECT DISTINCT t2_id + t2_int, t2_int FROM t2 WHERE t1.t1_id + 1 > t2.t2_id * 2)";
let msg = format!("Creating logical plan for '{sql}'");
let dataframe = ctx.sql(&("explain ".to_owned() + sql)).await.expect(&msg);
let plan = dataframe.into_optimized_plan()?;

let expected = vec![
"Explain [plan_type:Utf8, plan:Utf8]",
" LeftAnti Join: Filter: CAST(t1.t1_id AS Int64) + Int64(1) > CAST(t2.t2_id AS Int64) * Int64(2) [t1_id:UInt32;N, t1_name:Utf8;N, t1_int:UInt32;N]",
" TableScan: t1 projection=[t1_id, t1_name, t1_int] [t1_id:UInt32;N, t1_name:Utf8;N, t1_int:UInt32;N]",
" Aggregate: groupBy=[[t2.t2_id]], aggr=[[]] [t2_id:UInt32;N]",
" TableScan: t2 projection=[t2_id] [t2_id:UInt32;N]",
];
let formatted = plan.display_indent_schema().to_string();
let actual: Vec<&str> = formatted.trim().lines().collect();
assert_eq!(
expected, actual,
"\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
);
let expected = vec![
"+-------+---------+--------+",
"| t1_id | t1_name | t1_int |",
"+-------+---------+--------+",
"| 11 | a | 1 |",
"+-------+---------+--------+",
];

let results = execute_to_batches(&ctx, sql).await;
assert_batches_sorted_eq!(expected, &results);
}

Ok(())
}

#[tokio::test]
async fn exists_distinct_subquery_to_join_with_literal() -> Result<()> {
let test_repartition_joins = vec![true, false];
for repartition_joins in test_repartition_joins {
let ctx = create_join_context("t1_id", "t2_id", repartition_joins)?;

// `1` is in the subquery project.
let sql = "SELECT * FROM t1 WHERE NOT EXISTS(SELECT DISTINCT 1, t2_int FROM t2 WHERE t1.t1_id + 1 > t2.t2_id * 2)";
let msg = format!("Creating logical plan for '{sql}'");
let dataframe = ctx.sql(&("explain ".to_owned() + sql)).await.expect(&msg);
let plan = dataframe.into_optimized_plan()?;

let expected = vec![
"Explain [plan_type:Utf8, plan:Utf8]",
" LeftAnti Join: Filter: CAST(t1.t1_id AS Int64) + Int64(1) > CAST(t2.t2_id AS Int64) * Int64(2) [t1_id:UInt32;N, t1_name:Utf8;N, t1_int:UInt32;N]",
" TableScan: t1 projection=[t1_id, t1_name, t1_int] [t1_id:UInt32;N, t1_name:Utf8;N, t1_int:UInt32;N]",
" Aggregate: groupBy=[[t2.t2_id]], aggr=[[]] [t2_id:UInt32;N]",
" TableScan: t2 projection=[t2_id] [t2_id:UInt32;N]",
];
let formatted = plan.display_indent_schema().to_string();
let actual: Vec<&str> = formatted.trim().lines().collect();
assert_eq!(
expected, actual,
"\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
);
let expected = vec![
"+-------+---------+--------+",
"| t1_id | t1_name | t1_int |",
"+-------+---------+--------+",
"| 11 | a | 1 |",
"+-------+---------+--------+",
];

let results = execute_to_batches(&ctx, sql).await;
assert_batches_sorted_eq!(expected, &results);
}

Ok(())
}
199 changes: 137 additions & 62 deletions datafusion/optimizer/src/decorrelate_where_exists.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,14 +16,16 @@
// under the License.

use crate::optimizer::ApplyOrder;
use crate::utils::{conjunction, extract_join_filters, split_conjunction};
use crate::utils::{
collect_subquery_cols, conjunction, extract_join_filters, split_conjunction,
};
use crate::{OptimizerConfig, OptimizerRule};
use datafusion_common::{Column, DataFusionError, Result};
use datafusion_expr::{
logical_plan::{Filter, JoinType, Subquery},
logical_plan::{Distinct, Filter, JoinType, Subquery},
Expr, LogicalPlan, LogicalPlanBuilder,
};
use std::collections::BTreeSet;

use std::sync::Arc;

/// Optimizer rule for rewriting subquery filters to joins
Expand Down Expand Up @@ -142,69 +144,70 @@ fn optimize_exists(
query_info: &SubqueryInfo,
outer_input: &LogicalPlan,
) -> Result<Option<LogicalPlan>> {
let maybe_subqury_filter = match query_info.query.subquery.as_ref() {
LogicalPlan::Distinct(subqry_distinct) => match subqry_distinct.input.as_ref() {
LogicalPlan::Projection(subqry_proj) => &subqry_proj.input,
_ => {
return Ok(None);
}
},
LogicalPlan::Projection(subqry_proj) => &subqry_proj.input,
_ => {
// Subquery currently only supports distinct or projection
return Ok(None);
}
}
.as_ref();
let subquery = query_info.query.subquery.as_ref();
if let Some((join_filter, optimized_subquery)) = optimize_subquery(subquery)? {
// join our sub query into the main plan
let join_type = match query_info.negated {
true => JoinType::LeftAnti,
false => JoinType::LeftSemi,
};

let new_plan = LogicalPlanBuilder::from(outer_input.clone())
.join(
optimized_subquery,
join_type,
(Vec::<Column>::new(), Vec::<Column>::new()),
Some(join_filter),
)?
.build()?;

// extract join filters
let (join_filters, subquery_input) = extract_join_filters(maybe_subqury_filter)?;
// cannot optimize non-correlated subquery
if join_filters.is_empty() {
return Ok(None);
Ok(Some(new_plan))
} else {
Ok(None)
}

let input_schema = subquery_input.schema();
let subquery_cols: BTreeSet<Column> =
join_filters
.iter()
.try_fold(BTreeSet::new(), |mut cols, expr| {
let using_cols: Vec<Column> = expr
.to_columns()?
}
/// Optimize the subquery and extract the possible join filter.
/// This function can't optimize non-correlated subquery, and will return None.
fn optimize_subquery(subquery: &LogicalPlan) -> Result<Option<(Expr, LogicalPlan)>> {
match subquery {
LogicalPlan::Distinct(subqry_distinct) => {
let distinct_input = &subqry_distinct.input;
let optimized_plan =
optimize_subquery(distinct_input)?.map(|(filters, right)| {
(
filters,
LogicalPlan::Distinct(Distinct {
input: Arc::new(right),
}),
)
});
Ok(optimized_plan)
}
Copy link
Contributor Author

@ygf11 ygf11 Mar 3, 2023

Choose a reason for hiding this comment

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

I am not sure the behavior of Distinct is correct, so do not handle Aggregate here.

Copy link
Contributor

@mingmwang mingmwang Mar 8, 2023 •
edited
Loading

Choose a reason for hiding this comment

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

I am not sure the behavior of Distinct is correct, so do not handle Aggregate here.

I will take a closer look at this PR tomorrow. If you do not know how to handle Aggregate here, you can just leave it here and only handle the Distinct case.

A simplest approach is checking whether there are out reference used by Aggregate expressions, if there are, return Err and else add the correlated columns to be part of the group by columns. (Only allow out reference columns referred in Filter expressions or Join expressions, this will limit the supported cases of Subqueries, but more safe.)
For example if the original inner aggregate is group by inner_a and there is correlation condition like outer_b = inner_b, them add the inner_b to be part of the group by conditions. It is not a perfect solution and sometimes might cause some bug.

SparkSQL has the similar logic in the rule pullOutCorrelatedPredicates

      case a @ Aggregate(grouping, expressions, child) =>
        val referencesToAdd = missingReferences(a)
        if (referencesToAdd.nonEmpty) {
          Aggregate(grouping ++ referencesToAdd, expressions ++ referencesToAdd, child)
        } else {
          a
        }

The latest SparkSQL's implementation also has bug here (need to differ the original Aggregate is Scalar Aggregate or Vector Aggregate)

SQL to reproduce the bug, you can have a try on both PostgreSQL and SparkSQL

CREATE TABLE t1 (id INT,name String);
CREATE TABLE t2 (id INT,name String);
insert into t1 values (11, "a"), (11, "a"), (22, "b"), (33, "c"), (44, "d"), (null, "e");
insert into t2 values (111, "z"), (111, "z"), (222, "y"), (444, "x"), (555, "w"), (null, "v");

-- Should output  all t1
SELECT t1.id, t1.name FROM t1 WHERE EXISTS (select 0);

-- Should output all t1
SELECT t1.id, t1.name FROM t1 WHERE EXISTS (SELECT count(*) FROM t2 WHERE t2.id = t1.id);

-- Should output all t1
SELECT t1.id, t1.name FROM t1 WHERE EXISTS (SELECT count(*) FROM t2 WHERE t2.id = t1.id having count(*) = 0);

PostgreSQL does not support decorrelate Subqueries which include Aggregates. SparkSQL supports some cases but not all the cases and the implementation has bug.

I'm going to implement a general subquery decorrelation rule based on the two well known papers.
#5492

Orthogonal Optimization of Subqueries and Aggregation
https://dl.acm.org/doi/10.1145/375663.375748

Unnesting Arbitrary Queries
https://cs.emis.de/LNI/Proceedings/Proceedings241/383.pdf

Copy link
Member

@jackwener jackwener Mar 9, 2023

Choose a reason for hiding this comment

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

I'm going to implement a general subquery decorrelation rule based on the two well known papers.

Great job !

may related #5368

LogicalPlan::Projection(projection) => {
// extract join filters
let (join_filters, subquery_input) = extract_join_filters(&projection.input)?;
// cannot optimize non-correlated subquery
if join_filters.is_empty() {
return Ok(None);
}
let input_schema = subquery_input.schema();
let project_exprs: Vec<Expr> =
collect_subquery_cols(&join_filters, input_schema.clone())?
.into_iter()
.filter(|col| input_schema.has_column(col))
.collect::<_>();

cols.extend(using_cols);
Result::<_, DataFusionError>::Ok(cols)
.map(Expr::Column)
.collect();
let right = LogicalPlanBuilder::from(subquery_input)
.project(project_exprs)?
.build()?;

// join_filters is not empty.
let join_filter = conjunction(join_filters).ok_or_else(|| {
DataFusionError::Internal("join filters should not be empty".to_string())
})?;

let projection_exprs: Vec<Expr> =
subquery_cols.into_iter().map(Expr::Column).collect();

let right = LogicalPlanBuilder::from(subquery_input)
.project(projection_exprs)?
.build()?;

let join_filter = conjunction(join_filters);

// join our sub query into the main plan
let join_type = match query_info.negated {
true => JoinType::LeftAnti,
false => JoinType::LeftSemi,
};

// TODO: add Distinct if the original plan is a Distinct.
let new_plan = LogicalPlanBuilder::from(outer_input.clone())
.join(
right,
join_type,
(Vec::<Column>::new(), Vec::<Column>::new()),
join_filter,
)?
.build()?;

Ok(Some(new_plan))
Ok(Some((join_filter, right)))
}
_ => Ok(None),
}
}

struct SubqueryInfo {
Expand Down Expand Up @@ -670,4 +673,76 @@ mod tests {

assert_plan_eq(&plan, expected)
}

#[test]
fn exists_distinct_subquery() -> Result<()> {
let table_scan = test_table_scan()?;
let subquery_scan = test_table_scan_with_name("sq")?;
let subquery = LogicalPlanBuilder::from(subquery_scan)
.filter((lit(1u32) + col("sq.a")).gt(col("test.a") * lit(2u32)))?
.project(vec![col("sq.c")])?
.distinct()?
.build()?;
let plan = LogicalPlanBuilder::from(table_scan)
.filter(exists(Arc::new(subquery)))?
.project(vec![col("test.b")])?
.build()?;

let expected = "Projection: test.b [b:UInt32]\
\n LeftSemi Join: Filter: UInt32(1) + sq.a > test.a * UInt32(2) [a:UInt32, b:UInt32, c:UInt32]\
\n TableScan: test [a:UInt32, b:UInt32, c:UInt32]\
\n Distinct: [a:UInt32]\
\n Projection: sq.a [a:UInt32]\
\n TableScan: sq [a:UInt32, b:UInt32, c:UInt32]";

assert_plan_eq(&plan, expected)
}

#[test]
fn exists_distinct_expr_subquery() -> Result<()> {
let table_scan = test_table_scan()?;
let subquery_scan = test_table_scan_with_name("sq")?;
let subquery = LogicalPlanBuilder::from(subquery_scan)
.filter((lit(1u32) + col("sq.a")).gt(col("test.a") * lit(2u32)))?
.project(vec![col("sq.b") + col("sq.c")])?
.distinct()?
.build()?;
let plan = LogicalPlanBuilder::from(table_scan)
.filter(exists(Arc::new(subquery)))?
.project(vec![col("test.b")])?
.build()?;

let expected = "Projection: test.b [b:UInt32]\
\n LeftSemi Join: Filter: UInt32(1) + sq.a > test.a * UInt32(2) [a:UInt32, b:UInt32, c:UInt32]\
\n TableScan: test [a:UInt32, b:UInt32, c:UInt32]\
\n Distinct: [a:UInt32]\
\n Projection: sq.a [a:UInt32]\
\n TableScan: sq [a:UInt32, b:UInt32, c:UInt32]";

assert_plan_eq(&plan, expected)
}
Copy link
Contributor Author

@ygf11 ygf11 Mar 3, 2023

Choose a reason for hiding this comment

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

I find the behaviors of Distinct in postgres and spark are not same.

  • For postgres, it will not add back the distinct to the optimized result.
  • For spark, it will add back the distinct, and will keep the unused project exprs(sq.b + sq.c in the above).

I think sq.b + sq.c is not used any more, so remove it.

I am not sure if the way of this pr is appropriate.
I want to hear your advice @mingmwang @alamb @jackwener

Copy link
Contributor

@alamb alamb Mar 7, 2023

Choose a reason for hiding this comment

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

For postgres, it will not add back the distinct to the optimized result.

I think it likely depends on how the join operators are implemented

# current logical plan:
Projection: t1.t1_id                                                                                                                
  LeftSemi Join:  Filter: t2.t2_id > t1.t1_id                                                                                                               
    TableScan: t1 projection=[t1_id]                                                                                                       
    Projection: t2.t2_id # DISTINCT is not added back                                                                                                                                    
      TableScan: t2 projection=[t2_id]

If the LeftSemiJoin is going to read its build side into (effectively) a HashSet then it doesn't really matter if the input is deduplicated prior to input.

However, since there is no equality predicate (the predicate is >) , i am not sure what our LeftSemiJoin will do (maybe it will buffer the entire build side?) in which case it might help performance to put a DISTINCT on he output of the Projection to reduce the cardinality that the join buffers.

Copy link
Member

@jackwener jackwener Mar 9, 2023

Choose a reason for hiding this comment

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

I think sq.b + sq.c is not used any more, so remove it.

This is a good job.

For postgres, it will not add back the distinct to the optimized result.

I'm also not sure about it

Copy link
Contributor

@mingmwang mingmwang Mar 10, 2023

Choose a reason for hiding this comment

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

I find the behaviors of Distinct in postgres and spark are not same.

  • For postgres, it will not add back the distinct to the optimized result.
  • For spark, it will add back the distinct, and will keep the unused project exprs(sq.b + sq.c in the above).

@ygf11 @jackwener @alamb
I guess why Postgres remove the Distinct in the final result is maybe it is related to how it implements NAAJ(Null Aware Anti Join) and to make sure the result is correct.

You can test try those SQLs in Postgres:

CREATE TABLE t1 (id integer,name varchar(100));
CREATE TABLE t2 (id integer,name varchar(100));

explain  
SELECT t1.id, t1.name FROM t1 WHERE EXISTS (SELECT distinct * FROM t2 WHERE t2.id = t1.id);

QUERY PLAN
Hash Join (cost=18.50..34.32 rows=160 width=222)
Hash Cond: (t1.id = t2.id)
-> Seq Scan on t1 (cost=0.00..13.20 rows=320 width=222)
-> Hash (cost=16.00..16.00 rows=200 width=4)
-> HashAggregate (cost=14.00..16.00 rows=200 width=4)
Group Key: t2.id
-> Seq Scan on t2 (cost=0.00..13.20 rows=320 width=4)

explain  
SELECT t1.id, t1.name FROM t1 WHERE NOT EXISTS (SELECT distinct t2.name FROM t2 WHERE t2.id = t1.id);

QUERY PLAN
--
Hash Anti Join (cost=17.20..35.82 rows=160 width=222)
Hash Cond: (t1.id = t2.id)
-> Seq Scan on t1 (cost=0.00..13.20 rows=320 width=222)
-> Hash (cost=13.20..13.20 rows=320 width=4)
-> Seq Scan on t2 (cost=0.00..13.20 rows=320 width=4)

explain  
SELECT t1.id, t1.name FROM t1 WHERE t1.id in (SELECT distinct t2.id FROM t2);

QUERY PLAN
--
Hash Join (cost=20.50..34.56 rows=320 width=222)
Hash Cond: (t1.id = t2.id)
-> Seq Scan on t1 (cost=0.00..13.20 rows=320 width=222)
-> Hash (cost=18.00..18.00 rows=200 width=4)
-> HashAggregate (cost=14.00..16.00 rows=200 width=4)
Group Key: t2.id
-> Seq Scan on t2 (cost=0.00..13.20 rows=320 width=4)

explain  
SELECT t1.id, t1.name FROM t1 WHERE t1.id not in (SELECT distinct t2.id FROM t2);

QUERY PLAN
--
Seq Scan on t1 (cost=16.50..30.50 rows=160 width=222)
Filter: (NOT (hashed SubPlan 1))
SubPlan 1
-> HashAggregate (cost=14.00..16.00 rows=200 width=4)
Group Key: t2.id
-> Seq Scan on t2 (cost=0.00..13.20 rows=320 width=4

You can see that for IN/EXISTS subqueries, there are Aggregates. But for NOT EXISTS, the Aggregateis removed, and forNOT IN`, the subquery is kept and not decorrelated.

I think in DataFusion, I like the implementation in this PR and the generated plan is more consistent. For the correctness of NAAJ, we have another ticket to track this and we can make the Hash Join itself Null aware.


#[test]
fn exists_distinct_subquery_with_literal() -> Result<()> {
let table_scan = test_table_scan()?;
let subquery_scan = test_table_scan_with_name("sq")?;
let subquery = LogicalPlanBuilder::from(subquery_scan)
.filter((lit(1u32) + col("sq.a")).gt(col("test.a") * lit(2u32)))?
.project(vec![lit(1u32), col("sq.c")])?
.distinct()?
.build()?;
let plan = LogicalPlanBuilder::from(table_scan)
.filter(exists(Arc::new(subquery)))?
.project(vec![col("test.b")])?
.build()?;

let expected = "Projection: test.b [b:UInt32]\
\n LeftSemi Join: Filter: UInt32(1) + sq.a > test.a * UInt32(2) [a:UInt32, b:UInt32, c:UInt32]\
\n TableScan: test [a:UInt32, b:UInt32, c:UInt32]\
\n Distinct: [a:UInt32]\
\n Projection: sq.a [a:UInt32]\
\n TableScan: sq [a:UInt32, b:UInt32, c:UInt32]";

assert_plan_eq(&plan, expected)
}
}
21 changes: 6 additions & 15 deletions datafusion/optimizer/src/decorrelate_where_in.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,12 @@

use crate::alias::AliasGenerator;
use crate::optimizer::ApplyOrder;
use crate::utils::{conjunction, extract_join_filters, only_or_err, split_conjunction};
use crate::utils::{
collect_subquery_cols, conjunction, extract_join_filters, only_or_err,
split_conjunction,
};
use crate::{OptimizerConfig, OptimizerRule};
use datafusion_common::{context, Column, DataFusionError, Result};
use datafusion_common::{context, Column, Result};
use datafusion_expr::expr_rewriter::{replace_col, unnormalize_col};
use datafusion_expr::logical_plan::{JoinType, Projection, Subquery};
use datafusion_expr::{Expr, Filter, LogicalPlan, LogicalPlanBuilder};
Expand Down Expand Up @@ -159,19 +162,7 @@ fn optimize_where_in(
// replace qualified name with subquery alias.
let subquery_alias = alias.next("__correlated_sq");
let input_schema = subquery_input.schema();
let mut subquery_cols: BTreeSet<Column> =
join_filters
.iter()
.try_fold(BTreeSet::new(), |mut cols, expr| {
let using_cols: Vec<Column> = expr
.to_columns()?
.into_iter()
.filter(|col| input_schema.has_column(col))
.collect::<_>();

cols.extend(using_cols);
Result::<_, DataFusionError>::Ok(cols)
})?;
let mut subquery_cols = collect_subquery_cols(&join_filters, input_schema.clone())?;
let join_filter = conjunction(join_filters).map_or(Ok(None), |filter| {
replace_qualified_name(filter, &subquery_cols, &subquery_alias).map(Option::Some)
})?;
Expand Down
Loading