alopex_dataframe/lazy/

optimizer.rs

use std::collections::HashSet;

use crate::expr::{Expr as E, Operator};
use crate::lazy::{LogicalPlan, ProjectionKind};
use crate::Expr;

/// Optimizer that rewrites `LogicalPlan` (e.g. predicate/projection pushdown).
pub struct Optimizer;

impl Optimizer {
    /// Optimize a `LogicalPlan` and return the rewritten plan.
    pub fn optimize(plan: &LogicalPlan) -> LogicalPlan {
        let plan = predicate_pushdown(plan.clone());
        projection_pushdown(plan)
    }
}

fn predicate_pushdown(plan: LogicalPlan) -> LogicalPlan {
    match plan {
        LogicalPlan::Filter { input, predicate } => {
            let input = predicate_pushdown(*input);

            match input {
                LogicalPlan::Filter {
                    input: inner,
                    predicate: inner_predicate,
                } => {
                    let combined = and_expr(inner_predicate, predicate);
                    predicate_pushdown(LogicalPlan::Filter {
                        input: inner,
                        predicate: combined,
                    })
                }
                LogicalPlan::Projection { input, exprs, kind } => {
                    if can_push_filter_through_projection(&predicate, &exprs, &kind) {
                        predicate_pushdown(LogicalPlan::Projection {
                            input: Box::new(LogicalPlan::Filter { input, predicate }),
                            exprs,
                            kind,
                        })
                    } else {
                        LogicalPlan::Filter {
                            input: Box::new(LogicalPlan::Projection { input, exprs, kind }),
                            predicate,
                        }
                    }
                }
                LogicalPlan::CsvScan {
                    path,
                    predicate: existing,
                    projection,
                } => LogicalPlan::CsvScan {
                    path,
                    predicate: Some(match existing {
                        Some(existing) => and_expr(existing, predicate),
                        None => predicate,
                    }),
                    projection,
                },
                LogicalPlan::ParquetScan {
                    path,
                    predicate: existing,
                    projection,
                } => LogicalPlan::ParquetScan {
                    path,
                    predicate: Some(match existing {
                        Some(existing) => and_expr(existing, predicate),
                        None => predicate,
                    }),
                    projection,
                },
                other => LogicalPlan::Filter {
                    input: Box::new(other),
                    predicate,
                },
            }
        }
        LogicalPlan::Projection { input, exprs, kind } => LogicalPlan::Projection {
            input: Box::new(predicate_pushdown(*input)),
            exprs,
            kind,
        },
        LogicalPlan::Concat { inputs, schema } => {
            let mut flattened = Vec::new();
            for input in inputs {
                match predicate_pushdown(input) {
                    LogicalPlan::Concat {
                        inputs: nested,
                        schema: nested_schema,
                    } if nested_schema == schema => flattened.extend(nested),
                    input => flattened.push(input),
                }
            }
            LogicalPlan::Concat {
                inputs: flattened,
                schema,
            }
        }
        LogicalPlan::Aggregate {
            input,
            group_by,
            aggs,
        } => LogicalPlan::Aggregate {
            input: Box::new(predicate_pushdown(*input)),
            group_by,
            aggs,
        },
        LogicalPlan::Join {
            left,
            right,
            keys,
            how,
        } => LogicalPlan::Join {
            left: Box::new(predicate_pushdown(*left)),
            right: Box::new(predicate_pushdown(*right)),
            keys,
            how,
        },
        LogicalPlan::Sort { input, options } => LogicalPlan::Sort {
            input: Box::new(predicate_pushdown(*input)),
            options,
        },
        LogicalPlan::Slice {
            input,
            offset,
            len,
            from_end,
        } => LogicalPlan::Slice {
            input: Box::new(predicate_pushdown(*input)),
            offset,
            len,
            from_end,
        },
        LogicalPlan::Unique { input, subset } => LogicalPlan::Unique {
            input: Box::new(predicate_pushdown(*input)),
            subset,
        },
        LogicalPlan::FillNull { input, fill } => LogicalPlan::FillNull {
            input: Box::new(predicate_pushdown(*input)),
            fill,
        },
        LogicalPlan::DropNulls { input, subset } => LogicalPlan::DropNulls {
            input: Box::new(predicate_pushdown(*input)),
            subset,
        },
        LogicalPlan::NullCount { input } => LogicalPlan::NullCount {
            input: Box::new(predicate_pushdown(*input)),
        },
        LogicalPlan::Explode { input, column } => LogicalPlan::Explode {
            input: Box::new(predicate_pushdown(*input)),
            column,
        },
        LogicalPlan::Implode { input } => LogicalPlan::Implode {
            input: Box::new(predicate_pushdown(*input)),
        },
        other => other,
    }
}

