Skip to content

Give example of making the plugin cross platform (rust and python) #63

@deanm0000

Description

@deanm0000

By way of example:

pub trait SumI64Trait {
    fn sum_i64(self, other: Expr)-> Expr
}

impl SumI64Trait for Expr {
    fn sum_i64(self, other: Expr) -> Expr {
        as_struct(vec![self, other]).map(
            |s| {
                let ca=s.struct_().unwrap();
                let inputs=ca.fields_as_series();
                let inputs_slice = inputs.as_slice();
                let out = sum_i64_common(inputs_slice).unwrap().into_column();
                Ok(Some(out))
            },
            GetOutput::from_type(DataType::Int64)
        )
    }
}

fn sum_i64_common(inputs: &[Series]) -> PolarsResult<Series> {
    let left: &Int64Chunked = inputs[0].i64()?;
    let right: &Int64Chunked = inputs[1].i64()?;
    // Note: there's a faster way of summing two columns, see
    // section 7.
    let out: Int64Chunked = broadcast_binary_elementwise(
        left,
        right,
        |left: Option<i64>, right: Option<i64>| match (left, right) {
            (Some(left), Some(right)) => Some(left + right),
            _ => None,
        },
    );
    Ok(out.into_series())
}


#[polars_expr(output_type=Int64)]
fn sum_i64_expr(inputs: &[Series]) -> PolarsResult<Series> {
    sum_i64_common(inputs)
}

That way people can use the plugin on the rust side with col("a").sum_i64(col("b"))

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions