proof-of-sql-planner 0.130.2

use super::{PlannerError, PlannerResult};
use crate::expr_to_proof_expr;
use datafusion::{
    logical_expr::{
        expr::{AggregateFunction, AggregateFunctionDefinition},
        Expr,
    },
    physical_plan,
};
use proof_of_sql::{
    base::database::{ColumnType, LiteralValue},
    sql::proof_exprs::DynProofExpr,
};
use sqlparser::ast::Ident;

/// An aggregate function we support
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum AggregateFunc {
    /// Sum
    Sum,
    /// Count
    Count,
}

/// Convert an [`AggregateFunction`] to a [`DynProofExpr`]
///
/// TODO: Some moderate changes are necessary once we upgrade `DataFusion` to 46.0.0
pub(crate) fn aggregate_function_to_proof_expr(
    function: &AggregateFunction,
    schema: &[(Ident, ColumnType)],
) -> PlannerResult<(AggregateFunc, DynProofExpr)> {
    match function {
        AggregateFunction {
            distinct: false,
            filter: None,
            order_by: None,
            args,
            func_def: AggregateFunctionDefinition::BuiltIn(op),
            ..
        } if args.len() == 1 => {
            let arg = &args[0];
            match (op, arg) {
                (physical_plan::aggregates::AggregateFunction::Count, &Expr::Wildcard { .. }) => {
                    // Special case for COUNT(*)
                    let proof_expr = DynProofExpr::new_literal(LiteralValue::BigInt(1));
                    Ok((AggregateFunc::Count, proof_expr))
                }
                (physical_plan::aggregates::AggregateFunction::Sum, _) => {
                    Ok((AggregateFunc::Sum, expr_to_proof_expr(arg, schema)?))
                }
                (physical_plan::aggregates::AggregateFunction::Count, _) => {
                    Ok((AggregateFunc::Count, expr_to_proof_expr(arg, schema)?))
                }
                _ => Err(PlannerError::UnsupportedAggregateOperation { op: op.clone() }),
            }
        }
        _ => Err(PlannerError::UnsupportedAggregateFunction {
            function: function.clone(),
        })?,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::df_util::*;
    use proof_of_sql::base::database::{ColumnRef, ColumnType, TableRef};

    // AggregateFunction to DynProofExpr
    #[test]
    fn we_can_convert_an_aggregate_function_to_proof_expr() {
        let expr = df_column("table", "a");
        let schema: Vec<(Ident, ColumnType)> = vec![("a".into(), ColumnType::BigInt)];
        for (function, operator) in &[
            (
                physical_plan::aggregates::AggregateFunction::Sum,
                AggregateFunc::Sum,
            ),
            (
                physical_plan::aggregates::AggregateFunction::Count,
                AggregateFunc::Count,
            ),
        ] {
            let function = AggregateFunction::new(
                function.clone(),
                vec![expr.clone()],
                false,
                None,
                None,
                None,
            );
            assert_eq!(
                aggregate_function_to_proof_expr(&function, &schema).unwrap(),
                (
                    *operator,
                    DynProofExpr::new_column(ColumnRef::new(
                        TableRef::from_names(None, "table"),
                        "a".into(),
                        ColumnType::BigInt
                    ))
                )
            );
        }
    }

    #[test]
    fn we_can_convert_count_star_to_count_one() {
        use proof_of_sql::base::database::LiteralValue;

        let wildcard_expr = Expr::Wildcard { qualifier: None };
        let schema: Vec<(Ident, ColumnType)> = vec![("a".into(), ColumnType::BigInt)];
        let function = AggregateFunction::new(
            physical_plan::aggregates::AggregateFunction::Count,
            vec![wildcard_expr],
            false,
            None,
            None,
            None,
        );
        assert_eq!(
            aggregate_function_to_proof_expr(&function, &schema).unwrap(),
            (
                AggregateFunc::Count,
                DynProofExpr::new_literal(LiteralValue::BigInt(1))
            )
        );
    }

    #[test]
    fn we_cannot_convert_an_aggregate_function_to_pair_if_unsupported() {
        let expr = df_column("table", "a");
        let schema = vec![("a".into(), ColumnType::BigInt)];
        let function = AggregateFunction::new(
            physical_plan::aggregates::AggregateFunction::RegrIntercept,
            vec![expr.clone()],
            false,
            None,
            None,
            None,
        );
        assert!(matches!(
            aggregate_function_to_proof_expr(&function, &schema),
            Err(PlannerError::UnsupportedAggregateOperation { .. })
        ));
    }

    #[test]
    fn we_cannot_convert_an_aggregate_function_to_pair_if_too_many_or_no_exprs() {
        let expr = df_column("table", "a");
        let schema = vec![("a".into(), ColumnType::BigInt)];
        // Too many exprs
        let function = AggregateFunction::new(
            physical_plan::aggregates::AggregateFunction::Sum,
            vec![expr.clone(); 2],
            false,
            None,
            None,
            None,
        );
        assert!(matches!(
            aggregate_function_to_proof_expr(&function, &schema),
            Err(PlannerError::UnsupportedAggregateFunction { .. })
        ));

        // No exprs
        let function = AggregateFunction::new(
            physical_plan::aggregates::AggregateFunction::Sum,
            Vec::<_>::new(),
            false,
            None,
            None,
            None,
        );
        assert!(matches!(
            aggregate_function_to_proof_expr(&function, &schema),
            Err(PlannerError::UnsupportedAggregateFunction { .. })
        ));
    }

    #[test]
    fn we_cannot_convert_an_aggregate_function_to_pair_if_unsupported_options() {
        // No distinct, filter, or order_by

        // Distinct
        let expr = df_column("table", "a");
        let schema = vec![("a".into(), ColumnType::BigInt)];
        let function = AggregateFunction::new(
            physical_plan::aggregates::AggregateFunction::Count,
            vec![expr.clone()],
            true,
            None,
            None,
            None,
        );
        assert!(matches!(
            aggregate_function_to_proof_expr(&function, &schema),
            Err(PlannerError::UnsupportedAggregateFunction { .. })
        ));

        // Filter
        let function = AggregateFunction::new(
            physical_plan::aggregates::AggregateFunction::Count,
            vec![expr.clone()],
            false,
            Some(Box::new(expr.clone())),
            None,
            None,
        );
        assert!(matches!(
            aggregate_function_to_proof_expr(&function, &schema),
            Err(PlannerError::UnsupportedAggregateFunction { .. })
        ));

        // OrderBy
        let function = AggregateFunction::new(
            physical_plan::aggregates::AggregateFunction::Count,
            vec![expr.clone()],
            false,
            None,
            Some(vec![expr.clone()]),
            None,
        );
        assert!(matches!(
            aggregate_function_to_proof_expr(&function, &schema),
            Err(PlannerError::UnsupportedAggregateFunction { .. })
        ));
    }
}