Implement smoke 'benchmark'

Signed-off-by: Abhi Agarwal <[email protected]>

Author: abhiaagarwal

crates/benchmarks/benches/smoke.rs

@@ -0,0 +1,35 @@
+use delta_benchmarks::{run_smoke_once, SmokeParams};
+use divan::{AllocProfiler, Bencher};
+use url::Url;
+
+fn main() {
+    divan::main();
+}
+
+#[global_allocator]
+static ALLOC: AllocProfiler = AllocProfiler::system();
+
+type Runtime = tokio::runtime::Runtime;
+
+fn bench_smoke(bencher: Bencher, params: &SmokeParams) {
+    let rt = Runtime::new().expect("create tokio runtime");
+    bencher
+        .with_inputs(|| tempfile::tempdir().expect("create temp dir"))
+        .bench_local_values(|tmp_dir| {
+            let table_url = Url::from_directory_path(tmp_dir.path()).expect("tmp dir url");
+            rt.block_on(async {
+                run_smoke_once(&table_url, params).await.expect("smoke run");
+            });
+            drop(tmp_dir);
+        });
+}
+
+#[divan::bench(args = [
+    SmokeParams { rows: 2 },
+    SmokeParams { rows: 10 },
+    SmokeParams { rows: 100 },
+    SmokeParams { rows: 1_000 },
+])]
+fn smoke(bencher: Bencher, params: &SmokeParams) {
+    bench_smoke(bencher, params);
+}

crates/benchmarks/src/lib.rs

@@ -1,184 +1,7 @@
-use std::path::Path;
+pub mod merge;
+pub mod smoke;
 
