vortex_layout/plan/plans/
row_idx.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Display;
5use std::fmt::Formatter;
6
7use vortex_array::EmptyMetadata;
8use vortex_array::dtype::DType;
9use vortex_array::dtype::FieldName;
10use vortex_array::dtype::Nullability;
11use vortex_array::dtype::PType;
12use vortex_array::dtype::StructFields;
13use vortex_array::expr::BoundExpression;
14use vortex_array::expr::transform::partition_bound;
15use vortex_array::expr::traversal::NodeExt;
16use vortex_array::expr::traversal::Transformed;
17use vortex_array::expr::traversal::TraversalOrder;
18use vortex_array::scalar_fn::fns::pack::Pack as PackFn;
19use vortex_error::VortexResult;
20use vortex_error::vortex_ensure_eq;
21use vortex_error::vortex_err;
22use vortex_session::registry::CachedId;
23
24use crate::layouts::row_idx::RowIdx as RowIdxFn;
25use crate::plan::EvalPlan;
26use crate::plan::PackPlan;
27use crate::plan::Plan;
28use crate::plan::PlanChildren;
29use crate::plan::PlanId;
30use crate::plan::PlanParts;
31use crate::plan::PlanRef;
32use crate::plan::PlanVTable;
33use crate::plan::check_child_count;
34use crate::plan::plans::pack::rewrite_partition_root;
35
36const ROW_IDX_PARTITION_NAME: &str = "row_idx";
37const CHILD_PARTITION_NAME: &str = "child";
38
39/// Generates global row indices for the current execution row domain.
40#[derive(Clone, Debug)]
41pub struct RowIdx;
42
43/// Operator-specific data for a [`RowIdx`] plan.
44#[derive(Clone, Debug)]
45pub struct RowIdxData;
46
47/// A childless source of global row indices supplied by its execution row domain.
48pub type RowIdxPlan = Plan<RowIdx>;
49
50impl RowIdxPlan {
51    /// Creates a row-index source covering `row_count` rows.
52    pub fn new(row_count: u64) -> Self {
53        PlanParts {
54            vtable: RowIdx,
55            dtype: row_idx_dtype(),
56            row_count,
57            children: PlanChildren::default(),
58            data: RowIdxData,
59        }
60        .into_typed()
61    }
62}
63
64/// Returns the dtype of a generated row index.
65pub fn row_idx_dtype() -> DType {
66    DType::Primitive(PType::U64, Nullability::NonNullable)
67}
68
69impl PlanVTable for RowIdx {
70    type PlanData = RowIdxData;
71    type Metadata = EmptyMetadata;
72
73    fn id(&self) -> PlanId {
74        static ID: CachedId = CachedId::new("vortex.plan.row_idx");
75        *ID
76    }
77
78    fn metadata(_plan: &Plan<Self>) -> Option<Self::Metadata> {
79        Some(EmptyMetadata)
80    }
81
82    fn with_children(
83        _plan: &Plan<Self>,
84        children: &PlanChildren,
85        _data: &mut Self::PlanData,
86    ) -> VortexResult<()> {
87        check_child_count("RowIdx", children, 0)
88    }
89}
90
91/// Plans an expression over a data source and its global row-index domain.
92///
93/// Expressions that only use data bypass the row-index source, expressions that only use
94/// `#row_idx` bypass `child`, and mixed expressions combine independently planned branches with a
95/// [`PackPlan`]. The file row offset is supplied by the execution row domain.
96pub fn plan_row_idx_expression(
97    expression: BoundExpression,
98    child: PlanRef,
99) -> VortexResult<PlanRef> {
100    let partitioned = partition_bound(expression.clone(), |node| {
101        if node
102            .as_scalar()
103            .is_some_and(|scalar_fn| scalar_fn.is::<RowIdxFn>())
104        {
105            vec![RowIdxExpressionPartition::RowIdx]
106        } else if node.is_root() {
107            vec![RowIdxExpressionPartition::Child]
108        } else {
109            vec![]
110        }
111    })?;
112
113    if partitioned.partition_annotations.is_empty() {
114        let row_domain = PackPlan::try_new(
115            StructFields::empty(),
116            Nullability::NonNullable,
117            child.row_count(),
118            Vec::new(),
119            None,
120        )?
121        .into_plan();
122        return Ok(EvalPlan::try_new(expression, row_domain)?.into_plan());
123    }
124
125    if partitioned.partition_annotations.len() == 1 {
126        return match partitioned.partition_annotations[0] {
127            RowIdxExpressionPartition::RowIdx => {
128                let expression = replace_row_idx(expression)?;
129                let values = RowIdxPlan::new(child.row_count()).into_plan();
130                Ok(EvalPlan::try_new(expression, values)?.into_plan())
131            }
132            RowIdxExpressionPartition::Child => {
133                Ok(EvalPlan::try_new(expression, child)?.into_plan())
134            }
135        };
136    }
137
138    vortex_ensure_eq!(
139        partitioned.partition_annotations.len(),
140        2,
141        "Row-index expression produced more than two partitions"
142    );
143    let row_idx_index = partitioned
144        .partition_annotations
145        .iter()
146        .position(|partition| *partition == RowIdxExpressionPartition::RowIdx)
147        .ok_or_else(|| vortex_err!("Row-index expression has no row-index partition"))?;
148    let child_index = partitioned
149        .partition_annotations
150        .iter()
151        .position(|partition| *partition == RowIdxExpressionPartition::Child)
152        .ok_or_else(|| vortex_err!("Row-index expression has no data partition"))?;
153
154    let row_idx_partition = &partitioned.partitions[row_idx_index];
155    let child_partition = &partitioned.partitions[child_index];
156    let (Some(row_idx_pack), Some(child_pack)) = (
157        row_idx_partition
158            .as_scalar()
159            .and_then(|scalar_fn| scalar_fn.as_opt::<PackFn>()),
160        child_partition
161            .as_scalar()
162            .and_then(|scalar_fn| scalar_fn.as_opt::<PackFn>()),
163    ) else {
164        return Err(vortex_err!(
165            "Row-index expression partitions must be struct packs"
166        ));
167    };
168    let row_idx_partition_name = partitioned.partition_names[row_idx_index].clone();
169    let child_partition_name = partitioned.partition_names[child_index].clone();
170    let mut collapsed = Vec::with_capacity(2);
171
172    let row_idx_expression = if row_idx_partition.children().len() == 1 {
173        let Some(value_name) = row_idx_pack.names.get(0) else {
174            return Err(vortex_err!("Row-index expression partition is empty"));
175        };
176        collapsed.push((row_idx_partition_name, value_name.clone()));
177        row_idx_partition.children()[0].clone()
178    } else {
179        row_idx_partition.clone()
180    };
181    let child_expression = if child_partition.children().len() == 1 {
182        let Some(value_name) = child_pack.names.get(0) else {
183            return Err(vortex_err!("Data expression partition is empty"));
184        };
185        collapsed.push((child_partition_name, value_name.clone()));
186        child_partition.children()[0].clone()
187    } else {
188        child_partition.clone()
189    };
190
191    let row_count = child.row_count();
192    let row_idx_expression = replace_row_idx(row_idx_expression)?;
193    let row_idx_plan =
194        EvalPlan::try_new(row_idx_expression, RowIdxPlan::new(row_count).into_plan())?.into_plan();
195    let child_plan = EvalPlan::try_new(child_expression, child)?.into_plan();
196    let fields = StructFields::from_iter([
197        (ROW_IDX_PARTITION_NAME, row_idx_plan.dtype().clone()),
198        (CHILD_PARTITION_NAME, child_plan.dtype().clone()),
199    ]);
200    let partitions = PackPlan::try_new(
201        fields,
202        Nullability::NonNullable,
203        row_count,
204        vec![row_idx_plan, child_plan],
205        None,
206    )?;
207    let residual =
208        rewrite_partition_root(partitioned.root, partitions.dtype().clone(), &collapsed)?;
209
210    Ok(EvalPlan::try_new(residual, partitions.into_plan())?.into_plan())
211}
212
213fn replace_row_idx(expression: BoundExpression) -> VortexResult<BoundExpression> {
214    Ok(expression
215        .transform_down(|node| {
216            if node
217                .as_scalar()
218                .is_some_and(|scalar_fn| scalar_fn.is::<RowIdxFn>())
219            {
220                Ok(Transformed {
221                    value: BoundExpression::new_root(row_idx_dtype()),
222                    changed: true,
223                    order: TraversalOrder::Skip,
224                })
225            } else {
226                Ok(Transformed::no(node))
227            }
228        })?
229        .into_inner())
230}
231
232#[derive(Clone, Copy, PartialEq, Eq, Hash)]
233enum RowIdxExpressionPartition {
234    RowIdx,
235    Child,
236}
237
238impl RowIdxExpressionPartition {
239    fn name(self) -> &'static str {
240        match self {
241            Self::RowIdx => ROW_IDX_PARTITION_NAME,
242            Self::Child => CHILD_PARTITION_NAME,
243        }
244    }
245}
246
247impl Display for RowIdxExpressionPartition {
248    fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
249        formatter.write_str(self.name())
250    }
251}
252
253impl From<RowIdxExpressionPartition> for FieldName {
254    fn from(partition: RowIdxExpressionPartition) -> Self {
255        FieldName::from(partition.name())
256    }
257}
vortex_layout/plan/plans/row_idx.rs

vortex_layout/plan/plans/
row_idx.rs
