refactor: make random table self-contained

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

@@ -38,10 +38,7 @@ mod record_batch_test_accessor_test;
 #[cfg(any(test, feature = "test"))]
 mod test_accessor_utility;
 #[cfg(any(test, feature = "test"))]
-pub use test_accessor_utility::{
-    make_random_test_accessor_data, make_random_test_accessor_owned_table,
-    RandomTestAccessorDescriptor,
-};
+pub use test_accessor_utility::{make_random_test_accessor_data, RandomTestAccessorDescriptor};
 
 mod owned_column;
 pub use owned_column::OwnedColumn;

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

@@ -1,5 +1,4 @@
-use super::{OwnedColumn, OwnedTable};
-use crate::base::{database::ColumnType, scalar::Scalar, time::timestamp::PoSQLTimeUnit};
+use crate::base::{database::ColumnType, time::timestamp::PoSQLTimeUnit};
 use arrow::{
     array::{
         Array, BooleanArray, Decimal128Array, Decimal256Array, Int16Array, Int32Array, Int64Array,
@@ -145,57 +144,6 @@ pub fn make_random_test_accessor_data(
     RecordBatch::try_new(schema, columns).unwrap()
 }
 
-/// Generate a OwnedTable with random data
-///
-/// Currently, this mirrors [make_random_test_accessor_data] and is intended to replace it.
-pub fn make_random_test_accessor_owned_table<S: Scalar>(
-    rng: &mut StdRng,
-    cols: &[(&str, ColumnType)],
-    descriptor: &RandomTestAccessorDescriptor,
-) -> OwnedTable<S> {
-    let n = Uniform::new(descriptor.min_rows, descriptor.max_rows + 1).sample(rng);
-    let dist = Uniform::new(descriptor.min_value, descriptor.max_value + 1);
-
-    OwnedTable::try_from_iter(cols.iter().map(|(col_name, col_type)| {
-        let values = dist.sample_iter(&mut *rng).take(n);
-        (
-            col_name.parse().unwrap(),
-            match col_type {
-                ColumnType::Boolean => OwnedColumn::Boolean(values.map(|x| x % 2 != 0).collect()),
-                ColumnType::SmallInt => {
-                    OwnedColumn::SmallInt(
-                        values
-                            .map(|x| ((x >> 48) as i16)) // Shift right to align the lower 16 bits
-                            .collect(),
-                    )
-                }
-                ColumnType::Int => {
-                    OwnedColumn::Int(
-                        values
-                            .map(|x| ((x >> 32) as i32)) // Shift right to align the lower 32 bits
-                            .collect(),
-                    )
-                }
-                ColumnType::BigInt => OwnedColumn::BigInt(values.collect()),
-                ColumnType::Int128 => OwnedColumn::Int128(values.map(|x| x as i128).collect()),
-                ColumnType::Decimal75(precision, scale) => {
-                    OwnedColumn::Decimal75(*precision, *scale, values.map(Into::into).collect())
-                }
-                ColumnType::VarChar => OwnedColumn::VarChar(
-                    values
-                        .map(|v| "s".to_owned() + &v.to_string()[..])
-                        .collect(),
-                ),
-                ColumnType::Scalar => OwnedColumn::Scalar(values.map(Into::into).collect()),
-                ColumnType::TimestampTZ(tu, tz) => {
-                    OwnedColumn::TimestampTZ(*tu, *tz, values.collect())
-                }
-            },
-        )
-    }))
-    .unwrap()
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;

crates/proof-of-sql/src/sql/ast/and_expr_test.rs

@@ -2,10 +2,7 @@ use super::{test_utility::*, ProvableExpr};
 use crate::{
     base::{
         commitment::InnerProductProof,
-        database::{
-            make_random_test_accessor_owned_table, owned_table_utility::*, Column, ColumnType,
-            OwnedTableTestAccessor, RandomTestAccessorDescriptor, TestAccessor,
-        },
+        database::{owned_table_utility::*, Column, OwnedTableTestAccessor, TestAccessor},
     },
     sql::{
         ast::{
@@ -75,24 +72,24 @@ fn we_can_prove_a_simple_and_query_with_128_bits() {
 }
 
 fn test_random_tables_with_given_offset(offset: usize) {
-    let descr = RandomTestAccessorDescriptor {
-        min_rows: 1,
-        max_rows: 20,
-        min_value: -3,
-        max_value: 3,
-    };
+    let dist = Uniform::new(-3, 4);
     let mut rng = StdRng::from_seed([0u8; 32]);
-    let cols = [
-        ("a", ColumnType::BigInt),
-        ("b", ColumnType::VarChar),
-        ("c", ColumnType::BigInt),
-        ("d", ColumnType::VarChar),
-    ];
     for _ in 0..20 {
-        let data = make_random_test_accessor_owned_table(&mut rng, &cols, &descr);
-        let filter_val1 = Uniform::new(descr.min_value, descr.max_value + 1).sample(&mut rng);
-        let filter_val1 = format!("s{filter_val1}");
-        let filter_val2 = Uniform::new(descr.min_value, descr.max_value + 1).sample(&mut rng);
+        let n = Uniform::new(1, 21).sample(&mut rng);
+        let data = owned_table([
+            bigint("a", dist.sample_iter(&mut rng).take(n)),
+            varchar(
+                "b",
+                dist.sample_iter(&mut rng).take(n).map(|v| format!("s{v}")),
+            ),
+            bigint("c", dist.sample_iter(&mut rng).take(n)),
+            varchar(
+                "d",
+                dist.sample_iter(&mut rng).take(n).map(|v| format!("s{v}")),
+            ),
+        ]);
+        let filter_val1 = format!("s{}", dist.sample(&mut rng));
+        let filter_val2 = dist.sample(&mut rng);
 
         let t = "sxt.t".parse().unwrap();
         let accessor = OwnedTableTestAccessor::<InnerProductProof>::new_from_table(