26694011222

Committed 30 May 2026 08:22PM UTC coverage: 94.726% (-0.009%) from 94.735%

Build # 26694011222

Build Type

push

github

Committed by

web-flow

Commit Message

perf: Use BufWriter when writing json (#1594)

* perf: Use BufWriter when writing json

Writing without a BufWriter is going to generate a syscall for every
emitted JSON token.

* fix: flush buffer after writing

* buffer other calls in the library

---------

Co-authored-by: Ivan Carvalho <ivancarvalho@gatech.edu>

Coverage Stats

29 of 32 new or added lines in 4 files covered. (90.63%)

1 existing line in 1 file now uncovered.

19148 of 20214 relevant lines covered (94.73%)

920306.39 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

85.59

/src/json/node_link_data.rs

// Licensed under the Apache License, Version 2.0 (the "License"); you may
// not use this file except in compliance with the License. You may obtain
// a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations
// under the License.

use std::collections::BTreeMap;
use std::fs::File;
use std::io::BufWriter;
use std::io::Write as _;

use hashbrown::HashMap;
use indexmap::IndexSet;

use serde::{Deserialize, Serialize};

use pyo3::IntoPyObjectExt;
use pyo3::Python;
use pyo3::prelude::*;

use petgraph::EdgeType;
use petgraph::visit::EdgeRef;
use petgraph::visit::IntoEdgeReferences;

use crate::JSONSerializationError;
use crate::NodeIndex;
use crate::StablePyGraph;

#[derive(Serialize)]
pub struct Graph {
    pub directed: bool,
    pub multigraph: bool,
    pub attrs: Option<BTreeMap<String, String>>,
    pub nodes: Vec<Node>,
    pub links: Vec<Link>,
}

#[derive(Deserialize)]
pub struct GraphInput {
    pub directed: bool,
    pub multigraph: bool,
    pub attrs: Option<BTreeMap<String, String>>,
    pub nodes: Vec<NodeInput>,
    pub links: Vec<LinkInput>,
}

#[derive(Serialize)]
pub struct Node {
    id: usize,
    data: Option<BTreeMap<String, String>>,
}

#[derive(Deserialize)]
pub struct NodeInput {
    id: Option<usize>,
    data: Option<BTreeMap<String, String>>,
}

#[derive(Deserialize)]
pub struct LinkInput {
    source: usize,
    target: usize,
    #[allow(dead_code)]
    id: Option<usize>,
    data: Option<BTreeMap<String, String>>,
}

#[derive(Serialize)]
pub struct Link {
    source: usize,
    target: usize,
    id: usize,
    data: Option<BTreeMap<String, String>>,
}

#[allow(clippy::too_many_arguments)]
pub fn parse_node_link_data<Ty: EdgeType>(
    py: &Python,
    graph: GraphInput,
    out_graph: &mut StablePyGraph<Ty>,
    node_attrs: Option<Py<PyAny>>,
    edge_attrs: Option<Py<PyAny>>,
) -> PyResult<bool> {
    let mut id_mapping: HashMap<usize, NodeIndex> = HashMap::with_capacity(graph.nodes.len());

    // Check if nodes have explicit IDs that need preservation
    let preserve_ids = graph.nodes.iter().any(|n| n.id.is_some());

    let node_removed = if !preserve_ids {
        // No explicit IDs, just add nodes sequentially (legacy behavior)
        for node in graph.nodes {
            let payload = match node.data {
                Some(data) => match node_attrs {
                    Some(ref callback) => callback.call1(*py, (data,))?,
                    None => data.into_py_any(*py)?,
                },
                None => py.None(),
            };
            let id = out_graph.add_node(payload);
            id_mapping.insert(id.index(), id);
        }
        false
    } else {
        // Find the maximum node ID to determine how many placeholder nodes we need
        let max_id = graph.nodes.iter().filter_map(|n| n.id).max().unwrap_or(0);

        // Create placeholder nodes up to max_id
        let mut tmp_nodes: IndexSet<NodeIndex> = (0..=max_id)
            .map(|_| out_graph.add_node(py.None()))
            .collect();

        // Replace placeholder nodes with actual data and track which to keep
        for node in graph.nodes {
            let payload = match node.data {
                Some(data) => match node_attrs {
                    Some(ref callback) => callback.call1(*py, (data,))?,
                    None => data.into_py_any(*py)?,
                },
                None => py.None(),
            };
            let node_id = node.id.ok_or_else(|| {
                pyo3::exceptions::PyValueError::new_err(
                    "All nodes must have explicit IDs when any node has an ID",
                )
            })?;
            let idx = NodeIndex::new(node_id);

            // Replace the placeholder with actual data
            if let Some(weight) = out_graph.node_weight_mut(idx) {
                *weight = payload;
            }

            id_mapping.insert(node_id, idx);
            tmp_nodes.swap_remove(&idx);
        }

        // Track if we're removing any nodes (indicates gaps in indices)
        let has_gaps = !tmp_nodes.is_empty();

        // Remove remaining placeholder nodes
        for tmp_node in tmp_nodes {
            out_graph.remove_node(tmp_node);
        }

        has_gaps
    };

    for edge in graph.links {
        let data = match edge.data {
            Some(data) => match edge_attrs {
                Some(ref callback) => callback.call1(*py, (data,))?,
                None => data.into_py_any(*py)?,
            },
            None => py.None(),
        };
        out_graph.add_edge(id_mapping[&edge.source], id_mapping[&edge.target], data);
    }
    Ok(node_removed)
}

#[allow(clippy::too_many_arguments)]
pub fn node_link_data<Ty: EdgeType>(
    py: Python,
    graph: &StablePyGraph<Ty>,
    multigraph: bool,
    attrs: &Py<PyAny>,
    path: Option<String>,
    graph_attrs: Option<Py<PyAny>>,
    node_attrs: Option<Py<PyAny>>,
    edge_attrs: Option<Py<PyAny>>,
) -> PyResult<Option<String>> {
    let attr_callable =
        |attrs: &Py<PyAny>, obj: &Py<PyAny>| -> PyResult<BTreeMap<String, String>> {
            let res = attrs.call1(py, (obj,))?;
            res.extract(py)
        };
    let mut nodes: Vec<Node> = Vec::with_capacity(graph.node_count());
    for n in graph.node_indices() {
        let data = match node_attrs {
            Some(ref callback) => Some(attr_callable(callback, &graph[n])?),
            None => None,
        };
        nodes.push(Node {
            id: n.index(),
            data,
        });
    }
    let mut links: Vec<Link> = Vec::with_capacity(graph.edge_count());
    for e in graph.edge_references() {
        let data = match edge_attrs {
            Some(ref callback) => Some(attr_callable(callback, e.weight())?),
            None => None,
        };
        links.push(Link {
            source: e.source().index(),
            target: e.target().index(),
            id: e.id().index(),
            data,
        });
    }

    let graph_attrs = match graph_attrs {
        Some(ref callback) => Some(attr_callable(callback, attrs)?),
        None => None,
    };

    let output_struct = Graph {
        directed: graph.is_directed(),
        multigraph,
        attrs: graph_attrs,
        nodes,
        links,
    };
    match path {
        None => match serde_json::to_string(&output_struct) {
            Ok(v) => Ok(Some(v)),
            Err(e) => Err(JSONSerializationError::new_err(format!("JSON Error: {e}"))),
        },
        Some(filename) => {
            let file = File::create(filename)?;
            let mut buffer = BufWriter::new(file);
            let res = match serde_json::to_writer(&mut buffer, &output_struct) {
                Ok(_) => {
                    buffer.flush()?;
                    Ok(None)
                }
                Err(e) => Err(JSONSerializationError::new_err(format!("JSON Error: {e}"))),
            }?;

            Ok(res)
        }
    }
}

1	// Licensed under the Apache License, Version 2.0 (the "License"); you may
2	// not use this file except in compliance with the License. You may obtain
3	// a copy of the License at
4	//
5	// http://www.apache.org/licenses/LICENSE-2.0
6	//
7	// Unless required by applicable law or agreed to in writing, software
8	// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
9	// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
10	// License for the specific language governing permissions and limitations
11	// under the License.
12
13	use std::collections::BTreeMap;
14	use std::fs::File;
15	use std::io::BufWriter;
16	use std::io::Write as _;
17
18	use hashbrown::HashMap;
19	use indexmap::IndexSet;
20
21	use serde::{Deserialize, Serialize};
22
23	use pyo3::IntoPyObjectExt;
24	use pyo3::Python;
25	use pyo3::prelude::*;
26
27	use petgraph::EdgeType;
28	use petgraph::visit::EdgeRef;
29	use petgraph::visit::IntoEdgeReferences;
30
31	use crate::JSONSerializationError;
32	use crate::NodeIndex;
33	use crate::StablePyGraph;
34
35	#[derive(Serialize)]
36	pub struct Graph {
37	pub directed: bool,
38	pub multigraph: bool,
39	pub attrs: Option<BTreeMap<String, String>>,
40	pub nodes: Vec<Node>,
41	pub links: Vec<Link>,
42	}
43
44	#[derive(Deserialize)]
45	pub struct GraphInput {
46	pub directed: bool,
47	pub multigraph: bool,
48	pub attrs: Option<BTreeMap<String, String>>,
49	pub nodes: Vec<NodeInput>,
50	pub links: Vec<LinkInput>,
51	}
52
53	#[derive(Serialize)]
54	pub struct Node {
55	id: usize,
56	data: Option<BTreeMap<String, String>>,
57	}
58
59	#[derive(Deserialize)]
60	pub struct NodeInput {
61	id: Option<usize>,
62	data: Option<BTreeMap<String, String>>,
63	}
64
65	#[derive(Deserialize)]
66	pub struct LinkInput {
67	source: usize,
68	target: usize,
69	#[allow(dead_code)]
70	id: Option<usize>,
71	data: Option<BTreeMap<String, String>>,
72	}
73
74	#[derive(Serialize)]
75	pub struct Link {
76	source: usize,
77	target: usize,
78	id: usize,
79	data: Option<BTreeMap<String, String>>,
80	}
81
82	#[allow(clippy::too_many_arguments)]
83	pub fn parse_node_link_data<Ty: EdgeType>(	26✔
84	py: &Python,	26✔
85	graph: GraphInput,	26✔
86	out_graph: &mut StablePyGraph<Ty>,	26✔
87	node_attrs: Option<Py<PyAny>>,	26✔
88	edge_attrs: Option<Py<PyAny>>,	26✔
89	) -> PyResult<bool> {	26✔
90	let mut id_mapping: HashMap<usize, NodeIndex> = HashMap::with_capacity(graph.nodes.len());	26✔
91
92	// Check if nodes have explicit IDs that need preservation
93	let preserve_ids = graph.nodes.iter().any(\|n\| n.id.is_some());	26✔
94
95	let node_removed = if !preserve_ids {	26✔
96	// No explicit IDs, just add nodes sequentially (legacy behavior)
97	for node in graph.nodes {	×
98	let payload = match node.data {	×
99	Some(data) => match node_attrs {	×
100	Some(ref callback) => callback.call1(*py, (data,))?,	×
101	None => data.into_py_any(*py)?,	×
102	},
103	None => py.None(),	×
104	};
105	let id = out_graph.add_node(payload);	×
106	id_mapping.insert(id.index(), id);	×
107	}
108	false	×
109	} else {
110	// Find the maximum node ID to determine how many placeholder nodes we need
111	let max_id = graph.nodes.iter().filter_map(\|n\| n.id).max().unwrap_or(0);	26✔
112
113	// Create placeholder nodes up to max_id
114	let mut tmp_nodes: IndexSet<NodeIndex> = (0..=max_id)	26✔
115	.map(\|_\| out_graph.add_node(py.None()))	3,856✔
116	.collect();	26✔
117
118	// Replace placeholder nodes with actual data and track which to keep
119	for node in graph.nodes {	3,840✔
120	let payload = match node.data {	3,840✔
121	Some(data) => match node_attrs {	3,544✔
122	Some(ref callback) => callback.call1(*py, (data,))?,	×
123	None => data.into_py_any(*py)?,	3,544✔
124	},
125	None => py.None(),	296✔
126	};
127	let node_id = node.id.ok_or_else(\|\| {	3,840✔
128	pyo3::exceptions::PyValueError::new_err(	×
129	"All nodes must have explicit IDs when any node has an ID",
130	)
131	})?;	×
132	let idx = NodeIndex::new(node_id);	3,840✔
133
134	// Replace the placeholder with actual data
135	if let Some(weight) = out_graph.node_weight_mut(idx) {	3,840✔
136	*weight = payload;	3,840✔
137	}	3,840✔
138
139	id_mapping.insert(node_id, idx);	3,840✔
140	tmp_nodes.swap_remove(&idx);	3,840✔
141	}
142
143	// Track if we're removing any nodes (indicates gaps in indices)
144	let has_gaps = !tmp_nodes.is_empty();	26✔
145
146	// Remove remaining placeholder nodes
147	for tmp_node in tmp_nodes {	26✔
148	out_graph.remove_node(tmp_node);	16✔
149	}	16✔
150
151	has_gaps	26✔
152	};
153
154	for edge in graph.links {	4,458✔
155	let data = match edge.data {	4,458✔
156	Some(data) => match edge_attrs {	4,184✔
157	Some(ref callback) => callback.call1(*py, (data,))?,	×
158	None => data.into_py_any(*py)?,	4,184✔
159	},
160	None => py.None(),	274✔
161	};
162	out_graph.add_edge(id_mapping[&edge.source], id_mapping[&edge.target], data);	4,458✔
163	}
164	Ok(node_removed)	26✔
165	}	26✔
166
167	#[allow(clippy::too_many_arguments)]
168	pub fn node_link_data<Ty: EdgeType>(	58✔
169	py: Python,	58✔
170	graph: &StablePyGraph<Ty>,	58✔
171	multigraph: bool,	58✔
172	attrs: &Py<PyAny>,	58✔
173	path: Option<String>,	58✔
174	graph_attrs: Option<Py<PyAny>>,	58✔
175	node_attrs: Option<Py<PyAny>>,	58✔
176	edge_attrs: Option<Py<PyAny>>,	58✔
177	) -> PyResult<Option<String>> {	58✔
178	let attr_callable =	58✔
179	\|attrs: &Py<PyAny>, obj: &Py<PyAny>\| -> PyResult<BTreeMap<String, String>> {	7,788✔
180	let res = attrs.call1(py, (obj,))?;	7,788✔
181	res.extract(py)	7,788✔
182	};	7,788✔
183	let mut nodes: Vec<Node> = Vec::with_capacity(graph.node_count());	58✔
184	for n in graph.node_indices() {	3,868✔
185	let data = match node_attrs {	3,868✔
186	Some(ref callback) => Some(attr_callable(callback, &graph[n])?),	3,568✔
187	None => None,	300✔
188	};
189	nodes.push(Node {	3,868✔
190	id: n.index(),	3,868✔
191	data,	3,868✔
192	});	3,868✔
193	}
194	let mut links: Vec<Link> = Vec::with_capacity(graph.edge_count());	58✔
195	for e in graph.edge_references() {	4,474✔
196	let data = match edge_attrs {	4,474✔
197	Some(ref callback) => Some(attr_callable(callback, e.weight())?),	4,200✔
198	None => None,	274✔
199	};
200	links.push(Link {	4,474✔
201	source: e.source().index(),	4,474✔
202	target: e.target().index(),	4,474✔
203	id: e.id().index(),	4,474✔
204	data,	4,474✔
205	});	4,474✔
206	}
207
208	let graph_attrs = match graph_attrs {	58✔
209	Some(ref callback) => Some(attr_callable(callback, attrs)?),	20✔
210	None => None,	38✔
211	};
212
213	let output_struct = Graph {	54✔
214	directed: graph.is_directed(),	54✔
215	multigraph,	54✔
216	attrs: graph_attrs,	54✔
217	nodes,	54✔
218	links,	54✔
219	};	54✔
220	match path {	54✔
221	None => match serde_json::to_string(&output_struct) {	38✔
222	Ok(v) => Ok(Some(v)),	38✔
223	Err(e) => Err(JSONSerializationError::new_err(format!("JSON Error: {e}"))),	×
224	},
225	Some(filename) => {	16✔
226	let file = File::create(filename)?;	16✔
227	let mut buffer = BufWriter::new(file);	12✔
228	let res = match serde_json::to_writer(&mut buffer, &output_struct) {	12✔
229	Ok(_) => {
230	buffer.flush()?;	12✔
231	Ok(None)	12✔
232	}
UNCOV 233	Err(e) => Err(JSONSerializationError::new_err(format!("JSON Error: {e}"))),	×
NEW 234	}?;	×
235
236	Ok(res)	12✔
237	}
238	}
239	}	58✔

Qiskit / rustworkx / 26694011222

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous