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

Code before:

#[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),
    };

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

#[cfg(feature = "dtype-datetime")]
// Determines the output data type, including user-specified t tz, and interval.
pub(super) fn build_datetime_supertype(
    default: DataType,
    tu: &Option<TimeUnit>,
    tz: &Option<TimeZone>,
    interval: &Duration,
) -> PolarsResult<DataType> {
    let mut dtype_out = match (&default, tu) {
        (DataType::Date, time_unit) => {
            if let Some(tu) = time_unit {
                DataType::Datetime(*tu, None)
            } else if interval.nanoseconds() % 1_000 != 0 {
                DataType::Datetime(TimeUnit::Nanoseconds, None)
            } else {
                // No datatype, use microseconds
                DataType::Datetime(TimeUnit::Microseconds, None)
            }
        },
        // overwrite nothing, keep as-is
        (DataType::Datetime(_, _), None) => default,
        // overwrite time unit, keep timezone
        (DataType::Datetime(_, tz), Some(tu)) => DataType::Datetime(*tu, tz.clone()),
        (dt, _) => {
            polars_bail!(InvalidOperation: "expected a temporal datatype, got {}", dt)
        },
    };

    // Overwrite time zone, if specified

Review comment (suggestion):
```suggestion
    start_end_supertype: DataType,
```

Code after:
}

#[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,
                options,
            }
        },
        dt => polars_bail!(ComputeError: "cannot convert time zone of dtype {:?}", dt),
    };

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

// Determines the output data type, including user-specified t tz, and interval.
#[cfg(feature = "dtype-datetime")]
pub(super) fn temporal_range_output_type(
    start_end_supertype: DataType,
    tu: &Option<TimeUnit>,
    tz: &Option<TimeZone>,
    interval: &Duration,
) -> PolarsResult<DataType> {
    let mut dtype_out = match (&start_end_supertype, tu) {
        (DataType::Date, time_unit) => {
            if let Some(tu) = time_unit {
                DataType::Datetime(*tu, None)
            } else if interval.nanoseconds() % 1_000 != 0 {
                DataType::Datetime(TimeUnit::Nanoseconds, None)
            } else {
                // No datatype, use microseconds
                DataType::Datetime(TimeUnit::Microseconds, None)
            }
        },
        // overwrite nothing, keep as-is
        (DataType::Datetime(_, _), None) => start_end_supertype,
        // overwrite time unit, keep timezone
        (DataType::Datetime(_, tz), Some(tu)) => DataType::Datetime(*tu, tz.clone()),