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;
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!(
139        partitioned.partition_annotations.len() == 2,
140        "Row-index expression produced more than two partitions"
141    );
142    let row_idx_index = partitioned
143        .partition_annotations
144        .iter()
145        .position(|partition| *partition == RowIdxExpressionPartition::RowIdx)
146        .ok_or_else(|| vortex_err!("Row-index expression has no row-index partition"))?;
147    let child_index = partitioned
148        .partition_annotations
149        .iter()
150        .position(|partition| *partition == RowIdxExpressionPartition::Child)
151        .ok_or_else(|| vortex_err!("Row-index expression has no data partition"))?;
152
153    let row_idx_partition = &partitioned.partitions[row_idx_index];
154    let child_partition = &partitioned.partitions[child_index];
155    let (Some(row_idx_pack), Some(child_pack)) = (
156        row_idx_partition
157            .as_scalar()
158            .and_then(|scalar_fn| scalar_fn.as_opt::<PackFn>()),
159        child_partition
160            .as_scalar()
161            .and_then(|scalar_fn| scalar_fn.as_opt::<PackFn>()),
162    ) else {
163        return Err(vortex_err!(
164            "Row-index expression partitions must be struct packs"
165        ));
166    };
167    let row_idx_partition_name = partitioned.partition_names[row_idx_index].clone();
168    let child_partition_name = partitioned.partition_names[child_index].clone();
169    let mut collapsed = Vec::with_capacity(2);
170
171    let row_idx_expression = if row_idx_partition.children().len() == 1 {
172        let Some(value_name) = row_idx_pack.names.get(0) else {
173            return Err(vortex_err!("Row-index expression partition is empty"));
174        };
175        collapsed.push((row_idx_partition_name, value_name.clone()));
176        row_idx_partition.children()[0].clone()
177    } else {
178        row_idx_partition.clone()
179    };
180    let child_expression = if child_partition.children().len() == 1 {
181        let Some(value_name) = child_pack.names.get(0) else {
182            return Err(vortex_err!("Data expression partition is empty"));
183        };
184        collapsed.push((child_partition_name, value_name.clone()));
185        child_partition.children()[0].clone()
186    } else {
187        child_partition.clone()
188    };
189
190    let row_count = child.row_count();
191    let row_idx_expression = replace_row_idx(row_idx_expression)?;
192    let row_idx_plan =
193        EvalPlan::try_new(row_idx_expression, RowIdxPlan::new(row_count).into_plan())?.into_plan();
194    let child_plan = EvalPlan::try_new(child_expression, child)?.into_plan();
195    let fields = StructFields::from_iter([
196        (ROW_IDX_PARTITION_NAME, row_idx_plan.dtype().clone()),
197        (CHILD_PARTITION_NAME, child_plan.dtype().clone()),
198    ]);
199    let partitions = PackPlan::try_new(
200        fields,
201        Nullability::NonNullable,
202        row_count,
203        vec![row_idx_plan, child_plan],
204        None,
205    )?;
206    let residual =
207        rewrite_partition_root(partitioned.root, partitions.dtype().clone(), &collapsed)?;
208
209    Ok(EvalPlan::try_new(residual, partitions.into_plan())?.into_plan())
210}
211
212fn replace_row_idx(expression: BoundExpression) -> VortexResult<BoundExpression> {
213    Ok(expression
214        .transform_down(|node| {
215            if node
216                .as_scalar()
217                .is_some_and(|scalar_fn| scalar_fn.is::<RowIdxFn>())
218            {
219                Ok(Transformed {
220                    value: BoundExpression::new_root(row_idx_dtype()),
221                    changed: true,
222                    order: TraversalOrder::Skip,
223                })
224            } else {
225                Ok(Transformed::no(node))
226            }
227        })?
228        .into_inner())
229}
230
231#[derive(Clone, Copy, PartialEq, Eq, Hash)]
232enum RowIdxExpressionPartition {
233    RowIdx,
234    Child,
235}
236
237impl RowIdxExpressionPartition {
238    fn name(self) -> &'static str {
239        match self {
240            Self::RowIdx => ROW_IDX_PARTITION_NAME,
241            Self::Child => CHILD_PARTITION_NAME,
242        }
243    }
244}
245
246impl Display for RowIdxExpressionPartition {
247    fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
248        formatter.write_str(self.name())
249    }
250}
251
252impl From<RowIdxExpressionPartition> for FieldName {
253    fn from(partition: RowIdxExpressionPartition) -> Self {
254        FieldName::from(partition.name())
255    }
256}
vortex_layout/plan/plans/row_idx.rs

vortex_layout/plan/plans/
row_idx.rs
