fix: RunEndBool stats and slice accounts for offsets (#1428)

robert3005 · web-flow · commit e7cbad6ea02b · 2024-11-21T07:46:35.000-05:00
diff --git a/encodings/runend-bool/benches/ree_bool_compress.rs b/encodings/runend-bool/benches/ree_bool_compress.rs
@@ -8,7 +8,10 @@ use itertools::Itertools;
 use rand::distributions::Open01;
 use rand::prelude::StdRng;
 use rand::{Rng, SeedableRng};
-use vortex_runend_bool::compress::{runend_bool_decode_slice, runend_bool_encode_slice};
+use vortex_error::VortexExpect;
+use vortex_runend_bool::compress::{
+    runend_bool_decode_slice, runend_bool_encode_slice, trimmed_ends_iter,
+};
 
 fn compress_compare(c: &mut Criterion) {
     compress_compare_param(c, 0.);
@@ -38,8 +41,15 @@ fn compress_compare_param(c: &mut Criterion, sel_fac: f32) {
     });
 
     let (ends, start) = runend_bool_encode_slice(&boolbuf);
+    let length = ends.last().copied().vortex_expect("must have one elem") as usize;
     group.bench_function("ree bool decompress", |b| {
-        b.iter(|| black_box(runend_bool_decode_slice(&ends, start, 0, ends.len())));
+        b.iter(|| {
+            black_box(runend_bool_decode_slice(
+                trimmed_ends_iter(&ends, 0, length),
+                start,
+                length,
+            ))
+        });
     });
     group.finish()
 }
