algeon-datafusion 0.1.2

use std::{any::Any, sync::Arc};

use arrow::datatypes::{DataType, TimeUnit};
use datafusion::{
    common::ScalarValue,
    functions::{
        core::nullif::NullIfFunc,
        datetime::{date_part::DatePartFunc, date_trunc::DateTruncFunc, to_timestamp::ToTimestampSecondsFunc},
        math::{self, abs::AbsFunc, ceil::CeilFunc, floor::FloorFunc, round::RoundFunc},
        regex::regexpreplace::RegexpReplaceFunc,
        string::{
            btrim::BTrimFunc, concat::ConcatFunc, contains::ContainsFunc, ends_with::EndsWithFunc, lower::LowerFunc,
            ltrim::LtrimFunc, octet_length::OctetLengthFunc, rtrim::RtrimFunc, starts_with::StartsWithFunc,
            upper::UpperFunc,
        },
        unicode::{character_length::CharacterLengthFunc, substr::SubstrFunc},
    },
    functions_aggregate::{
        average::Avg,
        count::Count,
        median::Median,
        min_max::{Max, Min},
        stddev::{Stddev, StddevPop},
        sum::Sum,
        variance::{VariancePopulation, VarianceSample},
    },
    logical_expr::{AggregateUDF, Operator, ScalarUDF},
    physical_expr::{PhysicalExpr, ScalarFunctionExpr, expressions::Literal},
};
use query_engine::rapids_backend::rapids_backend;
use query_planning::{
    datatypes::{
        is_narrow_decimal_carry_type, is_supported_arrow_cast, is_supported_list_carry_type, is_timestamp_with_timezone,
    },
    expr::{AggregateFunction as NativeAggregateFunction, UnaryMathOp},
    validation::{is_integral_type, is_string_type},
};

use super::{ExpressionCapability, ExpressionRejection, ExpressionRejectionSubject};
use crate::native::literal::parse_yyyy_mm_dd_date32;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(in crate::native) enum SupportedBinaryOperator {
    Add,
    Subtract,
    Multiply,
    Divide,
    Modulo,
    BitwiseAnd,
    BitwiseOr,
    BitwiseXor,
    ShiftLeft,
    ShiftRight,
    Eq,
    NotEq,
    Lt,
    LtEq,
    Gt,
    GtEq,
    And,
    Or,
    StringConcat,
    IsDistinctFrom,
    IsNotDistinctFrom,
}

impl SupportedBinaryOperator {
    pub(in crate::native) fn from_datafusion(op: &Operator) -> Result<Self, ExpressionRejection> {
        match op {
            Operator::Plus => Ok(Self::Add),
            Operator::Minus => Ok(Self::Subtract),
            Operator::Multiply => Ok(Self::Multiply),
            Operator::Divide => Ok(Self::Divide),
            Operator::Modulo => Ok(Self::Modulo),
            Operator::BitwiseAnd => Ok(Self::BitwiseAnd),
            Operator::BitwiseOr => Ok(Self::BitwiseOr),
            Operator::BitwiseXor => Ok(Self::BitwiseXor),
            Operator::BitwiseShiftLeft => Ok(Self::ShiftLeft),
            Operator::BitwiseShiftRight => Ok(Self::ShiftRight),
            Operator::Eq => Ok(Self::Eq),
            Operator::NotEq => Ok(Self::NotEq),
            Operator::Lt => Ok(Self::Lt),
            Operator::LtEq => Ok(Self::LtEq),
            Operator::Gt => Ok(Self::Gt),
            Operator::GtEq => Ok(Self::GtEq),
            Operator::And => Ok(Self::And),
            Operator::Or => Ok(Self::Or),
            Operator::StringConcat => Ok(Self::StringConcat),
            Operator::IsDistinctFrom => Ok(Self::IsDistinctFrom),
            Operator::IsNotDistinctFrom => Ok(Self::IsNotDistinctFrom),
            _ => Err(ExpressionRejection::unsupported_expression(ExpressionCapability::BinaryOperator)),
        }
    }

    pub(super) const fn rejects_timestamp_timezone(self) -> bool {
        matches!(
            self,
            Self::Add
                | Self::Subtract
                | Self::Multiply
                | Self::Divide
                | Self::Modulo
                | Self::Eq
                | Self::NotEq
                | Self::Lt
                | Self::LtEq
                | Self::Gt
                | Self::GtEq
                | Self::IsDistinctFrom
                | Self::IsNotDistinctFrom
        )
    }

