Skip to main content

vortex_layout/plan/physical/plans/
segment_scan.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::ops::Range;
5
6use vortex_array::EmptyMetadata;
7use vortex_array::dtype::DType;
8use vortex_buffer::ByteBuffer;
9use vortex_error::VortexResult;
10use vortex_mask::Mask;
11use vortex_session::registry::CachedId;
12use vortex_session::registry::ReadContext;
13
14use crate::plan::physical::Plan;
15use crate::plan::physical::PlanChildren;
16use crate::plan::physical::PlanId;
17use crate::plan::physical::PlanParts;
18use crate::plan::physical::PlanVTable;
19use crate::plan::physical::check_child_count;
20use crate::plan::physical::pipeline::Chain;
21use crate::plan::physical::pipeline::Compiler;
22use crate::plan::physical::pipeline::Reach;
23use crate::segments::SegmentId;
24
25/// Reads one serialized array segment.
26#[derive(Clone, Debug)]
27pub struct SegmentScan;
28
29/// Data needed to read and decode a single segment.
30#[derive(Clone, Debug)]
31pub struct SegmentScanData {
32    segment_id: SegmentId,
33    array_ctx: ReadContext,
34    array_tree: Option<ByteBuffer>,
35}
36
37/// A plan that reads one serialized array segment.
38pub type SegmentScanPlan = Plan<SegmentScan>;
39
40impl SegmentScanPlan {
41    /// Creates a segment scan over `segment_id`.
42    pub fn new(
43        dtype: DType,
44        row_count: u64,
45        segment_id: SegmentId,
46        array_ctx: ReadContext,
47        array_tree: Option<ByteBuffer>,
48    ) -> Self {
49        PlanParts {
50            vtable: SegmentScan,
51            dtype,
52            row_count,
53            children: PlanChildren::default(),
54            data: SegmentScanData {
55                segment_id,
56                array_ctx,
57                array_tree,
58            },
59        }
60        .into_typed()
61    }
62
63    /// Returns the segment this plan reads.
64    pub fn segment_id(&self) -> SegmentId {
65        self.data().segment_id
66    }
67
68    /// Returns the read context for the serialized array.
69    pub fn array_ctx(&self) -> &ReadContext {
70        &self.data().array_ctx
71    }
72
73    /// Returns the serialized array encoding tree, when it is stored out of line.
74    pub fn array_tree(&self) -> Option<&ByteBuffer> {
75        self.data().array_tree.as_ref()
76    }
77}
78
79impl PlanVTable for SegmentScan {
80    type PlanData = SegmentScanData;
81    type Metadata = EmptyMetadata;
82
83    fn id(&self) -> PlanId {
84        static ID: CachedId = CachedId::new("vortex.plan.segment_scan");
85        *ID
86    }
87
88    fn metadata(_plan: &Plan<Self>) -> Option<Self::Metadata> {
89        // The segment ID and read context are not yet covered by a metadata codec.
90        None
91    }
92
93    fn with_children(
94        _plan: &Plan<Self>,
95        children: &PlanChildren,
96        _data: &mut Self::PlanData,
97    ) -> VortexResult<()> {
98        check_child_count("SegmentScan", children, 0)?;
99        Ok(())
100    }
101
102    fn compile(
103        plan: &Plan<Self>,
104        rows: Range<u64>,
105        mask: &Mask,
106        compiler: &mut Compiler<'_>,
107    ) -> VortexResult<Option<Chain>> {
108        // A scan keeps the selected rows itself, as every plan produces only those.
109        compiler.scan(plan, rows, Some(mask.clone()))
110    }
111
112    fn reach(
113        plan: &Plan<Self>,
114        rows: Range<u64>,
115        at: &Reach,
116        visit: &mut dyn FnMut(SegmentId, Range<u64>),
117    ) -> VortexResult<()> {
118        visit(plan.segment_id(), at.root(&rows));
119        Ok(())
120    }
121}