Add Datafusion solution [updated]

.gitignore

@@ -30,3 +30,5 @@ run.out
 clickhouse/etc_sudoers.bak
 workdir/
 timeout-exit-codes.out
+*/target
+*.lock

.history

@@ -0,0 +1,18 @@
+#V2
+\\h
+\\h Create table as
+\\h create table
+CREATE EXTERNAL TABLE x STORED AS CSV LOCATION "data/J1_1e7_NA_0_0.csv";
+SELECT * FROM x LIMIT 5;
+\\q
+CREATE EXTERNAL TABLE x STORED AS CSV WITH HEADER ROW LOCATION "data/J1_1e7_NA_0_0.csv"\n;
+select * from x limit 5;
+show x\n;
+show columns from x;
+\\q
+select 2^2;
+select 2 * 2;
+\\q
+select 2^2;
+select power(2,2);
+\\q

datafusion/Cargo.toml

@@ -0,0 +1,22 @@
+[package]
+name = "db-benchmark"
+version = "0.1.0"
+edition = "2018"
+
+[dependencies]
+datafusion = { git = "https://github.com/apache/arrow-datafusion.git", features = ["simd"]}
+tokio = { version = "^1.0", features = ["macros", "rt", "rt-multi-thread"] }
+snmalloc-rs = {version = "0.2", features= ["cache-friendly"]}
+num_cpus = "1.0"
+
+[profile.release]
+lto = true
+codegen-units = 1
+
+[[bin]]
+name = "groupby"
+path = "src/bin/groupby.rs"
+
+[[bin]]
+name = "join"
+path = "src/bin/join.rs"

datafusion/exec.sh

@@ -0,0 +1,4 @@
+#!/bin/bash
+set -e
+
+RUSTFLAGS='-C target-cpu=native' cargo +nightly run --release

datafusion/groupby-datafusion.py

