## Coverage Report

#### Created: 2025-11-17 14:14

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

    /Users/andrewlamb/Software/arrow-rs/arrow-select/src/window.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
    //! Defines windowing functions, like `shift`ing
    19
    20
    use crate::concat::concat;
    21
    use arrow_array::{Array, ArrayRef, make_array, new_null_array};
    22
    use arrow_schema::ArrowError;
    23
    use num_traits::abs;
    24
    25
    /// Shifts array by defined number of items (to left or right)
    26
    /// A positive value for `offset` shifts the array to the right
    27
    /// a negative value shifts the array to the left.
    28
    /// # Examples
    29
    /// ```
    30
    /// # use arrow_array::Int32Array;
    31
    /// # use arrow_select::window::shift;
    32
    ///
    33
    /// let a: Int32Array = vec![Some(1), None, Some(4)].into();
    34
    ///
    35
    /// // shift array 1 element to the right
    36
    /// let res = shift(&a, 1).unwrap();
    37
    /// let expected: Int32Array = vec![None, Some(1), None].into();
    38
    /// assert_eq!(res.as_ref(), &expected);
    39
    ///
    40
    /// // shift array 1 element to the left
    41
    /// let res = shift(&a, -1).unwrap();
    42
    /// let expected: Int32Array = vec![None, Some(4), None].into();
    43
    /// assert_eq!(res.as_ref(), &expected);
    44
    ///
    45
    /// // shift array 0 element, although not recommended
    46
    /// let res = shift(&a, 0).unwrap();
    47
    /// let expected: Int32Array = vec![Some(1), None, Some(4)].into();
    48
    /// assert_eq!(res.as_ref(), &expected);
    49
    ///
    50
    /// // shift array 3 element to the right
    51
    /// let res = shift(&a, 3).unwrap();
    52
    /// let expected: Int32Array = vec![None, None, None].into();
    53
    /// assert_eq!(res.as_ref(), &expected);
    54
    /// ```
    55
    12
    pub fn shift(array: &dyn Array, offset: i64) -> Result<ArrayRef, ArrowError> {
    56
    12
        let value_len = array.len() as i64;
    57
    12
        if offset == 0 {
    58
    1
            Ok(make_array(array.to_data()))
    59
    11
        } else if offset == i64::MIN || abs10(offset) >= value_len {
    60
    5
            Ok(new_null_array(array.data_type(), array.len()))
    61
        } else {
    62
            // Concatenate both arrays, add nulls after if shift > 0 else before
    63
    6
            if offset > 0 {
    64
    3
                let length = array.len() - offset as usize;
    65
    3
                let slice = array.slice(0, length);
    66
    67
                // Generate array with remaining `null` items
    68
    3
                let null_arr = new_null_array(array.data_type(), offset as usize);
    69
    3
                concat(&[null_arr.as_ref(), slice.as_ref()])
    70
            } else {
    71
    3
                let offset = -offset as usize;
    72
    3
                let length = array.len() - offset;
    73
    3
                let slice = array.slice(offset, length);
    74
    75
                // Generate array with remaining `null` items
    76
    3
                let null_arr = new_null_array(array.data_type(), offset);
    77
    3
                concat(&[slice.as_ref(), null_arr.as_ref()])
    78
            }
    79
        }
    80
    12
    }
    81
    82
    #[cfg(test)]
    83
    mod tests {
    84
        use super::*;
    85
        use arrow_array::{Float64Array, Int32Array, Int32DictionaryArray};
    86
    87
        #[test]
    88
    1
        fn test_shift_neg() {
    89
    1
            let a: Int32Array = vec![Some(1), None, Some(4)].into();
    90
    1
            let res = shift(&a, -1).unwrap();
    91
    1
            let expected: Int32Array = vec![None, Some(4), None].into();
    92
    1
            assert_eq!(res.as_ref(), &expected);
    93
    1
        }
    94
    95
        #[test]
    96
    1
        fn test_shift_pos() {
    97
    1
            let a: Int32Array = vec![Some(1), None, Some(4)].into();
    98
    1
            let res = shift(&a, 1).unwrap();
    99
    1
            let expected: Int32Array = vec![None, Some(1), None].into();
    100
    1
            assert_eq!(res.as_ref(), &expected);
    101
    1
        }
    102
    103
        #[test]
    104
    1
        fn test_shift_neg_float64() {
    105
    1
            let a: Float64Array = vec![Some(1.), None, Some(4.)].into();
    106
    1
            let res = shift(&a, -1).unwrap();
    107
    1
            let expected: Float64Array = vec![None, Some(4.), None].into();
    108
    1
            assert_eq!(res.as_ref(), &expected);
    109
    1
        }
    110
    111
        #[test]
    112
    1
        fn test_shift_pos_float64() {
    113
    1
            let a: Float64Array = vec![Some(1.), None, Some(4.)].into();
    114
    1
            let res = shift(&a, 1).unwrap();
    115
    1
            let expected: Float64Array = vec![None, Some(1.), None].into();
    116
    1
            assert_eq!(res.as_ref(), &expected);
    117
    1
        }
    118
    119
        #[test]
    120
    1
        fn test_shift_neg_int32_dict() {
    121
    1
            let a: Int32DictionaryArray = [Some("alpha"), None, Some("beta"), Some("alpha")]
    122
    1
                .iter()
    123
    1
                .copied()
    124
    1
                .collect();
    125
    1
            let res = shift(&a, -1).unwrap();
    126
    1
            let expected: Int32DictionaryArray = [None, Some("beta"), Some("alpha"), None]
    127
    1
                .iter()
    128
    1
                .copied()
    129
    1
                .collect();
    130
    1
            assert_eq!(res.as_ref(), &expected);
    131
    1
        }
    132
    133
        #[test]
    134
    1
        fn test_shift_pos_int32_dict() {
    135
    1
            let a: Int32DictionaryArray = [Some("alpha"), None, Some("beta"), Some("alpha")]
    136
    1
                .iter()
    137
    1
                .copied()
    138
    1
                .collect();
    139
    1
            let res = shift(&a, 1).unwrap();
    140
    1
            let expected: Int32DictionaryArray = [None, Some("alpha"), None, Some("beta")]
    141
    1
                .iter()
    142
    1
                .copied()
    143
    1
                .collect();
    144
    1
            assert_eq!(res.as_ref(), &expected);
    145
    1
        }
    146
    147
        #[test]
    148
    1
        fn test_shift_nil() {
    149
    1
            let a: Int32Array = vec![Some(1), None, Some(4)].into();
    150
    1
            let res = shift(&a, 0).unwrap();
    151
    1
            let expected: Int32Array = vec![Some(1), None, Some(4)].into();
    152
    1
            assert_eq!(res.as_ref(), &expected);
    153
    1
        }
    154
    155
        #[test]
    156
    1
        fn test_shift_boundary_pos() {
    157
    1
            let a: Int32Array = vec![Some(1), None, Some(4)].into();
    158
    1
            let res = shift(&a, 3).unwrap();
    159
    1
            let expected: Int32Array = vec![None, None, None].into();
    160
    1
            assert_eq!(res.as_ref(), &expected);
    161
    1
        }
    162
    163
        #[test]
    164
    1
        fn test_shift_boundary_neg() {
    165
    1
            let a: Int32Array = vec![Some(1), None, Some(4)].into();
    166
    1
            let res = shift(&a, -3).unwrap();
    167
    1
            let expected: Int32Array = vec![None, None, None].into();
    168
    1
            assert_eq!(res.as_ref(), &expected);
    169
    1
        }
    170
    171
        #[test]
    172
    1
        fn test_shift_boundary_neg_min() {
    173
    1
            let a: Int32Array = vec![Some(1), None, Some(4)].into();
    174
    1
            let res = shift(&a, i64::MIN).unwrap();
    175
    1
            let expected: Int32Array = vec![None, None, None].into();
    176
    1
            assert_eq!(res.as_ref(), &expected);
    177
    1
        }
    178
    179
        #[test]
    180
    1
        fn test_shift_large_pos() {
    181
    1
            let a: Int32Array = vec![Some(1), None, Some(4)].into();
    182
    1
            let res = shift(&a, 1000).unwrap();
    183
    1
            let expected: Int32Array = vec![None, None, None].into();
    184
    1
            assert_eq!(res.as_ref(), &expected);
    185
    1
        }
    186
    187
        #[test]
    188
    1
        fn test_shift_large_neg() {
    189
    1
            let a: Int32Array = vec![Some(1), None, Some(4)].into();
    190
    1
            let res = shift(&a, -1000).unwrap();
    191
    1
            let expected: Int32Array = vec![None, None, None].into();
    192
    1
            assert_eq!(res.as_ref(), &expected);
    193
    1
        }
    194
    }