fn and_expr(left: Expr, right: Expr) -> Expr {
    let mut conjuncts = Vec::new();
    conjuncts.extend(flatten_and(left));
    conjuncts.extend(flatten_and(right));
    build_and(conjuncts)
}

fn flatten_and(expr: Expr) -> Vec<Expr> {
    match expr {
        E::BinaryOp {
            left,
            op: Operator::And,
            right,
        } => {
            let mut out = flatten_and(*left);
            out.extend(flatten_and(*right));
            out
        }
        other => vec![other],
    }
}

fn build_and(mut conjuncts: Vec<Expr>) -> Expr {
    let first = conjuncts
        .pop()
        .expect("build_and must be called with non-empty conjuncts");
    conjuncts
        .into_iter()
        .rev()
        .fold(first, |acc, expr| E::BinaryOp {
            left: Box::new(expr),
            op: Operator::And,
            right: Box::new(acc),
        })
}

fn can_push_filter_through_projection(
    predicate: &Expr,
    exprs: &[Expr],
    kind: &ProjectionKind,
) -> bool {
    let referenced = referenced_columns(predicate);

    match kind {
        ProjectionKind::Select => match projection_select_output_columns(exprs) {
            OutputColumns::Some(cols) => referenced.is_subset(&cols),
            OutputColumns::All | OutputColumns::Unknown => false,
        },
        ProjectionKind::WithColumns => {
            let assigned = projection_assigned_columns(exprs);
            !referenced.iter().any(|c| assigned.contains(c))
        }
    }
}

fn referenced_columns(expr: &Expr) -> HashSet<String> {
    let mut out = HashSet::new();
    collect_referenced_columns(expr, &mut out);
    out
}

fn collect_referenced_columns(expr: &Expr, out: &mut HashSet<String>) {
    match expr {
        E::Column(name) => {
            out.insert(name.clone());
        }
        E::Alias { expr, .. } => collect_referenced_columns(expr, out),
        E::UnaryOp { expr, .. } => collect_referenced_columns(expr, out),
        E::BinaryOp { left, right, .. } => {
            collect_referenced_columns(left, out);
            collect_referenced_columns(right, out);
        }
        E::Agg { expr, .. } => collect_referenced_columns(expr, out),
        E::Function { input, .. } => collect_referenced_columns(input, out),
        E::ConcatStr { inputs, .. } => {
            for input in inputs {
                collect_referenced_columns(input, out);
            }
        }
        E::Literal(_) | E::Wildcard => {}
    }
}

enum OutputColumns {
    All,
    Some(HashSet<String>),
    Unknown,
}

fn projection_select_output_columns(exprs: &[Expr]) -> OutputColumns {
    let mut cols = HashSet::new();
    for expr in exprs {
        match expr {
            E::Wildcard => return OutputColumns::All,
            E::Alias { name, .. } => {
                cols.insert(name.clone());
            }
            E::Column(name) => {
                cols.insert(name.clone());
            }
            _ => return OutputColumns::Unknown,
        }
    }
    OutputColumns::Some(cols)
}