@@ -0,0 +1,187 @@
+#!/usr/bin/env python
+
+print("# groupby-datafusion.py", flush=True)
+
+import os
+import gc
+import timeit
+import datafusion as df
+from datafusion import functions as f
+from datafusion import col
+import pyarrow as pa
+from pyarrow import csv as pacsv
+
+exec(open("./_helpers/helpers.py").read())
+
+def ans_shape(batches):
+    rows, cols = 0, 0
+    for batch in batches:
+        rows += batch.num_rows
+        if cols == 0:
+            cols = batch.num_columns
+        else:
+            assert(cols == batch.num_columns)
+
+    return rows, cols
+
+# ver = df.__version__
+ver = "6.0.0"
+git = ""
+task = "groupby"
+solution = "datafusion"
+fun = ".groupby"
+cache = "TRUE"
+on_disk = "FALSE"
+
+data_name = os.environ["SRC_DATANAME"]
+src_grp = os.path.join("data", data_name + ".csv")
+print("loading dataset %s" % data_name, flush=True)
+
+data = pacsv.read_csv(src_grp)
+
+ctx = df.ExecutionContext()
+ctx.register_record_batches("x", [data.to_batches()])
+
+in_rows = data.num_rows
+
+task_init = timeit.default_timer()
+
+question = "sum v1 by id1" # q1
+gc.collect()
+t_start = timeit.default_timer()
+ans = ctx.sql("SELECT id1, SUM(v1) AS v1 FROM x GROUP BY id1").collect()
+t = timeit.default_timer() - t_start
+print(t)
+shape = ans_shape(ans)
+print(shape)
+t_start = timeit.default_timer()
+df = ctx.create_dataframe([ans])
+chk = df.aggregate([], [f.sum(col("v1"))]).collect()[0].column(0)[0]
+chkt = timeit.default_timer() - t_start
+m = memory_usage()
+write_log(task=task, data=data_name, in_rows=in_rows, question=question, out_rows=shape[0], out_cols=shape[1], solution=solution, version=ver, git=git, fun=fun, run=1, time_sec=t, mem_gb=m, cache=cache, chk=make_chk([chk]), chk_time_sec=chkt, on_disk=on_disk)
+del ans
+gc.collect()
+
+question = "sum v1 by id1:id2" # q2
+gc.collect()
+t_start = timeit.default_timer()
+ans = ctx.sql("SELECT id1, id2, SUM(v1) AS v1 FROM x GROUP BY id1, id2").collect()
+t = timeit.default_timer() - t_start
+print(t)
+shape = ans_shape(ans)
+print(shape)
+t_start = timeit.default_timer()
+df = ctx.create_dataframe([ans])
+chk = df.aggregate([], [f.sum(col("v1"))]).collect()[0].column(0)[0]
+chkt = timeit.default_timer() - t_start
+m = memory_usage()
+write_log(task=task, data=data_name, in_rows=in_rows, question=question, out_rows=shape[0], out_cols=shape[1], solution=solution, version=ver, git=git, fun=fun, run=1, time_sec=t, mem_gb=m, cache=cache, chk=make_chk([chk]), chk_time_sec=chkt, on_disk=on_disk)
+del ans
+gc.collect()
+
+question = "sum v1 mean v3 by id3" # q3
+gc.collect()
+t_start = timeit.default_timer()
+ans = ctx.sql("SELECT id3, SUM(v1) AS v1, AVG(v3) AS v3 FROM x GROUP BY id3").collect()
+t = timeit.default_timer() - t_start
+print(t)
+shape = ans_shape(ans)
+print(shape)
+t_start = timeit.default_timer()
+df = ctx.create_dataframe([ans])
+chk = df.aggregate([], [f.sum(col("v1")), f.sum(col("v3"))]).collect()[0].to_pandas().to_numpy()[0]
+chkt = timeit.default_timer() - t_start
+m = memory_usage()
+write_log(task=task, data=data_name, in_rows=in_rows, question=question, out_rows=shape[0], out_cols=shape[1], solution=solution, version=ver, git=git, fun=fun, run=1, time_sec=t, mem_gb=m, cache=cache, chk=make_chk([chk]), chk_time_sec=chkt, on_disk=on_disk)
+del ans
+gc.collect()
+
+question = "mean v1:v3 by id4" # q4
+gc.collect()
+t_start = timeit.default_timer()
+ans = ctx.sql("SELECT id4, AVG(v1) AS v1, AVG(v2) AS v2, AVG(v3) AS v3 FROM x GROUP BY id4").collect()
+t = timeit.default_timer() - t_start
+print(t)
+shape = ans_shape(ans)
+print(shape)
+t_start = timeit.default_timer()
+df = ctx.create_dataframe([ans])
+chk = df.aggregate([], [f.sum(col("v1")), f.sum(col("v2")), f.sum(col("v3"))]).collect()[0].to_pandas().to_numpy()[0]
+chkt = timeit.default_timer() - t_start
+m = memory_usage()
+write_log(task=task, data=data_name, in_rows=in_rows, question=question, out_rows=shape[0], out_cols=shape[1], solution=solution, version=ver, git=git, fun=fun, run=1, time_sec=t, mem_gb=m, cache=cache, chk=make_chk([chk]), chk_time_sec=chkt, on_disk=on_disk)
+del ans
+gc.collect()
+
+question = "sum v1:v3 by id6" # q5
+gc.collect()
+t_start = timeit.default_timer()
+ans = ctx.sql("SELECT id6, SUM(v1) AS v1, SUM(v2) AS v2, SUM(v3) AS v3 FROM x GROUP BY id6").collect()
+t = timeit.default_timer() - t_start
+print(t)
+shape = ans_shape(ans)
+print(shape)
+t_start = timeit.default_timer()
+df = ctx.create_dataframe([ans])
+chk = df.aggregate([], [f.sum(col("v1")), f.sum(col("v2")), f.sum(col("v3"))]).collect()[0].to_pandas().to_numpy()[0]
+chkt = timeit.default_timer() - t_start
+m = memory_usage()
+write_log(task=task, data=data_name, in_rows=in_rows, question=question, out_rows=shape[0], out_cols=shape[1], solution=solution, version=ver, git=git, fun=fun, run=1, time_sec=t, mem_gb=m, cache=cache, chk=make_chk([chk]), chk_time_sec=chkt, on_disk=on_disk)
+del ans
+gc.collect()
+
+question = "max v1 - min v2 by id3" # q7
+gc.collect()
+t_start = timeit.default_timer()
+ans = ctx.sql("SELECT id3, MAX(v1) - MIN(v2) AS range_v1_v2 FROM x GROUP BY id3").collect()
+t = timeit.default_timer() - t_start
+print(t)
+shape = ans_shape(ans)
+print(shape)
+t_start = timeit.default_timer()
+df = ctx.create_dataframe([ans])
+chk = df.aggregate([], [f.sum(col("range_v1_v2"))]).collect()[0].column(0)[0]
+chkt = timeit.default_timer() - t_start
+m = memory_usage()
+write_log(task=task, data=data_name, in_rows=in_rows, question=question, out_rows=shape[0], out_cols=shape[1], solution=solution, version=ver, git=git, fun=fun, run=1, time_sec=t, mem_gb=m, cache=cache, chk=make_chk([chk]), chk_time_sec=chkt, on_disk=on_disk)
+del ans
+gc.collect()
+
+question = "largest two v3 by id6" # q8
+gc.collect()
+t_start = timeit.default_timer()
+ans = ctx.sql("SELECT id6, v3 from (SELECT id6, v3, row_number() OVER (PARTITION BY id6 ORDER BY v3 DESC) AS row FROM x) t WHERE row <= 2").collect()
+t = timeit.default_timer() - t_start
+print(t)
+shape = ans_shape(ans)
+print(shape)
+t_start = timeit.default_timer()
+df = ctx.create_dataframe([ans])
+chk = df.aggregate([], [f.sum(col("v3"))]).collect()[0].column(0)[0]
+chkt = timeit.default_timer() - t_start
+m = memory_usage()
+write_log(task=task, data=data_name, in_rows=in_rows, question=question, out_rows=shape[0], out_cols=shape[1], solution=solution, version=ver, git=git, fun=fun, run=1, time_sec=t, mem_gb=m, cache=cache, chk=make_chk([chk]), chk_time_sec=chkt, on_disk=on_disk)
+del ans
+gc.collect()
+
+question = "sum v3 count by id1:id6" # q10
+gc.collect()
+t_start = timeit.default_timer()
+ans = ctx.sql("SELECT id1, id2, id3, id4, id5, id6, SUM(v3) as v3, COUNT(*) AS cnt FROM x GROUP BY id1, id2, id3, id4, id5, id6").collect()
+t = timeit.default_timer() - t_start
+print(t)
+shape = ans_shape(ans)
+print(shape)
+t_start = timeit.default_timer()
+df = ctx.create_dataframe([ans])
+chk = df.aggregate([], [f.sum(col("v3")), f.sum(col("cnt"))]).collect()[0].to_pandas().to_numpy()[0]
+chkt = timeit.default_timer() - t_start
+m = memory_usage()
+write_log(task=task, data=data_name, in_rows=in_rows, question=question, out_rows=shape[0], out_cols=shape[1], solution=solution, version=ver, git=git, fun=fun, run=1, time_sec=t, mem_gb=m, cache=cache, chk=make_chk([chk]), chk_time_sec=chkt, on_disk=on_disk)
+del ans
+gc.collect()
+
+print("grouping finished, took %0.fs" % (timeit.default_timer() - task_init), flush=True)
+
+exit(0)

