hamelin_datafusion 0.7.5

//! DataFusion translations for conditional functions.

use datafusion::logical_expr::expr::Case as DFCase;
use datafusion::logical_expr::Expr as DFExpr;
use datafusion_functions::core::expr_ext::FieldAccessor;
use datafusion_functions::core::expr_fn as core_fn;

use hamelin_lib::func::defs::{Case, Coalesce, If2, If3};

use super::DataFusionTranslationRegistry;

pub fn register(registry: &mut DataFusionTranslationRegistry) {
    // if(condition, then) -> CASE WHEN condition THEN then END
    registry.register::<If2>(|mut params| {
        let condition = params.take()?.expr;
        let then = params.take()?.expr;
        Ok(DFExpr::Case(DFCase {
            expr: None,
            when_then_expr: vec![(Box::new(condition), Box::new(then))],
            else_expr: None,
        }))
    });

    // if(condition, then, else) -> CASE WHEN condition THEN then ELSE else END
    registry.register::<If3>(|mut params| {
        let condition = params.take()?.expr;
        let then = params.take()?.expr;
        let else_expr = params.take()?.expr;
        Ok(DFExpr::Case(DFCase {
            expr: None,
            when_then_expr: vec![(Box::new(condition), Box::new(then))],
            else_expr: Some(Box::new(else_expr)),
        }))
    });

    // case((cond1, val1), (cond2, val2), ...) -> CASE WHEN cond1 THEN val1 WHEN cond2 THEN val2 ... END
    registry.register::<Case>(|params| {
        // Each parameter is a tuple (condition, value), represented as a struct
        // We need to extract the condition and value from each tuple
        let mut when_then_expr = Vec::new();

        for tuple_translation in params.into_iter() {
            // The tuple is represented as a struct with fields c0 (condition) and c1 (value)
            // Extract using get_field
            let tuple_expr = tuple_translation.expr;
            let condition = tuple_expr.clone().field("c0");
            let value = tuple_expr.field("c1");
            when_then_expr.push((Box::new(condition), Box::new(value)));
        }

        Ok(DFExpr::Case(DFCase {
            expr: None,
            when_then_expr,
            else_expr: None,
        }))
    });

    // coalesce(val1, val2, ...) -> COALESCE(val1, val2, ...)
    registry.register::<Coalesce>(|params| {
        let args: Vec<DFExpr> = params.into_iter().map(|t| t.expr).collect();
        Ok(core_fn::coalesce(args))
    });
}