diff --git a/vortex-geo/src/lib.rs b/vortex-geo/src/lib.rs index 652311be3b2..3a0c52f1eb9 100644 --- a/vortex-geo/src/lib.rs +++ b/vortex-geo/src/lib.rs @@ -20,6 +20,7 @@ use crate::extension::Polygon; use crate::extension::Rect; use crate::extension::WellKnownBinary; use crate::prune::GeoDistancePrune; +use crate::prune::GeoIntersectsPrune; use crate::scalar_fn::contains::GeoContains; use crate::scalar_fn::distance::GeoDistance; use crate::scalar_fn::intersects::GeoIntersects; @@ -70,6 +71,7 @@ pub fn initialize(session: &VortexSession) { // geometry columns. session.aggregate_fns().register(GeometryAabb); - // Register the spatial pruning rule that uses that AABB. + // Register the spatial pruning rules that use that AABB. session.stats().register_rewrite(GeoDistancePrune); + session.stats().register_rewrite(GeoIntersectsPrune); } diff --git a/vortex-geo/src/prune.rs b/vortex-geo/src/prune/distance.rs similarity index 54% rename from vortex-geo/src/prune.rs rename to vortex-geo/src/prune/distance.rs index 6b86c57c3e2..380571c11f9 100644 --- a/vortex-geo/src/prune.rs +++ b/vortex-geo/src/prune/distance.rs @@ -1,52 +1,36 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Chunk pruning for spatial filters, using the per-chunk [`GeometryAabb`] axis-aligned -//! bounding box (AABB). +//! `ST_Distance(geom, const) radius` pruning. -use geo::BoundingRect; use geo::Rect as GeoRect; -use vortex_array::VortexSessionExecute; -use vortex_array::aggregate_fn::AggregateFnVTableExt; -use vortex_array::aggregate_fn::EmptyOptions; use vortex_array::expr::Expression; -use vortex_array::expr::case_when; -use vortex_array::expr::checked_add; -use vortex_array::expr::ext_storage; -use vortex_array::expr::get_item; use vortex_array::expr::gt; use vortex_array::expr::gt_eq; -use vortex_array::expr::is_root; use vortex_array::expr::lit; use vortex_array::expr::lt; use vortex_array::expr::lt_eq; -use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_array::scalar_fn::ScalarFnVTableExt; use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::literal::Literal; use vortex_array::scalar_fn::fns::operators::Operator; use vortex_array::stats::rewrite::StatsRewriteCtx; use vortex_array::stats::rewrite::StatsRewriteRule; -use vortex_array::stats::stat; -use vortex_error::VortexExpect; use vortex_error::VortexResult; -use crate::aggregate_fn::GeometryAabb; -use crate::extension::is_native_geometry; -use crate::extension::single_geometry; +use super::aabb_stat; +use super::geometry_and_constant; +use super::max_dist_sq; +use super::min_dist_sq; +use super::query_aabb; use crate::scalar_fn::distance::GeoDistance; -/// Prunes chunks for `ST_Distance(geom, const) r` filters using the chunk's [`GeometryAabb`] -/// bounding box. Register it with `crate::initialize`; a chunk written without the `GeometryAabb` -/// statistic is scanned rather than skipped. +/// Prunes chunks for `ST_Distance(geom, const) r` filters. /// /// All four comparisons prune: `<= r` / `< r` skip a chunk whose box is wholly beyond `r` (box /// min-distance); `>= r` / `> r` skip one wholly within `r` (box max-distance). `==` / `!=` don't -/// prune. To add another spatial predicate, write a sibling [`StatsRewriteRule`] from the -/// `geometry_and_constant` + `distance_prune_proof` helpers; no new statistic or file-format change -/// is needed. +/// prune. #[derive(Debug)] pub struct GeoDistancePrune; @@ -78,15 +62,12 @@ impl StatsRewriteRule for GeoDistancePrune { if distance.as_opt::().is_none() { return Ok(None); } - let Some((geom, constant)) = geometry_and_constant(distance, ctx)? else { - return Ok(None); - }; let Some(radius) = expr.child(1).as_opt::() else { return Ok(None); }; - // Casts any primitive radius (integer literals included); it fails only for a null or - // non-primitive literal, where falling through means "scan the chunk", which is always - // sound. + // Casts any primitive radius (filters arrive uncoerced, so integer literals are + // legitimate); a null or non-primitive (e.g. extension-typed) literal has no value to + // reason about, decline and scan, never error. let Ok(radius) = f64::try_from(radius) else { return Ok(None); }; @@ -96,49 +77,20 @@ impl StatsRewriteRule for GeoDistancePrune { return Ok(None); } - // Reduce `const` (any geometry type) to its AABB. Every row sits in the chunk AABB - // and `const` in this box, so the box-to-box distance bounds the true distance soundly for - // any geometry types. - let mut exec = ctx.session().create_execution_ctx(); - let Some(query) = single_geometry(constant, &mut exec)?.bounding_rect() else { + let Some((geom, constant)) = geometry_and_constant(distance, ctx)? else { + return Ok(None); + }; + let Some(query) = query_aabb(constant, ctx)? else { return Ok(None); }; Ok(distance_prune_proof(geom, query, op, radius)) } } -/// Shared AABB-pruning helper: split a symmetric geo predicate's operands into the scope-rooted -/// geometry column and the constant's scalar, or `None` when the expression doesn't have that -/// shape or the geometry's dtype has no [`GeometryAabb`] support. Symmetric only - an asymmetric -/// predicate that needs to know *which* operand is the column must recover the role separately. -fn geometry_and_constant<'a>( - expr: &'a Expression, - ctx: &StatsRewriteCtx<'_>, -) -> VortexResult> { - // The predicate is symmetric, so the geometry column (scope root) and the constant may be on - // either side. - let (lhs, rhs) = (expr.child(0), expr.child(1)); - let (geom, constant) = if is_root(lhs) { - (lhs, rhs) - } else if is_root(rhs) { - (rhs, lhs) - } else { - return Ok(None); - }; - - // A `GeometryAabb` stat reference only binds for dtypes it supports; anything else (e.g. a - // WKB column) must fall through to the scan. - if !is_native_geometry(&ctx.return_dtype(geom)?) { - return Ok(None); - } - - Ok(constant.as_opt::().map(|scalar| (geom, scalar))) -} - -/// Build the prune proof for `ST_Distance(geom, const) radius` from the chunk's bounding-box -/// stat, or `None` when this operator/radius cannot prune. `<=` / `<` prune a chunk whose box is -/// wholly beyond `radius` (min box-distance); `>=` / `>` prune one wholly within it (max -/// box-distance). Every row sits in the box, so those bounds prove the whole chunk one-sided. +/// Build the prune proof for `ST_Distance(geom, const) radius`, or `None` when this +/// operator/radius cannot prune. The box-to-box distance bounds every row's true distance for any +/// geometry types: `<=` / `<` prune a chunk whose box is wholly beyond `radius` (min +/// box-distance); `>=` / `>` prune one wholly within it (max box-distance). /// /// A distance is always `>= 0`, which decides the degenerate radii up front. fn distance_prune_proof( @@ -157,9 +109,8 @@ fn distance_prune_proof( Operator::Gt if radius < 0.0 => return None, _ => {} } - // The stat is read through `ext_storage`/`get_item`, which propagate a missing stat (null) to - // "keep the chunk". Compared squared to avoid a `sqrt`; all operands are `>= 0`. - let aabb = ext_storage(stat(geom.clone(), GeometryAabb.bind(EmptyOptions))); + // Compared squared to avoid a `sqrt`; all operands are `>= 0`. + let aabb = aabb_stat(geom); let r2 = lit(radius * radius); Some(match op { // Beyond the threshold: even the nearest the box can be exceeds `r`. @@ -172,112 +123,41 @@ fn distance_prune_proof( }) } -/// Squared minimum distance between the chunk box `aabb` and the query box - a lower bound on every -/// row's distance. `dx^2 + dy^2`, each axis gap clamped at zero (zero when the intervals overlap). -fn min_dist_sq(aabb: &Expression, query: GeoRect) -> Expression { - let field = |name: &str| get_item(name, aabb.clone()); - // max(0, q_lo - aabb_hi, aabb_lo - q_hi): positive only when the intervals are separated. - let gap = |q_lo: f64, q_hi: f64, lo: Expression, hi: Expression| { - maximum( - lit(0.0), - maximum( - binop(Operator::Sub, lit(q_lo), hi), - binop(Operator::Sub, lo, lit(q_hi)), - ), - ) - }; - let dx = gap(query.min().x, query.max().x, field("xmin"), field("xmax")); - let dy = gap(query.min().y, query.max().y, field("ymin"), field("ymax")); - checked_add(square(dx), square(dy)) -} - -/// Squared maximum distance between the chunk box `aabb` and the query box - an upper bound on every -/// row's distance. `Dx^2 + Dy^2`, each axis span the full extent of the two intervals' union. -fn max_dist_sq(aabb: &Expression, query: GeoRect) -> Expression { - let field = |name: &str| get_item(name, aabb.clone()); - // max(q_hi, aabb_hi) - min(q_lo, aabb_lo): the farthest two points of the boxes can be on an axis. - // The (nullable) AABB field is passed as the second arg so `case_when`'s else branch carries the - // nullability - a missing stat then propagates null through to "keep the chunk". - let span = |q_lo: f64, q_hi: f64, lo: Expression, hi: Expression| { - binop( - Operator::Sub, - maximum(lit(q_hi), hi), - minimum(lit(q_lo), lo), - ) - }; - let dx = span(query.min().x, query.max().x, field("xmin"), field("xmax")); - let dy = span(query.min().y, query.max().y, field("ymin"), field("ymax")); - checked_add(square(dx), square(dy)) -} - -/// `a b` as a binary-operator expression. -fn binop(op: Operator, a: Expression, b: Expression) -> Expression { - Binary - .try_new_expr(op, [a, b]) - .vortex_expect("binary expression") -} - -/// `e * e` as an expression. -fn square(e: Expression) -> Expression { - binop(Operator::Mul, e.clone(), e) -} - -/// `max(a, b)` as an expression. -fn maximum(a: Expression, b: Expression) -> Expression { - case_when(gt(a.clone(), b.clone()), a, b) -} - -/// `min(a, b)` as an expression. -fn minimum(a: Expression, b: Expression) -> Expression { - case_when(lt(a.clone(), b.clone()), a, b) -} - #[cfg(test)] mod tests { - use std::sync::Arc; - use rstest::rstest; - use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::aggregate_fn::AggregateFnVTableExt; - use vortex_array::aggregate_fn::EmptyOptions as AggregateEmptyOptions; - use vortex_array::arrays::ExtensionArray; - use vortex_array::arrays::PrimitiveArray; - use vortex_array::arrays::StructArray; use vortex_array::dtype::DType; - use vortex_array::dtype::FieldNames; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::dtype::extension::ExtDType; use vortex_array::expr::Expression; use vortex_array::expr::gt_eq; use vortex_array::expr::lit; use vortex_array::expr::lt_eq; use vortex_array::expr::root; + use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::ScalarFnVTableExt; use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::operators::Operator; use vortex_array::stats::rewrite::StatsRewriteCtx; use vortex_array::stats::rewrite::StatsRewriteRule; - use vortex_array::validity::Validity; use vortex_error::VortexResult; - use vortex_layout::layouts::zoned::zone_map::ZoneMap; use super::GeoDistancePrune; - use crate::aggregate_fn::GeometryAabb; - use crate::extension::GeoMetadata; - use crate::extension::Rect; + use crate::prune::test_harness::aabb_zone_map; + use crate::prune::test_harness::empty_zone_map; use crate::scalar_fn::distance::GeoDistance; use crate::test_harness::geo_session; use crate::test_harness::point_column; /// Run the rule against `GeoDistance(root, origin) radius`, operands swapped when - /// `geom_first` is false. + /// `geom_first` is false. The radius is any literal scalar, matching the uncoerced filter + /// expressions the rule sees in production. fn falsify_distance( operator: Operator, geom_first: bool, - radius: f64, + radius: impl Into, ) -> VortexResult> { let session = geo_session(); let mut ctx = session.create_execution_ctx(); @@ -290,7 +170,7 @@ mod tests { [lit(origin), root()] }; let distance = GeoDistance.new_expr(EmptyOptions, operands); - let predicate = Binary.new_expr(operator, [distance, lit(radius)]); + let predicate = Binary.new_expr(operator, [distance, lit(radius.into())]); GeoDistancePrune.falsify(&predicate, &StatsRewriteCtx::new(&session, &scope)) } @@ -312,10 +192,12 @@ mod tests { Ok(()) } - /// Distance is symmetric: `GeoDistance(const, geom) <= r` falsifies just like the geom-first form. - #[test] - fn falsifies_with_constant_as_left_operand() -> VortexResult<()> { - assert!(falsify_distance(Operator::Lte, false, 0.5)?.is_some()); + /// Distance is symmetric: both operand orders produce a proof. + #[rstest] + #[case(true)] + #[case(false)] + fn falsifies_either_operand_order(#[case] geom_first: bool) -> VortexResult<()> { + assert!(falsify_distance(Operator::Lte, geom_first, 0.5)?.is_some()); Ok(()) } @@ -333,6 +215,34 @@ mod tests { Ok(()) } + /// Filter expressions arrive uncoerced, so `distance <= 10` may carry an integer literal - + /// it casts and prunes like an f64 radius. + #[test] + fn integer_radius_prunes() -> VortexResult<()> { + assert!(falsify_distance(Operator::Lte, true, 10i64)?.is_some()); + Ok(()) + } + + /// A null radius has no value to reason about; the rule declines and the chunk is scanned. + #[test] + fn null_radius_never_prunes() -> VortexResult<()> { + let radius = Scalar::null(DType::Primitive(PType::F64, Nullability::Nullable)); + assert!(falsify_distance(Operator::Lte, true, radius)?.is_none()); + Ok(()) + } + + /// An extension-typed radius passes `Binary`'s typecheck (extension operands are exempt) but + /// has no numeric value - the rule declines rather than erroring, and the chunk is scanned. + #[test] + fn extension_radius_never_prunes() -> VortexResult<()> { + let session = geo_session(); + let mut ctx = session.create_execution_ctx(); + + let geometry = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; + assert!(falsify_distance(Operator::Lte, true, geometry)?.is_none()); + Ok(()) + } + /// A scope dtype without `GeometryAabb` support gets no proof - the stat reference would /// fail to bind at prune time. #[test] @@ -369,30 +279,11 @@ mod tests { let mut ctx = session.create_execution_ctx(); let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); - let aabb_fn = GeometryAabb.bind(AggregateEmptyOptions); - - // Two chunks' AABBs, stored as the native `geoarrow.box` stat column with default - // (unreferenced) metadata to match the aggregate's return dtype: chunk 0 near the origin - // (0,0..1,1), chunk 1 far away (100,100..101,101). - let ord = |a: f64, b: f64| PrimitiveArray::from_iter([a, b]).into_array(); - let boxes = StructArray::try_new( - ["xmin", "ymin", "xmax", "ymax"].into(), - vec![ - ord(0.0, 100.0), - ord(0.0, 100.0), - ord(1.0, 101.0), - ord(1.0, 101.0), - ], - 2, - Validity::AllValid, - )? - .into_array(); - let box_dtype = - ExtDType::::try_new(GeoMetadata::default(), boxes.dtype().clone())?.erased(); - let aabbs = ExtensionArray::try_new(box_dtype, boxes)?.into_array(); - let zone_array = StructArray::from_fields(&[(aabb_fn.to_string().as_str(), aabbs)])?; - let zone_map = - ZoneMap::try_new(point_dtype.clone(), zone_array, Arc::new([aabb_fn]), 1, 2)?; + // Chunk 0 near the origin (0,0..1,1), chunk 1 far away (100,100..101,101). + let zone_map = aabb_zone_map( + &point_dtype, + &[[0.0, 0.0, 1.0, 1.0], [100.0, 100.0, 101.0, 101.0]], + )?; let origin = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; let distance = GeoDistance.new_expr(EmptyOptions, [root(), lit(origin)]); @@ -408,31 +299,16 @@ mod tests { } /// The true-distance prune skips a chunk that is *diagonally* farther than `r`, even though - /// neither axis alone exceeds `r` - the case a per-axis box-overlap test would wrongly keep. + /// neither axis alone exceeds `r`, the case a per-axis box-overlap test would wrongly keep. #[test] fn prunes_diagonally_distant_chunk() -> VortexResult<()> { let session = geo_session(); let mut ctx = session.create_execution_ctx(); let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); - let aabb_fn = GeometryAabb.bind(AggregateEmptyOptions); - // One chunk, AABB (0.8,0.8)..(0.9,0.9): each axis is only 0.8 from the origin (<= r = 1), but // the near corner is sqrt(0.8^2 + 0.8^2) ~= 1.13 away (> 1), so no point in the box is within 1. - let ord = |a: f64| PrimitiveArray::from_iter([a]).into_array(); - let boxes = StructArray::try_new( - ["xmin", "ymin", "xmax", "ymax"].into(), - vec![ord(0.8), ord(0.8), ord(0.9), ord(0.9)], - 1, - Validity::AllValid, - )? - .into_array(); - let box_dtype = - ExtDType::::try_new(GeoMetadata::default(), boxes.dtype().clone())?.erased(); - let aabbs = ExtensionArray::try_new(box_dtype, boxes)?.into_array(); - let zone_array = StructArray::from_fields(&[(aabb_fn.to_string().as_str(), aabbs)])?; - let zone_map = - ZoneMap::try_new(point_dtype.clone(), zone_array, Arc::new([aabb_fn]), 1, 1)?; + let zone_map = aabb_zone_map(&point_dtype, &[[0.8, 0.8, 0.9, 0.9]])?; let origin = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; let distance = GeoDistance.new_expr(EmptyOptions, [root(), lit(origin)]); @@ -451,74 +327,51 @@ mod tests { Ok(()) } - /// Backward compat: a zone map written without the `GeometryAabb` stat (an older file) keeps - /// every zone - the missing stat binds to null and `null_as_false` retains the zone. + /// A `>= r` filter prunes a chunk lying wholly *within* `r` (every row nearer than `r`, so none + /// satisfy `>= r`) via the box max-distance, while a chunk beyond `r` is kept. #[test] - fn missing_aabb_stat_keeps_all_zones() -> VortexResult<()> { + fn prunes_within_chunk_for_far_filter() -> VortexResult<()> { let session = geo_session(); let mut ctx = session.create_execution_ctx(); let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); - let zone_map = ZoneMap::try_new( - point_dtype.clone(), - StructArray::try_new(FieldNames::empty(), vec![], 2, Validity::NonNullable)?, - Arc::new([]), - 1, - 2, + // Chunk 0 (AABB 0,0..0.5,0.5, farthest corner ~= 0.707) is entirely within 2 of the origin; + // chunk 1 (100,100..101,101) is entirely beyond it. + let zone_map = aabb_zone_map( + &point_dtype, + &[[0.0, 0.0, 0.5, 0.5], [100.0, 100.0, 101.0, 101.0]], )?; let origin = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; let distance = GeoDistance.new_expr(EmptyOptions, [root(), lit(origin)]); - let proof = lt_eq(distance, lit(0.5f64)) + let proof = gt_eq(distance, lit(2.0f64)) .falsify(&point_dtype, &session)? .expect("distance filter should be falsifiable"); + // Chunk 0 (within 2) is pruned for `>= 2`; chunk 1 (beyond 2) is kept. let mask = zone_map.prune(&proof, &session)?; - assert_eq!(mask.iter().collect::>(), vec![false, false]); + assert_eq!(mask.iter().collect::>(), vec![true, false]); Ok(()) } - /// A `>= r` filter prunes a chunk lying wholly *within* `r` (every row nearer than `r`, so none - /// satisfy `>= r`) via the box max-distance, while a chunk beyond `r` is kept. + /// Backward compat: a zone map written without the `GeometryAabb` stat (an older file) keeps + /// every zone, the missing stat binds to null and `null_as_false` retains the zone. #[test] - fn prunes_within_chunk_for_far_filter() -> VortexResult<()> { + fn missing_aabb_stat_keeps_all_zones() -> VortexResult<()> { let session = geo_session(); let mut ctx = session.create_execution_ctx(); let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); - let aabb_fn = GeometryAabb.bind(AggregateEmptyOptions); - - // Chunk 0 (AABB 0,0..0.5,0.5, farthest corner ~= 0.707) is entirely within 2 of the origin; - // chunk 1 (100,100..101,101) is entirely beyond it. - let ord = |a: f64, b: f64| PrimitiveArray::from_iter([a, b]).into_array(); - let boxes = StructArray::try_new( - ["xmin", "ymin", "xmax", "ymax"].into(), - vec![ - ord(0.0, 100.0), - ord(0.0, 100.0), - ord(0.5, 101.0), - ord(0.5, 101.0), - ], - 2, - Validity::AllValid, - )? - .into_array(); - let box_dtype = - ExtDType::::try_new(GeoMetadata::default(), boxes.dtype().clone())?.erased(); - let aabbs = ExtensionArray::try_new(box_dtype, boxes)?.into_array(); - let zone_array = StructArray::from_fields(&[(aabb_fn.to_string().as_str(), aabbs)])?; - let zone_map = - ZoneMap::try_new(point_dtype.clone(), zone_array, Arc::new([aabb_fn]), 1, 2)?; + let zone_map = empty_zone_map(&point_dtype)?; let origin = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; let distance = GeoDistance.new_expr(EmptyOptions, [root(), lit(origin)]); - let proof = gt_eq(distance, lit(2.0f64)) + let proof = lt_eq(distance, lit(0.5f64)) .falsify(&point_dtype, &session)? .expect("distance filter should be falsifiable"); - // Chunk 0 (within 2) is pruned for `>= 2`; chunk 1 (beyond 2) is kept. let mask = zone_map.prune(&proof, &session)?; - assert_eq!(mask.iter().collect::>(), vec![true, false]); + assert_eq!(mask.iter().collect::>(), vec![false, false]); Ok(()) } } diff --git a/vortex-geo/src/prune/intersects.rs b/vortex-geo/src/prune/intersects.rs new file mode 100644 index 00000000000..9af6c5b4d9e --- /dev/null +++ b/vortex-geo/src/prune/intersects.rs @@ -0,0 +1,166 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! `ST_Intersects(geom, const)` pruning. + +use vortex_array::expr::Expression; +use vortex_array::expr::gt; +use vortex_array::expr::lit; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::stats::rewrite::StatsRewriteCtx; +use vortex_array::stats::rewrite::StatsRewriteRule; +use vortex_error::VortexResult; + +use super::aabb_stat; +use super::geometry_and_constant; +use super::min_dist_sq; +use super::query_aabb; +use crate::scalar_fn::intersects::GeoIntersects; + +/// Prunes chunks for `ST_Intersects(geom, const)` filters: a chunk whose box is strictly +/// separated from the constant's bounding box cannot contain an intersecting row. +/// +/// Only the positive form prunes. `NOT ST_Intersects` cannot: it would need every row to provably +/// intersect the constant, and box overlap never proves geometry intersection. +#[derive(Debug)] +pub struct GeoIntersectsPrune; + +impl StatsRewriteRule for GeoIntersectsPrune { + fn scalar_fn_id(&self) -> ScalarFnId { + // Unlike the distance rule, the boolean predicate is itself the expression root. + GeoIntersects.id() + } + + fn falsify( + &self, + expr: &Expression, + ctx: &StatsRewriteCtx<'_>, + ) -> VortexResult> { + let Some((geom, constant)) = geometry_and_constant(expr, ctx)? else { + return Ok(None); + }; + let Some(query) = query_aabb(constant, ctx)? else { + return Ok(None); + }; + // Disjoint iff the minimum box-to-box distance is positive. Strictly (`gt`, not `gt_eq`): + // boxes that merely touch must scan, since touching geometries do intersect. + Ok(Some(gt(min_dist_sq(&aabb_stat(geom), query), lit(0.0)))) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::VortexSessionExecute; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::expr::Expression; + use vortex_array::expr::lit; + use vortex_array::expr::root; + use vortex_array::scalar_fn::EmptyOptions; + use vortex_array::scalar_fn::ScalarFnVTableExt; + use vortex_array::stats::rewrite::StatsRewriteCtx; + use vortex_array::stats::rewrite::StatsRewriteRule; + use vortex_error::VortexResult; + + use super::GeoIntersectsPrune; + use crate::prune::test_harness::aabb_zone_map; + use crate::prune::test_harness::empty_zone_map; + use crate::scalar_fn::intersects::GeoIntersects; + use crate::test_harness::geo_session; + use crate::test_harness::point_column; + + /// Run the intersects rule against `GeoIntersects(root, point(1.0, 0.5))`, operands swapped + /// when `geom_first` is false. + fn falsify_intersects(geom_first: bool) -> VortexResult> { + let session = geo_session(); + let mut ctx = session.create_execution_ctx(); + + let scope = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let query = point_column(vec![1.0], vec![0.5])?.execute_scalar(0, &mut ctx)?; + let operands = if geom_first { + [root(), lit(query)] + } else { + [lit(query), root()] + }; + let predicate = GeoIntersects.new_expr(EmptyOptions, operands); + GeoIntersectsPrune.falsify(&predicate, &StatsRewriteCtx::new(&session, &scope)) + } + + /// Intersects is symmetric: both operand orders produce a proof. + #[rstest] + #[case(true)] + #[case(false)] + fn falsifies_either_operand_order(#[case] geom_first: bool) -> VortexResult<()> { + assert!(falsify_intersects(geom_first)?.is_some()); + Ok(()) + } + + /// A scope dtype without `GeometryAabb` support gets no proof, the stat reference would + /// fail to bind at prune time. + #[test] + fn unsupported_scope_is_not_pruned() -> VortexResult<()> { + let session = geo_session(); + let mut ctx = session.create_execution_ctx(); + + let scope = DType::Primitive(PType::F64, Nullability::NonNullable); + let query = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; + let predicate = GeoIntersects.new_expr(EmptyOptions, [root(), lit(query)]); + + let ctx = StatsRewriteCtx::new(&session, &scope); + assert!(GeoIntersectsPrune.falsify(&predicate, &ctx)?.is_none()); + Ok(()) + } + + /// End-to-end: a zone strictly separated from the query is skipped; zones containing or merely + /// touching the query must scan, touching geometries intersect under OGC semantics. + #[test] + fn prunes_disjoint_keeps_touching_and_containing() -> VortexResult<()> { + let session = geo_session(); + let mut ctx = session.create_execution_ctx(); + + let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + // Query point (1.0, 0.5): zone 0 contains it, zone 1 only touches it at `x == 1`, zone 2 + // is strictly separated. + let zone_map = aabb_zone_map( + &point_dtype, + &[ + [0.5, 0.0, 1.5, 1.0], + [-1.0, 0.0, 1.0, 1.0], + [3.0, 0.0, 4.0, 1.0], + ], + )?; + + let query = point_column(vec![1.0], vec![0.5])?.execute_scalar(0, &mut ctx)?; + let predicate = GeoIntersects.new_expr(EmptyOptions, [root(), lit(query)]); + let proof = predicate + .falsify(&point_dtype, &session)? + .expect("intersects filter should be falsifiable"); + + let mask = zone_map.prune(&proof, &session)?; + assert_eq!(mask.iter().collect::>(), vec![false, false, true]); + Ok(()) + } + + /// Backward compat: a zone map written without the `GeometryAabb` stat keeps every zone. + #[test] + fn missing_aabb_stat_keeps_all_zones() -> VortexResult<()> { + let session = geo_session(); + let mut ctx = session.create_execution_ctx(); + + let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + let zone_map = empty_zone_map(&point_dtype)?; + + let query = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; + let proof = GeoIntersects + .new_expr(EmptyOptions, [root(), lit(query)]) + .falsify(&point_dtype, &session)? + .expect("intersects filter should be falsifiable"); + + let mask = zone_map.prune(&proof, &session)?; + assert_eq!(mask.iter().collect::>(), vec![false, false]); + Ok(()) + } +} diff --git a/vortex-geo/src/prune/mod.rs b/vortex-geo/src/prune/mod.rs new file mode 100644 index 00000000000..00200f1500c --- /dev/null +++ b/vortex-geo/src/prune/mod.rs @@ -0,0 +1,164 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Chunk pruning for spatial filters, using the per-chunk [`GeometryAabb`] axis-aligned +//! bounding box (AABB). +//! +//! One module per predicate. To prune a new one: recover the geometry column and the constant +//! (symmetric predicates can use `geometry_and_constant`), reduce the constant to its box with +//! `query_aabb`, build the proof over `aabb_stat` from the box-relation builders (`min_dist_sq`, +//! `max_dist_sq`), and register the rule in [`crate::initialize`]. No new statistic or +//! file-format change is needed. + +mod distance; +mod intersects; +#[cfg(test)] +mod test_harness; + +pub use distance::GeoDistancePrune; +use geo::BoundingRect; +use geo::Rect as GeoRect; +pub use intersects::GeoIntersectsPrune; +use vortex_array::VortexSessionExecute; +use vortex_array::aggregate_fn::AggregateFnVTableExt; +use vortex_array::aggregate_fn::EmptyOptions; +use vortex_array::expr::Expression; +use vortex_array::expr::case_when; +use vortex_array::expr::checked_add; +use vortex_array::expr::ext_storage; +use vortex_array::expr::get_item; +use vortex_array::expr::gt; +use vortex_array::expr::is_root; +use vortex_array::expr::lit; +use vortex_array::expr::lt; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::scalar_fn::fns::binary::Binary; +use vortex_array::scalar_fn::fns::literal::Literal; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::stats::rewrite::StatsRewriteCtx; +use vortex_array::stats::stat; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; + +use crate::aggregate_fn::GeometryAabb; +use crate::extension::is_native_geometry; +use crate::extension::single_geometry; + +/// Splits a symmetric two-operand geo predicate into the scope-rooted geometry column and the +/// constant operand's scalar. +/// +/// `None` means the rule must decline: the expression doesn't have the `f(column, constant)` +/// shape (in either operand order), or the column's dtype carries no [`GeometryAabb`] statistic. +/// An asymmetric predicate (e.g. a future contains) must recover which operand is the column +/// itself instead of calling this. +fn geometry_and_constant<'a>( + expr: &'a Expression, + ctx: &StatsRewriteCtx<'_>, +) -> VortexResult> { + // The predicate is symmetric, so the column (scope root) and the constant may be on either + // side. + let (lhs, rhs) = (expr.child(0), expr.child(1)); + let (geom, constant) = if is_root(lhs) { + (lhs, rhs) + } else if is_root(rhs) { + (rhs, lhs) + } else { + return Ok(None); + }; + + // A `GeometryAabb` stat reference only binds for dtypes it supports; anything else (e.g. a + // WKB column) must fall through to the scan. + if !is_native_geometry(&ctx.return_dtype(geom)?) { + return Ok(None); + } + + Ok(constant.as_opt::().map(|scalar| (geom, scalar))) +} + +/// The 2D bounding box of a constant geometry of any type, or `None` for one without an extent +/// (e.g. an empty geometry). +/// +/// Prove claims against this box rather than the constant itself: the constant lies inside it, +/// so whatever holds for the box holds for the geometry. +fn query_aabb(constant: &Scalar, ctx: &StatsRewriteCtx<'_>) -> VortexResult>> { + // Decoding the constant into a concrete geometry runs through the compute stack, which needs + // an execution context. + let mut exec = ctx.session().create_execution_ctx(); + Ok(single_geometry(constant, &mut exec)?.bounding_rect()) +} + +/// The chunk's AABB statistic, as the storage struct with `xmin`/`ymin`/`xmax`/`ymax` fields. +/// +/// A chunk written without the statistic reads as null here; every proof built on top must let +/// that null propagate to its root, where the zone map keeps the chunk. +fn aabb_stat(geom: &Expression) -> Expression { + // `ext_storage` unwraps the native `geoarrow.box` stat value to its backing struct, so + // proofs can `get_item` the coordinate fields. + ext_storage(stat(geom.clone(), GeometryAabb.bind(EmptyOptions))) +} + +/// Lower bound on every row's squared distance to the query AABB: zero when the boxes overlap or +/// touch, positive iff they are strictly separated. +/// +/// Prunes "near" predicates: `min_dist_sq > r^2` proves every row is farther than `r`. +fn min_dist_sq(aabb: &Expression, query: GeoRect) -> Expression { + let field = |name: &str| get_item(name, aabb.clone()); + // Per axis: gap = max(0, q_lo - aabb_hi, aabb_lo - q_hi), positive only when the intervals + // are separated. The nearest two points of the boxes are one axis-gap apart per axis, so the + // squared distance is gap_x^2 + gap_y^2 (squared throughout to avoid a sqrt). + let gap = |q_lo: f64, q_hi: f64, lo: Expression, hi: Expression| { + maximum( + lit(0.0), + maximum( + binop(Operator::Sub, lit(q_lo), hi), + binop(Operator::Sub, lo, lit(q_hi)), + ), + ) + }; + let dx = gap(query.min().x, query.max().x, field("xmin"), field("xmax")); + let dy = gap(query.min().y, query.max().y, field("ymin"), field("ymax")); + checked_add(square(dx), square(dy)) +} + +/// Upper bound on every row's squared distance to the query AABB. +/// +/// Prunes "far" predicates: `max_dist_sq < r^2` proves every row is within `r`. +fn max_dist_sq(aabb: &Expression, query: GeoRect) -> Expression { + let field = |name: &str| get_item(name, aabb.clone()); + // Per axis: span = max(q_hi, aabb_hi) - min(q_lo, aabb_lo), the farthest two points of the + // boxes can be apart. The nullable AABB field is the second `maximum`/`minimum` argument so + // that `case_when`'s else branch carries the nullability - a missing stat propagates null. + let span = |q_lo: f64, q_hi: f64, lo: Expression, hi: Expression| { + binop( + Operator::Sub, + maximum(lit(q_hi), hi), + minimum(lit(q_lo), lo), + ) + }; + let dx = span(query.min().x, query.max().x, field("xmin"), field("xmax")); + let dy = span(query.min().y, query.max().y, field("ymin"), field("ymax")); + checked_add(square(dx), square(dy)) +} + +/// `a b`. +fn binop(op: Operator, a: Expression, b: Expression) -> Expression { + Binary + .try_new_expr(op, [a, b]) + .vortex_expect("binary expression") +} + +/// `e * e`. +fn square(e: Expression) -> Expression { + binop(Operator::Mul, e.clone(), e) +} + +/// `max(a, b)`. +fn maximum(a: Expression, b: Expression) -> Expression { + case_when(gt(a.clone(), b.clone()), a, b) +} + +/// `min(a, b)`. +fn minimum(a: Expression, b: Expression) -> Expression { + case_when(lt(a.clone(), b.clone()), a, b) +} diff --git a/vortex-geo/src/prune/test_harness.rs b/vortex-geo/src/prune/test_harness.rs new file mode 100644 index 00000000000..fb3e12ef9d5 --- /dev/null +++ b/vortex-geo/src/prune/test_harness.rs @@ -0,0 +1,59 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Shared zone-map fixtures for the pruning-rule tests; each rule module owns its own cases. + +use std::sync::Arc; + +use vortex_array::IntoArray; +use vortex_array::aggregate_fn::AggregateFnVTableExt; +use vortex_array::aggregate_fn::EmptyOptions; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldNames; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; +use vortex_layout::layouts::zoned::zone_map::ZoneMap; + +use crate::aggregate_fn::GeometryAabb; +use crate::extension::GeoMetadata; +use crate::extension::Rect; + +/// A single-column zone map holding one native-box AABB stat row (`[xmin, ymin, xmax, ymax]`) +/// per zone, with default (unreferenced) metadata to match the aggregate's return dtype. +pub(super) fn aabb_zone_map(point_dtype: &DType, boxes: &[[f64; 4]]) -> VortexResult { + let aabb_fn = GeometryAabb.bind(EmptyOptions); + let col = |i: usize| PrimitiveArray::from_iter(boxes.iter().map(move |b| b[i])).into_array(); + let storage = StructArray::try_new( + ["xmin", "ymin", "xmax", "ymax"].into(), + vec![col(0), col(1), col(2), col(3)], + boxes.len(), + Validity::AllValid, + )? + .into_array(); + let box_dtype = + ExtDType::::try_new(GeoMetadata::default(), storage.dtype().clone())?.erased(); + let aabbs = ExtensionArray::try_new(box_dtype, storage)?.into_array(); + let zone_array = StructArray::from_fields(&[(aabb_fn.to_string().as_str(), aabbs)])?; + ZoneMap::try_new( + point_dtype.clone(), + zone_array, + Arc::new([aabb_fn]), + 1, + boxes.len() as u64, + ) +} + +/// A two-zone zone map with no stat columns at all, as written before the AABB stat existed. +pub(super) fn empty_zone_map(point_dtype: &DType) -> VortexResult { + ZoneMap::try_new( + point_dtype.clone(), + StructArray::try_new(FieldNames::empty(), vec![], 2, Validity::NonNullable)?, + Arc::new([]), + 1, + 2, + ) +}