Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Add binary_elementwise_into_string_amortized #17892

Closed
wants to merge 1 commit into from

Conversation

condekind
Copy link
Contributor

We already have apply_into_string_amortized and binary_elementwise. The contributed function is a specialized version that combines the logic of these two. Here's an example of how it can simplify some scenarios. This code is from a plugin that performs reverse geocoding:

#[polars_expr(output_type=String)]
fn reverse_geocode(inputs: &[Series]) -> PolarsResult<Series> {
    let lat = inputs[0].f64()?;
    let lon = inputs[1].f64()?;
    let geocoder = ReverseGeocoder::new();

    let (lhs, rhs) = align_chunks_binary(lat, lon);
    let chunks = lhs
        .downcast_iter()
        .zip(rhs.downcast_iter())
        .map(|(lat_arr, lon_arr)| {
            let mut mutarr = MutablePlString::with_capacity(lat_arr.len());

            for (lat_opt_val, lon_opt_val) in lat_arr.iter().zip(lon_arr.iter()) {
                match (lat_opt_val, lon_opt_val) {
                    (Some(lat_val), Some(lon_val)) => {
                        let res = &geocoder.search((*lat_val, *lon_val)).record.name;
                        mutarr.push(Some(res))
                    }
                    _ => mutarr.push_null(),
                }
            }

            mutarr.freeze().boxed()
        })
        .collect();
    let out: StringChunked = unsafe { ChunkedArray::from_chunks("placeholder", chunks) };
    Ok(out.into_series())
}

Simplified version with binary_elementwise_into_string_amortized:

#[polars_expr(output_type=String)]
fn reverse_geocode(inputs: &[Series]) -> PolarsResult<Series> {
    let lat = inputs[0].f64()?;
    let lon = inputs[1].f64()?;
    let geocoder = ReverseGeocoder::new();


    let out: StringChunked =
        binary_elementwise_into_string_amortized(lat, lon, |lhs, rhs, output| {
            let _ = write!(output, "{}", &geocoder.search((lhs, rhs)).record.name);
        });

    Ok(out.into_series())
}

@condekind condekind changed the title feat: add binary_elementwise_into_string_amortized feat: Add binary_elementwise_into_string_amortized Jul 26, 2024
@github-actions github-actions bot added enhancement New feature or an improvement of an existing feature python Related to Python Polars rust Related to Rust Polars and removed title needs formatting labels Jul 26, 2024
Copy link
Collaborator

@MarcoGorelli MarcoGorelli left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks @condekind - yup, this is exactly what I was suggesting, nice one

The function's currently unused in Polars - to prevent someone accidentally removing it, perhaps add a little docstring to it, in which you mention that it's intended as a utility for plugin authors?

@condekind
Copy link
Contributor Author

Closing this since I didn't realize I opened the PR from the main branch. Will reopen from a dedicated branch and add the suggested docstrings, thank you @MarcoGorelli!

@condekind condekind closed this Jul 26, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or an improvement of an existing feature python Related to Python Polars rust Related to Rust Polars
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants