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!