From 7f6e3889018e27bf26cd6a78c2ab047a1b7af04a Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Thu, 2 Jul 2026 12:59:13 -0400 Subject: [PATCH 1/4] feat: support ST_DWithin pushdown in vortex MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Single-table `ST_DWithin(col, literal, radius)` filters over native-geometry Vortex scans now push down as `vortex.geo.distance(..) <= radius`, evaluated natively via a `GeoDistance` point fast path — with no query rewriting; any query using `ST_DWithin` benefits. Genuine `ST_DWithin` cannot push as-is: spatial's bind folds the radius into opaque bind data, while its SPATIAL_JOIN optimization requires exactly that folded 2-argument form. The pieces: - `duckdb_vx_register_st_dwithin_override` shadows `ST_DWithin` in the user catalog with a copy whose bind is cleared, so bound calls keep the radius as `children[2]`. Unpushed occurrences still execute correctly through spatial's own 3-column code path. - The expression converter lowers the 3-argument `st_dwithin` (and `st_distance`) to `GeoDistance`, guarded by the scan's fields: geometry columns must be native (`vortex.geo.wkb` columns fall back to DuckDB). `can_push_expression` dry-runs the lowering so DuckDB never installs an ExpressionFilter the scan would later drop. - `RestoreStDWithin`, in the existing Vortex optimizer pass, rebinds every remaining 3-argument `st_dwithin` (join conditions, unpushed filters) through spatial's original entry, restoring the folded form its spatial-join optimization requires. duckdb-bench registers the override once per connection, after `LOAD spatial`; it is a no-op when spatial is absent. SF1.0: Q1 47ms (parquet) / 16ms (vortex WKB) / 5.4ms (native, pushed); Q8 keeps SPATIAL_JOIN at ~89ms on the native lane. Signed-off-by: Nemo Yu --- benchmarks/duckdb-bench/src/lib.rs | 10 + vortex-duckdb/cpp/expr.cpp | 51 +++++ vortex-duckdb/cpp/include/expr.h | 7 + .../cpp/include/scalar_fn_pushdown.hpp | 8 + vortex-duckdb/cpp/scalar_fn_pushdown.cpp | 47 ++++ vortex-duckdb/cpp/vortex_duckdb.cpp | 3 + vortex-duckdb/src/convert/expr.rs | 200 +++++++++++++++--- vortex-duckdb/src/duckdb/database.rs | 10 + vortex-duckdb/src/table_function.rs | 2 +- vortex-geo/src/extension/mod.rs | 111 ++++++++++ vortex-geo/src/scalar_fn/distance.rs | 88 +++++++- 11 files changed, 502 insertions(+), 35 deletions(-) diff --git a/benchmarks/duckdb-bench/src/lib.rs b/benchmarks/duckdb-bench/src/lib.rs index bf64f123956..11be0dd80e5 100644 --- a/benchmarks/duckdb-bench/src/lib.rs +++ b/benchmarks/duckdb-bench/src/lib.rs @@ -78,6 +78,11 @@ impl DuckClient { for stmt in &statements { self.connection().query(stmt)?; } + // After `LOAD spatial`, shadow `ST_DWithin` so radius filters push. No-op without it. + self.db + .as_ref() + .vortex_expect("DuckClient database accessed after close") + .register_st_dwithin_override()?; self.init_sql = statements; Ok(()) } @@ -127,6 +132,11 @@ impl DuckClient { .vortex_expect("connection just opened") .query(stmt)?; } + // Re-shadow `ST_DWithin` against the fresh instance. + self.db + .as_ref() + .vortex_expect("database just opened") + .register_st_dwithin_override()?; Ok(()) } diff --git a/vortex-duckdb/cpp/expr.cpp b/vortex-duckdb/cpp/expr.cpp index afe2573adc2..dbd45999422 100644 --- a/vortex-duckdb/cpp/expr.cpp +++ b/vortex-duckdb/cpp/expr.cpp @@ -11,6 +11,17 @@ #include "duckdb/planner/expression/bound_operator_expression.hpp" #include "duckdb/planner/expression/bound_conjunction_expression.hpp" +#include "duckdb/catalog/catalog.hpp" +#include "duckdb/catalog/catalog_entry/scalar_function_catalog_entry.hpp" +#include "duckdb/main/capi/capi_internal.hpp" +#include "duckdb/main/client_context.hpp" +#include "duckdb/main/connection.hpp" +#include "duckdb/main/database_manager.hpp" +#include "duckdb/parser/parsed_data/create_scalar_function_info.hpp" +#include "duckdb/transaction/meta_transaction.hpp" + +#include + using namespace duckdb; extern "C" const char *duckdb_vx_sfunc_name(duckdb_vx_sfunc ffi_func) { @@ -21,6 +32,46 @@ extern "C" const char *duckdb_vx_sfunc_name(duckdb_vx_sfunc ffi_func) { return func->name.c_str(); } +extern "C" duckdb_state duckdb_vx_register_st_dwithin_override(duckdb_database ffi_db) { + if (!ffi_db) { + return DuckDBError; + } + const DatabaseWrapper &wrapper = *reinterpret_cast(ffi_db); + try { + Connection conn(*wrapper.database->instance); + ClientContext &context = *conn.context; + context.RunFunctionInTransaction([&]() { + auto &system = Catalog::GetSystemCatalog(context); + auto &entry = + system.GetEntry(context, DEFAULT_SCHEMA, "st_dwithin"); + ScalarFunctionSet set("st_dwithin"); + for (const auto &overload : entry.functions.functions) { + ScalarFunction copy = overload; + // Keep the radius as children[2]; spatial's bind folds it into private bind data. + copy.bind = nullptr; + set.AddFunction(copy); + } + CreateScalarFunctionInfo info(std::move(set)); + info.on_conflict = OnCreateConflict::REPLACE_ON_CONFLICT; + // `internal` entries are only accepted by the system catalog. + info.internal = false; + // The user catalog binds ahead of the system catalog, shadowing spatial's entry; + // `RestoreStDWithin` rebinds unpushed calls through the original. + auto &catalog = + Catalog::GetCatalog(context, DatabaseManager::GetDefaultDatabase(context)); + // Durable catalogs require the modified mark; scalar function entries are never + // persisted, so this is metadata-only. + MetaTransaction::Get(context).ModifyDatabase(catalog.GetAttached(), + DatabaseModificationType()); + catalog.CreateFunction(context, info); + }); + } catch (const std::exception &) { + // No `spatial` loaded, so there is no `ST_DWithin` to override; nothing to register. + return DuckDBSuccess; + } + return DuckDBSuccess; +} + extern "C" const char *duckdb_vx_expr_to_string(duckdb_vx_expr ffi_expr) { if (!ffi_expr) { return nullptr; diff --git a/vortex-duckdb/cpp/include/expr.h b/vortex-duckdb/cpp/include/expr.h index 5b7997596d6..d71e4e270a7 100644 --- a/vortex-duckdb/cpp/include/expr.h +++ b/vortex-duckdb/cpp/include/expr.h @@ -13,6 +13,13 @@ typedef struct duckdb_vx_sfunc_ *duckdb_vx_sfunc; const char *duckdb_vx_sfunc_name(duckdb_vx_sfunc ffi_func); +/// Shadow the spatial extension's `ST_DWithin` in the user catalog with a copy whose bind keeps +/// the radius as the third argument (spatial's bind folds it into private bind data), so radius +/// predicates can push into Vortex scans. Unpushed occurrences still execute via spatial's own +/// 3-column code path, and the Vortex optimizer pass rebinds them (e.g. join conditions) back to +/// spatial's folded form, which its spatial-join optimization requires. No-op without `spatial`. +duckdb_state duckdb_vx_register_st_dwithin_override(duckdb_database ffi_db); + typedef struct duckdb_vx_expr_ *duckdb_vx_expr; /// Return the string representation of the expression. Must be freed with `duckdb_free`. diff --git a/vortex-duckdb/cpp/include/scalar_fn_pushdown.hpp b/vortex-duckdb/cpp/include/scalar_fn_pushdown.hpp index ef590c96dcc..00a47216672 100644 --- a/vortex-duckdb/cpp/include/scalar_fn_pushdown.hpp +++ b/vortex-duckdb/cpp/include/scalar_fn_pushdown.hpp @@ -77,6 +77,14 @@ using Projections = unordered_map; LogicalOperatorPtr TryPushdownScalarFunctions(ClientContext &context, LogicalOperatorPtr plan); +/* + * Rebind remaining 3-argument `st_dwithin` calls (bound via the Vortex override; see expr.h) + * back through spatial's original entry, whose folded 2-argument form its spatial-join + * optimization requires. Runs after DuckDB's optimizers, when pushable filters are already + * absorbed into Vortex scans. + */ +void RestoreStDWithin(ClientContext &context, LogicalOperator &plan); + /** * Collect fn(col) expressions i.e. expressions where a single function (not * a function chain) wraps a single bound column. If "col" is used without diff --git a/vortex-duckdb/cpp/scalar_fn_pushdown.cpp b/vortex-duckdb/cpp/scalar_fn_pushdown.cpp index 057920f79f7..0f66033e574 100644 --- a/vortex-duckdb/cpp/scalar_fn_pushdown.cpp +++ b/vortex-duckdb/cpp/scalar_fn_pushdown.cpp @@ -1,6 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors #include "duckdb/catalog/catalog.hpp" +#include "duckdb/catalog/catalog_entry/scalar_function_catalog_entry.hpp" +#include "duckdb/common/string_util.hpp" +#include "duckdb/function/function_binder.hpp" #include "duckdb/planner/operator/logical_projection.hpp" #include "scalar_fn_pushdown.hpp" #include "table_function.hpp" @@ -231,3 +234,47 @@ ScalarFnReplace::ScalarFnReplace(Analyses &analyses, const Projections &projecti TableColumnStorageIndex GetAnalysis::StorageIndex(TableColumnScanIndex idx) const { return get.GetColumnIds()[idx].GetPrimaryIndex(); } + +namespace { + +// See RestoreStDWithin: rebinding through spatial's own entry lets spatial build its own bind +// data, so nothing here depends on its internals. +class StDWithinRestore final : public LogicalOperatorVisitor { +public: + explicit StDWithinRestore(ClientContext &context) : context(context) { + } + + ExpressionPtr VisitReplace(BoundFunctionExpression &expr, ExpressionPtr *) override { + if (expr.children.size() != 3 || !StringUtil::CIEquals(expr.function.name, "st_dwithin")) { + return nullptr; // Not the override's shape: keep it and descend into its children. + } + // The system catalog holds spatial's original; the user-catalog override cannot shadow + // this lookup. + auto original = Catalog::GetSystemCatalog(context).GetEntry( + context, DEFAULT_SCHEMA, "st_dwithin", OnEntryNotFound::RETURN_NULL); + if (!original) { + return nullptr; + } + vector children; + children.reserve(expr.children.size()); + for (const auto &child : expr.children) { + children.push_back(child->Copy()); + } + ErrorData error; + FunctionBinder binder(context); + auto bound = binder.BindScalarFunction(*original, std::move(children), error); + if (!bound) { + return nullptr; // No matching overload: keep the executable 3-argument form. + } + return bound; + } + +private: + ClientContext &context; +}; + +} // namespace + +void RestoreStDWithin(ClientContext &context, LogicalOperator &plan) { + StDWithinRestore(context).VisitOperator(plan); +} diff --git a/vortex-duckdb/cpp/vortex_duckdb.cpp b/vortex-duckdb/cpp/vortex_duckdb.cpp index 091f98a1703..0a0cee3bc6b 100644 --- a/vortex-duckdb/cpp/vortex_duckdb.cpp +++ b/vortex-duckdb/cpp/vortex_duckdb.cpp @@ -269,6 +269,9 @@ extern "C" duckdb_blob duckdb_vx_value_get_geometry(duckdb_value value) { static void VortexOptimizeFunction(OptimizerExtensionInput &input, unique_ptr &plan) { plan = TryPushdownScalarFunctions(input.context, std::move(plan)); + // Runs before spatial's own optimizer pass: this extension registers at database open, + // spatial's at LOAD, and DuckDB invokes optimizer extensions in registration order. + RestoreStDWithin(input.context, *plan); } struct VortexOptimizerExtension final : OptimizerExtension { diff --git a/vortex-duckdb/src/convert/expr.rs b/vortex-duckdb/src/convert/expr.rs index 387b644fe30..ae0f7787ad8 100644 --- a/vortex-duckdb/src/convert/expr.rs +++ b/vortex-duckdb/src/convert/expr.rs @@ -28,6 +28,7 @@ use vortex::expr::not; use vortex::expr::or_collect; use vortex::expr::root; use vortex::scalar::Scalar; +use vortex::scalar_fn::EmptyOptions; use vortex::scalar_fn::ScalarFnVTableExt; use vortex::scalar_fn::fns::between::Between; use vortex::scalar_fn::fns::between::BetweenOptions; @@ -37,6 +38,12 @@ use vortex::scalar_fn::fns::like::Like; use vortex::scalar_fn::fns::like::LikeOptions; use vortex::scalar_fn::fns::literal::Literal; use vortex::scalar_fn::fns::operators::Operator; +use vortex_geo::extension::MultiPolygon; +use vortex_geo::extension::Point; +use vortex_geo::extension::Polygon; +use vortex_geo::extension::WellKnownBinary; +use vortex_geo::extension::native_geometry_scalar_from_wkb; +use vortex_geo::scalar_fn::distance::GeoDistance; use crate::cpp::DUCKDB_TYPE; use crate::cpp::DUCKDB_VX_EXPR_TYPE; @@ -73,15 +80,134 @@ fn build_list_length(expr: Expression, nullability: Nullability) -> Expression { cast(list_length(expr), DType::Primitive(PType::I64, nullability)) } +/// Read an `f64` from a constant expression (the `ST_DWithin` radius); `None` for non-constants. +fn from_bound_f64(value: &duckdb::ExpressionRef) -> VortexResult> { + match value.as_class().vortex_expect("unknown class") { + BoundConstant(constant) => Ok(Some(f64::try_from(&Scalar::try_from(constant.value)?)?)), + _ => Ok(None), + } +} + +/// Context threaded through expression conversion. +#[derive(Clone, Copy)] +struct ConvertCtx<'a> { + /// Substituted for `BoundRef` references when converting scan-scoped table filters. + col_sub: Option<&'a Expression>, + /// The scan's fields, when known; geo lowering requires them to verify a geometry column. + fields: Option<&'a [DuckdbField]>, +} + +/// Whether `name` is a native geometry column of the scan. The pushed `GeoDistance` cannot +/// evaluate `vortex.geo.wkb` columns, which also surface to DuckDB as `GEOMETRY`. +fn is_native_geo_column(fields: Option<&[DuckdbField]>, name: &str) -> bool { + fields + .into_iter() + .flatten() + .filter(|field| field.name == name) + .any(|field| match field.dtype.as_extension_opt() { + Some(ext) => ext.is::() || ext.is::() || ext.is::(), + None => false, + }) +} + +/// Lower a geo operand: a `GEOMETRY` literal arrives as WKB, decoded once to its native type so the +/// pushed `GeoDistance` stays native; a column must be native geometry. `None` skips the push. +fn geo_operand( + value: &duckdb::ExpressionRef, + ctx: ConvertCtx<'_>, +) -> VortexResult> { + match value.as_class() { + Some(BoundConstant(constant)) => { + let scalar = Scalar::try_from(constant.value)?; + let DType::Extension(ext_dtype) = scalar.dtype() else { + return Ok(None); + }; + if !ext_dtype.is::() { + return Ok(None); + } + let storage = scalar.as_extension().to_storage_scalar(); + let Some(buf) = storage.as_binary_opt().and_then(|b| b.value()) else { + return Ok(None); + }; + Ok(native_geometry_scalar_from_wkb(buf.as_slice())?.map(lit)) + } + Some(BoundColumnRef(col_ref)) + if is_native_geo_column(ctx.fields, col_ref.name.as_ref()) => + { + try_from_expression_inner(value, ctx) + } + _ => Ok(None), + } +} + +/// Lower geo UDFs to native Vortex geo ops so the work runs in the scan. `None` otherwise. +fn try_from_geo_function( + name: &str, + func: &BoundFunction, + ctx: ConvertCtx<'_>, +) -> VortexResult> { + // Catch-all for every bound function: reject non-geo names before touching the children. + if !is_geo_function(name) { + debug!("bound function {name}"); + return Ok(None); + } + let children: Vec<_> = func.children().collect(); + let expr = match name.to_ascii_lowercase().as_str() { + // The Vortex override keeps the radius as `children[2]`; see + // `duckdb_vx_register_st_dwithin_override`. + "st_dwithin" => { + if children.len() != 3 { + return Ok(None); + } + let Some(a) = geo_operand(children[0], ctx)? else { + return Ok(None); + }; + let Some(b) = geo_operand(children[1], ctx)? else { + return Ok(None); + }; + // A non-constant radius is left for DuckDB to evaluate. + let Some(distance) = from_bound_f64(children[2])? else { + return Ok(None); + }; + let geo_distance = GeoDistance.new_expr(EmptyOptions, [a, b]); + Binary.new_expr(Operator::Lte, [geo_distance, lit(distance)]) + } + "st_distance" => { + if children.len() != 2 { + return Ok(None); + } + let Some(a) = geo_operand(children[0], ctx)? else { + return Ok(None); + }; + let Some(b) = geo_operand(children[1], ctx)? else { + return Ok(None); + }; + GeoDistance.new_expr(EmptyOptions, [a, b]) + } + _ => return Ok(None), + }; + + Ok(Some(expr)) +} + +/// Geo UDFs that `try_from_geo_function` lowers - shared with `can_push_expression` so the pushable +/// and lowered sets can't drift. +fn is_geo_function(name: &str) -> bool { + matches!( + name.to_ascii_lowercase().as_str(), + "st_distance" | "st_dwithin" + ) +} + fn try_from_bound_function( func: &BoundFunction, - col_sub: Option<&Expression>, + ctx: ConvertCtx<'_>, ) -> VortexResult> { let expr = match func.scalar_function.name() { "strlen" => { let children: Vec<_> = func.children().collect(); vortex_ensure!(children.len() == 1); - let Some(col) = try_from_expression_inner(children[0], col_sub)? else { + let Some(col) = try_from_expression_inner(children[0], ctx)? else { return Ok(None); }; let col = byte_length(col); @@ -95,7 +221,7 @@ fn try_from_bound_function( "struct_extract" => { let children: Vec<_> = func.children().collect(); vortex_ensure!(children.len() == 2); - let Some(child) = try_from_expression_inner(children[0], col_sub)? else { + let Some(child) = try_from_expression_inner(children[0], ctx)? else { return Ok(None); }; let field = from_bound_str(children[1])?; @@ -104,10 +230,10 @@ fn try_from_bound_function( like @ ("~~" | "!~~") => { let children: Vec<_> = func.children().collect(); vortex_ensure!(children.len() == 2); - let Some(string) = try_from_expression_inner(children[0], col_sub)? else { + let Some(string) = try_from_expression_inner(children[0], ctx)? else { return Ok(None); }; - let Some(target) = try_from_expression_inner(children[1], col_sub)? else { + let Some(target) = try_from_expression_inner(children[1], ctx)? else { return Ok(None); }; let opts = LikeOptions { @@ -119,7 +245,7 @@ fn try_from_bound_function( matchers @ ("contains" | "prefix" | "suffix") => { let children: Vec<_> = func.children().collect(); vortex_ensure!(children.len() == 2); - let Some(value) = try_from_expression_inner(children[0], col_sub)? else { + let Some(value) = try_from_expression_inner(children[0], ctx)? else { return Ok(None); }; let pattern = from_bound_str(children[1])?; @@ -137,7 +263,7 @@ fn try_from_bound_function( if children.len() != 1 { return Ok(None); } - let Some(col) = try_from_expression_inner(children[0], col_sub)? else { + let Some(col) = try_from_expression_inner(children[0], ctx)? else { return Ok(None); }; @@ -151,7 +277,7 @@ fn try_from_bound_function( let child = children[0]; if returns_a_list(child) { - let Some(col) = try_from_expression_inner(child, col_sub)? else { + let Some(col) = try_from_expression_inner(child, ctx)? else { return Ok(None); }; @@ -162,10 +288,8 @@ fn try_from_bound_function( return Ok(None); } } - _ => { - debug!("bound function {}", func.scalar_function.name()); - return Ok(None); - } + // Geo UDFs are handled here; non-geo names return `None` inside. + name => return try_from_geo_function(name, func, ctx), }; Ok(Some(expr)) @@ -173,15 +297,30 @@ fn try_from_bound_function( pub fn try_from_bound_expression( value: &duckdb::ExpressionRef, + fields: &[DuckdbField], ) -> VortexResult> { - try_from_expression_inner(value, None) + try_from_expression_inner( + value, + ConvertCtx { + col_sub: None, + fields: Some(fields), + }, + ) } pub(super) fn try_from_bound_expression_with_col_sub( value: &duckdb::ExpressionRef, col_sub: &Expression, ) -> VortexResult> { - try_from_expression_inner(value, Some(col_sub)) + // No fields: scan-time table filters never carry geo functions, because + // `can_push_expression` refuses them. + try_from_expression_inner( + value, + ConvertCtx { + col_sub: Some(col_sub), + fields: None, + }, + ) } fn is_supported_length_alias(func: &BoundFunction) -> bool { @@ -227,6 +366,9 @@ pub fn can_push_expression(value: &duckdb::ExpressionRef) -> bool { || name == "strlen" || name == "array_length" || (matches!(name, "len" | "length") && is_supported_length_alias(&func)) + // Geo functions are absent on purpose: they push only via + // `pushdown_complex_filter`, which has the scan's fields to verify the geometry + // columns are native. } ExpressionClass::BoundOperator(op) => { if !matches!( @@ -284,7 +426,7 @@ pub fn try_from_projection_expression( // can_push_expression fn try_from_expression_inner( value: &duckdb::ExpressionRef, - col_sub: Option<&Expression>, + ctx: ConvertCtx<'_>, ) -> VortexResult> { let Some(value) = value.as_class() else { debug!( @@ -295,7 +437,7 @@ fn try_from_expression_inner( }; Ok(Some(match value { BoundRef => { - let Some(col) = col_sub else { + let Some(col) = ctx.col_sub else { vortex_bail!("BoundRef requested but no column supplied"); }; col.clone() @@ -305,23 +447,23 @@ fn try_from_expression_inner( BoundComparison(compare) => { let operator: Operator = compare.op.try_into()?; - let Some(left) = try_from_expression_inner(compare.left, col_sub)? else { + let Some(left) = try_from_expression_inner(compare.left, ctx)? else { return Ok(None); }; - let Some(right) = try_from_expression_inner(compare.right, col_sub)? else { + let Some(right) = try_from_expression_inner(compare.right, ctx)? else { return Ok(None); }; Binary.new_expr(operator, [left, right]) } BoundBetween(between) => { - let Some(array) = try_from_expression_inner(between.input, col_sub)? else { + let Some(array) = try_from_expression_inner(between.input, ctx)? else { return Ok(None); }; - let Some(lower) = try_from_expression_inner(between.lower, col_sub)? else { + let Some(lower) = try_from_expression_inner(between.lower, ctx)? else { return Ok(None); }; - let Some(upper) = try_from_expression_inner(between.upper, col_sub)? else { + let Some(upper) = try_from_expression_inner(between.upper, ctx)? else { return Ok(None); }; Between.new_expr( @@ -346,7 +488,7 @@ fn try_from_expression_inner( | DUCKDB_VX_EXPR_TYPE::DUCKDB_VX_EXPR_TYPE_OPERATOR_IS_NOT_NULL => { let children: Vec<_> = operator.children().collect(); vortex_ensure!(children.len() == 1); - let Some(child) = try_from_expression_inner(children[0], col_sub)? else { + let Some(child) = try_from_expression_inner(children[0], ctx)? else { return Ok(None); }; match operator.op { @@ -359,10 +501,10 @@ fn try_from_expression_inner( } } DUCKDB_VX_EXPR_TYPE::DUCKDB_VX_EXPR_TYPE_COMPARE_IN => { - return try_from_compare_in(operator, col_sub, false); + return try_from_compare_in(operator, ctx, false); } DUCKDB_VX_EXPR_TYPE::DUCKDB_VX_EXPR_TYPE_COMPARE_NOT_IN => { - return try_from_compare_in(operator, col_sub, true); + return try_from_compare_in(operator, ctx, true); } _ => { debug!(op=?operator.op, "cannot push down operator"); @@ -370,12 +512,12 @@ fn try_from_expression_inner( } }, ExpressionClass::BoundFunction(func) => { - return try_from_bound_function(&func, col_sub); + return try_from_bound_function(&func, ctx); } BoundConjunction(conj) => { let Some(children) = conj .children() - .map(|c| try_from_expression_inner(c, col_sub)) + .map(|c| try_from_expression_inner(c, ctx)) .collect::>>>()? else { return Ok(None); @@ -395,13 +537,13 @@ fn try_from_expression_inner( fn try_from_compare_in( operator: BoundOperator, - col_sub: Option<&Expression>, + ctx: ConvertCtx<'_>, not_in: bool, ) -> VortexResult> { // First child is element, rest form the list. let children: Vec<_> = operator.children().collect(); assert!(children.len() >= 2); - let Some(element) = try_from_expression_inner(children[0], col_sub)? else { + let Some(element) = try_from_expression_inner(children[0], ctx)? else { return Ok(None); }; @@ -409,7 +551,7 @@ fn try_from_compare_in( .iter() .skip(1) .map(|c| { - let Some(value) = try_from_expression_inner(c, col_sub)? else { + let Some(value) = try_from_expression_inner(c, ctx)? else { return Ok(None); }; Ok(Some( diff --git a/vortex-duckdb/src/duckdb/database.rs b/vortex-duckdb/src/duckdb/database.rs index ab86503b291..e7033b75ee2 100644 --- a/vortex-duckdb/src/duckdb/database.rs +++ b/vortex-duckdb/src/duckdb/database.rs @@ -90,4 +90,14 @@ impl DatabaseRef { ); Ok(()) } + + /// Shadow `ST_DWithin` with a radius-visible copy so its filters push into the Vortex scan + /// (see `duckdb_vx_register_st_dwithin_override`). No-op when `spatial` is not loaded. + pub fn register_st_dwithin_override(&self) -> VortexResult<()> { + duckdb_try!( + unsafe { cpp::duckdb_vx_register_st_dwithin_override(self.as_ptr()) }, + "Failed to register the ST_DWithin override" + ); + Ok(()) + } } diff --git a/vortex-duckdb/src/table_function.rs b/vortex-duckdb/src/table_function.rs index 5151e47a464..098d8bf30ee 100644 --- a/vortex-duckdb/src/table_function.rs +++ b/vortex-duckdb/src/table_function.rs @@ -391,7 +391,7 @@ pub fn pushdown_complex_filter( ) -> VortexResult { debug!(%expr, "pushing down expression"); - let Some(expr) = try_from_bound_expression(expr)? else { + let Some(expr) = try_from_bound_expression(expr, &bind_data.column_fields)? else { debug!(%expr, "failed to push down expression"); return Ok(false); }; diff --git a/vortex-geo/src/extension/mod.rs b/vortex-geo/src/extension/mod.rs index 37f903aa0ca..4e4896aa5e5 100644 --- a/vortex-geo/src/extension/mod.rs +++ b/vortex-geo/src/extension/mod.rs @@ -10,11 +10,19 @@ mod wkb; use std::fmt::Display; use std::sync::Arc; +use ::wkb::reader::GeometryType; +use arrow_array::BinaryArray; use geo_types::Geometry; +use geoarrow::array::GenericWkbArray; use geoarrow::array::GeoArrowArray; +use geoarrow::datatypes::CoordType; use geoarrow::datatypes::Crs; +use geoarrow::datatypes::Dimension; use geoarrow::datatypes::GeoArrowType; use geoarrow::datatypes::Metadata; +use geoarrow::datatypes::MultiPolygonType; +use geoarrow::datatypes::PointType; +use geoarrow::datatypes::PolygonType; use geoarrow::datatypes::WkbType; use geoarrow_cast::cast::cast; pub use multipolygon::*; @@ -27,6 +35,8 @@ use vortex_array::arrays::ConstantArray; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::arrow::FromArrowArray; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::dtype::extension::ExtVTable; use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -72,6 +82,63 @@ pub(crate) fn single_geometry( .ok_or_else(|| vortex_err!("geo: constant operand decoded to no geometry")) } +/// Decode a WKB geometry literal (DuckDB's wire form for `GEOMETRY` constants) to its native +/// `Point`/`Polygon`/`MultiPolygon` scalar. `None` for unsupported types. Plan-time, one value only. +pub fn native_geometry_scalar_from_wkb(bytes: &[u8]) -> VortexResult> { + let metadata = geoarrow_metadata(&GeoMetadata::default()); + let binary = BinaryArray::from(vec![Some(bytes)]); + let wkb = GenericWkbArray::::try_from(( + &binary as &dyn arrow_array::Array, + WkbType::new(Arc::clone(&metadata)), + )) + .map_err(|e| vortex_err!("failed to read WKB literal: {e}"))?; + + // Cast the WKB value to `target`, import its native storage as a Vortex array. + let to_storage = |target: &GeoArrowType| -> VortexResult { + let native = + cast(&wkb, target).map_err(|e| vortex_err!("failed to cast WKB literal: {e}"))?; + ArrayRef::from_arrow(native.to_array_ref().as_ref(), false) + }; + + let scalar = match Wkb::try_from_bytes(bytes)?.geometry_type() { + GeometryType::Point => { + let target = GeoArrowType::Point( + PointType::new(Dimension::XY, metadata).with_coord_type(CoordType::Separated), + ); + geo_ext_scalar(Point, to_storage(&target)?)? + } + GeometryType::Polygon => { + let target = GeoArrowType::Polygon( + PolygonType::new(Dimension::XY, metadata).with_coord_type(CoordType::Separated), + ); + geo_ext_scalar(Polygon, to_storage(&target)?)? + } + GeometryType::MultiPolygon => { + let target = GeoArrowType::MultiPolygon( + MultiPolygonType::new(Dimension::XY, metadata) + .with_coord_type(CoordType::Separated), + ); + geo_ext_scalar(MultiPolygon, to_storage(&target)?)? + } + _ => return Ok(None), + }; + Ok(Some(scalar)) +} + +/// Wrap cast-from-WKB `storage` in its `vtable` extension type and pull out the single scalar. +// `scalar_at` is deprecated for `execute_scalar`, but there is no execution context at plan time. +#[allow(deprecated)] +fn geo_ext_scalar>( + vtable: V, + storage: ArrayRef, +) -> VortexResult { + let ext = ExtDType::try_with_vtable(vtable, GeoMetadata::default(), storage.dtype().clone())? + .erased(); + ExtensionArray::try_new(ext, storage)? + .into_array() + .scalar_at(0) +} + /// Extension metadata that is common to all the geospatial extension types. /// /// Currently, this is just the coordinate reference system (CRS). @@ -130,7 +197,13 @@ pub(crate) fn geo_metadata_from_arrow(metadata: &Metadata) -> GeoMetadata { #[cfg(test)] mod tests { use prost::Message; + use vortex_array::dtype::DType; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + use super::Point; + use super::Polygon; + use super::native_geometry_scalar_from_wkb; use crate::extension::GeoMetadata; #[test] @@ -145,4 +218,42 @@ mod tests { let decoded = GeoMetadata::decode(bytes.as_slice()).unwrap(); assert_eq!(decoded, meta); } + + /// A little-endian WKB `POINT` literal decodes to the native `Point` extension scalar. + #[test] + fn decodes_wkb_point_to_native() -> VortexResult<()> { + let mut wkb = vec![1u8]; // little-endian byte order + wkb.extend_from_slice(&1u32.to_le_bytes()); // geometry type: point + wkb.extend_from_slice(&1.0f64.to_le_bytes()); // x + wkb.extend_from_slice(&2.0f64.to_le_bytes()); // y + + let scalar = native_geometry_scalar_from_wkb(&wkb)?.expect("a point scalar"); + let DType::Extension(ext) = scalar.dtype() else { + panic!("expected an extension dtype, got {}", scalar.dtype()); + }; + assert!(ext.is::()); + Ok(()) + } + + /// A little-endian WKB `POLYGON` literal decodes to the native `Polygon` extension scalar. + #[test] + fn decodes_wkb_polygon_to_native() -> VortexResult<()> { + let ring = [(0.0, 0.0), (1.0, 0.0), (0.0, 1.0), (0.0, 0.0)]; + let mut wkb = vec![1u8]; // little-endian byte order + wkb.extend_from_slice(&3u32.to_le_bytes()); // geometry type: polygon + wkb.extend_from_slice(&1u32.to_le_bytes()); // one ring + let ring_len = u32::try_from(ring.len()).map_err(|e| vortex_err!("{e}"))?; + wkb.extend_from_slice(&ring_len.to_le_bytes()); + for (x, y) in ring { + wkb.extend_from_slice(&f64::to_le_bytes(x)); + wkb.extend_from_slice(&f64::to_le_bytes(y)); + } + + let scalar = native_geometry_scalar_from_wkb(&wkb)?.expect("a polygon scalar"); + let DType::Extension(ext) = scalar.dtype() else { + panic!("expected an extension dtype, got {}", scalar.dtype()); + }; + assert!(ext.is::()); + Ok(()) + } } diff --git a/vortex-geo/src/scalar_fn/distance.rs b/vortex-geo/src/scalar_fn/distance.rs index 6531c1dd8f2..b894946eecb 100644 --- a/vortex-geo/src/scalar_fn/distance.rs +++ b/vortex-geo/src/scalar_fn/distance.rs @@ -10,8 +10,12 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::Constant; use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::ScalarFnArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; @@ -28,6 +32,8 @@ use vortex_error::vortex_ensure; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use crate::extension::Point; +use crate::extension::coordinate::coordinate_from_struct; use crate::extension::geometries; use crate::extension::single_geometry; @@ -101,14 +107,25 @@ impl ScalarFnVTable for GeoDistance { (Some(query), None) => distances_to_constant(&b, query.scalar(), ctx), (None, Some(query)) => distances_to_constant(&a, query.scalar(), ctx), (None, None) => { - let ag = geometries(&a, ctx)?; - let bg = geometries(&b, ctx)?; vortex_ensure!( - ag.len() == bg.len(), + a.len() == b.len(), "geo distance: operand length mismatch {} vs {}", - ag.len(), - bg.len() + a.len(), + b.len() ); + // Fast path: two Point columns, distance straight over their `x`/`y` f64 buffers. + if is_nonnull_point(a.dtype()) && is_nonnull_point(b.dtype()) { + let (xa, ya) = point_xy(&a, ctx)?; + let (xb, yb) = point_xy(&b, ctx)?; + return Ok(point_distances( + xa.as_slice::().iter().copied(), + ya.as_slice::().iter().copied(), + xb.as_slice::().iter().copied(), + yb.as_slice::().iter().copied(), + )); + } + let ag = geometries(&a, ctx)?; + let bg = geometries(&b, ctx)?; let distances = ag.iter().zip(&bg).map(|(x, y)| Euclidean.distance(x, y)); Ok(PrimitiveArray::from_iter(distances).into_array()) } @@ -123,12 +140,73 @@ fn distances_to_constant( query: &Scalar, ctx: &mut ExecutionCtx, ) -> VortexResult { + // Fast path: Point column vs constant Point, `x`/`y` f64 buffers, broadcasting the constant. + if is_nonnull_point(operand.dtype()) && is_point(query.dtype()) { + let q = coordinate_from_struct(&query.as_extension().to_storage_scalar())?; + let (xs, ys) = point_xy(operand, ctx)?; + return Ok(point_distances( + xs.as_slice::().iter().copied(), + ys.as_slice::().iter().copied(), + std::iter::repeat(q.x), + std::iter::repeat(q.y), + )); + } + let query = single_geometry(query, ctx)?; let geoms = geometries(operand, ctx)?; let distances = geoms.iter().map(|g| Euclidean.distance(g, &query)); Ok(PrimitiveArray::from_iter(distances).into_array()) } +/// Extract the `x` and `y` `f64` columns from a native `Point` operand, for the columnar fast paths. +fn point_xy( + operand: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult<(PrimitiveArray, PrimitiveArray)> { + let storage = operand + .clone() + .execute::(ctx)? + .storage_array() + .clone() + .execute::(ctx)?; + let xs = storage + .unmasked_field_by_name("x")? + .clone() + .execute::(ctx)?; + let ys = storage + .unmasked_field_by_name("y")? + .clone() + .execute::(ctx)?; + Ok((xs, ys)) +} + +/// Per-row planar distance `sqrt(dx^2 + dy^2)` over two `(x, y)` f64 streams; a constant side is fed +/// as `repeat(c)`. +fn point_distances( + xa: impl Iterator, + ya: impl Iterator, + xb: impl Iterator, + yb: impl Iterator, +) -> ArrayRef { + let distances = xa.zip(ya).zip(xb.zip(yb)).map(|((xa, ya), (xb, yb))| { + let (dx, dy) = (xa - xb, ya - yb); + (dx * dx + dy * dy).sqrt() + }); + PrimitiveArray::from_iter(distances).into_array() +} + +/// Whether `dtype` is the native `Point` extension (eligible for the columnar fast path). +fn is_point(dtype: &DType) -> bool { + dtype + .as_extension_opt() + .is_some_and(|ext| ext.is::()) +} + +/// A non-nullable native `Point`, a column operand the fast path can read straight from `x`/`y`. +fn is_nonnull_point(dtype: &DType) -> bool { + is_point(dtype) && !dtype.is_nullable() +} + #[cfg(test)] mod tests { use vortex_array::ArrayRef; From 666e132b9436c38f2fc9715d1de10b7ebb72e7d5 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Thu, 2 Jul 2026 14:02:36 -0400 Subject: [PATCH 2/4] chore: address PR review comments - Move the ST_DWithin restore pass to the optimizer extension's pre_optimize_function. DuckDB runs all extensions' pre-optimize hooks before any post-optimize pass, so this guarantees the restore precedes spatial's spatial-join optimization without relying on registration order. Since it now runs before filter pushdown, restrict it to join conditions so filters keep the visible radius they need to push. - Compare the function name with `==`; DuckDB identifiers are already case-insensitive. - Terser doc comment on `duckdb_vx_register_st_dwithin_override`. - Drop a debug print from the geo function converter. Signed-off-by: Nemo Yu --- vortex-duckdb/cpp/include/expr.h | 7 ++----- .../cpp/include/scalar_fn_pushdown.hpp | 6 +++--- vortex-duckdb/cpp/scalar_fn_pushdown.cpp | 20 +++++++++++++++++-- vortex-duckdb/cpp/vortex_duckdb.cpp | 9 ++++++--- vortex-duckdb/src/convert/expr.rs | 1 - 5 files changed, 29 insertions(+), 14 deletions(-) diff --git a/vortex-duckdb/cpp/include/expr.h b/vortex-duckdb/cpp/include/expr.h index d71e4e270a7..792ebf7aa10 100644 --- a/vortex-duckdb/cpp/include/expr.h +++ b/vortex-duckdb/cpp/include/expr.h @@ -13,11 +13,8 @@ typedef struct duckdb_vx_sfunc_ *duckdb_vx_sfunc; const char *duckdb_vx_sfunc_name(duckdb_vx_sfunc ffi_func); -/// Shadow the spatial extension's `ST_DWithin` in the user catalog with a copy whose bind keeps -/// the radius as the third argument (spatial's bind folds it into private bind data), so radius -/// predicates can push into Vortex scans. Unpushed occurrences still execute via spatial's own -/// 3-column code path, and the Vortex optimizer pass rebinds them (e.g. join conditions) back to -/// spatial's folded form, which its spatial-join optimization requires. No-op without `spatial`. +/// Shadow `ST_DWithin` with a copy that keeps the radius as the third argument, so +/// radius filters can push into Vortex scans. duckdb_state duckdb_vx_register_st_dwithin_override(duckdb_database ffi_db); typedef struct duckdb_vx_expr_ *duckdb_vx_expr; diff --git a/vortex-duckdb/cpp/include/scalar_fn_pushdown.hpp b/vortex-duckdb/cpp/include/scalar_fn_pushdown.hpp index 00a47216672..3853b66e8eb 100644 --- a/vortex-duckdb/cpp/include/scalar_fn_pushdown.hpp +++ b/vortex-duckdb/cpp/include/scalar_fn_pushdown.hpp @@ -78,10 +78,10 @@ using Projections = unordered_map; LogicalOperatorPtr TryPushdownScalarFunctions(ClientContext &context, LogicalOperatorPtr plan); /* - * Rebind remaining 3-argument `st_dwithin` calls (bound via the Vortex override; see expr.h) + * Rebind 3-argument `st_dwithin` join conditions (bound via the Vortex override; see expr.h) * back through spatial's original entry, whose folded 2-argument form its spatial-join - * optimization requires. Runs after DuckDB's optimizers, when pushable filters are already - * absorbed into Vortex scans. + * optimization requires. Runs in the pre-optimize hook, which is guaranteed to precede every + * extension's post-optimize pass; filters are left untouched so they can push into Vortex scans. */ void RestoreStDWithin(ClientContext &context, LogicalOperator &plan); diff --git a/vortex-duckdb/cpp/scalar_fn_pushdown.cpp b/vortex-duckdb/cpp/scalar_fn_pushdown.cpp index 0f66033e574..f836bdf0405 100644 --- a/vortex-duckdb/cpp/scalar_fn_pushdown.cpp +++ b/vortex-duckdb/cpp/scalar_fn_pushdown.cpp @@ -2,7 +2,6 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors #include "duckdb/catalog/catalog.hpp" #include "duckdb/catalog/catalog_entry/scalar_function_catalog_entry.hpp" -#include "duckdb/common/string_util.hpp" #include "duckdb/function/function_binder.hpp" #include "duckdb/planner/operator/logical_projection.hpp" #include "scalar_fn_pushdown.hpp" @@ -244,8 +243,25 @@ class StDWithinRestore final : public LogicalOperatorVisitor { explicit StDWithinRestore(ClientContext &context) : context(context) { } + // Restore join conditions, filters must keep the radius visible so + // DuckDB's filter pushdown can offer them to Vortex scans. + void VisitOperator(LogicalOperator &op) override { + using enum LogicalOperatorType; + switch (op.type) { + case LOGICAL_COMPARISON_JOIN: + case LOGICAL_ANY_JOIN: + case LOGICAL_DELIM_JOIN: + case LOGICAL_ASOF_JOIN: + VisitOperatorExpressions(op); + break; + default: + break; + } + VisitOperatorChildren(op); + } + ExpressionPtr VisitReplace(BoundFunctionExpression &expr, ExpressionPtr *) override { - if (expr.children.size() != 3 || !StringUtil::CIEquals(expr.function.name, "st_dwithin")) { + if (expr.children.size() != 3 || expr.function.name != "st_dwithin") { return nullptr; // Not the override's shape: keep it and descend into its children. } // The system catalog holds spatial's original; the user-catalog override cannot shadow diff --git a/vortex-duckdb/cpp/vortex_duckdb.cpp b/vortex-duckdb/cpp/vortex_duckdb.cpp index 0a0cee3bc6b..86537368841 100644 --- a/vortex-duckdb/cpp/vortex_duckdb.cpp +++ b/vortex-duckdb/cpp/vortex_duckdb.cpp @@ -269,13 +269,16 @@ extern "C" duckdb_blob duckdb_vx_value_get_geometry(duckdb_value value) { static void VortexOptimizeFunction(OptimizerExtensionInput &input, unique_ptr &plan) { plan = TryPushdownScalarFunctions(input.context, std::move(plan)); - // Runs before spatial's own optimizer pass: this extension registers at database open, - // spatial's at LOAD, and DuckDB invokes optimizer extensions in registration order. +} + +static void VortexPreOptimizeFunction(OptimizerExtensionInput &input, + unique_ptr &plan) { RestoreStDWithin(input.context, *plan); } struct VortexOptimizerExtension final : OptimizerExtension { - inline VortexOptimizerExtension() : OptimizerExtension(VortexOptimizeFunction, nullptr, {}) { + inline VortexOptimizerExtension() + : OptimizerExtension(VortexOptimizeFunction, VortexPreOptimizeFunction, {}) { } }; diff --git a/vortex-duckdb/src/convert/expr.rs b/vortex-duckdb/src/convert/expr.rs index ae0f7787ad8..41a236cd3f9 100644 --- a/vortex-duckdb/src/convert/expr.rs +++ b/vortex-duckdb/src/convert/expr.rs @@ -148,7 +148,6 @@ fn try_from_geo_function( ) -> VortexResult> { // Catch-all for every bound function: reject non-geo names before touching the children. if !is_geo_function(name) { - debug!("bound function {name}"); return Ok(None); } let children: Vec<_> = func.children().collect(); From 1e6f08e220117fd1f96ac951ee9873a43a4a5977 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Thu, 2 Jul 2026 17:29:20 -0400 Subject: [PATCH 3/4] feat(vortex-geo): geometry bounding-box zone-map statistic + distance-filter pruning Add a `GeometryBounds` aggregate that computes a per-chunk 2D minimum bounding rectangle (`Struct`) for native geometry columns and stores it as a zone-map statistic, plus a `GeoDistanceBoundsPrune` stats-rewrite rule that skips a chunk when its MBR is disjoint from a `GeoDistance(geom, const) <= r` query box. - vortex-array: aggregates self-declare as default zone stats via a `zone_stat_default` vtable/plugin hook; the zoned writer discovers them from the session registry in deterministic (id-sorted) order. - vortex-geo: `GeometryBounds` covers Point/Polygon/MultiPolygon; the prune rule guards on `is_native_geometry`, ignores a NaN radius, and only handles the near forms `<=` / `<`. Missing stats bind to null, so older files degrade to no pruning. Signed-off-by: Nemo Yu --- Cargo.lock | 3 + vortex-array/src/aggregate_fn/plugin.rs | 11 + vortex-array/src/aggregate_fn/session.rs | 15 + vortex-array/src/aggregate_fn/vtable.rs | 6 + vortex-geo/Cargo.toml | 1 + vortex-geo/src/aggregate_fn/bounds.rs | 437 ++++++++++++++++++++++ vortex-geo/src/aggregate_fn/mod.rs | 8 + vortex-geo/src/extension/mod.rs | 33 ++ vortex-geo/src/extension/multipolygon.rs | 19 + vortex-geo/src/extension/point.rs | 9 + vortex-geo/src/extension/polygon.rs | 23 ++ vortex-geo/src/lib.rs | 12 + vortex-geo/src/prune.rs | 329 ++++++++++++++++ vortex-geo/src/test_harness.rs | 106 +++++- vortex-layout/src/layouts/zoned/writer.rs | 18 +- 15 files changed, 1017 insertions(+), 13 deletions(-) create mode 100644 vortex-geo/src/aggregate_fn/bounds.rs create mode 100644 vortex-geo/src/aggregate_fn/mod.rs create mode 100644 vortex-geo/src/prune.rs diff --git a/Cargo.lock b/Cargo.lock index d99d2124e45..c973fa7dd7d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9897,6 +9897,8 @@ dependencies = [ "bzip2", "clap", "futures", + "geo", + "geo-traits", "geoarrow", "geoarrow-cast", "get_dir", @@ -10417,6 +10419,7 @@ dependencies = [ "rstest", "vortex-array", "vortex-error", + "vortex-layout", "vortex-session", "wkb", ] diff --git a/vortex-array/src/aggregate_fn/plugin.rs b/vortex-array/src/aggregate_fn/plugin.rs index b7ff8b893ac..d07008a7338 100644 --- a/vortex-array/src/aggregate_fn/plugin.rs +++ b/vortex-array/src/aggregate_fn/plugin.rs @@ -10,6 +10,7 @@ use crate::aggregate_fn::AggregateFn; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnVTable; +use crate::dtype::DType; /// Reference-counted pointer to an aggregate function plugin. pub type AggregateFnPluginRef = Arc; @@ -28,6 +29,12 @@ pub trait AggregateFnPlugin: 'static + Send + Sync { /// Deserialize an aggregate function from serialized metadata. fn deserialize(&self, metadata: &[u8], session: &VortexSession) -> VortexResult; + + /// The default per-chunk zone statistic to store for a column of `input_dtype`, or `None` if + /// this aggregate isn't one. + fn zone_stat_default(&self, _input_dtype: &DType) -> Option { + None + } } impl std::fmt::Debug for dyn AggregateFnPlugin { @@ -51,4 +58,8 @@ impl AggregateFnPlugin for V { let options = AggregateFnVTable::deserialize(self, metadata, session)?; Ok(AggregateFn::new(self.clone(), options).erased()) } + + fn zone_stat_default(&self, input_dtype: &DType) -> Option { + AggregateFnVTable::zone_stat_default(self, input_dtype) + } } diff --git a/vortex-array/src/aggregate_fn/session.rs b/vortex-array/src/aggregate_fn/session.rs index c496f7a9bf8..0528e88b4e6 100644 --- a/vortex-array/src/aggregate_fn/session.rs +++ b/vortex-array/src/aggregate_fn/session.rs @@ -10,6 +10,7 @@ use vortex_session::SessionVar; use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnPluginRef; +use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::fns::all_nan::AllNan; use crate::aggregate_fn::fns::all_non_distinct::AllNonDistinct; @@ -44,6 +45,7 @@ use crate::arrays::chunked::compute::aggregate::ChunkedArrayAggregate; use crate::arrays::dict::compute::is_constant::DictIsConstantKernel; use crate::arrays::dict::compute::is_sorted::DictIsSortedKernel; use crate::arrays::dict::compute::min_max::DictMinMaxKernel; +use crate::dtype::DType; /// Session state for aggregate functions and encoding-specific aggregate kernels. /// @@ -134,6 +136,19 @@ impl AggregateFnSession { self.registry.insert(id, pluginref); } + /// The default per-chunk zone statistics for a column of `input_dtype`, collected from every + /// registered aggregate's `zone_stat_default`. + pub fn zone_stat_defaults(&self, input_dtype: &DType) -> Vec { + self.registry.read(|registry| { + let mut fns: Vec = registry + .values() + .filter_map(|plugin| plugin.zone_stat_default(input_dtype)) + .collect(); + fns.sort_by_key(|f| f.id()); + fns + }) + } + /// Returns the aggregate kernel registered for `array_id` and `agg_fn_id`, if any. /// /// Lookup first checks for a kernel registered for the exact aggregate function, then falls diff --git a/vortex-array/src/aggregate_fn/vtable.rs b/vortex-array/src/aggregate_fn/vtable.rs index 49b28dd26d7..3f0a4cf567c 100644 --- a/vortex-array/src/aggregate_fn/vtable.rs +++ b/vortex-array/src/aggregate_fn/vtable.rs @@ -93,6 +93,12 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync { /// Returns `None` if the aggregate function cannot be applied to the input dtype. fn return_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option; + /// If this aggregate should be computed as a default zone statistic for `input_dtype`, return + /// the bound aggregate to store. Default: not a zone-map default. + fn zone_stat_default(&self, _input_dtype: &DType) -> Option { + None + } + /// DType of the intermediate partial accumulator state. /// /// Use a struct dtype when multiple fields are needed diff --git a/vortex-geo/Cargo.toml b/vortex-geo/Cargo.toml index 2f0583b49e6..068831b6a62 100644 --- a/vortex-geo/Cargo.toml +++ b/vortex-geo/Cargo.toml @@ -29,6 +29,7 @@ wkb = { workspace = true } [dev-dependencies] rstest = { workspace = true } +vortex-layout = { workspace = true } [lints] workspace = true diff --git a/vortex-geo/src/aggregate_fn/bounds.rs b/vortex-geo/src/aggregate_fn/bounds.rs new file mode 100644 index 00000000000..4b5bdadb59f --- /dev/null +++ b/vortex-geo/src/aggregate_fn/bounds.rs @@ -0,0 +1,437 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! An aggregate computing the 2D minimum bounding rectangle of a native geometry column as +//! `Struct`. Stored as a zone statistic, it lets spatial filters prune +//! chunks whose bounding box cannot intersect the query region. + +use vortex_array::ArrayRef; +use vortex_array::Columnar; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::aggregate_fn::AggregateFnId; +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::AggregateFnVTableExt; +use vortex_array::aggregate_fn::EmptyOptions; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::dtype::StructFields; +use vortex_array::scalar::Scalar; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::extension::coordinates; +use crate::extension::is_native_geometry; + +/// Aggregate computing the 2D minimum bounding rectangle of a native geometry column. +#[derive(Clone, Debug)] +pub struct GeometryBounds; + +/// An axis-aligned bounding box. +#[derive(Clone, Copy)] +struct Bbox { + xmin: f64, + ymin: f64, + xmax: f64, + ymax: f64, +} + +impl Bbox { + /// The smallest box containing both `self` and `other`. + fn union(self, other: Bbox) -> Bbox { + Bbox { + xmin: self.xmin.min(other.xmin), + ymin: self.ymin.min(other.ymin), + xmax: self.xmax.max(other.xmax), + ymax: self.ymax.max(other.ymax), + } + } +} + +/// Partial MBR accumulator: the union of every bounding box seen so far, or `None` when empty. +pub struct BoundsPartial { + bbox: Option, +} + +impl BoundsPartial { + fn merge(&mut self, other: Bbox) { + self.bbox = Some(match self.bbox { + Some(cur) => cur.union(other), + None => other, + }); + } +} + +/// `Struct` of `f64`. Nullable so an empty group is a null MBR; the fields +/// are nullable too, so the pruning proof's `get_item` keeps a consistent nullable dtype. +fn bounds_dtype() -> DType { + let coord = DType::Primitive(PType::F64, Nullability::Nullable); + DType::Struct( + StructFields::from_iter([ + ("xmin", coord.clone()), + ("ymin", coord.clone()), + ("xmax", coord.clone()), + ("ymax", coord), + ]), + Nullability::Nullable, + ) +} + +/// The bounding box of the coordinate slices, or `None` for an empty chunk. +fn bounds_of(xs: &[f64], ys: &[f64]) -> Option { + if xs.is_empty() { + return None; + } + let min_max = |vals: &[f64]| { + vals.iter() + .fold((f64::INFINITY, f64::NEG_INFINITY), |(lo, hi), &v| { + (lo.min(v), hi.max(v)) + }) + }; + let (xmin, xmax) = min_max(xs); + let (ymin, ymax) = min_max(ys); + Some(Bbox { + xmin, + ymin, + xmax, + ymax, + }) +} + +impl AggregateFnVTable for GeometryBounds { + type Options = EmptyOptions; + type Partial = BoundsPartial; + + fn id(&self) -> AggregateFnId { + static ID: CachedId = CachedId::new("vortex.geo.bounds"); + *ID + } + + // Serializable so the zoned writer can persist this as a per-chunk stat. No options to encode. + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { + Ok(Some(vec![])) + } + + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) + } + + fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option { + is_native_geometry(input_dtype).then(bounds_dtype) + } + + fn zone_stat_default(&self, input_dtype: &DType) -> Option { + is_native_geometry(input_dtype).then(|| self.bind(EmptyOptions)) + } + + fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { + self.return_dtype(options, input_dtype) + } + + fn empty_partial( + &self, + _options: &Self::Options, + _input_dtype: &DType, + ) -> VortexResult { + Ok(BoundsPartial { bbox: None }) + } + + fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { + if other.is_null() { + return Ok(()); + } + let fields = other.as_struct(); + let read = |name: &str| -> VortexResult { + f64::try_from( + &fields + .field(name) + .ok_or_else(|| vortex_err!("bounds missing {name}"))?, + ) + }; + partial.merge(Bbox { + xmin: read("xmin")?, + ymin: read("ymin")?, + xmax: read("xmax")?, + ymax: read("ymax")?, + }); + Ok(()) + } + + fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + Ok(match partial.bbox { + Some(b) => Scalar::struct_( + bounds_dtype(), + vec![ + Scalar::primitive(b.xmin, Nullability::Nullable), + Scalar::primitive(b.ymin, Nullability::Nullable), + Scalar::primitive(b.xmax, Nullability::Nullable), + Scalar::primitive(b.ymax, Nullability::Nullable), + ], + ), + None => Scalar::null(bounds_dtype()), + }) + } + + fn reset(&self, partial: &mut Self::Partial) { + partial.bbox = None; + } + + fn is_saturated(&self, _partial: &Self::Partial) -> bool { + // A bounding box can always grow, so it is never saturated. + false + } + + fn accumulate( + &self, + partial: &mut Self::Partial, + batch: &Columnar, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + let array = match batch { + Columnar::Canonical(canonical) => canonical.clone().into_array(), + Columnar::Constant(constant) => constant.clone().into_array(), + }; + let coords = coordinates(&array, ctx)?; + let xs = coords + .unmasked_field_by_name("x")? + .clone() + .execute::(ctx)?; + let ys = coords + .unmasked_field_by_name("y")? + .clone() + .execute::(ctx)?; + if let Some(bbox) = bounds_of(xs.as_slice::(), ys.as_slice::()) { + partial.merge(bbox); + } + Ok(()) + } + + fn finalize(&self, partials: ArrayRef) -> VortexResult { + // The stored partial is already the MBR struct, so finalizing is the identity. + Ok(partials) + } + + fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { + self.to_scalar(partial) + } +} + +#[cfg(test)] +mod tests { + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::Accumulator; + use vortex_array::aggregate_fn::AggregateFnVTable; + use vortex_array::aggregate_fn::DynAccumulator; + use vortex_array::aggregate_fn::EmptyOptions; + use vortex_array::aggregate_fn::session::AggregateFnSessionExt; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::scalar::Scalar; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + + use super::Bbox; + use super::BoundsPartial; + use super::GeometryBounds; + use super::bounds_dtype; + use crate::test_harness::multipolygon_column; + use crate::test_harness::point_column; + use crate::test_harness::polygon_column; + + /// The aggregate must be serializable so the zoned writer can persist its zone-stat descriptor. + #[test] + fn serializes_for_zone_storage() -> VortexResult<()> { + let session = vortex_array::array_session(); + let metadata = GeometryBounds + .serialize(&EmptyOptions)? + .expect("GeometryBounds must be serializable to be stored as a zone statistic"); + GeometryBounds.deserialize(&metadata, &session)?; + Ok(()) + } + + /// The MBR result's corners as `(xmin, ymin, xmax, ymax)`. + fn mbr(result: &Scalar) -> VortexResult<(f64, f64, f64, f64)> { + let fields = result.as_struct(); + let read = |name: &str| -> VortexResult { + f64::try_from( + &fields + .field(name) + .ok_or_else(|| vortex_err!("missing {name}"))?, + ) + }; + Ok((read("xmin")?, read("ymin")?, read("xmax")?, read("ymax")?)) + } + + /// The MBR of a Point column is the min/max of its coordinates, accumulated across batches. + #[test] + fn point_bounds_across_batches() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let mut acc = Accumulator::try_new(GeometryBounds, EmptyOptions, dtype)?; + + acc.accumulate(&point_column(vec![1.0, 3.0], vec![2.0, 4.0])?, &mut ctx)?; + acc.accumulate(&point_column(vec![-1.0], vec![5.0])?, &mut ctx)?; + + assert_eq!(mbr(&acc.finish()?)?, (-1.0, 2.0, 3.0, 5.0)); + Ok(()) + } + + /// The MBR of a Polygon column unions every ring vertex — exercising the `List>` + /// unwrap, not just the bare Point struct. + #[test] + fn polygon_bounds_union_all_vertices() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + // Two rectangles: (0,0)-(2,3) and (5,5)-(7,8). The chunk MBR is their union: (0,0)-(7,8). + let polygons = polygon_column(vec![ + vec![vec![(0.0, 0.0), (2.0, 0.0), (2.0, 3.0), (0.0, 3.0)]], + vec![vec![(5.0, 5.0), (7.0, 5.0), (7.0, 8.0), (5.0, 8.0)]], + ])?; + let dtype = polygons.dtype().clone(); + let mut acc = Accumulator::try_new(GeometryBounds, EmptyOptions, dtype)?; + acc.accumulate(&polygons, &mut ctx)?; + + assert_eq!(mbr(&acc.finish()?)?, (0.0, 0.0, 7.0, 8.0)); + Ok(()) + } + + /// The MBR of a MultiPolygon column unions every vertex of every polygon's rings — exercising + /// the triple-`List` unwrap. + #[test] + fn multipolygon_bounds_union_all_vertices() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + // Multipolygon 0: squares (0,0)-(1,1) and (4,4)-(5,5); multipolygon 1: square (-3,7)-(-2,9). + let multipolygons = multipolygon_column(vec![ + vec![ + vec![vec![(0.0, 0.0), (1.0, 0.0), (1.0, 1.0), (0.0, 1.0)]], + vec![vec![(4.0, 4.0), (5.0, 4.0), (5.0, 5.0), (4.0, 5.0)]], + ], + vec![vec![vec![ + (-3.0, 7.0), + (-2.0, 7.0), + (-2.0, 9.0), + (-3.0, 9.0), + ]]], + ])?; + let dtype = multipolygons.dtype().clone(); + let mut acc = Accumulator::try_new(GeometryBounds, EmptyOptions, dtype)?; + acc.accumulate(&multipolygons, &mut ctx)?; + + assert_eq!(mbr(&acc.finish()?)?, (-3.0, 0.0, 5.0, 9.0)); + Ok(()) + } + + /// `combine_partials` unions partial boxes — the path the zoned writer takes when a zone's + /// array is chunked. + #[test] + fn combine_partials_unions_boxes() -> VortexResult<()> { + let bbox = |xmin, ymin, xmax, ymax| BoundsPartial { + bbox: Some(Bbox { + xmin, + ymin, + xmax, + ymax, + }), + }; + let mut partial = BoundsPartial { bbox: None }; + GeometryBounds.combine_partials( + &mut partial, + GeometryBounds.to_scalar(&bbox(0.0, 0.0, 1.0, 1.0))?, + )?; + GeometryBounds.combine_partials( + &mut partial, + GeometryBounds.to_scalar(&bbox(5.0, -2.0, 7.0, 3.0))?, + )?; + assert_eq!( + mbr(&GeometryBounds.to_scalar(&partial)?)?, + (0.0, -2.0, 7.0, 3.0) + ); + Ok(()) + } + + /// A null partial (an empty group's MBR) is a no-op in `combine_partials`. + #[test] + fn combine_partials_ignores_null() -> VortexResult<()> { + let mut partial = BoundsPartial { + bbox: Some(Bbox { + xmin: 0.0, + ymin: 0.0, + xmax: 1.0, + ymax: 1.0, + }), + }; + GeometryBounds.combine_partials(&mut partial, Scalar::null(bounds_dtype()))?; + assert_eq!( + mbr(&GeometryBounds.to_scalar(&partial)?)?, + (0.0, 0.0, 1.0, 1.0) + ); + Ok(()) + } + + /// All-NaN coordinates fold to an inverted box (min > max). Sound to store: the pruning proof + /// then skips the chunk, and NaN-coordinate rows can never satisfy `distance <= r` anyway. + #[test] + fn all_nan_coordinates_yield_inverted_box() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let column = point_column(vec![f64::NAN, f64::NAN], vec![f64::NAN, f64::NAN])?; + let mut acc = Accumulator::try_new(GeometryBounds, EmptyOptions, column.dtype().clone())?; + acc.accumulate(&column, &mut ctx)?; + + let (xmin, ymin, xmax, ymax) = mbr(&acc.finish()?)?; + assert!(xmin > xmax && ymin > ymax); + Ok(()) + } + + /// An empty group yields a null MBR. + #[test] + fn empty_group_is_null() -> VortexResult<()> { + let dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let mut acc = Accumulator::try_new(GeometryBounds, EmptyOptions, dtype)?; + assert!(acc.finish()?.is_null()); + Ok(()) + } + + /// After `initialize`, the registry yields a default zone statistic for geometry columns (so the + /// zoned writer stores it) but none for ordinary numeric columns. + #[test] + fn registered_as_geometry_zone_default() -> VortexResult<()> { + let session = vortex_array::array_session(); + crate::initialize(&session); + + let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + assert!( + !session + .aggregate_fns() + .zone_stat_defaults(&point_dtype) + .is_empty(), + "a geometry zone-stat default should be discovered for Point columns" + ); + let i32_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + assert!( + session + .aggregate_fns() + .zone_stat_defaults(&i32_dtype) + .is_empty(), + "no geometry zone-stat default should apply to numeric columns" + ); + Ok(()) + } +} diff --git a/vortex-geo/src/aggregate_fn/mod.rs b/vortex-geo/src/aggregate_fn/mod.rs new file mode 100644 index 00000000000..499355beaac --- /dev/null +++ b/vortex-geo/src/aggregate_fn/mod.rs @@ -0,0 +1,8 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Aggregate functions over geometry columns. + +mod bounds; + +pub use bounds::*; diff --git a/vortex-geo/src/extension/mod.rs b/vortex-geo/src/extension/mod.rs index 4e4896aa5e5..e3a050bc64e 100644 --- a/vortex-geo/src/extension/mod.rs +++ b/vortex-geo/src/extension/mod.rs @@ -33,8 +33,10 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::StructArray; use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::arrow::FromArrowArray; +use vortex_array::dtype::DType; use vortex_array::dtype::extension::ExtDType; use vortex_array::dtype::extension::ExtVTable; use vortex_array::scalar::Scalar; @@ -43,6 +45,37 @@ use vortex_error::vortex_bail; use vortex_error::vortex_err; pub use wkb::*; +/// Whether `dtype` is a native geometry extension. +pub(crate) fn is_native_geometry(dtype: &DType) -> bool { + dtype + .as_extension_opt() + .is_some_and(|ext| ext.is::() || ext.is::() || ext.is::()) +} + +/// The flat coordinate `Struct` of a native geometry column. +pub(crate) fn coordinates(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let Some(ext) = array.dtype().as_extension_opt() else { + vortex_bail!( + "geo: operand is not a geometry extension type, was {}", + array.dtype() + ); + }; + let storage = array + .clone() + .execute::(ctx)? + .storage_array() + .clone(); + if ext.is::() { + point_coordinates(&storage, ctx) + } else if ext.is::() { + polygon_coordinates(&storage, ctx) + } else if ext.is::() { + multipolygon_coordinates(&storage, ctx) + } else { + vortex_bail!("geo: unsupported geometry extension {}", array.dtype()) + } +} + /// Decode a native geometry column to `geo_types`. A non-geometry operand is an error. pub(crate) fn geometries( array: &ArrayRef, diff --git a/vortex-geo/src/extension/multipolygon.rs b/vortex-geo/src/extension/multipolygon.rs index 67dac438970..3af3ec5adc7 100644 --- a/vortex-geo/src/extension/multipolygon.rs +++ b/vortex-geo/src/extension/multipolygon.rs @@ -20,10 +20,13 @@ use geoarrow::datatypes::CoordType; use geoarrow::datatypes::MultiPolygonType; use prost::Message; use vortex_array::ArrayRef; +use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::StructArray; use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::arrays::listview::ListViewArrayExt; use vortex_array::arrow::ArrowExport; use vortex_array::arrow::ArrowExportVTable; use vortex_array::arrow::ArrowImport; @@ -53,6 +56,7 @@ use super::coordinate::coordinate_storage_dtype; use super::geo_metadata_from_arrow; use super::geoarrow_metadata; use super::geoarrow_to_wkb; +use super::polygon_coordinates; /// A multipolygon (`geoarrow.multipolygon`); a single `Polygon` is a one-element multipolygon. #[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] @@ -117,6 +121,21 @@ fn multipolygon_type(geo_metadata: &GeoMetadata, dimension: Dimension) -> MultiP MultiPolygonType::new(dimension.into(), geoarrow_metadata(geo_metadata)) } +/// The coordinate `Struct` of `MultiPolygon` storage, flattening all three `List` +/// levels to every vertex of every ring of every polygon. +pub(crate) fn multipolygon_coordinates( + storage: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let polygons = storage + .clone() + .execute::(ctx)? + .into_listview() + .elements() + .clone(); + polygon_coordinates(&polygons, ctx) +} + /// Decode storage to `geo_types` for the geo scalar functions (CRS is irrelevant to planar ops). pub(crate) fn multipolygon_geometries( storage: &ArrayRef, diff --git a/vortex-geo/src/extension/point.rs b/vortex-geo/src/extension/point.rs index 6371105ca25..7c4470ed3a7 100644 --- a/vortex-geo/src/extension/point.rs +++ b/vortex-geo/src/extension/point.rs @@ -22,6 +22,7 @@ use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::StructArray; use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::arrow::ArrowExport; use vortex_array::arrow::ArrowExportVTable; @@ -120,6 +121,14 @@ impl PointData { } } +/// The coordinate `Struct` of `Point` storage. +pub(crate) fn point_coordinates( + storage: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + storage.clone().execute::(ctx) +} + /// 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 { diff --git a/vortex-geo/src/extension/polygon.rs b/vortex-geo/src/extension/polygon.rs index aeda30931a2..7cc43aeb2d4 100644 --- a/vortex-geo/src/extension/polygon.rs +++ b/vortex-geo/src/extension/polygon.rs @@ -20,10 +20,13 @@ use geoarrow::datatypes::CoordType; use geoarrow::datatypes::PolygonType; use prost::Message; use vortex_array::ArrayRef; +use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::StructArray; use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::arrays::listview::ListViewArrayExt; use vortex_array::arrow::ArrowExport; use vortex_array::arrow::ArrowExportVTable; use vortex_array::arrow::ArrowImport; @@ -114,6 +117,26 @@ fn polygon_type(geo_metadata: &GeoMetadata, dimension: Dimension) -> PolygonType PolygonType::new(dimension.into(), geoarrow_metadata(geo_metadata)) } +/// The coordinate `Struct` of `Polygon` storage, flattening both `List` levels of +/// `List>` to every vertex of every ring. +pub(crate) fn polygon_coordinates( + storage: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let rings = storage + .clone() + .execute::(ctx)? + .into_listview() + .elements() + .clone(); + let coords = rings + .execute::(ctx)? + .into_listview() + .elements() + .clone(); + coords.execute::(ctx) +} + /// Decode `Polygon` storage (`List>`) to `geo_types` polygons, for the geo scalar /// functions. CRS does not affect planar geometry ops, so default metadata is used. pub(crate) fn polygon_geometries( diff --git a/vortex-geo/src/lib.rs b/vortex-geo/src/lib.rs index 2cc8004efc5..4405e20c3d6 100644 --- a/vortex-geo/src/lib.rs +++ b/vortex-geo/src/lib.rs @@ -3,18 +3,24 @@ use std::sync::Arc; +use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::arrow::ArrowSessionExt; use vortex_array::dtype::session::DTypeSessionExt; use vortex_array::scalar_fn::session::ScalarFnSessionExt; +use vortex_array::stats::session::StatsSessionExt; use vortex_session::VortexSession; +use crate::aggregate_fn::GeometryBounds; use crate::extension::MultiPolygon; use crate::extension::Point; use crate::extension::Polygon; use crate::extension::WellKnownBinary; +use crate::prune::GeoDistanceBoundsPrune; use crate::scalar_fn::distance::GeoDistance; +pub mod aggregate_fn; pub mod extension; +pub mod prune; pub mod scalar_fn; #[cfg(test)] mod test_harness; @@ -39,4 +45,10 @@ pub fn initialize(session: &VortexSession) { // Register the geometry scalar functions. session.scalar_fns().register(GeoDistance); + + // The bounding-box aggregate; self-declares as a per-chunk zone stat for geometry columns. + session.aggregate_fns().register(GeometryBounds); + + // Register the spatial pruning rule that uses that bounding box. + session.stats().register_rewrite(GeoDistanceBoundsPrune); } diff --git a/vortex-geo/src/prune.rs b/vortex-geo/src/prune.rs new file mode 100644 index 00000000000..11023637724 --- /dev/null +++ b/vortex-geo/src/prune.rs @@ -0,0 +1,329 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Stats-rewrite pruning for spatial filters, backed by the per-chunk [`GeometryBounds`] MBR. +//! +//! [`GeoDistanceBoundsPrune`] falsifies `ST_Distance(geom, const) <= r` (or `< r`): any row within +//! `r` of the constant must lie inside the constant's bounding box expanded by `r`, so if a chunk's +//! geometry MBR is disjoint from that expanded box, no row in the chunk can match and the chunk is +//! skipped. +//! +//! # Limitations +//! +//! Only the "near" forms `<= r` / `< r` are handled — the predicates a radius/within search uses. +//! They prune via the MBR's *lower* distance bound (the nearest point of the box to `const`). Every +//! other comparison falls through to `None`, leaving the chunk to be scanned — correct, just not +//! pruned: +//! +//! - `> r` / `>= r` are *soundly* prunable via the symmetric upper bound (the farthest corner of the +//! MBR being `<= r` proves every geometry is within `r`), but are intentionally omitted: "far from" +//! filters are rare and rarely selective, so the prune would almost never fire. +//! - `== r` would need both bounds at once and is not a realistic query. +//! - `!= r` is unprunable: a bounding box cannot prove every row sits at exactly distance `r`. + +use geo::BoundingRect; +use vortex_array::VortexSessionExecute; +use vortex_array::aggregate_fn::AggregateFnVTableExt; +use vortex_array::aggregate_fn::EmptyOptions; +use vortex_array::expr::Expression; +use vortex_array::expr::get_item; +use vortex_array::expr::gt; +use vortex_array::expr::is_root; +use vortex_array::expr::lit; +use vortex_array::expr::lt; +use vortex_array::expr::or; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::fns::binary::Binary; +use vortex_array::scalar_fn::fns::literal::Literal; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::stats::rewrite::StatsRewriteCtx; +use vortex_array::stats::rewrite::StatsRewriteRule; +use vortex_array::stats::stat; +use vortex_error::VortexResult; + +use crate::aggregate_fn::GeometryBounds; +use crate::extension::is_native_geometry; +use crate::extension::single_geometry; +use crate::scalar_fn::distance::GeoDistance; + +/// Prunes chunks for `GeoDistance(geom, const) <= r` / `< r` using the chunk's [`GeometryBounds`] +/// MBR. Registered against the comparison's scalar-function id, since the comparison — not +/// `GeoDistance` — is the predicate root. +#[derive(Debug)] +pub struct GeoDistanceBoundsPrune; + +impl StatsRewriteRule for GeoDistanceBoundsPrune { + fn scalar_fn_id(&self) -> ScalarFnId { + Binary.id() + } + + fn falsify( + &self, + expr: &Expression, + ctx: &StatsRewriteCtx<'_>, + ) -> VortexResult> { + // Only the near forms `<= r` / `< r` prune (see the module Limitations). + match expr.as_::() { + Operator::Lte | Operator::Lt => {} + _ => return Ok(None), + } + let distance = expr.child(0); + let threshold = expr.child(1); + + // The left operand must be `GeoDistance(geom, const)`. + if distance.as_opt::().is_none() { + return Ok(None); + } + + // Distance is symmetric, so the geometry column (scope root) and constant may be on either + // side. + let (lhs, rhs) = (distance.child(0), distance.child(1)); + let (geom, constant) = if is_root(lhs) { + (lhs, rhs) + } else if is_root(rhs) { + (rhs, lhs) + } else { + return Ok(None); + }; + + // The stat reference below only binds for dtypes `GeometryBounds` supports; anything else + // (e.g. a WKB column) must fall through to the scan. + if !is_native_geometry(&ctx.return_dtype(geom)?) { + return Ok(None); + } + + let (Some(constant), Some(radius)) = + (constant.as_opt::(), threshold.as_opt::()) + else { + return Ok(None); + }; + let Ok(radius) = f64::try_from(radius) else { + return Ok(None); + }; + // A NaN radius would poison the expanded box into pruning every zone, while the scan's + // total-order compare treats `distance <= NaN` as true for every row. + if radius.is_nan() { + return Ok(None); + } + + // Bounding box of the constant geometry, expanded by the radius. + let mut exec = ctx.session().create_execution_ctx(); + let Some(rect) = single_geometry(constant, &mut exec)?.bounding_rect() else { + return Ok(None); + }; + let (xmin, xmax) = (rect.min().x - radius, rect.max().x + radius); + let (ymin, ymax) = (rect.min().y - radius, rect.max().y + radius); + + // Prune when the chunk MBR is disjoint from the expanded box on any axis: no row can match. + let mbr = stat(geom.clone(), GeometryBounds.bind(EmptyOptions)); + let proof = or( + or( + lt(get_item("xmax", mbr.clone()), lit(xmin)), + gt(get_item("xmin", mbr.clone()), lit(xmax)), + ), + or( + lt(get_item("ymax", mbr.clone()), lit(ymin)), + gt(get_item("ymin", mbr), lit(ymax)), + ), + ); + Ok(Some(proof)) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::AggregateFnVTableExt; + use vortex_array::aggregate_fn::EmptyOptions as AggregateEmptyOptions; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::StructArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::FieldNames; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::expr::Expression; + use vortex_array::expr::lit; + use vortex_array::expr::lt_eq; + use vortex_array::expr::root; + use vortex_array::scalar_fn::EmptyOptions; + use vortex_array::scalar_fn::ScalarFnVTableExt; + use vortex_array::scalar_fn::fns::binary::Binary; + use vortex_array::scalar_fn::fns::operators::Operator; + use vortex_array::stats::rewrite::StatsRewriteCtx; + use vortex_array::stats::rewrite::StatsRewriteRule; + use vortex_array::validity::Validity; + use vortex_error::VortexResult; + use vortex_layout::layouts::zoned::zone_map::ZoneMap; + + use super::GeoDistanceBoundsPrune; + use crate::aggregate_fn::GeometryBounds; + use crate::scalar_fn::distance::GeoDistance; + use crate::test_harness::point_column; + + /// Run the rule against `GeoDistance(root, origin) radius`, operands swapped when + /// `geom_first` is false. + fn falsify_distance( + operator: Operator, + geom_first: bool, + radius: f64, + ) -> VortexResult> { + let session = vortex_array::array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + + let scope = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let origin = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; + let operands = if geom_first { + [root(), lit(origin)] + } else { + [lit(origin), root()] + }; + let distance = GeoDistance.new_expr(EmptyOptions, operands); + let predicate = Binary.new_expr(operator, [distance, lit(radius)]); + + GeoDistanceBoundsPrune.falsify(&predicate, &StatsRewriteCtx::new(&session, &scope)) + } + + /// Only the upper-bounded "near" forms (`<=`/`<`) are pruned; the rest are left to the scan. + #[rstest] + #[case(Operator::Lte, true)] + #[case(Operator::Lt, true)] + #[case(Operator::Gt, false)] + #[case(Operator::Gte, false)] + #[case(Operator::Eq, false)] + #[case(Operator::NotEq, false)] + fn prunes_only_near_distance( + #[case] operator: Operator, + #[case] prunes: bool, + ) -> VortexResult<()> { + assert_eq!(falsify_distance(operator, true, 0.5)?.is_some(), prunes); + Ok(()) + } + + /// Distance is symmetric: `GeoDistance(const, geom) <= r` falsifies just like the geom-first form. + #[test] + fn falsifies_with_constant_as_left_operand() -> VortexResult<()> { + assert!(falsify_distance(Operator::Lte, false, 0.5)?.is_some()); + Ok(()) + } + + /// A NaN radius must not prune — the scan's total-order compare treats `dist <= NaN` as true. + #[test] + fn nan_radius_never_prunes() -> VortexResult<()> { + assert!(falsify_distance(Operator::Lte, true, f64::NAN)?.is_none()); + Ok(()) + } + + /// A negative radius prunes every zone — vacuously sound: `distance <= r < 0` matches no row. + #[test] + fn negative_radius_prunes_vacuously() -> VortexResult<()> { + assert!(falsify_distance(Operator::Lte, true, -0.5)?.is_some()); + Ok(()) + } + + /// A scope dtype without `GeometryBounds` support gets no proof — the stat reference would + /// fail to bind at prune time. + #[test] + fn unsupported_scope_is_not_pruned() -> VortexResult<()> { + let session = vortex_array::array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + + let scope = DType::Primitive(PType::F64, Nullability::NonNullable); + let origin = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; + let distance = GeoDistance.new_expr(EmptyOptions, [root(), lit(origin)]); + let predicate = lt_eq(distance, lit(0.5f64)); + + let ctx = StatsRewriteCtx::new(&session, &scope); + assert!(GeoDistanceBoundsPrune.falsify(&predicate, &ctx)?.is_none()); + Ok(()) + } + + /// A comparison that does not wrap `GeoDistance` is left untouched. + #[test] + fn ignores_non_distance_comparison() -> VortexResult<()> { + let session = vortex_array::array_session(); + crate::initialize(&session); + let scope = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + + let predicate = lt_eq(lit(1.0f64), lit(2.0f64)); + let ctx = StatsRewriteCtx::new(&session, &scope); + assert!(GeoDistanceBoundsPrune.falsify(&predicate, &ctx)?.is_none()); + Ok(()) + } + + /// End-to-end over a hand-built zone map: the far chunk is skipped, the near one kept. + #[test] + fn prunes_far_chunk_keeps_near() -> VortexResult<()> { + let session = vortex_array::array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + + let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let bounds_fn = GeometryBounds.bind(AggregateEmptyOptions); + + // Two chunks: chunk 0 near the origin (MBR 0,0..1,1), chunk 1 far away (MBR 100,100..101,101). + let coord = + |a: f64, b: f64| PrimitiveArray::from_option_iter([Some(a), Some(b)]).into_array(); + let mbrs = StructArray::try_new( + ["xmin", "ymin", "xmax", "ymax"].into(), + vec![ + coord(0.0, 100.0), + coord(0.0, 100.0), + coord(1.0, 101.0), + coord(1.0, 101.0), + ], + 2, + Validity::AllValid, + )? + .into_array(); + let zone_array = StructArray::from_fields(&[(bounds_fn.to_string().as_str(), mbrs)])?; + let zone_map = + ZoneMap::try_new(point_dtype.clone(), zone_array, Arc::new([bounds_fn]), 1, 2)?; + + let origin = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; + let distance = GeoDistance.new_expr(EmptyOptions, [root(), lit(origin)]); + let predicate = lt_eq(distance, lit(0.5f64)); + let proof = predicate + .falsify(&point_dtype, &session)? + .expect("distance filter should be falsifiable"); + + // `true` means the zone is pruned: chunk 0 (near origin) is kept, chunk 1 (far) is skipped. + let mask = zone_map.prune(&proof, &session)?; + assert_eq!(mask.iter().collect::>(), vec![false, true]); + Ok(()) + } + + /// Backward compat: a zone map written without the `GeometryBounds` stat (an older file) keeps + /// every zone — the missing stat binds to null and `null_as_false` retains the zone. + #[test] + fn missing_bounds_stat_keeps_all_zones() -> VortexResult<()> { + let session = vortex_array::array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + + let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let zone_map = ZoneMap::try_new( + point_dtype.clone(), + StructArray::try_new(FieldNames::empty(), vec![], 2, Validity::NonNullable)?, + Arc::new([]), + 1, + 2, + )?; + + let origin = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; + let distance = GeoDistance.new_expr(EmptyOptions, [root(), lit(origin)]); + let proof = lt_eq(distance, lit(0.5f64)) + .falsify(&point_dtype, &session)? + .expect("distance filter should be falsifiable"); + + let mask = zone_map.prune(&proof, &session)?; + assert_eq!(mask.iter().collect::>(), vec![false, false]); + Ok(()) + } +} diff --git a/vortex-geo/src/test_harness.rs b/vortex-geo/src/test_harness.rs index 2e9e7f43c27..78e528350a9 100644 --- a/vortex-geo/src/test_harness.rs +++ b/vortex-geo/src/test_harness.rs @@ -6,31 +6,121 @@ use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::ListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; +use vortex_array::dtype::Nullability; use vortex_array::dtype::extension::ExtDType; use vortex_array::scalar::Scalar; +use vortex_array::validity::Validity; use vortex_error::VortexResult; +use vortex_error::vortex_err; use crate::extension::GeoMetadata; +use crate::extension::MultiPolygon; use crate::extension::Point; +use crate::extension::Polygon; use crate::extension::coordinate::Coordinate; +use crate::extension::coordinate::Dimension; use crate::extension::coordinate::coordinate_from_struct; +use crate::extension::multipolygon_storage_dtype; +use crate::extension::polygon_storage_dtype; -/// A `Point` column (CRS `EPSG:4326`) over the given x/y coordinates. -pub(crate) fn point_column(xs: Vec, ys: Vec) -> VortexResult { - let storage = StructArray::from_fields(&[ +/// The WGS 84 (`EPSG:4326`) metadata tagged onto test geometry columns. +fn wgs84() -> GeoMetadata { + GeoMetadata { + crs: Some("EPSG:4326".to_string()), + } +} + +/// A coordinate `Struct` over the parallel x/y buffers. +fn xy_struct(xs: Vec, ys: Vec) -> VortexResult { + Ok(StructArray::from_fields(&[ ("x", PrimitiveArray::from_iter(xs).into_array()), ("y", PrimitiveArray::from_iter(ys).into_array()), ])? - .into_array(); - let metadata = GeoMetadata { - crs: Some("EPSG:4326".to_string()), - }; - let dtype = ExtDType::::try_new(metadata, storage.dtype().clone())?; + .into_array()) +} + +/// A `Point` column (CRS `EPSG:4326`) over the given x/y coordinates. +pub(crate) fn point_column(xs: Vec, ys: Vec) -> VortexResult { + let storage = xy_struct(xs, ys)?; + let dtype = ExtDType::::try_new(wgs84(), storage.dtype().clone())?; Ok(ExtensionArray::new(dtype.erased(), storage).into_array()) } +/// A list offset as `i32`. +fn offset(n: usize) -> VortexResult { + i32::try_from(n).map_err(|_| vortex_err!("geometry offset overflow")) +} + +/// Wrap `values` in a non-nullable `List` with the given offsets. +fn list(values: ArrayRef, offsets: Vec) -> VortexResult { + Ok(ListArray::try_new( + values, + PrimitiveArray::from_iter(offsets).into_array(), + Validity::NonNullable, + )? + .into_array()) +} + +/// A `Polygon` column (CRS `EPSG:4326`). Each polygon is a list of rings; each ring a list of +/// `(x, y)` vertices. Stored as `List>>`. +pub(crate) fn polygon_column(polygons: Vec>>) -> VortexResult { + let (mut xs, mut ys) = (Vec::new(), Vec::new()); + let mut ring_offsets = vec![0i32]; + let mut polygon_offsets = vec![0i32]; + for rings in &polygons { + for ring in rings { + for &(x, y) in ring { + xs.push(x); + ys.push(y); + } + ring_offsets.push(offset(xs.len())?); + } + polygon_offsets.push(offset(ring_offsets.len() - 1)?); + } + + let storage = list(list(xy_struct(xs, ys)?, ring_offsets)?, polygon_offsets)?; + let dtype = ExtDType::::try_new( + wgs84(), + polygon_storage_dtype(Dimension::Xy, Nullability::NonNullable), + )?; + Ok(ExtensionArray::try_new(dtype.erased(), storage)?.into_array()) +} + +/// One multipolygon: polygons → rings → `(x, y)` vertices. +pub(crate) type MultiPolygonRings = Vec>>; + +/// A `MultiPolygon` column (CRS `EPSG:4326`), stored as `List>>>`. +pub(crate) fn multipolygon_column(multipolygons: Vec) -> VortexResult { + let (mut xs, mut ys) = (Vec::new(), Vec::new()); + let mut ring_offsets = vec![0i32]; + let mut polygon_offsets = vec![0i32]; + let mut multipolygon_offsets = vec![0i32]; + for polygons in &multipolygons { + for rings in polygons { + for ring in rings { + for &(x, y) in ring { + xs.push(x); + ys.push(y); + } + ring_offsets.push(offset(xs.len())?); + } + polygon_offsets.push(offset(ring_offsets.len() - 1)?); + } + multipolygon_offsets.push(offset(polygon_offsets.len() - 1)?); + } + + let rings = list(xy_struct(xs, ys)?, ring_offsets)?; + let storage = list(list(rings, polygon_offsets)?, multipolygon_offsets)?; + let dtype = ExtDType::::try_new( + wgs84(), + multipolygon_storage_dtype(Dimension::Xy, Nullability::NonNullable), + )?; + Ok(ExtensionArray::try_new(dtype.erased(), storage)?.into_array()) +} + /// Decode a [`Coordinate`] from an extension-typed point scalar (unwrapped to its coordinate /// storage) or a bare coordinate `Struct` scalar — used to read back a single point in assertions. pub(crate) fn coordinate_from_scalar(scalar: &Scalar) -> VortexResult { diff --git a/vortex-layout/src/layouts/zoned/writer.rs b/vortex-layout/src/layouts/zoned/writer.rs index 6ed7b541bcb..c6ab64a06fb 100644 --- a/vortex-layout/src/layouts/zoned/writer.rs +++ b/vortex-layout/src/layouts/zoned/writer.rs @@ -26,6 +26,7 @@ use vortex_array::aggregate_fn::fns::min::Min; use vortex_array::aggregate_fn::fns::nan_count::NanCount; use vortex_array::aggregate_fn::fns::null_count::NullCount; use vortex_array::aggregate_fn::fns::sum::Sum; +use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::dtype::DType; use vortex_error::VortexError; use vortex_error::VortexResult; @@ -109,7 +110,7 @@ impl LayoutStrategy for ZonedStrategy { .options .aggregate_fns .clone() - .unwrap_or_else(|| default_zoned_aggregate_fns(stream.dtype())); + .unwrap_or_else(|| default_zoned_aggregate_fns(stream.dtype(), session)); let compute_session = session.clone(); let stats_accumulator = Arc::new(Mutex::new(AggregateStatsAccumulator::new( @@ -192,7 +193,7 @@ impl LayoutStrategy for ZonedStrategy { } } -fn default_zoned_aggregate_fns(dtype: &DType) -> Arc<[AggregateFnRef]> { +fn default_zoned_aggregate_fns(dtype: &DType, session: &VortexSession) -> Arc<[AggregateFnRef]> { let (max, min) = match dtype { DType::Utf8(_) | DType::Binary(_) => ( BoundedMax.bind(BoundedMaxOptions { @@ -218,6 +219,9 @@ fn default_zoned_aggregate_fns(dtype: &DType) -> Arc<[AggregateFnRef]> { aggregate_fns.push(NanCount.bind(EmptyOptions)); aggregate_fns.push(NullCount.bind(EmptyOptions)); + // Stats from geo extension types are discovered from the registry at runtime instead. + aggregate_fns.extend(session.aggregate_fns().zone_stat_defaults(dtype)); + aggregate_fns.into() } @@ -237,7 +241,10 @@ mod tests { #[test] fn default_aggregates_bound_variable_length_min_max() { - let aggregate_fns = default_zoned_aggregate_fns(&DType::Utf8(Nullability::NonNullable)); + let aggregate_fns = default_zoned_aggregate_fns( + &DType::Utf8(Nullability::NonNullable), + &vortex_array::array_session(), + ); assert_eq!( aggregate_fns[0].as_::().max_bytes, @@ -251,7 +258,8 @@ mod tests { #[test] fn default_aggregates_keep_fixed_width_min_max_exact() { - let aggregate_fns = default_zoned_aggregate_fns(&PType::I32.into()); + let aggregate_fns = + default_zoned_aggregate_fns(&PType::I32.into(), &vortex_array::array_session()); assert!(aggregate_fns[0].is::()); assert!(aggregate_fns[1].is::()); @@ -263,7 +271,7 @@ mod tests { let dtype = DType::Extension( Timestamp::new(TimeUnit::Microseconds, Nullability::Nullable).erased(), ); - let aggregate_fns = default_zoned_aggregate_fns(&dtype); + let aggregate_fns = default_zoned_aggregate_fns(&dtype, &vortex_array::array_session()); assert!( aggregate_fns From 22d686fc55cb19d800a4f37ac0c15b1c7c305ae8 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Thu, 2 Jul 2026 17:29:20 -0400 Subject: [PATCH 4/4] feat(vortex-bench): spatially sort SpatialBench source parquet (Z-order) Morton-sort each generated table by its geometry column's bounding-box center so every lane (parquet, vortex-WKB, vortex-geo-native) reads spatially-clustered data and the geometry zone-map prune can actually skip chunks. Idempotent via a parquet marker; stale derived vortex files from pre-sort parquet are deleted so the existence-keyed conversions regenerate. Signed-off-by: Nemo Yu --- vortex-bench/Cargo.toml | 2 + vortex-bench/src/spatialbench/benchmark.rs | 10 + vortex-bench/src/spatialbench/datagen/mod.rs | 2 + .../src/spatialbench/datagen/spatial_sort.rs | 253 ++++++++++++++++++ 4 files changed, 267 insertions(+) create mode 100644 vortex-bench/src/spatialbench/datagen/spatial_sort.rs diff --git a/vortex-bench/Cargo.toml b/vortex-bench/Cargo.toml index 069f33a57e9..c64e2551664 100644 --- a/vortex-bench/Cargo.toml +++ b/vortex-bench/Cargo.toml @@ -34,6 +34,8 @@ async-trait = { workspace = true } bzip2 = { workspace = true } clap = { workspace = true, features = ["derive"] } futures = { workspace = true } +geo = { workspace = true } +geo-traits = { workspace = true } geoarrow = { workspace = true } geoarrow-cast = { workspace = true } get_dir = { workspace = true } diff --git a/vortex-bench/src/spatialbench/benchmark.rs b/vortex-bench/src/spatialbench/benchmark.rs index d7b9d94061d..690891fba44 100644 --- a/vortex-bench/src/spatialbench/benchmark.rs +++ b/vortex-bench/src/spatialbench/benchmark.rs @@ -95,6 +95,16 @@ impl Benchmark for SpatialBenchBenchmark { crate::conversions::add_geoparquet_metadata(&zone_file, &geo).await?; } } + + // Cluster the source parquet along a spatial curve so all lanes read the same layout and + // the geometry zone-map prune can skip chunks. + let derived_dirs = [ + base_data_dir.join(Format::OnDiskVortex.name()), + base_data_dir.join(Format::VortexCompact.name()), + base_data_dir.join(NATIVE_DIR), + ]; + datagen::spatially_sort_tables(&base_data_dir.join(Format::Parquet.name()), &derived_dirs) + .await?; Ok(()) } diff --git a/vortex-bench/src/spatialbench/datagen/mod.rs b/vortex-bench/src/spatialbench/datagen/mod.rs index 7808e06cbc3..4c0d4bd2f81 100644 --- a/vortex-bench/src/spatialbench/datagen/mod.rs +++ b/vortex-bench/src/spatialbench/datagen/mod.rs @@ -6,9 +6,11 @@ //! source of truth for the base tables both stages share. pub mod native; +pub mod spatial_sort; pub mod table; pub mod wkb; pub use native::write_native_vortex; +pub use spatial_sort::spatially_sort_tables; pub use table::Table; pub use wkb::generate_tables; diff --git a/vortex-bench/src/spatialbench/datagen/spatial_sort.rs b/vortex-bench/src/spatialbench/datagen/spatial_sort.rs new file mode 100644 index 00000000000..798a65065c9 --- /dev/null +++ b/vortex-bench/src/spatialbench/datagen/spatial_sort.rs @@ -0,0 +1,253 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Spatial clustering of the source parquet, in place, so every downstream lane (parquet, +//! vortex-WKB, `vortex-geo-native`) reads the same layout. +//! +//! The geometry zone-map prune skips a chunk only when its bounding box is disjoint from the query +//! region. In generation order every chunk's box spans the whole map, so nothing prunes; ordering +//! rows on a Z-order (Morton) curve of each geometry's bounding-box center gives each chunk a +//! compact box instead. The center works uniformly for points, polygons, and multipolygons. +//! +//! Every table with a geometry column is sorted by its first one, each part independently — global +//! order is unnecessary for per-chunk pruning. A parquet marker makes it idempotent; derived vortex +//! files from pre-sort parquet are deleted so the existence-keyed conversions regenerate. + +use std::path::Path; +use std::path::PathBuf; +use std::sync::Arc; + +use anyhow::Context; +use arrow_array::Array; +use arrow_array::ArrayRef; +use arrow_array::RecordBatch; +use arrow_array::UInt64Array; +use arrow_schema::DataType; +use arrow_select::concat::concat_batches; +use arrow_select::take::take; +use futures::TryStreamExt; +use geo::BoundingRect; +use geo::Geometry; +use geo_traits::to_geo::ToGeoGeometry; +use geoarrow::array::GenericWkbArray; +use geoarrow::array::GeoArrowArrayAccessor; +use geoarrow::array::WkbViewArray; +use geoarrow::datatypes::Crs; +use geoarrow::datatypes::Metadata; +use geoarrow::datatypes::WkbType; +use parquet::arrow::AsyncArrowWriter; +use parquet::arrow::ParquetRecordBatchStreamBuilder; +use parquet::basic::Compression; +use parquet::file::metadata::KeyValue; +use parquet::file::properties::WriterProperties; +use tokio::fs::File as TokioFile; +use tracing::info; + +use super::table::Table; +use super::wkb::geo_parquet_metadata; + +/// Parquet metadata marker: this file is already spatially sorted (makes the step idempotent). +const SORTED_KEY: &str = "vortex_spatial_sorted"; + +fn geo_metadata() -> Arc { + Arc::new(Metadata::new(Crs::default(), None)) +} + +/// Spatially sort every table with a geometry column, in place. +pub async fn spatially_sort_tables( + parquet_dir: &Path, + derived_dirs: &[PathBuf], +) -> anyhow::Result<()> { + for table in Table::ALL.into_iter().filter(|table| table.is_generated()) { + sort_table_parquet(table, parquet_dir, derived_dirs).await?; + } + Ok(()) +} + +/// Sort every unsorted `{table}_*.parquet` part by its first geometry column's bounding-box center. +async fn sort_table_parquet( + table: Table, + parquet_dir: &Path, + derived_dirs: &[PathBuf], +) -> anyhow::Result<()> { + let Some(geom) = table.geometry_columns().first() else { + return Ok(()); + }; + + let pattern = parquet_dir.join(format!("{}_*.parquet", table.name())); + let mut files: Vec = + glob::glob(&pattern.to_string_lossy())?.collect::>()?; + files.sort(); + let mut pending = Vec::new(); + for file in files { + if !is_sorted(&file).await? { + pending.push(file); + } + } + if pending.is_empty() { + return Ok(()); + } + + // Delete before rewriting: an interrupted run leaves unsorted parts unmarked and re-deletes. + for dir in derived_dirs { + let stale = dir.join(format!("{}_*.vortex", table.name())); + for file in glob::glob(&stale.to_string_lossy())?.flatten() { + tokio::fs::remove_file(&file).await?; + info!(path = %file.display(), "removed stale derived file from pre-sort parquet"); + } + } + + for file in pending { + sort_part(&file, table, geom.name).await?; + } + Ok(()) +} + +/// Whether `path` already carries the [`SORTED_KEY`] marker. +async fn is_sorted(path: &Path) -> anyhow::Result { + let builder = ParquetRecordBatchStreamBuilder::new(TokioFile::open(path).await?).await?; + Ok(builder + .metadata() + .file_metadata() + .key_value_metadata() + .is_some_and(|kvs| kvs.iter().any(|kv| kv.key == SORTED_KEY))) +} + +/// Rewrite one parquet part with its rows in Z-order of `geom_col`. +async fn sort_part(path: &Path, table: Table, geom_col: &str) -> anyhow::Result<()> { + let builder = ParquetRecordBatchStreamBuilder::new(TokioFile::open(path).await?).await?; + let schema = Arc::clone(builder.schema()); + let mut reader = builder.build()?; + let mut batches = Vec::new(); + while let Some(batch) = reader.try_next().await? { + batches.push(batch); + } + if batches.iter().all(|batch| batch.num_rows() == 0) { + return Ok(()); + } + let batch = concat_batches(&schema, &batches)?; + drop(batches); + + let geom_idx = schema.index_of(geom_col)?; + let (xs, ys) = wkb_centers(batch.column(geom_idx).as_ref()) + .with_context(|| format!("decoding geometry column {geom_col} of {}", path.display()))?; + let indices = morton_sort_indices(&xs, &ys); + + let columns: Vec = batch + .columns() + .iter() + .map(|column| take(column.as_ref(), &indices, None)) + .collect::>()?; + let sorted = RecordBatch::try_new(Arc::clone(&schema), columns)?; + let num_rows = sorted.num_rows(); + + let tmp_path = path.with_extension("parquet.sorttmp"); + let props = WriterProperties::builder() + .set_compression(Compression::SNAPPY) + .build(); + let mut writer = + AsyncArrowWriter::try_new(TokioFile::create(&tmp_path).await?, schema, Some(props))?; + writer.write(&sorted).await?; + // A fresh write drops metadata: re-tag geo so DuckDB reads `GEOMETRY`, and add the sorted marker. + if let Some(geo) = geo_parquet_metadata(table) { + writer.append_key_value_metadata(KeyValue::new("geo".to_string(), Some(geo))); + } + writer.append_key_value_metadata(KeyValue::new( + SORTED_KEY.to_string(), + Some(format!("morton:{geom_col}")), + )); + writer.close().await?; + tokio::fs::rename(&tmp_path, path).await?; + + info!( + path = %path.display(), + rows = num_rows, + column = geom_col, + "spatially sorted parquet (morton z-order)" + ); + Ok(()) +} + +/// The bounding-box center `(x, y)` of every geometry in a WKB column, whatever its geometry type. +fn wkb_centers(column: &dyn Array) -> anyhow::Result<(Vec, Vec)> { + let wkb_type = WkbType::new(geo_metadata()); + // Expanded per concrete WKB array type. + macro_rules! centers { + ($array:expr) => {{ + let mut xy = Vec::new(); + for item in $array.iter() { + xy.push(match item { + Some(geometry) => bbox_center(&geometry?.to_geometry()), + None => (f64::NAN, f64::NAN), + }); + } + Ok(xy.into_iter().unzip()) + }}; + } + match column.data_type() { + DataType::Binary => centers!(GenericWkbArray::::try_from((column, wkb_type))?), + DataType::LargeBinary => centers!(GenericWkbArray::::try_from((column, wkb_type))?), + DataType::BinaryView => centers!(WkbViewArray::try_from((column, wkb_type))?), + other => anyhow::bail!("unsupported WKB column type {other}"), + } +} + +/// The center of a geometry's bounding box, or `NaN` for an empty geometry. +fn bbox_center(geometry: &Geometry) -> (f64, f64) { + geometry + .bounding_rect() + .map(|r| ((r.min().x + r.max().x) / 2.0, (r.min().y + r.max().y) / 2.0)) + .unwrap_or((f64::NAN, f64::NAN)) +} + +/// Row indices ordering `(x, y)` along a 32-bit-per-axis Z-order curve, with each axis +/// quantized over the data's own extent. +fn morton_sort_indices(xs: &[f64], ys: &[f64]) -> UInt64Array { + let (mut xmin, mut xmax) = (f64::INFINITY, f64::NEG_INFINITY); + let (mut ymin, mut ymax) = (f64::INFINITY, f64::NEG_INFINITY); + for (&x, &y) in xs.iter().zip(ys) { + if !x.is_nan() && !y.is_nan() { + xmin = xmin.min(x); + xmax = xmax.max(x); + ymin = ymin.min(y); + ymax = ymax.max(y); + } + } + + let keys: Vec = xs + .iter() + .zip(ys) + .map(|(&x, &y)| { + if x.is_nan() || y.is_nan() { + 0 + } else { + interleave(quantize(x, xmin, xmax)) | (interleave(quantize(y, ymin, ymax)) << 1) + } + }) + .collect(); + + let mut order: Vec = (0..xs.len()).collect(); + order.sort_unstable_by_key(|&i| keys[i]); + UInt64Array::from_iter_values(order.into_iter().map(|i| i as u64)) +} + +/// Map `value` within `[min, max]` to the full `u32` range. +fn quantize(value: f64, min: f64, max: f64) -> u32 { + let range = max - min; + if range <= 0.0 { + return 0; + } + let t = ((value - min) / range).clamp(0.0, 1.0); + (t * f64::from(u32::MAX)) as u32 +} + +/// Spread a 32-bit value's bits into the even positions of a `u64`. +fn interleave(value: u32) -> u64 { + let mut n = u64::from(value); + n = (n | (n << 16)) & 0x0000_FFFF_0000_FFFF; + n = (n | (n << 8)) & 0x00FF_00FF_00FF_00FF; + n = (n | (n << 4)) & 0x0F0F_0F0F_0F0F_0F0F; + n = (n | (n << 2)) & 0x3333_3333_3333_3333; + n = (n | (n << 1)) & 0x5555_5555_5555_5555; + n +}