vortex_layout/plan/physical/
children.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt;
5use std::sync::Arc;
6
7use once_cell::sync::OnceCell;
8use vortex_error::VortexResult;
9use vortex_error::vortex_err;
10
11use crate::plan::physical::PlanRef;
12
13type ChildInitializer = dyn Fn(usize) -> VortexResult<PlanRef> + 'static + Send + Sync;
14
15/// Ordered plan children that may be initialized one slot at a time.
16///
17/// Eagerly constructed operators store already-filled slots. Layout lowering instead installs an
18/// initializer that owns the source layout and lowers each child on first access.
19#[derive(Clone)]
20pub struct PlanChildren {
21    initializer: Option<Arc<ChildInitializer>>,
22    cache: Arc<[OnceCell<PlanRef>]>,
23}
24
25impl PlanChildren {
26    /// Creates lazy child slots backed by `initializer`.
27    pub(crate) fn lazy(
28        len: usize,
29        initializer: impl Fn(usize) -> VortexResult<PlanRef> + 'static + Send + Sync,
30    ) -> Self {
31        Self {
32            initializer: Some(Arc::new(initializer)),
33            cache: (0..len).map(|_| OnceCell::new()).collect::<Vec<_>>().into(),
34        }
35    }
36
37    /// Returns the number of children without initializing any slot.
38    pub fn len(&self) -> usize {
39        self.cache.len()
40    }
41
42    /// Returns whether there are no children.
43    pub fn is_empty(&self) -> bool {
44        self.cache.is_empty()
45    }
46
47    /// Returns a child, initializing and caching its slot on first access.
48    pub fn get(&self, index: usize) -> VortexResult<Option<PlanRef>> {
49        self.get_ref(index).map(|child| child.cloned())
50    }
51
52    /// Borrows a child, initializing and caching its slot on first access.
53    pub(crate) fn get_ref(&self, index: usize) -> VortexResult<Option<&PlanRef>> {
54        let Some(cell) = self.cache.get(index) else {
55            return Ok(None);
56        };
57        if let Some(child) = cell.get() {
58            return Ok(Some(child));
59        }
60
61        let initializer = self
62            .initializer
63            .as_ref()
64            .ok_or_else(|| vortex_err!("Plan child {index} was not initialized"))?;
65        Ok(Some(cell.get_or_try_init(|| initializer(index))?))
66    }
67
68    /// Iterates over the children in logical order, initializing slots as they are visited.
69    pub fn iter(&self) -> impl ExactSizeIterator<Item = VortexResult<PlanRef>> + '_ {
70        (0..self.len()).map(|index| {
71            self.get(index)?
72                .ok_or_else(|| vortex_err!("Plan child {index} is absent"))
73        })
74    }
75
76    /// Borrows the children in logical order, initializing slots as they are visited.
77    pub(crate) fn iter_refs(&self) -> impl ExactSizeIterator<Item = VortexResult<&PlanRef>> + '_ {
78        (0..self.len()).map(|index| {
79            self.get_ref(index)?
80                .ok_or_else(|| vortex_err!("Plan child {index} is absent"))
81        })
82    }
83
84    /// Materializes all children into an eager vector.
85    pub fn to_vec(&self) -> VortexResult<Vec<PlanRef>> {
86        self.iter().collect()
87    }
88}
89
90impl From<Vec<PlanRef>> for PlanChildren {
91    fn from(children: Vec<PlanRef>) -> Self {
92        let cache = children
93            .into_iter()
94            .map(OnceCell::with_value)
95            .collect::<Vec<_>>()
96            .into();
97        Self {
98            initializer: None,
99            cache,
100        }
101    }
102}
103
104impl<const N: usize> From<[PlanRef; N]> for PlanChildren {
105    fn from(children: [PlanRef; N]) -> Self {
106        Vec::from(children).into()
107    }
108}
109
110impl Default for PlanChildren {
111    fn default() -> Self {
112        Vec::new().into()
113    }
114}
115
116impl fmt::Debug for PlanChildren {
117    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
118        formatter
119            .debug_struct("PlanChildren")
120            .field("len", &self.len())
121            .field(
122                "initialized",
123                &self
124                    .cache
125                    .iter()
126                    .filter(|slot| slot.get().is_some())
127                    .count(),
128            )
129            .finish()
130    }
131}
vortex_layout/plan/physical/children.rs

vortex_layout/plan/physical/
children.rs
