diff --git a/benchmarks/duckdb-bench/src/lib.rs b/benchmarks/duckdb-bench/src/lib.rs index 11be0dd80e5..4f3bdd388fc 100644 --- a/benchmarks/duckdb-bench/src/lib.rs +++ b/benchmarks/duckdb-bench/src/lib.rs @@ -78,11 +78,12 @@ impl DuckClient { for stmt in &statements { self.connection().query(stmt)?; } - // After `LOAD spatial`, shadow `ST_DWithin` so radius filters push. No-op without it. + // After `LOAD spatial`, shadow the overridden spatial functions so that their filters + // push down. No-op without it. self.db .as_ref() .vortex_expect("DuckClient database accessed after close") - .register_st_dwithin_override()?; + .register_spatial_overrides()?; self.init_sql = statements; Ok(()) } @@ -132,11 +133,11 @@ impl DuckClient { .vortex_expect("connection just opened") .query(stmt)?; } - // Re-shadow `ST_DWithin` against the fresh instance. + // Re-shadow the overridden spatial functions against the fresh instance. self.db .as_ref() .vortex_expect("database just opened") - .register_st_dwithin_override()?; + .register_spatial_overrides()?; Ok(()) } diff --git a/vortex-duckdb/build.rs b/vortex-duckdb/build.rs index c6bc02f95de..c15009a92c3 100644 --- a/vortex-duckdb/build.rs +++ b/vortex-duckdb/build.rs @@ -27,12 +27,13 @@ const DEFAULT_DUCKDB_VERSION: &str = "1.5.3"; const BUILD_ARTIFACTS: [&str; 3] = ["libduckdb.dylib", "libduckdb.so", "libduckdb_static.a"]; -const SOURCE_FILES: [&str; 10] = [ +const SOURCE_FILES: [&str; 11] = [ "cpp/vortex_duckdb.cpp", "cpp/copy_function.cpp", "cpp/expr.cpp", "cpp/optimizer.cpp", "cpp/scalar_fn_pushdown.cpp", + "cpp/spatial_overrides.cpp", "cpp/cast_pushdown.cpp", "cpp/aggregate_fn_pushdown.cpp", "cpp/table_filter.cpp", @@ -178,9 +179,10 @@ const DUCKDB_C_API_FUNCTIONS: [&str; 133] = [ "duckdb_vector_size", ]; -const DUCKDB_C_API_HEADERS: [&str; 6] = [ +const DUCKDB_C_API_HEADERS: [&str; 7] = [ "cpp/include/vortex_duckdb.h", "cpp/include/expr.h", + "cpp/include/spatial_overrides.h", "cpp/include/table_filter.h", "cpp/include/vector.h", "cpp/include/copy_function.h", diff --git a/vortex-duckdb/cpp/expr.cpp b/vortex-duckdb/cpp/expr.cpp index 819133f7873..54ef37cd959 100644 --- a/vortex-duckdb/cpp/expr.cpp +++ b/vortex-duckdb/cpp/expr.cpp @@ -14,19 +14,6 @@ #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) { @@ -37,52 +24,6 @@ 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_agg_func_name(duckdb_vx_agg_func ffi) { D_ASSERT(ffi); return reinterpret_cast(ffi)->name.c_str(); diff --git a/vortex-duckdb/cpp/include/expr.h b/vortex-duckdb/cpp/include/expr.h index 886c1d21170..704c8406c90 100644 --- a/vortex-duckdb/cpp/include/expr.h +++ b/vortex-duckdb/cpp/include/expr.h @@ -14,11 +14,6 @@ typedef struct duckdb_vx_agg_func_ *duckdb_vx_agg_func; 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; const char *duckdb_vx_agg_func_name(duckdb_vx_agg_func func); diff --git a/vortex-duckdb/cpp/include/scalar_fn_pushdown.hpp b/vortex-duckdb/cpp/include/scalar_fn_pushdown.hpp index c94fab80408..19bad361cb0 100644 --- a/vortex-duckdb/cpp/include/scalar_fn_pushdown.hpp +++ b/vortex-duckdb/cpp/include/scalar_fn_pushdown.hpp @@ -35,5 +35,3 @@ struct ScalarFnReplace final : LogicalOperatorVisitor { ExpressionPtr VisitReplace(BoundColumnRefExpression &expr, ExpressionPtr *ptr) override; ExpressionPtr VisitReplace(BoundFunctionExpression &expr, ExpressionPtr *ptr) override; }; - -void RestoreStDWithin(ClientContext &context, LogicalOperator &plan); diff --git a/vortex-duckdb/cpp/include/spatial_overrides.h b/vortex-duckdb/cpp/include/spatial_overrides.h new file mode 100644 index 00000000000..e5bbe4dde04 --- /dev/null +++ b/vortex-duckdb/cpp/include/spatial_overrides.h @@ -0,0 +1,19 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#pragma once + +#include "duckdb.h" + +#ifdef __cplusplus /* If compiled as C++, use C ABI */ +extern "C" { +#endif + +/// Shadow the spatial functions that block filter pushdown with pushable copies; see +/// `SPATIAL_OVERRIDES` in spatial_overrides.cpp for the list. Call after `LOAD spatial`; +/// does nothing when spatial is not loaded. +duckdb_state duckdb_vx_register_spatial_overrides(duckdb_database ffi_db); + +#ifdef __cplusplus +} +#endif diff --git a/vortex-duckdb/cpp/include/spatial_overrides.hpp b/vortex-duckdb/cpp/include/spatial_overrides.hpp new file mode 100644 index 00000000000..ce49d82fef8 --- /dev/null +++ b/vortex-duckdb/cpp/include/spatial_overrides.hpp @@ -0,0 +1,32 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#pragma once +#include "duckdb/main/client_context.hpp" +#include "duckdb/planner/expression/bound_function_expression.hpp" +#include "duckdb/planner/logical_operator_visitor.hpp" + +using namespace duckdb; + +/* + * Vortex shadows some spatial (`ST_*`) scalar functions with tweaked copies, because as + * spatial registers them their filters can never push into Vortex scans. + * + * `duckdb_vx_register_spatial_overrides` (spatial_overrides.h) installs the copies after + * `LOAD spatial`. `RestoreSpatialOverrides` below hands join conditions back to spatial's + * originals: Vortex never pushes join conditions, and spatial's join machinery expects its + * own functions. + */ + +// Rebinds overridden spatial calls in join conditions back to spatial's original, so spatial's +// own machinery handles joins. Filters are left untouched: they keep the override and push to +// Vortex scans. +struct SpatialOverrideRestore final : LogicalOperatorVisitor { + ClientContext &context; + + explicit SpatialOverrideRestore(ClientContext &context); + void VisitOperator(LogicalOperator &op) override; + unique_ptr VisitReplace(BoundFunctionExpression &expr, unique_ptr *ptr) override; +}; + +/// Rebind overridden spatial calls in join conditions to spatial's originals. +void RestoreSpatialOverrides(ClientContext &context, LogicalOperator &plan); diff --git a/vortex-duckdb/cpp/scalar_fn_pushdown.cpp b/vortex-duckdb/cpp/scalar_fn_pushdown.cpp index b3986ef49cc..bcf23496443 100644 --- a/vortex-duckdb/cpp/scalar_fn_pushdown.cpp +++ b/vortex-duckdb/cpp/scalar_fn_pushdown.cpp @@ -2,9 +2,6 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors #include "scalar_fn_pushdown.hpp" -#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 "duckdb/planner/operator/logical_get.hpp" @@ -122,67 +119,3 @@ ScalarFnCollect::ScalarFnCollect(Analyses &analyses, const Projections &projecti ScalarFnReplace::ScalarFnReplace(Analyses &analyses, const Projections &projections) : analyses(analyses), projections(projections) { } - -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/spatial_overrides.cpp b/vortex-duckdb/cpp/spatial_overrides.cpp new file mode 100644 index 00000000000..28f095a7b89 --- /dev/null +++ b/vortex-duckdb/cpp/spatial_overrides.cpp @@ -0,0 +1,157 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#include "spatial_overrides.hpp" + +#include "spatial_overrides.h" + +#include "duckdb/catalog/catalog.hpp" +#include "duckdb/catalog/catalog_entry/scalar_function_catalog_entry.hpp" +#include "duckdb/common/error_data.hpp" +#include "duckdb/function/function_binder.hpp" +#include "duckdb/function/scalar_function.hpp" +#include "duckdb/logging/logger.hpp" +#include "duckdb/main/capi/capi_internal.hpp" +#include "duckdb/main/connection.hpp" +#include "duckdb/main/database_manager.hpp" +#include "duckdb/parser/parsed_data/create_scalar_function_info.hpp" +#include "duckdb/planner/expression/bound_function_expression.hpp" +#include "duckdb/planner/logical_operator_visitor.hpp" +#include "duckdb/transaction/meta_transaction.hpp" + +#include +#include + +/// A spatial function Vortex shadows so that its filters can push into Vortex scans. +struct SpatialOverride { + const char *name; + idx_t arity; + void (*tweak)(ScalarFunction &); +}; + +/// Drop spatial's bind so the filter keeps the radius visible as `children[2]`. +static void DropBind(ScalarFunction &fn) { + fn.bind = nullptr; +} + +/// Clear the error mode so the filter pushes through view projections. +static void ClearErrorMode(ScalarFunction &fn) { + fn.SetErrorMode(FunctionErrors::CANNOT_ERROR); +} + +static constexpr SpatialOverride SPATIAL_OVERRIDES[] = { + {"st_dwithin", 3, DropBind}, + {"st_intersects", 2, ClearErrorMode}, + {"st_contains", 2, ClearErrorMode}, + {"st_within", 2, ClearErrorMode}, +}; + +/// Apply one override, later calls to the function bind to the pushable copy instead of +/// spatial's original. +static void RegisterSpatialOverride(ClientContext &context, const SpatialOverride &fn_override) { + auto &system = Catalog::GetSystemCatalog(context); + auto entry = system.GetEntry(context, + DEFAULT_SCHEMA, + fn_override.name, + OnEntryNotFound::RETURN_NULL); + if (!entry) { + return; + } + ScalarFunctionSet set(fn_override.name); + for (const auto &overload : entry->functions.functions) { + ScalarFunction copy = overload; + fn_override.tweak(copy); + set.AddFunction(copy); + } + CreateScalarFunctionInfo info(std::move(set)); + info.on_conflict = OnCreateConflict::REPLACE_ON_CONFLICT; + + info.internal = false; + // Register in the default database's catalog: unqualified calls resolve there before the + // system catalog, so the copy shadows spatial's entry. + auto &catalog = Catalog::GetCatalog(context, DatabaseManager::GetDefaultDatabase(context)); + // A file-backed catalog rejects writes unless the database is marked modified; the mark is + // harmless here because function entries are never persisted to disk. + MetaTransaction::Get(context).ModifyDatabase(catalog.GetAttached(), DatabaseModificationType()); + catalog.CreateFunction(context, info); +} + +/// Apply every override in `SPATIAL_OVERRIDES` in one transaction. +extern "C" duckdb_state duckdb_vx_register_spatial_overrides(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([&]() { + for (const auto &fn_override : SPATIAL_OVERRIDES) { + RegisterSpatialOverride(context, fn_override); + } + }); + } catch (const std::exception &e) { + ErrorData data(e); + DUCKDB_LOG_ERROR(db, "Failed to register the spatial overrides:\t" + data.Message()); + return DuckDBError; + } + return DuckDBSuccess; +} + +SpatialOverrideRestore::SpatialOverrideRestore(ClientContext &context) : context(context) { +} + +void SpatialOverrideRestore::VisitOperator(LogicalOperator &op) { + 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); +} + +unique_ptr SpatialOverrideRestore::VisitReplace(BoundFunctionExpression &expr, + unique_ptr *) { + const bool overridden = + std::any_of(std::begin(SPATIAL_OVERRIDES), + std::end(SPATIAL_OVERRIDES), + [&](const SpatialOverride &o) { + return expr.function.name == o.name && expr.children.size() == o.arity; + }); + if (!overridden) { + return nullptr; // Not an overridden call: leave it as is. + } + // Spatial's original lives in the system catalog, where the override cannot shadow it. + auto original = + Catalog::GetSystemCatalog(context).GetEntry(context, + DEFAULT_SCHEMA, + expr.function.name, + OnEntryNotFound::RETURN_NULL); + if (!original) { + return nullptr; + } + // Rebind a copy of the call's arguments through the original function. + 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: the override call still executes, keep it. + } + return bound; +} + +void RestoreSpatialOverrides(ClientContext &context, LogicalOperator &plan) { + SpatialOverrideRestore(context).VisitOperator(plan); +} diff --git a/vortex-duckdb/cpp/vortex_duckdb.cpp b/vortex-duckdb/cpp/vortex_duckdb.cpp index 8b023c960dd..05ee9098583 100644 --- a/vortex-duckdb/cpp/vortex_duckdb.cpp +++ b/vortex-duckdb/cpp/vortex_duckdb.cpp @@ -5,6 +5,7 @@ #include "data.hpp" #include "error.hpp" #include "scalar_fn_pushdown.hpp" +#include "spatial_overrides.hpp" #include "cast_pushdown.hpp" #include "vortex_duckdb.h" @@ -276,7 +277,7 @@ static void VortexOptimizeFunction(OptimizerExtensionInput &input, unique_ptr &plan) { - RestoreStDWithin(input.context, *plan); + RestoreSpatialOverrides(input.context, *plan); } struct VortexOptimizerExtension final : OptimizerExtension { diff --git a/vortex-duckdb/src/convert/expr.rs b/vortex-duckdb/src/convert/expr.rs index 9acd763bf30..c8c40e20af6 100644 --- a/vortex-duckdb/src/convert/expr.rs +++ b/vortex-duckdb/src/convert/expr.rs @@ -60,7 +60,9 @@ 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::contains::GeoContains; use vortex_geo::scalar_fn::distance::GeoDistance; +use vortex_geo::scalar_fn::intersects::GeoIntersects; use crate::convert::dtype::FromLogicalType; use crate::cpp::DUCKDB_TYPE; @@ -116,13 +118,14 @@ struct ConvertCtx<'a> { 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`. +/// Whether `name` is a non-nullable native geometry column of the scan. The pushed geo kernels +/// reject nullable operands and 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) + .filter(|field| field.name == name && !field.dtype.is_nullable()) .any(|field| match field.dtype.as_extension_opt() { Some(ext) => { ext.is::() @@ -166,48 +169,72 @@ fn geo_operand( } } +/// Lower all geometry operands of a geo function. Returns `None`, skipping the push, when any +/// operand is neither a constant geometry nor a native geometry column. +fn geo_operands( + children: &[&duckdb::ExpressionRef], + ctx: ConvertCtx<'_>, +) -> VortexResult>> { + children + .iter() + .map(|child| geo_operand(child, ctx)) + .collect() +} + /// 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`. + // Spatial's own st_dwithin folds the radius into bind data; the override + // (cpp/spatial_overrides.cpp) keeps it visible here as `children[2]`. "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 { + let Some(operands) = geo_operands(&children[..2], 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(ScalarEmptyOptions, [a, b]); + let geo_distance = GeoDistance.new_expr(ScalarEmptyOptions, operands); 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 { + let Some(operands) = geo_operands(&children, ctx)? else { return Ok(None); }; - let Some(b) = geo_operand(children[1], ctx)? else { + GeoDistance.new_expr(ScalarEmptyOptions, operands) + } + "st_intersects" => { + if children.len() != 2 { + return Ok(None); + } + let Some(operands) = geo_operands(&children, ctx)? else { return Ok(None); }; - GeoDistance.new_expr(ScalarEmptyOptions, [a, b]) + GeoIntersects.new_expr(ScalarEmptyOptions, operands) + } + containment @ ("st_contains" | "st_within") => { + if children.len() != 2 { + return Ok(None); + } + let Some(mut operands) = geo_operands(&children, ctx)? else { + return Ok(None); + }; + // `st_within(a, b)` is `st_contains(b, a)`; both lower to the contains kernel. + if containment == "st_within" { + operands.swap(0, 1); + } + GeoContains.new_expr(ScalarEmptyOptions, operands) } _ => return Ok(None), }; @@ -215,15 +242,6 @@ fn try_from_geo_function( 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, ctx: ConvertCtx<'_>, diff --git a/vortex-duckdb/src/duckdb/database.rs b/vortex-duckdb/src/duckdb/database.rs index e7033b75ee2..1d258f05751 100644 --- a/vortex-duckdb/src/duckdb/database.rs +++ b/vortex-duckdb/src/duckdb/database.rs @@ -91,12 +91,13 @@ 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<()> { + /// Shadow the spatial functions that block filter pushdown with pushable copies; see + /// cpp/spatial_overrides.cpp for the list. Call after `LOAD spatial`; does nothing when + /// spatial is not loaded. + pub fn register_spatial_overrides(&self) -> VortexResult<()> { duckdb_try!( - unsafe { cpp::duckdb_vx_register_st_dwithin_override(self.as_ptr()) }, - "Failed to register the ST_DWithin override" + unsafe { cpp::duckdb_vx_register_spatial_overrides(self.as_ptr()) }, + "Failed to register the spatial function overrides" ); Ok(()) } diff --git a/vortex-duckdb/src/e2e_test/geo_pushdown_test.rs b/vortex-duckdb/src/e2e_test/geo_pushdown_test.rs new file mode 100644 index 00000000000..f058de20d47 --- /dev/null +++ b/vortex-duckdb/src/e2e_test/geo_pushdown_test.rs @@ -0,0 +1,160 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Pushdown tests for the geo scalar functions: every lowered filter must reach the Vortex +//! scan on a direct file scan and through a view. + +use num_traits::AsPrimitive; +use rstest::rstest; +use tempfile::NamedTempFile; +use vortex::array::IntoArray; +use vortex::array::arrays::PrimitiveArray; +use vortex::array::arrays::StructArray; +use vortex::file::WriteOptionsSessionExt; +use vortex::io::runtime::BlockingRuntime; +use vortex_array::arrays::ExtensionArray; +use vortex_array::dtype::extension::ExtDType; +use vortex_geo::extension::GeoMetadata; +use vortex_geo::extension::Point; + +use crate::RUNTIME; +use crate::SESSION; +use crate::cpp::duckdb_string_t; +use crate::duckdb::Connection; +use crate::duckdb::Database; + +/// The test points; the filters below state how many of them they match. +const POINTS: [(f64, f64); 5] = [ + (1.0, 1.0), + (4.0, 4.0), + (10.0, 10.0), + (-1.0, 5.0), + (2.0, 3.0), +]; + +/// Matches three of [`POINTS`]: two inside the polygon and one on its boundary. +const ST_INTERSECTS_FILTER: &str = + "ST_Intersects(geometry, ST_GeomFromText('POLYGON((0 0, 4 0, 4 4, 0 4, 0 0))'))"; + +/// Matches two of [`POINTS`]: those closer than `3` to `(1, 1)`. +const ST_DISTANCE_FILTER: &str = "ST_Distance(geometry, ST_GeomFromText('POINT (1 1)')) < 3.0"; + +/// Matches the same two points as [`ST_DISTANCE_FILTER`], phrased as `ST_DWithin`. +const ST_DWITHIN_FILTER: &str = "ST_DWithin(geometry, ST_GeomFromText('POINT (1 1)'), 3.0)"; + +/// Matches two of [`POINTS`]: the boundary point of [`ST_INTERSECTS_FILTER`]'s polygon is not +/// contained (OGC contains excludes the boundary). +const ST_CONTAINS_FILTER: &str = + "ST_Contains(ST_GeomFromText('POLYGON((0 0, 4 0, 4 4, 0 4, 0 0))'), geometry)"; + +/// Matches the same two points as [`ST_CONTAINS_FILTER`], phrased as `ST_Within`. +const ST_WITHIN_FILTER: &str = + "ST_Within(geometry, ST_GeomFromText('POLYGON((0 0, 4 0, 4 4, 0 4, 0 0))'))"; + +/// A vortex file whose single column `geometry` holds [`POINTS`] as native `Point`s. +fn native_point_file() -> NamedTempFile { + RUNTIME.block_on(async { + let xs = PrimitiveArray::from_iter(POINTS.map(|(x, _)| x)).into_array(); + let ys = PrimitiveArray::from_iter(POINTS.map(|(_, y)| y)).into_array(); + let storage = StructArray::from_fields(&[("x", xs), ("y", ys)]) + .unwrap() + .into_array(); + let dtype = + ExtDType::::try_new(GeoMetadata { crs: None }, storage.dtype().clone()).unwrap(); + let points = ExtensionArray::new(dtype.erased(), storage).into_array(); + + let file = NamedTempFile::with_suffix(".vortex").unwrap(); + let table = StructArray::from_fields(&[("geometry", points)]).unwrap(); + let mut writer = async_fs::File::create(&file).await.unwrap(); + SESSION + .write_options() + .write(&mut writer, table.into_array().to_array_stream()) + .await + .unwrap(); + file + }) +} + +/// An in-memory database with the Vortex extension initialized, `spatial` loaded, and the +/// spatial overrides registered. +fn spatial_database() -> (Database, Connection) { + let db = Database::open_in_memory().unwrap(); + db.register_vortex_scan_replacement().unwrap(); + crate::initialize(&db).unwrap(); + let conn = db.connect().unwrap(); + conn.query("INSTALL spatial; LOAD spatial;").unwrap(); + // Must follow `LOAD spatial`: the overrides copy spatial's catalog entries. + db.register_spatial_overrides().unwrap(); + (db, conn) +} + +/// Read back the single `i64` of a one-row, one-column query. +fn query_i64(conn: &Connection, query: &str) -> i64 { + let result = conn.query(query).unwrap(); + let chunk = result.into_iter().next().unwrap(); + chunk + .get_vector(0) + .as_slice_with_len::(chunk.len().as_())[0] +} + +/// The `EXPLAIN` physical plan of `query` as one string. +fn explain_plan(conn: &Connection, query: &str) -> String { + let explain = conn.query(&format!("EXPLAIN {query}")).unwrap(); + let mut plan = String::new(); + for mut chunk in explain { + let len = chunk.len().as_(); + let vec = chunk.get_vector_mut(1); + for value in unsafe { vec.as_slice_mut::(len) } { + let slice: &[u8] = unsafe { + std::slice::from_raw_parts( + crate::cpp::duckdb_string_t_data(&raw mut *value) as _, + crate::cpp::duckdb_string_t_length(*value) as usize, + ) + }; + plan.push_str(&String::from_utf8_lossy(slice)); + } + } + plan +} + +/// Assert that filtering `table` with `filter` counts `expected` points and leaves no `FILTER` +/// operator in the plan, i.e. the predicate ran inside the scan. +fn assert_pushed(conn: &Connection, table: &str, filter: &str, expected: i64) { + let query = format!("SELECT count(*) FROM {table} WHERE {filter}"); + assert_eq!(query_i64(conn, &query), expected); + let plan = explain_plan(conn, &query); + assert!(!plan.contains("FILTER"), "filter was not pushed:\n{plan}"); +} + +/// Every lowered geo filter pushes on a direct file scan. +#[rstest] +#[case::st_intersects(ST_INTERSECTS_FILTER, 3)] +#[case::st_distance(ST_DISTANCE_FILTER, 2)] +#[case::st_dwithin(ST_DWITHIN_FILTER, 2)] +#[case::st_contains(ST_CONTAINS_FILTER, 2)] +#[case::st_within(ST_WITHIN_FILTER, 2)] +fn geo_filter_pushes_on_file_scan(#[case] filter: &str, #[case] expected: i64) { + let file = native_point_file(); + let (_db, conn) = spatial_database(); + let table = format!("'{}'", file.path().to_string_lossy()); + assert_pushed(&conn, &table, filter, expected); +} + +/// Every lowered geo filter pushes through a view; without the overrides, DuckDB would keep +/// `ST_Intersects` (can-throw) above the view's projection and hide `ST_DWithin`'s radius. +#[rstest] +#[case::st_intersects(ST_INTERSECTS_FILTER, 3)] +#[case::st_distance(ST_DISTANCE_FILTER, 2)] +#[case::st_dwithin(ST_DWITHIN_FILTER, 2)] +#[case::st_contains(ST_CONTAINS_FILTER, 2)] +#[case::st_within(ST_WITHIN_FILTER, 2)] +fn geo_filter_pushes_through_view(#[case] filter: &str, #[case] expected: i64) { + let file = native_point_file(); + let (_db, conn) = spatial_database(); + conn.query(&format!( + "CREATE VIEW points_v AS SELECT * FROM read_vortex('{}')", + file.path().to_string_lossy() + )) + .unwrap(); + assert_pushed(&conn, "points_v", filter, expected); +} diff --git a/vortex-duckdb/src/e2e_test/mod.rs b/vortex-duckdb/src/e2e_test/mod.rs index 34182bd133d..5d41df51220 100644 --- a/vortex-duckdb/src/e2e_test/mod.rs +++ b/vortex-duckdb/src/e2e_test/mod.rs @@ -1,5 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +#[cfg(test)] +mod geo_pushdown_test; #[cfg(test)] mod vortex_scan_test; diff --git a/vortex-geo/Cargo.toml b/vortex-geo/Cargo.toml index bfb6a639742..36a26bfd835 100644 --- a/vortex-geo/Cargo.toml +++ b/vortex-geo/Cargo.toml @@ -30,6 +30,7 @@ wkb = { workspace = true } [dev-dependencies] rstest = { workspace = true } +vortex-array = { workspace = true, features = ["_test-harness"] } [lints] workspace = true diff --git a/vortex-geo/src/extension/mod.rs b/vortex-geo/src/extension/mod.rs index 73f6b6317b0..663d10820e0 100644 --- a/vortex-geo/src/extension/mod.rs +++ b/vortex-geo/src/extension/mod.rs @@ -43,15 +43,45 @@ use vortex_array::IntoArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::dtype::DType; use vortex_array::dtype::extension::ExtDType; use vortex_array::dtype::extension::ExtVTable; use vortex_array::scalar::Scalar; use vortex_arrow::FromArrowArray; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; use vortex_error::vortex_err; pub use wkb::*; +/// Whether `dtype` is one of the native geometry extension types the geo kernels operate on. +fn is_native_geometry(dtype: &DType) -> bool { + dtype.as_extension_opt().is_some_and(|ext| { + ext.is::() + || ext.is::() + || ext.is::() + || ext.is::() + || ext.is::() + || ext.is::() + }) +} + +/// Validate the operands of a geo scalar function: each must be a native geometry type (so the +/// kernel can decode it) and non-nullable (geometry arrays never carry nulls). +pub(crate) fn validate_geometry_operands(dtypes: &[DType]) -> VortexResult<()> { + for dtype in dtypes { + vortex_ensure!( + is_native_geometry(dtype), + "geo: operand {dtype} is not a native geometry type" + ); + vortex_ensure!( + !dtype.is_nullable(), + "geo: nullable operand {dtype} is unsupported" + ); + } + Ok(()) +} + /// 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/lib.rs b/vortex-geo/src/lib.rs index 4ddb8c25c3b..ad0c3c503e2 100644 --- a/vortex-geo/src/lib.rs +++ b/vortex-geo/src/lib.rs @@ -15,7 +15,9 @@ use crate::extension::MultiPolygon; use crate::extension::Point; use crate::extension::Polygon; use crate::extension::WellKnownBinary; +use crate::scalar_fn::contains::GeoContains; use crate::scalar_fn::distance::GeoDistance; +use crate::scalar_fn::intersects::GeoIntersects; pub mod extension; pub mod scalar_fn; @@ -50,5 +52,7 @@ pub fn initialize(session: &VortexSession) { session.arrow().register_importer(Arc::new(MultiPolygon)); // Register the geometry scalar functions. + session.scalar_fns().register(GeoContains); session.scalar_fns().register(GeoDistance); + session.scalar_fns().register(GeoIntersects); } diff --git a/vortex-geo/src/scalar_fn/contains.rs b/vortex-geo/src/scalar_fn/contains.rs new file mode 100644 index 00000000000..5638c224e20 --- /dev/null +++ b/vortex-geo/src/scalar_fn/contains.rs @@ -0,0 +1,307 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! `ST_Contains`: OGC containment test between two native geometries. + +use geo::Contains; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::TypedScalarFnInstance; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::extension::geometries; +use crate::extension::single_geometry; +use crate::extension::validate_geometry_operands; + +/// OGC `ST_Contains` between two native geometry operands, each a column or a constant +/// literal: true where operand `b` lies completely inside operand `a` (boundary contact alone +/// does not count). Containment is not symmetric; the operand order is significant. +#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] +pub struct GeoContains; + +impl GeoContains { + /// A lazy `ScalarFnArray` computing per-row whether operand `a` contains operand `b`; + /// either may be constant. The output length is taken from `a`. + pub fn try_new_array(a: ArrayRef, b: ArrayRef) -> VortexResult { + ScalarFnArray::try_new( + TypedScalarFnInstance::new(GeoContains, EmptyOptions).erased(), + vec![a, b], + ) + } +} + +impl ScalarFnVTable for GeoContains { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.geo.contains"); + *ID + } + + fn serialize(&self, _: &Self::Options) -> VortexResult>> { + Ok(Some(vec![])) + } + + fn deserialize(&self, _: &[u8], _: &VortexSession) -> VortexResult { + Ok(EmptyOptions) + } + + fn arity(&self, _: &Self::Options) -> Arity { + Arity::Exact(2) + } + + fn child_name(&self, _: &Self::Options, child_idx: usize) -> ChildName { + match child_idx { + 0 => ChildName::from("a"), + 1 => ChildName::from("b"), + _ => unreachable!("contains has exactly two children"), + } + } + + fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { + validate_geometry_operands(dtypes)?; + Ok(DType::Bool(Nullability::NonNullable)) + } + + fn execute( + &self, + _: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let a = args.get(0)?; + let b = args.get(1)?; + // Containment is not symmetric: `a` is always the container and `b` the contained. + match (a.as_opt::(), b.as_opt::()) { + (Some(qa), Some(qb)) => { + let ga = single_geometry(qa.scalar(), ctx)?; + let gb = single_geometry(qb.scalar(), ctx)?; + Ok(ConstantArray::new( + Scalar::bool(ga.contains(&gb), Nullability::NonNullable), + a.len(), + ) + .into_array()) + } + (Some(qa), None) => constant_contains_column(qa.scalar(), &b, ctx), + (None, Some(qb)) => column_contains_constant(&a, qb.scalar(), ctx), + (None, None) => { + vortex_ensure_eq!( + a.len(), + b.len(), + "geo contains: operand length mismatch {} vs {}", + a.len(), + b.len() + ); + let ag = geometries(&a, ctx)?; + let bg = geometries(&b, ctx)?; + let hits = ag.iter().zip(&bg).map(|(x, y)| x.contains(y)); + Ok(BoolArray::from_iter(hits).into_array()) + } + } + } +} + +/// Whether the constant `container` contains each row of `contained`. +fn constant_contains_column( + container: &Scalar, + contained: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let container = single_geometry(container, ctx)?; + let geoms = geometries(contained, ctx)?; + let hits = geoms.iter().map(|g| container.contains(g)); + Ok(BoolArray::from_iter(hits).into_array()) +} + +/// Whether each row of `container` contains the constant `contained`. +fn column_contains_constant( + container: &ArrayRef, + contained: &Scalar, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let contained = single_geometry(contained, ctx)?; + let geoms = geometries(container, ctx)?; + let hits = geoms.iter().map(|g| g.contains(&contained)); + Ok(BoolArray::from_iter(hits).into_array()) +} + +#[cfg(test)] +mod tests { + use geo_types::Geometry; + use geo_types::LineString; + use geo_types::Point; + use geo_types::Polygon; + use rstest::rstest; + use vortex_array::ArrayRef; + use vortex_array::Canonical; + use vortex_array::ExecutionCtx; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ConstantArray; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::scalar_fn::EmptyOptions; + use vortex_array::scalar_fn::ScalarFnVTable; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + use wkb::writer::WriteOptions; + + use super::GeoContains; + use crate::test_harness::point_column; + + /// A rectangle polygon with corners `(x0, y0)` and `(x1, y1)`, no holes. + fn rect_polygon(x0: f64, y0: f64, x1: f64, y1: f64) -> Polygon { + Polygon::new( + LineString::from(vec![(x0, y0), (x1, y0), (x1, y1), (x0, y1), (x0, y0)]), + vec![], + ) + } + + /// A constant column of length `len`, every row the native form of `geometry`. + fn geometry_constant(geometry: &Geometry, len: usize) -> VortexResult { + let mut buf = Vec::new(); + wkb::writer::write_geometry(&mut buf, geometry, &WriteOptions::default()) + .map_err(|e| vortex_err!("writing WKB failed: {e}"))?; + let scalar = crate::extension::native_geometry_scalar_from_wkb(&buf)? + .ok_or_else(|| vortex_err!("unsupported geometry type"))?; + Ok(ConstantArray::new(scalar, len).into_array()) + } + + /// Materialize `array` so it is no longer a `Constant`, forcing the non-constant kernel + /// paths. + fn materialize(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + Ok(array.execute::(ctx)?.into_array()) + } + + /// Execute `GeoContains(a, b)` and assert the per-row verdicts equal `expected`. + fn assert_contains( + a: ArrayRef, + b: ArrayRef, + expected: impl IntoIterator, + ) -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let contains = GeoContains::try_new_array(a, b)?.into_array(); + assert_arrays_eq!(contains, BoolArray::from_iter(expected), &mut ctx); + Ok(()) + } + + // The tests cover each `execute` dispatch arm in match order, then the edge cases. + + /// Constant vs constant: a polygon contains a nested polygon but not a partially + /// overlapping or disjoint one; every output row carries the same verdict. + #[rstest] + #[case::nested(rect_polygon(1.0, 1.0, 3.0, 3.0), true)] + #[case::overlapping(rect_polygon(2.0, 2.0, 6.0, 6.0), false)] + #[case::disjoint(rect_polygon(20.0, 20.0, 24.0, 24.0), false)] + fn constant_vs_constant_polygons( + #[case] other: Polygon, + #[case] expected: bool, + ) -> VortexResult<()> { + let container = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 3)?; + let other = geometry_constant(&Geometry::Polygon(other), 3)?; + assert_contains(container, other, [expected; 3]) + } + + /// Partially overlapping polygons contain each other in neither direction. + #[test] + fn overlapping_polygons_contain_neither_way() -> VortexResult<()> { + let a = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 2)?; + let b = geometry_constant(&Geometry::Polygon(rect_polygon(2.0, 2.0, 6.0, 6.0)), 2)?; + assert_contains(a.clone(), b.clone(), [false; 2])?; + assert_contains(b, a, [false; 2]) + } + + /// Containment is not symmetric: a polygon contains an interior point, but the point does + /// not contain the polygon. + #[test] + fn contains_is_asymmetric() -> VortexResult<()> { + let polygon = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 2)?; + let point = geometry_constant(&Geometry::Point(Point::new(2.0, 2.0)), 2)?; + assert_contains(polygon.clone(), point.clone(), [true; 2])?; + assert_contains(point, polygon, [false; 2]) + } + + /// Constant polygon vs point column: a strictly interior point is contained; points outside + /// or exactly on the boundary are not (OGC contains excludes the boundary). + #[test] + fn constant_polygon_vs_point_column() -> VortexResult<()> { + let container = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 3)?; + let points = point_column(vec![2.0, 10.0, 0.0], vec![2.0, 10.0, 2.0])?; + assert_contains(container, points, [true, false, false]) + } + + /// Polygon column vs constant point: only the polygon around the point contains it. + #[test] + fn polygon_column_vs_constant_point() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let around = materialize( + geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 2)?, + &mut ctx, + )?; + let away = materialize( + geometry_constant(&Geometry::Polygon(rect_polygon(20.0, 20.0, 24.0, 24.0)), 2)?, + &mut ctx, + )?; + let point = geometry_constant(&Geometry::Point(Point::new(2.0, 2.0)), 2)?; + + assert_contains(around, point.clone(), [true; 2])?; + assert_contains(away, point, [false; 2]) + } + + /// Column vs column pairs rows: each polygon row is tested against the point row at the + /// same position. + #[test] + fn polygon_column_vs_point_column() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let polygons = materialize( + geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 2)?, + &mut ctx, + )?; + let points = point_column(vec![2.0, 10.0], vec![2.0, 10.0])?; + assert_contains(polygons, points, [true, false]) + } + + /// Geometry arrays are never nullable, so a nullable operand dtype is rejected. + #[test] + fn nullable_operand_is_rejected() -> VortexResult<()> { + let dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let result = GeoContains.return_dtype(&EmptyOptions, &[dtype.as_nullable(), dtype]); + assert!(result.is_err()); + Ok(()) + } + + /// A non-geometry operand dtype is rejected up front, before execution. + #[test] + fn non_geometry_operand_is_rejected() -> VortexResult<()> { + let geo = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let numeric = DType::Primitive(PType::I32, Nullability::NonNullable); + let result = GeoContains.return_dtype(&EmptyOptions, &[geo, numeric]); + assert!(result.is_err()); + Ok(()) + } +} diff --git a/vortex-geo/src/scalar_fn/distance.rs b/vortex-geo/src/scalar_fn/distance.rs index b894946eecb..2e196706ec2 100644 --- a/vortex-geo/src/scalar_fn/distance.rs +++ b/vortex-geo/src/scalar_fn/distance.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! `ST_Distance`: planar (Euclidean) distance between two native geometries via the `geo` crate. +//! `ST_Distance`: planar (Euclidean) distance between two native geometries. use geo::Distance; use geo::Euclidean; @@ -10,12 +10,8 @@ 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,17 +24,16 @@ use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; use vortex_error::VortexResult; -use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; 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; +use crate::extension::validate_geometry_operands; -/// Planar (Euclidean) `ST_Distance` (no geodesic correction) between two native geometry operands. -/// Each is a column or a constant literal; `geo` computes the distance between each pair. +/// Planar (Euclidean) `ST_Distance` (no geodesic correction) between two native geometry +/// operands, each a column or a constant literal. #[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] pub struct GeoDistance; @@ -81,7 +76,8 @@ impl ScalarFnVTable for GeoDistance { } } - fn return_dtype(&self, _: &Self::Options, _: &[DType]) -> VortexResult { + fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { + validate_geometry_operands(dtypes)?; Ok(DType::Primitive(PType::F64, Nullability::NonNullable)) } @@ -107,23 +103,13 @@ 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) => { - vortex_ensure!( - a.len() == b.len(), + vortex_ensure_eq!( + a.len(), + b.len(), "geo distance: operand length mismatch {} vs {}", 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)); @@ -133,80 +119,19 @@ impl ScalarFnVTable for GeoDistance { } } -/// Distance from each row of `operand` to a constant `query` geometry, decoded once and broadcast. -/// Distance is symmetric, so this serves a constant on either side. +/// Distance from each row of `operand` to the constant `query` geometry. Distance is symmetric, +/// so this serves a constant on either side. fn distances_to_constant( operand: &ArrayRef, 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; @@ -215,6 +140,11 @@ mod tests { use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::scalar_fn::EmptyOptions; + use vortex_array::scalar_fn::ScalarFnVTable; use vortex_error::VortexResult; use super::GeoDistance; @@ -296,4 +226,23 @@ mod tests { assert_eq!(distances(distance, &mut ctx)?, vec![5.0, 5.0, 5.0]); Ok(()) } + + /// Geometry arrays are never nullable, so a nullable operand dtype is rejected. + #[test] + fn nullable_operand_is_rejected() -> VortexResult<()> { + let dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let result = GeoDistance.return_dtype(&EmptyOptions, &[dtype.as_nullable(), dtype]); + assert!(result.is_err()); + Ok(()) + } + + /// A non-geometry operand dtype is rejected up front, before execution. + #[test] + fn non_geometry_operand_is_rejected() -> VortexResult<()> { + let geo = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let numeric = DType::Primitive(PType::I32, Nullability::NonNullable); + let result = GeoDistance.return_dtype(&EmptyOptions, &[geo, numeric]); + assert!(result.is_err()); + Ok(()) + } } diff --git a/vortex-geo/src/scalar_fn/intersects.rs b/vortex-geo/src/scalar_fn/intersects.rs new file mode 100644 index 00000000000..801745dc277 --- /dev/null +++ b/vortex-geo/src/scalar_fn/intersects.rs @@ -0,0 +1,334 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! `ST_Intersects`: OGC intersection test between two native geometries. + +use geo::Intersects; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::TypedScalarFnInstance; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::extension::geometries; +use crate::extension::single_geometry; +use crate::extension::validate_geometry_operands; + +/// OGC `ST_Intersects` (not disjoint; boundary contact counts) between two native geometry +/// operands, each a column or a constant literal. +#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] +pub struct GeoIntersects; + +impl GeoIntersects { + /// A lazy `ScalarFnArray` computing per-row whether operands `a` and `b` intersect; either may + /// be constant. The output length is taken from `a`. + pub fn try_new_array(a: ArrayRef, b: ArrayRef) -> VortexResult { + ScalarFnArray::try_new( + TypedScalarFnInstance::new(GeoIntersects, EmptyOptions).erased(), + vec![a, b], + ) + } +} + +impl ScalarFnVTable for GeoIntersects { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.geo.intersects"); + *ID + } + + fn serialize(&self, _: &Self::Options) -> VortexResult>> { + Ok(Some(vec![])) + } + + fn deserialize(&self, _: &[u8], _: &VortexSession) -> VortexResult { + Ok(EmptyOptions) + } + + fn arity(&self, _: &Self::Options) -> Arity { + Arity::Exact(2) + } + + fn child_name(&self, _: &Self::Options, child_idx: usize) -> ChildName { + match child_idx { + 0 => ChildName::from("a"), + 1 => ChildName::from("b"), + _ => unreachable!("intersects has exactly two children"), + } + } + + fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { + validate_geometry_operands(dtypes)?; + Ok(DType::Bool(Nullability::NonNullable)) + } + + fn execute( + &self, + _: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let a = args.get(0)?; + let b = args.get(1)?; + match (a.as_opt::(), b.as_opt::()) { + (Some(qa), Some(qb)) => { + let ga = single_geometry(qa.scalar(), ctx)?; + let gb = single_geometry(qb.scalar(), ctx)?; + Ok(ConstantArray::new( + Scalar::bool(ga.intersects(&gb), Nullability::NonNullable), + a.len(), + ) + .into_array()) + } + (Some(query), None) => intersects_constant(&b, query.scalar(), ctx), + (None, Some(query)) => intersects_constant(&a, query.scalar(), ctx), + (None, None) => { + vortex_ensure_eq!( + a.len(), + b.len(), + "geo intersects: operand length mismatch {} vs {}", + a.len(), + b.len() + ); + let ag = geometries(&a, ctx)?; + let bg = geometries(&b, ctx)?; + let hits = ag.iter().zip(&bg).map(|(x, y)| x.intersects(y)); + Ok(BoolArray::from_iter(hits).into_array()) + } + } + } +} + +/// Whether each row of `operand` intersects the constant `query` geometry. Intersection is +/// symmetric, so this serves a constant on either side. +fn intersects_constant( + operand: &ArrayRef, + query: &Scalar, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let query = single_geometry(query, ctx)?; + let geoms = geometries(operand, ctx)?; + let hits = geoms.iter().map(|g| g.intersects(&query)); + Ok(BoolArray::from_iter(hits).into_array()) +} + +#[cfg(test)] +mod tests { + use geo_types::Coord; + use geo_types::Geometry; + use geo_types::LineString; + use geo_types::MultiPolygon; + use geo_types::Polygon; + use rstest::rstest; + use vortex_array::ArrayRef; + use vortex_array::Canonical; + use vortex_array::ExecutionCtx; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ConstantArray; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::scalar_fn::EmptyOptions; + use vortex_array::scalar_fn::ScalarFnVTable; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + use wkb::writer::WriteOptions; + + use super::GeoIntersects; + use crate::test_harness::point_column; + + /// A rectangle polygon with corners `(x0, y0)` and `(x1, y1)`, no holes. + fn rect_polygon(x0: f64, y0: f64, x1: f64, y1: f64) -> Polygon { + Polygon::new( + LineString::from(vec![(x0, y0), (x1, y0), (x1, y1), (x0, y1), (x0, y0)]), + vec![], + ) + } + + /// The query region shared by the point tests: a `10x10` square with a `4..6` square hole. + fn donut() -> Geometry { + Geometry::Polygon(Polygon::new( + LineString::from(vec![(0.0, 0.0), (10.0, 0.0), (10.0, 10.0), (0.0, 10.0)]), + vec![LineString::from(vec![ + (4.0, 4.0), + (6.0, 4.0), + (6.0, 6.0), + (4.0, 6.0), + ])], + )) + } + + /// Probes for [`donut`], one per verdict class: interior, exterior, on the outer boundary, + /// and inside the hole. + fn donut_probes() -> VortexResult { + point_column(vec![2.0, 20.0, 0.0, 5.0], vec![2.0, 20.0, 5.0, 5.0]) + } + + /// [`donut_probes`]'s verdicts: interior and boundary contact intersect; exterior and + /// in-hole do not. + const DONUT_EXPECTED: [bool; 4] = [true, false, true, false]; + + /// A constant column of length `len`, every row the native form of `geometry`. + fn geometry_constant(geometry: &Geometry, len: usize) -> VortexResult { + let mut buf = Vec::new(); + wkb::writer::write_geometry(&mut buf, geometry, &WriteOptions::default()) + .map_err(|e| vortex_err!("writing WKB failed: {e}"))?; + let scalar = crate::extension::native_geometry_scalar_from_wkb(&buf)? + .ok_or_else(|| vortex_err!("unsupported geometry type"))?; + Ok(ConstantArray::new(scalar, len).into_array()) + } + + /// Materialize `array` so it is no longer a `Constant`, forcing the non-constant kernel paths. + fn materialize(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + Ok(array.execute::(ctx)?.into_array()) + } + + /// Execute `GeoIntersects(a, b)` and assert the per-row verdicts equal `expected`. + fn assert_intersects( + a: ArrayRef, + b: ArrayRef, + expected: impl IntoIterator, + ) -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let intersects = GeoIntersects::try_new_array(a, b)?.into_array(); + assert_arrays_eq!(intersects, BoolArray::from_iter(expected), &mut ctx); + Ok(()) + } + + // The tests cover each `execute` dispatch arm in match order, then the edge cases. + + /// Constant vs constant: overlapping and touching (edge or corner) polygons intersect, + /// disjoint ones do not; every output row carries the same verdict. + #[rstest] + #[case::overlapping(rect_polygon(2.0, 2.0, 6.0, 6.0), true)] + #[case::touching_edge(rect_polygon(4.0, 0.0, 8.0, 4.0), true)] + #[case::touching_corner(rect_polygon(4.0, 4.0, 8.0, 8.0), true)] + #[case::disjoint(rect_polygon(20.0, 20.0, 24.0, 24.0), false)] + fn constant_vs_constant_polygons( + #[case] other: Polygon, + #[case] expected: bool, + ) -> VortexResult<()> { + let a = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 3)?; + let b = geometry_constant(&Geometry::Polygon(other), 3)?; + assert_intersects(a, b, [expected; 3]) + } + + /// Point column vs constant polygon, with the constant on either side since intersection + /// is symmetric: the [`DONUT_EXPECTED`] verdicts. + #[test] + fn point_column_vs_constant_polygon() -> VortexResult<()> { + assert_intersects( + donut_probes()?, + geometry_constant(&donut(), 4)?, + DONUT_EXPECTED, + )?; + assert_intersects( + geometry_constant(&donut(), 4)?, + donut_probes()?, + DONUT_EXPECTED, + ) + } + + /// Point column vs constant multipolygon: a point in either part intersects, a point in + /// neither does not. + #[test] + fn point_column_vs_constant_multipolygon() -> VortexResult<()> { + let query = Geometry::MultiPolygon(MultiPolygon::new(vec![ + rect_polygon(0.0, 0.0, 2.0, 2.0), + rect_polygon(10.0, 10.0, 12.0, 12.0), + ])); + let points = point_column(vec![1.0, 11.0, 5.0], vec![1.0, 11.0, 5.0])?; + assert_intersects(points, geometry_constant(&query, 3)?, [true, true, false]) + } + + /// Polygon column vs constant: the same pairs and verdicts as + /// [`constant_vs_constant_polygons`]. + #[rstest] + #[case::overlapping(rect_polygon(2.0, 2.0, 6.0, 6.0), true)] + #[case::touching_edge(rect_polygon(4.0, 0.0, 8.0, 4.0), true)] + #[case::disjoint(rect_polygon(20.0, 20.0, 24.0, 24.0), false)] + fn polygon_column_vs_constant( + #[case] other: Polygon, + #[case] expected: bool, + ) -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let column = materialize(geometry_constant(&Geometry::Polygon(other), 2)?, &mut ctx)?; + let query = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 2)?; + assert_intersects(column, query, [expected; 2]) + } + + /// Point column vs point column: rows intersect exactly at equal coordinates. + #[test] + fn point_column_vs_point_column() -> VortexResult<()> { + let a = point_column(vec![0.0, 1.0], vec![0.0, 1.0])?; + let b = point_column(vec![0.0, 2.0], vec![0.0, 1.0])?; + assert_intersects(a, b, [true, false]) + } + + /// Point column vs polygon column pairwise: agrees with point column vs constant on the + /// same data. + #[test] + fn columns_agree_with_constant() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let constant = geometry_constant(&donut(), 4)?; + let column = materialize(constant.clone(), &mut ctx)?; + + let against_constant = + GeoIntersects::try_new_array(donut_probes()?, constant)?.into_array(); + let pairwise = GeoIntersects::try_new_array(donut_probes()?, column)?.into_array(); + + assert_arrays_eq!(against_constant, pairwise, &mut ctx); + Ok(()) + } + + /// An empty constant geometry intersects nothing. + #[test] + fn empty_constant_intersects_nothing() -> VortexResult<()> { + let empty = Geometry::LineString(LineString::new(Vec::::new())); + let points = point_column(vec![0.0, 1.0], vec![0.0, 1.0])?; + assert_intersects(points, geometry_constant(&empty, 2)?, [false, false]) + } + + /// Geometry arrays are never nullable, so a nullable operand dtype is rejected. + #[test] + fn nullable_operand_is_rejected() -> VortexResult<()> { + let dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let result = GeoIntersects.return_dtype(&EmptyOptions, &[dtype.as_nullable(), dtype]); + assert!(result.is_err()); + Ok(()) + } + + /// A non-geometry operand dtype is rejected up front, before execution. + #[test] + fn non_geometry_operand_is_rejected() -> VortexResult<()> { + let geo = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let numeric = DType::Primitive(PType::I32, Nullability::NonNullable); + let result = GeoIntersects.return_dtype(&EmptyOptions, &[geo, numeric]); + assert!(result.is_err()); + Ok(()) + } +} diff --git a/vortex-geo/src/scalar_fn/mod.rs b/vortex-geo/src/scalar_fn/mod.rs index 385208f1991..2246f82621b 100644 --- a/vortex-geo/src/scalar_fn/mod.rs +++ b/vortex-geo/src/scalar_fn/mod.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Geometry scalar functions over the [`Point`](crate::extension::Point) type. +//! Geometry scalar functions over the native geometry extension types. +pub mod contains; pub mod distance; +pub mod intersects;