algeon-datafusion 0.1.2

//! Operator-level admission checks shared by normalization and lowering.
//!
//! The checks decide whether a DataFusion aggregate, sort ordering, or
//! expression tree has a native counterpart, and return a `LoweringRejection`
//! that names the function and defect when it does not. They read plan
//! structure and Arrow types; expression semantics are proven by
//! `expression_contract`.

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

use arrow::datatypes::{DataType, SchemaRef};
use cudf::interop::arrow_host::is_supported_arrow_type;
use datafusion::{
    physical_expr::{
        LexOrdering, PhysicalExpr, ScalarFunctionExpr,
        aggregate::AggregateFunctionExpr,
        expressions::{
            BinaryExpr, CaseExpr, CastExpr, Column, DynamicFilterPhysicalExpr, InListExpr, IsNotNullExpr, IsNullExpr,
            LikeExpr, Literal, NegativeExpr, NotExpr, TryCastExpr,
        },
    },
    physical_plan::{
        ExecutionPlan,
        aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy},
        filter::FilterExec,
        projection::ProjectionExec,
    },
};
use query_planning::{
    capability::is_cudf_median_input_type,
    datatypes::is_timestamp_with_timezone,
    expr::AggregateFunction as NativeAggregateFunction,
    validation::{AggregateArgumentUse, is_validity_count_input_type},
};

use super::{
    LoweringRejection,
    expression_contract::{
        is_list_carry_type, is_narrow_decimal_carry, native_aggregate_function_kind, resolve_column,
        validate_physical_expr_contract,
    },
};
use crate::{
    capability::CapabilityReason,
    native::ExpressionRejection,
    report::{AggregateDefect, PlanBlocker},
};

/// A rejection of `reason` that names the aggregate function `function` and
/// what is unsupported about its use.
pub(super) fn aggregate_blocked(
    reason: CapabilityReason,
    function: &str,
    defect: AggregateDefect,
) -> LoweringRejection {
    LoweringRejection::blocked_by(reason, PlanBlocker::Aggregate { function: function.to_string(), defect })
}

pub(super) fn validate_schema(schema: SchemaRef) -> Result<(), CapabilityReason> {
    for field in schema.fields() {
        if !is_supported_arrow_type(field.data_type()) {
            return Err(CapabilityReason::UnsupportedType);
        }
    }
    Ok(())
}

