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