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

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

//! Retrieval-tree optimization, active-statement checks and physical execution.

use super::driver::{context::PhysicalDriverContext, PhysicalRetrievalDriver};
use super::{OperatorOutput, PlanExecutor};
use uqa_core::PostingList;
use uqa_operators::{OperatorTree, TextScoringMode};
use uqa_sql::{SQLError, SQLParam};

type DriverResult<T> = Result<T, SQLError>;

pub trait RetrievalPlanOptimizer: Sync {
    fn optimize(&self, table: &str, tree: &OperatorTree) -> DriverResult<OperatorTree>;
}

pub trait RetrievalTransactionState: Sync {
    fn transaction_depth(&self) -> usize;
}

#[derive(Clone, Copy)]
pub struct TreeExecutionContext<'a> {
    pub driver: PhysicalDriverContext<'a>,
    pub optimizer: &'a dyn RetrievalPlanOptimizer,
    pub transaction: &'a dyn RetrievalTransactionState,
}

/// Execute below an existing statement boundary; calibration must already have a transaction.
pub fn execute_tree(
    context: &TreeExecutionContext<'_>,
    table: &str,
    signal_table: &str,
    params: &[SQLParam],
    tree: &OperatorTree,
) -> DriverResult<OperatorOutput> {
    if context.transaction.transaction_depth() == 0 && tree_may_persist_calibration(tree) {
        return Err(SQLError::Internal(
            "calibrating operator execution requires an active statement transaction".into(),
        ));
    }
    optimize_and_execute_tree(context, table, signal_table, params, tree)
}

pub fn optimize_and_execute_tree(
    context: &TreeExecutionContext<'_>,
    table: &str,
    signal_table: &str,
    params: &[SQLParam],
    tree: &OperatorTree,
) -> DriverResult<OperatorOutput> {
    validate_text_top_k_placement(tree)?;
    let optimized = context.optimizer.optimize(table, tree)?;
    execute_physical_tree(context, table, signal_table, params, &optimized)
}

pub fn execute_preoptimized_tree(
    context: &TreeExecutionContext<'_>,
    table: &str,
    signal_table: &str,
    params: &[SQLParam],
    tree: &OperatorTree,
) -> DriverResult<OperatorOutput> {
    if context.transaction.transaction_depth() == 0 && tree_may_persist_calibration(tree) {
        return Err(SQLError::Internal(
            "calibrating operator execution requires an active statement transaction".into(),
        ));
    }
    execute_physical_tree(context, table, signal_table, params, tree)
}

fn execute_physical_tree(
    context: &TreeExecutionContext<'_>,
    table: &str,
    signal_table: &str,
    params: &[SQLParam],
    tree: &OperatorTree,
) -> DriverResult<OperatorOutput> {
    validate_text_top_k_placement(tree)?;
    let driver = PhysicalRetrievalDriver::new(context.driver, table, signal_table, params);
    let mut executor = PlanExecutor::new(&driver);
    executor.execute(tree)
}

fn validate_text_top_k_placement(tree: &OperatorTree) -> DriverResult<()> {
    let root_is_physical_text = matches!(tree, OperatorTree::Term { top_k: Some(_), .. });
    let mut physical_text_nodes = 0_usize;
    tree.visit(&mut |node| {
        if matches!(node, OperatorTree::Term { top_k: Some(_), .. }) {
            physical_text_nodes += 1;
        }
    });
    if physical_text_nodes == usize::from(root_is_physical_text) {
        Ok(())
    } else {
        Err(SQLError::Internal(
            "physical text top-k is valid only as the root retrieval leaf".into(),
        ))
    }
}

pub fn tree_may_persist_calibration(tree: &OperatorTree) -> bool {
    let mut may_persist = false;
    tree.visit(&mut |node| {
        may_persist |= matches!(
            node,
            OperatorTree::BayesianScore { .. }
                | OperatorTree::Term {
                    scoring: Some(TextScoringMode::BayesianBM25),
                    ..
                }
                | OperatorTree::Phrase {
                    scoring: Some(TextScoringMode::BayesianBM25),
                    ..
                }
                | OperatorTree::BayesianMatchWithPrior { .. }
                | OperatorTree::MultiFieldSearch { .. }
        );
    });
    may_persist
}

pub fn expect_posting_output(output: OperatorOutput, context: &str) -> DriverResult<PostingList> {
    match output {
        OperatorOutput::Posting(result) => Ok(result),
        OperatorOutput::Graph(result) => Ok(result.to_posting_list()),
        OperatorOutput::Generalized(_) => Err(SQLError::TypeMismatch(format!(
            "{context} requires single-document rows, but the physical plan produced join tuples"
        ))),
    }
}

#[cfg(test)]
mod tests;