vortex_layout/plan/physical/plans/
share.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::borrow::Cow;
5use std::ops::Range;
6
7use vortex_array::ArrayRef;
8use vortex_array::EmptyMetadata;
9use vortex_array::IntoArray;
10use vortex_error::VortexResult;
11use vortex_error::vortex_bail;
12use vortex_error::vortex_err;
13use vortex_mask::Mask;
14use vortex_session::registry::CachedId;
15
16use crate::arrays::shared::Shared;
17use crate::plan::physical::Plan;
18use crate::plan::physical::PlanChildren;
19use crate::plan::physical::PlanId;
20use crate::plan::physical::PlanParts;
21use crate::plan::physical::PlanRef;
22use crate::plan::physical::PlanVTable;
23use crate::plan::physical::check_child_count;
24use crate::plan::physical::operators::OnceSource;
25use crate::plan::physical::operators::UNBOUNDED;
26use crate::plan::physical::operators::WholeSource;
27use crate::plan::physical::pipeline::Chain;
28use crate::plan::physical::pipeline::Compiler;
29use crate::plan::physical::pipeline::Reach;
30use crate::segments::SegmentId;
31
32/// Produces its child's whole value once per split, as a
33/// [`SharedArray`](crate::arrays::shared::SharedArray) that the split's later scans reuse.
34///
35/// Dictionary values sit under a share. An expression the optimizer pushes onto the values stays
36/// above it, so the expression runs over values canonicalized once, and the projection reading the
37/// same dictionary finds them already canonical, rather than each evaluating over the encoded
38/// values.
39#[derive(Clone, Debug)]
40pub struct Share;
41
42/// Operator-specific data for a [`Share`] plan.
43#[derive(Clone, Debug, Default)]
44pub struct ShareData;
45
46/// A plan that shares its child's value between executions.
47pub type SharePlan = Plan<Share>;
48
49impl SharePlan {
50    /// Creates a share of `child`.
51    pub fn new(child: PlanRef) -> Self {
52        PlanParts {
53            vtable: Share,
54            dtype: child.dtype().clone(),
55            row_count: child.row_count(),
56            children: vec![child].into(),
57            data: ShareData,
58        }
59        .into_typed()
60    }
61
62    /// Returns the shared plan.
63    pub fn child_plan(&self) -> VortexResult<PlanRef> {
64        self.child_required(0)
65    }
66}
67
68impl PlanVTable for Share {
69    type PlanData = ShareData;
70    type Metadata = EmptyMetadata;
71
72    fn id(&self) -> PlanId {
73        static ID: CachedId = CachedId::new("vortex.plan.share");
74        *ID
75    }
76
77    fn metadata(_plan: &Plan<Self>) -> Option<Self::Metadata> {
78        Some(EmptyMetadata)
79    }
80
81    fn with_children(
82        plan: &Plan<Self>,
83        children: &PlanChildren,
84        _data: &mut Self::PlanData,
85    ) -> VortexResult<()> {
86        check_child_count("Share", children, 1)?;
87        let child = children
88            .get(0)?
89            .ok_or_else(|| vortex_err!("Share child is absent"))?;
90        if child.dtype() != plan.dtype() || child.row_count() != plan.row_count() {
91            vortex_bail!("Share child does not match the share's dtype and row count");
92        }
93        Ok(())
94    }
95
96    fn child_name(_plan: &Plan<Self>, index: usize) -> Cow<'_, str> {
97        if index == 0 {
98            Cow::Borrowed("child")
99        } else {
100            Cow::Owned(format!("child[{index}]"))
101        }
102    }
103
104    /// A share whose value an earlier scan of the split produced emits the selected rows of it;
105    /// otherwise it collects its child whole, keeps it for the split, and emits the selected
106    /// rows.
107    fn compile(
108        plan: &Plan<Self>,
109        rows: Range<u64>,
110        mask: &Mask,
111        compiler: &mut Compiler<'_>,
112    ) -> VortexResult<Option<Chain>> {
113        if mask.all_false() {
114            return Ok(None);
115        }
116        let cache = compiler.values().clone();
117        let shared = plan.clone().into_plan();
118        if let Some(value) = cache.get(&shared) {
119            return Ok(Some(Chain::new(OnceSource::new(selected_rows(
120                &value, &rows, mask,
121            )?))));
122        }
123        let child = plan.child_plan()?;
124        let len = usize::try_from(plan.row_count())?;
125        let Some(chain) = compiler.compile(&child, 0..plan.row_count(), &Mask::new_true(len))?
126        else {
127            return Ok(None);
128        };
129        let mask = mask.clone();
130        let source = WholeSource::new(plan.dtype().clone(), move |whole, _cx| {
131            let value = cache.insert(&shared, Shared::new(whole).into_array());
132            selected_rows(&value, &rows, &mask)
133        });
134        Ok(Some(compiler.join_with(vec![chain], |_| UNBOUNDED, source)))
135    }
136
137    fn reach(
138        plan: &Plan<Self>,
139        rows: Range<u64>,
140        at: &Reach,
141        visit: &mut dyn FnMut(SegmentId, Range<u64>),
142    ) -> VortexResult<()> {
143        plan.child_plan()?
144            .reach(0..plan.row_count(), &at.fixed(&rows), visit)
145    }
146}
147
148/// The rows of `rows` that `mask` selects, from a whole shared value.
149fn selected_rows(value: &ArrayRef, rows: &Range<u64>, mask: &Mask) -> VortexResult<ArrayRef> {
150    let array = value.slice(usize::try_from(rows.start)?..usize::try_from(rows.end)?)?;
151    if mask.all_true() {
152        return Ok(array);
153    }
154    array.filter(mask.clone())
155}
vortex_layout/plan/physical/plans/share.rs

vortex_layout/plan/physical/plans/
share.rs
