vortex_geo/prune/
intersects.rs1use 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#[derive(Debug)]
27pub struct GeoIntersectsPrune;
28
29impl StatsRewriteRule for GeoIntersectsPrune {
30 fn scalar_fn_id(&self) -> ScalarFnId {
31 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 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 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 #[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 #[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 #[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 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 #[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}