uqa-execution 0.3.0

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
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143

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

//! Owned physical query results and public result conversion.

use crate::physical::physical_exec_error;
use uqa_core::Value;
use uqa_sql::{ResultRow, SQLError, SQLResult};

mod exists;
pub use exists::collect_exists_key_operator;

pub enum QueryRows {
    Rows {
        named: Vec<ResultRow>,
        positional: Option<Vec<Vec<Value>>>,
    },
    SharedSpill(crate::SharedSpill),
    ExistsKeySet(crate::CanonicalRowHashSet),
}

pub struct QueryOutput {
    pub columns: Vec<String>,
    pub column_types: Vec<Option<uqa_sql::ast::ColumnType>>,
    /// Physical columns include internal row metadata that is available to a parent query block but never exposed through [`SQLResult`].
    pub internal_columns: Vec<String>,
    pub internal_types: Vec<Option<uqa_sql::ast::ColumnType>>,
    pub rows: QueryRows,
}

impl QueryOutput {
    pub fn into_cursor(self) -> Result<crate::query::cursor::SQLCursor, SQLError> {
        match self.rows {
            QueryRows::SharedSpill(rows) => {
                crate::query::cursor::SQLCursor::from_spill(self.columns, self.column_types, rows)
            }
            QueryRows::Rows { .. } | QueryRows::ExistsKeySet(_) => Err(SQLError::Internal(
                "cursor query unexpectedly used unbounded row materialization".into(),
            )),
        }
    }

    pub fn into_sql_result(self) -> Result<SQLResult, SQLError> {
        let (rows, positional_rows) = match self.rows {
            QueryRows::Rows { named, positional } => (named, positional),
            QueryRows::SharedSpill(rows) => {
                let mut scan = crate::SharedSpillScan::new(rows);
                (
                    crate::physical::run_to_rows(&mut scan)
                        .map_err(physical_exec_error)?
                        .1,
                    None,
                )
            }
            QueryRows::ExistsKeySet(_) => {
                return Err(SQLError::Internal(
                    "EXISTS key-set output cannot become a SQL result".into(),
                ));
            }
        };
        Ok(SQLResult::from_typed_rows_with_positions(
            self.columns,
            self.column_types,
            rows,
            positional_rows,
        ))
    }

    pub fn into_operator<'a>(self) -> Box<dyn crate::PhysicalOperator + 'a> {
        match self.rows {
            QueryRows::Rows { named, .. } => Box::new(crate::TableScan::from_typed_rows(
                self.internal_columns,
                self.internal_types,
                named,
            )),
            QueryRows::SharedSpill(rows) => Box::new(crate::SharedSpillScan::new(rows)),
            QueryRows::ExistsKeySet(_) => {
                panic!("EXISTS key-set output cannot become a physical operator")
            }
        }
    }

    pub fn into_public_operator<'a>(self) -> Box<dyn crate::PhysicalOperator + 'a> {
        let columns = self.columns.clone();
        let public_width = columns.len();
        let operator = self.into_operator();
        let positions = columns
            .into_iter()
            .enumerate()
            .map(|(position, column)| (column, position))
            .collect();
        debug_assert!(operator.row_schema().len() >= public_width);
        Box::new(crate::ColumnSelection::with_positions(operator, positions))
    }

    pub fn into_subquery_result(self) -> Result<crate::SubqueryResult, SQLError> {
        let QueryOutput {
            columns,
            column_types: _,
            internal_columns,
            internal_types,
            rows,
        } = self;
        let rows: Box<dyn Iterator<Item = Result<crate::OwnedPhysicalRow, SQLError>> + Send> =
            match rows {
                QueryRows::Rows { named, .. } => {
                    let schema = crate::RowSchema::with_types(internal_columns, internal_types);
                    Box::new(named.into_iter().map(move |row| {
                        Ok(crate::OwnedPhysicalRow::new(
                            schema.clone(),
                            crate::PhysicalRow::from_result_row(&schema, row),
                        ))
                    }))
                }
                QueryRows::SharedSpill(rows) => Box::new(
                    rows.read_rows()
                        .map_err(physical_exec_error)?
                        .map(|row| row.map_err(physical_exec_error)),
                ),
                QueryRows::ExistsKeySet(_) => {
                    return Err(SQLError::Internal(
                        "EXISTS key-set output cannot become a scalar subquery result".into(),
                    ));
                }
            };
        Ok(crate::SubqueryResult { columns, rows })
    }
}

pub fn query_output_shared(
    output: QueryOutput,
    label: &str,
) -> Result<crate::SharedSpill, SQLError> {
    let QueryRows::SharedSpill(rows) = output.rows else {
        return Err(SQLError::Internal(format!(
            "{label} execution returned in-memory rows at an internal streaming boundary"
        )));
    };
    Ok(rows)
}