uqa_sql/semantics/partition/
tree.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! The partitions of a partitioned table in the order `PostgreSQL` visits them when it derives objects on them.
8
9use std::collections::BTreeSet;
10
11use super::{partition_bound_order, PartitionContext};
12use crate::SQLError;
13
14/// A partition and the partitioned table it belongs to.
15#[derive(Debug, Clone, PartialEq, Eq)]
16pub struct PartitionTreeNode {
17    pub table: String,
18    pub object_id: [u8; 16],
19    pub parent: String,
20    pub parent_object_id: [u8; 16],
21}
22
23fn object_id(context: &PartitionContext<'_>, table: &str) -> Result<[u8; 16], SQLError> {
24    context
25        .catalog
26        .try_table_object_id(table)
27        .map_err(|error| SQLError::Internal(format!("read table identity: {error}")))?
28        .ok_or_else(|| SQLError::UnknownTable(table.to_string()))
29}
30
31/// The partitions below `root`, and `root` itself first when `include_root` names a partition, each before its own partitions, with siblings in partition order.
32pub fn partition_tree(
33    context: &PartitionContext<'_>,
34    root: &str,
35    include_root: bool,
36) -> Result<Vec<PartitionTreeNode>, SQLError> {
37    let mut output = Vec::new();
38    let mut visited = BTreeSet::new();
39    let parent = if include_root {
40        let hierarchy = context
41            .catalog
42            .try_table_hierarchy(root)
43            .map_err(|error| SQLError::Internal(format!("read partition hierarchy: {error}")))?;
44        let parent = hierarchy
45            .parents
46            .first()
47            .filter(|_| hierarchy.is_partition())
48            .cloned()
49            .ok_or_else(|| SQLError::Internal(format!("`{root}` is not a partition")))?;
50        let parent_object_id = object_id(context, &parent)?;
51        Some((parent, parent_object_id))
52    } else {
53        None
54    };
55    visit(context, root, parent, &mut visited, &mut output)?;
56    Ok(output)
57}
58
59fn visit(
60    context: &PartitionContext<'_>,
61    table: &str,
62    parent: Option<(String, [u8; 16])>,
63    visited: &mut BTreeSet<String>,
64    output: &mut Vec<PartitionTreeNode>,
65) -> Result<(), SQLError> {
66    if !visited.insert(table.to_string()) {
67        return Err(SQLError::Internal(format!(
68            "partition hierarchy cycle reaches `{table}`"
69        )));
70    }
71    let table_object_id = object_id(context, table)?;
72    if let Some((parent, parent_object_id)) = parent {
73        output.push(PartitionTreeNode {
74            table: table.to_string(),
75            object_id: table_object_id,
76            parent,
77            parent_object_id,
78        });
79    }
80    let hierarchy = context
81        .catalog
82        .try_table_hierarchy(table)
83        .map_err(|error| SQLError::Internal(format!("read partition hierarchy: {error}")))?;
84    if hierarchy.partition_spec.is_none() {
85        return Ok(());
86    }
87    let mut partitions = Vec::new();
88    for child in context.catalog.direct_hierarchy_children(table)? {
89        let bound = context
90            .catalog
91            .try_table_hierarchy(&child)
92            .map_err(|error| SQLError::Internal(format!("read partition hierarchy: {error}")))?
93            .partition_bound
94            .ok_or_else(|| {
95                SQLError::Internal(format!("partition `{child}` of `{table}` has no bound"))
96            })?;
97        partitions.push((child, bound));
98    }
99    for child in partition_bound_order(context, partitions)? {
100        visit(
101            context,
102            &child,
103            Some((table.to_string(), table_object_id)),
104            visited,
105            output,
106        )?;
107    }
108    Ok(())
109}
uqa_sql/semantics/partition/tree.rs

uqa_sql/semantics/partition/
tree.rs
