vortex_layout/plan/
lower.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Test support for constructing physical plans from stored layout trees.
5//!
6//! This module is only used to build physical-plan fixtures for tests. It is not a production
7//! planning API.
8
9use vortex_error::VortexResult;
10use vortex_error::vortex_bail;
11use vortex_error::vortex_err;
12
13use crate::LayoutRef;
14use crate::layouts::chunked::Chunked;
15use crate::layouts::chunked::ChunkedLayout;
16use crate::layouts::dict::Dict;
17use crate::layouts::dict::DictLayout;
18use crate::layouts::flat::Flat;
19use crate::layouts::flat::FlatLayout;
20use crate::layouts::list::ELEMENTS_CHILD_INDEX;
21use crate::layouts::list::List;
22use crate::layouts::list::ListLayout;
23use crate::layouts::list::OFFSETS_CHILD_INDEX;
24use crate::layouts::list::VALIDITY_CHILD_INDEX;
25use crate::layouts::struct_::Struct;
26use crate::layouts::struct_::StructLayout;
27use crate::plan::ConcatPlan;
28use crate::plan::ListPackPlan;
29use crate::plan::PackPlan;
30use crate::plan::PlanChildren;
31use crate::plan::PlanRef;
32use crate::plan::SegmentScanPlan;
33use crate::plan::TakePlan;
34
35/// Constructs a physical-plan fixture from `layout` for tests.
36///
37/// The root operator is built immediately. Its child container owns a hidden clone of the source
38/// layout and lowers each child independently on first access.
39pub fn lower(layout: &LayoutRef) -> VortexResult<PlanRef> {
40    if let Some(layout) = layout.as_opt::<Flat>() {
41        return Ok(lower_flat(layout).into_plan());
42    }
43    if let Some(layout) = layout.as_opt::<Chunked>() {
44        return Ok(lower_chunked(layout)?.into_plan());
45    }
46    if let Some(layout) = layout.as_opt::<Struct>() {
47        return Ok(lower_struct(layout)?.into_plan());
48    }
49    if let Some(layout) = layout.as_opt::<Dict>() {
50        return Ok(lower_dict(layout)?.into_plan());
51    }
52    if let Some(layout) = layout.as_opt::<List>() {
53        return Ok(lower_list(layout)?.into_plan());
54    }
55    vortex_bail!(
56        "No physical plan implementation for layout '{}'",
57        layout.encoding_id()
58    )
59}
60
61fn lower_flat(layout: &FlatLayout) -> SegmentScanPlan {
62    SegmentScanPlan::new(
63        layout.dtype().clone(),
64        layout.row_count(),
65        layout.segment_id(),
66        layout.array_ctx().clone(),
67        layout.array_tree().cloned(),
68    )
69}
70
71fn lower_chunked(layout: &ChunkedLayout) -> VortexResult<ConcatPlan> {
72    let mut row_offsets = Vec::with_capacity(layout.nchildren());
73    let mut row_count = 0u64;
74    for index in 0..layout.nchildren() {
75        row_offsets.push(row_count);
76        row_count = row_count
77            .checked_add(layout.child_row_count(index))
78            .ok_or_else(|| vortex_err!("Chunked row count overflow"))?;
79    }
80    // SAFETY: Chunked layout construction validates that every child has the parent dtype and
81    // that their row counts sum to the parent; offsets were computed with checked addition above.
82    Ok(unsafe {
83        ConcatPlan::from_children_unchecked(
84            layout.dtype().clone(),
85            layout.row_count(),
86            row_offsets.into(),
87            lazy_children(layout.to_layout(), (0..layout.nchildren()).collect()),
88        )
89    })
90}
91
92fn lower_struct(layout: &StructLayout) -> VortexResult<PackPlan> {
93    // Struct layout slot 0 is validity and field i is slot i + 1. The plan puts validity last so
94    // field indices are identical to their plan-child indices.
95    let fields = layout.struct_fields().clone();
96    let mut slots = (1..=fields.nfields()).collect::<Vec<_>>();
97    if layout.dtype().is_nullable() {
98        slots.push(0);
99    }
100    // SAFETY: Struct layout construction validates child count and row counts, and its logical
101    // slot schema defines every field dtype plus the optional non-nullable boolean validity dtype.
102    Ok(unsafe {
103        PackPlan::from_children_unchecked(
104            fields,
105            layout.dtype().nullability(),
106            layout.row_count(),
107            lazy_children(layout.to_layout(), slots),
108        )
109    })
110}
111
112fn lower_dict(layout: &DictLayout) -> VortexResult<TakePlan> {
113    // Dict serialization stores values before codes; the plan order is deliberately codes,
114    // values because that is the optimizer-facing logical shape.
115    // SAFETY: Dict layout construction validates its values and codes slots. The plan reorders
116    // those slots to [codes, values] while preserving the derived output dtype and row domain.
117    Ok(unsafe {
118        TakePlan::from_children_unchecked(
119            layout.dtype().clone(),
120            layout.row_count(),
121            lazy_children(layout.to_layout(), vec![1, 0]),
122        )
123    })
124}
125
126fn lower_list(layout: &ListLayout) -> VortexResult<ListPackPlan> {
127    let mut slots = vec![ELEMENTS_CHILD_INDEX, OFFSETS_CHILD_INDEX];
128    if layout.dtype().is_nullable() {
129        slots.push(VALIDITY_CHILD_INDEX);
130    }
131    // SAFETY: List layout construction validates its element, offsets, and optional validity
132    // child shapes before this plan preserves them in the same logical order.
133    Ok(unsafe {
134        ListPackPlan::from_children_unchecked(
135            layout.dtype().clone(),
136            layout.row_count(),
137            lazy_children(layout.to_layout(), slots),
138        )
139    })
140}
141
142fn lazy_children(layout: LayoutRef, slots: Vec<usize>) -> PlanChildren {
143    PlanChildren::lazy(slots.len(), move |index| {
144        let slot = slots
145            .get(index)
146            .copied()
147            .ok_or_else(|| vortex_err!("Missing plan child slot {index}"))?;
148        let child = layout
149            .slot(slot)?
150            .ok_or_else(|| vortex_err!("Layout child slot {slot} is absent"))?;
151        lower(&child)
152    })
153}
vortex_layout/plan/lower.rs

vortex_layout/plan/
lower.rs
