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..22b520ec7a0 100644 --- a/vortex-duckdb/cpp/expr.cpp +++ b/vortex-duckdb/cpp/expr.cpp @@ -11,6 +11,19 @@ #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/common/error_data.hpp" +#include "duckdb/logging/logger.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 +34,52 @@ 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); + DatabaseInstance &db = *wrapper.database->instance; + try { + Connection conn(db); + ClientContext &context = *conn.context; + context.RunFunctionInTransaction([&]() { + auto &system = Catalog::GetSystemCatalog(context); + auto entry = system.GetEntry(context, + DEFAULT_SCHEMA, + "st_dwithin", + OnEntryNotFound::RETURN_NULL); + if (!entry) { + // No `spatial` loaded, so there is no `ST_DWithin` to override. + return; + } + 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 &e) { + ErrorData data(e); + DUCKDB_LOG_ERROR(db, "Failed to register the ST_DWithin override:\t" + data.Message()); + return DuckDBError; + } + 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..adc706540bc 100644 --- a/vortex-duckdb/cpp/include/expr.h +++ b/vortex-duckdb/cpp/include/expr.h @@ -13,6 +13,11 @@ typedef struct duckdb_vx_sfunc_ *duckdb_vx_sfunc; const char *duckdb_vx_sfunc_name(duckdb_vx_sfunc ffi_func); +/// Shadow `ST_DWithin` with a copy that keeps the radius as the third argument, so +/// radius filters can push into Vortex scans. See `RestoreStDWithin` in +/// scalar_fn_pushdown.hpp for the override/restore example. +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..24c0fd8bd1a 100644 --- a/vortex-duckdb/cpp/include/scalar_fn_pushdown.hpp +++ b/vortex-duckdb/cpp/include/scalar_fn_pushdown.hpp @@ -77,6 +77,21 @@ using Projections = unordered_map; LogicalOperatorPtr TryPushdownScalarFunctions(ClientContext &context, LogicalOperatorPtr plan); +/* + * We override spatial's `ST_DWithin` with a copy that keeps the radius as a plain third + * argument (see expr.h); the original folds it into bind data at bind time. We need the + * radius visible to push the filter into a Vortex scan, but spatial's join optimizer only + * recognizes the folded 2-argument form. So this pass rebinds join conditions through + * spatial's original function and leaves filters alone: + * + * FILTER st_dwithin(t.geom, 'POINT(0 0)', 10.0) -- untouched, pushed into the scan + * JOIN ON st_dwithin(a.geom, b.geom, 10.0) -- rebound: st_dwithin(a.geom, b.geom) + * with the radius in bind data + * + * Runs in the pre-optimize hook, before any extension's optimizer pass. + */ +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..8bac6d86749 100644 --- a/vortex-duckdb/cpp/scalar_fn_pushdown.cpp +++ b/vortex-duckdb/cpp/scalar_fn_pushdown.cpp @@ -1,6 +1,8 @@ // 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/function/function_binder.hpp" #include "duckdb/planner/operator/logical_projection.hpp" #include "scalar_fn_pushdown.hpp" #include "table_function.hpp" @@ -231,3 +233,67 @@ 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) { + } + + // 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 || 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..bafa777a061 100644 --- a/vortex-duckdb/cpp/vortex_duckdb.cpp +++ b/vortex-duckdb/cpp/vortex_duckdb.cpp @@ -271,8 +271,13 @@ static void VortexOptimizeFunction(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 387b644fe30..32778870845 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,133 @@ 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. + 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) { + 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 +220,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 +229,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 +244,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 +262,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 +276,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 +287,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 +296,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 +365,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 +425,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 +436,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 +446,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 +487,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 +500,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 +511,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 +536,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 +550,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;