vortex_layout/plan/physical/plans/
legacy_stats.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_array::expr::stats::Stat;
11use vortex_error::VortexResult;
12use vortex_error::vortex_bail;
13use vortex_error::vortex_err;
14use vortex_mask::Mask;
15use vortex_session::registry::CachedId;
16
17use crate::layouts::zoned::legacy_stats_to_aggregates;
18use crate::plan::physical::EvalPlan;
19use crate::plan::physical::Plan;
20use crate::plan::physical::PlanChildren;
21use crate::plan::physical::PlanId;
22use crate::plan::physical::PlanParts;
23use crate::plan::physical::PlanRef;
24use crate::plan::physical::PlanVTable;
25use crate::plan::physical::Zoned;
26use crate::plan::physical::ZonedPlan;
27use crate::plan::physical::check_child_count;
28use crate::plan::physical::optimizer::PlanParentReduceRule;
29use crate::plan::physical::optimizer::PlanReduceRule;
30use crate::plan::physical::pipeline::Chain;
31use crate::plan::physical::pipeline::Compiler;
32use crate::plan::physical::pipeline::Reach;
33use crate::segments::SegmentId;
34
35const DATA: usize = 0;
36const ZONES: usize = 1;
37
38/// Data alongside a legacy `vortex.stats` zone table, whose fields are named by [`Stat`] rather
39/// than by aggregate function.
40///
41/// The optimizer rewrites it into a [`Zoned`] plan over the zone table converted to the aggregate
42/// schema, or into its data alone when a [`Zoned`] plan already summarises the same data. Compiled
43/// unoptimized, it reads its data and prunes nothing.
44#[derive(Clone, Debug)]
45pub struct LegacyStats;
46
47/// Legacy-stats-plan-specific data.
48#[derive(Clone, Debug)]
49pub struct LegacyStatsData {
50    zone_len: u64,
51    /// The stats the zone table stores.
52    stats: Arc<[Stat]>,
53}
54
55/// A plan that pairs data with a legacy stats zone table.
56pub type LegacyStatsPlan = Plan<LegacyStats>;
57
58impl LegacyStatsPlan {
59    pub(crate) fn from_children(
60        dtype: DType,
61        row_count: u64,
62        children: PlanChildren,
63        zone_len: u64,
64        stats: Arc<[Stat]>,
65    ) -> Self {
66        PlanParts {
67            vtable: LegacyStats,
68            dtype,
69            row_count,
70            children,
71            data: LegacyStatsData { zone_len, stats },
72        }
73        .into_typed()
74    }
75
76    /// Creates a legacy stats plan over `data` summarised by `zones` of `zone_len` rows each,
77    /// whose fields hold `stats`.
78    pub fn new(data: PlanRef, zones: PlanRef, zone_len: u64, stats: Arc<[Stat]>) -> Self {
79        let dtype = data.dtype().clone();
80        let row_count = data.row_count();
81        Self::from_children(dtype, row_count, vec![data, zones].into(), zone_len, stats)
82    }
83
84    /// Returns the number of rows each zone summarises.
85    pub fn zone_len(&self) -> u64 {
86        self.data().zone_len
87    }
88
89    /// Returns the stats the zone table stores.
90    pub fn stats(&self) -> &Arc<[Stat]> {
91        &self.data().stats
92    }
93
94    /// Returns the plan producing the summarised data.
95    pub fn data_plan(&self) -> VortexResult<PlanRef> {
96        self.child_required(DATA)
97    }
98
99    /// Returns the plan producing the legacy zone table.
100    pub fn zones_plan(&self) -> VortexResult<PlanRef> {
101        self.child_required(ZONES)
102    }
103}
104
105impl PlanVTable for LegacyStats {
106    type PlanData = LegacyStatsData;
107    type Metadata = EmptyMetadata;
108
109    fn id(&self) -> PlanId {
110        static ID: CachedId = CachedId::new("vortex.plan.legacy_stats");
111        *ID
112    }
113
114    fn metadata(_plan: &Plan<Self>) -> Option<Self::Metadata> {
115        None
116    }
117
118    fn with_children(
119        plan: &Plan<Self>,
120        children: &PlanChildren,
121        _data: &mut Self::PlanData,
122    ) -> VortexResult<()> {
123        check_child_count("LegacyStats", children, 2)?;
124        let data = children
125            .get(DATA)?
126            .ok_or_else(|| vortex_err!("LegacyStats data child is absent"))?;
127        if data.dtype() != plan.dtype() || data.row_count() != plan.row_count() {
128            vortex_bail!("LegacyStats data child shape does not match the plan output");
129        }
130        Ok(())
131    }
132
133    fn child_name(_plan: &Plan<Self>, index: usize) -> Cow<'_, str> {
134        match index {
135            DATA => Cow::Borrowed("data"),
136            ZONES => Cow::Borrowed("zones"),
137            _ => Cow::Owned(format!("child[{index}]")),
138        }
139    }
140
141    fn compile(
142        plan: &Plan<Self>,
143        rows: Range<u64>,
144        mask: &Mask,
145        compiler: &mut Compiler<'_>,
146    ) -> VortexResult<Option<Chain>> {
147        compiler.compile(&plan.data_plan()?, rows, mask)
148    }
149
150    fn reach(
151        plan: &Plan<Self>,
152        rows: Range<u64>,
153        at: &Reach,
154        visit: &mut dyn FnMut(SegmentId, Range<u64>),
155    ) -> VortexResult<()> {
156        plan.data_plan()?.reach(rows, at, visit)
157    }
158}
159
160/// Rewrites a legacy stats plan into a [`Zoned`] plan over its zone table converted to the
161/// aggregate schema.
162///
163/// A plan whose data is already [`Zoned`] keeps only that data: legacy stats are used only when no
164/// aggregate zone map summarises the same rows. A zero zone length, as older files may record, or
165/// a table with no stat an aggregate can stand for, leaves only the data too.
166#[derive(Debug)]
167pub(crate) struct LegacyStatsRule;
168
169impl PlanReduceRule<LegacyStats> for LegacyStatsRule {
170    fn reduce(&self, plan: &Plan<LegacyStats>) -> VortexResult<Option<PlanRef>> {
171        let data = plan.data_plan()?;
172        if plan.zone_len() == 0 || data.is::<Zoned>() {
173            return Ok(Some(data));
174        }
175        let zones = plan.zones_plan()?;
176        let (aggregate_fns, conversion) =
177            legacy_stats_to_aggregates(plan.dtype(), zones.dtype(), plan.stats());
178        if aggregate_fns.is_empty() {
179            return Ok(Some(data));
180        }
181        let zones = EvalPlan::try_new(conversion.bind(zones.dtype())?, zones)?.into_plan();
182        let zoned = ZonedPlan::new(data, zones, plan.zone_len(), aggregate_fns);
183        Ok(Some(zoned.into_plan()))
184    }
185}
186
187/// Drops legacy stats beneath a [`Zoned`] plan: its aggregate zone map already summarises the
188/// same rows.
189#[derive(Debug)]
190pub(crate) struct ZonedLegacyStatsRule;
191
192impl PlanParentReduceRule<LegacyStats> for ZonedLegacyStatsRule {
193    type Parent = Zoned;
194
195    fn reduce_parent(
196        &self,
197        child: &LegacyStatsPlan,
198        parent: &ZonedPlan,
199        child_idx: usize,
200    ) -> VortexResult<Option<PlanRef>> {
201        if child_idx != DATA || parent.is_pruning() {
202            return Ok(None);
203        }
204        let children = vec![child.data_plan()?, parent.zones_plan()?];
205        Ok(Some(parent.clone().into_plan().with_children(children)?))
206    }
207}
vortex_layout/plan/physical/plans/legacy_stats.rs

vortex_layout/plan/physical/plans/
legacy_stats.rs
