vortex_layout/plan/
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::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        let Some(cell) = self.cache.get(index) else {
50            return Ok(None);
51        };
52        if let Some(child) = cell.get() {
53            return Ok(Some(child.clone()));
54        }
55
56        let initializer = self
57            .initializer
58            .as_ref()
59            .ok_or_else(|| vortex_err!("Plan child {index} was not initialized"))?;
60        Ok(Some(cell.get_or_try_init(|| initializer(index))?.clone()))
61    }
62
63    /// Iterates over the children in logical order, initializing slots as they are visited.
64    pub fn iter(&self) -> impl ExactSizeIterator<Item = VortexResult<PlanRef>> + '_ {
65        (0..self.len()).map(|index| {
66            self.get(index)?
67                .ok_or_else(|| vortex_err!("Plan child {index} is absent"))
68        })
69    }
70
71    /// Materializes all children into an eager vector.
72    pub fn to_vec(&self) -> VortexResult<Vec<PlanRef>> {
73        self.iter().collect()
74    }
75}
76
77impl From<Vec<PlanRef>> for PlanChildren {
78    fn from(children: Vec<PlanRef>) -> Self {
79        let cache = children
80            .into_iter()
81            .map(OnceCell::with_value)
82            .collect::<Vec<_>>()
83            .into();
84        Self {
85            initializer: None,
86            cache,
87        }
88    }
89}
90
91impl<const N: usize> From<[PlanRef; N]> for PlanChildren {
92    fn from(children: [PlanRef; N]) -> Self {
93        Vec::from(children).into()
94    }
95}
96
97impl Default for PlanChildren {
98    fn default() -> Self {
99        Vec::new().into()
100    }
101}
102
103impl fmt::Debug for PlanChildren {
104    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
105        formatter
106            .debug_struct("PlanChildren")
107            .field("len", &self.len())
108            .field(
109                "initialized",
110                &self
111                    .cache
112                    .iter()
113                    .filter(|slot| slot.get().is_some())
114                    .count(),
115            )
116            .finish()
117    }
118}
vortex_layout/plan/children.rs

vortex_layout/plan/
children.rs