diff --git a/encodings/runend-bool/src/array.rs b/encodings/runend-bool/src/array.rs
@@ -17,7 +17,7 @@ use vortex_dtype::{match_each_integer_ptype, match_each_unsigned_integer_ptype,
 use vortex_error::{vortex_bail, VortexExpect as _, VortexResult};
 use vortex_scalar::Scalar;
 
-use crate::compress::{runend_bool_decode_slice, runend_bool_encode_slice};
+use crate::compress::{runend_bool_decode_slice, runend_bool_encode_slice, trimmed_ends_iter};
 
 impl_encoding!("vortex.runendbool", ids::RUN_END_BOOL, RunEndBool);
 
@@ -164,7 +164,7 @@ pub(crate) fn decode_runend_bool(
     length: usize,
 ) -> VortexResult<BoolArray> {
     match_each_integer_ptype!(run_ends.ptype(), |$E| {
-        let bools = runend_bool_decode_slice::<$E>(run_ends.maybe_null_slice(), start, offset, length);
+        let bools = runend_bool_decode_slice(trimmed_ends_iter(run_ends.maybe_null_slice::<$E>(), offset, length), start, length);
         Ok(BoolArray::try_new(bools, validity)?)
     })
 }
@@ -179,8 +179,14 @@ pub(crate) fn value_at_index(idx: usize, start: bool) -> bool {
 
 impl BoolArrayTrait for RunEndBoolArray {
     fn invert(&self) -> VortexResult<ArrayData> {
-        RunEndBoolArray::try_new(self.ends(), !self.start(), self.validity())
-            .map(|a| a.into_array())
+        RunEndBoolArray::with_offset_and_size(
+            self.ends(),
+            !self.start(),
+            self.validity(),
+            self.len(),
+            self.offset(),
+        )
+        .map(|a| a.into_array())
     }
 }
 
@@ -229,19 +235,19 @@ impl ArrayStatisticsCompute for RunEndBoolArray {
             Stat::NullCount => Some(self.validity().null_count(self.len())?.into()),
             Stat::TrueCount => {
                 let pends = self.ends().into_primitive()?;
-                let mut true_count: u64 = 0;
-                let mut prev_end: u64 = 0;
+                let mut true_count: usize = 0;
+                let mut prev_end: usize = 0;
                 let mut include = self.start();
                 match_each_unsigned_integer_ptype!(pends.ptype(), |$P| {
-                    for end in pends.maybe_null_slice::<$P>() {
+                    for end in trimmed_ends_iter(pends.maybe_null_slice::<$P>(), self.offset(), self.len()) {
                         if include {
-                            true_count += (*end as u64 - prev_end);
+                            true_count += end - prev_end;
                         }
                         include = !include;
-                        prev_end = *end as u64;
+                        prev_end = end;
                     }
                 });
-                Some(true_count.into())
+                Some((true_count as u64).into())
             }
             _ => None,
         };
@@ -259,7 +265,7 @@ mod test {
 
     use itertools::Itertools as _;
     use rstest::rstest;
-    use vortex_array::array::BoolArray;
+    use vortex_array::array::{BoolArray, PrimitiveArray};
     use vortex_array::compute::unary::scalar_at;
     use vortex_array::compute::{slice, take, TakeOptions};
     use vortex_array::stats::{ArrayStatistics as _, ArrayStatisticsCompute};
@@ -408,4 +414,16 @@ mod test {
 
         assert_eq!(arr.statistics().compute_run_count(), Some(ends_len));
     }
+
+    #[test]
+    fn sliced_true_count() {
+        let arr = RunEndBoolArray::try_new(
+            PrimitiveArray::from(vec![5u32, 7, 10]).into_array(),
+            true,
+            Validity::NonNullable,
+        )
+        .unwrap();
+        let sliced = slice(&arr, 4, 8).unwrap();
+        assert_eq!(sliced.statistics().compute_true_count().unwrap(), 2);
+    }
 }
diff --git a/encodings/runend-bool/src/compress.rs b/encodings/runend-bool/src/compress.rs
@@ -2,7 +2,7 @@ use std::cmp::min;
 
 use arrow_buffer::buffer::BooleanBuffer;
 use arrow_buffer::BooleanBufferBuilder;
-use num_traits::{AsPrimitive, FromPrimitive};
+use num_traits::AsPrimitive;
 use vortex_dtype::NativePType;
 use vortex_error::{vortex_panic, VortexExpect as _};
 
@@ -35,12 +35,12 @@ pub fn runend_bool_encode_slice(elements: &BooleanBuffer) -> (Vec<u64>, bool) {
     (ends, first_bool)
 }
 
-pub fn runend_bool_decode_slice<E: NativePType + AsPrimitive<usize> + FromPrimitive + Ord>(
+#[inline]
+pub fn trimmed_ends_iter<E: NativePType + AsPrimitive<usize> + Ord>(
     run_ends: &[E],
-    start: bool,
     offset: usize,
     length: usize,
-) -> BooleanBuffer {
+) -> impl Iterator<Item = usize> + use<'_, E> {
     let offset_e = E::from_usize(offset).unwrap_or_else(|| {
         vortex_panic!(
             "offset {} cannot be converted to {}",
@@ -55,14 +55,22 @@ pub fn runend_bool_decode_slice<E: NativePType + AsPrimitive<usize> + FromPrimit
             std::any::type_name::<E>()
         )
     });
-    let trimmed_ends = run_ends
+    run_ends
         .iter()
-        .map(|v| *v - offset_e)
-        .map(|v| min(v, length_e));
+        .copied()
+        .map(move |v| v - offset_e)
+        .map(move |v| min(v, length_e))
+        .map(|v| v.as_())
+}
 
+pub fn runend_bool_decode_slice(
+    run_ends_iter: impl Iterator<Item = usize>,
+    start: bool,
+    length: usize,
+) -> BooleanBuffer {
     let mut decoded = BooleanBufferBuilder::new(length);
-    for (idx, end) in trimmed_ends.enumerate() {
-        decoded.append_n(end.as_() - decoded.len(), value_at_index(idx, start));
+    for (idx, end) in run_ends_iter.enumerate() {
+        decoded.append_n(end - decoded.len(), value_at_index(idx, start));
     }
     BooleanBuffer::from(decoded)
 }
@@ -78,7 +86,7 @@ mod test {
     use vortex_array::validity::Validity;
     use vortex_array::IntoArrayVariant;
 
-    use crate::compress::{runend_bool_decode_slice, runend_bool_encode_slice};
+    use crate::compress::{runend_bool_decode_slice, runend_bool_encode_slice, trimmed_ends_iter};
     use crate::decode_runend_bool;
 
     #[test]
@@ -107,27 +115,30 @@ mod test {
     fn encode_decode_bool() {
         let input = [true, true, false, true, true, false];
         let (ends, start) = runend_bool_encode_slice(&BooleanBuffer::from(input.as_slice()));
+        let ends_iter = trimmed_ends_iter(ends.as_slice(), 0, input.len());
 
-        let decoded = runend_bool_decode_slice(ends.as_slice(), start, 0, input.len());
+        let decoded = runend_bool_decode_slice(ends_iter, start, input.len());
         assert_eq!(decoded, BooleanBuffer::from(input.as_slice()))
     }
 
     #[test]
     fn encode_decode_bool_false_start() {
         let input = [false, false, true, true, false, true, true, false];
         let (ends, start) = runend_bool_encode_slice(&BooleanBuffer::from(input.as_slice()));
+        let ends_iter = trimmed_ends_iter(ends.as_slice(), 0, input.len());
 
-        let decoded = runend_bool_decode_slice(ends.as_slice(), start, 0, input.len());
+        let decoded = runend_bool_decode_slice(ends_iter, start, input.len());
         assert_eq!(decoded, BooleanBuffer::from(input.as_slice()))
     }
 
     #[test]
     fn encode_decode_bool_false_start_long() {
-        let input = vec![true; 1024];
-        // input.extend([false, true, false, true].as_slice());
+        let mut input = vec![true; 1024];
+        input.extend([false, true, false, true].as_slice());
         let (ends, start) = runend_bool_encode_slice(&BooleanBuffer::from(input.as_slice()));
+        let ends_iter = trimmed_ends_iter(ends.as_slice(), 0, input.len());
 
-        let decoded = runend_bool_decode_slice(ends.as_slice(), start, 0, input.len());
+        let decoded = runend_bool_decode_slice(ends_iter, start, input.len());
         assert_eq!(decoded, BooleanBuffer::from(input.as_slice()))
     }
 
@@ -136,8 +147,9 @@ mod test {
         let mut rng = StdRng::seed_from_u64(4352);
         let input = (0..1024 * 4).map(|_x| rng.gen::<bool>()).collect_vec();
         let (ends, start) = runend_bool_encode_slice(&BooleanBuffer::from(input.as_slice()));
+        let ends_iter = trimmed_ends_iter(ends.as_slice(), 0, input.len());
 
-        let decoded = runend_bool_decode_slice(ends.as_slice(), start, 0, input.len());
+        let decoded = runend_bool_decode_slice(ends_iter, start, input.len());
         assert_eq!(decoded, BooleanBuffer::from(input.as_slice()))
     }
 
diff --git a/encodings/runend-bool/src/compute.rs b/encodings/runend-bool/src/compute.rs
@@ -43,9 +43,9 @@ impl TakeFn<RunEndBoolArray> for RunEndBoolEncoding {
         let primitive_indices = indices.clone().into_primitive()?;
         let physical_indices = match_each_integer_ptype!(primitive_indices.ptype(), |$P| {
             primitive_indices
-                .maybe_null_slice::<$P>()
-                .iter()
-                .map(|idx| *idx as usize)
+                .into_maybe_null_slice::<$P>()
+                .into_iter()
+                .map(|idx| idx as usize)
                 .map(|idx| {
                     if idx >= array.len() {
                         vortex_bail!(OutOfBounds: idx, 0, array.len())
@@ -72,7 +72,7 @@ impl SliceFn<RunEndBoolArray> for RunEndBoolEncoding {
             value_at_index(slice_begin, array.start()),
             array.validity().slice(slice_begin, slice_end + 1)?,
             stop - start,
-            start,
+            start + array.offset(),
         )?
         .into_array())
     }
