proof_of_sql/sql/proof_plans/
table_exec.rs

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