hamelin_lib 0.19.0

//! Aggregate functions

use std::any::TypeId;

use crate::func::def::{FunctionDef, ParameterBindingProvider, Parameters, SpecialPosition};
use crate::types::array::Array;
use crate::types::map::Map;
use crate::types::matcher::{
    numeric_or_interval_matcher, AnyMatcher, ExactMatcher, MapKeyMatcher, NumericMatcher, OrMatcher,
};
use crate::types::{Type, BOOLEAN, DOUBLE, INT, INTERVAL, STRING, TIMESTAMP};

/// count() - Count all rows (COUNT(*)).
#[derive(Default, Clone)]
pub struct CountStar;

impl FunctionDef for CountStar {
    fn name(&self) -> &'static str {
        "count"
    }

    fn return_type(&self, _bindings: &dyn ParameterBindingProvider) -> anyhow::Result<Type> {
        Ok(INT)
    }

    fn type_id(&self) -> TypeId {
        TypeId::of::<Self>()
    }

    fn special_position(&self) -> Option<SpecialPosition> {
        Some(SpecialPosition::Agg)
    }
}

/// count(x) - Count non-null values.
#[derive(Default, Clone)]
pub struct CountAny;

impl FunctionDef for CountAny {
    fn name(&self) -> &'static str {
        "count"
    }

    fn parameters(&self) -> Parameters {
        Parameters::new().with("x", AnyMatcher::default())
    }

    fn return_type(&self, _bindings: &dyn ParameterBindingProvider) -> anyhow::Result<Type> {
        Ok(INT)
    }

    fn type_id(&self) -> TypeId {
        TypeId::of::<Self>()
    }

    fn special_position(&self) -> Option<SpecialPosition> {
        Some(SpecialPosition::Agg)
    }
}

/// count_distinct(x) - Count distinct non-null values.
#[derive(Default, Clone)]
pub struct CountDistinct;

impl FunctionDef for CountDistinct {
    fn name(&self) -> &'static str {
        "count_distinct"
    }

    fn parameters(&self) -> Parameters {
        Parameters::new().with("x", AnyMatcher::default())
    }

    fn return_type(&self, _bindings: &dyn ParameterBindingProvider) -> anyhow::Result<Type> {
        Ok(INT)
    }

    fn type_id(&self) -> TypeId {
        TypeId::of::<Self>()
    }

    fn special_position(&self) -> Option<SpecialPosition> {
        Some(SpecialPosition::Agg)
    }
}

/// approx_distinct(x) - Approximate count distinct.
#[derive(Default, Clone)]
pub struct ApproxDistinct;

impl FunctionDef for ApproxDistinct {
    fn name(&self) -> &'static str {
        "approx_distinct"
    }

    fn parameters(&self) -> Parameters {
        Parameters::new().with("x", AnyMatcher::default())
    }

    fn return_type(&self, _bindings: &dyn ParameterBindingProvider) -> anyhow::Result<Type> {
        Ok(INT)
    }

    fn type_id(&self) -> TypeId {
        TypeId::of::<Self>()
    }

    fn special_position(&self) -> Option<SpecialPosition> {
        Some(SpecialPosition::Agg)
    }
}

/// count_if(condition) - Count rows where condition is true.
#[derive(Default, Clone)]
pub struct CountIf;

impl FunctionDef for CountIf {
    fn name(&self) -> &'static str {
        "count_if"
    }

    fn parameters(&self) -> Parameters {
        Parameters::new().with("condition", ExactMatcher::of(BOOLEAN))
    }

    fn return_type(&self, _bindings: &dyn ParameterBindingProvider) -> anyhow::Result<Type> {
        Ok(INT)
    }

    fn type_id(&self) -> TypeId {
        TypeId::of::<Self>()
    }

    fn special_position(&self) -> Option<SpecialPosition> {
        Some(SpecialPosition::Agg)
    }
}

/// sum(x) - Sum of numeric values.
#[derive(Default, Clone)]
pub struct Sum;

impl FunctionDef for Sum {
    fn name(&self) -> &'static str {
        "sum"
    }

    fn parameters(&self) -> Parameters {
        Parameters::new().with("x", numeric_or_interval_matcher())
    }

    fn return_type(&self, bindings: &dyn ParameterBindingProvider) -> anyhow::Result<Type> {
        // Return type matches input type
        Ok(bindings.get_by_index(0)?.typ().clone())
    }

    fn type_id(&self) -> TypeId {
        TypeId::of::<Self>()
    }

    fn special_position(&self) -> Option<SpecialPosition> {
        Some(SpecialPosition::Agg)
    }
}

/// avg(x) - Average of numeric values.
#[derive(Default, Clone)]
pub struct Avg;

