// Copyright (c) 2023-2026 ParadeDB, Inc.
//
// This file is part of ParadeDB - Postgres for Search and Analytics
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.

//! DataFusion plan builder for aggregate-on-join queries.
//!
//! Builds a DataFusion logical plan from a [`RelNode`] join tree and a
//! [`JoinAggregateTargetList`], producing: scan(s) → join → aggregate.
//!
//! Key difference from JoinScan's plan builder: no CTID columns, no late
//! materialization, no SegmentedTopK — aggregates run entirely on fast fields
//! and the result is aggregate rows, not individual tuples.

use super::join_targetlist::AggOrderByEntry;
use crate::index::fast_fields_helper::WhichFastField;
use crate::index::reader::index::SearchIndexManifest;
use crate::postgres::customscan::aggregatescan::join_targetlist::{
    AggKind, JoinAggregateEntry, JoinAggregateTargetList,
};
use crate::postgres::customscan::aggregatescan::privdat::{CompareOp, DataFusionTopK, FilterExpr};
use crate::postgres::customscan::datafusion::numeric_agg::{
    numeric64_avg_udaf, numeric64_sum_udaf, numeric_bytes_avg_udaf, numeric_bytes_sum_udaf,
};
use crate::postgres::customscan::datafusion::translator::{
    apply_join_level_filter, build_join_df, make_col, make_source_col, ColumnMapper,
    PredicateTranslator,
};
use crate::postgres::customscan::joinscan::build::{
    JoinLevelSearchPredicate, JoinSource, RelNode, RelationAlias,
};
use crate::postgres::customscan::joinscan::scan_state::{
    create_datafusion_session_context, register_source_table, SessionContextProfile,
};
use crate::scan::PgSearchTableProvider;
use crate::schema::SearchFieldType;
use datafusion::common::{DataFusionError, Result};
use datafusion::functions_aggregate::array_agg::array_agg_udaf;
use datafusion::functions_aggregate::count::count_udaf;
use datafusion::functions_aggregate::expr_fn::{
    array_agg, avg, bool_and, bool_or, count, max, min, stddev, stddev_pop, sum, var_pop,
    var_sample,
};
use datafusion::functions_aggregate::string_agg::string_agg_udaf;
use datafusion::logical_expr::expr::{AggregateFunction, Sort};
use datafusion::logical_expr::{lit, Expr};
use datafusion::prelude::{DataFrame, SessionContext};
use futures::future::{FutureExt, LocalBoxFuture};
use pgrx::pg_sys;

/// Creates a DataFusion [`SessionContext`] for aggregate-on-join workloads.
///
/// Thin wrapper around the shared
/// [`create_datafusion_session_context`] with the
/// [`SessionContextProfile::Aggregate`] profile. Kept as a named function so
/// the call sites in `aggregatescan/mod.rs` remain stable; if more aggregate-
/// specific session setup ever appears, this is the place to put it.
pub fn create_aggregate_session_context() -> SessionContext {
    create_datafusion_session_context(SessionContextProfile::Aggregate)
}

