## Coverage Report

#### Created: 2024-10-13 08:39

    /Users/andrewlamb/Software/datafusion/datafusion/physical-plan/src/aggregates/group_values/bytes_view.rs

    Line
    Count
    Source (jump to first uncovered line)
    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 crate::aggregates::group_values::GroupValues;
    19
    use arrow_array::{Array, ArrayRef, RecordBatch};
    20
    use datafusion_expr::EmitTo;
    21
    use datafusion_physical_expr::binary_map::OutputType;
    22
    use datafusion_physical_expr_common::binary_view_map::ArrowBytesViewMap;
    23
    24
    /// A [`GroupValues`] storing single column of Utf8View/BinaryView values
    25
    ///
    26
    /// This specialization is significantly faster than using the more general
    27
    /// purpose `Row`s format
    28
    pub struct GroupValuesBytesView {
    29
        /// Map string/binary values to group index
    30
        map: ArrowBytesViewMap<usize>,
    31
        /// The total number of groups so far (used to assign group_index)
    32
        num_groups: usize,
    33
    }
    34
    35
    impl GroupValuesBytesView {
    36
    0
        pub fn new(output_type: OutputType) -> Self {
    37
    0
            Self {
    38
    0
                map: ArrowBytesViewMap::new(output_type),
    39
    0
                num_groups: 0,
    40
    0
            }
    41
    0
        }
    42
    }
    43
    44
    impl GroupValues for GroupValuesBytesView {
    45
    0
        fn intern(
    46
    0
            &mut self,
    47
    0
            cols: &[ArrayRef],
    48
    0
            groups: &mut Vec<usize>,
    49
    0
        ) -> datafusion_common::Result<()> {
    50
    0
            assert_eq!(cols.len(), 1);
    51
    52
            // look up / add entries in the table
    53
    0
            let arr = &cols[0];
    54
    0
    55
    0
            groups.clear();
    56
    0
            self.map.insert_if_new(
    57
    0
                arr,
    58
    0
                // called for each new group
    59
    0
                |_value| {
    60
    0
                    // assign new group index on each insert
    61
    0
                    let group_idx = self.num_groups;
    62
    0
                    self.num_groups += 1;
    63
    0
                    group_idx
    64
    0
                },
    65
    0
                // called for each group
    66
    0
                |group_idx| {
    67
    0
                    groups.push(group_idx);
    68
    0
                },
    69
    0
            );
    70
    0
    71
    0
            // ensure we assigned a group to for each row
    72
    0
            assert_eq!(groups.len(), arr.len());
    73
    0
            Ok(())
    74
    0
        }
    75
    76
    0
        fn size(&self) -> usize {
    77
    0
            self.map.size() + std::mem::size_of::<Self>()
    78
    0
        }
    79
    80
    0
        fn is_empty(&self) -> bool {
    81
    0
            self.num_groups == 0
    82
    0
        }
    83
    84
    0
        fn len(&self) -> usize {
    85
    0
            self.num_groups
    86
    0
        }
    87
    88
    0
        fn emit(&mut self, emit_to: EmitTo) -> datafusion_common::Result<Vec<ArrayRef>> {
    89
    0
            // Reset the map to default, and convert it into a single array
    90
    0
            let map_contents = self.map.take().into_state();
    91
    92
    0
            let group_values = match emit_to {
    93
                EmitTo::All => {
    94
    0
                    self.num_groups -= map_contents.len();
    95
    0
                    map_contents
    96
                }
    97
    0
                EmitTo::First(n) if n == self.len() => {
    98
    0
                    self.num_groups -= map_contents.len();
    99
    0
                    map_contents
    100
                }
    101
    0
                EmitTo::First(n) => {
    102
    0
                    // if we only wanted to take the first n, insert the rest back
    103
    0
                    // into the map we could potentially avoid this reallocation, at
    104
    0
                    // the expense of much more complex code.
    105
    0
                    // see https://github.com/apache/datafusion/issues/9195
    106
    0
                    let emit_group_values = map_contents.slice(0, n);
    107
    0
                    let remaining_group_values =
    108
    0
                        map_contents.slice(n, map_contents.len() - n);
    109
    0
    110
    0
                    self.num_groups = 0;
    111
    0
                    let mut group_indexes = vec![];
    112
    0
                    self.intern(&[remaining_group_values], &mut group_indexes)?;
    113
    114
                    // Verify that the group indexes were assigned in the correct order
    115
    0
                    assert_eq!(0, group_indexes[0]);
    116
    117
    0
                    emit_group_values
    118
                }
    119
            };
    120
    121
    0
            Ok(vec![group_values])
    122
    0
        }
    123
    124
    0
        fn clear_shrink(&mut self, _batch: &RecordBatch) {
    125
    0
            // in theory we could potentially avoid this reallocation and clear the
    126
    0
            // contents of the maps, but for now we just reset the map from the beginning
    127
    0
            self.map.take();
    128
    0
        }
    129
    }