algeon-query-planning 0.1.2

//! Rewrite planning and the node rewriter for the late materialization pass.
//!
//! `shared_aggregate_family_rewrite_plan` decides, without touching the plan,
//! how one family is shared and which predicate each branch keeps; every
//! `Err` is the report reason for that family. `SharedSubplanMaterializationRewriter`
//! then rebuilds the plan: the representative subtree is cloned once per
//! family with widened sources, and each member aggregate is replaced by its
//! branch node over that clone.

use std::{
    collections::{BTreeMap, BTreeSet},
    sync::Arc,
};

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

use super::{
    diagnostics::{SharedSubplanMaterializationDiagnostics, SharedSubplanMaterializationRejectedReason},
    expr_helpers::{
        collect_expr_column_names, collect_join_expr_columns, conjunction, disjunction, field_with_name,
        predicate_references_only_source_filterable_columns, scalar_count_star_aggregate_expr, scalar_count_star_expr,
        scalar_filtered_aggregate_streaming_supported, scalar_filtered_aggregate_supported,
        source_reader_predicate_filter_is_elidable,
    },
    lineage::{
        join_output_schema, projection_output_preserves_column, propagatable_materialized_field,
        scalar_input_expr_lineage_signature, scalar_input_expr_name_collides_with_unresolved_output,
        subtree_can_propagate_materialized_column, unique_join_input_that_can_propagate_column,
        unique_join_input_that_can_propagate_materialized_column,
    },
    signature::{
        AggregateMaterializationSignature, SourceColumnMaterialization, local_source_materialization_signature,
    },
    source_widening::{
        matching_source_entry, merge_source_columns, source_column_names_are_unique_across_sources,
        source_materialization_sources_are_unique, widen_local_parquet_source,
    },
};
use crate::{
    expr::{AggregateExpr, AggregateSpec, BinaryOp, ColumnRef, Expr, LiteralValue},
    optimizer::rewrite as optimizer_rewrite,
    plan::{
        AggregateNode, FilterNode, JoinOutputColumn, JoinSide, NodeId, PlanFacts, PlanNode, ProjectionNode,
        ProjectionOutput, QueryNode, QueryPlan,
    },
    source::SourceNode,
};

/// Shape signature of a scalar family's aggregate input, or `None` when it has
/// an operator the scalar rewrite cannot clone with extra columns (a limit,
/// aggregate, window, union or GPU function). Identity projections, filters
/// that repeat the reader predicate, and, when allowed, column projections are
/// left out so branches that differ only in those compare equal.
pub(super) fn scalar_materialization_input_shape(
    plan: &QueryPlan,
    node_id: NodeId,
    allow_column_projection_elision: bool,
) -> Option<String> {
    scalar_materialization_input_shape_inner(plan, node_id, true, allow_column_projection_elision)
}

fn scalar_materialization_input_shape_inner(
    plan: &QueryPlan,
    node_id: NodeId,
    allow_count_star_projection_elision: bool,
    allow_column_projection_elision: bool,
) -> Option<String> {
    let node = plan.node(node_id)?;
    match &node.kind {
        PlanNode::Source(source) => match source.as_ref() {
            SourceNode::LocalParquet(source) => {
                let signature = local_source_materialization_signature(source);
                Some(format!("local_parquet({})", signature.rowset_fingerprint))
            },
            _ => None,
        },
        PlanNode::Filter(filter) if source_reader_predicate_filter_is_elidable(plan, filter) => {
            scalar_materialization_input_shape_inner(
                plan,
                filter.input,
                allow_count_star_projection_elision,
                allow_column_projection_elision,
            )
        },
        PlanNode::Filter(filter) => Some(format!(
            "filter(predicate={};input={})",
            filter.predicate,
            scalar_materialization_input_shape_inner(plan, filter.input, false, allow_column_projection_elision)?
        )),
        PlanNode::Projection(projection) if identity_projection_is_elidable(projection) => {
            scalar_materialization_input_shape_inner(
                plan,
                projection.input,
                allow_count_star_projection_elision,
                allow_column_projection_elision,
            )
        },
        PlanNode::Projection(projection) if allow_column_projection_elision => {
            scalar_materialization_input_shape_inner(
                plan,
                projection.input,
                allow_count_star_projection_elision,
                allow_column_projection_elision,
            )
        },
        PlanNode::Projection(projection)
            if allow_count_star_projection_elision && count_star_projection_is_elidable(projection) =>
        {
            scalar_materialization_input_shape_inner(plan, projection.input, true, allow_column_projection_elision)
        },
        PlanNode::Projection(projection) => {
            let outputs = projection
                .outputs
                .iter()
                .map(|output| {
                    format!(
                        "{}:{:?}:nullable={}={}",
                        output.field.name(),
                        output.field.data_type(),
                        output.field.is_nullable(),
                        output.expr
                    )
                })
                .collect::<Vec<_>>()
                .join(",");
            Some(format!(
                "projection(outputs=[{outputs}];input={})",
                scalar_materialization_input_shape_inner(
                    plan,
                    projection.input,
                    false,
                    allow_column_projection_elision
                )?
            ))
        },
        PlanNode::Join(join) => Some(format!(
            "join(spec={:?};input0={};input1={})",
            join.spec,
            scalar_materialization_input_shape_inner(plan, join.inputs[0], false, allow_column_projection_elision)?,
            scalar_materialization_input_shape_inner(plan, join.inputs[1], false, allow_column_projection_elision)?
        )),
        PlanNode::Sort(sort) => Some(format!(
            "sort(spec={:?};input={})",
            sort.spec,
            scalar_materialization_input_shape_inner(plan, sort.input, false, allow_column_projection_elision)?
        )),
        PlanNode::Limit(_) => None,
        PlanNode::Aggregate(_)
        | PlanNode::Window(_)
        | PlanNode::Union(_)
        | PlanNode::EdgeNormalize(_)
        | PlanNode::GraphAlgorithm(_)
        | PlanNode::VectorAlgorithm(_)
        | PlanNode::MlAlgorithm(_) => None,
    }
}