/// Build the complete DataFusion logical plan for an aggregate-on-join query:
/// scan(s) → join → aggregate [→ sort → limit].
#[allow(clippy::too_many_arguments)]
pub async fn build_join_aggregate_plan(
    plan: &RelNode,
    targetlist: &JoinAggregateTargetList,
    topk: Option<&DataFusionTopK>,
    join_level_predicates: &[JoinLevelSearchPredicate],
    custom_exprs: *mut pg_sys::List,
    custom_scan_tlist: *mut pg_sys::List,
    having_filter: Option<&FilterExpr>,
    ctx: &SessionContext,
    expr_context: Option<*mut pg_sys::ExprContext>,
    planstate: Option<*mut pg_sys::PlanState>,
    mpp_manifests: Option<&[SearchIndexManifest]>,
) -> Result<(datafusion::logical_expr::LogicalPlan, Vec<usize>)> {
    // Step 1: Build the join DataFrame from the RelNode tree
    let df = build_relnode_df(
        ctx,
        plan,
        join_level_predicates,
        custom_exprs,
        custom_scan_tlist,
        expr_context,
        planstate,
        mpp_manifests,
    )
    .await?;

    // Step 2: Build GROUP BY expressions
    // DataFusion deduplicates grouping expressions that resolve to the same
    // column name (e.g. metadata.brand). We must track which DataFusion output
    // column index corresponds to each of our original targetlist.group_columns.
    let mut group_exprs = Vec::new();
    let mut field_to_df_idx = crate::api::HashMap::default();
    let mut group_df_indices = Vec::with_capacity(targetlist.group_columns.len());

    for gc in &targetlist.group_columns {
        // Dedup key by (plan_position, field_name): plan_position is the
        // unique source identity; field_name distinguishes columns within
        // a source. Keying by rti would collapse rti-aliased sources from
        // sub-PlannerInfos into one DataFusion column.
        let entry = field_to_df_idx.entry((gc.plan_position, gc.field_name.clone()));
        let df_idx = match entry {
            std::collections::hash_map::Entry::Vacant(v) => {
                let df_idx = group_exprs.len();
                v.insert(df_idx);
                group_exprs.push(make_plan_position_col(
                    plan,
                    gc.plan_position,
                    &gc.field_name,
                ));
                df_idx
            }
            std::collections::hash_map::Entry::Occupied(o) => *o.get(),
        };
        group_df_indices.push(df_idx);
    }

    // Step 3: Build aggregate expressions
    let agg_exprs: Vec<Expr> = targetlist
        .aggregates
        .iter()
        .enumerate()
        .map(|(i, agg)| {
            let agg_expr = match agg.agg_kind {
                AggKind::CountStar => Ok(count(lit(1))),
                AggKind::Count => agg_field_col(agg, plan).map(count),
                AggKind::CountDistinct => {
                    let col_exprs = agg_field_cols(agg, plan)?;
                    Ok(Expr::AggregateFunction(AggregateFunction::new_udf(
                        count_udaf(),
                        col_exprs,
                        true,   // distinct
                        None,   // filter
                        vec![], // order_by
                        None,   // null_treatment
                    )))
                }
                // NUMERIC SUM/AVG route to scaled-Int64 or decimal-bytes UDAFs.
                // The Numeric64 UDAFs take the scale as a plan literal so it
                // survives plan serialization for parallel and MPP execution;
                // decimal-bytes values are self-describing.
                AggKind::Sum => agg_field_col(agg, plan).map(|col| match agg.numeric {
                    None => sum(col),
                    Some(SearchFieldType::Numeric64(_, scale)) => {
                        numeric64_sum_udaf().call(vec![col, lit(scale as i32)])
                    }
                    Some(_) => numeric_bytes_sum_udaf().call(vec![col]),
                }),
                AggKind::Avg => agg_field_col(agg, plan).map(|col| match agg.numeric {
                    None => avg(col),
                    Some(SearchFieldType::Numeric64(_, scale)) => {
                        numeric64_avg_udaf().call(vec![col, lit(scale as i32)])
                    }
                    Some(_) => numeric_bytes_avg_udaf().call(vec![col]),
                }),
                AggKind::Min => agg_field_col(agg, plan).map(min),
                AggKind::Max => agg_field_col(agg, plan).map(max),
                AggKind::StddevSamp => agg_field_col(agg, plan).map(stddev),
                AggKind::StddevPop => agg_field_col(agg, plan).map(stddev_pop),
                AggKind::VarSamp => agg_field_col(agg, plan).map(var_sample),
                AggKind::VarPop => agg_field_col(agg, plan).map(var_pop),
                AggKind::BoolAnd => agg_field_col(agg, plan).map(bool_and),
                AggKind::BoolOr => agg_field_col(agg, plan).map(bool_or),
                AggKind::ArrayAgg => {
                    let col_expr = agg_field_col(agg, plan)?;
                    if agg.order_by.is_empty() {
                        Ok(array_agg(col_expr))
                    } else {
                        Ok(Expr::AggregateFunction(AggregateFunction::new_udf(
                            array_agg_udaf(),
                            vec![col_expr],
                            false,
                            None,
                            agg_order_by_exprs(&agg.order_by, plan),
                            None,
                        )))
                    }
                }
                AggKind::StringAgg(ref sep) => {
                    let col_expr = agg_field_col(agg, plan)?;
                    let sep_lit = lit(sep.clone());
                    if agg.order_by.is_empty() {
                        Ok(datafusion::functions_aggregate::string_agg::string_agg(
                            col_expr, sep_lit,
                        ))
                    } else {
                        Ok(Expr::AggregateFunction(AggregateFunction::new_udf(
                            string_agg_udaf(),
                            vec![col_expr, sep_lit],
                            false,
                            None,
                            agg_order_by_exprs(&agg.order_by, plan),
                            None,
                        )))
                    }
                }
            }?;
            // Apply DISTINCT flag for non-CountDistinct aggregates.
            // CountDistinct already sets distinct=true via new_udf above.
            let agg_expr = if agg.distinct
                && !matches!(agg.agg_kind, AggKind::CountDistinct | AggKind::CountStar)
            {
                with_distinct(agg_expr)
            } else {
                agg_expr
            };
            // Apply per-aggregate FILTER clause if present.
            let agg_expr = if let Some(ref filter_expr) = agg.filter {
                let filter_ctx = FilterExprExecContext {
                    targetlist: None,
                    plan: Some(plan),
                };
                let df_filter = filter_expr.to_datafusion(&filter_ctx).ok_or_else(|| {
                    DataFusionError::Internal(
                        "Failed to translate aggregate FILTER clause to DataFusion".to_string(),
                    )
                })?;
                with_filter(agg_expr, df_filter)
            } else {
                agg_expr
            };
            // Alias for stable reference
            Ok(agg_expr.alias(format!("agg_{}", i)))
        })
        .collect::<Result<Vec<Expr>>>()?;

    // Step 4: Apply aggregate
    let mut df = df.aggregate(group_exprs, agg_exprs)?;

    // Step 4.5: Apply HAVING filter (post-aggregate)
    if let Some(having) = having_filter {
        let having_ctx = FilterExprExecContext {
            targetlist: Some(targetlist),
            plan: None,
        };
        let expr = having.to_datafusion(&having_ctx).ok_or_else(|| {
            DataFusionError::Internal(
                "Failed to translate HAVING clause to DataFusion expression".to_string(),
            )
        })?;
        df = df.filter(expr)?;
    }

    // Step 5: If TopK is requested, add sort + limit so DataFusion handles
    // it internally. DataFusion's built-in TopKAggregation optimizer rule
    // can then push the limit into AggregateExec for group-key and MIN/MAX
    // ordering. For COUNT/SUM/AVG ordering, SortExec(fetch=K) uses a
    // bounded TopK heap.
    if let Some(topk) = topk {
        let sort_col_name = topk.sort_target.resolve_sort_col_name(targetlist, plan);
        let sort_expr = datafusion::prelude::col(&sort_col_name)
            .sort(topk.direction.is_asc(), topk.direction.is_nulls_first());
        let df = df.sort(vec![sort_expr])?;
        let df = df.limit(0, Some(topk.k))?;
        return Ok((df.into_optimized_plan()?, group_df_indices));
    }

    Ok((df.into_optimized_plan()?, group_df_indices))
}

