Add fuzzing infrastructure for distributed DataFusion

This change introduces a new `fuzz` binary that enables fuzzing of distributed
DataFusion against a localhost cluster. The fuzzing framework is designed to
be extensible, allowing us to define custom workloads and oracles.

This change defines a TPC-DS workload which uses the duckdb data generator and 99
TPC-DS queries from their github repo. It also defines one oracle, the `SingleNodeOracle`
which validates the set of rows produced by distributed datafusion against the set of
rows produced by single node datafusion (if a query errors in both, then this is not
counted as a failure). You can run this using the following command:
```
RUST_LOG=info cargo run --features integration --bin fuzz -- tpcds --force-regenerate
```

Fuzzing also uses randomized cluster configurations using a deterministic seed.

Next steps:
- Add ordering oracle to validate ORDER BY correctness
  - Idea: Inspect the ordering properties in the logical plan and assert this
    on the RecordBatches
- Observability
  - Log stats on queries that were invalid (ie. failed to execute on single node df)
  - Log localhost cluster ports
- Add metrics oracle to validate output_rows metric (ensuring metrics are
  working correctly)
- Set up nightly github actions workflow to run the fuzzer automatically
  - Ensure that the data is available to be downloaded so we can reproduce any failures
    locally

Author: jayshrivastava

.gitignore

@@ -1,4 +1,5 @@
 /.idea
 /target
 /benchmarks/data/
-testdata/tpch/data/
+testdata/tpch/data/
+testdata/tpcds/data/

Cargo.lock

@@ -10,6 +10,11 @@ name = "datafusion-distributed"
 version = "0.1.0"
 edition = "2024"
 
+[[bin]]
+name = "fuzz"
+path = "src/bin/fuzz.rs"
+required-features = ["integration"]
+
 [dependencies]
 chrono = { version = "0.4.42" }
 datafusion = { workspace = true }
@@ -30,8 +35,12 @@ delegate = "0.13.4"
 dashmap = "6.1.0"
 prost = "0.13.5"
 rand = "0.8.5"
+rand_chacha = "0.3.1"
+base64 = "0.22"
 object_store = "0.12.3"
 bytes = "1.10.1"
+log = "0.4"
+env_logger = "0.11"
 
 # integration_tests deps
 insta = { version = "1.43.1", features = ["filters"], optional = true }
@@ -42,6 +51,7 @@ arrow = { version = "56.1.0", optional = true }
 tokio-stream = { version = "0.1.17", optional = true }
 hyper-util = { version = "0.1.16", optional = true }
 pin-project = "1.1.10"
