8. I'd like to have an argument, please¶
Say you want to rewrite
def add_suffix(s, *, suffix):
return s + suffix
s.map_elements(lambda x: add_suffix(x, suffix='-billy'))
We've covered passing in extra columns, but...how about passing extra keyword arguments?
We'll do this with kwargs
. In minimal_plugin/__init__.py
, add the
following:
def add_suffix(expr: IntoExprColumn, *, suffix: str) -> pl.Expr:
return register_plugin_function(
args=[expr],
plugin_path=LIB,
function_name="add_suffix",
is_elementwise=True,
kwargs={"suffix": suffix},
)
In src/expressions.rs
, we'll then first have to define a struct to hold
our keyword-arguments:
Then, we can just pass an argument of this type to a add_suffix
function,
which is going to be very similar to the good version of pig_latinnify
:
#[polars_expr(output_type=String)]
fn add_suffix(inputs: &[Series], kwargs: AddSuffixKwargs) -> PolarsResult<Series> {
let s = &inputs[0];
let ca = s.str()?;
let out = ca.apply_into_string_amortized(|value, output| {
write!(output, "{}{}", value, kwargs.suffix).unwrap();
});
Ok(out.into_series())
}
To see it in action, compile with maturin develop
(or maturin develop --release
if you're
benchmarking), and then you should be able to put
import polars as pl
import minimal_plugin as mp
df = pl.DataFrame({'a': ['bob', 'billy']})
print(df.with_columns(mp.add_suffix('a', suffix='-billy')))
run.py
, and run it to get
You can add multiple keyword-arguments in the same function, just make sure to
include them in the struct which you define on the Rust side.