Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

REE infers length from ends array #147

Merged
merged 2 commits into from
Apr 3, 2024
Merged

REE infers length from ends array #147

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3451907
REE infers length from ends array
robert3005 Mar 26, 2024
14115da
Merge remote-tracking branch 'origin/develop' into rk/ree-length
lwwmanning Apr 3, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 0 additions & 2 deletions vortex-ree/src/compress.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,6 @@ impl EncodingCompression for REEEncoding {
compressed_ends,
compressed_values,
ctx.compress_validity(primitive_array.validity())?,
array.len(),
)
.into_array())
}
Expand Down Expand Up @@ -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(
Expand Down
25 changes: 7 additions & 18 deletions vortex-ree/src/ree.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand All @@ -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);
}
Expand All @@ -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,
Expand All @@ -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")),
}
Expand Down Expand Up @@ -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!(
Expand All @@ -236,7 +227,6 @@ mod test {
vec![2u32, 5, 10].into_array(),
vec![1i32, 2, 3].into_array(),
None,
10,
)
.slice(3, 8)
.unwrap();
Expand All @@ -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>(),
Expand Down
5 changes: 1 addition & 4 deletions vortex-ree/src/serde.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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())?;
Expand All @@ -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())
}
}

Expand Down Expand Up @@ -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();
Expand Down