flat

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<R>,
-    pub(crate) selection: &'a [usize],
-}
-
-impl<'a, R> DictionaryArrayView<'a, R>
-where
-    R: ReservationTracker,
-{
-    pub fn try_from_array(array: &'a Array<R>) -> Result<Self> {
-        let selection = array.buffer.try_as_slice::<PhysicalDictionary>()?;
-        let dict_buffer = array
-            .buffer
-            .secondary_buffers()
-            .try_as_dictionary_buffer()?;
-
-        Ok(DictionaryArrayView {
-            validity: &dict_buffer.validity,
-            array_buffer: &dict_buffer.buffer,
-            selection,
-        })
-    }
-}

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<O::MutableStorage<'a>>),
     {
         if array.is_dictionary() {
-            let view = DictionaryArrayView::try_from_array(array)?;
-            return Self::execute_dictionary::<S, _, _>(view, selection, out, out_validity, op);
+            let view = FlatArrayView::from_array(array)?;
+            return Self::execute_flat::<S, _, _>(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<Item = usize>,
         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::<PhysicalI32>(&NopReservationTracker, 3).unwrap();
+        let mut validity = Validity::new_all_valid(3);
+
+        let flat = FlatArrayView::from_array(&array).unwrap();
+
+        UnaryExecutor::execute_flat::<PhysicalI32, PhysicalI32, _>(
+            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::<PhysicalI32>().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());

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<R>,
+    pub(crate) selection: FlatSelection<'a>,
+}
+
+impl<'a, R> FlatArrayView<'a, R>
+where
+    R: ReservationTracker,
+{
+    pub fn from_array(array: &'a Array<R>) -> Result<Self> {
+        if array.is_dictionary() {
+            let selection = array.buffer.try_as_slice::<PhysicalDictionary>()?;
+            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<usize> {
+        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<Self::Item> {
+        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)
+    }
+}

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;