1use crate::error::{DecodeError, Result};
8use crate::primitives::{RawJson, Reader, Writer};
9
10pub const SSG2_MAGIC: &[u8; 4] = b"SSG2";
11pub const SSG2_FORMAT_VERSION: u16 = 1;
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15pub enum ColumnType {
16 String,
17 Integer,
18 Float,
19 Boolean,
20 Json,
21 Bytes,
22 BlobRef,
25 Crdt,
29}
30
31impl ColumnType {
32 pub fn from_tag(tag: u8) -> Option<Self> {
33 match tag {
34 1 => Some(ColumnType::String),
35 2 => Some(ColumnType::Integer),
36 3 => Some(ColumnType::Float),
37 4 => Some(ColumnType::Boolean),
38 5 => Some(ColumnType::Json),
39 6 => Some(ColumnType::Bytes),
40 7 => Some(ColumnType::BlobRef),
41 8 => Some(ColumnType::Crdt),
42 _ => None,
43 }
44 }
45
46 pub fn tag(self) -> u8 {
47 match self {
48 ColumnType::String => 1,
49 ColumnType::Integer => 2,
50 ColumnType::Float => 3,
51 ColumnType::Boolean => 4,
52 ColumnType::Json => 5,
53 ColumnType::Bytes => 6,
54 ColumnType::BlobRef => 7,
55 ColumnType::Crdt => 8,
56 }
57 }
58
59 pub fn name(self) -> &'static str {
60 match self {
61 ColumnType::String => "string",
62 ColumnType::Integer => "integer",
63 ColumnType::Float => "float",
64 ColumnType::Boolean => "boolean",
65 ColumnType::Json => "json",
66 ColumnType::Bytes => "bytes",
67 ColumnType::BlobRef => "blob_ref",
68 ColumnType::Crdt => "crdt",
69 }
70 }
71
72 pub fn from_name(name: &str) -> Option<Self> {
75 Some(match name {
76 "string" => ColumnType::String,
77 "integer" => ColumnType::Integer,
78 "float" => ColumnType::Float,
79 "boolean" => ColumnType::Boolean,
80 "json" => ColumnType::Json,
81 "bytes" => ColumnType::Bytes,
82 "blob_ref" => ColumnType::BlobRef,
83 "crdt" => ColumnType::Crdt,
84 _ => return None,
85 })
86 }
87}
88
89#[derive(Debug, Clone, PartialEq)]
90pub struct Column {
91 pub name: String,
92 pub ty: ColumnType,
93 pub nullable: bool,
94}
95
96#[derive(Debug, Clone, PartialEq)]
98pub enum ColumnValue {
99 String(String),
100 Integer(i64),
101 Float(f64),
102 Boolean(bool),
103 Json(RawJson),
104 Bytes(Vec<u8>),
105 BlobRef(RawJson),
107 Crdt(Vec<u8>),
110}
111
112pub type Row = Vec<Option<ColumnValue>>;
114
115#[derive(Debug, Clone, PartialEq)]
119pub enum SparseSlot {
120 Absent,
121 Null,
122 Value(ColumnValue),
123}
124
125pub type SparseRow = Vec<SparseSlot>;
127
128#[derive(Debug, Clone, PartialEq)]
130pub struct SegmentRow {
131 pub server_version: i64,
132 pub values: Row,
133}
134
135#[derive(Debug, Clone, PartialEq)]
136pub struct RowsSegment {
137 pub table: String,
138 pub schema_version: i32,
139 pub columns: Vec<Column>,
140 pub blocks: Vec<Vec<SegmentRow>>,
143}
144
145fn read_column_value(r: &mut Reader<'_>, col: &Column) -> Result<ColumnValue> {
147 Ok(match col.ty {
148 ColumnType::String => ColumnValue::String(r.str(&col.name)?),
149 ColumnType::Integer => ColumnValue::Integer(r.i64(&col.name)?),
150 ColumnType::Float => ColumnValue::Float(r.f64(&col.name)?),
151 ColumnType::Boolean => ColumnValue::Boolean(r.bool(&col.name)?),
152 ColumnType::Json => ColumnValue::Json(r.json(&col.name)?),
153 ColumnType::Bytes => ColumnValue::Bytes(r.bytes(&col.name)?),
154 ColumnType::BlobRef => ColumnValue::BlobRef(r.blob_ref(&col.name)?),
155 ColumnType::Crdt => ColumnValue::Crdt(r.bytes(&col.name)?),
157 })
158}
159
160fn write_column_value(w: &mut Writer, value: &ColumnValue) {
162 match value {
163 ColumnValue::String(s) => w.str(s),
164 ColumnValue::Integer(v) => w.i64(*v),
165 ColumnValue::Float(v) => w.f64(*v),
166 ColumnValue::Boolean(v) => w.bool(*v),
167 ColumnValue::Json(j) => w.str(&j.0),
168 ColumnValue::Bytes(b) => w.bytes(b),
169 ColumnValue::BlobRef(j) => w.str(&j.0),
170 ColumnValue::Crdt(b) => w.bytes(b),
171 }
172}
173
174fn bitmap_bit(bitmap: &[u8], index: usize) -> bool {
176 bitmap[index / 8] & (1 << (index % 8)) != 0
177}
178
179fn set_bitmap_bit(bitmap: &mut [u8], index: usize) {
180 bitmap[index / 8] |= 1 << (index % 8);
181}
182
183pub fn decode_row(r: &mut Reader<'_>, columns: &[Column]) -> Result<Row> {
186 let bitmap_len = columns.len().div_ceil(8);
187 let bitmap = r.take(bitmap_len, "row null bitmap")?.to_vec();
188 for bit in columns.len()..bitmap_len * 8 {
190 if bitmap[bit / 8] & (1 << (bit % 8)) != 0 {
191 return Err(DecodeError::invalid(
192 "row null bitmap has a set padding bit (non-canonical encoding)",
193 ));
194 }
195 }
196 let mut row: Row = Vec::with_capacity(columns.len());
197 for (i, col) in columns.iter().enumerate() {
198 let is_null = bitmap[i / 8] & (1 << (i % 8)) != 0;
199 if is_null {
200 if !col.nullable {
201 return Err(DecodeError::invalid(format!(
202 "null bit set for non-nullable column {:?}",
203 col.name
204 )));
205 }
206 row.push(None);
207 continue;
208 }
209 row.push(Some(read_column_value(r, col)?));
210 }
211 Ok(row)
212}
213
214pub fn encode_row(w: &mut Writer, columns: &[Column], row: &Row) {
216 let bitmap_len = columns.len().div_ceil(8);
217 let mut bitmap = vec![0u8; bitmap_len];
218 for (i, value) in row.iter().enumerate() {
219 if value.is_none() {
220 bitmap[i / 8] |= 1 << (i % 8);
221 }
222 }
223 w.raw(&bitmap);
224 for value in row.iter().flatten() {
225 write_column_value(w, value);
226 }
227}
228
229pub fn decode_sparse_row(
238 columns: &[Column],
239 primary_key_index: usize,
240 bytes: &[u8],
241) -> Result<SparseRow> {
242 let mut r = Reader::new(bytes);
243 let presence_len = columns.len().div_ceil(8);
244 let presence = r.take(presence_len, "sparse row presence bitmap")?.to_vec();
245 for bit in columns.len()..presence_len * 8 {
246 if bitmap_bit(&presence, bit) {
247 return Err(DecodeError::invalid(
248 "sparse row presence bitmap has a set padding bit (non-canonical encoding)",
249 ));
250 }
251 }
252 let primary_key = columns.get(primary_key_index);
253 if primary_key.is_none() || !bitmap_bit(&presence, primary_key_index) {
254 return Err(DecodeError::invalid(format!(
255 "sparse row is missing the primary-key column {}",
256 primary_key.map_or(primary_key_index.to_string(), |c| format!("{:?}", c.name))
257 )));
258 }
259 let present_count: usize = presence.iter().map(|b| b.count_ones() as usize).sum();
260 let null_len = present_count.div_ceil(8);
261 let nulls = r.take(null_len, "sparse row null bitmap")?.to_vec();
262 for bit in present_count..null_len * 8 {
263 if bitmap_bit(&nulls, bit) {
264 return Err(DecodeError::invalid(
265 "sparse row null bitmap sets a null bit for an absent column",
266 ));
267 }
268 }
269 let mut row: SparseRow = Vec::with_capacity(columns.len());
270 let mut present_index = 0usize;
271 for (i, col) in columns.iter().enumerate() {
272 if !bitmap_bit(&presence, i) {
273 row.push(SparseSlot::Absent);
274 continue;
275 }
276 let is_null = bitmap_bit(&nulls, present_index);
277 present_index += 1;
278 if !is_null {
279 row.push(SparseSlot::Value(read_column_value(&mut r, col)?));
280 continue;
281 }
282 if !col.nullable {
283 return Err(DecodeError::invalid(format!(
284 "null bit set for non-nullable column {:?}",
285 col.name
286 )));
287 }
288 row.push(SparseSlot::Null);
289 }
290 if !r.is_empty() {
291 return Err(DecodeError::invalid(
292 "trailing bytes after sparse row payload",
293 ));
294 }
295 Ok(row)
296}
297
298pub fn encode_sparse_row(columns: &[Column], primary_key_index: usize, row: &SparseRow) -> Vec<u8> {
302 assert_eq!(
303 row.len(),
304 columns.len(),
305 "sparse row length must match the column table"
306 );
307 let mut presence = vec![0u8; columns.len().div_ceil(8)];
308 let mut present_count = 0usize;
309 for (i, slot) in row.iter().enumerate() {
310 if matches!(slot, SparseSlot::Absent) {
311 continue;
312 }
313 set_bitmap_bit(&mut presence, i);
314 present_count += 1;
315 }
316 let primary_key = columns.get(primary_key_index).unwrap_or_else(|| {
317 panic!("primary-key column index {primary_key_index} is outside the column table")
318 });
319 assert!(
320 bitmap_bit(&presence, primary_key_index),
321 "sparse row is missing the primary-key column {:?}",
322 primary_key.name
323 );
324 let mut nulls = vec![0u8; present_count.div_ceil(8)];
325 let mut present_index = 0usize;
326 for (i, slot) in row.iter().enumerate() {
327 match slot {
328 SparseSlot::Absent => continue,
329 SparseSlot::Null => {
330 let col = &columns[i];
331 assert!(col.nullable, "column {:?} is not nullable", col.name);
332 set_bitmap_bit(&mut nulls, present_index);
333 }
334 SparseSlot::Value(_) => {}
335 }
336 present_index += 1;
337 }
338 let mut w = Writer::new();
339 w.raw(&presence);
340 w.raw(&nulls);
341 for slot in row {
342 if let SparseSlot::Value(value) = slot {
343 write_column_value(&mut w, value);
344 }
345 }
346 w.into_bytes()
347}
348
349pub fn decode_rows_segment(bytes: &[u8]) -> Result<RowsSegment> {
352 let mut r = Reader::new(bytes);
353 let magic = r.take(4, "SSG2 magic")?;
354 if magic != SSG2_MAGIC {
355 return Err(DecodeError::invalid("bad SSG2 magic"));
356 }
357 let format_version = r.u16("formatVersion")?;
358 if format_version != SSG2_FORMAT_VERSION {
359 return Err(DecodeError::invalid(format!(
360 "unsupported SSG2 formatVersion {format_version}"
361 )));
362 }
363 let flags = r.u16("flags")?;
364 if flags != 0 {
365 return Err(DecodeError::invalid(format!(
366 "SSG2 flags must be 0, got 0x{flags:04x}"
367 )));
368 }
369
370 let table = r.str("table")?;
371 let schema_version = r.i32("schemaVersion")?;
372 let column_count = r.u16("column count")? as usize;
373 let mut columns = Vec::with_capacity(column_count);
374 for _ in 0..column_count {
375 let name = r.str("column name")?;
376 let tag = r.u8("column type")?;
377 let ty = ColumnType::from_tag(tag).ok_or_else(|| {
378 DecodeError::invalid(format!("unknown column type tag {tag} for column {name:?}"))
379 })?;
380 let col_flags = r.u8("column flags")?;
381 if col_flags & !0x01 != 0 {
382 return Err(DecodeError::invalid(format!(
383 "reserved column flag bits set for column {name:?} (0x{col_flags:02x})"
384 )));
385 }
386 columns.push(Column {
387 name,
388 ty,
389 nullable: col_flags & 0x01 != 0,
390 });
391 }
392
393 let mut blocks: Vec<Vec<SegmentRow>> = Vec::new();
394 loop {
395 let row_count = r.u32("block rowCount")? as usize;
396 if row_count == 0 {
397 if !r.is_empty() {
399 return Err(DecodeError::invalid(
400 "trailing bytes after the SSG2 end-of-segment marker",
401 ));
402 }
403 break;
404 }
405 let byte_length = r.u32("block byteLength")? as usize;
406 let row_bytes = r.take(byte_length, "block rows")?;
407 let mut rr = Reader::new(row_bytes);
408 let mut rows = Vec::with_capacity(row_count.min(4096));
409 for _ in 0..row_count {
410 let server_version = rr.i64("row serverVersion")?;
412 if server_version < 1 {
413 return Err(DecodeError::invalid(format!(
414 "row serverVersion must be >= 1, got {server_version}"
415 )));
416 }
417 rows.push(SegmentRow {
418 server_version,
419 values: decode_row(&mut rr, &columns)?,
420 });
421 }
422 if !rr.is_empty() {
423 return Err(DecodeError::invalid(
424 "block rows do not consume exactly byteLength bytes",
425 ));
426 }
427 blocks.push(rows);
428 }
429 Ok(RowsSegment {
430 table,
431 schema_version,
432 columns,
433 blocks,
434 })
435}
436
437pub fn encode_rows_segment(seg: &RowsSegment) -> Vec<u8> {
439 let mut w = Writer::new();
440 w.raw(SSG2_MAGIC);
441 w.u16(SSG2_FORMAT_VERSION);
442 w.u16(0); w.str(&seg.table);
444 w.i32(seg.schema_version);
445 w.u16(seg.columns.len() as u16);
446 for col in &seg.columns {
447 w.str(&col.name);
448 w.u8(col.ty.tag());
449 w.u8(u8::from(col.nullable));
450 }
451 for block in &seg.blocks {
452 let mut bw = Writer::new();
453 for row in block {
454 bw.i64(row.server_version);
455 encode_row(&mut bw, &seg.columns, &row.values);
456 }
457 let row_bytes = bw.into_bytes();
458 w.u32(block.len() as u32);
459 w.u32(row_bytes.len() as u32);
460 w.raw(&row_bytes);
461 }
462 w.u32(0); w.into_bytes()
464}
465
466#[cfg(test)]
467mod tests {
468 use super::*;
469
470 fn sparse_columns() -> Vec<Column> {
473 vec![
474 Column {
475 name: "id".to_owned(),
476 ty: ColumnType::String,
477 nullable: false,
478 },
479 Column {
480 name: "title".to_owned(),
481 ty: ColumnType::String,
482 nullable: true,
483 },
484 Column {
485 name: "body".to_owned(),
486 ty: ColumnType::String,
487 nullable: false,
488 },
489 Column {
490 name: "count".to_owned(),
491 ty: ColumnType::Integer,
492 nullable: true,
493 },
494 Column {
495 name: "score".to_owned(),
496 ty: ColumnType::Float,
497 nullable: false,
498 },
499 Column {
500 name: "done".to_owned(),
501 ty: ColumnType::Boolean,
502 nullable: false,
503 },
504 Column {
505 name: "meta".to_owned(),
506 ty: ColumnType::Json,
507 nullable: true,
508 },
509 Column {
510 name: "blob".to_owned(),
511 ty: ColumnType::Bytes,
512 nullable: true,
513 },
514 Column {
515 name: "doc".to_owned(),
516 ty: ColumnType::Crdt,
517 nullable: true,
518 },
519 ]
520 }
521
522 fn partial_row() -> SparseRow {
523 vec![
524 SparseSlot::Value(ColumnValue::String("n-1".to_owned())),
525 SparseSlot::Null,
526 SparseSlot::Absent,
527 SparseSlot::Absent,
528 SparseSlot::Absent,
529 SparseSlot::Absent,
530 SparseSlot::Value(ColumnValue::Json(RawJson("{\"tags\": [\"a\"]}".to_owned()))),
531 SparseSlot::Absent,
532 SparseSlot::Absent,
533 ]
534 }
535
536 fn decode_error(columns: &[Column], bytes: &[u8]) -> String {
537 let error =
538 decode_sparse_row(columns, 0, bytes).expect_err("sparse payload must be rejected");
539 assert_eq!(error.code, crate::error::ErrorCode::InvalidRequest);
540 error.detail
541 }
542
543 #[test]
544 fn partial_row_round_trips_with_a_present_null() {
545 let columns = sparse_columns();
546 let row = partial_row();
547 let bytes = encode_sparse_row(&columns, 0, &row);
548 assert_eq!(bytes[0], 0b0100_0011); assert_eq!(bytes[1], 0x00); assert_eq!(bytes[2], 0b0000_0010); let decoded = decode_sparse_row(&columns, 0, &bytes).unwrap();
552 assert_eq!(decoded, row);
553 assert_eq!(encode_sparse_row(&columns, 0, &decoded), bytes);
554 }
555
556 #[test]
557 fn full_row_still_leads_with_the_presence_bitmap() {
558 let columns = sparse_columns();
559 let row: SparseRow = vec![
560 SparseSlot::Value(ColumnValue::String("n-2".to_owned())),
561 SparseSlot::Null,
562 SparseSlot::Value(ColumnValue::String("body".to_owned())),
563 SparseSlot::Value(ColumnValue::Integer(-7)),
564 SparseSlot::Value(ColumnValue::Float(0.125)),
565 SparseSlot::Value(ColumnValue::Boolean(false)),
566 SparseSlot::Value(ColumnValue::Json(RawJson("{\"k\":true}".to_owned()))),
567 SparseSlot::Value(ColumnValue::Bytes(vec![0, 1, 254, 255])),
568 SparseSlot::Value(ColumnValue::Crdt(vec![0x01, 0x02, 0xfe])),
569 ];
570 let bytes = encode_sparse_row(&columns, 0, &row);
571 assert_eq!(&bytes[..2], &[0xff, 0x01]);
572 assert_eq!(decode_sparse_row(&columns, 0, &bytes).unwrap(), row);
573 let full_row: Row = row
576 .iter()
577 .map(|slot| match slot {
578 SparseSlot::Absent => None,
579 SparseSlot::Null => None,
580 SparseSlot::Value(value) => Some(value.clone()),
581 })
582 .collect();
583 let mut w = Writer::new();
584 encode_row(&mut w, &columns, &full_row);
585 assert_ne!(w.into_bytes(), bytes);
586 }
587
588 #[test]
589 fn every_column_count_boundary_round_trips() {
590 for count in [1usize, 7, 8, 9] {
591 let columns: Vec<Column> = (0..count)
592 .map(|i| Column {
593 name: format!("c{i}"),
594 ty: ColumnType::String,
595 nullable: true,
596 })
597 .collect();
598 let row: SparseRow = (0..count)
599 .map(|i| {
600 if i > 0 && i % 3 == 0 {
601 SparseSlot::Absent
602 } else if i == 1 {
603 SparseSlot::Null
604 } else {
605 SparseSlot::Value(ColumnValue::String(format!("v{i}")))
606 }
607 })
608 .collect();
609 let bytes = encode_sparse_row(&columns, 0, &row);
610 let pad_start = count % 8;
611 if pad_start != 0 {
612 let last = bytes[count.div_ceil(8) - 1];
613 for bit in pad_start..8 {
614 assert_eq!((last >> bit) & 1, 0, "column count {count}");
615 }
616 }
617 let decoded = decode_sparse_row(&columns, 0, &bytes).unwrap();
618 assert_eq!(decoded, row, "column count {count}");
619 assert_eq!(encode_sparse_row(&columns, 0, &decoded), bytes);
620 }
621 }
622
623 #[test]
624 fn decode_rejects_presence_padding_bits() {
625 let columns = sparse_columns();
627 let detail = decode_error(&columns, &[0b0000_0001, 0b0000_0010]);
628 assert!(
629 detail.contains("presence bitmap has a set padding bit"),
630 "{detail}"
631 );
632 }
633
634 #[test]
635 fn decode_rejects_a_null_bit_for_an_absent_column() {
636 let columns = sparse_columns();
638 let detail = decode_error(&columns, &[0b0000_0001, 0x00, 0b0000_1000]);
639 assert!(detail.contains("null bit for an absent column"), "{detail}");
640 }
641
642 #[test]
643 fn decode_rejects_a_null_bit_for_a_non_nullable_column() {
644 let columns = sparse_columns();
646 let detail = decode_error(&columns, &[0b0000_0011, 0x00, 0b0000_0001]);
647 assert!(
648 detail.contains("null bit set for non-nullable column \"id\""),
649 "{detail}"
650 );
651 }
652
653 #[test]
654 fn decode_rejects_an_absent_primary_key() {
655 let columns = sparse_columns();
656 let mut w = Writer::new();
657 w.raw(&[0b0000_0010, 0x00]); w.u8(0x00); w.str("unnamed");
660 let detail = decode_error(&columns, &w.into_bytes());
661 assert!(
662 detail.contains("missing the primary-key column \"id\""),
663 "{detail}"
664 );
665 }
666
667 #[test]
668 fn decode_rejects_trailing_bytes() {
669 let columns = sparse_columns();
670 let mut bytes = encode_sparse_row(&columns, 0, &partial_row());
671 bytes.push(0x00);
672 let detail = decode_error(&columns, &bytes);
673 assert!(detail.contains("trailing bytes"), "{detail}");
674 }
675
676 #[test]
677 #[should_panic(expected = "is missing the primary-key column \"id\"")]
678 fn encode_rejects_a_missing_primary_key() {
679 let columns = sparse_columns();
680 let row: SparseRow = vec![SparseSlot::Absent; columns.len()];
681 encode_sparse_row(&columns, 0, &row);
682 }
683
684 #[test]
685 #[should_panic(expected = "is not nullable")]
686 fn encode_rejects_a_null_in_a_non_nullable_column() {
687 let columns = sparse_columns();
688 let mut row = partial_row();
689 row[0] = SparseSlot::Null;
690 encode_sparse_row(&columns, 0, &row);
691 }
692}