fn projection_assigned_columns(exprs: &[Expr]) -> HashSet<String> {
    let mut cols = HashSet::new();
    for expr in exprs {
        match expr {
            E::Alias { name, .. } => {
                cols.insert(name.clone());
            }
            E::Column(name) => {
                cols.insert(name.clone());
            }
            _ => {}
        }
    }
    cols
}

fn projection_pushdown(plan: LogicalPlan) -> LogicalPlan {
    projection_pushdown_inner(plan, RequiredColumns::All).0
}

#[derive(Debug, Clone)]
enum RequiredColumns {
    All,
    Some(HashSet<String>),
}

impl RequiredColumns {
    fn union(self, other: Self) -> Self {
        match (self, other) {
            (RequiredColumns::All, _) | (_, RequiredColumns::All) => RequiredColumns::All,
            (RequiredColumns::Some(mut a), RequiredColumns::Some(b)) => {
                a.extend(b);
                RequiredColumns::Some(a)
            }
        }
    }
}

fn projection_pushdown_inner(
    plan: LogicalPlan,
    required: RequiredColumns,
) -> (LogicalPlan, RequiredColumns) {
    match plan {
        LogicalPlan::Projection { input, exprs, kind } => match kind {
            ProjectionKind::Select => {
                let input_required = required_columns_for_select(&exprs);
                let (new_input, _) = projection_pushdown_inner(*input, input_required);
                (
                    LogicalPlan::Projection {
                        input: Box::new(new_input),
                        exprs,
                        kind: ProjectionKind::Select,
                    },
                    required,
                )
            }
            ProjectionKind::WithColumns => {
                // Conservative: keep required columns and any inputs to compute overwritten/new columns
                let mut needed = HashSet::new();
                if let RequiredColumns::Some(ref req) = required {
                    needed.extend(req.clone());
                }
                for expr in &exprs {
                    match expr {
                        E::Alias { expr, .. } => needed.extend(referenced_columns(expr)),
                        E::Column(_) => {}
                        _ => {}
                    }
                }
                let (new_input, _) =
                    projection_pushdown_inner(*input, RequiredColumns::Some(needed));
                (
                    LogicalPlan::Projection {
                        input: Box::new(new_input),
                        exprs,
                        kind: ProjectionKind::WithColumns,
                    },
                    required,
                )
            }
        },
        LogicalPlan::Filter { input, predicate } => {
            let input_required = required
                .clone()
                .union(RequiredColumns::Some(referenced_columns(&predicate)));
            let (new_input, _) = projection_pushdown_inner(*input, input_required);
            (
                LogicalPlan::Filter {
                    input: Box::new(new_input),
                    predicate,
                },
                required,
            )
        }
        // Keep the validated concat schema intact.  Pushing a downstream projection into
        // individual inputs would require rewriting this node's schema as well as proving that
        // every input still produces the identical narrowed schema.  That optimisation is not
        // needed for v0.8 correctness, so recurse without narrowing each input.
        LogicalPlan::Concat { inputs, schema } => (
            LogicalPlan::Concat {
                inputs: inputs
                    .into_iter()
                    .map(|input| projection_pushdown_inner(input, RequiredColumns::All).0)
                    .collect(),
                schema,
            },
            required,
        ),
        LogicalPlan::Aggregate {
            input,
            group_by,
            aggs,
        } => {
            let mut needed = HashSet::new();
            for e in group_by.iter().chain(aggs.iter()) {
                needed.extend(referenced_columns(e));
            }
            let (new_input, _) = projection_pushdown_inner(*input, RequiredColumns::Some(needed));
            (
                LogicalPlan::Aggregate {
                    input: Box::new(new_input),
                    group_by,
                    aggs,
                },
                required,
            )
        }
        LogicalPlan::Join {
            left,
            right,
            keys,
            how,
        } => {
            let (new_left, _) = projection_pushdown_inner(*left, RequiredColumns::All);
            let (new_right, _) = projection_pushdown_inner(*right, RequiredColumns::All);
            (
                LogicalPlan::Join {
                    left: Box::new(new_left),
                    right: Box::new(new_right),
                    keys,
                    how,
                },
                required,
            )
        }
        LogicalPlan::Sort { input, options } => {
            let (new_input, _) = projection_pushdown_inner(*input, required.clone());
            (
                LogicalPlan::Sort {
                    input: Box::new(new_input),
                    options,
                },
                required,
            )
        }
        LogicalPlan::Slice {
            input,
            offset,
            len,
            from_end,
        } => {
            let (new_input, _) = projection_pushdown_inner(*input, required.clone());
            (
                LogicalPlan::Slice {
                    input: Box::new(new_input),
                    offset,
                    len,
                    from_end,
                },
                required,
            )
        }
        LogicalPlan::Unique { input, subset } => {
            let (new_input, _) = projection_pushdown_inner(*input, required.clone());
            (
                LogicalPlan::Unique {
                    input: Box::new(new_input),
                    subset,
                },
                required,
            )
        }
        LogicalPlan::FillNull { input, fill } => {
            let (new_input, _) = projection_pushdown_inner(*input, required.clone());
            (
                LogicalPlan::FillNull {
                    input: Box::new(new_input),
                    fill,
                },
                required,
            )
        }
        LogicalPlan::DropNulls { input, subset } => {
            let (new_input, _) = projection_pushdown_inner(*input, required.clone());
            (
                LogicalPlan::DropNulls {
                    input: Box::new(new_input),
                    subset,
                },
                required,
            )
        }
        LogicalPlan::NullCount { input } => {
            let (new_input, _) = projection_pushdown_inner(*input, RequiredColumns::All);
            (
                LogicalPlan::NullCount {
                    input: Box::new(new_input),
                },
                RequiredColumns::All,
            )
        }
        LogicalPlan::Explode { input, column } => {
            let (new_input, _) = projection_pushdown_inner(*input, RequiredColumns::All);
            (
                LogicalPlan::Explode {
                    input: Box::new(new_input),
                    column,
                },
                RequiredColumns::All,
            )
        }
        LogicalPlan::Implode { input } => {
            let (new_input, _) = projection_pushdown_inner(*input, RequiredColumns::All);
            (
                LogicalPlan::Implode {
                    input: Box::new(new_input),
                },
                RequiredColumns::All,
            )
        }
        LogicalPlan::CsvScan {
            path,
            predicate,
            projection,
        } => {
            let mut needed = match required {
                RequiredColumns::All => None,
                RequiredColumns::Some(s) => Some(s),
            };
            if let Some(pred) = &predicate {
                let cols = referenced_columns(pred);
                needed = Some(match needed {
                    Some(mut s) => {
                        s.extend(cols);
                        s
                    }
                    None => cols,
                });
            }
            (
                LogicalPlan::CsvScan {
                    path,
                    predicate,
                    projection: merge_projection(projection, needed),
                },
                RequiredColumns::All,
            )
        }
        LogicalPlan::ParquetScan {
            path,
            predicate,
            projection,
        } => {
            let mut needed = match required {
                RequiredColumns::All => None,
                RequiredColumns::Some(s) => Some(s),
            };
            if let Some(pred) = &predicate {
                let cols = referenced_columns(pred);
                needed = Some(match needed {
                    Some(mut s) => {
                        s.extend(cols);
                        s
                    }
                    None => cols,
                });
            }
            (
                LogicalPlan::ParquetScan {
                    path,
                    predicate,
                    projection: merge_projection(projection, needed),
                },
                RequiredColumns::All,
            )
        }
        other => (other, RequiredColumns::All),
    }
}

