pola-rs / pyo3-polars

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

Kernel panics when passing list(categorical) as input #54

Closed ion-elgreco closed 7 months ago

ion-elgreco commented 8 months ago

Passing a list(categorical) to a registered plugin expr will instantly kill the kernel:

Assume this simple expression

#[polars_expr(output_type=Float64)]
fn test(inputs: &[Series]) -> PolarsResult<Series> {
    Ok(Series::from_vec("TEST", vec![0.0]))
}

Then this plugin code:

@pl.api.register_expr_namespace("testing")
class TestSpace:
    def __init__(self, expr: pl.Expr):
        self._expr = expr

    def test(self) -> pl.Expr:
        """test"""
        return self._expr.register_plugin(
            lib=lib,
            symbol="test",
            is_elementwise=True,
        )

Now exucting this expression on list(categorical) will faill.

df = pl.DataFrame(
    {"x": [["1"]],},
    schema={"x": pl.List(pl.Categorical)},
)

df.with_columns(
    pl.col("x").testing.test(),
)

image

However this runs fine:

df = pl.DataFrame(
    {"x": [["1"]]},
    schema={"x": pl.List(pl.Utf8)},
)

df.with_columns(
    pl.col("x").dist_list.test(),
)
shape: (1, 1)
┌─────┐
│ x   │
│ --- │
│ f64 │
╞═════╡
│ 0.0 │
└─────┘