From fb5483931fdc5ec38055928ace2506ef47828759 Mon Sep 17 00:00:00 2001 From: Sean Smith Date: Fri, 20 Dec 2024 17:05:34 -0600 Subject: [PATCH] flat --- crates/rayexec_bullet/src/exp/array.rs | 26 ---- .../src/exp/executors/scalar/unary.rs | 37 ++++-- crates/rayexec_bullet/src/exp/flat_array.rs | 117 ++++++++++++++++++ crates/rayexec_bullet/src/exp/mod.rs | 1 + 4 files changed, 148 insertions(+), 33 deletions(-) create mode 100644 crates/rayexec_bullet/src/exp/flat_array.rs diff --git a/crates/rayexec_bullet/src/exp/array.rs b/crates/rayexec_bullet/src/exp/array.rs index 229cf89ff..b8fb6ddd1 100644 --- a/crates/rayexec_bullet/src/exp/array.rs +++ b/crates/rayexec_bullet/src/exp/array.rs @@ -140,29 +140,3 @@ where } } } - -#[derive(Debug)] -pub struct DictionaryArrayView<'a, R: ReservationTracker = NopReservationTracker> { - pub(crate) validity: &'a Validity, - pub(crate) array_buffer: &'a ArrayBuffer, - pub(crate) selection: &'a [usize], -} - -impl<'a, R> DictionaryArrayView<'a, R> -where - R: ReservationTracker, -{ - pub fn try_from_array(array: &'a Array) -> Result { - let selection = array.buffer.try_as_slice::()?; - let dict_buffer = array - .buffer - .secondary_buffers() - .try_as_dictionary_buffer()?; - - Ok(DictionaryArrayView { - validity: &dict_buffer.validity, - array_buffer: &dict_buffer.buffer, - selection, - }) - } -} diff --git a/crates/rayexec_bullet/src/exp/executors/scalar/unary.rs b/crates/rayexec_bullet/src/exp/executors/scalar/unary.rs index 89e82d285..970322724 100644 --- a/crates/rayexec_bullet/src/exp/executors/scalar/unary.rs +++ b/crates/rayexec_bullet/src/exp/executors/scalar/unary.rs @@ -1,11 +1,12 @@ use rayexec_error::Result; use crate::compute::util::IntoExactSizedIterator; -use crate::exp::array::{Array, DictionaryArrayView}; +use crate::exp::array::Array; use crate::exp::buffer::addressable::{AddressableStorage, MutableAddressableStorage}; use crate::exp::buffer::physical_type::{MutablePhysicalStorage, PhysicalStorage}; use crate::exp::buffer::ArrayBuffer; use crate::exp::executors::OutputBuffer; +use crate::exp::flat_array::FlatArrayView; use crate::exp::validity::Validity; #[derive(Debug, Clone)] @@ -26,8 +27,8 @@ impl UnaryExecutor { for<'a> Op: FnMut(&S::StorageType, OutputBuffer>), { if array.is_dictionary() { - let view = DictionaryArrayView::try_from_array(array)?; - return Self::execute_dictionary::(view, selection, out, out_validity, op); + let view = FlatArrayView::from_array(array)?; + return Self::execute_flat::(view, selection, out, out_validity, op); } let input = S::get_storage(array.buffer())?; @@ -58,8 +59,8 @@ impl UnaryExecutor { Ok(()) } - pub fn execute_dictionary<'a, S, O, Op>( - array: DictionaryArrayView<'a>, + pub fn execute_flat<'a, S, O, Op>( + array: FlatArrayView<'a>, selection: impl IntoExactSizedIterator, out: &mut ArrayBuffer, out_validity: &mut Validity, @@ -77,7 +78,7 @@ impl UnaryExecutor { if validity.all_valid() { for (output_idx, input_idx) in selection.into_iter().enumerate() { - let selected_idx = array.selection[input_idx]; + let selected_idx = array.selection.get(input_idx).unwrap(); op( input.get(selected_idx).unwrap(), @@ -86,7 +87,7 @@ impl UnaryExecutor { } } else { for (output_idx, input_idx) in selection.into_iter().enumerate() { - let selected_idx = array.selection[input_idx]; + let selected_idx = array.selection.get(input_idx).unwrap(); if validity.is_valid(selected_idx) { op( @@ -160,6 +161,28 @@ mod tests { assert_eq!(&[3, 4, 5], out_slice); } + #[test] + fn int32_inc_by_2_using_flat_view() { + let array = Array::new(DataType::Int32, Int32Builder::from_iter([1, 2, 3]).unwrap()); + let mut out = ArrayBuffer::with_len::(&NopReservationTracker, 3).unwrap(); + let mut validity = Validity::new_all_valid(3); + + let flat = FlatArrayView::from_array(&array).unwrap(); + + UnaryExecutor::execute_flat::( + flat, + 0..3, + &mut out, + &mut validity, + |&v, buf| buf.put(&(v + 2)), + ) + .unwrap(); + assert!(validity.all_valid()); + + let out_slice = out.try_as_slice::().unwrap(); + assert_eq!(&[3, 4, 5], out_slice); + } + #[test] fn int32_inc_by_2_in_place() { let mut array = Array::new(DataType::Int32, Int32Builder::from_iter([1, 2, 3]).unwrap()); diff --git a/crates/rayexec_bullet/src/exp/flat_array.rs b/crates/rayexec_bullet/src/exp/flat_array.rs new file mode 100644 index 000000000..31b736abd --- /dev/null +++ b/crates/rayexec_bullet/src/exp/flat_array.rs @@ -0,0 +1,117 @@ +use rayexec_error::Result; + +use super::array::Array; +use super::buffer::physical_type::PhysicalDictionary; +use super::buffer::reservation::{NopReservationTracker, ReservationTracker}; +use super::buffer::ArrayBuffer; +use super::validity::Validity; + +/// A view on top of normal arrays flattening some parts of the nested +/// structure. +#[derive(Debug)] +pub struct FlatArrayView<'a, R: ReservationTracker = NopReservationTracker> { + pub(crate) validity: &'a Validity, + pub(crate) array_buffer: &'a ArrayBuffer, + pub(crate) selection: FlatSelection<'a>, +} + +impl<'a, R> FlatArrayView<'a, R> +where + R: ReservationTracker, +{ + pub fn from_array(array: &'a Array) -> Result { + if array.is_dictionary() { + let selection = array.buffer.try_as_slice::()?; + let dict_buffer = array + .buffer + .secondary_buffers() + .try_as_dictionary_buffer()?; + + Ok(FlatArrayView { + validity: &dict_buffer.validity, + array_buffer: &dict_buffer.buffer, + selection: FlatSelection::selection(selection), + }) + } else { + Ok(FlatArrayView { + validity: &array.validity, + array_buffer: &array.buffer, + selection: FlatSelection::linear(array.len()), + }) + } + } +} + +#[derive(Debug, Clone, Copy)] +pub enum FlatSelection<'a> { + /// Represents a linear selection. + /// + /// '0..len' + Linear(usize), + /// Represents the true location to use for some index. + Selection(&'a [usize]), +} + +impl<'a> FlatSelection<'a> { + pub fn linear(len: usize) -> Self { + Self::Linear(len) + } + + pub fn selection(sel: &'a [usize]) -> Self { + Self::Selection(sel) + } + + pub fn iter(&self) -> FlatSelectionIter { + FlatSelectionIter { idx: 0, sel: *self } + } + + pub fn len(&self) -> usize { + match self { + Self::Linear(len) => *len, + Self::Selection(sel) => sel.len(), + } + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + #[inline] + pub fn get(&self, idx: usize) -> Option { + match self { + Self::Linear(len) => { + if idx >= *len { + None + } else { + Some(idx) + } + } + Self::Selection(sel) => sel.get(idx).copied(), + } + } +} + +#[derive(Debug)] +pub struct FlatSelectionIter<'a> { + idx: usize, + sel: FlatSelection<'a>, +} + +impl<'a> Iterator for FlatSelectionIter<'a> { + type Item = usize; + + fn next(&mut self) -> Option { + if self.idx >= self.sel.len() { + return None; + } + + let v = match self.sel { + FlatSelection::Linear(_) => self.idx, + FlatSelection::Selection(sel) => sel[self.idx], + }; + + self.idx += 1; + + Some(v) + } +} diff --git a/crates/rayexec_bullet/src/exp/mod.rs b/crates/rayexec_bullet/src/exp/mod.rs index 548af4bfa..740962263 100644 --- a/crates/rayexec_bullet/src/exp/mod.rs +++ b/crates/rayexec_bullet/src/exp/mod.rs @@ -2,4 +2,5 @@ pub mod array; pub mod batch; pub mod buffer; pub mod executors; +pub mod flat_array; pub mod validity;