pola-rs / pyo3-polars

Plugins/extension for Polars
MIT License
232 stars 38 forks source link

How to pass kwarg to func output type #58

Closed ion-elgreco closed 6 months ago

ion-elgreco commented 7 months ago

We added a function in polars-xdt which changes the timezone of the datetime, this know at hand since we pass the time-zone through the kwargs, see some example code here to mimic the behaviour:

#[derive(Deserialize)]
pub struct OutputKwargs{
    to_tz: std::string::String,
}

pub fn output(input_fields: &[Field]) -> PolarsResult<Field> {
    let field = input_fields[0].clone();
    let dtype = match field.dtype {
        DataType::Datetime(unit, _) => DataType::Datetime(unit, "SET TZ HERE from KWARGS"),
        _ => polars_bail!(InvalidOperation:
            "dtype '{}' not supported", field.dtype
        ),
    };
    Ok(Field::new(&field.name, dtype))
}
#[polars_expr(output_type_func=output)]
fn function(inputs: &[Series], kwargs: OutputKwargs) -> PolarsResult<Series> {
    let s1 = &inputs[0];
    let ca = s1.datetime().unwrap();
    let converted = ca.clone().convert_time_zone(kwargs.to_tz)?;

    Ok(converted.clone().into_series())
}

I would like to pass the kargs.to_tz to the from_local_datetime_output, is this possible? If not how should I handle the func output type in that case.

Do you have any ideas @ritchie46?

ritchie46 commented 7 months ago

You will have to compile to expressions. In your register_plugin you will then know which expression to call.

ion-elgreco commented 7 months ago

@ritchie46 so I would have to create an expression with a func_output_field for each timezone?

ritchie46 commented 7 months ago

Oh.. it is different for every timezone? Hmm.. have to think about this.

ion-elgreco commented 7 months ago

Yes, the user provides the timezone, and then the function returns the datetime with this timezone