1use hawdb_core::{LabelId, Value};
18use hawdb_plan_cypher::SortDirection;
19use hawdb_storage::{NodeRecord, RelRecord};
20use std::cmp::Ordering;
21use std::collections::{BTreeMap, BTreeSet};
22
23#[derive(Debug, Clone, PartialEq, Eq)]
24pub struct Binding {
25 pub values: BTreeMap<String, Value>,
26 pub nodes: BTreeMap<String, NodeRecord>,
27 pub relationships: BTreeMap<String, RelRecord>,
28}
29
30impl Binding {
31 pub fn values(values: BTreeMap<String, Value>) -> Self {
32 Self {
33 values,
34 nodes: BTreeMap::new(),
35 relationships: BTreeMap::new(),
36 }
37 }
38
39 pub fn scalar(name: impl Into<String>, value: Value) -> Self {
40 Self::values(BTreeMap::from([(name.into(), value)]))
41 }
42}
43
44#[derive(Debug, Clone, PartialEq, Eq)]
45pub struct TopNBinding {
46 pub sort_values: Vec<(Value, SortDirection)>,
47 pub ordinal: u64,
48 pub binding: Binding,
49}
50
51impl TopNBinding {
52 pub fn memory_bytes(&self) -> usize {
53 binding_memory_bytes(&self.binding).saturating_add(self.sort_values.iter().fold(
54 std::mem::size_of::<Vec<(Value, SortDirection)>>(),
55 |total, (value, _)| total.saturating_add(value_memory_bytes(value)),
56 ))
57 }
58}
59
60impl Ord for TopNBinding {
61 fn cmp(&self, other: &Self) -> Ordering {
62 for ((left, direction), (right, other_direction)) in
63 self.sort_values.iter().zip(&other.sort_values)
64 {
65 debug_assert_eq!(direction, other_direction);
66 let ordering = match direction {
67 SortDirection::Asc => left.cmp(right),
68 SortDirection::Desc => left.cmp(right).reverse(),
69 };
70 if ordering != Ordering::Equal {
71 return ordering;
72 }
73 }
74 self.ordinal.cmp(&other.ordinal)
75 }
76}
77
78impl PartialOrd for TopNBinding {
79 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
80 Some(self.cmp(other))
81 }
82}
83
84pub fn map_payload_bytes(values: &BTreeMap<String, Value>) -> usize {
85 values.iter().fold(0usize, |total, (name, value)| {
86 total
87 .saturating_add(name.len())
88 .saturating_add(value_payload_bytes(value))
89 })
90}
91
92pub fn map_memory_bytes(values: &BTreeMap<String, Value>) -> usize {
93 map_memory_bytes_with_entries(values.iter().map(|(name, value)| (name.as_str(), value)))
94}
95
96fn map_memory_bytes_with_entries<'a>(values: impl Iterator<Item = (&'a str, &'a Value)>) -> usize {
97 std::mem::size_of::<BTreeMap<String, Value>>().saturating_add(values.fold(
98 0usize,
99 |total, (name, value)| {
100 total
101 .saturating_add(std::mem::size_of::<(String, Value)>() * 3)
102 .saturating_add(name.len())
103 .saturating_add(value_memory_bytes(value))
104 },
105 ))
106}
107
108pub(crate) fn projected_node_binding_memory_bytes(
111 variable: &str,
112 node: &NodeRecord,
113 required_properties: &BTreeSet<String>,
114) -> usize {
115 std::mem::size_of::<Binding>()
116 .saturating_add(variable.len())
117 .saturating_add(std::mem::size_of_val(&node.id))
118 .saturating_add(
119 node.labels
120 .len()
121 .saturating_mul(std::mem::size_of::<LabelId>()),
122 )
123 .saturating_add(map_memory_bytes_with_entries(
124 required_properties.iter().filter_map(|property| {
125 node.properties
126 .get(property)
127 .map(|value| (property.as_str(), value))
128 }),
129 ))
130 .saturating_add(std::mem::size_of::<usize>() * 6)
131}
132
133pub fn value_memory_bytes(value: &Value) -> usize {
134 std::mem::size_of::<Value>().saturating_add(match value {
135 Value::Null | Value::Bool(_) | Value::Int(_) | Value::Float(_) => 0,
136 Value::String(value) => value.len(),
137 Value::Binary(value) => value.len(),
138 Value::Uuid(_) => 16,
139 Value::List(values) => values
140 .iter()
141 .fold(std::mem::size_of::<Vec<Value>>(), |total, value| {
142 total.saturating_add(value_memory_bytes(value))
143 }),
144 Value::Map(values) => map_memory_bytes(values),
145 })
146}
147
148pub fn binding_payload_bytes(binding: &Binding) -> usize {
149 map_payload_bytes(&binding.values).saturating_add(graph_binding_payload_bytes(binding))
150}
151
152fn graph_binding_payload_bytes(binding: &Binding) -> usize {
153 graph_binding_bytes(binding, map_payload_bytes)
154}
155
156fn graph_binding_bytes(
157 binding: &Binding,
158 property_bytes: fn(&BTreeMap<String, Value>) -> usize,
159) -> usize {
160 graph_binding_parts_bytes(
161 binding
162 .nodes
163 .iter()
164 .map(|(name, node)| (name.as_str(), node)),
165 binding
166 .relationships
167 .iter()
168 .map(|(name, row)| (name.as_str(), row)),
169 property_bytes,
170 )
171}
172
173fn graph_binding_parts_bytes<'a>(
174 nodes: impl Iterator<Item = (&'a str, &'a NodeRecord)>,
175 relationships: impl Iterator<Item = (&'a str, &'a RelRecord)>,
176 property_bytes: fn(&BTreeMap<String, Value>) -> usize,
177) -> usize {
178 nodes
179 .fold(0usize, |total, (name, node)| {
180 total
181 .saturating_add(name.len())
182 .saturating_add(std::mem::size_of_val(&node.id))
183 .saturating_add(
184 node.labels
185 .len()
186 .saturating_mul(std::mem::size_of::<LabelId>()),
187 )
188 .saturating_add(property_bytes(&node.properties))
189 })
190 .saturating_add(relationships.fold(0usize, |total, (name, row)| {
191 total
192 .saturating_add(name.len())
193 .saturating_add(std::mem::size_of_val(&row.id))
194 .saturating_add(std::mem::size_of_val(&row.source))
195 .saturating_add(std::mem::size_of_val(&row.target))
196 .saturating_add(std::mem::size_of_val(&row.rel_type))
197 .saturating_add(property_bytes(&row.properties))
198 }))
199}
200
201pub(crate) fn binding_payload_bytes_with_parts<'a>(
202 values: impl Iterator<Item = (&'a str, &'a Value)>,
203 nodes: impl Iterator<Item = (&'a str, &'a NodeRecord)>,
204 relationships: impl Iterator<Item = (&'a str, &'a RelRecord)>,
205) -> usize {
206 values
207 .fold(0usize, |total, (name, value)| {
208 total
209 .saturating_add(name.len())
210 .saturating_add(value_payload_bytes(value))
211 })
212 .saturating_add(graph_binding_parts_bytes(
213 nodes,
214 relationships,
215 map_payload_bytes,
216 ))
217}
218
219pub fn binding_memory_bytes(binding: &Binding) -> usize {
220 binding_memory_bytes_with_values(
221 binding,
222 binding
223 .values
224 .iter()
225 .map(|(name, value)| (name.as_str(), value)),
226 )
227}
228
229pub(crate) fn binding_memory_bytes_replacing_value(
232 binding: &Binding,
233 name: &str,
234 value: &Value,
235) -> usize {
236 binding_memory_bytes_with_values(
237 binding,
238 binding
239 .values
240 .iter()
241 .filter(|(key, _)| key.as_str() != name)
242 .map(|(key, value)| (key.as_str(), value))
243 .chain(std::iter::once((name, value))),
244 )
245}
246
247pub(crate) fn binding_memory_bytes_with_values<'a>(
250 binding: &'a Binding,
251 values: impl Iterator<Item = (&'a str, &'a Value)>,
252) -> usize {
253 binding_memory_bytes_with_parts(
254 values,
255 binding
256 .nodes
257 .iter()
258 .map(|(name, node)| (name.as_str(), node)),
259 binding
260 .relationships
261 .iter()
262 .map(|(name, row)| (name.as_str(), row)),
263 )
264}
265
266pub(crate) fn binding_memory_bytes_with_parts<'a>(
267 values: impl Iterator<Item = (&'a str, &'a Value)>,
268 nodes: impl Iterator<Item = (&'a str, &'a NodeRecord)>,
269 relationships: impl Iterator<Item = (&'a str, &'a RelRecord)>,
270) -> usize {
271 let entries = std::cell::Cell::new(0usize);
272 let graph_bytes = graph_binding_parts_bytes(
273 nodes.inspect(|_| entries.set(entries.get().saturating_add(1))),
274 relationships.inspect(|_| entries.set(entries.get().saturating_add(1))),
275 map_memory_bytes,
276 );
277 std::mem::size_of::<Binding>()
278 .saturating_add(graph_bytes)
279 .saturating_add(
280 entries
281 .get()
282 .saturating_mul(std::mem::size_of::<usize>() * 6),
283 )
284 .saturating_add(values.fold(0usize, |total, (name, value)| {
285 total
286 .saturating_add(name.len())
287 .saturating_add(match value {
288 Value::List(_) | Value::Map(_) => {
289 value_memory_bytes(value).saturating_sub(std::mem::size_of::<Value>())
290 }
291 _ => value_payload_bytes(value),
292 })
293 .saturating_add(std::mem::size_of::<usize>() * 6)
294 }))
295}
296
297pub fn node_memory_bytes(node: &NodeRecord) -> usize {
298 std::mem::size_of::<NodeRecord>()
299 .saturating_add(
300 node.labels
301 .len()
302 .saturating_mul(std::mem::size_of::<LabelId>() * 3),
303 )
304 .saturating_add(map_memory_bytes(&node.properties))
305}
306
307pub fn relationship_memory_bytes(relationship: &RelRecord) -> usize {
308 std::mem::size_of::<RelRecord>().saturating_add(map_memory_bytes(&relationship.properties))
309}
310
311pub fn value_payload_bytes(value: &Value) -> usize {
312 match value {
313 Value::Null => 0,
314 Value::Bool(_) => std::mem::size_of::<bool>(),
315 Value::Int(_) => std::mem::size_of::<i64>(),
316 Value::Float(_) => std::mem::size_of::<f64>(),
317 Value::String(value) => value.len(),
318 Value::Binary(value) => value.len(),
319 Value::Uuid(_) => 16,
320 Value::List(values) => values.iter().fold(0usize, |total, value| {
321 total.saturating_add(value_payload_bytes(value))
322 }),
323 Value::Map(values) => map_payload_bytes(values),
324 }
325}
326
327#[cfg(test)]
328mod tests {
329 use super::*;
330 use hawdb_core::RelTypeId;
331 use hawdb_storage::{NodeId, RelId};
332 use std::collections::BTreeSet;
333
334 #[test]
335 fn value_only_constructors_leave_graph_bindings_empty() {
336 let values = Binding::values(BTreeMap::from([("left".to_string(), Value::Int(1))]));
337 assert_eq!(values.values["left"], Value::Int(1));
338 assert!(values.nodes.is_empty());
339 assert!(values.relationships.is_empty());
340
341 let scalar = Binding::scalar("right", Value::Bool(true));
342 assert_eq!(scalar.values["right"], Value::Bool(true));
343 assert!(scalar.nodes.is_empty());
344 assert!(scalar.relationships.is_empty());
345 }
346
347 #[test]
348 fn binding_accounting_includes_nested_values_and_graph_records() {
349 let binding = Binding {
350 values: BTreeMap::from([(
351 "nested".to_string(),
352 Value::List(vec![Value::Int(1), Value::String("two".to_string())]),
353 )]),
354 nodes: BTreeMap::from([(
355 "n".to_string(),
356 NodeRecord {
357 id: NodeId(7),
358 labels: BTreeSet::from([LabelId(3)]),
359 properties: BTreeMap::new(),
360 },
361 )]),
362 relationships: BTreeMap::from([(
363 "r".to_string(),
364 RelRecord {
365 id: RelId(11),
366 source: NodeId(7),
367 target: NodeId(8),
368 rel_type: RelTypeId(4),
369 properties: BTreeMap::new(),
370 },
371 )]),
372 };
373
374 assert!(binding_payload_bytes(&binding) > "nested".len() + "two".len());
375 assert!(binding_memory_bytes(&binding) > binding_payload_bytes(&binding));
376 }
377
378 #[test]
379 fn replacement_admission_matches_final_rows_without_mutating_input() {
380 let column = hawdb_plan_cypher::SCORING_RERANK_SCORE_COLUMN;
381 for previous in [
382 None,
383 Some(Value::Null),
384 Some(Value::Float(-1.0)),
385 Some(Value::List(vec![Value::Map(BTreeMap::from([(
386 "nested".into(),
387 Value::String("x".repeat(1_024)),
388 )]))])),
389 ] {
390 let mut input = Binding::scalar("score", Value::Float(0.5));
391 if let Some(previous) = previous {
392 input.values.insert(column.into(), previous);
393 }
394 let original = input.clone();
395 let estimate =
396 binding_memory_bytes_replacing_value(&input, column, &Value::Float(0.25));
397 assert_eq!(input, original, "admission must not alter the candidate");
398 let mut output = input;
399 output.values.insert(column.into(), Value::Float(0.25));
400 assert_eq!(estimate, binding_memory_bytes(&output));
401 }
402 }
403}