Repository: pola-rs/polars
File: crates/polars-plan/src/plans/conversion/type_coercion/datetime.rs (Rust)

Code before:
        expr.set_dtype(to_type.clone());
    }
    Ok(())
}

#[cfg(all(feature = "dtype-datetime", feature = "timezones"))]
pub(super) fn replace_tz(
    e: &mut ExprIR,
    dtype: &DataType,
    expr_arena: &mut Arena<AExpr>,
) -> PolarsResult<()> {
    let replacement_expr = match &dtype {
        // Wrap our node in a ReplaceTimezone node
        &DataType::Datetime(_, tz @ Some(_)) => AExpr::Function {
            input: {
                // We must add the ambiguous input argument.
                let scalar = Scalar::new(DataType::String, AnyValue::from("raise"));
                let node = expr_arena.add(AExpr::Literal(LiteralValue::Scalar(scalar)));
                let ambiguous = ExprIR::from_node(node, expr_arena);
                vec![e.clone(), ambiguous]
            },
            function: IRFunctionExpr::TemporalExpr(IRTemporalFunction::ReplaceTimeZone(
                tz.clone(),
                NonExistent::Raise,
            )),
            options: FunctionOptions::elementwise(),
        },
        dt => polars_bail!(ComputeError: "cannot replace time zone of dtype {:?}", dt),
    };

    e.set_node(expr_arena.add(replacement_expr));
    e.set_dtype(dtype.clone());
    Ok(())
}

#[cfg(all(feature = "dtype-datetime", feature = "timezones"))]
pub(super) fn convert_tz(
    e: &mut ExprIR,
    dtype: &DataType,
    expr_arena: &mut Arena<AExpr>,
) -> PolarsResult<()> {
    let replacement_expr = match &dtype {
        // Wrap our node in a ReplaceTimezone node
        &DataType::Datetime(_, Some(tz)) => AExpr::Function {
            input: vec![e.clone()],
            function: IRFunctionExpr::TemporalExpr(IRTemporalFunction::ConvertTimeZone(tz.clone())),
            options: FunctionOptions::elementwise(),
        },
        dt => polars_bail!(ComputeError: "cannot convert time zone of dtype {:?}", dt),
    };


Review comment (suggestion):
This should initialize with `function.function_options()`, with a `let function = IRFunctionExpr::TemporalExpr(..)` above.


Code after:
}

#[cfg(all(feature = "dtype-datetime", feature = "timezones"))]
pub(super) fn replace_tz(
    e: &mut ExprIR,
    dtype: &DataType,
    expr_arena: &mut Arena<AExpr>,
) -> PolarsResult<()> {
    let replacement_expr = match &dtype {
        &DataType::Datetime(_, tz @ Some(_)) => {
            // Wrap our node in a ReplaceTimezone node
            let function = IRFunctionExpr::TemporalExpr(IRTemporalFunction::ReplaceTimeZone(
                tz.clone(),
                NonExistent::Raise,
            ));
            let options = function.function_options();
            AExpr::Function {
                input: {
                    // We must add the ambiguous input argument.
                    let scalar = Scalar::new(DataType::String, AnyValue::from("raise"));
                    let node = expr_arena.add(AExpr::Literal(LiteralValue::Scalar(scalar)));
                    let ambiguous = ExprIR::from_node(node, expr_arena);
                    vec![e.clone(), ambiguous]
                },
                function,
                options,
            }
        },
        dt => polars_bail!(ComputeError: "cannot replace time zone of dtype {:?}", dt),
    };

    e.set_node(expr_arena.add(replacement_expr));
    e.set_dtype(dtype.clone());
    Ok(())
}

#[cfg(all(feature = "dtype-datetime", feature = "timezones"))]
pub(super) fn convert_tz(
    e: &mut ExprIR,
    dtype: &DataType,
    expr_arena: &mut Arena<AExpr>,
) -> PolarsResult<()> {
    let replacement_expr = match &dtype {
        // Wrap our node in a ReplaceTimezone node
        &DataType::Datetime(_, Some(tz)) => {
            let function =
                IRFunctionExpr::TemporalExpr(IRTemporalFunction::ConvertTimeZone(tz.clone()));
            let options = function.function_options();
            AExpr::Function {
                input: vec![e.clone()],
                function,