algeon-query-planning 0.1.2

//! Mirrors post-aggregate group-key filters onto dimension reader predicates.
//!
//! The rule matches a `Filter` conjunct whose columns trace, through renaming
//! projections, a grouping-only aggregate, the preserved side of an inner or left
//! join, and sorts, to file columns of one local Parquet source, crossing at
//! least one aggregate and no node with more than one parent. In the rewrite, the
//! conjunct, rewritten to source column names, is conjoined onto that source's
//! reader predicate; the original filter stays in place. The copy changes no
//! result because a group-key value that fails the filter above the aggregate
//! fails it below as well, a preserved join side carries its input's own values,
//! and the filter above still runs. Only conjuncts that crossed an aggregate are
//! seeded; a filter that reaches a source without one is left to the frontend's
//! own predicate placement. Sources are refined in place
//! (`RewrittenNodes::refined_sources`), so routes bound to them survive.

use super::{
    AggregateNode, AnalysisContext, BinaryOp, ColumnRef, Expr, JoinNode, JoinSide, JoinType, LocalParquetSource,
    NodeId, PlanNode, ProjectionNode, QueryNode, QueryPlan, RewrittenNodes, SourceColumnRef, SourceNode,
    mutable_local_parquet_source, predicate_columns, rewrite,
};
use crate::error::Result;

/// Seeds every eligible conjunct and reports whether any source changed. A
/// conjunct already present on the source is not added twice, so a second run
/// over the same plan reports no change.
pub(super) fn seed_reader_predicates_from_group_key_filters(
    plan: QueryPlan,
    ctx: AnalysisContext<'_>,
) -> Result<(QueryPlan, bool)> {
    let mut nodes = plan.nodes().to_vec();
    let parent_counts = rewrite::parent_counts(&plan);
    let mut changed = false;

    for node in plan.nodes() {
        let PlanNode::Filter(filter) = &node.kind else {
            continue;
        };
        for conjunct in predicate_conjuncts(&filter.predicate) {
            changed |= push_filter_conjunct_to_reader_source(
                &plan,
                &parent_counts,
                &mut nodes,
                filter.input,
                conjunct.clone(),
                false,
                false,
                ctx,
            );
        }
    }

    if !changed {
        return Ok((plan, false));
    }
    let root = plan.root_id();
    Ok((plan.rewrite(RewrittenNodes::refined_sources(nodes, root))?, true))
}

/// Returns a copy of `plan` with eligible post-aggregate group-key filters
/// mirrored onto local Parquet dimension-reader predicates.
///
/// The original post-aggregate filter remains in the plan. This helper is used
/// by adapter-side object-store pruning, where row-group statistics are not
/// available inside `query-engine` but the seeded dimension predicate is
/// still needed to derive a key range.
pub fn seed_dimension_derived_reader_predicates(plan: QueryPlan, ctx: AnalysisContext<'_>) -> Result<QueryPlan> {
    Ok(seed_reader_predicates_from_group_key_filters(plan, ctx)?.0)
}

fn predicate_conjuncts(expr: &Expr) -> Vec<&Expr> {
    match expr {
        Expr::Binary { left, op: BinaryOp::And, right, .. } => {
            let mut conjuncts = predicate_conjuncts(left);
            conjuncts.extend(predicate_conjuncts(right));
            conjuncts
        },
        _ => vec![expr],
    }
}