    pub(super) const fn rejects_list_carry(self) -> bool {
        self.rejects_timestamp_timezone()
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(in crate::native) enum ScalarFunctionKind {
    Substring,
    Extract,
    DateTrunc,
    ToTimestampSeconds,
    Lower,
    Upper,
    Trim,
    LeftTrim,
    RightTrim,
    Contains,
    StartsWith,
    EndsWith,
    StringLength,
    OctetLength,
    RegexpReplace,
    Concat,
    Abs,
    UnaryMath(UnaryMathOp),
    Round,
    NullIf,
}

impl ScalarFunctionKind {
    /// Identifies a scalar function by its implementation type. A user UDF
    /// registered under a builtin name must not inherit the builtin's
    /// validation and lowering.
    pub(in crate::native) fn from_scalar_function(function: &ScalarFunctionExpr) -> Result<Self, ExpressionRejection> {
        let implementation = function.fun().inner().as_ref() as &dyn Any;
        if implementation.is::<SubstrFunc>() {
            Ok(Self::Substring)
        } else if implementation.is::<DatePartFunc>() {
            Ok(Self::Extract)
        } else if implementation.is::<DateTruncFunc>() {
            Ok(Self::DateTrunc)
        } else if implementation.is::<ToTimestampSecondsFunc>() {
            Ok(Self::ToTimestampSeconds)
        } else if implementation.is::<LowerFunc>() {
            Ok(Self::Lower)
        } else if implementation.is::<UpperFunc>() {
            Ok(Self::Upper)
        } else if implementation.is::<BTrimFunc>() {
            Ok(Self::Trim)
        } else if implementation.is::<LtrimFunc>() {
            Ok(Self::LeftTrim)
        } else if implementation.is::<RtrimFunc>() {
            Ok(Self::RightTrim)
        } else if implementation.is::<ContainsFunc>() {
            Ok(Self::Contains)
        } else if implementation.is::<StartsWithFunc>() {
            Ok(Self::StartsWith)
        } else if implementation.is::<EndsWithFunc>() {
            Ok(Self::EndsWith)
        } else if implementation.is::<CharacterLengthFunc>() {
            Ok(Self::StringLength)
        } else if implementation.is::<OctetLengthFunc>() {
            Ok(Self::OctetLength)
        } else if implementation.is::<RegexpReplaceFunc>() {
            Ok(Self::RegexpReplace)
        } else if implementation.is::<ConcatFunc>() {
            Ok(Self::Concat)
        } else if implementation.is::<AbsFunc>() {
            Ok(Self::Abs)
        } else if function.name().eq_ignore_ascii_case("sqrt") && is_builtin_scalar_impl(function.fun(), &math::sqrt())
        {
            Ok(Self::UnaryMath(UnaryMathOp::Sqrt))
        } else if function.name().eq_ignore_ascii_case("ln") && is_builtin_scalar_impl(function.fun(), &math::ln()) {
            Ok(Self::UnaryMath(UnaryMathOp::Ln))
        } else if function.name().eq_ignore_ascii_case("exp") && is_builtin_scalar_impl(function.fun(), &math::exp()) {
            Ok(Self::UnaryMath(UnaryMathOp::Exp))
        } else if implementation.is::<FloorFunc>() {
            Ok(Self::UnaryMath(UnaryMathOp::Floor))
        } else if implementation.is::<CeilFunc>() {
            Ok(Self::UnaryMath(UnaryMathOp::Ceil))
        } else if implementation.is::<RoundFunc>() {
            Ok(Self::Round)
        } else if implementation.is::<NullIfFunc>() {
            Ok(Self::NullIf)
        } else {
            Err(ExpressionRejection::unsupported_expression(ExpressionCapability::ScalarFunction)
                .with_subject(ExpressionRejectionSubject::ScalarFunction(function.name().to_string())))
        }
    }
}

/// Compares a scalar implementation against a freshly built builtin by
/// `TypeId`. Used only for builtins whose implementation types are private to
/// `datafusion-functions` (`sqrt`, `ln`, `exp`); everything else downcasts to
/// a named type above.
fn is_builtin_scalar_impl(fun: &ScalarUDF, builtin: &ScalarUDF) -> bool {
    (fun.inner().as_ref() as &dyn Any).type_id() == (builtin.inner().as_ref() as &dyn Any).type_id()
}

/// Maps an aggregate function to its native lowering by implementation type,
/// not by name: a user UDAF registered under a builtin name resolves to
/// `None` and falls back to CPU.
pub(in crate::native) fn native_aggregate_function_kind(fun: &AggregateUDF) -> Option<NativeAggregateFunction> {
    let implementation = fun.inner().as_ref() as &dyn Any;
    if implementation.is::<Sum>() {
        Some(NativeAggregateFunction::Sum)
    } else if implementation.is::<Min>() {
        Some(NativeAggregateFunction::Min)
    } else if implementation.is::<Max>() {
        Some(NativeAggregateFunction::Max)
    } else if implementation.is::<Count>() {
        Some(NativeAggregateFunction::Count)
    } else if implementation.is::<Avg>() {
        Some(NativeAggregateFunction::Avg)
    } else if implementation.is::<Median>() {
        Some(NativeAggregateFunction::Median)
    } else if implementation.is::<Stddev>() {
        Some(NativeAggregateFunction::StddevSamp)
    } else if implementation.is::<StddevPop>() {
        Some(NativeAggregateFunction::StddevPop)
    } else if implementation.is::<VarianceSample>() {
        Some(NativeAggregateFunction::VarSamp)
    } else if implementation.is::<VariancePopulation>() {
        Some(NativeAggregateFunction::VarPop)
    } else {
        None
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(in crate::native) enum CastCapability {
    Native,
    LiteralStringDate32,
    Decimal128ToFloat32ViaFloat64,
}

pub(in crate::native) fn is_list_carry_type(data_type: &DataType) -> bool {
    is_supported_list_carry_type(data_type)
}

pub(in crate::native) fn is_narrow_decimal_carry(data_type: &DataType) -> bool {
    is_narrow_decimal_carry_type(data_type)
}

pub(in crate::native) fn is_supported_native_cast(source_type: &DataType, target_type: &DataType) -> bool {
    is_supported_arrow_cast(source_type, target_type, rapids_backend())
}

pub(in crate::native) fn cast_capability(
    expr: &Arc<dyn PhysicalExpr>,
    source_type: &DataType,
    target_type: &DataType,
) -> Result<CastCapability, ExpressionRejection> {
    if is_supported_native_cast(source_type, target_type) {
        return Ok(CastCapability::Native);
    }
    if literal_string_date32_cast_is_supported(expr, target_type) {
        return Ok(CastCapability::LiteralStringDate32);
    }
    if is_decimal128_to_float32_cast_via_float64(source_type, target_type) {
        return Ok(CastCapability::Decimal128ToFloat32ViaFloat64);
    }
    Err(ExpressionRejection::unsupported_type(ExpressionCapability::CastPair)
        .with_subject(ExpressionRejectionSubject::Cast { source: source_type.clone(), target: target_type.clone() }))
}

/// Arrow casts `Decimal128` to `Float32` by dividing in `Float64` and then
/// rounding to `Float32`. A direct libcudf fixed-point cast divides in
/// `Float32`, which can round differently, so the lowered expression composes
/// the two supported casts.
pub(in crate::native) fn is_decimal128_to_float32_cast_via_float64(
    source_type: &DataType,
    target_type: &DataType,
) -> bool {
    matches!((source_type, target_type), (DataType::Decimal128(_, _), DataType::Float32))
        && is_supported_native_cast(source_type, &DataType::Float64)
        && is_supported_native_cast(&DataType::Float64, target_type)
}

pub(in crate::native) fn is_supported_try_cast_pair(source_type: &DataType, target_type: &DataType) -> bool {
    is_string_type(source_type) && is_integral_type(target_type)
}

pub(in crate::native) fn is_supported_to_timestamp_seconds_cast(
    source_type: &DataType,
    target_type: &DataType,
) -> bool {
    matches!(target_type, DataType::Timestamp(TimeUnit::Second, None))
        && matches!(
            source_type,
            DataType::Int8
                | DataType::Int16
                | DataType::Int32
                | DataType::Int64
                | DataType::UInt8
                | DataType::UInt16
                | DataType::UInt32
        )
        && is_supported_arrow_cast(source_type, target_type, rapids_backend())
}

pub(in crate::native) fn is_numeric_or_decimal_type(data_type: &DataType) -> bool {
    matches!(
        data_type,
        DataType::Int8
            | DataType::Int16
            | DataType::Int32
            | DataType::Int64
            | DataType::UInt8
            | DataType::UInt16
            | DataType::UInt32
            | DataType::UInt64
            | DataType::Float32
            | DataType::Float64
            | DataType::Decimal128(_, _)
    )
}

pub(in crate::native) fn compatible_types(left: &DataType, right: &DataType) -> bool {
    !is_timestamp_with_timezone(left)
        && !is_timestamp_with_timezone(right)
        && !is_list_carry_type(left)
        && !is_list_carry_type(right)
        && !is_narrow_decimal_carry(left)
        && !is_narrow_decimal_carry(right)
        && (left == right || is_supported_native_cast(left, right) || is_supported_native_cast(right, left))
}

fn literal_string_date32_cast_is_supported(expr: &Arc<dyn PhysicalExpr>, target_type: &DataType) -> bool {
    matches!(target_type, DataType::Date32)
        && expr
            .downcast_ref::<Literal>()
            .and_then(|literal| match literal.value() {
                ScalarValue::Utf8(Some(value))
                | ScalarValue::LargeUtf8(Some(value))
                | ScalarValue::Utf8View(Some(value)) => parse_yyyy_mm_dd_date32(value),
                _ => None,
            })
            .is_some()
}