Push down filter plan for non-unnest column

[jayzhan211]

Which issue does this PR close?

Follow on #10991

Closes #11016

Rationale for this change

In #10991, if we see unnest column, we skip the push down optimization.
In this PR, we select the non-unnest columns to push down, even there is unnest columns besides.

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

[alamb]

Thank you very much @jayzhan211 -- looks great to me

I also updated the ticket description to reflect that this PR would close #11016 as far as I understand

cc @jonahgao

[alamb]

Can you please also update the comment a few lines above to reflect the issue is fixed?

# Could push the filter (column1 = 2) down below unnest
# https://github.com/apache/datafusion/issues/11016

[jonahgao]

Looks good to me!

I am considering whether it is beneficial to wait for the next rewrite pass to handle the unnest input, because we already have the capability to push down filters below projections. This might simplify the implementation, and we don't need to assume that the unnest input is a projection.

            LogicalPlan::Unnest(mut unnest) => {
                let predicates = split_conjunction_owned(filter.predicate.clone());
                let mut non_unnest_predicates = vec![];
                let mut unnest_predicates = vec![];
                for predicate in predicates {
                    // collect all the Expr::Column in predicate recursively
                    let mut accum: HashSet<Column> = HashSet::new();
                    expr_to_columns(&predicate, &mut accum)?;

                    if unnest.exec_columns.iter().any(|c| accum.contains(c)) {
                        unnest_predicates.push(predicate);
                    } else {
                        non_unnest_predicates.push(predicate);
                    }
                }
                // Unnest predicates should not be pushed down.
                // If no non-unnest predicates exist, early return
                if non_unnest_predicates.is_empty() {
                    filter.input = Arc::new(LogicalPlan::Unnest(unnest));
                    return Ok(Transformed::no(LogicalPlan::Filter(filter)));
                }
                let new_filter = LogicalPlan::Filter(Filter::try_new(
                    conjunction(non_unnest_predicates).unwrap(),
                    unnest.input.clone(),
                )?);
                unnest.input = Arc::new(new_filter);
                if unnest_predicates.is_empty() {
                    Ok(Transformed::yes(LogicalPlan::Unnest(unnest)))
                } else {
                    filter.predicate = conjunction(unnest_predicates).unwrap();
                    filter.input = Arc::new(LogicalPlan::Unnest(unnest));
                    Ok(Transformed::no(LogicalPlan::Filter(filter)))
                }
            }

[jayzhan211]

Looks good to me!

I am considering whether it is beneficial to wait for the next rewrite pass to handle the unnest input, because we already have the capability to push down filters below projections. This might simplify the implementation, and we don't need to assume that the unnest input is a projection.

            LogicalPlan::Unnest(mut unnest) => {
                let predicates = split_conjunction_owned(filter.predicate.clone());
                let mut non_unnest_predicates = vec![];
                let mut unnest_predicates = vec![];
                for predicate in predicates {
                    // collect all the Expr::Column in predicate recursively
                    let mut accum: HashSet<Column> = HashSet::new();
                    expr_to_columns(&predicate, &mut accum)?;

                    if unnest.exec_columns.iter().any(|c| accum.contains(c)) {
                        unnest_predicates.push(predicate);
                    } else {
                        non_unnest_predicates.push(predicate);
                    }
                }
                // Unnest predicates should not be pushed down.
                // If no non-unnest predicates exist, early return
                if non_unnest_predicates.is_empty() {
                    filter.input = Arc::new(LogicalPlan::Unnest(unnest));
                    return Ok(Transformed::no(LogicalPlan::Filter(filter)));
                }
                let new_filter = LogicalPlan::Filter(Filter::try_new(
                    conjunction(non_unnest_predicates).unwrap(),
                    unnest.input.clone(),
                )?);
                unnest.input = Arc::new(new_filter);
                if unnest_predicates.is_empty() {
                    Ok(Transformed::yes(LogicalPlan::Unnest(unnest)))
                } else {
                    filter.predicate = conjunction(unnest_predicates).unwrap();
                    filter.input = Arc::new(LogicalPlan::Unnest(unnest));
                    Ok(Transformed::no(LogicalPlan::Filter(filter)))
                }
            }

For projection plan, we check if predicates contains volatile expression, so we can't directly assign it as the new input. Not sure what is the case for other plan, they may have their own special case handling.
Also, I think currently we only run the PushDownFilter pass once 🤔 unlike SimplifyExpressions

I see. It might worth to try.

[jonahgao]

I think the optimizer already handles recursion, so it might not be necessary to call rewrite here.

[jayzhan211]

Agree!

[jonahgao]

How about directly cloning the input since it is an Arc? This way, we don't need to construct a Default LogicalPlan either.

[jayzhan211]

We have a plan to convert Arc<LogicalPlan> to Box<LogicalPlan>, but before that we need to make sure there is no clone exist, so I think we should try our best to avoid clone. It is also the reason why we have rewrite in Logical optimizer rule that takes LogicalPlan instead of Arc<LogicalPlan>

Ref: #9637 (comment)

[jonahgao]

Since there's a purpose for doing this, make sense to me. Thanks @jayzhan211

[jonahgao]

LGTM! Thanks @jayzhan211

[alamb]

👍

[jayzhan211]

Thanks @jonahgao and @alamb