proof_of_sql/sql/proof_plans/
table_exec.rs

1use crate::{
2    base::{
3        database::{
4            ColumnField, ColumnRef, LiteralValue, OwnedTable, Table, TableEvaluation, TableRef,
5        },
6        map::{indexset, IndexMap, IndexSet},
7        proof::{PlaceholderResult, ProofError},
8        scalar::Scalar,
9    },
10    sql::proof::{
11        FinalRoundBuilder, FirstRoundBuilder, ProofPlan, ProverEvaluate, VerificationBuilder,
12    },
13    utils::log,
14};
15use alloc::vec::Vec;
16use bumpalo::Bump;
17use serde::{Deserialize, Serialize};
18use sqlparser::ast::Ident;
19
20/// Source [`ProofPlan`] for (sub)queries with table source such as `SELECT col from tab;`
21/// Inspired by `DataFusion` data source [`ExecutionPlan`]s such as [`ArrowExec`] and [`CsvExec`].
22/// Note that we only need to load the columns we use.
23#[derive(Debug, PartialEq, Serialize, Deserialize, Clone)]
24pub struct TableExec {
25    /// Table reference
26    table_ref: TableRef,
27    /// Schema of the table
28    schema: Vec<ColumnField>,
29}
30
31impl TableExec {
32    /// Creates a new [`TableExec`].
33    #[must_use]
34    pub fn new(table_ref: TableRef, schema: Vec<ColumnField>) -> Self {
35        Self { table_ref, schema }
36    }
37
38    /// Get the table reference
39    #[must_use]
40    pub fn table_ref(&self) -> &TableRef {
41        &self.table_ref
42    }
43
44    /// Get the schema
45    #[must_use]
46    pub fn schema(&self) -> &[ColumnField] {
47        &self.schema
48    }
49}
50
51impl ProofPlan for TableExec {
52    #[expect(unused_variables)]
53    fn verifier_evaluate<S: Scalar>(
54        &self,
55        builder: &mut impl VerificationBuilder<S>,
56        accessor: &IndexMap<TableRef, IndexMap<Ident, S>>,
57        _result: Option<&OwnedTable<S>>,
58        chi_eval_map: &IndexMap<TableRef, (S, usize)>,
59        params: &[LiteralValue],
60    ) -> Result<TableEvaluation<S>, ProofError> {
61        let column_evals = self
62            .schema
63            .iter()
64            .map(|field| {
65                *accessor
66                    .get(self.table_ref())
67                    .expect("Table does not exist")
68                    .get(&field.name())
69                    .expect("Column does not exist")
70            })
71            .collect::<Vec<_>>();
72        let chi_eval = *chi_eval_map
73            .get(&self.table_ref)
74            .expect("Chi eval not found");
75        Ok(TableEvaluation::new(column_evals, chi_eval))
76    }
77
78    fn get_column_result_fields(&self) -> Vec<ColumnField> {
79        self.schema.clone()
80    }
81
82    fn get_column_references(&self) -> IndexSet<ColumnRef> {
83        self.schema
84            .iter()
85            .map(|field| ColumnRef::new(self.table_ref.clone(), field.name(), field.data_type()))
86            .collect()
87    }
88
89    fn get_table_references(&self) -> IndexSet<TableRef> {
90        indexset! {self.table_ref.clone()}
91    }
92}
93
94impl ProverEvaluate for TableExec {
95    #[tracing::instrument(name = "TableExec::first_round_evaluate", level = "debug", skip_all)]
96    fn first_round_evaluate<'a, S: Scalar>(
97        &self,
98        _builder: &mut FirstRoundBuilder<'a, S>,
99        _alloc: &'a Bump,
100        table_map: &IndexMap<TableRef, Table<'a, S>>,
101        _params: &[LiteralValue],
102    ) -> PlaceholderResult<Table<'a, S>> {
103        log::log_memory_usage("Start");
104
105        let first_round_table = table_map
106            .get(&self.table_ref)
107            .expect("Table not found")
108            .clone();
109
110        log::log_memory_usage("End");
111
112        Ok(first_round_table)
113    }
114
115    #[tracing::instrument(name = "TableExec::final_round_evaluate", level = "debug", skip_all)]
116    #[expect(unused_variables)]
117    fn final_round_evaluate<'a, S: Scalar>(
118        &self,
119        builder: &mut FinalRoundBuilder<'a, S>,
120        alloc: &'a Bump,
121        table_map: &IndexMap<TableRef, Table<'a, S>>,
122        _params: &[LiteralValue],
123    ) -> PlaceholderResult<Table<'a, S>> {
124        log::log_memory_usage("Start");
125
126        let final_round_table = table_map
127            .get(&self.table_ref)
128            .expect("Table not found")
129            .clone();
130
131        log::log_memory_usage("End");
132
133        Ok(final_round_table)
134    }
135}