## 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/functions-aggregate-common/src/merge_arrays.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 arrow::compute::SortOptions;
    19
    use datafusion_common::utils::compare_rows;
    20
    use datafusion_common::{exec_err, ScalarValue};
    21
    use std::cmp::Ordering;
    22
    use std::collections::{BinaryHeap, VecDeque};
    23
    24
    /// This is a wrapper struct to be able to correctly merge `ARRAY_AGG` data from
    25
    /// multiple partitions using `BinaryHeap`. When used inside `BinaryHeap`, this
    26
    /// struct returns smallest `CustomElement`, where smallest is determined by
    27
    /// `ordering` values (`Vec<ScalarValue>`) according to `sort_options`.
    28
    #[derive(Debug, PartialEq, Eq)]
    29
    struct CustomElement<'a> {
    30
        /// Stores the partition this entry came from
    31
        branch_idx: usize,
    32
        /// Values to merge
    33
        value: ScalarValue,
    34
        // Comparison "key"
    35
        ordering: Vec<ScalarValue>,
    36
        /// Options defining the ordering semantics
    37
        sort_options: &'a [SortOptions],
    38
    }
    39
    40
    impl<'a> CustomElement<'a> {
    41
    0
        fn new(
    42
    0
            branch_idx: usize,
    43
    0
            value: ScalarValue,
    44
    0
            ordering: Vec<ScalarValue>,
    45
    0
            sort_options: &'a [SortOptions],
    46
    0
        ) -> Self {
    47
    0
            Self {
    48
    0
                branch_idx,
    49
    0
                value,
    50
    0
                ordering,
    51
    0
                sort_options,
    52
    0
            }
    53
    0
        }
    54
    55
    0
        fn ordering(
    56
    0
            &self,
    57
    0
            current: &[ScalarValue],
    58
    0
            target: &[ScalarValue],
    59
    0
        ) -> datafusion_common::Result<Ordering> {
    60
            // Calculate ordering according to `sort_options`
    61
    0
            compare_rows(current, target, self.sort_options)
    62
    0
        }
    63
    }
    64
    65
    // Overwrite ordering implementation such that
    66
    // - `self.ordering` values are used for comparison,
    67
    // - When used inside `BinaryHeap` it is a min-heap.
    68
    impl Ord for CustomElement<'_> {
    69
    0
        fn cmp(&self, other: &Self) -> Ordering {
    70
            // TODO Ord/PartialOrd is not consistent with PartialEq; PartialOrd contract is violated
    71
            // Compares according to custom ordering
    72
    0
            self.ordering(&self.ordering, &other.ordering)
    73
                // Convert max heap to min heap
    74
    0
                .map(|ordering| ordering.reverse())
    75
                // This function return error, when `self.ordering` and `other.ordering`
    76
                // have different types (such as one is `ScalarValue::Int64`, other is `ScalarValue::Float32`)
    77
                // Here this case won't happen, because data from each partition will have same type
    78
    0
                .unwrap()
    79
    0
        }
    80
    }
    81
    82
    impl PartialOrd for CustomElement<'_> {
    83
    0
        fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
    84
    0
            Some(self.cmp(other))
    85
    0
        }
    86
    }
    87
    88
    /// This functions merges `values` array (`&[Vec<ScalarValue>]`) into single array `Vec<ScalarValue>`
    89
    /// Merging done according to ordering values stored inside `ordering_values` (`&[Vec<Vec<ScalarValue>>]`)
    90
    /// Inner `Vec<ScalarValue>` in the `ordering_values` can be thought as ordering information for
    91
    /// each `ScalarValue` in the `values` array.
    92
    /// Desired ordering specified by `sort_options` argument (Should have same size with inner `Vec<ScalarValue>`
    93
    /// of the `ordering_values` array).
    94
    ///
    95
    /// As an example
    96
    /// values can be \[
    97
    ///      \[1, 2, 3, 4, 5\],
    98
    ///      \[1, 2, 3, 4\],
    99
    ///      \[1, 2, 3, 4, 5, 6\],
    100
    /// \]
    101
    /// In this case we will be merging three arrays (doesn't have to be same size)
    102
    /// and produce a merged array with size 15 (sum of 5+4+6)
    103
    /// Merging will be done according to ordering at `ordering_values` vector.
    104
    /// As an example `ordering_values` can be [
    105
    ///      \[(1, a), (2, b), (3, b), (4, a), (5, b) \],
    106
    ///      \[(1, a), (2, b), (3, b), (4, a) \],
    107
    ///      \[(1, b), (2, c), (3, d), (4, e), (5, a), (6, b) \],
    108
    /// ]
    109
    /// For each ScalarValue in the `values` we have a corresponding `Vec<ScalarValue>` (like timestamp of it)
    110
    /// for the example above `sort_options` will have size two, that defines ordering requirement of the merge.
    111
    /// Inner `Vec<ScalarValue>`s of the `ordering_values` will be compared according `sort_options` (Their sizes should match)
    112
    0
    pub fn merge_ordered_arrays(
    113
    0
        // We will merge values into single `Vec<ScalarValue>`.
    114
    0
        values: &mut [VecDeque<ScalarValue>],
    115
    0
        // `values` will be merged according to `ordering_values`.
    116
    0
        // Inner `Vec<ScalarValue>` can be thought as ordering information for the
    117
    0
        // each `ScalarValue` in the values`.
    118
    0
        ordering_values: &mut [VecDeque<Vec<ScalarValue>>],
    119
    0
        // Defines according to which ordering comparisons should be done.
    120
    0
        sort_options: &[SortOptions],
    121
    0
    ) -> datafusion_common::Result<(Vec<ScalarValue>, Vec<Vec<ScalarValue>>)> {
    122
        // Keep track of the most recent data of each branch, in a binary heap data structure.
    123
    0
        let mut heap = BinaryHeap::<CustomElement>::new();
    124
    125
    0
        if values.len() != ordering_values.len() {
    126
    0
            return exec_err!(
    127
    0
                "Expects values and ordering_values to have same size but got {} and {}",
    128
    0
                values.len(),
    129
    0
                ordering_values.len()
    130
            );
    131
    0
        }
    132
    0
        if let Some((idx, (values, ordering_values))) = values
    133
    0
            .iter()
    134
    0
            .zip(ordering_values.iter())
    135
    0
            .enumerate()
    136
    0
            .find(|(_, (vals, ordering_vals))| vals.len() != ordering_vals.len())
    137
        {
    138
    0
            return exec_err!(
    139
    0
                "Expects values elements and ordering_values elements to have same size but got {} and {} at index {}",
    140
    0
                values.len(),
    141
    0
                ordering_values.len(),
    142
                idx
    143
            );
    144
    0
        }
    145
    0
        let n_branch = values.len();
    146
    0
        let mut merged_values = vec![];
    147
    0
        let mut merged_orderings = vec![];
    148
        // Continue iterating the loop until consuming data of all branches.
    149
        loop {
    150
    0
            let minimum = if let Some(minimum) = heap.pop() {
    151
    0
                minimum
    152
            } else {
    153
                // Heap is empty, fill it with the next entries from each branch.
    154
    0
                for branch_idx in 0..n_branch {
    155
    0
                    if let Some(orderings) = ordering_values[branch_idx].pop_front() {
    156
    0
                        // Their size should be same, we can safely .unwrap here.
    157
    0
                        let value = values[branch_idx].pop_front().unwrap();
    158
    0
                        // Push the next element to the heap:
    159
    0
                        heap.push(CustomElement::new(
    160
    0
                            branch_idx,
    161
    0
                            value,
    162
    0
                            orderings,
    163
    0
                            sort_options,
    164
    0
                        ));
    165
    0
                    }
    166
                    // If None, we consumed this branch, skip it.
    167
                }
    168
    169
                // Now we have filled the heap, get the largest entry (this will be
    170
                // the next element in merge).
    171
    0
                if let Some(minimum) = heap.pop() {
    172
    0
                    minimum
    173
                } else {
    174
                    // Heap is empty, this means that all indices are same with
    175
                    // `end_indices`. We have consumed all of the branches, merge
    176
                    // is completed, exit from the loop:
    177
    0
                    break;
    178
                }
    179
            };
    180
            let CustomElement {
    181
    0
                branch_idx,
    182
    0
                value,
    183
    0
                ordering,
    184
                ..
    185
    0
            } = minimum;
    186
            // Add minimum value in the heap to the result
    187
    0
            merged_values.push(value);
    188
    0
            merged_orderings.push(ordering);
    189
    190
            // If there is an available entry, push next entry in the most
    191
            // recently consumed branch to the heap.
    192
    0
            if let Some(orderings) = ordering_values[branch_idx].pop_front() {
    193
    0
                // Their size should be same, we can safely .unwrap here.
    194
    0
                let value = values[branch_idx].pop_front().unwrap();
    195
    0
                // Push the next element to the heap:
    196
    0
                heap.push(CustomElement::new(
    197
    0
                    branch_idx,
    198
    0
                    value,
    199
    0
                    orderings,
    200
    0
                    sort_options,
    201
    0
                ));
    202
    0
            }
    203
        }
    204
    205
    0
        Ok((merged_values, merged_orderings))
    206
    0
    }
    207
    208
    #[cfg(test)]
    209
    mod tests {
    210
        use super::*;
    211
    212
        use std::collections::VecDeque;
    213
        use std::sync::Arc;
    214
    215
        use arrow::array::{ArrayRef, Int64Array};
    216
    217
        use datafusion_common::utils::get_row_at_idx;
    218
        use datafusion_common::{Result, ScalarValue};
    219
    220
        #[test]
    221
        fn test_merge_asc() -> Result<()> {
    222
            let lhs_arrays: Vec<ArrayRef> = vec![
    223
                Arc::new(Int64Array::from(vec![0, 0, 1, 1, 2])),
    224
                Arc::new(Int64Array::from(vec![0, 1, 2, 3, 4])),
    225
            ];
    226
            let n_row = lhs_arrays[0].len();
    227
            let lhs_orderings = (0..n_row)
    228
                .map(|idx| get_row_at_idx(&lhs_arrays, idx))
    229
                .collect::<Result<VecDeque<_>>>()?;
    230
    231
            let rhs_arrays: Vec<ArrayRef> = vec![
    232
                Arc::new(Int64Array::from(vec![0, 0, 1, 1, 2])),
    233
                Arc::new(Int64Array::from(vec![0, 1, 2, 3, 4])),
    234
            ];
    235
            let n_row = rhs_arrays[0].len();
    236
            let rhs_orderings = (0..n_row)
    237
                .map(|idx| get_row_at_idx(&rhs_arrays, idx))
    238
                .collect::<Result<VecDeque<_>>>()?;
    239
            let sort_options = vec![
    240
                SortOptions {
    241
                    descending: false,
    242
                    nulls_first: false,
    243
                },
    244
                SortOptions {
    245
                    descending: false,
    246
                    nulls_first: false,
    247
                },
    248
            ];
    249
    250
            let lhs_vals_arr = Arc::new(Int64Array::from(vec![0, 1, 2, 3, 4])) as ArrayRef;
    251
            let lhs_vals = (0..lhs_vals_arr.len())
    252
                .map(|idx| ScalarValue::try_from_array(&lhs_vals_arr, idx))
    253
                .collect::<Result<VecDeque<_>>>()?;
    254
    255
            let rhs_vals_arr = Arc::new(Int64Array::from(vec![0, 1, 2, 3, 4])) as ArrayRef;
    256
            let rhs_vals = (0..rhs_vals_arr.len())
    257
                .map(|idx| ScalarValue::try_from_array(&rhs_vals_arr, idx))
    258
                .collect::<Result<VecDeque<_>>>()?;
    259
            let expected =
    260
                Arc::new(Int64Array::from(vec![0, 0, 1, 1, 2, 2, 3, 3, 4, 4])) as ArrayRef;
    261
            let expected_ts = vec![
    262
                Arc::new(Int64Array::from(vec![0, 0, 0, 0, 1, 1, 1, 1, 2, 2])) as ArrayRef,
    263
                Arc::new(Int64Array::from(vec![0, 0, 1, 1, 2, 2, 3, 3, 4, 4])) as ArrayRef,
    264
            ];
    265
    266
            let (merged_vals, merged_ts) = merge_ordered_arrays(
    267
                &mut [lhs_vals, rhs_vals],
    268
                &mut [lhs_orderings, rhs_orderings],
    269
                &sort_options,
    270
            )?;
    271
            let merged_vals = ScalarValue::iter_to_array(merged_vals.into_iter())?;
    272
            let merged_ts = (0..merged_ts[0].len())
    273
                .map(|col_idx| {
    274
                    ScalarValue::iter_to_array(
    275
                        (0..merged_ts.len())
    276
                            .map(|row_idx| merged_ts[row_idx][col_idx].clone()),
    277
                    )
    278
                })
    279
                .collect::<Result<Vec<_>>>()?;
    280
    281
            assert_eq!(&merged_vals, &expected);
    282
            assert_eq!(&merged_ts, &expected_ts);
    283
    284
            Ok(())
    285
        }
    286
    287
        #[test]
    288
        fn test_merge_desc() -> Result<()> {
    289
            let lhs_arrays: Vec<ArrayRef> = vec![
    290
                Arc::new(Int64Array::from(vec![2, 1, 1, 0, 0])),
    291
                Arc::new(Int64Array::from(vec![4, 3, 2, 1, 0])),
    292
            ];
    293
            let n_row = lhs_arrays[0].len();
    294
            let lhs_orderings = (0..n_row)
    295
                .map(|idx| get_row_at_idx(&lhs_arrays, idx))
    296
                .collect::<Result<VecDeque<_>>>()?;
    297
    298
            let rhs_arrays: Vec<ArrayRef> = vec![
    299
                Arc::new(Int64Array::from(vec![2, 1, 1, 0, 0])),
    300
                Arc::new(Int64Array::from(vec![4, 3, 2, 1, 0])),
    301
            ];
    302
            let n_row = rhs_arrays[0].len();
    303
            let rhs_orderings = (0..n_row)
    304
                .map(|idx| get_row_at_idx(&rhs_arrays, idx))
    305
                .collect::<Result<VecDeque<_>>>()?;
    306
            let sort_options = vec![
    307
                SortOptions {
    308
                    descending: true,
    309
                    nulls_first: false,
    310
                },
    311
                SortOptions {
    312
                    descending: true,
    313
                    nulls_first: false,
    314
                },
    315
            ];
    316
    317
            // Values (which will be merged) doesn't have to be ordered.
    318
            let lhs_vals_arr = Arc::new(Int64Array::from(vec![0, 1, 2, 1, 2])) as ArrayRef;
    319
            let lhs_vals = (0..lhs_vals_arr.len())
    320
                .map(|idx| ScalarValue::try_from_array(&lhs_vals_arr, idx))
    321
                .collect::<Result<VecDeque<_>>>()?;
    322
    323
            let rhs_vals_arr = Arc::new(Int64Array::from(vec![0, 1, 2, 1, 2])) as ArrayRef;
    324
            let rhs_vals = (0..rhs_vals_arr.len())
    325
                .map(|idx| ScalarValue::try_from_array(&rhs_vals_arr, idx))
    326
                .collect::<Result<VecDeque<_>>>()?;
    327
            let expected =
    328
                Arc::new(Int64Array::from(vec![0, 0, 1, 1, 2, 2, 1, 1, 2, 2])) as ArrayRef;
    329
            let expected_ts = vec![
    330
                Arc::new(Int64Array::from(vec![2, 2, 1, 1, 1, 1, 0, 0, 0, 0])) as ArrayRef,
    331
                Arc::new(Int64Array::from(vec![4, 4, 3, 3, 2, 2, 1, 1, 0, 0])) as ArrayRef,
    332
            ];
    333
            let (merged_vals, merged_ts) = merge_ordered_arrays(
    334
                &mut [lhs_vals, rhs_vals],
    335
                &mut [lhs_orderings, rhs_orderings],
    336
                &sort_options,
    337
            )?;
    338
            let merged_vals = ScalarValue::iter_to_array(merged_vals.into_iter())?;
    339
            let merged_ts = (0..merged_ts[0].len())
    340
                .map(|col_idx| {
    341
                    ScalarValue::iter_to_array(
    342
                        (0..merged_ts.len())
    343
                            .map(|row_idx| merged_ts[row_idx][col_idx].clone()),
    344
                    )
    345
                })
    346
                .collect::<Result<Vec<_>>>()?;
    347
    348
            assert_eq!(&merged_vals, &expected);
    349
            assert_eq!(&merged_ts, &expected_ts);
    350
            Ok(())
    351
        }
    352
    }