uqa-execution 0.4.5

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
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124

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

//! Missing generated fields borrow only their selected expression inputs.

use super::RowLayout;
use crate::query::generated::{
    evaluate_generated_expression_with_control, prepare_generated_column_with_control,
    GeneratedControl,
};
use crate::query::table_snapshot::documents::projection::visit_source_projection;
use uqa_core::{
    memory::{BudgetedMap, BudgetedVec, MemoryReservation, ProductionControl},
    DocId, Value,
};
use uqa_sql::{ast::ColumnDef, expr::RowLookup, schema::ColumnTypeSchema, SQLError};
use uqa_storage::{DocumentStore, StorageBackendResult};

struct ProjectedInput<'a> {
    slots: &'a BudgetedMap<&'a str, usize>,
    values: &'a [&'a Value],
}

impl RowLookup for ProjectedInput<'_> {
    fn column(&self, name: &str) -> Option<&Value> {
        self.slots.get(name).map(|index| self.values[*index])
    }

    fn qualified_column(&self, _: &str, _: &str) -> Option<&Value> {
        None
    }
}

impl RowLayout {
    pub(super) fn generated_field(
        &self,
        source: &dyn DocumentStore,
        id: DocId,
        column: &ColumnDef,
        memory: &mut MemoryReservation,
    ) -> StorageBackendResult<Option<Value>> {
        assert!(self.control.memory().shares_allowance(memory.budget()));
        self.control.check()?;
        let production = ProductionControl::new(
            self.control.memory(),
            self.control.cancellation(),
            self.control.cancellation(),
        );
        let present = source.contains_doc_id(id)?;
        self.control.check()?;
        if !present {
            return Ok(None);
        }
        let schema = ColumnTypeSchema::new(&self.columns);
        let expression = prepare_generated_column_with_control(
            &schema,
            column.generated.as_ref().expect("selected generated field"),
            &GeneratedControl {
                budget: self.control.memory(),
                original: self.control.cancellation(),
                invoking: self.control.cancellation(),
            },
        )
        .map_err(Self::error)?;
        let mut names = BudgetedVec::new(self.control.memory());
        let mut slots = BudgetedMap::new(self.control.memory());
        let projectable = expression.scalar.try_visit_columns(&mut |name| {
            self.control.check()?;
            if !slots.contains_key(name) {
                names.reserve(1)?;
                slots.insert(name, names.len())?;
                names.push(name)?;
            }
            Ok::<_, uqa_storage::StorageBackendError>(())
        })?;
        if !projectable {
            return Err(Self::error(SQLError::Internal(
                "validated generated expression requires a relational row shape".into(),
            )));
        }
        let projection = self.generated_inputs(&names)?;
        let mut output = None;
        let read = visit_source_projection(
            &self.control,
            source,
            &[id],
            &projection,
            &[],
            &mut |_, present, values| {
                if present {
                    let row = ProjectedInput {
                        slots: &slots,
                        values,
                    };
                    output = Some(
                        evaluate_generated_expression_with_control(&expression.scalar, &row, &production)
                            .and_then(|value| {
                                uqa_sql::assignment::conversion::convert_value_to_column_type_with_control(
                                    value, &column.ty, &production,
                                )
                            })
                            .map_err(Self::error),
                    );
                }
                false
            },
        );
        // Keep the evaluator's first typed failure if the provider observes a later cancellation while returning.
        let output = output.transpose()?;
        read?;
        self.control.check()?;
        Ok(output.map(|value| {
            let (value, allocation) = value
                .into_budgeted()
                .expect("controlled generated field retains its allowance")
                .into_parts();
            memory.absorb(allocation);
            value
        }))
    }
}