utils

Author: jayshrivastava

src/test_utils/tpcds.rs

@@ -1,5 +1,6 @@
 
 use std::path::Path;
+use std::process::Command;
 use datafusion::{
     execution::context::SessionContext,
     prelude::ParquetReadOptions,
@@ -170,6 +171,94 @@ pub async fn register_available_tpcds_tables(
     Ok((registered_tables, missing_tables))
 }
 
+/// Generate TPCDS data by calling the generate.sh script
+/// 
+/// # Arguments
+/// * `scale_factor` - Optional scale factor for data generation (e.g., "1", "10", "100")
+/// 
+/// # Returns
+/// * `Result<()>` - Ok if generation successful, error otherwise
+pub fn generate_tpcds_data(scale_factor: Option<&str>) -> Result<()> {
+    let tpcds_dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("testdata/tpcds");
+    let generate_script = tpcds_dir.join("generate.sh");
+    
+    if !generate_script.exists() {
+        return Err(DataFusionError::Execution(format!(
+            "TPCDS generate script not found: {}",
+            generate_script.display()
+        )));
+    }
+
+    let mut cmd = Command::new("bash");
+    cmd.arg(&generate_script);
+    cmd.current_dir(&tpcds_dir);
+    
+    // Add scale factor if provided
+    if let Some(scale) = scale_factor {
+        cmd.arg(scale);
+    }
+
+    let output = cmd.output().map_err(|e| {
+        DataFusionError::Execution(format!(
+            "Failed to execute TPCDS generate script: {}",
+            e
+        ))
+    })?;
+
+    if !output.status.success() {
+        let stderr = String::from_utf8_lossy(&output.stderr);
+        let stdout = String::from_utf8_lossy(&output.stdout);
+        return Err(DataFusionError::Execution(format!(
+            "TPCDS generate script failed with status: {}\nstdout: {}\nstderr: {}",
+            output.status, stdout, stderr
+        )));
+    }
+
+    println!("TPCDS data generation completed successfully");
+    if !output.stdout.is_empty() {
+        println!("Output: {}", String::from_utf8_lossy(&output.stdout));
+    }
+
+    Ok(())
+}
+
+/// Ensure TPCDS data is available, generating it if necessary
+/// 
+/// # Arguments
+/// * `scale_factor` - Optional scale factor for data generation
+/// * `force_regenerate` - If true, regenerate data even if it already exists
+/// 
+/// # Returns
+/// * `Result<()>` - Ok if data is available, error otherwise
+pub fn ensure_tpcds_data(scale_factor: Option<&str>, force_regenerate: bool) -> Result<()> {
+    let data_dir = get_data_dir();
+    
+    // Check if data already exists (basic check for a few key tables)
+    let key_tables = ["date_dim", "item", "customer"];
+    let mut tables_exist = true;
+    
+    if !force_regenerate {
+        for table in &key_tables {
+            let table_path = data_dir.join(table);
+            if !table_path.exists() {
+                tables_exist = false;
+                break;
+            }
+        }
+    } else {
+        tables_exist = false;
+    }
+    
+    if !tables_exist {
+        println!("Generating TPCDS data (scale factor: {:?})", scale_factor);
+        generate_tpcds_data(scale_factor)?;
+    } else {
+        println!("TPCDS data already exists, skipping generation");
+    }
+    
+    Ok(())
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;