From b5ea3f65e6f5a8797a83d21949ed3efd859f725c Mon Sep 17 00:00:00 2001
From: Andrew Lamb <andrew@nerdnetworks.org>
Date: Thu, 9 Sep 2021 16:27:49 -0400
Subject: [PATCH] fix: Handle slices in unary kernel (#739) (#755)

Co-authored-by: Ben Chambers <35960+bjchambers@users.noreply.github.com>
---
 arrow/src/buffer/immutable.rs      |  2 +-
 arrow/src/compute/kernels/arity.rs | 24 +++++++++++++++++++++++-
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/arrow/src/buffer/immutable.rs b/arrow/src/buffer/immutable.rs
index c00af6ef1573..f0aefd9b94b7 100644
--- a/arrow/src/buffer/immutable.rs
+++ b/arrow/src/buffer/immutable.rs
@@ -184,7 +184,7 @@ impl Buffer {
     /// If the offset is byte-aligned the returned buffer is a shallow clone,
     /// otherwise a new buffer is allocated and filled with a copy of the bits in the range.
     pub fn bit_slice(&self, offset: usize, len: usize) -> Self {
-        if offset % 8 == 0 && len % 8 == 0 {
+        if offset % 8 == 0 {
             return self.slice(offset / 8);
         }
 
diff --git a/arrow/src/compute/kernels/arity.rs b/arrow/src/compute/kernels/arity.rs
index 4aa7f3d6e5dc..d7beae605993 100644
--- a/arrow/src/compute/kernels/arity.rs
+++ b/arrow/src/compute/kernels/arity.rs
@@ -30,7 +30,10 @@ fn into_primitive_array_data<I: ArrowPrimitiveType, O: ArrowPrimitiveType>(
         O::DATA_TYPE,
         array.len(),
         None,
-        array.data_ref().null_buffer().cloned(),
+        array
+            .data_ref()
+            .null_buffer()
+            .map(|b| b.bit_slice(array.offset(), array.len())),
         0,
         vec![buffer],
         vec![],
@@ -72,3 +75,22 @@ where
     let data = into_primitive_array_data::<_, O>(array, buffer);
     PrimitiveArray::<O>::from(data)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::array::{as_primitive_array, Float64Array};
+
+    #[test]
+    fn test_unary_f64_slice() {
+        let input =
+            Float64Array::from(vec![Some(5.1f64), None, Some(6.8), None, Some(7.2)]);
+        let input_slice = input.slice(1, 4);
+        let input_slice: &Float64Array = as_primitive_array(&input_slice);
+        let result = unary(input_slice, |n| n.round());
+        assert_eq!(
+            result,
+            Float64Array::from(vec![None, Some(7.0), None, Some(7.0)])
+        )
+    }
+}