extension.rs

use std::sync::Arc;

use polars_core::error::PolarsResult;
use polars_core::prelude::*;
use polars_plan::dsl::{ColumnsUdf, SpecialEq};
use polars_plan::plans::IRExtensionFunction;

pub fn function_expr_to_udf(func: IRExtensionFunction) -> SpecialEq<Arc<dyn ColumnsUdf>> {
    use IRExtensionFunction::*;
    match func {
        To(dtype) => map!(ext_to, dtype.clone()),
        Storage => map!(ext_storage),
    }
}

fn ext_to(s: &Column, dtype: DataType) -> PolarsResult<Column> {
    let DataType::Extension(typ, storage) = &dtype else {
        polars_bail!(ComputeError: "ext.to() requires an Extension dtype")
    };

    Ok(s.apply_unary_elementwise(|s| {
        assert!(*s.dtype() == **storage);
        s.clone().into_extension(typ.clone())
    }))
}

fn ext_storage(s: &Column) -> PolarsResult<Column> {
    Ok(s.apply_unary_elementwise(|s| s.to_storage().clone()))
}