impl FunctionDef for Avg {
    fn name(&self) -> &'static str {
        "avg"
    }

    fn parameters(&self) -> Parameters {
        Parameters::new().with("x", numeric_or_interval_matcher())
    }

    fn return_type(&self, bindings: &dyn ParameterBindingProvider) -> anyhow::Result<Type> {
        match bindings.get_by_index(0)?.typ() {
            Type::Int => Ok(DOUBLE),
            other => Ok(other.clone()),
        }
    }

    fn type_id(&self) -> TypeId {
        TypeId::of::<Self>()
    }

    fn special_position(&self) -> Option<SpecialPosition> {
        Some(SpecialPosition::Agg)
    }
}

/// stddev(x) - Standard deviation of numeric values.
#[derive(Default, Clone)]
pub struct Stddev;

impl FunctionDef for Stddev {
    fn name(&self) -> &'static str {
        "stddev"
    }

    fn parameters(&self) -> Parameters {
        Parameters::new().with("x", NumericMatcher::default())
    }

    fn return_type(&self, _bindings: &dyn ParameterBindingProvider) -> anyhow::Result<Type> {
        Ok(DOUBLE)
    }

    fn type_id(&self) -> TypeId {
        TypeId::of::<Self>()
    }

    fn special_position(&self) -> Option<SpecialPosition> {
        Some(SpecialPosition::Agg)
    }
}

/// approx_percentile(x, percentile) - Approximate percentile.
#[derive(Default, Clone)]
pub struct ApproxPercentile;

impl FunctionDef for ApproxPercentile {
    fn name(&self) -> &'static str {
        "approx_percentile"
    }

    fn parameters(&self) -> Parameters {
        Parameters::new()
            .with("x", NumericMatcher::default())
            .with("percentile", NumericMatcher::default())
    }

    fn return_type(&self, bindings: &dyn ParameterBindingProvider) -> anyhow::Result<Type> {
        match bindings.get_by_name("x")?.typ() {
            Type::Int => Ok(INT.into()),
            _ => Ok(DOUBLE.clone()),
        }
    }

    fn type_id(&self) -> TypeId {
        TypeId::of::<Self>()
    }

    fn special_position(&self) -> Option<SpecialPosition> {
        Some(SpecialPosition::Agg)
    }
}

/// min(x) - Minimum value (aggregate).
#[derive(Default, Clone)]
pub struct AggMin;

impl FunctionDef for AggMin {
    fn name(&self) -> &'static str {
        "min"
    }

    fn parameters(&self) -> Parameters {
        Parameters::new().with(
            "x",
            OrMatcher::default()
                .with(NumericMatcher::default())
                .with(ExactMatcher::of(STRING))
                .with(ExactMatcher::of(TIMESTAMP))
                .with(ExactMatcher::of(INTERVAL)),
        )
    }

    fn return_type(&self, bindings: &dyn ParameterBindingProvider) -> anyhow::Result<Type> {
        Ok(bindings.get_by_index(0)?.typ().clone())
    }

    fn type_id(&self) -> TypeId {
        TypeId::of::<Self>()
    }

    fn special_position(&self) -> Option<SpecialPosition> {
        Some(SpecialPosition::Agg)
    }
}

/// max(x) - Maximum value (aggregate).
#[derive(Default, Clone)]
pub struct AggMax;

impl FunctionDef for AggMax {
    fn name(&self) -> &'static str {
        "max"
    }

    fn parameters(&self) -> Parameters {
        Parameters::new().with(
            "x",
            OrMatcher::default()
                .with(NumericMatcher::default())
                .with(ExactMatcher::of(STRING))
                .with(ExactMatcher::of(TIMESTAMP))
                .with(ExactMatcher::of(INTERVAL)),
        )
    }

    fn return_type(&self, bindings: &dyn ParameterBindingProvider) -> anyhow::Result<Type> {
        Ok(bindings.get_by_index(0)?.typ().clone())
    }

    fn type_id(&self) -> TypeId {
        TypeId::of::<Self>()
    }

    fn special_position(&self) -> Option<SpecialPosition> {
        Some(SpecialPosition::Agg)
    }
}

/// any_value(x) - Return any value from the group.
#[derive(Default, Clone)]
pub struct AnyValue;

impl FunctionDef for AnyValue {
    fn name(&self) -> &'static str {
        "any_value"
    }

    fn parameters(&self) -> Parameters {
        Parameters::new().with("x", AnyMatcher::default())
    }

    fn return_type(&self, bindings: &dyn ParameterBindingProvider) -> anyhow::Result<Type> {
        Ok(bindings.get_by_index(0)?.typ().clone())
    }

    fn type_id(&self) -> TypeId {
        TypeId::of::<Self>()
    }

    fn special_position(&self) -> Option<SpecialPosition> {
        Some(SpecialPosition::Agg)
    }
}

/// array_agg(x) - Aggregate values into an array.
#[derive(Default, Clone)]
pub struct ArrayAgg;