#[expect(
    clippy::too_many_arguments,
    reason = "each parameter is independent recursive-walk state (plan, parent_counts, the rewritten node buffer, \
              cursor, carried predicate, two crossed-boundary flags, and the backend context); grouping them into a \
              struct would not reduce the walk's real state count"
)]
fn push_filter_conjunct_to_reader_source(
    plan: &QueryPlan,
    parent_counts: &[usize],
    nodes: &mut [QueryNode],
    node_id: NodeId,
    predicate: Expr,
    crossed_aggregate: bool,
    crossed_shared_path: bool,
    ctx: AnalysisContext<'_>,
) -> bool {
    let Some(node) = plan.node(node_id) else {
        return false;
    };
    // Past a node with a second parent the conjunct describes only one of the
    // source's consumers; the flag is sticky so the source is never seeded.
    let crossed_shared_path = crossed_shared_path || parent_counts[node_id.index()] > 1;
    match &node.kind {
        PlanNode::Source(source) => match source.as_ref() {
            SourceNode::LocalParquet(source)
                if crossed_aggregate
                    && !crossed_shared_path
                    && seed_predicate_is_supported(source, &predicate, ctx) =>
            {
                let Some(source) = mutable_local_parquet_source(nodes, node_id) else {
                    return false;
                };
                if source.reader_predicate.as_ref() == Some(&predicate) {
                    return false;
                }
                source.reader_predicate = Some(match source.reader_predicate.take() {
                    Some(existing) => Expr::binary(existing, BinaryOp::And, predicate),
                    None => predicate,
                });
                source.refresh_diagnostics();
                true
            },
            _ => false,
        },
        PlanNode::Filter(filter) => push_filter_conjunct_to_reader_source(
            plan,
            parent_counts,
            nodes,
            filter.input,
            predicate,
            crossed_aggregate,
            crossed_shared_path,
            ctx,
        ),
        PlanNode::Projection(projection) => {
            remap_projection_predicate(&predicate, projection).is_some_and(|predicate| {
                push_filter_conjunct_to_reader_source(
                    plan,
                    parent_counts,
                    nodes,
                    projection.input,
                    predicate,
                    crossed_aggregate,
                    crossed_shared_path,
                    ctx,
                )
            })
        },
        PlanNode::Aggregate(aggregate) => {
            remap_aggregate_group_predicate(&predicate, aggregate).is_some_and(|predicate| {
                push_filter_conjunct_to_reader_source(
                    plan,
                    parent_counts,
                    nodes,
                    aggregate.input,
                    predicate,
                    true,
                    crossed_shared_path,
                    ctx,
                )
            })
        },
        PlanNode::Join(join) => remap_join_output_predicate(&predicate, join).is_some_and(|(side, predicate)| {
            let input = match side {
                JoinSide::Left => join.inputs[0],
                JoinSide::Right => join.inputs[1],
            };
            push_filter_conjunct_to_reader_source(
                plan,
                parent_counts,
                nodes,
                input,
                predicate,
                crossed_aggregate,
                crossed_shared_path,
                ctx,
            )
        }),
        PlanNode::Sort(sort) => push_filter_conjunct_to_reader_source(
            plan,
            parent_counts,
            nodes,
            sort.input,
            predicate,
            crossed_aggregate,
            crossed_shared_path,
            ctx,
        ),
        PlanNode::Limit(_)
        | PlanNode::Window(_)
        | PlanNode::Union(_)
        | PlanNode::EdgeNormalize(_)
        | PlanNode::GraphAlgorithm(_)
        | PlanNode::VectorAlgorithm(_)
        | PlanNode::MlAlgorithm(_) => false,
    }
}

fn seed_predicate_is_supported(source: &LocalParquetSource, predicate: &Expr, ctx: AnalysisContext<'_>) -> bool {
    if !predicate_is_reader_visible(source, predicate) {
        return false;
    }
    let combined = match source.reader_predicate.as_ref() {
        Some(existing) => Expr::binary(existing.clone(), BinaryOp::And, predicate.clone()),
        None => predicate.clone(),
    };
    source.reader_predicate_is_supported(&combined, ctx)
}

/// Whether every column the predicate reads is a physical file column. Partition
/// and missing columns are synthesized after the read and are not visible to the
/// Parquet reader predicate.
fn predicate_is_reader_visible(source: &LocalParquetSource, predicate: &Expr) -> bool {
    let columns = predicate_columns(predicate);
    !columns.is_empty()
        && columns.iter().all(|column| matches!(source.source_column_ref(column), Some(SourceColumnRef::File(_))))
}

fn remap_projection_predicate(predicate: &Expr, projection: &ProjectionNode) -> Option<Expr> {
    remap_unqualified_columns(predicate, &mut |column| {
        let output = projection.outputs.iter().find(|output| output.field.name() == column.name())?;
        simple_projection_expr(&output.expr)
    })
}

fn remap_aggregate_group_predicate(predicate: &Expr, aggregate: &AggregateNode) -> Option<Expr> {
    if aggregate.spec.grouping_sets.is_some() {
        return None;
    }
    remap_unqualified_columns(predicate, &mut |column| {
        aggregate
            .spec
            .group_exprs
            .iter()
            .find(|expr| aggregate_group_output_name(expr) == Some(column.name()))
            .and_then(simple_projection_expr)
    })
}

fn remap_join_output_predicate(predicate: &Expr, join: &JoinNode) -> Option<(JoinSide, Expr)> {
    if !matches!(join.spec.join_type, JoinType::Inner | JoinType::Left) || join.spec.output_projection.is_empty() {
        return None;
    }
    let mut side = None;
    let remapped = remap_unqualified_columns(predicate, &mut |column| {
        let output = join.spec.output_projection.iter().find(|output| output.output_name == column.name())?;
        // A left join's right-side columns are null for unmatched rows, so a
        // predicate over them cannot be evaluated on the right input alone.
        if join.spec.join_type == JoinType::Left && output.side != JoinSide::Left {
            return None;
        }
        match side {
            Some(existing) if existing != output.side => return None,
            Some(_) => {},
            None => side = Some(output.side),
        }
        Some(Expr::column(output.column.name()))
    })?;
    side.map(|side| (side, remapped))
}

// This helper is deliberately narrower than the similarly named helper in
// `join_filter_reassociation`. Dimension pruning must reject expression
// families it cannot prove safe to push into source pruning.
pub(super) fn remap_unqualified_columns(
    expr: &Expr,
    map_column: &mut impl FnMut(&ColumnRef) -> Option<Expr>,
) -> Option<Expr> {
    match expr {
        Expr::Column(column) if column.qualifier().is_none() => map_column(column),
        Expr::Column(_) => None,
        Expr::Literal(value) => Some(Expr::Literal(value.clone())),
        Expr::Alias { expr, alias } => Some(remap_unqualified_columns(expr, map_column)?.alias(alias.clone())),
        Expr::Cast { expr, target_type } => {
            Some(Expr::cast(remap_unqualified_columns(expr, map_column)?, target_type.clone()))
        },
        Expr::TryCast { expr, target_type } => {
            Some(Expr::try_cast(remap_unqualified_columns(expr, map_column)?, target_type.clone()))
        },
        Expr::Binary { left, op, right, output_type } => Some(Expr::Binary {
            left: Box::new(remap_unqualified_columns(left, map_column)?),
            op: *op,
            right: Box::new(remap_unqualified_columns(right, map_column)?),
            output_type: output_type.clone(),
        }),
        Expr::Not(expr) => Some(Expr::logical_not(remap_unqualified_columns(expr, map_column)?)),
        Expr::IsNull(expr) => Some(Expr::is_null(remap_unqualified_columns(expr, map_column)?)),
        Expr::IsNotNull(expr) => Some(Expr::is_not_null(remap_unqualified_columns(expr, map_column)?)),
        Expr::InList { expr, list, negated } => Some(Expr::InList {
            expr: Box::new(remap_unqualified_columns(expr, map_column)?),
            list: list.iter().map(|expr| remap_unqualified_columns(expr, map_column)).collect::<Option<Vec<_>>>()?,
            negated: *negated,
        }),
        Expr::Like { .. }
        | Expr::Substring { .. }
        | Expr::Lower(_)
        | Expr::Upper(_)
        | Expr::Trim { .. }
        | Expr::LTrim { .. }
        | Expr::RTrim { .. }
        | Expr::StringContains { .. }
        | Expr::StartsWith { .. }
        | Expr::EndsWith { .. }
        | Expr::StringLength { .. }
        | Expr::RegexpReplace { .. }
        | Expr::Concat(_)
        | Expr::Abs(_)
        | Expr::UnaryMath { .. }
        | Expr::Round { .. }
        | Expr::IsDistinctFrom { .. }
        | Expr::IsNotDistinctFrom { .. }
        | Expr::Case { .. }
        | Expr::Extract { .. }
        | Expr::DateTrunc { .. } => None,
    }
}

fn simple_projection_expr(expr: &Expr) -> Option<Expr> {
    match expr {
        Expr::Column(column) => Some(Expr::column(column.name())),
        Expr::Alias { expr, .. } => simple_projection_expr(expr),
        _ => None,
    }
}

pub(super) fn aggregate_group_output_name(expr: &Expr) -> Option<&str> {
    match expr {
        Expr::Column(column) => Some(column.display_name()),
        Expr::Alias { alias, .. } => Some(alias.as_str()),
        _ => None,
    }
}