Interested in adding a section on using rust classes in the kwargs?

[thomasfrederikhoeck]

I have a use-case where I want to be able to use some logic in rust to create custom class and us this as a kwarg in my plugin.

I took some time to figure out how to wire it together and I thought it might be interesting to have this available for others (I have written with a least one person in the polars Discord about it so there is some interest). And given that this is the defacto guide on plugins I thought this would be the right place.

Below is an example of how it can be used where I have used bincode for the serde.

The impl serde::Serialize and __setstate__ are not needed bot kept for completness.

use std::collections::HashMap;

use polars::prelude::*;
use pyo3::prelude::*;
use pyo3_polars::derive::polars_expr;
use serde::{Deserialize, Serialize};

/// Creates a sample lookup table for testing purposes.
/// # Returns
/// A `MyClass` instance containing a HashMap where each key `i` maps to
/// a vector `[i*3, i*3+1, i*3+2]`

#[pyfunction]
pub fn get_lookup_table(len: usize) -> MyClass {
    let lookup_table: HashMap<i64, Vec<i64>> = (0..len as i64)
        .map(|i| {
            let start = i * 3;
            (i, vec![start, start + 1, start + 2])
        })
        .collect();
    MyClass { lookup_table }
}

/// A lookup table class exposed to Python via PyO3.
///
/// This class wraps a `HashMap<i64, Vec<i64>>` and is designed to be:
/// 1. Picklable in Python (via `__getstate__`/`__setstate__`)
/// 2. Passable through Polars plugin kwargs (via custom serde implementation)
///
/// # Serialization Strategy
///
/// Polars plugins use `serde_pickle` to deserialize kwargs on the Rust side.
/// When Python pickles the kwargs dict, objects with `__getstate__` return bytes.
/// The challenge is that `serde_pickle` expects to deserialize these bytes
/// back into the original struct.
///
/// To solve this, `MyClass` implements custom `Serialize`/`Deserialize` that:
/// - **Serialize**: Encodes the data as raw bytes using `serializer.serialize_bytes()`
/// - **Deserialize**: Expects raw bytes via `deserializer.deserialize_bytes()`
///
/// This matches how Polars' `DataFrame` handles serde (see `polars-core/src/serde/df.rs`),
/// allowing custom types to pass through the plugin kwargs system.
///
/// The `impl_polars_kwarg_class!` macro generates all the boilerplate for this.
#[pyclass(module = "mylib._internal")]
#[derive(Clone)]
pub struct MyClass {
    pub(crate) lookup_table: HashMap<i64, Vec<i64>>,
}

impl serde::Serialize for MyClass {
    fn serialize<S>(&self,serializer:S) -> Result<S::Ok,S::Error>where S:serde::Serializer,{
        let bytes = bincode::serialize(&self.lookup_table).map_err(serde::ser::Error::custom)?;
        serializer.serialize_bytes(&bytes)
    }

    }
impl <'de>serde::Deserialize<'de>for MyClass {
    fn deserialize<D>(deserializer:D) -> Result<Self,D::Error>where D:serde::Deserializer<'de>,{
        struct FieldVisitor;
        
        impl <'de>serde::de::Visitor<'de>for FieldVisitor {
            type Value = MyClass;
            fn expecting(&self,formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
                formatter.write_str(concat!("a byte array containing bincode-serialized ",stringify!(MyClass)," data"))
            }
            fn visit_bytes<E>(self,v: &[u8]) -> Result<Self::Value,E>where E:serde::de::Error,{
                let lookup_table:HashMap<i64,Vec<i64>>  = bincode::deserialize(v).map_err(serde::de::Error::custom)?;
                Ok(MyClass {
                    lookup_table
                })
            }
            fn visit_byte_buf<E>(self,v:Vec<u8>) -> Result<Self::Value,E>where E:serde::de::Error,{
                self.visit_bytes(&v)
            }
        
            }
        deserializer.deserialize_bytes(FieldVisitor)
    }

    }
#[pyo3::pymethods]
impl MyClass {
    pub fn new(lookup_table:HashMap<i64,Vec<i64>>) -> Self {
        MyClass {
            lookup_table
        }
    }
    pub fn __setstate__(&mut self,state: &pyo3::Bound<'_,pyo3::types::PyBytes>,) -> pyo3::PyResult<()>{
        self.lookup_table = bincode::deserialize(state.as_bytes()).unwrap();
        Ok(())
    }

    pub fn __getstate__<'py>(&self,py:pyo3::Python<'py>,) -> pyo3::PyResult<pyo3::Bound<'py,pyo3::types::PyBytes>>{
        Ok(pyo3::types::PyBytes::new(py, &bincode::serialize(&self.lookup_table).unwrap(),))
    }
    pub fn __getnewargs__(&self) -> pyo3::PyResult<(HashMap<i64,Vec<i64>> ,)>{
        Ok((self.lookup_table.clone(),))
    }

    }

/// Kwargs struct for `my_class_function` Polars plugin.
///
/// This struct defines the keyword arguments that can be passed to the
/// `my_class_function` expression plugin. Polars serializes kwargs as pickle
/// on the Python side and deserializes them via serde_pickle on the Rust side.
///
/// The field names must match the Python kwargs exactly.
#[derive(Deserialize, Serialize)]
pub struct KwargsMyClass {
    pub my_class: MyClass,
}


/// Polars expression plugin that uses a `MyClass` lookup table.
///
/// This function demonstrates passing a custom Rust struct through Polars
/// plugin kwargs. It retrieves a value from the lookup table and adds it
/// to each element of the input column.
///
/// # Arguments
/// * `inputs` - Array of input Series (expects one i64 column)
/// * `kwargs` - Contains the `MyClass` instance with the lookup table
///
/// # Returns
/// A new i64 Series where each value is `input[i] + lookup_table[0][0]`
#[polars_expr(output_type=Int64)]
fn my_class_function(inputs: &[Series], kwargs: KwargsMyClass) -> PolarsResult<Series> {
    let input: &Int64Chunked = inputs[0].i64()?;
    let variable_0 = *kwargs
        .my_class
        .lookup_table
        .get(&0)
        .ok_or(PolarsError::ComputeError(
            "Key 0 not found in variable_0.".into(),
        ))?
        .get(0)
        .unwrap();

    let value: ChunkedArray<Int64Type> = unary_elementwise_values(input, |k| k + variable_0);

    Ok(value.into_series())
}

With the following in src/lib.rs

#[pymodule]
fn _internal(_py: Python, m: &Bound<PyModule>) -> PyResult<()> {
    m.add("__version__", env!("CARGO_PKG_VERSION"))?;
    m.add_class::<MyClass>()?;
    m.add_function(wrap_pyfunction!(get_lookup_table, m)?)?;
    Ok(())
}

[MarcoGorelli]

hmm yeah sure, perhaps https://marcogorelli.github.io/polars-plugins-tutorial/arguments/ could become a section and this could be a subsection of it in its own page?

[thomasfrederikhoeck]

@MarcoGorelli Coolios - I will find some time in the comming weeks and make a PR :-)

[marinegor]

@thomasfrederikhoeck thanks a lot, I can confirm that it works for my toy usecase: https://github.com/marinegor/polars_counter (commit 82dd8a2ea7de44f53d20d596931281a56ffddd36 if someone interested).

only thing I struggled with was finding appropriate bincode version (I have bincode = {version = "<2.0.0"} in Cargo.toml), but the rest worked smoothly!