REE infers length from ends array (#147)

fixes #19
spiraldb · Apr 3, 2024 · 1a5c2fe · 1a5c2fe
1 parent 221e2fb
commit 1a5c2fe
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 24 deletions.
diff --git a/vortex-ree/src/compress.rs b/vortex-ree/src/compress.rs
@@ -59,7 +59,6 @@ impl EncodingCompression for REEEncoding {
             compressed_ends,
             compressed_values,
             ctx.compress_validity(primitive_array.validity())?,
-            array.len(),
         )
         .into_array())
     }
@@ -195,7 +194,6 @@ mod test {
             vec![2u32, 5, 10].into_array(),
             vec![1i32, 2, 3].into_array(),
             Some(validity),
-            10,
         );
 
         let decoded = ree_decode(

diff --git a/vortex-ree/src/ree.rs b/vortex-ree/src/ree.rs
@@ -3,6 +3,7 @@ use std::sync::{Arc, RwLock};
 use vortex::array::validity::Validity;
 use vortex::array::{check_slice_bounds, Array, ArrayKind, ArrayRef};
 use vortex::compress::EncodingCompression;
+use vortex::compute::scalar_at::scalar_at;
 use vortex::compute::search_sorted::SearchSortedSide;
 use vortex::compute::ArrayCompute;
 use vortex::encoding::{Encoding, EncodingId, EncodingRef};
@@ -26,21 +27,16 @@ pub struct REEArray {
 }
 
 impl REEArray {
-    pub fn new(
-        ends: ArrayRef,
-        values: ArrayRef,
-        validity: Option<Validity>,
-        length: usize,
-    ) -> Self {
-        Self::try_new(ends, values, validity, length).unwrap()
+    pub fn new(ends: ArrayRef, values: ArrayRef, validity: Option<Validity>) -> Self {
+        Self::try_new(ends, values, validity).unwrap()
     }
 
     pub fn try_new(
         ends: ArrayRef,
         values: ArrayRef,
         validity: Option<Validity>,
-        length: usize,
     ) -> VortexResult<Self> {
+        let length: usize = scalar_at(ends.as_ref(), ends.len() - 1)?.try_into()?;
         if let Some(v) = &validity {
             assert_eq!(v.len(), length);
         }
@@ -53,7 +49,6 @@ impl REEArray {
             vortex_bail!("Ends array must be strictly sorted",);
         }
 
-        // TODO(ngates): https://github.com/fulcrum-so/spiral/issues/873
         Ok(Self {
             ends,
             values,
@@ -76,13 +71,10 @@ impl REEArray {
         match ArrayKind::from(array) {
             ArrayKind::Primitive(p) => {
                 let (ends, values) = ree_encode(p);
-                Ok(REEArray::new(
-                    ends.into_array(),
-                    values.into_array(),
-                    p.validity(),
-                    p.len(),
+                Ok(
+                    REEArray::new(ends.into_array(), values.into_array(), p.validity())
+                        .into_array(),
                 )
-                .into_array())
             }
             _ => Err(vortex_err!("REE can only encode primitive arrays")),
         }
@@ -213,7 +205,6 @@ mod test {
             vec![2u32, 5, 10].into_array(),
             vec![1i32, 2, 3].into_array(),
             None,
-            10,
         );
         assert_eq!(arr.len(), 10);
         assert_eq!(
@@ -236,7 +227,6 @@ mod test {
             vec![2u32, 5, 10].into_array(),
             vec![1i32, 2, 3].into_array(),
             None,
-            10,
         )
         .slice(3, 8)
         .unwrap();
@@ -258,7 +248,6 @@ mod test {
             vec![2u32, 5, 10].into_array(),
             vec![1i32, 2, 3].into_array(),
             None,
-            10,
         );
         assert_eq!(
             flatten_primitive(&arr).unwrap().typed_data::<i32>(),

diff --git a/vortex-ree/src/serde.rs b/vortex-ree/src/serde.rs
@@ -6,7 +6,6 @@ use crate::{REEArray, REEEncoding};
 
 impl ArraySerde for REEArray {
     fn write(&self, ctx: &mut WriteCtx) -> VortexResult<()> {
-        ctx.write_usize(self.len())?;
         ctx.write_validity(self.validity())?;
         // TODO(robert): Stop writing this
         ctx.dtype(self.ends().dtype())?;
@@ -21,12 +20,11 @@ impl ArraySerde for REEArray {
 
 impl EncodingSerde for REEEncoding {
     fn read(&self, ctx: &mut ReadCtx) -> VortexResult<ArrayRef> {
-        let len = ctx.read_usize()?;
         let validity = ctx.read_validity()?;
         let ends_dtype = ctx.dtype()?;
         let ends = ctx.with_schema(&ends_dtype).read()?;
         let values = ctx.read()?;
-        Ok(REEArray::new(ends, values, validity, len).into_array())
+        Ok(REEArray::new(ends, values, validity).into_array())
     }
 }
 
@@ -57,7 +55,6 @@ mod test {
             vec![0u8, 9, 20, 32, 49].into_array(),
             vec![-7i64, -13, 17, 23].into_array(),
             None,
-            49,
         );
         let read_arr = roundtrip_array(&arr).unwrap();
         let read_ree = read_arr.as_ree();