fn identity_projection_is_elidable(projection: &ProjectionNode) -> bool {
    !projection.outputs.is_empty()
        && projection.outputs.iter().all(|output| projection_output_preserves_column(output, output.field.name()))
}

fn count_star_projection_is_elidable(projection: &ProjectionNode) -> bool {
    projection.outputs.is_empty()
        || projection.outputs.iter().all(|output| {
            scalar_count_star_expr(&output.expr) || projection_output_preserves_column(output, output.field.name())
        })
}

/// How one family is shared. The representative's subtree is the one cloned;
/// the other members contribute only their predicates and columns.
#[derive(Debug, Clone)]
pub(super) struct SharedAggregateFamilyRewrite {
    family_key: String,
    representative_aggregate_node_id: NodeId,
    /// Reader predicate the cloned source gets, by source fingerprint. Absent
    /// for a source whose predicate every member shares.
    widened_source_predicates: BTreeMap<String, Expr>,
    /// Columns the cloned source must read beyond the representative's, by
    /// source fingerprint. Grouped families never widen columns.
    widened_source_columns: BTreeMap<String, BTreeMap<String, SourceColumnMaterialization>>,
    kind: SharedAggregateFamilyRewriteKind,
}

#[derive(Debug, Clone)]
enum SharedAggregateFamilyRewriteKind {
    /// Each member becomes a filter over the shared aggregate; `None` means
    /// the member reads the shared aggregate directly.
    Grouped { branch_filters: BTreeMap<NodeId, Option<Expr>> },
    /// Each member becomes a projection that renames its `FILTER`ed outputs
    /// of the one shared scalar aggregate.
    ScalarFiltered { required_aggregate_input_columns: BTreeSet<String>, branches: Vec<ScalarFilteredAggregateBranch> },
}

#[derive(Debug, Clone)]
struct ScalarFilteredAggregateBranch {
    aggregate_node_id: NodeId,
    outputs: Vec<ScalarFilteredAggregateOutput>,
}

#[derive(Debug, Clone)]
struct ScalarFilteredAggregateOutput {
    original_output_field: Field,
    materialized_alias: String,
    aggregate_expr: AggregateExpr,
}

/// Sources whose reader predicate differs between members and so must be
/// widened. A source read identically by every member keeps its predicate.
fn widened_source_fingerprints(family: &[AggregateMaterializationSignature]) -> BTreeSet<String> {
    let mut predicates_by_source = BTreeMap::<String, BTreeSet<String>>::new();
    for signature in family {
        for source in &signature.source_predicates {
            predicates_by_source
                .entry(source.abstract_fingerprint.clone())
                .or_default()
                .insert(source.predicate.as_ref().map_or_else(|| "none".to_string(), ToString::to_string));
        }
    }
    predicates_by_source
        .into_iter()
        .filter_map(|(fingerprint, predicates)| (predicates.len() > 1).then_some(fingerprint))
        .collect()
}

/// Plans how `family` is shared, or names the reason it cannot be. The
/// reason is a report contract; the check order fixes which one a family with
/// several defects reports.
pub(super) fn shared_aggregate_family_rewrite_plan(
    plan: &QueryPlan,
    family: &[AggregateMaterializationSignature],
) -> Result<SharedAggregateFamilyRewrite, SharedSubplanMaterializationRejectedReason> {
    let representative =
        family.first().ok_or(SharedSubplanMaterializationRejectedReason::PredicateNotPostAggregateFilterable)?;
    if representative.group_exprs.is_empty() {
        return shared_scalar_filtered_aggregate_rewrite_plan(plan, family);
    }
    if family.iter().any(|signature| !source_materialization_sources_are_unique(signature)) {
        return Err(SharedSubplanMaterializationRejectedReason::PredicateColumnPropagationUnsupported);
    }
    let widened_source_fingerprints = widened_source_fingerprints(family);
    if widened_source_fingerprints.is_empty() {
        return Ok(SharedAggregateFamilyRewrite {
            family_key: representative.family_key.clone(),
            representative_aggregate_node_id: representative.aggregate_node_id,
            widened_source_predicates: BTreeMap::new(),
            widened_source_columns: BTreeMap::new(),
            kind: SharedAggregateFamilyRewriteKind::Grouped {
                branch_filters: family.iter().map(|signature| (signature.aggregate_node_id, None)).collect(),
            },
        });
    }

    let mut widened_source_predicates = BTreeMap::new();
    let mut branch_filter_parts =
        family.iter().map(|signature| (signature.aggregate_node_id, Vec::new())).collect::<BTreeMap<_, _>>();
    for fingerprint in &widened_source_fingerprints {
        let factored = factored_grouped_source_predicate(family, fingerprint)?;
        widened_source_predicates.insert(fingerprint.clone(), factored.widened_source_predicate);
        for (node_id, predicate) in factored.branch_predicates {
            if let Some(predicate) = predicate {
                branch_filter_parts.entry(node_id).or_default().push(predicate);
            }
        }
    }
    let branch_filters = branch_filter_parts
        .into_iter()
        .map(|(node_id, predicates)| (node_id, conjunction(predicates)))
        .collect::<BTreeMap<_, _>>();

    Ok(SharedAggregateFamilyRewrite {
        family_key: representative.family_key.clone(),
        representative_aggregate_node_id: representative.aggregate_node_id,
        widened_source_predicates,
        widened_source_columns: BTreeMap::new(),
        kind: SharedAggregateFamilyRewriteKind::Grouped { branch_filters },
    })
}

struct FactoredGroupedSourcePredicate {
    widened_source_predicate: Expr,
    branch_predicates: BTreeMap<NodeId, Option<Expr>>,
}

/// Splits the members' reader predicates on one source into the conjuncts
/// they all share and the conjuncts that vary.
///
/// The shared source keeps the common conjuncts and reads the disjunction of
/// the varying ones, so it produces every row any member needs and no row
/// none of them needs. Each member re-applies its own varying conjuncts after
/// the aggregate, which is correct only when they read group columns the
/// source emits unchanged; any other column is a `PredicateNotPostAggregateFilterable`
/// rejection. A member with no varying conjunct needs every row the common
/// part admits, so the disjunction is then omitted.
fn factored_grouped_source_predicate(
    family: &[AggregateMaterializationSignature],
    fingerprint: &str,
) -> Result<FactoredGroupedSourcePredicate, SharedSubplanMaterializationRejectedReason> {
    let mut branches = Vec::with_capacity(family.len());
    for signature in family {
        let source = signature
            .source_predicates
            .iter()
            .find(|source| source.abstract_fingerprint == fingerprint)
            .ok_or(SharedSubplanMaterializationRejectedReason::PredicateNotPostAggregateFilterable)?;
        let predicate = source
            .predicate
            .clone()
            .ok_or(SharedSubplanMaterializationRejectedReason::PredicateNotPostAggregateFilterable)?;
        let conjuncts = split_conjuncts(predicate);
        let keys = conjuncts.iter().map(ToString::to_string).collect::<BTreeSet<_>>();
        branches.push((signature, conjuncts, keys));
    }

    let mut common_keys = branches
        .first()
        .map(|(_, _, keys)| keys.clone())
        .ok_or(SharedSubplanMaterializationRejectedReason::PredicateNotPostAggregateFilterable)?;
    for (_, _, keys) in branches.iter().skip(1) {
        common_keys = common_keys.intersection(keys).cloned().collect();
    }

    let mut common_seen = BTreeSet::new();
    let common_predicates = branches[0]
        .1
        .iter()
        .filter(|expr| common_keys.contains(&expr.to_string()))
        .filter(|expr| common_seen.insert(expr.to_string()))
        .cloned()
        .collect::<Vec<_>>();

    let mut branch_predicates = BTreeMap::new();
    let mut varying_predicates = Vec::new();
    let mut branch_without_varying_predicate = false;
    for (signature, conjuncts, _) in branches {
        let varying = conjuncts.into_iter().filter(|expr| !common_keys.contains(&expr.to_string())).collect::<Vec<_>>();
        let branch_predicate = conjunction(varying);
        if let Some(predicate) = &branch_predicate {
            if !predicate_references_only_source_filterable_columns(predicate, signature, fingerprint) {
                return Err(SharedSubplanMaterializationRejectedReason::PredicateNotPostAggregateFilterable);
            }
            varying_predicates.push(predicate.clone());
        } else {
            branch_without_varying_predicate = true;
        }
        branch_predicates.insert(signature.aggregate_node_id, branch_predicate);
    }

    let mut widened_predicates = common_predicates;
    if !branch_without_varying_predicate {
        let varying_union = disjunction(varying_predicates)
            .ok_or(SharedSubplanMaterializationRejectedReason::PredicateNotPostAggregateFilterable)?;
        widened_predicates.push(varying_union);
    }
    let widened_source_predicate = conjunction(widened_predicates)
        .ok_or(SharedSubplanMaterializationRejectedReason::PredicateNotPostAggregateFilterable)?;

    Ok(FactoredGroupedSourcePredicate { widened_source_predicate, branch_predicates })
}

fn split_conjuncts(predicate: Expr) -> Vec<Expr> {
    match predicate {
        Expr::Binary { left, op: BinaryOp::And, right, output_type: None } => {
            let mut conjuncts = split_conjuncts(*left);
            conjuncts.extend(split_conjuncts(*right));
            conjuncts
        },
        other => vec![other],
    }
}

