chore: Add more array methods (#3399)

Split out from mutable array stuff

crates/rayexec_execution/src/arrays/array/physical_type.rs

@@ -209,7 +209,7 @@ impl ProtoConv for PhysicalType {
 }
 
 /// Represents an in-memory array that can be indexed into to retrieve values.
-pub trait Addressable: Debug {
+pub trait Addressable<'a>: Debug {
     /// The type that get's returned.
     type T: Send + Debug + ?Sized;
 
@@ -220,10 +220,10 @@ pub trait Addressable: Debug {
     }
 
     /// Get a value at the given index.
-    fn get(&self, idx: usize) -> Option<&Self::T>;
+    fn get(&self, idx: usize) -> Option<&'a Self::T>;
 }
 
-impl<T> Addressable for &[T]
+impl<'a, T> Addressable<'a> for &'a [T]
 where
     T: Debug + Send,
 {
@@ -233,7 +233,7 @@ where
         (**self).len()
     }
 
-    fn get(&self, idx: usize) -> Option<&Self::T> {
+    fn get(&self, idx: usize) -> Option<&'a Self::T> {
         (**self).get(idx)
     }
 }
@@ -355,7 +355,7 @@ pub trait PhysicalStorage: Debug + Sync + Send + Clone + Copy + 'static {
     type StorageType: Sync + Send + ?Sized;
 
     /// The type of the addressable storage.
-    type Addressable<'a>: Addressable<T = Self::StorageType>;
+    type Addressable<'a>: Addressable<'a, T = Self::StorageType>;
 
     /// Get addressable storage for indexing into the array.
     fn get_addressable<B: BufferManager>(buffer: &ArrayBuffer<B>) -> Result<Self::Addressable<'_>>;
@@ -744,6 +744,16 @@ impl PhysicalStorage for PhysicalList {
     }
 }
 
+impl MutablePhysicalStorage for PhysicalList {
+    type AddressableMut<'a> = &'a mut [Self::StorageType];
+
+    fn get_addressable_mut<B: BufferManager>(
+        buffer: &mut ArrayBuffer<B>,
+    ) -> Result<Self::AddressableMut<'_>> {
+        buffer.try_as_slice_mut::<Self>()
+    }
+}
+
 /// Dictionary arrays have the selection vector as the primary data buffer.
 #[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
 pub struct PhysicalDictionary;

crates/rayexec_execution/src/arrays/array/string_view.rs

@@ -8,14 +8,14 @@ pub struct StringViewAddressable<'a> {
     pub(crate) heap: &'a StringViewHeap,
 }
 
-impl Addressable for StringViewAddressable<'_> {
+impl<'a> Addressable<'a> for StringViewAddressable<'a> {
     type T = str;
 
     fn len(&self) -> usize {
         self.metadata.len()
     }
 
-    fn get(&self, idx: usize) -> Option<&Self::T> {
+    fn get(&self, idx: usize) -> Option<&'a Self::T> {
         let m = self.metadata.get(idx)?;
         let bs = self.heap.get(m)?;
         Some(unsafe { std::str::from_utf8_unchecked(bs) })
@@ -54,14 +54,14 @@ pub struct BinaryViewAddressable<'a> {
     pub(crate) heap: &'a StringViewHeap,
 }
 
-impl Addressable for BinaryViewAddressable<'_> {
+impl<'a> Addressable<'a> for BinaryViewAddressable<'a> {
     type T = [u8];
 
     fn len(&self) -> usize {
         self.metadata.len()
     }
 
-    fn get(&self, idx: usize) -> Option<&Self::T> {
+    fn get(&self, idx: usize) -> Option<&'a Self::T> {
         let m = self.metadata.get(idx)?;
         self.heap.get(m)
     }

crates/rayexec_execution/src/arrays/compute/cast/array.rs

@@ -88,7 +88,7 @@ pub fn cast_array(arr: &Array, to: DataType, behavior: CastFailBehavior) -> Resu
     let arr = match arr.datatype() {
         DataType::Null => {
             // Can cast NULL to anything else.
-            let data = to.physical_type()?.zeroed_array_data(arr.logical_len());
+            let data = to.physical_type().zeroed_array_data(arr.logical_len());
             let validity = Bitmap::new_with_all_false(arr.logical_len());
             Array::new_with_validity_and_array_data(to, validity, data)
         }

crates/rayexec_execution/src/arrays/datatype.rs

@@ -1,6 +1,6 @@
 use std::fmt;
 
-use rayexec_error::{not_implemented, OptionExt, RayexecError, Result, ResultExt};
+use rayexec_error::{OptionExt, RayexecError, Result, ResultExt};
 use rayexec_proto::ProtoConv;
 use serde::{Deserialize, Serialize};
 
@@ -445,8 +445,8 @@ impl DataType {
         }
     }
 
