hamelin_datafusion 0.7.8

//! DataFusion translations for interval arithmetic operations.
//!
//! DataFusion doesn't support interval * number or interval / number directly,
//! so we convert to millis/months, do the arithmetic, and convert back.

use datafusion::logical_expr::{BinaryExpr, Expr as DFExpr, Operator as DFOperator};

use hamelin_lib::func::defs::{
    CalendarIntervalDivideInt, CalendarIntervalMinusCalendarInterval, CalendarIntervalMultiplyInt,
    CalendarIntervalPlusCalendarInterval, CalendarIntervalPlusTimestamp,
    IntMultiplyCalendarInterval, IntervalDivideNumeric, IntervalMinusInterval,
    IntervalMultiplyNumeric, IntervalPlusInterval, IntervalPlusTimestamp, NumericMultiplyInterval,
    TimestampMinusCalendarInterval, TimestampMinusInterval, TimestampMinusTimestamp,
    TimestampPlusCalendarInterval, TimestampPlusInterval,
};

use super::DataFusionTranslationRegistry;
use crate::udf::{from_millis_udf, from_months_udf, to_millis_udf, to_months_udf};

pub fn register(registry: &mut DataFusionTranslationRegistry) {
    // Interval + Interval -> left + right
    registry.register::<IntervalPlusInterval>(|mut params| {
        let left = params.take_by_name("left")?.expr;
        let right = params.take_by_name("right")?.expr;
        Ok(DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(left),
            DFOperator::Plus,
            Box::new(right),
        )))
    });

    // Interval - Interval -> left - right
    registry.register::<IntervalMinusInterval>(|mut params| {
        let left = params.take_by_name("left")?.expr;
        let right = params.take_by_name("right")?.expr;
        Ok(DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(left),
            DFOperator::Minus,
            Box::new(right),
        )))
    });

    // Interval * Numeric -> from_millis(to_millis(interval) * n)
    // DataFusion doesn't support interval * number directly
    registry.register::<IntervalMultiplyNumeric>(|mut params| {
        let left = params.take_by_name("left")?.expr;
        let right = params.take_by_name("right")?.expr;

        // to_millis(interval) * n
        let millis = to_millis_udf().call(vec![left]);
        let product = DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(millis),
            DFOperator::Multiply,
            Box::new(right),
        ));

        // from_millis(product)
        Ok(from_millis_udf().call(vec![product]))
    });

    // Numeric * Interval -> from_millis(n * to_millis(interval))
    registry.register::<NumericMultiplyInterval>(|mut params| {
        let left = params.take_by_name("left")?.expr;
        let right = params.take_by_name("right")?.expr;

        // n * to_millis(interval)
        let millis = to_millis_udf().call(vec![right]);
        let product = DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(left),
            DFOperator::Multiply,
            Box::new(millis),
        ));

        // from_millis(product)
        Ok(from_millis_udf().call(vec![product]))
    });

    // Interval / Numeric -> from_millis(to_millis(interval) / n)
    registry.register::<IntervalDivideNumeric>(|mut params| {
        let left = params.take_by_name("left")?.expr;
        let right = params.take_by_name("right")?.expr;

        // to_millis(interval) / n
        let millis = to_millis_udf().call(vec![left]);
        let quotient = DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(millis),
            DFOperator::Divide,
            Box::new(right),
        ));

        // from_millis(quotient)
        Ok(from_millis_udf().call(vec![quotient]))
    });

    // CalendarInterval + CalendarInterval -> left + right
    registry.register::<CalendarIntervalPlusCalendarInterval>(|mut params| {
        let left = params.take_by_name("left")?.expr;
        let right = params.take_by_name("right")?.expr;
        Ok(DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(left),
            DFOperator::Plus,
            Box::new(right),
        )))
    });

    // CalendarInterval - CalendarInterval -> left - right
    registry.register::<CalendarIntervalMinusCalendarInterval>(|mut params| {
        let left = params.take_by_name("left")?.expr;
        let right = params.take_by_name("right")?.expr;
        Ok(DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(left),
            DFOperator::Minus,
            Box::new(right),
        )))
    });

    // CalendarInterval * Int -> from_months(to_months(interval) * n)
    registry.register::<CalendarIntervalMultiplyInt>(|mut params| {
        let left = params.take_by_name("left")?.expr;
        let right = params.take_by_name("right")?.expr;

        // to_months(interval) * n
        let months = to_months_udf().call(vec![left]);
        let product = DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(months),
            DFOperator::Multiply,
            Box::new(right),
        ));

        // from_months(product)
        Ok(from_months_udf().call(vec![product]))
    });

    // Int * CalendarInterval -> from_months(n * to_months(interval))
    registry.register::<IntMultiplyCalendarInterval>(|mut params| {
        let left = params.take_by_name("left")?.expr;
        let right = params.take_by_name("right")?.expr;

        // n * to_months(interval)
        let months = to_months_udf().call(vec![right]);
        let product = DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(left),
            DFOperator::Multiply,
            Box::new(months),
        ));

        // from_months(product)
        Ok(from_months_udf().call(vec![product]))
    });

    // CalendarInterval / Int -> from_months(to_months(interval) / n)
    registry.register::<CalendarIntervalDivideInt>(|mut params| {
        let left = params.take_by_name("left")?.expr;
        let right = params.take_by_name("right")?.expr;

        // to_months(interval) / n
        let months = to_months_udf().call(vec![left]);
        let quotient = DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(months),
            DFOperator::Divide,
            Box::new(right),
        ));

        // from_months(quotient)
        Ok(from_months_udf().call(vec![quotient]))
    });

    // Timestamp + Interval -> timestamp + interval
    registry.register::<TimestampPlusInterval>(|mut params| {
        let left = params.take_by_name("left")?.expr;
        let right = params.take_by_name("right")?.expr;
        Ok(DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(left),
            DFOperator::Plus,
            Box::new(right),
        )))
    });

    // Interval + Timestamp -> interval + timestamp
    registry.register::<IntervalPlusTimestamp>(|mut params| {
        let left = params.take_by_name("left")?.expr;
        let right = params.take_by_name("right")?.expr;
        Ok(DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(left),
            DFOperator::Plus,
            Box::new(right),
        )))
    });

    // Timestamp - Interval -> timestamp - interval
    registry.register::<TimestampMinusInterval>(|mut params| {
        let left = params.take_by_name("left")?.expr;
        let right = params.take_by_name("right")?.expr;
        Ok(DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(left),
            DFOperator::Minus,
            Box::new(right),
        )))
    });

    // Timestamp - Timestamp -> timestamp - timestamp (yields interval)
    registry.register::<TimestampMinusTimestamp>(|mut params| {
        let left = params.take_by_name("left")?.expr;
        let right = params.take_by_name("right")?.expr;
        Ok(DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(left),
            DFOperator::Minus,
            Box::new(right),
        )))
    });

    // Timestamp + CalendarInterval -> timestamp + interval
    registry.register::<TimestampPlusCalendarInterval>(|mut params| {
        let left = params.take_by_name("left")?.expr;
        let right = params.take_by_name("right")?.expr;
        Ok(DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(left),
            DFOperator::Plus,
            Box::new(right),
        )))
    });

    // CalendarInterval + Timestamp -> interval + timestamp
    registry.register::<CalendarIntervalPlusTimestamp>(|mut params| {
        let left = params.take_by_name("left")?.expr;
        let right = params.take_by_name("right")?.expr;
        Ok(DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(left),
            DFOperator::Plus,
            Box::new(right),
        )))
    });

    // Timestamp - CalendarInterval -> timestamp - interval
    registry.register::<TimestampMinusCalendarInterval>(|mut params| {
        let left = params.take_by_name("left")?.expr;
        let right = params.take_by_name("right")?.expr;
        Ok(DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(left),
            DFOperator::Minus,
            Box::new(right),
        )))
    });
}