## Coverage Report

#### Created: 2025-11-22 06:00

next uncovered line (L), next uncovered region (R), next uncovered branch (B)

    /Users/andrewlamb/Software/datafusion/datafusion/physical-expr/src/statistics/stats_solver.rs

    Line
    Count
    Source
    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
    use std::sync::Arc;
    19
    20
    use crate::expressions::Literal;
    21
    use crate::intervals::cp_solver::PropagationResult;
    22
    use crate::physical_expr::PhysicalExpr;
    23
    use crate::utils::{build_dag, ExprTreeNode};
    24
    25
    use arrow::datatypes::{DataType, Schema};
    26
    use datafusion_common::{Result, ScalarValue};
    27
    use datafusion_expr::statistics::Distribution;
    28
    use datafusion_expr_common::interval_arithmetic::Interval;
    29
    30
    use petgraph::adj::DefaultIx;
    31
    use petgraph::prelude::Bfs;
    32
    use petgraph::stable_graph::{NodeIndex, StableGraph};
    33
    use petgraph::visit::DfsPostOrder;
    34
    use petgraph::Outgoing;
    35
    36
    /// This object implements a directed acyclic expression graph (DAEG) that
    37
    /// is used to compute statistics/distributions for expressions hierarchically.
    38
    #[derive(Clone, Debug)]
    39
    pub struct ExprStatisticsGraph {
    40
        graph: StableGraph<ExprStatisticsGraphNode, usize>,
    41
        root: NodeIndex,
    42
    }
    43
    44
    /// This is a node in the DAEG; it encapsulates a reference to the actual
    45
    /// [`PhysicalExpr`] as well as its statistics/distribution.
    46
    #[derive(Clone, Debug)]
    47
    pub struct ExprStatisticsGraphNode {
    48
        expr: Arc<dyn PhysicalExpr>,
    49
        dist: Distribution,
    50
    }
    51
    52
    impl ExprStatisticsGraphNode {
    53
        /// Constructs a new DAEG node based on the given interval with a
    54
        /// `Uniform` distribution.
    55
    0
        fn new_uniform(expr: Arc<dyn PhysicalExpr>, interval: Interval) -> Result<Self> {
    56
    0
            Distribution::new_uniform(interval)
    57
    0
                .map(|dist| ExprStatisticsGraphNode { expr, dist })
    58
    0
        }
    59
    60
        /// Constructs a new DAEG node with a `Bernoulli` distribution having an
    61
        /// unknown success probability.
    62
    0
        fn new_bernoulli(expr: Arc<dyn PhysicalExpr>) -> Result<Self> {
    63
    0
            Distribution::new_bernoulli(ScalarValue::Float64(None))
    64
    0
                .map(|dist| ExprStatisticsGraphNode { expr, dist })
    65
    0
        }
    66
    67
        /// Constructs a new DAEG node with a `Generic` distribution having no
    68
        /// definite summary statistics.
    69
    0
        fn new_generic(expr: Arc<dyn PhysicalExpr>, dt: &DataType) -> Result<Self> {
    70
    0
            let interval = Interval::make_unbounded(dt)?;
    71
    0
            let dist = Distribution::new_from_interval(interval)?;
    72
    0
            Ok(ExprStatisticsGraphNode { expr, dist })
    73
    0
        }
    74
    75
        /// Get the [`Distribution`] object representing the statistics of the
    76
        /// expression.
    77
    0
        pub fn distribution(&self) -> &Distribution {
    78
    0
            &self.dist
    79
    0
        }
    80
    81
        /// This function creates a DAEG node from DataFusion's [`ExprTreeNode`]
    82
        /// object. Literals are created with `Uniform` distributions with a
    83
        /// definite, singleton interval. Expressions with a `Boolean` data type
    84
        /// result in a`Bernoulli` distribution with an unknown success probability.
    85
        /// Any other expression starts with an `Unknown` distribution with an
    86
        /// indefinite range (i.e. `[-∞, ∞]`).
    87
    0
        pub fn make_node(node: &ExprTreeNode<NodeIndex>, schema: &Schema) -> Result<Self> {
    88
    0
            let expr = Arc::clone(&node.expr);
    89
    0
            if let Some(literal) = expr.as_any().downcast_ref::<Literal>() {
    90
    0
                let value = literal.value();
    91
    0
                Interval::try_new(value.clone(), value.clone())
    92
    0
                    .and_then(|interval| Self::new_uniform(expr, interval))
    93
            } else {
    94
    0
                expr.data_type(schema).and_then(|dt| {
    95
    0
                    if dt.eq(&DataType::Boolean) {
    96
    0
                        Self::new_bernoulli(expr)
    97
                    } else {
    98
    0
                        Self::new_generic(expr, &dt)
    99
                    }
    100
    0
                })
    101
            }
    102
    0
        }
    103
    }
    104
    105
    impl ExprStatisticsGraph {
    106
    0
        pub fn try_new(expr: Arc<dyn PhysicalExpr>, schema: &Schema) -> Result<Self> {
    107
            // Build the full graph:
    108
    0
            let (root, graph) = build_dag(expr, &|node| {
    109
    0
                ExprStatisticsGraphNode::make_node(node, schema)
    110
    0
            })?;
    111
    0
            Ok(Self { graph, root })
    112
    0
        }
    113
    114
        /// This function assigns given distributions to expressions in the DAEG.
    115
        /// The argument `assignments` associates indices of sought expressions
    116
        /// with their corresponding new distributions.
    117
    0
        pub fn assign_statistics(&mut self, assignments: &[(usize, Distribution)]) {
    118
    0
            for (index, stats) in assignments {
    119
    0
                let node_index = NodeIndex::from(*index as DefaultIx);
    120
    0
                self.graph[node_index].dist = stats.clone();
    121
    0
            }
    122
    0
        }
    123
    124
        /// Computes statistics/distributions for an expression via a bottom-up
    125
        /// traversal.
    126
    0
        pub fn evaluate_statistics(&mut self) -> Result<&Distribution> {
    127
    0
            let mut dfs = DfsPostOrder::new(&self.graph, self.root);
    128
    0
            while let Some(idx) = dfs.next(&self.graph) {
    129
    0
                let neighbors = self.graph.neighbors_directed(idx, Outgoing);
    130
    0
                let mut children_statistics = neighbors
    131
    0
                    .map(|child| self.graph[child].distribution())
    132
    0
                    .collect::<Vec<_>>();
    133
                // Note that all distributions are assumed to be independent.
    134
    0
                if !children_statistics.is_empty() {
    135
                    // Reverse to align with `PhysicalExpr`'s children:
    136
    0
                    children_statistics.reverse();
    137
    0
                    self.graph[idx].dist = self.graph[idx]
    138
    0
                        .expr
    139
    0
                        .evaluate_statistics(&children_statistics)?;
    140
    0
                }
    141
            }
    142
    0
            Ok(self.graph[self.root].distribution())
    143
    0
        }
    144
    145
        /// Runs a propagation mechanism in a top-down manner to update statistics
    146
        /// of leaf nodes.
    147
    0
        pub fn propagate_statistics(
    148
    0
            &mut self,
    149
    0
            given_stats: Distribution,
    150
    0
        ) -> Result<PropagationResult> {
    151
            // Adjust the root node with the given statistics:
    152
    0
            let root_range = self.graph[self.root].dist.range()?;
    153
    0
            let given_range = given_stats.range()?;
    154
    0
            if let Some(interval) = root_range.intersect(&given_range)? {
    155
    0
                if interval != root_range {
    156
                    // If the given statistics enable us to obtain a more precise
    157
                    // range for the root, update it:
    158
    0
                    let subset = root_range.contains(given_range)?;
    159
    0
                    self.graph[self.root].dist = if subset == Interval::TRUE {
    160
                        // Given statistics is strictly more informative, use it as is:
    161
    0
                        given_stats
    162
                    } else {
    163
                        // Intersecting ranges gives us a more precise range:
    164
    0
                        Distribution::new_from_interval(interval)?
    165
                    };
    166
    0
                }
    167
            } else {
    168
    0
                return Ok(PropagationResult::Infeasible);
    169
            }
    170
    171
    0
            let mut bfs = Bfs::new(&self.graph, self.root);
    172
    173
    0
            while let Some(node) = bfs.next(&self.graph) {
    174
    0
                let neighbors = self.graph.neighbors_directed(node, Outgoing);
    175
    0
                let mut children = neighbors.collect::<Vec<_>>();
    176
                // If the current expression is a leaf, its statistics is now final.
    177
                // So, just continue with the propagation procedure:
    178
    0
                if children.is_empty() {
    179
    0
                    continue;
    180
    0
                }
    181
                // Reverse to align with `PhysicalExpr`'s children:
    182
    0
                children.reverse();
    183
    0
                let children_stats = children
    184
    0
                    .iter()
    185
    0
                    .map(|child| self.graph[*child].distribution())
    186
    0
                    .collect::<Vec<_>>();
    187
    0
                let node_statistics = self.graph[node].distribution();
    188
    0
                let propagated_statistics = self.graph[node]
    189
    0
                    .expr
    190
    0
                    .propagate_statistics(node_statistics, &children_stats)?;
    191
    0
                if let Some(propagated_stats) = propagated_statistics {
    192
    0
                    for (child_idx, stats) in children.into_iter().zip(propagated_stats) {
    193
    0
                        self.graph[child_idx].dist = stats;
    194
    0
                    }
    195
                } else {
    196
                    // The constraint is infeasible, report:
    197
    0
                    return Ok(PropagationResult::Infeasible);
    198
                }
    199
            }
    200
    0
            Ok(PropagationResult::Success)
    201
    0
        }
    202
    }
    203
    204
    #[cfg(test)]
    205
    mod tests {
    206
        use std::sync::Arc;
    207
    208
        use crate::expressions::{binary, try_cast, Column};
    209
        use crate::intervals::cp_solver::PropagationResult;
    210
        use crate::statistics::stats_solver::ExprStatisticsGraph;
    211
    212
        use arrow::datatypes::{DataType, Field, Schema};
    213
        use datafusion_common::{Result, ScalarValue};
    214
        use datafusion_expr_common::interval_arithmetic::Interval;
    215
        use datafusion_expr_common::operator::Operator;
    216
        use datafusion_expr_common::statistics::Distribution;
    217
        use datafusion_expr_common::type_coercion::binary::BinaryTypeCoercer;
    218
        use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
    219
    220
    0
        pub fn binary_expr(
    221
    0
            left: Arc<dyn PhysicalExpr>,
    222
    0
            op: Operator,
    223
    0
            right: Arc<dyn PhysicalExpr>,
    224
    0
            schema: &Schema,
    225
    0
        ) -> Result<Arc<dyn PhysicalExpr>> {
    226
    0
            let left_type = left.data_type(schema)?;
    227
    0
            let right_type = right.data_type(schema)?;
    228
    0
            let binary_type_coercer = BinaryTypeCoercer::new(&left_type, &op, &right_type);
    229
    0
            let (lhs, rhs) = binary_type_coercer.get_input_types()?;
    230
    231
    0
            let left_expr = try_cast(left, schema, lhs)?;
    232
    0
            let right_expr = try_cast(right, schema, rhs)?;
    233
    0
            binary(left_expr, op, right_expr, schema)
    234
    0
        }
    235
    236
        #[test]
    237
    0
        fn test_stats_integration() -> Result<()> {
    238
    0
            let schema = &Schema::new(vec![
    239
    0
                Field::new("a", DataType::Float64, false),
    240
    0
                Field::new("b", DataType::Float64, false),
    241
    0
                Field::new("c", DataType::Float64, false),
    242
    0
                Field::new("d", DataType::Float64, false),
    243
    0
            ]);
    244
    245
    0
            let a = Arc::new(Column::new("a", 0)) as _;
    246
    0
            let b = Arc::new(Column::new("b", 1)) as _;
    247
    0
            let c = Arc::new(Column::new("c", 2)) as _;
    248
    0
            let d = Arc::new(Column::new("d", 3)) as _;
    249
    250
    0
            let left = binary_expr(a, Operator::Plus, b, schema)?;
    251
    0
            let right = binary_expr(c, Operator::Minus, d, schema)?;
    252
    0
            let expr = binary_expr(left, Operator::Eq, right, schema)?;
    253
    254
    0
            let mut graph = ExprStatisticsGraph::try_new(expr, schema)?;
    255
            // 2, 5 and 6 are BinaryExpr
    256
    0
            graph.assign_statistics(&[
    257
                (
    258
                    0usize,
    259
    0
                    Distribution::new_uniform(Interval::make(Some(0.), Some(1.))?)?,
    260
                ),
    261
                (
    262
                    1usize,
    263
    0
                    Distribution::new_uniform(Interval::make(Some(0.), Some(2.))?)?,
    264
                ),
    265
                (
    266
                    3usize,
    267
    0
                    Distribution::new_uniform(Interval::make(Some(1.), Some(3.))?)?,
    268
                ),
    269
                (
    270
                    4usize,
    271
    0
                    Distribution::new_uniform(Interval::make(Some(1.), Some(5.))?)?,
    272
                ),
    273
            ]);
    274
    0
            let ev_stats = graph.evaluate_statistics()?;
    275
    0
            assert_eq!(
    276
                ev_stats,
    277
    0
                &Distribution::new_bernoulli(ScalarValue::Float64(None))?
    278
            );
    279
    280
    0
            let one = ScalarValue::new_one(&DataType::Float64)?;
    281
    0
            assert_eq!(
    282
    0
                graph.propagate_statistics(Distribution::new_bernoulli(one)?)?,
    283
                PropagationResult::Success
    284
            );
    285
    0
            Ok(())
    286
    0
        }
    287
    }