From c1f78c389b5e2a38420616d0c6b487eca3b42810 Mon Sep 17 00:00:00 2001 From: Ian Joiner <14581281+iajoiner@users.noreply.github.com> Date: Wed, 13 Nov 2024 16:24:51 -0500 Subject: [PATCH] feat: add `Table::columns()`, `ProvableQueryResult::new_from_table` to simplify the code --- .../proof-of-sql/src/base/database/table.rs | 4 ++ .../src/sql/proof/provable_query_result.rs | 9 +++- .../proof-of-sql/src/sql/proof/query_proof.rs | 8 +--- .../src/sql/proof_plans/filter_exec_test.rs | 48 +++++-------------- .../sql/proof_plans/projection_exec_test.rs | 36 ++++---------- .../src/sql/proof_plans/table_exec.rs | 4 ++ 6 files changed, 38 insertions(+), 71 deletions(-) diff --git a/crates/proof-of-sql/src/base/database/table.rs b/crates/proof-of-sql/src/base/database/table.rs index 9f8685111..52db2e1f5 100644 --- 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 { self.table.keys() } + /// Returns the columns of this table as an Iterator + pub fn columns(&self) -> impl Iterator> { + 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 index b45d1dd18..b1db29c5d 100644 --- 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(table: &Table) -> Self { + let columns = table.columns().copied().collect::>(); + 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 index 9fba7b5bc..99db705e9 100644 --- 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 QueryProof { // Evaluate query result let result_table = expr.result_evaluate(&alloc, accessor); - let result_cols = result_table - .inner_table() - .values() - .cloned() - .collect::>(); - 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 index 1e8ee53b5..bfe530754 100644 --- 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::>(); 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 = - ProvableQueryResult::new(result_table.num_rows() as u64, &result_cols) - .to_owned_table(fields) - .unwrap(); + let res: OwnedTable = ProvableQueryResult::new_from_table(&result_table) + .to_owned_table(fields) + .unwrap(); let expected: OwnedTable = 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::>(); 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 = - ProvableQueryResult::new(result_table.num_rows() as u64, &result_cols) - .to_owned_table(fields) - .unwrap(); + let res: OwnedTable = ProvableQueryResult::new_from_table(&result_table) + .to_owned_table(fields) + .unwrap(); let expected: OwnedTable = 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::>(); let mut builder = FirstRoundBuilder::new(); expr.first_round_evaluate(&mut builder); let fields = &[]; - let res: OwnedTable = - ProvableQueryResult::new(result_table.num_rows() as u64, &result_cols) - .to_owned_table(fields) - .unwrap(); + let res: OwnedTable = 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::>(); 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 = - ProvableQueryResult::new(result_table.num_rows() as u64, &result_cols) - .to_owned_table(fields) - .unwrap(); + let res: OwnedTable = ProvableQueryResult::new_from_table(&result_table) + .to_owned_table(fields) + .unwrap(); let expected: OwnedTable = 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 index 372d22ec8..efca22a9c 100644 --- 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::>(); 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 = - ProvableQueryResult::new(result_table.num_rows() as u64, &result_cols) - .to_owned_table(fields) - .unwrap(); + let res: OwnedTable = ProvableQueryResult::new_from_table(&result_table) + .to_owned_table(fields) + .unwrap(); let expected: OwnedTable = 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::>(); let mut builder = FirstRoundBuilder::new(); expr.first_round_evaluate(&mut builder); let fields = &[]; - let res: OwnedTable = - ProvableQueryResult::new(result_table.num_rows() as u64, &result_cols) - .to_owned_table(fields) - .unwrap(); + let res: OwnedTable = 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::>(); 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 = - ProvableQueryResult::new(result_table.num_rows() as u64, &result_cols) - .to_owned_table(fields) - .unwrap(); + let res: OwnedTable = ProvableQueryResult::new_from_table(&result_table) + .to_owned_table(fields) + .unwrap(); let expected: OwnedTable = 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 index d3346120b..2091b3390 100644 --- 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) -> Table<'a, S> { Table::<'a, S>::try_from_iter(self.schema.iter().map(|field| { (