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

exoware_sql/
types.rs
