Search code examples
rustrust-polars

How to use types with lifetimes in Expr::apply closures?


The following function will fail to compile because foo needs to be 'static to work in the closure. However it is clear in this example that foo will outlive the aggregation just fine. Since we cannot used scoped threads here, is there a way in either safe or unsafe rust to allow you to use such types?

use polars::prelude::*;

struct Foo<'a> {
    a: &'a str,
}

fn process<'a>(df: DataFrame, foo: Foo<'a>) -> DataFrame {
    let foo = Arc::new(foo);

    df.lazy()
        .group_by(["idx"])
        .agg([col("col1").apply(move |s| agg_fn(s, foo.clone()), Default::default())])
        .collect()
        .unwrap()
}

fn agg_fn(_s: Series, _foo: Arc<Foo>) -> PolarsResult<Option<Series>> {
    unimplemented!()
}

fn main() {
    let foo = Foo { a: "hello" };
    let df = df!(
        "idx" => [1, 1, 3],
        "col1" => [5.0, 3.0, 1.0],
    )
    .unwrap();

    let result = process(df, foo);
    println!("Result: {result:#}");
}

Solution

  • Polars Expr doesn't have a lifetime so anything you put inside must be 'static, so there is no direct solution to your problem.

    The Arc trick doesn't work because you end up having a Arc<Foo<'a>> that is no better than the original one for your problem.

    The easiest fix would be to avoid the lifetime in Foo in the first place. In your particular example the inner value is a &'static str but I guess that is just an example artifact. If you could that value into a Box or Arc instead of the whole Foo`:

    struct Foo {
        a: Arc<str>,
    }
    

    Then the lifetime will just disappear.

    There is a playground with a simplified version of your problem and this solution.

    If you like to live risky, you can lie to the compiler and get away with it, sometimes. Just don't complain when your program eventually blows up in flames (it will, I know ;-).

    I cannot recommend this line of work, because you are breaking the rules, but here it is anyways:

    impl<'a> Foo<'a> {
        unsafe fn transmute_static(x: Self) -> Foo<'static> {
            std::mem::transmute(x)
        }
    }
    

    And now you can call it instead of using an Arc:

       //let foo = Arc::new(foo);
       let foo = unsafe { Foo::transmute_static(foo) };
    

    It this transmute sound? Well, I think the jury is still out, and people smarter than me don't seem to agree in that point. But for now it seems to work. Here is a playground with a simplified version of your issue and this hacky solution.