datafusion/setup-datafusion.sh

@@ -0,0 +1,5 @@
+#!/bin/bash
+set -e
+
+curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
+

datafusion/src/bin/groupby.rs

@@ -0,0 +1,99 @@
+use datafusion::error::Result;
+use datafusion::prelude::*;
+use datafusion::{
+    arrow::datatypes::{DataType, Field, Schema},
+    datasource::MemTable,
+};
+use std::time::Instant;
+use std::{env, sync::Arc};
+
+#[global_allocator]
+static ALLOC: snmalloc_rs::SnMalloc = snmalloc_rs::SnMalloc;
+
+async fn exec_query(ctx: &mut ExecutionContext, query: &str, name: &str) -> Result<()> {
+    let start = Instant::now();
+
+    let ans = ctx.sql(query).await?.collect().await?;
+
+    // TODO: print details
+
+    println!("{} took {} ms", name, start.elapsed().as_millis());
+
+    Ok(())
+}
+#[tokio::main]
+async fn main() -> Result<()> {
+    let batch_size = 65536;
+    let partition_size = num_cpus::get();
+    let cfg = ExecutionConfig::new()
+        .with_target_partitions(partition_size)
+        .with_batch_size(batch_size);
+    let mut ctx = ExecutionContext::with_config(cfg);
+    let data = format!("../data/{}.csv", env::var("SRC_DATANAME").unwrap());
+
+    let schema = Schema::new(vec![
+        Field::new("id1", DataType::Utf8, false),
+        Field::new("id2", DataType::Utf8, false),
+        Field::new("id3", DataType::Utf8, false),
+        Field::new("id4", DataType::Int32, false),
+        Field::new("id5", DataType::Int32, false),
+        Field::new("id6", DataType::Int32, false),
+        Field::new("v1", DataType::Int32, false),
+        Field::new("v2", DataType::Int32, false),
+        Field::new("v3", DataType::Float64, false),
+    ]);
+    let options = CsvReadOptions::new().schema(&schema).has_header(true);
+
+    let df = ctx.read_csv(&data, options).await?;
+    let batches = df.collect_partitioned().await?;
+    let memtbl = MemTable::try_new(Arc::new(schema), batches)?;
+    ctx.register_table("tbl", Arc::new(memtbl))?;
+
+    exec_query(
+        &mut ctx,
+        "SELECT id1, SUM(v1) AS v1 FROM tbl GROUP BY id1",
+        "q1",
+    )
+    .await?;
+    exec_query(
+        &mut ctx,
+        "SELECT id1, id2, SUM(v1) AS v1 FROM tbl GROUP BY id1, id2",
+        "q2",
+    )
+    .await?;
+    exec_query(
+        &mut ctx,
+        "SELECT id3, SUM(v1) AS v1, AVG(v3) AS v3 FROM tbl GROUP BY id3",
+        "q3",
+    )
+    .await?;
+    exec_query(
+        &mut ctx,
+        "SELECT id4, AVG(v1) AS v1, AVG(v2) AS v2, AVG(v3) AS v3 FROM tbl GROUP BY id4",
+        "q4",
+    )
+    .await?;
+    exec_query(
+        &mut ctx,
+        "SELECT id6, SUM(v1) AS v1, SUM(v2) AS v2, SUM(v3) AS v3 FROM tbl GROUP BY id6",
+        "q5",
+    )
+    .await?;
+    exec_query(
+        &mut ctx,
+        "SELECT id3, MAX(v1) - MIN(v2) AS range_v1_v2 FROM tbl GROUP BY id3",
+        "q7",
+    )
+    .await?;
+
+    exec_query(
+        &mut ctx,
+        "select id6, v3 from (select id6, v3, row_number() over (partition by id6 order by v3 desc) as row from tbl) t where row <= 2",
+        "q8",
+    )
+    .await?;
+
+    exec_query(&mut ctx, "SELECT id1, id2, id3, id4, id5, id6, SUM(v3) as v3, COUNT(*) AS cnt FROM tbl GROUP BY id1, id2, id3, id4, id5, id6", "q10").await?;
+
+    Ok(())
+}