exoware_sql/
types.rs

1use std::collections::{HashMap, HashSet};
2use std::sync::Arc;
3
4use datafusion::arrow::datatypes::{i256, DataType, Field, Schema, SchemaRef, TimeUnit};
5use exoware_sdk::keys::{Key, Prefix};
6use exoware_sdk::PrefixedStoreClient;
7
8use crate::codec::{primary_key_prefix, secondary_index_prefix};
9
10/// Every table/index family is named by a single packed byte
11/// `(table_prefix << 4) | discriminator`: the high nibble is the table prefix
12/// and the low nibble is the family discriminator. Discriminator `0x0` is the
13/// primary-row family and `0x1..=0xF` names secondary index slot
14/// `discriminator - 1`. All families of table `t` occupy the contiguous byte
15/// range `[t << 4, (t << 4) | 0xF]`; these ranges are pairwise disjoint within
16/// and across tables, and the primary discriminator (`0`) sorts below every
17/// index family of the same table.
18pub(crate) const FAMILY_PREFIX_LEN: usize = 1;
19/// Discriminator nibble for a table's primary-row family.
20pub(crate) const PRIMARY_FAMILY_DISCRIMINATOR: u8 = 0x00;
21/// Byte offset of the first payload byte within a primary/index key, measured
22/// from the store-stripped key (i.e. after the family prefix). Both families
23/// share the same one-byte prefix, so both offsets are `FAMILY_PREFIX_LEN`.
24pub(crate) const PRIMARY_KEY_BYTE_OFFSET: usize = FAMILY_PREFIX_LEN;
25pub(crate) const INDEX_KEY_BYTE_OFFSET: usize = FAMILY_PREFIX_LEN;
26/// Capacity caps for the packed one-byte family prefix. The table prefix
27/// occupies the high nibble and the discriminator the low nibble, giving up to
28/// 16 tables and 15 secondary index slots (discriminator `0x0` is reserved for
29/// the primary family).
30pub(crate) const MAX_TABLES: usize = 16;
31pub(crate) const MAX_INDEX_SPECS: usize = 15;
32pub(crate) const STRING_KEY_INLINE_LIMIT: usize = 15;
33pub(crate) const STRING_KEY_TERMINATOR: u8 = 0x00;
34pub(crate) const STRING_KEY_ESCAPE_PREFIX: u8 = 0x01;
35pub(crate) const STRING_KEY_ESCAPE_FF: u8 = 0x02;
36pub(crate) const PAGE_SIZE: usize = 1_000;
37pub(crate) const BATCH_FLUSH_ROWS: usize = 2_048;
38pub(crate) const INDEX_BACKFILL_FLUSH_ENTRIES: usize = 4_096;
39
40#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
41pub struct IndexBackfillReport {
42    pub scanned_rows: u64,
43    pub indexes_backfilled: usize,
44    pub index_entries_written: u64,
45}
46
47#[derive(Debug, Clone, PartialEq, Eq)]
48pub struct IndexBackfillOptions {
49    pub row_batch_size: usize,
50    pub start_from_primary_key: Option<Key>,
51}
52
53impl Default for IndexBackfillOptions {
54    fn default() -> Self {
55        Self {
56            row_batch_size: PAGE_SIZE,
57            start_from_primary_key: None,
58        }
59    }
60}
61
62#[derive(Debug, Clone, PartialEq, Eq)]
63pub enum IndexBackfillEvent {
64    Started {
65        table_name: String,
66        indexes_backfilled: usize,
67        row_batch_size: usize,
68        start_cursor: Key,
69    },
70    Progress {
71        scanned_rows: u64,
72        index_entries_written: u64,
73        last_scanned_primary_key: Key,
74        next_cursor: Option<Key>,
75    },
76    Completed {
77        report: IndexBackfillReport,
78    },
79}
80
81#[derive(Debug, Clone, Copy, PartialEq, Eq)]
82pub(crate) enum ListElementKind {
83    Int64,
84    Float64,
85    Boolean,
86    Utf8,
87}
88
89#[derive(Debug, Clone, Copy, PartialEq, Eq)]
90pub(crate) enum ColumnKind {
91    Int64,
92    UInt64,
93    Float64,
94    Boolean,
95    Utf8,
96    Date32,
97    Date64,
98    Timestamp,
99    Decimal128,
100    Decimal256,
101    FixedSizeBinary(usize),
102    Binary,
103    List(ListElementKind),
104}
105
106impl ColumnKind {
107    pub(crate) fn from_data_type(data_type: &DataType) -> Result<Self, String> {
108        match data_type {
109            DataType::Int64 => Ok(Self::Int64),
110            DataType::UInt64 => Ok(Self::UInt64),
111            DataType::Float64 => Ok(Self::Float64),
112            DataType::Boolean => Ok(Self::Boolean),
113            DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => Ok(Self::Utf8),
114            DataType::Date32 => Ok(Self::Date32),
115            DataType::Date64 => Ok(Self::Date64),
116            DataType::Timestamp(_, _) => Ok(Self::Timestamp),
117            DataType::Decimal128(_, _) => Ok(Self::Decimal128),
118            DataType::Decimal256(_, _) => Ok(Self::Decimal256),
119            DataType::FixedSizeBinary(n) => Ok(Self::FixedSizeBinary(*n as usize)),
120            DataType::Binary | DataType::LargeBinary | DataType::BinaryView => Ok(Self::Binary),
121            DataType::List(field) | DataType::LargeList(field) => {
122                let inner = Self::from_data_type(field.data_type())?;
123                let elem = match inner {
124                    Self::Int64 => ListElementKind::Int64,
125                    Self::Float64 => ListElementKind::Float64,
126                    Self::Boolean => ListElementKind::Boolean,
127                    Self::Utf8 => ListElementKind::Utf8,
128                    _ => {
129                        return Err(format!(
130                            "unsupported list element type {:?}; \
131                             list elements must be Int64, Float64, Boolean, or Utf8",
132                            field.data_type()
133                        ))
134                    }
135                };
136                Ok(Self::List(elem))
137            }
138            other => Err(format!(
139                "unsupported column type {other:?}; supported: \
140                 Int64, UInt64, Float64, Boolean, Utf8, Date32, Date64, Timestamp, \
141                 Decimal128, Decimal256, FixedSizeBinary, Binary, List"
142            )),
143        }
144    }
145
146    pub(crate) fn fixed_key_width(self) -> Option<usize> {
147        match self {
148            Self::Int64 => Some(8),
149            Self::UInt64 => Some(8),
150            Self::Float64 => Some(8),
151            Self::Boolean => Some(1),
152            Self::Utf8 => None,
153            Self::Date32 => Some(4),
154            Self::Date64 => Some(8),
155            Self::Timestamp => Some(8),
156            Self::Decimal128 => Some(16),
157            Self::Decimal256 => Some(32),
158            Self::FixedSizeBinary(n) => Some(n),
159            Self::Binary => None,
160            Self::List(_) => None,
161        }
162    }
163
164    pub(crate) fn key_width(self) -> usize {
165        self.fixed_key_width()
166            .unwrap_or(STRING_KEY_INLINE_LIMIT + 1)
167    }
168
169    pub(crate) fn indexable(self) -> bool {
170        !matches!(self, Self::List(_) | Self::Binary)
171    }
172}
173
174#[derive(Debug, Clone)]
175pub struct TableColumnConfig {
176    pub name: String,
177    pub data_type: DataType,
178    pub nullable: bool,
179}
180
181impl TableColumnConfig {
182    pub fn new(name: impl Into<String>, data_type: DataType, nullable: bool) -> Self {
183        Self {
184            name: name.into(),
185            data_type,
186            nullable,
187        }
188    }
189}
190
191#[derive(Debug, Clone, Copy, PartialEq, Eq)]
192pub enum IndexLayout {
193    Lexicographic,
194    ZOrder,
195}
196
197#[derive(Debug, Clone)]
198pub struct IndexSpec {
199    name: String,
200    key_columns: Vec<String>,
201    cover_columns: Vec<String>,
202    layout: IndexLayout,
203}
204
205impl IndexSpec {
206    #[cfg(test)]
207    pub(crate) fn new(name: impl Into<String>, key_columns: Vec<String>) -> Result<Self, String> {
208        Self::lexicographic(name, key_columns)
209    }
210
211    pub fn lexicographic(
212        name: impl Into<String>,
213        key_columns: Vec<String>,
214    ) -> Result<Self, String> {
215        let name = name.into();
216        if name.trim().is_empty() {
217            return Err("index name must not be empty".to_string());
218        }
219        if key_columns.is_empty() {
220            return Err("key_columns must not be empty".to_string());
221        }
222        Ok(Self {
223            name,
224            key_columns,
225            cover_columns: Vec::new(),
226            layout: IndexLayout::Lexicographic,
227        })
228    }
229
230    pub fn z_order(name: impl Into<String>, key_columns: Vec<String>) -> Result<Self, String> {
231        Self::lexicographic(name, key_columns).map(|spec| spec.with_layout(IndexLayout::ZOrder))
232    }
233
234    pub fn with_cover_columns(mut self, cover_columns: Vec<String>) -> Self {
235        self.cover_columns = cover_columns;
236        self
237    }
238
239    pub fn with_layout(mut self, layout: IndexLayout) -> Self {
240        self.layout = layout;
241        self
242    }
243
244    pub fn name(&self) -> &str {
245        &self.name
246    }
247
248    pub fn key_columns(&self) -> &[String] {
249        &self.key_columns
250    }
251
252    pub fn cover_columns(&self) -> &[String] {
253        &self.cover_columns
254    }
255
256    pub fn layout(&self) -> &IndexLayout {
257        &self.layout
258    }
259}
260
261pub fn default_orders_index_specs() -> Vec<IndexSpec> {
262    vec![IndexSpec::lexicographic(
263        "region_customer",
264        vec!["region".to_string(), "customer_id".to_string()],
265    )
266    .expect("default orders index must be valid")]
267}
268
269#[derive(Debug, Clone)]
270pub(crate) struct KvTableConfig {
271    pub(crate) table_prefix: u8,
272    pub(crate) columns: Vec<TableColumnConfig>,
273    pub(crate) primary_key_columns: Vec<String>,
274    pub(crate) index_specs: Vec<IndexSpec>,
275}
276
277impl KvTableConfig {
278    pub(crate) fn new(
279        table_prefix: u8,
280        columns: Vec<TableColumnConfig>,
281        primary_key_columns: Vec<String>,
282        index_specs: Vec<IndexSpec>,
283    ) -> Result<Self, String> {
284        if usize::from(table_prefix) >= MAX_TABLES {
285            return Err(format!(
286                "table prefix {table_prefix} exceeds max {} for key layout",
287                MAX_TABLES - 1
288            ));
289        }
290        if columns.is_empty() {
291            return Err("table config requires at least one column".to_string());
292        }
293        if primary_key_columns.is_empty() {
294            return Err("primary key must have at least one column".to_string());
295        }
296
297        let mut seen = HashSet::new();
298        let mut col_kinds = HashMap::new();
299        for col in &columns {
300            if col.name.trim().is_empty() {
301                return Err("column name must not be empty".to_string());
302            }
303            if !seen.insert(col.name.clone()) {
304                return Err(format!("duplicate column '{}'", col.name));
305            }
306            let kind = ColumnKind::from_data_type(&col.data_type)?;
307            col_kinds.insert(col.name.clone(), kind);
308        }
309
310        let mut total_pk_width = 0usize;
311        for pk_col in &primary_key_columns {
312            let kind = col_kinds
313                .get(pk_col)
314                .ok_or_else(|| format!("primary key column '{pk_col}' not found"))?;
315            match kind {
316                ColumnKind::Int64
317                | ColumnKind::UInt64
318                | ColumnKind::Utf8
319                | ColumnKind::FixedSizeBinary(_) => {}
320                _ => {
321                    return Err(format!(
322                        "primary key column '{pk_col}' must be Int64, UInt64, Utf8, or FixedSizeBinary"
323                    ));
324                }
325            }
326            total_pk_width += kind.key_width();
327        }
328        if total_pk_width > primary_key_prefix(table_prefix)?.max_payload_len() {
329            return Err(format!(
330                "composite primary key is too wide ({total_pk_width} bytes) for key payload"
331            ));
332        }
333
334        Ok(Self {
335            table_prefix,
336            columns,
337            primary_key_columns,
338            index_specs,
339        })
340    }
341
342    pub(crate) fn to_schema(&self) -> SchemaRef {
343        Arc::new(Schema::new(
344            self.columns
345                .iter()
346                .map(|col| {
347                    let dt = match &col.data_type {
348                        DataType::Timestamp(_, tz) => {
349                            DataType::Timestamp(TimeUnit::Microsecond, tz.clone())
350                        }
351                        DataType::LargeList(field) => DataType::List(field.clone()),
352                        other => other.clone(),
353                    };
354                    Field::new(&col.name, dt, col.nullable)
355                })
356                .collect::<Vec<_>>(),
357        ))
358    }
359}
360
361#[derive(Debug, Clone)]
362pub(crate) struct ResolvedColumn {
363    pub(crate) name: String,
364    pub(crate) kind: ColumnKind,
365    pub(crate) nullable: bool,
366}
367
368#[derive(Debug, Clone)]
369pub(crate) struct ResolvedIndexSpec {
370    pub(crate) id: u8,
371    pub(crate) prefix: Prefix,
372    pub(crate) name: String,
373    pub(crate) layout: IndexLayout,
374    pub(crate) key_columns: Vec<usize>,
375    pub(crate) value_column_mask: Vec<bool>,
376    pub(crate) key_columns_width: usize,
377}
378
379#[derive(Debug, Clone)]
380pub(crate) struct TableModel {
381    pub(crate) table_prefix: u8,
382    pub(crate) primary_key_prefix: Prefix,
383    pub(crate) schema: SchemaRef,
384    pub(crate) columns: Vec<ResolvedColumn>,
385    pub(crate) columns_by_name: HashMap<String, usize>,
386    pub(crate) primary_key_indices: Vec<usize>,
387    pub(crate) primary_key_kinds: Vec<ColumnKind>,
388    pub(crate) primary_key_width: usize,
389}
390
391impl TableModel {
392    pub(crate) fn from_config(config: &KvTableConfig) -> Result<Self, String> {
393        let schema = config.to_schema();
394        let mut columns = Vec::with_capacity(config.columns.len());
395        let mut columns_by_name = HashMap::with_capacity(config.columns.len());
396
397        for (idx, col) in config.columns.iter().enumerate() {
398            let kind = ColumnKind::from_data_type(&col.data_type)?;
399            columns.push(ResolvedColumn {
400                name: col.name.clone(),
401                kind,
402                nullable: col.nullable,
403            });
404            columns_by_name.insert(col.name.clone(), idx);
405        }
406
407        let mut primary_key_indices = Vec::with_capacity(config.primary_key_columns.len());
408        let mut primary_key_kinds = Vec::with_capacity(config.primary_key_columns.len());
409        let mut primary_key_width = 0usize;
410        for pk_col in &config.primary_key_columns {
411            let idx = *columns_by_name
412                .get(pk_col)
413                .ok_or_else(|| format!("primary key column '{pk_col}' not found"))?;
414            let kind = columns[idx].kind;
415            primary_key_indices.push(idx);
416            primary_key_kinds.push(kind);
417            primary_key_width += kind.key_width();
418        }
419
420        Ok(Self {
421            table_prefix: config.table_prefix,
422            primary_key_prefix: primary_key_prefix(config.table_prefix)?,
423            schema,
424            columns,
425            columns_by_name,
426            primary_key_indices,
427            primary_key_kinds,
428            primary_key_width,
429        })
430    }
431
432    /// Whether a column index is part of the primary key.
433    pub(crate) fn is_pk_column(&self, col_idx: usize) -> bool {
434        self.primary_key_indices.contains(&col_idx)
435    }
436
437    pub(crate) fn pk_position(&self, col_idx: usize) -> Option<usize> {
438        self.primary_key_indices
439            .iter()
440            .position(|&idx| idx == col_idx)
441    }
442
443    pub(crate) fn resolve_index_specs(
444        &self,
445        specs: &[IndexSpec],
446    ) -> Result<Vec<ResolvedIndexSpec>, String> {
447        let mut out = Vec::with_capacity(specs.len());
448        let mut names = HashSet::new();
449
450        for (idx, spec) in specs.iter().enumerate() {
451            if !names.insert(spec.name.clone()) {
452                return Err(format!("duplicate index name '{}'", spec.name));
453            }
454
455            let id = u8::try_from(idx + 1).map_err(|_| {
456                format!("too many index specs for key layout (max {MAX_INDEX_SPECS})")
457            })?;
458            if usize::from(id) > MAX_INDEX_SPECS {
459                return Err(format!(
460                    "too many index specs for key layout (max {MAX_INDEX_SPECS})"
461                ));
462            }
463            let mut key_columns = Vec::with_capacity(spec.key_columns.len());
464            let mut key_columns_width = 0usize;
465            let mut value_column_mask = vec![false; self.columns.len()];
466            for col_name in &spec.key_columns {
467                let Some(col_idx) = self.columns_by_name.get(col_name).copied() else {
468                    return Err(format!(
469                        "index '{}' references unknown column '{}'",
470                        spec.name, col_name
471                    ));
472                };
473                if !self.columns[col_idx].kind.indexable() {
474                    return Err(format!(
475                        "index '{}' references non-indexable column '{}'",
476                        spec.name, col_name
477                    ));
478                }
479                if self.columns[col_idx].nullable {
480                    return Err(format!(
481                        "index '{}' references nullable column '{}'; \
482                         nullable columns cannot be used in index keys",
483                        spec.name, col_name
484                    ));
485                }
486                // Z-order stored keys interleave the actual encoded bytes, but
487                // decode and bound construction assume the reserved fixed key
488                // width, so a variable-width z-order key column silently
489                // matches no rows.
490                if spec.layout == IndexLayout::ZOrder
491                    && self.columns[col_idx].kind.fixed_key_width().is_none()
492                {
493                    return Err(format!(
494                        "index '{}' z-order key column '{}' must have a \
495                         fixed-width kind; variable-width kinds (e.g. Utf8) \
496                         cannot be used in z-order index keys",
497                        spec.name, col_name
498                    ));
499                }
500                key_columns.push(col_idx);
501                key_columns_width += self.columns[col_idx].kind.key_width();
502                if !self.is_pk_column(col_idx) {
503                    value_column_mask[col_idx] = true;
504                }
505            }
506
507            for col_name in &spec.cover_columns {
508                let Some(col_idx) = self.columns_by_name.get(col_name).copied() else {
509                    return Err(format!(
510                        "index '{}' cover list references unknown column '{}'",
511                        spec.name, col_name
512                    ));
513                };
514                if self.is_pk_column(col_idx) {
515                    return Err(format!(
516                        "index '{}' cover column '{}' is a primary key column; \
517                         PK columns are always available from key bytes",
518                        spec.name, col_name
519                    ));
520                }
521                if !value_column_mask[col_idx] {
522                    value_column_mask[col_idx] = true;
523                }
524            }
525            let prefix = secondary_index_prefix(self.table_prefix, id)?;
526            if key_columns_width + self.primary_key_width > prefix.max_payload_len() {
527                return Err(format!(
528                    "index '{}' key layout too wide for key payload",
529                    spec.name
530                ));
531            }
532
533            out.push(ResolvedIndexSpec {
534                id,
535                prefix,
536                name: spec.name.clone(),
537                layout: spec.layout,
538                key_columns,
539                value_column_mask,
540                key_columns_width,
541            });
542        }
543
544        Ok(out)
545    }
546
547    pub(crate) fn column(&self, index: usize) -> &ResolvedColumn {
548        &self.columns[index]
549    }
550}
551
552#[derive(Debug, Clone)]
553pub enum CellValue {
554    Null,
555    Int64(i64),
556    UInt64(u64),
557    Float64(f64),
558    Boolean(bool),
559    Date32(i32),
560    Date64(i64),
561    Timestamp(i64),
562    Decimal128(i128),
563    Decimal256(i256),
564    Utf8(String),
565    FixedBinary(Vec<u8>),
566    Binary(Vec<u8>),
567    List(Vec<CellValue>),
568}
569
570#[derive(Debug, Clone)]
571pub(crate) struct KvRow {
572    pub(crate) values: Vec<CellValue>,
573}
574
575impl KvRow {
576    pub(crate) fn primary_key_values(&self, model: &TableModel) -> Vec<&CellValue> {
577        model
578            .primary_key_indices
579            .iter()
580            .map(|&idx| &self.values[idx])
581            .collect()
582    }
583
584    pub(crate) fn value_at(&self, idx: usize) -> &CellValue {
585        &self.values[idx]
586    }
587}
588
589#[derive(Debug, Clone, Default)]
590pub(crate) struct DecodedIndexEntry {
591    pub(crate) primary_key: Key,
592    pub(crate) primary_key_values: Vec<CellValue>,
593    pub(crate) values: HashMap<usize, CellValue>,
594}
595
596#[derive(Debug, Clone, PartialEq)]
597pub(crate) struct KeyRange {
598    pub(crate) start: Key,
599    pub(crate) end: Key,
600}
601
602#[derive(Debug, Clone)]
603pub(crate) struct IndexPlan {
604    pub(crate) spec_idx: usize,
605    pub(crate) ranges: Vec<KeyRange>,
606    pub(crate) constrained_prefix_len: usize,
607    pub(crate) constrained_column_count: usize,
608}
609
610#[derive(Debug, Clone)]
611pub(crate) struct KvTable {
612    pub(crate) client: PrefixedStoreClient,
613    pub(crate) model: Arc<TableModel>,
614    pub(crate) index_specs: Arc<Vec<ResolvedIndexSpec>>,
615}
616
617impl KvTable {
618    pub(crate) fn new(client: PrefixedStoreClient, config: KvTableConfig) -> Result<Self, String> {
619        let model = Arc::new(TableModel::from_config(&config)?);
620        let index_specs = Arc::new(model.resolve_index_specs(&config.index_specs)?);
621        Ok(Self {
622            client,
623            model,
624            index_specs,
625        })
626    }
627}
exoware_sql/types.rs

exoware_sql/
types.rs
