vortex_layout/plan/physical/plans/
zoned.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::borrow::Cow;
5use std::fmt;
6use std::ops::Range;
7use std::sync::Arc;
8
9use parking_lot::RwLock;
10use vortex_array::ArrayRef;
11use vortex_array::EmptyMetadata;
12use vortex_array::IntoArray;
13use vortex_array::VortexSessionExecute;
14use vortex_array::aggregate_fn::AggregateFnRef;
15use vortex_array::arrays::BoolArray;
16use vortex_array::arrays::ConstantArray;
17use vortex_array::arrays::StructArray;
18use vortex_array::dtype::DType;
19use vortex_array::expr::BoundExpression;
20use vortex_array::expr::traversal::NodeExt;
21use vortex_array::expr::traversal::Transformed;
22use vortex_array::expr::traversal::TraversalOrder;
23use vortex_array::scalar::Scalar;
24use vortex_array::scalar_fn::fns::dynamic::DynamicExprUpdates;
25use vortex_array::scalar_fn::fns::stat::StatFn;
26use vortex_array::validity::Validity;
27use vortex_buffer::BitBufferMut;
28use vortex_error::VortexExpect;
29use vortex_error::VortexResult;
30use vortex_error::vortex_bail;
31use vortex_error::vortex_err;
32use vortex_mask::Mask;
33use vortex_session::VortexSession;
34use vortex_session::registry::CachedId;
35
36use crate::layouts::zoned::zone_map::ZoneMap;
37use crate::plan::physical::Eval;
38use crate::plan::physical::Plan;
39use crate::plan::physical::PlanChildren;
40use crate::plan::physical::PlanId;
41use crate::plan::physical::PlanParts;
42use crate::plan::physical::PlanRef;
43use crate::plan::physical::PlanVTable;
44use crate::plan::physical::check_child_count;
45use crate::plan::physical::operators::OnceSource;
46use crate::plan::physical::operators::UNBOUNDED;
47use crate::plan::physical::operators::WholeSource;
48use crate::plan::physical::optimizer::PlanParentReduceRule;
49use crate::plan::physical::pipeline::Chain;
50use crate::plan::physical::pipeline::Compiler;
51use crate::plan::physical::pipeline::Reach;
52use crate::segments::SegmentId;
53
54const DATA: usize = 0;
55const ZONES: usize = 1;
56
57#[derive(Clone, Debug)]
58struct ZonedPruningState {
59    expression: BoundExpression,
60    /// The dtype of the column the zones summarise, which the expression is bound to.
61    column_dtype: DType,
62    /// Which zones the expression proves can be pruned, shared by every execution of this plan.
63    pruned_zones: Arc<PrunedZones>,
64}
65
66/// Which zones a pruning expression proves can be pruned, shared by every execution of a plan.
67///
68/// The zone table is read once and kept. An expression with dynamic comparisons, such as a top-N
69/// bound an engine tightens during the scan, is proven again over the kept zones whenever one of
70/// them changes, as the default scan's zoned reader does.
71pub(crate) struct PrunedZones {
72    updates: Option<DynamicExprUpdates>,
73    /// The zones, and the pruned zones as of an update version.
74    state: RwLock<Option<(ZoneMap, u64, Mask)>>,
75}
76
77impl PrunedZones {
78    fn new(expression: &BoundExpression) -> Self {
79        Self {
80            updates: DynamicExprUpdates::new(expression),
81            state: RwLock::new(None),
82        }
83    }
84
85    fn version(&self) -> u64 {
86        self.updates.as_ref().map_or(0, DynamicExprUpdates::version)
87    }
88
89    /// The pruned zones, proven again if a dynamic comparison changed, or `None` before any
90    /// execution read the zones.
91    pub(crate) fn get(
92        &self,
93        expression: &BoundExpression,
94        session: &VortexSession,
95    ) -> VortexResult<Option<Mask>> {
96        let version = self.version();
97        {
98            let state = self.state.read();
99            match &*state {
100                None => return Ok(None),
101                Some((_, proven, pruned)) if *proven >= version => return Ok(Some(pruned.clone())),
102                Some(_) => {}
103            }
104        }
105        let mut state = self.state.write();
106        let Some((zones, proven, pruned)) = state.as_mut() else {
107            return Ok(None);
108        };
109        if *proven < version {
110            *pruned = zones.prune(expression, session)?;
111            *proven = version;
112        }
113        Ok(Some(pruned.clone()))
114    }
115
116    /// Proves `expression` over `zones` and keeps both, unless another execution did first.
117    pub(crate) fn init(
118        &self,
119        zones: ZoneMap,
120        expression: &BoundExpression,
121        session: &VortexSession,
122    ) -> VortexResult<Mask> {
123        let version = self.version();
124        let pruned = zones.prune(expression, session)?;
125        let mut state = self.state.write();
126        if state.is_none() {
127            *state = Some((zones, version, pruned.clone()));
128        }
129        Ok(pruned)
130    }
131}
132
133impl fmt::Debug for PrunedZones {
134    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
135        f.debug_struct("PrunedZones")
136            .field("dynamic", &self.updates.is_some())
137            .field("read", &self.state.read().is_some())
138            .finish()
139    }
140}
141
142/// Zoned-plan-specific data.
143#[derive(Clone, Debug)]
144pub struct ZonedData {
145    zone_len: u64,
146    /// The aggregate functions whose results the zone table stores, in field order.
147    aggregate_fns: Arc<[AggregateFnRef]>,
148    pruning: Option<ZonedPruningState>,
149}
150
151/// Reads data alongside the zone statistics summarising it.
152///
153/// Its zone table has the aggregate schema of `vortex.zoned` layouts; legacy `vortex.stats`
154/// layouts reach it through [`LegacyStats`](crate::plan::physical::LegacyStats), which converts
155/// their tables. An expression containing abstract statistic functions can rewrite it into a
156/// pruning plan that retains only the zone-statistics child.
157#[derive(Clone, Debug)]
158pub struct Zoned;
159
160/// A plan that pairs data with its zone statistics or represents a zone-backed pruning proof.
161pub type ZonedPlan = Plan<Zoned>;
162
163impl ZonedPlan {
164    pub(crate) fn from_children(
165        dtype: DType,
166        row_count: u64,
167        children: PlanChildren,
168        zone_len: u64,
169        aggregate_fns: Arc<[AggregateFnRef]>,
170    ) -> Self {
171        PlanParts {
172            vtable: Zoned,
173            dtype,
174            row_count,
175            children,
176            data: ZonedData {
177                zone_len,
178                aggregate_fns,
179                pruning: None,
180            },
181        }
182        .into_typed()
183    }
184
185    /// Creates a zoned plan over `data` summarised by `zones` of `zone_len` rows each, whose
186    /// fields hold the results of `aggregate_fns`.
187    pub fn new(
188        data: PlanRef,
189        zones: PlanRef,
190        zone_len: u64,
191        aggregate_fns: Arc<[AggregateFnRef]>,
192    ) -> Self {
193        let dtype = data.dtype().clone();
194        let row_count = data.row_count();
195        Self::from_children(
196            dtype,
197            row_count,
198            vec![data, zones].into(),
199            zone_len,
200            aggregate_fns,
201        )
202    }
203
204    /// Returns the number of rows each zone summarises.
205    pub fn zone_len(&self) -> u64 {
206        self.data().zone_len
207    }
208
209    /// Returns the aggregate functions whose results the zone table stores.
210    pub fn aggregate_fns(&self) -> &Arc<[AggregateFnRef]> {
211        &self.data().aggregate_fns
212    }
213
214    /// Returns the dtype of the column the zones summarise, when this is a pruning plan.
215    pub fn pruning_column_dtype(&self) -> Option<&DType> {
216        self.data()
217            .pruning
218            .as_ref()
219            .map(|state| &state.column_dtype)
220    }
221
222    /// Returns the cache of pruned zones shared by every execution of this pruning plan.
223    pub(crate) fn pruned_zones(&self) -> Option<&Arc<PrunedZones>> {
224        self.data()
225            .pruning
226            .as_ref()
227            .map(|state| &state.pruned_zones)
228    }
229
230    /// Returns the plan producing the summarised data, unless this is a pruning plan.
231    pub fn data_plan(&self) -> VortexResult<Option<PlanRef>> {
232        if self.is_pruning() {
233            return Ok(None);
234        }
235        self.child(DATA)
236    }
237
238    /// Returns the plan producing the zone statistics.
239    pub fn zones_plan(&self) -> VortexResult<PlanRef> {
240        let index = if self.is_pruning() { 0 } else { ZONES };
241        self.child_required(index)
242    }
243
244    /// Returns whether this plan represents a zone-backed pruning proof.
245    pub fn is_pruning(&self) -> bool {
246        self.data().pruning.is_some()
247    }
248
249    /// Returns the abstract pruning proof carried by this plan, when present.
250    pub fn pruning_expression(&self) -> Option<&BoundExpression> {
251        self.data().pruning.as_ref().map(|state| &state.expression)
252    }
253
254    fn with_pruning(&self, expression: BoundExpression) -> VortexResult<Option<Self>> {
255        if self.data().zone_len == 0 || self.is_pruning() {
256            return Ok(None);
257        }
258        let mut data = self.data().clone();
259        data.pruning = Some(ZonedPruningState {
260            expression: expression.clone(),
261            column_dtype: self.dtype().clone(),
262            pruned_zones: Arc::new(PrunedZones::new(&expression)),
263        });
264        Ok(Some(
265            PlanParts {
266                vtable: Zoned,
267                dtype: expression.dtype().clone(),
268                row_count: self.row_count(),
269                children: vec![self.zones_plan()?].into(),
270                data,
271            }
272            .into_typed(),
273        ))
274    }
275}
276
277impl PlanVTable for Zoned {
278    type PlanData = ZonedData;
279    type Metadata = EmptyMetadata;
280
281    fn id(&self) -> PlanId {
282        static ID: CachedId = CachedId::new("vortex.plan.zoned");
283        *ID
284    }
285
286    fn fmt(plan: &Plan<Self>, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
287        if let Some(expression) = plan.pruning_expression() {
288            write!(formatter, " prune={expression}")?;
289        }
290        Ok(())
291    }
292
293    fn metadata(_plan: &Plan<Self>) -> Option<Self::Metadata> {
294        None
295    }
296
297    fn with_children(
298        plan: &Plan<Self>,
299        children: &PlanChildren,
300        data: &mut Self::PlanData,
301    ) -> VortexResult<()> {
302        if plan.is_pruning() {
303            check_child_count("Zoned pruning", children, 1)?;
304            // New zones may prove different zones, so the proof of the old ones is not kept.
305            if let Some(state) = &mut data.pruning {
306                state.pruned_zones = Arc::new(PrunedZones::new(&state.expression));
307            }
308            return Ok(());
309        }
310
311        check_child_count("Zoned", children, 2)?;
312        let data = children
313            .get(DATA)?
314            .ok_or_else(|| vortex_error::vortex_err!("Zoned data child is absent"))?;
315        if data.dtype() != plan.dtype() || data.row_count() != plan.row_count() {
316            vortex_error::vortex_bail!("Zoned data child shape does not match the plan output");
317        }
318        Ok(())
319    }
320
321    fn child_name(plan: &Plan<Self>, index: usize) -> Cow<'_, str> {
322        if plan.is_pruning() {
323            return if index == 0 {
324                Cow::Borrowed("zones")
325            } else {
326                Cow::Owned(format!("child[{index}]"))
327            };
328        }
329        match index {
330            DATA => Cow::Borrowed("data"),
331            ZONES => Cow::Borrowed("zones"),
332            _ => Cow::Owned(format!("child[{index}]")),
333        }
334    }
335
336    /// A data plan compiles as its data child. A pruning plan whose zones an earlier execution
337    /// proved emits the expanded result; otherwise it collects the whole zone table, proves the
338    /// expression over it, keeps the result, and emits it expanded to the selected rows.
339    fn compile(
340        plan: &Plan<Self>,
341        rows: Range<u64>,
342        mask: &Mask,
343        compiler: &mut Compiler<'_>,
344    ) -> VortexResult<Option<Chain>> {
345        if !plan.is_pruning() {
346            let Some(data) = plan.data_plan()? else {
347                vortex_bail!("Zoned plan has no data child");
348            };
349            return compiler.compile(&data, rows, mask);
350        }
351        if mask.all_false() {
352            return Ok(None);
353        }
354        let expression = plan
355            .pruning_expression()
356            .ok_or_else(|| vortex_err!("Zoned pruning plan has no expression"))?;
357        if let Some(pruned) = pruned_zones(plan)?.get(expression, compiler.session())? {
358            return Ok(Some(Chain::new(OnceSource::new(expand_zones(
359                plan, &rows, mask, &pruned,
360            )?))));
361        }
362        let zones = plan.zones_plan()?;
363        let count = usize::try_from(zones.row_count())?;
364        let Some(chain) = compiler.compile(&zones, 0..zones.row_count(), &Mask::new_true(count))?
365        else {
366            vortex_bail!("Zoned pruning plan has an empty zone table");
367        };
368        let zoned = plan.clone();
369        let mask = mask.clone();
370        let source = WholeSource::new(zones.dtype().clone(), move |table, cx| {
371            let pruned = prune_zones(&zoned, cx.session(), table)?;
372            expand_zones(&zoned, &rows, &mask, &pruned)
373        });
374        Ok(Some(compiler.join_with(vec![chain], |_| UNBOUNDED, source)))
375    }
376
377    fn reach(
378        plan: &Plan<Self>,
379        rows: Range<u64>,
380        at: &Reach,
381        visit: &mut dyn FnMut(SegmentId, Range<u64>),
382    ) -> VortexResult<()> {
383        if !plan.is_pruning() {
384            return match plan.data_plan()? {
385                Some(data) => data.reach(rows, at, visit),
386                None => Ok(()),
387            };
388        }
389        let zones = plan.zones_plan()?;
390        zones.reach(0..zones.row_count(), &at.fixed(&rows), visit)
391    }
392}
393
394/// The per-zone results kept on a pruning `plan`.
395fn pruned_zones(plan: &ZonedPlan) -> VortexResult<&Arc<PrunedZones>> {
396    plan.pruned_zones()
397        .ok_or_else(|| vortex_err!("Zone pruning needs a pruning plan"))
398}
399
400/// Builds the zone map from the zone table and proves the expression over it, one value per
401/// zone, keeping both on the plan.
402fn prune_zones(plan: &ZonedPlan, session: &VortexSession, table: ArrayRef) -> VortexResult<Mask> {
403    let table = table.execute::<StructArray>(&mut session.create_execution_ctx())?;
404    let (Some(expression), Some(column_dtype)) =
405        (plan.pruning_expression(), plan.pruning_column_dtype())
406    else {
407        vortex_bail!("Zone pruning needs a pruning plan");
408    };
409    let zone_map = ZoneMap::try_new(
410        column_dtype.clone(),
411        table,
412        Arc::clone(plan.aggregate_fns()),
413        plan.zone_len(),
414        plan.row_count(),
415    )?;
416    pruned_zones(plan)?.init(zone_map, expression, session)
417}
418
419/// Expands the per-zone result to one value per row of `rows` that `mask` selects.
420///
421/// Every row of a zone shares the zone's value, so the rows are filled a zone at a time and the
422/// selection is applied once at the end.
423fn expand_zones(
424    plan: &ZonedPlan,
425    rows: &Range<u64>,
426    mask: &Mask,
427    pruned: &Mask,
428) -> VortexResult<ArrayRef> {
429    let zone_len = plan.zone_len();
430    let start_zone = usize::try_from(rows.start / zone_len)?;
431    let end_zone = usize::try_from(rows.end.div_ceil(zone_len))?;
432    let selected_zones = pruned.slice(start_zone..end_zone);
433    if selected_zones.all_true() || selected_zones.all_false() {
434        let scalar = Scalar::from(selected_zones.all_true()).cast(plan.dtype())?;
435        return Ok(ConstantArray::new(scalar, mask.true_count()).into_array());
436    }
437    let mut bits = BitBufferMut::with_capacity(mask.len());
438    let mut row = rows.start;
439    while row < rows.end {
440        let zone = usize::try_from(row / zone_len).vortex_expect("zone index must fit in usize");
441        let zone_end = ((row / zone_len) + 1) * zone_len;
442        let end = zone_end.min(rows.end);
443        bits.append_n(
444            pruned.value(zone),
445            usize::try_from(end - row).vortex_expect("zone rows must fit in usize"),
446        );
447        row = end;
448    }
449    let validity = if plan.dtype().is_nullable() {
450        Validity::AllValid
451    } else {
452        Validity::NonNullable
453    };
454    let array = BoolArray::try_new(bits.freeze(), validity)?.into_array();
455    if mask.all_true() {
456        return Ok(array);
457    }
458    array.filter(mask.clone())
459}
460
461/// Rewrites an abstract statistic expression over a zoned plan into its pruning state.
462#[derive(Debug)]
463pub(crate) struct ExpressionZonedRule;
464
465impl PlanParentReduceRule<Zoned> for ExpressionZonedRule {
466    type Parent = Eval;
467
468    fn reduce_parent(
469        &self,
470        child: &ZonedPlan,
471        parent: &Plan<Eval>,
472        _child_idx: usize,
473    ) -> VortexResult<Option<PlanRef>> {
474        let mut contains_stat = false;
475        let mut contains_root = false;
476        parent.expression().clone().transform_down(|expression| {
477            if expression
478                .as_scalar()
479                .is_some_and(|scalar_fn| scalar_fn.is::<StatFn>())
480            {
481                contains_stat = true;
482                return Ok(Transformed {
483                    value: expression,
484                    order: TraversalOrder::Skip,
485                    changed: false,
486                });
487            }
488            contains_root |= expression.is_root();
489            Ok(Transformed::no(expression))
490        })?;
491        if !parent.dtype().is_boolean() || !contains_stat || contains_root {
492            return Ok(None);
493        }
494
495        Ok(child
496            .with_pruning(parent.expression().clone())?
497            .map(Plan::into_plan))
498    }
499}
vortex_layout/plan/physical/plans/zoned.rs

vortex_layout/plan/physical/plans/
zoned.rs
