refactor!: replace DataAccessor with Table in ProofExpr && remove input_length from ProofPlan::result_evaluate

crates/proof-of-sql/src/base/database/accessor.rs

@@ -1,6 +1,7 @@
 use crate::base::{
     commitment::Commitment,
-    database::{Column, ColumnRef, ColumnType, TableRef},
+    database::{Column, ColumnRef, ColumnType, Table, TableOptions, TableRef},
+    map::{IndexMap, IndexSet},
     scalar::Scalar,
 };
 use alloc::vec::Vec;
@@ -85,6 +86,28 @@ pub trait CommitmentAccessor<C: Commitment>: MetadataAccessor {
 pub trait DataAccessor<S: Scalar>: MetadataAccessor {
     /// Return the data span in the table (not the full-table data)
     fn get_column(&self, column: ColumnRef) -> Column<S>;
+
+    /// Creates a new [`Table`] from a [`TableRef`] and [`ColumnRef`]s.
+    ///
+    /// Columns are retrieved from the [`DataAccessor`] using the provided [`TableRef`] and [`ColumnRef`]s.
+    /// The only reason why [`table_ref` is needed is because `column_refs` can be empty.
+    /// # Panics
+    /// Column length mismatches can occur in theory. In practice, this should not happen.
+    fn get_table(&self, table_ref: TableRef, column_refs: &IndexSet<ColumnRef>) -> Table<S> {
+        if column_refs.is_empty() {
+            let input_length = self.get_length(table_ref);
+            Table::<S>::try_new_with_options(
+                IndexMap::default(),
+                TableOptions::new(Some(input_length)),
+            )
+        } else {
+            Table::<S>::try_from_iter(column_refs.into_iter().map(|column_ref| {
+                let column = self.get_column(*column_ref);
+                (column_ref.column_id(), column)
+            }))
+        }
+        .expect("Failed to create table from table and column references")
+    }
 }
 
 /// Access tables and their schemas in a database.

crates/proof-of-sql/src/base/database/mod.rs

@@ -65,9 +65,9 @@ mod owned_table_test;
 pub mod owned_table_utility;
 
 mod table;
-pub use table::Table;
 #[cfg(test)]
 pub(crate) use table::TableError;
+pub use table::{Table, TableOptions};
 #[cfg(test)]
 mod table_test;
 pub mod table_utility;

crates/proof-of-sql/src/base/database/table.rs

@@ -3,12 +3,36 @@ use crate::base::{map::IndexMap, scalar::Scalar};
 use proof_of_sql_parser::Identifier;
 use snafu::Snafu;
 
+/// Options for creating a table.
+/// Inspired by [`RecordBatchOptions`](https://docs.rs/arrow/latest/arrow/record_batch/struct.RecordBatchOptions.html)
+#[derive(Debug, Default, Clone, Copy)]
+pub struct TableOptions {
+    /// The number of rows in the table. Mostly useful for tables without columns.
+    pub row_count: Option<usize>,
+}
+
+impl TableOptions {
+    /// Creates a new [`TableOptions`].
+    #[must_use]
+    pub fn new(row_count: Option<usize>) -> Self {
+        Self { row_count }
+    }
+}
+
 /// An error that occurs when working with tables.
 #[derive(Snafu, Debug, PartialEq, Eq)]
 pub enum TableError {
     /// The columns have different lengths.
     #[snafu(display("Columns have different lengths"))]
     ColumnLengthMismatch,
+
+    /// At least one column has length different from the provided row count.
+    #[snafu(display("Column has length different from the provided row count"))]
+    ColumnLengthMismatchWithSpecifiedRowCount,
+
+    /// The table is empty and there is no specified row count.
+    #[snafu(display("Table is empty and no row count is specified"))]
+    EmptyTableWithoutSpecifiedRowCount,
 }
 /// A table of data, with schema included. This is simply a map from `Identifier` to `Column`,
 /// where columns order matters.
@@ -18,35 +42,64 @@ pub enum TableError {
 #[derive(Debug, Clone, Eq)]
 pub struct Table<'a, S: Scalar> {
     table: IndexMap<Identifier, Column<'a, S>>,
+    row_count: usize,
 }
 impl<'a, S: Scalar> Table<'a, S> {
-    /// Creates a new [`Table`].
+    /// Creates a new [`Table`] with the given columns and default [`TableOptions`].
     pub fn try_new(table: IndexMap<Identifier, Column<'a, S>>) -> Result<Self, TableError> {
-        if table.is_empty() {
-            return Ok(Self { table });
-        }
-        let num_rows = table[0].len();
-        if table.values().any(|column| column.len() != num_rows) {
-            Err(TableError::ColumnLengthMismatch)
-        } else {
-            Ok(Self { table })
+        Self::try_new_with_options(table, TableOptions::default())
+    }
+
+    /// Creates a new [`Table`] with the given columns and with [`TableOptions`].
+    pub fn try_new_with_options(
+        table: IndexMap<Identifier, Column<'a, S>>,
+        options: TableOptions,
+    ) -> Result<Self, TableError> {
+        match (table.is_empty(), options.row_count) {
+            (true, None) => Err(TableError::EmptyTableWithoutSpecifiedRowCount),
+            (true, Some(row_count)) => Ok(Self { table, row_count }),
+            (false, None) => {
+                let row_count = table[0].len();
+                if table.values().any(|column| column.len() != row_count) {
+                    Err(TableError::ColumnLengthMismatch)
+                } else {
+                    Ok(Self { table, row_count })
+                }
+            }
+            (false, Some(row_count)) => {
+                if table.values().any(|column| column.len() != row_count) {
+                    Err(TableError::ColumnLengthMismatchWithSpecifiedRowCount)
+                } else {
+                    Ok(Self { table, row_count })
+                }
+            }
         }
     }
-    /// Creates a new [`Table`].
+
+    /// Creates a new [`Table`] from an iterator of `(Identifier, Column)` pairs with default [`TableOptions`].
     pub fn try_from_iter<T: IntoIterator<Item = (Identifier, Column<'a, S>)>>(
         iter: T,
     ) -> Result<Self, TableError> {
-        Self::try_new(IndexMap::from_iter(iter))
+        Self::try_from_iter_with_options(iter, TableOptions::default())
     }
+
+    /// Creates a new [`Table`] from an iterator of `(Identifier, Column)` pairs with [`TableOptions`].
+    pub fn try_from_iter_with_options<T: IntoIterator<Item = (Identifier, Column<'a, S>)>>(
+        iter: T,
+        options: TableOptions,
+    ) -> Result<Self, TableError> {
+        Self::try_new_with_options(IndexMap::from_iter(iter), options)
+    }
+
     /// Number of columns in the table.
     #[must_use]
     pub fn num_columns(&self) -> usize {
         self.table.len()
     }
-    /// Number of rows in the table. For an empty table, this will return `None`.
+    /// Number of rows in the table.
     #[must_use]
-    pub fn num_rows(&self) -> Option<usize> {
-        (!self.table.is_empty()).then(|| self.table[0].len())
+    pub fn num_rows(&self) -> usize {
+        self.row_count
     }
     /// Whether the table has no columns.
     #[must_use]

crates/proof-of-sql/src/base/database/table_test.rs

@@ -1,6 +1,6 @@
 use crate::base::{
-    database::{table_utility::*, Column, Table, TableError},
-    map::IndexMap,
+    database::{table_utility::*, Column, Table, TableError, TableOptions},
+    map::{indexmap, IndexMap},
     scalar::test_scalar::TestScalar,
 };
 use bumpalo::Bump;
@@ -10,13 +10,104 @@ use proof_of_sql_parser::{
 };
 
 #[test]
-fn we_can_create_a_table_with_no_columns() {
-    let table = Table::<TestScalar>::try_new(IndexMap::default()).unwrap();
+fn we_can_create_a_table_with_no_columns_specifying_row_count() {
+    let table =
+        Table::<TestScalar>::try_new_with_options(IndexMap::default(), TableOptions::new(Some(1)))
+            .unwrap();
     assert_eq!(table.num_columns(), 0);
-    assert_eq!(table.num_rows(), None);
+    assert_eq!(table.num_rows(), 1);
+
+    let table =
+        Table::<TestScalar>::try_new_with_options(IndexMap::default(), TableOptions::new(Some(0)))
+            .unwrap();
+    assert_eq!(table.num_columns(), 0);
+    assert_eq!(table.num_rows(), 0);
+}
+
+#[test]
+fn we_can_create_a_table_with_default_options() {
+    let table = Table::<TestScalar>::try_new(indexmap! {
+        "a".parse().unwrap() => Column::BigInt(&[0, 1]),
+        "b".parse().unwrap() => Column::Int128(&[0, 1]),
+    })
+    .unwrap();
+    assert_eq!(table.num_columns(), 2);
+    assert_eq!(table.num_rows(), 2);
+
+    let table = Table::<TestScalar>::try_new(indexmap! {
+        "a".parse().unwrap() => Column::BigInt(&[]),
+        "b".parse().unwrap() => Column::Int128(&[]),
+    })
+    .unwrap();
+    assert_eq!(table.num_columns(), 2);
+    assert_eq!(table.num_rows(), 0);
+}
+
+#[test]
+fn we_can_create_a_table_with_specified_row_count() {
+    let table = Table::<TestScalar>::try_new_with_options(
+        indexmap! {
+            "a".parse().unwrap() => Column::BigInt(&[0, 1]),
+            "b".parse().unwrap() => Column::Int128(&[0, 1]),
+        },
+        TableOptions::new(Some(2)),
+    )
+    .unwrap();
+    assert_eq!(table.num_columns(), 2);
+    assert_eq!(table.num_rows(), 2);
+
+    let table = Table::<TestScalar>::try_new_with_options(
+        indexmap! {
+            "a".parse().unwrap() => Column::BigInt(&[]),
+            "b".parse().unwrap() => Column::Int128(&[]),
+        },
+        TableOptions::new(Some(0)),
+    )
+    .unwrap();
+    assert_eq!(table.num_columns(), 2);
+    assert_eq!(table.num_rows(), 0);
+}
+
+#[test]
+fn we_cannot_create_a_table_with_differing_column_lengths() {
+    assert!(matches!(
+        Table::<TestScalar>::try_from_iter([
+            ("a".parse().unwrap(), Column::BigInt(&[0])),
+            ("b".parse().unwrap(), Column::BigInt(&[])),
+        ]),
+        Err(TableError::ColumnLengthMismatch)
+    ));
 }
+
 #[test]
-fn we_can_create_an_empty_table() {
+fn we_cannot_create_a_table_with_column_length_different_from_specified_row_count() {
+    assert!(matches!(
+        Table::<TestScalar>::try_from_iter_with_options(
+            [
+                ("a".parse().unwrap(), Column::BigInt(&[0])),
+                ("b".parse().unwrap(), Column::BigInt(&[1])),
+            ],
+            TableOptions::new(Some(0))
+        ),
+        Err(TableError::ColumnLengthMismatchWithSpecifiedRowCount)
+    ));
+}
+
+#[test]
+fn we_cannot_create_a_table_with_no_columns_without_specified_row_count() {
+    assert!(matches!(
+        Table::<TestScalar>::try_from_iter_with_options([], TableOptions::new(None)),
+        Err(TableError::EmptyTableWithoutSpecifiedRowCount)
+    ));
+
+    assert!(matches!(
+        Table::<TestScalar>::try_new(IndexMap::default()),
+        Err(TableError::EmptyTableWithoutSpecifiedRowCount)
+    ));
+}
+
+#[test]
+fn we_can_create_an_empty_table_with_some_columns() {
     let alloc = Bump::new();
     let borrowed_table = table::<TestScalar>([
         borrowed_bigint("bigint", [0; 0], &alloc),
@@ -193,14 +284,3 @@ fn we_get_inequality_between_tables_with_differing_data() {
 
     assert_ne!(table_a, table_b);
 }
-
-#[test]
-fn we_cannot_create_a_table_with_differing_column_lengths() {
-    assert!(matches!(
-        Table::<TestScalar>::try_from_iter([
-            ("a".parse().unwrap(), Column::BigInt(&[0])),
-            ("b".parse().unwrap(), Column::BigInt(&[])),
-        ]),
-        Err(TableError::ColumnLengthMismatch)
-    ));
-}

crates/proof-of-sql/src/base/database/table_test_accessor.rs

@@ -111,12 +111,7 @@ impl<CP: CommitmentEvaluationProof> MetadataAccessor for TableTestAccessor<'_, C
     ///
     /// Will panic if the `table_ref` is not found in `self.tables`, indicating that an invalid reference was provided.
     fn get_length(&self, table_ref: TableRef) -> usize {
-        self.tables
-            .get(&table_ref)
-            .unwrap()
-            .0
-            .num_rows()
-            .unwrap_or(0)
+        self.tables.get(&table_ref).unwrap().0.num_rows()
     }
     ///
     /// # Panics

crates/proof-of-sql/src/sql/proof/proof_plan.rs

@@ -35,10 +35,9 @@ pub trait ProofPlan: Debug + Send + Sync + ProverEvaluate {
 
 #[enum_dispatch::enum_dispatch(DynProofPlan)]
 pub trait ProverEvaluate {
-    /// Evaluate the query and modify `FirstRoundBuilder` to track the result of the query.
+    /// Evaluate the query and return the result.
     fn result_evaluate<'a, S: Scalar>(
         &self,
-        input_length: usize,
         alloc: &'a Bump,
         accessor: &'a dyn DataAccessor<S>,
     ) -> Vec<Column<'a, S>>;

crates/proof-of-sql/src/sql/proof/query_proof.rs

@@ -76,7 +76,7 @@ impl<CP: CommitmentEvaluationProof> QueryProof<CP> {
         let alloc = Bump::new();
 
         // Evaluate query result
-        let result_cols = expr.result_evaluate(range_length, &alloc, accessor);
+        let result_cols = expr.result_evaluate(&alloc, accessor);
         let output_length = result_cols.first().map_or(0, Column::len);
         let provable_result = ProvableQueryResult::new(output_length as u64, &result_cols);
 

crates/proof-of-sql/src/sql/proof/query_proof_test.rs

@@ -42,7 +42,6 @@ impl Default for TrivialTestProofPlan {
 impl ProverEvaluate for TrivialTestProofPlan {
     fn result_evaluate<'a, S: Scalar>(
         &self,
-        _input_length: usize,
         alloc: &'a Bump,
         _accessor: &'a dyn DataAccessor<S>,
     ) -> Vec<Column<'a, S>> {
@@ -212,7 +211,6 @@ impl Default for SquareTestProofPlan {
 impl ProverEvaluate for SquareTestProofPlan {
     fn result_evaluate<'a, S: Scalar>(
         &self,
-        _table_length: usize,
         alloc: &'a Bump,
         _accessor: &'a dyn DataAccessor<S>,
     ) -> Vec<Column<'a, S>> {
@@ -390,7 +388,6 @@ impl Default for DoubleSquareTestProofPlan {
 impl ProverEvaluate for DoubleSquareTestProofPlan {
     fn result_evaluate<'a, S: Scalar>(
         &self,
-        _input_length: usize,
         alloc: &'a Bump,
         _accessor: &'a dyn DataAccessor<S>,
     ) -> Vec<Column<'a, S>> {
@@ -598,7 +595,6 @@ struct ChallengeTestProofPlan {}
 impl ProverEvaluate for ChallengeTestProofPlan {
     fn result_evaluate<'a, S: Scalar>(
         &self,
-        _input_length: usize,
         _alloc: &'a Bump,
         _accessor: &'a dyn DataAccessor<S>,
     ) -> Vec<Column<'a, S>> {

crates/proof-of-sql/src/sql/proof/verifiable_query_result_test.rs

@@ -27,7 +27,6 @@ pub(super) struct EmptyTestQueryExpr {
 impl ProverEvaluate for EmptyTestQueryExpr {
     fn result_evaluate<'a, S: Scalar>(
         &self,
-        _input_length: usize,
         alloc: &'a Bump,
         _accessor: &'a dyn DataAccessor<S>,
     ) -> Vec<Column<'a, S>> {