vortex_layout/plan/plans/
take.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::borrow::Cow;
5
6use vortex_array::EmptyMetadata;
7use vortex_array::dtype::DType;
8use vortex_array::expr::ExactBoundExpr;
9use vortex_array::expr::label_bound_tree;
10use vortex_error::VortexResult;
11use vortex_session::registry::CachedId;
12
13use crate::plan::Eval;
14use crate::plan::EvalPlan;
15use crate::plan::Plan;
16use crate::plan::PlanChildren;
17use crate::plan::PlanId;
18use crate::plan::PlanParts;
19use crate::plan::PlanRef;
20use crate::plan::PlanVTable;
21use crate::plan::check_child_count;
22use crate::plan::optimizer::PlanParentReduceRule;
23
24const CODES: usize = 0;
25const VALUES: usize = 1;
26
27/// Indexes `values` by `codes`, with children ordered as `[codes, values]`.
28#[derive(Clone, Debug)]
29pub struct Take;
30
31/// A plan that indexes one child by another.
32pub type TakePlan = Plan<Take>;
33
34impl TakePlan {
35    /// Creates a take from potentially unresolved children without validation.
36    ///
37    /// # Safety
38    ///
39    /// `children` must be `[codes, values]`; `codes` must have `row_count` rows; and `dtype` must
40    /// be the values dtype unioned with the codes nullability.
41    pub(crate) unsafe fn from_children_unchecked(
42        dtype: DType,
43        row_count: u64,
44        children: PlanChildren,
45    ) -> Self {
46        PlanParts {
47            vtable: Take,
48            dtype,
49            row_count,
50            children,
51            data: (),
52        }
53        .into_typed()
54    }
55
56    /// Creates a take of `values` at `codes`.
57    ///
58    /// The row domain is that of `codes`, and the output dtype is that of `values`.
59    pub fn new(codes: PlanRef, values: PlanRef) -> Self {
60        let dtype = values
61            .dtype()
62            .union_nullability(codes.dtype().nullability());
63        let row_count = codes.row_count();
64        // SAFETY: Parent metadata is derived from the ordered children immediately above.
65        unsafe { Self::from_children_unchecked(dtype, row_count, vec![codes, values].into()) }
66    }
67
68    /// Returns the plan producing indices.
69    pub fn codes(&self) -> VortexResult<PlanRef> {
70        self.child_required(CODES)
71    }
72
73    /// Returns the plan producing the values being indexed.
74    pub fn values(&self) -> VortexResult<PlanRef> {
75        self.child_required(VALUES)
76    }
77}
78
79impl PlanVTable for Take {
80    type PlanData = ();
81    type Metadata = EmptyMetadata;
82
83    fn id(&self) -> PlanId {
84        static ID: CachedId = CachedId::new("vortex.plan.take");
85        *ID
86    }
87
88    fn metadata(_plan: &Plan<Self>) -> Option<Self::Metadata> {
89        Some(EmptyMetadata)
90    }
91
92    fn with_children(
93        plan: &Plan<Self>,
94        children: &PlanChildren,
95        _data: &mut Self::PlanData,
96    ) -> VortexResult<()> {
97        check_child_count("Take", children, 2)?;
98        let codes = children
99            .get(CODES)?
100            .ok_or_else(|| vortex_error::vortex_err!("Take codes child is absent"))?;
101        let values = children
102            .get(VALUES)?
103            .ok_or_else(|| vortex_error::vortex_err!("Take values child is absent"))?;
104        let dtype = values
105            .dtype()
106            .union_nullability(codes.dtype().nullability());
107        if codes.row_count() != plan.row_count() || &dtype != plan.dtype() {
108            vortex_error::vortex_bail!("Take child shape does not match the plan output");
109        }
110        Ok(())
111    }
112
113    fn child_name(_plan: &Plan<Self>, index: usize) -> Cow<'_, str> {
114        match index {
115            CODES => Cow::Borrowed("codes"),
116            VALUES => Cow::Borrowed("values"),
117            _ => Cow::Owned(format!("child[{index}]")),
118        }
119    }
120}
121
122/// Pushes a strict, infallible boolean expression onto the dictionary values of a [`Take`].
123#[derive(Debug)]
124pub(crate) struct ExpressionTakeRule;
125
126impl PlanParentReduceRule<Take> for ExpressionTakeRule {
127    type Parent = Eval;
128
129    fn reduce_parent(
130        &self,
131        child: &Plan<Take>,
132        parent: &Plan<Eval>,
133        _child_idx: usize,
134    ) -> VortexResult<Option<PlanRef>> {
135        let expression = parent.expression();
136        if !expression.dtype().is_boolean() {
137            return Ok(None);
138        }
139        // Evaluating over values rather than codes is only sound when the expression reads the
140        // root, is strict, and cannot fail: otherwise per-row behaviour is not preserved.
141        let labels = label_bound_tree(
142            expression,
143            |node| match node.as_scalar() {
144                Some(scalar_fn) => (
145                    false,
146                    scalar_fn.signature().is_strict(),
147                    scalar_fn.signature().is_infallible(),
148                ),
149                None => (true, true, true),
150            },
151            |acc, &child| (acc.0 | child.0, acc.1 & child.1, acc.2 & child.2),
152        );
153        let (references_root, is_strict, is_infallible) = labels
154            .get(&ExactBoundExpr(expression.clone()))
155            .copied()
156            .unwrap_or((false, false, false));
157        if !references_root || !is_strict || !is_infallible {
158            return Ok(None);
159        }
160
161        let values = EvalPlan::try_new(expression.clone(), child.values()?)?.into_plan();
162        Ok(Some(TakePlan::new(child.codes()?, values).into_plan()))
163    }
164}
vortex_layout/plan/plans/take.rs

vortex_layout/plan/plans/
take.rs