pub(super) fn validate_aggregate(aggregate: &AggregateExec, allow_partial: bool) -> Result<(), LoweringRejection> {
    if !(matches!(aggregate.mode(), AggregateMode::Single | AggregateMode::SinglePartitioned)
        || allow_partial && matches!(aggregate.mode(), AggregateMode::Partial))
    {
        return Err(CapabilityReason::UnsupportedOperator.into());
    }
    let group_by = aggregate.group_expr();
    if group_by.has_grouping_set() {
        validate_grouping_sets(group_by)?;
    } else if !group_by.expr().is_empty()
        && (group_by.groups().len() != 1 || group_by.groups()[0].iter().any(|is_null| *is_null))
    {
        return Err(CapabilityReason::UnsupportedOperator.into());
    }
    if let Some(median) = median_with_floating_statistical(aggregate) {
        return Err(aggregate_blocked(
            CapabilityReason::UnsupportedOperator,
            median.fun().name(),
            AggregateDefect::MedianWithStatistical,
        ));
    }
    for (index, aggregate_expr) in aggregate.aggr_expr().iter().enumerate() {
        let kind = native_aggregate_function_kind(aggregate_expr.fun());
        let function = aggregate_expr.fun().name();
        if aggregate_expr.is_distinct() && !matches!(kind, Some(NativeAggregateFunction::Count)) {
            return Err(aggregate_blocked(CapabilityReason::UnsupportedOperator, function, AggregateDefect::Distinct));
        }
        if !aggregate_expr.order_bys().is_empty() {
            return Err(aggregate_blocked(CapabilityReason::UnsupportedExpression, function, AggregateDefect::OrderBy));
        }
        if kind.is_none() {
            return Err(aggregate_blocked(
                CapabilityReason::UnsupportedOperator,
                function,
                AggregateDefect::UnsupportedFunction,
            ));
        }
        if matches!(aggregate.mode(), AggregateMode::Partial) && matches!(kind, Some(NativeAggregateFunction::Median)) {
            return Err(aggregate_blocked(
                CapabilityReason::UnsupportedOperator,
                function,
                AggregateDefect::PartialStage,
            ));
        }
        let expressions = aggregate_expr.expressions();
        let [arg] = expressions.as_slice() else {
            return Err(aggregate_blocked(
                CapabilityReason::UnsupportedExpression,
                function,
                AggregateDefect::ArgumentCount(expressions.len()),
            ));
        };
        let validity_only = kind.is_some_and(|function| {
            AggregateArgumentUse {
                function,
                distinct: aggregate_expr.is_distinct(),
                filtered: aggregate.filter_expr().get(index).is_some_and(Option::is_some),
                grouped: !group_by.expr().is_empty(),
                column_argument: arg.downcast_ref::<Column>().is_some(),
            }
            .reads_only_top_level_validity()
        });
        if validity_only {
            validate_validity_count_argument_expr(arg, &aggregate.input().schema())?;
        } else {
            validate_physical_expr_contract(arg, &aggregate.input().schema())?;
            validate_operator_value_expr(arg, &aggregate.input().schema())?;
        }
        if matches!(kind, Some(NativeAggregateFunction::Median)) {
            validate_median_input_type(arg, &aggregate.input().schema())?;
        }
    }
    validate_aggregate_filters(aggregate)?;
    for (group_expr, _alias) in group_by.expr() {
        validate_operator_key_expr(group_expr, &aggregate.input().schema())?;
    }
    Ok(())
}

/// The `MEDIAN` aggregate of `aggregate` when it also holds a floating-point
/// statistical aggregate.
fn median_with_floating_statistical(aggregate: &AggregateExec) -> Option<&Arc<AggregateFunctionExpr>> {
    let median = aggregate
        .aggr_expr()
        .iter()
        .find(|expr| matches!(native_aggregate_function_kind(expr.fun()), Some(NativeAggregateFunction::Median)))?;
    aggregate_has_floating_statistical(aggregate).then_some(median)
}

fn first_floating_statistical(aggregate: &AggregateExec) -> Option<&Arc<AggregateFunctionExpr>> {
    aggregate.aggr_expr().iter().find(|expr| {
        matches!(
            native_aggregate_function_kind(expr.fun()),
            Some(
                NativeAggregateFunction::StddevSamp
                    | NativeAggregateFunction::StddevPop
                    | NativeAggregateFunction::VarSamp
                    | NativeAggregateFunction::VarPop
            )
        )
    })
}

fn aggregate_has_floating_statistical(aggregate: &AggregateExec) -> bool {
    first_floating_statistical(aggregate).is_some()
}

/// The first floating statistical aggregate in `plan` whose input DataFusion
/// reads through more than one partition. DataFusion's own result then
/// depends on the order partitions deliver rows, which no native evaluation
/// order reproduces.
pub(crate) fn floating_statistical_aggregate_input_is_partitioned(plan: &dyn ExecutionPlan) -> Option<PlanBlocker> {
    if let Some(aggregate) = plan.downcast_ref::<AggregateExec>()
        && let Some(statistical) = first_floating_statistical(aggregate)
        && subtree_has_multiple_partitions(aggregate.input().as_ref())
    {
        return Some(PlanBlocker::Aggregate {
            function: statistical.fun().name().to_string(),
            defect: AggregateDefect::PartitionedInput,
        });
    }
    plan.children().into_iter().find_map(|child| floating_statistical_aggregate_input_is_partitioned(child.as_ref()))
}

