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    ast::{ColumnDef, Expr, TableHierarchy},
11    type_resolution::FunctionTypeResolver,
12    ResultRow, RowSchema, SQLError, SQLParam,
13};
14use std::cmp::Ordering;
15use uqa_core::Value;
16
17/// Hierarchy definitions needed to validate and route a row; no storage writes are exposed.
18pub trait PartitionCatalog {
19    fn try_table_hierarchy(&self, table: &str) -> Result<TableHierarchy, String>;
20    fn direct_hierarchy_children(&self, parent: &str) -> Result<Vec<String>, SQLError>;
21    fn try_resolve_table_name(&self, name: &str) -> Result<Option<String>, String>;
22    fn try_describe_table(&self, table: &str) -> Result<Option<Vec<ColumnDef>>, String>;
23    /// The object identity of a table, which survives renames.
24    fn try_table_object_id(&self, table: &str) -> Result<Option<[u8; 16]>, String>;
25}
26
27/// Evaluate declared partition keys and bounds with the caller's expression scope.
28pub trait PartitionExpressions {
29    fn evaluate_bound(&self, expression: &Expr, params: &[SQLParam]) -> Result<Value, SQLError>;
30    fn evaluate_row(
31        &self,
32        expression: &Expr,
33        row: &ResultRow,
34        schema: &RowSchema,
35        params: &[SQLParam],
36    ) -> Result<Value, SQLError>;
37}
38
39#[derive(Clone, Copy)]
40pub struct PartitionContext<'a> {
41    pub catalog: &'a dyn PartitionCatalog,
42    pub expressions: &'a dyn PartitionExpressions,
43    pub types: &'a dyn FunctionTypeResolver,
44}
45
46mod hash;
47
48pub fn validate_hash_partition_spec(
49    context: &PartitionContext<'_>,
50    spec: &crate::ast::PartitionSpec,
51    columns: &[crate::ast::ColumnDef],
52) -> Result<(), SQLError> {
53    hash::validate_partition_spec(context.types, spec, columns)
54}
55
56pub fn validate_new_partition_bound(
57    context: &PartitionContext<'_>,
58    parent: &str,
59    bound: &crate::ast::PartitionBound,
60) -> Result<(), SQLError> {
61    let hierarchy = context
62        .catalog
63        .try_table_hierarchy(parent)
64        .map_err(|error| SQLError::Internal(format!("read parent partition metadata: {error}")))?;
65    let spec = hierarchy
66        .partition_spec
67        .as_ref()
68        .ok_or_else(|| SQLError::Routine {
69            sqlstate: "42809".into(),
70            message: format!("relation \"{parent}\" is not partitioned"),
71        })?;
72    validate_partition_bound_width(spec, bound)?;
73    if let crate::ast::PartitionBound::Hash { modulus, remainder } = bound {
74        hash::validate_bound(*modulus, *remainder)?;
75        let mut existing_moduli = Vec::new();
76        for sibling in context.catalog.direct_hierarchy_children(parent)? {
77            let sibling_hierarchy = context
78                .catalog
79                .try_table_hierarchy(&sibling)
80                .map_err(|error| SQLError::Internal(format!("read sibling partition: {error}")))?;
81            match sibling_hierarchy.partition_bound.as_ref() {
82                Some(crate::ast::PartitionBound::Hash { modulus, remainder }) => {
83                    hash::validate_bound(*modulus, *remainder)?;
84                    existing_moduli.push(*modulus);
85                }
86                Some(crate::ast::PartitionBound::Default) => {
87                    return Err(SQLError::Internal(format!(
88                        "HASH-partitioned table `{parent}` has a default partition"
89                    )))
90                }
91                Some(_) => {
92                    return Err(SQLError::Internal(
93                        "partition siblings use different bound strategies".into(),
94                    ))
95                }
96                None => {}
97            }
98        }
99        hash::validate_modulus_chain(*modulus, existing_moduli)?;
100    }
101    if let crate::ast::PartitionBound::Range { lower, upper } = bound {
102        if compare_partition_points(context, lower, upper)? != Ordering::Less {
103            return Err(invalid_partition_bound(
104                "empty range bound specified for partition",
105            ));
106        }
107    }
108    for sibling in context.catalog.direct_hierarchy_children(parent)? {
109        let sibling_hierarchy = context
110            .catalog
111            .try_table_hierarchy(&sibling)
112            .map_err(|error| SQLError::Internal(format!("read sibling partition: {error}")))?;
113        let Some(sibling_bound) = sibling_hierarchy.partition_bound.as_ref() else {
114            continue;
115        };
116        if partition_bounds_overlap(context, bound, sibling_bound)? {
117            return Err(invalid_partition_bound(format!(
118                "partition would overlap partition \"{sibling}\""
119            )));
120        }
121    }
122    Ok(())
123}
124
125/// Test one stored row against a prospective direct-child bound before the
126/// hierarchy edge is installed. DEFAULT means no existing non-default sibling
127/// accepts the row, matching the routing decision the new edge will expose.
128pub fn prospective_partition_bound_accepts_document(
129    context: &PartitionContext<'_>,
130    parent: &str,
131    bound: &crate::ast::PartitionBound,
132    document: &ResultRow,
133) -> Result<bool, SQLError> {
134    let hierarchy = context
135        .catalog
136        .try_table_hierarchy(parent)
137        .map_err(|error| SQLError::Internal(format!("read parent partition metadata: {error}")))?;
138    let spec = hierarchy
139        .partition_spec
140        .as_ref()
141        .ok_or_else(|| SQLError::Routine {
142            sqlstate: "42809".into(),
143            message: format!("relation \"{parent}\" is not partitioned"),
144        })?;
145    let (keys, row_hash) = partition_key_values_and_hash(context, parent, spec, document)?;
146    if !matches!(bound, crate::ast::PartitionBound::Default) {
147        return partition_bound_matches(context, bound, &keys, &[], row_hash);
148    }
149    for sibling in context.catalog.direct_hierarchy_children(parent)? {
150        let sibling_hierarchy = context
151            .catalog
152            .try_table_hierarchy(&sibling)
153            .map_err(|error| SQLError::Internal(format!("read child partition: {error}")))?;
154        let Some(sibling_bound) = sibling_hierarchy.partition_bound.as_ref() else {
155            continue;
156        };
157        if matches!(sibling_bound, crate::ast::PartitionBound::Default) {
158            continue;
159        }
160        if partition_bound_matches(context, sibling_bound, &keys, &[], row_hash)? {
161            return Ok(false);
162        }
163    }
164    Ok(true)
165}
166
167/// Evaluate a retained detached-partition CHECK without requiring a live
168/// parent edge. This is also the exact predicate used for prospective ATTACH
169/// row scans.
170pub fn partition_constraint_accepts_document(
171    context: &PartitionContext<'_>,
172    table: &str,
173    spec: &crate::ast::PartitionSpec,
174    bound: &crate::ast::PartitionBound,
175    document: &ResultRow,
176) -> Result<bool, SQLError> {
177    let (keys, row_hash) = partition_key_values_and_hash(context, table, spec, document)?;
178    partition_bound_matches(context, bound, &keys, &[], row_hash)
179}
180
181fn partition_key_values_and_hash(
182    context: &PartitionContext<'_>,
183    table: &str,
184    spec: &crate::ast::PartitionSpec,
185    document: &ResultRow,
186) -> Result<(Vec<Value>, Option<u64>), SQLError> {
187    let (keys, definitions) = evaluate_partition_keys(context, table, &spec.keys, document, &[])?;
188    let row_hash = (spec.strategy == crate::ast::PartitionStrategy::Hash)
189        .then(|| hash::row_hash(context.types, spec, &definitions, &keys))
190        .transpose()?;
191    Ok((keys, row_hash))
192}
193
194fn validate_partition_bound_width(
195    spec: &crate::ast::PartitionSpec,
196    bound: &crate::ast::PartitionBound,
197) -> Result<(), SQLError> {
198    use crate::ast::{PartitionBound, PartitionStrategy};
199    match (spec.strategy, bound) {
200        (_, PartitionBound::Default) => Ok(()),
201        (PartitionStrategy::List, PartitionBound::List(_)) if spec.keys.len() != 1 => {
202            Err(invalid_partition_bound(
203                "cannot use list partition bounds with more than one partition key",
204            ))
205        }
206        (PartitionStrategy::List, PartitionBound::List(_)) => Ok(()),
207        (PartitionStrategy::Range, PartitionBound::Range { lower, upper })
208            if lower.len() != spec.keys.len() || upper.len() != spec.keys.len() =>
209        {
210            Err(invalid_partition_bound(
211                "partition bound has the wrong number of columns",
212            ))
213        }
214        (PartitionStrategy::Range, PartitionBound::Range { .. })
215        | (PartitionStrategy::Hash, PartitionBound::Hash { .. }) => Ok(()),
216        (strategy, _) => Err(invalid_partition_bound(format!(
217            "invalid bound specification for a {} partitioned table",
218            match strategy {
219                PartitionStrategy::List => "list",
220                PartitionStrategy::Range => "range",
221                PartitionStrategy::Hash => "hash",
222            }
223        ))),
224    }
225}
226
227fn partition_bounds_overlap(
228    context: &PartitionContext<'_>,
229    left: &crate::ast::PartitionBound,
230    right: &crate::ast::PartitionBound,
231) -> Result<bool, SQLError> {
232    use crate::ast::PartitionBound;
233    match (left, right) {
234        (PartitionBound::Default, PartitionBound::Default) => Ok(true),
235        (PartitionBound::Default, _) | (_, PartitionBound::Default) => Ok(false),
236        (PartitionBound::List(left), PartitionBound::List(right)) => {
237            let left = evaluate_bound_values(context, left)?;
238            let right = evaluate_bound_values(context, right)?;
239            Ok(left.iter().any(|value| right.contains(value)))
240        }
241        (
242            PartitionBound::Range {
243                lower: left_lower,
244                upper: left_upper,
245            },
246            PartitionBound::Range {
247                lower: right_lower,
248                upper: right_upper,
249            },
250        ) => Ok(
251            compare_partition_points(context, left_lower, right_upper)? == Ordering::Less
252                && compare_partition_points(context, right_lower, left_upper)? == Ordering::Less,
253        ),
254        (
255            PartitionBound::Hash {
256                modulus: left_modulus,
257                remainder: left_remainder,
258            },
259            PartitionBound::Hash {
260                modulus: right_modulus,
261                remainder: right_remainder,
262            },
263        ) => hash::bounds_overlap(
264            *left_modulus,
265            *left_remainder,
266            *right_modulus,
267            *right_remainder,
268        ),
269        _ => Err(SQLError::Internal(
270            "partition siblings use different bound strategies".into(),
271        )),
272    }
273}
274
275fn evaluate_bound_values(
276    context: &PartitionContext<'_>,
277    expressions: &[crate::ast::Expr],
278) -> Result<Vec<Value>, SQLError> {
279    expressions
280        .iter()
281        .map(|expression| context.expressions.evaluate_bound(expression, &[]))
282        .collect()
283}
284
285fn compare_partition_points(
286    context: &PartitionContext<'_>,
287    left: &[crate::ast::PartitionRangeDatum],
288    right: &[crate::ast::PartitionRangeDatum],
289) -> Result<Ordering, SQLError> {
290    if left.len() != right.len() {
291        return Err(invalid_partition_bound(
292            "partition range points have different widths",
293        ));
294    }
295    for (left, right) in left.iter().zip(right) {
296        let ordering = match (left, right) {
297            (
298                crate::ast::PartitionRangeDatum::MinValue,
299                crate::ast::PartitionRangeDatum::MinValue,
300            )
301            | (
302                crate::ast::PartitionRangeDatum::MaxValue,
303                crate::ast::PartitionRangeDatum::MaxValue,
304            ) => Ordering::Equal,
305            (crate::ast::PartitionRangeDatum::MinValue, _)
306            | (_, crate::ast::PartitionRangeDatum::MaxValue) => Ordering::Less,
307            (crate::ast::PartitionRangeDatum::MaxValue, _)
308            | (_, crate::ast::PartitionRangeDatum::MinValue) => Ordering::Greater,
309            (
310                crate::ast::PartitionRangeDatum::Value(left),
311                crate::ast::PartitionRangeDatum::Value(right),
312            ) => context
313                .expressions
314                .evaluate_bound(left, &[])?
315                .cmp(&context.expressions.evaluate_bound(right, &[])?),
316        };
317        if ordering != Ordering::Equal {
318            return Ok(ordering);
319        }
320    }
321    Ok(Ordering::Equal)
322}
323
324fn invalid_partition_bound(message: impl Into<String>) -> SQLError {
325    SQLError::Routine {
326        sqlstate: "42P17".into(),
327        message: message.into(),
328    }
329}
330
331/// 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.
332#[derive(Debug, Clone, PartialEq)]
333pub enum PartitionRejection {
334    /// 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`.
335    Constraint { relation: String },
336    /// No partition of `relation` accepts the row, whose partition key has the values `keys`: `no partition of relation "x" found for row`.
337    NoPartition { relation: String, keys: Vec<Value> },
338}
339
340/// Where a row that a statement writes through a relation is stored.
341#[derive(Debug, Clone, PartialEq)]
342pub enum PartitionRoute {
343    /// The relation that stores the row.
344    Target(String),
345    /// Routing rejects the row.
346    Rejected(PartitionRejection),
347}
348
349/// 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`].
350pub fn route_partition_insert(
351    context: &PartitionContext<'_>,
352    requested_table: &str,
353    document: &ResultRow,
354    params: &[SQLParam],
355    include_descendants: bool,
356) -> Result<PartitionRoute, SQLError> {
357    let table = context
358        .catalog
359        .try_resolve_table_name(requested_table)
360        .map_err(|error| SQLError::Internal(format!("resolve INSERT table: {error}")))?
361        .ok_or_else(|| SQLError::UnknownTable(requested_table.to_string()))?;
362    let hierarchy = context
363        .catalog
364        .try_table_hierarchy(&table)
365        .map_err(|error| SQLError::Internal(format!("read partition metadata: {error}")))?;
366    let Some(spec) = hierarchy.partition_spec.as_ref() else {
367        return Ok(PartitionRoute::Target(table));
368    };
369    if !include_descendants {
370        return Err(SQLError::Routine {
371            sqlstate: "42809".into(),
372            message: format!("cannot insert into partitioned table \"{requested_table}\""),
373        });
374    }
375    if !partition_constraint_accepts_row(context, &table, document, params)? {
376        return Ok(PartitionRoute::Rejected(PartitionRejection::Constraint {
377            relation: table,
378        }));
379    }
380    route_partition_tree(context, &table, spec, document, params)
381}
382
383/// 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.
384pub fn partition_constraint_accepts_row(
385    context: &PartitionContext<'_>,
386    table: &str,
387    document: &ResultRow,
388    params: &[SQLParam],
389) -> Result<bool, SQLError> {
390    let mut child = table.to_string();
391    let mut visited = std::collections::BTreeSet::new();
392    loop {
393        if !visited.insert(child.clone()) {
394            return Err(SQLError::Internal(format!(
395                "partition hierarchy cycle reaches `{child}`"
396            )));
397        }
398        let hierarchy = context
399            .catalog
400            .try_table_hierarchy(&child)
401            .map_err(|error| SQLError::Internal(format!("read partition metadata: {error}")))?;
402        if hierarchy.partition_bound.is_none() {
403            return Ok(true);
404        }
405        let parent = hierarchy.parents.first().ok_or_else(|| {
406            SQLError::Internal(format!("partition `{child}` has no parent relation"))
407        })?;
408        let parent_hierarchy = context
409            .catalog
410            .try_table_hierarchy(parent)
411            .map_err(|error| {
412                SQLError::Internal(format!("read parent partition metadata: {error}"))
413            })?;
414        let spec = parent_hierarchy.partition_spec.as_ref().ok_or_else(|| {
415            SQLError::Internal(format!("partition parent `{parent}` has no partition key"))
416        })?;
417        match select_direct_partition(context, parent, spec, document, params)? {
418            DirectPartition::Child(selected) if selected == child => {}
419            DirectPartition::Child(_) | DirectPartition::None { .. } => return Ok(false),
420        }
421        child.clone_from(parent);
422    }
423}
424
425fn route_partition_tree(
426    context: &PartitionContext<'_>,
427    table: &str,
428    spec: &crate::ast::PartitionSpec,
429    document: &ResultRow,
430    params: &[SQLParam],
431) -> Result<PartitionRoute, SQLError> {
432    let child = match select_direct_partition(context, table, spec, document, params)? {
433        DirectPartition::Child(child) => child,
434        DirectPartition::None { keys } => {
435            return Ok(PartitionRoute::Rejected(PartitionRejection::NoPartition {
436                relation: table.to_string(),
437                keys,
438            }))
439        }
440    };
441    let hierarchy = context
442        .catalog
443        .try_table_hierarchy(&child)
444        .map_err(|error| SQLError::Internal(format!("read partition metadata: {error}")))?;
445    match hierarchy.partition_spec.as_ref() {
446        Some(spec) => route_partition_tree(context, &child, spec, document, params),
447        None => Ok(PartitionRoute::Target(child)),
448    }
449}
450
451/// The direct partition of a partitioned table that accepts a row, or the row's partition key values when none does.
452enum DirectPartition {
453    Child(String),
454    None { keys: Vec<Value> },
455}
456
457fn select_direct_partition(
458    context: &PartitionContext<'_>,
459    parent: &str,
460    spec: &crate::ast::PartitionSpec,
461    document: &ResultRow,
462    params: &[SQLParam],
463) -> Result<DirectPartition, SQLError> {
464    let (keys, definitions) =
465        evaluate_partition_keys(context, parent, &spec.keys, document, params)?;
466    let row_hash = (spec.strategy == crate::ast::PartitionStrategy::Hash)
467        .then(|| hash::row_hash(context.types, spec, &definitions, &keys))
468        .transpose()?;
469    let mut default = None;
470    for child in context.catalog.direct_hierarchy_children(parent)? {
471        let child_hierarchy = context
472            .catalog
473            .try_table_hierarchy(&child)
474            .map_err(|error| SQLError::Internal(format!("read child partition: {error}")))?;
475        let Some(bound) = child_hierarchy.partition_bound.as_ref() else {
476            continue;
477        };
478        if matches!(bound, crate::ast::PartitionBound::Default) {
479            if default.replace(child).is_some() {
480                return Err(SQLError::Internal(format!(
481                    "partitioned table `{parent}` has more than one default partition"
482                )));
483            }
484            continue;
485        }
486        if partition_bound_matches(context, bound, &keys, params, row_hash)? {
487            return Ok(DirectPartition::Child(child));
488        }
489    }
490    Ok(default.map_or(DirectPartition::None { keys }, DirectPartition::Child))
491}
492
493fn evaluate_partition_keys(
494    context: &PartitionContext<'_>,
495    table: &str,
496    expressions: &[crate::ast::Expr],
497    document: &ResultRow,
498    params: &[SQLParam],
499) -> Result<(Vec<Value>, Vec<crate::ast::ColumnDef>), SQLError> {
500    let definitions = context
501        .catalog
502        .try_describe_table(table)
503        .map_err(|error| SQLError::Internal(format!("read partition row type: {error}")))?
504        .ok_or_else(|| SQLError::UnknownTable(table.to_string()))?;
505    let schema = crate::RowSchema::with_types(
506        definitions
507            .iter()
508            .map(|definition| definition.name.clone())
509            .collect(),
510        definitions
511            .iter()
512            .map(|definition| Some(definition.ty.clone()))
513            .collect(),
514    );
515    let values = expressions
516        .iter()
517        .map(|expression| {
518            context
519                .expressions
520                .evaluate_row(expression, document, &schema, params)
521        })
522        .collect::<Result<Vec<_>, _>>()?;
523    Ok((values, definitions))
524}
525
526fn partition_bound_matches(
527    context: &PartitionContext<'_>,
528    bound: &crate::ast::PartitionBound,
529    keys: &[Value],
530    params: &[SQLParam],
531    row_hash: Option<u64>,
532) -> Result<bool, SQLError> {
533    use crate::ast::PartitionBound;
534    match bound {
535        PartitionBound::Default => Ok(true),
536        PartitionBound::List(values) => {
537            let [key] = keys else {
538                return Err(SQLError::Internal(
539                    "LIST partition has more than one partition key".into(),
540                ));
541            };
542            for expression in values {
543                if context.expressions.evaluate_bound(expression, params)? == *key {
544                    return Ok(true);
545                }
546            }
547            Ok(false)
548        }
549        PartitionBound::Range { lower, upper } => {
550            if keys.iter().any(|value| matches!(value, Value::Null)) {
551                return Ok(false);
552            }
553            Ok(
554                compare_key_to_bound(context, keys, lower, params)? != Ordering::Less
555                    && compare_key_to_bound(context, keys, upper, params)? == Ordering::Less,
556            )
557        }
558        PartitionBound::Hash { modulus, remainder } => hash::bound_matches(
559            row_hash.ok_or_else(|| {
560                SQLError::Internal("HASH partition bound has no computed row hash".into())
561            })?,
562            *modulus,
563            *remainder,
564        ),
565    }
566}
567
568fn compare_key_to_bound(
569    context: &PartitionContext<'_>,
570    keys: &[Value],
571    bound: &[crate::ast::PartitionRangeDatum],
572    params: &[SQLParam],
573) -> Result<Ordering, SQLError> {
574    if keys.len() != bound.len() {
575        return Err(SQLError::Internal(format!(
576            "partition key width {} differs from bound width {}",
577            keys.len(),
578            bound.len()
579        )));
580    }
581    for (key, datum) in keys.iter().zip(bound) {
582        let ordering = match datum {
583            crate::ast::PartitionRangeDatum::MinValue => Ordering::Greater,
584            crate::ast::PartitionRangeDatum::MaxValue => Ordering::Less,
585            crate::ast::PartitionRangeDatum::Value(expression) => {
586                key.cmp(&context.expressions.evaluate_bound(expression, params)?)
587            }
588        };
589        if ordering != Ordering::Equal {
590            return Ok(ordering);
591        }
592    }
593    Ok(Ordering::Equal)
594}
595
596mod identity;
597mod order;
598pub use order::partition_bound_order;
599mod tree;
600pub use identity::{
601    foreign_key_scan_tables, partition_ancestor_tables, partition_hierarchy_root,
602    partition_identity_owner,
603};
604pub use tree::{partition_tree, PartitionTreeNode};