feat: wire vortex-native into duckdb-bench

Signed-off-by: Nemo Yu <zyu379@wisc.edu>

Author: HarukiMoriarty

Cargo.lock

@@ -111,10 +111,9 @@ pub fn format_to_df_format(format: Format) -> Arc<dyn FileFormat> {
         Format::Csv => Arc::new(CsvFormat::default()) as _,
         Format::Arrow => Arc::new(ArrowFormat),
         Format::Parquet => Arc::new(ParquetFormat::new()),
-        Format::OnDiskVortex | Format::VortexCompact => Arc::new(VortexFormat::new_with_options(
-            SESSION.clone(),
-            vortex_table_options(),
-        )),
+        Format::OnDiskVortex | Format::VortexCompact | Format::VortexNative => Arc::new(
+            VortexFormat::new_with_options(SESSION.clone(), vortex_table_options()),
+        ),
         Format::OnDiskDuckDB | Format::Lance => {
             unimplemented!("Format {format} cannot be turned into a DataFusion `FileFormat`")
         }

benchmarks/datafusion-bench/src/lib.rs

@@ -169,7 +169,10 @@ impl DuckClient {
         file_format: Format,
     ) -> Result<()> {
         let object_type = match file_format {
-            Format::Parquet | Format::OnDiskVortex | Format::VortexCompact => "VIEW",
+            Format::Parquet
+            | Format::OnDiskVortex
+            | Format::VortexCompact
+            | Format::VortexNative => "VIEW",
             Format::OnDiskDuckDB => "TABLE",
             Format::Lance => {
                 anyhow::bail!(

benchmarks/duckdb-bench/src/lib.rs

@@ -142,6 +142,7 @@ fn main() -> anyhow::Result<()> {
                     // OnDiskDuckDB tables are created during register_tables by loading from Parquet
                     _ => {}
                 }
+                benchmark.prepare_format(format, &base_path).await?;
             }
 
             anyhow::Ok(())
@@ -197,7 +198,8 @@ fn main() -> anyhow::Result<()> {
             if !args.reuse {
                 ctx.reopen()?;
             }
-            ctx.execute_query_result(query)
+            let query = benchmark.query_for_format(query, format);
+            ctx.execute_query_result(&query)
         },
     )?;
 

benchmarks/duckdb-bench/src/main.rs

@@ -24,6 +24,7 @@ vortex = { workspace = true, features = [
     "zstd",
 ] }
 vortex-tensor = { workspace = true } # TODO(connor): In the future, this might be inside vortex.
+vortex-geo = { workspace = true }
 
 anyhow = { workspace = true }
 arrow-array = { workspace = true }
@@ -33,6 +34,8 @@ async-trait = { workspace = true }
 bzip2 = { workspace = true }
 clap = { workspace = true, features = ["derive"] }
 futures = { workspace = true }
+geoarrow = { workspace = true }
+geoarrow-cast = { workspace = true }
 get_dir = { workspace = true }
 glob = { workspace = true }
 humansize = { workspace = true }

vortex-bench/Cargo.toml

@@ -3,6 +3,8 @@
 
 //! Core benchmark trait and types.
 
+use std::path::Path;
+
 use arrow_schema::Schema;
 use glob::Pattern;
 use url::Url;
@@ -33,6 +35,11 @@ pub trait Benchmark: Send + Sync {
     /// Get all available queries for this benchmark
     fn queries(&self) -> anyhow::Result<Vec<(usize, String)>>;
 
+    /// Adapt a query to a specific storage `format` before execution. Default: unchanged.
+    fn query_for_format(&self, query: &str, _format: Format) -> String {
+        query.to_string()
+    }
+
     /// SQL an `engine` must run before this benchmark's queries (e.g. loading engine
     /// extensions). Runners replay these after every (re)open. Default: none.
     fn engine_init_sql(&self, _engine: Engine) -> Vec<String> {
@@ -47,6 +54,13 @@ pub trait Benchmark: Send + Sync {
     /// call this method to ensure base data exists, then perform their own format conversion.
     async fn generate_base_data(&self) -> anyhow::Result<()>;
 
+    /// Prepare benchmark- and format-specific data beyond the Parquet base that
+    /// [`Benchmark::generate_base_data`] produced. Called once per requested format, after the base
+    /// data exists. Default: nothing.
+    async fn prepare_format(&self, _format: Format, _base_path: &Path) -> anyhow::Result<()> {
+        Ok(())
+    }
+
     /// Get expected row counts for validation (optional)
     /// If None, no validation will be performed
     fn expected_row_counts(&self) -> Option<Vec<usize>> {
@@ -80,4 +94,10 @@ pub trait Benchmark: Send + Sync {
         _ = format;
         None
     }
+
+    /// SQL projection substituted into `SELECT {..} FROM read_<fmt>(..)` when registering
+    /// `table_name` as a DuckDB view. Defaults to `*`.
+    fn view_projection(&self, _table_name: &str, _format: Format) -> String {
+        "*".to_string()
+    }
 }

vortex-bench/src/benchmark.rs

@@ -76,8 +76,11 @@ use vortex::session::VortexSession;
 #[global_allocator]
 static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
 
-pub static SESSION: LazyLock<VortexSession> =
-    LazyLock::new(|| VortexSession::default().with_tokio());
+pub static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
+    let session = VortexSession::default().with_tokio();
+    vortex_geo::initialize(&session);
+    session
+});
 
 #[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, Serialize, Deserialize)]
 pub struct Target {
@@ -146,6 +149,9 @@ pub enum Format {
     #[clap(name = "vortex-compact")]
     #[serde(rename = "vortex-compact")]
     VortexCompact,
+    #[clap(name = "vortex-native")]
+    #[serde(rename = "vortex-native")]
+    VortexNative,
     #[clap(name = "duckdb")]
     #[serde(rename = "duckdb")]
     OnDiskDuckDB,
@@ -185,6 +191,7 @@ impl Format {
             Format::Parquet => "parquet",
             Format::OnDiskVortex => "vortex-file-compressed",
             Format::VortexCompact => "vortex-compact",
+            Format::VortexNative => "vortex-native",
             Format::OnDiskDuckDB => "duckdb",
             Format::Lance => "lance",
         }
@@ -197,6 +204,7 @@ impl Format {
             Format::Parquet => "parquet",
             Format::OnDiskVortex => "vortex",
             Format::VortexCompact => "vortex",
+            Format::VortexNative => "vortex",
             Format::OnDiskDuckDB => "duckdb",
             Format::Lance => "lance",
         }
@@ -451,8 +459,16 @@ where
             object_type.to_lowercase()
         );
 
+        let projection = benchmark.view_projection(name, load_format);
+        // SpatialBench's native and WKB lanes both register `trip` from the same db path but with different casts —
+        // so always replace views (cheap, metadata-only). Tables hold materialized data: keep them.
+        let create = if object_type == "VIEW" {
+            format!("CREATE OR REPLACE VIEW {name}")
+        } else {
+            format!("CREATE {object_type} IF NOT EXISTS {name}")
+        };
         sql_statements.push(format!(
-            "CREATE {object_type} IF NOT EXISTS {name} AS SELECT * FROM read_{extension}('{base_dir}/{pattern}');\n",
+            "{create} AS SELECT {projection} FROM read_{extension}('{base_dir}/{pattern}');\n",
         ));
     }
 

vortex-bench/src/lib.rs

@@ -4,6 +4,7 @@
 //! SpatialBench benchmark implementation
 
 use std::fs;
+use std::path::Path;
 
 use url::Url;
 
@@ -13,9 +14,13 @@ use crate::Engine;
 use crate::Format;
 use crate::TableSpec;
 use crate::spatialbench::datagen;
+use crate::spatialbench::datagen::Table;
 use crate::utils::file::resolve_data_url;
 use crate::workspace_root;
 
+/// Data-dir subfolder for the native-geometry Vortex files (the `vortex-native` lane).
+pub const NATIVE_DIR: &str = "vortex-native";
+
 /// SpatialBench geospatial benchmark (Apache Sedona): a `trip` point table, `building` polygons, and
 /// a `customer` attribute table, queried with spatial filters and joins. `zone` polygons are sourced
 /// externally and registered when present. See <https://sedona.apache.org/spatialbench/>.
@@ -34,6 +39,21 @@ impl SpatialBenchBenchmark {
             scale_factor,
         })
     }
+
+    /// Tables to materialize and register: the in-process base tables (`trip`, `building`,
+    /// `customer`) plus the externally-sourced `zone` when its parquet is present. Shared by native
+    /// data-gen and table registration so both lanes cover the same set.
+    fn base_tables(&self) -> Vec<Table> {
+        let mut tables = vec![Table::Trip, Table::Building, Table::Customer];
+        let zone_present = match self.data_url.to_file_path() {
+            Ok(base) => zone_parquet_present(&base.join(Format::Parquet.name())),
+            Err(()) => true,
+        };
+        if zone_present {
+            tables.push(Table::Zone);
+        }
+        tables
+    }
 }
 
 #[async_trait::async_trait]
@@ -58,6 +78,16 @@ impl Benchmark for SpatialBenchBenchmark {
             .collect())
     }
 
