base_d/encoders/algorithms/schema/
binary_packer.rs1use crate::encoders::algorithms::schema::types::{
2 FLAG_HAS_NULLS, FLAG_HAS_ROOT_KEY, FieldType, IntermediateRepresentation, SchemaValue,
3};
4
5pub fn pack(ir: &IntermediateRepresentation) -> Vec<u8> {
7 let mut buffer = Vec::new();
8
9 pack_header(&mut buffer, ir);
11
12 pack_values(&mut buffer, ir);
14
15 buffer
16}
17
18fn pack_header(buffer: &mut Vec<u8>, ir: &IntermediateRepresentation) {
20 let header = &ir.header;
21
22 buffer.push(header.flags);
24
25 if header.has_flag(FLAG_HAS_ROOT_KEY) && header.root_key.is_some() {
27 let key = header.root_key.as_ref().unwrap();
28 encode_varint(buffer, key.len() as u64);
29 buffer.extend_from_slice(key.as_bytes());
30 }
31
32 encode_varint(buffer, header.row_count as u64);
34
35 encode_varint(buffer, header.fields.len() as u64);
37
38 pack_field_types(buffer, ir);
40
41 for field in &header.fields {
43 encode_varint(buffer, field.name.len() as u64);
44 buffer.extend_from_slice(field.name.as_bytes());
45 }
46
47 if header.has_flag(FLAG_HAS_NULLS) && header.null_bitmap.is_some() {
49 let bitmap = header.null_bitmap.as_ref().unwrap();
50 buffer.extend_from_slice(bitmap);
51 }
52}
53
54fn pack_field_types(buffer: &mut Vec<u8>, ir: &IntermediateRepresentation) {
56 let mut type_buffer = Vec::new();
57 let mut nibble_count = 0;
58
59 for field in &ir.header.fields {
60 pack_field_type_recursive(&mut type_buffer, &field.field_type, &mut nibble_count);
61 }
62
63 encode_varint(buffer, type_buffer.len() as u64);
65 buffer.extend_from_slice(&type_buffer);
66}
67
68fn pack_field_type_recursive(
70 buffer: &mut Vec<u8>,
71 field_type: &FieldType,
72 nibble_count: &mut usize,
73) {
74 let tag = field_type.type_tag();
75
76 if (*nibble_count).is_multiple_of(2) {
78 buffer.push(tag);
80 } else {
81 let last_idx = buffer.len() - 1;
83 buffer[last_idx] |= tag << 4;
84 }
85 *nibble_count += 1;
86
87 if let FieldType::Array(element_type) = field_type {
89 pack_field_type_recursive(buffer, element_type, nibble_count);
90 }
91}
92
93fn pack_values(buffer: &mut Vec<u8>, ir: &IntermediateRepresentation) {
95 for value in &ir.values {
96 pack_value(buffer, value);
97 }
98}
99
100fn pack_value(buffer: &mut Vec<u8>, value: &SchemaValue) {
102 match value {
103 SchemaValue::U64(v) => encode_varint(buffer, *v),
104 SchemaValue::I64(v) => encode_signed_varint(buffer, *v),
105 SchemaValue::F64(v) => buffer.extend_from_slice(&v.to_le_bytes()),
106 SchemaValue::String(s) => {
107 encode_varint(buffer, s.len() as u64);
108 buffer.extend_from_slice(s.as_bytes());
109 }
110 SchemaValue::Bool(b) => buffer.push(if *b { 1 } else { 0 }),
111 SchemaValue::Null => {} SchemaValue::Array(arr) => {
113 encode_varint(buffer, arr.len() as u64);
114 for item in arr {
115 pack_value(buffer, item);
116 }
117 }
118 }
119}
120
121pub(crate) fn encode_varint(buffer: &mut Vec<u8>, mut value: u64) {
123 loop {
124 let mut byte = (value & 0x7F) as u8;
125 value >>= 7;
126 if value != 0 {
127 byte |= 0x80; }
129 buffer.push(byte);
130 if value == 0 {
131 break;
132 }
133 }
134}
135
136pub(crate) fn encode_signed_varint(buffer: &mut Vec<u8>, value: i64) {
138 let encoded = ((value << 1) ^ (value >> 63)) as u64;
139 encode_varint(buffer, encoded);
140}
141
142#[cfg(test)]
143mod tests {
144 use super::*;
145 use crate::encoders::algorithms::schema::types::{FieldDef, SchemaHeader};
146
147 #[test]
148 fn test_encode_varint() {
149 let mut buf = Vec::new();
150 encode_varint(&mut buf, 0);
151 assert_eq!(buf, vec![0]);
152
153 buf.clear();
154 encode_varint(&mut buf, 1);
155 assert_eq!(buf, vec![1]);
156
157 buf.clear();
158 encode_varint(&mut buf, 127);
159 assert_eq!(buf, vec![127]);
160
161 buf.clear();
162 encode_varint(&mut buf, 128);
163 assert_eq!(buf, vec![0x80, 0x01]);
164
165 buf.clear();
166 encode_varint(&mut buf, 16383);
167 assert_eq!(buf, vec![0xFF, 0x7F]);
168
169 buf.clear();
170 encode_varint(&mut buf, 16384);
171 assert_eq!(buf, vec![0x80, 0x80, 0x01]);
172 }
173
174 #[test]
175 fn test_encode_signed_varint() {
176 let mut buf = Vec::new();
177 encode_signed_varint(&mut buf, 0);
178 assert_eq!(buf, vec![0]);
179
180 buf.clear();
181 encode_signed_varint(&mut buf, -1);
182 assert_eq!(buf, vec![1]);
183
184 buf.clear();
185 encode_signed_varint(&mut buf, 1);
186 assert_eq!(buf, vec![2]);
187
188 buf.clear();
189 encode_signed_varint(&mut buf, -64);
190 assert_eq!(buf, vec![127]);
191
192 buf.clear();
193 encode_signed_varint(&mut buf, 64);
194 assert_eq!(buf, vec![128, 1]);
195 }
196
197 #[test]
198 fn test_pack_simple_ir() {
199 let fields = vec![
200 FieldDef::new("id", FieldType::U64),
201 FieldDef::new("name", FieldType::String),
202 ];
203 let header = SchemaHeader::new(1, fields);
204
205 let values = vec![
206 SchemaValue::U64(42),
207 SchemaValue::String("Alice".to_string()),
208 ];
209
210 let ir = IntermediateRepresentation::new(header, values).unwrap();
211 let packed = pack(&ir);
212
213 assert!(!packed.is_empty());
215
216 assert_eq!(packed[0], 0);
218 }
219
220 #[test]
221 fn test_pack_with_root_key() {
222 let mut header = SchemaHeader::new(1, vec![FieldDef::new("id", FieldType::U64)]);
223 header.root_key = Some("users".to_string());
224 header.set_flag(FLAG_HAS_ROOT_KEY);
225
226 let values = vec![SchemaValue::U64(42)];
227 let ir = IntermediateRepresentation::new(header, values).unwrap();
228 let packed = pack(&ir);
229
230 assert_eq!(packed[0] & FLAG_HAS_ROOT_KEY, FLAG_HAS_ROOT_KEY);
232 }
233
234 #[test]
235 fn test_pack_field_types() {
236 let fields = vec![
237 FieldDef::new("a", FieldType::U64), FieldDef::new("b", FieldType::I64), FieldDef::new("c", FieldType::String), ];
241 let header = SchemaHeader::new(1, fields);
242 let values = vec![
243 SchemaValue::U64(1),
244 SchemaValue::I64(-1),
245 SchemaValue::String("x".to_string()),
246 ];
247
248 let ir = IntermediateRepresentation::new(header, values).unwrap();
249 let packed = pack(&ir);
250
251 assert!(!packed.is_empty());
255 }
256
257 #[test]
258 fn test_pack_values() {
259 let mut buffer = Vec::new();
260
261 pack_value(&mut buffer, &SchemaValue::U64(42));
262 assert_eq!(buffer, vec![42]);
263
264 buffer.clear();
265 pack_value(&mut buffer, &SchemaValue::Bool(true));
266 assert_eq!(buffer, vec![1]);
267
268 buffer.clear();
269 pack_value(&mut buffer, &SchemaValue::String("hi".to_string()));
270 assert_eq!(buffer, vec![2, b'h', b'i']);
271 }
272
273 #[test]
274 fn test_pack_array() {
275 let mut buffer = Vec::new();
276 let array = SchemaValue::Array(vec![SchemaValue::U64(1), SchemaValue::U64(2)]);
277 pack_value(&mut buffer, &array);
278
279 assert_eq!(buffer, vec![2, 1, 2]);
281 }
282}