Add test comparing distributed + single node execution on TPCH data

* add TPCH queries from DataFusion repo (see tests/fixtures/tpch/queries)

* add lazy TPCH parquet data generator using github.com/clflushopt/tpchgen-rs

* write consistency test that runs queries twice (distributed and single-node)
  and asserts match

Author: jayshrivastava

.gitignore

@@ -1,2 +1,3 @@
 /.idea
-/target
+/target
+tests/fixtures/tpch/data/

Cargo.lock

@@ -25,3 +25,7 @@ object_store = "0.12.3"
 
 [dev-dependencies]
 insta = { version = "1.43.1", features = ["filters"] }
+tpchgen = { git = "https://github.com/clflushopt/tpchgen-rs" }
+tpchgen-arrow = { git = "https://github.com/clflushopt/tpchgen-rs" }
+parquet = "55.2.0"
+arrow = "55.2.0"

Cargo.toml

@@ -5,15 +5,36 @@ use datafusion::{
     catalog::{MemTable, TableProvider},
 };
 
+use std::fs;
+
+use tokio::sync::OnceCell;
+
+use arrow::record_batch::RecordBatch;
+use parquet::{arrow::arrow_writer::ArrowWriter, file::properties::WriterProperties};
+use tpchgen::generators::{
+    CustomerGenerator, LineItemGenerator, NationGenerator, OrderGenerator, PartGenerator,
+    PartSuppGenerator, RegionGenerator, SupplierGenerator,
+};
+use tpchgen_arrow::{
+    CustomerArrow, LineItemArrow, NationArrow, OrderArrow, PartArrow, PartSuppArrow, RegionArrow,
+    SupplierArrow,
+};
+
+const QUERIES_DIR: &str = "tests/fixtures/tpch/queries";
+pub const DATA_DIR: &str = "tests/fixtures/tpch/data"; // mirroed in .gitignore
+pub const NUM_QUERIES: u8 = 22; // number of queries in the TPCH benchmark numbered from 1 to 22
+
+const SCALE_FACTOR: f64 = 0.001;
+
 pub fn tpch_table(name: &str) -> Arc<dyn TableProvider> {
     let schema = Arc::new(get_tpch_table_schema(name));
     Arc::new(MemTable::try_new(schema, vec![]).unwrap())
 }
 
 pub fn tpch_query(num: u8) -> String {
-    // read the query from the test/tpch/queries/ directory and return it
-    let query_path = format!("testing/tpch/queries/q{}.sql", num);
-    std::fs::read_to_string(query_path)
+    // read the query from the queries directory in the fixtures dir and return it
+    let query_path = format!("{}/q{}.sql", QUERIES_DIR, num);
+    fs::read_to_string(query_path)
         .unwrap_or_else(|_| panic!("Failed to read TPCH query file: q{}.sql", num))
         .trim()
         .to_string()
@@ -113,3 +134,61 @@ pub fn get_tpch_table_schema(table: &str) -> Schema {
         _ => unimplemented!(),
     }
 }
+
+// generate_table creates a parquet file in the DATA_DIR directory from an arrow RecordBatch row
+// source.
+fn generate_table<A>(mut data_source: A, table_name: &str) -> Result<(), Box<dyn std::error::Error>>
+where
+    A: Iterator<Item = RecordBatch>,
+{
+    let output_path = format!("{}/{}.parquet", DATA_DIR, table_name);
+
+    if let Some(first_batch) = data_source.next() {
+        let file = fs::File::create(&output_path)?;
+        let props = WriterProperties::builder().build();
+        let mut writer = ArrowWriter::try_new(file, first_batch.schema(), Some(props))?;
+
+        writer.write(&first_batch)?;
+
+        while let Some(batch) = data_source.next() {
+            writer.write(&batch)?;
+        }
+
+        writer.close()?;
+    }
+
+    println!("Generated {} table: {}", table_name, output_path);
+    Ok(())
+}
+
+macro_rules! must_generate_tpch_table {
+    ($generator:ident, $arrow:ident, $name:literal) => {
+        generate_table(
+            // TODO: Consider adjusting the partitions and batch sizes.
+            $arrow::new($generator::new(SCALE_FACTOR, 1, 1)).with_batch_size(1000),
+            $name,
+        )
+        .expect(concat!("Failed to generate ", $name, " table"));
+    };
+}
+
+// INIT_TPCH_TABLES is ensures that TPC-H tables are generated only once.
+static INIT_TPCH_TABLES: OnceCell<()> = OnceCell::const_new();
+
+// generate_tpch_data generates all TPC-H tables in the DATA_DIR directory.
+pub async fn generate_tpch_data() {
+    INIT_TPCH_TABLES
+        .get_or_init(|| async {
+            fs::create_dir_all(DATA_DIR).expect("Failed to create data directory");
+
+            must_generate_tpch_table!(RegionGenerator, RegionArrow, "region");
+            must_generate_tpch_table!(NationGenerator, NationArrow, "nation");
+            must_generate_tpch_table!(CustomerGenerator, CustomerArrow, "customer");
+            must_generate_tpch_table!(SupplierGenerator, SupplierArrow, "supplier");
+            must_generate_tpch_table!(PartGenerator, PartArrow, "part");
+            must_generate_tpch_table!(PartSuppGenerator, PartSuppArrow, "partsupp");
+            must_generate_tpch_table!(OrderGenerator, OrderArrow, "orders");
+            must_generate_tpch_table!(LineItemGenerator, LineItemArrow, "lineitem");
+        })
+        .await;
+}