+    /// On the `vortex-native` lane, geometry columns surface as `GEOMETRY`, so drop the
+    /// `ST_GeomFromWKB(..)` wrappers and let DuckDB's `spatial` extension evaluate the `ST_*`
+    /// predicates directly on the native geometry.
+    fn query_for_format(&self, query: &str, format: Format) -> String {
+        match format {
+            Format::VortexNative => strip_wkb_wrappers(query),
+            _ => query.to_string(),
+        }
+    }
+
     async fn generate_base_data(&self) -> anyhow::Result<()> {
         if self.data_url.scheme() != "file" {
             return Ok(());
@@ -66,14 +96,37 @@ impl Benchmark for SpatialBenchBenchmark {
             .data_url
             .to_file_path()
             .map_err(|_| anyhow::anyhow!("Invalid file URL: {}", self.data_url.as_str()))?;
-        datagen::generate_tables(&self.scale_factor, base_data_dir).await?;
+        datagen::generate_tables(&self.scale_factor, base_data_dir.clone()).await?;
+        Ok(())
+    }
+
+    /// The `vortex-native` lane decodes each table's WKB geometry to native GeoArrow once, into the
+    /// `vortex-native` dir, so its queries read DuckDB `GEOMETRY` directly. Idempotent.
+    async fn prepare_format(&self, format: Format, base_path: &Path) -> anyhow::Result<()> {
+        if format == Format::VortexNative {
+            let parquet_dir = base_path.join(Format::Parquet.name());
+            let native_dir = base_path.join(NATIVE_DIR);
+            for table in self.base_tables() {
+                datagen::write_native_vortex(table, &parquet_dir, &native_dir).await?;
+            }
+        }
         Ok(())
     }
 
     fn data_url(&self) -> &Url {
         &self.data_url
     }
 
+    /// The `vortex-native` lane reads the native-geometry Vortex dir; every other format reads its
+    /// own `{format}` subfolder.
+    fn format_path(&self, format: Format, base_url: &Url) -> anyhow::Result<Url> {
+        let dir = match format {
+            Format::VortexNative => NATIVE_DIR,
+            other => other.name(),
+        };
+        Ok(base_url.join(&format!("{dir}/"))?)
+    }
+
     fn expected_row_counts(&self) -> Option<Vec<usize>> {
         // Indexed by `query_idx` (1-based), so index 0 is a dummy and Q1's count is at index 1 (TPC-H
         // convention). Only SF1.0 and SF10.0 are validated (like TPC-H); other scale factors return
@@ -101,22 +154,32 @@ impl Benchmark for SpatialBenchBenchmark {
         format!("spatialbench(sf={})", self.scale_factor)
     }
 
+    /// Both lanes register the same tables (WKB reads `parquet`/`vortex`, native reads
+    /// `vortex-native`); `zone` is externally sourced and optional, registered only when present.
     fn table_specs(&self) -> Vec<TableSpec> {
-        let mut specs = vec![
-            TableSpec::new("trip", None),
-            TableSpec::new("building", None),
-            TableSpec::new("customer", None),
-        ];
-        // `zone` is externally sourced and optional; register it only when present so queries that
-        // don't need it don't fail on the missing glob.
-        let zone_present = match self.data_url.to_file_path() {
-            Ok(base) => zone_parquet_present(&base.join(Format::Parquet.name())),
-            Err(()) => true,
-        };
-        if zone_present {
-            specs.push(TableSpec::new("zone", None));
+        self.base_tables()
+            .iter()
+            .map(|table| TableSpec::new(table.name(), None))
+            .collect()
+    }
+
+    /// DuckDB's view star-expansion drops native `GEOMETRY` columns down to `BLOB`, so `ST_*` fail to
+    /// bind. Re-cast every geometry column back to `GEOMETRY` in the view's projection.
+    fn view_projection(&self, table_name: &str, format: Format) -> String {
+        if format == Format::VortexNative
+            && let Some(table) = Table::from_name(table_name)
+        {
+            let geometry_columns = table.geometry_columns();
+            if !geometry_columns.is_empty() {
+                let casts = geometry_columns
+                    .iter()
+                    .map(|column| format!("{name}::GEOMETRY AS {name}", name = column.name))
+                    .collect::<Vec<_>>()
+                    .join(", ");
+                return format!("* REPLACE ({casts})");
+            }
         }
-        specs
+        "*".to_string()
     }
 
     /// Scope each table to its own `{table}_*.{ext}` files; the default globs every file in the
@@ -141,8 +204,33 @@ impl Benchmark for SpatialBenchBenchmark {
 
 /// Whether an externally-sourced `zone_*.parquet` exists under `parquet_dir` (generated by the
 /// upstream `spatialbench-cli`; see the module docs).
-fn zone_parquet_present(parquet_dir: &std::path::Path) -> bool {
+fn zone_parquet_present(parquet_dir: &Path) -> bool {
     glob::glob(&parquet_dir.join("zone_*.parquet").to_string_lossy())
         .map(|mut paths| paths.next().is_some())
         .unwrap_or(false)
 }
+
+/// Strip `ST_GeomFromWKB(<inner>)` → `<inner>` so the native lane reads the already-`GEOMETRY`
+/// column directly. Assumes the wrapped expression contains no inner `)` (true for our column refs).
+fn strip_wkb_wrappers(sql: &str) -> String {
+    const OPEN: &str = "ST_GeomFromWKB(";
+    let mut out = String::with_capacity(sql.len());
+    let mut rest = sql;
+    while let Some(pos) = rest.find(OPEN) {
+        out.push_str(&rest[..pos]);
+        let after = &rest[pos + OPEN.len()..];
+        match after.find(')') {
+            Some(close) => {
+                out.push_str(&after[..close]);
+                rest = &after[close + 1..];
+            }
+            // Unbalanced wrapper: emit it verbatim and stop rewriting.
+            None => {
+                out.push_str(OPEN);
+                rest = after;
+            }
+        }
+    }
+    out.push_str(rest);
+    out
+}

vortex-bench/src/spatialbench/benchmark.rs

@@ -1,11 +1,14 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-//! SpatialBench data preparation. [`wkb`] generates the canonical WKB base tables (Parquet + Vortex);
-//! the [`table`] catalog is the single source of truth for the base tables.
+//! SpatialBench data preparation. [`wkb`] generates the canonical WKB base tables; [`native`] derives
+//! native-geometry Vortex files from them for `points=native`. The [`table`] catalog is the single
+//! source of truth for the base tables both stages share.
 
+pub mod native;
 pub mod table;
 pub mod wkb;
 
+pub use native::write_native_vortex;
 pub use table::Table;
 pub use wkb::generate_tables;