Skip to main content

uqa_planner/statement_planning/
prepared.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Custom-versus-generic plan selection and type-preserving parameter specialization.
8
9#[derive(Clone, Copy)]
10pub struct PreparedPlanUsage {
11    pub has_parameters: bool,
12    pub custom_plans: i64,
13    pub total_custom_cost: f64,
14}
15
16pub fn choose_custom_plan(usage: PreparedPlanUsage, mode: &str, generic_cost: Option<f64>) -> bool {
17    if !usage.has_parameters {
18        return false;
19    }
20    match mode {
21        "force_generic_plan" => false,
22        "force_custom_plan" => true,
23        _ if usage.custom_plans < 5 => true,
24        _ => generic_cost
25            .is_some_and(|cost| cost >= usage.total_custom_cost / usage.custom_plans as f64),
26    }
27}
28
29pub fn specialize_parameters(plan: &mut crate::UnifiedPlan, parameters: &[uqa_sql::SQLParam]) {
30    use uqa_sql::SQLParam;
31    use uqa_sql::ScalarExpr;
32    plan.rewrite_scalar_expressions(&mut |expression| {
33        let ScalarExpr::Param(index) = expression else {
34            return;
35        };
36        let Some(parameter) = index.checked_sub(1).and_then(|index| parameters.get(index)) else {
37            return;
38        };
39        *expression = match parameter {
40            SQLParam::TypedScalar { value, ty } => ScalarExpr::TypedLiteral {
41                value: value.clone(),
42                ty: ty.sql_name(),
43                bound_type: Some(ty.clone()),
44                parameter_index: Some(*index),
45            },
46            SQLParam::Scalar(value) => ScalarExpr::Literal(value.clone()),
47            SQLParam::Vector(_) | SQLParam::Tensor(_) => return,
48        };
49    });
50}
51
52pub mod selection;
53
54#[cfg(test)]
55mod tests;