4
\$\begingroup\$

I am working on converting some python code over to Rust, and I have come across a bit of a peculiarity in the way that my code is behaving. Namely, the module that I have written in Rust is much slower than the same code written in Python. I originally thought that this was due to the fact that the module that I was writing had a lot of overhead from converting the Python list object to a Rust vector but now I am not so sure. Even when I scale this for large grid graphs (400x400 or larger), the overhead seems to just scale with it, so there might be something else wrong with the code. Here is the bit that appears to be causing the issue:

use pyo3::prelude::*;
use pyo3::wrap_pyfunction;
use rand::seq::SliceRandom;
use rand::thread_rng;
use std::collections::HashMap;
use pyo3::types::PyList;


fn current_component(n: usize, component_merge_dict: &mut HashMap<usize, usize>) -> usize
{
    let mut nodeid = n;
    let mut nodeids_to_update = Vec::with_capacity(component_merge_dict.len());
    while let Some(&next_nodeid) = component_merge_dict.get(&nodeid)
    {
        if nodeid == next_nodeid { break; }
        nodeids_to_update.push(nodeid);
        nodeid = next_nodeid;
    }
    for nid in nodeids_to_update
    {
        component_merge_dict.insert(nid, nodeid);
    }
    nodeid
}

#[pyfunction]
fn rand_kruskal_memo(_py: Python, 
                py_node_list: &PyList, 
                py_edge_list: &PyList) -> PyResult<Vec<((i32, i32), (i32, i32))>>
{
    let node_list: Vec<(i32, i32)> = py_node_list.extract()?;
    let edge_list: Vec<((i32, i32), (i32, i32))> = py_edge_list.extract()?;


    let mut tree_edge_list: Vec<((i32, i32), (i32, i32))> = Vec::with_capacity(node_list.len() - 1);

    let mut edge_indices: Vec<usize> = (0..edge_list.len()).collect();

    edge_indices.shuffle(&mut rand::thread_rng());
    
    let mut nodes_to_components_dict: HashMap<&(i32, i32), usize> = HashMap::new();
    let mut component_merge_dict: HashMap<usize, usize> = HashMap::new();

    node_list.iter().enumerate().into_iter().for_each(|(index, node)| {
        nodes_to_components_dict.insert(node, index);
        component_merge_dict.insert(index, index);
    });
    

    let mut num_components: usize = node_list.len();

    let mut curr: usize = 0;

    while num_components > 1
    {
        let this_edge: ((i32, i32), (i32, i32)) = edge_list[edge_indices[curr]];
        curr += 1;

        let component_1: &usize = nodes_to_components_dict.get(&this_edge.0).unwrap();
        let component_num_1: usize = current_component(*component_1, &mut component_merge_dict);
        let component_2: &usize = nodes_to_components_dict.get(&this_edge.1).unwrap();
        let component_num_2: usize = current_component(*component_2, &mut component_merge_dict);

        if component_num_1 != component_num_2
        {
            component_merge_dict.insert(component_num_1, component_num_2);
            
            tree_edge_list.push(this_edge);

            num_components -= 1;
        }

    }

    Ok(tree_edge_list)
}

#[pymodule]
fn rusty_tree(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(rand_kruskal_memo, m)?)?;
    Ok(())
}

I have also implemented a Union-Find method that seems to be suffering from the same issue, and I don't understand it. I also implemented both of these methods in C++ using Pybind11, and that code does not seem to have the conversion overhead issue that I am seeing here, so I am a bit confused as to what is going on. Admittedly, I am mainly a Python and C++ developer, so it is possible that I am just not accustomed to working with Rust quite yet and there is a simple fix that I am not aware of. Regardless, if someone wouldn't mind going over this and either telling me where I went wrong or telling me a more efficient way to define the bindings between Rust and python objects, I would greatly appreciate it. Thank you!

\$\endgroup\$
1

1 Answer 1

7
\$\begingroup\$

Rust defaults to a safer but slower hash function implementation. This page discusses the issue. Basically, Rust's function is designed to insure that you don't have too many collisions, even if a hostile party is trying to mess with your program. However, most of us aren't really in that situation and can benefit from a faster hash function. Python, on the other hand, relies really heavily on its hash function and makes sure that it is really fast. I suspect that your code spends almost all of its time looking up hashes and thus the hash lookup dominates the differences between Python and Rust.

You can solve this by using crates like rustc-hash, fnv, and ahash which provide alternatives to the standard hashmap which use faster hash functions. Alternately, you can try to restructure the algorithm to index into Vecs which will be much faster.

Also, make sure you are compiling in release mode.

\$\endgroup\$
4
  • \$\begingroup\$ Ah, that hash alternative was exactly what I was looking for. This thing was actually slowing down other sections of my code that I didn't include here as well. Thank you! \$\endgroup\$
    – peabody
    Jul 13, 2023 at 17:18
  • \$\begingroup\$ AFAIK, Python also uses SipHash in dict, just like Rust. It switched to it early on to prevent Hash Flooding in web frameworks such as Django. It may be a different configuration, and it may be useful to switch Rust code to another hash function regardless -- I personally like the fxhash crate -- but it doesn't explain why the Rust code would suffer unduly from SipHash when the Python code doesn't. \$\endgroup\$ Jul 14, 2023 at 7:11
  • 1
    \$\begingroup\$ @MatthieuM. Based on what I see here peps.python.org/pep-0456 it looks like Python doesn't use siphash for everything, just bytes and strings. \$\endgroup\$ Jul 14, 2023 at 13:46
  • 1
    \$\begingroup\$ @WinstonEwert: TIL! Clever of them, Hash Flooding was an issue in web framework parsing HTTP parameters & HTTP headers in hash-maps, in which case the keys would be bytes or strings. Seems fraught with peril to ignore other cases, but that PEP fixed the issue for 99.9% of affected users I'd bet. \$\endgroup\$ Jul 14, 2023 at 14:17

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.