/// Plans a scalar filtered family: one aggregate whose expressions are every
/// member's aggregates with that member's reader predicate as a `FILTER`
/// clause, over the shared source read with the disjunction of those
/// predicates.
///
/// The filter columns must reach the aggregate input, so every projection and
/// join on the way is checked for the ability to carry them, and the widened
/// column set must be unambiguous across sources. Count-star arguments are
/// rewritten to a literal because the shared input no longer has the
/// member's count-preserving column.
fn shared_scalar_filtered_aggregate_rewrite_plan(
    plan: &QueryPlan,
    family: &[AggregateMaterializationSignature],
) -> Result<SharedAggregateFamilyRewrite, SharedSubplanMaterializationRejectedReason> {
    let representative =
        family.first().ok_or(SharedSubplanMaterializationRejectedReason::UnsupportedScalarAggregateShape)?;
    let representative_node = plan
        .node(representative.aggregate_node_id)
        .ok_or(SharedSubplanMaterializationRejectedReason::UnsupportedScalarAggregateShape)?;
    let PlanNode::Aggregate(representative_aggregate) = &representative_node.kind else {
        return Err(SharedSubplanMaterializationRejectedReason::UnsupportedScalarAggregateShape);
    };
    if !scalar_filtered_aggregate_supported(&representative_aggregate.spec) {
        return Err(SharedSubplanMaterializationRejectedReason::UnsupportedScalarAggregateShape);
    }
    if !scalar_filtered_aggregate_streaming_supported(&representative_aggregate.spec) {
        return Err(SharedSubplanMaterializationRejectedReason::UnsupportedScalarAggregateShape);
    }
    if family.iter().any(|signature| !source_materialization_sources_are_unique(signature)) {
        return Err(SharedSubplanMaterializationRejectedReason::PredicateColumnPropagationUnsupported);
    }

    for signature in family {
        let node = plan
            .node(signature.aggregate_node_id)
            .ok_or(SharedSubplanMaterializationRejectedReason::UnsupportedScalarAggregateShape)?;
        let PlanNode::Aggregate(aggregate) = &node.kind else {
            return Err(SharedSubplanMaterializationRejectedReason::UnsupportedScalarAggregateShape);
        };
        if !scalar_filtered_aggregate_supported(&aggregate.spec) {
            return Err(SharedSubplanMaterializationRejectedReason::UnsupportedScalarAggregateShape);
        }
        if !scalar_filtered_aggregate_streaming_supported(&aggregate.spec) {
            return Err(SharedSubplanMaterializationRejectedReason::UnsupportedScalarAggregateShape);
        }
        for aggregate_expr in &aggregate.spec.aggr_exprs {
            if scalar_count_star_aggregate_expr(plan, aggregate.input, &aggregate_expr.expr) {
                continue;
            }
            let Some(branch_signature) =
                scalar_input_expr_lineage_signature(plan, aggregate.input, &aggregate_expr.expr)
            else {
                return Err(SharedSubplanMaterializationRejectedReason::PredicateColumnPropagationUnsupported);
            };
            match scalar_input_expr_lineage_signature(plan, representative_aggregate.input, &aggregate_expr.expr) {
                Some(representative_signature) if branch_signature != representative_signature => {
                    return Err(SharedSubplanMaterializationRejectedReason::PredicateColumnPropagationUnsupported);
                },
                Some(_) => {},
                None => {
                    if scalar_input_expr_name_collides_with_unresolved_output(
                        plan,
                        representative_aggregate.input,
                        &aggregate_expr.expr,
                    ) {
                        return Err(SharedSubplanMaterializationRejectedReason::PredicateColumnPropagationUnsupported);
                    }
                },
            }
        }
    }

    let widened_source_fingerprints = widened_source_fingerprints(family);
    if widened_source_fingerprints.is_empty() {
        return Err(SharedSubplanMaterializationRejectedReason::UnsupportedScalarAggregateShape);
    }

    let mut widened_source_predicates = BTreeMap::new();
    let mut required_source_columns = BTreeMap::<String, BTreeSet<String>>::new();
    let mut widened_source_columns = BTreeMap::<String, BTreeMap<String, SourceColumnMaterialization>>::new();
    for fingerprint in &widened_source_fingerprints {
        let mut predicates = Vec::new();
        let mut seen = BTreeSet::new();
        let mut required_columns = BTreeSet::new();
        for signature in family {
            if let Some(columns) = signature.source_columns.get(fingerprint) {
                merge_source_columns(widened_source_columns.entry(fingerprint.clone()).or_default(), columns)?;
            }
            let source = signature
                .source_predicates
                .iter()
                .find(|source| &source.abstract_fingerprint == fingerprint)
                .ok_or(SharedSubplanMaterializationRejectedReason::PredicateNotPostAggregateFilterable)?;
            let predicate = source
                .predicate
                .clone()
                .ok_or(SharedSubplanMaterializationRejectedReason::PredicateNotPostAggregateFilterable)?;
            collect_expr_column_names(&predicate, &mut required_columns);
            if seen.insert(predicate.to_string()) {
                predicates.push(predicate);
            }
        }
        required_source_columns.insert(fingerprint.clone(), required_columns);
        widened_source_predicates.insert(
            fingerprint.clone(),
            disjunction(predicates)
                .ok_or(SharedSubplanMaterializationRejectedReason::PredicateNotPostAggregateFilterable)?,
        );
    }

    let mut required_aggregate_input_columns = BTreeSet::new();
    for columns in required_source_columns.values() {
        for column in columns {
            if !subtree_can_propagate_materialized_column(
                plan,
                representative_aggregate.input,
                column,
                &widened_source_columns,
            ) {
                return Err(SharedSubplanMaterializationRejectedReason::PredicateColumnPropagationUnsupported);
            }
            required_aggregate_input_columns.insert(column.clone());
        }
    }
    for signature in family {
        let node = plan
            .node(signature.aggregate_node_id)
            .ok_or(SharedSubplanMaterializationRejectedReason::UnsupportedScalarAggregateShape)?;
        let PlanNode::Aggregate(aggregate) = &node.kind else {
            return Err(SharedSubplanMaterializationRejectedReason::UnsupportedScalarAggregateShape);
        };
        for aggregate_expr in &aggregate.spec.aggr_exprs {
            if scalar_count_star_aggregate_expr(plan, aggregate.input, &aggregate_expr.expr) {
                continue;
            }
            collect_expr_column_names(&aggregate_expr.expr, &mut required_aggregate_input_columns);
        }
        for source in &signature.source_predicates {
            if let Some(columns) = signature.source_columns.get(&source.abstract_fingerprint) {
                merge_source_columns(
                    widened_source_columns.entry(source.abstract_fingerprint.clone()).or_default(),
                    columns,
                )?;
            }
        }
    }
    if !source_column_names_are_unique_across_sources(&widened_source_columns) {
        return Err(SharedSubplanMaterializationRejectedReason::PredicateColumnPropagationUnsupported);
    }
    for column in &required_aggregate_input_columns {
        if !subtree_can_propagate_materialized_column(
            plan,
            representative_aggregate.input,
            column,
            &widened_source_columns,
        ) {
            return Err(SharedSubplanMaterializationRejectedReason::PredicateColumnPropagationUnsupported);
        }
    }

    let mut branches = Vec::with_capacity(family.len());
    for signature in family {
        let node = plan
            .node(signature.aggregate_node_id)
            .ok_or(SharedSubplanMaterializationRejectedReason::UnsupportedScalarAggregateShape)?;
        let PlanNode::Aggregate(aggregate) = &node.kind else {
            return Err(SharedSubplanMaterializationRejectedReason::UnsupportedScalarAggregateShape);
        };
        if aggregate.spec.aggr_exprs.len() != node.output_schema().fields().len() {
            return Err(SharedSubplanMaterializationRejectedReason::UnsupportedScalarAggregateShape);
        }
        let mut predicates = Vec::new();
        for source in &signature.source_predicates {
            if !widened_source_fingerprints.contains(&source.abstract_fingerprint) {
                continue;
            }
            predicates.push(
                source
                    .predicate
                    .clone()
                    .ok_or(SharedSubplanMaterializationRejectedReason::PredicateNotPostAggregateFilterable)?,
            );
        }
        let filter = conjunction(predicates)
            .ok_or(SharedSubplanMaterializationRejectedReason::PredicateNotPostAggregateFilterable)?;
        let outputs = aggregate
            .spec
            .aggr_exprs
            .iter()
            .zip(node.output_schema().fields())
            .enumerate()
            .map(|(index, (aggregate_expr, field))| {
                let count_star = scalar_count_star_aggregate_expr(plan, aggregate.input, &aggregate_expr.expr);
                let mut aggregate_expr = aggregate_expr.clone();
                aggregate_expr.alias =
                    format!("__algeon_native_shared_scalar_{}_{}", signature.aggregate_node_id.index(), index);
                aggregate_expr.filter = Some(filter.clone());
                if count_star {
                    aggregate_expr.expr = Expr::literal(LiteralValue::Int64(1));
                }
                ScalarFilteredAggregateOutput {
                    original_output_field: field.as_ref().clone(),
                    materialized_alias: aggregate_expr.alias.clone(),
                    aggregate_expr,
                }
            })
            .collect();
        branches.push(ScalarFilteredAggregateBranch { aggregate_node_id: signature.aggregate_node_id, outputs });
    }

    Ok(SharedAggregateFamilyRewrite {
        family_key: representative.family_key.clone(),
        representative_aggregate_node_id: representative.aggregate_node_id,
        widened_source_predicates,
        widened_source_columns,
        kind: SharedAggregateFamilyRewriteKind::ScalarFiltered { required_aggregate_input_columns, branches },
    })
}

/// Rebuilds a plan with each applied family's members replaced.
///
/// `memo` maps the source nodes that were cloned unchanged to their images
/// and becomes the rewrite origin. Replaced aggregates and the nodes of a
/// cloned materialization are kept out of it on purpose: a route or advisory
/// bound to a member aggregate describes that member alone, and `rewrite`
/// discards side state whose node has no image.
pub(super) struct SharedSubplanMaterializationRewriter<'a> {
    original: &'a QueryPlan,
    families_by_key: BTreeMap<String, SharedAggregateFamilyRewrite>,
    /// Member aggregate to the key of the family that replaces it.
    replacements: BTreeMap<NodeId, String>,
    /// The one shared node per family, created on first use.
    materialized_nodes: BTreeMap<String, NodeId>,
    /// Branch node per replaced member, created on first use so a member
    /// consumed twice gets one image.
    branch_nodes: BTreeMap<NodeId, NodeId>,
    pub(super) nodes: Vec<QueryNode>,
    pub(super) memo: BTreeMap<NodeId, NodeId>,
    pub(super) diagnostics: SharedSubplanMaterializationDiagnostics,
    pub(super) changed: bool,
}

impl<'a> SharedSubplanMaterializationRewriter<'a> {
    pub(super) fn new(
        original: &'a QueryPlan,
        families: Vec<SharedAggregateFamilyRewrite>,
        diagnostics: SharedSubplanMaterializationDiagnostics,
    ) -> Self {
        let mut replacements = BTreeMap::new();
        let mut families_by_key = BTreeMap::new();
        for family in families {
            match &family.kind {
                SharedAggregateFamilyRewriteKind::Grouped { branch_filters } => {
                    for aggregate_node_id in branch_filters.keys() {
                        replacements.insert(*aggregate_node_id, family.family_key.clone());
                    }
                },
                SharedAggregateFamilyRewriteKind::ScalarFiltered { branches, .. } => {
                    for branch in branches {
                        replacements.insert(branch.aggregate_node_id, family.family_key.clone());
                    }
                },
            }
            families_by_key.insert(family.family_key.clone(), family);
        }
        Self {
            original,
            families_by_key,
            replacements,
            materialized_nodes: BTreeMap::new(),
            branch_nodes: BTreeMap::new(),
            nodes: Vec::new(),
            memo: BTreeMap::new(),
            diagnostics,
            changed: false,
        }
    }

