proof_of_sql/sql/proof_plans/
table_exec.rs

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