Split tpch tests

Author: gabotechs

.github/workflows/ci.yml

@@ -38,7 +38,7 @@ jobs:
     steps:
       - uses: actions/checkout@v4
       - uses: ./.github/actions/setup
-      - run: cargo test --features tpch --test tpch_validation_test
+      - run: cargo test --features tpch --test 'tpch_*'
 
   format-check:
     runs-on: ubuntu-latest

.gitignore

@@ -1,4 +1,5 @@
 /.idea
 /target
 /benchmarks/data/
-testdata/tpch/data/
+testdata/tpch/*
+!testdata/tpch/queries

tests/tpch_correctness_test.rs

@@ -0,0 +1,219 @@
+#[cfg(all(feature = "integration", feature = "tpch", test))]
+mod tests {
+    use datafusion::physical_plan::execute_stream;
+    use datafusion::prelude::SessionContext;
+    use datafusion_distributed::DefaultSessionBuilder;
+    use datafusion_distributed::test_utils::localhost::start_localhost_context;
+    use datafusion_distributed::test_utils::tpch;
+    use futures::TryStreamExt;
+    use std::error::Error;
+    use std::fmt::Display;
+    use std::fs;
+    use std::path::Path;
+    use tokio::sync::OnceCell;
+
+    const PARTITIONS: usize = 6;
+    const TPCH_SCALE_FACTOR: f64 = 1.0;
+    const TPCH_DATA_PARTS: i32 = 16;
+
+    #[tokio::test]
+    async fn test_tpch_1() -> Result<(), Box<dyn Error>> {
+        test_tpch_query(1).await
+    }
+
+    #[tokio::test]
+    async fn test_tpch_2() -> Result<(), Box<dyn Error>> {
+        test_tpch_query(2).await
+    }
+
+    #[tokio::test]
+    async fn test_tpch_3() -> Result<(), Box<dyn Error>> {
+        test_tpch_query(3).await
+    }
+
+    #[tokio::test]
+    async fn test_tpch_4() -> Result<(), Box<dyn Error>> {
+        test_tpch_query(4).await
+    }
+
+    #[tokio::test]
+    async fn test_tpch_5() -> Result<(), Box<dyn Error>> {
+        test_tpch_query(5).await
+    }
+
+    #[tokio::test]
+    async fn test_tpch_6() -> Result<(), Box<dyn Error>> {
+        test_tpch_query(6).await
+    }
+
+    #[tokio::test]
+    async fn test_tpch_7() -> Result<(), Box<dyn Error>> {
+        test_tpch_query(7).await
+    }
+
+    #[tokio::test]
+    async fn test_tpch_8() -> Result<(), Box<dyn Error>> {
+        test_tpch_query(8).await
+    }
+
+    #[tokio::test]
+    async fn test_tpch_9() -> Result<(), Box<dyn Error>> {
+        test_tpch_query(9).await
+    }
+
+    #[tokio::test]
+    async fn test_tpch_10() -> Result<(), Box<dyn Error>> {
+        test_tpch_query(10).await
+    }
+
+    #[tokio::test]
+    async fn test_tpch_11() -> Result<(), Box<dyn Error>> {
+        test_tpch_query(11).await
+    }
+
+    #[tokio::test]
+    async fn test_tpch_12() -> Result<(), Box<dyn Error>> {
+        test_tpch_query(12).await
+    }
+
+    #[tokio::test]
+    async fn test_tpch_13() -> Result<(), Box<dyn Error>> {
+        test_tpch_query(13).await
+    }
+
+    #[tokio::test]
+    async fn test_tpch_14() -> Result<(), Box<dyn Error>> {
+        test_tpch_query(14).await
+    }
+
+    #[tokio::test]
+    async fn test_tpch_15() -> Result<(), Box<dyn Error>> {
+        test_tpch_query(15).await
+    }
+
+    #[tokio::test]
+    async fn test_tpch_16() -> Result<(), Box<dyn Error>> {
+        test_tpch_query(16).await
+    }
+
+    #[tokio::test]
+    async fn test_tpch_17() -> Result<(), Box<dyn Error>> {
+        test_tpch_query(17).await
+    }
+
+    #[tokio::test]
+    async fn test_tpch_18() -> Result<(), Box<dyn Error>> {
+        test_tpch_query(18).await
+    }
+
+    #[tokio::test]
+    async fn test_tpch_19() -> Result<(), Box<dyn Error>> {
+        test_tpch_query(19).await
+    }
+
+    #[tokio::test]
+    async fn test_tpch_20() -> Result<(), Box<dyn Error>> {
+        test_tpch_query(20).await
+    }
+
+    #[tokio::test]
+    async fn test_tpch_21() -> Result<(), Box<dyn Error>> {
+        test_tpch_query(21).await
+    }
+
+    #[tokio::test]
+    async fn test_tpch_22() -> Result<(), Box<dyn Error>> {
+        test_tpch_query(22).await
+    }
+
+    async fn test_tpch_query(query_id: u8) -> Result<(), Box<dyn Error>> {
+        let (ctx, _guard) = start_localhost_context(4, DefaultSessionBuilder).await;
+        let results_d = run_tpch_query(ctx, query_id).await?;
+        let results_s = run_tpch_query(SessionContext::new(), query_id).await?;
+
+        assert_eq!(
+            results_d.to_string(),
+            results_s.to_string(),
+            "Query {query_id} results differ between executions",
+        );
+        Ok(())
+    }
+
+    // test_non_distributed_consistency runs each TPC-H query twice - once in a distributed manner
+    // and once in a non-distributed manner. For each query, it asserts that the results are identical.
+    async fn run_tpch_query(
+        ctx: SessionContext,
+        query_id: u8,
+    ) -> Result<impl Display, Box<dyn Error>> {
+        let data_dir = ensure_tpch_data(TPCH_SCALE_FACTOR, TPCH_DATA_PARTS).await;
+        let sql = get_test_tpch_query(query_id);
+        ctx.state_ref()
+            .write()
+            .config_mut()
+            .options_mut()
+            .execution
+            .target_partitions = PARTITIONS;
+
+        // Register tables for first context
+        for table_name in [
+            "lineitem", "orders", "part", "partsupp", "customer", "nation", "region", "supplier",
+        ] {
+            let query_path = data_dir.join(table_name);
+            ctx.register_parquet(
+                table_name,
+                query_path.to_string_lossy().as_ref(),
+                datafusion::prelude::ParquetReadOptions::default(),
+            )
+            .await?;
+        }
+
+        // Query 15 has three queries in it, one creating the view, the second
+        // executing, which we want to capture the output of, and the third
+        // tearing down the view
+        let stream = if query_id == 15 {
+            let queries: Vec<&str> = sql
+                .split(';')
+                .map(str::trim)
+                .filter(|s| !s.is_empty())
+                .collect();
+
+            ctx.sql(queries[0]).await?.collect().await?;
+            let df = ctx.sql(queries[1]).await?;
+            let plan = df.create_physical_plan().await?;
+            let stream = execute_stream(plan.clone(), ctx.task_ctx())?;
+            ctx.sql(queries[2]).await?.collect().await?;
+
+            stream
+        } else {
+            let df = ctx.sql(&sql).await?;
+            let plan = df.create_physical_plan().await?;
+            execute_stream(plan.clone(), ctx.task_ctx())?
+        };
+
+        let batches = stream.try_collect::<Vec<_>>().await?;
+
+        Ok(arrow::util::pretty::pretty_format_batches(&batches)?)
+    }
+
+    pub fn get_test_tpch_query(num: u8) -> String {
+        let queries_dir = Path::new(env!("CARGO_MANIFEST_DIR")).join("testdata/tpch/queries");
+        tpch::tpch_query_from_dir(&queries_dir, num)
+    }
+
+    // OnceCell to ensure TPCH tables are generated only once for tests
+    static INIT_TEST_TPCH_TABLES: OnceCell<()> = OnceCell::const_new();
+
+    // ensure_tpch_data initializes the TPCH data on disk.
+    pub async fn ensure_tpch_data(sf: f64, parts: i32) -> std::path::PathBuf {
+        let data_dir =
+            Path::new(env!("CARGO_MANIFEST_DIR")).join(format!("testdata/tpch/correctness_sf{sf}"));
+        INIT_TEST_TPCH_TABLES
+            .get_or_init(|| async {
+                if !fs::exists(&data_dir).unwrap() {
+                    tpch::generate_tpch_data(&data_dir, sf, parts);
+                }
+            })
+            .await;
+        data_dir
+    }
+}