datafusion_physical_plan/
statistics.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Statistics computation for physical plans.
19//!
20//! [`StatisticsArgs`] provides external context to
21//! [`ExecutionPlan::statistics_from_inputs`].
22
23use crate::ExecutionPlan;
24use datafusion_common::{
25    Result, Statistics, assert_eq_or_internal_err, assert_or_internal_err,
26};
27use std::cell::RefCell;
28use std::collections::HashMap;
29use std::rc::Rc;
30use std::sync::Arc;
31
32/// Per-call memoization cache for statistics computation.
33///
34/// Keyed by `(plan node pointer address, partition)`. Shared across
35/// a single statistics walk via [`StatisticsContext`].
36///
37/// The pointer-based key is safe within a single synchronous walk:
38/// all `Arc<dyn ExecutionPlan>` nodes are held by the plan tree for
39/// the duration of the walk, so addresses cannot be reused.
40#[derive(Debug, Default)]
41struct StatsCache(HashMap<(usize, Option<usize>), Arc<Statistics>>);
42
43impl StatsCache {
44    fn get(
45        &self,
46        plan: &dyn ExecutionPlan,
47        partition: Option<usize>,
48    ) -> Option<&Arc<Statistics>> {
49        let key = (
50            plan as *const dyn ExecutionPlan as *const () as usize,
51            partition,
52        );
53        self.0.get(&key)
54    }
55
56    fn insert(
57        &mut self,
58        plan: &dyn ExecutionPlan,
59        partition: Option<usize>,
60        stats: Arc<Statistics>,
61    ) {
62        let key = (
63            plan as *const dyn ExecutionPlan as *const () as usize,
64            partition,
65        );
66        self.0.insert(key, stats);
67    }
68}
69
70/// Arguments passed to [`ExecutionPlan::statistics_from_inputs`] carrying
71/// external information that operators can use when computing their
72/// statistics.
73#[derive(Debug, Default, Clone)]
74pub struct StatisticsArgs {
75    partition: Option<usize>,
76}
77
78impl StatisticsArgs {
79    /// Creates new statistics arguments.
80    ///
81    /// By default the partition is set to `None` (statistics should be computed
82    /// for the entire plan).
83    pub fn new() -> Self {
84        Default::default()
85    }
86
87    /// Set the partition to compute statistics
88    ///
89    /// * `None` means statistics should be computed for the entire plan.
90    /// * `Some(idx)` means statistics should be computed for the specified
91    ///   partition index.
92    pub fn set_partition(&mut self, partition: Option<usize>) {
93        self.partition = partition;
94    }
95
96    /// Builder Style API for [`Self::set_partition`]
97    pub fn with_partition(mut self, partition: Option<usize>) -> Self {
98        self.set_partition(partition);
99        self
100    }
101
102    /// Return the partition to compute statistics
103    pub fn partition(&self) -> Option<usize> {
104        self.partition
105    }
106}
107
108/// Directive returned by [`ExecutionPlan::child_stats_requests`] describing
109/// how the [`StatisticsContext`] should obtain each child's statistics.
110#[derive(Debug, Clone, Copy, PartialEq, Eq)]
111pub enum ChildStats {
112    /// Compute the child's statistics at this partition (`None` = overall).
113    At(Option<usize>),
114    /// Skip this child; the parent does not need its statistics. A placeholder
115    /// [`Statistics::new_unknown`] is supplied in its slot.
116    Skip,
117}
118
119/// Owns the bottom-up traversal and per-walk memoization cache for statistics
120/// computation. Call [`StatisticsContext::compute`] to walk a plan tree.
121pub struct StatisticsContext {
122    cache: Rc<RefCell<StatsCache>>,
123}
124
125impl Default for StatisticsContext {
126    fn default() -> Self {
127        Self::new()
128    }
129}
130
131impl StatisticsContext {
132    /// Creates a context with an empty cache.
133    pub fn new() -> Self {
134        Self {
135            cache: Rc::new(RefCell::new(StatsCache::default())),
136        }
137    }
138
139    /// Clears the memoization cache.
140    ///
141    /// The cache is keyed by raw plan-node pointers, which are only stable
142    /// while the current plan tree is alive. Reset between optimizer passes
143    /// (which rewrite the plan) when reusing one context across them, so stale
144    /// pointer keys cannot collide.
145    pub fn reset_cache(&self) {
146        self.cache.borrow_mut().0.clear();
147    }
148
149    /// Computes statistics for `plan`, resolving children first and passing
150    /// the results to [`ExecutionPlan::statistics_from_inputs`].
151    ///
152    /// When `args.partition()` is `Some(idx)`, `idx` is validated against the
153    /// plan's partition count.
154    pub fn compute(
155        &self,
156        plan: &dyn ExecutionPlan,
157        args: &StatisticsArgs,
158    ) -> Result<Arc<Statistics>> {
159        let partition = args.partition();
160
161        if let Some(idx) = partition {
162            let partition_count = plan.properties().partitioning.partition_count();
163            assert_or_internal_err!(
164                idx < partition_count,
165                "Invalid partition index: {}, the partition count is {}",
166                idx,
167                partition_count
168            );
169        }
170
171        if let Some(cached) = self.cache.borrow().get(plan, partition) {
172            return Ok(Arc::clone(cached));
173        }
174
175        let children = plan.children();
176        let requests = plan.child_stats_requests(partition);
177        assert_eq_or_internal_err!(
178            requests.len(),
179            children.len(),
180            "{} child_stats_requests returned {} entries for {} children",
181            plan.name(),
182            requests.len(),
183            children.len()
184        );
185        let child_stats = children
186            .iter()
187            .zip(requests)
188            .map(|(child, directive)| match directive {
189                ChildStats::At(p) => {
190                    self.compute(child.as_ref(), &StatisticsArgs::new().with_partition(p))
191                }
192                ChildStats::Skip => {
193                    Ok(Arc::new(Statistics::new_unknown(child.schema().as_ref())))
194                }
195            })
196            .collect::<Result<Vec<_>>>()?;
197
198        let result = plan.statistics_from_inputs(&child_stats, args)?;
199        self.cache
200            .borrow_mut()
201            .insert(plan, partition, Arc::clone(&result));
202        Ok(result)
203    }
204}
205
206#[cfg(all(test, feature = "test_utils"))]
207mod tests {
208    use super::*;
209    use crate::coalesce_partitions::CoalescePartitionsExec;
210    use crate::test::exec::StatisticsExec;
211    use arrow::datatypes::{DataType, Field, Schema};
212    use datafusion_common::{ColumnStatistics, stats::Precision};
213
214    fn make_stats_leaf(num_rows: usize) -> Arc<dyn ExecutionPlan> {
215        let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
216        let col_stats = vec![ColumnStatistics {
217            null_count: Precision::Exact(0),
218            max_value: Precision::Absent,
219            min_value: Precision::Absent,
220            sum_value: Precision::Absent,
221            distinct_count: Precision::Absent,
222            byte_size: Precision::Absent,
223        }];
224        Arc::new(StatisticsExec::new(
225            Statistics {
226                num_rows: Precision::Exact(num_rows),
227                total_byte_size: Precision::Absent,
228                column_statistics: col_stats,
229            },
230            schema,
231        ))
232    }
233
234    #[test]
235    fn coalesce_returns_overall_stats_for_any_partition() {
236        let leaf = make_stats_leaf(100);
237        let plan: Arc<dyn ExecutionPlan> = Arc::new(CoalescePartitionsExec::new(leaf));
238
239        let ctx = StatisticsContext::new();
240        let stats = ctx
241            .compute(
242                plan.as_ref(),
243                &StatisticsArgs::new().with_partition(Some(0)),
244            )
245            .unwrap();
246        assert_eq!(stats.num_rows, Precision::Exact(100));
247
248        let stats_none = ctx.compute(plan.as_ref(), &StatisticsArgs::new()).unwrap();
249        assert_eq!(stats_none.num_rows, Precision::Exact(100));
250    }
251
252    #[test]
253    fn context_caches_within_walk() {
254        let leaf = make_stats_leaf(42);
255        let ctx = StatisticsContext::new();
256        let args = StatisticsArgs::new();
257
258        let s1 = ctx.compute(leaf.as_ref(), &args).unwrap();
259        assert!(!ctx.cache.borrow().0.is_empty());
260
261        let s2 = ctx.compute(leaf.as_ref(), &args).unwrap();
262        assert!(Arc::ptr_eq(&s1, &s2));
263    }
264
265    #[test]
266    fn reset_cache_clears_entries() {
267        let leaf = make_stats_leaf(10);
268        let ctx = StatisticsContext::new();
269        let _ = ctx.compute(leaf.as_ref(), &StatisticsArgs::new()).unwrap();
270        assert!(!ctx.cache.borrow().0.is_empty());
271        ctx.reset_cache();
272        assert!(ctx.cache.borrow().0.is_empty());
273    }
274}
datafusion_physical_plan/statistics.rs

datafusion_physical_plan/
statistics.rs