    pub(super) fn rewrite_subtree(&mut self, old_id: NodeId) -> NodeId {
        if self.replacements.contains_key(&old_id) {
            self.changed = true;
            return self.branch_node(old_id);
        }
        if let Some(new_id) = self.memo.get(&old_id) {
            return *new_id;
        }

        let old_node = self.original.node(old_id).expect("optimizer input node must exist");
        let new_kind = self.clone_kind_with_rewritten_inputs(&old_node.kind);
        let new_id = self.push_node(new_kind, old_node.facts.clone());
        self.memo.insert(old_id, new_id);
        new_id
    }

    fn branch_node(&mut self, old_aggregate_id: NodeId) -> NodeId {
        if let Some(new_id) = self.branch_nodes.get(&old_aggregate_id) {
            return *new_id;
        }
        let family_key = self.replacements.get(&old_aggregate_id).expect("replacement family must exist").clone();
        let materialized_id = self.materialized_node(&family_key);
        let family = self.families_by_key.get(&family_key).expect("replacement family must exist").clone();
        let new_id = match &family.kind {
            SharedAggregateFamilyRewriteKind::Grouped { branch_filters } => {
                let predicate = branch_filters.get(&old_aggregate_id).cloned().flatten();
                if let Some(predicate) = predicate {
                    let facts =
                        self.nodes.get(materialized_id.index()).expect("materialized node must exist").facts.clone();
                    self.push_node(PlanNode::Filter(FilterNode { input: materialized_id, predicate }), facts)
                } else {
                    materialized_id
                }
            },
            SharedAggregateFamilyRewriteKind::ScalarFiltered { branches, .. } => {
                let branch = branches
                    .iter()
                    .find(|branch| branch.aggregate_node_id == old_aggregate_id)
                    .expect("scalar filtered branch must exist");
                let outputs = branch
                    .outputs
                    .iter()
                    .map(|output| {
                        ProjectionOutput::new(
                            Expr::column(output.materialized_alias.clone())
                                .alias(output.original_output_field.name().clone()),
                            output.original_output_field.clone(),
                        )
                    })
                    .collect();
                let projection = ProjectionNode::new(materialized_id, outputs);
                let facts = PlanFacts::new(Arc::clone(projection.output_schema()));
                self.push_node(PlanNode::Projection(projection), facts)
            },
        };
        self.branch_nodes.insert(old_aggregate_id, new_id);
        new_id
    }

    fn materialized_node(&mut self, family_key: &str) -> NodeId {
        if let Some(new_id) = self.materialized_nodes.get(family_key) {
            return *new_id;
        }
        let family = self.families_by_key.get(family_key).expect("materialization family must exist").clone();
        let new_id = self.clone_materialized_subtree(family.representative_aggregate_node_id, &family);
        self.materialized_nodes.insert(family_key.to_string(), new_id);
        new_id
    }

