vortex_layout/plan/plans/
concat.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::borrow::Cow;
5use std::sync::Arc;
6
7use vortex_array::EmptyMetadata;
8use vortex_array::dtype::DType;
9use vortex_error::VortexResult;
10use vortex_error::vortex_bail;
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::optimizer::PlanParentReduceRule;
22
23/// Concatenates its children row-wise.
24#[derive(Clone, Debug)]
25pub struct Concat;
26
27/// Row offsets of each concatenated child.
28#[derive(Clone, Debug)]
29pub struct ConcatData {
30    row_offsets: Arc<[u64]>,
31}
32
33/// A plan that concatenates its children row-wise.
34pub type ConcatPlan = Plan<Concat>;
35
36impl ConcatPlan {
37    /// Creates a concatenation from potentially unresolved children without validation.
38    ///
39    /// # Safety
40    ///
41    /// Every child must have `dtype`; `row_offsets` must contain the cumulative row offset of
42    /// every child; and the sum of all child row counts must equal `row_count` without overflow.
43    pub(crate) unsafe fn from_children_unchecked(
44        dtype: DType,
45        row_count: u64,
46        row_offsets: Arc<[u64]>,
47        children: PlanChildren,
48    ) -> Self {
49        PlanParts {
50            vtable: Concat,
51            dtype,
52            row_count,
53            children,
54            data: ConcatData { row_offsets },
55        }
56        .into_typed()
57    }
58
59    /// Creates a concatenation over `children`.
60    ///
61    /// Every child must produce `dtype`, and the row domain is the sum of the child row counts.
62    pub fn try_new(dtype: DType, children: Vec<PlanRef>) -> VortexResult<Self> {
63        let mut row_offsets = Vec::with_capacity(children.len());
64        let mut row_count = 0u64;
65        for child in &children {
66            if child.dtype() != &dtype {
67                vortex_bail!(
68                    "Concat child dtype {} does not match {dtype}",
69                    child.dtype()
70                );
71            }
72            row_offsets.push(row_count);
73            row_count = row_count
74                .checked_add(child.row_count())
75                .ok_or_else(|| vortex_error::vortex_err!("Concat row count overflow"))?;
76        }
77        // SAFETY: Child dtypes and the checked cumulative row metadata were validated above.
78        Ok(unsafe {
79            Self::from_children_unchecked(dtype, row_count, row_offsets.into(), children.into())
80        })
81    }
82
83    /// Returns the first row of each child within this plan's row domain.
84    pub fn row_offsets(&self) -> &[u64] {
85        &self.data().row_offsets
86    }
87}
88
89impl PlanVTable for Concat {
90    type PlanData = ConcatData;
91    type Metadata = EmptyMetadata;
92
93    fn id(&self) -> PlanId {
94        static ID: CachedId = CachedId::new("vortex.plan.concat");
95        *ID
96    }
97
98    fn metadata(_plan: &Plan<Self>) -> Option<Self::Metadata> {
99        // Row offsets are derived from the children, so nothing needs storing.
100        Some(EmptyMetadata)
101    }
102
103    fn with_children(
104        plan: &Plan<Self>,
105        children: &PlanChildren,
106        data: &mut Self::PlanData,
107    ) -> VortexResult<()> {
108        if children.len() != plan.children().len() {
109            vortex_bail!(
110                "Concat expects {} children but got {}",
111                plan.children().len(),
112                children.len()
113            );
114        }
115
116        let mut row_offsets = Vec::with_capacity(children.len());
117        let mut row_count = 0u64;
118        for child in children.iter() {
119            let child = child?;
120            if child.dtype() != plan.dtype() {
121                vortex_bail!(
122                    "Concat child dtype {} does not match {}",
123                    child.dtype(),
124                    plan.dtype()
125                );
126            }
127            row_offsets.push(row_count);
128            row_count = row_count
129                .checked_add(child.row_count())
130                .ok_or_else(|| vortex_error::vortex_err!("Concat row count overflow"))?;
131        }
132        if row_count != plan.row_count() {
133            vortex_bail!(
134                "Concat children have {row_count} rows but the plan has {}",
135                plan.row_count()
136            );
137        }
138        data.row_offsets = row_offsets.into();
139        Ok(())
140    }
141
142    fn child_name(_plan: &Plan<Self>, index: usize) -> Cow<'_, str> {
143        Cow::Owned(format!("chunks[{index}]"))
144    }
145}
146
147/// Pushes an expression into every chunk of a [`Concat`].
148#[derive(Debug)]
149pub(crate) struct ExpressionConcatRule;
150
151impl PlanParentReduceRule<Concat> for ExpressionConcatRule {
152    type Parent = Eval;
153
154    fn reduce_parent(
155        &self,
156        child: &Plan<Concat>,
157        parent: &Plan<Eval>,
158        _child_idx: usize,
159    ) -> VortexResult<Option<PlanRef>> {
160        let expression = parent.expression();
161        let chunks = child
162            .children()
163            .iter()
164            .map(|chunk| Ok(EvalPlan::try_new(expression.clone(), chunk?)?.into_plan()))
165            .collect::<VortexResult<Vec<_>>>()?;
166        Ok(Some(
167            ConcatPlan::try_new(expression.dtype().clone(), chunks)?.into_plan(),
168        ))
169    }
170}
vortex_layout/plan/plans/concat.rs

vortex_layout/plan/plans/
concat.rs
