Skip to main content

nodedb_query/msgpack_scan/
group_key.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Binary group key construction for GROUP BY on raw MessagePack documents.
4//!
5//! Builds a deterministic string key from field values extracted directly
6//! from msgpack bytes, avoiding full document decode.
7
8use crate::msgpack_scan::field::extract_field;
9use crate::msgpack_scan::index::FieldIndex;
10use crate::msgpack_scan::reader::{read_f64, read_i64, read_null, read_str};
11
12/// Build a GROUP BY key string from raw msgpack bytes.
13///
14/// Format: `[val1,val2,...]` where values are formatted as JSON literals.
15/// This is compatible with the legacy `sonic_rs::to_string(&key_parts)` format,
16/// so the result-construction code can parse it back with `sonic_rs::from_str`.
17pub fn build_group_key(doc: &[u8], group_fields: &[String]) -> String {
18    if group_fields.is_empty() {
19        return "__all__".to_string();
20    }
21
22    let mut key_buf = String::new();
23    key_buf.push('[');
24    for (i, field) in group_fields.iter().enumerate() {
25        if i > 0 {
26            key_buf.push(',');
27        }
28        append_field_value(&mut key_buf, doc, field);
29    }
30    key_buf.push(']');
31    key_buf
32}
33
34/// Build a GROUP BY key using a pre-built `FieldIndex` for O(1) lookups.
35pub fn build_group_key_indexed(doc: &[u8], group_fields: &[String], idx: &FieldIndex) -> String {
36    if group_fields.is_empty() {
37        return "__all__".to_string();
38    }
39
40    let mut key_buf = String::new();
41    key_buf.push('[');
42    for (i, field) in group_fields.iter().enumerate() {
43        if i > 0 {
44            key_buf.push(',');
45        }
46        let range = idx.get(field);
47        append_field_value_range(&mut key_buf, doc, range);
48    }
49    key_buf.push(']');
50    key_buf
51}
52
53/// Append a single field's value to the key buffer as a JSON literal.
54fn append_field_value(buf: &mut String, doc: &[u8], field: &str) {
55    let Some((start, end)) = extract_field(doc, 0, field) else {
56        buf.push_str("null");
57        return;
58    };
59
60    if read_null(doc, start) {
61        buf.push_str("null");
62    } else if let Some(s) = read_str(doc, start) {
63        buf.push('"');
64        buf.push_str(s);
65        buf.push('"');
66    } else if let Some(n) = read_i64(doc, start) {
67        use std::fmt::Write;
68        let _ = write!(buf, "{n}");
69    } else if let Some(n) = read_f64(doc, start) {
70        use std::fmt::Write;
71        let _ = write!(buf, "{n}");
72    } else {
73        // Complex value (array/map/bin) — hex-encode raw bytes as key.
74        let bytes = &doc[start..end];
75        for b in bytes {
76            use std::fmt::Write;
77            let _ = write!(buf, "{b:02x}");
78        }
79    }
80}
81
82/// Append a field value from a pre-resolved range.
83fn append_field_value_range(buf: &mut String, doc: &[u8], range: Option<(usize, usize)>) {
84    let Some((start, end)) = range else {
85        buf.push_str("null");
86        return;
87    };
88
89    if read_null(doc, start) {
90        buf.push_str("null");
91    } else if let Some(s) = read_str(doc, start) {
92        buf.push('"');
93        buf.push_str(s);
94        buf.push('"');
95    } else if let Some(n) = read_i64(doc, start) {
96        use std::fmt::Write;
97        let _ = write!(buf, "{n}");
98    } else if let Some(n) = read_f64(doc, start) {
99        use std::fmt::Write;
100        let _ = write!(buf, "{n}");
101    } else {
102        let bytes = &doc[start..end];
103        for b in bytes {
104            use std::fmt::Write;
105            let _ = write!(buf, "{b:02x}");
106        }
107    }
108}
109
110#[cfg(test)]
111mod tests {
112    use super::*;
113    use serde_json::json;
114
115    fn encode(v: &serde_json::Value) -> Vec<u8> {
116        nodedb_types::json_msgpack::json_to_msgpack(v).expect("encode")
117    }
118
119    #[test]
120    fn single_string_field() {
121        let doc = encode(&json!({"name": "alice", "age": 30}));
122        let key = build_group_key(&doc, &["name".into()]);
123        assert_eq!(key, r#"["alice"]"#);
124    }
125
126    #[test]
127    fn single_int_field() {
128        let doc = encode(&json!({"status": 200}));
129        let key = build_group_key(&doc, &["status".into()]);
130        assert_eq!(key, "[200]");
131    }
132
133    #[test]
134    fn multiple_fields() {
135        let doc = encode(&json!({"city": "ny", "year": 2024}));
136        let key = build_group_key(&doc, &["city".into(), "year".into()]);
137        assert_eq!(key, r#"["ny",2024]"#);
138    }
139
140    #[test]
141    fn missing_field_is_null() {
142        let doc = encode(&json!({"x": 1}));
143        let key = build_group_key(&doc, &["missing".into()]);
144        assert_eq!(key, "[null]");
145    }
146
147    #[test]
148    fn empty_group_fields() {
149        let doc = encode(&json!({"x": 1}));
150        let key = build_group_key(&doc, &[]);
151        assert_eq!(key, "__all__");
152    }
153
154    #[test]
155    fn null_field_value() {
156        let doc = encode(&json!({"v": null}));
157        let key = build_group_key(&doc, &["v".into()]);
158        assert_eq!(key, "[null]");
159    }
160
161    #[test]
162    fn float_field() {
163        let doc = encode(&json!({"temp": 36.6}));
164        let key = build_group_key(&doc, &["temp".into()]);
165        assert_eq!(key, "[36.6]");
166    }
167}