/// Recursively lower a [`RelNode`] tree into a DataFusion [`DataFrame`].
///
/// Unlike JoinScan's `build_relnode_df`, this version:
/// - Does NOT include CTID columns (no heap fetch needed for aggregates)
/// - Does NOT handle LIMIT, ORDER BY, DISTINCT, or output projection
///   (those are handled by the aggregate layer above)
/// - Is single-threaded (no partitioning logic)
#[allow(clippy::too_many_arguments)]
fn build_relnode_df<'a>(
    ctx: &'a SessionContext,
    node: &'a RelNode,
    join_level_predicates: &'a [JoinLevelSearchPredicate],
    custom_exprs: *mut pg_sys::List,
    custom_scan_tlist: *mut pg_sys::List,
    expr_context: Option<*mut pg_sys::ExprContext>,
    planstate: Option<*mut pg_sys::PlanState>,
    mpp_manifests: Option<&'a [SearchIndexManifest]>,
) -> LocalBoxFuture<'a, Result<DataFrame>> {
    async move {
        match node {
            RelNode::Scan(source) => {
                let plan_position = source.plan_position;
                let df = build_source_df(
                    ctx,
                    source,
                    plan_position,
                    expr_context,
                    planstate,
                    mpp_manifests,
                )
                .await?;
                let alias =
                    RelationAlias::new(source.scan_info.alias.as_deref()).execution(plan_position);
                Ok(df.alias(&alias)?)
            }
            RelNode::Join(join) => {
                let left_df = build_relnode_df(
                    ctx,
                    &join.left,
                    join_level_predicates,
                    custom_exprs,
                    custom_scan_tlist,
                    expr_context,
                    planstate,
                    mpp_manifests,
                )
                .await?;
                let right_df = build_relnode_df(
                    ctx,
                    &join.right,
                    join_level_predicates,
                    custom_exprs,
                    custom_scan_tlist,
                    expr_context,
                    planstate,
                    mpp_manifests,
                )
                .await?;

                build_join_df(left_df, right_df, join)
            }
            RelNode::Filter(filter) => {
                let df = build_relnode_df(
                    ctx,
                    &filter.input,
                    join_level_predicates,
                    custom_exprs,
                    custom_scan_tlist,
                    expr_context,
                    planstate,
                    mpp_manifests,
                )
                .await?;

                let has_predicates = !join_level_predicates.is_empty() || !custom_exprs.is_null();

                if !has_predicates {
                    // No predicates to apply — pass through
                    return Ok(df);
                }

                let sources = filter.input.output_sources();

                // Translate custom_exprs (non-@@@ cross-table predicates) using
                // PredicateTranslator, mirroring JoinScan's scan_state.rs:562-576.
                // After setrefs, Vars in custom_exprs are INDEX_VAR references
                // that index into custom_scan_tlist. We need a mapper to resolve
                // them back to the correct DataFusion column names.
                let mut translated_exprs = Vec::new();
                if !custom_exprs.is_null() {
                    let mapper = AggregateIndexVarMapper {
                        sources: &sources,
                        custom_scan_tlist,
                    };
                    let translator =
                        PredicateTranslator::new(&sources).with_mapper(Box::new(mapper));
                    unsafe {
                        let expr_list = pgrx::PgList::<pg_sys::Node>::from_pg(custom_exprs);
                        for (i, expr_node) in expr_list.iter_ptr().enumerate() {
                            let expr = translator.translate(expr_node).ok_or_else(|| {
                                DataFusionError::Internal(format!(
                                    "Failed to translate aggregate custom expression at index {}",
                                    i
                                ))
                            })?;
                            translated_exprs.push(expr);
                        }
                    }
                }

                apply_join_level_filter(
                    df,
                    &filter.predicate,
                    &translated_exprs,
                    &sources,
                    /* handle_mark = */ false,
                )
            }
        }
    }
    .boxed_local()
}

