Skip to main content

vortex_layout/plan/physical/
lower.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Lowering stored layout trees into physical plans.
5
6use std::sync::Arc;
7
8use vortex_error::VortexResult;
9use vortex_error::vortex_bail;
10use vortex_error::vortex_err;
11
12use crate::LayoutRef;
13use crate::layouts::chunked::Chunked;
14use crate::layouts::chunked::ChunkedLayout;
15use crate::layouts::dict::Dict;
16use crate::layouts::dict::DictLayout;
17use crate::layouts::flat::Flat;
18use crate::layouts::flat::FlatLayout;
19use crate::layouts::list::ELEMENTS_CHILD_INDEX;
20use crate::layouts::list::List;
21use crate::layouts::list::ListLayout;
22use crate::layouts::list::OFFSETS_CHILD_INDEX;
23use crate::layouts::list::VALIDITY_CHILD_INDEX;
24use crate::layouts::struct_::Struct;
25use crate::layouts::struct_::StructLayout;
26use crate::layouts::zoned::LegacyStats;
27use crate::layouts::zoned::Zoned;
28use crate::plan::physical::ConcatPlan;
29use crate::plan::physical::FilterPlan;
30use crate::plan::physical::LegacyStatsPlan;
31use crate::plan::physical::ListPackPlan;
32use crate::plan::physical::PackPlan;
33use crate::plan::physical::PlanChildren;
34use crate::plan::physical::PlanRef;
35use crate::plan::physical::SegmentScanPlan;
36use crate::plan::physical::TakePlan;
37use crate::plan::physical::ZonedPlan;
38
39/// Lowers `layout` into a physical plan.
40///
41/// Zoned and legacy stats layouts keep their zone statistics: unoptimized they read only their
42/// data, and the optimizer rewrites a statistics predicate over them into a pruning plan. Those
43/// with no zones or no statistics lower to their data.
44///
45/// The root operator is built immediately. Its child container owns a hidden clone of the source
46/// layout and lowers each child independently on first access.
47///
48/// A flat layout lowers to a [`Filter`](crate::plan::physical::Filter) over its segment scan, so
49/// the scan returns only the rows it is executed with.
50pub fn lower(layout: &LayoutRef) -> VortexResult<PlanRef> {
51    if let Some(layout) = layout.as_opt::<Flat>() {
52        return Ok(FilterPlan::new(lower_flat(layout).into_plan()).into_plan());
53    }
54    if let Some(layout) = layout.as_opt::<Chunked>() {
55        return Ok(lower_chunked(layout)?.into_plan());
56    }
57    if let Some(layout) = layout.as_opt::<Struct>() {
58        return Ok(lower_struct(layout)?.into_plan());
59    }
60    if let Some(layout) = layout.as_opt::<Dict>() {
61        return Ok(lower_dict(layout)?.into_plan());
62    }
63    if let Some(layout) = layout.as_opt::<List>() {
64        return Ok(lower_list(layout)?.into_plan());
65    }
66    if layout.is::<Zoned>() || layout.is::<LegacyStats>() {
67        return lower_zoned(layout);
68    }
69    vortex_bail!(
70        "No physical plan implementation for layout '{}'",
71        layout.encoding_id()
72    )
73}
74
75fn lower_flat(layout: &FlatLayout) -> SegmentScanPlan {
76    SegmentScanPlan::new(
77        layout.dtype().clone(),
78        layout.row_count(),
79        layout.segment_id(),
80        layout.array_ctx().clone(),
81        layout.array_tree().cloned(),
82    )
83}
84
85fn lower_chunked(layout: &ChunkedLayout) -> VortexResult<ConcatPlan> {
86    let mut row_offsets = Vec::with_capacity(layout.nchildren());
87    let mut row_count = 0u64;
88    for index in 0..layout.nchildren() {
89        row_offsets.push(row_count);
90        row_count = row_count
91            .checked_add(layout.child_row_count(index))
92            .ok_or_else(|| vortex_err!("Chunked row count overflow"))?;
93    }
94    // SAFETY: Chunked layout construction validates that every child has the parent dtype and
95    // that their row counts sum to the parent; offsets were computed with checked addition above.
96    Ok(unsafe {
97        ConcatPlan::from_children_unchecked(
98            layout.dtype().clone(),
99            layout.row_count(),
100            row_offsets.into(),
101            lazy_children(layout.to_layout(), (0..layout.nchildren()).collect()),
102        )
103    })
104}
105
106fn lower_struct(layout: &StructLayout) -> VortexResult<PackPlan> {
107    // Struct layout slot 0 is validity and field i is slot i + 1. The plan puts validity last so
108    // field indices are identical to their plan-child indices.
109    let fields = layout.struct_fields().clone();
110    let mut slots = (1..=fields.nfields()).collect::<Vec<_>>();
111    if layout.dtype().is_nullable() {
112        slots.push(0);
113    }
114    // SAFETY: Struct layout construction validates child count and row counts, and its logical
115    // slot schema defines every field dtype plus the optional non-nullable boolean validity dtype.
116    Ok(unsafe {
117        PackPlan::from_children_unchecked(
118            fields,
119            layout.dtype().nullability(),
120            layout.row_count(),
121            lazy_children(layout.to_layout(), slots),
122        )
123    })
124}
125
126fn lower_dict(layout: &DictLayout) -> VortexResult<TakePlan> {
127    // Dict serialization stores values before codes; the plan order is deliberately codes,
128    // values because that is the optimizer-facing logical shape.
129    // SAFETY: Dict layout construction validates its values and codes slots. The plan reorders
130    // those slots to [codes, values] while preserving the derived output dtype and row domain.
131    Ok(unsafe {
132        TakePlan::from_children_unchecked(
133            layout.dtype().clone(),
134            layout.row_count(),
135            lazy_children(layout.to_layout(), vec![1, 0]),
136        )
137    })
138}
139
140fn lower_list(layout: &ListLayout) -> VortexResult<ListPackPlan> {
141    let mut slots = vec![ELEMENTS_CHILD_INDEX, OFFSETS_CHILD_INDEX];
142    if layout.dtype().is_nullable() {
143        slots.push(VALIDITY_CHILD_INDEX);
144    }
145    // SAFETY: List layout construction validates its element, offsets, and optional validity
146    // child shapes before this plan preserves them in the same logical order.
147    Ok(unsafe {
148        ListPackPlan::from_children_unchecked(
149            layout.dtype().clone(),
150            layout.row_count(),
151            lazy_children(layout.to_layout(), slots),
152        )
153    })
154}
155
156fn lazy_children(layout: LayoutRef, slots: Vec<usize>) -> PlanChildren {
157    PlanChildren::lazy(slots.len(), move |index| {
158        let slot = slots
159            .get(index)
160            .copied()
161            .ok_or_else(|| vortex_err!("Missing plan child slot {index}"))?;
162        let child = layout
163            .slot(slot)?
164            .ok_or_else(|| vortex_err!("Layout child slot {slot} is absent"))?;
165        lower(&child)
166    })
167}
168
169fn lower_zoned(layout: &LayoutRef) -> VortexResult<PlanRef> {
170    // Zoned and legacy stats layouts share a child shape: transparent data, auxiliary zones.
171    let metadata = if let Some(layout) = layout.as_opt::<Zoned>() {
172        layout.data()
173    } else if let Some(layout) = layout.as_opt::<LegacyStats>() {
174        layout.data()
175    } else {
176        vortex_bail!("Zoned plan requires a zoned layout")
177    };
178    // Without zones, or with a zone table holding no statistics, nothing can be pruned.
179    let no_stats = metadata
180        .stats_table_dtype()
181        .as_struct_fields_opt()
182        .is_none_or(|fields| fields.nfields() == 0);
183    if metadata.zone_len() == 0 || no_stats {
184        let data = layout
185            .slot(0)?
186            .ok_or_else(|| vortex_err!("Zoned layout has no data child"))?;
187        return lower(&data);
188    }
189    let children = lazy_children(Arc::clone(layout), vec![0, 1]);
190    let zone_len = u64::try_from(metadata.zone_len())?;
191    // A legacy table is converted to the aggregate schema when the plan is optimized.
192    Ok(match metadata.legacy_stats() {
193        Some(stats) => LegacyStatsPlan::from_children(
194            layout.dtype().clone(),
195            layout.row_count(),
196            children,
197            zone_len,
198            Arc::clone(stats),
199        )
200        .into_plan(),
201        None => ZonedPlan::from_children(
202            layout.dtype().clone(),
203            layout.row_count(),
204            children,
205            zone_len,
206            metadata.aggregate_fns(),
207        )
208        .into_plan(),
209    })
210}