tests/common/tpch.rs

@@ -0,0 +1,21 @@
+select
+    l_returnflag,
+    l_linestatus,
+    sum(l_quantity) as sum_qty,
+    sum(l_extendedprice) as sum_base_price,
+    sum(l_extendedprice * (1 - l_discount)) as sum_disc_price,
+    sum(l_extendedprice * (1 - l_discount) * (1 + l_tax)) as sum_charge,
+    avg(l_quantity) as avg_qty,
+    avg(l_extendedprice) as avg_price,
+    avg(l_discount) as avg_disc,
+    count(*) as count_order
+from
+    lineitem
+where
+        l_shipdate <= date '1998-09-02'
+group by
+    l_returnflag,
+    l_linestatus
+order by
+    l_returnflag,
+    l_linestatus;

tests/fixtures/tpch/queries/q1.sql

@@ -0,0 +1,31 @@
+select
+    c_custkey,
+    c_name,
+    sum(l_extendedprice * (1 - l_discount)) as revenue,
+    c_acctbal,
+    n_name,
+    c_address,
+    c_phone,
+    c_comment
+from
+    customer,
+    orders,
+    lineitem,
+    nation
+where
+        c_custkey = o_custkey
+  and l_orderkey = o_orderkey
+  and o_orderdate >= date '1993-10-01'
+  and o_orderdate < date '1994-01-01'
+  and l_returnflag = 'R'
+  and c_nationkey = n_nationkey
+group by
+    c_custkey,
+    c_name,
+    c_acctbal,
+    c_phone,
+    n_name,
+    c_address,
+    c_comment
+order by
+    revenue desc;

tests/fixtures/tpch/queries/q10.sql

@@ -0,0 +1,27 @@
+select
+    ps_partkey,
+    sum(ps_supplycost * ps_availqty) as value
+from
+    partsupp,
+    supplier,
+    nation
+where
+    ps_suppkey = s_suppkey
+  and s_nationkey = n_nationkey
+  and n_name = 'GERMANY'
+group by
+    ps_partkey having
+    sum(ps_supplycost * ps_availqty) > (
+    select
+    sum(ps_supplycost * ps_availqty) * 0.0001
+    from
+    partsupp,
+    supplier,
+    nation
+    where
+    ps_suppkey = s_suppkey
+                  and s_nationkey = n_nationkey
+                  and n_name = 'GERMANY'
+    )
+order by
+    value desc;

tests/fixtures/tpch/queries/q11.sql

@@ -0,0 +1,30 @@
+select
+    l_shipmode,
+    sum(case
+            when o_orderpriority = '1-URGENT'
+                or o_orderpriority = '2-HIGH'
+                then 1
+            else 0
+        end) as high_line_count,
+    sum(case
+            when o_orderpriority <> '1-URGENT'
+                and o_orderpriority <> '2-HIGH'
+                then 1
+            else 0
+        end) as low_line_count
+from
+    lineitem
+        join
+    orders
+    on
+            l_orderkey = o_orderkey
+where
+        l_shipmode in ('MAIL', 'SHIP')
+  and l_commitdate < l_receiptdate
+  and l_shipdate < l_commitdate
+  and l_receiptdate >= date '1994-01-01'
+  and l_receiptdate < date '1995-01-01'
+group by
+    l_shipmode
+order by
+    l_shipmode;

tests/fixtures/tpch/queries/q12.sql

@@ -0,0 +1,20 @@
+select
+    c_count,
+    count(*) as custdist
+from
+    (
+        select
+            c_custkey,
+            count(o_orderkey)
+        from
+            customer left outer join orders on
+                        c_custkey = o_custkey
+                    and o_comment not like '%special%requests%'
+        group by
+            c_custkey
+    ) as c_orders (c_custkey, c_count)
+group by
+    c_count
+order by
+    custdist desc,
+    c_count desc;

tests/fixtures/tpch/queries/q13.sql

@@ -0,0 +1,13 @@
+select
+            100.00 * sum(case
+                             when p_type like 'PROMO%'
+                                 then l_extendedprice * (1 - l_discount)
+                             else 0
+            end) / sum(l_extendedprice * (1 - l_discount)) as promo_revenue
+from
+    lineitem,
+    part
+where
+        l_partkey = p_partkey
+  and l_shipdate >= date '1995-09-01'
+  and l_shipdate < date '1995-10-01';