vortex_layout/plan/physical/plans/
concat.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::borrow::Cow;
5use std::ops::Range;
6use std::sync::Arc;
7
8use vortex_array::EmptyMetadata;
9use vortex_array::dtype::DType;
10use vortex_error::VortexResult;
11use vortex_error::vortex_bail;
12use vortex_mask::Mask;
13use vortex_session::registry::CachedId;
14
15use crate::plan::physical::Eval;
16use crate::plan::physical::EvalPlan;
17use crate::plan::physical::Plan;
18use crate::plan::physical::PlanChildren;
19use crate::plan::physical::PlanId;
20use crate::plan::physical::PlanParts;
21use crate::plan::physical::PlanRef;
22use crate::plan::physical::PlanVTable;
23use crate::plan::physical::operators::ConcatSource;
24use crate::plan::physical::optimizer::PlanParentReduceRule;
25use crate::plan::physical::pipeline::Chain;
26use crate::plan::physical::pipeline::Compiler;
27use crate::plan::physical::pipeline::Reach;
28use crate::plan::physical::pipeline::overlapping;
29use crate::plan::physical::selection::selected_slice;
30use crate::segments::SegmentId;
31
32/// Concatenates its children row-wise.
33#[derive(Clone, Debug)]
34pub struct Concat;
35
36/// Row offsets of each concatenated child.
37#[derive(Clone, Debug)]
38pub struct ConcatData {
39    row_offsets: Arc<[u64]>,
40}
41
42/// A plan that concatenates its children row-wise.
43pub type ConcatPlan = Plan<Concat>;
44
45impl ConcatPlan {
46    /// Creates a concatenation from potentially unresolved children without validation.
47    ///
48    /// # Safety
49    ///
50    /// Every child must have `dtype`; `row_offsets` must contain the cumulative row offset of
51    /// every child; and the sum of all child row counts must equal `row_count` without overflow.
52    pub(crate) unsafe fn from_children_unchecked(
53        dtype: DType,
54        row_count: u64,
55        row_offsets: Arc<[u64]>,
56        children: PlanChildren,
57    ) -> Self {
58        PlanParts {
59            vtable: Concat,
60            dtype,
61            row_count,
62            children,
63            data: ConcatData { row_offsets },
64        }
65        .into_typed()
66    }
67
68    /// Creates a concatenation over `children`.
69    ///
70    /// Every child must produce `dtype`, and the row domain is the sum of the child row counts.
71    pub fn try_new(dtype: DType, children: Vec<PlanRef>) -> VortexResult<Self> {
72        let mut row_offsets = Vec::with_capacity(children.len());
73        let mut row_count = 0u64;
74        for child in &children {
75            if child.dtype() != &dtype {
76                vortex_bail!(
77                    "Concat child dtype {} does not match {dtype}",
78                    child.dtype()
79                );
80            }
81            row_offsets.push(row_count);
82            row_count = row_count
83                .checked_add(child.row_count())
84                .ok_or_else(|| vortex_error::vortex_err!("Concat row count overflow"))?;
85        }
86        // SAFETY: Child dtypes and the checked cumulative row metadata were validated above.
87        Ok(unsafe {
88            Self::from_children_unchecked(dtype, row_count, row_offsets.into(), children.into())
89        })
90    }
91
92    /// Returns the first row of each child within this plan's row domain.
93    pub fn row_offsets(&self) -> &[u64] {
94        &self.data().row_offsets
95    }
96}
97
98impl PlanVTable for Concat {
99    type PlanData = ConcatData;
100    type Metadata = EmptyMetadata;
101
102    fn id(&self) -> PlanId {
103        static ID: CachedId = CachedId::new("vortex.plan.concat");
104        *ID
105    }
106
107    fn metadata(_plan: &Plan<Self>) -> Option<Self::Metadata> {
108        // Row offsets are derived from the children, so nothing needs storing.
109        Some(EmptyMetadata)
110    }
111
112    fn with_children(
113        plan: &Plan<Self>,
114        children: &PlanChildren,
115        data: &mut Self::PlanData,
116    ) -> VortexResult<()> {
117        if children.len() != plan.children().len() {
118            vortex_bail!(
119                "Concat expects {} children but got {}",
120                plan.children().len(),
121                children.len()
122            );
123        }
124
125        let mut row_offsets = Vec::with_capacity(children.len());
126        let mut row_count = 0u64;
127        for child in children.iter() {
128            let child = child?;
129            if child.dtype() != plan.dtype() {
130                vortex_bail!(
131                    "Concat child dtype {} does not match {}",
132                    child.dtype(),
133                    plan.dtype()
134                );
135            }
136            row_offsets.push(row_count);
137            row_count = row_count
138                .checked_add(child.row_count())
139                .ok_or_else(|| vortex_error::vortex_err!("Concat row count overflow"))?;
140        }
141        if row_count != plan.row_count() {
142            vortex_bail!(
143                "Concat children have {row_count} rows but the plan has {}",
144                plan.row_count()
145            );
146        }
147        data.row_offsets = row_offsets.into();
148        Ok(())
149    }
150
151    fn child_name(_plan: &Plan<Self>, index: usize) -> Cow<'_, str> {
152        Cow::Owned(format!("chunks[{index}]"))
153    }
154
155    fn compile(
156        plan: &Plan<Self>,
157        rows: Range<u64>,
158        mask: &Mask,
159        compiler: &mut Compiler<'_>,
160    ) -> VortexResult<Option<Chain>> {
161        let capacity = if mask.all_true() {
162            overlapping(plan, &rows).size_hint().1.unwrap_or(0)
163        } else {
164            0
165        };
166        let mut chains = Vec::with_capacity(capacity);
167        for (index, start, local) in overlapping(plan, &rows) {
168            let Some((span, local_mask)) = selected_slice(
169                mask,
170                usize::try_from(local.start - rows.start)?
171                    ..usize::try_from(local.end - rows.start)?,
172            ) else {
173                continue;
174            };
175            let chunk = plan.child_required(index)?;
176            if let Some(chain) = compiler.compile(
177                &chunk,
178                rows.start + span.start as u64 - start..rows.start + span.end as u64 - start,
179                &local_mask,
180            )? {
181                chains.push(chain);
182            }
183        }
184        Ok(match chains.len() {
185            0 => None,
186            // Rows inside one chunk are that chunk's rows: no concatenation is built.
187            1 => chains.pop(),
188            count => Some(compiler.join(chains, ConcatSource::new(count))),
189        })
190    }
191
192    fn reach(
193        plan: &Plan<Self>,
194        rows: Range<u64>,
195        at: &Reach,
196        visit: &mut dyn FnMut(SegmentId, Range<u64>),
197    ) -> VortexResult<()> {
198        for (index, start, local) in overlapping(plan, &rows) {
199            plan.child_required(index)?.reach(
200                local.start - start..local.end - start,
201                &at.shift(start),
202                visit,
203            )?;
204        }
205        Ok(())
206    }
207}
208
209/// Pushes an expression into every chunk of a [`Concat`].
210#[derive(Debug)]
211pub(crate) struct ExpressionConcatRule;
212
213impl PlanParentReduceRule<Concat> for ExpressionConcatRule {
214    type Parent = Eval;
215
216    fn reduce_parent(
217        &self,
218        child: &Plan<Concat>,
219        parent: &Plan<Eval>,
220        _child_idx: usize,
221    ) -> VortexResult<Option<PlanRef>> {
222        let expression = parent.expression();
223        let chunks = child
224            .children()
225            .iter()
226            .map(|chunk| Ok(EvalPlan::try_new(expression.clone(), chunk?)?.into_plan()))
227            .collect::<VortexResult<Vec<_>>>()?;
228        Ok(Some(
229            ConcatPlan::try_new(expression.dtype().clone(), chunks)?.into_plan(),
230        ))
231    }
232}
vortex_layout/plan/physical/plans/concat.rs

vortex_layout/plan/physical/plans/
concat.rs
