Skip to content

Commit

Permalink
add tests
Browse files Browse the repository at this point in the history
Signed-off-by: remzi <13716567376yh@gmail.com>
  • Loading branch information
HaoYang670 committed Mar 26, 2022
1 parent 1d5221f commit aa6a138
Showing 1 changed file with 66 additions and 4 deletions.
70 changes: 66 additions & 4 deletions arrow/src/compute/kernels/length.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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> {
Expand Down Expand Up @@ -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<()> {
Expand Down Expand Up @@ -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>>)> {
Expand Down Expand Up @@ -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() {
Expand All @@ -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())?;
Expand All @@ -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);
Expand Down

0 comments on commit aa6a138

Please sign in to comment.