-use deltalake_core::datafusion::functions::expr_fn;
-use deltalake_core::kernel::engine::arrow_conversion::TryIntoKernel;
-use deltalake_core::kernel::{StructField, StructType};
-use deltalake_core::operations::merge::MergeBuilder;
-use deltalake_core::{arrow, DeltaResult};
-use deltalake_core::{
-    datafusion::{
-        logical_expr::{cast, lit},
-        prelude::{DataFrame, ParquetReadOptions, SessionContext},
-    },
-    DeltaOps, DeltaTable, DeltaTableError,
+pub use merge::{
+    merge_delete, merge_insert, merge_upsert, prepare_source_and_table, MergeOp, MergePerfParams,
 };
-use tempfile::TempDir;
-use url::Url;
-
-pub type MergeOp = fn(DataFrame, DeltaTable) -> Result<MergeBuilder, DeltaTableError>;
-
-#[derive(Debug, Clone)]
-pub struct MergePerfParams {
-    pub sample_matched_rows: f32,
-    pub sample_not_matched_rows: f32,
-}
-
-pub fn merge_upsert(source: DataFrame, table: DeltaTable) -> Result<MergeBuilder, DeltaTableError> {
-    deltalake_core::DeltaOps(table)
-        .merge(source, "source.wr_item_sk = target.wr_item_sk and source.wr_order_number = target.wr_order_number")
-        .with_source_alias("source")
-        .with_target_alias("target")
-        .when_matched_update(|update| {
-            update
-                .update("wr_returned_date_sk", "source.wr_returned_date_sk")
-                .update("wr_returned_time_sk", "source.wr_returned_time_sk")
-                .update("wr_item_sk", "source.wr_item_sk")
-                .update("wr_refunded_customer_sk", "source.wr_refunded_customer_sk")
-                .update("wr_refunded_cdemo_sk", "source.wr_refunded_cdemo_sk")
-                .update("wr_refunded_hdemo_sk", "source.wr_refunded_hdemo_sk")
-                .update("wr_refunded_addr_sk", "source.wr_refunded_addr_sk")
-                .update("wr_returning_customer_sk", "source.wr_returning_customer_sk")
-                .update("wr_returning_cdemo_sk", "source.wr_returning_cdemo_sk")
-                .update("wr_returning_hdemo_sk", "source.wr_returning_hdemo_sk")
-                .update("wr_returning_addr_sk", "source.wr_returning_addr_sk")
-                .update("wr_web_page_sk", "source.wr_web_page_sk")
-                .update("wr_reason_sk", "source.wr_reason_sk")
-                .update("wr_order_number", "source.wr_order_number")
-                .update("wr_return_quantity", "source.wr_return_quantity")
-                .update("wr_return_amt", "source.wr_return_amt")
-                .update("wr_return_tax", "source.wr_return_tax")
-                .update("wr_return_amt_inc_tax", "source.wr_return_amt_inc_tax")
-                .update("wr_fee", "source.wr_fee")
-                .update("wr_return_ship_cost", "source.wr_return_ship_cost")
-                .update("wr_refunded_cash", "source.wr_refunded_cash")
-                .update("wr_reversed_charge", "source.wr_reversed_charge")
-                .update("wr_account_credit", "source.wr_account_credit")
-                .update("wr_net_loss", "source.wr_net_loss")
-        })?
-        .when_not_matched_insert(|insert| {
-            insert
-                .set("wr_returned_date_sk", "source.wr_returned_date_sk")
-                .set("wr_returned_time_sk", "source.wr_returned_time_sk")
-                .set("wr_item_sk", "source.wr_item_sk")
-                .set("wr_refunded_customer_sk", "source.wr_refunded_customer_sk")
-                .set("wr_refunded_cdemo_sk", "source.wr_refunded_cdemo_sk")
-                .set("wr_refunded_hdemo_sk", "source.wr_refunded_hdemo_sk")
-                .set("wr_refunded_addr_sk", "source.wr_refunded_addr_sk")
-                .set("wr_returning_customer_sk", "source.wr_returning_customer_sk")
-                .set("wr_returning_cdemo_sk", "source.wr_returning_cdemo_sk")
-                .set("wr_returning_hdemo_sk", "source.wr_returning_hdemo_sk")
-                .set("wr_returning_addr_sk", "source.wr_returning_addr_sk")
-                .set("wr_web_page_sk", "source.wr_web_page_sk")
-                .set("wr_reason_sk", "source.wr_reason_sk")
-                .set("wr_order_number", "source.wr_order_number")
-                .set("wr_return_quantity", "source.wr_return_quantity")
-                .set("wr_return_amt", "source.wr_return_amt")
-                .set("wr_return_tax", "source.wr_return_tax")
-                .set("wr_return_amt_inc_tax", "source.wr_return_amt_inc_tax")
-                .set("wr_fee", "source.wr_fee")
-                .set("wr_return_ship_cost", "source.wr_return_ship_cost")
-                .set("wr_refunded_cash", "source.wr_refunded_cash")
-                .set("wr_reversed_charge", "source.wr_reversed_charge")
-                .set("wr_account_credit", "source.wr_account_credit")
-                .set("wr_net_loss", "source.wr_net_loss")
-        })
-}
-
-pub fn merge_insert(source: DataFrame, table: DeltaTable) -> Result<MergeBuilder, DeltaTableError> {
-    deltalake_core::DeltaOps(table)
-        .merge(source, "source.wr_item_sk = target.wr_item_sk and source.wr_order_number = target.wr_order_number")
-        .with_source_alias("source")
-        .with_target_alias("target")
-        .when_not_matched_insert(|insert| {
-            insert
-                .set("wr_returned_date_sk", "source.wr_returned_date_sk")
-                .set("wr_returned_time_sk", "source.wr_returned_time_sk")
-                .set("wr_item_sk", "source.wr_item_sk")
-                .set("wr_refunded_customer_sk", "source.wr_refunded_customer_sk")
-                .set("wr_refunded_cdemo_sk", "source.wr_refunded_cdemo_sk")
-                .set("wr_refunded_hdemo_sk", "source.wr_refunded_hdemo_sk")
-                .set("wr_refunded_addr_sk", "source.wr_refunded_addr_sk")
-                .set("wr_returning_customer_sk", "source.wr_returning_customer_sk")
-                .set("wr_returning_cdemo_sk", "source.wr_returning_cdemo_sk")
-                .set("wr_returning_hdemo_sk", "source.wr_returning_hdemo_sk")
-                .set("wr_returning_addr_sk", "source.wr_returning_addr_sk")
-                .set("wr_web_page_sk", "source.wr_web_page_sk")
-                .set("wr_reason_sk", "source.wr_reason_sk")
-                .set("wr_order_number", "source.wr_order_number")
-                .set("wr_return_quantity", "source.wr_return_quantity")
-                .set("wr_return_amt", "source.wr_return_amt")
-                .set("wr_return_tax", "source.wr_return_tax")
-                .set("wr_return_amt_inc_tax", "source.wr_return_amt_inc_tax")
-                .set("wr_fee", "source.wr_fee")
-                .set("wr_return_ship_cost", "source.wr_return_ship_cost")
-                .set("wr_refunded_cash", "source.wr_refunded_cash")
-                .set("wr_reversed_charge", "source.wr_reversed_charge")
-                .set("wr_account_credit", "source.wr_account_credit")
-                .set("wr_net_loss", "source.wr_net_loss")
-        })
-}
-
-pub fn merge_delete(source: DataFrame, table: DeltaTable) -> Result<MergeBuilder, DeltaTableError> {
-    deltalake_core::DeltaOps(table)
-        .merge(source, "source.wr_item_sk = target.wr_item_sk and source.wr_order_number = target.wr_order_number")
-        .with_source_alias("source")
-        .with_target_alias("target")
-        .when_matched_delete(|delete| delete)
-}
-
-/// Prepare source DataFrame and target Delta table from DuckDB-generated TPC-DS parquet.
-/// Creates a temporary Delta table from web_returns.parquet as the target.
-/// Returns (source_df, target_table) for benchmarking.
-pub async fn prepare_source_and_table(
-    params: &MergePerfParams,
-    tmp_dir: &TempDir,
-    parquet_dir: &Path,
-) -> DeltaResult<(DataFrame, DeltaTable)> {
-    let ctx = SessionContext::new();
-
-    let parquet_path = parquet_dir
-        .join("web_returns.parquet")
-        .to_str()
-        .unwrap()
-        .to_owned();
-
-    let parquet_df = ctx
-        .read_parquet(&parquet_path, ParquetReadOptions::default())
-        .await?;
-    let temp_table_url = Url::from_directory_path(tmp_dir).unwrap();
-
-    let schema = parquet_df.schema();
-    let delta_schema: StructType = schema.as_arrow().try_into_kernel().unwrap();
-
-    let batches = parquet_df.collect().await?;
-    let fields: Vec<StructField> = delta_schema.fields().cloned().collect();
-    let table = DeltaOps::try_from_uri(temp_table_url)
-        .await?
-        .create()
-        .with_columns(fields)
-        .await?;
-
-    let table = DeltaOps(table).write(batches).await?;
-
-    // Now prepare source DataFrame with sampling
-    let source = ctx
-        .read_parquet(&parquet_path, ParquetReadOptions::default())
-        .await?;
-
-    // Split matched and not-matched portions
-    let matched = source
-        .clone()
-        .filter(expr_fn::random().lt_eq(lit(params.sample_matched_rows)))?;
-
-    let rand = cast(
-        expr_fn::random() * lit(u32::MAX),
-        arrow::datatypes::DataType::Int64,
-    );
-    let not_matched = source
-        .filter(expr_fn::random().lt_eq(lit(params.sample_not_matched_rows)))?
-        .with_column("wr_item_sk", rand.clone())?
-        .with_column("wr_order_number", rand)?;
-
-    let source = matched.union(not_matched)?;
-    Ok((source, table))
-}
+pub use smoke::{run_smoke_once, SmokeParams};

crates/benchmarks/src/main.rs

@@ -1,10 +1,12 @@
 use std::{path::PathBuf, time::Instant};
 
-use clap::{Parser, ValueEnum};
+use clap::{Parser, Subcommand, ValueEnum};
 
 use delta_benchmarks::{
-    merge_delete, merge_insert, merge_upsert, prepare_source_and_table, MergeOp, MergePerfParams,
+    merge_delete, merge_insert, merge_upsert, prepare_source_and_table, run_smoke_once, MergeOp,
+    MergePerfParams, SmokeParams,
 };
+use deltalake_core::ensure_table_uri;
 
 #[derive(Copy, Clone, Debug, ValueEnum)]
 enum OpKind {
@@ -14,53 +16,103 @@ enum OpKind {
 }
 
 #[derive(Parser, Debug)]
-#[command(about = "Run a merge benchmark with configurable parameters")]
+#[command(about = "Run delta-rs benchmarks")]
 struct Cli {
-    /// Operation to benchmark
-    #[arg(value_enum)]
-    op: OpKind,
+    #[command(subcommand)]
+    command: Command,
+}
+
+#[derive(Debug, Subcommand)]
+enum Command {
+    /// Run a merge benchmark with configurable parameters
+    Merge {
+        /// Operation to benchmark
+        #[arg(value_enum)]
+        op: OpKind,
+
+        /// Fraction of rows that match an existing key (0.0-1.0)
+        #[arg(long, default_value_t = 0.01)]
+        matched: f32,
+
+        /// Fraction of rows that do not match (0.0-1.0)
+        #[arg(long, default_value_t = 0.10)]
+        not_matched: f32,
+    },
 
-    /// Fraction of rows that match an existing key (0.0-1.0)
-    #[arg(long, default_value_t = 0.01)]
-    matched: f32,
+    /// Run the smoke workload to validate delta-rs read/write operations
+    Smoke {
+        /// Number of rows to write into the smoke table
+        #[arg(long, default_value_t = 2)]
+        rows: usize,
 
-    /// Fraction of rows that do not match (0.0-1.0)
-    #[arg(long, default_value_t = 0.10)]
-    not_matched: f32,
+        /// Optional table path to reuse for the smoke run (defaults to a temporary directory)
+        #[arg(long)]
+        table_path: Option<PathBuf>,
+    },
 }
 
 #[tokio::main]
-async fn main() {
+async fn main() -> anyhow::Result<()> {
     let cli = Cli::parse();
 
-    let op_fn: MergeOp = match cli.op {
-        OpKind::Upsert => merge_upsert,
-        OpKind::Delete => merge_delete,
-        OpKind::Insert => merge_insert,
-    };
+    match cli.command {
+        Command::Merge {
+            op,
+            matched,
+            not_matched,
+        } => {
+            let op_fn: MergeOp = match op {
+                OpKind::Upsert => merge_upsert,
+                OpKind::Delete => merge_delete,
+                OpKind::Insert => merge_insert,
+            };
+
+            let params = MergePerfParams {
+                sample_matched_rows: matched,
+                sample_not_matched_rows: not_matched,
+            };
+
+            let tmp_dir = tempfile::tempdir()?;
+
+            let parquet_dir = PathBuf::from(
+                std::env::var("TPCDS_PARQUET_DIR")
+                    .unwrap_or_else(|_| "crates/benchmarks/data/tpcds_parquet".to_string()),
+            );
 
-    let params = MergePerfParams {
-        sample_matched_rows: cli.matched,
-        sample_not_matched_rows: cli.not_matched,
-    };
+            let (source, table) = prepare_source_and_table(&params, &tmp_dir, &parquet_dir).await?;
 
-    let tmp_dir = tempfile::tempdir().expect("create tmp dir");
+            let start = Instant::now();
+            let (_table, metrics) = op_fn(source, table)?.await?;
+            let elapsed = start.elapsed();
 
-    let parquet_dir = PathBuf::from(
-        std::env::var("TPCDS_PARQUET_DIR")
-            .unwrap_or_else(|_| "crates/benchmarks/data/tpcds_parquet".to_string()),
-    );
+            println!(
+                "merge_duration_ms={} metrics={:?}",
+                elapsed.as_millis(),
+                metrics
+            );
+        }
+        Command::Smoke { rows, table_path } => {
+            let params = SmokeParams { rows };
+            let (table_url, _guard) = match table_path {
+                Some(path) => (ensure_table_uri(path.to_string_lossy().as_ref())?, None),
+                None => {
+                    let dir = tempfile::tempdir()?;
+                    let url = ensure_table_uri(dir.path().to_string_lossy().as_ref())?;
+                    (url, Some(dir))
+                }
+            };
 
-    let (source, table) = prepare_source_and_table(&params, &tmp_dir, &parquet_dir)
-        .await
-        .expect("prepare inputs");
+            let start = Instant::now();
+            run_smoke_once(&table_url, &params).await?;
+            let elapsed = start.elapsed();
 
-    let start = Instant::now();
-    let (_table, metrics) = op_fn(source, table)
-        .expect("build merge")
-        .await
-        .expect("execute merge");
-    let elapsed = start.elapsed();
+            println!(
+                "smoke_duration_ms={} table_uri={}",
+                elapsed.as_millis(),
+                table_url
+            );
+        }
+    }
 
-    println!("duration_ms={} metrics={:?}", elapsed.as_millis(), metrics)
+    Ok(())
 }