/// Maps INDEX_VAR references (from setrefs-transformed custom_exprs) back to
/// DataFusion column names. In the aggregate scan, custom_scan_tlist mirrors
/// the plan's targetlist (plus any Vars we added for predicates), and INDEX_VAR
/// varattno indexes into it. We resolve each Var by looking up the original
/// (rti, attno) from custom_scan_tlist and finding the corresponding source.
struct AggregateIndexVarMapper<'a> {
    sources: &'a [&'a JoinSource],
    custom_scan_tlist: *mut pg_sys::List,
}

impl<'a> ColumnMapper for AggregateIndexVarMapper<'a> {
    fn map_var(&self, varno: pg_sys::Index, varattno: pg_sys::AttrNumber) -> Option<Expr> {
        let (rti, attno) = if varno == pg_sys::INDEX_VAR as pg_sys::Index {
            // INDEX_VAR: look up the original Var from custom_scan_tlist.
            // varattno is 1-indexed into the target list.
            unsafe {
                let tlist = pgrx::PgList::<pg_sys::TargetEntry>::from_pg(self.custom_scan_tlist);
                let idx = (varattno - 1) as usize;
                let te = tlist.get_ptr(idx)?;
                if (*(*te).expr).type_ != pg_sys::NodeTag::T_Var {
                    return None;
                }
                let var = (*te).expr as *mut pg_sys::Var;
                ((*var).varno as pg_sys::Index, (*var).varattno)
            }
        } else {
            (varno, varattno)
        };

        let source = self.sources.iter().find(|s| s.contains_rti(rti))?;
        let field_name = source.column_name(attno)?;
        Some(make_source_col(source, &field_name))
    }
}