    /// Clones the representative's subtree for `family`. Sources that match a
    /// widened entry get the widened predicate and columns, and a scalar
    /// family's root becomes the combined `FILTER`ed aggregate.
    fn clone_materialized_subtree(&mut self, old_id: NodeId, family: &SharedAggregateFamilyRewrite) -> NodeId {
        let old_node = self.original.node(old_id).expect("optimizer input node must exist");
        if old_id == family.representative_aggregate_node_id
            && let SharedAggregateFamilyRewriteKind::ScalarFiltered {
                required_aggregate_input_columns, branches, ..
            } = &family.kind
        {
            let PlanNode::Aggregate(aggregate) = &old_node.kind else {
                panic!("scalar filtered materialization root must be an aggregate");
            };
            let input = self.clone_materialized_subtree_with_required_columns(
                aggregate.input,
                family,
                required_aggregate_input_columns,
            );
            let spec = AggregateSpec::new(
                Vec::new(),
                branches
                    .iter()
                    .flat_map(|branch| branch.outputs.iter().map(|output| output.aggregate_expr.clone()))
                    .collect(),
            );
            let output_schema = Arc::new(Schema::new(
                branches
                    .iter()
                    .flat_map(|branch| {
                        branch.outputs.iter().map(|output| {
                            Field::new(
                                output.materialized_alias.clone(),
                                output.original_output_field.data_type().clone(),
                                output.original_output_field.is_nullable(),
                            )
                        })
                    })
                    .collect::<Vec<_>>(),
            ));
            return self.push_node(PlanNode::Aggregate(AggregateNode { input, spec }), PlanFacts::new(output_schema));
        }
        let new_facts = old_node.facts.clone();
        let new_kind = match &old_node.kind {
            PlanNode::Source(source) => match source.as_ref() {
                SourceNode::LocalParquet(source) => {
                    let mut source = source.clone();
                    let signature = local_source_materialization_signature(&source);
                    let mut source_facts = None;
                    if let Some(columns) = matching_source_entry(&family.widened_source_columns, &signature)
                        && widen_local_parquet_source(&mut source, columns)
                    {
                        source_facts = Some(PlanFacts::new(Arc::clone(source.output_schema())));
                    }
                    if let Some(predicate) = matching_source_entry(&family.widened_source_predicates, &signature) {
                        source.reader_predicate = Some(predicate.clone());
                        source.refresh_diagnostics();
                    }
                    let new_kind = PlanNode::Source(Box::new(SourceNode::LocalParquet(source)));
                    return self.push_node(new_kind, source_facts.unwrap_or_else(|| old_node.facts.clone()));
                },
                _ => PlanNode::Source(source.clone()),
            },
            _ => self.clone_kind_with_input_mapper(&old_node.kind, |this, input| {
                this.clone_materialized_subtree(input, family)
            }),
        };
        self.push_node(new_kind, new_facts)
    }

