diff --git a/crates/proof-of-sql/src/base/database/mod.rs b/crates/proof-of-sql/src/base/database/mod.rs index b45f3ea6a..eb726b930 100644 --- a/crates/proof-of-sql/src/base/database/mod.rs +++ b/crates/proof-of-sql/src/base/database/mod.rs @@ -64,6 +64,14 @@ pub(crate) use owned_table::OwnedTableError; mod owned_table_test; pub mod owned_table_utility; +mod table; +pub use table::Table; +#[cfg(test)] +pub(crate) use table::TableError; +#[cfg(test)] +mod table_test; +pub mod table_utility; + /// TODO: add docs pub(crate) mod expression_evaluation; mod expression_evaluation_error; @@ -87,6 +95,11 @@ pub use owned_table_test_accessor::OwnedTableTestAccessor; #[cfg(all(test, feature = "blitzar"))] mod owned_table_test_accessor_test; +mod table_test_accessor; +pub use table_test_accessor::TableTestAccessor; +#[cfg(all(test, feature = "blitzar"))] +mod table_test_accessor_test; + /// TODO: add docs pub(crate) mod filter_util; #[cfg(test)] diff --git a/crates/proof-of-sql/src/base/database/table.rs b/crates/proof-of-sql/src/base/database/table.rs new file mode 100644 index 000000000..e570e8adc --- /dev/null +++ b/crates/proof-of-sql/src/base/database/table.rs @@ -0,0 +1,93 @@ +use super::Column; +use crate::base::{map::IndexMap, scalar::Scalar}; +use proof_of_sql_parser::Identifier; +use snafu::Snafu; + +/// 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, +} +/// A table of data, with schema included. This is simply a map from `Identifier` to `Column`, +/// where columns order matters. +/// This is primarily used as an internal result that is used before +/// converting to the final result in either Arrow format or JSON. +/// This is the analog of an arrow [`RecordBatch`](arrow::record_batch::RecordBatch). +#[derive(Debug, Clone, Eq)] +pub struct Table<'a, S: Scalar> { + table: IndexMap>, +} +impl<'a, S: Scalar> Table<'a, S> { + /// Creates a new [`Table`]. + pub fn try_new(table: IndexMap>) -> Result { + 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 }) + } + } + /// Creates a new [`Table`]. + pub fn try_from_iter)>>( + iter: T, + ) -> Result { + Self::try_new(IndexMap::from_iter(iter)) + } + /// 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`. + #[must_use] + pub fn num_rows(&self) -> Option { + (!self.table.is_empty()).then(|| self.table[0].len()) + } + /// Whether the table has no columns. + #[must_use] + pub fn is_empty(&self) -> bool { + self.table.is_empty() + } + /// Returns the columns of this table as an `IndexMap` + #[must_use] + pub fn into_inner(self) -> IndexMap> { + self.table + } + /// Returns the columns of this table as an `IndexMap` + #[must_use] + pub fn inner_table(&self) -> &IndexMap> { + &self.table + } + /// Returns the columns of this table as an Iterator + pub fn column_names(&self) -> impl Iterator { + self.table.keys() + } +} + +// Note: we modify the default PartialEq for IndexMap to also check for column ordering. +// This is to align with the behaviour of a `RecordBatch`. +impl PartialEq for Table<'_, S> { + fn eq(&self, other: &Self) -> bool { + self.table == other.table + && self + .table + .keys() + .zip(other.table.keys()) + .all(|(a, b)| a == b) + } +} + +#[cfg(test)] +impl<'a, S: Scalar> core::ops::Index<&str> for Table<'a, S> { + type Output = Column<'a, S>; + fn index(&self, index: &str) -> &Self::Output { + self.table + .get(&index.parse::().unwrap()) + .unwrap() + } +} diff --git a/crates/proof-of-sql/src/base/database/table_test.rs b/crates/proof-of-sql/src/base/database/table_test.rs new file mode 100644 index 000000000..515de66fa --- /dev/null +++ b/crates/proof-of-sql/src/base/database/table_test.rs @@ -0,0 +1,206 @@ +use crate::base::{ + database::{table_utility::*, Column, Table, TableError}, + map::IndexMap, + scalar::test_scalar::TestScalar, +}; +use bumpalo::Bump; +use proof_of_sql_parser::{ + posql_time::{PoSQLTimeUnit, PoSQLTimeZone}, + Identifier, +}; + +#[test] +fn we_can_create_a_table_with_no_columns() { + let table = Table::::try_new(IndexMap::default()).unwrap(); + assert_eq!(table.num_columns(), 0); + assert_eq!(table.num_rows(), None); +} +#[test] +fn we_can_create_an_empty_table() { + let alloc = Bump::new(); + let borrowed_table = table::([ + borrowed_bigint("bigint", [0; 0], &alloc), + borrowed_int128("decimal", [0; 0], &alloc), + borrowed_varchar("varchar", ["0"; 0], &alloc), + borrowed_scalar("scalar", [0; 0], &alloc), + borrowed_boolean("boolean", [true; 0], &alloc), + ]); + let mut table = IndexMap::default(); + table.insert(Identifier::try_new("bigint").unwrap(), Column::BigInt(&[])); + table.insert(Identifier::try_new("decimal").unwrap(), Column::Int128(&[])); + table.insert( + Identifier::try_new("varchar").unwrap(), + Column::VarChar((&[], &[])), + ); + table.insert(Identifier::try_new("scalar").unwrap(), Column::Scalar(&[])); + table.insert( + Identifier::try_new("boolean").unwrap(), + Column::Boolean(&[]), + ); + assert_eq!(borrowed_table.into_inner(), table); +} + +#[test] +fn we_can_create_a_table_with_data() { + let alloc = Bump::new(); + + let borrowed_table = table::([ + borrowed_bigint( + "bigint", + [0_i64, 1, 2, 3, 4, 5, 6, i64::MIN, i64::MAX], + &alloc, + ), + borrowed_int128( + "decimal", + [0_i128, 1, 2, 3, 4, 5, 6, i128::MIN, i128::MAX], + &alloc, + ), + borrowed_varchar( + "varchar", + ["0", "1", "2", "3", "4", "5", "6", "7", "8"], + &alloc, + ), + borrowed_scalar("scalar", [0, 1, 2, 3, 4, 5, 6, 7, 8], &alloc), + borrowed_boolean( + "boolean", + [true, false, true, false, true, false, true, false, true], + &alloc, + ), + borrowed_timestamptz( + "time_stamp", + PoSQLTimeUnit::Second, + PoSQLTimeZone::Utc, + [0_i64, 1, 2, 3, 4, 5, 6, i64::MIN, i64::MAX], + &alloc, + ), + ]); + + let mut expected_table = IndexMap::default(); + + let time_stamp_data = alloc.alloc_slice_copy(&[0_i64, 1, 2, 3, 4, 5, 6, i64::MIN, i64::MAX]); + expected_table.insert( + Identifier::try_new("time_stamp").unwrap(), + Column::TimestampTZ(PoSQLTimeUnit::Second, PoSQLTimeZone::Utc, time_stamp_data), + ); + + let bigint_data = alloc.alloc_slice_copy(&[0_i64, 1, 2, 3, 4, 5, 6, i64::MIN, i64::MAX]); + expected_table.insert( + Identifier::try_new("bigint").unwrap(), + Column::BigInt(bigint_data), + ); + + let decimal_data = alloc.alloc_slice_copy(&[0_i128, 1, 2, 3, 4, 5, 6, i128::MIN, i128::MAX]); + expected_table.insert( + Identifier::try_new("decimal").unwrap(), + Column::Int128(decimal_data), + ); + + let varchar_data: Vec<&str> = ["0", "1", "2", "3", "4", "5", "6", "7", "8"] + .iter() + .map(|&s| alloc.alloc_str(s) as &str) + .collect(); + let varchar_str_slice = alloc.alloc_slice_clone(&varchar_data); + let varchar_scalars: Vec = varchar_data.iter().map(Into::into).collect(); + let varchar_scalars_slice = alloc.alloc_slice_clone(&varchar_scalars); + expected_table.insert( + Identifier::try_new("varchar").unwrap(), + Column::VarChar((varchar_str_slice, varchar_scalars_slice)), + ); + + let scalar_data: Vec = (0..=8).map(TestScalar::from).collect(); + let scalar_slice = alloc.alloc_slice_copy(&scalar_data); + expected_table.insert( + Identifier::try_new("scalar").unwrap(), + Column::Scalar(scalar_slice), + ); + + let boolean_data = + alloc.alloc_slice_copy(&[true, false, true, false, true, false, true, false, true]); + expected_table.insert( + Identifier::try_new("boolean").unwrap(), + Column::Boolean(boolean_data), + ); + + assert_eq!(borrowed_table.into_inner(), expected_table); +} + +#[test] +fn we_get_inequality_between_tables_with_differing_column_order() { + let alloc = Bump::new(); + + let table_a: Table<'_, TestScalar> = table([ + borrowed_bigint("a", [0; 0], &alloc), + borrowed_int128("b", [0; 0], &alloc), + borrowed_varchar("c", ["0"; 0], &alloc), + borrowed_boolean("d", [false; 0], &alloc), + borrowed_timestamptz( + "time_stamp", + PoSQLTimeUnit::Second, + PoSQLTimeZone::Utc, + [0_i64; 0], + &alloc, + ), + ]); + + let table_b: Table<'_, TestScalar> = table([ + borrowed_boolean("d", [false; 0], &alloc), + borrowed_int128("b", [0; 0], &alloc), + borrowed_bigint("a", [0; 0], &alloc), + borrowed_varchar("c", ["0"; 0], &alloc), + borrowed_timestamptz( + "time_stamp", + PoSQLTimeUnit::Second, + PoSQLTimeZone::Utc, + [0_i64; 0], + &alloc, + ), + ]); + + assert_ne!(table_a, table_b); +} + +#[test] +fn we_get_inequality_between_tables_with_differing_data() { + let alloc = Bump::new(); + + let table_a: Table<'_, TestScalar> = table([ + borrowed_bigint("a", [0], &alloc), + borrowed_int128("b", [0], &alloc), + borrowed_varchar("c", ["0"], &alloc), + borrowed_boolean("d", [true], &alloc), + borrowed_timestamptz( + "time_stamp", + PoSQLTimeUnit::Second, + PoSQLTimeZone::Utc, + [1_625_072_400], + &alloc, + ), + ]); + + let table_b: Table<'_, TestScalar> = table([ + borrowed_bigint("a", [1], &alloc), + borrowed_int128("b", [0], &alloc), + borrowed_varchar("c", ["0"], &alloc), + borrowed_boolean("d", [true], &alloc), + borrowed_timestamptz( + "time_stamp", + PoSQLTimeUnit::Second, + PoSQLTimeZone::Utc, + [1_625_076_000], + &alloc, + ), + ]); + + assert_ne!(table_a, table_b); +} + +#[test] +fn we_cannot_create_a_table_with_differing_column_lengths() { + assert!(matches!( + Table::::try_from_iter([ + ("a".parse().unwrap(), Column::BigInt(&[0])), + ("b".parse().unwrap(), Column::BigInt(&[])), + ]), + Err(TableError::ColumnLengthMismatch) + )); +} diff --git a/crates/proof-of-sql/src/base/database/table_test_accessor.rs b/crates/proof-of-sql/src/base/database/table_test_accessor.rs new file mode 100644 index 000000000..705c19c6c --- /dev/null +++ b/crates/proof-of-sql/src/base/database/table_test_accessor.rs @@ -0,0 +1,175 @@ +use super::{ + Column, ColumnRef, ColumnType, CommitmentAccessor, DataAccessor, MetadataAccessor, + SchemaAccessor, Table, TableRef, TestAccessor, +}; +use crate::base::{ + commitment::{CommitmentEvaluationProof, VecCommitmentExt}, + map::IndexMap, +}; +use alloc::vec::Vec; +use proof_of_sql_parser::Identifier; + +/// A test accessor that uses [`Table`] as the underlying table type. +/// Note: this is intended for testing and examples. It is not optimized for performance, so should not be used for benchmarks or production use-cases. +pub struct TableTestAccessor<'a, CP: CommitmentEvaluationProof> { + tables: IndexMap, usize)>, + setup: Option>, +} + +impl Default for TableTestAccessor<'_, CP> { + fn default() -> Self { + Self { + tables: IndexMap::default(), + setup: None, + } + } +} + +impl Clone for TableTestAccessor<'_, CP> { + fn clone(&self) -> Self { + Self { + tables: self.tables.clone(), + setup: self.setup, + } + } +} + +impl<'a, CP: CommitmentEvaluationProof> TestAccessor for TableTestAccessor<'a, CP> { + type Table = Table<'a, CP::Scalar>; + + fn new_empty() -> Self { + TableTestAccessor::default() + } + + fn add_table(&mut self, table_ref: TableRef, data: Self::Table, table_offset: usize) { + self.tables.insert(table_ref, (data, table_offset)); + } + /// + /// # Panics + /// + /// Will panic if the `table_ref` is not found in `self.tables`, indicating + /// that an invalid reference was provided. + fn get_column_names(&self, table_ref: TableRef) -> Vec<&str> { + self.tables + .get(&table_ref) + .unwrap() + .0 + .column_names() + .map(proof_of_sql_parser::Identifier::as_str) + .collect() + } + + /// + /// # Panics + /// + /// Will panic if the `table_ref` is not found in `self.tables`, indicating that an invalid reference was provided. + fn update_offset(&mut self, table_ref: TableRef, new_offset: usize) { + self.tables.get_mut(&table_ref).unwrap().1 = new_offset; + } +} + +/// +/// # Panics +/// +/// Will panic if the `column.table_ref()` is not found in `self.tables`, or if +/// the `column.column_id()` is not found in the inner table for that reference, +/// indicating that an invalid column reference was provided. +impl<'a, CP: CommitmentEvaluationProof> DataAccessor for TableTestAccessor<'a, CP> { + fn get_column(&self, column: ColumnRef) -> Column<'a, CP::Scalar> { + *self + .tables + .get(&column.table_ref()) + .unwrap() + .0 + .inner_table() + .get(&column.column_id()) + .unwrap() + } +} + +/// +/// # Panics +/// +/// Will panic if the `column.table_ref()` is not found in `self.tables`, or if the `column.column_id()` is not found in the inner table for that reference,indicating that an invalid column reference was provided. +impl CommitmentAccessor + for TableTestAccessor<'_, CP> +{ + fn get_commitment(&self, column: ColumnRef) -> CP::Commitment { + let (table, offset) = self.tables.get(&column.table_ref()).unwrap(); + let borrowed_column = table.inner_table().get(&column.column_id()).unwrap(); + Vec::::from_columns_with_offset( + [borrowed_column], + *offset, + self.setup.as_ref().unwrap(), + )[0] + .clone() + } +} +impl MetadataAccessor for TableTestAccessor<'_, CP> { + /// + /// # Panics + /// + /// 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) + } + /// + /// # Panics + /// + /// Will panic if the `table_ref` is not found in `self.tables`, indicating that an invalid reference was provided. + fn get_offset(&self, table_ref: TableRef) -> usize { + self.tables.get(&table_ref).unwrap().1 + } +} +impl SchemaAccessor for TableTestAccessor<'_, CP> { + fn lookup_column(&self, table_ref: TableRef, column_id: Identifier) -> Option { + Some( + self.tables + .get(&table_ref)? + .0 + .inner_table() + .get(&column_id)? + .column_type(), + ) + } + /// + /// # Panics + /// + /// Will panic if the `table_ref` is not found in `self.tables`, indicating that an invalid reference was provided. + fn lookup_schema(&self, table_ref: TableRef) -> Vec<(Identifier, ColumnType)> { + self.tables + .get(&table_ref) + .unwrap() + .0 + .inner_table() + .iter() + .map(|(&id, col)| (id, col.column_type())) + .collect() + } +} + +impl<'a, CP: CommitmentEvaluationProof> TableTestAccessor<'a, CP> { + /// Create a new empty test accessor with the given setup. + pub fn new_empty_with_setup(setup: CP::ProverPublicSetup<'a>) -> Self { + let mut res = Self::new_empty(); + res.setup = Some(setup); + res + } + + /// Create a new test accessor containing the provided table. + pub fn new_from_table( + table_ref: TableRef, + table: Table<'a, CP::Scalar>, + offset: usize, + setup: CP::ProverPublicSetup<'a>, + ) -> Self { + let mut res = Self::new_empty_with_setup(setup); + res.add_table(table_ref, table, offset); + res + } +} diff --git a/crates/proof-of-sql/src/base/database/table_test_accessor_test.rs b/crates/proof-of-sql/src/base/database/table_test_accessor_test.rs new file mode 100644 index 000000000..08f913faa --- /dev/null +++ b/crates/proof-of-sql/src/base/database/table_test_accessor_test.rs @@ -0,0 +1,305 @@ +use super::{ + Column, ColumnRef, ColumnType, CommitmentAccessor, DataAccessor, MetadataAccessor, + SchemaAccessor, TableTestAccessor, TestAccessor, +}; +use crate::base::{ + commitment::{ + naive_commitment::NaiveCommitment, test_evaluation_proof::TestEvaluationProof, Commitment, + CommittableColumn, + }, + database::table_utility::*, + scalar::test_scalar::TestScalar, +}; +use bumpalo::Bump; +use proof_of_sql_parser::posql_time::{PoSQLTimeUnit, PoSQLTimeZone}; + +#[test] +fn we_can_query_the_length_of_a_table() { + let alloc = Bump::new(); + let mut accessor = TableTestAccessor::::new_empty_with_setup(()); + let table_ref_1 = "sxt.test".parse().unwrap(); + let table_ref_2 = "sxt.test2".parse().unwrap(); + + let data1 = table([ + borrowed_bigint("a", [1, 2, 3], &alloc), + borrowed_bigint("b", [4, 5, 6], &alloc), + ]); + accessor.add_table(table_ref_1, data1, 0_usize); + + assert_eq!(accessor.get_length(table_ref_1), 3); + + let data2 = table([ + borrowed_bigint("a", [1, 2, 3, 4], &alloc), + borrowed_bigint("b", [4, 5, 6, 5], &alloc), + ]); + accessor.add_table(table_ref_2, data2, 0_usize); + + assert_eq!(accessor.get_length(table_ref_1), 3); + assert_eq!(accessor.get_length(table_ref_2), 4); +} + +#[test] +fn we_can_access_the_columns_of_a_table() { + let alloc = Bump::new(); + let mut accessor = TableTestAccessor::::new_empty_with_setup(()); + let table_ref_1 = "sxt.test".parse().unwrap(); + let table_ref_2 = "sxt.test2".parse().unwrap(); + + let data1 = table([ + borrowed_bigint("a", [1, 2, 3], &alloc), + borrowed_bigint("b", [4, 5, 6], &alloc), + ]); + accessor.add_table(table_ref_1, data1, 0_usize); + + let column = ColumnRef::new(table_ref_1, "b".parse().unwrap(), ColumnType::BigInt); + match accessor.get_column(column) { + Column::BigInt(col) => assert_eq!(col.to_vec(), vec![4, 5, 6]), + _ => panic!("Invalid column type"), + }; + + let data2 = table([ + borrowed_bigint("a", [1, 2, 3, 4], &alloc), + borrowed_bigint("b", [4, 5, 6, 5], &alloc), + borrowed_int128("c128", [1, 2, 3, 4], &alloc), + borrowed_varchar("varchar", ["a", "bc", "d", "e"], &alloc), + borrowed_scalar("scalar", [1, 2, 3, 4], &alloc), + borrowed_boolean("boolean", [true, false, true, false], &alloc), + borrowed_timestamptz( + "time", + PoSQLTimeUnit::Second, + PoSQLTimeZone::Utc, + [4, 5, 6, 5], + &alloc, + ), + ]); + accessor.add_table(table_ref_2, data2, 0_usize); + + let column = ColumnRef::new(table_ref_1, "a".parse().unwrap(), ColumnType::BigInt); + match accessor.get_column(column) { + Column::BigInt(col) => assert_eq!(col.to_vec(), vec![1, 2, 3]), + _ => panic!("Invalid column type"), + }; + + let column = ColumnRef::new(table_ref_2, "b".parse().unwrap(), ColumnType::BigInt); + match accessor.get_column(column) { + Column::BigInt(col) => assert_eq!(col.to_vec(), vec![4, 5, 6, 5]), + _ => panic!("Invalid column type"), + }; + + let column = ColumnRef::new(table_ref_2, "c128".parse().unwrap(), ColumnType::Int128); + match accessor.get_column(column) { + Column::Int128(col) => assert_eq!(col.to_vec(), vec![1, 2, 3, 4]), + _ => panic!("Invalid column type"), + }; + + let col_slice: Vec<_> = vec!["a", "bc", "d", "e"]; + let col_scalars: Vec<_> = ["a", "bc", "d", "e"] + .iter() + .map(core::convert::Into::into) + .collect(); + let column = ColumnRef::new(table_ref_2, "varchar".parse().unwrap(), ColumnType::VarChar); + match accessor.get_column(column) { + Column::VarChar((col, scals)) => { + assert_eq!(col.to_vec(), col_slice); + assert_eq!(scals.to_vec(), col_scalars); + } + _ => panic!("Invalid column type"), + }; + + let column = ColumnRef::new(table_ref_2, "scalar".parse().unwrap(), ColumnType::Scalar); + match accessor.get_column(column) { + Column::Scalar(col) => assert_eq!( + col.to_vec(), + vec![ + TestScalar::from(1), + TestScalar::from(2), + TestScalar::from(3), + TestScalar::from(4) + ] + ), + _ => panic!("Invalid column type"), + }; + + let column = ColumnRef::new(table_ref_2, "boolean".parse().unwrap(), ColumnType::Boolean); + match accessor.get_column(column) { + Column::Boolean(col) => assert_eq!(col.to_vec(), vec![true, false, true, false]), + _ => panic!("Invalid column type"), + }; + + let column = ColumnRef::new( + table_ref_2, + "time".parse().unwrap(), + ColumnType::TimestampTZ(PoSQLTimeUnit::Second, PoSQLTimeZone::Utc), + ); + match accessor.get_column(column) { + Column::TimestampTZ(_, _, col) => assert_eq!(col.to_vec(), vec![4, 5, 6, 5]), + _ => panic!("Invalid column type"), + }; +} + +#[test] +fn we_can_access_the_commitments_of_table_columns() { + let alloc = Bump::new(); + let mut accessor = TableTestAccessor::::new_empty_with_setup(()); + let table_ref_1 = "sxt.test".parse().unwrap(); + let table_ref_2 = "sxt.test2".parse().unwrap(); + + let data1 = table([ + borrowed_bigint("a", [1, 2, 3], &alloc), + borrowed_bigint("b", [4, 5, 6], &alloc), + ]); + accessor.add_table(table_ref_1, data1, 0_usize); + + let column = ColumnRef::new(table_ref_1, "b".parse().unwrap(), ColumnType::BigInt); + assert_eq!( + accessor.get_commitment(column), + NaiveCommitment::compute_commitments( + &[CommittableColumn::from(&[4i64, 5, 6][..])], + 0_usize, + &() + )[0] + ); + + let data2 = table([ + borrowed_bigint("a", [1, 2, 3, 4], &alloc), + borrowed_bigint("b", [4, 5, 6, 5], &alloc), + ]); + accessor.add_table(table_ref_2, data2, 0_usize); + + let column = ColumnRef::new(table_ref_1, "a".parse().unwrap(), ColumnType::BigInt); + assert_eq!( + accessor.get_commitment(column), + NaiveCommitment::compute_commitments( + &[CommittableColumn::from(&[1i64, 2, 3][..])], + 0_usize, + &() + )[0] + ); + + let column = ColumnRef::new(table_ref_2, "b".parse().unwrap(), ColumnType::BigInt); + assert_eq!( + accessor.get_commitment(column), + NaiveCommitment::compute_commitments( + &[CommittableColumn::from(&[4i64, 5, 6, 5][..])], + 0_usize, + &() + )[0] + ); +} + +#[test] +fn we_can_access_the_type_of_table_columns() { + let alloc = Bump::new(); + let mut accessor = TableTestAccessor::::new_empty_with_setup(()); + let table_ref_1 = "sxt.test".parse().unwrap(); + let table_ref_2 = "sxt.test2".parse().unwrap(); + + let data1 = table([ + borrowed_bigint("a", [1, 2, 3], &alloc), + borrowed_bigint("b", [4, 5, 6], &alloc), + ]); + accessor.add_table(table_ref_1, data1, 0_usize); + + let column = ColumnRef::new(table_ref_1, "b".parse().unwrap(), ColumnType::BigInt); + assert_eq!( + accessor.lookup_column(column.table_ref(), column.column_id()), + Some(ColumnType::BigInt) + ); + + let column = ColumnRef::new(table_ref_1, "c".parse().unwrap(), ColumnType::BigInt); + assert!(accessor + .lookup_column(column.table_ref(), column.column_id()) + .is_none()); + + let data2 = table([ + borrowed_bigint("a", [1, 2, 3, 4], &alloc), + borrowed_bigint("b", [4, 5, 6, 5], &alloc), + ]); + accessor.add_table(table_ref_2, data2, 0_usize); + + let column = ColumnRef::new(table_ref_1, "a".parse().unwrap(), ColumnType::BigInt); + assert_eq!( + accessor.lookup_column(column.table_ref(), column.column_id()), + Some(ColumnType::BigInt) + ); + + let column = ColumnRef::new(table_ref_2, "b".parse().unwrap(), ColumnType::BigInt); + assert_eq!( + accessor.lookup_column(column.table_ref(), column.column_id()), + Some(ColumnType::BigInt) + ); + + let column = ColumnRef::new(table_ref_2, "c".parse().unwrap(), ColumnType::BigInt); + assert!(accessor + .lookup_column(column.table_ref(), column.column_id()) + .is_none()); +} + +#[test] +fn we_can_access_schema_and_column_names() { + let alloc = Bump::new(); + let mut accessor = TableTestAccessor::::new_empty_with_setup(()); + let table_ref_1 = "sxt.test".parse().unwrap(); + + let data1 = table([ + borrowed_bigint("a", [1, 2, 3], &alloc), + borrowed_varchar("b", ["x", "y", "z"], &alloc), + ]); + accessor.add_table(table_ref_1, data1, 0_usize); + + assert_eq!( + accessor.lookup_schema(table_ref_1), + vec![ + ("a".parse().unwrap(), ColumnType::BigInt), + ("b".parse().unwrap(), ColumnType::VarChar) + ] + ); + assert_eq!(accessor.get_column_names(table_ref_1), vec!["a", "b"]); +} + +#[test] +fn we_can_correctly_update_offsets() { + let alloc = Bump::new(); + let mut accessor1 = TableTestAccessor::::new_empty_with_setup(()); + let table_ref = "sxt.test".parse().unwrap(); + + let data = table([ + borrowed_bigint("a", [1, 2, 3], &alloc), + borrowed_bigint("b", [123, 5, 123], &alloc), + ]); + accessor1.add_table(table_ref, data.clone(), 0_usize); + + let offset = 123; + let mut accessor2 = TableTestAccessor::::new_empty_with_setup(()); + accessor2.add_table(table_ref, data, offset); + + let column = ColumnRef::new(table_ref, "a".parse().unwrap(), ColumnType::BigInt); + assert_ne!( + accessor1.get_commitment(column), + accessor2.get_commitment(column) + ); + let column = ColumnRef::new(table_ref, "b".parse().unwrap(), ColumnType::BigInt); + assert_ne!( + accessor1.get_commitment(column), + accessor2.get_commitment(column) + ); + + assert_eq!(accessor1.get_offset(table_ref), 0); + assert_eq!(accessor2.get_offset(table_ref), offset); + + accessor1.update_offset(table_ref, offset); + + let column = ColumnRef::new(table_ref, "a".parse().unwrap(), ColumnType::BigInt); + assert_eq!( + accessor1.get_commitment(column), + accessor2.get_commitment(column) + ); + let column = ColumnRef::new(table_ref, "b".parse().unwrap(), ColumnType::BigInt); + assert_eq!( + accessor1.get_commitment(column), + accessor2.get_commitment(column) + ); + + assert_eq!(accessor1.get_offset(table_ref), offset); + assert_eq!(accessor2.get_offset(table_ref), offset); +} diff --git a/crates/proof-of-sql/src/base/database/table_utility.rs b/crates/proof-of-sql/src/base/database/table_utility.rs new file mode 100644 index 000000000..5273f04a8 --- /dev/null +++ b/crates/proof-of-sql/src/base/database/table_utility.rs @@ -0,0 +1,338 @@ +//! Utility functions for creating [`Table`]s and [`Column`]s. +//! These functions are primarily intended for use in tests. +//! +//! # Example +//! ``` +//! use bumpalo::Bump; +//! use proof_of_sql::base::{database::table_utility::*, scalar::Curve25519Scalar}; +//! let alloc = Bump::new(); +//! let result = table::([ +//! borrowed_bigint("a", [1, 2, 3], &alloc), +//! borrowed_boolean("b", [true, false, true], &alloc), +//! borrowed_int128("c", [1, 2, 3], &alloc), +//! borrowed_scalar("d", [1, 2, 3], &alloc), +//! borrowed_varchar("e", ["a", "b", "c"], &alloc), +//! borrowed_decimal75("f", 12, 1, [1, 2, 3], &alloc), +//! ]); +//! ``` +use super::{Column, Table}; +use crate::base::scalar::Scalar; +use alloc::{string::String, vec::Vec}; +use bumpalo::Bump; +use core::ops::Deref; +use proof_of_sql_parser::{ + posql_time::{PoSQLTimeUnit, PoSQLTimeZone}, + Identifier, +}; + +/// Creates an [`Table`] from a list of `(Identifier, Column)` pairs. +/// This is a convenience wrapper around [`Table::try_from_iter`] primarily for use in tests and +/// intended to be used along with the other methods in this module (e.g. [bigint], [boolean], etc). +/// The function will panic under a variety of conditions. See [`Table::try_from_iter`] for more details. +/// +/// # Example +/// ``` +/// use bumpalo::Bump; +/// use proof_of_sql::base::{database::table_utility::*, scalar::Curve25519Scalar}; +/// let alloc = Bump::new(); +/// let result = table::([ +/// borrowed_bigint("a", [1, 2, 3], &alloc), +/// borrowed_boolean("b", [true, false, true], &alloc), +/// borrowed_int128("c", [1, 2, 3], &alloc), +/// borrowed_scalar("d", [1, 2, 3], &alloc), +/// borrowed_varchar("e", ["a", "b", "c"], &alloc), +/// borrowed_decimal75("f", 12, 1, [1, 2, 3], &alloc), +/// ]); +/// ``` +/// +/// # Panics +/// - Panics if converting the iterator into an `Table<'a, S>` fails. +pub fn table<'a, S: Scalar>( + iter: impl IntoIterator)>, +) -> Table<'a, S> { + Table::try_from_iter(iter).unwrap() +} + +/// Creates a (Identifier, `Column`) pair for a tinyint column. +/// This is primarily intended for use in conjunction with [`table`]. +/// # Example +/// ``` +/// use bumpalo::Bump; +/// use proof_of_sql::base::{database::table_utility::*, scalar::Curve25519Scalar}; +/// let alloc = Bump::new(); +/// let result = table::([ +/// borrowed_tinyint("a", [1_i8, 2, 3], &alloc), +/// ]); +///``` +/// # Panics +/// - Panics if `name.parse()` fails to convert the name into an `Identifier`. +pub fn borrowed_tinyint( + name: impl Deref, + data: impl IntoIterator>, + alloc: &Bump, +) -> (Identifier, Column<'_, S>) { + let transformed_data: Vec = data.into_iter().map(Into::into).collect(); + let alloc_data = alloc.alloc_slice_copy(&transformed_data); + (name.parse().unwrap(), Column::TinyInt(alloc_data)) +} + +/// Creates a `(Identifier, Column)` pair for a smallint column. +/// This is primarily intended for use in conjunction with [`table`]. +/// +/// # Example +/// ```rust +/// use bumpalo::Bump; +/// use proof_of_sql::base::{database::table_utility::*, scalar::Curve25519Scalar}; +/// let alloc = Bump::new(); +/// let result = table::([ +/// borrowed_smallint("a", [1_i16, 2, 3], &alloc), +/// ]); +/// ``` +/// +/// # Panics +/// - Panics if `name.parse()` fails to convert the name into an `Identifier`. +pub fn borrowed_smallint( + name: impl Deref, + data: impl IntoIterator>, + alloc: &Bump, +) -> (Identifier, Column<'_, S>) { + let transformed_data: Vec = data.into_iter().map(Into::into).collect(); + let alloc_data = alloc.alloc_slice_copy(&transformed_data); + (name.parse().unwrap(), Column::SmallInt(alloc_data)) +} + +/// Creates a `(Identifier, Column)` pair for an int column. +/// This is primarily intended for use in conjunction with [`table`]. +/// +/// # Example +/// ```rust +/// use bumpalo::Bump; +/// use proof_of_sql::base::{database::table_utility::*, scalar::Curve25519Scalar}; +/// let alloc = Bump::new(); +/// let result = table::([ +/// borrowed_int("a", [1, 2, 3], &alloc), +/// ]); +/// ``` +/// +/// # Panics +/// - Panics if `name.parse()` fails to convert the name into an `Identifier`. +pub fn borrowed_int( + name: impl Deref, + data: impl IntoIterator>, + alloc: &Bump, +) -> (Identifier, Column<'_, S>) { + let transformed_data: Vec = data.into_iter().map(Into::into).collect(); + let alloc_data = alloc.alloc_slice_copy(&transformed_data); + (name.parse().unwrap(), Column::Int(alloc_data)) +} + +/// Creates a `(Identifier, Column)` pair for a bigint column. +/// This is primarily intended for use in conjunction with [`table`]. +/// +/// # Example +/// ```rust +/// use bumpalo::Bump; +/// use proof_of_sql::base::{database::table_utility::*, scalar::Curve25519Scalar}; +/// let alloc = Bump::new(); +/// let result = table::([ +/// borrowed_bigint("a", [1, 2, 3], &alloc), +/// ]); +/// ``` +/// +/// # Panics +/// - Panics if `name.parse()` fails to convert the name into an `Identifier`. +pub fn borrowed_bigint( + name: impl Deref, + data: impl IntoIterator>, + alloc: &Bump, +) -> (Identifier, Column<'_, S>) { + let transformed_data: Vec = data.into_iter().map(Into::into).collect(); + let alloc_data = alloc.alloc_slice_copy(&transformed_data); + (name.parse().unwrap(), Column::BigInt(alloc_data)) +} + +/// Creates a `(Identifier, Column)` pair for a boolean column. +/// This is primarily intended for use in conjunction with [`table`]. +/// +/// # Example +/// ``` +/// use bumpalo::Bump; +/// use proof_of_sql::base::{database::table_utility::*, scalar::Curve25519Scalar}; +/// let alloc = Bump::new(); +/// let result = table::([ +/// borrowed_boolean("a", [true, false, true], &alloc), +/// ]); +/// ``` +/// +/// # Panics +/// - Panics if `name.parse()` fails to convert the name into an `Identifier`. +pub fn borrowed_boolean( + name: impl Deref, + data: impl IntoIterator>, + alloc: &Bump, +) -> (Identifier, Column<'_, S>) { + let transformed_data: Vec = data.into_iter().map(Into::into).collect(); + let alloc_data = alloc.alloc_slice_copy(&transformed_data); + (name.parse().unwrap(), Column::Boolean(alloc_data)) +} + +/// Creates a `(Identifier, Column)` pair for an int128 column. +/// This is primarily intended for use in conjunction with [`table`]. +/// +/// # Example +/// ``` +/// use bumpalo::Bump; +/// use proof_of_sql::base::{database::table_utility::*, scalar::Curve25519Scalar}; +/// let alloc = Bump::new(); +/// let result = table::([ +/// borrowed_int128("a", [1, 2, 3], &alloc), +/// ]); +/// ``` +/// +/// # Panics +/// - Panics if `name.parse()` fails to convert the name into an `Identifier`. +pub fn borrowed_int128( + name: impl Deref, + data: impl IntoIterator>, + alloc: &Bump, +) -> (Identifier, Column<'_, S>) { + let transformed_data: Vec = data.into_iter().map(Into::into).collect(); + let alloc_data = alloc.alloc_slice_copy(&transformed_data); + (name.parse().unwrap(), Column::Int128(alloc_data)) +} + +/// Creates a `(Identifier, Column)` pair for a scalar column. +/// This is primarily intended for use in conjunction with [`table`]. +/// +/// # Example +/// ``` +/// use bumpalo::Bump; +/// use proof_of_sql::base::{database::table_utility::*, scalar::Curve25519Scalar}; +/// let alloc = Bump::new(); +/// let result = table::([ +/// borrowed_scalar("a", [1, 2, 3], &alloc), +/// ]); +/// ``` +/// +/// # Panics +/// - Panics if `name.parse()` fails to convert the name into an `Identifier`. +pub fn borrowed_scalar( + name: impl Deref, + data: impl IntoIterator>, + alloc: &Bump, +) -> (Identifier, Column<'_, S>) { + let transformed_data: Vec = data.into_iter().map(Into::into).collect(); + let alloc_data = alloc.alloc_slice_copy(&transformed_data); + (name.parse().unwrap(), Column::Scalar(alloc_data)) +} + +/// Creates a `(Identifier, Column)` pair for a varchar column. +/// This is primarily intended for use in conjunction with [`table`]. +/// # Example +/// ``` +/// use bumpalo::Bump; +/// use proof_of_sql::base::{database::table_utility::*, scalar::Curve25519Scalar}; +/// let alloc = Bump::new(); +/// let result = table::([ +/// borrowed_varchar("a", ["a", "b", "c"], &alloc), +/// ]); +/// ``` +/// +/// # Panics +/// - Panics if `name.parse()` fails to convert the name into an `Identifier`. +pub fn borrowed_varchar<'a, S: Scalar>( + name: impl Deref, + data: impl IntoIterator>, + alloc: &'a Bump, +) -> (Identifier, Column<'a, S>) { + let strings: Vec<&'a str> = data + .into_iter() + .map(|item| { + let string = item.into(); + alloc.alloc_str(&string) as &'a str + }) + .collect(); + let alloc_strings = alloc.alloc_slice_clone(&strings); + let scalars: Vec = strings.iter().map(|s| (*s).into()).collect(); + let alloc_scalars = alloc.alloc_slice_copy(&scalars); + ( + name.parse().unwrap(), + Column::VarChar((alloc_strings, alloc_scalars)), + ) +} + +/// Creates a `(Identifier, Column)` pair for a decimal75 column. +/// This is primarily intended for use in conjunction with [`table`]. +/// # Example +/// ``` +/// use bumpalo::Bump; +/// use proof_of_sql::base::{database::table_utility::*, scalar::Curve25519Scalar}; +/// let alloc = Bump::new(); +/// let result = table::([ +/// borrowed_decimal75("a", 12, 1, [1, 2, 3], &alloc), +/// ]); +/// ``` +/// +/// # Panics +/// - Panics if `name.parse()` fails to convert the name into an `Identifier`. +/// - Panics if creating the `Precision` from the specified precision value fails. +pub fn borrowed_decimal75( + name: impl Deref, + precision: u8, + scale: i8, + data: impl IntoIterator>, + alloc: &Bump, +) -> (Identifier, Column<'_, S>) { + let transformed_data: Vec = data.into_iter().map(Into::into).collect(); + let alloc_data = alloc.alloc_slice_copy(&transformed_data); + ( + name.parse().unwrap(), + Column::Decimal75( + crate::base::math::decimal::Precision::new(precision).unwrap(), + scale, + alloc_data, + ), + ) +} + +/// Creates a `(Identifier, Column)` pair for a timestamp column. +/// This is primarily intended for use in conjunction with [`table`]. +/// +/// # Parameters +/// - `name`: The name of the column. +/// - `time_unit`: The time unit of the timestamps. +/// - `timezone`: The timezone for the timestamps. +/// - `data`: The data for the column, provided as an iterator over `i64` values representing time since the unix epoch. +/// - `alloc`: The bump allocator to use for allocating the column data. +/// +/// # Example +/// ``` +/// use bumpalo::Bump; +/// use proof_of_sql::base::{database::table_utility::*, +/// scalar::Curve25519Scalar, +/// }; +/// use proof_of_sql_parser::{ +/// posql_time::{PoSQLTimeZone, PoSQLTimeUnit}}; +/// +/// let alloc = Bump::new(); +/// let result = table::([ +/// borrowed_timestamptz("event_time", PoSQLTimeUnit::Second, PoSQLTimeZone::Utc, vec![1625072400, 1625076000, 1625079600], &alloc), +/// ]); +/// ``` +/// +/// # Panics +/// - Panics if `name.parse()` fails to convert the name into an `Identifier`. +pub fn borrowed_timestamptz( + name: impl Deref, + time_unit: PoSQLTimeUnit, + timezone: PoSQLTimeZone, + data: impl IntoIterator, + alloc: &Bump, +) -> (Identifier, Column<'_, S>) { + let vec_data: Vec = data.into_iter().collect(); + let alloc_data = alloc.alloc_slice_copy(&vec_data); + ( + name.parse().unwrap(), + Column::TimestampTZ(time_unit, timezone, alloc_data), + ) +}