Using tpch script from datafusion-benchmarks

Author: edmondop

.github/workflows/rust.yml

@@ -2,24 +2,36 @@ name: Rust
 
 on:
   push:
-    branches: [ "main" ]
   pull_request:
-    branches: [ "main" ]
 
 env:
   CARGO_TERM_COLOR: always
+  PYTHON_VERSION: 3.9
+  TPCH_SAMPLING_RATE: "1" # (1/100)
+  TPCH_TEST_PARTITIONS: "2"
+  TPCH_DATA_PATH: "data"
 
 jobs:
   build:
-
     runs-on: ubuntu-latest
 
     steps:
-    - uses: actions/checkout@v3
-    - name: Install protobuf compiler
-      shell: bash
-      run: sudo apt-get install protobuf-compiler
-    - name: Build Rust code
-      run: cargo build --verbose
-    - name: Run tests
-      run: cargo test --verbose
+      - uses: actions/checkout@v3
+      - name: Install protobuf compiler
+        shell: bash
+        run: sudo apt-get install protobuf-compiler
+      - name: Build Rust code
+        run: cargo build --verbose
+      - name: Set up Python
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ env.PYTHON_VERSION }}
+      - name: Install test dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -r tpch/requirements.txt
+      - name: Generate test data
+        run: |
+          ./scripts/gen-test-data.sh
+      - name: Run tests
+        run: cargo test --verbose

.gitignore

@@ -5,3 +5,4 @@ venv
 *.so
 *.log
 results-sf*
+data

Cargo.toml

@@ -45,6 +45,10 @@ uuid = "1.2"
 rustc_version = "0.4.0"
 tonic-build = { version = "0.8", default-features = false, features = ["transport", "prost"] }
 
+[dev-dependencies]
+anyhow = "1.0.89"
+pretty_assertions = "1.4.0"
+
 [lib]
 name = "datafusion_ray"
 crate-type = ["cdylib", "rlib"]
@@ -54,4 +58,4 @@ name = "datafusion_ray._datafusion_ray_internal"
 
 [profile.release]
 codegen-units = 1
-lto = true
+lto = true

scripts/gen-test-data.sh

@@ -0,0 +1,6 @@
+#!/bin/bash
+
+set -e
+mkdir -p data
+python -m tpch.tpchgen generate --scale-factor "$TPCH_SAMPLING_RATE" --partitions "$TPCH_TEST_PARTITIONS"
+python -m tpch.tpchgen convert --partitions "$TPCH_TEST_PARTITIONS:lua require('telescope').extensions.live_grep_args.live_grep_args()"

src/planner.rs

@@ -264,125 +264,124 @@ mod test {
     use super::*;
     use datafusion::physical_plan::displayable;
     use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext};
-    use std::fs;
+    use pretty_assertions::assert_eq;
     use std::path::Path;
+    use std::{env, fs};
+    type TestResult<T> = std::result::Result<T, anyhow::Error>;
 
     #[tokio::test]
-    async fn test_q1() -> Result<()> {
+    async fn test_q1() -> TestResult<()> {
         do_test(1).await
     }
 
     #[tokio::test]
-    async fn test_q2() -> Result<()> {
+    async fn test_q2() -> TestResult<()> {
         do_test(2).await
     }
 
     #[tokio::test]
-    async fn test_q3() -> Result<()> {
+    async fn test_q3() -> TestResult<()> {
         do_test(3).await
     }
 
     #[tokio::test]
-    async fn test_q4() -> Result<()> {
+    async fn test_q4() -> TestResult<()> {
         do_test(4).await
     }
 
     #[tokio::test]
-    async fn test_q5() -> Result<()> {
+    async fn test_q5() -> TestResult<()> {
         do_test(5).await
     }
 
     #[tokio::test]
-    async fn test_q6() -> Result<()> {
+    async fn test_q6() -> TestResult<()> {
         do_test(6).await
     }
 
     #[tokio::test]
-    async fn test_q7() -> Result<()> {
+    async fn test_q7() -> TestResult<()> {
         do_test(7).await
     }
 
     #[tokio::test]
-    async fn test_q8() -> Result<()> {
+    async fn test_q8() -> TestResult<()> {
         do_test(8).await
     }
 
     #[tokio::test]
-    async fn test_q9() -> Result<()> {
+    async fn test_q9() -> TestResult<()> {
         do_test(9).await
     }
 
     #[tokio::test]
-    async fn test_q10() -> Result<()> {
+    async fn test_q10() -> TestResult<()> {
         do_test(10).await
     }
 
     #[tokio::test]
-    async fn test_q11() -> Result<()> {
+    async fn test_q11() -> TestResult<()> {
         do_test(11).await
     }
 
     #[tokio::test]
-    async fn test_q12() -> Result<()> {
+    async fn test_q12() -> TestResult<()> {
         do_test(12).await
     }
 
     #[tokio::test]
-    async fn test_q13() -> Result<()> {
+    async fn test_q13() -> TestResult<()> {
         do_test(13).await
     }
 
     #[tokio::test]
-    async fn test_q14() -> Result<()> {
+    async fn test_q14() -> TestResult<()> {
         do_test(14).await
     }
 
     #[ignore]
     #[tokio::test]
-    async fn test_q15() -> Result<()> {
+    async fn test_q15() -> TestResult<()> {
         do_test(15).await
     }
 
     #[tokio::test]
-    async fn test_q16() -> Result<()> {
+    async fn test_q16() -> TestResult<()> {
         do_test(16).await
     }
 
     #[tokio::test]
-    async fn test_q17() -> Result<()> {
+    async fn test_q17() -> TestResult<()> {
         do_test(17).await
     }
 
     #[tokio::test]
-    async fn test_q18() -> Result<()> {
+    async fn test_q18() -> TestResult<()> {
         do_test(18).await
     }
 
     #[tokio::test]
-    async fn test_q19() -> Result<()> {
+    async fn test_q19() -> TestResult<()> {
         do_test(19).await
     }
 
     #[tokio::test]
-    async fn test_q20() -> Result<()> {
+    async fn test_q20() -> TestResult<()> {
         do_test(20).await
     }
 
     #[tokio::test]
-    async fn test_q21() -> Result<()> {
+    async fn test_q21() -> TestResult<()> {
         do_test(21).await
     }
 
     #[tokio::test]
-    async fn test_q22() -> Result<()> {
+    async fn test_q22() -> TestResult<()> {
         do_test(22).await
     }
 
-    async fn do_test(n: u8) -> Result<()> {
-        let data_path = "/mnt/bigdata/tpch/sf10-parquet";
-        if !Path::new(&data_path).exists() {
-            return Ok(());
-        }
+    async fn do_test(n: u8) -> TestResult<()> {
+        let data_path = env::var("TPCH_DATA_PATH")?;
         let file = format!("testdata/queries/q{n}.sql");
         let sql = fs::read_to_string(&file)?;
         let config = SessionConfig::new().with_target_partitions(4);