feat: add Table::columns(), ProvableQueryResult::new_from_table t…

…o simplify the code
spaceandtimelabs · Nov 14, 2024 · c1f78c3 · c1f78c3
1 parent 513d374
commit c1f78c3
Show file tree

Hide file tree

Showing 6 changed files with 38 additions and 71 deletions.
diff --git a/crates/proof-of-sql/src/base/database/table.rs b/crates/proof-of-sql/src/base/database/table.rs
@@ -120,6 +120,10 @@ impl<'a, S: Scalar> Table<'a, S> {
     pub fn column_names(&self) -> impl Iterator<Item = &Identifier> {
         self.table.keys()
     }
+    /// Returns the columns of this table as an Iterator
+    pub fn columns(&self) -> impl Iterator<Item = &Column<'a, S>> {
+        self.table.values()
+    }
 }
 
 // Note: we modify the default PartialEq for IndexMap to also check for column ordering.

diff --git a/crates/proof-of-sql/src/sql/proof/provable_query_result.rs b/crates/proof-of-sql/src/sql/proof/provable_query_result.rs
@@ -1,6 +1,6 @@
 use super::{decode_and_convert, decode_multiple_elements, ProvableResultColumn, QueryError};
 use crate::base::{
-    database::{Column, ColumnField, ColumnType, OwnedColumn, OwnedTable},
+    database::{Column, ColumnField, ColumnType, OwnedColumn, OwnedTable, Table},
     polynomial::compute_evaluation_vector,
     scalar::Scalar,
 };
@@ -55,6 +55,13 @@ impl ProvableQueryResult {
         }
     }
 
+    /// Form intermediate query result from a table
+    #[must_use]
+    pub fn new_from_table<S: Scalar>(table: &Table<S>) -> Self {
+        let columns = table.columns().copied().collect::<Vec<_>>();
+        Self::new(table.num_rows() as u64, &columns)
+    }
+
     /// Form intermediate query result from index rows and result columns
     /// # Panics
     ///

diff --git a/crates/proof-of-sql/src/sql/proof/query_proof.rs b/crates/proof-of-sql/src/sql/proof/query_proof.rs
@@ -77,13 +77,7 @@ impl<CP: CommitmentEvaluationProof> QueryProof<CP> {
 
         // Evaluate query result
         let result_table = expr.result_evaluate(&alloc, accessor);
-        let result_cols = result_table
-            .inner_table()
-            .values()
-            .cloned()
-            .collect::<Vec<_>>();
-        let provable_result =
-            ProvableQueryResult::new(result_table.num_rows() as u64, &result_cols);
+        let provable_result = ProvableQueryResult::new_from_table(&result_table);
 
         // Prover First Round
         let mut first_round_builder = FirstRoundBuilder::new();

diff --git a/crates/proof-of-sql/src/sql/proof_plans/filter_exec_test.rs b/crates/proof-of-sql/src/sql/proof_plans/filter_exec_test.rs
@@ -195,11 +195,6 @@ fn we_can_get_an_empty_result_from_a_basic_filter_on_an_empty_table_using_result
     );
     let alloc = Bump::new();
     let result_table = expr.result_evaluate(&alloc, &accessor);
-    let result_cols = result_table
-        .inner_table()
-        .values()
-        .cloned()
-        .collect::<Vec<_>>();
     let mut builder = FirstRoundBuilder::new();
     expr.first_round_evaluate(&mut builder);
     let fields = &[
@@ -211,10 +206,9 @@ fn we_can_get_an_empty_result_from_a_basic_filter_on_an_empty_table_using_result
             ColumnType::Decimal75(Precision::new(75).unwrap(), 0),
         ),
     ];