fn subtree_has_multiple_partitions(plan: &dyn ExecutionPlan) -> bool {
    plan.properties().output_partitioning().partition_count() > 1
        || plan.children().into_iter().any(|child| subtree_has_multiple_partitions(child.as_ref()))
}

pub(super) fn validate_aggregate_for_native_lowering(
    aggregate: &AggregateExec,
    allow_partial: bool,
) -> Result<(), LoweringRejection> {
    if aggregate.group_expr().has_grouping_set() {
        validate_grouping_set_aggregate(aggregate, allow_partial)
    } else {
        validate_aggregate(aggregate, allow_partial)
    }
}

fn validate_grouping_set_aggregate(aggregate: &AggregateExec, allow_partial: bool) -> Result<(), LoweringRejection> {
    if !(matches!(aggregate.mode(), AggregateMode::Single | AggregateMode::SinglePartitioned)
        || allow_partial && matches!(aggregate.mode(), AggregateMode::Partial))
    {
        return Err(CapabilityReason::UnsupportedOperator.into());
    }
    let group_by = aggregate.group_expr();
    if !group_by.has_grouping_set()
        || group_by.expr().is_empty()
        || group_by.null_expr().len() != group_by.expr().len()
        || group_by.groups().is_empty()
        || group_by.groups().iter().any(|group| group.len() != group_by.expr().len())
    {
        return Err(CapabilityReason::UnsupportedOperator.into());
    }
    for aggregate_expr in aggregate.aggr_expr() {
        let function = aggregate_expr.fun().name();
        if aggregate_expr.is_distinct() {
            return Err(aggregate_blocked(
                CapabilityReason::UnsupportedExpression,
                function,
                AggregateDefect::Distinct,
            ));
        }
        if !aggregate_expr.order_bys().is_empty() {
            return Err(aggregate_blocked(CapabilityReason::UnsupportedExpression, function, AggregateDefect::OrderBy));
        }
        let kind = native_aggregate_function_kind(aggregate_expr.fun());
        if kind.is_none() {
            return Err(aggregate_blocked(
                CapabilityReason::UnsupportedOperator,
                function,
                AggregateDefect::UnsupportedFunction,
            ));
        }
        if matches!(kind, Some(NativeAggregateFunction::Median)) {
            return Err(aggregate_blocked(
                CapabilityReason::UnsupportedOperator,
                function,
                AggregateDefect::GroupingSets,
            ));
        }
        let expressions = aggregate_expr.expressions();
        let [arg] = expressions.as_slice() else {
            return Err(aggregate_blocked(
                CapabilityReason::UnsupportedExpression,
                function,
                AggregateDefect::ArgumentCount(expressions.len()),
            ));
        };
        validate_physical_expr_contract(arg, &aggregate.input().schema())?;
        validate_operator_value_expr(arg, &aggregate.input().schema())?;
    }
    validate_aggregate_filters(aggregate)?;
    for (group_expr, _alias) in group_by.expr() {
        validate_operator_key_expr(group_expr, &aggregate.input().schema())?;
    }
    for (null_expr, _alias) in group_by.null_expr() {
        validate_physical_expr_contract(null_expr, &aggregate.input().schema())?;
    }
    Ok(())
}

fn validate_aggregate_filters(aggregate: &AggregateExec) -> Result<(), LoweringRejection> {
    for filter in aggregate.filter_expr().iter().flatten() {
        validate_physical_expr_contract(filter, &aggregate.input().schema())?;
        let data_type = filter
            .data_type(aggregate.input().schema().as_ref())
            .map_err(|_error| CapabilityReason::UnsupportedExpression)?;
        if !matches!(data_type, DataType::Boolean) {
            return Err(ExpressionRejection::unsupported_value_type(filter, data_type).into());
        }
    }
    Ok(())
}

