vortex_layout/plan/physical/plans/
take.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::Canonical;
8use vortex_array::EmptyMetadata;
9use vortex_array::IntoArray;
10use vortex_array::dtype::DType;
11use vortex_array::expr::BoundExpression;
12use vortex_array::expr::BoundLabels;
13use vortex_array::expr::ExactBoundExpr;
14use vortex_array::expr::label_bound_tree;
15use vortex_array::scalar_fn::fns::dynamic::DynamicComparison;
16use vortex_array::scalar_fn::is_negative_cost;
17use vortex_error::VortexResult;
18use vortex_mask::Mask;
19use vortex_session::registry::CachedId;
20
21use crate::plan::physical::Eval;
22use crate::plan::physical::EvalPlan;
23use crate::plan::physical::Plan;
24use crate::plan::physical::PlanChildren;
25use crate::plan::physical::PlanId;
26use crate::plan::physical::PlanParts;
27use crate::plan::physical::PlanRef;
28use crate::plan::physical::PlanVTable;
29use crate::plan::physical::Share;
30use crate::plan::physical::check_child_count;
31use crate::plan::physical::operators::OnceSource;
32use crate::plan::physical::operators::TakeSource;
33use crate::plan::physical::optimizer::PlanParentReduceRule;
34use crate::plan::physical::pipeline::Chain;
35use crate::plan::physical::pipeline::Compiler;
36use crate::plan::physical::pipeline::Reach;
37use crate::segments::SegmentId;
38
39const CODES: usize = 0;
40const VALUES: usize = 1;
41
42/// Indexes `values` by `codes`, with children ordered as `[codes, values]`.
43#[derive(Clone, Debug)]
44pub struct Take;
45
46/// A plan that indexes one child by another.
47pub type TakePlan = Plan<Take>;
48
49impl TakePlan {
50    /// Creates a take from potentially unresolved children without validation.
51    ///
52    /// # Safety
53    ///
54    /// `children` must be `[codes, values]`; `codes` must have `row_count` rows; and `dtype` must
55    /// be the values dtype unioned with the codes nullability.
56    pub(crate) unsafe fn from_children_unchecked(
57        dtype: DType,
58        row_count: u64,
59        children: PlanChildren,
60    ) -> Self {
61        PlanParts {
62            vtable: Take,
63            dtype,
64            row_count,
65            children,
66            data: (),
67        }
68        .into_typed()
69    }
70
71    /// Creates a take of `values` at `codes`.
72    ///
73    /// The row domain is that of `codes`, and the output dtype is that of `values`.
74    pub fn new(codes: PlanRef, values: PlanRef) -> Self {
75        let dtype = values
76            .dtype()
77            .union_nullability(codes.dtype().nullability());
78        let row_count = codes.row_count();
79        PlanParts {
80            vtable: Take,
81            dtype,
82            row_count,
83            children: vec![codes, values].into(),
84            data: (),
85        }
86        .into_typed()
87    }
88
89    /// Returns the plan producing indices.
90    pub fn codes(&self) -> VortexResult<PlanRef> {
91        self.child_required(CODES)
92    }
93
94    /// Returns the plan producing the values being indexed.
95    pub fn values(&self) -> VortexResult<PlanRef> {
96        self.child_required(VALUES)
97    }
98}
99
100impl PlanVTable for Take {
101    type PlanData = ();
102    type Metadata = EmptyMetadata;
103
104    fn id(&self) -> PlanId {
105        static ID: CachedId = CachedId::new("vortex.plan.take");
106        *ID
107    }
108
109    fn metadata(_plan: &Plan<Self>) -> Option<Self::Metadata> {
110        Some(EmptyMetadata)
111    }
112
113    fn with_children(
114        plan: &Plan<Self>,
115        children: &PlanChildren,
116        _data: &mut Self::PlanData,
117    ) -> VortexResult<()> {
118        check_child_count("Take", children, 2)?;
119        let codes = children
120            .get(CODES)?
121            .ok_or_else(|| vortex_error::vortex_err!("Take codes child is absent"))?;
122        let values = children
123            .get(VALUES)?
124            .ok_or_else(|| vortex_error::vortex_err!("Take values child is absent"))?;
125        let dtype = values
126            .dtype()
127            .union_nullability(codes.dtype().nullability());
128        if codes.row_count() != plan.row_count() || &dtype != plan.dtype() {
129            vortex_error::vortex_bail!("Take child shape does not match the plan output");
130        }
131        Ok(())
132    }
133
134    fn child_name(_plan: &Plan<Self>, index: usize) -> Cow<'_, str> {
135        match index {
136            CODES => Cow::Borrowed("codes"),
137            VALUES => Cow::Borrowed("values"),
138            _ => Cow::Owned(format!("child[{index}]")),
139        }
140    }
141
142    /// A take whose values an earlier scan of the split produced looks its codes up in them;
143    /// otherwise it joins its codes with the whole of its values.
144    fn compile(
145        plan: &Plan<Self>,
146        rows: Range<u64>,
147        mask: &Mask,
148        compiler: &mut Compiler<'_>,
149    ) -> VortexResult<Option<Chain>> {
150        let Some(codes) = compiler.compile(&plan.codes()?, rows, mask)? else {
151            return Ok(None);
152        };
153        let cache = compiler.values().clone();
154        let values_plan = plan.values()?;
155        if let Some(values) = cache.get(&values_plan) {
156            let source = TakeSource::new(plan.clone(), Some(values), None);
157            return Ok(Some(compiler.join(vec![codes], source)));
158        }
159        let values = values_plan;
160        let values_rows = 0..values.row_count();
161        let len = usize::try_from(values_rows.end)?;
162        let values = match compiler.compile(&values, values_rows, &Mask::new_true(len))? {
163            Some(values) => values,
164            None => Chain::new(OnceSource::new(
165                Canonical::empty(values.dtype()).into_array(),
166            )),
167        };
168        let source = TakeSource::new(plan.clone(), None, Some(cache));
169        Ok(Some(compiler.join_with(
170            vec![codes, values],
171            TakeSource::capacity,
172            source,
173        )))
174    }
175
176    fn reach(
177        plan: &Plan<Self>,
178        rows: Range<u64>,
179        at: &Reach,
180        visit: &mut dyn FnMut(SegmentId, Range<u64>),
181    ) -> VortexResult<()> {
182        let values = plan.values()?;
183        let len = values.row_count();
184        values.reach(0..len, &at.fixed(&rows), visit)?;
185        plan.codes()?.reach(rows, at, visit)
186    }
187}
188
189/// Pushes an expression, or the part of it that reads the values, onto the dictionary values of
190/// a [`Take`].
191///
192/// Evaluating over values rather than codes is only sound for infallible expressions, since values
193/// no selected row references are evaluated too, and, when codes may be null, strict ones:
194/// otherwise per-row behaviour is not preserved.
195///
196/// A boolean expression is pushed whole onto the values as they are, so it stays above any
197/// [`Share`] of them. Otherwise the largest part holding every read of the values and built only
198/// from negative-cost functions, such as the byte length in `cast(byte_length($))`, is pushed
199/// beneath the share, since it is cheaper over the encoded values than canonicalizing them; the
200/// rest stays above the take.
201#[derive(Debug)]
202pub(crate) struct ExpressionTakeRule;
203
204/// Per expression node: whether it reads the root, and whether it is strict, infallible, and built
205/// only from negative-cost functions.
206type Labels = BoundLabels<(bool, bool, bool, bool)>;
207
208impl PlanParentReduceRule<Take> for ExpressionTakeRule {
209    type Parent = Eval;
210
211    fn reduce_parent(
212        &self,
213        child: &Plan<Take>,
214        parent: &Plan<Eval>,
215        _child_idx: usize,
216    ) -> VortexResult<Option<PlanRef>> {
217        let expression = parent.expression();
218        // A dynamic comparison stays above the take: evaluated per batch it reads the latest
219        // bound, and the dictionary batches the take emits evaluate it over their values anyway.
220        // Values it never reaches can always be kept between scans.
221        if expression.contains::<DynamicComparison>()? {
222            return Ok(None);
223        }
224        // Strictness keeps a null code null. Without null codes every row takes its value's result.
225        let codes_nullable = child.codes()?.dtype().is_nullable();
226        let labels = label_bound_tree(
227            expression,
228            |node| match node.as_scalar() {
229                Some(scalar_fn) => (
230                    false,
231                    scalar_fn.signature().is_strict() || !codes_nullable,
232                    scalar_fn.signature().is_infallible(),
233                    is_negative_cost(scalar_fn.id()),
234                ),
235                None => (true, true, true, true),
236            },
237            |acc, &child| {
238                (
239                    acc.0 | child.0,
240                    acc.1 & child.1,
241                    acc.2 & child.2,
242                    acc.3 & child.3,
243                )
244            },
245        );
246        let label = |node: &BoundExpression| {
247            labels
248                .get(&ExactBoundExpr(node.clone()))
249                .copied()
250                .unwrap_or((false, false, false, false))
251        };
252
253        let (references_root, is_strict, is_infallible, is_negative_cost) = label(expression);
254        if references_root && is_strict && is_infallible && !is_negative_cost {
255            if !expression.dtype().is_boolean() {
256                return Ok(None);
257            }
258            let values = child.values()?;
259            let Some(expression) = rebind_root(expression, child.dtype(), values.dtype()) else {
260                return Ok(None);
261            };
262            let values = EvalPlan::try_new(expression, values)?.into_plan();
263            let take = TakePlan::new(child.codes()?, values).into_plan();
264            if take.dtype() != parent.dtype() {
265                return Ok(None);
266            }
267            return Ok(Some(take));
268        }
269
270        let Some(inner) = negative_cost_part(expression, &labels) else {
271            return Ok(None);
272        };
273        let values = child.values()?;
274        let values = match values.as_opt::<Share>() {
275            Some(share) => share.child_plan()?,
276            None => values,
277        };
278        let Some(pushed) = rebind_root(&inner, child.dtype(), values.dtype()) else {
279            return Ok(None);
280        };
281        let values = EvalPlan::try_new(pushed, values)?.into_plan();
282        let take = TakePlan::new(child.codes()?, values).into_plan();
283        let outer = replace(
284            expression,
285            &inner,
286            &BoundExpression::new_root(take.dtype().clone()),
287        )?;
288        if outer.dtype() != expression.dtype() {
289            return Ok(None);
290        }
291        if outer.is_root() {
292            return Ok(Some(take));
293        }
294        Ok(Some(EvalPlan::try_new(outer, take)?.into_plan()))
295    }
296}
297
298/// The largest strict, infallible, negative-cost part of `expression` holding every read of the
299/// root, unless that is the bare root.
300fn negative_cost_part(expression: &BoundExpression, labels: &Labels) -> Option<BoundExpression> {
301    let (references_root, is_strict, is_infallible, is_negative_cost) =
302        labels.get(&ExactBoundExpr(expression.clone())).copied()?;
303    if !references_root {
304        return None;
305    }
306    if is_strict && is_infallible && is_negative_cost {
307        return (!expression.is_root()).then(|| expression.clone());
308    }
309    let mut reading = expression.children().iter().filter(|child| {
310        labels
311            .get(&ExactBoundExpr((*child).clone()))
312            .is_some_and(|label| label.0)
313    });
314    let child = reading.next()?;
315    if reading.next().is_some() {
316        return None;
317    }
318    negative_cost_part(child, labels)
319}
320
321/// `expression`, bound to the take's dtype `from`, bound instead to its values' dtype `to`, or
322/// `None` when it does not apply there.
323///
324/// Values are less nullable than the take when its codes may be null: the take's nulls come from
325/// its codes, which a strict expression keeps null.
326fn rebind_root(expression: &BoundExpression, from: &DType, to: &DType) -> Option<BoundExpression> {
327    if from == to {
328        return Some(expression.clone());
329    }
330    replace(
331        expression,
332        &BoundExpression::new_root(from.clone()),
333        &BoundExpression::new_root(to.clone()),
334    )
335    .ok()
336}
337
338/// `expression` with each occurrence of `needle` replaced by `replacement`.
339fn replace(
340    expression: &BoundExpression,
341    needle: &BoundExpression,
342    replacement: &BoundExpression,
343) -> VortexResult<BoundExpression> {
344    if ExactBoundExpr(expression.clone()) == ExactBoundExpr(needle.clone()) {
345        return Ok(replacement.clone());
346    }
347    if expression.children().is_empty() {
348        return Ok(expression.clone());
349    }
350    let children = expression
351        .children()
352        .iter()
353        .map(|child| replace(child, needle, replacement))
354        .collect::<VortexResult<Vec<_>>>()?;
355    expression.clone().with_children(children)
356}
vortex_layout/plan/physical/plans/take.rs

vortex_layout/plan/physical/plans/
take.rs
