hamelin_datafusion 0.7.8

//! DataFusion translations for JSON functions.

use datafusion::logical_expr::lit;
use datafusion_functions::core::expr_fn as core_fn;

use hamelin_lib::func::defs::{ParseJsonString, ParseJsonVariant, ToJsonString, Typeof};

use super::DataFusionTranslationRegistry;
use crate::struct_expansion::hamelin_type_to_arrow;

pub fn register(registry: &mut DataFusionTranslationRegistry) {
    // parse_json(json_string) -> hamelin_json_to_variant(json_string)
    // Converts a JSON string to a native Arrow Variant type
    registry.register::<ParseJsonString>(|mut params| {
        let json_string = params.take()?.expr;
        Ok(crate::udf::json_to_variant_udf().call(vec![json_string]))
    });

    // parse_json(variant) -> pass through (already a variant)
    registry.register::<ParseJsonVariant>(|mut params| {
        let variant = params.take()?.expr;
        Ok(variant)
    });

    // to_json_string(variant) -> hamelin_to_json_string(variant)
    registry.register::<ToJsonString>(|mut params| {
        let value = params.take()?.expr;
        Ok(crate::udf::variant_to_json_udf().call(vec![value]))
    });

    // typeof(x) -> named_struct("hamelin_type", <type_string>, "sql_type", <arrow_type_string>)
    // Returns a struct with the Hamelin type name and the corresponding Arrow/SQL type name.
    // This is evaluated at compile time - the type is known statically.
    registry.register::<Typeof>(|mut params| {
        let x = params.take()?;
        let hamelin_type = x.typ.as_ref();

        // Get the Hamelin type string representation
        let hamelin_type_str = hamelin_type.to_string();

        // Get the Arrow/DataFusion type string representation
        let arrow_type = hamelin_type_to_arrow(hamelin_type);
        let sql_type_str = format!("{}", arrow_type);

        // Build a struct with both type names as literal strings
        Ok(core_fn::named_struct(vec![
            lit("hamelin_type"),
            lit(hamelin_type_str),
            lit("sql_type"),
            lit(sql_type_str),
        ]))
    });
}