impl FunctionDef for ArrayAgg {
    fn name(&self) -> &'static str {
        "array_agg"
    }

    fn parameters(&self) -> Parameters {
        Parameters::new().with("x", AnyMatcher::default())
    }

    fn return_type(&self, bindings: &dyn ParameterBindingProvider) -> anyhow::Result<Type> {
        Ok(Array::new(bindings.get_by_index(0)?.typ().clone()).into())
    }

    fn type_id(&self) -> TypeId {
        TypeId::of::<Self>()
    }

    fn special_position(&self) -> Option<SpecialPosition> {
        Some(SpecialPosition::Agg)
    }

    fn sortable_input(&self) -> bool {
        true
    }
}

/// set_agg(x) - Aggregate distinct values into an array.
#[derive(Default, Clone)]
pub struct SetAgg;

impl FunctionDef for SetAgg {
    fn name(&self) -> &'static str {
        "set_agg"
    }

    fn parameters(&self) -> Parameters {
        Parameters::new().with("x", AnyMatcher::default())
    }

    fn return_type(&self, bindings: &dyn ParameterBindingProvider) -> anyhow::Result<Type> {
        Ok(Array::new(bindings.get_by_index(0)?.typ().clone()).into())
    }

    fn type_id(&self) -> TypeId {
        TypeId::of::<Self>()
    }

    fn special_position(&self) -> Option<SpecialPosition> {
        Some(SpecialPosition::Agg)
    }
}

/// map_agg(key, value) - Aggregate key-value pairs into a map.
#[derive(Default, Clone)]
pub struct MapAgg;

impl FunctionDef for MapAgg {
    fn name(&self) -> &'static str {
        "map_agg"
    }

    fn parameters(&self) -> Parameters {
        Parameters::new()
            .with("key", MapKeyMatcher::default())
            .with("value", AnyMatcher::default())
    }

    fn return_type(&self, bindings: &dyn ParameterBindingProvider) -> anyhow::Result<Type> {
        let key = bindings.get_by_index(0)?.typ().clone();
        let value = bindings.get_by_index(1)?.typ().clone();
        Ok(Map::new(key, value).into())
    }

    fn type_id(&self) -> TypeId {
        TypeId::of::<Self>()
    }

    fn special_position(&self) -> Option<SpecialPosition> {
        Some(SpecialPosition::Agg)
    }

    fn sortable_input(&self) -> bool {
        true
    }
}

/// multimap_agg(key, value) - Aggregate key-value pairs into a multimap.
#[derive(Default, Clone)]
pub struct MultimapAgg;

impl FunctionDef for MultimapAgg {
    fn name(&self) -> &'static str {
        "multimap_agg"
    }

    fn parameters(&self) -> Parameters {
        Parameters::new()
            .with("key", MapKeyMatcher::default())
            .with("value", AnyMatcher::default())
    }

    fn return_type(&self, bindings: &dyn ParameterBindingProvider) -> anyhow::Result<Type> {
        let key = bindings.get_by_index(0)?.typ().clone();
        let value = bindings.get_by_index(1)?.typ().clone();
        Ok(Map::new(key, Array::new(value).into()).into())
    }

    fn type_id(&self) -> TypeId {
        TypeId::of::<Self>()
    }

    fn special_position(&self) -> Option<SpecialPosition> {
        Some(SpecialPosition::Agg)
    }

    fn sortable_input(&self) -> bool {
        true
    }
}

/// any(x) - Aggregate boolean OR (any true).
#[derive(Default, Clone)]
pub struct AggAny;

impl FunctionDef for AggAny {
    fn name(&self) -> &'static str {
        "any"
    }

    fn parameters(&self) -> Parameters {
        Parameters::new().with("x", ExactMatcher::of(BOOLEAN))
    }

    fn return_type(&self, _bindings: &dyn ParameterBindingProvider) -> anyhow::Result<Type> {
        Ok(BOOLEAN)
    }

    fn type_id(&self) -> TypeId {
        TypeId::of::<Self>()
    }

    fn special_position(&self) -> Option<SpecialPosition> {
        Some(SpecialPosition::Agg)
    }
}

/// all(x) - Aggregate boolean AND (all true).
#[derive(Default, Clone)]
pub struct AggAll;

impl FunctionDef for AggAll {
    fn name(&self) -> &'static str {
        "all"
    }

    fn parameters(&self) -> Parameters {
        Parameters::new().with("x", ExactMatcher::of(BOOLEAN))
    }

    fn return_type(&self, _bindings: &dyn ParameterBindingProvider) -> anyhow::Result<Type> {
        Ok(BOOLEAN)
    }

    fn type_id(&self) -> TypeId {
        TypeId::of::<Self>()
    }

    fn special_position(&self) -> Option<SpecialPosition> {
        Some(SpecialPosition::Agg)
    }
}