uqa-execution 0.5.3

Volcano physical operators with row-batch pipelines
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109

//
// Unified Query Algebra
//
// Copyright (c) 2023-2026 Cognica, Inc.
//

//! Aggregate call metadata and comparison support retained across groups, frames and spill restoration.

use super::{AggregateAccumulator, AggregateStatePlan, SQLAggregateFunction};
use std::sync::Arc;
use uqa_sql::ast::ColumnType;
use uqa_sql::expr::enums::EnumComparisonState;

#[derive(Clone)]
pub struct AggregateAccumulatorTemplate {
    kind: AggregateKind,
    extrema: Option<Arc<EnumComparisonState>>,
    ordering: Option<Arc<[EnumComparisonState]>>,
}

#[derive(Clone)]
enum AggregateKind {
    Builtin(AggregateStatePlan),
    Registered(Arc<dyn SQLAggregateFunction>),
}

impl AggregateAccumulatorTemplate {
    pub(crate) fn builtin(name: &str, input_type: Option<&ColumnType>) -> Self {
        Self {
            kind: AggregateKind::Builtin(AggregateStatePlan::builtin_with_input_type(
                name, input_type,
            )),
            extrema: enum_extrema_state(name, input_type),
            ordering: None,
        }
    }

    pub(super) fn generic() -> Self {
        Self {
            kind: AggregateKind::Builtin(AggregateStatePlan::Generic),
            extrema: None,
            ordering: None,
        }
    }

    pub(super) fn registered(function: Arc<dyn SQLAggregateFunction>) -> Self {
        Self {
            kind: AggregateKind::Registered(function),
            extrema: None,
            ordering: None,
        }
    }

    pub(super) fn with_ordering(mut self, keys: usize) -> Self {
        if keys != 0 {
            self.ordering = Some((0..keys).map(|_| EnumComparisonState::default()).collect());
        }
        self
    }

    pub(super) fn state_plan(&self) -> Option<AggregateStatePlan> {
        match self.kind {
            AggregateKind::Builtin(plan) => Some(plan),
            AggregateKind::Registered(_) => None,
        }
    }

    pub(super) fn instantiate(&self, budget_bytes: usize) -> AggregateAccumulator {
        let mut accumulator = match &self.kind {
            AggregateKind::Builtin(plan) => {
                AggregateAccumulator::from_plan_with_budget(*plan, budget_bytes)
            }
            AggregateKind::Registered(function) => {
                AggregateAccumulator::registered_with_budget(Arc::clone(function), budget_bytes)
            }
        };
        self.restore_comparison(&mut accumulator);
        accumulator
    }

    pub(super) fn restore_comparison(&self, accumulator: &mut AggregateAccumulator) {
        accumulator.enum_comparison.clone_from(&self.extrema);
        accumulator
            .values
            .comparison_states
            .clone_from(&self.ordering);
        accumulator
            .distinct
            .values
            .comparison_states
            .clone_from(&self.ordering);
        accumulator
            .registered_ordered
            .comparison_states
            .clone_from(&self.ordering);
    }
}

pub(super) fn enum_extrema_state(
    name: &str,
    input_type: Option<&ColumnType>,
) -> Option<Arc<EnumComparisonState>> {
    if let Some(ColumnType::Domain { base, .. }) = input_type {
        return enum_extrema_state(name, Some(base));
    }
    ((name.eq_ignore_ascii_case("min") || name.eq_ignore_ascii_case("max"))
        && matches!(input_type, Some(ColumnType::Enum(_))))
    .then(|| Arc::new(EnumComparisonState::default()))
}