diff --git a/src/lib.rs b/src/lib.rs index 3b6cd2fc..805169f3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -40,6 +40,7 @@ pub mod result; pub mod stage; pub mod stage_reader; pub mod test_utils; +pub mod tpch_tests; pub mod util; pub mod vocab; diff --git a/src/tpch_tests.rs b/src/tpch_tests.rs new file mode 100644 index 00000000..de3e3165 --- /dev/null +++ b/src/tpch_tests.rs @@ -0,0 +1,241 @@ +#[cfg(test)] +mod tpch_serialization_tests { + use std::sync::Arc; + use std::path::Path; + use std::fs; + + use datafusion::{ + prelude::*, + execution::config::SessionConfig, + physical_plan::ExecutionPlan, + }; + use datafusion_proto::physical_plan::AsExecutionPlan; + + use crate::codec::DFRayCodec; + + /// Test serialization/deserialization round trip for a physical plan + async fn test_plan_serialization(physical_plan: Arc, query_name: &str) -> Result<(), Box> { + let codec = DFRayCodec {}; + + // Serialize to protobuf + let proto = datafusion_proto::protobuf::PhysicalPlanNode::try_from_physical_plan( + physical_plan.clone(), + &codec + )?; + + println!("āœ… Serialization successful for {}", query_name); + + // Create a new context for deserialization (simulating worker context) + let worker_ctx = create_tpch_context().await.expect("Failed to create worker context"); + + // Deserialize back to physical plan + let result = proto.try_into_physical_plan( + &worker_ctx, + worker_ctx.runtime_env().as_ref(), + &codec + ); + + match result { + Ok(deserialized_plan) => { + println!("āœ… Deserialization successful for {}", query_name); + + // Check if schemas match + let orig_schema = physical_plan.schema(); + let deser_schema = deserialized_plan.schema(); + + if orig_schema != deser_schema { + return Err(format!("Schema mismatch for {}: orig={:?}, deser={:?}", + query_name, orig_schema, deser_schema).into()); + } + + println!("āœ… Schema validation passed for {}", query_name); + Ok(()) + } + Err(e) => { + println!("āŒ Deserialization failed for {}: {}", query_name, e); + Err(e.into()) + } + } + } + + /// Create a SessionContext with all TPC-H tables registered + async fn create_tpch_context() -> Result> { + let config = SessionConfig::default(); + let ctx = SessionContext::new_with_config(config); + + // Check if TPC-H data is available + if !Path::new("/tmp/tpch_s1/").exists() { + return Err("TPC-H data not found at /tmp/tpch_s1/. Please ensure the data is available.".into()); + } + + // Define all TPC-H tables + let tables = vec![ + "customer", "lineitem", "nation", "orders", + "part", "partsupp", "region", "supplier" + ]; + + for table in tables { + let table_sql = format!( + "CREATE EXTERNAL TABLE {} STORED AS PARQUET LOCATION '/tmp/tpch_s1/{}.parquet'", + table, table + ); + + ctx.sql(&table_sql).await + .map_err(|e| format!("Failed to create table {}: {}", table, e))?; + } + + Ok(ctx) + } + + /// Read and parse a TPC-H query file + fn read_query_file(query_name: &str) -> Result> { + let path = format!("tpch/queries/{}.sql", query_name); + let content = fs::read_to_string(&path) + .map_err(|e| format!("Failed to read {}: {}", path, e))?; + Ok(content) + } + + /// Test a single TPC-H query for serialization bugs + async fn test_single_query(ctx: &SessionContext, query_name: &str) -> Result<(), Box> { + println!("\nšŸ” Testing query: {}", query_name); + + let sql = read_query_file(query_name)?; + + // Create logical plan + let logical_plan = ctx.state().create_logical_plan(&sql).await + .map_err(|e| format!("Failed to create logical plan for {}: {}", query_name, e))?; + + let logical_plan = ctx.state().optimize(&logical_plan) + .map_err(|e| format!("Failed to optimize logical plan for {}: {}", query_name, e))?; + + // Create physical plan + let physical_plan = ctx.state().create_physical_plan(&logical_plan).await + .map_err(|e| format!("Failed to create physical plan for {}: {}", query_name, e))?; + + // Test serialization/deserialization + test_plan_serialization(physical_plan, query_name).await?; + + Ok(()) + } + + #[tokio::test] + async fn test_all_tpch_queries_serialization() { + // Skip test if TPC-H data not available + if !Path::new("/tmp/tpch_s1/").exists() { + println!("āš ļø Skipping TPC-H tests - data not found at /tmp/tpch_s1/"); + return; + } + + let ctx = create_tpch_context().await.expect("Failed to create TPC-H context"); + + // List of all TPC-H queries + let queries = vec![ + "q1", "q2", "q3", "q4", "q5", "q6", "q7", "q8", "q9", "q10", + "q11", "q12", "q13", "q14", "q15", "q16", "q17", "q18", "q19", + "q20", "q21", "q22" + ]; + + let mut passed = 0; + let mut failed = 0; + let mut failed_queries = Vec::new(); + + for query in &queries { + match test_single_query(&ctx, query).await { + Ok(_) => { + println!("āœ… {} passed", query); + passed += 1; + } + Err(e) => { + println!("āŒ {} failed: {}", query, e); + failed += 1; + failed_queries.push((query, e.to_string())); + } + } + } + + println!("\nšŸ“Š Results Summary:"); + println!(" āœ… Passed: {}/{}", passed, queries.len()); + println!(" āŒ Failed: {}/{}", failed, queries.len()); + + if !failed_queries.is_empty() { + println!("\nāŒ Failed queries:"); + for (query, error) in &failed_queries { + println!(" {} - {}", query, error); + } + } + + // For now, we expect some queries to fail due to the bug + // Once we fix the serialization issues, we can change this to assert all pass + if failed > 0 { + println!("\nāš ļø Some queries failed serialization - this is expected until we fix the bugs"); + + // Let's specifically check if q16 fails with the expected error + let q16_failed = failed_queries.iter().any(|(query, error)| { + **query == "q16" && error.contains("data type inlist should be same") + }); + + if q16_failed { + println!("āœ… Confirmed: Q16 fails with the expected 'data type inlist should be same' error"); + } + } + } + + #[tokio::test] + async fn test_q16_specifically() { + // This test specifically focuses on the q16 bug we know about + if !Path::new("/tmp/tpch_s1/").exists() { + println!("āš ļø Skipping Q16 test - TPC-H data not found at /tmp/tpch_s1/"); + return; + } + + let ctx = create_tpch_context().await.expect("Failed to create TPC-H context"); + + match test_single_query(&ctx, "q16").await { + Ok(_) => { + println!("āœ… Q16 serialization works! The bug may have been fixed."); + } + Err(e) => { + if e.to_string().contains("data type inlist should be same") { + println!("āœ… Successfully reproduced Q16 serialization bug!"); + println!("Error: {}", e); + // This is expected for now + } else { + panic!("Q16 failed with unexpected error: {}", e); + } + } + } + } + + #[tokio::test] + async fn test_minimal_in_expression() { + // Test a minimal IN expression without TPC-H dependency + let ctx = SessionContext::new(); + + let sql = r#" + WITH sample_data AS ( + SELECT column1 as id, column2 as size_col FROM VALUES + (1, 14), (2, 6), (3, 5), (4, 31), (5, 49) + ) + SELECT * FROM sample_data WHERE size_col IN (14, 6, 5, 31, 49, 15, 41, 47) + "#; + + println!("šŸ” Testing minimal IN expression..."); + + let logical_plan = ctx.state().create_logical_plan(sql).await.expect("logical plan"); + let physical_plan = ctx.state().create_physical_plan(&logical_plan).await.expect("physical plan"); + + match test_plan_serialization(physical_plan, "minimal_in_expression").await { + Ok(_) => { + println!("āœ… Minimal IN expression serialization works"); + } + Err(e) => { + if e.to_string().contains("data type inlist should be same") { + println!("āœ… Reproduced the bug with minimal IN expression!"); + println!("Error: {}", e); + } else { + println!("āŒ Unexpected error: {}", e); + } + } + } + } +} \ No newline at end of file