Skip to main content

vortex_layout/layouts/zoned/
mod.rs

1//! Zoned layouts wrap a data layout with an auxiliary per-zone statistics layout.
2//!
3//! The zoned layout tree has exactly two children:
4//! - a transparent `data` child containing the underlying column data
5//! - an auxiliary `zones` child containing one row of aggregate statistics per zone
6//!
7//! Metadata stores the logical zone length in rows plus the aggregate functions present in the
8//! auxiliary table. During scans, pruning first evaluates a falsification predicate against the
9//! `zones` child and only forwards surviving rows to the underlying `data` child.
10
11// SPDX-License-Identifier: Apache-2.0
12// SPDX-FileCopyrightText: Copyright the Vortex contributors
13
14mod builder;
15mod pruning;
16mod reader;
17mod schema;
18pub mod writer;
19pub mod zone_map;
20
21use std::num::NonZeroUsize;
22use std::sync::Arc;
23
24pub(crate) use builder::AggregateStatsAccumulator;
25pub(crate) use builder::aggregate_partials;
26use prost::Message;
27pub use schema::MAX_IS_TRUNCATED;
28pub use schema::MIN_IS_TRUNCATED;
29use vortex_array::DeserializeMetadata;
30use vortex_array::SerializeMetadata;
31use vortex_array::aggregate_fn::AggregateFnRef;
32use vortex_array::dtype::DType;
33use vortex_array::dtype::TryFromBytes;
34use vortex_array::expr::stats::Stat;
35use vortex_array::stats::as_stat_bitset_bytes;
36use vortex_array::stats::stats_from_bitset_bytes;
37use vortex_error::VortexExpect;
38use vortex_error::VortexResult;
39use vortex_error::vortex_bail;
40use vortex_error::vortex_ensure;
41use vortex_error::vortex_ensure_eq;
42use vortex_error::vortex_panic;
43use vortex_session::VortexSession;
44use vortex_session::registry::CachedId;
45
46use crate::Layout;
47use crate::LayoutChildType;
48use crate::LayoutDeserializeArgs;
49use crate::LayoutId;
50use crate::LayoutParts;
51use crate::LayoutReaderContext;
52use crate::LayoutReaderRef;
53use crate::LayoutRef;
54use crate::VTable;
55use crate::children::OwnedLayoutChildren;
56use crate::layouts::zoned::reader::ZonedReader;
57use crate::layouts::zoned::schema::AggregateSpecProto;
58use crate::layouts::zoned::schema::aggregate_specs_from_fns;
59use crate::layouts::zoned::schema::aggregate_stats_table_dtype;
60use crate::layouts::zoned::schema::legacy_stats_table_dtype;
61use crate::layouts::zoned::schema::try_aggregate_fns_from_specs;
62use crate::segments::SegmentSource;
63
64/// Zoned layout vtable.
65#[derive(Clone, Debug)]
66pub struct Zoned;
67
68/// Legacy `vortex.stats` layout vtable.
69#[derive(Clone, Debug)]
70pub struct LegacyStats;
71
72/// Backwards-compatible plugin names.
73pub use LegacyStats as LegacyStatsLayoutEncoding;
74pub use Zoned as ZonedLayoutEncoding;
75
76/// Zoned-layout-specific data.
77#[derive(Clone, Debug)]
78pub struct ZonedData {
79    zone_len: usize,
80    zone_map_schema: ZoneMapSchema,
81    stats_table_dtype: DType,
82}
83
84/// A layout annotating a data child with per-zone statistics.
85pub type ZonedLayout = Layout<Zoned>;
86/// A legacy `vortex.stats` layout using the same runtime data.
87pub type LegacyStatsLayout = Layout<LegacyStats>;
88
89impl VTable for Zoned {
90    type LayoutData = ZonedData;
91    type Metadata = ZonedMetadata;
92
93    fn id(&self) -> LayoutId {
94        static ID: CachedId = CachedId::new("vortex.zoned");
95        *ID
96    }
97
98    fn metadata(layout: &Layout<Self>) -> Self::Metadata {
99        ZonedMetadata {
100            zone_len: u32::try_from(layout.zone_len).vortex_expect("Invalid zone length"),
101            aggregate_specs: match &layout.zone_map_schema {
102                ZoneMapSchema::AggregateFns(aggregate_fns) => {
103                    aggregate_specs_from_fns(aggregate_fns).vortex_expect(
104                        "aggregate functions should be validated as serializable during build",
105                    )
106                }
107                ZoneMapSchema::LegacyStats(_) => {
108                    vortex_panic!("Cannot serialize legacy stats schema as vortex.zoned")
109                }
110            },
111        }
112    }
113
114    fn deserialize(
115        &self,
116        args: &LayoutDeserializeArgs<'_>,
117        metadata: &ZonedMetadata,
118    ) -> VortexResult<Self::LayoutData> {
119        vortex_ensure_eq!(
120            args.children.nchildren(),
121            2,
122            "ZonedLayout expects exactly 2 children (data, zones)"
123        );
124        vortex_ensure_eq!(args.children.child_row_count(0), args.row_count);
125        let Some(aggregate_fns) =
126            try_aggregate_fns_from_specs(&metadata.aggregate_specs, args.session)?
127        else {
128            args.children.child(0, args.dtype)?;
129            return Ok(ZonedData {
130                zone_len: 0,
131                zone_map_schema: ZoneMapSchema::AggregateFns(Arc::new([])),
132                stats_table_dtype: aggregate_stats_table_dtype(args.dtype, &[]),
133            });
134        };
135        aggregate_specs_from_fns(&aggregate_fns)?;
136        let stats_table_dtype = aggregate_stats_table_dtype(args.dtype, &aggregate_fns);
137        args.children.child(0, args.dtype)?;
138        args.children.child(1, &stats_table_dtype)?;
139        Ok(ZonedData {
140            zone_len: metadata.zone_len as usize,
141            zone_map_schema: ZoneMapSchema::AggregateFns(aggregate_fns),
142            stats_table_dtype,
143        })
144    }
145
146    fn child_dtype(layout: &Layout<Self>, idx: usize) -> VortexResult<DType> {
147        match idx {
148            0 => Ok(layout.dtype().clone()),
149            1 => Ok(layout.stats_table_dtype.clone()),
150            _ => vortex_bail!("Invalid child index: {idx}"),
151        }
152    }
153
154    fn child_type(_layout: &Layout<Self>, idx: usize) -> LayoutChildType {
155        match idx {
156            0 => LayoutChildType::Transparent("data".into()),
157            1 => LayoutChildType::Auxiliary("zones".into()),
158            _ => vortex_panic!("Invalid child index: {}", idx),
159        }
160    }
161
162    fn new_reader(
163        layout: &Layout<Self>,
164        name: Arc<str>,
165        segment_source: Arc<dyn SegmentSource>,
166        session: &VortexSession,
167        ctx: &LayoutReaderContext,
168    ) -> VortexResult<LayoutReaderRef> {
169        layout.zoned_reader(name, segment_source, session, ctx)
170    }
171}
172
173impl VTable for LegacyStats {
174    type LayoutData = ZonedData;
175    type Metadata = LegacyStatsMetadata;
176
177    fn id(&self) -> LayoutId {
178        static ID: CachedId = CachedId::new("vortex.stats");
179        *ID
180    }
181
182    fn metadata(layout: &Layout<Self>) -> Self::Metadata {
183        LegacyStatsMetadata {
184            zone_len: u32::try_from(layout.zone_len).vortex_expect("Invalid zone length"),
185            zone_map_schema: layout.zone_map_schema.clone(),
186        }
187    }
188
189    fn deserialize(
190        &self,
191        args: &LayoutDeserializeArgs<'_>,
192        metadata: &LegacyStatsMetadata,
193    ) -> VortexResult<Self::LayoutData> {
194        vortex_ensure_eq!(
195            args.children.nchildren(),
196            2,
197            "LegacyStatsLayout expects exactly 2 children (data, zones)"
198        );
199        let stats_table_dtype = match &metadata.zone_map_schema {
200            ZoneMapSchema::LegacyStats(stats) => legacy_stats_table_dtype(args.dtype, stats),
201            ZoneMapSchema::AggregateFns(aggregate_fns) => {
202                aggregate_stats_table_dtype(args.dtype, aggregate_fns)
203            }
204        };
205        args.children.child(0, args.dtype)?;
206        args.children.child(1, &stats_table_dtype)?;
207        Ok(ZonedData {
208            zone_len: metadata.zone_len as usize,
209            zone_map_schema: metadata.zone_map_schema.clone(),
210            stats_table_dtype,
211        })
212    }
213
214    fn child_dtype(layout: &Layout<Self>, idx: usize) -> VortexResult<DType> {
215        match idx {
216            0 => Ok(layout.dtype().clone()),
217            1 => Ok(layout.stats_table_dtype.clone()),
218            _ => vortex_bail!("Invalid child index: {idx}"),
219        }
220    }
221
222    fn child_type(_layout: &Layout<Self>, idx: usize) -> LayoutChildType {
223        match idx {
224            0 => LayoutChildType::Transparent("data".into()),
225            1 => LayoutChildType::Auxiliary("zones".into()),
226            _ => vortex_panic!("Invalid child index: {idx}"),
227        }
228    }
229
230    fn new_reader(
231        layout: &Layout<Self>,
232        name: Arc<str>,
233        segment_source: Arc<dyn SegmentSource>,
234        session: &VortexSession,
235        ctx: &LayoutReaderContext,
236    ) -> VortexResult<LayoutReaderRef> {
237        layout.zoned_reader(name, segment_source, session, ctx)
238    }
239}
240
241impl LegacyStatsLayout {
242    fn zoned_reader(
243        &self,
244        name: Arc<str>,
245        segment_source: Arc<dyn SegmentSource>,
246        session: &VortexSession,
247        ctx: &LayoutReaderContext,
248    ) -> VortexResult<LayoutReaderRef> {
249        if self.zone_len == 0 {
250            return self
251                .slot(0)?
252                .vortex_expect("ZonedLayout always has a data child")
253                .new_reader(name, segment_source, session, ctx);
254        }
255
256        Ok(Arc::new(ZonedReader::try_new(
257            self.clone(),
258            self.nzones(),
259            name,
260            segment_source,
261            session.clone(),
262            ctx.clone(),
263        )?))
264    }
265
266    pub fn nzones(&self) -> usize {
267        usize::try_from(self.children().child_row_count(1))
268            .vortex_expect("Invalid number of zones, cannot handle more than usize zones")
269    }
270
271    /// Returns display names for the zone-map aggregates stored by this layout.
272    pub fn present_aggregates(&self) -> Arc<[String]> {
273        present_aggregates(&self.zone_map_schema)
274    }
275}
276
277#[derive(Clone, Debug, PartialEq, Eq)]
278pub(crate) enum ZoneMapSchema {
279    LegacyStats(Arc<[Stat]>),
280    AggregateFns(Arc<[AggregateFnRef]>),
281}
282
283impl ZonedLayout {
284    /// Create a zoned layout from a data child, a zone-map child, a zone length, and the aggregate
285    /// functions stored in the zone map.
286    pub fn try_new(
287        data: LayoutRef,
288        zones: LayoutRef,
289        zone_len: NonZeroUsize,
290        aggregate_fns: Arc<[AggregateFnRef]>,
291    ) -> VortexResult<Self> {
292        let expected_dtype = aggregate_stats_table_dtype(data.dtype(), &aggregate_fns);
293        if zones.dtype() != &expected_dtype {
294            vortex_bail!("Invalid zone map layout: zones dtype does not match expected dtype");
295        }
296
297        // Verify that every aggregate is serializable.
298        aggregate_specs_from_fns(&aggregate_fns)?;
299
300        let dtype = data.dtype().clone();
301        let row_count = data.row_count();
302        Ok(LayoutParts::new(
303            Zoned,
304            dtype,
305            row_count,
306            Vec::new(),
307            OwnedLayoutChildren::layout_children(vec![data, zones]),
308            ZonedData {
309                zone_len: zone_len.get(),
310                zone_map_schema: ZoneMapSchema::AggregateFns(aggregate_fns),
311                stats_table_dtype: expected_dtype,
312            },
313        )
314        .into_typed())
315    }
316
317    pub fn zone_len(&self) -> usize {
318        self.zone_len
319    }
320
321    /// Returns display names for the zone-map aggregates stored by this layout.
322    pub fn present_aggregates(&self) -> Arc<[String]> {
323        present_aggregates(&self.zone_map_schema)
324    }
325
326    /// Builds a reader for a zoned layout, bypassing [`ZonedReader`] when the zone map is empty
327    /// (`zone_len == 0`) and reading the data child directly, since there is nothing to prune with.
328    /// This covers both legacy zero-length zones and layouts whose aggregates the session cannot
329    /// reconstruct.
330    fn zoned_reader(
331        &self,
332        name: Arc<str>,
333        segment_source: Arc<dyn SegmentSource>,
334        session: &VortexSession,
335        ctx: &LayoutReaderContext,
336    ) -> VortexResult<LayoutReaderRef> {
337        if self.zone_len == 0 {
338            return self
339                .slot(0)?
340                .vortex_expect("ZonedLayout always has a data child")
341                .new_reader(name, segment_source, session, ctx);
342        }
343
344        Ok(Arc::new(ZonedReader::try_new(
345            self.clone(),
346            self.nzones(),
347            name,
348            segment_source,
349            session.clone(),
350            ctx.clone(),
351        )?))
352    }
353
354    pub fn nzones(&self) -> usize {
355        usize::try_from(self.children().child_row_count(1))
356            .vortex_expect("Invalid number of zones, cannot handle more than usize zones")
357    }
358}
359
360impl ZonedData {
361    fn aggregate_fns(&self) -> Arc<[AggregateFnRef]> {
362        match &self.zone_map_schema {
363            ZoneMapSchema::LegacyStats(stats) => stats
364                .iter()
365                .filter_map(Stat::aggregate_fn)
366                .collect::<Vec<_>>()
367                .into(),
368            ZoneMapSchema::AggregateFns(aggregate_fns) => Arc::clone(aggregate_fns),
369        }
370    }
371}
372
373fn present_aggregates(schema: &ZoneMapSchema) -> Arc<[String]> {
374    match schema {
375        ZoneMapSchema::LegacyStats(stats) => stats
376            .iter()
377            .filter_map(Stat::aggregate_fn)
378            .map(|aggregate_fn| aggregate_fn.to_string())
379            .collect::<Vec<_>>()
380            .into(),
381        ZoneMapSchema::AggregateFns(aggregate_fns) => aggregate_fns
382            .iter()
383            .map(ToString::to_string)
384            .collect::<Vec<_>>()
385            .into(),
386    }
387}
388
389/// Serialized zoned-layout metadata.
390///
391/// `zone_len` is the logical row length of each zone. `aggregate_specs` is the ordered list of
392/// aggregate functions stored in the auxiliary stats-table child.
393#[derive(Debug, PartialEq, Eq, Clone)]
394pub struct ZonedMetadata {
395    pub(super) zone_len: u32,
396    pub(super) aggregate_specs: Arc<[AggregateSpecProto]>,
397}
398
399/// Serialized metadata for legacy `vortex.stats` layouts.
400#[derive(Debug, PartialEq, Eq, Clone)]
401pub struct LegacyStatsMetadata {
402    pub(super) zone_len: u32,
403    pub(crate) zone_map_schema: ZoneMapSchema,
404}
405
406const ZONED_METADATA_PROTO_VERSION: u8 = 1;
407
408#[derive(Clone, PartialEq, Message)]
409struct ZonedMetadataProto {
410    #[prost(uint32, tag = "1")]
411    zone_len: u32,
412    #[prost(message, repeated, tag = "2")]
413    aggregate_specs: Vec<AggregateSpecProto>,
414}
415
416impl DeserializeMetadata for ZonedMetadata {
417    type Output = Self;
418
419    fn deserialize(metadata: &[u8]) -> VortexResult<Self::Output> {
420        let Some((&version, proto_bytes)) = metadata.split_first() else {
421            vortex_bail!("Zoned metadata missing protobuf version");
422        };
423
424        vortex_ensure!(
425            version == ZONED_METADATA_PROTO_VERSION,
426            "Unsupported zoned metadata version: {}",
427            version
428        );
429        vortex_ensure!(!proto_bytes.is_empty(), "Zoned metadata missing protobuf");
430
431        let proto = ZonedMetadataProto::decode(proto_bytes)?;
432        Ok(Self {
433            zone_len: proto.zone_len,
434            aggregate_specs: proto.aggregate_specs.into(),
435        })
436    }
437}
438
439impl SerializeMetadata for ZonedMetadata {
440    fn serialize(self) -> Vec<u8> {
441        let proto = ZonedMetadataProto {
442            zone_len: self.zone_len,
443            aggregate_specs: self.aggregate_specs.to_vec(),
444        };
445        let mut metadata = vec![ZONED_METADATA_PROTO_VERSION];
446        metadata.extend(proto.encode_to_vec());
447        metadata
448    }
449}
450
451impl DeserializeMetadata for LegacyStatsMetadata {
452    type Output = Self;
453
454    fn deserialize(metadata: &[u8]) -> VortexResult<Self::Output> {
455        vortex_ensure!(
456            metadata.len() >= 4,
457            "Legacy zoned metadata must contain at least 4 bytes for zone length, got {}",
458            metadata.len()
459        );
460
461        // Backward compat: older files may encode `zone_len == 0`. Preserve the raw metadata on
462        // read and let the reader disable zoned pruning for those layouts instead of rejecting
463        // deserialization outright.
464        let zone_len = u32::try_from_le_bytes(&metadata[0..4])?;
465        let present_stats: Arc<[Stat]> = stats_from_bitset_bytes(&metadata[4..]).into();
466
467        Ok(Self {
468            zone_len,
469            zone_map_schema: ZoneMapSchema::LegacyStats(present_stats),
470        })
471    }
472}
473
474impl SerializeMetadata for LegacyStatsMetadata {
475    fn serialize(self) -> Vec<u8> {
476        match self.zone_map_schema {
477            ZoneMapSchema::LegacyStats(stats) => {
478                let mut metadata = self.zone_len.to_le_bytes().to_vec();
479                metadata.extend(as_stat_bitset_bytes(&stats));
480                metadata
481            }
482            ZoneMapSchema::AggregateFns(_) => {
483                vortex_panic!("Cannot serialize aggregate specs as legacy stats metadata")
484            }
485        }
486    }
487}
488
489#[cfg(test)]
490mod tests {
491    use std::panic;
492
493    use rstest::rstest;
494    use vortex_array::aggregate_fn::AggregateFnRef;
495    use vortex_array::aggregate_fn::AggregateFnVTableExt;
496    use vortex_array::aggregate_fn::NumericalAggregateOpts;
497    use vortex_array::aggregate_fn::fns::bounded_max::BoundedMax;
498    use vortex_array::aggregate_fn::fns::bounded_max::BoundedMaxOptions;
499    use vortex_array::aggregate_fn::fns::max::Max;
500    use vortex_array::aggregate_fn::fns::min::Min;
501    use vortex_array::aggregate_fn::session::AggregateFnSession;
502    use vortex_array::dtype::DType;
503    use vortex_array::dtype::Nullability;
504    use vortex_array::dtype::PType;
505    use vortex_array::stats::as_stat_bitset_bytes;
506    use vortex_session::VortexSession;
507    use vortex_session::registry::ReadContext;
508
509    use super::*;
510    use crate::LayoutBuildContext;
511    use crate::children::OwnedLayoutChildren;
512    use crate::layouts::flat::FlatLayout;
513    use crate::segments::SegmentId;
514
515    fn aggregate_spec(aggregate_fn: AggregateFnRef) -> AggregateSpecProto {
516        AggregateSpecProto::try_from_aggregate_fn(&aggregate_fn).unwrap()
517    }
518
519    #[rstest]
520    #[case(ZonedMetadata {
521            zone_len: u32::MAX,
522            aggregate_specs: Arc::new([]),
523        })]
524    #[case::min_max(ZonedMetadata {
525            zone_len: 314,
526            aggregate_specs: Arc::new([
527                aggregate_spec(Max.bind(NumericalAggregateOpts::skip_nans())),
528                aggregate_spec(Min.bind(NumericalAggregateOpts::skip_nans())),
529            ]),
530        })]
531    fn test_metadata_serialization(#[case] metadata: ZonedMetadata) {
532        let serialized = metadata.clone().serialize();
533        assert_eq!(serialized[0], ZONED_METADATA_PROTO_VERSION);
534        let deserialized = ZonedMetadata::deserialize(&serialized).unwrap();
535        assert_eq!(deserialized, metadata);
536    }
537
538    #[test]
539    fn test_metadata_serialization_preserves_aggregate_options() -> VortexResult<()> {
540        let aggregate_fn = BoundedMax.bind(BoundedMaxOptions {
541            // SAFETY: 128 is non-zero.
542            max_bytes: unsafe { NonZeroUsize::new_unchecked(128) },
543        });
544        let metadata = ZonedMetadata {
545            zone_len: 314,
546            aggregate_specs: Arc::new([AggregateSpecProto::try_from_aggregate_fn(&aggregate_fn)?]),
547        };
548
549        let deserialized = ZonedMetadata::deserialize(&metadata.serialize())?;
550        let session = VortexSession::empty().with::<AggregateFnSession>();
551        let aggregate_fns = try_aggregate_fns_from_specs(&deserialized.aggregate_specs, &session)?
552            .expect("known aggregates resolve");
553
554        assert_eq!(aggregate_fns.as_ref(), std::slice::from_ref(&aggregate_fn));
555        Ok(())
556    }
557
558    #[test]
559    fn test_deserialize_legacy_stat_bitset_as_legacy_stats() {
560        let mut serialized = u32::MAX.to_le_bytes().to_vec();
561        serialized.extend(as_stat_bitset_bytes(&[
562            Stat::IsStrictSorted,
563            Stat::IsSorted,
564            Stat::Max,
565        ]));
566        let deserialized = LegacyStatsMetadata::deserialize(&serialized).unwrap();
567        let ZoneMapSchema::LegacyStats(legacy_stats) = deserialized.zone_map_schema else {
568            panic!("legacy bitset metadata should deserialize as legacy stats");
569        };
570
571        assert!(legacy_stats.is_sorted());
572        assert_eq!(
573            legacy_stats.as_ref(),
574            &[Stat::IsSorted, Stat::IsStrictSorted, Stat::Max]
575        );
576    }
577
578    #[rstest]
579    #[case::empty(vec![])]
580    #[case::unsupported_version(vec![0])]
581    #[case::missing_proto(vec![ZONED_METADATA_PROTO_VERSION])]
582    #[case::malformed_proto(vec![ZONED_METADATA_PROTO_VERSION, 0])]
583    fn test_deserialize_short_metadata_errors(#[case] metadata: Vec<u8>) {
584        assert!(ZonedMetadata::deserialize(&metadata).is_err());
585    }
586
587    #[test]
588    fn test_deserialize_short_metadata_returns_error_not_panic() {
589        let result = panic::catch_unwind(|| ZonedMetadata::deserialize(&[]));
590        assert!(
591            result.is_ok(),
592            "deserialize should return an error, not panic"
593        );
594        assert!(result.unwrap().is_err());
595    }
596
597    #[test]
598    fn test_deserialize_zero_zone_len_is_allowed_for_backcompat() {
599        let metadata = 0u32.to_le_bytes();
600        let deserialized = LegacyStatsMetadata::deserialize(&metadata).unwrap();
601        assert_eq!(deserialized.zone_len, 0);
602        let ZoneMapSchema::LegacyStats(legacy_stats) = deserialized.zone_map_schema else {
603            panic!("legacy bitset metadata should deserialize as legacy stats");
604        };
605        assert!(legacy_stats.is_empty());
606    }
607
608    #[test]
609    fn test_build_allows_zero_zone_len_for_backcompat() -> VortexResult<()> {
610        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
611        let read_ctx = ReadContext::new([]);
612        let children = OwnedLayoutChildren::layout_children(vec![
613            FlatLayout::new(0, dtype.clone(), SegmentId::from(0), read_ctx.clone()).into_layout(),
614            FlatLayout::new(
615                0,
616                legacy_stats_table_dtype(&dtype, &[]),
617                SegmentId::from(1),
618                read_ctx,
619            )
620            .into_layout(),
621        ]);
622        let session = vortex_array::array_session();
623        let build_read_ctx = ReadContext::new([]);
624        let build_ctx = LayoutBuildContext {
625            session: &session,
626            array_read_ctx: &build_read_ctx,
627        };
628
629        let layout = <LegacyStats as VTable>::build(
630            &LegacyStatsLayoutEncoding,
631            &dtype,
632            0,
633            &LegacyStatsMetadata {
634                zone_len: 0,
635                zone_map_schema: ZoneMapSchema::LegacyStats(Arc::new([])),
636            },
637            vec![],
638            children.as_ref(),
639            &build_ctx,
640        )?;
641
642        assert_eq!(layout.zone_len, 0);
643        Ok(())
644    }
645
646    #[test]
647    fn test_build_rejects_invalid_child_count() {
648        let metadata = ZonedMetadata {
649            zone_len: 3,
650            aggregate_specs: Arc::new([]),
651        };
652        let children = OwnedLayoutChildren::layout_children(vec![]);
653        let session = vortex_array::array_session();
654        let build_read_ctx = ReadContext::new([]);
655        let build_ctx = LayoutBuildContext {
656            session: &session,
657            array_read_ctx: &build_read_ctx,
658        };
659
660        let result = <Zoned as VTable>::build(
661            &ZonedLayoutEncoding,
662            &DType::Primitive(PType::I32, Nullability::NonNullable),
663            0,
664            &metadata,
665            vec![],
666            children.as_ref(),
667            &build_ctx,
668        );
669
670        assert!(result.is_err());
671    }
672
673    #[test]
674    fn test_build_unknown_aggregate_disables_pruning_when_allowed() -> VortexResult<()> {
675        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
676        let read_ctx = ReadContext::new([]);
677        let children = OwnedLayoutChildren::layout_children(vec![
678            FlatLayout::new(0, dtype.clone(), SegmentId::from(0), read_ctx.clone()).into_layout(),
679            FlatLayout::new(0, dtype.clone(), SegmentId::from(1), read_ctx).into_layout(),
680        ]);
681        let session = vortex_array::array_session();
682        session.allow_unknown();
683        let build_read_ctx = ReadContext::new([]);
684        let build_ctx = LayoutBuildContext {
685            session: &session,
686            array_read_ctx: &build_read_ctx,
687        };
688
689        let metadata = ZonedMetadata {
690            zone_len: 8,
691            aggregate_specs: Arc::new([AggregateSpecProto::new_unknown("vortex.test.unknown")]),
692        };
693
694        let layout = <Zoned as VTable>::build(
695            &ZonedLayoutEncoding,
696            &dtype,
697            0,
698            &metadata,
699            vec![],
700            children.as_ref(),
701            &build_ctx,
702        )?;
703
704        // An unknown aggregate disables zoned pruning (`zone_len == 0`) while leaving the data
705        // child readable, so the index is ignored rather than turned into a hard read error.
706        assert_eq!(layout.zone_len, 0);
707        Ok(())
708    }
709
710    #[test]
711    fn test_build_unknown_aggregate_errors_without_allow_unknown() {
712        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
713        let read_ctx = ReadContext::new([]);
714        let children = OwnedLayoutChildren::layout_children(vec![
715            FlatLayout::new(0, dtype.clone(), SegmentId::from(0), read_ctx.clone()).into_layout(),
716            FlatLayout::new(0, dtype.clone(), SegmentId::from(1), read_ctx).into_layout(),
717        ]);
718        let session = vortex_array::array_session();
719        let build_read_ctx = ReadContext::new([]);
720        let build_ctx = LayoutBuildContext {
721            session: &session,
722            array_read_ctx: &build_read_ctx,
723        };
724
725        let metadata = ZonedMetadata {
726            zone_len: 8,
727            aggregate_specs: Arc::new([AggregateSpecProto::new_unknown("vortex.test.unknown")]),
728        };
729
730        let result = <Zoned as VTable>::build(
731            &ZonedLayoutEncoding,
732            &dtype,
733            0,
734            &metadata,
735            vec![],
736            children.as_ref(),
737            &build_ctx,
738        );
739
740        assert!(result.is_err());
741    }
742}