fn required_columns_for_select(exprs: &[Expr]) -> RequiredColumns {
    let mut needed = HashSet::new();
    for expr in exprs {
        match expr {
            E::Wildcard => return RequiredColumns::All,
            other => needed.extend(referenced_columns(other)),
        }
    }
    RequiredColumns::Some(needed)
}

fn merge_projection(
    existing: Option<Vec<String>>,
    needed: Option<HashSet<String>>,
) -> Option<Vec<String>> {
    let Some(needed) = needed else {
        return existing;
    };

    let mut out = Vec::new();
    let mut seen = HashSet::new();

    if let Some(existing) = existing {
        for c in existing {
            if seen.insert(c.clone()) {
                out.push(c);
            }
        }
    }

    for c in needed {
        if seen.insert(c.clone()) {
            out.push(c);
        }
    }

    Some(out)
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use arrow::datatypes::{DataType, Field, Schema};

    use super::Optimizer;
    use crate::expr::{col, lit};
    use crate::lazy::LogicalPlan;
    use crate::lazy::ProjectionKind;

    #[test]
    fn predicate_pushdown_moves_filter_into_scan() {
        let plan = LogicalPlan::Filter {
            input: Box::new(LogicalPlan::CsvScan {
                path: "data.csv".into(),
                predicate: None,
                projection: None,
            }),
            predicate: col("a").gt(lit(1_i64)),
        };

        let optimized = Optimizer::optimize(&plan);
        match optimized {
            LogicalPlan::CsvScan { predicate, .. } => assert!(predicate.is_some()),
            other => panic!("expected CsvScan, got {other:?}"),
        }
    }

    #[test]
    fn predicate_pushdown_combines_multiple_filters_with_and() {
        let plan = LogicalPlan::Filter {
            input: Box::new(LogicalPlan::Filter {
                input: Box::new(LogicalPlan::CsvScan {
                    path: "data.csv".into(),
                    predicate: None,
                    projection: None,
                }),
                predicate: col("a").gt(lit(1_i64)),
            }),
            predicate: col("b").lt(lit(10_i64)),
        };

        let optimized = Optimizer::optimize(&plan);
        match optimized {
            LogicalPlan::CsvScan {
                predicate: Some(p), ..
            } => {
                let s = format!("{p:?}");
                assert!(s.contains("And"));
            }
            other => panic!("expected CsvScan with predicate, got {other:?}"),
        }
    }

    #[test]
    fn predicate_pushdown_does_not_cross_select_when_column_not_selected() {
        let plan = LogicalPlan::Filter {
            input: Box::new(LogicalPlan::Projection {
                input: Box::new(LogicalPlan::CsvScan {
                    path: "data.csv".into(),
                    predicate: None,
                    projection: None,
                }),
                exprs: vec![col("a")],
                kind: ProjectionKind::Select,
            }),
            predicate: col("b").gt(lit(1_i64)),
        };

        let optimized = Optimizer::optimize(&plan);
        assert!(matches!(optimized, LogicalPlan::Filter { .. }));
    }

    #[test]
    fn projection_pushdown_sets_scan_projection() {
        let plan = LogicalPlan::Projection {
            input: Box::new(LogicalPlan::CsvScan {
                path: "data.csv".into(),
                predicate: None,
                projection: None,
            }),
            exprs: vec![col("a"), col("b")],
            kind: ProjectionKind::Select,
        };

        let optimized = Optimizer::optimize(&plan);
        let s = optimized.display();
        assert!(s.contains("projection"));
    }

    #[test]
    fn concat_optimizer_flattens_only_compatible_nested_inputs_in_order() {
        let schema = Arc::new(Schema::new(vec![Field::new(
            "value",
            DataType::Int64,
            true,
        )]));
        let scan = |path: &str| LogicalPlan::CsvScan {
            path: path.into(),
            predicate: None,
            projection: None,
        };
        let plan = LogicalPlan::Concat {
            inputs: vec![
                LogicalPlan::Concat {
                    inputs: vec![scan("first.csv"), scan("second.csv")],
                    schema: Some(schema.clone()),
                },
                scan("third.csv"),
            ],
            schema: Some(schema),
        };

        let display = Optimizer::optimize(&plan).display();
        assert!(display.contains("concat inputs=3"));
        assert!(display.find("first.csv").unwrap() < display.find("second.csv").unwrap());
        assert!(display.find("second.csv").unwrap() < display.find("third.csv").unwrap());
    }
}