skade-flight/examples/flight_bench.rs
//! skade-flight bench — what a Power BI refresh costs, measured over real
//! TCP with the stock Flight SQL client.
//!
//! Arms (each one `BenchResult` JSON line on stdout — the shared nordisk
//! envelope; a human table on stderr):
//!
//! flight.plan_latency GetFlightInfo (plan only) per op, µs
//! flight.sql_info GetSqlInfo round trip per op, µs
//! flight.get_tables GetTables round trip per op, µs
//! flight.scan_full SELECT * — rows/s and MB/s on the wire
//! flight.scan_projected two of six columns — rows/s
//! flight.filter_pushdown WHERE id % 1000 = 0 — base rows/s, rows out
//! flight.aggregate count/sum/avg over the table — base rows/s
//! flight.group_by GROUP BY a 100-bucket key — base rows/s
//! flight.concurrent_scans 8 clients × SELECT * — aggregate rows/s
//! skade-flight.capabilities static machine facts (no timed work)
//!
//! Env: SKADE_FLIGHT_BENCH_ROWS (default 2_000_000), SKADE_FLIGHT_BENCH_CLIENTS (8),
//! SKADE_FLIGHT_BENCH_PLAN_ITERS (200).
//! Run: cargo run --release --example flight_bench
use std::sync::Arc;
use std::time::{Duration, Instant};
use anyhow::Result;
use arrow_array::{Float64Array, Int32Array, Int64Array, RecordBatch, StringArray, TimestampMicrosecondArray};
use arrow_schema::{DataType, Field, Schema, TimeUnit};
use futures::TryStreamExt;
use nornir_bench::{BenchResult, machine_capabilities};
use serde_json::{Map, json};
use skade_flight::arrow_flight::sql::client::FlightSqlServiceClient;
use skade_flight::arrow_flight::sql::{CommandGetTables, SqlInfo};
use skade_flight::{Auth, Engine};
use tonic::transport::Channel;
fn env_usize(key: &str, default: usize) -> usize {
std::env::var(key).ok().and_then(|v| v.parse().ok()).unwrap_or(default)
}
const BATCH: usize = 65_536;
fn schema() -> Schema {
Schema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new("bucket", DataType::Int32, false),
Field::new("f", DataType::Float64, false),
Field::new("name", DataType::Utf8, false),
Field::new("ts", DataType::Timestamp(TimeUnit::Microsecond, None), false),
Field::new("g", DataType::Float64, true),
])
}
fn batch(schema: &Arc<Schema>, start: i64, n: usize) -> RecordBatch {
let ids: Vec<i64> = (start..start + n as i64).collect();
let bucket: Vec<i32> = ids.iter().map(|i| (i % 100) as i32).collect();
let f: Vec<f64> = ids.iter().map(|i| (*i as f64) * 0.5).collect();
let name: Vec<String> = ids.iter().map(|i| format!("row-{i}")).collect();
let ts: Vec<i64> = ids.iter().map(|i| 1_700_000_000_000_000 + i * 1_000).collect();
let g: Vec<Option<f64>> = ids.iter().map(|i| if i % 7 == 0 { None } else { Some(*i as f64) }).collect();
RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int64Array::from(ids)),
Arc::new(Int32Array::from(bucket)),
Arc::new(Float64Array::from(f)),
Arc::new(StringArray::from(name)),
Arc::new(TimestampMicrosecondArray::from(ts)),
Arc::new(Float64Array::from(g)),
],
)
.unwrap()
}
struct Arm {
name: &'static str,
metrics: Map<String, serde_json::Value>,
}
impl Arm {
fn result(&self) -> BenchResult {
BenchResult::measured(format!("flight.{}", self.name), self.metrics.clone())
}
}
fn rate_arm(name: &'static str, secs: f64, base_rows: usize, rows_out: usize, bytes: usize) -> Arm {
let mut m = Map::new();
m.insert("rows_per_sec".into(), json!(base_rows as f64 / secs));
m.insert("rows_out".into(), json!(rows_out as f64));
m.insert("mb_per_sec".into(), json!(bytes as f64 / 1e6 / secs));
m.insert("elapsed_s".into(), json!(secs));
Arm { name, metrics: m }
}
fn latency_arm(name: &'static str, iters: usize, total: Duration) -> Arm {
let mut m = Map::new();
m.insert("per_op_us".into(), json!(total.as_secs_f64() * 1e6 / iters as f64));
m.insert("ops_per_sec".into(), json!(iters as f64 / total.as_secs_f64()));
m.insert("iters".into(), json!(iters as f64));
Arm { name, metrics: m }
}
/// Run a statement, drain the stream, return (rows, wire bytes, elapsed).
async fn scan(c: &mut FlightSqlServiceClient<Channel>, sql: &str) -> Result<(usize, usize, Duration)> {
let t0 = Instant::now();
let info = c.execute(sql.to_string(), None).await?;
let ticket = info.endpoint[0].ticket.clone().expect("ticket");
let mut stream = c.do_get(ticket).await?;
let (mut rows, mut bytes) = (0usize, 0usize);
while let Some(b) = stream.try_next().await? {
rows += b.num_rows();
bytes += b.get_array_memory_size();
}
Ok((rows, bytes, t0.elapsed()))
}
#[tokio::main(flavor = "multi_thread")]
async fn main() -> Result<()> {
let rows = env_usize("SKADE_FLIGHT_BENCH_ROWS", 2_000_000);
let clients = env_usize("SKADE_FLIGHT_BENCH_CLIENTS", 8);
let plan_iters = env_usize("SKADE_FLIGHT_BENCH_PLAN_ITERS", 200);
// ── warehouse ──────────────────────────────────────────────────────────
let dir = tempfile::tempdir()?;
let wh = skade::Warehouse::open(dir.path().join("wh")).await?;
let schema = Arc::new(schema());
let mut table = wh.create_table("t", &schema).await?;
let t0 = Instant::now();
let mut batches = Vec::with_capacity(rows / BATCH + 1);
let mut start = 0i64;
while (start as usize) < rows {
let n = BATCH.min(rows - start as usize);
batches.push(batch(&schema, start, n));
start += n as i64;
}
table.append_parallel(&batches).await?;
eprintln!("# wrote {rows} rows in {} batches: {:.2}s", batches.len(), t0.elapsed().as_secs_f64());
drop(batches);
// ── server ─────────────────────────────────────────────────────────────
let mut eopts = skade_flight::engine::EngineOptions::default();
if let Some(b) = std::env::var("SKADE_FLIGHT_BENCH_BATCH").ok().and_then(|v| v.parse().ok()) {
eopts.batch_size = b;
}
let engine = Arc::new(Engine::with_engine_options(wh, eopts).await?);
let engine_for_local = engine.clone();
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
let addr = listener.local_addr()?;
let _server = tokio::spawn(skade_flight::serve_on(engine, Arc::new(Auth::open()), listener, Default::default(), futures::future::pending()));
let endpoint = format!("http://{addr}");
// HTTP/2 flow control: the SENDER is bounded by the receiver's advertised
// window. tonic's client default is h2's 64 KiB, which on loopback caps a
// stream near 1.2 GB/s. SKADE_FLIGHT_BENCH_WINDOW_MB (default 8) opens it,
// which is what a real client (gRPC C++ with BDP probing) converges to.
let win_mb: u32 = std::env::var("SKADE_FLIGHT_BENCH_WINDOW_MB").ok().and_then(|v| v.parse().ok()).unwrap_or(8);
let channel = mk_endpoint(endpoint.clone(), win_mb).connect().await?;
let mut c = FlightSqlServiceClient::new(channel);
let mut arms: Vec<Arm> = Vec::new();
// plan latency
let t0 = Instant::now();
for _ in 0..plan_iters {
c.execute("SELECT id, f FROM t WHERE id > 5 ORDER BY id LIMIT 10".into(), None).await?;
}
arms.push(latency_arm("plan_latency", plan_iters, t0.elapsed()));
// sql_info
let t0 = Instant::now();
for _ in 0..plan_iters {
let info = c.get_sql_info(vec![SqlInfo::FlightSqlServerName]).await?;
let _: Vec<RecordBatch> = c.do_get(info.endpoint[0].ticket.clone().unwrap()).await?.try_collect().await?;
}
arms.push(latency_arm("sql_info", plan_iters, t0.elapsed()));
// get_tables
let t0 = Instant::now();
for _ in 0..plan_iters {
let info = c.get_tables(CommandGetTables { include_schema: true, ..Default::default() }).await?;
let _: Vec<RecordBatch> = c.do_get(info.endpoint[0].ticket.clone().unwrap()).await?.try_collect().await?;
}
arms.push(latency_arm("get_tables", plan_iters, t0.elapsed()));
// engine only, no wire: where does the time go?
for (name, sql) in [("engine_scan_full", "SELECT * FROM t"), ("engine_scan_projected", "SELECT id, f FROM t")] {
let t0 = Instant::now();
let mut s = engine_for_local.execute_stream(sql).await?;
let (mut r, mut b) = (0usize, 0usize);
while let Some(batch) = s.try_next().await? {
r += batch.num_rows();
b += batch.get_array_memory_size();
}
arms.push(rate_arm(name, t0.elapsed().as_secs_f64(), rows, r, b));
}
// raw frames, no client decode: the server + wire alone
{
let t0 = Instant::now();
let info = c.execute("SELECT * FROM t".into(), None).await?;
let ticket = info.endpoint[0].ticket.clone().expect("ticket");
let mut raw = c.inner_mut().do_get(ticket).await?.into_inner();
let (mut frames, mut bytes) = (0usize, 0usize);
while let Some(fd) = raw.try_next().await? {
frames += 1;
bytes += fd.data_header.len() + fd.data_body.len();
}
let mut arm = rate_arm("scan_full_raw", t0.elapsed().as_secs_f64(), rows, rows, bytes);
arm.metrics.insert("frames".into(), json!(frames as f64));
arms.push(arm);
}
// scans
let (r, b, d) = scan(&mut c, "SELECT * FROM t").await?;
assert_eq!(r, rows);
arms.push(rate_arm("scan_full", d.as_secs_f64(), rows, r, b));
let (r, b, d) = scan(&mut c, "SELECT id, f FROM t").await?;
arms.push(rate_arm("scan_projected", d.as_secs_f64(), rows, r, b));
let (r, b, d) = scan(&mut c, "SELECT id FROM t WHERE id % 1000 = 0").await?;
arms.push(rate_arm("filter_pushdown", d.as_secs_f64(), rows, r, b));
let (r, b, d) = scan(&mut c, "SELECT count(*), sum(f), avg(g) FROM t").await?;
arms.push(rate_arm("aggregate", d.as_secs_f64(), rows, r, b));
let (r, b, d) = scan(&mut c, "SELECT bucket, count(*), avg(f) FROM t GROUP BY bucket").await?;
arms.push(rate_arm("group_by", d.as_secs_f64(), rows, r, b));
// concurrent scans
let t0 = Instant::now();
let mut handles = Vec::new();
for _ in 0..clients {
let ep = endpoint.clone();
handles.push(tokio::spawn(async move {
let ch = mk_endpoint(ep, win_mb).connect().await.unwrap();
let mut c = FlightSqlServiceClient::new(ch);
scan(&mut c, "SELECT * FROM t").await.unwrap()
}));
}
let (mut r_sum, mut b_sum) = (0, 0);
for h in handles {
let (r, b, _) = h.await?;
r_sum += r;
b_sum += b;
}
let d = t0.elapsed();
let mut arm = rate_arm("concurrent_scans", d.as_secs_f64(), r_sum, r_sum, b_sum);
arm.metrics.insert("clients".into(), json!(clients as f64));
arms.push(arm);
// ── emit ───────────────────────────────────────────────────────────────
println!("{}", serde_json::to_string(&machine_capabilities("skade-flight.capabilities"))?);
for a in &arms {
println!("{}", serde_json::to_string(&a.result())?);
}
eprintln!("# {rows} rows · {clients} clients");
eprintln!("# {:<24} {:>14} {:>12} {:>10}", "arm", "rows/s|per_op_us", "MB/s|ops/s", "rows_out");
for a in &arms {
let m = &a.metrics;
let (x, y) = if let Some(v) = m.get("rows_per_sec") {
(v.as_f64().unwrap_or(0.0), m["mb_per_sec"].as_f64().unwrap_or(0.0))
} else {
(m["per_op_us"].as_f64().unwrap_or(0.0), m["ops_per_sec"].as_f64().unwrap_or(0.0))
};
let out = m.get("rows_out").and_then(|v| v.as_f64()).unwrap_or(0.0);
eprintln!("# {:<24} {:>14.0} {:>12.1} {:>10.0}", a.name, x, y, out);
}
Ok(())
}
fn mk_endpoint(ep: String, win_mb: u32) -> tonic::transport::Endpoint {
let e = Channel::from_shared(ep).expect("endpoint");
let fs: u32 = std::env::var("SKADE_FLIGHT_BENCH_FRAME_KB").ok().and_then(|v| v.parse().ok()).unwrap_or(1024);
let e = if fs == 0 { e } else { e.max_frame_size((fs << 10).clamp(16 * 1024, (1 << 24) - 1)) };
if win_mb == 0 { e } else { e.initial_stream_window_size(win_mb << 20).initial_connection_window_size(win_mb << 21) }
}