nodedb_columnar/reader/
cell.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! One decoded segment cell as a `Value`.
4//!
5//! The rule is the memtable's read rule (`ColumnData::get_value`): a cell
6//! reads as the same `Value` before and after its memtable flushes. The
7//! declared column type decides the variant, because a segment block records
8//! only its physical layout.
9
10use nodedb_types::columnar::ColumnType;
11use nodedb_types::value::Value;
12use nodedb_types::value_from_msgpack;
13
14use crate::error::ColumnarError;
15
16use super::types::DecodedColumn;
17
18/// Width of a decimal, UUID or ULID cell.
19const ID_WIDTH: usize = 16;
20
21/// Width of one vector element: a little-endian `f32`.
22const F32_WIDTH: usize = 4;
23
24/// Row `row` of `col` as the `Value` the memtable reads for a cell of
25/// `declared` type.
26///
27/// A null cell or a row past the column is `Value::Null`. A cell whose bytes
28/// do not hold a value of `declared` type is a corrupt segment and an error.
29pub fn decoded_cell_value(
30    col: &DecodedColumn,
31    row: usize,
32    declared: &ColumnType,
33) -> Result<Value, ColumnarError> {
34    let is_valid = |valid: &[bool]| valid.get(row).copied().unwrap_or(false);
35    match col {
36        // A time column decodes as `Int64`. The declared type types the cell.
37        DecodedColumn::Int64 { values, valid } | DecodedColumn::Timestamp { values, valid } => {
38            Ok(match values.get(row) {
39                Some(&v) if is_valid(valid) => declared.time_cell(v),
40                _ => Value::Null,
41            })
42        }
43        DecodedColumn::Float64 { values, valid } => Ok(match values.get(row) {
44            Some(&v) if is_valid(valid) => Value::Float(v),
45            _ => Value::Null,
46        }),
47        DecodedColumn::Bool { values, valid } => Ok(match values.get(row) {
48            Some(&v) if is_valid(valid) => Value::Bool(v),
49            _ => Value::Null,
50        }),
51        DecodedColumn::DictEncoded {
52            ids,
53            dictionary,
54            valid,
55        } => {
56            let Some(&id) = ids.get(row).filter(|_| is_valid(valid)) else {
57                return Ok(Value::Null);
58            };
59            let text = usize::try_from(id)
60                .ok()
61                .and_then(|id| dictionary.get(id))
62                .ok_or_else(|| {
63                    corruption(format!(
64                        "dictionary ID {id} is outside a dictionary of {} entries",
65                        dictionary.len()
66                    ))
67                })?;
68            Ok(Value::String(text.clone()))
69        }
70        DecodedColumn::Binary {
71            data,
72            offsets,
73            valid,
74        } => {
75            if !is_valid(valid) {
76                return Ok(Value::Null);
77            }
78            let bytes = cell_bytes(data, offsets, row)?;
79            binary_cell_value(bytes, declared)
80        }
81    }
82}
83
84/// The bytes of row `row` of a binary column.
85fn cell_bytes<'a>(data: &'a [u8], offsets: &[u32], row: usize) -> Result<&'a [u8], ColumnarError> {
86    let bound = |i: usize| offsets.get(i).and_then(|&o| usize::try_from(o).ok());
87    bound(row)
88        .zip(row.checked_add(1).and_then(bound))
89        .and_then(|(start, end)| data.get(start..end))
90        .ok_or_else(|| {
91            corruption(format!(
92                "row {row} has no byte range in a binary column of {} offsets and {} bytes",
93                offsets.len(),
94                data.len()
95            ))
96        })
97}
98
99/// A binary cell as the `Value` the memtable reads for `declared`.
100fn binary_cell_value(bytes: &[u8], declared: &ColumnType) -> Result<Value, ColumnarError> {
101    Ok(match declared {
102        ColumnType::String
103        | ColumnType::Regex
104        | ColumnType::SparseVector
105        | ColumnType::Geometry => Value::String(utf8(bytes, declared)?.to_owned()),
106        ColumnType::Bytes
107        | ColumnType::Array
108        | ColumnType::Set
109        | ColumnType::Range
110        | ColumnType::Record => Value::Bytes(bytes.to_vec()),
111        ColumnType::Json if bytes.is_empty() => Value::Null,
112        ColumnType::Json => value_from_msgpack(bytes)
113            .map_err(|e| corruption(format!("JSON cell is not MessagePack: {e}")))?,
114        ColumnType::Decimal { .. } => Value::Decimal(rust_decimal::Decimal::deserialize(id_cell(
115            bytes, declared,
116        )?)),
117        ColumnType::Uuid => {
118            Value::Uuid(uuid::Uuid::from_bytes(id_cell(bytes, declared)?).to_string())
119        }
120        ColumnType::Ulid => {
121            Value::Ulid(ulid::Ulid::from_bytes(id_cell(bytes, declared)?).to_string())
122        }
123        ColumnType::Vector(dim) => vector_cell(bytes, *dim)?,
124        ColumnType::Int64
125        | ColumnType::Float64
126        | ColumnType::Bool
127        | ColumnType::Timestamp
128        | ColumnType::Timestamptz
129        | ColumnType::SystemTimestamp
130        | ColumnType::Duration => {
131            return Err(corruption(format!(
132                "a {declared} column has no binary cells"
133            )));
134        }
135        // `ColumnType` is `#[non_exhaustive]`. The memtable stores a type it
136        // does not know as raw bytes.
137        _ => Value::Bytes(bytes.to_vec()),
138    })
139}
140
141fn utf8<'a>(bytes: &'a [u8], declared: &ColumnType) -> Result<&'a str, ColumnarError> {
142    std::str::from_utf8(bytes).map_err(|e| corruption(format!("{declared} cell is not UTF-8: {e}")))
143}
144
145fn id_cell(bytes: &[u8], declared: &ColumnType) -> Result<[u8; ID_WIDTH], ColumnarError> {
146    <[u8; ID_WIDTH]>::try_from(bytes).map_err(|_| {
147        corruption(format!(
148            "{declared} cell holds {} bytes, not {ID_WIDTH}",
149            bytes.len()
150        ))
151    })
152}
153
154/// A packed little-endian `f32` vector cell as an array of floats.
155fn vector_cell(bytes: &[u8], dim: u32) -> Result<Value, ColumnarError> {
156    let expected = usize::try_from(dim)
157        .ok()
158        .and_then(|d| d.checked_mul(F32_WIDTH));
159    if expected != Some(bytes.len()) {
160        return Err(corruption(format!(
161            "VECTOR({dim}) cell holds {} bytes",
162            bytes.len()
163        )));
164    }
165    let (chunks, _) = bytes.as_chunks::<F32_WIDTH>();
166    Ok(Value::Array(
167        chunks
168            .iter()
169            .map(|c| Value::Float(f64::from(f32::from_le_bytes(*c))))
170            .collect(),
171    ))
172}
173
174fn corruption(reason: String) -> ColumnarError {
175    ColumnarError::Corruption {
176        segment_id: None,
177        reason,
178        offset: None,
179    }
180}
181
182#[cfg(test)]
183mod tests {
184    use std::collections::HashMap;
185
186    use nodedb_types::columnar::{ColumnDef, ColumnType, ColumnarSchema};
187    use nodedb_types::value::Value;
188    use nodedb_types::{NdbDateTime, NdbDuration};
189
190    use super::decoded_cell_value;
191    use crate::memtable::ColumnarMemtable;
192    use crate::reader::SegmentReader;
193    use crate::test_support::test_memory;
194    use crate::writer::{PROFILE_PLAIN, SegmentWriter};
195
196    /// Rows past one block, so the second block's offsets are covered.
197    const ROWS: usize = 1500;
198
199    const BASE_MICROS: i64 = 1_700_000_000_000_000;
200
201    fn every_type_schema() -> ColumnarSchema {
202        ColumnarSchema::new(vec![
203            ColumnDef::required("id", ColumnType::Int64).with_primary_key(),
204            ColumnDef::nullable("f", ColumnType::Float64),
205            ColumnDef::nullable("b", ColumnType::Bool),
206            ColumnDef::nullable("s", ColumnType::String),
207            ColumnDef::nullable("tag", ColumnType::String),
208            ColumnDef::nullable("raw", ColumnType::Bytes),
209            ColumnDef::nullable("ts", ColumnType::Timestamp),
210            ColumnDef::nullable("tz", ColumnType::Timestamptz),
211            ColumnDef::nullable("sys", ColumnType::SystemTimestamp),
212            ColumnDef::nullable("dec", ColumnType::Decimal(None)),
213            ColumnDef::nullable("u", ColumnType::Uuid),
214            ColumnDef::nullable("ul", ColumnType::Ulid),
215            ColumnDef::nullable("geo", ColumnType::Geometry),
216            ColumnDef::nullable("vec", ColumnType::Vector(3)),
217            ColumnDef::nullable("j", ColumnType::Json),
218            ColumnDef::nullable("dur", ColumnType::Duration),
219            ColumnDef::nullable("re", ColumnType::Regex),
220            ColumnDef::nullable("arr", ColumnType::Array),
221            ColumnDef::nullable("sparse", ColumnType::SparseVector),
222        ])
223        .expect("valid schema")
224    }
225
226    fn ulid_text(i: usize) -> String {
227        let mut bytes = [0u8; 16];
228        bytes[..8].copy_from_slice(&(i as u64 + 1).to_be_bytes());
229        bytes[15] = 7;
230        ulid::Ulid::from_bytes(bytes).to_string()
231    }
232
233    fn json_cell(i: usize) -> Value {
234        if i.is_multiple_of(5) {
235            // JSON text of a string scalar.
236            Value::String("\"scalar\"".into())
237        } else {
238            Value::Object(HashMap::from([
239                ("n".to_string(), Value::Integer(i as i64)),
240                ("tag".to_string(), Value::String("x".into())),
241            ]))
242        }
243    }
244
245    fn row(i: usize) -> Vec<Value> {
246        let mut values = vec![Value::Integer(i as i64)];
247        if i % 11 == 4 {
248            values.extend(std::iter::repeat_n(Value::Null, 18));
249            return values;
250        }
251        values.extend([
252            Value::Float(i as f64 * 0.5),
253            Value::Bool(i.is_multiple_of(2)),
254            Value::String(format!("row-{i}")),
255            Value::String(["alpha", "beta", "gamma"][i % 3].into()),
256            Value::Bytes(vec![i as u8, 0xFF, 0x00]),
257            Value::NaiveDateTime(NdbDateTime::from_micros(BASE_MICROS + i as i64)),
258            Value::DateTime(NdbDateTime::from_micros(BASE_MICROS - i as i64)),
259            Value::Integer(i as i64 * 1_000),
260            Value::Decimal(rust_decimal::Decimal::new(i as i64 * 25, 2)),
261            Value::Uuid(uuid::Uuid::from_u128(i as u128 + 1).to_string()),
262            Value::Ulid(ulid_text(i)),
263            Value::String(format!("POINT({i} 1)")),
264            Value::Array(vec![
265                Value::Float(i as f64),
266                Value::Float(0.5),
267                Value::Float(-1.25),
268            ]),
269            json_cell(i),
270            Value::Duration(NdbDuration::from_micros(i as i64 * 3)),
271            Value::Regex(format!("^r{i}$")),
272            Value::Array(vec![Value::Integer(i as i64), Value::String("a".into())]),
273            Value::String(format!("{{{i}: 0.5}}")),
274        ]);
275        values
276    }
277
278    /// Every column type reads from a flushed segment as the same `Value` the
279    /// live memtable reads for the same cell, nulls and the second block
280    /// included.
281    #[test]
282    fn a_flushed_cell_reads_as_its_memtable_value() {
283        let schema = every_type_schema();
284        let mut mt = ColumnarMemtable::new(&schema);
285        for i in 0..ROWS {
286            mt.append_row(&row(i)).expect("append");
287        }
288        let expected: Vec<Vec<Value>> = (0..ROWS)
289            .map(|i| mt.get_row(i).expect("read").expect("memtable row"))
290            .collect();
291
292        let (schema, columns, row_count) = mt.drain_optimized();
293        let segment = SegmentWriter::new(PROFILE_PLAIN, test_memory())
294            .write_segment(&schema, &columns, row_count, None)
295            .expect("write segment");
296        let reader = SegmentReader::open(&segment).expect("open segment");
297        let indices: Vec<usize> = (0..schema.columns.len()).collect();
298        let decoded = reader.read_columns(&indices, &[]).expect("read columns");
299
300        for (i, memtable_row) in expected.iter().enumerate() {
301            for ((col, def), want) in decoded.iter().zip(&schema.columns).zip(memtable_row) {
302                let got = decoded_cell_value(col, i, &def.column_type).expect("decodable cell");
303                assert_eq!(&got, want, "row {i} column '{}'", def.name);
304            }
305        }
306
307        // The memtable shapes the comparison relies on.
308        let first = &expected[1];
309        assert_eq!(first[10], Value::Uuid(uuid::Uuid::from_u128(2).to_string()));
310        assert_eq!(first[11], Value::Ulid(ulid_text(1)));
311        assert_eq!(first[12], Value::String("POINT(1 1)".into()));
312        assert_eq!(first[15], Value::Integer(3));
313        assert_eq!(first[16], Value::String("^r1$".into()));
314        assert_eq!(expected[4][10], Value::Null);
315    }
316
317    /// A cell whose bytes do not hold its declared type is an error, not a
318    /// default value.
319    #[test]
320    fn a_corrupt_binary_cell_is_an_error() {
321        let col = crate::reader::DecodedColumn::Binary {
322            data: vec![1, 2, 3],
323            offsets: vec![0, 3],
324            valid: vec![true],
325        };
326        for ty in [
327            ColumnType::Uuid,
328            ColumnType::Ulid,
329            ColumnType::Decimal(None),
330            ColumnType::Vector(2),
331        ] {
332            assert!(
333                decoded_cell_value(&col, 0, &ty).is_err(),
334                "{ty} must reject a 3-byte cell"
335            );
336        }
337        // 0xC1 is the one byte MessagePack never uses.
338        let not_msgpack = crate::reader::DecodedColumn::Binary {
339            data: vec![0xC1],
340            offsets: vec![0, 1],
341            valid: vec![true],
342        };
343        assert!(decoded_cell_value(&not_msgpack, 0, &ColumnType::Json).is_err());
344        let not_utf8 = crate::reader::DecodedColumn::Binary {
345            data: vec![0xFF],
346            offsets: vec![0, 1],
347            valid: vec![true],
348        };
349        assert!(decoded_cell_value(&not_utf8, 0, &ColumnType::String).is_err());
350        assert_eq!(
351            decoded_cell_value(&not_utf8, 0, &ColumnType::Bytes).expect("bytes"),
352            Value::Bytes(vec![0xFF])
353        );
354    }
355}
nodedb_columnar/reader/cell.rs

nodedb_columnar/reader/
cell.rs
