[Python] Using multithreading in Polars Expression Plugins

Discussão em 'Python' iniciado por Stack, Março 3, 2025.

  1. Stack

    Stack Membro Participativo

    I want to write a custom Polars Expression Plugins (https://docs.pola.rs/user-guide/plugins/expr_plugins/) for doing some compute heavy calculations. I could see a big andvantage in making them multi-threaded (especially because they operate element-wise i.e rows are independent).

    In this issue it is mentioned that:


    But each expression individually will be single-threaded. If you want it parallel you need to implement that yourself, but then you risk bad interactions with Polars' own multithreading
    ​

    Do you know if there is a smart/recomended way to do this? Could one use the same thread pool (or similar) as polars is using to avoid starving ressources?

    I have tried digging around in the polars code and can see the POOL object but I'm unsure how to go about it. Could you provide an example of how one could modify the code in the example below to achieve this?

    // src/expressions.rs
    use polars::prelude::*;
    use pyo3_polars::derive::polars_expr;
    use std::fmt::Write;

    fn pig_latin_str(value: &str, output: &mut String) {
    if let Some(first_char) = value.chars().next() {
    write!(output, "{}{}ay", &value[1..], first_char).unwrap()
    }
    }

    #[polars_expr(output_type=String)]
    fn pig_latinnify(inputs: &[Series]) -> PolarsResult<Series> {
    let ca = inputs[0].str()?;
    let out: StringChunked = ca.apply_into_string_amortized(pig_latin_str);
    Ok(out.into_series())
    }

    Continue reading...

Compartilhe esta Página