uqa_sql/semantics/
cte_strategy.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Shared CTE scheduling decisions for execution and traversal analysis.
8
9use crate::{
10    ast::CteMaterialization,
11    plan::{CtePlan, QueryPlan},
12    semantics::{
13        cte_definition_references, cte_references_own_name, ordered_plan_ctes,
14        reachable_plan_cte_names, single_reference_plan_cte_names,
15        volatility::{query_contains_volatile_function, VolatilityCatalog},
16    },
17    SQLError,
18};
19use std::collections::{BTreeMap, BTreeSet};
20
21pub struct ScheduledCte<'a> {
22    pub plan: &'a CtePlan,
23    pub deferred: bool,
24}
25
26pub fn schedule_plan_ctes<'a>(
27    catalog: &dyn VolatilityCatalog,
28    plan: &'a QueryPlan,
29) -> Result<Vec<ScheduledCte<'a>>, SQLError> {
30    let ordered = ordered_plan_ctes(plan)?;
31    let reachable = reachable_plan_cte_names(plan);
32    let single_reference = single_reference_plan_cte_names(plan);
33    Ok(ordered
34        .into_iter()
35        .filter(|cte| reachable.contains(&cte.name))
36        .map(|cte| ScheduledCte {
37            plan: cte,
38            deferred: !cte.body.modifies_data()
39                && !cte_references_own_name(cte)
40                && match cte.materialization {
41                    CteMaterialization::Default => single_reference.contains(&cte.name),
42                    CteMaterialization::Materialized => false,
43                    CteMaterialization::NotMaterialized => true,
44                }
45                && matches!(
46                    cte.body
47                        .query()
48                        .map_or(Ok(true), |query| query_contains_volatile_function(
49                            catalog, query
50                        )),
51                    Ok(false)
52                ),
53        })
54        .collect())
55}
56
57/// The materialized WITH items of a statement, split by when they run. `PostgreSQL` runs a data-modifying item while its reader draws rows from it, and once the primary query has finished `ExecPostprocessPlan` runs every item no reader drew from to completion, the last defined first.
58pub struct StatementCteOrder<'a> {
59    /// The items that run before the primary query: the ones it reads, directly or through other items, and the queries that no postponed item reads.
60    pub primary: Vec<&'a CtePlan>,
61    /// The items that run once the primary query has finished, each data-modifying item after the items it reads.
62    pub postponed: Vec<&'a CtePlan>,
63}
64
65/// Split `scheduled`, the items of `ctes` that a statement materializes in dependency order, into the items that run before its primary query and the items that run after it. `primary_references` names the items the primary query reads.
66pub fn order_statement_ctes<'a>(
67    ctes: &'a [CtePlan],
68    scheduled: Vec<&'a CtePlan>,
69    primary_references: &BTreeSet<String>,
70) -> StatementCteOrder<'a> {
71    if !scheduled.iter().any(|cte| cte.body.modifies_data()) {
72        return StatementCteOrder {
73            primary: scheduled,
74            postponed: Vec::new(),
75        };
76    }
77    let references = ctes
78        .iter()
79        .enumerate()
80        .map(|(index, cte)| (cte.name.as_str(), cte_definition_references(ctes, index)))
81        .collect::<BTreeMap<_, _>>();
82    // The items `roots` read, directly or through the items they read, with the roots themselves.
83    let read_through = |roots: Vec<String>| {
84        let mut reached = BTreeSet::new();
85        let mut pending = roots;
86        while let Some(name) = pending.pop() {
87            if let Some(names) = references.get(name.as_str()) {
88                pending.extend(
89                    names
90                        .iter()
91                        .filter(|name| !reached.contains(*name))
92                        .cloned(),
93                );
94            }
95            reached.insert(name);
96        }
97        reached
98    };
99    let read = read_through(primary_references.iter().cloned().collect());
100    let mut postponed_names = BTreeSet::new();
101    let mut postponed = Vec::new();
102    for item in ctes
103        .iter()
104        .rev()
105        .filter(|cte| cte.body.modifies_data() && !read.contains(&cte.name))
106    {
107        let needed = read_through(vec![item.name.clone()]);
108        for cte in &scheduled {
109            if needed.contains(&cte.name)
110                && !read.contains(&cte.name)
111                && postponed_names.insert(cte.name.clone())
112            {
113                postponed.push(*cte);
114            }
115        }
116    }
117    let primary = scheduled
118        .into_iter()
119        .filter(|cte| !postponed_names.contains(&cte.name))
120        .collect();
121    StatementCteOrder { primary, postponed }
122}
123
124#[cfg(test)]
125mod tests;
uqa_sql/semantics/cte_strategy.rs

uqa_sql/semantics/
cte_strategy.rs