fn validate_grouping_sets(group_by: &PhysicalGroupBy) -> Result<(), CapabilityReason> {
    let group_expr_count = group_by.expr().len();
    if group_expr_count == 0 || group_expr_count > 64 {
        return Err(CapabilityReason::UnsupportedOperator);
    }
    if group_by.null_expr().len() != group_expr_count {
        return Err(CapabilityReason::UnsupportedOperator);
    }
    if group_by.groups().is_empty() || group_by.groups().iter().any(|group| group.len() != group_expr_count) {
        return Err(CapabilityReason::UnsupportedOperator);
    }
    Ok(())
}

pub(super) fn validate_sort_ordering(
    ordering: &LexOrdering,
    input_schema: &SchemaRef,
) -> Result<(), LoweringRejection> {
    if ordering.is_empty() {
        return Err(CapabilityReason::UnsupportedExpression.into());
    }
    for sort_expr in ordering {
        validate_operator_key_expr(&sort_expr.expr, input_schema)?;
    }
    Ok(())
}

fn validate_operator_key_expr(expr: &Arc<dyn PhysicalExpr>, input_schema: &SchemaRef) -> Result<(), LoweringRejection> {
    validate_physical_expr_contract(expr, input_schema)?;
    validate_operator_value_expr(expr, input_schema)
}

fn validate_operator_value_expr(
    expr: &Arc<dyn PhysicalExpr>,
    input_schema: &SchemaRef,
) -> Result<(), LoweringRejection> {
    let data_type = expr.data_type(input_schema.as_ref()).map_err(|_error| CapabilityReason::UnsupportedExpression)?;
    if is_timestamp_with_timezone(&data_type) || is_list_carry_type(&data_type) || is_narrow_decimal_carry(&data_type) {
        Err(ExpressionRejection::unsupported_value_type(expr, data_type).into())
    } else {
        Ok(())
    }
}

/// A validity-only `COUNT` still proves the argument expression against the
/// native expression contract, but reads no values, so list carry types pass.
fn validate_validity_count_argument_expr(
    expr: &Arc<dyn PhysicalExpr>,
    input_schema: &SchemaRef,
) -> Result<(), LoweringRejection> {
    validate_physical_expr_contract(expr, input_schema)?;
    let data_type = expr.data_type(input_schema.as_ref()).map_err(|_error| CapabilityReason::UnsupportedExpression)?;
    if is_validity_count_input_type(&data_type) {
        Ok(())
    } else {
        Err(ExpressionRejection::unsupported_value_type(expr, data_type).into())
    }
}

fn validate_median_input_type(expr: &Arc<dyn PhysicalExpr>, input_schema: &SchemaRef) -> Result<(), LoweringRejection> {
    let data_type = expr.data_type(input_schema.as_ref()).map_err(|_error| CapabilityReason::UnsupportedExpression)?;
    if is_cudf_median_input_type(&data_type) {
        Ok(())
    } else {
        Err(ExpressionRejection::unsupported_value_type(expr, data_type).into())
    }
}

pub(super) fn native_zero_column_count_source_is_safe(
    filter: Option<&FilterExec>,
    pre_aggregate_projection: Option<&ProjectionExec>,
    aggregate: &AggregateExec,
) -> bool {
    filter.is_none() && pre_aggregate_projection.is_none() && aggregate_is_count_star_only(aggregate)
}

pub(super) fn aggregate_is_count_star_only(aggregate: &AggregateExec) -> bool {
    if !aggregate.group_expr().expr().is_empty()
        || !aggregate.filter_expr().iter().all(Option::is_none)
        || aggregate.aggr_expr().len() != 1
    {
        return false;
    }
    let expr = &aggregate.aggr_expr()[0];
    if !matches!(native_aggregate_function_kind(expr.fun()), Some(NativeAggregateFunction::Count))
        || expr.is_distinct()
        || !expr.order_bys().is_empty()
    {
        return false;
    }
    let expressions = expr.expressions();
    let [arg] = expressions.as_slice() else {
        return false;
    };
    if arg.downcast_ref::<Literal>().is_none_or(|literal| literal.value().is_null()) {
        return false;
    }
    let mut required_columns = BTreeSet::new();
    collect_required_columns(arg, &aggregate.input().schema(), &mut required_columns).is_ok()
        && required_columns.is_empty()
}

