Duckdb run over public_bi (#3143)

Author: joseph-isaacs

bench-vortex/scripts/public_pi_run.sh

@@ -0,0 +1,17 @@
+#!/bin/bash
+
+set -Eeu -o pipefail -x
+
+# List files in the current directory
+readonly dir=$(dirname ${BASH_SOURCE[0]})
+readonly files=$(ls $dir/../public_bi/benchmark)
+
+for file in $files; do
+    echo "Running public BI: $file"
+      
+    file_lowercase=$(echo "$file" | tr '[:upper:]' '[:lower:]')
+
+    cargo run --profile bench --bin public_bi -- --targets=datafusion:vortex,duckdb:vortex -d $file_lowercase -i1
+
+    echo ""
+done

bench-vortex/src/bin/public_bi.rs

@@ -6,8 +6,8 @@ use bench_vortex::metrics::MetricsSetExt;
 use bench_vortex::public_bi::{FileType, PBI_DATASETS, PBIDataset};
 use bench_vortex::utils::constants::STORAGE_NVME;
 use bench_vortex::utils::new_tokio_runtime;
-use bench_vortex::{Engine, Format, Target, default_env_filter, df};
-use clap::Parser;
+use bench_vortex::{Format, Target, default_env_filter, df};
+use clap::{Parser, value_parser};
 use indicatif::ProgressBar;
 use itertools::Itertools;
 use tracing::info_span;
@@ -22,10 +22,20 @@ struct Args {
     iterations: usize,
     #[arg(short, long)]
     threads: Option<usize>,
-    #[arg(long, value_delimiter = ',', value_enum, default_values_t = vec![Format::Parquet, Format::OnDiskVortex])]
-    formats: Vec<Format>,
+    #[arg(long, value_delimiter = ',', value_parser = value_parser!(Target),
+        default_values = vec![
+            "datafusion:parquet",
+            "datafusion:vortex",
+            "duckdb:parquet",
+            "duckdb:vortex",
+            "duckdb:duckdb"
+        ]
+    )]
+    targets: Vec<Target>,
     #[arg(short, long)]
     verbose: bool,
+    #[arg(long)]
+    display_metrics: bool,
     #[arg(long, default_value_t, value_enum)]
     display_format: DisplayFormat,
     #[arg(long, default_value_t = false)]
@@ -76,12 +86,6 @@ fn main() -> anyhow::Result<()> {
 
     let runtime = new_tokio_runtime(args.threads);
 
-    let targets = args
-        .formats
-        .iter()
-        .map(|f| Target::new(Engine::DataFusion, *f))
-        .collect_vec();
-
     let pbi_dataset = PBI_DATASETS.get(args.dataset);
     let queries = match args.queries.clone() {
         None => pbi_dataset.queries()?,
@@ -92,7 +96,7 @@ fn main() -> anyhow::Result<()> {
             .collect(),
     };
 
-    let progress_bar = ProgressBar::new((queries.len() * args.formats.len()) as u64);
+    let progress_bar = ProgressBar::new((queries.len() * args.targets.len()) as u64);
     let mut all_measurements = Vec::default();
     let mut metrics = Vec::new();
 
@@ -101,7 +105,7 @@ fn main() -> anyhow::Result<()> {
     // download csvs, unzip, convert to parquet, and convert that to vortex
     runtime.block_on(dataset.write_as_vortex());
 
-    for target in &targets {
+    for target in &args.targets {
         let format = target.format();
         let session =
             df::get_session_context(args.emulate_object_store, args.disable_datafusion_cache);
@@ -164,21 +168,23 @@ fn main() -> anyhow::Result<()> {
 
     match args.display_format {
         DisplayFormat::Table => {
-            for (query, format, metric_sets) in metrics {
-                println!("\nmetrics for query={query}, {format}:");
-                for (idx, metric_set) in metric_sets.into_iter().enumerate() {
-                    println!("scan[{idx}]:");
-                    for m in metric_set
-                        .timestamps_removed()
-                        .aggregate()
-                        .sorted_for_display()
-                        .iter()
-                    {
-                        println!("{}", m);
+            if args.display_metrics {
+                for (query, format, metric_sets) in metrics {
+                    println!("\nmetrics for query={query}, {format}:");
+                    for (idx, metric_set) in metric_sets.into_iter().enumerate() {
+                        println!("scan[{idx}]:");
+                        for m in metric_set
+                            .timestamps_removed()
+                            .aggregate()
+                            .sorted_for_display()
+                            .iter()
+                        {
+                            println!("{}", m);
+                        }
                     }
                 }
             }
-            render_table(all_measurements, &targets)
+            render_table(all_measurements, &args.targets)
         }
         DisplayFormat::GhJson => print_measurements_json(all_measurements),
     }

bench-vortex/src/public_bi.rs

@@ -42,7 +42,7 @@ pub static PBI_DATASETS: LazyLock<PBIDatasets> = LazyLock::new(|| {
 });
 
 #[derive(Copy, Clone, Debug, Eq, PartialEq, Hash, ValueEnum)]
-#[clap(rename_all = "PascalCase")]
+#[clap(rename_all = "LowerCase")]
 pub enum PBIDataset {
     Arade,
     Bimbo,
@@ -312,7 +312,9 @@ impl PBIData {
                     public_bi_csv_to_parquet_file(table, csv, &output_path).await
                 })
                 .await
-                .vortex_expect("failed to create parquet file");
+                .vortex_expect(
+                    "failed to create parquet file, either the file or duckdb is missing",
+                );
                 let pq_size = parquet_file.metadata().unwrap().size();
                 info!(
                     "Parquet size: {}, {}B",

vortex-array/src/arrays/extension/compute/to_arrow.rs

@@ -24,7 +24,9 @@ impl ToArrowFn<&ExtensionArray> for ExtensionEncoding {
         // NOTE(ngates): this is really gross... but I guess it's ok given how tightly integrated
         //  we are with Arrow.
         if is_temporal_ext_type(array.id()) {
-            temporal_to_arrow(TemporalArray::try_from(array.to_array())?).map(Some)
+            // TODO(joe): push this cast into `temporal_to_arrow`
+            let arrow = temporal_to_arrow(TemporalArray::try_from(array.to_array())?)?;
+            Ok(arrow_cast::cast(&arrow, data_type).map(Some)?)
         } else {
             // Convert storage array directly into arrow, losing type information
             // that will let us round-trip.