Skip to main content

vortex_geo/prune/
intersects.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! `ST_Intersects(geom, const)` pruning.
5
6use vortex_array::expr::Expression;
7use vortex_array::expr::gt;
8use vortex_array::expr::lit;
9use vortex_array::scalar_fn::ScalarFnId;
10use vortex_array::scalar_fn::ScalarFnVTable;
11use vortex_array::stats::rewrite::StatsRewriteCtx;
12use vortex_array::stats::rewrite::StatsRewriteRule;
13use vortex_error::VortexResult;
14
15use super::aabb_stat;
16use super::geometry_and_constant;
17use super::min_dist_sq;
18use super::query_aabb;
19use crate::scalar_fn::intersects::GeoIntersects;
20
21/// Prunes chunks for `ST_Intersects(geom, const)` filters: a chunk whose box is strictly
22/// separated from the constant's bounding box cannot contain an intersecting row.
23///
24/// Only the positive form prunes. `NOT ST_Intersects` cannot: it would need every row to provably
25/// intersect the constant, and box overlap never proves geometry intersection.
26#[derive(Debug)]
27pub struct GeoIntersectsPrune;
28
29impl StatsRewriteRule for GeoIntersectsPrune {
30    fn scalar_fn_id(&self) -> ScalarFnId {
31        // Unlike the distance rule, the boolean predicate is itself the expression root.
32        GeoIntersects.id()
33    }
34
35    fn falsify(
36        &self,
37        expr: &Expression,
38        ctx: &StatsRewriteCtx<'_>,
39    ) -> VortexResult<Option<Expression>> {
40        let Some((geom, constant)) = geometry_and_constant(expr, ctx)? else {
41            return Ok(None);
42        };
43        let Some(query) = query_aabb(constant, ctx)? else {
44            return Ok(None);
45        };
46        // Disjoint iff the minimum box-to-box distance is positive. Strictly (`gt`, not `gt_eq`):
47        // boxes that merely touch must scan, since touching geometries do intersect.
48        Ok(Some(gt(min_dist_sq(&aabb_stat(geom), query), lit(0.0))))
49    }
50}
51
52#[cfg(test)]
53mod tests {
54    use rstest::rstest;
55    use vortex_array::VortexSessionExecute;
56    use vortex_array::dtype::DType;
57    use vortex_array::dtype::Nullability;
58    use vortex_array::dtype::PType;
59    use vortex_array::expr::Expression;
60    use vortex_array::expr::lit;
61    use vortex_array::expr::root;
62    use vortex_array::scalar_fn::EmptyOptions;
63    use vortex_array::scalar_fn::ScalarFnVTableExt;
64    use vortex_array::stats::rewrite::StatsRewriteCtx;
65    use vortex_array::stats::rewrite::StatsRewriteRule;
66    use vortex_error::VortexResult;
67
68    use super::GeoIntersectsPrune;
69    use crate::prune::test_harness::aabb_zone_map;
70    use crate::prune::test_harness::empty_zone_map;
71    use crate::scalar_fn::intersects::GeoIntersects;
72    use crate::test_harness::geo_session;
73    use crate::test_harness::point_column;
74
75    /// Run the intersects rule against `GeoIntersects(root, point(1.0, 0.5))`, operands swapped
76    /// when `geom_first` is false.
77    fn falsify_intersects(geom_first: bool) -> VortexResult<Option<Expression>> {
78        let session = geo_session();
79        let mut ctx = session.create_execution_ctx();
80
81        let scope = point_column(vec![0.0], vec![0.0])?.dtype().clone();
82        let query = point_column(vec![1.0], vec![0.5])?.execute_scalar(0, &mut ctx)?;
83        let operands = if geom_first {
84            [root(), lit(query)]
85        } else {
86            [lit(query), root()]
87        };
88        let predicate = GeoIntersects.new_expr(EmptyOptions, operands);
89        GeoIntersectsPrune.falsify(&predicate, &StatsRewriteCtx::new(&session, &scope))
90    }
91
92    /// Intersects is symmetric: both operand orders produce a proof.
93    #[rstest]
94    #[case(true)]
95    #[case(false)]
96    fn falsifies_either_operand_order(#[case] geom_first: bool) -> VortexResult<()> {
97        assert!(falsify_intersects(geom_first)?.is_some());
98        Ok(())
99    }
100
101    /// A scope dtype without `GeometryAabb` support gets no proof, the stat reference would
102    /// fail to bind at prune time.
103    #[test]
104    fn unsupported_scope_is_not_pruned() -> VortexResult<()> {
105        let session = geo_session();
106        let mut ctx = session.create_execution_ctx();
107
108        let scope = DType::Primitive(PType::F64, Nullability::NonNullable);
109        let query = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?;
110        let predicate = GeoIntersects.new_expr(EmptyOptions, [root(), lit(query)]);
111
112        let ctx = StatsRewriteCtx::new(&session, &scope);
113        assert!(GeoIntersectsPrune.falsify(&predicate, &ctx)?.is_none());
114        Ok(())
115    }
116
117    /// End-to-end: a zone strictly separated from the query is skipped; zones containing or merely
118    /// touching the query must scan, touching geometries intersect under OGC semantics.
119    #[test]
120    fn prunes_disjoint_keeps_touching_and_containing() -> VortexResult<()> {
121        let session = geo_session();
122        let mut ctx = session.create_execution_ctx();
123
124        let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone();
125        // Query point (1.0, 0.5): zone 0 contains it, zone 1 only touches it at `x == 1`, zone 2
126        // is strictly separated.
127        let zone_map = aabb_zone_map(
128            &point_dtype,
129            &[
130                [0.5, 0.0, 1.5, 1.0],
131                [-1.0, 0.0, 1.0, 1.0],
132                [3.0, 0.0, 4.0, 1.0],
133            ],
134        )?;
135
136        let query = point_column(vec![1.0], vec![0.5])?.execute_scalar(0, &mut ctx)?;
137        let predicate = GeoIntersects.new_expr(EmptyOptions, [root(), lit(query)]);
138        let proof = predicate
139            .falsify(&point_dtype, &session)?
140            .expect("intersects filter should be falsifiable");
141
142        let mask = zone_map.prune(&proof, &session)?;
143        assert_eq!(mask.iter().collect::<Vec<bool>>(), vec![false, false, true]);
144        Ok(())
145    }
146
147    /// Backward compat: a zone map written without the `GeometryAabb` stat keeps every zone.
148    #[test]
149    fn missing_aabb_stat_keeps_all_zones() -> VortexResult<()> {
150        let session = geo_session();
151        let mut ctx = session.create_execution_ctx();
152
153        let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone();
154        let zone_map = empty_zone_map(&point_dtype)?;
155
156        let query = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?;
157        let proof = GeoIntersects
158            .new_expr(EmptyOptions, [root(), lit(query)])
159            .falsify(&point_dtype, &session)?
160            .expect("intersects filter should be falsifiable");
161
162        let mask = zone_map.prune(&proof, &session)?;
163        assert_eq!(mask.iter().collect::<Vec<bool>>(), vec![false, false]);
164        Ok(())
165    }
166}