arrow/src/compute/kernels/window.rs

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! Defines windowing functions, like `shift`ing

use crate::array::{Array, ArrayRef};
use crate::error::Result;
use crate::{
    array::{make_array, new_null_array},
    compute::concat,
};
use num::{abs, clamp};

/// Shifts array by defined number of items (to left or right)
/// A positive value for `offset` shifts the array to the right
/// a negative value shifts the array to the left.
/// # Examples
/// ```
/// use arrow::array::Int32Array;
/// use arrow::error::Result;
/// use arrow::compute::shift;
///
/// let a: Int32Array = vec![Some(1), None, Some(4)].into();
///
/// // shift array 1 element to the right
/// let res = shift(&a, 1).unwrap();
/// let expected: Int32Array = vec![None, Some(1), None].into();
/// assert_eq!(res.as_ref(), &expected);
///
/// // shift array 1 element to the left
/// let res = shift(&a, -1).unwrap();
/// let expected: Int32Array = vec![None, Some(4), None].into();
/// assert_eq!(res.as_ref(), &expected);
///
/// // shift array 0 element, although not recommended
/// let res = shift(&a, 0).unwrap();
/// let expected: Int32Array = vec![Some(1), None, Some(4)].into();
/// assert_eq!(res.as_ref(), &expected);
///
/// // shift array 3 element tot he right
/// let res = shift(&a, 3).unwrap();
/// let expected: Int32Array = vec![None, None, None].into();
/// assert_eq!(res.as_ref(), &expected);
/// ```
pub fn shift(array: &Array, offset: i64) -> Result<ArrayRef> {
    let value_len = array.len() as i64;
    if offset == 0 {
        Ok(make_array(array.data_ref().clone()))
    } else if offset == i64::MIN || abs(offset) >= value_len {
        Ok(new_null_array(array.data_type(), array.len()))
    } else {
        let slice_offset = clamp(-offset, 0, value_len) as usize;
        let length = array.len() - abs(offset) as usize;
        let slice = array.slice(slice_offset, length);

        // Generate array with remaining `null` items
        let nulls = abs(offset) as usize;
        let null_arr = new_null_array(array.data_type(), nulls);

        // Concatenate both arrays, add nulls after if shift > 0 else before
        if offset > 0 {
            concat(&[null_arr.as_ref(), slice.as_ref()])
        } else {
            concat(&[slice.as_ref(), null_arr.as_ref()])
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::array::{Float64Array, Int32Array, Int32DictionaryArray};

    #[test]
    fn test_shift_neg() {
        let a: Int32Array = vec![Some(1), None, Some(4)].into();
        let res = shift(&a, -1).unwrap();
        let expected: Int32Array = vec![None, Some(4), None].into();
        assert_eq!(res.as_ref(), &expected);
    }

    #[test]
    fn test_shift_pos() {
        let a: Int32Array = vec![Some(1), None, Some(4)].into();
        let res = shift(&a, 1).unwrap();
        let expected: Int32Array = vec![None, Some(1), None].into();
        assert_eq!(res.as_ref(), &expected);
    }

    #[test]
    fn test_shift_neg_float64() {
        let a: Float64Array = vec![Some(1.), None, Some(4.)].into();
        let res = shift(&a, -1).unwrap();
        let expected: Float64Array = vec![None, Some(4.), None].into();
        assert_eq!(res.as_ref(), &expected);
    }

    #[test]
    fn test_shift_pos_float64() {
        let a: Float64Array = vec![Some(1.), None, Some(4.)].into();
        let res = shift(&a, 1).unwrap();
        let expected: Float64Array = vec![None, Some(1.), None].into();
        assert_eq!(res.as_ref(), &expected);
    }

    #[test]
    fn test_shift_neg_int32_dict() {
        let a: Int32DictionaryArray = [Some("alpha"), None, Some("beta"), Some("alpha")]
            .iter()
            .copied()
            .collect();
        let res = shift(&a, -1).unwrap();
        let expected: Int32DictionaryArray = [None, Some("beta"), Some("alpha"), None]
            .iter()
            .copied()
            .collect();
        assert_eq!(res.as_ref(), &expected);
    }

    #[test]
    fn test_shift_pos_int32_dict() {
        let a: Int32DictionaryArray = [Some("alpha"), None, Some("beta"), Some("alpha")]
            .iter()
            .copied()
            .collect();
        let res = shift(&a, 1).unwrap();
        let expected: Int32DictionaryArray = [None, Some("alpha"), None, Some("beta")]
            .iter()
            .copied()
            .collect();
        assert_eq!(res.as_ref(), &expected);
    }

    #[test]
    fn test_shift_nil() {
        let a: Int32Array = vec![Some(1), None, Some(4)].into();
        let res = shift(&a, 0).unwrap();
        let expected: Int32Array = vec![Some(1), None, Some(4)].into();
        assert_eq!(res.as_ref(), &expected);
    }

    #[test]
    fn test_shift_boundary_pos() {
        let a: Int32Array = vec![Some(1), None, Some(4)].into();
        let res = shift(&a, 3).unwrap();
        let expected: Int32Array = vec![None, None, None].into();
        assert_eq!(res.as_ref(), &expected);
    }

    #[test]
    fn test_shift_boundary_neg() {
        let a: Int32Array = vec![Some(1), None, Some(4)].into();
        let res = shift(&a, -3).unwrap();
        let expected: Int32Array = vec![None, None, None].into();
        assert_eq!(res.as_ref(), &expected);
    }

    #[test]
    fn test_shift_boundary_neg_min() {
        let a: Int32Array = vec![Some(1), None, Some(4)].into();
        let res = shift(&a, i64::MIN).unwrap();
        let expected: Int32Array = vec![None, None, None].into();
        assert_eq!(res.as_ref(), &expected);
    }

    #[test]
    fn test_shift_large_pos() {
        let a: Int32Array = vec![Some(1), None, Some(4)].into();
        let res = shift(&a, 1000).unwrap();
        let expected: Int32Array = vec![None, None, None].into();
        assert_eq!(res.as_ref(), &expected);
    }

    #[test]
    fn test_shift_large_neg() {
        let a: Int32Array = vec![Some(1), None, Some(4)].into();
        let res = shift(&a, -1000).unwrap();
        let expected: Int32Array = vec![None, None, None].into();
        assert_eq!(res.as_ref(), &expected);
    }
}