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    table_ref: TableRef,
26    /// Schema of the table
27    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    /// Get the table reference
38    #[must_use]
39    pub fn table_ref(&self) -> &TableRef {
40        &self.table_ref
41    }
42
43    /// Get the schema
44    #[must_use]
45    pub fn schema(&self) -> &[ColumnField] {
46        &self.schema
47    }
48}
49
50impl ProofPlan for TableExec {
51    #[expect(unused_variables)]
52    fn verifier_evaluate<S: Scalar>(
53        &self,
54        builder: &mut impl VerificationBuilder<S>,
55        accessor: &IndexMap<ColumnRef, S>,
56        _result: Option<&OwnedTable<S>>,
57        chi_eval_map: &IndexMap<TableRef, S>,
58        params: &[LiteralValue],
59    ) -> Result<TableEvaluation<S>, ProofError> {
60        let column_evals = self
61            .schema
62            .iter()
63            .map(|field| {
64                let column_ref =
65                    ColumnRef::new(self.table_ref.clone(), field.name(), field.data_type());
66                *accessor.get(&column_ref).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}