test_utils/tpcds

Author: jayshrivastava

src/test_utils/mod.rs

@@ -6,4 +6,5 @@ pub mod mock_exec;
 pub mod parquet;
 pub mod plans;
 pub mod session_context;
+pub mod tpcds;
 pub mod tpch;

src/test_utils/tpcds.rs

@@ -0,0 +1,227 @@
+
+use std::path::Path;
+use datafusion::{
+    execution::context::SessionContext,
+    prelude::ParquetReadOptions,
+    error::{DataFusionError, Result},
+};
+
+pub fn get_data_dir() -> std::path::PathBuf {
+    std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("testdata/tpcds/data")
+}
+
+pub fn get_queries_dir() -> std::path::PathBuf {
+    std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("testdata/tpcds/queries")
+}
+
+/// All TPCDS table names
+pub const TPCDS_TABLES: &[&str] = &[
+    "call_center",
+    "catalog_page", 
+    "catalog_returns",
+    "catalog_sales",
+    "customer",
+    "customer_address",
+    "customer_demographics", 
+    "date_dim",
+    "household_demographics",
+    "income_band",
+    "inventory",
+    "item",
+    "promotion",
+    "reason",
+    "ship_mode",
+    "store",
+    "store_returns",
+    "store_sales",
+    "time_dim",
+    "warehouse",
+    "web_page",
+    "web_returns", 
+    "web_sales",
+    "web_site"
+];
+
+/// Register a single TPCDS table from parquet files in the data directory
+/// 
+/// # Arguments
+/// * `ctx` - The DataFusion SessionContext to register the table with
+/// * `table_name` - The name of the table to register
+/// * `data_dir` - Optional path to the TPCDS data directory. If None, uses get_data_dir()
+/// 
+/// # Returns
+/// * `Result<()>` - Ok if successful, error if table directory is missing or registration fails
+pub async fn register_tpcds_table(
+    ctx: &SessionContext,
+    table_name: &str,
+    data_dir: Option<&Path>,
+) -> Result<()> {
+    let default_data_dir = get_data_dir();
+    let data_path = data_dir.unwrap_or(&default_data_dir);
+    let table_path = data_path.join(table_name);
+    
+    if !table_path.exists() {
+        return Err(DataFusionError::Execution(format!(
+            "TPCDS table directory not found: {}",
+            table_path.display()
+        )));
+    }
+    
+    // Check if this is a partitioned table or single files
+    let entries = std::fs::read_dir(&table_path)
+        .map_err(|e| DataFusionError::Execution(format!(
+            "Failed to read table directory {}: {}", 
+            table_path.display(), e
+        )))?;
+        
+    let mut has_partitions = false;
+    let mut has_direct_files = false;
+    
+    for entry in entries {
+        let entry = entry.map_err(|e| DataFusionError::Execution(format!(
+            "Failed to read directory entry: {}", e
+        )))?;
+        let path = entry.path();
+        if path.is_dir() {
+            has_partitions = true;
+        } else if path.extension().and_then(|s| s.to_str()) == Some("parquet") {
+            has_direct_files = true;
+        }
+    }
+    
+    let table_path_str = table_path.to_string_lossy();
+    
+    if has_partitions && !has_direct_files {
+        // Register as partitioned table
+        ctx.register_parquet(table_name, &table_path_str, ParquetReadOptions::default())
+            .await
+    } else if has_direct_files && !has_partitions {
+        // Register single parquet files
+        ctx.register_parquet(table_name, &table_path_str, ParquetReadOptions::default())
+            .await  
+    } else if has_partitions && has_direct_files {
+        return Err(DataFusionError::Execution(format!(
+            "Mixed directory structure not supported for table: {}", table_name
+        )));
+    } else {
+        return Err(DataFusionError::Execution(format!(
+            "No parquet files found in table directory: {}", table_path.display()
+        )));
+    }
+}
+
+/// Register all available TPCDS tables from the data directory
+/// 
+/// # Arguments
+/// * `ctx` - The DataFusion SessionContext to register tables with
+/// * `data_dir` - Optional path to the TPCDS data directory. If None, uses get_data_dir()
+/// 
+/// # Returns
+/// * `Result<Vec<String>>` - List of successfully registered table names, or error if any required table is missing
+pub async fn register_all_tpcds_tables(
+    ctx: &SessionContext,
+    data_dir: Option<&Path>,
+) -> Result<Vec<String>> {
+    let mut registered_tables = Vec::new();
+    
+    for &table_name in TPCDS_TABLES {
+        match register_tpcds_table(ctx, table_name, data_dir).await {
+            Ok(_) => {
+                registered_tables.push(table_name.to_string());
+            }
+            Err(e) => {
+                return Err(DataFusionError::Execution(format!(
+                    "Failed to register required TPCDS table '{}': {}", 
+                    table_name, e
+                )));
+            }
+        }
+    }
+    
+    Ok(registered_tables)
+}
+
+/// Register only available TPCDS tables from the data directory, skipping missing ones
+/// 
+/// # Arguments
+/// * `ctx` - The DataFusion SessionContext to register tables with  
+/// * `data_dir` - Optional path to the TPCDS data directory. If None, uses get_data_dir()
+/// 
+/// # Returns
+/// * `Result<(Vec<String>, Vec<String>)>` - Tuple of (successfully registered tables, missing tables)
+pub async fn register_available_tpcds_tables(
+    ctx: &SessionContext,
+    data_dir: Option<&Path>, 
+) -> Result<(Vec<String>, Vec<String>)> {
+    let mut registered_tables = Vec::new();
+    let mut missing_tables = Vec::new();
+    
+    for &table_name in TPCDS_TABLES {
+        match register_tpcds_table(ctx, table_name, data_dir).await {
+            Ok(_) => {
+                registered_tables.push(table_name.to_string());
+            }
+            Err(_) => {
+                missing_tables.push(table_name.to_string());
+            }
+        }
+    }
+    
+    Ok((registered_tables, missing_tables))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use datafusion::execution::context::SessionContext;
+
+    #[tokio::test]
+    async fn test_register_available_tpcds_tables() {
+        let ctx = SessionContext::new();
+        
+        // Test with the actual data directory
+        let (registered, missing) = register_available_tpcds_tables(&ctx, None)
+            .await
+            .expect("Should successfully register available tables");
+        
+        // We should have at least some tables registered
+        assert!(!registered.is_empty(), "Expected at least some TPCDS tables to be available");
+        
+        // Test that we can query a registered table
+        if registered.contains(&"date_dim".to_string()) {
+            let result = ctx.sql("SELECT COUNT(*) as count FROM date_dim").await;
+            assert!(result.is_ok(), "Should be able to query registered table date_dim");
+            
+            let df = result.unwrap();
+            let batches = df.collect().await;
+            assert!(batches.is_ok(), "Should be able to collect results from date_dim");
+            assert!(!batches.unwrap().is_empty(), "date_dim should have some data");
+        }
+        
+        println!("Registered tables: {:?}", registered);
+        println!("Missing tables: {:?}", missing);
+    }
+    
+    #[tokio::test] 
+    async fn test_register_single_tpcds_table() {
+        let ctx = SessionContext::new();
+        
+        // Test registering a single table that should exist
+        match register_tpcds_table(&ctx, "date_dim", None).await {
+            Ok(_) => {
+                // Table registered successfully, try to query it
+                let result = ctx.sql("SELECT COUNT(*) FROM date_dim").await;
+                assert!(result.is_ok(), "Should be able to query registered date_dim table");
+            }
+            Err(e) => {
+                // If table doesn't exist, that's also acceptable for this test
+                println!("date_dim table not available: {}", e);
+            }
+        }
+        
+        // Test with a non-existent table - should return error
+        let result = register_tpcds_table(&ctx, "nonexistent_table", None).await;
+        assert!(result.is_err(), "Should return error for non-existent table");
+    }
+}
+