From b36255172bb08c9c50ee6913a8b654af866b3f45 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Fri, 26 Jun 2026 13:33:32 -0400 Subject: [PATCH 1/7] feat: wire in the vx-benchmark Signed-off-by: Nemo Yu --- bench-orchestrator/bench_orchestrator/cli.py | 8 ++- .../bench_orchestrator/config.py | 27 +++++++- bench-orchestrator/tests/test_config.py | 26 ++++++++ vortex-bench/src/conversions.rs | 62 +++++++++++++++++- vortex-bench/src/spatialbench/datagen/wkb.rs | 63 +++++++++++++++++++ 5 files changed, 177 insertions(+), 9 deletions(-) diff --git a/bench-orchestrator/bench_orchestrator/cli.py b/bench-orchestrator/bench_orchestrator/cli.py index a9e37e70309..b9d7f9bb2ef 100644 --- a/bench-orchestrator/bench_orchestrator/cli.py +++ b/bench-orchestrator/bench_orchestrator/cli.py @@ -85,9 +85,11 @@ def run_ref_auto_complete() -> list[str]: return list(map(lambda x: x.run_id, ResultStore().list_runs(limit=None))) -def targets_from_axes(engine: str, format: str) -> tuple[list[BenchmarkTarget], list[str]]: +def targets_from_axes( + engine: str, format: str, benchmark: Benchmark | None = None +) -> tuple[list[BenchmarkTarget], list[str]]: """Resolve legacy engine/format axes into explicit benchmark targets.""" - return resolve_axis_targets(parse_engines(engine), parse_formats(format)) + return resolve_axis_targets(parse_engines(engine), parse_formats(format), benchmark) def backends_for_engines(engines: list[Engine]) -> list[Engine]: @@ -260,7 +262,7 @@ def run( targets = parse_targets_json(targets_json) warnings: list[str] = [] else: - targets, warnings = targets_from_axes(engine, format) + targets, warnings = targets_from_axes(engine, format, benchmark) except ValueError as exc: console.print(f"[red]{exc}[/red]") raise typer.Exit(1) from exc diff --git a/bench-orchestrator/bench_orchestrator/config.py b/bench-orchestrator/bench_orchestrator/config.py index c597e84c6be..e358bf18f01 100644 --- a/bench-orchestrator/bench_orchestrator/config.py +++ b/bench-orchestrator/bench_orchestrator/config.py @@ -52,6 +52,7 @@ class Benchmark(Enum): POLARSIGNALS = "polarsignals" PUBLIC_BI = "public-bi" STATPOPGEN = "statpopgen" + SPATIALBENCH = "spatialbench" # Engine to supported formats mapping. @@ -72,6 +73,19 @@ class Benchmark(Enum): Engine.LANCE: [Format.LANCE], } +# Engines each benchmark can run on. Benchmarks default to *every* engine; list one here only to +# restrict it. SpatialBench's queries use DuckDB-specific `ST_*` spatial SQL that DataFusion has no +# functions for yet. +BENCHMARK_ENGINES: dict[Benchmark, frozenset[Engine]] = { + Benchmark.SPATIALBENCH: frozenset({Engine.DUCKDB}), +} + + +def engines_for_benchmark(benchmark: Benchmark) -> frozenset[Engine]: + """Return the engines `benchmark` supports, defaulting to every engine when unrestricted.""" + return BENCHMARK_ENGINES.get(benchmark, frozenset(Engine)) + + T = TypeVar("T") @@ -175,13 +189,16 @@ def parse_formats_json(value: str) -> list[Format]: def resolve_axis_targets( - engines: Iterable[Engine], formats: Iterable[Format] + engines: Iterable[Engine], formats: Iterable[Format], benchmark: Benchmark | None = None ) -> tuple[list[BenchmarkTarget], list[str]]: """Expand engine/format axes into supported explicit targets.""" warnings: list[str] = [] targets: list[BenchmarkTarget] = [] for engine in engines: + if benchmark is not None and engine not in engines_for_benchmark(benchmark): + warnings.append(f"Benchmark {benchmark.value} does not support engine {engine.value}") + continue for fmt in formats: target = BenchmarkTarget(engine=engine, format=fmt).normalized() if not target.is_supported(): @@ -200,7 +217,9 @@ def group_targets_by_backend(targets: Iterable[BenchmarkTarget]) -> dict[Engine, return groups -def validate_targets(targets: Iterable[BenchmarkTarget], options: dict[str, str]) -> list[str]: +def validate_targets( + targets: Iterable[BenchmarkTarget], options: dict[str, str], benchmark: Benchmark | None = None +) -> list[str]: """Validate explicit targets against benchmark runner constraints.""" errors: list[str] = [] @@ -208,6 +227,8 @@ def validate_targets(targets: Iterable[BenchmarkTarget], options: dict[str, str] for target in normalized_targets: if not target.is_supported(): errors.append(f"Format {target.format.value} is not supported by engine {target.engine.value}") + if benchmark is not None and target.engine not in engines_for_benchmark(benchmark): + errors.append(f"Benchmark {benchmark.value} does not support engine {target.engine.value}") if options.get("remote-data-dir") and any(target.format == Format.LANCE for target in normalized_targets): errors.append("Lance format is not supported for remote storage benchmarks.") @@ -242,7 +263,7 @@ def backends(self) -> list[Engine]: def validate(self) -> list[str]: """Validate the configuration and return any errors.""" - return validate_targets(self.targets, self.options) + return validate_targets(self.targets, self.options, self.benchmark) @dataclass diff --git a/bench-orchestrator/tests/test_config.py b/bench-orchestrator/tests/test_config.py index c7e2d6bb291..f900048f87b 100644 --- a/bench-orchestrator/tests/test_config.py +++ b/bench-orchestrator/tests/test_config.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright the Vortex contributors from bench_orchestrator.config import ( + Benchmark, BenchmarkTarget, Engine, Format, @@ -39,6 +40,31 @@ def test_resolve_axis_targets_filters_unsupported_combinations() -> None: assert warnings == ["Format arrow is not supported by engine duckdb"] +def test_resolve_axis_targets_skips_engines_a_benchmark_cannot_run() -> None: + # SpatialBench is DuckDB-only (ST_* spatial SQL), so the DataFusion axis is dropped with a warning. + targets, warnings = resolve_axis_targets( + [Engine.DATAFUSION, Engine.DUCKDB], + [Format.PARQUET, Format.VORTEX], + Benchmark.SPATIALBENCH, + ) + + assert targets == [ + BenchmarkTarget(engine=Engine.DUCKDB, format=Format.PARQUET), + BenchmarkTarget(engine=Engine.DUCKDB, format=Format.VORTEX), + ] + assert warnings == ["Benchmark spatialbench does not support engine datafusion"] + + +def test_validate_targets_rejects_engine_a_benchmark_cannot_run() -> None: + errors = validate_targets( + [BenchmarkTarget(engine=Engine.DATAFUSION, format=Format.PARQUET)], + {}, + Benchmark.SPATIALBENCH, + ) + + assert errors == ["Benchmark spatialbench does not support engine datafusion"] + + def test_validate_targets_rejects_remote_lance() -> None: errors = validate_targets( [BenchmarkTarget(engine=Engine.DATAFUSION, format=Format.LANCE)], diff --git a/vortex-bench/src/conversions.rs b/vortex-bench/src/conversions.rs index ad84ef61b1d..6fa12c21b76 100644 --- a/vortex-bench/src/conversions.rs +++ b/vortex-bench/src/conversions.rs @@ -4,6 +4,7 @@ use std::fs; use std::path::Path; use std::path::PathBuf; +use std::sync::Arc; use futures::StreamExt; use futures::TryStreamExt; @@ -33,14 +34,22 @@ use vortex::array::arrow::FromArrowArray; use vortex::array::builders::builder_with_capacity; use vortex::array::stream::ArrayStreamAdapter; use vortex::array::stream::ArrayStreamExt; +use vortex::compressor::BtrBlocksCompressorBuilder; use vortex::dtype::DType; +use vortex::dtype::FieldPath; use vortex::dtype::StructFields; use vortex::dtype::arrow::FromArrowType; use vortex::dtype::extension::ExtDType; use vortex::dtype::extension::ExtDTypeRef; use vortex::error::VortexResult; use vortex::error::vortex_err; +use vortex::file::VortexWriteOptions; use vortex::file::WriteOptionsSessionExt; +use vortex::file::WriteStrategyBuilder; +use vortex::layout::LayoutStrategy; +use vortex::layout::layouts::chunked::writer::ChunkedLayoutStrategy; +use vortex::layout::layouts::compressed::CompressingStrategy; +use vortex::layout::layouts::flat::writer::FlatLayoutStrategy; use vortex::session::VortexSession; use vortex::utils::aliases::hash_set::HashSet; use vortex_geo::extension::GeoMetadata; @@ -142,8 +151,7 @@ pub async fn convert_parquet_file_to_vortex( .open(output_path) .await?; - compaction - .apply_options(SESSION.write_options()) + write_options_for(compaction, &dtype, is_spatialbench(parquet_path)) .write( &mut output_file, ArrayStreamExt::boxed(ArrayStreamAdapter::new(dtype, stream)), @@ -153,6 +161,54 @@ pub async fn convert_parquet_file_to_vortex( Ok(()) } +/// Whether `path` points at SpatialBench data. +fn is_spatialbench(path: &Path) -> bool { + path.components() + .any(|component| component.as_os_str() == "spatialbench") +} + +/// Vortex write options for converting `dtype`-shaped data. +/// +/// For SpatialBench (`skip_binary_dict`), the geometry blobs are large and +/// unique, so the dictionary builder balloons memory (tens of GB) for zero gain. +fn write_options_for( + compaction: CompactionStrategy, + dtype: &DType, + skip_binary_dict: bool, +) -> VortexWriteOptions { + let binary_fields: Vec<_> = match dtype { + DType::Struct(fields, _) if skip_binary_dict => fields + .names() + .iter() + .zip(fields.fields()) + .filter(|(_, field)| matches!(field, DType::Binary(_))) + .map(|(name, _)| name.clone()) + .collect(), + _ => Vec::new(), + }; + if binary_fields.is_empty() { + return compaction.apply_options(SESSION.write_options()); + } + + let mut builder = WriteStrategyBuilder::default(); + if matches!(compaction, CompactionStrategy::Compact) { + builder = + builder.with_btrblocks_builder(BtrBlocksCompressorBuilder::default().with_compact()); + } + for name in binary_fields { + builder = builder.with_field_writer(FieldPath::from_name(name), no_dict_layout()); + } + SESSION.write_options().with_strategy(builder.build()) +} + +/// A chunked + compressed layout that skips dictionary encoding for opaque `Binary` blobs. +fn no_dict_layout() -> Arc { + Arc::new(CompressingStrategy::new( + ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()), + BtrBlocksCompressorBuilder::default().build(), + )) +} + /// Convert all Parquet files in a directory to Vortex format. /// /// This function reads Parquet files from `{input_path}/parquet/` and writes Vortex files to @@ -251,7 +307,7 @@ pub async fn add_geoparquet_metadata(parquet_path: &Path, geo_json: &str) -> any return Ok(()); } - let schema = std::sync::Arc::clone(builder.schema()); + let schema = Arc::clone(builder.schema()); let mut reader = builder.build()?; let tmp_path = parquet_path.with_extension("parquet.tmp"); diff --git a/vortex-bench/src/spatialbench/datagen/wkb.rs b/vortex-bench/src/spatialbench/datagen/wkb.rs index 422505a3623..c8d5e968c6f 100644 --- a/vortex-bench/src/spatialbench/datagen/wkb.rs +++ b/vortex-bench/src/spatialbench/datagen/wkb.rs @@ -5,9 +5,11 @@ //! Geometry is emitted as WKB, which DuckDB reads directly as `GEOMETRY` via `ST_GeomFromWKB`. use std::fs; +use std::path::Path; use std::path::PathBuf; use std::sync::Arc; +use anyhow::Context; use anyhow::Result; // spatialbench emits arrow-56 batches, so they must be written with its matching arrow-56 // parquet crate, not the workspace's arrow-58 one. The parquet file itself is version-neutral. @@ -23,7 +25,9 @@ use spatialbench_parquet::basic::Compression; use spatialbench_parquet::file::properties::WriterProperties; use spatialbench_parquet::format::KeyValue; use tokio::fs::File as TokioFile; +use tokio::process::Command; use tracing::info; +use tracing::warn; use super::table::Table; use crate::Format; @@ -109,6 +113,65 @@ pub async fn generate_tables(scale_factor: &str, output_dir: PathBuf) -> Result< } } + // `zone` isn't generated in-process (`Table::is_generated` is false); it comes from the upstream + // CLI. Best-effort: a missing/failed CLI shouldn't block the zone-free queries, so warn and + // carry on. + if let Err(e) = generate_zone(scale_factor, &parquet_dir).await { + warn!( + error = %e, + "zone table not generated — SpatialBench queries Q2/Q4/Q6/Q10/Q11 need it. Install the \ + upstream generator (`cargo install --path /spatialbench-cli`) or \ + set SPATIALBENCH_CLI to its binary, then re-run." + ); + } + + Ok(()) +} + +/// Generate the externally-sourced `zone` table by shelling out to the upstream `spatialbench-cli`. +async fn generate_zone(scale_factor: f64, parquet_dir: &Path) -> Result<()> { + if parquet_dir.join("zone_0.parquet").exists() { + return Ok(()); + } + let cli = std::env::var("SPATIALBENCH_CLI").unwrap_or_else(|_| "spatialbench-cli".to_string()); + + // Generate into a scratch dir so the CLI's `zone.parquet` name can't collide with the base + // tables, then move the produced parts into place as `zone_{part}.parquet`. + // Start from an empty scratch dir (clear any leftover from an interrupted run). + let scratch = parquet_dir.join(".zone-scratch"); + fs::remove_dir_all(&scratch).ok(); + fs::create_dir_all(&scratch)?; + + info!( + scale_factor, + cli, "Generating SpatialBench zone table via spatialbench-cli" + ); + let status = Command::new(&cli) + .arg("-s") + .arg(scale_factor.to_string()) + .args(["-T", "zone", "-f", "parquet", "-o"]) + .arg(&scratch) + .status() + .await + .with_context(|| format!("failed to spawn `{cli}` (is it installed / on PATH?)"))?; + anyhow::ensure!( + status.success(), + "`{cli}` exited with {status} while generating zone" + ); + + // The CLI writes `zone.parquet` (single part) or `zone/zone.N.parquet`. + let mut produced: Vec = glob::glob(&scratch.join("**/*.parquet").to_string_lossy())? + .collect::>()?; + produced.sort(); + anyhow::ensure!( + !produced.is_empty(), + "`{cli}` produced no zone parquet under {}", + scratch.display() + ); + for (part_idx, src) in produced.iter().enumerate() { + fs::rename(src, parquet_dir.join(format!("zone_{part_idx}.parquet")))?; + } + fs::remove_dir_all(&scratch).ok(); Ok(()) } From 1a8001be8a6194f93461393a37e758817f775ac9 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Wed, 1 Jul 2026 10:09:54 -0400 Subject: [PATCH 2/7] feat(vortex-bench): store SpatialBench geometry as little-endian WKB DuckDB's GEOMETRY only accepts little-endian (NDR) WKB, but the externally sourced zone table (Overture Maps via spatialbench-cli) is big-endian, so the vortex lane failed spatial queries with "Only little-endian WKB is supported". Re-encode geometry columns to little-endian during the parquet->vortex conversion so the vortex file stores canonical little-endian WKB; columns that are already little-endian pass through without a copy. Also drop the best-effort warn around zone generation so data-generation failures propagate. Signed-off-by: Nemo Yu --- Cargo.lock | 1 + vortex-bench/Cargo.toml | 1 + vortex-bench/src/conversions.rs | 65 ++++++++++++++++---- vortex-bench/src/spatialbench/datagen/wkb.rs | 13 +--- 4 files changed, 57 insertions(+), 23 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8de4efe05fc..ddf5b6e939a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9919,6 +9919,7 @@ dependencies = [ "vortex", "vortex-geo", "vortex-tensor", + "wkb", ] [[package]] diff --git a/vortex-bench/Cargo.toml b/vortex-bench/Cargo.toml index 9cd7fc92d30..f655b5213bb 100644 --- a/vortex-bench/Cargo.toml +++ b/vortex-bench/Cargo.toml @@ -69,6 +69,7 @@ tracing-subscriber = { workspace = true, features = [ ] } url = { workspace = true } uuid = { workspace = true, features = ["v4"] } +wkb = { workspace = true } [dev-dependencies] insta = { workspace = true } diff --git a/vortex-bench/src/conversions.rs b/vortex-bench/src/conversions.rs index 6fa12c21b76..b7c367ef8f6 100644 --- a/vortex-bench/src/conversions.rs +++ b/vortex-bench/src/conversions.rs @@ -23,12 +23,14 @@ use tracing::info; use tracing::trace; use vortex::VortexSessionDefault; use vortex::array::ArrayRef; +use vortex::array::ExecutionCtx; use vortex::array::IntoArray; use vortex::array::VortexSessionExecute; use vortex::array::arrays::ChunkedArray; use vortex::array::arrays::ExtensionArray; use vortex::array::arrays::Struct; use vortex::array::arrays::StructArray; +use vortex::array::arrays::VarBinViewArray; use vortex::array::arrays::struct_::StructArrayExt; use vortex::array::arrow::FromArrowArray; use vortex::array::builders::builder_with_capacity; @@ -54,6 +56,10 @@ use vortex::session::VortexSession; use vortex::utils::aliases::hash_set::HashSet; use vortex_geo::extension::GeoMetadata; use vortex_geo::extension::WellKnownBinary; +use wkb::Endianness; +use wkb::reader::read_wkb; +use wkb::writer::WriteOptions; +use wkb::writer::write_geometry; use crate::CompactionStrategy; use crate::Format; @@ -370,7 +376,8 @@ fn tag_geo_dtype(dtype: DType, geo_columns: &HashSet) -> VortexResult
) -> VortexResult { if geo_columns.is_empty() { return Ok(chunk); @@ -381,17 +388,51 @@ fn tag_geo_array(chunk: ArrayRef, geo_columns: &HashSet) -> VortexResult let names = struct_array.names().clone(); let validity = struct_array.struct_validity(); let len = struct_array.len(); - let tagged = names - .iter() - .zip(struct_array.iter_unmasked_fields()) - .map(|(name, field)| { - if geo_columns.contains(name.as_ref()) && field.dtype().is_binary() { - let ext = wkb_ext_dtype(field.dtype())?; - Ok(ExtensionArray::try_new(ext, field.clone())?.into_array()) - } else { - Ok(field.clone()) + let mut ctx = SESSION.create_execution_ctx(); + let mut tagged = Vec::with_capacity(names.len()); + for (name, field) in names.iter().zip(struct_array.iter_unmasked_fields()) { + if geo_columns.contains(name.as_ref()) && field.dtype().is_binary() { + let ext = wkb_ext_dtype(field.dtype())?; + let little_endian = wkb_field_to_little_endian(field, &mut ctx)?; + tagged.push(ExtensionArray::try_new(ext, little_endian)?.into_array()); + } else { + tagged.push(field.clone()); + } + } + Ok(StructArray::try_new(names, tagged, len, validity)?.into_array()) +} + +/// Re-encode a binary `field`'s WKB values as little-endian (NDR), the only byte order DuckDB's +/// `GEOMETRY` accepts. Byte order is uniform within a column, so the array is returned untouched (no +/// copy) unless its first non-null value is big-endian; only then is each value re-encoded. +fn wkb_field_to_little_endian(field: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let values = field.clone().execute::(ctx)?; + let len = values.len(); + let validity = values.validity()?.execute_mask(len, ctx)?; + let is_big_endian = (0..len) + .find(|&i| validity.value(i)) + .is_some_and(|i| values.bytes_at(i).as_slice().first() == Some(&0x00)); + if !is_big_endian { + return Ok(field.clone()); + } + let little_endian: Vec>> = (0..len) + .map(|i| { + if !validity.value(i) { + return Ok(None); } + let wkb = values.bytes_at(i); + let geometry = read_wkb(wkb.as_slice()).map_err(|e| vortex_err!("invalid WKB: {e}"))?; + let mut encoded = Vec::with_capacity(wkb.len()); + write_geometry( + &mut encoded, + &geometry, + &WriteOptions { + endianness: Endianness::LittleEndian, + }, + ) + .map_err(|e| vortex_err!("re-encoding WKB as little-endian: {e}"))?; + Ok(Some(encoded)) }) - .collect::>>()?; - Ok(StructArray::try_new(names, tagged, len, validity)?.into_array()) + .collect::>()?; + Ok(VarBinViewArray::from_iter_nullable_bin(little_endian).into_array()) } diff --git a/vortex-bench/src/spatialbench/datagen/wkb.rs b/vortex-bench/src/spatialbench/datagen/wkb.rs index c8d5e968c6f..ca29b59e91b 100644 --- a/vortex-bench/src/spatialbench/datagen/wkb.rs +++ b/vortex-bench/src/spatialbench/datagen/wkb.rs @@ -27,7 +27,6 @@ use spatialbench_parquet::format::KeyValue; use tokio::fs::File as TokioFile; use tokio::process::Command; use tracing::info; -use tracing::warn; use super::table::Table; use crate::Format; @@ -114,16 +113,8 @@ pub async fn generate_tables(scale_factor: &str, output_dir: PathBuf) -> Result< } // `zone` isn't generated in-process (`Table::is_generated` is false); it comes from the upstream - // CLI. Best-effort: a missing/failed CLI shouldn't block the zone-free queries, so warn and - // carry on. - if let Err(e) = generate_zone(scale_factor, &parquet_dir).await { - warn!( - error = %e, - "zone table not generated — SpatialBench queries Q2/Q4/Q6/Q10/Q11 need it. Install the \ - upstream generator (`cargo install --path /spatialbench-cli`) or \ - set SPATIALBENCH_CLI to its binary, then re-run." - ); - } + // `spatialbench-cli` (install it, or set `SPATIALBENCH_CLI` to its binary). + generate_zone(scale_factor, &parquet_dir).await?; Ok(()) } From c282925febb6ed137abbcb6a6d52669fd5676383 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Wed, 1 Jul 2026 10:33:18 -0400 Subject: [PATCH 3/7] refactor(vortex-bench): reuse idempotent_async for zone generation The zone table's data generation hand-rolled idempotency: an existence check plus a scratch dir cleared with remove_dir_all to survive interrupted runs. Reuse the shared idempotent_async helper instead, which skips when zone_0.parquet exists and writes via a unique temp path + atomic rename, so interrupted runs leave nothing to clean up. Zone is a single part (the CLI runs without --parts), so it is handled as one file. Signed-off-by: Nemo Yu --- vortex-bench/src/spatialbench/datagen/wkb.rs | 79 +++++++++----------- 1 file changed, 37 insertions(+), 42 deletions(-) diff --git a/vortex-bench/src/spatialbench/datagen/wkb.rs b/vortex-bench/src/spatialbench/datagen/wkb.rs index ca29b59e91b..e0160ea6f9c 100644 --- a/vortex-bench/src/spatialbench/datagen/wkb.rs +++ b/vortex-bench/src/spatialbench/datagen/wkb.rs @@ -119,50 +119,45 @@ pub async fn generate_tables(scale_factor: &str, output_dir: PathBuf) -> Result< Ok(()) } -/// Generate the externally-sourced `zone` table by shelling out to the upstream `spatialbench-cli`. +/// Generate the externally-sourced `zone` table by shelling out to the upstream +/// `spatialbench-cli`. async fn generate_zone(scale_factor: f64, parquet_dir: &Path) -> Result<()> { - if parquet_dir.join("zone_0.parquet").exists() { - return Ok(()); - } let cli = std::env::var("SPATIALBENCH_CLI").unwrap_or_else(|_| "spatialbench-cli".to_string()); - - // Generate into a scratch dir so the CLI's `zone.parquet` name can't collide with the base - // tables, then move the produced parts into place as `zone_{part}.parquet`. - // Start from an empty scratch dir (clear any leftover from an interrupted run). - let scratch = parquet_dir.join(".zone-scratch"); - fs::remove_dir_all(&scratch).ok(); - fs::create_dir_all(&scratch)?; - - info!( - scale_factor, - cli, "Generating SpatialBench zone table via spatialbench-cli" - ); - let status = Command::new(&cli) - .arg("-s") - .arg(scale_factor.to_string()) - .args(["-T", "zone", "-f", "parquet", "-o"]) - .arg(&scratch) - .status() - .await - .with_context(|| format!("failed to spawn `{cli}` (is it installed / on PATH?)"))?; - anyhow::ensure!( - status.success(), - "`{cli}` exited with {status} while generating zone" - ); - - // The CLI writes `zone.parquet` (single part) or `zone/zone.N.parquet`. - let mut produced: Vec = glob::glob(&scratch.join("**/*.parquet").to_string_lossy())? - .collect::>()?; - produced.sort(); - anyhow::ensure!( - !produced.is_empty(), - "`{cli}` produced no zone parquet under {}", - scratch.display() - ); - for (part_idx, src) in produced.iter().enumerate() { - fs::rename(src, parquet_dir.join(format!("zone_{part_idx}.parquet")))?; - } - fs::remove_dir_all(&scratch).ok(); + idempotent_async(parquet_dir.join("zone_0.parquet"), |zone_path| async move { + info!( + scale_factor, + cli, "Generating SpatialBench zone table via spatialbench-cli" + ); + // The CLI writes a fixed `zone.parquet` name into an output directory, so + // generate into a scratch dir and move the produced parquet onto `zone_path`. + let scratch = zone_path.with_extension("zone-scratch"); + fs::create_dir_all(&scratch)?; + let status = Command::new(&cli) + .arg("-s") + .arg(scale_factor.to_string()) + .args(["-T", "zone", "-f", "parquet", "-o"]) + .arg(&scratch) + .status() + .await + .with_context(|| format!("failed to spawn `{cli}` (is it installed / on PATH?)"))?; + anyhow::ensure!( + status.success(), + "`{cli}` exited with {status} while generating zone" + ); + let produced = glob::glob(&scratch.join("**/*.parquet").to_string_lossy())? + .flatten() + .next() + .with_context(|| { + format!( + "`{cli}` produced no zone parquet under {}", + scratch.display() + ) + })?; + fs::rename(&produced, &zone_path)?; + fs::remove_dir_all(&scratch).ok(); + Ok(()) + }) + .await?; Ok(()) } From fd5e606dba6aeb773c587fa1ad8488c312b6b486 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Fri, 26 Jun 2026 17:26:16 -0400 Subject: [PATCH 4/7] feat: support geo multipolygon Signed-off-by: Nemo Yu --- Cargo.lock | 14 + Cargo.toml | 1 + vortex-geo/Cargo.toml | 1 + vortex-geo/src/extension/mod.rs | 4 + vortex-geo/src/extension/multipolygon.rs | 372 +++++++++++++++++++++++ vortex-geo/src/extension/point.rs | 52 +++- vortex-geo/src/extension/polygon.rs | 48 ++- vortex-geo/src/lib.rs | 4 + 8 files changed, 482 insertions(+), 14 deletions(-) create mode 100644 vortex-geo/src/extension/multipolygon.rs diff --git a/Cargo.lock b/Cargo.lock index 0f70678a143..7979d78234f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4238,6 +4238,19 @@ dependencies = [ "wkt 0.14.0", ] +[[package]] +name = "geoarrow-cast" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41c308d653690a4e8ef3cbba69696056bd819e624766ece66d64cc26a638acc1" +dependencies = [ + "arrow-schema 58.3.0", + "geo-traits", + "geoarrow-array", + "geoarrow-schema", + "wkt 0.14.0", +] + [[package]] name = "geoarrow-schema" version = "0.8.0" @@ -10397,6 +10410,7 @@ dependencies = [ "geo-traits", "geo-types", "geoarrow", + "geoarrow-cast", "prost 0.14.4", "rstest", "vortex-array", diff --git a/Cargo.toml b/Cargo.toml index bf5057432a6..ba54a50a2ab 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -165,6 +165,7 @@ geo = "0.31.0" geo-traits = "0.3.0" geo-types = "0.7.19" geoarrow = "0.8.0" +geoarrow-cast = "0.8.0" get_dir = "0.5.0" glob = "0.3.2" goldenfile = "1" diff --git a/vortex-geo/Cargo.toml b/vortex-geo/Cargo.toml index e2f7e4dc10f..2f0583b49e6 100644 --- a/vortex-geo/Cargo.toml +++ b/vortex-geo/Cargo.toml @@ -20,6 +20,7 @@ geo = { workspace = true } geo-traits = { workspace = true } geo-types = { workspace = true } geoarrow = { workspace = true } +geoarrow-cast = { workspace = true } prost = { workspace = true } vortex-array = { workspace = true } vortex-error = { workspace = true } diff --git a/vortex-geo/src/extension/mod.rs b/vortex-geo/src/extension/mod.rs index 684c83bade0..5cccc489297 100644 --- a/vortex-geo/src/extension/mod.rs +++ b/vortex-geo/src/extension/mod.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors pub(crate) mod coordinate; +mod multipolygon; mod point; mod polygon; mod wkb; @@ -12,6 +13,7 @@ use std::sync::Arc; use geo_types::Geometry; use geoarrow::datatypes::Crs; use geoarrow::datatypes::Metadata; +pub use multipolygon::*; pub use point::*; pub use polygon::*; use vortex_array::ArrayRef; @@ -46,6 +48,8 @@ pub(crate) fn geometries( point_geometries(&storage, ctx) } else if ext.is::() { polygon_geometries(&storage, ctx) + } else if ext.is::() { + multipolygon_geometries(&storage, ctx) } else { vortex_bail!("geo: unsupported geometry extension {}", array.dtype()) } diff --git a/vortex-geo/src/extension/multipolygon.rs b/vortex-geo/src/extension/multipolygon.rs new file mode 100644 index 00000000000..82fe081b316 --- /dev/null +++ b/vortex-geo/src/extension/multipolygon.rs @@ -0,0 +1,372 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The [`MultiPolygon`] extension type (`vortex.geo.multipolygon`), stored as +//! `List>>>` (polygons → rings → coordinates) and tagged with +//! [`GeoMetadata`]. A single `Polygon` is a one-element multipolygon. + +use std::sync::Arc; + +use arrow_array::ArrayRef as ArrowArrayRef; +use arrow_schema::DataType; +use arrow_schema::Field; +use arrow_schema::extension::ExtensionType; +use geo_traits::to_geo::ToGeoGeometry; +use geo_types::Geometry; +use geoarrow::array::GeoArrowArray; +use geoarrow::array::GeoArrowArrayAccessor; +use geoarrow::array::IntoArrow; +use geoarrow::array::MultiPolygonArray; +use geoarrow::datatypes::CoordType; +use geoarrow::datatypes::GeoArrowType; +use geoarrow::datatypes::MultiPolygonType; +use geoarrow::datatypes::WkbType; +use geoarrow_cast::cast::cast; +use prost::Message; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::arrow::ArrowExport; +use vortex_array::arrow::ArrowExportVTable; +use vortex_array::arrow::ArrowImport; +use vortex_array::arrow::ArrowImportVTable; +use vortex_array::arrow::ArrowSession; +use vortex_array::arrow::ArrowSessionExt; +use vortex_array::arrow::FromArrowArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::arrow::FromArrowType; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::dtype::extension::ExtId; +use vortex_array::dtype::extension::ExtVTable; +use vortex_array::scalar::ScalarValue; +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_session::registry::CachedId; +use vortex_session::registry::Id; + +use super::GeoMetadata; +use super::coordinate::Dimension; +use super::coordinate::coordinate_dimension; +use super::coordinate::coordinate_storage_dtype; +use super::geo_metadata_from_arrow; +use super::geoarrow_metadata; + +/// A multipolygon (`geoarrow.multipolygon`); a single `Polygon` is a one-element multipolygon. +#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] +pub struct MultiPolygon; + +impl ExtVTable for MultiPolygon { + type Metadata = GeoMetadata; + // No cheap owned value like Point's `Coordinate`; expose the raw storage scalar. + type NativeValue<'a> = &'a ScalarValue; + + fn id(&self) -> ExtId { + ExtId::new_static("vortex.geo.multipolygon") + } + + fn serialize_metadata(&self, metadata: &Self::Metadata) -> VortexResult> { + Ok(metadata.encode_to_vec()) + } + + fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult { + Ok(GeoMetadata::decode(metadata)?) + } + + fn validate_dtype(ext_dtype: &ExtDType) -> VortexResult<()> { + multipolygon_dimension(ext_dtype.storage_dtype()).map(|_| ()) + } + + fn unpack_native<'a>( + _ext_dtype: &'a ExtDType, + storage_value: &'a ScalarValue, + ) -> VortexResult<&'a ScalarValue> { + Ok(storage_value) + } +} + +/// Storage `List>>`: polygons → rings → coordinates. +pub(crate) fn multipolygon_storage_dtype(dim: Dimension, nullability: Nullability) -> DType { + let coords = coordinate_storage_dtype(dim, Nullability::NonNullable); + let ring = DType::List(Arc::new(coords), Nullability::NonNullable); + let polygon = DType::List(Arc::new(ring), Nullability::NonNullable); + DType::List(Arc::new(polygon), nullability) +} + +/// Validate `dtype` is `List>>` and return its [`Dimension`]. +pub(crate) fn multipolygon_dimension(dtype: &DType) -> VortexResult { + let DType::List(polygon, _) = dtype else { + vortex_bail!("multipolygon storage must be a List of polygons, was {dtype}"); + }; + let DType::List(ring, _) = polygon.as_ref() else { + vortex_bail!("multipolygon polygon storage must be a List of rings, was {polygon}"); + }; + let DType::List(coords, _) = ring.as_ref() else { + vortex_bail!("multipolygon ring storage must be a List of coordinates, was {ring}"); + }; + coordinate_dimension(coords) +} + +static ARROW_MULTIPOLYGON: CachedId = CachedId::new(MultiPolygonType::NAME); + +/// The `geoarrow.multipolygon` type for `dimension`, with separated (struct) coordinates. +fn multipolygon_type(geo_metadata: &GeoMetadata, dimension: Dimension) -> MultiPolygonType { + MultiPolygonType::new(dimension.into(), geoarrow_metadata(geo_metadata)) +} + +/// Decode storage to `geo_types` for the geo scalar functions (CRS is irrelevant to planar ops). +pub(crate) fn multipolygon_geometries( + storage: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult>> { + multipolygon_array(storage, ctx)? + .iter() + .map(|geometry| -> VortexResult> { + Ok(geometry + .ok_or_else(|| vortex_err!("geo: null geometry is not supported"))? + .map_err(|e| vortex_err!("geo: geometry access failed: {e}"))? + .to_geometry()) + }) + .collect() +} + +/// Build a geoarrow `MultiPolygonArray` from the `MultiPolygon` storage. +fn multipolygon_array( + storage: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let multipolygon_type = multipolygon_type( + &GeoMetadata::default(), + multipolygon_dimension(storage.dtype())?, + ); + let session = ctx.session().clone(); + let arrow = session.arrow().execute_arrow(storage.clone(), None, ctx)?; + MultiPolygonArray::try_from((arrow.as_ref(), multipolygon_type)) + .map_err(|e| vortex_err!("failed to construct MultiPolygonArray: {e}")) +} + +/// A validated `MultiPolygon` array (`try_from` checks the extension type). +pub struct MultiPolygonData(ExtensionArray); + +impl TryFrom for MultiPolygonData { + type Error = VortexError; + + fn try_from(ext: ExtensionArray) -> Result { + vortex_ensure!( + ext.ext_dtype().is::(), + "expected a MultiPolygon extension array" + ); + Ok(MultiPolygonData(ext)) + } +} + +impl MultiPolygonData { + /// Serialize multipolygons to WKB (a view array) via geoarrow's cast — the form DuckDB + /// `GEOMETRY` takes. + pub fn to_wkb(&self, ctx: &mut ExecutionCtx) -> VortexResult { + let multipolygons = multipolygon_array(&self.0.storage_array().clone(), ctx)?; + let wkb_type = + GeoArrowType::WkbView(WkbType::new(geoarrow_metadata(&GeoMetadata::default()))); + let wkb = cast(&multipolygons, &wkb_type) + .map_err(|e| vortex_err!("failed to cast multipolygons to WKB: {e}"))?; + ArrayRef::from_arrow(wkb.to_array_ref().as_ref(), false) + } +} + +impl ArrowExportVTable for MultiPolygon { + fn arrow_ext_id(&self) -> Id { + *ARROW_MULTIPOLYGON + } + + fn vortex_id(&self) -> Id { + self.id() + } + + fn to_arrow_field( + &self, + name: &str, + dtype: &DType, + session: &ArrowSession, + ) -> VortexResult> { + let ext_type = dtype.as_extension(); + let geo_metadata = ext_type.metadata::(); + let dimension = multipolygon_dimension(ext_type.storage_dtype())?; + + let mut field = session.to_arrow_field(name, ext_type.storage_dtype())?; + field.try_with_extension_type(multipolygon_type(geo_metadata, dimension))?; + + Ok(Some(field)) + } + + fn execute_arrow( + &self, + array: ArrayRef, + target: &Field, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let is_multipolygon = array + .dtype() + .as_extension_opt() + .map(|ext| ext.is::()) + .unwrap_or(false); + if !is_multipolygon { + return Ok(ArrowExport::Unsupported(array)); + } + + let Ok(multipolygon_meta) = target.try_extension_type::() else { + return Ok(ArrowExport::Unsupported(array)); + }; + if multipolygon_meta.coord_type() != CoordType::Separated { + return Ok(ArrowExport::Unsupported(array)); + } + + let executed = array.execute::(ctx)?; + let storage = executed.storage_array().clone(); + + let storage_field = Field::new( + String::new(), + target.data_type().clone(), + target.is_nullable(), + ); + let session = ctx.session().clone(); + let arrow_storage = session + .arrow() + .execute_arrow(storage, Some(&storage_field), ctx)?; + + // Round-trip through GeoArrow's multipolygon array; `into_arrow` is concrete, so wrap in `Arc`. + let multipolygons = + MultiPolygonArray::try_from((arrow_storage.as_ref(), multipolygon_meta)) + .map_err(|e| vortex_err!("failed to construct MultiPolygonArray: {e}"))?; + + Ok(ArrowExport::Exported(Arc::new(multipolygons.into_arrow()))) + } +} + +impl ArrowImportVTable for MultiPolygon { + fn arrow_ext_id(&self) -> Id { + *ARROW_MULTIPOLYGON + } + + /// Import a `geoarrow.multipolygon` field (matched by GeoArrow name). Accepts the full + /// `MultiPolygonType`, or a metadata-less literal (name only), inferring the dimension. + fn from_arrow_field(&self, field: &Field) -> VortexResult> { + let (dimension, metadata) = + if let Ok(multipolygon_meta) = field.try_extension_type::() { + vortex_ensure!( + multipolygon_meta.coord_type() == CoordType::Separated, + "geoarrow.multipolygon with interleaved coordinates is not supported; \ + re-encode with separated (struct) coordinates" + ); + ( + multipolygon_meta.dimension().into(), + geo_metadata_from_arrow(multipolygon_meta.metadata()), + ) + } else { + // Literal: peel the three `List` layers to the coordinate struct and read its + // dimension from the field names (the canonical check rejects nullable coordinates). + if field.extension_type_name() != Some(MultiPolygonType::NAME) { + return Ok(None); + } + let DType::List(polygon, _) = DType::from_arrow(field) else { + return Ok(None); + }; + let DType::List(ring, _) = polygon.as_ref() else { + return Ok(None); + }; + let DType::List(coords, _) = ring.as_ref() else { + return Ok(None); + }; + let DType::Struct(fields, _) = coords.as_ref() else { + return Ok(None); + }; + let Ok(dimension) = Dimension::from_field_names(fields.names()) else { + return Ok(None); + }; + (dimension, GeoMetadata::default()) + }; + + let storage_dtype = multipolygon_storage_dtype(dimension, field.is_nullable().into()); + Ok(Some(DType::Extension( + ExtDType::try_with_vtable(MultiPolygon, metadata, storage_dtype)?.erased(), + ))) + } + + fn from_arrow_array( + &self, + array: ArrowArrayRef, + field: &Field, + dtype: &DType, + ) -> VortexResult { + let Some(ext_dtype) = dtype.as_extension_opt() else { + return Ok(ArrowImport::Unsupported(array)); + }; + if !ext_dtype.is::() + || field.try_extension_type::().is_err() + || !matches!(array.data_type(), DataType::List(_)) + { + return Ok(ArrowImport::Unsupported(array)); + } + + let storage = ArrayRef::from_arrow(array.as_ref(), field.is_nullable())?; + Ok(ArrowImport::Imported( + ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), + )) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use rstest::rstest; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::dtype::extension::ExtDType; + use vortex_error::VortexResult; + + use super::MultiPolygon; + use super::multipolygon_storage_dtype; + use crate::extension::GeoMetadata; + use crate::extension::coordinate::Dimension; + use crate::extension::coordinate::coordinate_storage_dtype; + + fn geo_meta() -> GeoMetadata { + GeoMetadata { + crs: Some("EPSG:4326".to_string()), + } + } + + /// `MultiPolygon` accepts the canonical `List>>` storage of every + /// dimension. + #[rstest] + #[case::xy(Dimension::Xy)] + #[case::xyz(Dimension::Xyz)] + #[case::xym(Dimension::Xym)] + #[case::xyzm(Dimension::Xyzm)] + fn multipolygon_validates_every_dimension(#[case] dim: Dimension) -> VortexResult<()> { + let storage = multipolygon_storage_dtype(dim, Nullability::NonNullable); + ExtDType::::try_new(geo_meta(), storage)?; + Ok(()) + } + + /// Non-multipolygon storage is rejected at dtype construction: a bare struct (point) and a + /// double list (polygon) both fail. + #[test] + fn multipolygon_rejects_invalid_storage() -> VortexResult<()> { + let primitive = DType::Primitive(PType::F64, Nullability::NonNullable); + assert!(ExtDType::::try_new(geo_meta(), primitive).is_err()); + + // A double list (polygon) is not a multipolygon. + let coords = coordinate_storage_dtype(Dimension::Xy, Nullability::NonNullable); + let ring = DType::List(Arc::new(coords), Nullability::NonNullable); + let polygon = DType::List(Arc::new(ring), Nullability::NonNullable); + assert!(ExtDType::::try_new(geo_meta(), polygon).is_err()); + Ok(()) + } +} diff --git a/vortex-geo/src/extension/point.rs b/vortex-geo/src/extension/point.rs index 389da97ad2b..4d33d5d6c77 100644 --- a/vortex-geo/src/extension/point.rs +++ b/vortex-geo/src/extension/point.rs @@ -12,11 +12,15 @@ use arrow_schema::Field; use arrow_schema::extension::ExtensionType; use geo_traits::to_geo::ToGeoGeometry; use geo_types::Geometry; +use geoarrow::array::GeoArrowArray; use geoarrow::array::GeoArrowArrayAccessor; use geoarrow::array::IntoArrow; use geoarrow::array::PointArray; use geoarrow::datatypes::CoordType; +use geoarrow::datatypes::GeoArrowType; use geoarrow::datatypes::PointType; +use geoarrow::datatypes::WkbType; +use geoarrow_cast::cast::cast; use prost::Message; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; @@ -37,6 +41,7 @@ use vortex_array::dtype::extension::ExtId; use vortex_array::dtype::extension::ExtVTable; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; +use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_err; @@ -97,20 +102,51 @@ fn point_type(geo_metadata: &GeoMetadata, dimension: Dimension) -> PointType { PointType::new(dimension.into(), geoarrow_metadata(geo_metadata)) } -/// Decode `Point` storage to `geo_types` points, for the geo scalar functions. -pub(crate) fn point_geometries( - storage: &ArrayRef, - ctx: &mut ExecutionCtx, -) -> VortexResult>> { +pub struct PointData(ExtensionArray); + +impl TryFrom for PointData { + type Error = VortexError; + + fn try_from(ext: ExtensionArray) -> Result { + vortex_ensure!( + ext.ext_dtype().is::(), + "expected a Point extension array" + ); + Ok(PointData(ext)) + } +} + +impl PointData { + /// Serialize points to WKB (a view array) via geoarrow's cast — the form DuckDB `GEOMETRY` takes. + pub fn to_wkb(&self, ctx: &mut ExecutionCtx) -> VortexResult { + let points = point_array(&self.0.storage_array().clone(), ctx)?; + let wkb_type = + GeoArrowType::WkbView(WkbType::new(geoarrow_metadata(&GeoMetadata::default()))); + let wkb = cast(&points, &wkb_type) + .map_err(|e| vortex_err!("failed to cast points to WKB: {e}"))?; + ArrayRef::from_arrow(wkb.to_array_ref().as_ref(), false) + } +} + +/// Build a geoarrow `PointArray` from a `Point`'s `Struct` storage, shared by WKB export +/// and `geo_types` decoding. +fn point_array(storage: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { let point_type = point_type( &GeoMetadata::default(), coordinate_dimension(storage.dtype())?, ); let session = ctx.session().clone(); let arrow = session.arrow().execute_arrow(storage.clone(), None, ctx)?; - let points = PointArray::try_from((arrow.as_ref(), point_type)) - .map_err(|e| vortex_err!("failed to construct PointArray: {e}"))?; - points + PointArray::try_from((arrow.as_ref(), point_type)) + .map_err(|e| vortex_err!("failed to construct PointArray: {e}")) +} + +/// Decode `Point` storage to `geo_types` points, for the geo scalar functions. +pub(crate) fn point_geometries( + storage: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult>> { + point_array(storage, ctx)? .iter() .map(|geometry| -> VortexResult> { Ok(geometry diff --git a/vortex-geo/src/extension/polygon.rs b/vortex-geo/src/extension/polygon.rs index 83ec6601158..9b3c74bc442 100644 --- a/vortex-geo/src/extension/polygon.rs +++ b/vortex-geo/src/extension/polygon.rs @@ -13,11 +13,15 @@ use arrow_schema::Field; use arrow_schema::extension::ExtensionType; use geo_traits::to_geo::ToGeoGeometry; use geo_types::Geometry; +use geoarrow::array::GeoArrowArray; use geoarrow::array::GeoArrowArrayAccessor; use geoarrow::array::IntoArrow; use geoarrow::array::PolygonArray; use geoarrow::datatypes::CoordType; +use geoarrow::datatypes::GeoArrowType; use geoarrow::datatypes::PolygonType; +use geoarrow::datatypes::WkbType; +use geoarrow_cast::cast::cast; use prost::Message; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; @@ -38,6 +42,7 @@ use vortex_array::dtype::extension::ExtDType; use vortex_array::dtype::extension::ExtId; use vortex_array::dtype::extension::ExtVTable; use vortex_array::scalar::ScalarValue; +use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -118,12 +123,7 @@ pub(crate) fn polygon_geometries( storage: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult>> { - let polygon_type = polygon_type(&GeoMetadata::default(), polygon_dimension(storage.dtype())?); - let session = ctx.session().clone(); - let arrow = session.arrow().execute_arrow(storage.clone(), None, ctx)?; - let polygons = PolygonArray::try_from((arrow.as_ref(), polygon_type)) - .map_err(|e| vortex_err!("failed to construct PolygonArray: {e}"))?; - polygons + polygon_array(storage, ctx)? .iter() .map(|geometry| -> VortexResult> { Ok(geometry @@ -134,6 +134,42 @@ pub(crate) fn polygon_geometries( .collect() } +/// Build a geoarrow `PolygonArray` from a `Polygon`'s `List>` storage. +fn polygon_array(storage: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let polygon_type = polygon_type(&GeoMetadata::default(), polygon_dimension(storage.dtype())?); + let session = ctx.session().clone(); + let arrow = session.arrow().execute_arrow(storage.clone(), None, ctx)?; + PolygonArray::try_from((arrow.as_ref(), polygon_type)) + .map_err(|e| vortex_err!("failed to construct PolygonArray: {e}")) +} + +/// A validated `Polygon` array (`try_from` checks the extension type). +pub struct PolygonData(ExtensionArray); + +impl TryFrom for PolygonData { + type Error = VortexError; + + fn try_from(ext: ExtensionArray) -> Result { + vortex_ensure!( + ext.ext_dtype().is::(), + "expected a Polygon extension array" + ); + Ok(PolygonData(ext)) + } +} + +impl PolygonData { + /// Serialize polygons to WKB (a view array) via geoarrow's cast — the form DuckDB `GEOMETRY` takes. + pub fn to_wkb(&self, ctx: &mut ExecutionCtx) -> VortexResult { + let polygons = polygon_array(&self.0.storage_array().clone(), ctx)?; + let wkb_type = + GeoArrowType::WkbView(WkbType::new(geoarrow_metadata(&GeoMetadata::default()))); + let wkb = cast(&polygons, &wkb_type) + .map_err(|e| vortex_err!("failed to cast polygons to WKB: {e}"))?; + ArrayRef::from_arrow(wkb.to_array_ref().as_ref(), false) + } +} + impl ArrowExportVTable for Polygon { fn arrow_ext_id(&self) -> Id { *ARROW_POLYGON diff --git a/vortex-geo/src/lib.rs b/vortex-geo/src/lib.rs index 951d93b7b4f..2cc8004efc5 100644 --- a/vortex-geo/src/lib.rs +++ b/vortex-geo/src/lib.rs @@ -8,6 +8,7 @@ use vortex_array::dtype::session::DTypeSessionExt; use vortex_array::scalar_fn::session::ScalarFnSessionExt; use vortex_session::VortexSession; +use crate::extension::MultiPolygon; use crate::extension::Point; use crate::extension::Polygon; use crate::extension::WellKnownBinary; @@ -32,6 +33,9 @@ pub fn initialize(session: &VortexSession) { session.dtypes().register(Polygon); session.arrow().register_exporter(Arc::new(Polygon)); session.arrow().register_importer(Arc::new(Polygon)); + session.dtypes().register(MultiPolygon); + session.arrow().register_exporter(Arc::new(MultiPolygon)); + session.arrow().register_importer(Arc::new(MultiPolygon)); // Register the geometry scalar functions. session.scalar_fns().register(GeoDistance); From 790a0d73e6df39d9679aa5fe4059c157def1c2bd Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Fri, 26 Jun 2026 17:47:35 -0400 Subject: [PATCH 5/7] test: add tests for multipolygon Signed-off-by: Nemo Yu --- vortex-geo/src/tests/mod.rs | 1 + vortex-geo/src/tests/multipolygon.rs | 94 ++++++++++++++++++++++++++++ 2 files changed, 95 insertions(+) create mode 100644 vortex-geo/src/tests/multipolygon.rs diff --git a/vortex-geo/src/tests/mod.rs b/vortex-geo/src/tests/mod.rs index 546de758eba..87b25ed1293 100644 --- a/vortex-geo/src/tests/mod.rs +++ b/vortex-geo/src/tests/mod.rs @@ -4,6 +4,7 @@ //! Arrow interop tests for the geospatial extension types, exercising the session wiring set up //! by [`crate::initialize`]. +mod multipolygon; mod point; mod wkb; diff --git a/vortex-geo/src/tests/multipolygon.rs b/vortex-geo/src/tests/multipolygon.rs new file mode 100644 index 00000000000..38f2543a96a --- /dev/null +++ b/vortex-geo/src/tests/multipolygon.rs @@ -0,0 +1,94 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Arrow interop for the `vortex.geo.multipolygon` extension type (`geoarrow.multipolygon`). + +use std::sync::Arc; + +use arrow_schema::DataType; +use arrow_schema::Field; +use arrow_schema::extension::ExtensionType as _; +use geoarrow::datatypes::CoordType; +use geoarrow::datatypes::Crs; +use geoarrow::datatypes::Dimension as GeoArrowDimension; +use geoarrow::datatypes::Metadata; +use geoarrow::datatypes::MultiPolygonType; +use vortex_array::arrow::ArrowSessionExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_error::VortexResult; + +use super::SESSION; +use crate::extension::MultiPolygon; + +/// A `geoarrow.multipolygon` Arrow field with separated (struct) XY coordinates. +fn multipolygon_field(name: &str, nullable: bool, crs: Option<&str>) -> Field { + let crs = crs + .map(|crs| Crs::from_unknown_crs_type(crs.to_string())) + .unwrap_or_default(); + let metadata = Arc::new(Metadata::new(crs, None)); + MultiPolygonType::new(GeoArrowDimension::XY, metadata).to_field(name, nullable) +} + +/// An imported `geoarrow.multipolygon` field maps to the MultiPolygon extension dtype, recovering the +/// CRS, the `List>>>` storage, and nullability. +#[test] +fn import_field_recovers_extension() -> VortexResult<()> { + let field = multipolygon_field("geom", true, Some("EPSG:4326")); + let dtype = SESSION.arrow().from_arrow_field(&field)?; + + let DType::Extension(ext) = &dtype else { + panic!("expected Extension dtype, got {dtype}"); + }; + assert!(ext.is::()); + assert_eq!( + ext.metadata::().crs.as_deref(), + Some("EPSG:4326") + ); + + // Storage peels three List layers (multipolygon → polygons → rings) to the coordinate struct. + let DType::List(polygons, nullability) = ext.storage_dtype() else { + panic!("expected List storage, got {}", ext.storage_dtype()); + }; + assert_eq!(*nullability, Nullability::Nullable); + let DType::List(rings, _) = polygons.as_ref() else { + panic!("expected List of polygons"); + }; + let DType::List(coords, _) = rings.as_ref() else { + panic!("expected List of rings"); + }; + let DType::Struct(fields, _) = coords.as_ref() else { + panic!("expected coordinate Struct"); + }; + let names: Vec<&str> = fields.names().iter().map(|n| n.as_ref()).collect(); + assert_eq!(names, vec!["x", "y"]); + Ok(()) +} + +/// A field with interleaved (`FixedSizeList`) coordinates fails to import. +#[test] +fn import_interleaved_field_fails() { + let multipolygon_type = MultiPolygonType::new(GeoArrowDimension::XY, Default::default()) + .with_coord_type(CoordType::Interleaved); + let field = multipolygon_type.to_field("geom", false); + assert!(SESSION.arrow().from_arrow_field(&field).is_err()); +} + +/// A field imported to the MultiPolygon dtype and exported back carries the `geoarrow.multipolygon` +/// extension over its `List` storage. +#[test] +fn export_field_carries_extension() -> VortexResult<()> { + let imported = + SESSION + .arrow() + .from_arrow_field(&multipolygon_field("geom", false, Some("EPSG:4326")))?; + let field = SESSION.arrow().to_arrow_field("geom", &imported)?; + + assert_eq!(field.extension_type_name(), Some(MultiPolygonType::NAME)); + assert!( + matches!(field.data_type(), DataType::List(_)), + "expected List storage, got {}", + field.data_type() + ); + Ok(()) +} From 42877b3dcf0a0e08951ad057c63e087fe38a94e7 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Wed, 1 Jul 2026 11:45:54 -0400 Subject: [PATCH 6/7] fix(vortex-geo): cache MultiPolygon extension id via CachedId `ExtVTable::id` interned the id with `ExtId::new_static` on every call, tripping the `disallowed_methods` clippy lint added on the base branch (#8617) since it grabs the interner lock per call. Match the `Point` and `Polygon` types by caching the id in a `static CachedId`. Signed-off-by: Nemo Yu --- vortex-geo/src/extension/multipolygon.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vortex-geo/src/extension/multipolygon.rs b/vortex-geo/src/extension/multipolygon.rs index 82fe081b316..c26b59012cc 100644 --- a/vortex-geo/src/extension/multipolygon.rs +++ b/vortex-geo/src/extension/multipolygon.rs @@ -67,7 +67,8 @@ impl ExtVTable for MultiPolygon { type NativeValue<'a> = &'a ScalarValue; fn id(&self) -> ExtId { - ExtId::new_static("vortex.geo.multipolygon") + static ID: CachedId = CachedId::new("vortex.geo.multipolygon"); + *ID } fn serialize_metadata(&self, metadata: &Self::Metadata) -> VortexResult> { @@ -238,7 +239,6 @@ impl ArrowExportVTable for MultiPolygon { .arrow() .execute_arrow(storage, Some(&storage_field), ctx)?; - // Round-trip through GeoArrow's multipolygon array; `into_arrow` is concrete, so wrap in `Arc`. let multipolygons = MultiPolygonArray::try_from((arrow_storage.as_ref(), multipolygon_meta)) .map_err(|e| vortex_err!("failed to construct MultiPolygonArray: {e}"))?; From 80e39d54ad3ae6adc946e30254e39dfddaf1c97e Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Wed, 1 Jul 2026 12:09:10 -0400 Subject: [PATCH 7/7] refactor(vortex-geo): extract shared geoarrow_to_wkb helper The `to_wkb` methods on `PointData`, `PolygonData`, and `MultiPolygonData` were byte-identical except for the geometry array type. Hoist the WKB cast into a single `geoarrow_to_wkb(&dyn GeoArrowArray)` in the extension module so each `to_wkb` is a one-liner, drop the redundant `storage_array().clone()` (the `*_array` helpers take `&ArrayRef`), and remove the now-unused `GeoArrowType`/`WkbType`/`cast`/`GeoArrowArray` imports from the three types. Signed-off-by: Nemo Yu --- vortex-geo/src/extension/mod.rs | 14 ++++++++++++++ vortex-geo/src/extension/multipolygon.rs | 15 +++------------ vortex-geo/src/extension/point.rs | 14 +++----------- vortex-geo/src/extension/polygon.rs | 14 +++----------- 4 files changed, 23 insertions(+), 34 deletions(-) diff --git a/vortex-geo/src/extension/mod.rs b/vortex-geo/src/extension/mod.rs index 5cccc489297..37f903aa0ca 100644 --- a/vortex-geo/src/extension/mod.rs +++ b/vortex-geo/src/extension/mod.rs @@ -11,8 +11,12 @@ use std::fmt::Display; use std::sync::Arc; use geo_types::Geometry; +use geoarrow::array::GeoArrowArray; use geoarrow::datatypes::Crs; +use geoarrow::datatypes::GeoArrowType; use geoarrow::datatypes::Metadata; +use geoarrow::datatypes::WkbType; +use geoarrow_cast::cast::cast; pub use multipolygon::*; pub use point::*; pub use polygon::*; @@ -22,6 +26,7 @@ use vortex_array::IntoArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::arrow::FromArrowArray; use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -99,6 +104,15 @@ pub(crate) fn geoarrow_metadata(geo_metadata: &GeoMetadata) -> Arc { )) } +/// Serialize a native geometry array to WKB (a `WkbView` array) via geoarrow's cast. +/// Shared by the `to_wkb` methods on the geometry extension types. +pub(crate) fn geoarrow_to_wkb(geo_array: &dyn GeoArrowArray) -> VortexResult { + let wkb_type = GeoArrowType::WkbView(WkbType::new(geoarrow_metadata(&GeoMetadata::default()))); + let wkb = cast(geo_array, &wkb_type) + .map_err(|e| vortex_err!("failed to cast geometry to WKB: {e}"))?; + ArrayRef::from_arrow(wkb.to_array_ref().as_ref(), false) +} + /// Recover [`GeoMetadata`] from GeoArrow metadata. pub(crate) fn geo_metadata_from_arrow(metadata: &Metadata) -> GeoMetadata { let crs = metadata.crs().crs_value().map(|value| { diff --git a/vortex-geo/src/extension/multipolygon.rs b/vortex-geo/src/extension/multipolygon.rs index c26b59012cc..67dac438970 100644 --- a/vortex-geo/src/extension/multipolygon.rs +++ b/vortex-geo/src/extension/multipolygon.rs @@ -13,15 +13,11 @@ use arrow_schema::Field; use arrow_schema::extension::ExtensionType; use geo_traits::to_geo::ToGeoGeometry; use geo_types::Geometry; -use geoarrow::array::GeoArrowArray; use geoarrow::array::GeoArrowArrayAccessor; use geoarrow::array::IntoArrow; use geoarrow::array::MultiPolygonArray; use geoarrow::datatypes::CoordType; -use geoarrow::datatypes::GeoArrowType; use geoarrow::datatypes::MultiPolygonType; -use geoarrow::datatypes::WkbType; -use geoarrow_cast::cast::cast; use prost::Message; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; @@ -56,6 +52,7 @@ use super::coordinate::coordinate_dimension; use super::coordinate::coordinate_storage_dtype; use super::geo_metadata_from_arrow; use super::geoarrow_metadata; +use super::geoarrow_to_wkb; /// A multipolygon (`geoarrow.multipolygon`); a single `Polygon` is a one-element multipolygon. #[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] @@ -167,15 +164,9 @@ impl TryFrom for MultiPolygonData { } impl MultiPolygonData { - /// Serialize multipolygons to WKB (a view array) via geoarrow's cast — the form DuckDB - /// `GEOMETRY` takes. + /// Serialize multipolygons to WKB (a view array) — the form DuckDB `GEOMETRY` takes. pub fn to_wkb(&self, ctx: &mut ExecutionCtx) -> VortexResult { - let multipolygons = multipolygon_array(&self.0.storage_array().clone(), ctx)?; - let wkb_type = - GeoArrowType::WkbView(WkbType::new(geoarrow_metadata(&GeoMetadata::default()))); - let wkb = cast(&multipolygons, &wkb_type) - .map_err(|e| vortex_err!("failed to cast multipolygons to WKB: {e}"))?; - ArrayRef::from_arrow(wkb.to_array_ref().as_ref(), false) + geoarrow_to_wkb(&multipolygon_array(self.0.storage_array(), ctx)?) } } diff --git a/vortex-geo/src/extension/point.rs b/vortex-geo/src/extension/point.rs index 4d33d5d6c77..6371105ca25 100644 --- a/vortex-geo/src/extension/point.rs +++ b/vortex-geo/src/extension/point.rs @@ -12,15 +12,11 @@ use arrow_schema::Field; use arrow_schema::extension::ExtensionType; use geo_traits::to_geo::ToGeoGeometry; use geo_types::Geometry; -use geoarrow::array::GeoArrowArray; use geoarrow::array::GeoArrowArrayAccessor; use geoarrow::array::IntoArrow; use geoarrow::array::PointArray; use geoarrow::datatypes::CoordType; -use geoarrow::datatypes::GeoArrowType; use geoarrow::datatypes::PointType; -use geoarrow::datatypes::WkbType; -use geoarrow_cast::cast::cast; use prost::Message; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; @@ -56,6 +52,7 @@ use super::coordinate::coordinate_from_struct; use super::coordinate::coordinate_storage_dtype; use super::geo_metadata_from_arrow; use super::geoarrow_metadata; +use super::geoarrow_to_wkb; /// A single location: `geoarrow.point`, stored as `Struct` of non-nullable `f64`. #[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] @@ -117,14 +114,9 @@ impl TryFrom for PointData { } impl PointData { - /// Serialize points to WKB (a view array) via geoarrow's cast — the form DuckDB `GEOMETRY` takes. + /// Serialize points to WKB (a view array) — the form DuckDB `GEOMETRY` takes. pub fn to_wkb(&self, ctx: &mut ExecutionCtx) -> VortexResult { - let points = point_array(&self.0.storage_array().clone(), ctx)?; - let wkb_type = - GeoArrowType::WkbView(WkbType::new(geoarrow_metadata(&GeoMetadata::default()))); - let wkb = cast(&points, &wkb_type) - .map_err(|e| vortex_err!("failed to cast points to WKB: {e}"))?; - ArrayRef::from_arrow(wkb.to_array_ref().as_ref(), false) + geoarrow_to_wkb(&point_array(self.0.storage_array(), ctx)?) } } diff --git a/vortex-geo/src/extension/polygon.rs b/vortex-geo/src/extension/polygon.rs index 9b3c74bc442..aeda30931a2 100644 --- a/vortex-geo/src/extension/polygon.rs +++ b/vortex-geo/src/extension/polygon.rs @@ -13,15 +13,11 @@ use arrow_schema::Field; use arrow_schema::extension::ExtensionType; use geo_traits::to_geo::ToGeoGeometry; use geo_types::Geometry; -use geoarrow::array::GeoArrowArray; use geoarrow::array::GeoArrowArrayAccessor; use geoarrow::array::IntoArrow; use geoarrow::array::PolygonArray; use geoarrow::datatypes::CoordType; -use geoarrow::datatypes::GeoArrowType; use geoarrow::datatypes::PolygonType; -use geoarrow::datatypes::WkbType; -use geoarrow_cast::cast::cast; use prost::Message; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; @@ -56,6 +52,7 @@ use super::coordinate::coordinate_dimension; use super::coordinate::coordinate_storage_dtype; use super::geo_metadata_from_arrow; use super::geoarrow_metadata; +use super::geoarrow_to_wkb; /// A polygon: `geoarrow.polygon`, stored as `List>>` (rings of vertices). #[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] @@ -159,14 +156,9 @@ impl TryFrom for PolygonData { } impl PolygonData { - /// Serialize polygons to WKB (a view array) via geoarrow's cast — the form DuckDB `GEOMETRY` takes. + /// Serialize polygons to WKB (a view array) — the form DuckDB `GEOMETRY` takes. pub fn to_wkb(&self, ctx: &mut ExecutionCtx) -> VortexResult { - let polygons = polygon_array(&self.0.storage_array().clone(), ctx)?; - let wkb_type = - GeoArrowType::WkbView(WkbType::new(geoarrow_metadata(&GeoMetadata::default()))); - let wkb = cast(&polygons, &wkb_type) - .map_err(|e| vortex_err!("failed to cast polygons to WKB: {e}"))?; - ArrayRef::from_arrow(wkb.to_array_ref().as_ref(), false) + geoarrow_to_wkb(&polygon_array(self.0.storage_array(), ctx)?) } }