1use crate::error::{DecodeError, Result};
7use crate::primitives::{RawJson, Reader, Writer};
8
9pub const SSG2_MAGIC: &[u8; 4] = b"SSG2";
10pub const SSG2_FORMAT_VERSION: u16 = 1;
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub enum ColumnType {
15 String,
16 Integer,
17 Float,
18 Boolean,
19 Json,
20 Bytes,
21 BlobRef,
24 Crdt,
28}
29
30impl ColumnType {
31 pub fn from_tag(tag: u8) -> Option<Self> {
32 match tag {
33 1 => Some(ColumnType::String),
34 2 => Some(ColumnType::Integer),
35 3 => Some(ColumnType::Float),
36 4 => Some(ColumnType::Boolean),
37 5 => Some(ColumnType::Json),
38 6 => Some(ColumnType::Bytes),
39 7 => Some(ColumnType::BlobRef),
40 8 => Some(ColumnType::Crdt),
41 _ => None,
42 }
43 }
44
45 pub fn tag(self) -> u8 {
46 match self {
47 ColumnType::String => 1,
48 ColumnType::Integer => 2,
49 ColumnType::Float => 3,
50 ColumnType::Boolean => 4,
51 ColumnType::Json => 5,
52 ColumnType::Bytes => 6,
53 ColumnType::BlobRef => 7,
54 ColumnType::Crdt => 8,
55 }
56 }
57
58 pub fn name(self) -> &'static str {
59 match self {
60 ColumnType::String => "string",
61 ColumnType::Integer => "integer",
62 ColumnType::Float => "float",
63 ColumnType::Boolean => "boolean",
64 ColumnType::Json => "json",
65 ColumnType::Bytes => "bytes",
66 ColumnType::BlobRef => "blob_ref",
67 ColumnType::Crdt => "crdt",
68 }
69 }
70}
71
72#[derive(Debug, Clone, PartialEq)]
73pub struct Column {
74 pub name: String,
75 pub ty: ColumnType,
76 pub nullable: bool,
77}
78
79#[derive(Debug, Clone, PartialEq)]
81pub enum ColumnValue {
82 String(String),
83 Integer(i64),
84 Float(f64),
85 Boolean(bool),
86 Json(RawJson),
87 Bytes(Vec<u8>),
88 BlobRef(RawJson),
90 Crdt(Vec<u8>),
93}
94
95pub type Row = Vec<Option<ColumnValue>>;
97
98#[derive(Debug, Clone, PartialEq)]
100pub struct SegmentRow {
101 pub server_version: i64,
102 pub values: Row,
103}
104
105#[derive(Debug, Clone, PartialEq)]
106pub struct RowsSegment {
107 pub table: String,
108 pub schema_version: i32,
109 pub columns: Vec<Column>,
110 pub blocks: Vec<Vec<SegmentRow>>,
113}
114
115pub fn decode_row(r: &mut Reader<'_>, columns: &[Column]) -> Result<Row> {
118 let bitmap_len = columns.len().div_ceil(8);
119 let bitmap = r.take(bitmap_len, "row null bitmap")?.to_vec();
120 for bit in columns.len()..bitmap_len * 8 {
122 if bitmap[bit / 8] & (1 << (bit % 8)) != 0 {
123 return Err(DecodeError::invalid(
124 "row null bitmap has a set padding bit (non-canonical encoding)",
125 ));
126 }
127 }
128 let mut row: Row = Vec::with_capacity(columns.len());
129 for (i, col) in columns.iter().enumerate() {
130 let is_null = bitmap[i / 8] & (1 << (i % 8)) != 0;
131 if is_null {
132 if !col.nullable {
133 return Err(DecodeError::invalid(format!(
134 "null bit set for non-nullable column {:?}",
135 col.name
136 )));
137 }
138 row.push(None);
139 continue;
140 }
141 let value = match col.ty {
142 ColumnType::String => ColumnValue::String(r.str(&col.name)?),
143 ColumnType::Integer => ColumnValue::Integer(r.i64(&col.name)?),
144 ColumnType::Float => ColumnValue::Float(r.f64(&col.name)?),
145 ColumnType::Boolean => ColumnValue::Boolean(r.bool(&col.name)?),
146 ColumnType::Json => ColumnValue::Json(r.json(&col.name)?),
147 ColumnType::Bytes => ColumnValue::Bytes(r.bytes(&col.name)?),
148 ColumnType::BlobRef => ColumnValue::BlobRef(r.blob_ref(&col.name)?),
149 ColumnType::Crdt => ColumnValue::Crdt(r.bytes(&col.name)?),
151 };
152 row.push(Some(value));
153 }
154 Ok(row)
155}
156
157pub fn encode_row(w: &mut Writer, columns: &[Column], row: &Row) {
159 let bitmap_len = columns.len().div_ceil(8);
160 let mut bitmap = vec![0u8; bitmap_len];
161 for (i, value) in row.iter().enumerate() {
162 if value.is_none() {
163 bitmap[i / 8] |= 1 << (i % 8);
164 }
165 }
166 w.raw(&bitmap);
167 for value in row.iter().flatten() {
168 match value {
169 ColumnValue::String(s) => w.str(s),
170 ColumnValue::Integer(v) => w.i64(*v),
171 ColumnValue::Float(v) => w.f64(*v),
172 ColumnValue::Boolean(v) => w.bool(*v),
173 ColumnValue::Json(j) => w.str(&j.0),
174 ColumnValue::Bytes(b) => w.bytes(b),
175 ColumnValue::BlobRef(j) => w.str(&j.0),
176 ColumnValue::Crdt(b) => w.bytes(b),
177 }
178 }
179}
180
181pub fn decode_rows_segment(bytes: &[u8]) -> Result<RowsSegment> {
184 let mut r = Reader::new(bytes);
185 let magic = r.take(4, "SSG2 magic")?;
186 if magic != SSG2_MAGIC {
187 return Err(DecodeError::invalid("bad SSG2 magic"));
188 }
189 let format_version = r.u16("formatVersion")?;
190 if format_version != SSG2_FORMAT_VERSION {
191 return Err(DecodeError::invalid(format!(
192 "unsupported SSG2 formatVersion {format_version}"
193 )));
194 }
195 let flags = r.u16("flags")?;
196 if flags != 0 {
197 return Err(DecodeError::invalid(format!(
198 "SSG2 flags must be 0, got 0x{flags:04x}"
199 )));
200 }
201
202 let table = r.str("table")?;
203 let schema_version = r.i32("schemaVersion")?;
204 let column_count = r.u16("column count")? as usize;
205 let mut columns = Vec::with_capacity(column_count);
206 for _ in 0..column_count {
207 let name = r.str("column name")?;
208 let tag = r.u8("column type")?;
209 let ty = ColumnType::from_tag(tag).ok_or_else(|| {
210 DecodeError::invalid(format!("unknown column type tag {tag} for column {name:?}"))
211 })?;
212 let col_flags = r.u8("column flags")?;
213 if col_flags & !0x01 != 0 {
214 return Err(DecodeError::invalid(format!(
215 "reserved column flag bits set for column {name:?} (0x{col_flags:02x})"
216 )));
217 }
218 columns.push(Column {
219 name,
220 ty,
221 nullable: col_flags & 0x01 != 0,
222 });
223 }
224
225 let mut blocks: Vec<Vec<SegmentRow>> = Vec::new();
226 loop {
227 let row_count = r.u32("block rowCount")? as usize;
228 if row_count == 0 {
229 if !r.is_empty() {
231 return Err(DecodeError::invalid(
232 "trailing bytes after the SSG2 end-of-segment marker",
233 ));
234 }
235 break;
236 }
237 let byte_length = r.u32("block byteLength")? as usize;
238 let row_bytes = r.take(byte_length, "block rows")?;
239 let mut rr = Reader::new(row_bytes);
240 let mut rows = Vec::with_capacity(row_count.min(4096));
241 for _ in 0..row_count {
242 let server_version = rr.i64("row serverVersion")?;
244 if server_version < 1 {
245 return Err(DecodeError::invalid(format!(
246 "row serverVersion must be >= 1, got {server_version}"
247 )));
248 }
249 rows.push(SegmentRow {
250 server_version,
251 values: decode_row(&mut rr, &columns)?,
252 });
253 }
254 if !rr.is_empty() {
255 return Err(DecodeError::invalid(
256 "block rows do not consume exactly byteLength bytes",
257 ));
258 }
259 blocks.push(rows);
260 }
261 Ok(RowsSegment {
262 table,
263 schema_version,
264 columns,
265 blocks,
266 })
267}
268
269pub fn encode_rows_segment(seg: &RowsSegment) -> Vec<u8> {
271 let mut w = Writer::new();
272 w.raw(SSG2_MAGIC);
273 w.u16(SSG2_FORMAT_VERSION);
274 w.u16(0); w.str(&seg.table);
276 w.i32(seg.schema_version);
277 w.u16(seg.columns.len() as u16);
278 for col in &seg.columns {
279 w.str(&col.name);
280 w.u8(col.ty.tag());
281 w.u8(u8::from(col.nullable));
282 }
283 for block in &seg.blocks {
284 let mut bw = Writer::new();
285 for row in block {
286 bw.i64(row.server_version);
287 encode_row(&mut bw, &seg.columns, &row.values);
288 }
289 let row_bytes = bw.into_bytes();
290 w.u32(block.len() as u32);
291 w.u32(row_bytes.len() as u32);
292 w.raw(&row_bytes);
293 }
294 w.u32(0); w.into_bytes()
296}