fuzzer working

Author: jayshrivastava

Cargo.lock

@@ -0,0 +1,271 @@
+use clap::{Parser, ValueEnum};
+use datafusion::error::Result;
+use datafusion_distributed::test_utils::{
+    fuzz::{FuzzDB, SessionConfig},
+    tpcds::{ensure_tpcds_data, get_queries_dir, register_available_tpcds_tables},
+};
+use std::fs;
+use std::process;
+use std::time::Instant;
+
+#[derive(Parser)]
+#[command(author, version, about, long_about = None)]
+#[command(name = "fuzz")]
+#[command(about = "Fuzz test distributed DataFusion with various workloads")]
+struct Cli {
+    /// Number of workers to use (must be >= 1)
+    #[arg(short, long, default_value_t = 3)]
+    workers: usize,
+
+    /// Workload to run
+    #[arg(long, value_enum, default_value_t = Workload::Tpcds)]
+    workload: Workload,
+
+    /// Scale factor for TPCDS data generation
+    #[arg(short, long, default_value = "0.01")]
+    scale_factor: String,
+
+    /// Generate data even if it already exists
+    #[arg(long, default_value_t = false)]
+    force_regenerate: bool,
+
+    /// Run only specific queries (comma-separated list, e.g., "q1,q5,q10")
+    #[arg(long)]
+    queries: Option<String>,
+
+    /// Verbose output
+    #[arg(short, long, default_value_t = false)]
+    verbose: bool,
+}
+
+#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum, Debug)]
+enum Workload {
+    /// TPC-DS benchmark queries
+    Tpcds,
+}
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    let cli = Cli::parse();
+
+    // Validate arguments
+    if let Err(e) = validate_args(&cli) {
+        eprintln!("Error: {}", e);
+        process::exit(1);
+    }
+
+    println!("🚀 Starting DataFusion distributed fuzz testing");
+    println!("   Workers: {}", cli.workers);
+    println!("   Workload: {:?}", cli.workload);
+    println!("   Scale factor: {}", cli.scale_factor);
+    println!();
+
+    match cli.workload {
+        Workload::Tpcds => {
+            if let Err(e) = run_tpcds_workload(cli).await {
+                eprintln!("❌ Fuzz testing failed: {}", e);
+                process::exit(1);
+            }
+        }
+    }
+
+    println!("✅ All fuzz tests passed!");
+    Ok(())
+}
+
+/// Validate command line arguments
+fn validate_args(cli: &Cli) -> std::result::Result<(), String> {
+    if cli.workers < 1 {
+        return Err("Number of workers must be >= 1".to_string());
+    }
+
+    // Validate scale factor is numeric (decimal allowed)
+    if cli.scale_factor.parse::<f64>().is_err() {
+        return Err(format!("Scale factor '{}' is not a valid number", cli.scale_factor));
+    }
+
+    Ok(())
+}
+
+/// Run the TPC-DS workload
+async fn run_tpcds_workload(cli: Cli) -> Result<()> {
+    println!("📊 Running TPC-DS workload...");
+
+    // Ensure TPCDS data is available
+    println!("🔧 Ensuring TPC-DS data is available (scale factor: {})...", cli.scale_factor);
+    if let Err(e) = ensure_tpcds_data(Some(&cli.scale_factor), cli.force_regenerate) {
+        eprintln!("⚠️  Warning: Failed to ensure TPCDS data: {}", e);
+        println!("   Continuing anyway - will skip queries that require missing tables");
+    }
+
+    // Create FuzzDB
+    println!("⚙️  Setting up distributed session with {} workers...", cli.workers);
+    let config = SessionConfig {
+        num_workers: cli.workers,
+        tasks_per_file: 4,  // Use reasonable defaults
+        cardinality_task_count_factor: 4,
+        target_partitions: cli.workers * 2,
+    };
+
+    let fuzz_db = match FuzzDB::new_with_config(config, |ctx| {
+        let result = tokio::task::block_in_place(|| {
+            tokio::runtime::Handle::current().block_on(async {
+                let (registered_tables, missing_tables) = register_available_tpcds_tables(ctx, None).await?;
+                
+                if !missing_tables.is_empty() {
+                    println!("Missing TPCDS tables: {:?}", missing_tables);
+                }
+                
+                Ok::<Vec<String>, datafusion::error::DataFusionError>(registered_tables)
+            })
+        });
+        result
+    }).await {
+        Ok(db) => {
+            if db.registered_tables.is_empty() {
+                println!("⚠️  Warning: No tables were registered. Some queries may fail.");
+            } else {
+                println!("✅ Successfully registered {} TPCDS tables", db.registered_tables.len());
+                if cli.verbose {
+                    println!("   Registered tables: {:?}", db.registered_tables);
+                }
+            }
+            db
+        }
+        Err(e) => {
+            return Err(datafusion::error::DataFusionError::Execution(format!(
+                "Failed to create FuzzDB: {}", e
+            )));
+        }
+    };
+
+    // Discover queries
+    let queries = discover_tpcds_queries(&cli)?;
+    println!("🔍 Found {} TPC-DS queries to execute", queries.len());
+
+    if queries.is_empty() {
+        return Err(datafusion::error::DataFusionError::Execution(
+            "No queries found to execute".to_string()
+        ));
+    }
+
+    // Run each query
+    let mut successful = 0;
+    let mut failed = 0;
+    let start_time = Instant::now();
+
+    for (query_name, query_sql) in queries {
+        if cli.verbose {
+            println!();
+            println!("🔄 Executing query: {}", query_name);
+            println!("Executing fuzz query: {}", query_sql.trim());
+            println!();
+        }
+
+        let query_start = Instant::now();
+        match fuzz_db.run(&query_sql).await {
+            Ok(results) => {
+                let query_duration = query_start.elapsed();
+                let total_rows: usize = results.iter().map(|b| b.num_rows()).sum();
+                successful += 1;
+                println!("✅ {} completed successfully", query_name);
+                if cli.verbose {
+                    println!("   Duration: {:?}", query_duration);
+                    println!("   Results: {} rows in {} batch(es)", total_rows, results.len());
+                }
+            }
+            Err(e) => {
+                failed += 1;
+                println!("❌ {} failed: {}", query_name, e);
+                if cli.verbose {
+                    println!("   Error details: {:?}", e);
+                }
+            }
+        }
+    }
+
+    let total_duration = start_time.elapsed();
+    println!();
+    println!("📈 Summary:");
+    println!("   Total queries: {}", successful + failed);
+    println!("   Successful: {}", successful);
+    println!("   Failed: {}", failed);
+    println!("   Total duration: {:?}", total_duration);
+
+    if failed > 0 {
+        return Err(datafusion::error::DataFusionError::Execution(format!(
+            "{} out of {} queries failed", failed, successful + failed
+        )));
+    }
+
+    Ok(())
+}
+
+/// Discover TPC-DS queries from the queries directory
+fn discover_tpcds_queries(cli: &Cli) -> Result<Vec<(String, String)>> {
+    let queries_dir = get_queries_dir();
+    
+    if !queries_dir.exists() {
+        return Err(datafusion::error::DataFusionError::Execution(format!(
+            "TPC-DS queries directory not found: {}", queries_dir.display()
+        )));
+    }
+
+    let mut queries = Vec::new();
+    let entries = fs::read_dir(&queries_dir).map_err(|e| {
+        datafusion::error::DataFusionError::Execution(format!(
+            "Failed to read queries directory {}: {}", queries_dir.display(), e
+        ))
+    })?;
+
+    // Get list of specific queries if provided
+    let specific_queries: Option<std::collections::HashSet<String>> = cli.queries.as_ref().map(|q| {
+        q.split(',')
+            .map(|s| s.trim().to_lowercase())
+            .collect()
+    });
+
+    for entry in entries {
+        let entry = entry.map_err(|e| {
+            datafusion::error::DataFusionError::Execution(format!(
+                "Failed to read directory entry: {}", e
+            ))
+        })?;
+
+        let path = entry.path();
+        if path.extension().and_then(|s| s.to_str()) == Some("sql") {
+            if let Some(file_stem) = path.file_stem().and_then(|s| s.to_str()) {
+                let query_name = file_stem.to_lowercase();
+                
+                // Filter specific queries if requested
+                if let Some(ref specific) = specific_queries {
+                    if !specific.contains(&query_name) {
+                        continue;
+                    }
+                }
+
+                let query_sql = fs::read_to_string(&path).map_err(|e| {
+                    datafusion::error::DataFusionError::Execution(format!(
+                        "Failed to read query file {}: {}", path.display(), e
+                    ))
+                })?;
+
+                queries.push((query_name, query_sql));
+            }
+        }
+    }
+
+    // Sort queries for consistent execution order
+    queries.sort_by(|a, b| {
+        // Extract query number for natural sorting (q1, q2, ..., q10, etc.)
+        let extract_num = |name: &str| {
+            name.strip_prefix('q')
+                .and_then(|s| s.parse::<u32>().ok())
+                .unwrap_or(0)
+        };
+        
+        extract_num(&a.0).cmp(&extract_num(&b.0))
+    });
+
+    Ok(queries)
+}