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