-    pub fn physical_type(&self) -> Result<PhysicalType> {
-        Ok(match self {
+    pub fn physical_type(&self) -> PhysicalType {
+        match self {
             DataType::Null => PhysicalType::UntypedNull,
             DataType::Boolean => PhysicalType::Boolean,
             DataType::Int8 => PhysicalType::Int8,
@@ -470,9 +470,9 @@ impl DataType {
             DataType::Interval => PhysicalType::Interval,
             DataType::Utf8 => PhysicalType::Utf8,
             DataType::Binary => PhysicalType::Binary,
-            DataType::Struct(_) => not_implemented!("struct data type to physical type"),
+            DataType::Struct(_) => PhysicalType::Struct,
             DataType::List(_) => PhysicalType::List,
-        })
+        }
     }
 
     /// Return if this datatype is null.

crates/rayexec_execution/src/arrays/executor/scalar/fill.rs

@@ -153,7 +153,7 @@ pub(crate) fn concat_with_exact_total_len(arrays: &[&Array], total_len: usize) -
         None => return Err(RayexecError::new("Cannot concat zero arrays")),
     };
 
-    match datatype.physical_type()? {
+    match datatype.physical_type() {
         PhysicalType::UntypedNull => Ok(Array {
             datatype: datatype.clone(),
             selection2: None,
@@ -380,7 +380,7 @@ pub fn interleave(arrays: &[&Array], indices: &[(usize, usize)]) -> Result<Array
         None => return Err(RayexecError::new("Cannot interleave zero arrays")),
     };
 
-    match datatype.physical_type()? {
+    match datatype.physical_type() {
         PhysicalType::UntypedNull => Ok(Array {
             datatype: datatype.clone(),
             selection2: None,

crates/rayexec_execution/src/arrays/testutil.rs

@@ -68,7 +68,7 @@ pub fn assert_arrays_eq_sel(
     let flat1 = array1.flat_view().unwrap();
     let flat2 = array2.flat_view().unwrap();
 
-    match array1.datatype.physical_type().unwrap() {
+    match array1.datatype.physical_type() {
         PhysicalType::Boolean => {
             assert_arrays_eq_sel_inner::<PhysicalBool>(flat1, sel1, flat2, sel2)
         }

crates/rayexec_execution/src/functions/aggregate/builtin/first.rs

@@ -81,7 +81,7 @@ impl AggregateFunction for First {
 
         let datatype = inputs[0].datatype(table_list)?;
 
-        let function_impl: Box<dyn AggregateFunctionImpl> = match datatype.physical_type()? {
+        let function_impl: Box<dyn AggregateFunctionImpl> = match datatype.physical_type() {
             PhysicalType::UntypedNull => Box::new(FirstUntypedNullImpl),
             PhysicalType::Boolean => Box::new(FirstBoolImpl),
             PhysicalType::Float16 => Box::new(FirstPrimitiveImpl::<PhysicalF16, f16>::new(

crates/rayexec_execution/src/functions/aggregate/builtin/minmax.rs

@@ -81,7 +81,7 @@ impl AggregateFunction for Min {
 
         let datatype = inputs[0].datatype(table_list)?;
 
-        let function_impl: Box<dyn AggregateFunctionImpl> = match datatype.physical_type()? {
+        let function_impl: Box<dyn AggregateFunctionImpl> = match datatype.physical_type() {
             PhysicalType::UntypedNull => Box::new(MinMaxUntypedNull),
             PhysicalType::Boolean => Box::new(MinBoolImpl::new()),
             PhysicalType::Float16 => {
@@ -175,7 +175,7 @@ impl AggregateFunction for Max {
 
         let datatype = inputs[0].datatype(table_list)?;
 
-        let function_impl: Box<dyn AggregateFunctionImpl> = match datatype.physical_type()? {
+        let function_impl: Box<dyn AggregateFunctionImpl> = match datatype.physical_type() {
             PhysicalType::UntypedNull => Box::new(MinMaxUntypedNull),
             PhysicalType::Boolean => Box::new(MaxBoolImpl::new()),
             PhysicalType::Float16 => {

crates/rayexec_execution/src/functions/scalar/builtin/comparison.rs

@@ -585,7 +585,7 @@ fn new_comparison_impl<O: ComparisonOperation>(
             }
             (DataType::List(m1), DataType::List(m2)) if m1 == m2 => {
                 // TODO: We'll want to figure out casting for lists.
-                Box::new(ListComparisonImpl::<O>::new(m1.datatype.physical_type()?))
+                Box::new(ListComparisonImpl::<O>::new(m1.datatype.physical_type()))
             }
             (a, b) => return Err(invalid_input_types_error(func, &[a, b])),
         },