## 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-common/src/utils.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::physical_expr::PhysicalExpr;
    21
    use crate::tree_node::ExprContext;
    22
    23
    use arrow::array::{make_array, Array, ArrayRef, BooleanArray, MutableArrayData};
    24
    use arrow::compute::{and_kleene, is_not_null, SlicesIterator};
    25
    use arrow::record_batch::RecordBatch;
    26
    use datafusion_common::Result;
    27
    use datafusion_expr_common::sort_properties::ExprProperties;
    28
    29
    /// Represents a [`PhysicalExpr`] node with associated properties (order and
    30
    /// range) in a context where properties are tracked.
    31
    pub type ExprPropertiesNode = ExprContext<ExprProperties>;
    32
    33
    impl ExprPropertiesNode {
    34
        /// Constructs a new `ExprPropertiesNode` with unknown properties for a
    35
        /// given physical expression. This node initializes with default properties
    36
        /// and recursively applies this to all child expressions.
    37
    0
        pub fn new_unknown(expr: Arc<dyn PhysicalExpr>) -> Self {
    38
    0
            let children = expr
    39
    0
                .children()
    40
    0
                .into_iter()
    41
    0
                .cloned()
    42
    0
                .map(Self::new_unknown)
    43
    0
                .collect();
    44
    0
            Self {
    45
    0
                expr,
    46
    0
                data: ExprProperties::new_unknown(),
    47
    0
                children,
    48
    0
            }
    49
    0
        }
    50
    }
    51
    52
    /// Scatter `truthy` array by boolean mask. When the mask evaluates `true`, next values of `truthy`
    53
    /// are taken, when the mask evaluates `false` values null values are filled.
    54
    ///
    55
    /// # Arguments
    56
    /// * `mask` - Boolean values used to determine where to put the `truthy` values
    57
    /// * `truthy` - All values of this array are to scatter according to `mask` into final result.
    58
    0
    pub fn scatter(mask: &BooleanArray, truthy: &dyn Array) -> Result<ArrayRef> {
    59
    0
        let truthy = truthy.to_data();
    60
    61
        // update the mask so that any null values become false
    62
        // (SlicesIterator doesn't respect nulls)
    63
    0
        let mask = and_kleene(mask, &is_not_null(mask)?)?;
    64
    65
    0
        let mut mutable = MutableArrayData::new(vec![&truthy], true, mask.len());
    66
    67
        // the SlicesIterator slices only the true values. So the gaps left by this iterator we need to
    68
        // fill with falsy values
    69
    70
        // keep track of how much is filled
    71
    0
        let mut filled = 0;
    72
        // keep track of current position we have in truthy array
    73
    0
        let mut true_pos = 0;
    74
    75
    0
        SlicesIterator::new(&mask).for_each(|(start, end)| {
    76
            // the gap needs to be filled with nulls
    77
    0
            if start > filled {
    78
    0
                mutable.extend_nulls(start - filled);
    79
    0
            }
    80
            // fill with truthy values
    81
    0
            let len = end - start;
    82
    0
            mutable.extend(0, true_pos, true_pos + len);
    83
    0
            true_pos += len;
    84
    0
            filled = end;
    85
    0
        });
    86
        // the remaining part is falsy
    87
    0
        if filled < mask.len() {
    88
    0
            mutable.extend_nulls(mask.len() - filled);
    89
    0
        }
    90
    91
    0
        let data = mutable.freeze();
    92
    0
        Ok(make_array(data))
    93
    0
    }
    94
    95
    /// Evaluates expressions against a record batch.
    96
    /// This will convert the resulting ColumnarValues to ArrayRefs,
    97
    /// duplicating any ScalarValues that may have been returned,
    98
    /// and validating that the returned arrays all have the same
    99
    /// number of rows as the input batch.
    100
    #[inline]
    101
    pub fn evaluate_expressions_to_arrays<'a>(
    102
        exprs: impl IntoIterator<Item = &'a Arc<dyn PhysicalExpr>>,
    103
        batch: &RecordBatch,
    104
    ) -> Result<Vec<ArrayRef>> {
    105
        let num_rows = batch.num_rows();
    106
        exprs
    107
            .into_iter()
    108
            .map(|e| {
    109
                e.evaluate(batch)
    110
                    .and_then(|col| col.into_array_of_size(num_rows))
    111
            })
    112
            .collect::<Result<Vec<ArrayRef>>>()
    113
    }
    114
    115
    #[cfg(test)]
    116
    mod tests {
    117
        use std::sync::Arc;
    118
    119
        use arrow::array::Int32Array;
    120
    121
        use datafusion_common::cast::{as_boolean_array, as_int32_array};
    122
    123
        use super::*;
    124
    125
        #[test]
    126
        fn scatter_int() -> Result<()> {
    127
            let truthy = Arc::new(Int32Array::from(vec![1, 10, 11, 100]));
    128
            let mask = BooleanArray::from(vec![true, true, false, false, true]);
    129
    130
            // the output array is expected to be the same length as the mask array
    131
            let expected =
    132
                Int32Array::from_iter(vec![Some(1), Some(10), None, None, Some(11)]);
    133
            let result = scatter(&mask, truthy.as_ref())?;
    134
            let result = as_int32_array(&result)?;
    135
    136
            assert_eq!(&expected, result);
    137
            Ok(())
    138
        }
    139
    140
        #[test]
    141
        fn scatter_int_end_with_false() -> Result<()> {
    142
            let truthy = Arc::new(Int32Array::from(vec![1, 10, 11, 100]));
    143
            let mask = BooleanArray::from(vec![true, false, true, false, false, false]);
    144
    145
            // output should be same length as mask
    146
            let expected =
    147
                Int32Array::from_iter(vec![Some(1), None, Some(10), None, None, None]);
    148
            let result = scatter(&mask, truthy.as_ref())?;
    149
            let result = as_int32_array(&result)?;
    150
    151
            assert_eq!(&expected, result);
    152
            Ok(())
    153
        }
    154
    155
        #[test]
    156
        fn scatter_with_null_mask() -> Result<()> {
    157
            let truthy = Arc::new(Int32Array::from(vec![1, 10, 11]));
    158
            let mask: BooleanArray = vec![Some(false), None, Some(true), Some(true), None]
    159
                .into_iter()
    160
                .collect();
    161
    162
            // output should treat nulls as though they are false
    163
            let expected = Int32Array::from_iter(vec![None, None, Some(1), Some(10), None]);
    164
            let result = scatter(&mask, truthy.as_ref())?;
    165
            let result = as_int32_array(&result)?;
    166
    167
            assert_eq!(&expected, result);
    168
            Ok(())
    169
        }
    170
    171
        #[test]
    172
        fn scatter_boolean() -> Result<()> {
    173
            let truthy = Arc::new(BooleanArray::from(vec![false, false, false, true]));
    174
            let mask = BooleanArray::from(vec![true, true, false, false, true]);
    175
    176
            // the output array is expected to be the same length as the mask array
    177
            let expected = BooleanArray::from_iter(vec![
    178
                Some(false),
    179
                Some(false),
    180
                None,
    181
                None,
    182
                Some(false),
    183
            ]);
    184
            let result = scatter(&mask, truthy.as_ref())?;
    185
            let result = as_boolean_array(&result)?;
    186
    187
            assert_eq!(&expected, result);
    188
            Ok(())
    189
        }
    190
    }