Skip to main content

uqa_sql/semantics/
partition.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Declarative partition validation, value comparison, and routing semantics.
8
9use crate::{
10    assignment::AssignmentContext,
11    ast::{ColumnDef, Expr, TableHierarchy},
12    schema::SchemaExpressionCatalog,
13    type_resolution::FunctionTypeResolver,
14    ResultRow, RowSchema, SQLError, SQLParam,
15};
16use std::cmp::Ordering;
17use uqa_core::Value;
18
19/// Hierarchy definitions needed to validate and route a row; no storage writes are exposed.
20pub trait PartitionCatalog {
21    fn try_table_hierarchy(&self, table: &str) -> Result<TableHierarchy, String>;
22    fn direct_hierarchy_children(&self, parent: &str) -> Result<Vec<String>, SQLError>;
23    fn try_resolve_table_name(&self, name: &str) -> Result<Option<String>, String>;
24    fn try_describe_table(&self, table: &str) -> Result<Option<Vec<ColumnDef>>, String>;
25    /// The object identity of a table, which survives renames.
26    fn try_table_object_id(&self, table: &str) -> Result<Option<[u8; 16]>, String>;
27}
28
29/// Evaluate declared partition keys and bounds with the caller's expression scope.
30pub trait PartitionExpressions {
31    /// Reconstruct a stored key using its retained type and routine identities and current catalog names.
32    fn expression_text(&self, expression: &Expr) -> Result<String, SQLError>;
33    fn evaluate_bound(&self, expression: &Expr, params: &[SQLParam]) -> Result<Value, SQLError>;
34    fn evaluate_row(
35        &self,
36        expression: &Expr,
37        row: &ResultRow,
38        schema: &RowSchema,
39        params: &[SQLParam],
40    ) -> Result<Value, SQLError>;
41}
42
43#[derive(Clone, Copy)]
44pub struct PartitionContext<'a> {
45    pub catalog: &'a dyn PartitionCatalog,
46    pub expressions: &'a dyn PartitionExpressions,
47    pub types: &'a dyn FunctionTypeResolver,
48    /// Catalog-aware assignment coercion of bound values and output of key values.
49    pub assignment: &'a dyn AssignmentContext,
50    /// Aggregate, window and set-returning classification of bound and key expressions.
51    pub schema: &'a dyn SchemaExpressionCatalog,
52}
53
54mod admission;
55mod bounds;
56mod datum_text;
57mod hash;
58mod key;
59pub use key::key_type as partition_key_type;
60
61pub use admission::validate_new_partition_bound;
62pub use bounds::transform_partition_bound;
63pub use datum_text::{partition_datum_text, range_bound_text, stored_datum};
64
65pub fn validate_hash_partition_spec(
66    context: &PartitionContext<'_>,
67    spec: &crate::ast::PartitionSpec,
68    columns: &[crate::ast::ColumnDef],
69) -> Result<(), SQLError> {
70    hash::validate_partition_spec(context.types, spec, columns)
71}
72
73/// Test one stored row against a prospective direct-child bound before the
74/// hierarchy edge is installed. DEFAULT means no existing non-default sibling
75/// accepts the row, matching the routing decision the new edge will expose.
76pub fn prospective_partition_bound_accepts_document(
77    context: &PartitionContext<'_>,
78    parent: &str,
79    bound: &crate::ast::PartitionBound,
80    document: &ResultRow,
81) -> Result<bool, SQLError> {
82    let hierarchy = context
83        .catalog
84        .try_table_hierarchy(parent)
85        .map_err(|error| SQLError::Internal(format!("read parent partition metadata: {error}")))?;
86    let spec = hierarchy
87        .partition_spec
88        .as_ref()
89        .ok_or_else(|| SQLError::Routine {
90            sqlstate: "42809".into(),
91            message: format!("relation \"{parent}\" is not partitioned"),
92        })?;
93    let key = routing_key(context, parent, spec, document, &[])?;
94    if !matches!(bound, crate::ast::PartitionBound::Default) {
95        return partition_bound_matches(context, bound, &key, &[]);
96    }
97    for sibling in context.catalog.direct_hierarchy_children(parent)? {
98        let sibling_hierarchy = context
99            .catalog
100            .try_table_hierarchy(&sibling)
101            .map_err(|error| SQLError::Internal(format!("read child partition: {error}")))?;
102        let Some(sibling_bound) = sibling_hierarchy.partition_bound.as_ref() else {
103            continue;
104        };
105        if matches!(sibling_bound, crate::ast::PartitionBound::Default) {
106            continue;
107        }
108        if partition_bound_matches(context, sibling_bound, &key, &[])? {
109            return Ok(false);
110        }
111    }
112    Ok(true)
113}
114
115/// Evaluate a retained detached-partition CHECK without requiring a live
116/// parent edge. This is also the exact predicate used for prospective ATTACH
117/// row scans.
118pub fn partition_constraint_accepts_document(
119    context: &PartitionContext<'_>,
120    table: &str,
121    spec: &crate::ast::PartitionSpec,
122    bound: &crate::ast::PartitionBound,
123    document: &ResultRow,
124) -> Result<bool, SQLError> {
125    let key = routing_key(context, table, spec, document, &[])?;
126    partition_bound_matches(context, bound, &key, &[])
127}
128
129/// A row's partition key under one partitioned table, with the row hash of a HASH partition key.
130struct RoutingKey {
131    values: Vec<Value>,
132    hash: Option<u64>,
133}
134
135fn routing_key(
136    context: &PartitionContext<'_>,
137    table: &str,
138    spec: &crate::ast::PartitionSpec,
139    document: &ResultRow,
140    params: &[SQLParam],
141) -> Result<RoutingKey, SQLError> {
142    let (values, definitions) =
143        evaluate_partition_keys(context, table, &spec.keys, document, params)?;
144    let hash = (spec.strategy == crate::ast::PartitionStrategy::Hash)
145        .then(|| hash::row_hash(context, spec, &definitions, &values))
146        .transpose()?;
147    Ok(RoutingKey { values, hash })
148}
149
150/// Why partition routing rejects a row. The executor reports each with a description of the row, which depends on the statement and the role that writes it.
151#[derive(Debug, Clone, PartialEq)]
152pub enum PartitionRejection {
153    /// The row does not satisfy the partition constraint of `relation`, which a statement names and routes from: `new row for relation "x" violates partition constraint`.
154    Constraint { relation: String },
155    /// No partition of `relation` accepts the row, whose partition key has the values `keys`: `no partition of relation "x" found for row`.
156    NoPartition { relation: String, keys: Vec<Value> },
157}
158
159/// Where a row that a statement writes through a relation is stored.
160#[derive(Debug, Clone, PartialEq)]
161pub enum PartitionRoute {
162    /// The relation that stores the row.
163    Target(String),
164    /// Routing rejects the row.
165    Rejected(PartitionRejection),
166}
167
168/// Route a row that an `INSERT` writes through `requested_table`, as `ExecFindPartition` does: a partitioned table that is itself a partition first checks its own partition constraint, and each level then selects the partition whose bound accepts the row. A relation that is not partitioned stores the row itself, without a check: `ExecInsert` checks the partition constraint of a partition that the statement names after its other constraints, which the caller does with [`partition_constraint_accepts_row`].
169pub fn route_partition_insert(
170    context: &PartitionContext<'_>,
171    requested_table: &str,
172    document: &ResultRow,
173    params: &[SQLParam],
174    include_descendants: bool,
175) -> Result<PartitionRoute, SQLError> {
176    let table = context
177        .catalog
178        .try_resolve_table_name(requested_table)
179        .map_err(|error| SQLError::Internal(format!("resolve INSERT table: {error}")))?
180        .ok_or_else(|| SQLError::UnknownTable(requested_table.to_string()))?;
181    let hierarchy = context
182        .catalog
183        .try_table_hierarchy(&table)
184        .map_err(|error| SQLError::Internal(format!("read partition metadata: {error}")))?;
185    let Some(spec) = hierarchy.partition_spec.as_ref() else {
186        return Ok(PartitionRoute::Target(table));
187    };
188    if !include_descendants {
189        return Err(SQLError::Routine {
190            sqlstate: "42809".into(),
191            message: format!("cannot insert into partitioned table \"{requested_table}\""),
192        });
193    }
194    if !partition_constraint_accepts_row(context, &table, document, params)? {
195        return Ok(PartitionRoute::Rejected(PartitionRejection::Constraint {
196            relation: table,
197        }));
198    }
199    route_partition_tree(context, &table, spec, document, params)
200}
201
202/// Whether the row satisfies the partition constraint of `table`, which holds the bounds of the partition and of each of its partition ancestors (`ExecPartitionCheck`). A relation that is not a partition accepts every row.
203pub fn partition_constraint_accepts_row(
204    context: &PartitionContext<'_>,
205    table: &str,
206    document: &ResultRow,
207    params: &[SQLParam],
208) -> Result<bool, SQLError> {
209    let mut child = table.to_string();
210    let mut visited = std::collections::BTreeSet::new();
211    loop {
212        if !visited.insert(child.clone()) {
213            return Err(SQLError::Internal(format!(
214                "partition hierarchy cycle reaches `{child}`"
215            )));
216        }
217        let hierarchy = context
218            .catalog
219            .try_table_hierarchy(&child)
220            .map_err(|error| SQLError::Internal(format!("read partition metadata: {error}")))?;
221        if hierarchy.partition_bound.is_none() {
222            return Ok(true);
223        }
224        let parent = hierarchy.parents.first().ok_or_else(|| {
225            SQLError::Internal(format!("partition `{child}` has no parent relation"))
226        })?;
227        let parent_hierarchy = context
228            .catalog
229            .try_table_hierarchy(parent)
230            .map_err(|error| {
231                SQLError::Internal(format!("read parent partition metadata: {error}"))
232            })?;
233        let spec = parent_hierarchy.partition_spec.as_ref().ok_or_else(|| {
234            SQLError::Internal(format!("partition parent `{parent}` has no partition key"))
235        })?;
236        match select_direct_partition(context, parent, spec, document, params)? {
237            DirectPartition::Child(selected) if selected == child => {}
238            DirectPartition::Child(_) | DirectPartition::None { .. } => return Ok(false),
239        }
240        child.clone_from(parent);
241    }
242}
243
244fn route_partition_tree(
245    context: &PartitionContext<'_>,
246    table: &str,
247    spec: &crate::ast::PartitionSpec,
248    document: &ResultRow,
249    params: &[SQLParam],
250) -> Result<PartitionRoute, SQLError> {
251    let child = match select_direct_partition(context, table, spec, document, params)? {
252        DirectPartition::Child(child) => child,
253        DirectPartition::None { keys } => {
254            return Ok(PartitionRoute::Rejected(PartitionRejection::NoPartition {
255                relation: table.to_string(),
256                keys,
257            }))
258        }
259    };
260    let hierarchy = context
261        .catalog
262        .try_table_hierarchy(&child)
263        .map_err(|error| SQLError::Internal(format!("read partition metadata: {error}")))?;
264    match hierarchy.partition_spec.as_ref() {
265        Some(spec) => route_partition_tree(context, &child, spec, document, params),
266        None => Ok(PartitionRoute::Target(child)),
267    }
268}
269
270/// The direct partition of a partitioned table that accepts a row, or the row's partition key values when none does.
271enum DirectPartition {
272    Child(String),
273    None { keys: Vec<Value> },
274}
275
276fn select_direct_partition(
277    context: &PartitionContext<'_>,
278    parent: &str,
279    spec: &crate::ast::PartitionSpec,
280    document: &ResultRow,
281    params: &[SQLParam],
282) -> Result<DirectPartition, SQLError> {
283    let key = routing_key(context, parent, spec, document, params)?;
284    Ok(match select_partition(context, parent, &key, params)? {
285        Some(child) => DirectPartition::Child(child),
286        None => DirectPartition::None { keys: key.values },
287    })
288}
289
290fn select_partition(
291    context: &PartitionContext<'_>,
292    parent: &str,
293    key: &RoutingKey,
294    params: &[SQLParam],
295) -> Result<Option<String>, SQLError> {
296    let mut default = None;
297    for child in context.catalog.direct_hierarchy_children(parent)? {
298        let child_hierarchy = context
299            .catalog
300            .try_table_hierarchy(&child)
301            .map_err(|error| SQLError::Internal(format!("read child partition: {error}")))?;
302        let Some(bound) = child_hierarchy.partition_bound.as_ref() else {
303            continue;
304        };
305        if matches!(bound, crate::ast::PartitionBound::Default) {
306            if default.replace(child).is_some() {
307                return Err(SQLError::Internal(format!(
308                    "partitioned table `{parent}` has more than one default partition"
309                )));
310            }
311            continue;
312        }
313        if partition_bound_matches(context, bound, key, params)? {
314            return Ok(Some(child));
315        }
316    }
317    Ok(default)
318}
319
320fn evaluate_partition_keys(
321    context: &PartitionContext<'_>,
322    table: &str,
323    expressions: &[crate::ast::Expr],
324    document: &ResultRow,
325    params: &[SQLParam],
326) -> Result<(Vec<Value>, Vec<crate::ast::ColumnDef>), SQLError> {
327    let definitions = context
328        .catalog
329        .try_describe_table(table)
330        .map_err(|error| SQLError::Internal(format!("read partition row type: {error}")))?
331        .ok_or_else(|| SQLError::UnknownTable(table.to_string()))?;
332    let schema = crate::RowSchema::with_types(
333        definitions
334            .iter()
335            .map(|definition| definition.name.clone())
336            .collect(),
337        definitions
338            .iter()
339            .map(|definition| Some(definition.ty.clone()))
340            .collect(),
341    );
342    let values = expressions
343        .iter()
344        .map(|expression| {
345            context
346                .expressions
347                .evaluate_row(expression, document, &schema, params)
348        })
349        .collect::<Result<Vec<_>, _>>()?;
350    Ok((values, definitions))
351}
352
353fn partition_bound_matches(
354    context: &PartitionContext<'_>,
355    bound: &crate::ast::PartitionBound,
356    key: &RoutingKey,
357    params: &[SQLParam],
358) -> Result<bool, SQLError> {
359    use crate::ast::PartitionBound;
360    let keys = key.values.as_slice();
361    match bound {
362        PartitionBound::Default => Ok(true),
363        PartitionBound::List(values) => {
364            let [key] = keys else {
365                return Err(SQLError::Internal(
366                    "LIST partition has more than one partition key".into(),
367                ));
368            };
369            for expression in values {
370                if context
371                    .expressions
372                    .evaluate_bound(expression, params)?
373                    .cmp(key)
374                    == Ordering::Equal
375                {
376                    return Ok(true);
377                }
378            }
379            Ok(false)
380        }
381        PartitionBound::Range { lower, upper } => {
382            if keys.iter().any(|value| matches!(value, Value::Null)) {
383                return Ok(false);
384            }
385            Ok(
386                compare_key_to_bound(context, keys, lower, params)? != Ordering::Less
387                    && compare_key_to_bound(context, keys, upper, params)? == Ordering::Less,
388            )
389        }
390        PartitionBound::Hash { modulus, remainder } => hash::bound_matches(
391            key.hash.ok_or_else(|| {
392                SQLError::Internal("HASH partition bound has no computed row hash".into())
393            })?,
394            *modulus,
395            *remainder,
396        ),
397    }
398}
399
400fn compare_key_to_bound(
401    context: &PartitionContext<'_>,
402    keys: &[Value],
403    bound: &[crate::ast::PartitionRangeDatum],
404    params: &[SQLParam],
405) -> Result<Ordering, SQLError> {
406    if keys.len() != bound.len() {
407        return Err(SQLError::Internal(format!(
408            "partition key width {} differs from bound width {}",
409            keys.len(),
410            bound.len()
411        )));
412    }
413    for (key, datum) in keys.iter().zip(bound) {
414        let ordering = match datum {
415            crate::ast::PartitionRangeDatum::MinValue => Ordering::Greater,
416            crate::ast::PartitionRangeDatum::MaxValue => Ordering::Less,
417            crate::ast::PartitionRangeDatum::Value(expression) => {
418                key.cmp(&context.expressions.evaluate_bound(expression, params)?)
419            }
420        };
421        if ordering != Ordering::Equal {
422            return Ok(ordering);
423        }
424    }
425    Ok(Ordering::Equal)
426}
427
428mod identity;
429mod order;
430pub use order::partition_bound_order;
431mod tree;
432pub use identity::{
433    foreign_key_scan_tables, partition_ancestor_tables, partition_hierarchy_root,
434    partition_identity_owner,
435};
436pub use tree::{partition_tree, PartitionTreeNode};