/// Context for the **exec phase** — translating a [`FilterExpr`] IR into a
/// DataFusion [`Expr`].
///
/// HAVING provides `targetlist` for resolving `AggRef`/`GroupRef`;
/// FILTER provides `plan` (a `RelNode` tree) for resolving `ColumnRef`.
///
/// This is distinct from the build-phase context in `datafusion_build.rs`,
/// which carries raw planner `JoinAggSource`s instead of a `RelNode` tree.
struct FilterExprExecContext<'a> {
    targetlist: Option<&'a JoinAggregateTargetList>,
    plan: Option<&'a RelNode>,
}

impl FilterExpr {
    /// Translate this expression to a DataFusion `Expr`.
    ///
    /// Used for both HAVING (pass `targetlist`) and per-aggregate FILTER (pass `plan`).
    fn to_datafusion(&self, ctx: &FilterExprExecContext<'_>) -> Option<Expr> {
        use datafusion::logical_expr::Operator;

        match self {
            FilterExpr::AggRef(idx) => {
                let tl = ctx.targetlist?;
                if *idx < tl.aggregates.len() {
                    Some(datafusion::prelude::col(format!("agg_{}", idx)))
                } else {
                    None
                }
            }
            FilterExpr::GroupRef(field_name) => Some(datafusion::prelude::col(field_name.as_str())),
            FilterExpr::ColumnRef {
                plan_position,
                field_name,
                ..
            } => {
                let plan = ctx.plan?;
                Some(make_plan_position_col(plan, *plan_position, field_name))
            }
            FilterExpr::LitInt(v) => Some(lit(*v)),
            FilterExpr::LitFloat(v) => Some(lit(*v)),
            FilterExpr::LitBool(v) => Some(lit(*v)),
            FilterExpr::LitString(v) => Some(lit(v.clone())),
            FilterExpr::BinOp { left, op, right } => {
                let l = left.to_datafusion(ctx)?;
                let r = right.to_datafusion(ctx)?;
                let df_op = match op {
                    CompareOp::Eq => Operator::Eq,
                    CompareOp::NotEq => Operator::NotEq,
                    CompareOp::Lt => Operator::Lt,
                    CompareOp::LtEq => Operator::LtEq,
                    CompareOp::Gt => Operator::Gt,
                    CompareOp::GtEq => Operator::GtEq,
                };
                Some(Expr::BinaryExpr(datafusion::logical_expr::BinaryExpr::new(
                    Box::new(l),
                    df_op,
                    Box::new(r),
                )))
            }
            FilterExpr::And(children) => {
                let exprs: Vec<Expr> = children
                    .iter()
                    .map(|c| c.to_datafusion(ctx))
                    .collect::<Option<Vec<Expr>>>()?;
                let mut result = exprs.into_iter();
                let first = result.next()?;
                Some(result.fold(first, |acc, e| acc.and(e)))
            }
            FilterExpr::Or(children) => {
                let exprs: Vec<Expr> = children
                    .iter()
                    .map(|c| c.to_datafusion(ctx))
                    .collect::<Option<Vec<Expr>>>()?;
                let mut result = exprs.into_iter();
                let first = result.next()?;
                Some(result.fold(first, |acc, e| acc.or(e)))
            }
            FilterExpr::Not(inner) => {
                let e = inner.to_datafusion(ctx)?;
                Some(Expr::Not(Box::new(e)))
            }
            FilterExpr::IsNull(inner) => {
                let e = inner.to_datafusion(ctx)?;
                Some(e.is_null())
            }
            FilterExpr::IsNotNull(inner) => {
                let e = inner.to_datafusion(ctx)?;
                Some(e.is_not_null())
            }
        }
    }
}

