use std::{sync::Arc, time::Instant};
use arrow_array::{Int32Array, LargeBinaryArray, RecordBatch};
use arrow_schema::{DataType, Field, Schema};
use lancedb::{
Result, Table,
arrow::{SendableRecordBatchStream, SimpleRecordBatchStream},
connect,
dataloader::permutation::{
builder::{PermutationBuilder, ShuffleStrategy},
reader::PermutationReader,
split::{SplitSizes, SplitStrategy},
},
query::Select,
};
fn env_usize(key: &str, default: usize) -> usize {
std::env::var(key)
.ok()
.and_then(|v| v.parse().ok())
.unwrap_or(default)
}
async fn make_base_table(num_rows: usize, row_bytes: usize) -> Result<Table> {
let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int32, false),
Field::new("payload", DataType::LargeBinary, false),
]));
let payload = vec![0u8; row_bytes];
let ids: Int32Array = (0..num_rows as i32).collect();
let payloads: LargeBinaryArray = (0..num_rows).map(|_| Some(payload.as_slice())).collect();
let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(ids), Arc::new(payloads)])?;
let stream: SendableRecordBatchStream = Box::pin(SimpleRecordBatchStream::new(
futures::stream::once(std::future::ready(Ok(batch))),
schema,
));
let db = connect("memory:///").execute().await?;
db.create_table("base", stream).execute().await
}
async fn make_permutation_table(base: &Table, num_splits: usize) -> Result<Table> {
PermutationBuilder::new(base.clone())
.with_split_strategy(
SplitStrategy::Random {
seed: Some(42),
sizes: SplitSizes::Fixed(num_splits as u64),
clump_size: None,
},
None,
)
.with_shuffle_strategy(ShuffleStrategy::Random {
seed: Some(42),
clump_size: None,
})
.build()
.await
}
async fn run_hot_loop(
readers: &[PermutationReader],
chunk_size: usize,
steps: usize,
) -> Result<(usize, f64)> {
let n = readers.len();
let split_sizes: Vec<usize> = readers.iter().map(|r| r.count_rows() as usize).collect();
struct SplitBuf {
batch: Option<RecordBatch>,
row_in_batch: usize,
consumed: usize,
}
let mut bufs: Vec<SplitBuf> = (0..n)
.map(|_| SplitBuf {
batch: None,
row_in_batch: 0,
consumed: 0,
})
.collect();
for i in 0..n {
let fetch = chunk_size.min(split_sizes[i]);
if fetch > 0 {
let offsets: Vec<u64> = (0..fetch as u64).collect();
bufs[i].batch = Some(readers[i].take_offsets(&offsets, Select::All).await?);
}
}
let mut total_rows = 0usize;
let t0 = Instant::now();
'outer: for _step in 0..steps {
for i in 0..n {
if bufs[i].consumed >= split_sizes[i] {
break 'outer;
}
let need_refill = bufs[i]
.batch
.as_ref()
.map(|b| bufs[i].row_in_batch >= b.num_rows())
.unwrap_or(true);
if need_refill {
let start = bufs[i].consumed as u64;
let remaining = (split_sizes[i] - bufs[i].consumed) as u64;
let fetch = chunk_size.min(remaining as usize);
let offsets: Vec<u64> = (start..start + fetch as u64).collect();
bufs[i].batch = Some(readers[i].take_offsets(&offsets, Select::All).await?);
bufs[i].row_in_batch = 0;
}
bufs[i].row_in_batch += 1;
bufs[i].consumed += 1;
total_rows += 1;
}
}
Ok((total_rows, t0.elapsed().as_secs_f64()))
}
#[tokio::main]
async fn main() -> Result<()> {
let num_splits = env_usize("BENCH_NUM_SPLITS", 24);
let num_rows = env_usize("BENCH_NUM_ROWS", num_splits * 2048);
let steps = env_usize("BENCH_STEPS", 200);
let row_bytes = env_usize("BENCH_ROW_BYTES", 4096);
let single_chunk: Option<usize> = std::env::var("BENCH_CHUNK")
.ok()
.and_then(|v| v.parse().ok());
let do_profile = std::env::var("BENCH_PROFILE")
.map(|v| v == "1")
.unwrap_or(false);
assert_eq!(
num_rows % num_splits,
0,
"NUM_ROWS must be divisible by NUM_SPLITS"
);
println!("Benchmark config:");
println!(
" num_rows={} num_splits={} rows/split={} steps={} row_bytes={}",
num_rows,
num_splits,
num_rows / num_splits,
steps,
row_bytes,
);
println!(
" ~{:.1} MB total",
(num_rows * row_bytes) as f64 / (1024.0 * 1024.0)
);
println!();
print!("Building base table... ");
let _ = std::io::Write::flush(&mut std::io::stdout());
let base = make_base_table(num_rows, row_bytes).await?;
println!("done");
print!("Building permutation table... ");
let _ = std::io::Write::flush(&mut std::io::stdout());
let perm = make_permutation_table(&base, num_splits).await?;
println!("done");
print!("Building {} PermutationReaders... ", num_splits);
let _ = std::io::Write::flush(&mut std::io::stdout());
let base_inner = base.base_table().clone();
let perm_inner = perm.base_table().clone();
let mut readers = Vec::with_capacity(num_splits);
for split in 0..num_splits {
readers.push(
PermutationReader::try_from_tables(
base_inner.clone(),
perm_inner.clone(),
split as u64,
)
.await?,
);
}
println!("done ({} rows/split)", readers[0].count_rows());
println!();
let chunk_sizes: Vec<usize> = if let Some(c) = single_chunk {
vec![c]
} else {
vec![1, 4, 16, 64, 256, 1024, 4096, 16384]
};
if do_profile {
#[cfg(unix)]
{
let chunk = chunk_sizes[0];
println!("Profiling chunk={chunk} for {steps} steps...");
let _ = run_hot_loop(&readers, chunk, 1).await?;
let guard = pprof::ProfilerGuardBuilder::default()
.frequency(1000)
.build()
.unwrap();
let (rows, elapsed) = run_hot_loop(&readers, chunk, steps).await?;
if let Ok(report) = guard.report().build() {
let svg_path = "flamegraph.svg";
let file = std::fs::File::create(svg_path).unwrap();
report.flamegraph(file).unwrap();
println!("Flamegraph written to {svg_path}");
}
let rows_per_sec = rows as f64 / elapsed;
println!("chunk={chunk} {rows} rows {elapsed:.3}s {rows_per_sec:.0} rows/s");
}
#[cfg(not(unix))]
{
println!("Flamegraph profiling (BENCH_PROFILE=1) is not supported on this platform.");
println!("Run without BENCH_PROFILE to get throughput numbers.");
}
} else {
println!(
"{:>6} {:>7} {:>8} {:>11} {:>10}",
"chunk", "rows", "elapsed", "rows/s", "ms/step"
);
println!("{}", "-".repeat(52));
for &chunk in &chunk_sizes {
let _ = run_hot_loop(&readers, chunk, 1).await?;
let (rows, elapsed) = run_hot_loop(&readers, chunk, steps).await?;
let rows_per_sec = rows as f64 / elapsed;
let ms_per_step = elapsed / steps as f64 * 1000.0;
println!(
"{:>6} {:>7} {:>7.3}s {:>11.0} {:>9.1}ms",
chunk, rows, elapsed, rows_per_sec, ms_per_step,
);
}
}
println!("\nDone.");
Ok(())
}