feat: ArrayOperations infallible, eager validation + new_unchecked (#4177)

ArrayOperations currently return VortexResult<>, but they really should
just be infallible. A failed array op is generally indicative of
programmer or encoding error. There's really nothing interesting we can
do to handle an out-of-bounds slice() or scalar_at.

There's a lot that falls out of this, like fixing a bunch of tests,
tweaking our scalar value casting to return Option instead of Result,
etc.

---------

Signed-off-by: Andrew Duffy <[email protected]>

Author: a10y

Cargo.lock

@@ -94,13 +94,13 @@ fn random_access(
             .vortex_expect("could not get DOLocationID");
         for (idx, loc) in [90i32, 249, 230, 79, 239, 236].iter().enumerate() {
             assert_eq!(
-                pu_location_id.scalar_at(idx).vortex_expect("scalar_at"),
+                pu_location_id.scalar_at(idx),
                 Scalar::primitive(*loc, NonNullable)
             );
         }
         for (idx, loc) in [164i32, 231, 25, 224, 243, 239].iter().enumerate() {
             assert_eq!(
-                do_location_id.scalar_at(idx).vortex_expect("scalar_at"),
+                do_location_id.scalar_at(idx),
                 Scalar::primitive(*loc, NonNullable)
             );
         }

bench-vortex/src/bin/random_access.rs

@@ -10,8 +10,9 @@ use vortex_array::vtable::{
 };
 use vortex_array::{Array, ArrayRef, Canonical, EncodingId, EncodingRef, vtable};
 use vortex_dtype::{DType, PType};
-use vortex_error::{VortexResult, vortex_bail};
+use vortex_error::{VortexExpect, VortexResult, vortex_ensure};
 
+use crate::ALPFloat;
 use crate::alp::{Exponents, decompress};
 
 vtable!(ALP);
@@ -51,17 +52,150 @@ pub struct ALPArray {
 pub struct ALPEncoding;
 
 impl ALPArray {
-    // TODO(ngates): remove try_new and panic on wrong DType?
+    fn validate(
+        encoded: &dyn Array,
+        exponents: Exponents,
+        patches: Option<&Patches>,
+    ) -> VortexResult<()> {
+        vortex_ensure!(
+            matches!(
+                encoded.dtype(),
+                DType::Primitive(PType::I32 | PType::I64, _)
+            ),
+            "ALP encoded ints have invalid DType {}",
+            encoded.dtype(),
+        );
+
+        // Validate exponents are in-bounds for the float, and that patches have the proper
+        // length and type.
+        let Exponents { e, f } = exponents;
+        match encoded.dtype().as_ptype() {
+            PType::I32 => {
+                vortex_ensure!(exponents.e <= f32::MAX_EXPONENT, "e out of bounds: {e}");
+                vortex_ensure!(exponents.f <= f32::MAX_EXPONENT, "f out of bounds: {f}");
+                if let Some(patches) = patches {
+                    Self::validate_patches::<f32>(patches, encoded)?;
+                }
+            }
+            PType::I64 => {
+                vortex_ensure!(e <= f64::MAX_EXPONENT, "e out of bounds: {e}");
+                vortex_ensure!(f <= f64::MAX_EXPONENT, "f out of bounds: {f}");
+
+                if let Some(patches) = patches {
+                    Self::validate_patches::<f64>(patches, encoded)?;
+                }
+            }
+            _ => unreachable!(),
+        }
+
+        // Validate patches
+        if let Some(patches) = patches {
+            vortex_ensure!(
+                patches.array_len() == encoded.len(),
+                "patches array_len != encoded len: {} != {}",
+                patches.array_len(),
+                encoded.len()
+            );
+
+            // Verify that the patches DType are of the proper DType.
+        }
+
+        Ok(())
+    }
+
+    /// Validate that any patches provided are valid for the ALPArray.
+    fn validate_patches<T: ALPFloat>(patches: &Patches, encoded: &dyn Array) -> VortexResult<()> {
+        vortex_ensure!(
+            patches.array_len() == encoded.len(),
+            "patches array_len != encoded len: {} != {}",
+            patches.array_len(),
+            encoded.len()
+        );
+
+        let expected_type = DType::Primitive(T::PTYPE, encoded.dtype().nullability());
+        vortex_ensure!(
+            patches.dtype() == &expected_type,
+            "Expected patches type {expected_type}, actual {}",
+            patches.dtype(),
+        );
+
+        Ok(())
+    }
+}
+
+impl ALPArray {
+    /// Build a new `ALPArray` from components, panicking on validation failure.
+    ///
+    /// See [`ALPArray::try_new`] for reference on preconditions that must pass before
+    /// calling this method.
+    pub fn new(encoded: ArrayRef, exponents: Exponents, patches: Option<Patches>) -> Self {
+        Self::try_new(encoded, exponents, patches).vortex_expect("ALPArray new")
+    }
+
+    /// Build a new `ALPArray` from components:
+    ///
+    /// * `encoded` contains the ALP-encoded ints. Any null values are replaced with placeholders
+    /// * `exponents` are the ALP exponents, valid range depends on the data type
+    /// * `patches` are any patch values that don't cleanly encode using the ALP conversion function
+    ///
+    /// This method validates the inputs and will return an error if any validation fails.
+    ///
+    /// # Validation
+    ///
+    /// * The `encoded` array must be either `i32` or `i64`
+    ///     * If `i32`, any `patches` must have DType `f32` with same nullability
+    ///     * If `i64`, then `patches`must have DType `f64` with same nullability
+    /// * `exponents` must be in the valid range depending on if the ALPArray is of type `f32` or
+    ///   `f64`.
+    /// * `patches` must have an `array_len` equal to the length of `encoded`
+    ///
+    /// Any failure of these preconditions will result in an error being returned.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use vortex_alp::{ALPArray, Exponents};
+    /// # use vortex_array::IntoArray;
+    /// # use vortex_buffer::buffer;
+    ///
+    /// // Returns error because buffer has wrong PType.
+    /// let result = ALPArray::try_new(
+    ///     buffer![1i8].into_array(),
+    ///     Exponents { e: 1, f: 1 },
+    ///     None
+    /// );
+    /// assert!(result.is_err());
+    ///
+    /// // Returns error because Exponents are out of bounds for f32
+    /// let result = ALPArray::try_new(
+    ///     buffer![1i32, 2i32].into_array(),
+    ///     Exponents { e: 100, f: 100 },
+    ///     None
+    /// );
+    /// assert!(result.is_err());
+    ///
+    /// // Success!
+    /// let value = ALPArray::try_new(
+    ///     buffer![0i32].into_array(),
+    ///     Exponents { e: 1, f: 1 },
+    ///     None
+    /// ).unwrap();
+    ///
+    /// assert_eq!(value.scalar_at(0), 0f32.into());
+    /// ```
     pub fn try_new(
         encoded: ArrayRef,
         exponents: Exponents,
         patches: Option<Patches>,
     ) -> VortexResult<Self> {
+        Self::validate(&encoded, exponents, patches.as_ref())?;
+
         let dtype = match encoded.dtype() {
             DType::Primitive(PType::I32, nullability) => DType::Primitive(PType::F32, *nullability),
             DType::Primitive(PType::I64, nullability) => DType::Primitive(PType::F64, *nullability),
-            d => vortex_bail!(MismatchedTypes: "int32 or int64", d),
+            _ => unreachable!(),
         };
+
         Ok(Self {
             dtype,
             encoded,
@@ -71,6 +205,25 @@ impl ALPArray {
         })
     }
 
+    /// Build a new `ALPArray` from components without validation.
+    ///
+    /// See [`ALPArray::try_new`] for information about the preconditions that should be checked
+    /// **before** calling this method.
+    pub(crate) unsafe fn new_unchecked(
+        encoded: ArrayRef,
+        exponents: Exponents,
+        patches: Option<Patches>,
+        dtype: DType,
+    ) -> Self {
+        Self {
+            dtype,
+            encoded,
+            exponents,
+            patches,
+            stats_set: Default::default(),
+        }
+    }
+
     pub fn ptype(&self) -> PType {
         self.dtype.as_ptype()
     }

encodings/alp/src/alp/array.rs

@@ -43,7 +43,15 @@ pub fn alp_encode(parray: &PrimitiveArray, exponents: Option<Exponents>) -> Vort
         _ => vortex_bail!("ALP can only encode f32 and f64"),
     };
 
-    ALPArray::try_new(encoded, exponents, patches)
+    // SAFETY: alp_encode_components_typed must return well-formed components
+    unsafe {
+        Ok(ALPArray::new_unchecked(
+            encoded,
+            exponents,
+            patches,
+            parray.dtype().clone(),
+        ))
+    }
 }
 
 #[allow(clippy::cast_possible_truncation)]
@@ -65,7 +73,7 @@ where
 
     let validity = values.validity_mask()?;
     // exceptional_positions may contain exceptions at invalid positions (which contain garbage
-    // data). We remove invalid exceptional positions in order to keep the Patches small.
+    // data). We remove null exceptions in order to keep the Patches small.
     let (valid_exceptional_positions, valid_exceptional_values): (Buffer<u64>, Buffer<T>) =
         match validity {
             Mask::AllTrue(_) => (exceptional_positions, exceptional_values),
@@ -194,10 +202,10 @@ mod tests {
         assert_eq!(encoded.exponents(), Exponents { e: 16, f: 13 });
 
         let decoded = decompress(&encoded).unwrap();
-        assert_eq!(decoded.scalar_at(0).unwrap(), array.scalar_at(0).unwrap());
-        assert_eq!(decoded.scalar_at(1).unwrap(), array.scalar_at(1).unwrap());
+        assert_eq!(decoded.scalar_at(0), array.scalar_at(0));
+        assert_eq!(decoded.scalar_at(1), array.scalar_at(1));
         assert!(!decoded.is_valid(2).unwrap());
-        assert_eq!(decoded.scalar_at(3).unwrap(), array.scalar_at(3).unwrap());
+        assert_eq!(decoded.scalar_at(3), array.scalar_at(3));
     }
 
     #[test]
@@ -216,12 +224,12 @@ mod tests {
         assert_eq!(encoded.exponents(), Exponents { e: 16, f: 13 });
 
         for idx in 0..3 {
-            let s = encoded.scalar_at(idx).unwrap();
+            let s = encoded.scalar_at(idx);
             assert!(s.is_valid());
         }
 
         assert!(!encoded.is_valid(4).unwrap());
-        let s = encoded.scalar_at(4).unwrap();
+        let s = encoded.scalar_at(4);
         assert!(s.is_null());
 
         let _decoded = decompress(&encoded).unwrap();
@@ -249,9 +257,9 @@ mod tests {
             decompressed.as_slice::<f64>()
         );
         assert_eq!(original.validity(), decompressed.validity());
-        assert_eq!(original.scalar_at(0).unwrap(), Scalar::null_typed::<f64>());
-        assert_eq!(original.scalar_at(1).unwrap(), Scalar::null_typed::<f64>());
-        assert_eq!(original.scalar_at(2).unwrap(), Scalar::null_typed::<f64>());
+        assert_eq!(original.scalar_at(0), Scalar::null_typed::<f64>());
+        assert_eq!(original.scalar_at(1), Scalar::null_typed::<f64>());
+        assert_eq!(original.scalar_at(2), Scalar::null_typed::<f64>());
     }
 
     #[test]

encodings/alp/src/alp/compress.rs

@@ -22,10 +22,18 @@ impl CastKernel for ALPVTable {
                     .with_nullability(dtype.nullability()),
             )?;
 
-            Ok(Some(
-                ALPArray::try_new(new_encoded, array.exponents(), array.patches().cloned())?
+            // SAFETY: casting nullability doesn't alter the invariants
+            unsafe {
+                Ok(Some(
+                    ALPArray::new_unchecked(
+                        new_encoded,
+                        array.exponents(),
+                        array.patches().cloned(),
+                        dtype.clone(),
+                    )
                     .into_array(),
-            ))
+                ))
+            }
         } else {
             Ok(None)
         }

encodings/alp/src/alp/compute/cast.rs

@@ -16,10 +16,16 @@ impl FilterKernel for ALPVTable {
             .transpose()?
             .flatten();
 
-        Ok(
-            ALPArray::try_new(filter(array.encoded(), mask)?, array.exponents(), patches)?
-                .to_array(),
-        )
+        // SAFETY: filtering the values does not change correctness
+        unsafe {
+            Ok(ALPArray::new_unchecked(
+                filter(array.encoded(), mask)?,
+                array.exponents(),
+                patches,
+                array.dtype().clone(),
+            )
+            .to_array())
+        }
     }
 }
 

encodings/alp/src/alp/compute/filter.rs

@@ -24,7 +24,7 @@ impl MaskKernel for ALPVTable {
                 )
             })
             .transpose()?;
-        Ok(ALPArray::try_new(masked_encoded, array.exponents(), masked_patches)?.to_array())
+        Ok(ALPArray::new(masked_encoded, array.exponents(), masked_patches).to_array())
     }
 }
 

encodings/alp/src/alp/compute/mask.rs

@@ -2,7 +2,7 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use vortex_array::compute::{TakeKernel, TakeKernelAdapter, take};
-use vortex_array::{Array, ArrayRef, register_kernel};
+use vortex_array::{Array, ArrayRef, IntoArray, register_kernel};
 use vortex_error::VortexResult;
 
 use crate::{ALPArray, ALPVTable};
@@ -23,7 +23,7 @@ impl TakeKernel for ALPVTable {
                 )
             })
             .transpose()?;
-        Ok(ALPArray::try_new(taken_encoded, array.exponents(), taken_patches)?.to_array())
+        Ok(ALPArray::new(taken_encoded, array.exponents(), taken_patches).into_array())
     }
 }