Verifying Why Python Rust Module is Running Slow

Question

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!

Answer 1

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.