From eaefe40af1759f262814147c18e434011b1a1168 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=B6rn=20Horstmann?= Date: Fri, 17 Jun 2022 14:33:14 +0200 Subject: [PATCH] Use bit_slice in combine_option_bitmap --- arrow/src/compute/util.rs | 41 +++++++++++++++++++++++++++++++++++++-- 1 file changed, 39 insertions(+), 2 deletions(-) diff --git a/arrow/src/compute/util.rs b/arrow/src/compute/util.rs index 7d583726105..c8e68fbeb35 100644 --- a/arrow/src/compute/util.rs +++ b/arrow/src/compute/util.rs @@ -55,7 +55,9 @@ pub(super) fn combine_option_bitmap( Err(ArrowError::ComputeError( "Arrays must not be empty".to_string(), )), - |(buffer, offset)| Ok(buffer.map(|buffer| buffer.slice(offset))), + |(buffer, offset)| { + Ok(buffer.map(|buffer| buffer.bit_slice(offset, len_in_bits))) + }, ) } @@ -185,7 +187,7 @@ pub(super) mod tests { offset: usize, null_bit_buffer: Option, ) -> Arc { - let buffer = Buffer::from(&vec![11; len]); + let buffer = Buffer::from(&vec![11; len + offset]); Arc::new( ArrayData::try_new( @@ -256,6 +258,41 @@ pub(super) mod tests { ); } + #[test] + fn test_combine_option_bitmap_with_offsets() { + let none_bitmap = make_data_with_null_bit_buffer(8, 0, None); + let bitmap0 = + make_data_with_null_bit_buffer(8, 0, Some(Buffer::from([0b10101010]))); + let bitmap1 = + make_data_with_null_bit_buffer(8, 1, Some(Buffer::from([0b01010100, 0b1]))); + let bitmap2 = + make_data_with_null_bit_buffer(8, 2, Some(Buffer::from([0b10101000, 0b10]))); + assert_eq!( + Some(Buffer::from([0b10101010])), + combine_option_bitmap(&[&bitmap1], 8).unwrap() + ); + assert_eq!( + Some(Buffer::from([0b10101010])), + combine_option_bitmap(&[&bitmap2], 8).unwrap() + ); + assert_eq!( + Some(Buffer::from([0b10101010])), + combine_option_bitmap(&[&bitmap1, &none_bitmap], 8).unwrap() + ); + assert_eq!( + Some(Buffer::from([0b10101010])), + combine_option_bitmap(&[&none_bitmap, &bitmap2], 8).unwrap() + ); + assert_eq!( + Some(Buffer::from([0b10101010])), + combine_option_bitmap(&[&bitmap0, &bitmap1], 8).unwrap() + ); + assert_eq!( + Some(Buffer::from([0b10101010])), + combine_option_bitmap(&[&bitmap1, &bitmap2], 8).unwrap() + ); + } + #[test] fn test_compare_option_bitmap() { let none_bitmap = make_data_with_null_bit_buffer(8, 0, None);