1use std::sync::Arc;
7
8use vortex_array::ArrayRef;
9use vortex_array::IntoArray;
10use vortex_array::VortexSessionExecute;
11use vortex_array::aggregate_fn::fns::all_nan::AllNan;
12use vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan;
13use vortex_array::aggregate_fn::fns::all_non_null::AllNonNull;
14use vortex_array::aggregate_fn::fns::all_null::AllNull;
15use vortex_array::aggregate_fn::fns::nan_count::NanCount;
16use vortex_array::arrays::ConstantArray;
17use vortex_array::arrays::PrimitiveArray;
18use vortex_array::arrays::StructArray;
19use vortex_array::arrays::struct_::StructArrayExt;
20use vortex_array::dtype::DType;
21use vortex_array::dtype::Nullability;
22use vortex_array::expr::Expression;
23use vortex_array::expr::eq;
24use vortex_array::expr::get_item;
25use vortex_array::expr::is_root;
26use vortex_array::expr::lit;
27use vortex_array::expr::root;
28use vortex_array::expr::stats::Stat;
29use vortex_array::expr::traversal::NodeExt;
30use vortex_array::expr::traversal::Transformed;
31use vortex_array::scalar::Scalar;
32use vortex_array::scalar_fn::EmptyOptions;
33use vortex_array::scalar_fn::ScalarFnVTableExt;
34use vortex_array::scalar_fn::fns::stat::StatFn;
35use vortex_array::scalar_fn::internal::row_count::RowCount;
36use vortex_array::scalar_fn::internal::row_count::contains_row_count;
37use vortex_array::scalar_fn::internal::row_count::substitute_row_count;
38use vortex_array::validity::Validity;
39use vortex_buffer::buffer;
40use vortex_error::VortexResult;
41use vortex_error::vortex_bail;
42use vortex_mask::Mask;
43use vortex_runend::RunEnd;
44use vortex_session::VortexSession;
45
46use crate::layouts::zoned::schema::stats_table_dtype;
47
48#[derive(Clone)]
53pub struct ZoneMap {
54 column_dtype: DType,
56 array: StructArray,
58 zone_len: u64,
60 row_count: u64,
62}
63
64impl ZoneMap {
65 pub fn try_new(
68 column_dtype: DType,
69 array: StructArray,
70 stats: Arc<[Stat]>,
71 zone_len: u64,
72 row_count: u64,
73 ) -> VortexResult<Self> {
74 let expected_dtype = stats_table_dtype(&column_dtype, &stats);
75 if &expected_dtype != array.dtype() {
76 vortex_bail!("Array dtype does not match expected zone map dtype: {expected_dtype}");
77 }
78
79 Ok(unsafe { Self::new_unchecked(column_dtype, array, zone_len, row_count) })
81 }
82
83 pub(super) unsafe fn new_unchecked(
84 column_dtype: DType,
85 array: StructArray,
86 zone_len: u64,
87 row_count: u64,
88 ) -> Self {
89 Self {
90 column_dtype,
91 array,
92 zone_len,
93 row_count,
94 }
95 }
96
97 #[deprecated(note = "zone-map stats table dtypes are an internal layout detail")]
101 pub fn dtype_for_stats_table(column_dtype: &DType, present_stats: &[Stat]) -> DType {
102 stats_table_dtype(column_dtype, present_stats)
103 }
104
105 pub fn prune(&self, predicate: &Expression, session: &VortexSession) -> VortexResult<Mask> {
119 let mut ctx = session.create_execution_ctx();
120 let num_zones = self.array.len();
121 let predicate = self.lower_stats(predicate.clone())?;
122
123 let applied = self.array.clone().into_array().apply(&predicate)?;
124
125 if !contains_row_count(&applied) {
126 return applied.execute::<Mask>(&mut ctx);
127 }
128
129 let row_count_array = row_count_array(self.zone_len, self.row_count, num_zones)?;
130 let substituted = substitute_row_count(applied, &row_count_array)?;
131 substituted.execute::<Mask>(&mut ctx)
132 }
133
134 fn lower_stats(&self, predicate: Expression) -> VortexResult<Expression> {
135 predicate
139 .transform_down(|expr| {
140 if expr.is::<StatFn>() {
141 return self.lower_stat_fn(expr).map(Transformed::yes);
142 }
143
144 Ok(Transformed::no(expr))
145 })
146 .map(Transformed::into_inner)
147 }
148
149 fn lower_stat_fn(&self, expr: Expression) -> VortexResult<Expression> {
150 let options = expr.as_::<StatFn>();
154 let input = expr.child(0);
155 let input_dtype = input.return_dtype(&self.column_dtype)?;
156 let input_is_root = is_root(input);
157
158 if options.aggregate_fn().is::<AllNan>() {
159 if !has_nans(&input_dtype) {
160 return Ok(lit(false));
161 }
162 if !input_is_root {
163 return Ok(null_expr(DType::Bool(Nullability::NonNullable)));
164 }
165 return Ok(eq(self.stat_field_expr(Stat::NaNCount)?, row_count_expr()));
166 }
167
168 if options.aggregate_fn().is::<AllNonNan>() {
169 if !has_nans(&input_dtype) {
170 return Ok(lit(true));
171 }
172 if !input_is_root {
173 return Ok(null_expr(DType::Bool(Nullability::NonNullable)));
174 }
175 return Ok(eq(self.stat_field_expr(Stat::NaNCount)?, lit(0u64)));
176 }
177
178 if options.aggregate_fn().is::<NanCount>() && !has_nans(&input_dtype) {
179 return Ok(lit(0u64));
180 }
181
182 let Some(return_dtype) = options.aggregate_fn().return_dtype(&input_dtype) else {
188 return Ok(null_expr(input_dtype.as_nullable()));
189 };
190
191 if !input_is_root {
192 return Ok(null_expr(return_dtype));
193 }
194
195 if options.aggregate_fn().is::<AllNull>() {
196 return Ok(eq(self.stat_field_expr(Stat::NullCount)?, row_count_expr()));
197 }
198
199 if options.aggregate_fn().is::<AllNonNull>() {
200 return Ok(eq(self.stat_field_expr(Stat::NullCount)?, lit(0u64)));
201 }
202
203 let Some(stat) = Stat::from_aggregate_fn(options.aggregate_fn()) else {
204 return Ok(null_expr(return_dtype));
205 };
206
207 self.stat_field_expr(stat)
208 }
209
210 fn stat_field_expr(&self, stat: Stat) -> VortexResult<Expression> {
211 if self.array.unmasked_field_by_name_opt(stat.name()).is_some() {
212 return Ok(get_item(stat.name(), root()));
213 }
214
215 let Some(dtype) = stat.dtype(&self.column_dtype) else {
216 vortex_bail!(
217 "Stat {} does not support column dtype {}",
218 stat,
219 self.column_dtype
220 );
221 };
222 Ok(null_expr(dtype))
223 }
224}
225
226fn row_count_expr() -> Expression {
227 RowCount.new_expr(EmptyOptions, [])
228}
229
230fn null_expr(dtype: DType) -> Expression {
231 lit(Scalar::null(dtype.as_nullable()))
232}
233
234fn has_nans(dtype: &DType) -> bool {
235 matches!(dtype, DType::Primitive(ptype, _) if ptype.is_float())
236}
237
238fn row_count_array(zone_len: u64, row_count: u64, num_zones: usize) -> VortexResult<ArrayRef> {
244 if num_zones == 0 {
245 return Ok(ConstantArray::new(0u64, 0).into_array());
246 }
247
248 let last_zone_len = row_count - zone_len.saturating_mul((num_zones as u64) - 1);
249 if num_zones == 1 || last_zone_len == zone_len {
250 return Ok(ConstantArray::new(last_zone_len, num_zones).into_array());
251 }
252
253 let ends = unsafe {
254 PrimitiveArray::new_unchecked(
255 buffer![num_zones as u64 - 1, num_zones as u64],
256 Validity::NonNullable,
257 )
258 }
259 .into_array();
260 let values = unsafe {
261 PrimitiveArray::new_unchecked(buffer![zone_len, last_zone_len], Validity::NonNullable)
262 }
263 .into_array();
264
265 Ok(unsafe { RunEnd::new_unchecked(ends, values, 0, num_zones) }.into_array())
268}
269
270#[cfg(test)]
271mod tests {
272 use std::sync::Arc;
273
274 use vortex_array::IntoArray;
275 use vortex_array::arrays::BoolArray;
276 use vortex_array::arrays::PrimitiveArray;
277 use vortex_array::arrays::StructArray;
278 use vortex_array::assert_arrays_eq;
279 use vortex_array::dtype::DType;
280 use vortex_array::dtype::FieldNames;
281 use vortex_array::dtype::Nullability;
282 use vortex_array::dtype::PType;
283 use vortex_array::expr::Expression;
284 use vortex_array::expr::cast;
285 use vortex_array::expr::gt;
286 use vortex_array::expr::gt_eq;
287 use vortex_array::expr::is_not_null;
288 use vortex_array::expr::is_null;
289 use vortex_array::expr::lit;
290 use vortex_array::expr::lt;
291 use vortex_array::expr::not_eq;
292 use vortex_array::expr::root;
293 use vortex_array::expr::stats::Stat;
294 use vortex_array::stats::all_nan;
295 use vortex_array::stats::all_non_nan;
296 use vortex_array::stats::all_non_null;
297 use vortex_array::stats::all_null;
298 use vortex_array::validity::Validity;
299 use vortex_buffer::buffer;
300
301 use crate::layouts::zoned::zone_map::ZoneMap;
302 use crate::test::SESSION;
303
304 fn falsify(expr: &Expression, dtype: DType) -> Expression {
305 expr.falsify(&dtype, &SESSION).unwrap().unwrap()
306 }
307
308 #[test]
309 fn test_zone_map_prunes() {
310 let zone_map = ZoneMap::try_new(
322 PType::I32.into(),
323 StructArray::from_fields(&[
324 (
325 "max",
326 PrimitiveArray::new(buffer![5i32, 6i32, 7i32], Validity::AllValid).into_array(),
327 ),
328 (
329 "max_is_truncated",
330 BoolArray::from_iter([false, false, false]).into_array(),
331 ),
332 (
333 "min",
334 PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::AllValid).into_array(),
335 ),
336 (
337 "min_is_truncated",
338 BoolArray::from_iter([false, false, false]).into_array(),
339 ),
340 ])
341 .unwrap(),
342 Arc::new([Stat::Max, Stat::Min]),
343 3,
344 10,
345 )
346 .unwrap();
347
348 let expr = gt_eq(root(), lit(6i32));
351 let pruning_expr = falsify(&expr, PType::I32.into());
352 let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap();
353 assert_arrays_eq!(
354 mask.into_array(),
355 BoolArray::from_iter([true, false, false])
356 );
357
358 let expr = gt(root(), lit(5i32));
361 let pruning_expr = falsify(&expr, PType::I32.into());
362 let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap();
363 assert_arrays_eq!(
364 mask.into_array(),
365 BoolArray::from_iter([true, false, false])
366 );
367
368 let expr = lt(root(), lit(2i32));
371 let pruning_expr = falsify(&expr, PType::I32.into());
372 let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap();
373 assert_arrays_eq!(mask.into_array(), BoolArray::from_iter([false, true, true]));
374 }
375
376 #[test]
377 fn row_count_prunes_short_trailing_zone() {
378 let zone_map = ZoneMap::try_new(
379 PType::U64.into(),
380 StructArray::from_fields(&[(
381 "null_count",
382 PrimitiveArray::new(buffer![0u64, 0, 2], Validity::AllValid).into_array(),
383 )])
384 .unwrap(),
385 Arc::new([Stat::NullCount]),
386 4,
387 10,
388 )
389 .unwrap();
390
391 let expr = is_not_null(root());
392 let pruning_expr = falsify(&expr, PType::U64.into());
393
394 let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap();
395 assert_arrays_eq!(
396 mask.into_array(),
397 BoolArray::from_iter([false, false, true])
398 );
399 }
400
401 #[test]
402 fn row_count_substitution_handles_empty_zone_map() {
403 let zone_map = ZoneMap::try_new(
404 PType::U64.into(),
405 StructArray::from_fields(&[(
406 "null_count",
407 PrimitiveArray::new::<u64>(buffer![], Validity::AllValid).into_array(),
408 )])
409 .unwrap(),
410 Arc::new([Stat::NullCount]),
411 4,
412 0,
413 )
414 .unwrap();
415
416 let expr = is_not_null(root());
417 let pruning_expr = falsify(&expr, PType::U64.into());
418
419 let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap();
420 assert_eq!(mask.len(), 0);
421 }
422
423 #[test]
424 fn all_null_stat_fn_lowers_to_null_count_and_row_count() {
425 let zone_map = ZoneMap::try_new(
426 PType::U64.into(),
427 StructArray::from_fields(&[(
428 "null_count",
429 PrimitiveArray::new(buffer![0u64, 4, 2], Validity::AllValid).into_array(),
430 )])
431 .unwrap(),
432 Arc::new([Stat::NullCount]),
433 4,
434 10,
435 )
436 .unwrap();
437
438 let mask = zone_map.prune(&all_null(root()), &SESSION).unwrap();
439 assert_arrays_eq!(mask.into_array(), BoolArray::from_iter([false, true, true]));
440 }
441
442 #[test]
443 fn all_non_null_stat_fn_lowers_to_null_count() {
444 let zone_map = ZoneMap::try_new(
445 PType::U64.into(),
446 StructArray::from_fields(&[(
447 "null_count",
448 PrimitiveArray::new(buffer![0u64, 4, 2], Validity::AllValid).into_array(),
449 )])
450 .unwrap(),
451 Arc::new([Stat::NullCount]),
452 4,
453 10,
454 )
455 .unwrap();
456
457 let mask = zone_map.prune(&all_non_null(root()), &SESSION).unwrap();
458 assert_arrays_eq!(
459 mask.into_array(),
460 BoolArray::from_iter([true, false, false])
461 );
462 }
463
464 #[test]
465 fn non_float_nan_stat_fns_lower_to_constants() {
466 let zone_map = ZoneMap::try_new(
467 PType::I32.into(),
468 StructArray::try_new(FieldNames::empty(), vec![], 2, Validity::NonNullable).unwrap(),
469 Arc::new([]),
470 4,
471 8,
472 )
473 .unwrap();
474
475 let mask = zone_map.prune(&all_nan(root()), &SESSION).unwrap();
476 assert_arrays_eq!(mask.into_array(), BoolArray::from_iter([false, false]));
477
478 let mask = zone_map.prune(&all_non_nan(root()), &SESSION).unwrap();
479 assert_arrays_eq!(mask.into_array(), BoolArray::from_iter([true, true]));
480 }
481
482 #[test]
483 fn unavailable_stat_fn_lowers_to_unknown_mask() {
484 let zone_map = ZoneMap::try_new(
485 PType::U64.into(),
486 StructArray::try_new(FieldNames::empty(), vec![], 3, Validity::NonNullable).unwrap(),
487 Arc::new([]),
488 4,
489 10,
490 )
491 .unwrap();
492
493 let mask = zone_map.prune(&all_non_null(root()), &SESSION).unwrap();
494 assert_arrays_eq!(
495 mask.into_array(),
496 BoolArray::from_iter([false, false, false])
497 );
498
499 let expr = gt(root(), lit(5u64));
500 let pruning_expr = falsify(&expr, PType::U64.into());
501 let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap();
502 assert_arrays_eq!(
503 mask.into_array(),
504 BoolArray::from_iter([false, false, false])
505 );
506 }
507
508 #[test]
509 fn float_min_max_stat_fn_requires_nan_count() {
510 let zone_map = ZoneMap::try_new(
511 PType::F32.into(),
512 StructArray::from_fields(&[
513 (
514 "max",
515 PrimitiveArray::new(buffer![5.0f32, 6.0, 7.0], Validity::AllValid).into_array(),
516 ),
517 (
518 "max_is_truncated",
519 BoolArray::from_iter([false, false, false]).into_array(),
520 ),
521 ])
522 .unwrap(),
523 Arc::new([Stat::Max]),
524 4,
525 12,
526 )
527 .unwrap();
528
529 let expr = gt(root(), lit(5.0f32));
530 let pruning_expr = falsify(&expr, PType::F32.into());
531 let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap();
532 assert_arrays_eq!(
533 mask.into_array(),
534 BoolArray::from_iter([false, false, false])
535 );
536
537 let zone_map = ZoneMap::try_new(
538 PType::F32.into(),
539 StructArray::from_fields(&[
540 (
541 "max",
542 PrimitiveArray::new(buffer![5.0f32, 6.0, 7.0], Validity::AllValid).into_array(),
543 ),
544 (
545 "max_is_truncated",
546 BoolArray::from_iter([false, false, false]).into_array(),
547 ),
548 (
549 "nan_count",
550 PrimitiveArray::new(buffer![0u64, 0, 0], Validity::AllValid).into_array(),
551 ),
552 ])
553 .unwrap(),
554 Arc::new([Stat::Max, Stat::NaNCount]),
555 4,
556 12,
557 )
558 .unwrap();
559
560 let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap();
561 assert_arrays_eq!(
562 mask.into_array(),
563 BoolArray::from_iter([true, false, false])
564 );
565 }
566
567 #[test]
568 fn float_cast_min_max_stat_fn_uses_source_nan_count() {
569 let zone_map = ZoneMap::try_new(
570 PType::F32.into(),
571 StructArray::from_fields(&[
572 (
573 "max",
574 PrimitiveArray::new(buffer![5.0f32, 5.0], Validity::AllValid).into_array(),
575 ),
576 (
577 "max_is_truncated",
578 BoolArray::from_iter([false, false]).into_array(),
579 ),
580 (
581 "min",
582 PrimitiveArray::new(buffer![5.0f32, 5.0], Validity::AllValid).into_array(),
583 ),
584 (
585 "min_is_truncated",
586 BoolArray::from_iter([false, false]).into_array(),
587 ),
588 (
589 "nan_count",
590 PrimitiveArray::new(buffer![1u64, 0], Validity::AllValid).into_array(),
591 ),
592 ])
593 .unwrap(),
594 Arc::new([Stat::Max, Stat::Min, Stat::NaNCount]),
595 4,
596 8,
597 )
598 .unwrap();
599
600 let cast_dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
601 let expr = not_eq(cast(root(), cast_dtype), lit(5i32));
602 let pruning_expr = falsify(&expr, PType::F32.into());
603
604 let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap();
605 assert_arrays_eq!(mask.into_array(), BoolArray::from_iter([false, true]));
606 }
607
608 #[test]
609 fn unsupported_aggregate_input_dtype_lowers_to_unknown() {
610 let elem_dtype = Arc::new(DType::Primitive(PType::I32, Nullability::Nullable));
617 let column_dtype = DType::FixedSizeList(elem_dtype, 1, Nullability::Nullable);
618
619 let zone_map = ZoneMap::try_new(
620 column_dtype,
621 StructArray::try_new(FieldNames::empty(), vec![], 3, Validity::NonNullable).unwrap(),
622 Arc::new([]),
623 4,
624 10,
625 )
626 .unwrap();
627
628 let max_fn = Stat::Max
629 .aggregate_fn()
630 .expect("max should have an aggregate function");
631 let predicate = is_null(vortex_array::stats::stat(root(), max_fn));
632
633 let mask = zone_map.prune(&predicate, &SESSION).unwrap();
636 assert_arrays_eq!(mask.into_array(), BoolArray::from_iter([true, true, true]));
637 }
638
639 #[test]
640 fn row_count_prunes_all_null_uniform_zones() {
641 let zone_map = ZoneMap::try_new(
642 PType::U64.into(),
643 StructArray::from_fields(&[(
644 "null_count",
645 PrimitiveArray::new(buffer![0u64, 4, 0], Validity::AllValid).into_array(),
646 )])
647 .unwrap(),
648 Arc::new([Stat::NullCount]),
649 4,
650 12,
651 )
652 .unwrap();
653
654 let expr = is_not_null(root());
655 let pruning_expr = falsify(&expr, PType::U64.into());
656
657 let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap();
659 assert_arrays_eq!(
660 mask.into_array(),
661 BoolArray::from_iter([false, true, false])
662 );
663 }
664}