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

Add Take for REEArray #162

Merged
merged 2 commits into from
Mar 28, 2024
Merged

Add Take for REEArray #162

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d52a11c
Add take for REEArray
robert3005 Mar 28, 2024
5b51e9a
simpler
robert3005 Mar 28, 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
47 changes: 45 additions & 2 deletions vortex-ree/src/compute.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,10 @@
use vortex::array::Array;
use vortex::compute::flatten::{flatten, FlattenFn, FlattenedArray};
use vortex::array::primitive::PrimitiveArray;
use vortex::array::{Array, ArrayRef};
use vortex::compute::flatten::{flatten, flatten_primitive, FlattenFn, FlattenedArray};
use vortex::compute::scalar_at::{scalar_at, ScalarAtFn};
use vortex::compute::take::{take, TakeFn};
use vortex::compute::ArrayCompute;
use vortex::match_each_integer_ptype;
use vortex::scalar::Scalar;
use vortex::validity::ArrayValidity;
use vortex_error::{VortexError, VortexResult};
Expand All @@ -17,6 +20,10 @@ impl ArrayCompute for REEArray {
fn scalar_at(&self) -> Option<&dyn ScalarAtFn> {
Some(self)
}

fn take(&self) -> Option<&dyn TakeFn> {
Some(self)
}
}

impl FlattenFn for REEArray {
Expand Down Expand Up @@ -45,3 +52,39 @@ impl ScalarAtFn for REEArray {
scalar_at(self.values(), self.find_physical_index(index)?)
}
}

impl TakeFn for REEArray {
fn take(&self, indices: &dyn Array) -> VortexResult<ArrayRef> {
let primitive_indices = flatten_primitive(indices)?;
let physical_indices = match_each_integer_ptype!(primitive_indices.ptype(), |$P| {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this should just be a search_sorted? Possibly this? #164

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That said, an iterator would allow us to apply the offset inline without allocating a new buffer

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is just search sorted. That’s what find_physical_index does

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, but a lot of function calls, vs search_sorted_many

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure

primitive_indices
.typed_data::<$P>()
.iter()
.map(|idx| {
self.find_physical_index(*idx as usize)
.map(|loc| loc as u64)
})
.collect::<VortexResult<Vec<_>>>()?
});
take(self.values(), &PrimitiveArray::from(physical_indices))
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the same pattern that happens in Dict array, repeated value take on some dictionary.

}
}

#[cfg(test)]
mod test {
use vortex::array::downcast::DowncastArrayBuiltin;
use vortex::array::primitive::PrimitiveArray;
use vortex::compute::take::take;

use crate::REEArray;

#[test]
fn ree_take() {
let ree = REEArray::encode(&PrimitiveArray::from(vec![
1, 1, 1, 4, 4, 4, 2, 2, 5, 5, 5, 5,
]))
.unwrap();
let taken = take(&ree, &PrimitiveArray::from(vec![9, 8, 1, 3])).unwrap();
assert_eq!(taken.as_primitive().typed_data::<i32>(), &[5, 5, 1, 4]);
}
}