/// Build a DataFusion [`DataFrame`] for a single scan source.
///
/// Unlike JoinScan's `build_source_df`, this version:
/// - Does NOT include CTID or Score columns
/// - Is always single-threaded (no partitioning)
/// - Does NOT set up late materialization
async fn build_source_df(
    ctx: &SessionContext,
    source: &JoinSource,
    plan_position: usize,
    expr_context: Option<*mut pg_sys::ExprContext>,
    planstate: Option<*mut pg_sys::PlanState>,
    mpp_manifests: Option<&[SearchIndexManifest]>,
) -> Result<DataFrame> {
    let scan_info = source.scan_info.clone();

    // Each source that solves runtime PostgreSQL expressions needs its own
    // per-tuple memory context. SearchQueryInput::solve_postgres_expressions()
    // resets that context before replacing Param/PostgresExpression nodes. If
    // two providers share one context, solving the second source invalidates
    // the Const/expression nodes retained by the first source. Generic plans
    // that parameterize predicates on both join inputs then dereference stale
    // nodes during reader construction and abort the backend.
    let source_query = scan_info.mode.query().clone();
    let needs_runtime_context =
        source_query.has_postgres_expressions() || source_query.has_parameters();
    // `.or(expr_context)` can hand every source the same context, but only the
    // EXPLAIN-only rebuild reaches it. There it is safe: `PgSearchTableProvider::scan()` returns
    // "postgres expressions have not been solved: missing planstate" before it would call
    // `solve_postgres_expressions`, so the shared context is never reset.
    let source_expr_context = if needs_runtime_context {
        unsafe {
            planstate
                .and_then(|planstate| {
                    if planstate.is_null() || (*planstate).state.is_null() {
                        None
                    } else {
                        Some(pg_sys::CreateExprContext((*planstate).state))
                    }
                })
                .or(expr_context)
        }
    } else {
        expr_context
    };

    let alias = RelationAlias::new(scan_info.alias.as_deref()).execution(plan_position);

    // Use all fast fields from the source (the provider exposes them to DataFusion).
    // Include Named fields plus Ctid as a sentinel — DataFusion needs at least one
    // column to produce RecordBatches with row counts (important for COUNT(*)).
    let mut fields: Vec<WhichFastField> = source
        .scan_info
        .fields
        .iter()
        .filter_map(|f| match &f.field {
            WhichFastField::Named(..) | WhichFastField::Deferred(..) => Some(f.field.clone()),
            WhichFastField::Ctid
            | WhichFastField::Score
            | WhichFastField::Junk(_)
            | WhichFastField::TableOid
            | WhichFastField::DeferredCtid(_)
            | WhichFastField::MatchTag(_) => None,
        })
        .collect();

    // Always include Ctid so the provider schema is never empty
    if fields.is_empty() || !fields.iter().any(|f| matches!(f, WhichFastField::Ctid)) {
        fields.push(WhichFastField::Ctid);
    }

    // MPP-aware provider setup. Every source gets its segments sliced across PG
    // parallel workers via `parallel_state.checkout_segment_for_source(plan_position)`
    // when this is an MPP plan.
    let source_idx = mpp_manifests.map(|_| plan_position);
    let mut provider = PgSearchTableProvider::new(scan_info, fields, source_idx);
    // The leader claims segments out of the DSM pool the same manifests populate, so its own
    // reader is built from the source's manifest. This plan never crosses the codec that
    // injects the manifest for JoinScan, so do it here.
    if let Some(manifests) = mpp_manifests {
        let manifest = manifests.get(plan_position).unwrap_or_else(|| {
            panic!(
                "missing captured manifest for aggregate source at plan_position {plan_position}"
            )
        });
        provider.set_manifest(manifest.clone());
    }
    if let crate::scan::ScanMode::Tagged { local_queries, .. } = &source.scan_info.mode {
        for tq in local_queries {
            provider.add_match_tag_column(&tq.tag_name);
        }
    }
    // HeapFilter queries (e.g. `=` on a column indexed via a
    // `pdb.literal(...)` cast) compile to runtime Postgres expressions
    // that can only be evaluated with a live ExprContext + PlanState.
    // The provider's `scan()` reaches for them via
    // `init_postgres_expressions` / `solve_postgres_expressions` only
    // when needed, so threading them through here is what makes the
    // agg-on-join path match JoinScan and Base Scan.
    provider.set_expr_context(source_expr_context);
    provider.set_planstate(planstate);
    let df = register_source_table(ctx, alias.as_str(), provider).await?;

    // Select all fields from the provider schema using their qualified names.
    // This mirrors JoinScan's pattern and ensures column names are accessible
    // via make_col(alias, field_name) in join keys and aggregate expressions.
    let exprs: Vec<Expr> = df
        .schema()
        .fields()
        .iter()
        .map(|f| make_col(alias.as_str(), f.name()))
        .collect();

    if exprs.is_empty() {
        // No fields at all — this can happen for COUNT(*) where no columns are
        // referenced from this source. Return the raw DataFrame.
        Ok(df)
    } else {
        df.select(exprs)
    }
}

