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

exoware_sql/
types.rs