+clap = { version = "4.5", features = ["derive"] }
 
 [features]
 integration = [

Cargo.toml

@@ -0,0 +1,227 @@
+use arrow::util::pretty::pretty_format_batches;
+use async_trait::async_trait;
+use clap::{Parser, Subcommand};
+use datafusion::error::Result;
+use datafusion::prelude::SessionContext;
+use datafusion_distributed::test_utils::{
+    fuzz::{FuzzConfig, FuzzDB},
+    rand::rng,
+    tpcds::{generate_tpcds_data, queries, register_tables},
+};
+use log::{Level, debug, error, info, log_enabled};
+use rand::Rng;
+use std::process;
+
+#[derive(Parser)]
+#[command(author, version, about, long_about = None)]
+#[command(name = "fuzz")]
+#[command(about = "Fuzz test distributed datafusion with various workloads")]
+struct Cli {
+    #[command(subcommand)]
+    command: Commands,
+
+    /// Base64-encoded seed for random number generator. This seed is used to
+    /// randomize configuration options and not used to generate data.
+    #[arg(long)]
+    seed: Option<String>,
+}
+
+/// Trait for defining fuzz workloads
+#[async_trait]
+trait Workload {
+    /// Do any necessary setup for the workload such as registering tables
+    async fn setup(ctx: SessionContext) -> Result<()>;
+
+    /// Iterator of (query_name, query_sql) pairs  
+    async fn queries(&self) -> Result<Box<dyn Iterator<Item = (String, String)> + Send>>;
+
+    /// Get workload name for logging
+    fn name(&self) -> &'static str;
+}
+
+#[derive(Subcommand)]
+enum Commands {
+    /// Run TPC-DS benchmark queries
+    Tpcds {
+        /// 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>,
+    },
+}
+
+/// TPC-DS workload implementation
+#[derive(Clone)]
+struct TpcdsWorkload {
+    queries: Option<String>,
+}
+
+impl TpcdsWorkload {
+    fn try_new(
+        scale_factor: String,
+        force_regenerate: bool,
+        queries: Option<String>,
+    ) -> Result<Self> {
+        if force_regenerate {
+            info!(
+                "Generating TPC-DS data with scale factor {}...",
+                scale_factor
+            );
+            generate_tpcds_data(&scale_factor)?;
+            info!("Done");
+        } else {
+            info!("Using existing TPC-DS data");
+        }
+
+        Ok(Self { queries })
+    }
+}
+
+#[async_trait]
+impl Workload for TpcdsWorkload {
+    async fn setup(ctx: SessionContext) -> Result<()> {
+        info!("Registering TPC-DS tables...");
+        let registered_tables = register_tables(&ctx).await?;
+        info!("Done");
+        debug!("Registered tables: {:?}", registered_tables);
+        Ok(())
+    }
+
+    async fn queries(&self) -> Result<Box<dyn Iterator<Item = (String, String)> + Send>> {
+        let queries_vec = queries(self.queries.as_deref())?;
+        Ok(Box::new(queries_vec.into_iter()))
+    }
+
+    fn name(&self) -> &'static str {
+        "TPC-DS"
+    }
+}
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    env_logger::init();
+    let cli = Cli::parse();
+
+    match cli.command {
+        Commands::Tpcds {
+            scale_factor,
+            force_regenerate,
+            queries,
+        } => {
+            // Validate arguments
+            if let Err(e) = validate_tpcds_args(&scale_factor) {
+                error!("Error: {}", e);
+                process::exit(1);
+            }
+
+            info!("🌀 Starting Distributed DataFusion Fuzz Test");
+
+            let workload = TpcdsWorkload::try_new(scale_factor.clone(), force_regenerate, queries)?;
+
+            if let Err(e) = run_workload(workload, cli.seed).await {
+                error!("❌ Fuzz testing failed: {}", e);
+                process::exit(1);
+            }
+
+            info!("✅ All fuzz tests passed!");
+        }
+    }
+
+    Ok(())
+}
+
+/// Validate TPC-DS command line arguments
+fn validate_tpcds_args(scale_factor: &str) -> std::result::Result<(), String> {
+    // Validate scale factor is numeric (decimal allowed)
+    if scale_factor.parse::<f64>().is_err() {
+        return Err(format!(
+            "Scale factor '{}' is not a valid number",
+            scale_factor
+        ));
+    }
+
+    Ok(())
+}
+
+fn randomized_fuzz_config(seed: Option<String>) -> Result<(FuzzConfig, String)> {
+    let (mut rng, seed_b64) = rng(seed)?;
+
+    let config = FuzzConfig {
+        num_workers: rng.gen_range(2..=8),
+        files_per_task: rng.gen_range(1..=8),
+        cardinality_task_count_factor: rng.gen_range(1.0..=3.0),
+    };
+
+    Ok((config, seed_b64))
+}
+
+/// Run a workload using the generic workload trait
+async fn run_workload<W>(workload: W, seed: Option<String>) -> Result<()>
+where
+    W: Workload,
+{
+    info!("Starting workload: {}", workload.name());
+
+    let (config, seed_b64) = randomized_fuzz_config(seed)?;
+    info!("Fuzz config: {}", config);
+    info!("Seed: {}", seed_b64);
+
+    let fuzz_db = match FuzzDB::new(config, W::setup).await {
+        Ok(db) => db,
+        Err(e) => {
+            return Err(datafusion::error::DataFusionError::Execution(format!(
+                "Failed to create FuzzDB: {}",
+                e
+            )));
+        }
+    };
+
+    // Get queries from workload
+    let queries_iter = workload.queries().await?;
+
+    // Run each query
+    let mut successful = 0;
+    let mut failed = 0;
+
+    for (query_name, query_sql) in queries_iter {
+        debug!("Executing fuzz query: {}", query_sql.trim());
+
+        match fuzz_db.run(&query_sql).await {
+            Ok(results) => {
+                successful += 1;
+                info!("✅ {} completed successfully", query_name);
+                if log_enabled!(Level::Debug) {
+                    match results {
+                        Some(batches) => {
+                            debug!("{}", pretty_format_batches(&batches)?,);
+                        }
+                        None => {
+                            debug!("No results (query errored expectedly)");
+                        }
+                    }
+                }
+            }
+            Err(e) => {
+                failed += 1;
+                error!("❌ {} failed: {}", query_name, e);
+            }
+        }
+    }
+
+    if failed > 0 {
+        return Err(datafusion::error::DataFusionError::Execution(format!(
+            "{} out of {} queries failed",
+            failed,
+            successful + failed
+        )));
+    }
+
+    Ok(())
+}