working

Author: jayshrivastava

src/bin/fuzz.rs

@@ -51,6 +51,7 @@ enum Commands {
 }
 
 /// TPC-DS workload implementation
+#[derive(Clone)]
 struct TpcdsWorkload {
     scale_factor: String,
     force_regenerate: bool,
@@ -69,7 +70,7 @@ impl TpcdsWorkload {
 
 #[async_trait]
 impl Workload for TpcdsWorkload {
-    async fn setup(&self, ctx: &SessionContext) -> Result<Vec<String>> {
+    async fn setup(&self, ctx: &SessionContext) -> Result<()> {
         // Register tables
         info!("🔧 Registering TPC-DS tables...");
         let (registered_tables, missing_tables) = register_available_tpcds_tables(ctx, None).await?;
@@ -120,7 +121,7 @@ async fn main() -> Result<()> {
             info!("   Workload: {}", workload.name());
             info!("   Scale factor: {}", scale_factor);
 
-            if let Err(e) = run_workload(Box::new(workload)).await {
+            if let Err(e) = run_workload(workload).await {
                 error!("❌ Fuzz testing failed: {}", e);
                 process::exit(1);
             }
@@ -143,16 +144,20 @@ fn validate_tpcds_args(scale_factor: &str) -> std::result::Result<(), String> {
 }
 
 /// Run a workload using the generic workload trait
-async fn run_workload(workload: Box<dyn Workload + Send + Sync>) -> Result<()> {
+async fn run_workload(workload: TpcdsWorkload) -> Result<()> {
     info!("📊 Running {} workload...", workload.name());
 
     // Create FuzzDB with randomized session config
     info!("⚙️  Setting up distributed session with randomized configuration...");
 
-    let fuzz_db = match FuzzDB::new(|ctx| {
-        let ctx = ctx.clone();
-        let workload = &workload;
-        workload.setup(&ctx)
+    let fuzz_db = match FuzzDB::new({
+        let workload_clone = workload.clone();
+        move |ctx| {
+            let workload_inner = workload_clone.clone();
+            async move {
+                workload_inner.setup(&ctx).await
+            }
+        }
     }).await {
         Ok(db) => db,
         Err(e) => {

src/test_utils/fuzz.rs

@@ -9,7 +9,6 @@ use datafusion::{
     logical_expr::LogicalPlan,
 };
 use rand::Rng;
-use std::collections::HashSet;
 
 /// Fuzzing database with distributed session context and helper functions
 pub struct FuzzDB {
@@ -42,8 +41,11 @@ impl Default for SessionConfig {
 }
 
 impl FuzzDB {
-    /// Create a new FuzzDB with randomized session parameters
-    pub async fn new() -> Result<Self> 
+    /// Create a new FuzzDB with randomized session parameters and setup function
+    pub async fn new<F, Fut>(setup: F) -> Result<Self> 
+    where
+        F: Fn(SessionContext) -> Fut + Send + Sync,
+        Fut: std::future::Future<Output = Result<()>> + Send
     {
         let config = randomize_session_config();
         create_db(config, setup).await
@@ -54,12 +56,12 @@ impl FuzzDB {
 
         // Execute on distributed context
         let df = self.distributed_ctx.sql(query).await?;
-        let logical_plan = df.logical_plan().clone();
+        let _logical_plan = df.logical_plan().clone();
         let results = df.collect().await?;
 
         // Run oracles
-        let single_node_oracle = SingleNodeOracle::new(&self.single_node_ctx);
-        let ordering_oracle = OrderingOracle::new();
+        let _single_node_oracle = SingleNodeOracle::new(&self.single_node_ctx);
+        let _ordering_oracle = OrderingOracle::new();
 
         // Validate with SingleNodeOracle
         // single_node_oracle.validate(&self.distributed_ctx, query, &results).await?;
@@ -87,7 +89,10 @@ fn randomize_session_config() -> SessionConfig {
 }
 
 /// Create distributed session context with specified configuration
-async fn create_db<F, Fut>(config: SessionConfig) -> Result<FuzzDB>
+async fn create_db<F, Fut>(config: SessionConfig, setup: F) -> Result<FuzzDB>
+where
+    F: Fn(SessionContext) -> Fut + Send + Sync,
+    Fut: std::future::Future<Output = Result<()>> + Send
 {
     println!("Creating FuzzDB with {} workers", config.num_workers);
 
@@ -114,8 +119,8 @@ async fn create_db<F, Fut>(config: SessionConfig) -> Result<FuzzDB>
     let single_node_ctx = SessionContext::new();
 
     // Call the setup function to register tables on both contexts
-    setup(&distributed_ctx).await?;
-    setup(&single_node_ctx).await?;
+    setup(distributed_ctx.clone()).await?;
+    setup(single_node_ctx.clone()).await?;
 
     // Log worker configuration
     println!("Session configuration:");
@@ -297,8 +302,6 @@ fn records_equal_as_sets(left: &[RecordBatch], right: &[RecordBatch]) -> bool {
 
 /// Perform detailed comparison of record batches
 fn detailed_batch_comparison(left: &[RecordBatch], right: &[RecordBatch]) -> bool {
-    use std::collections::HashSet;
-    
     // Convert both sides to sets of string representations of rows
     let left_rows = batch_rows_to_strings(left);
     let right_rows = batch_rows_to_strings(right);
@@ -307,8 +310,8 @@ fn detailed_batch_comparison(left: &[RecordBatch], right: &[RecordBatch]) -> boo
         return false;
     }
 
-    let left_set: HashSet<_> = left_rows.into_iter().collect();
-    let right_set: HashSet<_> = right_rows.into_iter().collect();
+    let left_set: std::collections::HashSet<_> = left_rows.into_iter().collect();
+    let right_set: std::collections::HashSet<_> = right_rows.into_iter().collect();
 
     if left_set != right_set {
         println!("Row content differs between result sets");