-    let res: OwnedTable<Curve25519Scalar> =
-        ProvableQueryResult::new(result_table.num_rows() as u64, &result_cols)
-            .to_owned_table(fields)
-            .unwrap();
+    let res: OwnedTable<Curve25519Scalar> = ProvableQueryResult::new_from_table(&result_table)
+        .to_owned_table(fields)
+        .unwrap();
     let expected: OwnedTable<Curve25519Scalar> = owned_table([
         bigint("b", [0; 0]),
         int128("c", [0; 0]),
@@ -245,11 +239,6 @@ fn we_can_get_an_empty_result_from_a_basic_filter_using_result_evaluate() {
     );
     let alloc = Bump::new();
     let result_table = expr.result_evaluate(&alloc, &accessor);
-    let result_cols = result_table
-        .inner_table()
-        .values()
-        .cloned()
-        .collect::<Vec<_>>();
     let mut builder = FirstRoundBuilder::new();
     expr.first_round_evaluate(&mut builder);
     let fields = &[
@@ -261,10 +250,9 @@ fn we_can_get_an_empty_result_from_a_basic_filter_using_result_evaluate() {
             ColumnType::Decimal75(Precision::new(1).unwrap(), 0),
         ),
     ];
-    let res: OwnedTable<Curve25519Scalar> =
-        ProvableQueryResult::new(result_table.num_rows() as u64, &result_cols)
-            .to_owned_table(fields)
-            .unwrap();
+    let res: OwnedTable<Curve25519Scalar> = ProvableQueryResult::new_from_table(&result_table)
+        .to_owned_table(fields)
+        .unwrap();
     let expected: OwnedTable<Curve25519Scalar> = owned_table([
         bigint("b", [0; 0]),
         int128("c", [0; 0]),
@@ -291,18 +279,12 @@ fn we_can_get_no_columns_from_a_basic_filter_with_no_selected_columns_using_resu
     let expr = filter(cols_expr_plan(t, &[], &accessor), tab(t), where_clause);
     let alloc = Bump::new();
     let result_table = expr.result_evaluate(&alloc, &accessor);
-    let result_cols = result_table
-        .inner_table()
-        .values()
-        .cloned()
-        .collect::<Vec<_>>();
     let mut builder = FirstRoundBuilder::new();
     expr.first_round_evaluate(&mut builder);
     let fields = &[];
-    let res: OwnedTable<Curve25519Scalar> =
-        ProvableQueryResult::new(result_table.num_rows() as u64, &result_cols)
-            .to_owned_table(fields)
-            .unwrap();
+    let res: OwnedTable<Curve25519Scalar> = ProvableQueryResult::new_from_table(&result_table)
+        .to_owned_table(fields)
+        .unwrap();
     let expected = OwnedTable::try_new(IndexMap::default()).unwrap();
     assert_eq!(res, expected);
 }
@@ -327,11 +309,6 @@ fn we_can_get_the_correct_result_from_a_basic_filter_using_result_evaluate() {
     );
     let alloc = Bump::new();
     let result_table = expr.result_evaluate(&alloc, &accessor);
-    let result_cols = result_table
-        .inner_table()
-        .values()
-        .cloned()
-        .collect::<Vec<_>>();
     let mut builder = FirstRoundBuilder::new();
     expr.first_round_evaluate(&mut builder);
     let fields = &[
@@ -343,10 +320,9 @@ fn we_can_get_the_correct_result_from_a_basic_filter_using_result_evaluate() {
             ColumnType::Decimal75(Precision::new(1).unwrap(), 0),
         ),
     ];
-    let res: OwnedTable<Curve25519Scalar> =
-        ProvableQueryResult::new(result_table.num_rows() as u64, &result_cols)
-            .to_owned_table(fields)
-            .unwrap();
+    let res: OwnedTable<Curve25519Scalar> = ProvableQueryResult::new_from_table(&result_table)
+        .to_owned_table(fields)
+        .unwrap();
     let expected: OwnedTable<Curve25519Scalar> = owned_table([
         bigint("b", [3, 5]),
         int128("c", [3, 5]),

diff --git a/crates/proof-of-sql/src/sql/proof_plans/projection_exec_test.rs b/crates/proof-of-sql/src/sql/proof_plans/projection_exec_test.rs
@@ -169,11 +169,6 @@ fn we_can_get_an_empty_result_from_a_basic_projection_on_an_empty_table_using_re
         projection(cols_expr_plan(t, &["b", "c", "d", "e"], &accessor), tab(t));
     let alloc = Bump::new();
     let result_table = expr.result_evaluate(&alloc, &accessor);
-    let result_cols = result_table
-        .inner_table()
-        .values()
-        .cloned()
-        .collect::<Vec<_>>();
     let mut builder = FirstRoundBuilder::new();
     expr.first_round_evaluate(&mut builder);
     let fields = &[
@@ -185,10 +180,9 @@ fn we_can_get_an_empty_result_from_a_basic_projection_on_an_empty_table_using_re
             ColumnType::Decimal75(Precision::new(75).unwrap(), 0),
         ),
     ];
-    let res: OwnedTable<Curve25519Scalar> =
-        ProvableQueryResult::new(result_table.num_rows() as u64, &result_cols)
-            .to_owned_table(fields)
-            .unwrap();
+    let res: OwnedTable<Curve25519Scalar> = ProvableQueryResult::new_from_table(&result_table)
+        .to_owned_table(fields)
+        .unwrap();
     let expected: OwnedTable<Curve25519Scalar> = owned_table([
         bigint("b", [0; 0]),
         int128("c", [0; 0]),
@@ -214,18 +208,12 @@ fn we_can_get_no_columns_from_a_basic_projection_with_no_selected_columns_using_
     let expr: DynProofPlan = projection(cols_expr_plan(t, &[], &accessor), tab(t));
     let alloc = Bump::new();
     let result_table = expr.result_evaluate(&alloc, &accessor);
-    let result_cols = result_table
-        .inner_table()
-        .values()
-        .cloned()
-        .collect::<Vec<_>>();
     let mut builder = FirstRoundBuilder::new();
     expr.first_round_evaluate(&mut builder);
     let fields = &[];
-    let res: OwnedTable<Curve25519Scalar> =
-        ProvableQueryResult::new(result_table.num_rows() as u64, &result_cols)
-            .to_owned_table(fields)
-            .unwrap();
+    let res: OwnedTable<Curve25519Scalar> = ProvableQueryResult::new_from_table(&result_table)
+        .to_owned_table(fields)
+        .unwrap();
     let expected = OwnedTable::try_new(IndexMap::default()).unwrap();
     assert_eq!(res, expected);
 }
@@ -256,11 +244,6 @@ fn we_can_get_the_correct_result_from_a_basic_projection_using_result_evaluate()
     );
     let alloc = Bump::new();
     let result_table = expr.result_evaluate(&alloc, &accessor);
-    let result_cols = result_table
-        .inner_table()
-        .values()
-        .cloned()
-        .collect::<Vec<_>>();
     let mut builder = FirstRoundBuilder::new();
     expr.first_round_evaluate(&mut builder);
     let fields = &[
@@ -272,10 +255,9 @@ fn we_can_get_the_correct_result_from_a_basic_projection_using_result_evaluate()
             ColumnType::Decimal75(Precision::new(1).unwrap(), 0),
         ),
     ];
-    let res: OwnedTable<Curve25519Scalar> =
-        ProvableQueryResult::new(result_table.num_rows() as u64, &result_cols)
-            .to_owned_table(fields)
-            .unwrap();
+    let res: OwnedTable<Curve25519Scalar> = ProvableQueryResult::new_from_table(&result_table)
+        .to_owned_table(fields)
+        .unwrap();
     let expected: OwnedTable<Curve25519Scalar> = owned_table([
         bigint("b", [2, 3, 4, 5, 6]),
         int128("prod", [1, 4, 9, 16, 25]),

diff --git a/crates/proof-of-sql/src/sql/proof_plans/table_exec.rs b/crates/proof-of-sql/src/sql/proof_plans/table_exec.rs
@@ -33,6 +33,10 @@ impl TableExec {
     }
 
     /// Returns the entire table.
+    ///
+    /// # Panics
+    /// Panics if columns from the accessor have different lengths.
+    /// In practice, this should never happen.
     fn get_table<'a, S: Scalar>(&self, accessor: &'a dyn DataAccessor<S>) -> Table<'a, S> {
         Table::<'a, S>::try_from_iter(self.schema.iter().map(|field| {
             (