vortex_layout/layouts/row_idx/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4mod expr;
5
6use std::fmt::Display;
7use std::fmt::Formatter;
8use std::ops::BitAnd;
9use std::ops::Range;
10use std::sync::Arc;
11use std::sync::OnceLock;
12
13use Nullability::NonNullable;
14pub use expr::*;
15use futures::FutureExt;
16use futures::future::BoxFuture;
17use vortex_array::ArrayRef;
18use vortex_array::Canonical;
19use vortex_array::IntoArray;
20use vortex_array::MaskFuture;
21use vortex_array::VortexSessionExecute;
22use vortex_array::dtype::DType;
23use vortex_array::dtype::FieldMask;
24use vortex_array::dtype::FieldName;
25use vortex_array::dtype::Nullability;
26use vortex_array::dtype::PType;
27use vortex_array::expr::BoundExpression;
28use vortex_array::expr::ExactBoundExpr;
29use vortex_array::expr::transform::BoundPartitionedExpr;
30use vortex_array::expr::transform::partition_bound;
31use vortex_array::expr::traversal::NodeExt;
32use vortex_array::expr::traversal::Transformed;
33use vortex_array::expr::traversal::TraversalOrder;
34use vortex_array::scalar::PValue;
35use vortex_error::VortexExpect;
36use vortex_error::VortexResult;
37use vortex_mask::Mask;
38use vortex_sequence::Sequence;
39use vortex_sequence::SequenceArray;
40use vortex_session::VortexSession;
41use vortex_utils::aliases::dash_map::DashMap;
42
43use crate::ArrayFuture;
44use crate::LayoutReader;
45use crate::RowSplits;
46use crate::SplitRange;
47use crate::layouts::partitioned::BoundPartitionedExprEval;
48
49pub struct RowIdxLayoutReader {
50    name: Arc<str>,
51    row_offset: u64,
52    child: Arc<dyn LayoutReader>,
53    partition_cache: DashMap<ExactBoundExpr, Arc<OnceLock<Partitioning>>>,
54    session: VortexSession,
55}
56
57impl RowIdxLayoutReader {
58    pub fn new(row_offset: u64, child: Arc<dyn LayoutReader>, session: VortexSession) -> Self {
59        Self {
60            name: Arc::clone(child.name()),
61            row_offset,
62            child,
63            partition_cache: DashMap::with_hasher(Default::default()),
64            session,
65        }
66    }
67
68    fn partition_expr(&self, expr: &BoundExpression) -> VortexResult<Partitioning> {
69        let key = ExactBoundExpr(expr.clone());
70
71        // Check cache first with read-only lock.
72        if let Some(entry) = self.partition_cache.get(&key)
73            && let Some(partitioning) = entry.value().get()
74        {
75            return Ok(partitioning.clone());
76        }
77
78        let result = self.compute_partitioning(expr)?;
79
80        self.partition_cache
81            .entry(key)
82            .or_insert_with(|| Arc::new(OnceLock::new()))
83            .get_or_init(|| result.clone());
84
85        Ok(result)
86    }
87
88    fn compute_partitioning(&self, expr: &BoundExpression) -> VortexResult<Partitioning> {
89        // Partition the expression into row idx and child expressions.
90        let mut partitioned = partition_bound(expr.clone(), |expr: &BoundExpression| {
91            if expr
92                .as_scalar()
93                .is_some_and(|scalar_fn| scalar_fn.is::<RowIdx>())
94            {
95                vec![Partition::RowIdx]
96            } else if expr.is_root() {
97                vec![Partition::Child]
98            } else {
99                vec![]
100            }
101        })?;
102
103        // If there's only a single partition, we can directly return the expression.
104        if partitioned.partitions.len() == 1 {
105            return Ok(match &partitioned.partition_annotations[0] {
106                Partition::RowIdx => Partitioning::RowIdx(replace_row_idx(expr.clone())?),
107                Partition::Child => Partitioning::Child(expr.clone()),
108            });
109        }
110
111        // Replace the row_idx expression with the root expression in the row_idx partition.
112        let partitions = partitioned
113            .partitions
114            .iter()
115            .cloned()
116            .map(replace_row_idx)
117            .collect::<VortexResult<Vec<_>>>()?
118            .into_boxed_slice();
119        partitioned.replace_partitions(partitions)?;
120
121        Ok(Partitioning::Partitioned(Arc::new(partitioned)))
122    }
123}
124
125#[derive(Clone)]
126enum Partitioning {
127    // An expression that only references the row index (e.g., `row_idx == 5`).
128    RowIdx(BoundExpression),
129    // An expression that does not reference the row index.
130    Child(BoundExpression),
131    // Contains both the RowIdx and Child expressions, (e.g., `row_idx < child.some_field`).
132    Partitioned(Arc<BoundPartitionedExpr<Partition>>),
133}
134
135#[derive(Clone, PartialEq, Eq, Hash)]
136enum Partition {
137    RowIdx,
138    Child,
139}
140
141impl Partition {
142    pub fn name(&self) -> &str {
143        match self {
144            Partition::RowIdx => "row_idx",
145            Partition::Child => "child",
146        }
147    }
148}
149
150impl From<Partition> for FieldName {
151    fn from(value: Partition) -> Self {
152        FieldName::from(value.name())
153    }
154}
155
156impl Display for Partition {
157    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
158        write!(f, "{}", self.name())
159    }
160}
161
162impl LayoutReader for RowIdxLayoutReader {
163    fn name(&self) -> &Arc<str> {
164        &self.name
165    }
166
167    fn dtype(&self) -> &DType {
168        self.child.dtype()
169    }
170
171    fn row_count(&self) -> u64 {
172        self.child.row_count()
173    }
174
175    fn register_splits(
176        &self,
177        field_mask: &[FieldMask],
178        split_range: &SplitRange,
179        splits: &mut RowSplits,
180    ) -> VortexResult<()> {
181        self.child.register_splits(field_mask, split_range, splits)
182    }
183
184    fn pruning_evaluation(
185        &self,
186        row_range: &Range<u64>,
187        expr: &BoundExpression,
188        mask: Mask,
189    ) -> VortexResult<MaskFuture> {
190        Ok(match &self.partition_expr(expr)? {
191            Partitioning::RowIdx(expr) => row_idx_mask_future(
192                self.row_offset,
193                row_range,
194                expr.clone(),
195                MaskFuture::ready(mask),
196                self.session.clone(),
197            ),
198            Partitioning::Child(expr) => self.child.pruning_evaluation(row_range, expr, mask)?,
199            Partitioning::Partitioned(..) => MaskFuture::ready(mask),
200        })
201    }
202
203    fn filter_evaluation(
204        &self,
205        row_range: &Range<u64>,
206        expr: &BoundExpression,
207        mask: MaskFuture,
208    ) -> VortexResult<MaskFuture> {
209        match &self.partition_expr(expr)? {
210            // Since this is run during pruning, we skip re-evaluating the row index expression
211            // during the filter evaluation.
212            Partitioning::RowIdx(_) => Ok(mask),
213            Partitioning::Child(expr) => self.child.filter_evaluation(row_range, expr, mask),
214            Partitioning::Partitioned(p) => Arc::clone(p).into_mask_future(
215                mask,
216                |annotation, expr, mask| match annotation {
217                    Partition::RowIdx => Ok(row_idx_mask_future(
218                        self.row_offset,
219                        row_range,
220                        expr.clone(),
221                        mask,
222                        self.session.clone(),
223                    )),
224                    Partition::Child => self.child.filter_evaluation(row_range, expr, mask),
225                },
226                |annotation, expr, mask| match annotation {
227                    Partition::RowIdx => Ok(row_idx_array_future(
228                        self.row_offset,
229                        row_range,
230                        expr.clone(),
231                        mask,
232                        self.session.clone(),
233                    )),
234                    Partition::Child => self.child.projection_evaluation(row_range, expr, mask),
235                },
236                self.session.clone(),
237            ),
238        }
239    }
240
241    fn projection_evaluation(
242        &self,
243        row_range: &Range<u64>,
244        expr: &BoundExpression,
245        mask: MaskFuture,
246    ) -> VortexResult<BoxFuture<'static, VortexResult<ArrayRef>>> {
247        match &self.partition_expr(expr)? {
248            Partitioning::RowIdx(expr) => Ok(row_idx_array_future(
249                self.row_offset,
250                row_range,
251                expr.clone(),
252                mask,
253                self.session.clone(),
254            )),
255            Partitioning::Child(expr) => self.child.projection_evaluation(row_range, expr, mask),
256            Partitioning::Partitioned(p) => {
257                Arc::clone(p).into_array_future(mask, |annotation, expr, mask| match annotation {
258                    Partition::RowIdx => Ok(row_idx_array_future(
259                        self.row_offset,
260                        row_range,
261                        expr.clone(),
262                        mask,
263                        self.session.clone(),
264                    )),
265                    Partition::Child => self.child.projection_evaluation(row_range, expr, mask),
266                })
267            }
268        }
269    }
270
271    fn as_any(&self) -> &dyn std::any::Any {
272        self
273    }
274}
275
276fn replace_row_idx(expr: BoundExpression) -> VortexResult<BoundExpression> {
277    Ok(expr
278        .transform_down(|node| {
279            if node
280                .as_scalar()
281                .is_some_and(|scalar_fn| scalar_fn.is::<RowIdx>())
282            {
283                Ok(Transformed {
284                    value: BoundExpression::new_root(row_idx_dtype()),
285                    changed: true,
286                    order: TraversalOrder::Skip,
287                })
288            } else {
289                Ok(Transformed::no(node))
290            }
291        })?
292        .into_inner())
293}
294
295fn row_idx_dtype() -> DType {
296    DType::Primitive(PType::U64, NonNullable)
297}
298
299// Returns a SequenceArray representing the row indices for the given row range,
300fn idx_array(row_offset: u64, row_range: &Range<u64>) -> SequenceArray {
301    Sequence::try_new(
302        PValue::U64(row_offset + row_range.start),
303        PValue::U64(1),
304        PType::U64,
305        NonNullable,
306        usize::try_from(row_range.end - row_range.start)
307            .vortex_expect("Row range length must fit in usize"),
308    )
309    .vortex_expect("Failed to create row index array")
310}
311
312fn row_idx_mask_future(
313    row_offset: u64,
314    row_range: &Range<u64>,
315    expr: BoundExpression,
316    mask: MaskFuture,
317    session: VortexSession,
318) -> MaskFuture {
319    let row_range = row_range.clone();
320    MaskFuture::new(mask.len(), async move {
321        let array = idx_array(row_offset, &row_range).into_array();
322
323        let mut ctx = session.create_execution_ctx();
324        let result_mask = array
325            .apply_bound(&expr)?
326            .null_as_false()
327            .execute(&mut ctx)?;
328
329        Ok(result_mask.bitand(&mask.await?))
330    })
331}
332
333fn row_idx_array_future(
334    row_offset: u64,
335    row_range: &Range<u64>,
336    expr: BoundExpression,
337    mask: MaskFuture,
338    session: VortexSession,
339) -> ArrayFuture {
340    let row_range = row_range.clone();
341    async move {
342        let array = idx_array(row_offset, &row_range).into_array();
343        let filtered = array.filter(mask.await?)?;
344        let mut ctx = session.create_execution_ctx();
345        let array = filtered.execute::<Canonical>(&mut ctx)?.into_array();
346        array.apply_bound(&expr)
347    }
348    .boxed()
349}
350
351#[cfg(test)]
352mod tests {
353    use std::sync::Arc;
354
355    use vortex_array::ArrayContext;
356    use vortex_array::IntoArray as _;
357    use vortex_array::MaskFuture;
358    use vortex_array::VortexSessionExecute;
359    use vortex_array::arrays::BoolArray;
360    use vortex_array::assert_arrays_eq;
361    use vortex_array::expr::eq;
362    use vortex_array::expr::gt;
363    use vortex_array::expr::lit;
364    use vortex_array::expr::or;
365    use vortex_array::expr::root;
366    use vortex_buffer::buffer;
367    use vortex_io::runtime::single::block_on;
368    use vortex_io::session::RuntimeSessionExt;
369
370    use crate::LayoutReader;
371    use crate::LayoutStrategy;
372    use crate::layouts::flat::writer::FlatLayoutStrategy;
373    use crate::layouts::row_idx::RowIdxLayoutReader;
374    use crate::layouts::row_idx::row_idx;
375    use crate::segments::TestSegments;
376    use crate::sequence::SequenceId;
377    use crate::sequence::SequentialArrayStreamExt;
378    use crate::test::new_session;
379
380    #[test]
381    fn flat_expr_no_row_id() {
382        block_on(|handle| async {
383            let session = new_session().with_handle(handle);
384            let mut ctx = session.create_execution_ctx();
385            let array_ctx = ArrayContext::empty();
386            let segments = Arc::new(TestSegments::default());
387            let (ptr, eof) = SequenceId::root().split();
388            let array = buffer![1..=5].into_array();
389            let layout = FlatLayoutStrategy::default()
390                .write_stream(
391                    array_ctx.into(),
392                    Arc::<TestSegments>::clone(&segments),
393                    array.to_array_stream().sequenced(ptr),
394                    eof,
395                    &session,
396                )
397                .await
398                .unwrap();
399
400            let expr = eq(root(), lit(3i32));
401            let reader = RowIdxLayoutReader::new(
402                0,
403                layout
404                    .new_reader("".into(), segments, &session, &Default::default())
405                    .unwrap(),
406                session.clone(),
407            );
408            let expr = expr.bind(reader.dtype()).unwrap();
409            let result = reader
410                .projection_evaluation(
411                    &(0..layout.row_count()),
412                    &expr,
413                    MaskFuture::new_true(layout.row_count().try_into().unwrap()),
414                )
415                .unwrap()
416                .await
417                .unwrap();
418
419            assert_arrays_eq!(
420                result,
421                BoolArray::from_iter([false, false, true, false, false]),
422                &mut ctx
423            );
424        })
425    }
426
427    #[test]
428    fn flat_expr_row_id() {
429        block_on(|handle| async {
430            let session = new_session().with_handle(handle);
431            let mut ctx = session.create_execution_ctx();
432            let array_ctx = ArrayContext::empty();
433            let segments = Arc::new(TestSegments::default());
434            let (ptr, eof) = SequenceId::root().split();
435            let array = buffer![1..=5].into_array();
436            let layout = FlatLayoutStrategy::default()
437                .write_stream(
438                    array_ctx.into(),
439                    Arc::<TestSegments>::clone(&segments),
440                    array.to_array_stream().sequenced(ptr),
441                    eof,
442                    &session,
443                )
444                .await
445                .unwrap();
446
447            let expr = gt(row_idx(), lit(3u64));
448            let reader = RowIdxLayoutReader::new(
449                0,
450                layout
451                    .new_reader("".into(), segments, &session, &Default::default())
452                    .unwrap(),
453                session.clone(),
454            );
455            let expr = expr.bind(reader.dtype()).unwrap();
456            let result = reader
457                .projection_evaluation(
458                    &(0..layout.row_count()),
459                    &expr,
460                    MaskFuture::new_true(layout.row_count().try_into().unwrap()),
461                )
462                .unwrap()
463                .await
464                .unwrap();
465
466            assert_arrays_eq!(
467                result,
468                BoolArray::from_iter([false, false, false, false, true]),
469                &mut ctx
470            );
471        })
472    }
473
474    #[test]
475    fn flat_expr_or() {
476        block_on(|handle| async {
477            let session = new_session().with_handle(handle);
478            let mut ctx = session.create_execution_ctx();
479            let array_ctx = ArrayContext::empty();
480            let segments = Arc::new(TestSegments::default());
481            let (ptr, eof) = SequenceId::root().split();
482            let array = buffer![1..=5].into_array();
483            let layout = FlatLayoutStrategy::default()
484                .write_stream(
485                    array_ctx.into(),
486                    Arc::<TestSegments>::clone(&segments),
487                    array.to_array_stream().sequenced(ptr),
488                    eof,
489                    &session,
490                )
491                .await
492                .unwrap();
493
494            let expr = or(
495                eq(root(), lit(3i32)),
496                or(gt(row_idx(), lit(3u64)), eq(root(), lit(1i32))),
497            );
498
499            let reader = RowIdxLayoutReader::new(
500                0,
501                layout
502                    .new_reader("".into(), segments, &session, &Default::default())
503                    .unwrap(),
504                session.clone(),
505            );
506            let expr = expr.bind(reader.dtype()).unwrap();
507            let result = reader
508                .projection_evaluation(
509                    &(0..layout.row_count()),
510                    &expr,
511                    MaskFuture::new_true(layout.row_count().try_into().unwrap()),
512                )
513                .unwrap()
514                .await
515                .unwrap();
516
517            assert_arrays_eq!(
518                result,
519                BoolArray::from_iter([true, false, true, false, true]),
520                &mut ctx
521            );
522        })
523    }
524}
vortex_layout/layouts/row_idx/mod.rs

vortex_layout/layouts/row_idx/
mod.rs
