cobble 0.1.0

A flexible embedded key-value storage engine for distributed systems as well as single-node applications.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147

use crate::error::{Error, Result};
use bytes::{BufMut, Bytes, BytesMut};
use std::convert::TryInto;

const PARQUET_META_MAGIC: &[u8; 4] = b"pqm1";

#[derive(Clone, Debug, PartialEq, Eq)]
pub(crate) struct ParquetRowGroupRange {
    pub(crate) start_key: Vec<u8>,
    pub(crate) end_key: Vec<u8>,
}

#[derive(Clone, Debug)]
pub(crate) struct ParquetMeta {
    row_count: u64,
    row_groups: Vec<ParquetRowGroupRange>,
}

impl ParquetMeta {
    pub(crate) fn new(row_count: u64, row_groups: Vec<ParquetRowGroupRange>) -> Self {
        Self {
            row_count,
            row_groups,
        }
    }

    pub(crate) fn row_count(&self) -> u64 {
        self.row_count
    }

    pub(crate) fn row_groups(&self) -> &[ParquetRowGroupRange] {
        &self.row_groups
    }

    pub(crate) fn encode(self) -> Bytes {
        let mut out = BytesMut::with_capacity(
            4 + 8
                + 4
                + self
                    .row_groups
                    .iter()
                    .map(|group| 8 + group.start_key.len() + group.end_key.len())
                    .sum::<usize>(),
        );
        out.extend_from_slice(PARQUET_META_MAGIC);
        out.put_u64_le(self.row_count);
        out.put_u32_le(self.row_groups.len() as u32);
        for group in self.row_groups {
            out.put_u32_le(group.start_key.len() as u32);
            out.extend_from_slice(&group.start_key);
            out.put_u32_le(group.end_key.len() as u32);
            out.extend_from_slice(&group.end_key);
        }
        out.freeze()
    }

    fn decode(bytes: &[u8]) -> Result<Self> {
        if bytes.len() < 16 {
            return Err(Error::IoError(format!(
                "Invalid parquet meta size: expected at least 16, got {}",
                bytes.len()
            )));
        }
        if &bytes[..4] != PARQUET_META_MAGIC {
            return Err(Error::IoError("Invalid parquet meta magic".to_string()));
        }

        let row_count = u64::from_le_bytes(
            bytes[4..12]
                .try_into()
                .map_err(|_| Error::IoError("Invalid parquet meta row_count".to_string()))?,
        );
        let row_group_count = u32::from_le_bytes(
            bytes[12..16]
                .try_into()
                .map_err(|_| Error::IoError("Invalid parquet meta row_group_count".to_string()))?,
        ) as usize;

        let mut offset = 16usize;
        let mut row_groups = Vec::with_capacity(row_group_count);
        for _ in 0..row_group_count {
            if offset + 4 > bytes.len() {
                return Err(Error::IoError(
                    "Invalid parquet meta: missing start key length".to_string(),
                ));
            }
            let start_len =
                u32::from_le_bytes(bytes[offset..offset + 4].try_into().map_err(|_| {
                    Error::IoError("Invalid parquet meta start key length".to_string())
                })?) as usize;
            offset += 4;
            if offset + start_len > bytes.len() {
                return Err(Error::IoError(
                    "Invalid parquet meta: start key out of range".to_string(),
                ));
            }
            let start_key = bytes[offset..offset + start_len].to_vec();
            offset += start_len;

            if offset + 4 > bytes.len() {
                return Err(Error::IoError(
                    "Invalid parquet meta: missing end key length".to_string(),
                ));
            }
            let end_len =
                u32::from_le_bytes(bytes[offset..offset + 4].try_into().map_err(|_| {
                    Error::IoError("Invalid parquet meta end key length".to_string())
                })?) as usize;
            offset += 4;
            if offset + end_len > bytes.len() {
                return Err(Error::IoError(
                    "Invalid parquet meta: end key out of range".to_string(),
                ));
            }
            let end_key = bytes[offset..offset + end_len].to_vec();
            offset += end_len;

            row_groups.push(ParquetRowGroupRange { start_key, end_key });
        }
        if offset != bytes.len() {
            return Err(Error::IoError(
                "Invalid parquet meta: trailing bytes detected".to_string(),
            ));
        }
        Ok(Self {
            row_count,
            row_groups,
        })
    }
}

pub(crate) fn decode_meta(meta_bytes: Option<Bytes>) -> Result<Option<ParquetMeta>> {
    let Some(meta) = meta_bytes else {
        return Ok(None);
    };
    ParquetMeta::decode(meta.as_ref()).map(Some)
}

pub(crate) fn decode_meta_row_count(meta_bytes: Option<Bytes>) -> Result<Option<u64>> {
    Ok(decode_meta(meta_bytes)?.map(|meta| meta.row_count()))
}

pub(crate) fn decode_meta_row_group_ranges(
    meta_bytes: Option<Bytes>,
) -> Result<Option<Vec<ParquetRowGroupRange>>> {
    Ok(decode_meta(meta_bytes)?.map(|meta| meta.row_groups().to_vec()))
}