pub(super) fn collect_required_columns(
    expr: &Arc<dyn PhysicalExpr>,
    input_schema: &SchemaRef,
    required: &mut BTreeSet<String>,
) -> Result<(), CapabilityReason> {
    if let Some(column) = expr.downcast_ref::<Column>() {
        let column = resolve_column(column, input_schema).map_err(|rejection| rejection.capability_reason())?;
        required.insert(column.field().name().to_string());
        return Ok(());
    }
    if expr.downcast_ref::<Literal>().is_some() {
        return Ok(());
    }
    if let Some(binary) = expr.downcast_ref::<BinaryExpr>() {
        collect_required_columns(binary.left(), input_schema, required)?;
        collect_required_columns(binary.right(), input_schema, required)?;
        return Ok(());
    }
    if let Some(not) = expr.downcast_ref::<NotExpr>() {
        collect_required_columns(not.arg(), input_schema, required)?;
        return Ok(());
    }
    if let Some(negative) = expr.downcast_ref::<NegativeExpr>() {
        return collect_required_columns(negative.arg(), input_schema, required);
    }
    if let Some(is_null) = expr.downcast_ref::<IsNullExpr>() {
        collect_required_columns(is_null.arg(), input_schema, required)?;
        return Ok(());
    }
    if let Some(is_not_null) = expr.downcast_ref::<IsNotNullExpr>() {
        collect_required_columns(is_not_null.arg(), input_schema, required)?;
        return Ok(());
    }
    if let Some(cast) = expr.downcast_ref::<CastExpr>() {
        collect_required_columns(cast.expr(), input_schema, required)?;
        return Ok(());
    }
    if let Some(like) = expr.downcast_ref::<LikeExpr>() {
        collect_required_columns(like.expr(), input_schema, required)?;
        return Ok(());
    }
    // IN-list admission belongs to the expression contract, which rejects
    // the list during lowering with its own capability.
    if let Some(in_list) = expr.downcast_ref::<InListExpr>() {
        collect_required_columns(in_list.expr(), input_schema, required)?;
        for expr in in_list.list() {
            collect_required_columns(expr, input_schema, required)?;
        }
        return Ok(());
    }
    if let Some(case) = expr.downcast_ref::<CaseExpr>() {
        if let Some(expr) = case.expr() {
            collect_required_columns(expr, input_schema, required)?;
        }
        for (when, then) in case.when_then_expr() {
            collect_required_columns(when, input_schema, required)?;
            collect_required_columns(then, input_schema, required)?;
        }
        if let Some(else_expr) = case.else_expr() {
            collect_required_columns(else_expr, input_schema, required)?;
        }
        return Ok(());
    }
    if let Some(dynamic_filter) = expr.downcast_ref::<DynamicFilterPhysicalExpr>() {
        let current = dynamic_filter.current().map_err(|_error| CapabilityReason::UnsupportedExpression)?;
        return collect_required_columns(&current, input_schema, required);
    }
    if let Some(function) = expr.downcast_ref::<ScalarFunctionExpr>() {
        collect_required_columns_for_scalar_function(function, input_schema, required)?;
        return Ok(());
    }
    if let Some(try_cast) = expr.downcast_ref::<TryCastExpr>() {
        collect_required_columns(try_cast.expr(), input_schema, required)?;
        return Ok(());
    }
    Err(CapabilityReason::UnsupportedExpression)
}

fn collect_required_columns_for_scalar_function(
    function: &ScalarFunctionExpr,
    input_schema: &SchemaRef,
    required: &mut BTreeSet<String>,
) -> Result<(), CapabilityReason> {
    for arg in function.args() {
        collect_required_columns(arg, input_schema, required)?;
    }
    Ok(())
}