hamelin_datafusion 0.7.8

//! DataFusion translations for comparison operators.

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

use hamelin_lib::func::defs::{
    BooleanEq, BooleanNeq, IntervalEq, IntervalGt, IntervalGte, IntervalLt, IntervalLte,
    IntervalNeq, Is, IsNot, NumericEq, NumericGt, NumericGte, NumericLt, NumericLte, NumericNeq,
    StringEq, StringNeq, TimestampEq, TimestampGt, TimestampGte, TimestampLt, TimestampLte,
    TimestampNeq,
};

use super::DataFusionTranslationRegistry;

/// Helper to create a binary comparison expression.
fn binary_eq(left: DFExpr, right: DFExpr) -> DFExpr {
    DFExpr::BinaryExpr(BinaryExpr::new(
        Box::new(left),
        DFOperator::Eq,
        Box::new(right),
    ))
}

fn binary_neq(left: DFExpr, right: DFExpr) -> DFExpr {
    DFExpr::BinaryExpr(BinaryExpr::new(
        Box::new(left),
        DFOperator::NotEq,
        Box::new(right),
    ))
}

fn binary_lt(left: DFExpr, right: DFExpr) -> DFExpr {
    DFExpr::BinaryExpr(BinaryExpr::new(
        Box::new(left),
        DFOperator::Lt,
        Box::new(right),
    ))
}

fn binary_lte(left: DFExpr, right: DFExpr) -> DFExpr {
    DFExpr::BinaryExpr(BinaryExpr::new(
        Box::new(left),
        DFOperator::LtEq,
        Box::new(right),
    ))
}

fn binary_gt(left: DFExpr, right: DFExpr) -> DFExpr {
    DFExpr::BinaryExpr(BinaryExpr::new(
        Box::new(left),
        DFOperator::Gt,
        Box::new(right),
    ))
}

fn binary_gte(left: DFExpr, right: DFExpr) -> DFExpr {
    DFExpr::BinaryExpr(BinaryExpr::new(
        Box::new(left),
        DFOperator::GtEq,
        Box::new(right),
    ))
}

pub fn register(registry: &mut DataFusionTranslationRegistry) {
    // Timestamp comparisons
    registry.register::<TimestampEq>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;
        Ok(binary_eq(left, right))
    });

    registry.register::<TimestampNeq>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;
        Ok(binary_neq(left, right))
    });

    registry.register::<TimestampLt>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;
        Ok(binary_lt(left, right))
    });

    registry.register::<TimestampLte>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;
        Ok(binary_lte(left, right))
    });

    registry.register::<TimestampGt>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;
        Ok(binary_gt(left, right))
    });

    registry.register::<TimestampGte>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;
        Ok(binary_gte(left, right))
    });

    // Interval comparisons
    registry.register::<IntervalEq>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;
        Ok(binary_eq(left, right))
    });

    registry.register::<IntervalNeq>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;
        Ok(binary_neq(left, right))
    });

    registry.register::<IntervalLt>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;
        Ok(binary_lt(left, right))
    });

    registry.register::<IntervalLte>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;
        Ok(binary_lte(left, right))
    });

    registry.register::<IntervalGt>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;
        Ok(binary_gt(left, right))
    });

    registry.register::<IntervalGte>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;
        Ok(binary_gte(left, right))
    });

    // Boolean comparisons
    registry.register::<BooleanEq>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;
        Ok(binary_eq(left, right))
    });

    registry.register::<BooleanNeq>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;
        Ok(binary_neq(left, right))
    });

    // String comparisons
    registry.register::<StringEq>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;
        Ok(binary_eq(left, right))
    });

    registry.register::<StringNeq>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;
        Ok(binary_neq(left, right))
    });

    // Numeric comparisons
    registry.register::<NumericEq>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;
        Ok(binary_eq(left, right))
    });

    registry.register::<NumericNeq>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;
        Ok(binary_neq(left, right))
    });

    registry.register::<NumericLt>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;
        Ok(binary_lt(left, right))
    });

    registry.register::<NumericLte>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;
        Ok(binary_lte(left, right))
    });

    registry.register::<NumericGt>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;
        Ok(binary_gt(left, right))
    });

    registry.register::<NumericGte>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;
        Ok(binary_gte(left, right))
    });

    // IS / IS NOT (null checking)
    registry.register::<Is>(|mut params| {
        let left = params.take()?.expr;
        // Right is always NULL in IS expressions
        let _right = params.take()?;
        Ok(left.is_null())
    });

    registry.register::<IsNot>(|mut params| {
        let left = params.take()?.expr;
        // Right is always NULL in IS NOT expressions
        let _right = params.take()?;
        Ok(left.is_not_null())
    });
}