/// Build a DataFusion column expression for a targetlist ref by its
/// previously-resolved `plan_position`.
fn make_plan_position_col(plan: &RelNode, plan_position: usize, field_name: &str) -> Expr {
    let source = plan
        .source_at_plan_position(plan_position)
        .unwrap_or_else(|| panic!("no source at plan_position {plan_position}"));
    make_source_col(source, field_name)
}

/// Replace an `Expr::AggregateFunction` with the same call but `distinct=true`.
/// Non-aggregate-function expressions are returned unchanged.
fn with_distinct(expr: Expr) -> Expr {
    match expr {
        Expr::AggregateFunction(af) => Expr::AggregateFunction(AggregateFunction::new_udf(
            af.func,
            af.params.args,
            true,
            af.params.filter,
            af.params.order_by,
            af.params.null_treatment,
        )),
        other => other,
    }
}

/// Replace an `Expr::AggregateFunction` with the same call but `filter=Some(...)`.
/// Non-aggregate-function expressions are returned unchanged.
fn with_filter(expr: Expr, filter: Expr) -> Expr {
    match expr {
        Expr::AggregateFunction(af) => Expr::AggregateFunction(AggregateFunction::new_udf(
            af.func,
            af.params.args,
            af.params.distinct,
            Some(Box::new(filter)),
            af.params.order_by,
            af.params.null_treatment,
        )),
        other => other,
    }
}

/// Build a DataFusion column expression for an aggregate's first field reference.
fn agg_field_col(agg: &JoinAggregateEntry, plan: &RelNode) -> Result<Expr> {
    let r = agg.field_refs.first().ok_or_else(|| {
        DataFusionError::Internal("non-COUNT(*) aggregate must have a field reference".to_string())
    })?;
    Ok(make_plan_position_col(plan, r.plan_position, &r.field_name))
}

/// Convert aggregate ORDER BY entries to DataFusion `Sort` expressions.
fn agg_order_by_exprs(order_by: &[AggOrderByEntry], plan: &RelNode) -> Vec<Sort> {
    order_by
        .iter()
        .map(|entry| {
            Sort::new(
                make_plan_position_col(plan, entry.plan_position, &entry.field_name),
                entry.direction.is_asc(),
                entry.direction.is_nulls_first(),
            )
        })
        .collect()
}

/// Build DataFusion column expressions for all of an aggregate's field references.
/// Used for multi-column DISTINCT (e.g. `COUNT(DISTINCT col1, col2)`).
fn agg_field_cols(agg: &JoinAggregateEntry, plan: &RelNode) -> Result<Vec<Expr>> {
    agg.field_refs
        .iter()
        .map(|r| Ok(make_plan_position_col(plan, r.plan_position, &r.field_name)))
        .collect()
}
