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