From aa6a1388e1d8cc002375b8b7e3916a86287bcbac Mon Sep 17 00:00:00 2001
From: remzi <13716567376yh@gmail.com>
Date: Sat, 26 Mar 2022 11:09:27 +0800
Subject: [PATCH] add tests

Signed-off-by: remzi <13716567376yh@gmail.com>
---
 arrow/src/compute/kernels/length.rs | 70 +++++++++++++++++++++++++++--
 1 file changed, 66 insertions(+), 4 deletions(-)

diff --git a/arrow/src/compute/kernels/length.rs b/arrow/src/compute/kernels/length.rs
index 262f6f0e5bae..ea25e6787ac4 100644
--- a/arrow/src/compute/kernels/length.rs
+++ b/arrow/src/compute/kernels/length.rs
@@ -156,15 +156,17 @@ pub fn bit_length(array: &dyn Array) -> Result<ArrayRef> {
         DataType::LargeUtf8 => Ok(bit_length_string::<i64, Int64Type>(array)),
         DataType::Binary => Ok(bit_length_binary::<i32, Int32Type>(array)),
         DataType::LargeBinary => Ok(bit_length_binary::<i64, Int64Type>(array)),
-        _ => Err(ArrowError::ComputeError(format!(
+        other => Err(ArrowError::ComputeError(format!(
             "bit_length not supported for {:?}",
-            array.data_type()
+            other
         ))),
     }
 }
 
 #[cfg(test)]
 mod tests {
+    use crate::datatypes::{Float32Type, Int8Type};
+
     use super::*;
 
     fn double_vec<T: Clone>(v: Vec<T>) -> Vec<T> {
@@ -199,6 +201,20 @@ mod tests {
         }};
     }
 
+    macro_rules! length_list_helper {
+        ($offset_ty: ty, $result_ty: ty, $element_ty: ty, $value: expr, $expected: expr) => {{
+            let array =
+                GenericListArray::<$offset_ty>::from_iter_primitive::<$element_ty, _, _>(
+                    $value,
+                );
+            let result = length(&array)?;
+            let result = result.as_any().downcast_ref::<$result_ty>().unwrap();
+            let expected: $result_ty = $expected.into();
+            assert_eq!(expected.data(), result.data());
+            Ok(())
+        }};
+    }
+
     #[test]
     #[cfg_attr(miri, ignore)] // running forever
     fn length_test_string() -> Result<()> {
@@ -247,6 +263,28 @@ mod tests {
         length_binary_helper!(i64, Int64Array, length, value, result)
     }
 
+    #[test]
+    fn length_test_list() -> Result<()> {
+        let value = vec![
+            Some(vec![]),
+            Some(vec![Some(1), Some(2), Some(4)]),
+            Some(vec![Some(0)]),
+        ];
+        let result: Vec<i32> = vec![0, 3, 1];
+        length_list_helper!(i32, Int32Array, Int32Type, value, result)
+    }
+
+    #[test]
+    fn length_test_large_list() -> Result<()> {
+        let value = vec![
+            Some(vec![]),
+            Some(vec![Some(1.1), Some(2.2), Some(3.3)]),
+            Some(vec![None]),
+        ];
+        let result: Vec<i64> = vec![0, 3, 1];
+        length_list_helper!(i64, Int64Array, Float32Type, value, result)
+    }
+
     type OptionStr = Option<&'static str>;
 
     fn length_null_cases_string() -> Vec<(Vec<OptionStr>, usize, Vec<Option<i32>>)> {
@@ -310,6 +348,30 @@ mod tests {
         length_binary_helper!(i64, Int64Array, length, value, result)
     }
 
+    #[test]
+    fn length_null_list() -> Result<()> {
+        let value = vec![
+            Some(vec![]),
+            None,
+            Some(vec![Some(1), None, Some(2), Some(4)]),
+            Some(vec![Some(0)]),
+        ];
+        let result: Vec<Option<i32>> = vec![Some(0), None, Some(4), Some(1)];
+        length_list_helper!(i32, Int32Array, Int8Type, value, result)
+    }
+
+    #[test]
+    fn length_null_large_list() -> Result<()> {
+        let value = vec![
+            Some(vec![]),
+            None,
+            Some(vec![Some(1.1), None, Some(4.0)]),
+            Some(vec![Some(0.1)]),
+        ];
+        let result: Vec<Option<i64>> = vec![Some(0), None, Some(3), Some(1)];
+        length_list_helper!(i64, Int64Array, Float32Type, value, result)
+    }
+
     /// Tests that length is not valid for u64.
     #[test]
     fn length_wrong_type() {
@@ -320,7 +382,7 @@ mod tests {
 
     /// Tests with an offset
     #[test]
-    fn length_offsets() -> Result<()> {
+    fn length_offsets_string() -> Result<()> {
         let a = StringArray::from(vec![Some("hello"), Some(" "), Some("world"), None]);
         let b = a.slice(1, 3);
         let result = length(b.as_ref())?;
@@ -333,7 +395,7 @@ mod tests {
     }
 
     #[test]
-    fn binary_length_offsets() -> Result<()> {
+    fn length_offsets_binary() -> Result<()> {
         let value: Vec<Option<&[u8]>> =
             vec![Some(b"hello"), Some(b" "), Some(&[0xff, 0xf8]), None];
         let a = BinaryArray::from(value);