    /// Scalar-family clone that threads the columns the `FILTER` clauses need
    /// down to the sources and back up: each projection and join output is
    /// extended with the required columns its input can supply. Operators
    /// the plan check excluded are unreachable here, so they panic.
    fn clone_materialized_subtree_with_required_columns(
        &mut self,
        old_id: NodeId,
        family: &SharedAggregateFamilyRewrite,
        required_output_columns: &BTreeSet<String>,
    ) -> NodeId {
        let old_node = self.original.node(old_id).expect("optimizer input node must exist");
        match &old_node.kind {
            PlanNode::Source(source) => {
                let (new_kind, facts) = match source.as_ref() {
                    SourceNode::LocalParquet(source) => {
                        let mut source = source.clone();
                        let signature = local_source_materialization_signature(&source);
                        let mut facts = old_node.facts.clone();
                        if let Some(columns) = matching_source_entry(&family.widened_source_columns, &signature)
                            && widen_local_parquet_source(&mut source, columns)
                        {
                            facts = PlanFacts::new(Arc::clone(source.output_schema()));
                        }
                        if let Some(predicate) = matching_source_entry(&family.widened_source_predicates, &signature) {
                            source.reader_predicate = Some(predicate.clone());
                            source.refresh_diagnostics();
                        }
                        (PlanNode::Source(Box::new(SourceNode::LocalParquet(source))), facts)
                    },
                    _ => (PlanNode::Source(source.clone()), old_node.facts.clone()),
                };
                self.push_node(new_kind, facts)
            },
            PlanNode::Filter(filter) => {
                if source_reader_predicate_filter_is_elidable(self.original, filter) {
                    return self.clone_materialized_subtree_with_required_columns(
                        filter.input,
                        family,
                        required_output_columns,
                    );
                }
                let mut input_required = required_output_columns.clone();
                collect_expr_column_names(&filter.predicate, &mut input_required);
                let input =
                    self.clone_materialized_subtree_with_required_columns(filter.input, family, &input_required);
                let mut filter = filter.clone();
                filter.input = input;
                let input_schema = Arc::clone(self.nodes[input.index()].output_schema());
                self.push_node(PlanNode::Filter(filter), PlanFacts::new(input_schema))
            },
            PlanNode::Projection(projection) => {
                let mut outputs = projection.outputs.clone();
                let mut output_names =
                    outputs.iter().map(|output| output.field.name().clone()).collect::<BTreeSet<_>>();
                for column in required_output_columns {
                    if output_names.insert(column.clone()) {
                        let input_field = field_with_name(
                            self.original.node(projection.input).expect("projection input must exist").output_schema(),
                            column,
                        )
                        .cloned()
                        .or_else(|| {
                            propagatable_materialized_field(
                                self.original,
                                projection.input,
                                column,
                                &family.widened_source_columns,
                            )
                        })
                        .expect("required projection column must exist on input");
                        outputs.push(ProjectionOutput::new(Expr::column(column.clone()), input_field));
                    }
                }
                let mut input_required = BTreeSet::new();
                for output in &outputs {
                    collect_expr_column_names(&output.expr, &mut input_required);
                }
                let input =
                    self.clone_materialized_subtree_with_required_columns(projection.input, family, &input_required);
                let projection = ProjectionNode::new(input, outputs);
                let facts = PlanFacts::new(Arc::clone(projection.output_schema()));
                self.push_node(PlanNode::Projection(projection), facts)
            },
            PlanNode::Join(join) => {
                let mut output_projection = join.spec.output_projection.clone();
                let mut left_required = BTreeSet::new();
                let mut right_required = BTreeSet::new();
                for key in &join.spec.keys {
                    left_required.insert(key.left.name().to_string());
                    right_required.insert(key.right.name().to_string());
                }
                if let Some(predicate) = &join.spec.residual_predicate
                    && collect_join_expr_columns(predicate, &mut left_required, &mut right_required).is_err()
                {
                    // An unattributable residual column must not be pruned from
                    // either input; plan validation rejects the residual later.
                    for (input, required) in
                        [(join.inputs[0], &mut left_required), (join.inputs[1], &mut right_required)]
                    {
                        let schema = self.original.node(input).expect("join input must exist").output_schema();
                        required.extend(schema.fields().iter().map(|field| field.name().clone()));
                    }
                }
                for output in &output_projection {
                    match output.side {
                        JoinSide::Left => {
                            left_required.insert(output.column.name().to_string());
                        },
                        JoinSide::Right => {
                            right_required.insert(output.column.name().to_string());
                        },
                    }
                }

                for column in required_output_columns {
                    if let Some(output) = output_projection.iter().find(|output| output.output_name == *column) {
                        match output.side {
                            JoinSide::Left => {
                                left_required.insert(output.column.name().to_string());
                            },
                            JoinSide::Right => {
                                right_required.insert(output.column.name().to_string());
                            },
                        }
                        continue;
                    }
                    let side = unique_join_input_that_can_propagate_column(
                        self.original,
                        join.inputs[0],
                        join.inputs[1],
                        column,
                    )
                    .or_else(|| {
                        unique_join_input_that_can_propagate_materialized_column(
                            self.original,
                            join.inputs[0],
                            join.inputs[1],
                            column,
                            &family.widened_source_columns,
                        )
                    })
                    .expect("required join column must have one input owner");
                    match side {
                        JoinSide::Left => {
                            left_required.insert(column.clone());
                            output_projection.push(JoinOutputColumn::left(
                                ColumnRef::input(0, column.clone(), column.clone()),
                                column.clone(),
                            ));
                        },
                        JoinSide::Right => {
                            right_required.insert(column.clone());
                            output_projection.push(JoinOutputColumn::right(
                                ColumnRef::input(1, column.clone(), column.clone()),
                                column.clone(),
                            ));
                        },
                    }
                }

                let left =
                    self.clone_materialized_subtree_with_required_columns(join.inputs[0], family, &left_required);
                let right =
                    self.clone_materialized_subtree_with_required_columns(join.inputs[1], family, &right_required);
                let mut join = join.clone();
                join.inputs = [left, right];
                join.spec.output_projection = output_projection;
                let output_schema = join_output_schema(
                    self.nodes[left.index()].output_schema(),
                    self.nodes[right.index()].output_schema(),
                    &join.spec.output_projection,
                );
                self.push_node(PlanNode::Join(join), PlanFacts::new(output_schema))
            },
            PlanNode::Sort(sort) => {
                let mut input_required = required_output_columns.clone();
                for sort_expr in &sort.spec.expressions {
                    collect_expr_column_names(&sort_expr.expr, &mut input_required);
                }
                let input = self.clone_materialized_subtree_with_required_columns(sort.input, family, &input_required);
                let mut sort = sort.clone();
                sort.input = input;
                let mut facts = old_node.facts.clone();
                facts.output_schema = Arc::clone(self.nodes[input.index()].output_schema());
                self.push_node(PlanNode::Sort(sort), facts)
            },
            PlanNode::Limit(limit) => {
                let input =
                    self.clone_materialized_subtree_with_required_columns(limit.input, family, required_output_columns);
                let mut limit = limit.clone();
                limit.input = input;
                let mut facts = old_node.facts.clone();
                facts.output_schema = Arc::clone(self.nodes[input.index()].output_schema());
                self.push_node(PlanNode::Limit(limit), facts)
            },
            PlanNode::Aggregate(_)
            | PlanNode::Window(_)
            | PlanNode::Union(_)
            | PlanNode::EdgeNormalize(_)
            | PlanNode::GraphAlgorithm(_)
            | PlanNode::VectorAlgorithm(_)
            | PlanNode::MlAlgorithm(_) => {
                panic!("unsupported scalar filtered materialization subtree")
            },
        }
    }

    fn clone_kind_with_rewritten_inputs(&mut self, kind: &PlanNode) -> PlanNode {
        self.clone_kind_with_input_mapper(kind, |this, input| this.rewrite_subtree(input))
    }

    fn clone_kind_with_input_mapper(
        &mut self,
        kind: &PlanNode,
        mut rewrite_input: impl FnMut(&mut Self, NodeId) -> NodeId,
    ) -> PlanNode {
        optimizer_rewrite::clone_plan_node_with_inputs(kind, |input| rewrite_input(self, input))
    }

    fn push_node(&mut self, kind: PlanNode, facts: PlanFacts) -> NodeId {
        let id = NodeId(self.nodes.len());
        self.nodes.push(QueryNode { id, kind, facts });
        id
    }
}