diff --git a/Cargo.lock b/Cargo.lock index 4086fff65be..bc27facbcc3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -125,7 +125,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -136,7 +136,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -1848,7 +1848,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "117725a109d387c937a1533ce01b450cbde6b88abceea8473c4d7a85853cda3c" dependencies = [ "lazy_static", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -3739,7 +3739,7 @@ dependencies = [ "libc", "option-ext", "redox_users", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -3905,7 +3905,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -4310,8 +4310,8 @@ dependencies = [ "libc", "log", "rustversion", - "windows-link 0.1.3", - "windows-result 0.3.4", + "windows-link 0.2.1", + "windows-result 0.4.1", ] [[package]] @@ -4991,7 +4991,7 @@ dependencies = [ "portable-atomic", "portable-atomic-util", "serde_core", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -6294,7 +6294,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -8275,7 +8275,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.4.15", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -8288,7 +8288,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.11.0", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -9243,7 +9243,7 @@ dependencies = [ "getrandom 0.3.4", "once_cell", "rustix 1.1.3", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -11216,7 +11216,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] diff --git a/encodings/alp/src/alp/array.rs b/encodings/alp/src/alp/array.rs index df61d2018b2..06a22e7b57a 100644 --- a/encodings/alp/src/alp/array.rs +++ b/encodings/alp/src/alp/array.rs @@ -186,7 +186,11 @@ impl VTable for ALPVTable { ALPArray::new( array.encoded().slice(range.clone()), array.exponents(), - array.patches().and_then(|p| p.slice(range)), + array + .patches() + .map(|p| p.slice(range)) + .transpose()? + .flatten(), ) .into_array(), )) @@ -349,7 +353,7 @@ impl ALPArray { /// None /// ).unwrap(); /// - /// assert_eq!(value.scalar_at(0), 0f32.into()); + /// assert_eq!(value.scalar_at(0).unwrap(), 0f32.into()); /// ``` pub fn try_new( encoded: ArrayRef, @@ -696,7 +700,7 @@ mod tests { for idx in 0..slice_len { let expected_value = values[slice_start + idx]; - let result_valid = result_primitive.validity_mask().value(idx); + let result_valid = result_primitive.validity_mask().unwrap().value(idx); assert_eq!( result_valid, expected_value.is_some(), diff --git a/encodings/alp/src/alp/compress.rs b/encodings/alp/src/alp/compress.rs index c792bc8f37e..acf01e4041d 100644 --- a/encodings/alp/src/alp/compress.rs +++ b/encodings/alp/src/alp/compress.rs @@ -75,7 +75,7 @@ where let encoded_array = PrimitiveArray::new(encoded, values.validity().clone()).into_array(); - let validity = values.validity_mask(); + let validity = values.validity_mask()?; // exceptional_positions may contain exceptions at invalid positions (which contain garbage // data). We remove null exceptions in order to keep the Patches small. let (valid_exceptional_positions, valid_exceptional_values): (Buffer, Buffer) = diff --git a/encodings/alp/src/alp/ops.rs b/encodings/alp/src/alp/ops.rs index 2703f588fb1..d17b0dddf02 100644 --- a/encodings/alp/src/alp/ops.rs +++ b/encodings/alp/src/alp/ops.rs @@ -3,6 +3,7 @@ use vortex_array::vtable::OperationsVTable; use vortex_error::VortexExpect; +use vortex_error::VortexResult; use vortex_scalar::Scalar; use crate::ALPArray; @@ -11,16 +12,16 @@ use crate::ALPVTable; use crate::match_each_alp_float_ptype; impl OperationsVTable for ALPVTable { - fn scalar_at(array: &ALPArray, index: usize) -> Scalar { + fn scalar_at(array: &ALPArray, index: usize) -> VortexResult { if let Some(patches) = array.patches() - && let Some(patch) = patches.get_patched(index) + && let Some(patch) = patches.get_patched(index)? { - return patch.cast(array.dtype()).vortex_expect("cast failure"); + return patch.cast(array.dtype()); } - let encoded_val = array.encoded().scalar_at(index); + let encoded_val = array.encoded().scalar_at(index)?; - match_each_alp_float_ptype!(array.ptype(), |T| { + Ok(match_each_alp_float_ptype!(array.ptype(), |T| { let encoded_val: ::ALPInt = encoded_val .as_ref() .try_into() @@ -29,6 +30,6 @@ impl OperationsVTable for ALPVTable { ::decode_single(encoded_val, array.exponents()), array.dtype().nullability(), ) - }) + })) } } diff --git a/encodings/alp/src/alp_rd/array.rs b/encodings/alp/src/alp_rd/array.rs index c2a8e2c56c1..581c0ac884b 100644 --- a/encodings/alp/src/alp_rd/array.rs +++ b/encodings/alp/src/alp_rd/array.rs @@ -82,7 +82,9 @@ impl VTable for ALPRDVTable { fn slice(array: &Self::Array, range: std::ops::Range) -> VortexResult> { let left_parts_exceptions = array .left_parts_patches() - .and_then(|patches| patches.slice(range.clone())); + .map(|patches| patches.slice(range.clone())) + .transpose()? + .flatten(); // SAFETY: slicing components does not change the encoded values Ok(Some(unsafe { @@ -341,7 +343,7 @@ impl ALPRDArray { let left_parts_patches = left_parts_patches .map(|patches| { - if !patches.values().all_valid() { + if !patches.values().all_valid()? { vortex_bail!("patches must be all valid: {}", patches.values()); } // TODO(ngates): assert the DType, don't cast it. diff --git a/encodings/alp/src/alp_rd/ops.rs b/encodings/alp/src/alp_rd/ops.rs index 91e2ae0a71a..cbb702e9ee4 100644 --- a/encodings/alp/src/alp_rd/ops.rs +++ b/encodings/alp/src/alp_rd/ops.rs @@ -4,18 +4,20 @@ use vortex_array::Array; use vortex_array::vtable::OperationsVTable; use vortex_error::VortexExpect; +use vortex_error::VortexResult; use vortex_scalar::Scalar; use crate::ALPRDArray; use crate::ALPRDVTable; impl OperationsVTable for ALPRDVTable { - fn scalar_at(array: &ALPRDArray, index: usize) -> Scalar { + fn scalar_at(array: &ALPRDArray, index: usize) -> VortexResult { // The left value can either be a direct value, or an exception. // The exceptions array represents exception positions with non-null values. - let maybe_patched_value = array - .left_parts_patches() - .and_then(|patches| patches.get_patched(index)); + let maybe_patched_value = match array.left_parts_patches() { + Some(patches) => patches.get_patched(index)?, + None => None, + }; let left = match maybe_patched_value { Some(patched_value) => patched_value .as_primitive() @@ -24,7 +26,7 @@ impl OperationsVTable for ALPRDVTable { _ => { let left_code: u16 = array .left_parts() - .scalar_at(index) + .scalar_at(index)? .as_primitive() .as_::() .vortex_expect("left_code must be non-null"); @@ -33,10 +35,10 @@ impl OperationsVTable for ALPRDVTable { }; // combine left and right values - if array.is_f32() { + Ok(if array.is_f32() { let right: u32 = array .right_parts() - .scalar_at(index) + .scalar_at(index)? .as_primitive() .as_::() .vortex_expect("non-null"); @@ -45,13 +47,13 @@ impl OperationsVTable for ALPRDVTable { } else { let right: u64 = array .right_parts() - .scalar_at(index) + .scalar_at(index)? .as_primitive() .as_::() .vortex_expect("non-null"); let packed = f64::from_bits(((left as u64) << array.right_bit_width()) | right); Scalar::primitive(packed, array.dtype().nullability()) - } + }) } } diff --git a/encodings/bytebool/src/array.rs b/encodings/bytebool/src/array.rs index 813d1b0539b..1b70568d4a0 100644 --- a/encodings/bytebool/src/array.rs +++ b/encodings/bytebool/src/array.rs @@ -219,8 +219,11 @@ impl BaseArrayVTable for ByteBoolVTable { } impl OperationsVTable for ByteBoolVTable { - fn scalar_at(array: &ByteBoolArray, index: usize) -> Scalar { - Scalar::bool(array.buffer()[index] == 1, array.dtype().nullability()) + fn scalar_at(array: &ByteBoolArray, index: usize) -> VortexResult { + Ok(Scalar::bool( + array.buffer()[index] == 1, + array.dtype().nullability(), + )) } } @@ -275,14 +278,14 @@ mod tests { assert_eq!(v_len, arr.len()); for idx in 0..arr.len() { - assert!(arr.is_valid(idx)); + assert!(arr.is_valid(idx).unwrap()); } let v = vec![Some(true), None, Some(false)]; let arr = ByteBoolArray::from(v); - assert!(arr.is_valid(0)); - assert!(!arr.is_valid(1)); - assert!(arr.is_valid(2)); + assert!(arr.is_valid(0).unwrap()); + assert!(!arr.is_valid(1).unwrap()); + assert!(arr.is_valid(2).unwrap()); assert_eq!(arr.len(), 3); let v: Vec> = vec![None, None]; @@ -292,7 +295,7 @@ mod tests { assert_eq!(v_len, arr.len()); for idx in 0..arr.len() { - assert!(!arr.is_valid(idx)); + assert!(!arr.is_valid(idx).unwrap()); } assert_eq!(arr.len(), 2); } diff --git a/encodings/datetime-parts/src/canonical.rs b/encodings/datetime-parts/src/canonical.rs index cf503cefd93..9268238f1d4 100644 --- a/encodings/datetime-parts/src/canonical.rs +++ b/encodings/datetime-parts/src/canonical.rs @@ -96,7 +96,7 @@ pub fn decode_to_temporal( } Ok(TemporalArray::new_timestamp( - PrimitiveArray::new(values.freeze(), Validity::copy_from_array(array.as_ref())) + PrimitiveArray::new(values.freeze(), Validity::copy_from_array(array.as_ref())?) .into_array(), temporal_metadata.time_unit(), temporal_metadata.time_zone().map(ToString::to_string), @@ -147,7 +147,7 @@ mod test { .unwrap(); assert_eq!( - date_times.validity_mask(), + date_times.validity_mask().unwrap(), validity.to_mask(date_times.len()) ); diff --git a/encodings/datetime-parts/src/ops.rs b/encodings/datetime-parts/src/ops.rs index cd81925cbe1..8e3095d6c8b 100644 --- a/encodings/datetime-parts/src/ops.rs +++ b/encodings/datetime-parts/src/ops.rs @@ -6,6 +6,7 @@ use vortex_array::vtable::OperationsVTable; use vortex_dtype::DType; use vortex_dtype::datetime::TemporalMetadata; use vortex_error::VortexExpect; +use vortex_error::VortexResult; use vortex_error::vortex_panic; use vortex_scalar::Scalar; @@ -15,7 +16,7 @@ use crate::timestamp; use crate::timestamp::TimestampParts; impl OperationsVTable for DateTimePartsVTable { - fn scalar_at(array: &DateTimePartsArray, index: usize) -> Scalar { + fn scalar_at(array: &DateTimePartsArray, index: usize) -> VortexResult { let DType::Extension(ext) = array.dtype().clone() else { vortex_panic!( "DateTimePartsArray must have extension dtype, found {}", @@ -27,25 +28,25 @@ impl OperationsVTable for DateTimePartsVTable { vortex_panic!(ComputeError: "must decode TemporalMetadata from extension metadata"); }; - if !array.is_valid(index) { - return Scalar::null(DType::Extension(ext)); + if !array.is_valid(index)? { + return Ok(Scalar::null(DType::Extension(ext))); } let days: i64 = array .days() - .scalar_at(index) + .scalar_at(index)? .as_primitive() .as_::() .vortex_expect("days fits in i64"); let seconds: i64 = array .seconds() - .scalar_at(index) + .scalar_at(index)? .as_primitive() .as_::() .vortex_expect("seconds fits in i64"); let subseconds: i64 = array .subseconds() - .scalar_at(index) + .scalar_at(index)? .as_primitive() .as_::() .vortex_expect("subseconds fits in i64"); @@ -59,6 +60,6 @@ impl OperationsVTable for DateTimePartsVTable { temporal_metadata.time_unit(), ); - Scalar::extension(ext, Scalar::from(ts)) + Ok(Scalar::extension(ext, Scalar::from(ts))) } } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs index 0870ccd704e..ca7ff934be7 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs @@ -59,7 +59,7 @@ impl CompareKernel for DecimalBytePartsVTable { // (depending on the `sign`) than all values in MSP. // If the LHS or the RHS contain nulls, then we must fallback to the canonicalized // implementation which does null-checking instead. - if lhs.all_valid() && rhs.all_valid() { + if lhs.all_valid()? && rhs.all_valid()? { Ok(Some( ConstantArray::new( unconvertible_value(sign, operator, nullability), diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs index eaf21ad548c..94db99d61be 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -257,15 +257,18 @@ fn to_canonical_decimal( } impl OperationsVTable for DecimalBytePartsVTable { - fn scalar_at(array: &DecimalBytePartsArray, index: usize) -> Scalar { + fn scalar_at(array: &DecimalBytePartsArray, index: usize) -> VortexResult { // TODO(joe): support parts len != 1 - let scalar = array.msp.scalar_at(index); + let scalar = array.msp.scalar_at(index)?; // Note. values in msp, can only be signed integers upto size i64. let primitive_scalar = scalar.as_primitive(); // TODO(joe): extend this to support multiple parts. let value = primitive_scalar.as_::().vortex_expect("non-null"); - Scalar::new(array.dtype.clone(), DecimalValue::I64(value).into()) + Ok(Scalar::new( + array.dtype.clone(), + DecimalValue::I64(value).into(), + )) } } @@ -314,14 +317,14 @@ mod tests { .unwrap() .to_array(); - assert_eq!(Scalar::null(dtype.clone()), array.scalar_at(0)); + assert_eq!(Scalar::null(dtype.clone()), array.scalar_at(0).unwrap()); assert_eq!( Scalar::new(dtype.clone(), DecimalValue::I64(200).into()), - array.scalar_at(1) + array.scalar_at(1).unwrap() ); assert_eq!( Scalar::new(dtype, DecimalValue::I64(400).into()), - array.scalar_at(2) + array.scalar_at(2).unwrap() ); } } diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs index b9de830748f..441db3682c6 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs @@ -210,7 +210,7 @@ pub fn gather_patches( }; let array_len = parray.len(); - let validity_mask = parray.validity_mask(); + let validity_mask = parray.validity_mask()?; let patches = if array_len < u8::MAX as usize { match_each_integer_ptype!(parray.ptype(), |T| { @@ -310,7 +310,7 @@ fn bit_width_histogram_typed( |v: T| (8 * size_of::()) - (PrimInt::leading_zeros(v) as usize); let mut bit_widths = vec![0usize; size_of::() * 8 + 1]; - match array.validity_mask().bit_buffer() { + match array.validity_mask()?.bit_buffer() { AllOr::All => { // All values are valid. for v in array.as_slice::() { @@ -464,6 +464,7 @@ mod test { (0..(1 << 4)).collect::>(), compressed .validity_mask() + .unwrap() .to_bit_buffer() .set_indices() .collect::>() diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs index dcecb856383..b4edd177fd7 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs @@ -46,7 +46,10 @@ pub fn unpack_to_pvector(array: &BitPackedArray) -> PVectorMut

// SAFETY: `decode_into` initialized exactly `len` elements into the spare (existing) capacity. unsafe { elements.set_len(len) }; - let mut validity = array.validity_mask().into_mut(); + let mut validity = array + .validity_mask() + .vortex_expect("validity_mask") + .into_mut(); debug_assert_eq!(validity.len(), len); // TODO(connor): Implement a fused version of patching instead. @@ -97,7 +100,7 @@ pub(crate) fn unpack_into_primitive_builder( // SAFETY: We later initialize the the uninitialized range of values with `copy_from_slice`. unsafe { // Append a dense null Mask. - uninit_range.append_mask(array.validity_mask()); + uninit_range.append_mask(array.validity_mask()?); } // SAFETY: `decode_into` will initialize all values in this range. @@ -136,7 +139,7 @@ pub fn apply_patches_to_uninit_range_fn T>( let indices = patches.indices().clone().execute::(ctx)?; let values = patches.values().clone().execute::(ctx)?; - let validity = values.validity_mask(); + let validity = values.validity_mask()?; let values = values.as_slice::(); match_each_unsigned_integer_ptype!(indices.ptype(), |P| { @@ -415,11 +418,11 @@ mod tests { // Verify the validity mask was correctly applied. assert_eq!(result.len(), 5); - assert!(!result.scalar_at(0).is_null()); - assert!(result.scalar_at(1).is_null()); - assert!(!result.scalar_at(2).is_null()); - assert!(!result.scalar_at(3).is_null()); - assert!(result.scalar_at(4).is_null()); + assert!(!result.scalar_at(0).unwrap().is_null()); + assert!(result.scalar_at(1).unwrap().is_null()); + assert!(!result.scalar_at(2).unwrap().is_null()); + assert!(!result.scalar_at(3).unwrap().is_null()); + assert!(result.scalar_at(4).unwrap().is_null()); Ok(()) } diff --git a/encodings/fastlanes/src/bitpacking/compute/take.rs b/encodings/fastlanes/src/bitpacking/compute/take.rs index e279327786d..8b90466ac7b 100644 --- a/encodings/fastlanes/src/bitpacking/compute/take.rs +++ b/encodings/fastlanes/src/bitpacking/compute/take.rs @@ -224,11 +224,11 @@ mod test { .enumerate() .for_each(|(ti, i)| { assert_eq!( - u32::try_from(&packed.scalar_at(*i as usize)).unwrap(), + u32::try_from(&packed.scalar_at(*i as usize).unwrap()).unwrap(), values[*i as usize] ); assert_eq!( - u32::try_from(&taken.scalar_at(ti)).unwrap(), + u32::try_from(&taken.scalar_at(ti).unwrap()).unwrap(), values[*i as usize] ); }); @@ -263,7 +263,7 @@ mod test { taken_primitive, PrimitiveArray::from_option_iter([Some(1i32), Some(2), None, Some(4)]) ); - assert_eq!(taken_primitive.to_primitive().invalid_count(), 1); + assert_eq!(taken_primitive.to_primitive().invalid_count().unwrap(), 1); } #[rstest] diff --git a/encodings/fastlanes/src/bitpacking/vtable/kernels/filter.rs b/encodings/fastlanes/src/bitpacking/vtable/kernels/filter.rs index 7d6b00e2b01..e99360544a4 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/kernels/filter.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/kernels/filter.rs @@ -144,7 +144,7 @@ fn filter_primitive_without_patches( ) -> VortexResult> { let values = filter_with_indices(array, selection.indices()); let validity = array - .validity_mask() + .validity_mask()? .filter(&Mask::Values(selection.clone())) .into_mut(); diff --git a/encodings/fastlanes/src/bitpacking/vtable/mod.rs b/encodings/fastlanes/src/bitpacking/vtable/mod.rs index d4085c5443f..47c06f7fd78 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/mod.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/mod.rs @@ -299,7 +299,11 @@ impl VTable for BitPackedVTable { array.packed().slice(encoded_start..encoded_stop), array.dtype.clone(), array.validity().slice(range.clone()), - array.patches().and_then(|p| p.slice(range.clone())), + array + .patches() + .map(|p| p.slice(range.clone())) + .transpose()? + .flatten(), array.bit_width(), range.len(), offset as u16, diff --git a/encodings/fastlanes/src/bitpacking/vtable/operations.rs b/encodings/fastlanes/src/bitpacking/vtable/operations.rs index b3ec21345f7..a1c2f3f4cd4 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/operations.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/operations.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_array::vtable::OperationsVTable; +use vortex_error::VortexResult; use vortex_scalar::Scalar; use crate::BitPackedArray; @@ -9,14 +10,16 @@ use crate::BitPackedVTable; use crate::bitpack_decompress; impl OperationsVTable for BitPackedVTable { - fn scalar_at(array: &BitPackedArray, index: usize) -> Scalar { - if let Some(patches) = array.patches() - && let Some(patch) = patches.get_patched(index) - { - patch - } else { - bitpack_decompress::unpack_single(array, index) - } + fn scalar_at(array: &BitPackedArray, index: usize) -> VortexResult { + Ok( + if let Some(patches) = array.patches() + && let Some(patch) = patches.get_patched(index)? + { + patch + } else { + bitpack_decompress::unpack_single(array, index) + }, + ) } } @@ -175,7 +178,7 @@ mod test { .into_array() }; assert_eq!( - packed_array.scalar_at(1), + packed_array.scalar_at(1).unwrap(), Scalar::null(DType::Primitive(PType::U32, Nullability::Nullable)) ); } @@ -188,7 +191,10 @@ mod test { assert!(packed.patches().is_some()); let patches = packed.patches().unwrap().indices().clone(); - assert_eq!(usize::try_from(&patches.scalar_at(0)).unwrap(), 256); + assert_eq!( + usize::try_from(&patches.scalar_at(0).unwrap()).unwrap(), + 256 + ); let expected = PrimitiveArray::from_iter(values.iter().copied()); assert_arrays_eq!(packed, expected); diff --git a/encodings/fastlanes/src/delta/array/delta_decompress.rs b/encodings/fastlanes/src/delta/array/delta_decompress.rs index abd3968f439..b2deabaae54 100644 --- a/encodings/fastlanes/src/delta/array/delta_decompress.rs +++ b/encodings/fastlanes/src/delta/array/delta_decompress.rs @@ -30,7 +30,8 @@ pub fn delta_decompress( // TODO(connor): This is incorrect, we need to untranspose the validity!!! - let validity = Validity::from_mask(array.deltas().validity_mask(), array.dtype().nullability()); + let validity = + Validity::from_mask(array.deltas().validity_mask()?, array.dtype().nullability()); let validity = validity.slice(start..end); Ok(match_each_unsigned_integer_ptype!(deltas.ptype(), |T| { diff --git a/encodings/fastlanes/src/delta/vtable/operations.rs b/encodings/fastlanes/src/delta/vtable/operations.rs index 96947350fc0..08875a0e13b 100644 --- a/encodings/fastlanes/src/delta/vtable/operations.rs +++ b/encodings/fastlanes/src/delta/vtable/operations.rs @@ -3,13 +3,14 @@ use vortex_array::ToCanonical; use vortex_array::vtable::OperationsVTable; +use vortex_error::VortexResult; use vortex_scalar::Scalar; use super::DeltaVTable; use crate::DeltaArray; impl OperationsVTable for DeltaVTable { - fn scalar_at(array: &DeltaArray, index: usize) -> Scalar { + fn scalar_at(array: &DeltaArray, index: usize) -> VortexResult { let decompressed = array.slice(index..index + 1).to_primitive(); decompressed.scalar_at(0) } @@ -196,7 +197,7 @@ mod tests { let delta = DeltaArray::try_from_vec((0u32..2048).collect()) .unwrap() .into_array(); - delta.scalar_at(2048); + delta.scalar_at(2048).unwrap(); } #[test] fn test_scalar_at_jagged_array() { @@ -214,7 +215,7 @@ mod tests { let delta = DeltaArray::try_from_vec((0u32..2000).collect()) .unwrap() .into_array(); - delta.scalar_at(2000); + delta.scalar_at(2000).unwrap(); } #[rstest] diff --git a/encodings/fastlanes/src/for/array/for_compress.rs b/encodings/fastlanes/src/for/array/for_compress.rs index 7b4ef7c111c..4b20acc5b77 100644 --- a/encodings/fastlanes/src/for/array/for_compress.rs +++ b/encodings/fastlanes/src/for/array/for_compress.rs @@ -113,7 +113,7 @@ mod test { assert!(compressed.reference_scalar().dtype().is_signed_int()); assert!(compressed.encoded().dtype().is_signed_int()); - let encoded = compressed.encoded().scalar_at(0); + let encoded = compressed.encoded().scalar_at(0).unwrap(); assert_eq!(encoded, Scalar::from(0i32)); } @@ -175,7 +175,10 @@ mod test { .iter() .enumerate() .for_each(|(i, v)| { - assert_eq!(*v, i8::try_from(compressed.scalar_at(i).as_ref()).unwrap()); + assert_eq!( + *v, + i8::try_from(compressed.scalar_at(i).unwrap().as_ref()).unwrap() + ); }); assert_arrays_eq!(decompressed, array); Ok(()) diff --git a/encodings/fastlanes/src/for/array/for_decompress.rs b/encodings/fastlanes/src/for/array/for_decompress.rs index 6661b7005ed..4677c90a655 100644 --- a/encodings/fastlanes/src/for/array/for_decompress.rs +++ b/encodings/fastlanes/src/for/array/for_decompress.rs @@ -109,7 +109,7 @@ pub(crate) fn fused_decompress< let mut uninit_range = builder.uninit_range(bp.len()); unsafe { // Append a dense null Mask. - uninit_range.append_mask(bp.validity_mask()); + uninit_range.append_mask(bp.validity_mask()?); } // SAFETY: `decode_into` will initialize all values in this range. diff --git a/encodings/fastlanes/src/for/vtable/operations.rs b/encodings/fastlanes/src/for/vtable/operations.rs index 2dff803197d..b89282c289c 100644 --- a/encodings/fastlanes/src/for/vtable/operations.rs +++ b/encodings/fastlanes/src/for/vtable/operations.rs @@ -4,19 +4,20 @@ use vortex_array::vtable::OperationsVTable; use vortex_dtype::match_each_integer_ptype; use vortex_error::VortexExpect; +use vortex_error::VortexResult; use vortex_scalar::Scalar; use super::FoRVTable; use crate::FoRArray; impl OperationsVTable for FoRVTable { - fn scalar_at(array: &FoRArray, index: usize) -> Scalar { - let encoded_pvalue = array.encoded().scalar_at(index); + fn scalar_at(array: &FoRArray, index: usize) -> VortexResult { + let encoded_pvalue = array.encoded().scalar_at(index)?; let encoded_pvalue = encoded_pvalue.as_primitive(); let reference = array.reference_scalar(); let reference = reference.as_primitive(); - match_each_integer_ptype!(array.ptype(), |P| { + Ok(match_each_integer_ptype!(array.ptype(), |P| { encoded_pvalue .typed_value::

() .map(|v| { @@ -28,7 +29,7 @@ impl OperationsVTable for FoRVTable { }) .map(|v| Scalar::primitive::

(v, array.reference_scalar().dtype().nullability())) .unwrap_or_else(|| Scalar::null(array.reference_scalar().dtype().clone())) - }) + })) } } diff --git a/encodings/fastlanes/src/rle/array/mod.rs b/encodings/fastlanes/src/rle/array/mod.rs index f76273b18b2..38afab37d54 100644 --- a/encodings/fastlanes/src/rle/array/mod.rs +++ b/encodings/fastlanes/src/rle/array/mod.rs @@ -188,12 +188,14 @@ impl RLEArray { pub(crate) fn values_idx_offset(&self, chunk_idx: usize) -> usize { self.values_idx_offsets .scalar_at(chunk_idx) + .expect("index must be in bounds") .as_primitive() .as_::() .expect("index must be of type usize") - self .values_idx_offsets .scalar_at(0) + .expect("index must be in bounds") .as_primitive() .as_::() .expect("index must be of type usize") @@ -279,9 +281,9 @@ mod tests { assert_eq!(rle_array.len(), 3); assert_eq!(rle_array.values().len(), 2); - assert!(rle_array.is_valid(0)); - assert!(!rle_array.is_valid(1)); - assert!(rle_array.is_valid(2)); + assert!(rle_array.is_valid(0).unwrap()); + assert!(!rle_array.is_valid(1).unwrap()); + assert!(rle_array.is_valid(2).unwrap()); } #[test] @@ -315,10 +317,10 @@ mod tests { let valid_slice = rle_array.slice(0..3).to_primitive(); // TODO(joe): replace with compute null count - assert!(valid_slice.all_valid()); + assert!(valid_slice.all_valid().unwrap()); let mixed_slice = rle_array.slice(1..5); - assert!(!mixed_slice.all_valid()); + assert!(!mixed_slice.all_valid().unwrap()); } #[test] @@ -356,10 +358,10 @@ mod tests { .to_canonical() .unwrap() .into_primitive(); - assert!(invalid_slice.all_invalid()); + assert!(invalid_slice.all_invalid().unwrap()); let mixed_slice = rle_array.slice(1..4); - assert!(!mixed_slice.all_invalid()); + assert!(!mixed_slice.all_invalid().unwrap()); } #[test] @@ -392,7 +394,7 @@ mod tests { .unwrap(); let sliced_array = rle_array.slice(1..4); - let validity_mask = sliced_array.validity_mask(); + let validity_mask = sliced_array.validity_mask().unwrap(); let expected_mask = Validity::from_iter([false, true, false]).to_mask(3); assert_eq!(validity_mask.len(), expected_mask.len()); diff --git a/encodings/fastlanes/src/rle/array/rle_decompress.rs b/encodings/fastlanes/src/rle/array/rle_decompress.rs index 391c81414c9..c8666cc427f 100644 --- a/encodings/fastlanes/src/rle/array/rle_decompress.rs +++ b/encodings/fastlanes/src/rle/array/rle_decompress.rs @@ -102,6 +102,6 @@ where buffer .freeze() .slice(offset_within_chunk..(offset_within_chunk + array.len())), - Validity::copy_from_array(array.as_ref()), + Validity::copy_from_array(array.as_ref())?, )) } diff --git a/encodings/fastlanes/src/rle/vtable/operations.rs b/encodings/fastlanes/src/rle/vtable/operations.rs index 4a5797df683..77fec3c50b0 100644 --- a/encodings/fastlanes/src/rle/vtable/operations.rs +++ b/encodings/fastlanes/src/rle/vtable/operations.rs @@ -3,6 +3,7 @@ use vortex_array::vtable::OperationsVTable; use vortex_error::VortexExpect; +use vortex_error::VortexResult; use vortex_scalar::Scalar; use super::RLEVTable; @@ -10,9 +11,9 @@ use crate::FL_CHUNK_SIZE; use crate::RLEArray; impl OperationsVTable for RLEVTable { - fn scalar_at(array: &RLEArray, index: usize) -> Scalar { + fn scalar_at(array: &RLEArray, index: usize) -> VortexResult { let offset_in_chunk = array.offset(); - let chunk_relative_idx = array.indices().scalar_at(offset_in_chunk + index); + let chunk_relative_idx = array.indices().scalar_at(offset_in_chunk + index)?; let chunk_relative_idx = chunk_relative_idx .as_primitive() @@ -24,9 +25,9 @@ impl OperationsVTable for RLEVTable { let scalar = array .values() - .scalar_at(value_idx_offset + chunk_relative_idx); + .scalar_at(value_idx_offset + chunk_relative_idx)?; - Scalar::new(array.dtype().clone(), scalar.into_value()) + Ok(Scalar::new(array.dtype().clone(), scalar.into_value())) } } @@ -167,7 +168,12 @@ mod tests { for &idx in &[1023, 1024, 1025, 2047, 2048, 2049] { if idx < encoded.len() { let original_value = expected[idx]; - let encoded_value = encoded.scalar_at(idx).as_primitive().as_::().unwrap(); + let encoded_value = encoded + .scalar_at(idx) + .unwrap() + .as_primitive() + .as_::() + .unwrap(); assert_eq!(original_value, encoded_value, "Mismatch at index {}", idx); } } @@ -177,14 +183,14 @@ mod tests { #[should_panic] fn test_scalar_at_out_of_bounds() { let array = fixture::rle_array(); - array.scalar_at(1025); + array.scalar_at(1025).unwrap(); } #[test] #[should_panic] fn test_scalar_at_slice_out_of_bounds() { let array = fixture::rle_array().slice(0..1); - array.scalar_at(1); + array.scalar_at(1).unwrap(); } #[test] diff --git a/encodings/fsst/src/array.rs b/encodings/fsst/src/array.rs index 32e5aa83b78..5d9bf171dad 100644 --- a/encodings/fsst/src/array.rs +++ b/encodings/fsst/src/array.rs @@ -190,7 +190,7 @@ impl VTable for FSSTVTable { // from it instead. let (buffers, views) = fsst_decode_views(array, builder.completed_block_count(), ctx)?; - builder.push_buffer_and_adjusted_views(&buffers, &views, array.validity_mask()); + builder.push_buffer_and_adjusted_views(&buffers, &views, array.validity_mask()?); Ok(()) } diff --git a/encodings/fsst/src/compute/compare.rs b/encodings/fsst/src/compute/compare.rs index a1e8310c58f..5b6016d1b0f 100644 --- a/encodings/fsst/src/compute/compare.rs +++ b/encodings/fsst/src/compute/compare.rs @@ -80,7 +80,7 @@ fn compare_fsst_constant( return Ok(Some( BoolArray::from_bit_buffer( buffer, - Validity::copy_from_array(left.as_ref()) + Validity::copy_from_array(left.as_ref())? .union_nullability(right.dtype().nullability()), ) .into_array(), @@ -181,12 +181,12 @@ mod tests { ConstantArray::new(Scalar::null(DType::Utf8(Nullability::Nullable)), lhs.len()); let equals_null = compare(lhs.as_ref(), null_rhs.as_ref(), Operator::Eq).unwrap(); for idx in 0..lhs.len() { - assert!(equals_null.scalar_at(idx).is_null()); + assert!(equals_null.scalar_at(idx).unwrap().is_null()); } let noteq_null = compare(lhs.as_ref(), null_rhs.as_ref(), Operator::NotEq).unwrap(); for idx in 0..lhs.len() { - assert!(noteq_null.scalar_at(idx).is_null()); + assert!(noteq_null.scalar_at(idx).unwrap().is_null()); } } } diff --git a/encodings/fsst/src/ops.rs b/encodings/fsst/src/ops.rs index 570b15253b5..9b69b743a36 100644 --- a/encodings/fsst/src/ops.rs +++ b/encodings/fsst/src/ops.rs @@ -5,18 +5,19 @@ use vortex_array::arrays::varbin_scalar; use vortex_array::vtable::OperationsVTable; use vortex_buffer::ByteBuffer; use vortex_error::VortexExpect; +use vortex_error::VortexResult; use vortex_scalar::Scalar; use crate::FSSTArray; use crate::FSSTVTable; impl OperationsVTable for FSSTVTable { - fn scalar_at(array: &FSSTArray, index: usize) -> Scalar { - let compressed = array.codes().scalar_at(index); + fn scalar_at(array: &FSSTArray, index: usize) -> VortexResult { + let compressed = array.codes().scalar_at(index)?; let binary_datum = compressed.as_binary().value().vortex_expect("non-null"); let decoded_buffer = ByteBuffer::from(array.decompressor().decompress(binary_datum.as_slice())); - varbin_scalar(decoded_buffer, array.dtype()) + Ok(varbin_scalar(decoded_buffer, array.dtype())) } } diff --git a/encodings/pco/src/array.rs b/encodings/pco/src/array.rs index 87504c5b423..5b796df58f5 100644 --- a/encodings/pco/src/array.rs +++ b/encodings/pco/src/array.rs @@ -199,7 +199,7 @@ pub(crate) fn number_type_from_dtype(dtype: &DType) -> NumberType { } fn collect_valid(parray: &PrimitiveArray) -> VortexResult { - let mask = parray.validity_mask(); + let mask = parray.validity_mask()?; Ok(filter(&parray.to_array(), &mask)?.to_primitive()) } @@ -526,7 +526,7 @@ impl BaseArrayVTable for PcoVTable { } impl OperationsVTable for PcoVTable { - fn scalar_at(array: &PcoArray, index: usize) -> Scalar { + fn scalar_at(array: &PcoArray, index: usize) -> VortexResult { array._slice(index, index + 1).decompress().scalar_at(0) } } diff --git a/encodings/pco/src/compute/cast.rs b/encodings/pco/src/compute/cast.rs index 2be9e55f227..162852b1500 100644 --- a/encodings/pco/src/compute/cast.rs +++ b/encodings/pco/src/compute/cast.rs @@ -14,7 +14,7 @@ use crate::PcoVTable; impl CastKernel for PcoVTable { fn cast(&self, array: &PcoArray, dtype: &DType) -> VortexResult> { - if !dtype.is_nullable() || !array.all_valid() { + if !dtype.is_nullable() || !array.all_valid()? { // TODO(joe): fixme // We cannot cast to non-nullable since the validity containing nulls is used to decode // the PCO array, this would require rewriting tables. diff --git a/encodings/pco/src/test.rs b/encodings/pco/src/test.rs index cd80b283b13..cd8d4f7c182 100644 --- a/encodings/pco/src/test.rs +++ b/encodings/pco/src/test.rs @@ -128,9 +128,12 @@ fn test_validity_vtable() { Validity::Array(BoolArray::from_iter(mask_bools.clone()).to_array()), ); let compressed = PcoArray::from_primitive(&array, 3, 0).unwrap(); - assert_eq!(compressed.validity_mask(), Mask::from_iter(mask_bools)); assert_eq!( - compressed.slice(1..4).validity_mask(), + compressed.validity_mask().unwrap(), + Mask::from_iter(mask_bools) + ); + assert_eq!( + compressed.slice(1..4).validity_mask().unwrap(), Mask::from_iter(vec![true, true, false]) ); } diff --git a/encodings/runend/benches/run_end_null_count.rs b/encodings/runend/benches/run_end_null_count.rs index 2bd8bd01408..d5613b01c18 100644 --- a/encodings/runend/benches/run_end_null_count.rs +++ b/encodings/runend/benches/run_end_null_count.rs @@ -53,7 +53,7 @@ fn null_count_run_end(bencher: Bencher, (n, run_step, valid_density): (usize, us bencher .with_inputs(|| &array) - .bench_refs(|array| array.invalid_count()); + .bench_refs(|array| array.invalid_count().unwrap()); } fn fixture(n: usize, run_step: usize, valid_density: f64) -> RunEndArray { diff --git a/encodings/runend/src/array.rs b/encodings/runend/src/array.rs index 43d048a0988..61edf670e64 100644 --- a/encodings/runend/src/array.rs +++ b/encodings/runend/src/array.rs @@ -216,13 +216,13 @@ impl RunEndArray { // Validate the offset and length are valid for the given ends and values if offset != 0 && length != 0 { - let first_run_end: usize = ends.scalar_at(0).as_ref().try_into()?; + let first_run_end: usize = ends.scalar_at(0)?.as_ref().try_into()?; if first_run_end <= offset { vortex_bail!("First run end {first_run_end} must be bigger than offset {offset}"); } } - let last_run_end: usize = ends.scalar_at(ends.len() - 1).as_ref().try_into()?; + let last_run_end: usize = ends.scalar_at(ends.len() - 1)?.as_ref().try_into()?; let min_required_end = offset + length; if last_run_end < min_required_end { vortex_bail!("Last run end {last_run_end} must be >= offset+length {min_required_end}"); @@ -244,15 +244,19 @@ impl RunEndArray { /// # use vortex_array::arrays::BoolArray; /// # use vortex_array::IntoArray; /// # use vortex_buffer::buffer; + /// # use vortex_error::VortexResult; /// # use vortex_runend::RunEndArray; + /// # fn main() -> VortexResult<()> { /// let ends = buffer![2u8, 3u8].into_array(); /// let values = BoolArray::from_iter([false, true]).into_array(); /// let run_end = RunEndArray::new(ends, values); /// /// // Array encodes - /// assert_eq!(run_end.scalar_at(0), false.into()); - /// assert_eq!(run_end.scalar_at(1), false.into()); - /// assert_eq!(run_end.scalar_at(2), true.into()); + /// assert_eq!(run_end.scalar_at(0)?, false.into()); + /// assert_eq!(run_end.scalar_at(1)?, false.into()); + /// assert_eq!(run_end.scalar_at(2)?, true.into()); + /// # Ok(()) + /// # } /// ``` pub fn new(ends: ArrayRef, values: ArrayRef) -> Self { Self::try_new(ends, values).vortex_expect("RunEndArray new") @@ -292,7 +296,7 @@ impl RunEndArray { let length: usize = if ends.is_empty() { 0 } else { - ends.scalar_at(ends.len() - 1).as_ref().try_into()? + ends.scalar_at(ends.len() - 1)?.as_ref().try_into()? }; Self::try_new_offset_length(ends, values, 0, length) @@ -342,14 +346,15 @@ impl RunEndArray { } /// Convert the given logical index to an index into the `values` array - pub fn find_physical_index(&self, index: usize) -> usize { - self.ends() + pub fn find_physical_index(&self, index: usize) -> VortexResult { + Ok(self + .ends() .as_primitive_typed() .search_sorted( &PValue::from(index + self.offset()), SearchSortedSide::Right, - ) - .to_ends_index(self.ends().len()) + )? + .to_ends_index(self.ends().len())) } /// Run the array through run-end encoding. @@ -451,8 +456,8 @@ impl ValidityVTable for RunEndVTable { }) } - fn validity_mask(array: &RunEndArray) -> Mask { - match array.values().validity_mask() { + fn validity_mask(array: &RunEndArray) -> VortexResult { + Ok(match array.values().validity_mask()? { Mask::AllTrue(_) => Mask::AllTrue(array.len()), Mask::AllFalse(_) => Mask::AllFalse(array.len()), Mask::Values(values) => { @@ -469,7 +474,7 @@ impl ValidityVTable for RunEndVTable { }; Mask::from_buffer(ree_validity.to_bool().bit_buffer().clone()) } - } + }) } } diff --git a/encodings/runend/src/arrow.rs b/encodings/runend/src/arrow.rs index 766ab1d5dfb..08f0e879fa7 100644 --- a/encodings/runend/src/arrow.rs +++ b/encodings/runend/src/arrow.rs @@ -12,6 +12,7 @@ use vortex_array::search_sorted::SearchSortedSide; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_dtype::NativePType; +use vortex_error::VortexExpect; use vortex_scalar::PValue; use crate::RunEndArray; @@ -36,8 +37,10 @@ where let slice_begin = ends .as_primitive_typed() .search_sorted(&PValue::from(offset), SearchSortedSide::Right) + .vortex_expect("search_sorted on primitive ends") .to_ends_index(ends.len()); - let slice_end = find_slice_end_index(ends.as_ref(), offset + len); + let slice_end = find_slice_end_index(ends.as_ref(), offset + len) + .vortex_expect("find_slice_end_index on primitive ends"); ( ends.slice(slice_begin..slice_end), diff --git a/encodings/runend/src/compress.rs b/encodings/runend/src/compress.rs index cd5e8d7f52d..33fc2e0d723 100644 --- a/encodings/runend/src/compress.rs +++ b/encodings/runend/src/compress.rs @@ -176,7 +176,7 @@ pub fn runend_decode_primitive( runend_decode_typed_primitive( trimmed_ends_iter(ends.as_slice::(), offset, length), values.as_slice::

(), - values.validity_mask(), + values.validity_mask().vortex_expect("validity_mask"), values.dtype().nullability(), length, ) @@ -194,7 +194,7 @@ pub fn runend_decode_bools( runend_decode_typed_bool( trimmed_ends_iter(ends.as_slice::(), offset, length), values.bit_buffer(), - values.validity_mask(), + values.validity_mask().vortex_expect("validity_mask"), values.dtype().nullability(), length, ) diff --git a/encodings/runend/src/compute/cast.rs b/encodings/runend/src/compute/cast.rs index 0cf19f32ff7..281a3bbb49e 100644 --- a/encodings/runend/src/compute/cast.rs +++ b/encodings/runend/src/compute/cast.rs @@ -76,25 +76,25 @@ mod tests { // RunEnd encoding should expand to [100, 100, 100, 200, 200, 100, 100, 100, 300, 300] assert_eq!(decoded.len(), 10); assert_eq!( - TryInto::::try_into(decoded.scalar_at(0).as_ref()) + TryInto::::try_into(decoded.scalar_at(0).unwrap().as_ref()) .ok() .unwrap(), 100i64 ); assert_eq!( - TryInto::::try_into(decoded.scalar_at(3).as_ref()) + TryInto::::try_into(decoded.scalar_at(3).unwrap().as_ref()) .ok() .unwrap(), 200i64 ); assert_eq!( - TryInto::::try_into(decoded.scalar_at(5).as_ref()) + TryInto::::try_into(decoded.scalar_at(5).unwrap().as_ref()) .ok() .unwrap(), 100i64 ); assert_eq!( - TryInto::::try_into(decoded.scalar_at(8).as_ref()) + TryInto::::try_into(decoded.scalar_at(8).unwrap().as_ref()) .ok() .unwrap(), 300i64 diff --git a/encodings/runend/src/compute/take.rs b/encodings/runend/src/compute/take.rs index 75ea164f3f7..32375618ecb 100644 --- a/encodings/runend/src/compute/take.rs +++ b/encodings/runend/src/compute/take.rs @@ -65,21 +65,21 @@ pub fn take_indices_unchecked>( // TODO(joe): use the validity mask to skip search sorted. let physical_indices = match_each_integer_ptype!(ends.ptype(), |I| { let end_slices = ends.as_slice::(); - let buffer = Buffer::from_trusted_len_iter( - indices - .iter() - .map(|idx| idx.as_() + array.offset()) - .map(|idx| { - match ::from(idx) { - Some(idx) => end_slices.search_sorted(&idx, SearchSortedSide::Right), - None => { - // The idx is too large for I, therefore it's out of bounds. - SearchResult::NotFound(ends_len) - } + let physical_indices_vec: Vec = indices + .iter() + .map(|idx| idx.as_() + array.offset()) + .map(|idx| { + match ::from(idx) { + Some(idx) => end_slices.search_sorted(&idx, SearchSortedSide::Right), + None => { + // The idx is too large for I, therefore it's out of bounds. + Ok(SearchResult::NotFound(ends_len)) } - }) - .map(|result| result.to_ends_index(ends_len) as u64), - ); + } + }) + .map(|result| result.map(|r| r.to_ends_index(ends_len) as u64)) + .collect::>>()?; + let buffer = Buffer::from(physical_indices_vec); PrimitiveArray::new(buffer, validity.clone()) }); diff --git a/encodings/runend/src/kernel.rs b/encodings/runend/src/kernel.rs index efacdd07ae9..339048a6871 100644 --- a/encodings/runend/src/kernel.rs +++ b/encodings/runend/src/kernel.rs @@ -42,26 +42,26 @@ impl ExecuteParentKernel for RunEndSliceKernel { _child_idx: usize, ctx: &mut ExecutionCtx, ) -> VortexResult> { - let sliced = slice(array, parent.slice_range().clone()); + let sliced = slice(array, parent.slice_range().clone())?; sliced.execute::(ctx).map(Some) } } -fn slice(array: &RunEndArray, range: Range) -> ArrayRef { +fn slice(array: &RunEndArray, range: Range) -> VortexResult { let new_length = range.len(); - let slice_begin = array.find_physical_index(range.start); - let slice_end = crate::ops::find_slice_end_index(array.ends(), range.end + array.offset()); + let slice_begin = array.find_physical_index(range.start)?; + let slice_end = crate::ops::find_slice_end_index(array.ends(), range.end + array.offset())?; // If the sliced range contains only a single run, opt to return a ConstantArray. if slice_begin + 1 == slice_end { - let value = array.values().scalar_at(slice_begin); - return ConstantArray::new(value, new_length).into_array(); + let value = array.values().scalar_at(slice_begin)?; + return Ok(ConstantArray::new(value, new_length).into_array()); } // SAFETY: we maintain the ends invariant in our slice implementation - unsafe { + Ok(unsafe { RunEndArray::new_unchecked( array.ends().slice(slice_begin..slice_end), array.values().slice(slice_begin..slice_end), @@ -69,5 +69,5 @@ fn slice(array: &RunEndArray, range: Range) -> ArrayRef { new_length, ) .into_array() - } + }) } diff --git a/encodings/runend/src/ops.rs b/encodings/runend/src/ops.rs index aa0fc36ac1d..1f3e55fe2a7 100644 --- a/encodings/runend/src/ops.rs +++ b/encodings/runend/src/ops.rs @@ -6,6 +6,7 @@ use vortex_array::search_sorted::SearchResult; use vortex_array::search_sorted::SearchSorted; use vortex_array::search_sorted::SearchSortedSide; use vortex_array::vtable::OperationsVTable; +use vortex_error::VortexResult; use vortex_scalar::PValue; use vortex_scalar::Scalar; @@ -13,8 +14,8 @@ use crate::RunEndArray; use crate::RunEndVTable; impl OperationsVTable for RunEndVTable { - fn scalar_at(array: &RunEndArray, index: usize) -> Scalar { - array.values().scalar_at(array.find_physical_index(index)) + fn scalar_at(array: &RunEndArray, index: usize) -> VortexResult { + array.values().scalar_at(array.find_physical_index(index)?) } } @@ -22,11 +23,11 @@ impl OperationsVTable for RunEndVTable { /// /// If the index exists in the array we want to take that position (as we are searching from the right) /// otherwise we want to take the next one -pub(crate) fn find_slice_end_index(array: &dyn Array, index: usize) -> usize { +pub(crate) fn find_slice_end_index(array: &dyn Array, index: usize) -> VortexResult { let result = array .as_primitive_typed() - .search_sorted(&PValue::from(index), SearchSortedSide::Right); - match result { + .search_sorted(&PValue::from(index), SearchSortedSide::Right)?; + Ok(match result { SearchResult::Found(i) => i, SearchResult::NotFound(i) => { if i == array.len() { @@ -35,7 +36,7 @@ pub(crate) fn find_slice_end_index(array: &dyn Array, index: usize) -> usize { i + 1 } } - } + }) } #[cfg(test)] @@ -149,7 +150,8 @@ mod tests { fn ree_scalar_at_end() { let scalar = RunEndArray::encode(buffer![1, 1, 1, 4, 4, 4, 2, 2, 5, 5, 5, 5].into_array()) .unwrap() - .scalar_at(11); + .scalar_at(11) + .unwrap(); assert_eq!(scalar, 5.into()); } diff --git a/encodings/sequence/src/array.rs b/encodings/sequence/src/array.rs index c3c8776912a..fad4ead18d7 100644 --- a/encodings/sequence/src/array.rs +++ b/encodings/sequence/src/array.rs @@ -388,11 +388,11 @@ impl BaseArrayVTable for SequenceVTable { } impl OperationsVTable for SequenceVTable { - fn scalar_at(array: &SequenceArray, index: usize) -> Scalar { - Scalar::new( + fn scalar_at(array: &SequenceArray, index: usize) -> VortexResult { + Ok(Scalar::new( array.dtype().clone(), ScalarValue::from(array.index_value(index)), - ) + )) } } @@ -401,8 +401,8 @@ impl ValidityVTable for SequenceVTable { Ok(Validity::AllValid) } - fn validity_mask(array: &SequenceArray) -> Mask { - Mask::AllTrue(array.len()) + fn validity_mask(array: &SequenceArray) -> VortexResult { + Ok(Mask::AllTrue(array.len())) } } @@ -455,7 +455,8 @@ mod tests { fn test_sequence_scalar_at() { let scalar = SequenceArray::typed_new(2i64, 3, Nullability::NonNullable, 4) .unwrap() - .scalar_at(2); + .scalar_at(2) + .unwrap(); assert_eq!( scalar, diff --git a/encodings/sequence/src/compress.rs b/encodings/sequence/src/compress.rs index c24838b2cc3..9f09f11873f 100644 --- a/encodings/sequence/src/compress.rs +++ b/encodings/sequence/src/compress.rs @@ -25,7 +25,7 @@ pub fn sequence_encode(primitive_array: &PrimitiveArray) -> VortexResult VortexResult { - let mask = indices.validity_mask(); + let mask = indices.validity_mask()?; let indices = indices.to_primitive(); let result_nullability = array.dtype().nullability() | indices.dtype().nullability(); diff --git a/encodings/sparse/src/canonical.rs b/encodings/sparse/src/canonical.rs index 745c3674b8a..a1a2465243d 100644 --- a/encodings/sparse/src/canonical.rs +++ b/encodings/sparse/src/canonical.rs @@ -131,7 +131,10 @@ fn canonicalize_sparse_lists( let n_filled = array.len() - resolved_patches.num_patches(); let total_canonical_values = values.elements().len() + fill_value.len() * n_filled; - let validity = Validity::from_mask(array.validity_mask(), nullability); + let validity = Validity::from_mask( + array.validity_mask().vortex_expect("validity_mask"), + nullability, + ); match_each_integer_ptype!(indices.ptype(), |I| { match_smallest_offset_type!(total_canonical_values, |O| { @@ -180,7 +183,12 @@ fn canonicalize_sparse_lists_inner( if position_is_patched { // Set with the patch value. builder - .append_value(patch_values.scalar_at(patch_idx).as_list()) + .append_value( + patch_values + .scalar_at(patch_idx) + .vortex_expect("scalar_at") + .as_list(), + ) .vortex_expect("Failed to append sparse value"); patch_idx += 1; } else { @@ -201,7 +209,10 @@ fn canonicalize_sparse_fixed_size_list(array: &SparseArray, nullability: Nullabi let values = resolved_patches.values().to_fixed_size_list(); let fill_value = array.fill_scalar().as_list(); - let validity = Validity::from_mask(array.validity_mask(), nullability); + let validity = Validity::from_mask( + array.validity_mask().vortex_expect("validity_mask"), + nullability, + ); match_each_integer_ptype!(indices.ptype(), |I| { canonicalize_sparse_fixed_size_list_inner::( @@ -252,7 +263,7 @@ fn canonicalize_sparse_fixed_size_list_inner( let patch_list = values.fixed_size_list_elements_at(patch_idx); for i in 0..list_size as usize { builder - .append_scalar(&patch_list.scalar_at(i)) + .append_scalar(&patch_list.scalar_at(i).vortex_expect("scalar_at")) .vortex_expect("element dtype must match"); } } else { @@ -376,7 +387,10 @@ fn canonicalize_sparse_struct( unresolved_patches.offset(), unresolved_patches.indices(), &Validity::from_mask( - unresolved_patches.values().validity_mask(), + unresolved_patches + .values() + .validity_mask() + .vortex_expect("validity_mask"), Nullability::Nullable, ), ) @@ -442,7 +456,10 @@ fn canonicalize_varbin( let patches = array.resolved_patches(); let indices = patches.indices().to_primitive(); let values = patches.values().to_varbinview(); - let validity = Validity::from_mask(array.validity_mask(), dtype.nullability()); + let validity = Validity::from_mask( + array.validity_mask().vortex_expect("validity_mask"), + dtype.nullability(), + ); let len = array.len(); match_each_integer_ptype!(indices.ptype(), |I| { diff --git a/encodings/sparse/src/lib.rs b/encodings/sparse/src/lib.rs index 6fdfa6582cc..cde54d0cfc9 100644 --- a/encodings/sparse/src/lib.rs +++ b/encodings/sparse/src/lib.rs @@ -160,9 +160,7 @@ impl VTable for SparseVTable { } fn slice(array: &SparseArray, range: Range) -> VortexResult> { - let new_patches = array.patches().slice(range.clone()); - - let Some(new_patches) = new_patches else { + let Some(new_patches) = array.patches().slice(range.clone())? else { return Ok(Some( ConstantArray::new(array.fill_scalar().clone(), range.len()).into_array(), )); @@ -223,7 +221,7 @@ impl SparseArray { // Verify the indices are all in the valid range if !indices.is_empty() { - let last_index = usize::try_from(&indices.scalar_at(indices.len() - 1))?; + let last_index = usize::try_from(&indices.scalar_at(indices.len() - 1)?)?; vortex_ensure!( last_index < len, @@ -306,7 +304,7 @@ impl SparseArray { fill_value.dtype() ) } - let mask = array.validity_mask(); + let mask = array.validity_mask()?; if mask.all_false() { // Array is constant NULL @@ -432,16 +430,16 @@ impl ValidityVTable for SparseVTable { )) } - fn validity_mask(array: &SparseArray) -> Mask { + fn validity_mask(array: &SparseArray) -> VortexResult { let fill_is_valid = array.fill_scalar().is_valid(); - let values_validity = array.patches().values().validity_mask(); + let values_validity = array.patches().values().validity_mask()?; let len = array.len(); if matches!(values_validity, Mask::AllTrue(_)) && fill_is_valid { - return Mask::AllTrue(len); + return Ok(Mask::AllTrue(len)); } if matches!(values_validity, Mask::AllFalse(_)) && !fill_is_valid { - return Mask::AllFalse(len); + return Ok(Mask::AllFalse(len)); } let mut is_valid_buffer = if fill_is_valid { @@ -458,7 +456,7 @@ impl ValidityVTable for SparseVTable { patch_validity(&mut is_valid_buffer, indices, index_offset, values_validity); }); - Mask::from_buffer(is_valid_buffer.freeze()) + Ok(Mask::from_buffer(is_valid_buffer.freeze())) } } @@ -544,16 +542,16 @@ mod test { pub fn test_scalar_at() { let array = sparse_array(nullable_fill()); - assert_eq!(array.scalar_at(0), nullable_fill()); - assert_eq!(array.scalar_at(2), Scalar::from(Some(100_i32))); - assert_eq!(array.scalar_at(5), Scalar::from(Some(200_i32))); + assert_eq!(array.scalar_at(0).unwrap(), nullable_fill()); + assert_eq!(array.scalar_at(2).unwrap(), Scalar::from(Some(100_i32))); + assert_eq!(array.scalar_at(5).unwrap(), Scalar::from(Some(200_i32))); } #[test] #[should_panic(expected = "out of bounds")] fn test_scalar_at_oob() { let array = sparse_array(nullable_fill()); - array.scalar_at(10); + array.scalar_at(10).unwrap(); } #[test] @@ -567,26 +565,26 @@ mod test { .unwrap(); assert_eq!( - PrimitiveScalar::try_from(&arr.scalar_at(10)) + PrimitiveScalar::try_from(&arr.scalar_at(10).unwrap()) .unwrap() .typed_value::(), Some(1234) ); - assert!(arr.scalar_at(0).is_null()); - assert!(arr.scalar_at(99).is_null()); + assert!(arr.scalar_at(0).unwrap().is_null()); + assert!(arr.scalar_at(99).unwrap().is_null()); } #[test] pub fn scalar_at_sliced() { let sliced = sparse_array(nullable_fill()).slice(2..7); - assert_eq!(usize::try_from(&sliced.scalar_at(0)).unwrap(), 100); + assert_eq!(usize::try_from(&sliced.scalar_at(0).unwrap()).unwrap(), 100); } #[test] pub fn validity_mask_sliced_null_fill() { let sliced = sparse_array(nullable_fill()).slice(2..7); assert_eq!( - sliced.validity_mask(), + sliced.validity_mask().unwrap(), Mask::from_iter(vec![true, false, false, true, false]) ); } @@ -607,7 +605,7 @@ mod test { .slice(2..7); assert_eq!( - sliced.validity_mask(), + sliced.validity_mask().unwrap(), Mask::from_iter(vec![false, true, true, false, true]) ); } @@ -615,17 +613,28 @@ mod test { #[test] pub fn scalar_at_sliced_twice() { let sliced_once = sparse_array(nullable_fill()).slice(1..8); - assert_eq!(usize::try_from(&sliced_once.scalar_at(1)).unwrap(), 100); + assert_eq!( + usize::try_from(&sliced_once.scalar_at(1).unwrap()).unwrap(), + 100 + ); let sliced_twice = sliced_once.slice(1..6); - assert_eq!(usize::try_from(&sliced_twice.scalar_at(3)).unwrap(), 200); + assert_eq!( + usize::try_from(&sliced_twice.scalar_at(3).unwrap()).unwrap(), + 200 + ); } #[test] pub fn sparse_validity_mask() { let array = sparse_array(nullable_fill()); assert_eq!( - array.validity_mask().to_bit_buffer().iter().collect_vec(), + array + .validity_mask() + .unwrap() + .to_bit_buffer() + .iter() + .collect_vec(), [ false, false, true, false, false, true, false, false, true, false ] @@ -635,7 +644,7 @@ mod test { #[test] fn sparse_validity_mask_non_null_fill() { let array = sparse_array(non_nullable_fill()); - assert!(array.validity_mask().all_true()); + assert!(array.validity_mask().unwrap().all_true()); } #[test] @@ -670,7 +679,7 @@ mod test { .vortex_expect("SparseArray::encode should succeed for test data"); let canonical = sparse.to_primitive(); assert_eq!( - sparse.validity_mask(), + sparse.validity_mask().unwrap(), Mask::from_iter(vec![ true, true, false, true, false, true, false, true, true, false, true, false, ]) @@ -687,7 +696,7 @@ mod test { let values = PrimitiveArray::from_option_iter([Some(0i16), Some(1), None, None, Some(4)]) .into_array(); let array = SparseArray::try_new(indices, values, 10, Scalar::null_typed::()).unwrap(); - let actual = array.validity_mask(); + let actual = array.validity_mask().unwrap(); let expected = Mask::from_iter([ true, false, true, false, false, false, false, false, true, false, ]); diff --git a/encodings/sparse/src/ops.rs b/encodings/sparse/src/ops.rs index dccafedc1ff..8a0629895de 100644 --- a/encodings/sparse/src/ops.rs +++ b/encodings/sparse/src/ops.rs @@ -2,17 +2,18 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_array::vtable::OperationsVTable; +use vortex_error::VortexResult; use vortex_scalar::Scalar; use crate::SparseArray; use crate::SparseVTable; impl OperationsVTable for SparseVTable { - fn scalar_at(array: &SparseArray, index: usize) -> Scalar { - array + fn scalar_at(array: &SparseArray, index: usize) -> VortexResult { + Ok(array .patches() - .get_patched(index) - .unwrap_or_else(|| array.fill_scalar().clone()) + .get_patched(index)? + .unwrap_or_else(|| array.fill_scalar().clone())) } } diff --git a/encodings/zigzag/src/array.rs b/encodings/zigzag/src/array.rs index c6aed03c032..e85880cac85 100644 --- a/encodings/zigzag/src/array.rs +++ b/encodings/zigzag/src/array.rs @@ -179,14 +179,14 @@ impl BaseArrayVTable for ZigZagVTable { } impl OperationsVTable for ZigZagVTable { - fn scalar_at(array: &ZigZagArray, index: usize) -> Scalar { - let scalar = array.encoded().scalar_at(index); + fn scalar_at(array: &ZigZagArray, index: usize) -> VortexResult { + let scalar = array.encoded().scalar_at(index)?; if scalar.is_null() { - return scalar.reinterpret_cast(array.ptype()); + return Ok(scalar.reinterpret_cast(array.ptype())); } let pscalar = scalar.as_primitive(); - match_each_unsigned_integer_ptype!(pscalar.ptype(), |P| { + Ok(match_each_unsigned_integer_ptype!(pscalar.ptype(), |P| { Scalar::primitive( <

::Int>::decode( pscalar @@ -195,7 +195,7 @@ impl OperationsVTable for ZigZagVTable { ), array.dtype().nullability(), ) - }) + })) } } @@ -245,7 +245,10 @@ mod test { let sliced = zigzag.slice(0..2); let sliced = sliced.as_::(); - assert_eq!(sliced.scalar_at(sliced.len() - 1), Scalar::from(-5i32)); + assert_eq!( + sliced.scalar_at(sliced.len() - 1).unwrap(), + Scalar::from(-5i32) + ); assert_eq!( sliced.statistics().compute_min::(), diff --git a/encodings/zigzag/src/compute/mod.rs b/encodings/zigzag/src/compute/mod.rs index 8bfd79b1f57..24cc9d46353 100644 --- a/encodings/zigzag/src/compute/mod.rs +++ b/encodings/zigzag/src/compute/mod.rs @@ -99,7 +99,7 @@ mod tests { Validity::AllValid, ))?; assert_eq!( - zigzag.scalar_at(1), + zigzag.scalar_at(1)?, Scalar::primitive(-160, Nullability::Nullable) ); Ok(()) diff --git a/encodings/zstd/src/array.rs b/encodings/zstd/src/array.rs index f3f2bef268a..66fab25ac06 100644 --- a/encodings/zstd/src/array.rs +++ b/encodings/zstd/src/array.rs @@ -216,12 +216,12 @@ fn choose_max_dict_size(uncompressed_size: usize) -> usize { } fn collect_valid_primitive(parray: &PrimitiveArray) -> VortexResult { - let mask = parray.validity_mask(); + let mask = parray.validity_mask()?; Ok(filter(&parray.to_array(), &mask)?.to_primitive()) } fn collect_valid_vbv(vbv: &VarBinViewArray) -> VortexResult<(ByteBuffer, Vec)> { - let mask = vbv.validity_mask(); + let mask = vbv.validity_mask()?; let buffer_and_value_byte_indices = match mask.bit_buffer() { AllOr::None => (Buffer::empty(), Vec::new()), _ => { @@ -773,7 +773,7 @@ impl BaseArrayVTable for ZstdVTable { } impl OperationsVTable for ZstdVTable { - fn scalar_at(array: &ZstdArray, index: usize) -> Scalar { + fn scalar_at(array: &ZstdArray, index: usize) -> VortexResult { array._slice(index, index + 1).decompress().scalar_at(0) } } diff --git a/encodings/zstd/src/test.rs b/encodings/zstd/src/test.rs index 060c0bae72c..74f1306ca3b 100644 --- a/encodings/zstd/src/test.rs +++ b/encodings/zstd/src/test.rs @@ -84,7 +84,7 @@ fn test_zstd_with_validity_and_multi_frame() { let slice = compressed.slice(176..179); let primitive = slice.to_primitive(); assert_eq!( - TryInto::::try_into(primitive.scalar_at(1).as_ref()) + TryInto::::try_into(primitive.scalar_at(1).unwrap().as_ref()) .ok() .unwrap(), 177 @@ -125,9 +125,12 @@ fn test_validity_vtable() { Validity::Array(BoolArray::from_iter(mask_bools.clone()).to_array()), ); let compressed = ZstdArray::from_primitive(&array, 3, 0).unwrap(); - assert_eq!(compressed.validity_mask(), Mask::from_iter(mask_bools)); assert_eq!( - compressed.slice(1..4).validity_mask(), + compressed.validity_mask().unwrap(), + Mask::from_iter(mask_bools) + ); + assert_eq!( + compressed.slice(1..4).validity_mask().unwrap(), Mask::from_iter(vec![true, true, false]) ); } diff --git a/fuzz/fuzz_targets/file_io.rs b/fuzz/fuzz_targets/file_io.rs index 894c2a91b56..cf32514adce 100644 --- a/fuzz/fuzz_targets/file_io.rs +++ b/fuzz/fuzz_targets/file_io.rs @@ -116,7 +116,9 @@ fuzz_target!(|fuzz: FuzzFileAction| -> Corpus { .vortex_expect("compare operation should succeed in fuzz test") .to_bool(); let true_count = bool_result.bit_buffer().true_count(); - if true_count != expected_array.len() && (bool_result.all_valid() || expected_array.all_valid()) + if true_count != expected_array.len() + && (bool_result.all_valid().vortex_expect("all_valid") + || expected_array.all_valid().vortex_expect("all_valid")) { vortex_panic!( "Failed to match original array {}with{}", diff --git a/fuzz/src/array/cast.rs b/fuzz/src/array/cast.rs index 906a2e33322..2a7842a7cbc 100644 --- a/fuzz/src/array/cast.rs +++ b/fuzz/src/array/cast.rs @@ -40,7 +40,7 @@ pub fn cast_canonical_array(array: &ArrayRef, target: &DType) -> VortexResult>(), - Validity::from_mask(array.validity_mask(), target.nullability()), + Validity::from_mask(array.validity_mask()?, target.nullability()), ) .to_array() }) @@ -66,7 +66,7 @@ pub fn cast_canonical_array(array: &ArrayRef, target: &DType) -> VortexResult>(), - Validity::from_mask(array.validity_mask(), target.nullability()), + Validity::from_mask(array.validity_mask()?, target.nullability()), ) .to_array(), )), @@ -81,7 +81,7 @@ pub fn cast_canonical_array(array: &ArrayRef, target: &DType) -> VortexResult>(), - Validity::from_mask(array.validity_mask(), target.nullability()), + Validity::from_mask(array.validity_mask()?, target.nullability()), ) .to_array(), )) diff --git a/fuzz/src/array/compare.rs b/fuzz/src/array/compare.rs index 618b2647892..e90e3d6c13c 100644 --- a/fuzz/src/array/compare.rs +++ b/fuzz/src/array/compare.rs @@ -39,7 +39,13 @@ pub fn compare_canonical_array(array: &dyn Array, value: &Scalar, operator: Oper .to_bool() .bit_buffer() .iter() - .zip(array.validity_mask().to_bit_buffer().iter()) + .zip( + array + .validity_mask() + .vortex_expect("validity_mask") + .to_bit_buffer() + .iter(), + ) .map(|(b, v)| v.then_some(b)), bool, operator, @@ -58,7 +64,13 @@ pub fn compare_canonical_array(array: &dyn Array, value: &Scalar, operator: Oper .as_slice::

() .iter() .copied() - .zip(array.validity_mask().to_bit_buffer().iter()) + .zip( + array + .validity_mask() + .vortex_expect("validity_mask") + .to_bit_buffer() + .iter(), + ) .map(|(b, v)| v.then_some(NativeValue(b))), NativeValue(pval), operator, @@ -80,7 +92,13 @@ pub fn compare_canonical_array(array: &dyn Array, value: &Scalar, operator: Oper buf.as_slice() .iter() .copied() - .zip(array.validity_mask().to_bit_buffer().iter()) + .zip( + array + .validity_mask() + .vortex_expect("validity_mask") + .to_bit_buffer() + .iter(), + ) .map(|(b, v)| v.then_some(b)), dval, operator, @@ -115,7 +133,9 @@ pub fn compare_canonical_array(array: &dyn Array, value: &Scalar, operator: Oper ) }), DType::Struct(..) | DType::List(..) | DType::FixedSizeList(..) => { - let scalar_vals: Vec = (0..array.len()).map(|i| array.scalar_at(i)).collect(); + let scalar_vals: Vec = (0..array.len()) + .map(|i| array.scalar_at(i).vortex_expect("scalar_at")) + .collect(); BoolArray::from_iter( scalar_vals .iter() diff --git a/fuzz/src/array/fill_null.rs b/fuzz/src/array/fill_null.rs index 302dc1d0256..eff983580b8 100644 --- a/fuzz/src/array/fill_null.rs +++ b/fuzz/src/array/fill_null.rs @@ -184,6 +184,7 @@ fn fill_varbinview_array( if validity_bits.value(i) { array .scalar_at(i) + .vortex_expect("scalar_at") .as_utf8() .value() .vortex_expect("cannot have null valid value") @@ -217,6 +218,7 @@ fn fill_varbinview_array( if validity_bits.value(i) { array .scalar_at(i) + .vortex_expect("scalar_at") .as_binary() .value() .vortex_expect("cannot have null valid value") diff --git a/fuzz/src/array/filter.rs b/fuzz/src/array/filter.rs index 4d6fe1c0e03..21b93f92280 100644 --- a/fuzz/src/array/filter.rs +++ b/fuzz/src/array/filter.rs @@ -23,7 +23,7 @@ use crate::array::take_canonical_array_non_nullable_indices; pub fn filter_canonical_array(array: &dyn Array, filter: &[bool]) -> VortexResult { let validity = if array.dtype().is_nullable() { - let validity_buff = array.validity_mask().to_bit_buffer(); + let validity_buff = array.validity_mask()?.to_bit_buffer(); Validity::from_iter( filter .iter() diff --git a/fuzz/src/array/mask.rs b/fuzz/src/array/mask.rs index 694256a6de0..4b89636a1de 100644 --- a/fuzz/src/array/mask.rs +++ b/fuzz/src/array/mask.rs @@ -149,7 +149,7 @@ mod tests { assert_eq!(result.len(), 5); // All values should still be null for i in 0..5 { - assert!(!result.is_valid(i)); + assert!(!result.is_valid(i).unwrap()); } } @@ -229,9 +229,9 @@ mod tests { let result = mask_canonical_array(array.to_canonical().unwrap(), &mask).unwrap(); assert_eq!(result.len(), 3); - assert!(result.is_valid(0)); - assert!(!result.is_valid(1)); - assert!(result.is_valid(2)); + assert!(result.is_valid(0).unwrap()); + assert!(!result.is_valid(1).unwrap()); + assert!(result.is_valid(2).unwrap()); } #[test] @@ -245,9 +245,9 @@ mod tests { let result = mask_canonical_array(array.to_canonical().unwrap(), &mask).unwrap(); assert_eq!(result.len(), 3); - assert!(!result.is_valid(0)); - assert!(result.is_valid(1)); - assert!(!result.is_valid(2)); + assert!(!result.is_valid(0).unwrap()); + assert!(result.is_valid(1).unwrap()); + assert!(!result.is_valid(2).unwrap()); } #[test] @@ -269,9 +269,9 @@ mod tests { let result = mask_canonical_array(array.to_canonical().unwrap(), &mask).unwrap(); assert_eq!(result.len(), 3); - assert!(result.is_valid(0)); - assert!(!result.is_valid(1)); - assert!(result.is_valid(2)); + assert!(result.is_valid(0).unwrap()); + assert!(!result.is_valid(1).unwrap()); + assert!(result.is_valid(2).unwrap()); } #[test] diff --git a/fuzz/src/array/mod.rs b/fuzz/src/array/mod.rs index 60cb868d05e..62a765e15a6 100644 --- a/fuzz/src/array/mod.rs +++ b/fuzz/src/array/mod.rs @@ -232,7 +232,9 @@ impl<'a> Arbitrary<'a> for FuzzArrayAction { } let scalar = if u.arbitrary()? { - current_array.scalar_at(u.choose_index(current_array.len())?) + current_array + .scalar_at(u.choose_index(current_array.len())?) + .vortex_expect("scalar_at") } else { random_scalar(u, current_array.dtype())? }; @@ -271,7 +273,9 @@ impl<'a> Arbitrary<'a> for FuzzArrayAction { } ActionType::Compare => { let scalar = if u.arbitrary()? { - current_array.scalar_at(u.choose_index(current_array.len())?) + current_array + .scalar_at(u.choose_index(current_array.len())?) + .vortex_expect("scalar_at") } else { // We can compare arrays with different nullability let null: Nullability = u.arbitrary()?; @@ -330,7 +334,9 @@ impl<'a> Arbitrary<'a> for FuzzArrayAction { return Err(EmptyChoose); } let fill_value = if u.arbitrary()? && !current_array.is_empty() { - current_array.scalar_at(u.choose_index(current_array.len())?) + current_array + .scalar_at(u.choose_index(current_array.len())?) + .vortex_expect("scalar_at") } else { random_scalar( u, @@ -636,7 +642,7 @@ pub fn run_fuzz_action(fuzz_action: FuzzArrayAction) -> crate::error::VortexFuzz Action::ScalarAt(indices) => { let expected_scalars = expected.scalar_vec(); for (j, &idx) in indices.iter().enumerate() { - let scalar = current_array.scalar_at(idx); + let scalar = current_array.scalar_at(idx).vortex_expect("scalar_at"); assert_scalar_eq(&expected_scalars[j], &scalar, i)?; } } @@ -658,7 +664,9 @@ fn assert_search_sorted( use crate::error::Backtrace; use crate::error::VortexFuzzError; - let search_result = array.search_sorted(&s, side); + let search_result = array + .search_sorted(&s, side) + .map_err(|e| VortexFuzzError::VortexError(e, Backtrace::capture()))?; if search_result != expected { Err(VortexFuzzError::SearchSortedError( s, @@ -704,8 +712,8 @@ pub fn assert_array_eq( )); } for idx in 0..lhs.len() { - let l = lhs.scalar_at(idx); - let r = rhs.scalar_at(idx); + let l = lhs.scalar_at(idx).vortex_expect("scalar_at"); + let r = rhs.scalar_at(idx).vortex_expect("scalar_at"); if l != r { return Err(VortexFuzzError::ArrayNotEqual( diff --git a/fuzz/src/array/scalar_at.rs b/fuzz/src/array/scalar_at.rs index d5598e1235c..033e44a946d 100644 --- a/fuzz/src/array/scalar_at.rs +++ b/fuzz/src/array/scalar_at.rs @@ -18,7 +18,7 @@ use vortex_scalar::Scalar; /// This implementation manually extracts the scalar value from each canonical type /// without using the scalar_at method, to serve as an independent baseline for testing. pub fn scalar_at_canonical_array(canonical: Canonical, index: usize) -> VortexResult { - if canonical.as_ref().is_invalid(index) { + if canonical.as_ref().is_invalid(index)? { return Ok(Scalar::null(canonical.as_ref().dtype().clone())); } Ok(match canonical { diff --git a/fuzz/src/array/search_sorted.rs b/fuzz/src/array/search_sorted.rs index 0f0e7349e99..80e4d71d3a3 100644 --- a/fuzz/src/array/search_sorted.rs +++ b/fuzz/src/array/search_sorted.rs @@ -4,7 +4,6 @@ use std::cmp::Ordering; use std::fmt::Debug; -use itertools::Itertools; use vortex_array::Array; use vortex_array::ToCanonical; use vortex_array::accessor::ArrayAccessor; @@ -25,9 +24,9 @@ use vortex_scalar::Scalar; struct SearchNullableSlice(Vec>); impl IndexOrd> for SearchNullableSlice { - fn index_cmp(&self, idx: usize, elem: &Option) -> Option { + fn index_cmp(&self, idx: usize, elem: &Option) -> VortexResult> { // SAFETY: Used in search_sorted_by same as the standard library. The search_sorted ensures idx is in bounds - unsafe { self.0.get_unchecked(idx) }.partial_cmp(elem) + Ok(unsafe { self.0.get_unchecked(idx) }.partial_cmp(elem)) } fn index_len(&self) -> usize { @@ -38,15 +37,15 @@ impl IndexOrd> for SearchNullableSlice { struct SearchPrimitiveSlice(Vec>); impl IndexOrd> for SearchPrimitiveSlice { - fn index_cmp(&self, idx: usize, elem: &Option) -> Option { + fn index_cmp(&self, idx: usize, elem: &Option) -> VortexResult> { match elem { None => unreachable!("Can't search for None"), Some(v) => { // SAFETY: Used in search_sorted_by same as the standard library. The search_sorted ensures idx is in bounds - match unsafe { self.0.get_unchecked(idx) } { + Ok(match unsafe { self.0.get_unchecked(idx) } { None => Some(Ordering::Less), Some(i) => Some(i.total_compare(*v)), - } + }) } } } @@ -64,7 +63,7 @@ pub fn search_sorted_canonical_array( match array.dtype() { DType::Bool(_) => { let bool_array = array.to_bool(); - let validity = bool_array.validity_mask().to_bit_buffer(); + let validity = bool_array.validity_mask()?.to_bit_buffer(); let opt_values = bool_array .bit_buffer() .iter() @@ -72,11 +71,11 @@ pub fn search_sorted_canonical_array( .map(|(b, v)| v.then_some(b)) .collect::>(); let to_find = scalar.try_into()?; - Ok(SearchNullableSlice(opt_values).search_sorted(&Some(to_find), side)) + SearchNullableSlice(opt_values).search_sorted(&Some(to_find), side) } DType::Primitive(p, _) => { let primitive_array = array.to_primitive(); - let validity = primitive_array.validity_mask().to_bit_buffer(); + let validity = primitive_array.validity_mask()?.to_bit_buffer(); match_each_native_ptype!(p, |P| { let opt_values = primitive_array .as_slice::

() @@ -86,12 +85,12 @@ pub fn search_sorted_canonical_array( .map(|(b, v)| v.then_some(b)) .collect::>(); let to_find: P = scalar.try_into()?; - Ok(SearchPrimitiveSlice(opt_values).search_sorted(&Some(to_find), side)) + SearchPrimitiveSlice(opt_values).search_sorted(&Some(to_find), side) }) } DType::Decimal(d, _) => { let decimal_array = array.to_decimal(); - let validity = decimal_array.validity_mask().to_bit_buffer(); + let validity = decimal_array.validity_mask()?.to_bit_buffer(); match_each_decimal_value_type!(decimal_array.values_type(), |D| { let buf = decimal_array.buffer::(); let opt_values = buf @@ -111,7 +110,7 @@ pub fn search_sorted_canonical_array( }) .transpose()? .ok_or_else(|| vortex_err!("unexpected null scalar"))?; - Ok(SearchNullableSlice(opt_values).search_sorted(&Some(to_find), side)) + SearchNullableSlice(opt_values).search_sorted(&Some(to_find), side) }) } DType::Utf8(_) | DType::Binary(_) => { @@ -123,11 +122,13 @@ pub fn search_sorted_canonical_array( } else { ByteBuffer::try_from(scalar)?.to_vec() }; - Ok(SearchNullableSlice(opt_values).search_sorted(&Some(to_find), side)) + SearchNullableSlice(opt_values).search_sorted(&Some(to_find), side) } DType::Struct(..) | DType::List(..) | DType::FixedSizeList(..) => { - let scalar_vals = (0..array.len()).map(|i| array.scalar_at(i)).collect_vec(); - Ok(scalar_vals.search_sorted(&scalar.cast(array.dtype())?, side)) + let scalar_vals = (0..array.len()) + .map(|i| array.scalar_at(i)) + .collect::>>()?; + scalar_vals.search_sorted(&scalar.cast(array.dtype())?, side) } d @ (DType::Null | DType::Extension(_)) => { unreachable!("DType {d} not supported for fuzzing") diff --git a/fuzz/src/array/slice.rs b/fuzz/src/array/slice.rs index b3061312358..ea851b94f18 100644 --- a/fuzz/src/array/slice.rs +++ b/fuzz/src/array/slice.rs @@ -26,7 +26,7 @@ pub fn slice_canonical_array( stop: usize, ) -> VortexResult { let validity = if array.dtype().is_nullable() { - let bool_buff = array.validity_mask().to_bit_buffer(); + let bool_buff = array.validity_mask()?.to_bit_buffer(); Validity::from(bool_buff.slice(start..stop)) } else { Validity::NonNullable diff --git a/fuzz/src/array/sort.rs b/fuzz/src/array/sort.rs index b1ffaf2f28e..c2b12ad24cf 100644 --- a/fuzz/src/array/sort.rs +++ b/fuzz/src/array/sort.rs @@ -28,7 +28,7 @@ pub fn sort_canonical_array(array: &dyn Array) -> VortexResult { let mut opt_values = bool_array .bit_buffer() .iter() - .zip(bool_array.validity_mask().to_bit_buffer().iter()) + .zip(bool_array.validity_mask()?.to_bit_buffer().iter()) .map(|(b, v)| v.then_some(b)) .collect::>(); opt_values.sort(); @@ -41,7 +41,7 @@ pub fn sort_canonical_array(array: &dyn Array) -> VortexResult { .as_slice::

() .iter() .copied() - .zip(primitive_array.validity_mask().to_bit_buffer().iter()) + .zip(primitive_array.validity_mask()?.to_bit_buffer().iter()) .map(|(p, v)| v.then_some(p)) .collect::>(); sort_primitive_slice(&mut opt_values); @@ -56,7 +56,7 @@ pub fn sort_canonical_array(array: &dyn Array) -> VortexResult { .as_slice() .iter() .copied() - .zip(decimal_array.validity_mask().to_bit_buffer().iter()) + .zip(decimal_array.validity_mask()?.to_bit_buffer().iter()) .map(|(p, v)| v.then_some(p)) .collect::>(); opt_values.sort(); @@ -75,7 +75,8 @@ pub fn sort_canonical_array(array: &dyn Array) -> VortexResult { sort_indices.sort_by(|a, b| { array .scalar_at(*a) - .partial_cmp(&array.scalar_at(*b)) + .vortex_expect("scalar_at") + .partial_cmp(&array.scalar_at(*b).vortex_expect("scalar_at")) .vortex_expect("must be a valid comparison") }); take_canonical_array_non_nullable_indices(array, &sort_indices) diff --git a/fuzz/src/array/take.rs b/fuzz/src/array/take.rs index 6d4a3b7660b..c74be63dbf8 100644 --- a/fuzz/src/array/take.rs +++ b/fuzz/src/array/take.rs @@ -45,7 +45,7 @@ pub fn take_canonical_array( let nullable: Nullability = indices.contains(&None).into(); let validity = if array.dtype().is_nullable() || nullable == Nullability::Nullable { - let validity_idx = array.validity_mask().to_bit_buffer(); + let validity_idx = array.validity_mask()?.to_bit_buffer(); Validity::from_iter( indices @@ -131,7 +131,7 @@ pub fn take_canonical_array( if let Some(idx) = idx { builder.append_scalar( &array - .scalar_at(*idx) + .scalar_at(*idx)? .cast(&array.dtype().union_nullability(nullable)) .vortex_expect("cannot cast scalar nullability"), )?; diff --git a/vortex-array/benches/scalar_at_struct.rs b/vortex-array/benches/scalar_at_struct.rs index 69adb80415e..485ba84f339 100644 --- a/vortex-array/benches/scalar_at_struct.rs +++ b/vortex-array/benches/scalar_at_struct.rs @@ -48,7 +48,7 @@ fn scalar_at_struct_simple(bencher: Bencher) { .with_inputs(|| (&struct_array, &indices)) .bench_refs(|(array, indices)| { for &idx in indices.iter() { - divan::black_box(array.scalar_at(idx)); + divan::black_box(array.scalar_at(idx).unwrap()); } }); } @@ -83,7 +83,7 @@ fn scalar_at_struct_wide(bencher: Bencher) { .with_inputs(|| (&struct_array, &indices)) .bench_refs(|(array, indices)| { for &idx in indices.iter() { - divan::black_box(array.scalar_at(idx)); + divan::black_box(array.scalar_at(idx).unwrap()); } }); } diff --git a/vortex-array/benches/search_sorted.rs b/vortex-array/benches/search_sorted.rs index 1375f259b96..353dec65a3b 100644 --- a/vortex-array/benches/search_sorted.rs +++ b/vortex-array/benches/search_sorted.rs @@ -28,7 +28,7 @@ fn binary_search_vortex(bencher: Bencher) { let (sorted_array, target) = fixture(); bencher .with_inputs(|| (&sorted_array, &target)) - .bench_refs(|(array, target)| array.search_sorted(target, SearchSortedSide::Left)); + .bench_refs(|(array, target)| array.search_sorted(target, SearchSortedSide::Left).unwrap()); } fn fixture() -> (Vec, i32) { diff --git a/vortex-array/src/array/mod.rs b/vortex-array/src/array/mod.rs index d33adcbdff5..34a910f74c1 100644 --- a/vortex-array/src/array/mod.rs +++ b/vortex-array/src/array/mod.rs @@ -110,37 +110,37 @@ pub trait Array: /// Fetch the scalar at the given index. /// /// This method panics if the index is out of bounds for the array. - fn scalar_at(&self, index: usize) -> Scalar; + fn scalar_at(&self, index: usize) -> VortexResult; /// Returns whether the item at `index` is valid. - fn is_valid(&self, index: usize) -> bool; + fn is_valid(&self, index: usize) -> VortexResult; /// Returns whether the item at `index` is invalid. - fn is_invalid(&self, index: usize) -> bool; + fn is_invalid(&self, index: usize) -> VortexResult; /// Returns whether all items in the array are valid. /// /// This is usually cheaper than computing a precise `valid_count`, but may return false /// negatives. - fn all_valid(&self) -> bool; + fn all_valid(&self) -> VortexResult; /// Returns whether the array is all invalid. /// /// This is usually cheaper than computing a precise `invalid_count`, but may return false /// negatives. - fn all_invalid(&self) -> bool; + fn all_invalid(&self) -> VortexResult; /// Returns the number of valid elements in the array. - fn valid_count(&self) -> usize; + fn valid_count(&self) -> VortexResult; /// Returns the number of invalid elements in the array. - fn invalid_count(&self) -> usize; + fn invalid_count(&self) -> VortexResult; /// Returns the [`Validity`] of the array. fn validity(&self) -> VortexResult; /// Returns the canonical validity mask for the array. - fn validity_mask(&self) -> Mask; + fn validity_mask(&self) -> VortexResult; /// Returns the canonical representation of the array. fn to_canonical(&self) -> VortexResult; @@ -229,37 +229,37 @@ impl Array for Arc { } #[inline] - fn scalar_at(&self, index: usize) -> Scalar { + fn scalar_at(&self, index: usize) -> VortexResult { self.as_ref().scalar_at(index) } #[inline] - fn is_valid(&self, index: usize) -> bool { + fn is_valid(&self, index: usize) -> VortexResult { self.as_ref().is_valid(index) } #[inline] - fn is_invalid(&self, index: usize) -> bool { + fn is_invalid(&self, index: usize) -> VortexResult { self.as_ref().is_invalid(index) } #[inline] - fn all_valid(&self) -> bool { + fn all_valid(&self) -> VortexResult { self.as_ref().all_valid() } #[inline] - fn all_invalid(&self) -> bool { + fn all_invalid(&self) -> VortexResult { self.as_ref().all_invalid() } #[inline] - fn valid_count(&self) -> usize { + fn valid_count(&self) -> VortexResult { self.as_ref().valid_count() } #[inline] - fn invalid_count(&self) -> usize { + fn invalid_count(&self) -> VortexResult { self.as_ref().invalid_count() } @@ -269,7 +269,7 @@ impl Array for Arc { } #[inline] - fn validity_mask(&self) -> Mask { + fn validity_mask(&self) -> VortexResult { self.as_ref().validity_mask() } @@ -537,84 +537,76 @@ impl Array for ArrayAdapter { .optimize() } - fn scalar_at(&self, index: usize) -> Scalar { - assert!(index < self.len(), "index {index} out of bounds"); - if self.is_invalid(index) { - return Scalar::null(self.dtype().clone()); + fn scalar_at(&self, index: usize) -> VortexResult { + vortex_ensure!(index < self.len(), OutOfBounds: index, 0, self.len()); + if self.is_invalid(index)? { + return Ok(Scalar::null(self.dtype().clone())); } - let scalar = >::scalar_at(&self.0, index); - assert_eq!(self.dtype(), scalar.dtype(), "Scalar dtype mismatch"); - scalar - } - - fn is_valid(&self, index: usize) -> bool { - if index >= self.len() { - vortex_panic!(OutOfBounds: index, 0, self.len()); - } - match self - .validity() - .vortex_expect("Failed to get validity for is_valid") - { - Validity::NonNullable | Validity::AllValid => true, - Validity::AllInvalid => false, - Validity::Array(a) => a - .scalar_at(index) + let scalar = >::scalar_at(&self.0, index)?; + vortex_ensure!(self.dtype() == scalar.dtype(), "Scalar dtype mismatch"); + Ok(scalar) + } + + fn is_valid(&self, index: usize) -> VortexResult { + vortex_ensure!(index < self.len(), OutOfBounds: index, 0, self.len()); + match self.validity()? { + Validity::NonNullable | Validity::AllValid => Ok(true), + Validity::AllInvalid => Ok(false), + Validity::Array(a) => Ok(a + .scalar_at(index)? .as_bool() .value() - .vortex_expect("validity must be non-nullable"), + .vortex_expect("validity must be non-nullable")), } } - fn is_invalid(&self, index: usize) -> bool { - !self.is_valid(index) + fn is_invalid(&self, index: usize) -> VortexResult { + Ok(!self.is_valid(index)?) } - fn all_valid(&self) -> bool { - match self.validity().vortex_expect("Failed to get validity") { - Validity::NonNullable | Validity::AllValid => true, - Validity::AllInvalid => false, - Validity::Array(a) => a.statistics().compute_min::().unwrap_or(false), + fn all_valid(&self) -> VortexResult { + match self.validity()? { + Validity::NonNullable | Validity::AllValid => Ok(true), + Validity::AllInvalid => Ok(false), + Validity::Array(a) => Ok(a.statistics().compute_min::().unwrap_or(false)), } } - fn all_invalid(&self) -> bool { - match self.validity().vortex_expect("Failed to get validity") { - Validity::NonNullable | Validity::AllValid => false, - Validity::AllInvalid => true, - Validity::Array(a) => !a.statistics().compute_max::().unwrap_or(true), + fn all_invalid(&self) -> VortexResult { + match self.validity()? { + Validity::NonNullable | Validity::AllValid => Ok(false), + Validity::AllInvalid => Ok(true), + Validity::Array(a) => Ok(!a.statistics().compute_max::().unwrap_or(true)), } } - fn valid_count(&self) -> usize { + fn valid_count(&self) -> VortexResult { if let Some(Precision::Exact(invalid_count)) = self.statistics().get_as::(Stat::NullCount) { - return self.len() - invalid_count; + return Ok(self.len() - invalid_count); } - let count = match self - .validity() - .vortex_expect("Failed to get validity for valid_count") - { + let count = match self.validity()? { Validity::NonNullable | Validity::AllValid => self.len(), Validity::AllInvalid => 0, Validity::Array(a) => { - let sum = compute::sum(&a).vortex_expect("Failed to compute sum for valid count"); + let sum = compute::sum(&a)?; sum.as_primitive() .as_::() .vortex_expect("Sum must be non-nullable") } }; - assert!(count <= self.len(), "Valid count exceeds array length"); + vortex_ensure!(count <= self.len(), "Valid count exceeds array length"); self.statistics() .set(Stat::NullCount, Precision::exact(self.len() - count)); - count + Ok(count) } - fn invalid_count(&self) -> usize { - self.len() - self.valid_count() + fn invalid_count(&self) -> VortexResult { + Ok(self.len() - self.valid_count()?) } fn validity(&self) -> VortexResult { @@ -634,13 +626,11 @@ impl Array for ArrayAdapter { } } - fn validity_mask(&self) -> Mask { - match self.validity().vortex_expect("Failed to get validity") { - Validity::NonNullable | Validity::AllValid => Mask::new_true(self.len()), - Validity::AllInvalid => Mask::new_false(self.len()), - Validity::Array(a) => a - .try_to_mask_fill_null_false() - .vortex_expect("Failed to get validity mask"), + fn validity_mask(&self) -> VortexResult { + match self.validity()? { + Validity::NonNullable | Validity::AllValid => Ok(Mask::new_true(self.len())), + Validity::AllInvalid => Ok(Mask::new_false(self.len())), + Validity::Array(a) => a.try_to_mask_fill_null_false(), } } diff --git a/vortex-array/src/arrays/assertions.rs b/vortex-array/src/arrays/assertions.rs index 3f452e1b575..226470b9cbb 100644 --- a/vortex-array/src/arrays/assertions.rs +++ b/vortex-array/src/arrays/assertions.rs @@ -22,7 +22,7 @@ pub fn format_indices>(indices: I) -> impl Display #[macro_export] macro_rules! assert_nth_scalar { ($arr:expr, $n:expr, $expected:expr) => { - assert_eq!($arr.scalar_at($n), $expected.try_into().unwrap()); + assert_eq!($arr.scalar_at($n).unwrap(), $expected.try_into().unwrap()); }; } @@ -52,7 +52,7 @@ macro_rules! assert_arrays_eq { } let n = left.len(); let mismatched_indices = (0..n) - .filter(|i| left.scalar_at(*i) != right.scalar_at(*i)) + .filter(|i| left.scalar_at(*i).unwrap() != right.scalar_at(*i).unwrap()) .collect::>(); if mismatched_indices.len() != 0 { panic!( diff --git a/vortex-array/src/arrays/bool/array.rs b/vortex-array/src/arrays/bool/array.rs index 28850af39d8..d4f34e02868 100644 --- a/vortex-array/src/arrays/bool/array.rs +++ b/vortex-array/src/arrays/bool/array.rs @@ -42,7 +42,7 @@ use crate::validity::Validity; /// assert_eq!(sliced.len(), 2); /// /// // Access individual values -/// let value = array.scalar_at(0); +/// let value = array.scalar_at(0).unwrap(); /// assert_eq!(value, true.into()); /// ``` #[derive(Clone, Debug)] @@ -188,12 +188,14 @@ impl BoolArray { pub fn to_mask(&self) -> Mask { self.maybe_to_mask() + .vortex_expect("failed to check validity") .vortex_expect("cannot convert nullable boolean array to mask") } - pub fn maybe_to_mask(&self) -> Option { - self.all_valid() - .then(|| Mask::from_buffer(self.bit_buffer().clone())) + pub fn maybe_to_mask(&self) -> VortexResult> { + Ok(self + .all_valid()? + .then(|| Mask::from_buffer(self.bit_buffer().clone()))) } pub fn to_mask_fill_null_false(&self) -> Mask { @@ -206,7 +208,10 @@ impl BoolArray { } } // Extract a boolean buffer, treating null values to false - let buffer = match self.validity_mask() { + let buffer = match self + .validity_mask() + .unwrap_or_else(|_| Mask::new_true(self.len())) + { Mask::AllTrue(_) => self.bit_buffer().clone(), Mask::AllFalse(_) => return Mask::new_false(self.len()), Mask::Values(validity) => validity.bit_buffer() & self.bit_buffer(), @@ -270,7 +275,7 @@ mod tests { #[test] fn bool_array() { let arr = BoolArray::from_iter([true, false, true]); - let scalar = bool::try_from(&arr.scalar_at(0)).unwrap(); + let scalar = bool::try_from(&arr.scalar_at(0).unwrap()).unwrap(); assert!(scalar); } @@ -280,9 +285,9 @@ mod tests { assert!(matches!(arr.validity(), Validity::AllValid)); - let scalar = bool::try_from(&arr.scalar_at(0)).unwrap(); + let scalar = bool::try_from(&arr.scalar_at(0).unwrap()).unwrap(); assert!(scalar); - let scalar = bool::try_from(&arr.scalar_at(1)).unwrap(); + let scalar = bool::try_from(&arr.scalar_at(1).unwrap()).unwrap(); assert!(!scalar); } @@ -290,19 +295,19 @@ mod tests { fn test_bool_from_iter() { let arr = BoolArray::from_iter([Some(true), Some(true), None, Some(false), None]); - let scalar = bool::try_from(&arr.scalar_at(0)).unwrap(); + let scalar = bool::try_from(&arr.scalar_at(0).unwrap()).unwrap(); assert!(scalar); - let scalar = bool::try_from(&arr.scalar_at(1)).unwrap(); + let scalar = bool::try_from(&arr.scalar_at(1).unwrap()).unwrap(); assert!(scalar); - let scalar = arr.scalar_at(2); + let scalar = arr.scalar_at(2).unwrap(); assert!(scalar.is_null()); - let scalar = bool::try_from(&arr.scalar_at(3)).unwrap(); + let scalar = bool::try_from(&arr.scalar_at(3).unwrap()).unwrap(); assert!(!scalar); - let scalar = arr.scalar_at(4); + let scalar = arr.scalar_at(4).unwrap(); assert!(scalar.is_null()); } diff --git a/vortex-array/src/arrays/bool/compute/is_sorted.rs b/vortex-array/src/arrays/bool/compute/is_sorted.rs index c23de6e8635..33d05f6ab48 100644 --- a/vortex-array/src/arrays/bool/compute/is_sorted.rs +++ b/vortex-array/src/arrays/bool/compute/is_sorted.rs @@ -13,7 +13,7 @@ use crate::register_kernel; impl IsSortedKernel for BoolVTable { fn is_sorted(&self, array: &BoolArray) -> VortexResult> { - match array.validity_mask() { + match array.validity_mask()? { Mask::AllFalse(_) => Ok(Some(true)), Mask::AllTrue(_) => Ok(Some(array.bit_buffer().iter().is_sorted())), Mask::Values(mask_values) => { @@ -32,7 +32,7 @@ impl IsSortedKernel for BoolVTable { } fn is_strict_sorted(&self, array: &BoolArray) -> VortexResult> { - match array.validity_mask() { + match array.validity_mask()? { Mask::AllFalse(_) => Ok(Some(false)), Mask::AllTrue(_) => Ok(Some(array.bit_buffer().iter().is_strict_sorted())), Mask::Values(mask_values) => { diff --git a/vortex-array/src/arrays/bool/compute/min_max.rs b/vortex-array/src/arrays/bool/compute/min_max.rs index d42b5796769..f39990575ab 100644 --- a/vortex-array/src/arrays/bool/compute/min_max.rs +++ b/vortex-array/src/arrays/bool/compute/min_max.rs @@ -18,11 +18,11 @@ use crate::register_kernel; impl MinMaxKernel for BoolVTable { fn min_max(&self, array: &BoolArray) -> VortexResult> { - let mask = array.validity_mask(); - let true_non_null = match mask { + let mask = array.validity_mask()?; + let true_non_null = match &mask { Mask::AllTrue(_) => array.bit_buffer().clone(), Mask::AllFalse(_) => return Ok(None), - Mask::Values(ref v) => array.bit_buffer().bitand(v.bit_buffer()), + Mask::Values(v) => array.bit_buffer().bitand(v.bit_buffer()), }; // TODO(ngates): we should be able to bail out earlier as soon as we have one true and @@ -48,7 +48,7 @@ impl MinMaxKernel for BoolVTable { }; // If the non null true slice doesn't cover the whole array we need to check for valid false values - match mask { + match &mask { // if the mask is all true or all false we don't need to look for false values Mask::AllTrue(_) | Mask::AllFalse(_) => {} Mask::Values(v) => { diff --git a/vortex-array/src/arrays/bool/compute/sum.rs b/vortex-array/src/arrays/bool/compute/sum.rs index 752d0eec2cc..3ddbbcd5dd8 100644 --- a/vortex-array/src/arrays/bool/compute/sum.rs +++ b/vortex-array/src/arrays/bool/compute/sum.rs @@ -16,7 +16,7 @@ use crate::register_kernel; impl SumKernel for BoolVTable { fn sum(&self, array: &BoolArray, accumulator: &Scalar) -> VortexResult { - let true_count: Option = match array.validity_mask().bit_buffer() { + let true_count: Option = match array.validity_mask()?.bit_buffer() { AllOr::All => { // All-valid Some(array.bit_buffer().true_count() as u64) diff --git a/vortex-array/src/arrays/bool/compute/take.rs b/vortex-array/src/arrays/bool/compute/take.rs index 04ac3c23f5f..6132e5d0af0 100644 --- a/vortex-array/src/arrays/bool/compute/take.rs +++ b/vortex-array/src/arrays/bool/compute/take.rs @@ -25,7 +25,7 @@ use crate::vtable::ValidityHelper; impl TakeKernel for BoolVTable { fn take(&self, array: &BoolArray, indices: &dyn Array) -> VortexResult { - let indices_nulls_zeroed = match indices.validity_mask() { + let indices_nulls_zeroed = match indices.validity_mask()? { Mask::AllTrue(_) => indices.to_array(), Mask::AllFalse(_) => { return Ok(ConstantArray::new( diff --git a/vortex-array/src/arrays/bool/test_harness.rs b/vortex-array/src/arrays/bool/test_harness.rs index f1df7f8f2d2..958e8661b52 100644 --- a/vortex-array/src/arrays/bool/test_harness.rs +++ b/vortex-array/src/arrays/bool/test_harness.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_error::VortexExpect; use vortex_error::vortex_panic; use crate::arrays::BoolArray; @@ -8,6 +9,7 @@ use crate::arrays::BoolArray; impl BoolArray { pub fn opt_bool_vec(&self) -> Vec> { self.validity_mask() + .vortex_expect("Failed to get validity mask") .to_bit_buffer() .iter() .zip(self.bit_buffer().iter()) @@ -17,6 +19,7 @@ impl BoolArray { pub fn bool_vec(&self) -> Vec { self.validity_mask() + .vortex_expect("Failed to get validity mask") .to_bit_buffer() .iter() .zip(self.bit_buffer().iter()) diff --git a/vortex-array/src/arrays/bool/vtable/operations.rs b/vortex-array/src/arrays/bool/vtable/operations.rs index 61ac756e3b9..a39c3319906 100644 --- a/vortex-array/src/arrays/bool/vtable/operations.rs +++ b/vortex-array/src/arrays/bool/vtable/operations.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_error::VortexResult; use vortex_scalar::Scalar; use crate::arrays::BoolArray; @@ -8,8 +9,11 @@ use crate::arrays::BoolVTable; use crate::vtable::OperationsVTable; impl OperationsVTable for BoolVTable { - fn scalar_at(array: &BoolArray, index: usize) -> Scalar { - Scalar::bool(array.bit_buffer().value(index), array.dtype().nullability()) + fn scalar_at(array: &BoolArray, index: usize) -> VortexResult { + Ok(Scalar::bool( + array.bit_buffer().value(index), + array.dtype().nullability(), + )) } } diff --git a/vortex-array/src/arrays/chunked/array.rs b/vortex-array/src/arrays/chunked/array.rs index b1a7924ac8b..66a9381d16e 100644 --- a/vortex-array/src/arrays/chunked/array.rs +++ b/vortex-array/src/arrays/chunked/array.rs @@ -120,7 +120,7 @@ impl ChunkedArray { self.chunk_offsets.to_buffer() } - pub(crate) fn find_chunk_idx(&self, index: usize) -> (usize, usize) { + pub(crate) fn find_chunk_idx(&self, index: usize) -> VortexResult<(usize, usize)> { assert!(index <= self.len(), "Index out of bounds of the array"); let index = index as u64; @@ -128,14 +128,14 @@ impl ChunkedArray { // and take the last chunk (we subtract 1 since there's a leading 0) let index_chunk = self .chunk_offsets() - .search_sorted(&index, SearchSortedSide::Right) + .search_sorted(&index, SearchSortedSide::Right)? .to_ends_index(self.nchunks() + 1) .saturating_sub(1); let chunk_start = self.chunk_offsets()[index_chunk]; let index_in_chunk = usize::try_from(index - chunk_start).vortex_expect("Index is too large for usize"); - (index_chunk, index_in_chunk) + Ok((index_chunk, index_in_chunk)) } pub fn chunks(&self) -> &[ArrayRef] { @@ -348,8 +348,8 @@ mod test { ChunkedArray::try_new(chunks, DType::Primitive(PType::U64, Nullability::Nullable))?; // Should be all_valid since all non-empty chunks are all_valid - assert!(chunked.all_valid()); - assert!(!chunked.all_invalid()); + assert!(chunked.all_valid().unwrap()); + assert!(!chunked.all_invalid().unwrap()); Ok(()) } @@ -368,8 +368,8 @@ mod test { ChunkedArray::try_new(chunks, DType::Primitive(PType::U64, Nullability::Nullable))?; // Should be all_invalid since all non-empty chunks are all_invalid - assert!(!chunked.all_valid()); - assert!(chunked.all_invalid()); + assert!(!chunked.all_valid().unwrap()); + assert!(chunked.all_invalid().unwrap()); Ok(()) } @@ -388,8 +388,8 @@ mod test { ChunkedArray::try_new(chunks, DType::Primitive(PType::U64, Nullability::Nullable))?; // Should be neither all_valid nor all_invalid - assert!(!chunked.all_valid()); - assert!(!chunked.all_invalid()); + assert!(!chunked.all_valid().unwrap()); + assert!(!chunked.all_invalid().unwrap()); Ok(()) } diff --git a/vortex-array/src/arrays/chunked/compute/filter.rs b/vortex-array/src/arrays/chunked/compute/filter.rs index 14849413517..f77ffd6ef9d 100644 --- a/vortex-array/src/arrays/chunked/compute/filter.rs +++ b/vortex-array/src/arrays/chunked/compute/filter.rs @@ -92,10 +92,10 @@ pub(crate) fn chunk_filters( let mut chunk_filters = vec![ChunkFilter::None; array.nchunks()]; for (slice_start, slice_end) in slices { - let (start_chunk, start_idx) = find_chunk_idx(slice_start, &chunk_offsets); + let (start_chunk, start_idx) = find_chunk_idx(slice_start, &chunk_offsets)?; // NOTE: we adjust slice end back by one, in case it ends on a chunk boundary, we do not // want to index into the unused chunk. - let (end_chunk, end_idx) = find_chunk_idx(slice_end - 1, &chunk_offsets); + let (end_chunk, end_idx) = find_chunk_idx(slice_end - 1, &chunk_offsets)?; // Adjust back to an exclusive range let end_idx = end_idx + 1; @@ -153,7 +153,7 @@ fn filter_indices( let chunk_offsets = array.chunk_offsets(); for set_index in indices { - let (chunk_id, index) = find_chunk_idx(set_index, &chunk_offsets); + let (chunk_id, index) = find_chunk_idx(set_index, &chunk_offsets)?; if chunk_id != current_chunk_id { // Push the chunk we've accumulated. if !chunk_indices.is_empty() { @@ -188,9 +188,9 @@ fn filter_indices( /// Mirrors the find_chunk_idx method on ChunkedArray, but avoids all of the overhead /// from scalars, dtypes, and metadata cloning. -pub(crate) fn find_chunk_idx(idx: usize, chunk_ends: &[u64]) -> (usize, usize) { +pub(crate) fn find_chunk_idx(idx: usize, chunk_ends: &[u64]) -> VortexResult<(usize, usize)> { let chunk_id = chunk_ends - .search_sorted(&(idx as u64), SearchSortedSide::Right) + .search_sorted(&(idx as u64), SearchSortedSide::Right)? .to_ends_index(chunk_ends.len()) .saturating_sub(1); let chunk_begin: usize = chunk_ends[chunk_id] @@ -198,7 +198,7 @@ pub(crate) fn find_chunk_idx(idx: usize, chunk_ends: &[u64]) -> (usize, usize) { .vortex_expect("chunk end must fit in usize"); let chunk_offset = idx - chunk_begin; - (chunk_id, chunk_offset) + Ok((chunk_id, chunk_offset)) } #[cfg(test)] diff --git a/vortex-array/src/arrays/chunked/compute/is_constant.rs b/vortex-array/src/arrays/chunked/compute/is_constant.rs index e70e1405e87..b8909fe043b 100644 --- a/vortex-array/src/arrays/chunked/compute/is_constant.rs +++ b/vortex-array/src/arrays/chunked/compute/is_constant.rs @@ -32,7 +32,7 @@ impl IsConstantKernel for ChunkedVTable { Some(true) => {} } - let first_value = first_chunk.scalar_at(0).into_nullable(); + let first_value = first_chunk.scalar_at(0)?.into_nullable(); for chunk in chunks { match is_constant_opts(chunk, opts)? { @@ -42,7 +42,7 @@ impl IsConstantKernel for ChunkedVTable { Some(true) => {} } - if first_value != chunk.scalar_at(0).into_nullable() { + if first_value != chunk.scalar_at(0)?.into_nullable() { return Ok(Some(false)); } } diff --git a/vortex-array/src/arrays/chunked/compute/is_sorted.rs b/vortex-array/src/arrays/chunked/compute/is_sorted.rs index b767b1d9695..76577946cc0 100644 --- a/vortex-array/src/arrays/chunked/compute/is_sorted.rs +++ b/vortex-array/src/arrays/chunked/compute/is_sorted.rs @@ -36,8 +36,8 @@ fn is_sorted_impl( continue; } - let first = chunk.scalar_at(0); - let last = chunk.scalar_at(chunk.len() - 1); + let first = chunk.scalar_at(0)?; + let last = chunk.scalar_at(chunk.len() - 1)?; first_last.push((first, last)); } diff --git a/vortex-array/src/arrays/chunked/compute/mask.rs b/vortex-array/src/arrays/chunked/compute/mask.rs index 034c3f84297..20fb44c37ce 100644 --- a/vortex-array/src/arrays/chunked/compute/mask.rs +++ b/vortex-array/src/arrays/chunked/compute/mask.rs @@ -63,7 +63,7 @@ fn mask_indices( let chunk_offsets = array.chunk_offsets(); for &set_index in indices { - let (chunk_id, index) = find_chunk_idx(set_index, &chunk_offsets); + let (chunk_id, index) = find_chunk_idx(set_index, &chunk_offsets)?; if chunk_id != current_chunk_id { let chunk = array.chunk(current_chunk_id).clone(); let chunk_len = chunk.len(); diff --git a/vortex-array/src/arrays/chunked/compute/take.rs b/vortex-array/src/arrays/chunked/compute/take.rs index 009d7bad300..5facc27e503 100644 --- a/vortex-array/src/arrays/chunked/compute/take.rs +++ b/vortex-array/src/arrays/chunked/compute/take.rs @@ -30,7 +30,7 @@ impl TakeKernel for ChunkedVTable { // TODO(joe): Should we split this implementation based on indices nullability? let nullability = indices.dtype().nullability(); - let indices_mask = indices.validity_mask(); + let indices_mask = indices.validity_mask()?; let indices = indices.as_slice::(); let mut chunks = Vec::new(); @@ -38,10 +38,10 @@ impl TakeKernel for ChunkedVTable { let mut start = 0; let mut stop = 0; // We assume indices are non-empty as it's handled in the top-level `take` function - let mut prev_chunk_idx = array.find_chunk_idx(indices[0].try_into()?).0; + let mut prev_chunk_idx = array.find_chunk_idx(indices[0].try_into()?)?.0; for idx in indices { let idx = usize::try_from(*idx)?; - let (chunk_idx, idx_in_chunk) = array.find_chunk_idx(idx); + let (chunk_idx, idx_in_chunk) = array.find_chunk_idx(idx)?; if chunk_idx != prev_chunk_idx { // Start a new chunk diff --git a/vortex-array/src/arrays/chunked/tests.rs b/vortex-array/src/arrays/chunked/tests.rs index 0e856ee68b0..48f0571fc3f 100644 --- a/vortex-array/src/arrays/chunked/tests.rs +++ b/vortex-array/src/arrays/chunked/tests.rs @@ -191,6 +191,6 @@ pub fn pack_nested_lists() { let canon_values = chunked_list.unwrap().to_listview(); - assert_eq!(l1.scalar_at(0), canon_values.scalar_at(0)); - assert_eq!(l2.scalar_at(0), canon_values.scalar_at(1)); + assert_eq!(l1.scalar_at(0).unwrap(), canon_values.scalar_at(0).unwrap()); + assert_eq!(l2.scalar_at(0).unwrap(), canon_values.scalar_at(1).unwrap()); } diff --git a/vortex-array/src/arrays/chunked/vtable/canonical.rs b/vortex-array/src/arrays/chunked/vtable/canonical.rs index cc534e481c4..9142459c28a 100644 --- a/vortex-array/src/arrays/chunked/vtable/canonical.rs +++ b/vortex-array/src/arrays/chunked/vtable/canonical.rs @@ -38,7 +38,7 @@ pub(super) fn _canonicalize( DType::Struct(struct_dtype, _) => { let struct_array = pack_struct_chunks( array.chunks(), - Validity::copy_from_array(array.as_ref()), + Validity::copy_from_array(array.as_ref())?, struct_dtype, ctx, )?; @@ -46,7 +46,7 @@ pub(super) fn _canonicalize( } DType::List(elem_dtype, _) => Canonical::List(swizzle_list_chunks( array.chunks(), - Validity::copy_from_array(array.as_ref()), + Validity::copy_from_array(array.as_ref())?, elem_dtype, ctx, )?), @@ -256,7 +256,7 @@ mod tests { let canon_values = chunked_list.unwrap().to_listview(); - assert_eq!(l1.scalar_at(0), canon_values.scalar_at(0)); - assert_eq!(l2.scalar_at(0), canon_values.scalar_at(1)); + assert_eq!(l1.scalar_at(0).unwrap(), canon_values.scalar_at(0).unwrap()); + assert_eq!(l2.scalar_at(0).unwrap(), canon_values.scalar_at(1).unwrap()); } } diff --git a/vortex-array/src/arrays/chunked/vtable/mod.rs b/vortex-array/src/arrays/chunked/vtable/mod.rs index 7bc1df33f3b..6047ab13c35 100644 --- a/vortex-array/src/arrays/chunked/vtable/mod.rs +++ b/vortex-array/src/arrays/chunked/vtable/mod.rs @@ -208,8 +208,8 @@ impl VTable for ChunkedVTable { } } - let (offset_chunk, offset_in_first_chunk) = array.find_chunk_idx(range.start); - let (length_chunk, length_in_last_chunk) = array.find_chunk_idx(range.end); + let (offset_chunk, offset_in_first_chunk) = array.find_chunk_idx(range.start)?; + let (length_chunk, length_in_last_chunk) = array.find_chunk_idx(range.end)?; if length_chunk == offset_chunk { let chunk = array.chunk(offset_chunk); diff --git a/vortex-array/src/arrays/chunked/vtable/operations.rs b/vortex-array/src/arrays/chunked/vtable/operations.rs index 28b62a78849..3762dbab42f 100644 --- a/vortex-array/src/arrays/chunked/vtable/operations.rs +++ b/vortex-array/src/arrays/chunked/vtable/operations.rs @@ -1,15 +1,17 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_error::VortexResult; use vortex_scalar::Scalar; +use crate::Array; use crate::arrays::ChunkedArray; use crate::arrays::ChunkedVTable; use crate::vtable::OperationsVTable; impl OperationsVTable for ChunkedVTable { - fn scalar_at(array: &ChunkedArray, index: usize) -> Scalar { - let (chunk_index, chunk_offset) = array.find_chunk_idx(index); + fn scalar_at(array: &ChunkedArray, index: usize) -> VortexResult { + let (chunk_index, chunk_offset) = array.find_chunk_idx(index)?; array.chunk(chunk_index).scalar_at(chunk_offset) } } @@ -105,8 +107,8 @@ mod tests { DType::Primitive(PType::U64, Nullability::NonNullable), ) .unwrap(); - assert_eq!(array.scalar_at(0), 1u64.into()); - assert_eq!(array.scalar_at(1), 2u64.into()); + assert_eq!(array.scalar_at(0).unwrap(), 1u64.into()); + assert_eq!(array.scalar_at(1).unwrap(), 2u64.into()); } #[test] @@ -121,10 +123,10 @@ mod tests { DType::Primitive(PType::U64, Nullability::NonNullable), ) .unwrap(); - assert_eq!(array.scalar_at(0), 1u64.into()); - assert_eq!(array.scalar_at(1), 2u64.into()); - assert_eq!(array.scalar_at(2), 3u64.into()); - assert_eq!(array.scalar_at(3), 4u64.into()); + assert_eq!(array.scalar_at(0).unwrap(), 1u64.into()); + assert_eq!(array.scalar_at(1).unwrap(), 2u64.into()); + assert_eq!(array.scalar_at(2).unwrap(), 3u64.into()); + assert_eq!(array.scalar_at(3).unwrap(), 4u64.into()); } #[test] @@ -139,9 +141,9 @@ mod tests { DType::Primitive(PType::U64, Nullability::NonNullable), ) .unwrap(); - assert_eq!(array.scalar_at(0), 1u64.into()); - assert_eq!(array.scalar_at(1), 2u64.into()); - assert_eq!(array.scalar_at(2), 3u64.into()); - assert_eq!(array.scalar_at(3), 4u64.into()); + assert_eq!(array.scalar_at(0).unwrap(), 1u64.into()); + assert_eq!(array.scalar_at(1).unwrap(), 2u64.into()); + assert_eq!(array.scalar_at(2).unwrap(), 3u64.into()); + assert_eq!(array.scalar_at(3).unwrap(), 4u64.into()); } } diff --git a/vortex-array/src/arrays/chunked/vtable/validity.rs b/vortex-array/src/arrays/chunked/vtable/validity.rs index 0d3a0c5ddfc..fee95a4be20 100644 --- a/vortex-array/src/arrays/chunked/vtable/validity.rs +++ b/vortex-array/src/arrays/chunked/vtable/validity.rs @@ -53,7 +53,7 @@ impl ValidityVTable for ChunkedVTable { )) } - fn validity_mask(array: &ChunkedArray) -> Mask { + fn validity_mask(array: &ChunkedArray) -> VortexResult { array.chunks().iter().map(|a| a.validity_mask()).collect() } } diff --git a/vortex-array/src/arrays/constant/array.rs b/vortex-array/src/arrays/constant/array.rs index 0828170d842..b1a38d6c0d8 100644 --- a/vortex-array/src/arrays/constant/array.rs +++ b/vortex-array/src/arrays/constant/array.rs @@ -29,4 +29,8 @@ impl ConstantArray { pub fn scalar(&self) -> &Scalar { &self.scalar } + + pub fn into_parts(self) -> Scalar { + self.scalar + } } diff --git a/vortex-array/src/arrays/constant/compute/boolean.rs b/vortex-array/src/arrays/constant/compute/boolean.rs index 2c1978ac478..f6683007f51 100644 --- a/vortex-array/src/arrays/constant/compute/boolean.rs +++ b/vortex-array/src/arrays/constant/compute/boolean.rs @@ -104,10 +104,10 @@ mod test { fn test_or(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) { let r = or(&lhs, &rhs).unwrap().to_bool().into_array(); - let v0 = r.scalar_at(0).as_bool().value(); - let v1 = r.scalar_at(1).as_bool().value(); - let v2 = r.scalar_at(2).as_bool().value(); - let v3 = r.scalar_at(3).as_bool().value(); + let v0 = r.scalar_at(0).unwrap().as_bool().value(); + let v1 = r.scalar_at(1).unwrap().as_bool().value(); + let v2 = r.scalar_at(2).unwrap().as_bool().value(); + let v3 = r.scalar_at(3).unwrap().as_bool().value(); assert!(v0.unwrap()); assert!(v1.unwrap()); @@ -123,10 +123,10 @@ mod test { fn test_and(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) { let r = and(&lhs, &rhs).unwrap().to_bool().into_array(); - let v0 = r.scalar_at(0).as_bool().value(); - let v1 = r.scalar_at(1).as_bool().value(); - let v2 = r.scalar_at(2).as_bool().value(); - let v3 = r.scalar_at(3).as_bool().value(); + let v0 = r.scalar_at(0).unwrap().as_bool().value(); + let v1 = r.scalar_at(1).unwrap().as_bool().value(); + let v2 = r.scalar_at(2).unwrap().as_bool().value(); + let v3 = r.scalar_at(3).unwrap().as_bool().value(); assert!(v0.unwrap()); assert!(!v1.unwrap()); diff --git a/vortex-array/src/arrays/constant/compute/take.rs b/vortex-array/src/arrays/constant/compute/take.rs index 4e90863b0ed..24b60dd4fb9 100644 --- a/vortex-array/src/arrays/constant/compute/take.rs +++ b/vortex-array/src/arrays/constant/compute/take.rs @@ -18,7 +18,7 @@ use crate::validity::Validity; impl TakeKernel for ConstantVTable { fn take(&self, array: &ConstantArray, indices: &dyn Array) -> VortexResult { - match indices.validity_mask().bit_buffer() { + match indices.validity_mask()?.bit_buffer() { AllOr::All => { let scalar = Scalar::new( array @@ -95,7 +95,10 @@ mod tests { Validity::from_iter([false, true, false]) ) ); - assert_eq!(taken.validity_mask().indices(), AllOr::Some(valid_indices)); + assert_eq!( + taken.validity_mask().unwrap().indices(), + AllOr::Some(valid_indices) + ); } #[test] @@ -114,7 +117,7 @@ mod tests { taken.to_primitive(), PrimitiveArray::new(buffer![42i32, 42, 42], Validity::AllValid) ); - assert_eq!(taken.validity_mask().indices(), AllOr::All); + assert_eq!(taken.validity_mask().unwrap().indices(), AllOr::All); } #[rstest] diff --git a/vortex-array/src/arrays/constant/vtable/canonical.rs b/vortex-array/src/arrays/constant/vtable/canonical.rs index 679b336a630..fdd80e7d54d 100644 --- a/vortex-array/src/arrays/constant/vtable/canonical.rs +++ b/vortex-array/src/arrays/constant/vtable/canonical.rs @@ -347,7 +347,7 @@ mod tests { let const_null = ConstantArray::new(Scalar::null(DType::Null), 42); let actual = const_null.to_null(); assert_eq!(actual.len(), 42); - assert_eq!(actual.scalar_at(33), Scalar::null(DType::Null)); + assert_eq!(actual.scalar_at(33).unwrap(), Scalar::null(DType::Null)); } #[test] @@ -360,7 +360,7 @@ mod tests { assert_eq!(canonical.len(), 4); for i in 0..=3 { - assert_eq!(canonical.scalar_at(i), "four".into()); + assert_eq!(canonical.scalar_at(i).unwrap(), "four".into()); } } @@ -400,7 +400,7 @@ mod tests { let canonical_const = const_array.to_primitive(); // Verify the scalar value is preserved through canonicalization - assert_eq!(canonical_const.scalar_at(0), f16_scalar); + assert_eq!(canonical_const.scalar_at(0).unwrap(), f16_scalar); } #[test] @@ -482,7 +482,7 @@ mod tests { let struct_array = array.to_struct(); assert_eq!(struct_array.len(), 3); - assert_eq!(struct_array.valid_count(), 0); + assert_eq!(struct_array.valid_count().unwrap(), 0); let field = struct_array.field_by_name("non_null_field").unwrap(); @@ -606,10 +606,10 @@ mod tests { // Check elements are repeated correctly. let elements = canonical.elements().to_varbinview(); - assert_eq!(elements.scalar_at(0), "hello".into()); - assert_eq!(elements.scalar_at(1), "world".into()); - assert_eq!(elements.scalar_at(2), "hello".into()); - assert_eq!(elements.scalar_at(3), "world".into()); + assert_eq!(elements.scalar_at(0).unwrap(), "hello".into()); + assert_eq!(elements.scalar_at(1).unwrap(), "world".into()); + assert_eq!(elements.scalar_at(2).unwrap(), "hello".into()); + assert_eq!(elements.scalar_at(3).unwrap(), "world".into()); } #[test] @@ -653,12 +653,12 @@ mod tests { // Check elements including nulls. let elements = canonical.elements().to_primitive(); - assert_eq!(elements.scalar_at(0), Scalar::from(100i32)); + assert_eq!(elements.scalar_at(0).unwrap(), Scalar::from(100i32)); assert_eq!( - elements.scalar_at(1), + elements.scalar_at(1).unwrap(), Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)) ); - assert_eq!(elements.scalar_at(2), Scalar::from(200i32)); + assert_eq!(elements.scalar_at(2).unwrap(), Scalar::from(200i32)); // Check element validity. let element_validity = elements.validity(); @@ -699,11 +699,11 @@ mod tests { // Check pattern repeats correctly. for i in 0..1000 { let base = i * 5; - assert_eq!(elements.scalar_at(base), Scalar::from(1u8)); - assert_eq!(elements.scalar_at(base + 1), Scalar::from(2u8)); - assert_eq!(elements.scalar_at(base + 2), Scalar::from(3u8)); - assert_eq!(elements.scalar_at(base + 3), Scalar::from(4u8)); - assert_eq!(elements.scalar_at(base + 4), Scalar::from(5u8)); + assert_eq!(elements.scalar_at(base).unwrap(), Scalar::from(1u8)); + assert_eq!(elements.scalar_at(base + 1).unwrap(), Scalar::from(2u8)); + assert_eq!(elements.scalar_at(base + 2).unwrap(), Scalar::from(3u8)); + assert_eq!(elements.scalar_at(base + 3).unwrap(), Scalar::from(4u8)); + assert_eq!(elements.scalar_at(base + 4).unwrap(), Scalar::from(5u8)); } } } diff --git a/vortex-array/src/arrays/constant/vtable/operations.rs b/vortex-array/src/arrays/constant/vtable/operations.rs index fe632e50b9e..ccb5048bd53 100644 --- a/vortex-array/src/arrays/constant/vtable/operations.rs +++ b/vortex-array/src/arrays/constant/vtable/operations.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_error::VortexResult; use vortex_scalar::Scalar; use crate::arrays::ConstantArray; @@ -8,7 +9,7 @@ use crate::arrays::ConstantVTable; use crate::vtable::OperationsVTable; impl OperationsVTable for ConstantVTable { - fn scalar_at(array: &ConstantArray, _index: usize) -> Scalar { - array.scalar.clone() + fn scalar_at(array: &ConstantArray, _index: usize) -> VortexResult { + Ok(array.scalar.clone()) } } diff --git a/vortex-array/src/arrays/constant/vtable/validity.rs b/vortex-array/src/arrays/constant/vtable/validity.rs index 67d50b1f58f..c89ac0054da 100644 --- a/vortex-array/src/arrays/constant/vtable/validity.rs +++ b/vortex-array/src/arrays/constant/vtable/validity.rs @@ -18,10 +18,10 @@ impl ValidityVTable for ConstantVTable { }) } - fn validity_mask(array: &ConstantArray) -> Mask { - match array.scalar().is_null() { + fn validity_mask(array: &ConstantArray) -> VortexResult { + Ok(match array.scalar().is_null() { true => Mask::AllFalse(array.len), false => Mask::AllTrue(array.len), - } + }) } } diff --git a/vortex-array/src/arrays/datetime/test.rs b/vortex-array/src/arrays/datetime/test.rs index fbbe387821f..214dd2050c9 100644 --- a/vortex-array/src/arrays/datetime/test.rs +++ b/vortex-array/src/arrays/datetime/test.rs @@ -1,10 +1,18 @@ // SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors +// SPDX-FileCopyrightText: Copyright the Vortex contributorsuse vortex_dtype::Nullability; +use std::sync::Arc; use rstest::rstest; use vortex_buffer::buffer; +use vortex_dtype::DType; +use vortex_dtype::ExtDType; +use vortex_dtype::Nullability; +use vortex_dtype::PType; +use vortex_dtype::datetime::TIMESTAMP_ID; use vortex_dtype::datetime::TemporalMetadata; use vortex_dtype::datetime::TimeUnit; +use vortex_error::VortexResult; +use vortex_scalar::Scalar; use crate::IntoArray; use crate::ToCanonical; @@ -12,6 +20,8 @@ use crate::array::Array; use crate::arrays::PrimitiveArray; use crate::arrays::TemporalArray; use crate::assert_arrays_eq; +use crate::expr::root; +use crate::expr::*; use crate::validity::Validity; use crate::vtable::ValidityHelper; @@ -195,3 +205,32 @@ fn test_validity_preservation(#[case] validity: Validity) { &validity ); } + +#[test] +fn test222() -> VortexResult<()> { + // Write file with MILLISECONDS timestamps + let ts_array = PrimitiveArray::from_iter(vec![1704067200000i64, 1704153600000, 1704240000000]) + .into_array(); + let temporal = TemporalArray::new_timestamp(ts_array, TimeUnit::Milliseconds, None); + + // Read with SECONDS filter scalar + let seconds_ext_dtype = Arc::new(ExtDType::new( + TIMESTAMP_ID.clone(), + Arc::new(DType::Primitive(PType::I64, Nullability::Nullable)), + Some(TemporalMetadata::Timestamp(TimeUnit::Seconds, None).into()), + )); + let filter_expr = gt( + root(), + lit(Scalar::extension( + seconds_ext_dtype, + Scalar::from(1704153600i64), + )), + ); + + let _result = temporal.as_ref().apply(&filter_expr); + + // let err = result.is_err().unwrap(); + // println!("Expected error: {}", err); + + Ok(()) +} diff --git a/vortex-array/src/arrays/decimal/compute/cast.rs b/vortex-array/src/arrays/decimal/compute/cast.rs index 6dbc90f9011..8687b3c0655 100644 --- a/vortex-array/src/arrays/decimal/compute/cast.rs +++ b/vortex-array/src/arrays/decimal/compute/cast.rs @@ -359,7 +359,7 @@ mod tests { assert_eq!(casted.len(), 3); // Check validity is preserved - let mask = casted.validity_mask(); + let mask = casted.validity_mask().unwrap(); assert!(mask.value(0)); assert!(!mask.value(1)); assert!(mask.value(2)); diff --git a/vortex-array/src/arrays/decimal/compute/fill_null.rs b/vortex-array/src/arrays/decimal/compute/fill_null.rs index 8d8016886b6..ab7ae6284d2 100644 --- a/vortex-array/src/arrays/decimal/compute/fill_null.rs +++ b/vortex-array/src/arrays/decimal/compute/fill_null.rs @@ -88,7 +88,7 @@ mod tests { p.buffer::().as_slice(), vec![4200, 800, 4200, 1000, 4200] ); - assert!(p.validity_mask().all_true()); + assert!(p.validity_mask().unwrap().all_true()); } #[test] diff --git a/vortex-array/src/arrays/decimal/compute/is_sorted.rs b/vortex-array/src/arrays/decimal/compute/is_sorted.rs index 306663f5b96..3c224138dc2 100644 --- a/vortex-array/src/arrays/decimal/compute/is_sorted.rs +++ b/vortex-array/src/arrays/decimal/compute/is_sorted.rs @@ -36,7 +36,7 @@ fn compute_is_sorted(array: &DecimalArray, strict: bool) - where dyn Iterator: IsSortedIteratorExt, { - match array.validity_mask() { + match array.validity_mask()? { Mask::AllFalse(_) => Ok(!strict), Mask::AllTrue(_) => { let buf = array.buffer::(); diff --git a/vortex-array/src/arrays/decimal/compute/min_max.rs b/vortex-array/src/arrays/decimal/compute/min_max.rs index 6c4c6bf714f..da4483a3a69 100644 --- a/vortex-array/src/arrays/decimal/compute/min_max.rs +++ b/vortex-array/src/arrays/decimal/compute/min_max.rs @@ -33,7 +33,7 @@ fn compute_min_max_with_validity(array: &DecimalArray) -> VortexResult

(); BitBuffer::collect_bool(array.len(), |idx| { @@ -73,7 +73,7 @@ impl ValidityVTable for DictVTable { AllOr::None => Mask::AllFalse(array.len()), AllOr::Some(validity_buff) => { let primitive_codes = array.codes().to_primitive(); - let values_mask = array.values().validity_mask(); + let values_mask = array.values().validity_mask()?; let is_valid_buffer = match_each_integer_ptype!(primitive_codes.ptype(), |P| { let codes_slice = primitive_codes.as_slice::

(); #[allow(clippy::cast_possible_truncation)] @@ -83,6 +83,6 @@ impl ValidityVTable for DictVTable { }); Mask::from_buffer(is_valid_buffer) } - } + }) } } diff --git a/vortex-array/src/arrays/extension/array.rs b/vortex-array/src/arrays/extension/array.rs index 37ee215c777..9661854f329 100644 --- a/vortex-array/src/arrays/extension/array.rs +++ b/vortex-array/src/arrays/extension/array.rs @@ -83,7 +83,7 @@ use crate::stats::ArrayStats; /// assert_eq!(currency_array.id().as_ref(), "example.currency"); /// /// // Access maintains extension type information -/// let first_value = currency_array.scalar_at(0); +/// let first_value = currency_array.scalar_at(0).unwrap(); /// assert!(first_value.as_extension_opt().is_some()); /// ``` #[derive(Clone, Debug)] diff --git a/vortex-array/src/arrays/extension/vtable/operations.rs b/vortex-array/src/arrays/extension/vtable/operations.rs index ab89ef99bb3..eefc8d07d6f 100644 --- a/vortex-array/src/arrays/extension/vtable/operations.rs +++ b/vortex-array/src/arrays/extension/vtable/operations.rs @@ -1,14 +1,19 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_error::VortexResult; use vortex_scalar::Scalar; +use crate::Array; use crate::arrays::extension::ExtensionArray; use crate::arrays::extension::ExtensionVTable; use crate::vtable::OperationsVTable; impl OperationsVTable for ExtensionVTable { - fn scalar_at(array: &ExtensionArray, index: usize) -> Scalar { - Scalar::extension(array.ext_dtype().clone(), array.storage().scalar_at(index)) + fn scalar_at(array: &ExtensionArray, index: usize) -> VortexResult { + Ok(Scalar::extension( + array.ext_dtype().clone(), + array.storage().scalar_at(index)?, + )) } } diff --git a/vortex-array/src/arrays/filter/vtable.rs b/vortex-array/src/arrays/filter/vtable.rs index 34d14e4e0f3..c16a7b67a98 100644 --- a/vortex-array/src/arrays/filter/vtable.rs +++ b/vortex-array/src/arrays/filter/vtable.rs @@ -155,7 +155,7 @@ pub(super) fn execute_fast_path( } // All null - child has no valid values - if array.validity_mask().true_count() == 0 { + if array.validity_mask()?.true_count() == 0 { return Ok(Some( ConstantArray::new(Scalar::null(array.dtype().clone()), true_count) .into_array() @@ -190,7 +190,7 @@ impl BaseArrayVTable for FilterVTable { } impl OperationsVTable for FilterVTable { - fn scalar_at(array: &FilterArray, index: usize) -> Scalar { + fn scalar_at(array: &FilterArray, index: usize) -> VortexResult { let rank_idx = array.mask.rank(index); array.child.scalar_at(rank_idx) } @@ -201,8 +201,8 @@ impl ValidityVTable for FilterVTable { array.child.validity()?.filter(&array.mask) } - fn validity_mask(array: &FilterArray) -> Mask { - Filter::filter(&array.child.validity_mask(), &array.mask) + fn validity_mask(array: &FilterArray) -> VortexResult { + Ok(Filter::filter(&array.child.validity_mask()?, &array.mask)) } } diff --git a/vortex-array/src/arrays/fixed_size_list/compute/is_constant.rs b/vortex-array/src/arrays/fixed_size_list/compute/is_constant.rs index 8f9f18207d1..1786fe7a7c4 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/is_constant.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/is_constant.rs @@ -26,12 +26,12 @@ impl IsConstantKernel for FixedSizeListVTable { array.len() > 1, "precondition for `is_constant` is incorrect" ); - let first_scalar = array.scalar_at(0); // We checked the array length above. + let first_scalar = array.scalar_at(0)?; // We checked the array length above. // TODO(connor): There must be a more efficient way to do this. Each `scalar_at()` call // makes several allocations... for i in 1..array.len() { - let current_scalar = array.scalar_at(i); + let current_scalar = array.scalar_at(i)?; if current_scalar != first_scalar { return Ok(Some(false)); } diff --git a/vortex-array/src/arrays/fixed_size_list/compute/take.rs b/vortex-array/src/arrays/fixed_size_list/compute/take.rs index 02a28e6da15..68432e2d798 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/take.rs @@ -137,8 +137,8 @@ fn take_nullable_fsl( let indices: &[I] = indices_array.as_slice::(); let new_len = indices.len(); - let array_validity = array.validity_mask(); - let indices_validity = indices_array.validity_mask(); + let array_validity = array.validity_mask()?; + let indices_validity = indices_array.validity_mask()?; // We must use placeholder zeros for null lists to maintain the array length without // propagating nullability to the element array's take operation. diff --git a/vortex-array/src/arrays/fixed_size_list/tests/basic.rs b/vortex-array/src/arrays/fixed_size_list/tests/basic.rs index 6fe97c2ea12..e25fe1f0b98 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/basic.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/basic.rs @@ -36,24 +36,24 @@ fn test_basic_fixed_size_list() { // Check the actual values in each list. let first_list = fsl.fixed_size_list_elements_at(0); - assert_eq!(first_list.scalar_at(0), 1i32.into()); - assert_eq!(first_list.scalar_at(1), 2i32.into()); - assert_eq!(first_list.scalar_at(2), 3i32.into()); + assert_eq!(first_list.scalar_at(0).unwrap(), 1i32.into()); + assert_eq!(first_list.scalar_at(1).unwrap(), 2i32.into()); + assert_eq!(first_list.scalar_at(2).unwrap(), 3i32.into()); let second_list = fsl.fixed_size_list_elements_at(1); - assert_eq!(second_list.scalar_at(0), 4i32.into()); - assert_eq!(second_list.scalar_at(1), 5i32.into()); - assert_eq!(second_list.scalar_at(2), 6i32.into()); + assert_eq!(second_list.scalar_at(0).unwrap(), 4i32.into()); + assert_eq!(second_list.scalar_at(1).unwrap(), 5i32.into()); + assert_eq!(second_list.scalar_at(2).unwrap(), 6i32.into()); let third_list = fsl.fixed_size_list_elements_at(2); - assert_eq!(third_list.scalar_at(0), 7i32.into()); - assert_eq!(third_list.scalar_at(1), 8i32.into()); - assert_eq!(third_list.scalar_at(2), 9i32.into()); + assert_eq!(third_list.scalar_at(0).unwrap(), 7i32.into()); + assert_eq!(third_list.scalar_at(1).unwrap(), 8i32.into()); + assert_eq!(third_list.scalar_at(2).unwrap(), 9i32.into()); let fourth_list = fsl.fixed_size_list_elements_at(3); - assert_eq!(fourth_list.scalar_at(0), 10i32.into()); - assert_eq!(fourth_list.scalar_at(1), 11i32.into()); - assert_eq!(fourth_list.scalar_at(2), 12i32.into()); + assert_eq!(fourth_list.scalar_at(0).unwrap(), 10i32.into()); + assert_eq!(fourth_list.scalar_at(1).unwrap(), 11i32.into()); + assert_eq!(fourth_list.scalar_at(2).unwrap(), 12i32.into()); } #[test] @@ -65,7 +65,7 @@ fn test_scalar_at() { let fsl = FixedSizeListArray::new(elements.into_array(), list_size, Validity::NonNullable, len); // First list: [1, 2, 3]. - let first = fsl.scalar_at(0); + let first = fsl.scalar_at(0).unwrap(); assert_eq!( first, Scalar::fixed_size_list( @@ -77,12 +77,12 @@ fn test_scalar_at() { // Additionally check individual elements via fixed_size_list_at. let first_list = fsl.fixed_size_list_elements_at(0); - assert_eq!(first_list.scalar_at(0), 1i32.into()); - assert_eq!(first_list.scalar_at(1), 2i32.into()); - assert_eq!(first_list.scalar_at(2), 3i32.into()); + assert_eq!(first_list.scalar_at(0).unwrap(), 1i32.into()); + assert_eq!(first_list.scalar_at(1).unwrap(), 2i32.into()); + assert_eq!(first_list.scalar_at(2).unwrap(), 3i32.into()); // Second list: [4, 5, 6]. - let second = fsl.scalar_at(1); + let second = fsl.scalar_at(1).unwrap(); assert_eq!( second, Scalar::fixed_size_list( @@ -94,9 +94,9 @@ fn test_scalar_at() { // Additionally check individual elements via fixed_size_list_at. let second_list = fsl.fixed_size_list_elements_at(1); - assert_eq!(second_list.scalar_at(0), 4i32.into()); - assert_eq!(second_list.scalar_at(1), 5i32.into()); - assert_eq!(second_list.scalar_at(2), 6i32.into()); + assert_eq!(second_list.scalar_at(0).unwrap(), 4i32.into()); + assert_eq!(second_list.scalar_at(1).unwrap(), 5i32.into()); + assert_eq!(second_list.scalar_at(2).unwrap(), 6i32.into()); } #[test] @@ -110,14 +110,14 @@ fn test_fixed_size_list_at() { // Get the first list [1.0, 2.0]. let first_list = fsl.fixed_size_list_elements_at(0); assert_eq!(first_list.len(), list_size as usize); - assert_eq!(first_list.scalar_at(0), 1.0f64.into()); - assert_eq!(first_list.scalar_at(1), 2.0f64.into()); + assert_eq!(first_list.scalar_at(0).unwrap(), 1.0f64.into()); + assert_eq!(first_list.scalar_at(1).unwrap(), 2.0f64.into()); // Get the third list [5.0, 6.0]. let third_list = fsl.fixed_size_list_elements_at(2); assert_eq!(third_list.len(), list_size as usize); - assert_eq!(third_list.scalar_at(0), 5.0f64.into()); - assert_eq!(third_list.scalar_at(1), 6.0f64.into()); + assert_eq!(third_list.scalar_at(0).unwrap(), 5.0f64.into()); + assert_eq!(third_list.scalar_at(1).unwrap(), 6.0f64.into()); } #[test] diff --git a/vortex-array/src/arrays/fixed_size_list/tests/degenerate.rs b/vortex-array/src/arrays/fixed_size_list/tests/degenerate.rs index 03453d657d4..d6dd5d6d245 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/degenerate.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/degenerate.rs @@ -54,7 +54,7 @@ fn test_fsl_size_0_length_1_non_nullable() { assert_eq!(fsl.elements().len(), 0); // Get the single empty list. - let scalar = fsl.scalar_at(0); + let scalar = fsl.scalar_at(0).unwrap(); assert!(!scalar.is_null()); assert_eq!( scalar, @@ -80,7 +80,7 @@ fn test_fsl_size_0_huge_length_non_nullable() { assert_eq!(fsl.elements().len(), 0); // Spot check a few lists. - let scalar_first = fsl.scalar_at(0); + let scalar_first = fsl.scalar_at(0).unwrap(); assert!(!scalar_first.is_null()); assert_eq!( scalar_first, @@ -91,7 +91,7 @@ fn test_fsl_size_0_huge_length_non_nullable() { ) ); - let scalar_middle = fsl.scalar_at(500_000_000_000); + let scalar_middle = fsl.scalar_at(500_000_000_000).unwrap(); assert!(!scalar_middle.is_null()); assert_eq!( scalar_middle, @@ -102,7 +102,7 @@ fn test_fsl_size_0_huge_length_non_nullable() { ) ); - let scalar_end = fsl.scalar_at(999_999_999_999); + let scalar_end = fsl.scalar_at(999_999_999_999).unwrap(); assert!(!scalar_end.is_null()); assert_eq!( scalar_end, @@ -152,7 +152,7 @@ fn test_fsl_size_0_length_1_nullable_valid() { assert_eq!(fsl.elements().len(), 0); // Get the single empty list (should be valid). - let scalar = fsl.scalar_at(0); + let scalar = fsl.scalar_at(0).unwrap(); assert!(!scalar.is_null()); assert_eq!( scalar, @@ -175,7 +175,7 @@ fn test_fsl_size_0_length_1_nullable_null() { assert_eq!(fsl.elements().len(), 0); // The single list should be null. - let scalar = fsl.scalar_at(0); + let scalar = fsl.scalar_at(0).unwrap(); assert!(scalar.is_null()); } @@ -200,7 +200,7 @@ fn test_fsl_size_0_length_10_nullable_mixed() { true, false, true, true, false, false, true, false, true, true, ]; for i in 0..len { - let scalar = fsl.scalar_at(i); + let scalar = fsl.scalar_at(i).unwrap(); if expected_valid[i] { assert!(!scalar.is_null()); assert_eq!( @@ -239,7 +239,7 @@ fn test_fsl_size_0_nullable_elements() { // All lists should be empty but valid. for i in 0..len { - let scalar = fsl.scalar_at(i); + let scalar = fsl.scalar_at(i).unwrap(); assert!(!scalar.is_null()); } } diff --git a/vortex-array/src/arrays/fixed_size_list/tests/filter.rs b/vortex-array/src/arrays/fixed_size_list/tests/filter.rs index 2fd321b1cbe..107f4e389ca 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/filter.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/filter.rs @@ -78,7 +78,7 @@ fn test_filter_degenerate_list_size_zero( // Should return a ConstantArray of nulls when all are invalid. let filtered_const = filtered.as_::(); for i in 0..expected_len { - assert!(filtered_const.scalar_at(i).is_null()); + assert!(filtered_const.scalar_at(i).unwrap().is_null()); } } else { let filtered_fsl = filtered.as_::(); @@ -111,13 +111,13 @@ fn test_filter_with_nulls() { // First list should be [1, 2] and valid. let first = filtered_fsl.fixed_size_list_elements_at(0); - assert_eq!(first.scalar_at(0), 1i32.into()); - assert_eq!(first.scalar_at(1), 2i32.into()); + assert_eq!(first.scalar_at(0).unwrap(), 1i32.into()); + assert_eq!(first.scalar_at(1).unwrap(), 2i32.into()); // Second list should be [5, 6] and valid. let second = filtered_fsl.fixed_size_list_elements_at(1); - assert_eq!(second.scalar_at(0), 5i32.into()); - assert_eq!(second.scalar_at(1), 6i32.into()); + assert_eq!(second.scalar_at(0).unwrap(), 5i32.into()); + assert_eq!(second.scalar_at(1).unwrap(), 6i32.into()); } #[test] @@ -138,11 +138,11 @@ fn test_filter_all_null_array() { "All-null FSL should produce ConstantArray" ); assert!( - filtered_const.scalar_at(0).is_null(), + filtered_const.scalar_at(0).unwrap().is_null(), "Expected null at index 0" ); assert!( - filtered_const.scalar_at(1).is_null(), + filtered_const.scalar_at(1).unwrap().is_null(), "Expected null at index 1" ); } @@ -191,16 +191,16 @@ fn test_filter_nested_fixed_size_lists() { // Check the actual values. let inner_list_0 = filtered_inner.fixed_size_list_elements_at(0); - assert_eq!(inner_list_0.scalar_at(0), 7i32.into()); - assert_eq!(inner_list_0.scalar_at(1), 8i32.into()); + assert_eq!(inner_list_0.scalar_at(0).unwrap(), 7i32.into()); + assert_eq!(inner_list_0.scalar_at(1).unwrap(), 8i32.into()); let inner_list_1 = filtered_inner.fixed_size_list_elements_at(1); - assert_eq!(inner_list_1.scalar_at(0), 9i32.into()); - assert_eq!(inner_list_1.scalar_at(1), 10i32.into()); + assert_eq!(inner_list_1.scalar_at(0).unwrap(), 9i32.into()); + assert_eq!(inner_list_1.scalar_at(1).unwrap(), 10i32.into()); let inner_list_2 = filtered_inner.fixed_size_list_elements_at(2); - assert_eq!(inner_list_2.scalar_at(0), 11i32.into()); - assert_eq!(inner_list_2.scalar_at(1), 12i32.into()); + assert_eq!(inner_list_2.scalar_at(0).unwrap(), 11i32.into()); + assert_eq!(inner_list_2.scalar_at(1).unwrap(), 12i32.into()); } // Conformance tests using rstest for various array configurations. @@ -297,8 +297,8 @@ fn test_filter_all_null_various_list_sizes() { let filtered0 = filter(fsl0.as_ref(), &mask0).unwrap(); assert_eq!(filtered0.len(), 2); // Check that all elements are null (might be ConstantArray or FixedSizeListArray) - assert!(filtered0.scalar_at(0).is_null()); - assert!(filtered0.scalar_at(1).is_null()); + assert!(filtered0.scalar_at(0).unwrap().is_null()); + assert!(filtered0.scalar_at(1).unwrap().is_null()); // Case 2: list_size == 1 let elements1 = buffer![1i32, 2, 3].into_array(); @@ -307,8 +307,8 @@ fn test_filter_all_null_various_list_sizes() { let filtered1 = filter(fsl1.as_ref(), &mask1).unwrap(); assert_eq!(filtered1.len(), 2); // Check that all elements are null - assert!(filtered1.scalar_at(0).is_null()); - assert!(filtered1.scalar_at(1).is_null()); + assert!(filtered1.scalar_at(0).unwrap().is_null()); + assert!(filtered1.scalar_at(1).unwrap().is_null()); // Case 3: list_size == 10 (large) let elements10 = buffer![0..50i32].into_array(); @@ -317,8 +317,8 @@ fn test_filter_all_null_various_list_sizes() { let filtered10 = filter(fsl10.as_ref(), &mask10).unwrap(); assert_eq!(filtered10.len(), 5); // Check that all elements are null - assert!(filtered10.scalar_at(0).is_null()); - assert!(filtered10.scalar_at(4).is_null()); + assert!(filtered10.scalar_at(0).unwrap().is_null()); + assert!(filtered10.scalar_at(4).unwrap().is_null()); } // Note: test_filter_to_empty_degenerate has been consolidated into test_filter_degenerate_list_size_zero above. @@ -355,13 +355,22 @@ fn test_mask_expansion_threshold_boundary() { // Verify correct elements were kept. let first = filtered_fsl.fixed_size_list_elements_at(0); - assert_eq!(first.scalar_at(0), (5i32 * list_size as i32).into()); + assert_eq!( + first.scalar_at(0).unwrap(), + (5i32 * list_size as i32).into() + ); let second = filtered_fsl.fixed_size_list_elements_at(1); - assert_eq!(second.scalar_at(0), (25i32 * list_size as i32).into()); + assert_eq!( + second.scalar_at(0).unwrap(), + (25i32 * list_size as i32).into() + ); let third = filtered_fsl.fixed_size_list_elements_at(2); - assert_eq!(third.scalar_at(0), (75i32 * list_size as i32).into()); + assert_eq!( + third.scalar_at(0).unwrap(), + (75i32 * list_size as i32).into() + ); // Test with list_size == 7 (just below threshold). let list_size_7 = 7u32; @@ -411,16 +420,16 @@ fn test_filter_large_list_size() { // Check that the correct lists were kept (indices 1, 3, 4 from original). let list_0 = filtered_fsl.fixed_size_list_elements_at(0); - assert_eq!(list_0.scalar_at(0), 100i64.into()); // Start of original list 1. - assert_eq!(list_0.scalar_at(99), 199i64.into()); // End of original list 1. + assert_eq!(list_0.scalar_at(0).unwrap(), 100i64.into()); // Start of original list 1. + assert_eq!(list_0.scalar_at(99).unwrap(), 199i64.into()); // End of original list 1. let list_1 = filtered_fsl.fixed_size_list_elements_at(1); - assert_eq!(list_1.scalar_at(0), 300i64.into()); // Start of original list 3. - assert_eq!(list_1.scalar_at(99), 399i64.into()); // End of original list 3. + assert_eq!(list_1.scalar_at(0).unwrap(), 300i64.into()); // Start of original list 3. + assert_eq!(list_1.scalar_at(99).unwrap(), 399i64.into()); // End of original list 3. let list_2 = filtered_fsl.fixed_size_list_elements_at(2); - assert_eq!(list_2.scalar_at(0), 400i64.into()); // Start of original list 4. - assert_eq!(list_2.scalar_at(99), 499i64.into()); // End of original list 4. + assert_eq!(list_2.scalar_at(0).unwrap(), 400i64.into()); // Start of original list 4. + assert_eq!(list_2.scalar_at(99).unwrap(), 499i64.into()); // End of original list 4. // Test edge case: filter out all but one large list. let mask_single = Mask::from_iter([false, false, true, false, false]); @@ -433,7 +442,7 @@ fn test_filter_large_list_size() { // Verify it's the correct list (original list 2). let single_list = filtered_single_fsl.fixed_size_list_elements_at(0); - assert_eq!(single_list.scalar_at(0), 200i64.into()); - assert_eq!(single_list.scalar_at(50), 250i64.into()); - assert_eq!(single_list.scalar_at(99), 299i64.into()); + assert_eq!(single_list.scalar_at(0).unwrap(), 200i64.into()); + assert_eq!(single_list.scalar_at(50).unwrap(), 250i64.into()); + assert_eq!(single_list.scalar_at(99).unwrap(), 299i64.into()); } diff --git a/vortex-array/src/arrays/fixed_size_list/tests/nested.rs b/vortex-array/src/arrays/fixed_size_list/tests/nested.rs index 9a2180d10c2..913f963675b 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/nested.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/nested.rs @@ -72,7 +72,7 @@ fn test_fsl_of_fsl_basic() { // The first outer list should contain 3 inner lists. // We can check by slicing and examining scalars. - let first_scalar = outer_fsl.scalar_at(0); + let first_scalar = outer_fsl.scalar_at(0).unwrap(); assert!(!first_scalar.is_null()); // Check the actual values in the nested structure. @@ -83,22 +83,22 @@ fn test_fsl_of_fsl_basic() { let inner_list_0 = first_outer_list .to_fixed_size_list() .fixed_size_list_elements_at(0); - assert_eq!(inner_list_0.scalar_at(0), 1i32.into()); - assert_eq!(inner_list_0.scalar_at(1), 2i32.into()); + assert_eq!(inner_list_0.scalar_at(0).unwrap(), 1i32.into()); + assert_eq!(inner_list_0.scalar_at(1).unwrap(), 2i32.into()); // Check second inner list [3,4]. let inner_list_1 = first_outer_list .to_fixed_size_list() .fixed_size_list_elements_at(1); - assert_eq!(inner_list_1.scalar_at(0), 3i32.into()); - assert_eq!(inner_list_1.scalar_at(1), 4i32.into()); + assert_eq!(inner_list_1.scalar_at(0).unwrap(), 3i32.into()); + assert_eq!(inner_list_1.scalar_at(1).unwrap(), 4i32.into()); // Check third inner list [5,6]. let inner_list_2 = first_outer_list .to_fixed_size_list() .fixed_size_list_elements_at(2); - assert_eq!(inner_list_2.scalar_at(0), 5i32.into()); - assert_eq!(inner_list_2.scalar_at(1), 6i32.into()); + assert_eq!(inner_list_2.scalar_at(0).unwrap(), 5i32.into()); + assert_eq!(inner_list_2.scalar_at(1).unwrap(), 6i32.into()); // Second outer list contains: [[7,8], [9,10], [11,12]]. let second_outer_list = outer_fsl.fixed_size_list_elements_at(1); @@ -107,22 +107,22 @@ fn test_fsl_of_fsl_basic() { let inner_list_0 = second_outer_list .to_fixed_size_list() .fixed_size_list_elements_at(0); - assert_eq!(inner_list_0.scalar_at(0), 7i32.into()); - assert_eq!(inner_list_0.scalar_at(1), 8i32.into()); + assert_eq!(inner_list_0.scalar_at(0).unwrap(), 7i32.into()); + assert_eq!(inner_list_0.scalar_at(1).unwrap(), 8i32.into()); // Check second inner list [9,10]. let inner_list_1 = second_outer_list .to_fixed_size_list() .fixed_size_list_elements_at(1); - assert_eq!(inner_list_1.scalar_at(0), 9i32.into()); - assert_eq!(inner_list_1.scalar_at(1), 10i32.into()); + assert_eq!(inner_list_1.scalar_at(0).unwrap(), 9i32.into()); + assert_eq!(inner_list_1.scalar_at(1).unwrap(), 10i32.into()); // Check third inner list [11,12]. let inner_list_2 = second_outer_list .to_fixed_size_list() .fixed_size_list_elements_at(2); - assert_eq!(inner_list_2.scalar_at(0), 11i32.into()); - assert_eq!(inner_list_2.scalar_at(1), 12i32.into()); + assert_eq!(inner_list_2.scalar_at(0).unwrap(), 11i32.into()); + assert_eq!(inner_list_2.scalar_at(1).unwrap(), 12i32.into()); } #[test] @@ -167,13 +167,13 @@ fn test_fsl_of_fsl_with_nulls() { assert_eq!(outer_fsl.len(), outer_len); // First outer list is valid. - assert!(!outer_fsl.scalar_at(0).is_null()); + assert!(!outer_fsl.scalar_at(0).unwrap().is_null()); // Second outer list is null. - assert!(outer_fsl.scalar_at(1).is_null()); + assert!(outer_fsl.scalar_at(1).unwrap().is_null()); // Third outer list is valid. - assert!(!outer_fsl.scalar_at(2).is_null()); + assert!(!outer_fsl.scalar_at(2).unwrap().is_null()); } #[test] @@ -225,21 +225,21 @@ fn test_deeply_nested_fsl() { // First level-2 list: [[1,2],[3,4]]. let level1_0_0 = level2_0.to_fixed_size_list().fixed_size_list_elements_at(0); - assert_eq!(level1_0_0.scalar_at(0), 1i32.into()); - assert_eq!(level1_0_0.scalar_at(1), 2i32.into()); + assert_eq!(level1_0_0.scalar_at(0).unwrap(), 1i32.into()); + assert_eq!(level1_0_0.scalar_at(1).unwrap(), 2i32.into()); let level1_0_1 = level2_0.to_fixed_size_list().fixed_size_list_elements_at(1); - assert_eq!(level1_0_1.scalar_at(0), 3i32.into()); - assert_eq!(level1_0_1.scalar_at(1), 4i32.into()); + assert_eq!(level1_0_1.scalar_at(0).unwrap(), 3i32.into()); + assert_eq!(level1_0_1.scalar_at(1).unwrap(), 4i32.into()); // Second level-2 list: [[5,6],[7,8]]. let level1_1_0 = level2_1.to_fixed_size_list().fixed_size_list_elements_at(0); - assert_eq!(level1_1_0.scalar_at(0), 5i32.into()); - assert_eq!(level1_1_0.scalar_at(1), 6i32.into()); + assert_eq!(level1_1_0.scalar_at(0).unwrap(), 5i32.into()); + assert_eq!(level1_1_0.scalar_at(1).unwrap(), 6i32.into()); let level1_1_1 = level2_1.to_fixed_size_list().fixed_size_list_elements_at(1); - assert_eq!(level1_1_1.scalar_at(0), 7i32.into()); - assert_eq!(level1_1_1.scalar_at(1), 8i32.into()); + assert_eq!(level1_1_1.scalar_at(0).unwrap(), 7i32.into()); + assert_eq!(level1_1_1.scalar_at(1).unwrap(), 8i32.into()); } //////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/vortex-array/src/arrays/fixed_size_list/tests/nullability.rs b/vortex-array/src/arrays/fixed_size_list/tests/nullability.rs index 984bb1d8755..0012d5942dc 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/nullability.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/nullability.rs @@ -29,7 +29,7 @@ fn test_nullable_fsl_with_nulls() { assert_eq!(fsl.list_size(), list_size); // First list is valid: [1, 2]. - let first = fsl.scalar_at(0); + let first = fsl.scalar_at(0).unwrap(); assert!(!first.is_null()); assert_eq!( first, @@ -42,15 +42,15 @@ fn test_nullable_fsl_with_nulls() { // Check individual elements of the first list. let first_list = fsl.fixed_size_list_elements_at(0); - assert_eq!(first_list.scalar_at(0), 1i32.into()); - assert_eq!(first_list.scalar_at(1), 2i32.into()); + assert_eq!(first_list.scalar_at(0).unwrap(), 1i32.into()); + assert_eq!(first_list.scalar_at(1).unwrap(), 2i32.into()); // Second list is null. - let second = fsl.scalar_at(1); + let second = fsl.scalar_at(1).unwrap(); assert!(second.is_null()); // Third list is valid: [5, 6]. - let third = fsl.scalar_at(2); + let third = fsl.scalar_at(2).unwrap(); assert!(!third.is_null()); assert_eq!( third, @@ -63,11 +63,11 @@ fn test_nullable_fsl_with_nulls() { // Check individual elements of the third list. let third_list = fsl.fixed_size_list_elements_at(2); - assert_eq!(third_list.scalar_at(0), 5i32.into()); - assert_eq!(third_list.scalar_at(1), 6i32.into()); + assert_eq!(third_list.scalar_at(0).unwrap(), 5i32.into()); + assert_eq!(third_list.scalar_at(1).unwrap(), 6i32.into()); // Fourth list is null. - let fourth = fsl.scalar_at(3); + let fourth = fsl.scalar_at(3).unwrap(); assert!(fourth.is_null()); } @@ -91,7 +91,7 @@ fn test_nullable_elements_non_nullable_lists() { )); // First list: [Some(1), None, Some(3)]. - let first = fsl.scalar_at(0); + let first = fsl.scalar_at(0).unwrap(); assert!(!first.is_null()); assert_eq!( first, @@ -103,7 +103,7 @@ fn test_nullable_elements_non_nullable_lists() { ); // Second list: [Some(4), Some(5), None]. - let second = fsl.scalar_at(1); + let second = fsl.scalar_at(1).unwrap(); assert!(!second.is_null()); assert_eq!( second, @@ -129,7 +129,7 @@ fn test_nullable_elements_and_nullable_lists() { assert_eq!(fsl.len(), len); // First list is valid: [Some(10), None]. - let first = fsl.scalar_at(0); + let first = fsl.scalar_at(0).unwrap(); assert!(!first.is_null()); assert_eq!( first, @@ -142,15 +142,15 @@ fn test_nullable_elements_and_nullable_lists() { // Check individual elements of the first list. let first_list = fsl.fixed_size_list_elements_at(0); - assert_eq!(first_list.scalar_at(0), Some(10u16).into()); - assert_eq!(first_list.scalar_at(1), None::.into()); + assert_eq!(first_list.scalar_at(0).unwrap(), Some(10u16).into()); + assert_eq!(first_list.scalar_at(1).unwrap(), None::.into()); // Second list is null (but elements would be [Some(20), Some(30)]). - let second = fsl.scalar_at(1); + let second = fsl.scalar_at(1).unwrap(); assert!(second.is_null()); // Third list is valid: [None, None]. - let third = fsl.scalar_at(2); + let third = fsl.scalar_at(2).unwrap(); assert!(!third.is_null()); assert_eq!( third, @@ -163,8 +163,8 @@ fn test_nullable_elements_and_nullable_lists() { // Check individual elements of the third list. let third_list = fsl.fixed_size_list_elements_at(2); - assert_eq!(third_list.scalar_at(0), None::.into()); - assert_eq!(third_list.scalar_at(1), None::.into()); + assert_eq!(third_list.scalar_at(0).unwrap(), None::.into()); + assert_eq!(third_list.scalar_at(1).unwrap(), None::.into()); } #[test] @@ -181,7 +181,7 @@ fn test_alternating_nulls() { // Check alternating pattern. for i in 0..len { - let scalar = fsl.scalar_at(i); + let scalar = fsl.scalar_at(i).unwrap(); if i % 2 == 0 { assert!(!scalar.is_null()); let expected_value = u8::try_from(i + 1).unwrap(); @@ -211,7 +211,7 @@ fn test_validity_types() { { let fsl = FixedSizeListArray::new(elements.clone(), list_size, Validity::AllInvalid, len); for i in 0..len { - assert!(fsl.scalar_at(i).is_null()); + assert!(fsl.scalar_at(i).unwrap().is_null()); } } @@ -225,10 +225,10 @@ fn test_validity_types() { len, ); - assert!(!fsl.scalar_at(0).is_null()); - assert!(!fsl.scalar_at(1).is_null()); - assert!(fsl.scalar_at(2).is_null()); - assert!(!fsl.scalar_at(3).is_null()); + assert!(!fsl.scalar_at(0).unwrap().is_null()); + assert!(!fsl.scalar_at(1).unwrap().is_null()); + assert!(fsl.scalar_at(2).unwrap().is_null()); + assert!(!fsl.scalar_at(3).unwrap().is_null()); } } @@ -254,22 +254,22 @@ fn test_mixed_nullability_patterns() { let fsl = FixedSizeListArray::new(elements.into_array(), list_size, validity, len); // List 0: valid with [Some(1), None]. - let list0 = fsl.scalar_at(0); + let list0 = fsl.scalar_at(0).unwrap(); assert!(!list0.is_null()); // List 1: null. - let list1 = fsl.scalar_at(1); + let list1 = fsl.scalar_at(1).unwrap(); assert!(list1.is_null()); // List 2: valid with [Some(5), Some(6)]. - let list2 = fsl.scalar_at(2); + let list2 = fsl.scalar_at(2).unwrap(); assert!(!list2.is_null()); // List 3: valid with [Some(7), None]. - let list3 = fsl.scalar_at(3); + let list3 = fsl.scalar_at(3).unwrap(); assert!(!list3.is_null()); // List 4: valid with [None, Some(10)]. - let list4 = fsl.scalar_at(4); + let list4 = fsl.scalar_at(4).unwrap(); assert!(!list4.is_null()); } diff --git a/vortex-array/src/arrays/fixed_size_list/tests/take.rs b/vortex-array/src/arrays/fixed_size_list/tests/take.rs index d305fc0270c..b9aa8c74441 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/take.rs @@ -53,25 +53,25 @@ fn test_take_basic_smoke_test() { // First list should be the original third list [5, 6]. let first = result_fsl.fixed_size_list_elements_at(0); assert_eq!( - first.scalar_at(0), + first.scalar_at(0).unwrap(), 5i32.into(), "Wrong value at [2][0] after take" ); assert_eq!( - first.scalar_at(1), + first.scalar_at(1).unwrap(), 6i32.into(), "Wrong value at [2][1] after take" ); // Second list should be the original first list [1, 2]. let second = result_fsl.fixed_size_list_elements_at(1); - assert_eq!(second.scalar_at(0), 1i32.into()); - assert_eq!(second.scalar_at(1), 2i32.into()); + assert_eq!(second.scalar_at(0).unwrap(), 1i32.into()); + assert_eq!(second.scalar_at(1).unwrap(), 2i32.into()); // Third list should be the original second list [3, 4]. let third = result_fsl.fixed_size_list_elements_at(2); - assert_eq!(third.scalar_at(0), 3i32.into()); - assert_eq!(third.scalar_at(1), 4i32.into()); + assert_eq!(third.scalar_at(0).unwrap(), 3i32.into()); + assert_eq!(third.scalar_at(1).unwrap(), 4i32.into()); } // Parameterized test for FSL-specific degenerate (list_size=0) cases. @@ -118,7 +118,7 @@ fn test_take_degenerate_lists( // Check nullability of results. for (i, expected_null) in expected_nulls.iter().enumerate() { - assert_eq!(result_fsl.scalar_at(i).is_null(), *expected_null); + assert_eq!(result_fsl.scalar_at(i).unwrap().is_null(), *expected_null); } } @@ -139,13 +139,13 @@ fn test_take_large_list_size() { // First list should be [200..300]. let first = result_fsl.fixed_size_list_elements_at(0); for i in 0..100i32 { - assert_eq!(first.scalar_at(i as usize), (200 + i).into()); + assert_eq!(first.scalar_at(i as usize).unwrap(), (200 + i).into()); } // Second list should be [0..100]. let second = result_fsl.fixed_size_list_elements_at(1); for i in 0..100i32 { - assert_eq!(second.scalar_at(i as usize), i.into()); + assert_eq!(second.scalar_at(i as usize).unwrap(), i.into()); } } @@ -164,19 +164,19 @@ fn test_take_fsl_with_null_indices_preserves_elements() { assert_eq!(result_fsl.list_size(), 2); // First list should be [3, 4]. - assert!(!result_fsl.scalar_at(0).is_null()); + assert!(!result_fsl.scalar_at(0).unwrap().is_null()); let first = result_fsl.fixed_size_list_elements_at(0); - assert_eq!(first.scalar_at(0), 3i32.into()); - assert_eq!(first.scalar_at(1), 4i32.into()); + assert_eq!(first.scalar_at(0).unwrap(), 3i32.into()); + assert_eq!(first.scalar_at(1).unwrap(), 4i32.into()); // Second list should be null. - assert!(result_fsl.scalar_at(1).is_null()); + assert!(result_fsl.scalar_at(1).unwrap().is_null()); // Third list should be [1, 2]. - assert!(!result_fsl.scalar_at(2).is_null()); + assert!(!result_fsl.scalar_at(2).unwrap().is_null()); let third = result_fsl.fixed_size_list_elements_at(2); - assert_eq!(third.scalar_at(0), 1i32.into()); - assert_eq!(third.scalar_at(1), 2i32.into()); + assert_eq!(third.scalar_at(0).unwrap(), 1i32.into()); + assert_eq!(third.scalar_at(1).unwrap(), 2i32.into()); } // Parameterized test for nullable array scenarios that are specific to FSL's implementation. @@ -240,6 +240,6 @@ fn test_take_nullable_arrays_fsl_specific( // Check nullability of results. for (i, expected_null) in expected_nulls.iter().enumerate() { - assert_eq!(result_fsl.scalar_at(i).is_null(), *expected_null); + assert_eq!(result_fsl.scalar_at(i).unwrap().is_null(), *expected_null); } } diff --git a/vortex-array/src/arrays/fixed_size_list/vtable/operations.rs b/vortex-array/src/arrays/fixed_size_list/vtable/operations.rs index 0daac82189b..222194ea4f1 100644 --- a/vortex-array/src/arrays/fixed_size_list/vtable/operations.rs +++ b/vortex-array/src/arrays/fixed_size_list/vtable/operations.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_error::VortexResult; use vortex_scalar::Scalar; use crate::arrays::FixedSizeListArray; @@ -8,17 +9,19 @@ use crate::arrays::FixedSizeListVTable; use crate::vtable::OperationsVTable; impl OperationsVTable for FixedSizeListVTable { - fn scalar_at(array: &FixedSizeListArray, index: usize) -> Scalar { + fn scalar_at(array: &FixedSizeListArray, index: usize) -> VortexResult { // By the preconditions we know that the list scalar is not null. let list = array.fixed_size_list_elements_at(index); - let children_elements: Vec = (0..list.len()).map(|i| list.scalar_at(i)).collect(); + let children_elements: Vec = (0..list.len()) + .map(|i| list.scalar_at(i)) + .collect::>()?; debug_assert_eq!(children_elements.len(), array.list_size() as usize); - Scalar::fixed_size_list( + Ok(Scalar::fixed_size_list( list.dtype().clone(), children_elements, array.dtype().nullability(), - ) + )) } } diff --git a/vortex-array/src/arrays/list/array.rs b/vortex-array/src/arrays/list/array.rs index 606cd88e599..1122e1243a3 100644 --- a/vortex-array/src/arrays/list/array.rs +++ b/vortex-array/src/arrays/list/array.rs @@ -262,6 +262,7 @@ impl ListArray { .unwrap_or_else(|| { self.offsets() .scalar_at(index) + .vortex_expect("offsets must support scalar_at") .as_primitive() .as_::() .vortex_expect("index must fit in usize") @@ -311,7 +312,7 @@ impl ListArray { } let offsets = self.offsets(); - let first_offset = offsets.scalar_at(0); + let first_offset = offsets.scalar_at(0)?; let adjusted_offsets = sub_scalar(offsets, first_offset)?; Self::try_new(elements, adjusted_offsets, self.validity.clone()) diff --git a/vortex-array/src/arrays/list/compute/is_constant.rs b/vortex-array/src/arrays/list/compute/is_constant.rs index 643daffa1fb..8730d9212d8 100644 --- a/vortex-array/src/arrays/list/compute/is_constant.rs +++ b/vortex-array/src/arrays/list/compute/is_constant.rs @@ -57,11 +57,11 @@ impl IsConstantKernel for ListVTable { array.len() > 1, "precondition for `is_constant` is incorrect" ); - let first_scalar = array.scalar_at(0); // We checked the array length above. + let first_scalar = array.scalar_at(0)?; // We checked the array length above. // All lists have the same length, so compare the actual list contents. for i in 1..array.len() { - let current_scalar = array.scalar_at(i); + let current_scalar = array.scalar_at(i)?; if current_scalar != first_scalar { return Ok(Some(false)); } diff --git a/vortex-array/src/arrays/list/compute/take.rs b/vortex-array/src/arrays/list/compute/take.rs index 4d7ee4303f2..5bc6c309671 100644 --- a/vortex-array/src/arrays/list/compute/take.rs +++ b/vortex-array/src/arrays/list/compute/take.rs @@ -53,8 +53,8 @@ fn _take( array: &ListArray, indices_array: &PrimitiveArray, ) -> VortexResult { - let data_validity = array.validity_mask(); - let indices_validity = indices_array.validity_mask(); + let data_validity = array.validity_mask()?; + let indices_validity = indices_array.validity_mask()?; if !indices_validity.all_true() || !data_validity.all_true() { return _take_nullable::(array, indices_array); @@ -117,8 +117,8 @@ fn _take_nullable::with_capacity( Nullability::NonNullable, @@ -228,9 +228,9 @@ mod test { let element_dtype: Arc = Arc::new(I32.into()); - assert!(result.is_valid(0)); + assert!(result.is_valid(0).unwrap()); assert_eq!( - result.scalar_at(0), + result.scalar_at(0).unwrap(), Scalar::list( element_dtype.clone(), vec![0i32.into(), 5.into()], @@ -238,11 +238,11 @@ mod test { ) ); - assert!(result.is_invalid(1)); + assert!(result.is_invalid(1).unwrap()); - assert!(result.is_valid(2)); + assert!(result.is_valid(2).unwrap()); assert_eq!( - result.scalar_at(2), + result.scalar_at(2).unwrap(), Scalar::list( element_dtype.clone(), vec![3i32.into()], @@ -250,9 +250,9 @@ mod test { ) ); - assert!(result.is_valid(3)); + assert!(result.is_valid(3).unwrap()); assert_eq!( - result.scalar_at(3), + result.scalar_at(3).unwrap(), Scalar::list(element_dtype, vec![], Nullability::Nullable) ); } @@ -308,9 +308,9 @@ mod test { let element_dtype: Arc = Arc::new(I32.into()); - assert!(result.is_valid(0)); + assert!(result.is_valid(0).unwrap()); assert_eq!( - result.scalar_at(0), + result.scalar_at(0).unwrap(), Scalar::list( element_dtype.clone(), vec![3i32.into()], @@ -318,9 +318,9 @@ mod test { ) ); - assert!(result.is_valid(1)); + assert!(result.is_valid(1).unwrap()); assert_eq!( - result.scalar_at(1), + result.scalar_at(1).unwrap(), Scalar::list( element_dtype.clone(), vec![0i32.into(), 5.into()], @@ -328,9 +328,9 @@ mod test { ) ); - assert!(result.is_valid(2)); + assert!(result.is_valid(2).unwrap()); assert_eq!( - result.scalar_at(2), + result.scalar_at(2).unwrap(), Scalar::list(element_dtype, vec![], Nullability::NonNullable) ); } @@ -416,8 +416,8 @@ mod test { let result_view = result.to_listview(); assert_eq!(result_view.len(), 2); - assert!(result_view.is_valid(0)); - assert!(result_view.is_valid(1)); + assert!(result_view.is_valid(0).unwrap()); + assert!(result_view.is_valid(1).unwrap()); } #[test] @@ -437,9 +437,9 @@ mod test { let result_view = result.to_listview(); assert_eq!(result_view.len(), 3); - assert!(result_view.is_valid(0)); - assert!(result_view.is_invalid(1)); - assert!(result_view.is_valid(2)); + assert!(result_view.is_valid(0).unwrap()); + assert!(result_view.is_invalid(1).unwrap()); + assert!(result_view.is_valid(2).unwrap()); } /// Regression test for validity length mismatch bug. diff --git a/vortex-array/src/arrays/list/tests.rs b/vortex-array/src/arrays/list/tests.rs index 581341c4d59..af4e0d694bf 100644 --- a/vortex-array/src/arrays/list/tests.rs +++ b/vortex-array/src/arrays/list/tests.rs @@ -47,7 +47,7 @@ fn test_simple_list_array() { vec![1.into(), 2.into()], Nullability::Nullable ), - list.scalar_at(0) + list.scalar_at(0).unwrap() ); assert_eq!( Scalar::list( @@ -55,11 +55,11 @@ fn test_simple_list_array() { vec![3.into(), 4.into()], Nullability::Nullable ), - list.scalar_at(1) + list.scalar_at(1).unwrap() ); assert_eq!( Scalar::list(Arc::new(I32.into()), vec![5.into()], Nullability::Nullable), - list.scalar_at(2) + list.scalar_at(2).unwrap() ); } @@ -76,8 +76,14 @@ fn test_simple_list_array_from_iter() { .unwrap(); assert_eq!(list.len(), list_from_iter.len()); - assert_eq!(list.scalar_at(0), list_from_iter.scalar_at(0)); - assert_eq!(list.scalar_at(1), list_from_iter.scalar_at(1)); + assert_eq!( + list.scalar_at(0).unwrap(), + list_from_iter.scalar_at(0).unwrap() + ); + assert_eq!( + list.scalar_at(1).unwrap(), + list_from_iter.scalar_at(1).unwrap() + ); } #[test] @@ -213,10 +219,10 @@ fn test_list_filter_with_nulls() { assert_eq!(filtered_list.len(), 4); // Check validity of filtered array. - assert!(filtered_list.scalar_at(0).is_valid()); - assert!(!filtered_list.scalar_at(1).is_valid()); // Was null. - assert!(!filtered_list.scalar_at(2).is_valid()); // Was null. - assert!(filtered_list.scalar_at(3).is_valid()); + assert!(filtered_list.scalar_at(0).unwrap().is_valid()); + assert!(!filtered_list.scalar_at(1).unwrap().is_valid()); // Was null. + assert!(!filtered_list.scalar_at(2).unwrap().is_valid()); // Was null. + assert!(filtered_list.scalar_at(3).unwrap().is_valid()); } #[test] @@ -608,7 +614,7 @@ fn test_list_of_lists() { assert_arrays_eq!(inner, PrimitiveArray::from_iter([7])); // Test scalar conversion. - let scalar = list_of_lists.scalar_at(0); + let scalar = list_of_lists.scalar_at(0).unwrap(); assert!(matches!(scalar.dtype(), DType::List(_, _))); let list_scalar = scalar.as_list(); assert_eq!(list_scalar.len(), 2); @@ -653,11 +659,11 @@ fn test_list_of_lists_nullable_outer() { )); // First element should be [[1, 2], [3]]. - let first = list_of_lists.scalar_at(0); + let first = list_of_lists.scalar_at(0).unwrap(); assert!(!first.is_null()); // Second element should be null. - let second = list_of_lists.scalar_at(1); + let second = list_of_lists.scalar_at(1).unwrap(); assert!(second.is_null()); // Third element should be [[4, 5, 6]]. @@ -710,7 +716,7 @@ fn test_list_of_lists_nullable_inner() { assert_eq!(first_list.len(), 3); // Check that second inner list is null. - let second_inner = first_list.scalar_at(1); + let second_inner = first_list.scalar_at(1).unwrap(); assert!(second_inner.is_null()); } @@ -738,7 +744,7 @@ fn test_list_of_lists_both_nullable() { )); // First outer list should have 2 elements, second is null inner list. - let first_outer = list_of_lists.scalar_at(0); + let first_outer = list_of_lists.scalar_at(0).unwrap(); assert!(!first_outer.is_null()); let first_outer_array = list_of_lists.list_elements_at(0); let first_list = first_outer_array.as_::(); @@ -749,11 +755,11 @@ fn test_list_of_lists_both_nullable() { assert_eq!(first_inner.len(), 2); // Second inner list should be null. - let second_inner = first_list.scalar_at(1); + let second_inner = first_list.scalar_at(1).unwrap(); assert!(second_inner.is_null()); // Second outer list should be null. - let second_outer = list_of_lists.scalar_at(1); + let second_outer = list_of_lists.scalar_at(1).unwrap(); assert!(second_outer.is_null()); // Third outer list should have [3]. @@ -767,7 +773,7 @@ fn test_list_of_lists_both_nullable() { let fourth_outer = list_of_lists.list_elements_at(3); let fourth_list = fourth_outer.as_::(); assert_eq!(fourth_list.len(), 1); - let inner = fourth_list.scalar_at(0); + let inner = fourth_list.scalar_at(0).unwrap(); assert!(inner.is_null()); } @@ -894,5 +900,8 @@ fn test_recursive_compact_list_of_lists() { assert_eq!(recursive_flat_elements.len(), 7); // Verify data integrity is preserved - assert_eq!(non_recursive.scalar_at(0), recursive.scalar_at(0)); + assert_eq!( + non_recursive.scalar_at(0).unwrap(), + recursive.scalar_at(0).unwrap() + ); } diff --git a/vortex-array/src/arrays/list/vtable/operations.rs b/vortex-array/src/arrays/list/vtable/operations.rs index 582ee3f2503..5102feb465b 100644 --- a/vortex-array/src/arrays/list/vtable/operations.rs +++ b/vortex-array/src/arrays/list/vtable/operations.rs @@ -3,6 +3,7 @@ use std::sync::Arc; +use vortex_error::VortexResult; use vortex_scalar::Scalar; use crate::arrays::ListArray; @@ -10,15 +11,17 @@ use crate::arrays::ListVTable; use crate::vtable::OperationsVTable; impl OperationsVTable for ListVTable { - fn scalar_at(array: &ListArray, index: usize) -> Scalar { + fn scalar_at(array: &ListArray, index: usize) -> VortexResult { // By the preconditions we know that the list scalar is not null. let elems = array.list_elements_at(index); - let scalars: Vec = (0..elems.len()).map(|i| elems.scalar_at(i)).collect(); + let scalars: Vec = (0..elems.len()) + .map(|i| elems.scalar_at(i)) + .collect::>()?; - Scalar::list( + Ok(Scalar::list( Arc::new(elems.dtype().clone()), scalars, array.dtype().nullability(), - ) + )) } } diff --git a/vortex-array/src/arrays/listview/array.rs b/vortex-array/src/arrays/listview/array.rs index 1c9389a9984..bb0c095ad7d 100644 --- a/vortex-array/src/arrays/listview/array.rs +++ b/vortex-array/src/arrays/listview/array.rs @@ -371,6 +371,7 @@ impl ListViewArray { // Slow path: use `scalar_at` if we can't downcast directly to `PrimitiveArray`. self.offsets .scalar_at(index) + .vortex_expect("offsets must support scalar_at") .as_primitive() .as_::() .vortex_expect("offset must fit in usize") @@ -398,6 +399,7 @@ impl ListViewArray { // Slow path: use `scalar_at` if we can't downcast directly to `PrimitiveArray`. self.sizes .scalar_at(index) + .vortex_expect("sizes must support scalar_at") .as_primitive() .as_::() .vortex_expect("size must fit in usize") diff --git a/vortex-array/src/arrays/listview/compute/is_constant.rs b/vortex-array/src/arrays/listview/compute/is_constant.rs index c77c350c98a..5a064be752e 100644 --- a/vortex-array/src/arrays/listview/compute/is_constant.rs +++ b/vortex-array/src/arrays/listview/compute/is_constant.rs @@ -36,11 +36,11 @@ impl IsConstantKernel for ListViewVTable { array.len() > 1, "precondition for `is_constant` is incorrect" ); - let first_scalar = array.scalar_at(0); + let first_scalar = array.scalar_at(0)?; // Compare all other scalars to the first. for i in 1..array.len() { - if array.scalar_at(i) != first_scalar { + if array.scalar_at(i)? != first_scalar { return Ok(Some(false)); } } diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index cd76b6f2898..6acd4a1b67b 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -143,7 +143,7 @@ impl ListViewArray { let mut n_elements = NewOffset::zero(); for index in 0..len { - if !self.is_valid(index) { + if !self.is_valid(index)? { // For NULL lists, place them after the previous item's data to maintain the // no-overlap invariant for zero-copy to `ListArray` arrays. new_offsets.push(n_elements); @@ -391,7 +391,7 @@ mod tests { // Note that element at index 2 (97) is preserved as a gap. let all_elements = trimmed.elements().to_primitive(); - assert_eq!(all_elements.scalar_at(2), 97i32.into()); + assert_eq!(all_elements.scalar_at(2).unwrap(), 97i32.into()); Ok(()) } @@ -431,10 +431,10 @@ mod tests { let exact = rebuilt.rebuild(ListViewRebuildMode::MakeExact)?; // Verify the result is still valid - assert!(exact.is_valid(0)); - assert!(exact.is_valid(1)); - assert!(!exact.is_valid(2)); - assert!(!exact.is_valid(3)); + assert!(exact.is_valid(0).unwrap()); + assert!(exact.is_valid(1).unwrap()); + assert!(!exact.is_valid(2).unwrap()); + assert!(!exact.is_valid(3).unwrap()); // Verify data is preserved assert_arrays_eq!( diff --git a/vortex-array/src/arrays/listview/tests/basic.rs b/vortex-array/src/arrays/listview/tests/basic.rs index 2e5d24755f8..b8255e2a0ba 100644 --- a/vortex-array/src/arrays/listview/tests/basic.rs +++ b/vortex-array/src/arrays/listview/tests/basic.rs @@ -54,7 +54,7 @@ fn test_basic_listview_comprehensive() { ); // Test scalar_at which returns entire lists as Scalar values. - let first_scalar = listview.scalar_at(0); + let first_scalar = listview.scalar_at(0).unwrap(); assert_eq!( first_scalar, Scalar::list( @@ -138,9 +138,9 @@ fn test_from_list_array() -> VortexResult<()> { ); // Check validity is preserved. - assert!(list_view.is_valid(0)); - assert!(list_view.is_invalid(1)); - assert!(list_view.is_valid(2)); + assert!(list_view.is_valid(0).unwrap()); + assert!(list_view.is_invalid(1).unwrap()); + assert!(list_view.is_valid(2).unwrap()); // Check third list. assert_arrays_eq!( @@ -197,9 +197,18 @@ fn test_listview_with_constant_arrays(#[case] const_sizes: bool, #[case] const_o assert_eq!(listview.list_elements_at(2).len(), 3); } else if const_offsets { // All lists start at offset 0, different sizes (overlapping). - assert_eq!(listview.list_elements_at(0).scalar_at(0), 1i32.into()); - assert_eq!(listview.list_elements_at(1).scalar_at(0), 1i32.into()); - assert_eq!(listview.list_elements_at(2).scalar_at(0), 1i32.into()); + assert_eq!( + listview.list_elements_at(0).scalar_at(0).unwrap(), + 1i32.into() + ); + assert_eq!( + listview.list_elements_at(1).scalar_at(0).unwrap(), + 1i32.into() + ); + assert_eq!( + listview.list_elements_at(2).scalar_at(0).unwrap(), + 1i32.into() + ); } } diff --git a/vortex-array/src/arrays/listview/tests/filter.rs b/vortex-array/src/arrays/listview/tests/filter.rs index 0318869ebb8..5404aef9694 100644 --- a/vortex-array/src/arrays/listview/tests/filter.rs +++ b/vortex-array/src/arrays/listview/tests/filter.rs @@ -186,11 +186,21 @@ fn test_filter_extreme_offsets() { // Verify we can still read the correct values. let list0 = result_list.list_elements_at(0); assert_eq!( - list0.scalar_at(0).as_primitive().as_::().unwrap(), + list0 + .scalar_at(0) + .unwrap() + .as_primitive() + .as_::() + .unwrap(), 4999 ); assert_eq!( - list0.scalar_at(1).as_primitive().as_::().unwrap(), + list0 + .scalar_at(1) + .unwrap() + .as_primitive() + .as_::() + .unwrap(), 5000 ); diff --git a/vortex-array/src/arrays/listview/tests/nested.rs b/vortex-array/src/arrays/listview/tests/nested.rs index 162ef096a9f..8bb223d1459 100644 --- a/vortex-array/src/arrays/listview/tests/nested.rs +++ b/vortex-array/src/arrays/listview/tests/nested.rs @@ -70,12 +70,44 @@ fn test_listview_of_listview_with_overlapping() { let inner1 = first_outer_lv.list_elements_at(1); // inner[0] should be [1, 2, 3]. - assert_eq!(inner0.scalar_at(0).as_primitive().as_::().unwrap(), 1); - assert_eq!(inner0.scalar_at(2).as_primitive().as_::().unwrap(), 3); + assert_eq!( + inner0 + .scalar_at(0) + .unwrap() + .as_primitive() + .as_::() + .unwrap(), + 1 + ); + assert_eq!( + inner0 + .scalar_at(2) + .unwrap() + .as_primitive() + .as_::() + .unwrap(), + 3 + ); // inner[1] should be [3, 4, 5] - shares element 3 with inner[0]. - assert_eq!(inner1.scalar_at(0).as_primitive().as_::().unwrap(), 3); - assert_eq!(inner1.scalar_at(1).as_primitive().as_::().unwrap(), 4); + assert_eq!( + inner1 + .scalar_at(0) + .unwrap() + .as_primitive() + .as_::() + .unwrap(), + 3 + ); + assert_eq!( + inner1 + .scalar_at(1) + .unwrap() + .as_primitive() + .as_::() + .unwrap(), + 4 + ); // Test slicing the outer ListView. let sliced = outer_listview.slice(1..2); @@ -255,7 +287,15 @@ fn test_listview_zero_and_overlapping() { let inner1 = first_outer_lv.list_elements_at(1); assert_eq!(inner1.len(), 3); // [1, 2, 3] - assert_eq!(inner1.scalar_at(0).as_primitive().as_::().unwrap(), 1); + assert_eq!( + inner1 + .scalar_at(0) + .unwrap() + .as_primitive() + .as_::() + .unwrap(), + 1 + ); let inner2 = first_outer_lv.list_elements_at(2); assert_eq!(inner2.len(), 0); // Empty @@ -266,7 +306,15 @@ fn test_listview_zero_and_overlapping() { let inner3 = second_outer_lv.list_elements_at(0); assert_eq!(inner3.len(), 3); // [2, 3, 4] - assert_eq!(inner3.scalar_at(0).as_primitive().as_::().unwrap(), 2); + assert_eq!( + inner3 + .scalar_at(0) + .unwrap() + .as_primitive() + .as_::() + .unwrap(), + 2 + ); // Verify slicing works with empty lists. let sliced = outer_listview.slice(0..2); @@ -337,7 +385,7 @@ fn test_listview_of_struct_with_nulls() { assert_eq!(list1.len(), 3); // The middle element (struct[2]) should be null. - assert!(list1.scalar_at(1).is_null()); + assert!(list1.scalar_at(1).unwrap().is_null()); // Test slicing preserves null handling. let sliced = listview.slice(1..3); diff --git a/vortex-array/src/arrays/listview/tests/nullability.rs b/vortex-array/src/arrays/listview/tests/nullability.rs index f4b947f4ca6..e4d1eb67fb8 100644 --- a/vortex-array/src/arrays/listview/tests/nullability.rs +++ b/vortex-array/src/arrays/listview/tests/nullability.rs @@ -33,9 +33,9 @@ fn test_nullable_listview_comprehensive() { assert_eq!(listview.len(), 3); // Check validity. - assert!(listview.is_valid(0)); - assert!(listview.is_invalid(1)); - assert!(listview.is_valid(2)); + assert!(listview.is_valid(0).unwrap()); + assert!(listview.is_invalid(1).unwrap()); + assert!(listview.is_valid(2).unwrap()); // Check dtype reflects nullability. assert!(matches!( @@ -44,7 +44,7 @@ fn test_nullable_listview_comprehensive() { )); // Test scalar_at with nulls. - let first = listview.scalar_at(0); + let first = listview.scalar_at(0).unwrap(); assert!(!first.is_null()); assert_eq!( first, @@ -55,10 +55,10 @@ fn test_nullable_listview_comprehensive() { ) ); - let second = listview.scalar_at(1); + let second = listview.scalar_at(1).unwrap(); assert!(second.is_null()); - let third = listview.scalar_at(2); + let third = listview.scalar_at(2).unwrap(); assert!(!third.is_null()); assert_eq!( third, @@ -72,8 +72,8 @@ fn test_nullable_listview_comprehensive() { // list_elements_at still returns data even for null lists. let null_list_data = listview.list_elements_at(1); assert_eq!(null_list_data.len(), 2); - assert_eq!(null_list_data.scalar_at(0), 3i32.into()); - assert_eq!(null_list_data.scalar_at(1), 4i32.into()); + assert_eq!(null_list_data.scalar_at(0).unwrap(), 3i32.into()); + assert_eq!(null_list_data.scalar_at(1).unwrap(), 4i32.into()); } // Parameterized tests for different null patterns. @@ -90,7 +90,7 @@ fn test_nullable_patterns(#[case] validity: Validity, #[case] expected_validity: let listview = unsafe { ListViewArray::new_unchecked(elements, offsets, sizes, validity) }; for (i, &expected) in expected_validity.iter().enumerate() { - assert_eq!(listview.is_valid(i), expected); + assert_eq!(listview.is_valid(i).unwrap(), expected); } } @@ -112,22 +112,22 @@ fn test_nullable_elements() { // First list: [Some(1), None]. let first_list = listview.list_elements_at(0); assert_eq!(first_list.len(), 2); - assert!(!first_list.scalar_at(0).is_null()); - assert_eq!(first_list.scalar_at(0), 1i32.into()); - assert!(first_list.scalar_at(1).is_null()); + assert!(!first_list.scalar_at(0).unwrap().is_null()); + assert_eq!(first_list.scalar_at(0).unwrap(), 1i32.into()); + assert!(first_list.scalar_at(1).unwrap().is_null()); // Second list: [Some(3), None]. let second_list = listview.list_elements_at(1); - assert!(!second_list.scalar_at(0).is_null()); - assert_eq!(second_list.scalar_at(0), 3i32.into()); - assert!(second_list.scalar_at(1).is_null()); + assert!(!second_list.scalar_at(0).unwrap().is_null()); + assert_eq!(second_list.scalar_at(0).unwrap(), 3i32.into()); + assert!(second_list.scalar_at(1).unwrap().is_null()); // Third list: [Some(5), Some(6)]. let third_list = listview.list_elements_at(2); - assert!(!third_list.scalar_at(0).is_null()); - assert_eq!(third_list.scalar_at(0), 5i32.into()); - assert!(!third_list.scalar_at(1).is_null()); - assert_eq!(third_list.scalar_at(1), 6i32.into()); + assert!(!third_list.scalar_at(0).unwrap().is_null()); + assert_eq!(third_list.scalar_at(0).unwrap(), 5i32.into()); + assert!(!third_list.scalar_at(1).unwrap().is_null()); + assert_eq!(third_list.scalar_at(1).unwrap(), 6i32.into()); // Check dtype of elements. assert!(matches!( diff --git a/vortex-array/src/arrays/listview/tests/operations.rs b/vortex-array/src/arrays/listview/tests/operations.rs index 5ad8f41ec4e..67e89d3ed0b 100644 --- a/vortex-array/src/arrays/listview/tests/operations.rs +++ b/vortex-array/src/arrays/listview/tests/operations.rs @@ -58,8 +58,8 @@ fn test_slice_comprehensive() { for i in 0..4 { // Compare the sliced elements assert_eq!( - full_list.scalar_at(i), - listview.scalar_at(i), + full_list.scalar_at(i).unwrap(), + listview.scalar_at(i).unwrap(), "Mismatch at index {}", i ); @@ -147,8 +147,8 @@ fn test_slice_with_nulls() { let sliced_list = sliced.as_::(); assert_eq!(sliced_list.len(), 2); - assert!(sliced_list.is_invalid(0)); // Original index 1 was null. - assert!(sliced_list.is_valid(1)); // Original index 2 was valid. + assert!(sliced_list.is_invalid(0).unwrap()); // Original index 1 was null. + assert!(sliced_list.is_valid(1).unwrap()); // Original index 2 was valid. // Verify offsets and sizes are preserved. assert_eq!(sliced_list.offset_at(0), 2); @@ -272,8 +272,8 @@ fn test_cast_with_nulls() { assert_eq!(result.dtype(), &target_dtype); let result_list = result.to_listview(); - assert!(result_list.is_valid(0)); - assert!(result_list.is_invalid(1)); + assert!(result_list.is_valid(0).unwrap()); + assert!(result_list.is_invalid(1).unwrap()); } #[rstest] @@ -516,10 +516,10 @@ fn test_mask_preserves_structure() { let result_list = result.to_listview(); // Check validity: true in mask means null. - assert!(!result_list.is_valid(0)); // Masked. - assert!(result_list.is_valid(1)); // Not masked. - assert!(!result_list.is_valid(2)); // Masked. - assert!(!result_list.is_valid(3)); // Masked. + assert!(!result_list.is_valid(0).unwrap()); // Masked. + assert!(result_list.is_valid(1).unwrap()); // Not masked. + assert!(!result_list.is_valid(2).unwrap()); // Masked. + assert!(!result_list.is_valid(3).unwrap()); // Masked. // Offsets and sizes are preserved. assert_eq!(result_list.offset_at(0), 0); @@ -553,9 +553,9 @@ fn test_mask_with_existing_nulls() { let result_list = result.to_listview(); // Check combined validity: - assert!(result_list.is_valid(0)); // Was valid, mask is false -> valid. - assert!(!result_list.is_valid(1)); // Was invalid, mask is true -> invalid. - assert!(!result_list.is_valid(2)); // Was valid, mask is true -> invalid. + assert!(result_list.is_valid(0).unwrap()); // Was valid, mask is false -> valid. + assert!(!result_list.is_valid(1).unwrap()); // Was invalid, mask is true -> invalid. + assert!(!result_list.is_valid(2).unwrap()); // Was valid, mask is true -> invalid. } #[test] @@ -573,9 +573,9 @@ fn test_mask_with_gaps() { let result_list = result.to_listview(); assert_eq!(result_list.len(), 3); - assert!(!result_list.is_valid(0)); // Masked - assert!(result_list.is_valid(1)); // Not masked - assert!(result_list.is_valid(2)); // Not masked + assert!(!result_list.is_valid(0).unwrap()); // Masked + assert!(result_list.is_valid(1).unwrap()); // Not masked + assert!(result_list.is_valid(2).unwrap()); // Not masked // Offsets and sizes still preserved assert_eq!(result_list.offset_at(1), 4); @@ -605,9 +605,9 @@ fn test_mask_constant_arrays() { let result_list = result.to_listview(); assert_eq!(result_list.len(), 3); - assert!(result_list.is_valid(0)); - assert!(!result_list.is_valid(1)); // Masked - assert!(result_list.is_valid(2)); + assert!(result_list.is_valid(0).unwrap()); + assert!(!result_list.is_valid(1).unwrap()); // Masked + assert!(result_list.is_valid(2).unwrap()); // All offsets and sizes remain constant assert_eq!(result_list.offset_at(0), 1); diff --git a/vortex-array/src/arrays/listview/tests/take.rs b/vortex-array/src/arrays/listview/tests/take.rs index d844046cf92..d2aedf33387 100644 --- a/vortex-array/src/arrays/listview/tests/take.rs +++ b/vortex-array/src/arrays/listview/tests/take.rs @@ -176,11 +176,21 @@ fn test_take_extreme_offsets() { // Verify we can still read the correct values. let list0 = result_list.list_elements_at(0); assert_eq!( - list0.scalar_at(0).as_primitive().as_::().unwrap(), + list0 + .scalar_at(0) + .unwrap() + .as_primitive() + .as_::() + .unwrap(), 4999 ); assert_eq!( - list0.scalar_at(1).as_primitive().as_::().unwrap(), + list0 + .scalar_at(1) + .unwrap() + .as_primitive() + .as_::() + .unwrap(), 5000 ); } diff --git a/vortex-array/src/arrays/listview/vtable/operations.rs b/vortex-array/src/arrays/listview/vtable/operations.rs index ec58b5c636b..e5b6a3cf7d1 100644 --- a/vortex-array/src/arrays/listview/vtable/operations.rs +++ b/vortex-array/src/arrays/listview/vtable/operations.rs @@ -3,6 +3,7 @@ use std::sync::Arc; +use vortex_error::VortexResult; use vortex_scalar::Scalar; use crate::arrays::ListViewArray; @@ -10,15 +11,17 @@ use crate::arrays::ListViewVTable; use crate::vtable::OperationsVTable; impl OperationsVTable for ListViewVTable { - fn scalar_at(array: &ListViewArray, index: usize) -> Scalar { + fn scalar_at(array: &ListViewArray, index: usize) -> VortexResult { // By the preconditions we know that the list scalar is not null. let list = array.list_elements_at(index); - let children: Vec = (0..list.len()).map(|i| list.scalar_at(i)).collect(); + let children: Vec = (0..list.len()) + .map(|i| list.scalar_at(i)) + .collect::>()?; - Scalar::list( + Ok(Scalar::list( Arc::new(list.dtype().clone()), children, array.dtype.nullability(), - ) + )) } } diff --git a/vortex-array/src/arrays/masked/array.rs b/vortex-array/src/arrays/masked/array.rs index 650a9eedb31..abb51e0ea20 100644 --- a/vortex-array/src/arrays/masked/array.rs +++ b/vortex-array/src/arrays/masked/array.rs @@ -24,7 +24,7 @@ impl MaskedArray { vortex_bail!("MaskedArray must have nullable validity, got {validity:?}") } - if !child.all_valid() { + if !child.all_valid()? { vortex_bail!("MaskedArray children must not have nulls"); } diff --git a/vortex-array/src/arrays/masked/compute/compare.rs b/vortex-array/src/arrays/masked/compute/compare.rs index 3e73cd7926f..e4b11a421a0 100644 --- a/vortex-array/src/arrays/masked/compute/compare.rs +++ b/vortex-array/src/arrays/masked/compute/compare.rs @@ -119,7 +119,10 @@ mod tests { vec![false, true, false] ); assert_eq!(res.dtype().nullability(), Nullability::Nullable); - assert_eq!(res.validity_mask(), Mask::from_iter([false, true, false])); + assert_eq!( + res.validity_mask().unwrap(), + Mask::from_iter([false, true, false]) + ); } #[test] @@ -142,6 +145,9 @@ mod tests { ); assert_eq!(res.dtype().nullability(), Nullability::Nullable); // Validity is union of both: lhs=[T,T,F], rhs=[T,F,T] => result=[T,F,F] - assert_eq!(res.validity_mask(), Mask::from_iter([true, false, false])); + assert_eq!( + res.validity_mask().unwrap(), + Mask::from_iter([true, false, false]) + ); } } diff --git a/vortex-array/src/arrays/masked/compute/take.rs b/vortex-array/src/arrays/masked/compute/take.rs index 5a5f6307057..945a1a047ed 100644 --- a/vortex-array/src/arrays/masked/compute/take.rs +++ b/vortex-array/src/arrays/masked/compute/take.rs @@ -18,7 +18,7 @@ use crate::vtable::ValidityHelper; impl TakeKernel for MaskedVTable { fn take(&self, array: &MaskedArray, indices: &dyn Array) -> VortexResult { - let taken_child = if !indices.all_valid() { + let taken_child = if !indices.all_valid()? { // This is safe because we'll mask out these positions in the validity let filled_take = fill_null( indices, diff --git a/vortex-array/src/arrays/masked/tests.rs b/vortex-array/src/arrays/masked/tests.rs index b65085d5a6c..fb0512bcd9e 100644 --- a/vortex-array/src/arrays/masked/tests.rs +++ b/vortex-array/src/arrays/masked/tests.rs @@ -59,12 +59,12 @@ fn test_masked_child_with_validity() { let prim = array.to_primitive(); // Positions where validity is false should be null in masked_child. - assert_eq!(prim.valid_count(), 3); - assert!(prim.is_valid(0)); - assert!(!prim.is_valid(1)); - assert!(prim.is_valid(2)); - assert!(!prim.is_valid(3)); - assert!(prim.is_valid(4)); + assert_eq!(prim.valid_count().unwrap(), 3); + assert!(prim.is_valid(0).unwrap()); + assert!(!prim.is_valid(1).unwrap()); + assert!(prim.is_valid(2).unwrap()); + assert!(!prim.is_valid(3).unwrap()); + assert!(prim.is_valid(4).unwrap()); } #[test] @@ -74,7 +74,7 @@ fn test_masked_child_all_valid() { let array = MaskedArray::try_new(child.clone(), Validity::AllValid).unwrap(); assert_eq!(array.len(), 3); - assert_eq!(array.valid_count(), 3); + assert_eq!(array.valid_count().unwrap(), 3); assert_arrays_eq!( PrimitiveArray::from_option_iter([10i32, 20, 30].map(Some)), array diff --git a/vortex-array/src/arrays/masked/vtable/canonical.rs b/vortex-array/src/arrays/masked/vtable/canonical.rs index 66092e2d0be..74d9efe8c5a 100644 --- a/vortex-array/src/arrays/masked/vtable/canonical.rs +++ b/vortex-array/src/arrays/masked/vtable/canonical.rs @@ -53,12 +53,12 @@ mod tests { let prim = canonical.as_ref().to_primitive(); // Check that null positions match validity. - assert_eq!(prim.valid_count(), 3); - assert!(prim.is_valid(0)); - assert!(!prim.is_valid(1)); - assert!(prim.is_valid(2)); - assert!(!prim.is_valid(3)); - assert!(prim.is_valid(4)); + assert_eq!(prim.valid_count().unwrap(), 3); + assert!(prim.is_valid(0).unwrap()); + assert!(!prim.is_valid(1).unwrap()); + assert!(prim.is_valid(2).unwrap()); + assert!(!prim.is_valid(3).unwrap()); + assert!(prim.is_valid(4).unwrap()); Ok(()) } @@ -71,7 +71,7 @@ mod tests { .unwrap(); let canonical = array.to_canonical()?; - assert_eq!(canonical.as_ref().valid_count(), 3); + assert_eq!(canonical.as_ref().valid_count().unwrap(), 3); assert_eq!( canonical.as_ref().dtype().nullability(), Nullability::Nullable diff --git a/vortex-array/src/arrays/masked/vtable/operations.rs b/vortex-array/src/arrays/masked/vtable/operations.rs index ad213e6b4c3..ca9f4607b24 100644 --- a/vortex-array/src/arrays/masked/vtable/operations.rs +++ b/vortex-array/src/arrays/masked/vtable/operations.rs @@ -1,15 +1,17 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_error::VortexResult; use vortex_scalar::Scalar; +use crate::Array; use crate::arrays::MaskedVTable; use crate::arrays::masked::MaskedArray; use crate::vtable::OperationsVTable; impl OperationsVTable for MaskedVTable { - fn scalar_at(array: &MaskedArray, index: usize) -> Scalar { + fn scalar_at(array: &MaskedArray, index: usize) -> VortexResult { // Invalid indices are handled by the entrypoint function. - array.child.scalar_at(index).into_nullable() + Ok(array.child.scalar_at(index)?.into_nullable()) } } diff --git a/vortex-array/src/arrays/null/compute/cast.rs b/vortex-array/src/arrays/null/compute/cast.rs index bf72d156208..ceb35504e63 100644 --- a/vortex-array/src/arrays/null/compute/cast.rs +++ b/vortex-array/src/arrays/null/compute/cast.rs @@ -69,7 +69,7 @@ mod tests { // Verify all values are null for i in 0..5 { - assert!(result.scalar_at(i).is_null()); + assert!(result.scalar_at(i).unwrap().is_null()); } } diff --git a/vortex-array/src/arrays/null/compute/mod.rs b/vortex-array/src/arrays/null/compute/mod.rs index 237a558738b..316798ba98d 100644 --- a/vortex-array/src/arrays/null/compute/mod.rs +++ b/vortex-array/src/arrays/null/compute/mod.rs @@ -29,7 +29,7 @@ mod test { let sliced = nulls.slice(0..4).to_null(); assert_eq!(sliced.len(), 4); - assert!(matches!(sliced.validity_mask(), Mask::AllFalse(4))); + assert!(matches!(sliced.validity_mask().unwrap(), Mask::AllFalse(4))); } #[test] @@ -40,14 +40,14 @@ mod test { .to_null(); assert_eq!(taken.len(), 5); - assert!(matches!(taken.validity_mask(), Mask::AllFalse(5))); + assert!(matches!(taken.validity_mask().unwrap(), Mask::AllFalse(5))); } #[test] fn test_scalar_at_nulls() { let nulls = NullArray::new(10); - let scalar = nulls.scalar_at(0); + let scalar = nulls.scalar_at(0).unwrap(); assert!(scalar.is_null()); assert_eq!(scalar.dtype().clone(), DType::Null); } diff --git a/vortex-array/src/arrays/null/mod.rs b/vortex-array/src/arrays/null/mod.rs index 746a2b3d5f0..fcd2436c916 100644 --- a/vortex-array/src/arrays/null/mod.rs +++ b/vortex-array/src/arrays/null/mod.rs @@ -112,7 +112,7 @@ impl VTable for NullVTable { /// assert_eq!(sliced.len(), 2); /// /// // All elements are null -/// let scalar = array.scalar_at(0); +/// let scalar = array.scalar_at(0).unwrap(); /// assert!(scalar.is_null()); /// ``` #[derive(Clone, Debug)] @@ -166,8 +166,8 @@ impl VisitorVTable for NullVTable { } impl OperationsVTable for NullVTable { - fn scalar_at(_array: &NullArray, _index: usize) -> Scalar { - Scalar::null(DType::Null) + fn scalar_at(_array: &NullArray, _index: usize) -> VortexResult { + Ok(Scalar::null(DType::Null)) } } @@ -176,7 +176,7 @@ impl ValidityVTable for NullVTable { Ok(Validity::AllInvalid) } - fn validity_mask(array: &NullArray) -> Mask { - Mask::AllFalse(array.len) + fn validity_mask(array: &NullArray) -> VortexResult { + Ok(Mask::AllFalse(array.len)) } } diff --git a/vortex-array/src/arrays/primitive/array/conversion.rs b/vortex-array/src/arrays/primitive/array/conversion.rs index 8bfbd8a52d5..b97c177024b 100644 --- a/vortex-array/src/arrays/primitive/array/conversion.rs +++ b/vortex-array/src/arrays/primitive/array/conversion.rs @@ -171,11 +171,11 @@ mod tests { assert_eq!(result.len(), 5); assert_eq!(result.ptype(), PType::I32); - assert!(result.is_valid(0)); - assert!(!result.is_valid(1)); - assert!(result.is_valid(2)); - assert!(result.is_valid(3)); - assert!(!result.is_valid(4)); + assert!(result.is_valid(0).unwrap()); + assert!(!result.is_valid(1).unwrap()); + assert!(result.is_valid(2).unwrap()); + assert!(result.is_valid(3).unwrap()); + assert!(!result.is_valid(4).unwrap()); } #[test] diff --git a/vortex-array/src/arrays/primitive/array/mod.rs b/vortex-array/src/arrays/primitive/array/mod.rs index 2cb835f4101..7dbe7bfab0f 100644 --- a/vortex-array/src/arrays/primitive/array/mod.rs +++ b/vortex-array/src/arrays/primitive/array/mod.rs @@ -57,7 +57,7 @@ use crate::buffer::BufferHandle; /// let sliced = array.slice(1..3); /// /// // Access individual values -/// let value = sliced.scalar_at(0); +/// let value = sliced.scalar_at(0).unwrap(); /// assert_eq!(value, 2i32.into()); /// /// // Convert into a type-erased array that can be passed to compute functions. diff --git a/vortex-array/src/arrays/primitive/array/top_value.rs b/vortex-array/src/arrays/primitive/array/top_value.rs index 6b009827882..0b67de41027 100644 --- a/vortex-array/src/arrays/primitive/array/top_value.rs +++ b/vortex-array/src/arrays/primitive/array/top_value.rs @@ -23,12 +23,12 @@ impl PrimitiveArray { return Ok(None); } - if self.all_invalid() { + if self.all_invalid()? { return Ok(None); } match_each_native_ptype!(self.ptype(), |P| { - let (top, count) = typed_top_value(self.as_slice::

(), self.validity_mask()); + let (top, count) = typed_top_value(self.as_slice::

(), self.validity_mask()?); Ok(Some((top.into(), count))) }) } diff --git a/vortex-array/src/arrays/primitive/compute/cast.rs b/vortex-array/src/arrays/primitive/compute/cast.rs index ceadfbea78b..83ab791d532 100644 --- a/vortex-array/src/arrays/primitive/compute/cast.rs +++ b/vortex-array/src/arrays/primitive/compute/cast.rs @@ -46,7 +46,7 @@ impl CastKernel for PrimitiveVTable { })); } - let mask = array.validity_mask(); + let mask = array.validity_mask()?; // Otherwise, we need to cast the values one-by-one Ok(Some(match_each_native_ptype!(new_ptype, |T| { @@ -217,7 +217,7 @@ mod test { PrimitiveArray::from_option_iter([None, Some(0u32), Some(10)]) ); assert_eq!( - p.validity_mask(), + p.validity_mask().unwrap(), Mask::from(BitBuffer::from(vec![false, true, true])) ); } diff --git a/vortex-array/src/arrays/primitive/compute/fill_null.rs b/vortex-array/src/arrays/primitive/compute/fill_null.rs index 8a3f5013e66..3d6d13e4b3e 100644 --- a/vortex-array/src/arrays/primitive/compute/fill_null.rs +++ b/vortex-array/src/arrays/primitive/compute/fill_null.rs @@ -65,7 +65,7 @@ mod test { .unwrap() .to_primitive(); assert_arrays_eq!(p, PrimitiveArray::from_iter([42u8, 8, 42, 10, 42])); - assert!(p.validity_mask().all_true()); + assert!(p.validity_mask().unwrap().all_true()); } #[test] @@ -76,7 +76,7 @@ mod test { .unwrap() .to_primitive(); assert_arrays_eq!(p, PrimitiveArray::from_iter([255u8, 255, 255, 255, 255])); - assert!(p.validity_mask().all_true()); + assert!(p.validity_mask().unwrap().all_true()); } #[test] @@ -89,7 +89,7 @@ mod test { .unwrap() .to_primitive(); assert_arrays_eq!(p, PrimitiveArray::from_iter([8u8, 10, 12, 14, 16])); - assert!(p.validity_mask().all_true()); + assert!(p.validity_mask().unwrap().all_true()); } #[test] @@ -99,6 +99,6 @@ mod test { .unwrap() .to_primitive(); assert_arrays_eq!(p, PrimitiveArray::from_iter([8u8, 10, 12, 14, 16])); - assert!(p.validity_mask().all_true()); + assert!(p.validity_mask().unwrap().all_true()); } } diff --git a/vortex-array/src/arrays/primitive/compute/is_sorted.rs b/vortex-array/src/arrays/primitive/compute/is_sorted.rs index 5ff6a22909f..e00e68d704d 100644 --- a/vortex-array/src/arrays/primitive/compute/is_sorted.rs +++ b/vortex-array/src/arrays/primitive/compute/is_sorted.rs @@ -61,7 +61,7 @@ where } fn compute_is_sorted(array: &PrimitiveArray, strict: bool) -> VortexResult { - match array.validity_mask() { + match array.validity_mask()? { Mask::AllFalse(_) => Ok(!strict), Mask::AllTrue(_) => { let slice = array.as_slice::(); diff --git a/vortex-array/src/arrays/primitive/compute/min_max.rs b/vortex-array/src/arrays/primitive/compute/min_max.rs index 3ef56d1f121..1c3445c983a 100644 --- a/vortex-array/src/arrays/primitive/compute/min_max.rs +++ b/vortex-array/src/arrays/primitive/compute/min_max.rs @@ -33,7 +33,7 @@ where T: NativePType, PValue: From, { - Ok(match array.validity_mask() { + Ok(match array.validity_mask()? { Mask::AllTrue(_) => compute_min_max(array.as_slice::().iter()), Mask::AllFalse(_) => None, Mask::Values(v) => compute_min_max( diff --git a/vortex-array/src/arrays/primitive/compute/nan_count.rs b/vortex-array/src/arrays/primitive/compute/nan_count.rs index 5a8a4e43c39..65f8873d8af 100644 --- a/vortex-array/src/arrays/primitive/compute/nan_count.rs +++ b/vortex-array/src/arrays/primitive/compute/nan_count.rs @@ -15,7 +15,7 @@ use crate::register_kernel; impl NaNCountKernel for PrimitiveVTable { fn nan_count(&self, array: &PrimitiveArray) -> VortexResult { Ok(match_each_float_ptype!(array.ptype(), |F| { - compute_nan_count_with_validity(array.as_slice::(), array.validity_mask()) + compute_nan_count_with_validity(array.as_slice::(), array.validity_mask()?) })) } } diff --git a/vortex-array/src/arrays/primitive/compute/sum.rs b/vortex-array/src/arrays/primitive/compute/sum.rs index bd55e45fa86..bd3d861d4eb 100644 --- a/vortex-array/src/arrays/primitive/compute/sum.rs +++ b/vortex-array/src/arrays/primitive/compute/sum.rs @@ -21,7 +21,7 @@ use crate::register_kernel; impl SumKernel for PrimitiveVTable { fn sum(&self, array: &PrimitiveArray, accumulator: &Scalar) -> VortexResult { - let array_sum_scalar = match array.validity_mask().bit_buffer() { + let array_sum_scalar = match array.validity_mask()?.bit_buffer() { AllOr::All => { // All-valid match_each_native_ptype!( diff --git a/vortex-array/src/arrays/primitive/compute/take/mod.rs b/vortex-array/src/arrays/primitive/compute/take/mod.rs index b740d6cb03d..fd2ddb63cf3 100644 --- a/vortex-array/src/arrays/primitive/compute/take/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/take/mod.rs @@ -114,6 +114,7 @@ fn take_primitive_scalar(array: &[T], indices: mod test { use rstest::rstest; use vortex_buffer::buffer; + use vortex_error::VortexExpect; use vortex_scalar::Scalar; use crate::Array; @@ -143,11 +144,20 @@ mod test { Validity::Array(BoolArray::from_iter([true, true, false]).into_array()), ); let actual = take(values.as_ref(), indices.as_ref()).unwrap(); - assert_eq!(actual.scalar_at(0), Scalar::from(Some(1))); + assert_eq!( + actual.scalar_at(0).vortex_expect("no fail"), + Scalar::from(Some(1)) + ); // position 3 is null - assert_eq!(actual.scalar_at(1), Scalar::null_typed::()); + assert_eq!( + actual.scalar_at(1).vortex_expect("no fail"), + Scalar::null_typed::() + ); // the third index is null - assert_eq!(actual.scalar_at(2), Scalar::null_typed::()); + assert_eq!( + actual.scalar_at(2).vortex_expect("no fail"), + Scalar::null_typed::() + ); } #[rstest] diff --git a/vortex-array/src/arrays/primitive/tests.rs b/vortex-array/src/arrays/primitive/tests.rs index a0131537d81..8f46f74d4fd 100644 --- a/vortex-array/src/arrays/primitive/tests.rs +++ b/vortex-array/src/arrays/primitive/tests.rs @@ -24,11 +24,12 @@ fn test_search_sorted_primitive( #[case] value: i32, #[case] side: SearchSortedSide, #[case] expected: SearchResult, -) { +) -> vortex_error::VortexResult<()> { let res = array .as_primitive_typed() - .search_sorted(&Some(PValue::from(value)), side); + .search_sorted(&Some(PValue::from(value)), side)?; assert_eq!(res, expected); + Ok(()) } #[test] diff --git a/vortex-array/src/arrays/primitive/vtable/operations.rs b/vortex-array/src/arrays/primitive/vtable/operations.rs index 63a472ec47a..fde30d7cb63 100644 --- a/vortex-array/src/arrays/primitive/vtable/operations.rs +++ b/vortex-array/src/arrays/primitive/vtable/operations.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_dtype::match_each_native_ptype; +use vortex_error::VortexResult; use vortex_scalar::Scalar; use crate::arrays::PrimitiveArray; @@ -9,9 +10,9 @@ use crate::arrays::PrimitiveVTable; use crate::vtable::OperationsVTable; impl OperationsVTable for PrimitiveVTable { - fn scalar_at(array: &PrimitiveArray, index: usize) -> Scalar { - match_each_native_ptype!(array.ptype(), |T| { + fn scalar_at(array: &PrimitiveArray, index: usize) -> VortexResult { + Ok(match_each_native_ptype!(array.ptype(), |T| { Scalar::primitive(array.as_slice::()[index], array.dtype().nullability()) - }) + })) } } diff --git a/vortex-array/src/arrays/scalar_fn/rules.rs b/vortex-array/src/arrays/scalar_fn/rules.rs index a72a7767224..9e3cd96af43 100644 --- a/vortex-array/src/arrays/scalar_fn/rules.rs +++ b/vortex-array/src/arrays/scalar_fn/rules.rs @@ -78,7 +78,7 @@ impl ArrayReduceRule for ScalarFnConstantRule { if array.is_empty() { Ok(Some(Canonical::empty(array.dtype()).into_array())) } else { - let result = array.scalar_at(0); + let result = array.scalar_at(0)?; Ok(Some(ConstantArray::new(result, array.len).into_array())) } } diff --git a/vortex-array/src/arrays/scalar_fn/vtable/operations.rs b/vortex-array/src/arrays/scalar_fn/vtable/operations.rs index 7fde151e2d5..39c3f8b4625 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/operations.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/operations.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_error::VortexExpect; +use vortex_error::VortexResult; use vortex_scalar::Scalar; use crate::Array; @@ -16,12 +16,12 @@ use crate::expr::ExecutionResult; use crate::vtable::OperationsVTable; impl OperationsVTable for ScalarFnVTable { - fn scalar_at(array: &ScalarFnArray, index: usize) -> Scalar { + fn scalar_at(array: &ScalarFnArray, index: usize) -> VortexResult { let inputs: Vec<_> = array .children .iter() - .map(|child| ConstantArray::new(child.scalar_at(index), 1).into_array()) - .collect::<_>(); + .map(|child| Ok(ConstantArray::new(child.scalar_at(index)?, 1).into_array())) + .collect::>()?; let mut ctx = LEGACY_SESSION.create_execution_ctx(); let args = ExecutionArgs { @@ -30,10 +30,7 @@ impl OperationsVTable for ScalarFnVTable { ctx: &mut ctx, }; - let result = array - .scalar_fn - .execute(args) - .vortex_expect("todo vortex result return"); + let result = array.scalar_fn.execute(args)?; let scalar = match result { ExecutionResult::Array(arr) => { @@ -41,7 +38,7 @@ impl OperationsVTable for ScalarFnVTable { "Scalar function {} returned non-constant array from execution over all scalar inputs", array.scalar_fn, ); - arr.as_ref().scalar_at(0) + arr.as_ref().scalar_at(0)? } ExecutionResult::Scalar(constant) => constant.scalar().clone(), }; @@ -55,7 +52,7 @@ impl OperationsVTable for ScalarFnVTable { array.dtype ); - scalar + Ok(scalar) } } @@ -63,49 +60,49 @@ impl OperationsVTable for ScalarFnVTable { mod tests { use vortex_buffer::buffer; use vortex_error::VortexResult; - use vortex_scalar::Scalar; use crate::IntoArray; + use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::scalar_fn::array::ScalarFnArray; + use crate::assert_arrays_eq; use crate::expr::Operator; use crate::expr::ScalarFn; use crate::expr::binary::Binary; + use crate::validity::Validity; #[test] - fn test_scalar_at_add() -> VortexResult<()> { + fn test_scalar_fn_add() -> VortexResult<()> { let lhs = buffer![1i32, 2, 3].into_array(); let rhs = buffer![10i32, 20, 30].into_array(); let scalar_fn = ScalarFn::new(Binary, Operator::Add); let scalar_fn_array = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs], 3)?; - assert_eq!(scalar_fn_array.scalar_at(0), Scalar::from(11i32)); - assert_eq!(scalar_fn_array.scalar_at(1), Scalar::from(22i32)); - assert_eq!(scalar_fn_array.scalar_at(2), Scalar::from(33i32)); + let result = scalar_fn_array.to_canonical()?.into_array(); + let expected = buffer![11i32, 22, 33].into_array(); + assert_arrays_eq!(result, expected); Ok(()) } #[test] - fn test_scalar_at_mul() -> VortexResult<()> { + fn test_scalar_fn_mul() -> VortexResult<()> { let lhs = buffer![2i32, 3, 4].into_array(); let rhs = buffer![5i32, 6, 7].into_array(); let scalar_fn = ScalarFn::new(Binary, Operator::Mul); let scalar_fn_array = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs], 3)?; - assert_eq!(scalar_fn_array.scalar_at(0), Scalar::from(10i32)); - assert_eq!(scalar_fn_array.scalar_at(1), Scalar::from(18i32)); - assert_eq!(scalar_fn_array.scalar_at(2), Scalar::from(28i32)); + let result = scalar_fn_array.to_canonical()?.into_array(); + let expected = buffer![10i32, 18, 28].into_array(); + assert_arrays_eq!(result, expected); Ok(()) } #[test] - fn test_scalar_at_with_nullable() -> VortexResult<()> { - use crate::validity::Validity; - + fn test_scalar_fn_with_nullable() -> VortexResult<()> { let lhs = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::AllValid).into_array(); let rhs = PrimitiveArray::new( buffer![10i32, 20, 30], @@ -116,24 +113,28 @@ mod tests { let scalar_fn = ScalarFn::new(Binary, Operator::Add); let scalar_fn_array = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs], 3)?; - assert_eq!(scalar_fn_array.scalar_at(0), Scalar::from(11i32)); - assert!(scalar_fn_array.scalar_at(1).is_null()); - assert_eq!(scalar_fn_array.scalar_at(2), Scalar::from(33i32)); + let result = scalar_fn_array.to_canonical()?.into_array(); + let expected = PrimitiveArray::new( + buffer![11i32, 0, 33], + Validity::from_iter([true, false, true]), + ) + .into_array(); + assert_arrays_eq!(result, expected); Ok(()) } #[test] - fn test_scalar_at_comparison() -> VortexResult<()> { + fn test_scalar_fn_comparison() -> VortexResult<()> { let lhs = buffer![1i32, 5, 3].into_array(); let rhs = buffer![2i32, 5, 1].into_array(); let scalar_fn = ScalarFn::new(Binary, Operator::Eq); let scalar_fn_array = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs], 3)?; - assert_eq!(scalar_fn_array.scalar_at(0), Scalar::from(false)); - assert_eq!(scalar_fn_array.scalar_at(1), Scalar::from(true)); - assert_eq!(scalar_fn_array.scalar_at(2), Scalar::from(false)); + let result = scalar_fn_array.to_canonical()?.into_array(); + let expected = BoolArray::from_iter([false, true, false]).into_array(); + assert_arrays_eq!(result, expected); Ok(()) } diff --git a/vortex-array/src/arrays/scalar_fn/vtable/validity.rs b/vortex-array/src/arrays/scalar_fn/vtable/validity.rs index 4896575ee4d..97c9a7c2ba4 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/validity.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/validity.rs @@ -3,7 +3,6 @@ use std::sync::Arc; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_mask::Mask; @@ -78,19 +77,12 @@ impl ValidityVTable for ScalarFnVTable { Ok(Validity::Array(execute_expr(&validity_expr, array.len())?)) } - fn validity_mask(array: &ScalarFnArray) -> Mask { + fn validity_mask(array: &ScalarFnArray) -> VortexResult { let mut ctx = LEGACY_SESSION.create_execution_ctx(); - let output = array - .to_array() - .execute::(&mut ctx) - .vortex_expect("Validity mask computation should be fallible"); - match output { + let output = array.to_array().execute::(&mut ctx)?; + Ok(match output { CanonicalOutput::Constant(c) => Mask::new(array.len, c.scalar().is_valid()), - CanonicalOutput::Array(a) => a - .into_array() - .validity() - .vortex_expect("cannot fail") - .to_mask(array.len()), - } + CanonicalOutput::Array(a) => a.into_array().validity()?.to_mask(array.len()), + }) } } diff --git a/vortex-array/src/arrays/slice/vtable.rs b/vortex-array/src/arrays/slice/vtable.rs index d922043202e..95e2969d79f 100644 --- a/vortex-array/src/arrays/slice/vtable.rs +++ b/vortex-array/src/arrays/slice/vtable.rs @@ -157,7 +157,7 @@ impl BaseArrayVTable for SliceVTable { } impl OperationsVTable for SliceVTable { - fn scalar_at(array: &SliceArray, index: usize) -> Scalar { + fn scalar_at(array: &SliceArray, index: usize) -> VortexResult { array.child.scalar_at(array.range.start + index) } } @@ -167,8 +167,8 @@ impl ValidityVTable for SliceVTable { Ok(array.child.validity()?.slice(array.range.clone())) } - fn validity_mask(array: &SliceArray) -> Mask { - array.child.validity_mask().slice(array.range.clone()) + fn validity_mask(array: &SliceArray) -> VortexResult { + Ok(array.child.validity_mask()?.slice(array.range.clone())) } } diff --git a/vortex-array/src/arrays/struct_/array.rs b/vortex-array/src/arrays/struct_/array.rs index 7b39f30c7bb..4f057cbc081 100644 --- a/vortex-array/src/arrays/struct_/array.rs +++ b/vortex-array/src/arrays/struct_/array.rs @@ -61,11 +61,11 @@ use crate::vtable::ValidityHelper; /// ).unwrap(); /// /// // Row 0 is valid - returns a struct scalar with field values -/// let row0 = struct_array.scalar_at(0); +/// let row0 = struct_array.scalar_at(0).unwrap(); /// assert!(!row0.is_null()); /// /// // Row 1 is null at struct level - returns null even though fields have values -/// let row1 = struct_array.scalar_at(1); +/// let row1 = struct_array.scalar_at(1).unwrap(); /// assert!(row1.is_null()); /// ``` /// @@ -94,7 +94,7 @@ use crate::vtable::ValidityHelper; /// /// // field_by_name returns the FIRST "data" field /// let first_data = struct_array.field_by_name("data").unwrap(); -/// assert_eq!(first_data.scalar_at(0), 1i32.into()); +/// assert_eq!(first_data.scalar_at(0).unwrap(), 1i32.into()); /// ``` /// /// ## Field Operations diff --git a/vortex-array/src/arrays/struct_/vtable/operations.rs b/vortex-array/src/arrays/struct_/vtable/operations.rs index 9efbd216e89..f8e6aa22332 100644 --- a/vortex-array/src/arrays/struct_/vtable/operations.rs +++ b/vortex-array/src/arrays/struct_/vtable/operations.rs @@ -1,22 +1,21 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use itertools::Itertools; +use vortex_error::VortexResult; use vortex_scalar::Scalar; +use crate::Array; use crate::arrays::struct_::StructArray; use crate::arrays::struct_::StructVTable; use crate::vtable::OperationsVTable; impl OperationsVTable for StructVTable { - fn scalar_at(array: &StructArray, index: usize) -> Scalar { - Scalar::struct_( - array.dtype().clone(), - array - .fields() - .iter() - .map(|field| field.scalar_at(index)) - .collect_vec(), - ) + fn scalar_at(array: &StructArray, index: usize) -> VortexResult { + let field_scalars: VortexResult> = array + .fields() + .iter() + .map(|field| field.scalar_at(index)) + .collect(); + Ok(Scalar::struct_(array.dtype().clone(), field_scalars?)) } } diff --git a/vortex-array/src/arrays/varbin/array.rs b/vortex-array/src/arrays/varbin/array.rs index 18c8dba69b9..717069b8594 100644 --- a/vortex-array/src/arrays/varbin/array.rs +++ b/vortex-array/src/arrays/varbin/array.rs @@ -149,7 +149,7 @@ impl VarBinArray { } let last_offset = offsets - .scalar_at(offsets.len() - 1) + .scalar_at(offsets.len() - 1)? .as_primitive() .as_::() .ok_or_else(|| vortex_err!("Last offset must be convertible to usize"))?; @@ -287,6 +287,7 @@ impl VarBinArray { self.offsets() .scalar_at(index) + .vortex_expect("offsets must support scalar_at") .as_ref() .try_into() .vortex_expect("Failed to convert offset to usize") diff --git a/vortex-array/src/arrays/varbin/builder.rs b/vortex-array/src/arrays/varbin/builder.rs index ace5428ab72..2c81676be78 100644 --- a/vortex-array/src/arrays/varbin/builder.rs +++ b/vortex-array/src/arrays/varbin/builder.rs @@ -124,9 +124,9 @@ mod tests { assert_eq!(array.len(), 3); assert_eq!(array.dtype().nullability(), Nullable); assert_eq!( - array.scalar_at(0), + array.scalar_at(0).unwrap(), Scalar::utf8("hello".to_string(), Nullable) ); - assert!(array.scalar_at(1).is_null()); + assert!(array.scalar_at(1).unwrap().is_null()); } } diff --git a/vortex-array/src/arrays/varbin/compute/compare.rs b/vortex-array/src/arrays/varbin/compute/compare.rs index 769f45ffceb..ffd658b2726 100644 --- a/vortex-array/src/arrays/varbin/compute/compare.rs +++ b/vortex-array/src/arrays/varbin/compute/compare.rs @@ -170,7 +170,7 @@ mod test { .to_bool(); assert_eq!( - &result.validity_mask().to_bit_buffer(), + &result.validity_mask().unwrap().to_bit_buffer(), &BitBuffer::from_iter([true, false, true]) ); assert_eq!( @@ -194,7 +194,7 @@ mod test { .to_bool(); assert_eq!( - &result.validity_mask().to_bit_buffer(), + &result.validity_mask().unwrap().to_bit_buffer(), &BitBuffer::from_iter([false, false, true]) ); assert_eq!( diff --git a/vortex-array/src/arrays/varbin/compute/filter.rs b/vortex-array/src/arrays/varbin/compute/filter.rs index 883d3f15b44..32f964f3669 100644 --- a/vortex-array/src/arrays/varbin/compute/filter.rs +++ b/vortex-array/src/arrays/varbin/compute/filter.rs @@ -59,7 +59,7 @@ fn filter_select_var_bin_by_slice( offsets.as_slice::(), values.bytes().as_slice(), mask_slices, - values.validity_mask(), + values.validity_mask()?, selection_count, ) }) diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index caa32c451a0..0562a20209e 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -42,64 +42,64 @@ impl TakeKernel for VarBinVTable { offsets.as_slice::(), data.as_slice(), indices.as_slice::(), - array.validity_mask(), - indices.validity_mask(), + array.validity_mask()?, + indices.validity_mask()?, ), PType::U16 => take::( dtype, offsets.as_slice::(), data.as_slice(), indices.as_slice::(), - array.validity_mask(), - indices.validity_mask(), + array.validity_mask()?, + indices.validity_mask()?, ), PType::U32 => take::( dtype, offsets.as_slice::(), data.as_slice(), indices.as_slice::(), - array.validity_mask(), - indices.validity_mask(), + array.validity_mask()?, + indices.validity_mask()?, ), PType::U64 => take::( dtype, offsets.as_slice::(), data.as_slice(), indices.as_slice::(), - array.validity_mask(), - indices.validity_mask(), + array.validity_mask()?, + indices.validity_mask()?, ), PType::I8 => take::( dtype, offsets.as_slice::(), data.as_slice(), indices.as_slice::(), - array.validity_mask(), - indices.validity_mask(), + array.validity_mask()?, + indices.validity_mask()?, ), PType::I16 => take::( dtype, offsets.as_slice::(), data.as_slice(), indices.as_slice::(), - array.validity_mask(), - indices.validity_mask(), + array.validity_mask()?, + indices.validity_mask()?, ), PType::I32 => take::( dtype, offsets.as_slice::(), data.as_slice(), indices.as_slice::(), - array.validity_mask(), - indices.validity_mask(), + array.validity_mask()?, + indices.validity_mask()?, ), PType::I64 => take::( dtype, offsets.as_slice::(), data.as_slice(), indices.as_slice::(), - array.validity_mask(), - indices.validity_mask(), + array.validity_mask()?, + indices.validity_mask()?, ), _ => unreachable!("invalid PType for offsets"), } diff --git a/vortex-array/src/arrays/varbin/vtable/canonical.rs b/vortex-array/src/arrays/varbin/vtable/canonical.rs index d92068b028a..7e8a0f85459 100644 --- a/vortex-array/src/arrays/varbin/vtable/canonical.rs +++ b/vortex-array/src/arrays/varbin/vtable/canonical.rs @@ -70,7 +70,7 @@ mod tests { let canonical = varbin.to_varbinview(); assert_eq!(canonical.dtype(), &dtype); - assert!(!canonical.is_valid(0)); + assert!(!canonical.is_valid(0).unwrap()); // First value is inlined (12 bytes) assert!(canonical.views()[1].is_inlined()); diff --git a/vortex-array/src/arrays/varbin/vtable/operations.rs b/vortex-array/src/arrays/varbin/vtable/operations.rs index 25cd1255222..17f5ebdce1f 100644 --- a/vortex-array/src/arrays/varbin/vtable/operations.rs +++ b/vortex-array/src/arrays/varbin/vtable/operations.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_error::VortexResult; use vortex_scalar::Scalar; use crate::arrays::VarBinArray; @@ -9,7 +10,7 @@ use crate::arrays::varbin_scalar; use crate::vtable::OperationsVTable; impl OperationsVTable for VarBinVTable { - fn scalar_at(array: &VarBinArray, index: usize) -> Scalar { - varbin_scalar(array.bytes_at(index), array.dtype()) + fn scalar_at(array: &VarBinArray, index: usize) -> VortexResult { + Ok(varbin_scalar(array.bytes_at(index), array.dtype())) } } diff --git a/vortex-array/src/arrays/varbinview/compact.rs b/vortex-array/src/arrays/varbinview/compact.rs index 033bdf6082c..8f3808a2659 100644 --- a/vortex-array/src/arrays/varbinview/compact.rs +++ b/vortex-array/src/arrays/varbinview/compact.rs @@ -26,7 +26,7 @@ impl VarBinViewArray { /// well-utilized buffers unchanged. pub fn compact_buffers(&self) -> VortexResult { // If there is nothing to be gained by compaction, return the original array untouched. - if !self.should_compact() { + if !self.should_compact()? { return Ok(self.clone()); } @@ -34,35 +34,35 @@ impl VarBinViewArray { self.compact_with_threshold(1.0) } - fn should_compact(&self) -> bool { + fn should_compact(&self) -> VortexResult { let nbuffers = self.nbuffers(); // If the array is entirely inlined strings, do not attempt to compact. if nbuffers == 0 { - return false; + return Ok(false); } // These will fail to write, so in most cases we want to compact this. if nbuffers > u16::MAX as usize { - return true; + return Ok(true); } - let bytes_referenced: u64 = self.count_referenced_bytes(); + let bytes_referenced: u64 = self.count_referenced_bytes()?; let buffer_total_bytes: u64 = self.buffers.iter().map(|buf| buf.len() as u64).sum(); // If there is any wasted space, we want to repack. // This is very aggressive. - bytes_referenced < buffer_total_bytes || buffer_total_bytes == 0 + Ok(bytes_referenced < buffer_total_bytes || buffer_total_bytes == 0) } /// Iterates over all valid, non-inlined views, calling the provided /// closure for each one. #[inline(always)] - fn iter_valid_views(&self, mut f: F) + fn iter_valid_views(&self, mut f: F) -> VortexResult<()> where F: FnMut(&Ref), { - match self.validity_mask() { + match self.validity_mask()? { Mask::AllTrue(_) => { for &view in self.views().iter() { if !view.is_inlined() { @@ -79,17 +79,18 @@ impl VarBinViewArray { } } } + Ok(()) } /// Count the number of bytes addressed by the views, not including null /// values or any inlined strings. - fn count_referenced_bytes(&self) -> u64 { + fn count_referenced_bytes(&self) -> VortexResult { let mut total = 0u64; - self.iter_valid_views(|view| total += view.size as u64); - total + self.iter_valid_views(|view| total += view.size as u64)?; + Ok(total) } - pub(crate) fn buffer_utilizations(&self) -> Vec { + pub(crate) fn buffer_utilizations(&self) -> VortexResult> { let mut utilizations: Vec = self .buffers() .iter() @@ -101,9 +102,9 @@ impl VarBinViewArray { self.iter_valid_views(|view| { utilizations[view.buffer_index as usize].add(view.offset, view.size); - }); + })?; - utilizations + Ok(utilizations) } /// Returns a compacted copy of the input array using selective buffer compaction. @@ -413,7 +414,7 @@ mod tests { let taken_array = taken.as_::(); // Get buffer stats before compaction - let utils_before = taken_array.buffer_utilizations(); + let utils_before = taken_array.buffer_utilizations().unwrap(); let original_buffer_count = taken_array.nbuffers(); // Compact with a threshold that should trigger slicing @@ -459,9 +460,10 @@ mod tests { #[case] expected_bytes: u64, #[case] expected_utils: &[f64], ) { - assert_eq!(arr.count_referenced_bytes(), expected_bytes); + assert_eq!(arr.count_referenced_bytes().unwrap(), expected_bytes); let utils: Vec = arr .buffer_utilizations() + .unwrap() .iter() .map(|u| u.overall_utilization()) .collect(); diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index 05481fadf74..90475e6698b 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/vortex-array/src/arrays/varbinview/compute/take.rs @@ -30,12 +30,9 @@ impl TakeKernel for VarBinViewVTable { let validity = array.validity().take(indices)?; let indices = indices.to_primitive(); + let indices_mask = indices.validity_mask()?; let views_buffer = match_each_integer_ptype!(indices.ptype(), |I| { - take_views( - array.views(), - indices.as_slice::(), - &indices.validity_mask(), - ) + take_views(array.views(), indices.as_slice::(), &indices_mask) }); // SAFETY: taking all components at same indices maintains invariants diff --git a/vortex-array/src/arrays/varbinview/compute/zip.rs b/vortex-array/src/arrays/varbinview/compute/zip.rs index c1ae7f4f6e2..e5674632715 100644 --- a/vortex-array/src/arrays/varbinview/compute/zip.rs +++ b/vortex-array/src/arrays/varbinview/compute/zip.rs @@ -52,8 +52,8 @@ impl ZipKernel for VarBinViewVTable { let mut views_builder = BufferMut::::with_capacity(len); let mut validity_builder = LazyBitBufferBuilder::new(len); - let true_validity = if_true.validity_mask(); - let false_validity = if_false.validity_mask(); + let true_validity = if_true.validity_mask()?; + let false_validity = if_false.validity_mask()?; match mask.slices() { AllOr::All => push_range( diff --git a/vortex-array/src/arrays/varbinview/vtable/operations.rs b/vortex-array/src/arrays/varbinview/vtable/operations.rs index 4eb88976376..2288d4dc381 100644 --- a/vortex-array/src/arrays/varbinview/vtable/operations.rs +++ b/vortex-array/src/arrays/varbinview/vtable/operations.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_error::VortexResult; use vortex_scalar::Scalar; use crate::arrays::VarBinViewArray; @@ -9,7 +10,7 @@ use crate::arrays::varbin_scalar; use crate::vtable::OperationsVTable; impl OperationsVTable for VarBinViewVTable { - fn scalar_at(array: &VarBinViewArray, index: usize) -> Scalar { - varbin_scalar(array.bytes_at(index), array.dtype()) + fn scalar_at(array: &VarBinViewArray, index: usize) -> VortexResult { + Ok(varbin_scalar(array.bytes_at(index), array.dtype())) } } diff --git a/vortex-array/src/arrow/array.rs b/vortex-array/src/arrow/array.rs index 04cea366592..8df97b55ccf 100644 --- a/vortex-array/src/arrow/array.rs +++ b/vortex-array/src/arrow/array.rs @@ -12,7 +12,6 @@ use vortex_dtype::arrow::FromArrowType; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; -use vortex_error::vortex_panic; use vortex_mask::Mask; use vortex_scalar::Scalar; @@ -150,8 +149,8 @@ impl BaseArrayVTable for ArrowVTable { } impl OperationsVTable for ArrowVTable { - fn scalar_at(_array: &ArrowArray, _index: usize) -> Scalar { - vortex_panic!("Not supported") + fn scalar_at(_array: &ArrowArray, _index: usize) -> VortexResult { + vortex_bail!("ArrowArray does not support scalar_at") } } @@ -173,12 +172,12 @@ impl ValidityVTable for ArrowVTable { }) } - fn validity_mask(array: &ArrowArray) -> Mask { - array + fn validity_mask(array: &ArrowArray) -> VortexResult { + Ok(array .inner .logical_nulls() .map(|null_buffer| Mask::from_buffer(null_buffer.inner().clone().into())) - .unwrap_or_else(|| Mask::new_true(array.inner.len())) + .unwrap_or_else(|| Mask::new_true(array.inner.len()))) } } diff --git a/vortex-array/src/arrow/datum.rs b/vortex-array/src/arrow/datum.rs index 80fcb3c3ab3..e40882f2425 100644 --- a/vortex-array/src/arrow/datum.rs +++ b/vortex-array/src/arrow/datum.rs @@ -5,6 +5,7 @@ use arrow_array::Array as ArrowArray; use arrow_array::ArrayRef as ArrowArrayRef; use arrow_array::Datum as ArrowDatum; use arrow_schema::DataType; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_panic; @@ -101,5 +102,11 @@ where ); } - ConstantArray::new(array.scalar_at(0), len).into_array() + ConstantArray::new( + array + .scalar_at(0) + .vortex_expect("array of length 1 must support scalar_at(0)"), + len, + ) + .into_array() } diff --git a/vortex-array/src/arrow/executor/bool.rs b/vortex-array/src/arrow/executor/bool.rs index b7764b558d6..7b6be7893e1 100644 --- a/vortex-array/src/arrow/executor/bool.rs +++ b/vortex-array/src/arrow/executor/bool.rs @@ -13,11 +13,11 @@ use crate::arrays::BoolArray; use crate::arrow::null_buffer::to_null_buffer; /// Convert a canonical BoolArray directly to Arrow. -pub fn canonical_bool_to_arrow(array: &BoolArray) -> ArrowArrayRef { - Arc::new(ArrowBooleanArray::new( +pub fn canonical_bool_to_arrow(array: &BoolArray) -> VortexResult { + Ok(Arc::new(ArrowBooleanArray::new( array.bit_buffer().clone().into(), - to_null_buffer(array.validity_mask()), - )) + to_null_buffer(array.validity_mask()?), + ))) } pub(super) fn to_arrow_bool( @@ -25,5 +25,5 @@ pub(super) fn to_arrow_bool( ctx: &mut ExecutionCtx, ) -> VortexResult { let bool_array = array.execute::(ctx)?; - Ok(canonical_bool_to_arrow(&bool_array)) + canonical_bool_to_arrow(&bool_array) } diff --git a/vortex-array/src/arrow/executor/byte_view.rs b/vortex-array/src/arrow/executor/byte_view.rs index ec6aea2fe97..2171a558c30 100644 --- a/vortex-array/src/arrow/executor/byte_view.rs +++ b/vortex-array/src/arrow/executor/byte_view.rs @@ -21,7 +21,9 @@ use crate::builtins::ArrayBuiltins; use crate::vtable::ValidityHelper; /// Convert a canonical VarBinViewArray directly to Arrow. -pub fn canonical_varbinview_to_arrow(array: &VarBinViewArray) -> ArrowArrayRef { +pub fn canonical_varbinview_to_arrow( + array: &VarBinViewArray, +) -> VortexResult { let views = ScalarBuffer::::from(array.views().clone().into_byte_buffer().into_arrow_buffer()); let buffers: Vec<_> = array @@ -29,10 +31,12 @@ pub fn canonical_varbinview_to_arrow(array: &VarBinViewArray) - .iter() .map(|buffer| buffer.clone().into_arrow_buffer()) .collect(); - let nulls = to_null_buffer(array.validity_mask()); + let nulls = to_null_buffer(array.validity_mask()?); // SAFETY: our own VarBinView array is considered safe. - Arc::new(unsafe { GenericByteViewArray::::new_unchecked(views, buffers, nulls) }) + Ok(Arc::new(unsafe { + GenericByteViewArray::::new_unchecked(views, buffers, nulls) + })) } pub fn execute_varbinview_to_arrow( @@ -65,5 +69,5 @@ pub(super) fn to_arrow_byte_view( let array = array.cast(DType::from_arrow((&T::DATA_TYPE, Nullability::Nullable)))?; let varbinview = array.execute::(ctx)?; - Ok(canonical_varbinview_to_arrow::(&varbinview)) + canonical_varbinview_to_arrow::(&varbinview) } diff --git a/vortex-array/src/arrow/executor/decimal.rs b/vortex-array/src/arrow/executor/decimal.rs index 65385136090..f6b37657d06 100644 --- a/vortex-array/src/arrow/executor/decimal.rs +++ b/vortex-array/src/arrow/executor/decimal.rs @@ -41,7 +41,7 @@ pub(super) fn to_arrow_decimal( } fn to_arrow_decimal32(array: DecimalArray) -> VortexResult { - let null_buffer = to_null_buffer(array.validity_mask()); + let null_buffer = to_null_buffer(array.validity_mask()?); let buffer: Buffer = match array.values_type() { DecimalType::I8 => { Buffer::from_trusted_len_iter(array.buffer::().into_iter().map(|x| x.as_())) @@ -85,7 +85,7 @@ fn to_arrow_decimal32(array: DecimalArray) -> VortexResult { } fn to_arrow_decimal64(array: DecimalArray) -> VortexResult { - let null_buffer = to_null_buffer(array.validity_mask()); + let null_buffer = to_null_buffer(array.validity_mask()?); let buffer: Buffer = match array.values_type() { DecimalType::I8 => { Buffer::from_trusted_len_iter(array.buffer::().into_iter().map(|x| x.as_())) @@ -124,7 +124,7 @@ fn to_arrow_decimal64(array: DecimalArray) -> VortexResult { } fn to_arrow_decimal128(array: DecimalArray) -> VortexResult { - let null_buffer = to_null_buffer(array.validity_mask()); + let null_buffer = to_null_buffer(array.validity_mask()?); let buffer: Buffer = match array.values_type() { DecimalType::I8 => { Buffer::from_trusted_len_iter(array.buffer::().into_iter().map(|x| x.as_())) @@ -158,7 +158,7 @@ fn to_arrow_decimal128(array: DecimalArray) -> VortexResult { } fn to_arrow_decimal256(array: DecimalArray) -> VortexResult { - let null_buffer = to_null_buffer(array.validity_mask()); + let null_buffer = to_null_buffer(array.validity_mask()?); let buffer: Buffer = match array.values_type() { DecimalType::I8 => { Buffer::from_trusted_len_iter(array.buffer::().into_iter().map(|x| x.as_())) diff --git a/vortex-array/src/arrow/executor/list.rs b/vortex-array/src/arrow/executor/list.rs index adab81da60c..ec3d0c1e882 100644 --- a/vortex-array/src/arrow/executor/list.rs +++ b/vortex-array/src/arrow/executor/list.rs @@ -126,7 +126,7 @@ fn list_view_zctl( // For ZCTL, we know that we only care about the final size. let final_size = sizes - .scalar_at(sizes.len() - 1) + .scalar_at(sizes.len() - 1)? .cast(&DType::Primitive(O::PTYPE, Nullability::NonNullable))?; let final_size = final_size .as_primitive() diff --git a/vortex-array/src/arrow/executor/primitive.rs b/vortex-array/src/arrow/executor/primitive.rs index f621cca6af3..87a221fb8fb 100644 --- a/vortex-array/src/arrow/executor/primitive.rs +++ b/vortex-array/src/arrow/executor/primitive.rs @@ -18,14 +18,16 @@ use crate::arrow::null_buffer::to_null_buffer; use crate::builtins::ArrayBuiltins; /// Convert a canonical PrimitiveArray directly to Arrow. -pub fn canonical_primitive_to_arrow(array: PrimitiveArray) -> ArrowArrayRef +pub fn canonical_primitive_to_arrow( + array: PrimitiveArray, +) -> VortexResult where T::Native: NativePType, { - let validity = array.validity_mask(); + let validity = array.validity_mask()?; let null_buffer = to_null_buffer(validity); let buffer = array.into_buffer::().into_arrow_scalar_buffer(); - Arc::new(ArrowPrimitiveArray::::new(buffer, null_buffer)) + Ok(Arc::new(ArrowPrimitiveArray::::new(buffer, null_buffer))) } pub(super) fn to_arrow_primitive( @@ -38,5 +40,5 @@ where // We use nullable here so we can essentially ignore nullability during the cast. let array = array.cast(DType::Primitive(T::Native::PTYPE, Nullability::Nullable))?; let primitive = array.execute::(ctx)?; - Ok(canonical_primitive_to_arrow::(primitive)) + canonical_primitive_to_arrow::(primitive) } diff --git a/vortex-array/src/arrow/executor/temporal.rs b/vortex-array/src/arrow/executor/temporal.rs index 4677a62f48a..939ec613d61 100644 --- a/vortex-array/src/arrow/executor/temporal.rs +++ b/vortex-array/src/arrow/executor/temporal.rs @@ -135,7 +135,7 @@ where primitive.ptype() ); - let validity = primitive.validity_mask(); + let validity = primitive.validity_mask()?; let buffer = primitive.to_buffer::(); let values = buffer.into_arrow_scalar_buffer(); diff --git a/vortex-array/src/arrow/record_batch.rs b/vortex-array/src/arrow/record_batch.rs index 37149e30bc9..17a6f303293 100644 --- a/vortex-array/src/arrow/record_batch.rs +++ b/vortex-array/src/arrow/record_batch.rs @@ -27,7 +27,7 @@ impl TryFrom<&dyn Array> for RecordBatch { }; vortex_ensure!( - struct_array.all_valid(), + struct_array.all_valid()?, "RecordBatch can only be constructed from StructArray with no nulls" ); diff --git a/vortex-array/src/builders/bool.rs b/vortex-array/src/builders/bool.rs index b2e500e68e1..167f32621d3 100644 --- a/vortex-array/src/builders/bool.rs +++ b/vortex-array/src/builders/bool.rs @@ -7,6 +7,7 @@ use std::mem; use vortex_buffer::BitBufferMut; use vortex_dtype::DType; use vortex_dtype::Nullability; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_mask::Mask; @@ -117,7 +118,11 @@ impl ArrayBuilder for BoolBuilder { let bool_array = array.to_bool(); self.inner.append_buffer(bool_array.bit_buffer()); - self.nulls.append_validity_mask(bool_array.validity_mask()); + self.nulls.append_validity_mask( + bool_array + .validity_mask() + .vortex_expect("validity_mask in extend_from_array_unchecked"), + ); } fn reserve_exact(&mut self, additional: usize) { diff --git a/vortex-array/src/builders/decimal.rs b/vortex-array/src/builders/decimal.rs index 4f504d40cab..beefe43a3ab 100644 --- a/vortex-array/src/builders/decimal.rs +++ b/vortex-array/src/builders/decimal.rs @@ -197,8 +197,11 @@ impl ArrayBuilder for DecimalBuilder { .extend(decimal_array.buffer::().iter().copied()); }); - self.nulls - .append_validity_mask(decimal_array.validity_mask()); + self.nulls.append_validity_mask( + decimal_array + .validity_mask() + .vortex_expect("validity_mask in extend_from_array_unchecked"), + ); } fn reserve_exact(&mut self, additional: usize) { @@ -314,7 +317,7 @@ mod tests { let i128s = i128s.finish(); for i in 0..i8s.len() { - assert_eq!(i8s.scalar_at(i), i128s.scalar_at(i)); + assert_eq!(i8s.scalar_at(i).unwrap(), i128s.scalar_at(i).unwrap()); } } @@ -338,7 +341,7 @@ mod tests { // Test by taking a scalar from the array and appending it to a new builder. let mut builder2 = DecimalBuilder::new::(DecimalDType::new(10, 2), true.into()); for i in 0..array.len() { - let scalar = array.scalar_at(i); + let scalar = array.scalar_at(i).unwrap(); builder2.append_scalar(&scalar).unwrap(); } diff --git a/vortex-array/src/builders/fixed_size_list.rs b/vortex-array/src/builders/fixed_size_list.rs index 149db4f7c90..d2070e0c1b7 100644 --- a/vortex-array/src/builders/fixed_size_list.rs +++ b/vortex-array/src/builders/fixed_size_list.rs @@ -241,7 +241,11 @@ impl ArrayBuilder for FixedSizeListBuilder { } self.elements_builder.extend_from_array(fsl.elements()); - self.nulls.append_validity_mask(array.validity_mask()); + self.nulls.append_validity_mask( + array + .validity_mask() + .vortex_expect("validity_mask in extend_from_array_unchecked"), + ); } fn reserve_exact(&mut self, additional: usize) { @@ -771,7 +775,7 @@ mod tests { // Check actual values using scalar_at. - let scalar0 = array.scalar_at(0); + let scalar0 = array.scalar_at(0).unwrap(); let list0 = scalar0.as_list(); assert_eq!(list0.len(), 2); if let Some(list0_items) = list0.elements() { @@ -779,7 +783,7 @@ mod tests { assert_eq!(list0_items[1].as_primitive().typed_value::(), Some(2)); } - let scalar1 = array.scalar_at(1); + let scalar1 = array.scalar_at(1).unwrap(); let list1 = scalar1.as_list(); assert_eq!(list1.len(), 2); if let Some(list1_items) = list1.elements() { diff --git a/vortex-array/src/builders/list.rs b/vortex-array/src/builders/list.rs index b93aa403c86..1d1aa25f203 100644 --- a/vortex-array/src/builders/list.rs +++ b/vortex-array/src/builders/list.rs @@ -221,7 +221,11 @@ impl ArrayBuilder for ListBuilder { } // Append validity information. - self.nulls.append_validity_mask(array.validity_mask()); + self.nulls.append_validity_mask( + array + .validity_mask() + .vortex_expect("validity_mask in extend_from_array_unchecked"), + ); // Note that `ListViewArray` has `n` offsets and sizes, not `n+1` offsets like `ListArray`. let elements = list.elements(); @@ -498,10 +502,13 @@ mod tests { let canon_values = chunked_list.unwrap().to_listview(); assert_eq!( - one_trailing_unused_element.scalar_at(0), - canon_values.scalar_at(0) + one_trailing_unused_element.scalar_at(0).unwrap(), + canon_values.scalar_at(0).unwrap() + ); + assert_eq!( + second_array.scalar_at(0).unwrap(), + canon_values.scalar_at(1).unwrap() ); - assert_eq!(second_array.scalar_at(0), canon_values.scalar_at(1)); } #[test] @@ -530,7 +537,7 @@ mod tests { // Check actual values using scalar_at. - let scalar0 = array.scalar_at(0); + let scalar0 = array.scalar_at(0).unwrap(); let list0 = scalar0.as_list(); assert_eq!(list0.len(), 2); if let Some(list0_items) = list0.elements() { @@ -538,7 +545,7 @@ mod tests { assert_eq!(list0_items[1].as_primitive().typed_value::(), Some(2)); } - let scalar1 = array.scalar_at(1); + let scalar1 = array.scalar_at(1).unwrap(); let list1 = scalar1.as_list(); assert_eq!(list1.len(), 3); if let Some(list1_items) = list1.elements() { @@ -547,7 +554,7 @@ mod tests { assert_eq!(list1_items[2].as_primitive().typed_value::(), Some(5)); } - let scalar2 = array.scalar_at(2); + let scalar2 = array.scalar_at(2).unwrap(); let list2 = scalar2.as_list(); assert!(list2.is_null()); // This should be null. diff --git a/vortex-array/src/builders/listview.rs b/vortex-array/src/builders/listview.rs index 88e051235b6..88f4f3e4fbe 100644 --- a/vortex-array/src/builders/listview.rs +++ b/vortex-array/src/builders/listview.rs @@ -300,7 +300,9 @@ impl ArrayBuilder for ListViewBuilder { // to manually append each scalar. if !listview.is_zero_copy_to_list() { for i in 0..listview.len() { - let list = listview.scalar_at(i); + let list = listview + .scalar_at(i) + .vortex_expect("scalar_at failed in extend_from_array_unchecked"); self.append_scalar(&list) .vortex_expect("was unable to extend the `ListViewBuilder`") @@ -316,7 +318,11 @@ impl ArrayBuilder for ListViewBuilder { .vortex_expect("ListViewArray::rebuild(MakeExact) failed in extend_from_array"); debug_assert!(listview.is_zero_copy_to_list()); - self.nulls.append_validity_mask(array.validity_mask()); + self.nulls.append_validity_mask( + array + .validity_mask() + .vortex_expect("validity_mask in extend_from_array_unchecked"), + ); // Bulk append the new elements (which should have no gaps or overlaps). let old_elements_len = self.elements_builder.len(); diff --git a/vortex-array/src/builders/mod.rs b/vortex-array/src/builders/mod.rs index f9ce9e2e687..9307c651b02 100644 --- a/vortex-array/src/builders/mod.rs +++ b/vortex-array/src/builders/mod.rs @@ -22,10 +22,10 @@ //! //! let strings = builder.finish(); //! -//! assert_eq!(strings.scalar_at(0), "a".into()); -//! assert_eq!(strings.scalar_at(1), "b".into()); -//! assert_eq!(strings.scalar_at(2), "c".into()); -//! assert_eq!(strings.scalar_at(3), "d".into()); +//! assert_eq!(strings.scalar_at(0).unwrap(), "a".into()); +//! assert_eq!(strings.scalar_at(1).unwrap(), "b".into()); +//! assert_eq!(strings.scalar_at(2).unwrap(), "c".into()); +//! assert_eq!(strings.scalar_at(3).unwrap(), "d".into()); //! ``` use std::any::Any; @@ -238,10 +238,10 @@ pub trait ArrayBuilder: Send { /// /// let strings = builder.finish(); /// -/// assert_eq!(strings.scalar_at(0), "a".into()); -/// assert_eq!(strings.scalar_at(1), "b".into()); -/// assert_eq!(strings.scalar_at(2), "c".into()); -/// assert_eq!(strings.scalar_at(3), "d".into()); +/// assert_eq!(strings.scalar_at(0).unwrap(), "a".into()); +/// assert_eq!(strings.scalar_at(1).unwrap(), "b".into()); +/// assert_eq!(strings.scalar_at(2).unwrap(), "c".into()); +/// assert_eq!(strings.scalar_at(3).unwrap(), "d".into()); /// ``` pub fn builder_with_capacity(dtype: &DType, capacity: usize) -> Box { match dtype { diff --git a/vortex-array/src/builders/primitive.rs b/vortex-array/src/builders/primitive.rs index b3a6652e404..16e4f871622 100644 --- a/vortex-array/src/builders/primitive.rs +++ b/vortex-array/src/builders/primitive.rs @@ -8,6 +8,7 @@ use vortex_buffer::BufferMut; use vortex_dtype::DType; use vortex_dtype::NativePType; use vortex_dtype::Nullability; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_mask::Mask; @@ -175,7 +176,11 @@ impl ArrayBuilder for PrimitiveBuilder { ); self.values.extend_from_slice(array.as_slice::()); - self.nulls.append_validity_mask(array.validity_mask()); + self.nulls.append_validity_mask( + array + .validity_mask() + .vortex_expect("validity_mask in extend_from_array_unchecked"), + ); } fn reserve_exact(&mut self, additional: usize) { @@ -414,9 +419,9 @@ mod tests { let array = builder.finish_into_primitive(); assert_eq!(array.len(), 3); // Check validity using scalar_at - nulls will return is_null() = true. - assert!(!array.scalar_at(0).is_null()); - assert!(array.scalar_at(1).is_null()); - assert!(!array.scalar_at(2).is_null()); + assert!(!array.scalar_at(0).unwrap().is_null()); + assert!(array.scalar_at(1).unwrap().is_null()); + assert!(!array.scalar_at(2).unwrap().is_null()); } /// REGRESSION TEST: This test verifies that `append_mask` validates the mask length. @@ -504,13 +509,13 @@ mod tests { assert_eq!(array.as_slice::(), &[100, 200, 10, 20, 30]); // Check validity - the first two should be valid (from append_value). - assert!(!array.scalar_at(0).is_null()); // initial value 100 - assert!(!array.scalar_at(1).is_null()); // initial value 200 + assert!(!array.scalar_at(0).unwrap().is_null()); // initial value 100 + assert!(!array.scalar_at(1).unwrap().is_null()); // initial value 200 // Check the range items with modified validity. - assert!(!array.scalar_at(2).is_null()); // range index 0 - set to valid - assert!(array.scalar_at(3).is_null()); // range index 1 - left as null - assert!(!array.scalar_at(4).is_null()); // range index 2 - set to valid + assert!(!array.scalar_at(2).unwrap().is_null()); // range index 0 - set to valid + assert!(array.scalar_at(3).unwrap().is_null()); // range index 1 - left as null + assert!(!array.scalar_at(4).unwrap().is_null()); // range index 2 - set to valid } /// Test that creating a zero-length uninit range panics. diff --git a/vortex-array/src/builders/struct_.rs b/vortex-array/src/builders/struct_.rs index 18de99bce31..73e6ac5bb65 100644 --- a/vortex-array/src/builders/struct_.rs +++ b/vortex-array/src/builders/struct_.rs @@ -173,7 +173,11 @@ impl ArrayBuilder for StructBuilder { builder.extend_from_array(a.as_ref()); } - self.nulls.append_validity_mask(array.validity_mask()); + self.nulls.append_validity_mask( + array + .validity_mask() + .vortex_expect("validity_mask in extend_from_array_unchecked"), + ); } fn reserve_exact(&mut self, capacity: usize) { @@ -244,7 +248,7 @@ mod tests { let struct_ = builder.finish(); assert_eq!(struct_.len(), 3); assert_eq!(struct_.dtype(), &dtype); - assert_eq!(struct_.valid_count(), 1); + assert_eq!(struct_.valid_count().unwrap(), 1); } #[test] diff --git a/vortex-array/src/builders/tests.rs b/vortex-array/src/builders/tests.rs index 9b898054208..79db5a8c10b 100644 --- a/vortex-array/src/builders/tests.rs +++ b/vortex-array/src/builders/tests.rs @@ -97,8 +97,8 @@ fn test_append_zeros_matches_default_value(#[case] dtype: DType) { // Compare each element. for i in 0..num_elements { - let scalar_zeros = array_zeros.scalar_at(i); - let scalar_manual = array_manual.scalar_at(i); + let scalar_zeros = array_zeros.scalar_at(i).unwrap(); + let scalar_manual = array_manual.scalar_at(i).unwrap(); assert_eq!( scalar_zeros, scalar_manual, @@ -203,7 +203,7 @@ fn test_append_defaults_behavior(#[case] dtype: DType, #[case] should_be_null: b assert_eq!(array.len(), 3); for i in 0..3 { - let scalar = array.scalar_at(i); + let scalar = array.scalar_at(i).unwrap(); if should_be_null { assert!(scalar.is_null(), "Element at index {} should be null", i); } else { @@ -258,8 +258,8 @@ where // Compare each element. for i in 0..array_direct.len() { - let scalar_direct = array_direct.scalar_at(i); - let scalar_indirect = array_indirect.scalar_at(i); + let scalar_direct = array_direct.scalar_at(i).unwrap(); + let scalar_indirect = array_indirect.scalar_at(i).unwrap(); assert_eq!( scalar_direct, scalar_indirect, @@ -553,13 +553,13 @@ fn test_append_scalar_comprehensive(#[case] dtype: DType) { // Verify each scalar matches. for (i, expected_scalar) in scalars.iter().enumerate() { - let actual_scalar = array.scalar_at(i); + let actual_scalar = array.scalar_at(i).unwrap(); assert_scalars_equal(&actual_scalar, expected_scalar, &dtype, i); } // If nullable, verify the last element is null. if dtype.is_nullable() { - let null_scalar = array.scalar_at(num_elements); + let null_scalar = array.scalar_at(num_elements).unwrap(); assert!( null_scalar.is_null(), "Last element should be null for nullable dtype" @@ -704,16 +704,16 @@ fn test_append_scalar_mixed_nulls(#[case] dtype: DType) { assert_eq!(array.len(), 5); // Check the pattern. - assert!(!array.scalar_at(0).is_null()); - assert!(array.scalar_at(1).is_null()); - assert!(!array.scalar_at(2).is_null()); - assert!(array.scalar_at(3).is_null()); - assert!(!array.scalar_at(4).is_null()); + assert!(!array.scalar_at(0).unwrap().is_null()); + assert!(array.scalar_at(1).unwrap().is_null()); + assert!(!array.scalar_at(2).unwrap().is_null()); + assert!(array.scalar_at(3).unwrap().is_null()); + assert!(!array.scalar_at(4).unwrap().is_null()); // Verify non-null values match. - assert_scalars_equal(&array.scalar_at(0), &test_scalars[0], &dtype, 0); - assert_scalars_equal(&array.scalar_at(2), &test_scalars[1], &dtype, 2); - assert_scalars_equal(&array.scalar_at(4), &test_scalars[2], &dtype, 4); + assert_scalars_equal(&array.scalar_at(0).unwrap(), &test_scalars[0], &dtype, 0); + assert_scalars_equal(&array.scalar_at(2).unwrap(), &test_scalars[1], &dtype, 2); + assert_scalars_equal(&array.scalar_at(4).unwrap(), &test_scalars[2], &dtype, 4); } /// Test that `append_scalar` correctly rejects scalars with wrong dtype. @@ -764,7 +764,7 @@ fn test_append_scalar_repeated_same_instance() { // All values should be 42. for i in 0..5 { - let actual = array.scalar_at(i); + let actual = array.scalar_at(i).unwrap(); assert_eq!( actual.as_primitive().typed_value::(), Some(42), diff --git a/vortex-array/src/builders/varbinview.rs b/vortex-array/src/builders/varbinview.rs index 42c9b805e8b..384dbb10c33 100644 --- a/vortex-array/src/builders/varbinview.rs +++ b/vortex-array/src/builders/varbinview.rs @@ -278,7 +278,11 @@ impl ArrayBuilder for VarBinViewBuilder { let array = array.to_varbinview(); self.flush_in_progress(); - self.push_only_validity_mask(array.validity_mask()); + self.push_only_validity_mask( + array + .validity_mask() + .vortex_expect("validity_mask in extend_from_array_unchecked"), + ); let view_adjustment = self.completed @@ -294,7 +298,10 @@ impl ArrayBuilder for VarBinViewBuilder { .iter() .map(|view| adjustment.adjust_view(view)), ), - ViewAdjustment::Rewriting(adjustment) => match array.validity_mask() { + ViewAdjustment::Rewriting(adjustment) => match array + .validity_mask() + .vortex_expect("validity_mask in extend_from_array_unchecked") + { Mask::AllTrue(_) => { for (idx, &view) in array.views().iter().enumerate() { let new_view = self.push_view(view, &adjustment, &array, idx); @@ -587,7 +594,9 @@ impl BuffersWithOffsets { }; } - let buffer_utilizations = array.buffer_utilizations(); + let buffer_utilizations = array + .buffer_utilizations() + .vortex_expect("buffer_utilizations in BuffersWithOffsets::from_array"); let mut has_rewrite = false; let mut has_nonzero_offset = false; for utilization in buffer_utilizations.iter() { @@ -1010,7 +1019,7 @@ mod tests { assert_eq!(array.len(), 1); // Verify the value was stored correctly - let retrieved = array.scalar_at(0).as_binary().value().unwrap(); + let retrieved = array.scalar_at(0).unwrap().as_binary().value().unwrap(); assert_eq!(retrieved.len(), 8192); assert_eq!(retrieved.as_slice(), &large_value); } diff --git a/vortex-array/src/canonical_to_vector.rs b/vortex-array/src/canonical_to_vector.rs index 9e8bcd04d7a..cb8068b3ba5 100644 --- a/vortex-array/src/canonical_to_vector.rs +++ b/vortex-array/src/canonical_to_vector.rs @@ -51,11 +51,11 @@ impl Canonical { Ok(match self { Canonical::Null(a) => Vector::Null(NullVector::new(a.len())), Canonical::Bool(a) => { - Vector::Bool(BoolVector::new(a.bit_buffer().clone(), a.validity_mask())) + Vector::Bool(BoolVector::new(a.bit_buffer().clone(), a.validity_mask()?)) } Canonical::Primitive(a) => { let ptype = a.ptype(); - let validity = a.validity_mask(); + let validity = a.validity_mask()?; match_each_native_ptype!(ptype, |T| { let buffer = a.to_buffer::(); Vector::Primitive(PVector::::new(buffer, validity).into()) @@ -72,7 +72,7 @@ impl Canonical { match_each_decimal_value_type!(min_value_type, |E| { let decimal_dtype = a.decimal_dtype(); let buffer = a.buffer::(); - let validity_mask = a.validity_mask(); + let validity_mask = a.validity_mask()?; // Copy from D to E, possibly widening, possibly narrowing let values = Buffer::::from_trusted_len_iter(buffer.iter().map(|d| { @@ -97,7 +97,7 @@ impl Canonical { }) } Canonical::VarBinView(a) => { - let validity = a.validity_mask(); + let validity = a.validity_mask()?; match a.dtype() { DType::Utf8(_) => { let views = a.views().clone(); @@ -137,9 +137,9 @@ impl Canonical { match_each_native_ptype!(offsets_ptype, |O| { match_each_native_ptype!(sizes_ptype, |S| { let offsets_vec = - PVector::::new(offsets.to_buffer::(), offsets.validity_mask()); + PVector::::new(offsets.to_buffer::(), offsets.validity_mask()?); let sizes_vec = - PVector::::new(sizes.to_buffer::(), sizes.validity_mask()); + PVector::::new(sizes.to_buffer::(), sizes.validity_mask()?); Vector::List(unsafe { ListViewVector::new_unchecked( Arc::new(elements_vector), @@ -152,7 +152,7 @@ impl Canonical { }) } Canonical::FixedSizeList(a) => { - let validity = a.validity_mask(); + let validity = a.validity_mask()?; let list_size = a.list_size(); let elements_vector = a.elements().clone().execute::(ctx)?; Vector::FixedSizeList(unsafe { @@ -164,7 +164,7 @@ impl Canonical { }) } Canonical::Struct(a) => { - let validity = a.validity_mask(); + let validity = a.validity_mask()?; let mut fields = Vec::with_capacity(a.fields().len()); for f in a.fields().iter().cloned() { fields.push(f.execute::(ctx)?); diff --git a/vortex-array/src/compute/between.rs b/vortex-array/src/compute/between.rs index 28679bfdf24..8c667e52d2f 100644 --- a/vortex-array/src/compute/between.rs +++ b/vortex-array/src/compute/between.rs @@ -123,7 +123,7 @@ impl ComputeFnVTable for Between { // A quick check to see if either array might is a null constant array. // Note: Depends on returning early if array is empty for is_invalid check. - if (lower.is_invalid(0) || upper.is_invalid(0)) + if (lower.is_invalid(0)? || upper.is_invalid(0)?) && let (Some(c_lower), Some(c_upper)) = (lower.as_constant(), upper.as_constant()) && (c_lower.is_null() || c_upper.is_null()) { diff --git a/vortex-array/src/compute/boolean.rs b/vortex-array/src/compute/boolean.rs index 353f8afc2dc..5b142cf8e27 100644 --- a/vortex-array/src/compute/boolean.rs +++ b/vortex-array/src/compute/boolean.rs @@ -339,10 +339,10 @@ mod tests { let r = r.to_bool().into_array(); - let v0 = r.scalar_at(0).as_bool().value(); - let v1 = r.scalar_at(1).as_bool().value(); - let v2 = r.scalar_at(2).as_bool().value(); - let v3 = r.scalar_at(3).as_bool().value(); + let v0 = r.scalar_at(0).unwrap().as_bool().value(); + let v1 = r.scalar_at(1).unwrap().as_bool().value(); + let v2 = r.scalar_at(2).unwrap().as_bool().value(); + let v3 = r.scalar_at(3).unwrap().as_bool().value(); assert!(v0.unwrap()); assert!(v1.unwrap()); @@ -359,10 +359,10 @@ mod tests { fn test_and(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) { let r = and(&lhs, &rhs).unwrap().to_bool().into_array(); - let v0 = r.scalar_at(0).as_bool().value(); - let v1 = r.scalar_at(1).as_bool().value(); - let v2 = r.scalar_at(2).as_bool().value(); - let v3 = r.scalar_at(3).as_bool().value(); + let v0 = r.scalar_at(0).unwrap().as_bool().value(); + let v1 = r.scalar_at(1).unwrap().as_bool().value(); + let v2 = r.scalar_at(2).unwrap().as_bool().value(); + let v3 = r.scalar_at(3).unwrap().as_bool().value(); assert!(v0.unwrap()); assert!(!v1.unwrap()); diff --git a/vortex-array/src/compute/compare.rs b/vortex-array/src/compute/compare.rs index 1345291c926..701df2d6c1d 100644 --- a/vortex-array/src/compute/compare.rs +++ b/vortex-array/src/compute/compare.rs @@ -671,9 +671,9 @@ mod tests { // Compare two lists together let result = compare(list.as_ref(), list.as_ref(), Operator::Eq).unwrap(); - assert!(result.scalar_at(0).is_valid()); - assert!(result.scalar_at(1).is_valid()); - assert!(result.scalar_at(2).is_valid()); + assert!(result.scalar_at(0).unwrap().is_valid()); + assert!(result.scalar_at(1).unwrap().is_valid()); + assert!(result.scalar_at(2).unwrap().is_valid()); } #[test] diff --git a/vortex-array/src/compute/conformance/binary_numeric.rs b/vortex-array/src/compute/conformance/binary_numeric.rs index 42b99befb5e..121dfaf7777 100644 --- a/vortex-array/src/compute/conformance/binary_numeric.rs +++ b/vortex-array/src/compute/conformance/binary_numeric.rs @@ -45,7 +45,11 @@ use crate::compute::numeric::numeric; fn to_vec_of_scalar(array: &dyn Array) -> Vec { // Not fast, but obviously correct (0..array.len()) - .map(|index| array.scalar_at(index)) + .map(|index| { + array + .scalar_at(index) + .vortex_expect("scalar_at should succeed in conformance test") + }) .collect_vec() } diff --git a/vortex-array/src/compute/conformance/cast.rs b/vortex-array/src/compute/conformance/cast.rs index 2e918914952..161549ed7b0 100644 --- a/vortex-array/src/compute/conformance/cast.rs +++ b/vortex-array/src/compute/conformance/cast.rs @@ -58,7 +58,14 @@ fn test_cast_identity(array: &dyn Array) { // Verify values are unchanged for i in 0..array.len().min(10) { - assert_eq!(array.scalar_at(i), result.scalar_at(i)); + assert_eq!( + array + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"), + result + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test") + ); } } @@ -84,7 +91,12 @@ fn test_cast_from_null(array: &dyn Array) { // Verify all values are null for i in 0..array.len().min(10) { - assert!(result.scalar_at(i).is_null()); + assert!( + result + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test") + .is_null() + ); } } @@ -100,14 +112,25 @@ fn test_cast_from_null(array: &dyn Array) { } fn test_cast_to_non_nullable(array: &dyn Array) { - if array.invalid_count() == 0 { + if array + .invalid_count() + .vortex_expect("invalid_count should succeed in conformance test") + == 0 + { let non_nullable = cast(array, &array.dtype().as_nonnullable()) .vortex_expect("arrays without nulls can cast to non-nullable"); assert_eq!(non_nullable.dtype(), &array.dtype().as_nonnullable()); assert_eq!(non_nullable.len(), array.len()); for i in 0..array.len().min(10) { - assert_eq!(array.scalar_at(i), non_nullable.scalar_at(i)); + assert_eq!( + array + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"), + non_nullable + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test") + ); } let back_to_nullable = cast(&non_nullable, array.dtype()) @@ -116,7 +139,14 @@ fn test_cast_to_non_nullable(array: &dyn Array) { assert_eq!(back_to_nullable.len(), array.len()); for i in 0..array.len().min(10) { - assert_eq!(array.scalar_at(i), back_to_nullable.scalar_at(i)); + assert_eq!( + array + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"), + back_to_nullable + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test") + ); } } else { if &DType::Null == array.dtype() { @@ -142,7 +172,14 @@ fn test_cast_to_nullable(array: &dyn Array) { assert_eq!(nullable.len(), array.len()); for i in 0..array.len().min(10) { - assert_eq!(array.scalar_at(i), nullable.scalar_at(i)); + assert_eq!( + array + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"), + nullable + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test") + ); } let back = cast(&nullable, array.dtype()) @@ -151,7 +188,13 @@ fn test_cast_to_nullable(array: &dyn Array) { assert_eq!(back.len(), array.len()); for i in 0..array.len().min(10) { - assert_eq!(array.scalar_at(i), back.scalar_at(i)); + assert_eq!( + array + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"), + back.scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test") + ); } } @@ -223,10 +266,21 @@ fn test_cast_to_primitive(array: &dyn Array, target_ptype: PType, test_round_tri array.display_values(), ) }); - assert_eq!(array.validity_mask(), casted.validity_mask()); + assert_eq!( + array + .validity_mask() + .vortex_expect("validity_mask should succeed in conformance test"), + casted + .validity_mask() + .vortex_expect("validity_mask should succeed in conformance test") + ); for i in 0..array.len().min(10) { - let original = array.scalar_at(i); - let casted = casted.scalar_at(i); + let original = array + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"); + let casted = casted + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( original .cast(casted.dtype()) diff --git a/vortex-array/src/compute/conformance/consistency.rs b/vortex-array/src/compute/conformance/consistency.rs index 49bf54a4594..6c52af1834f 100644 --- a/vortex-array/src/compute/conformance/consistency.rs +++ b/vortex-array/src/compute/conformance/consistency.rs @@ -88,8 +88,12 @@ fn test_filter_take_consistency(array: &dyn Array) { ); for i in 0..filtered.len() { - let filtered_val = filtered.scalar_at(i); - let taken_val = taken.scalar_at(i); + let filtered_val = filtered + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"); + let taken_val = taken + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( filtered_val, taken_val, "Filter and take produced different values at index {i}. \ @@ -152,8 +156,12 @@ fn test_double_mask_consistency(array: &dyn Array) { ); for i in 0..double_masked.len() { - let double_val = double_masked.scalar_at(i); - let direct_val = directly_masked.scalar_at(i); + let double_val = double_masked + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"); + let direct_val = directly_masked + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( double_val, direct_val, "Sequential masking and combined masking produced different values at index {i}. \ @@ -197,8 +205,12 @@ fn test_filter_identity(array: &dyn Array) { ); for i in 0..len { - let original_val = array.scalar_at(i); - let filtered_val = filtered.scalar_at(i); + let original_val = array + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"); + let filtered_val = filtered + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( filtered_val, original_val, "Filtering with all-true mask should preserve all values. \ @@ -247,8 +259,12 @@ fn test_mask_identity(array: &dyn Array) { ); for i in 0..len { - let original_val = array.scalar_at(i); - let masked_val = masked.scalar_at(i); + let original_val = array + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"); + let masked_val = masked + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"); let expected_val = original_val.clone().into_nullable(); assert_eq!( masked_val, expected_val, @@ -298,8 +314,12 @@ fn test_slice_filter_consistency(array: &dyn Array) { ); for i in 0..filtered.len() { - let filtered_val = filtered.scalar_at(i); - let sliced_val = sliced.scalar_at(i); + let filtered_val = filtered + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"); + let sliced_val = sliced + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( filtered_val, sliced_val, "Filter with contiguous mask and slice produced different values at index {i}. \ @@ -345,8 +365,12 @@ fn test_take_slice_consistency(array: &dyn Array) { ); for i in 0..taken.len() { - let taken_val = taken.scalar_at(i); - let sliced_val = sliced.scalar_at(i); + let taken_val = taken + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"); + let sliced_val = sliced + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( taken_val, sliced_val, "Take with sequential indices and slice produced different values at index {i}. \ @@ -371,9 +395,30 @@ fn test_filter_preserves_order(array: &dyn Array) { // Verify the filtered array contains the right elements in order assert_eq!(filtered.len(), 3.min(len)); if len >= 4 { - assert_eq!(filtered.scalar_at(0), array.scalar_at(0)); - assert_eq!(filtered.scalar_at(1), array.scalar_at(2),); - assert_eq!(filtered.scalar_at(2), array.scalar_at(3)); + assert_eq!( + filtered + .scalar_at(0) + .vortex_expect("scalar_at should succeed in conformance test"), + array + .scalar_at(0) + .vortex_expect("scalar_at should succeed in conformance test") + ); + assert_eq!( + filtered + .scalar_at(1) + .vortex_expect("scalar_at should succeed in conformance test"), + array + .scalar_at(2) + .vortex_expect("scalar_at should succeed in conformance test") + ); + assert_eq!( + filtered + .scalar_at(2) + .vortex_expect("scalar_at should succeed in conformance test"), + array + .scalar_at(3) + .vortex_expect("scalar_at should succeed in conformance test") + ); } } @@ -390,7 +435,14 @@ fn test_take_repeated_indices(array: &dyn Array) { assert_eq!(taken.len(), 3); for i in 0..3 { - assert_eq!(taken.scalar_at(i), array.scalar_at(0),); + assert_eq!( + taken + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"), + array + .scalar_at(0) + .vortex_expect("scalar_at should succeed in conformance test") + ); } } @@ -418,7 +470,14 @@ fn test_mask_filter_null_consistency(array: &dyn Array) { assert_eq!(filtered.len(), direct_filtered.len()); for i in 0..filtered.len() { - assert_eq!(filtered.scalar_at(i), direct_filtered.scalar_at(i)); + assert_eq!( + filtered + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"), + direct_filtered + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test") + ); } } @@ -462,7 +521,14 @@ fn test_take_preserves_properties(array: &dyn Array) { assert_eq!(taken.len(), array.len()); assert_eq!(taken.dtype(), array.dtype()); for i in 0..len { - assert_eq!(taken.scalar_at(i), array.scalar_at(i),); + assert_eq!( + taken + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"), + array + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test") + ); } } @@ -509,8 +575,13 @@ fn test_nullable_indices_consistency(array: &dyn Array) { ); // Check first element (from index 0) - let expected_0 = array.scalar_at(0).into_nullable(); - let actual_0 = taken.scalar_at(0); + let expected_0 = array + .scalar_at(0) + .vortex_expect("scalar_at should succeed in conformance test") + .into_nullable(); + let actual_0 = taken + .scalar_at(0) + .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( actual_0, expected_0, "Take with nullable indices: element at position 0 should be from array index 0. \ @@ -518,15 +589,22 @@ fn test_nullable_indices_consistency(array: &dyn Array) { ); // Check second element (should be null) - let actual_1 = taken.scalar_at(1); + let actual_1 = taken + .scalar_at(1) + .vortex_expect("scalar_at should succeed in conformance test"); assert!( actual_1.is_null(), "Take with nullable indices: element at position 1 should be null, but got {actual_1:?}" ); // Check third element (from index 2) - let expected_2 = array.scalar_at(2).into_nullable(); - let actual_2 = taken.scalar_at(2); + let expected_2 = array + .scalar_at(2) + .vortex_expect("scalar_at should succeed in conformance test") + .into_nullable(); + let actual_2 = taken + .scalar_at(2) + .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( actual_2, expected_2, "Take with nullable indices: element at position 2 should be from array index 2. \ @@ -555,7 +633,14 @@ fn test_large_array_consistency(array: &dyn Array) { // Results should match assert_eq!(taken.len(), filtered.len()); for i in 0..taken.len() { - assert_eq!(taken.scalar_at(i), filtered.scalar_at(i),); + assert_eq!( + taken + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"), + filtered + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test") + ); } } @@ -591,7 +676,9 @@ fn test_comparison_inverse_consistency(array: &dyn Array) { let test_scalar = if len == 0 { return; } else { - array.scalar_at(len / 2) + array + .scalar_at(len / 2) + .vortex_expect("scalar_at should succeed in conformance test") }; // Test Eq vs NotEq @@ -610,8 +697,12 @@ fn test_comparison_inverse_consistency(array: &dyn Array) { ); for i in 0..inverted_eq.len() { - let inv_val = inverted_eq.scalar_at(i); - let neq_val = neq_result.scalar_at(i); + let inv_val = inverted_eq + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"); + let neq_val = neq_result + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( inv_val, neq_val, "At index {i}: NOT(Eq) should equal NotEq. \ @@ -629,8 +720,12 @@ fn test_comparison_inverse_consistency(array: &dyn Array) { invert(>_result).vortex_expect("invert should succeed in conformance test"); for i in 0..inverted_gt.len() { - let inv_val = inverted_gt.scalar_at(i); - let lte_val = lte_result.scalar_at(i); + let inv_val = inverted_gt + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"); + let lte_val = lte_result + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( inv_val, lte_val, "At index {i}: NOT(Gt) should equal Lte. \ @@ -648,8 +743,12 @@ fn test_comparison_inverse_consistency(array: &dyn Array) { invert(<_result).vortex_expect("invert should succeed in conformance test"); for i in 0..inverted_lt.len() { - let inv_val = inverted_lt.scalar_at(i); - let gte_val = gte_result.scalar_at(i); + let inv_val = inverted_lt + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"); + let gte_val = gte_result + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( inv_val, gte_val, "At index {i}: NOT(Lt) should equal Gte. \ @@ -690,7 +789,9 @@ fn test_comparison_symmetry_consistency(array: &dyn Array) { let test_scalar = if len == 2 { return; } else { - array.scalar_at(len / 2) + array + .scalar_at(len / 2) + .vortex_expect("scalar_at should succeed in conformance test") }; // Create a constant array with the test scalar for reverse comparison @@ -708,8 +809,12 @@ fn test_comparison_symmetry_consistency(array: &dyn Array) { ); for i in 0..arr_gt_scalar.len() { - let arr_gt = arr_gt_scalar.scalar_at(i); - let scalar_lt = scalar_lt_arr.scalar_at(i); + let arr_gt = arr_gt_scalar + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"); + let scalar_lt = scalar_lt_arr + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( arr_gt, scalar_lt, "At index {i}: (array > scalar) should equal (scalar < array). \ @@ -724,8 +829,12 @@ fn test_comparison_symmetry_consistency(array: &dyn Array) { compare(const_array.as_ref(), array, Operator::Eq), ) { for i in 0..arr_eq_scalar.len() { - let arr_eq = arr_eq_scalar.scalar_at(i); - let scalar_eq = scalar_eq_arr.scalar_at(i); + let arr_eq = arr_eq_scalar + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"); + let scalar_eq = scalar_eq_arr + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( arr_eq, scalar_eq, "At index {i}: (array == scalar) should equal (scalar == array). \ @@ -776,8 +885,12 @@ fn test_boolean_demorgan_consistency(array: &dyn Array) { ); for i in 0..not_a_and_b.len() { - let left = not_a_and_b.scalar_at(i); - let right = not_a_or_not_b.scalar_at(i); + let left = not_a_and_b + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"); + let right = not_a_or_not_b + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( left, right, "De Morgan's first law failed at index {i}: \ @@ -793,8 +906,12 @@ fn test_boolean_demorgan_consistency(array: &dyn Array) { and(¬_a, ¬_b).vortex_expect("and should succeed in conformance test"); for i in 0..not_a_or_b.len() { - let left = not_a_or_b.scalar_at(i); - let right = not_a_and_not_b.scalar_at(i); + let left = not_a_or_b + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"); + let right = not_a_and_not_b + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( left, right, "De Morgan's second law failed at index {i}: \ @@ -841,13 +958,16 @@ fn test_slice_aggregate_consistency(array: &dyn Array) { let canonical_sliced = canonical.as_ref().slice(start..end); // Test null count through invalid_count + let sliced_invalid_count = sliced + .invalid_count() + .vortex_expect("invalid_count should succeed in conformance test"); + let canonical_invalid_count = canonical_sliced + .invalid_count() + .vortex_expect("invalid_count should succeed in conformance test"); assert_eq!( - sliced.invalid_count(), - canonical_sliced.invalid_count(), + sliced_invalid_count, canonical_invalid_count, "null_count on sliced array should match canonical. \ - Sliced: {}, Canonical: {}", - sliced.invalid_count(), - canonical_sliced.invalid_count() + Sliced: {sliced_invalid_count}, Canonical: {canonical_invalid_count}", ); // Test sum for numeric types @@ -1063,10 +1183,15 @@ fn test_cast_slice_consistency(array: &dyn Array) { // Compare each value against the canonical form for i in 0..slice_then_cast.len() { - let slice_cast_val = slice_then_cast.scalar_at(i); + let slice_cast_val = slice_then_cast + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"); // Get the corresponding value from the canonical array (adjusted for slice offset) - let canonical_val = canonical.as_ref().scalar_at(start + i); + let canonical_val = canonical + .as_ref() + .scalar_at(start + i) + .vortex_expect("scalar_at should succeed in conformance test"); // Cast the canonical scalar to the target dtype let expected_val = match canonical_val.cast(&target_dtype) { @@ -1104,8 +1229,12 @@ fn test_cast_slice_consistency(array: &dyn Array) { ); for i in 0..slice_then_cast.len() { - let slice_cast_val = slice_then_cast.scalar_at(i); - let cast_slice_val = cast_then_slice.scalar_at(i); + let slice_cast_val = slice_then_cast + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"); + let cast_slice_val = cast_then_slice + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( slice_cast_val, cast_slice_val, "Slice-then-cast and cast-then-slice produced different values at index {i}. \ diff --git a/vortex-array/src/compute/conformance/filter.rs b/vortex-array/src/compute/conformance/filter.rs index 6db22fb7e5d..57fa8ba0ed5 100644 --- a/vortex-array/src/compute/conformance/filter.rs +++ b/vortex-array/src/compute/conformance/filter.rs @@ -91,7 +91,14 @@ fn test_selective_filter(array: &dyn Array) { // Verify correct elements are kept for (filtered_idx, i) in (0..len).step_by(2).enumerate() { - assert_eq!(filtered.scalar_at(filtered_idx), array.scalar_at(i)); + assert_eq!( + filtered + .scalar_at(filtered_idx) + .vortex_expect("scalar_at should succeed in conformance test"), + array + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test") + ); } // Test first and last only @@ -103,8 +110,22 @@ fn test_selective_filter(array: &dyn Array) { let filtered = filter(array, &mask).vortex_expect("filter should succeed in conformance test"); assert_eq!(filtered.len(), 2); - assert_eq!(filtered.scalar_at(0), array.scalar_at(0)); - assert_eq!(filtered.scalar_at(1), array.scalar_at(len - 1)); + assert_eq!( + filtered + .scalar_at(0) + .vortex_expect("scalar_at should succeed in conformance test"), + array + .scalar_at(0) + .vortex_expect("scalar_at should succeed in conformance test") + ); + assert_eq!( + filtered + .scalar_at(1) + .vortex_expect("scalar_at should succeed in conformance test"), + array + .scalar_at(len - 1) + .vortex_expect("scalar_at should succeed in conformance test") + ); } } @@ -120,7 +141,14 @@ fn test_single_element_filter(array: &dyn Array) { let mask = Mask::from_iter(mask_values); let filtered = filter(array, &mask).vortex_expect("filter should succeed in conformance test"); assert_eq!(filtered.len(), 1); - assert_eq!(filtered.scalar_at(0), array.scalar_at(0)); + assert_eq!( + filtered + .scalar_at(0) + .vortex_expect("scalar_at should succeed in conformance test"), + array + .scalar_at(0) + .vortex_expect("scalar_at should succeed in conformance test") + ); // Test selecting only the last element if len > 1 { @@ -130,7 +158,14 @@ fn test_single_element_filter(array: &dyn Array) { let filtered = filter(array, &mask).vortex_expect("filter should succeed in conformance test"); assert_eq!(filtered.len(), 1); - assert_eq!(filtered.scalar_at(0), array.scalar_at(len - 1)); + assert_eq!( + filtered + .scalar_at(0) + .vortex_expect("scalar_at should succeed in conformance test"), + array + .scalar_at(len - 1) + .vortex_expect("scalar_at should succeed in conformance test") + ); } } @@ -185,7 +220,14 @@ fn test_alternating_pattern_filter(array: &dyn Array) { let mut filtered_idx = 0; for (i, &keep) in pattern.iter().enumerate() { if keep { - assert_eq!(filtered.scalar_at(filtered_idx), array.scalar_at(i)); + assert_eq!( + filtered + .scalar_at(filtered_idx) + .vortex_expect("scalar_at should succeed in conformance test"), + array + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test") + ); filtered_idx += 1; } } diff --git a/vortex-array/src/compute/conformance/mask.rs b/vortex-array/src/compute/conformance/mask.rs index c3d4abc96e8..463a461b156 100644 --- a/vortex-array/src/compute/conformance/mask.rs +++ b/vortex-array/src/compute/conformance/mask.rs @@ -45,9 +45,21 @@ fn test_heterogenous_mask(array: &dyn Array) { // Verify masked elements are null and unmasked elements are preserved for (i, &masked_out) in mask_pattern.iter().enumerate() { if masked_out { - assert!(!masked.is_valid(i)); + assert!( + !masked + .is_valid(i) + .vortex_expect("is_valid should succeed in conformance test") + ); } else { - assert_eq!(masked.scalar_at(i), array.scalar_at(i).into_nullable()); + assert_eq!( + masked + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"), + array + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test") + .into_nullable() + ); } } } @@ -63,7 +75,15 @@ fn test_empty_mask(array: &dyn Array) { // All elements should be preserved for i in 0..len { - assert_eq!(masked.scalar_at(i), array.scalar_at(i).into_nullable()); + assert_eq!( + masked + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"), + array + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test") + .into_nullable() + ); } } @@ -78,7 +98,11 @@ fn test_full_mask(array: &dyn Array) { // All elements should be null for i in 0..len { - assert!(!masked.is_valid(i)); + assert!( + !masked + .is_valid(i) + .vortex_expect("is_valid should succeed in conformance test") + ); } } @@ -93,9 +117,21 @@ fn test_alternating_mask(array: &dyn Array) { for i in 0..len { if i % 2 == 0 { - assert!(!masked.is_valid(i)); + assert!( + !masked + .is_valid(i) + .vortex_expect("is_valid should succeed in conformance test") + ); } else { - assert_eq!(masked.scalar_at(i), array.scalar_at(i).into_nullable()); + assert_eq!( + masked + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"), + array + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test") + .into_nullable() + ); } } } @@ -115,13 +151,24 @@ fn test_sparse_mask(array: &dyn Array) { assert_eq!(masked.len(), array.len()); // Count how many elements are valid after masking - let valid_count = (0..len).filter(|&i| masked.is_valid(i)).count(); + let valid_count = (0..len) + .filter(|&i| { + masked + .is_valid(i) + .vortex_expect("is_valid should succeed in conformance test") + }) + .count(); // Count how many elements should be invalid: // - Elements that were masked (pattern[i] == true) // - Elements that were already invalid in the original array let expected_invalid_count = (0..len) - .filter(|&i| pattern[i] || !array.is_valid(i)) + .filter(|&i| { + pattern[i] + || !array + .is_valid(i) + .vortex_expect("is_valid should succeed in conformance test") + }) .count(); assert_eq!(valid_count, len - expected_invalid_count); @@ -137,10 +184,22 @@ fn test_single_element_mask(array: &dyn Array) { let mask_array = Mask::from_iter(pattern); let masked = mask(array, &mask_array).vortex_expect("mask should succeed in conformance test"); - assert!(!masked.is_valid(0)); + assert!( + !masked + .is_valid(0) + .vortex_expect("is_valid should succeed in conformance test") + ); for i in 1..len { - assert_eq!(masked.scalar_at(i), array.scalar_at(i).into_nullable()); + assert_eq!( + masked + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"), + array + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test") + .into_nullable() + ); } } @@ -162,11 +221,20 @@ fn test_double_mask(array: &dyn Array) { // Elements should be null if either mask is true for i in 0..len { if mask1_pattern[i] || mask2_pattern[i] { - assert!(!double_masked.is_valid(i)); + assert!( + !double_masked + .is_valid(i) + .vortex_expect("is_valid should succeed in conformance test") + ); } else { assert_eq!( - double_masked.scalar_at(i), - array.scalar_at(i).into_nullable() + double_masked + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"), + array + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test") + .into_nullable() ); } } @@ -193,9 +261,21 @@ fn test_nullable_mask_input(array: &dyn Array) { // Elements are masked only if the mask is true AND valid for i in 0..len { if bool_values[i] && validity_values[i] { - assert!(!masked.is_valid(i)); + assert!( + !masked + .is_valid(i) + .vortex_expect("is_valid should succeed in conformance test") + ); } else { - assert_eq!(masked.scalar_at(i), array.scalar_at(i).into_nullable()); + assert_eq!( + masked + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"), + array + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test") + .into_nullable() + ); } } } diff --git a/vortex-array/src/compute/conformance/take.rs b/vortex-array/src/compute/conformance/take.rs index 8a6f31a0847..1c9257a78a5 100644 --- a/vortex-array/src/compute/conformance/take.rs +++ b/vortex-array/src/compute/conformance/take.rs @@ -79,7 +79,14 @@ fn test_take_all(array: &dyn Array) { _ => { // For non-primitive types, check scalar values for i in 0..len { - assert_eq!(array.scalar_at(i), result.scalar_at(i)); + assert_eq!( + array + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"), + result + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test") + ); } } } @@ -110,8 +117,12 @@ fn test_take_selective(array: &dyn Array) { // Verify the taken elements for (result_idx, &original_idx) in indices.iter().enumerate() { assert_eq!( - array.scalar_at(original_idx as usize), - result.scalar_at(result_idx) + array + .scalar_at(original_idx as usize) + .vortex_expect("scalar_at should succeed in conformance test"), + result + .scalar_at(result_idx) + .vortex_expect("scalar_at should succeed in conformance test") ); } } @@ -123,8 +134,22 @@ fn test_take_first_and_last(array: &dyn Array) { take(array, indices.as_ref()).vortex_expect("take should succeed in conformance test"); assert_eq!(result.len(), 2); - assert_eq!(array.scalar_at(0), result.scalar_at(0)); - assert_eq!(array.scalar_at(len - 1), result.scalar_at(1)); + assert_eq!( + array + .scalar_at(0) + .vortex_expect("scalar_at should succeed in conformance test"), + result + .scalar_at(0) + .vortex_expect("scalar_at should succeed in conformance test") + ); + assert_eq!( + array + .scalar_at(len - 1) + .vortex_expect("scalar_at should succeed in conformance test"), + result + .scalar_at(1) + .vortex_expect("scalar_at should succeed in conformance test") + ); } #[allow(clippy::cast_possible_truncation)] @@ -154,12 +179,21 @@ fn test_take_with_nullable_indices(array: &dyn Array) { for (i, idx_opt) in indices_vec.iter().enumerate() { match idx_opt { Some(idx) => { - let expected = array.scalar_at(*idx as usize); - let actual = result.scalar_at(i); + let expected = array + .scalar_at(*idx as usize) + .vortex_expect("scalar_at should succeed in conformance test"); + let actual = result + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!(expected, actual); } None => { - assert!(result.scalar_at(i).is_null()); + assert!( + result + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test") + .is_null() + ); } } } @@ -176,9 +210,16 @@ fn test_take_repeated_indices(array: &dyn Array) { take(array, indices.as_ref()).vortex_expect("take should succeed in conformance test"); assert_eq!(result.len(), 3); - let first_elem = array.scalar_at(0); + let first_elem = array + .scalar_at(0) + .vortex_expect("scalar_at should succeed in conformance test"); for i in 0..3 { - assert_eq!(result.scalar_at(i), first_elem); + assert_eq!( + result + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test"), + first_elem + ); } } @@ -202,7 +243,14 @@ fn test_take_reverse(array: &dyn Array) { // Verify elements are in reverse order for i in 0..len { - assert_eq!(array.scalar_at(len - 1 - i), result.scalar_at(i)); + assert_eq!( + array + .scalar_at(len - 1 - i) + .vortex_expect("scalar_at should succeed in conformance test"), + result + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test") + ); } } @@ -215,7 +263,14 @@ fn test_take_single_middle(array: &dyn Array) { take(array, indices.as_ref()).vortex_expect("take should succeed in conformance test"); assert_eq!(result.len(), 1); - assert_eq!(array.scalar_at(middle_idx), result.scalar_at(0)); + assert_eq!( + array + .scalar_at(middle_idx) + .vortex_expect("scalar_at should succeed in conformance test"), + result + .scalar_at(0) + .vortex_expect("scalar_at should succeed in conformance test") + ); } #[allow(clippy::cast_possible_truncation)] @@ -238,7 +293,14 @@ fn test_take_random_unsorted(array: &dyn Array) { // Verify elements match for (i, &idx) in indices.iter().enumerate() { - assert_eq!(array.scalar_at(idx as usize), result.scalar_at(i)); + assert_eq!( + array + .scalar_at(idx as usize) + .vortex_expect("scalar_at should succeed in conformance test"), + result + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test") + ); } } @@ -256,7 +318,14 @@ fn test_take_contiguous_range(array: &dyn Array) { // Verify elements for i in 0..(end - start) { - assert_eq!(array.scalar_at(start + i), result.scalar_at(i)); + assert_eq!( + array + .scalar_at(start + i) + .vortex_expect("scalar_at should succeed in conformance test"), + result + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test") + ); } } @@ -284,7 +353,14 @@ fn test_take_mixed_repeated(array: &dyn Array) { // Verify elements for (i, &idx) in indices.iter().enumerate() { - assert_eq!(array.scalar_at(idx as usize), result.scalar_at(i)); + assert_eq!( + array + .scalar_at(idx as usize) + .vortex_expect("scalar_at should succeed in conformance test"), + result + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test") + ); } } @@ -308,6 +384,13 @@ fn test_take_large_indices(array: &dyn Array) { // Spot check a few elements for i in (0..num_indices).step_by(1000) { let expected_idx = indices[i] as usize; - assert_eq!(array.scalar_at(expected_idx), result.scalar_at(i)); + assert_eq!( + array + .scalar_at(expected_idx) + .vortex_expect("scalar_at should succeed in conformance test"), + result + .scalar_at(i) + .vortex_expect("scalar_at should succeed in conformance test") + ); } } diff --git a/vortex-array/src/compute/fill_null.rs b/vortex-array/src/compute/fill_null.rs index 0dc9298c49a..869cb1faf84 100644 --- a/vortex-array/src/compute/fill_null.rs +++ b/vortex-array/src/compute/fill_null.rs @@ -104,11 +104,11 @@ impl ComputeFnVTable for FillNull { ) -> VortexResult { let FillNullArgs { array, fill_value } = FillNullArgs::try_from(args)?; - if !array.dtype().is_nullable() || array.all_valid() { + if !array.dtype().is_nullable() || array.all_valid()? { return Ok(cast(array, fill_value.dtype())?.into()); } - if array.all_invalid() { + if array.all_invalid()? { return Ok(ConstantArray::new(fill_value.clone(), array.len()) .into_array() .into()); diff --git a/vortex-array/src/compute/filter.rs b/vortex-array/src/compute/filter.rs index 8d6251ed6c0..8bfb6cf9b1b 100644 --- a/vortex-array/src/compute/filter.rs +++ b/vortex-array/src/compute/filter.rs @@ -51,17 +51,21 @@ pub(crate) fn warm_up_vtable() -> usize { /// use vortex_array::{Array, IntoArray}; /// use vortex_array::arrays::{BoolArray, PrimitiveArray}; /// use vortex_array::compute::{ filter, mask}; +/// use vortex_error::VortexResult; /// use vortex_mask::Mask; /// use vortex_scalar::Scalar; /// +/// # fn main() -> VortexResult<()> { /// let array = /// PrimitiveArray::from_option_iter([Some(0i32), None, Some(1i32), None, Some(2i32)]); /// let mask = Mask::from_iter([true, false, false, false, true]); /// -/// let filtered = filter(array.as_ref(), &mask).unwrap(); +/// let filtered = filter(array.as_ref(), &mask)?; /// assert_eq!(filtered.len(), 2); -/// assert_eq!(filtered.scalar_at(0), Scalar::from(Some(0_i32))); -/// assert_eq!(filtered.scalar_at(1), Scalar::from(Some(2_i32))); +/// assert_eq!(filtered.scalar_at(0)?, Scalar::from(Some(0_i32))); +/// assert_eq!(filtered.scalar_at(1)?, Scalar::from(Some(2_i32))); +/// # Ok(()) +/// # } /// ``` /// /// # Panics @@ -106,7 +110,7 @@ impl ComputeFnVTable for Filter { } // If the entire array is null, then we only need to adjust the length of the array. - if array.validity_mask().true_count() == 0 { + if array.validity_mask()?.true_count() == 0 { return Ok( ConstantArray::new(Scalar::null(array.dtype().clone()), true_count) .into_array() @@ -126,7 +130,7 @@ impl ComputeFnVTable for Filter { // Otherwise, we can use scalar_at if the mask has length 1. if mask.true_count() == 1 { let idx = mask.first().vortex_expect("true_count == 1"); - return Ok(ConstantArray::new(array.scalar_at(idx), 1) + return Ok(ConstantArray::new(array.scalar_at(idx)?, 1) .into_array() .into()); } diff --git a/vortex-array/src/compute/is_constant.rs b/vortex-array/src/compute/is_constant.rs index ac73a1ea0a4..0de0dfada5f 100644 --- a/vortex-array/src/compute/is_constant.rs +++ b/vortex-array/src/compute/is_constant.rs @@ -141,12 +141,12 @@ fn is_constant_impl( return Ok(Some(true)); } - let all_invalid = array.all_invalid(); + let all_invalid = array.all_invalid()?; if all_invalid { return Ok(Some(true)); } - let all_valid = array.all_valid(); + let all_valid = array.all_valid()?; // If we have some nulls, array can't be constant if !all_valid && !all_invalid { diff --git a/vortex-array/src/compute/is_sorted.rs b/vortex-array/src/compute/is_sorted.rs index 6296ca4a939..3d47eeb4f1c 100644 --- a/vortex-array/src/compute/is_sorted.rs +++ b/vortex-array/src/compute/is_sorted.rs @@ -255,13 +255,13 @@ fn is_sorted_impl( // Enforce strictness before we even try to check if the array is sorted. if strict { - let invalid_count = array.invalid_count(); + let invalid_count = array.invalid_count()?; match invalid_count { // We can keep going 0 => {} // If we have a potential null value - it has to be the first one. 1 => { - if !array.is_invalid(0) { + if !array.is_invalid(0)? { return Ok(Some(false)); } } diff --git a/vortex-array/src/compute/list_contains.rs b/vortex-array/src/compute/list_contains.rs index 4189ff5b6e3..041be9d1048 100644 --- a/vortex-array/src/compute/list_contains.rs +++ b/vortex-array/src/compute/list_contains.rs @@ -138,7 +138,7 @@ impl ComputeFnVTable for ListContains { ); }; - if value.all_invalid() || array.all_invalid() { + if value.all_invalid()? || array.all_invalid()? { return Ok(Output::Array( ConstantArray::new( Scalar::null(DType::Bool(Nullability::Nullable)), @@ -260,7 +260,7 @@ fn list_contains_scalar( // If the list array is constant, we perform a single comparison. if array.len() > 1 && array.is::() { let contains = list_contains_scalar(&array.slice(0..1), value, nullability)?; - return Ok(ConstantArray::new(contains.scalar_at(0), array.len()).into_array()); + return Ok(ConstantArray::new(contains.scalar_at(0)?, array.len()).into_array()); } let list_array = array.to_listview(); diff --git a/vortex-array/src/compute/mask.rs b/vortex-array/src/compute/mask.rs index 474d4785073..02b9991bcf2 100644 --- a/vortex-array/src/compute/mask.rs +++ b/vortex-array/src/compute/mask.rs @@ -49,20 +49,24 @@ pub(crate) fn warm_up_vtable() -> usize { /// use vortex_array::IntoArray; /// use vortex_array::arrays::{BoolArray, PrimitiveArray}; /// use vortex_array::compute::{ mask}; +/// use vortex_error::VortexResult; /// use vortex_mask::Mask; /// use vortex_scalar::Scalar; /// +/// # fn main() -> VortexResult<()> { /// let array = /// PrimitiveArray::from_option_iter([Some(0i32), None, Some(1i32), None, Some(2i32)]); /// let mask_array = Mask::from_iter([true, false, false, false, true]); /// -/// let masked = mask(array.as_ref(), &mask_array).unwrap(); +/// let masked = mask(array.as_ref(), &mask_array)?; /// assert_eq!(masked.len(), 5); -/// assert!(!masked.is_valid(0)); -/// assert!(!masked.is_valid(1)); -/// assert_eq!(masked.scalar_at(2), Scalar::from(Some(1))); -/// assert!(!masked.is_valid(3)); -/// assert!(!masked.is_valid(4)); +/// assert!(!masked.is_valid(0)?); +/// assert!(!masked.is_valid(1)?); +/// assert_eq!(masked.scalar_at(2)?, Scalar::from(Some(1))); +/// assert!(!masked.is_valid(3)?); +/// assert!(!masked.is_valid(4)?); +/// # Ok(()) +/// # } /// ``` /// pub fn mask(array: &dyn Array, mask: &Mask) -> VortexResult { @@ -126,7 +130,7 @@ impl ComputeFnVTable for MaskFn { } // Do nothing if the array is already all nulls. - if array.all_invalid() { + if array.all_invalid()? { return Ok(array.to_array().into()); } diff --git a/vortex-array/src/compute/min_max.rs b/vortex-array/src/compute/min_max.rs index 31dd2131e0c..5b75fdb8d6a 100644 --- a/vortex-array/src/compute/min_max.rs +++ b/vortex-array/src/compute/min_max.rs @@ -146,7 +146,7 @@ fn min_max_impl( array: &dyn Array, kernels: &[ArcRef], ) -> VortexResult> { - if array.is_empty() || array.valid_count() == 0 { + if array.is_empty() || array.valid_count()? == 0 { return Ok(None); } diff --git a/vortex-array/src/compute/nan_count.rs b/vortex-array/src/compute/nan_count.rs index 361108076f3..d9be3e857d5 100644 --- a/vortex-array/src/compute/nan_count.rs +++ b/vortex-array/src/compute/nan_count.rs @@ -115,7 +115,7 @@ impl Kernel for NaNCountKernelAdapter { } fn nan_count_impl(array: &dyn Array, kernels: &[ArcRef]) -> VortexResult { - if array.is_empty() || array.valid_count() == 0 { + if array.is_empty() || array.valid_count()? == 0 { return Ok(0); } diff --git a/vortex-array/src/compute/numeric.rs b/vortex-array/src/compute/numeric.rs index c6c7045a382..fefcfb5d6b6 100644 --- a/vortex-array/src/compute/numeric.rs +++ b/vortex-array/src/compute/numeric.rs @@ -320,7 +320,7 @@ mod test { .to_primitive(); let actual = (0..result.len()) - .map(|index| result.scalar_at(index)) + .map(|index| result.scalar_at(index).unwrap()) .collect::>(); assert_eq!( actual, diff --git a/vortex-array/src/compute/sum.rs b/vortex-array/src/compute/sum.rs index a9b8caa6774..fcfccc44640 100644 --- a/vortex-array/src/compute/sum.rs +++ b/vortex-array/src/compute/sum.rs @@ -237,7 +237,7 @@ pub fn sum_impl( accumulator: &Scalar, kernels: &[ArcRef], ) -> VortexResult { - if array.is_empty() || array.all_invalid() || accumulator.is_null() { + if array.is_empty() || array.all_invalid()? || accumulator.is_null() { return Ok(accumulator.clone()); } diff --git a/vortex-array/src/compute/take.rs b/vortex-array/src/compute/take.rs index 42388cf17b2..73ddf1aaa36 100644 --- a/vortex-array/src/compute/take.rs +++ b/vortex-array/src/compute/take.rs @@ -81,7 +81,7 @@ impl ComputeFnVTable for Take { // TODO(ngates): if indices min is quite high, we could slice self and offset the indices // such that canonicalize does less work. - if indices.all_invalid() { + if indices.all_invalid()? { return Ok(ConstantArray::new( Scalar::null(array.dtype().as_nullable()), indices.len(), @@ -132,7 +132,7 @@ fn propagate_take_stats( indices: &dyn Array, ) -> VortexResult<()> { target.statistics().with_mut_typed_stats_set(|mut st| { - if indices.all_valid() { + if indices.all_valid().unwrap_or(false) { let is_constant = source.statistics().get_as::(Stat::IsConstant); if is_constant == Some(Precision::Exact(true)) { // Any combination of elements from a constant array is still const diff --git a/vortex-array/src/display/mod.rs b/vortex-array/src/display/mod.rs index 469439f1caa..f0cc385073f 100644 --- a/vortex-array/src/display/mod.rs +++ b/vortex-array/src/display/mod.rs @@ -238,7 +238,11 @@ impl dyn Array + '_ { write!( f, "{}", - (0..self.len()).map(|i| self.scalar_at(i)).format(sep) + (0..self.len()) + .map(|i| self + .scalar_at(i) + .map_or_else(|e| format!(""), |s| s.to_string())) + .format(sep) )?; write!(f, "]") } @@ -255,8 +259,10 @@ impl dyn Array + '_ { let DType::Struct(sf, _) = self.dtype() else { // For non-struct arrays, simply display a single column table without header. for row_idx in 0..self.len() { - let value = self.scalar_at(row_idx); - builder.push_record([value.to_string()]); + let value = self + .scalar_at(row_idx) + .map_or_else(|e| format!(""), |s| s.to_string()); + builder.push_record([value]); } let mut table = builder.build(); @@ -269,14 +275,16 @@ impl dyn Array + '_ { builder.push_record(sf.names().iter().map(|name| name.to_string())); for row_idx in 0..self.len() { - if !self.is_valid(row_idx) { + if !self.is_valid(row_idx).unwrap_or(false) { let null_row = vec!["null".to_string(); sf.names().len()]; builder.push_record(null_row); } else { let mut row = Vec::new(); for field_array in struct_.fields().iter() { - let value = field_array.scalar_at(row_idx); - row.push(value.to_string()); + let value = field_array + .scalar_at(row_idx) + .map_or_else(|e| format!(""), |s| s.to_string()); + row.push(value); } builder.push_record(row); } @@ -291,7 +299,7 @@ impl dyn Array + '_ { } for row_idx in 0..self.len() { - if !self.is_valid(row_idx) { + if !self.is_valid(row_idx).unwrap_or(false) { table.modify( (1 + row_idx, 0), tabled::settings::Span::column(sf.names().len() as isize), diff --git a/vortex-array/src/display/tree.rs b/vortex-array/src/display/tree.rs index d47acad0a37..fda1f369da1 100644 --- a/vortex-array/src/display/tree.rs +++ b/vortex-array/src/display/tree.rs @@ -43,12 +43,22 @@ impl fmt::Display for StatsDisplay<'_> { write!(f, "nulls={}", nc)?; } } else if self.0.dtype().is_nullable() { - if self.0.all_valid() { - sep(f)?; - f.write_str("all_valid")?; - } else if self.0.all_invalid() { - sep(f)?; - f.write_str("all_invalid")?; + match self.0.all_valid() { + Ok(true) => { + sep(f)?; + f.write_str("all_valid")?; + } + Ok(false) => { + if self.0.all_invalid().unwrap_or(false) { + sep(f)?; + f.write_str("all_invalid")?; + } + } + Err(e) => { + tracing::warn!("Failed to check validity: {e}"); + sep(f)?; + f.write_str("validity_failed")?; + } } } diff --git a/vortex-array/src/expr/exprs/binary.rs b/vortex-array/src/expr/exprs/binary.rs index 7f553cb69a1..fa290ec176e 100644 --- a/vortex-array/src/expr/exprs/binary.rs +++ b/vortex-array/src/expr/exprs/binary.rs @@ -662,7 +662,7 @@ mod tests { // Test using compare compute function directly let result_equal = compare(&lhs_struct, &rhs_struct_equal, compute::Operator::Eq).unwrap(); assert_eq!( - result_equal.scalar_at(0), + result_equal.scalar_at(0).vortex_expect("value"), Scalar::bool(true, Nullability::NonNullable), "Equal structs should be equal" ); @@ -670,7 +670,7 @@ mod tests { let result_different = compare(&lhs_struct, &rhs_struct_different, compute::Operator::Eq).unwrap(); assert_eq!( - result_different.scalar_at(0), + result_different.scalar_at(0).vortex_expect("value"), Scalar::bool(false, Nullability::NonNullable), "Different structs should not be equal" ); diff --git a/vortex-array/src/expr/exprs/get_item.rs b/vortex-array/src/expr/exprs/get_item.rs index fb00cbf653c..07dd13cd422 100644 --- a/vortex-array/src/expr/exprs/get_item.rs +++ b/vortex-array/src/expr/exprs/get_item.rs @@ -115,7 +115,7 @@ impl VTable for GetItem { match input.dtype().nullability() { Nullability::NonNullable => Ok(field), - Nullability::Nullable => mask(&field, &input.validity_mask().not()), + Nullability::Nullable => mask(&field, &input.validity_mask()?.not()), }? .execute(args.ctx) } @@ -296,7 +296,7 @@ mod tests { let get_item = get_item("a", root()); let item = get_item.evaluate(&st).unwrap(); assert_eq!( - item.scalar_at(0), + item.scalar_at(0).unwrap(), Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)) ); } diff --git a/vortex-array/src/expr/exprs/is_null.rs b/vortex-array/src/expr/exprs/is_null.rs index 3a62892af9c..8518e3b69e1 100644 --- a/vortex-array/src/expr/exprs/is_null.rs +++ b/vortex-array/src/expr/exprs/is_null.rs @@ -174,7 +174,7 @@ mod tests { for (i, expected_value) in expected.iter().enumerate() { assert_eq!( - result.scalar_at(i), + result.scalar_at(i).unwrap(), Scalar::bool(*expected_value, Nullability::NonNullable) ); } @@ -228,7 +228,7 @@ mod tests { for (i, expected_value) in expected.iter().enumerate() { assert_eq!( - result.scalar_at(i), + result.scalar_at(i).unwrap(), Scalar::bool(*expected_value, Nullability::NonNullable) ); } diff --git a/vortex-array/src/expr/exprs/list_contains.rs b/vortex-array/src/expr/exprs/list_contains.rs index fc01309e4c6..94a04707d71 100644 --- a/vortex-array/src/expr/exprs/list_contains.rs +++ b/vortex-array/src/expr/exprs/list_contains.rs @@ -217,9 +217,12 @@ mod tests { let expr = list_contains(root(), lit(1)); let item = expr.evaluate(&arr).unwrap(); - assert_eq!(item.scalar_at(0), Scalar::bool(true, Nullability::Nullable)); assert_eq!( - item.scalar_at(1), + item.scalar_at(0).unwrap(), + Scalar::bool(true, Nullability::Nullable) + ); + assert_eq!( + item.scalar_at(1).unwrap(), Scalar::bool(false, Nullability::Nullable) ); } @@ -231,8 +234,14 @@ mod tests { let expr = list_contains(root(), lit(2)); let item = expr.evaluate(&arr).unwrap(); - assert_eq!(item.scalar_at(0), Scalar::bool(true, Nullability::Nullable)); - assert_eq!(item.scalar_at(1), Scalar::bool(true, Nullability::Nullable)); + assert_eq!( + item.scalar_at(0).unwrap(), + Scalar::bool(true, Nullability::Nullable) + ); + assert_eq!( + item.scalar_at(1).unwrap(), + Scalar::bool(true, Nullability::Nullable) + ); } #[test] @@ -243,11 +252,11 @@ mod tests { let item = expr.evaluate(&arr).unwrap(); assert_eq!( - item.scalar_at(0), + item.scalar_at(0).unwrap(), Scalar::bool(false, Nullability::Nullable) ); assert_eq!( - item.scalar_at(1), + item.scalar_at(1).unwrap(), Scalar::bool(false, Nullability::Nullable) ); } @@ -265,9 +274,12 @@ mod tests { let expr = list_contains(root(), lit(2)); let item = expr.evaluate(&arr).unwrap(); - assert_eq!(item.scalar_at(0), Scalar::bool(true, Nullability::Nullable)); assert_eq!( - item.scalar_at(1), + item.scalar_at(0).unwrap(), + Scalar::bool(true, Nullability::Nullable) + ); + assert_eq!( + item.scalar_at(1).unwrap(), Scalar::bool(false, Nullability::Nullable) ); } @@ -285,8 +297,11 @@ mod tests { let expr = list_contains(root(), lit(2)); let item = expr.evaluate(&arr).unwrap(); - assert_eq!(item.scalar_at(0), Scalar::bool(true, Nullability::Nullable)); - assert!(!item.is_valid(1)); + assert_eq!( + item.scalar_at(0).unwrap(), + Scalar::bool(true, Nullability::Nullable) + ); + assert!(!item.is_valid(1).unwrap()); } #[test] diff --git a/vortex-array/src/mask.rs b/vortex-array/src/mask.rs index 284af06ce90..9dce50dbbe5 100644 --- a/vortex-array/src/mask.rs +++ b/vortex-array/src/mask.rs @@ -35,7 +35,7 @@ impl Executable for Mask { } CanonicalOutput::Array(a) => { let bool = a.into_array().execute::(ctx)?; - let mask = bool.validity_mask(); + let mask = bool.validity_mask()?; let bits = bool.into_bit_buffer(); // To handle nullable boolean arrays, we treat nulls as false in the mask. // TODO(ngates): is this correct? Feels like we should just force the caller to diff --git a/vortex-array/src/patches.rs b/vortex-array/src/patches.rs index 32e43a083c1..848db75bf75 100644 --- a/vortex-array/src/patches.rs +++ b/vortex-array/src/patches.rs @@ -175,8 +175,12 @@ impl Patches { "Patch indices must be shorter than the array length" ); assert!(!indices.is_empty(), "Patch indices must not be empty"); - let max = usize::try_from(&indices.scalar_at(indices.len() - 1)) - .vortex_expect("indices must be a number"); + let max = usize::try_from( + &indices + .scalar_at(indices.len() - 1) + .vortex_expect("indices must support scalar_at"), + ) + .vortex_expect("indices must be a number"); assert!( max - offset < array_len, "Patch indices {max:?}, offset {offset} are longer than the array length {array_len}" @@ -291,6 +295,7 @@ impl Patches { chunk_offsets .scalar_at(idx) + .vortex_expect("chunk_offsets must support scalar_at") .as_primitive() .as_::() .vortex_expect("chunk offset must be usize") @@ -357,10 +362,11 @@ impl Patches { } /// Get the patched value at a given index if it exists. - pub fn get_patched(&self, index: usize) -> Option { - self.search_index(index) + pub fn get_patched(&self, index: usize) -> VortexResult> { + self.search_index(index)? .to_found() .map(|patch_idx| self.values().scalar_at(patch_idx)) + .transpose() } /// Searches for `index` in the indices array. @@ -380,7 +386,7 @@ impl Patches { /// /// [`SearchResult::Found(patch_idx)`]: SearchResult::Found /// [`SearchResult::NotFound(insertion_point)`]: SearchResult::NotFound - pub fn search_index(&self, index: usize) -> SearchResult { + pub fn search_index(&self, index: usize) -> VortexResult { if self.chunk_offsets.is_some() { return self.search_index_chunked(index); } @@ -393,7 +399,10 @@ impl Patches { /// # Returns /// [`SearchResult::Found`] with the position if needle exists, or [`SearchResult::NotFound`] /// with the insertion point if not found. - fn search_index_binary_search(indices: &dyn Array, needle: usize) -> SearchResult { + fn search_index_binary_search( + indices: &dyn Array, + needle: usize, + ) -> VortexResult { if indices.is_canonical() { let primitive = indices.to_primitive(); match_each_integer_ptype!(primitive.ptype(), |T| { @@ -402,7 +411,7 @@ impl Patches { // // The needle is a non-negative integer (a usize); therefore, it must be larger // than all values in this array. - return SearchResult::NotFound(primitive.len()); + return Ok(SearchResult::NotFound(primitive.len())); }; return primitive .as_slice::() @@ -424,7 +433,7 @@ impl Patches { /// /// # Panics /// Panics if `chunk_offsets` or `offset_within_chunk` are not set. - fn search_index_chunked(&self, index: usize) -> SearchResult { + fn search_index_chunked(&self, index: usize) -> VortexResult { let Some(chunk_offsets) = &self.chunk_offsets else { vortex_panic!("chunk_offsets is required to be set") }; @@ -434,7 +443,7 @@ impl Patches { }; if index >= self.array_len() { - return SearchResult::NotFound(self.indices().len()); + return Ok(SearchResult::NotFound(self.indices().len())); } let chunk_idx = (index + self.offset % PATCH_CHUNK_SIZE) / PATCH_CHUNK_SIZE; @@ -458,12 +467,12 @@ impl Patches { }; let chunk_indices = self.indices.slice(patches_start_idx..patches_end_idx); - let result = Self::search_index_binary_search(&chunk_indices, index + self.offset); + let result = Self::search_index_binary_search(&chunk_indices, index + self.offset)?; - match result { + Ok(match result { SearchResult::Found(idx) => SearchResult::Found(patches_start_idx + idx), SearchResult::NotFound(idx) => SearchResult::NotFound(patches_start_idx + idx), - } + }) } /// Batch version of `search_index`. @@ -476,7 +485,7 @@ impl Patches { indices: &[T], chunk_offsets: &[O], index: T, - ) -> SearchResult + ) -> VortexResult where T: UnsignedPType, O: UnsignedPType, @@ -490,11 +499,11 @@ impl Patches { let chunk_idx = { let Ok(index) = usize::try_from(index) else { // If the needle cannot be converted to usize, it's larger than all values in this array. - return SearchResult::NotFound(indices.len()); + return Ok(SearchResult::NotFound(indices.len())); }; if index >= self.array_len() { - return SearchResult::NotFound(self.indices().len()); + return Ok(SearchResult::NotFound(self.indices().len())); } (index + self.offset % PATCH_CHUNK_SIZE) / PATCH_CHUNK_SIZE @@ -534,16 +543,16 @@ impl Patches { let Some(offset) = T::from(self.offset) else { // If the offset cannot be converted to T, it's larger than all values in this array. - return SearchResult::NotFound(indices.len()); + return Ok(SearchResult::NotFound(indices.len())); }; let chunk_indices = &indices[patches_start_idx..patches_end_idx]; - let result = chunk_indices.search_sorted(&(index + offset), SearchSortedSide::Left); + let result = chunk_indices.search_sorted(&(index + offset), SearchSortedSide::Left)?; - match result { + Ok(match result { SearchResult::Found(idx) => SearchResult::Found(patches_start_idx + idx), SearchResult::NotFound(idx) => SearchResult::NotFound(patches_start_idx + idx), - } + }) } /// Returns the minimum patch index @@ -551,6 +560,7 @@ impl Patches { let first = self .indices .scalar_at(0) + .vortex_expect("indices must support scalar_at") .as_primitive() .as_::() .vortex_expect("non-null"); @@ -562,6 +572,7 @@ impl Patches { let last = self .indices .scalar_at(self.indices.len() - 1) + .vortex_expect("indices must support scalar_at") .as_primitive() .as_::() .vortex_expect("non-null"); @@ -643,12 +654,12 @@ impl Patches { } /// Slice the patches by a range of the patched array. - pub fn slice(&self, range: Range) -> Option { - let slice_start_idx = self.search_index(range.start).to_index(); - let slice_end_idx = self.search_index(range.end).to_index(); + pub fn slice(&self, range: Range) -> VortexResult> { + let slice_start_idx = self.search_index(range.start)?.to_index(); + let slice_end_idx = self.search_index(range.end)?.to_index(); if slice_start_idx == slice_end_idx { - return None; + return Ok(None); } let values = self.values().slice(slice_start_idx..slice_end_idx); @@ -664,20 +675,21 @@ impl Patches { let offset_within_chunk = chunk_offsets.as_ref().map(|chunk_offsets| { let base_offset = chunk_offsets .scalar_at(0) + .vortex_expect("chunk_offsets must support scalar_at") .as_primitive() .as_::() .vortex_expect("chunk offset must be usize"); slice_start_idx - base_offset }); - Some(Self { + Ok(Some(Self { array_len: range.len(), offset: range.start + self.offset(), indices, values, chunk_offsets, offset_within_chunk, - }) + })) } // https://docs.google.com/spreadsheets/d/1D9vBZ1QJ6mwcIvV5wIL0hjGgVchcEnAyhvitqWu2ugU @@ -745,7 +757,7 @@ impl Patches { take_indices_with_search_fn( patch_indices_slice, take_slice, - take_indices.validity_mask(), + take_indices.validity_mask()?, include_nulls, |take_idx| { self.search_index_chunked_batch( @@ -754,24 +766,24 @@ impl Patches { take_idx, ) }, - ) + )? }) } else { take_indices_with_search_fn( patch_indices_slice, take_slice, - take_indices.validity_mask(), + take_indices.validity_mask()?, include_nulls, |take_idx| { let Some(offset) = ::from(self.offset) else { // If the offset cannot be converted to T, it's larger than all values in this array. - return SearchResult::NotFound(patch_indices_slice.len()); + return Ok(SearchResult::NotFound(patch_indices_slice.len())); }; patch_indices_slice .search_sorted(&(take_idx + offset), SearchSortedSide::Left) }, - ) + )? } }) }); @@ -1146,31 +1158,39 @@ fn filter_patches_with_mask( ))) } -fn take_indices_with_search_fn SearchResult>( +fn take_indices_with_search_fn< + I: UnsignedPType, + T: IntegerPType, + F: Fn(I) -> VortexResult, +>( indices: &[I], take_indices: &[T], take_validity: Mask, include_nulls: bool, search_fn: F, -) -> (BufferMut, BufferMut) { - take_indices - .iter() - .enumerate() - .filter_map(|(new_patch_idx, &take_idx)| { - if include_nulls && !take_validity.value(new_patch_idx) { - // For nulls, patch index doesn't matter - use 0 for consistency - Some((0u64, new_patch_idx as u64)) - } else { - let search_result = I::from(take_idx) - .map(&search_fn) - .unwrap_or(SearchResult::NotFound(indices.len())); +) -> VortexResult<(BufferMut, BufferMut)> { + let mut values_indices = BufferMut::with_capacity(take_indices.len()); + let mut new_indices = BufferMut::with_capacity(take_indices.len()); + + for (new_patch_idx, &take_idx) in take_indices.iter().enumerate() { + if include_nulls && !take_validity.value(new_patch_idx) { + // For nulls, patch index doesn't matter - use 0 for consistency + values_indices.push(0u64); + new_indices.push(new_patch_idx as u64); + } else { + let search_result = match I::from(take_idx) { + Some(idx) => search_fn(idx)?, + None => SearchResult::NotFound(indices.len()), + }; - search_result - .to_found() - .map(|patch_idx| (patch_idx as u64, new_patch_idx as u64)) + if let Some(patch_idx) = search_result.to_found() { + values_indices.push(patch_idx as u64); + new_indices.push(new_patch_idx as u64); } - }) - .unzip() + } + } + + Ok((values_indices, new_indices)) } #[cfg(test)] @@ -1232,7 +1252,7 @@ mod test { ); assert_arrays_eq!(primitive_indices, PrimitiveArray::from_iter([0u64])); assert_eq!( - primitive_values.validity_mask(), + primitive_values.validity_mask().unwrap(), Mask::from_iter(vec![true]) ); } @@ -1264,7 +1284,7 @@ mod test { assert_arrays_eq!(taken.indices(), PrimitiveArray::from_iter([0u64, 1])); assert_eq!( - primitive_values.validity_mask(), + primitive_values.validity_mask().unwrap(), Mask::from_iter([true, false]) ); } @@ -1367,7 +1387,7 @@ mod test { let patches = Patches::new(101, 0, indices, values, None); - let sliced = patches.slice(15..100).unwrap(); + let sliced = patches.slice(15..100).unwrap().unwrap(); assert_eq!(sliced.array_len(), 100 - 15); assert_arrays_eq!(sliced.values(), PrimitiveArray::from_iter([13531u32])); } @@ -1379,11 +1399,11 @@ mod test { let patches = Patches::new(101, 0, indices, values, None); - let sliced = patches.slice(15..100).unwrap(); + let sliced = patches.slice(15..100).unwrap().unwrap(); assert_eq!(sliced.array_len(), 100 - 15); assert_arrays_eq!(sliced.values(), PrimitiveArray::from_iter([13531u32])); - let doubly_sliced = sliced.slice(35..36).unwrap(); + let doubly_sliced = sliced.slice(35..36).unwrap().unwrap(); assert_arrays_eq!( doubly_sliced.values(), PrimitiveArray::from_iter([13531u32]) @@ -1423,9 +1443,9 @@ mod test { masked.values(), PrimitiveArray::from_iter([100i32, 200, 300]) ); - assert!(masked.values().is_valid(0)); - assert!(masked.values().is_valid(1)); - assert!(masked.values().is_valid(2)); + assert!(masked.values().is_valid(0).unwrap()); + assert!(masked.values().is_valid(1).unwrap()); + assert!(masked.values().is_valid(2).unwrap()); // Indices should remain unchanged assert_arrays_eq!(masked.indices(), PrimitiveArray::from_iter([2u64, 5, 8])); @@ -1499,9 +1519,12 @@ mod test { // Values should be the null and 300 let masked_values = masked.values().to_primitive(); assert_eq!(masked_values.len(), 2); - assert!(!masked_values.is_valid(0)); // the null value at index 5 - assert!(masked_values.is_valid(1)); // the 300 value at index 8 - assert_eq!(i32::try_from(&masked_values.scalar_at(1)).unwrap(), 300i32); + assert!(!masked_values.is_valid(0).unwrap()); // the null value at index 5 + assert!(masked_values.is_valid(1).unwrap()); // the 300 value at index 8 + assert_eq!( + i32::try_from(&masked_values.scalar_at(1).unwrap()).unwrap(), + 300i32 + ); } #[test] @@ -1569,7 +1592,7 @@ mod test { None, ); - let sliced = patches.slice(0..10).unwrap(); + let sliced = patches.slice(0..10).unwrap().unwrap(); assert_arrays_eq!(sliced.indices(), PrimitiveArray::from_iter([2u64, 5, 8])); assert_arrays_eq!( @@ -1589,7 +1612,7 @@ mod test { ); // Slice from 3 to 8 (includes patch at 5) - let sliced = patches.slice(3..8).unwrap(); + let sliced = patches.slice(3..8).unwrap().unwrap(); assert_arrays_eq!(sliced.indices(), PrimitiveArray::from_iter([5u64])); // Index stays the same assert_arrays_eq!(sliced.values(), PrimitiveArray::from_iter([200i32])); @@ -1608,7 +1631,7 @@ mod test { ); // Slice from 6 to 7 (no patches in this range) - let sliced = patches.slice(6..7); + let sliced = patches.slice(6..7).unwrap(); assert!(sliced.is_none()); } @@ -1623,7 +1646,7 @@ mod test { ); // Slice from 3 to 8 (includes patch at actual index 5) - let sliced = patches.slice(3..8).unwrap(); + let sliced = patches.slice(3..8).unwrap().unwrap(); assert_arrays_eq!(sliced.indices(), PrimitiveArray::from_iter([10u64])); // Index stays the same (offset + 5 = 10) assert_arrays_eq!(sliced.values(), PrimitiveArray::from_iter([200i32])); @@ -1641,9 +1664,18 @@ mod test { ); let values = patches.values().to_primitive(); - assert_eq!(i32::try_from(&values.scalar_at(0)).unwrap(), 100i32); - assert_eq!(i32::try_from(&values.scalar_at(1)).unwrap(), 200i32); - assert_eq!(i32::try_from(&values.scalar_at(2)).unwrap(), 300i32); + assert_eq!( + i32::try_from(&values.scalar_at(0).unwrap()).unwrap(), + 100i32 + ); + assert_eq!( + i32::try_from(&values.scalar_at(1).unwrap()).unwrap(), + 200i32 + ); + assert_eq!( + i32::try_from(&values.scalar_at(2).unwrap()).unwrap(), + 300i32 + ); } #[test] @@ -1671,15 +1703,15 @@ mod test { ); // Search for exact indices - assert_eq!(patches.search_index(2), SearchResult::Found(0)); - assert_eq!(patches.search_index(5), SearchResult::Found(1)); - assert_eq!(patches.search_index(8), SearchResult::Found(2)); + assert_eq!(patches.search_index(2).unwrap(), SearchResult::Found(0)); + assert_eq!(patches.search_index(5).unwrap(), SearchResult::Found(1)); + assert_eq!(patches.search_index(8).unwrap(), SearchResult::Found(2)); // Search for non-patch indices - assert_eq!(patches.search_index(0), SearchResult::NotFound(0)); - assert_eq!(patches.search_index(3), SearchResult::NotFound(1)); - assert_eq!(patches.search_index(6), SearchResult::NotFound(2)); - assert_eq!(patches.search_index(9), SearchResult::NotFound(3)); + assert_eq!(patches.search_index(0).unwrap(), SearchResult::NotFound(0)); + assert_eq!(patches.search_index(3).unwrap(), SearchResult::NotFound(1)); + assert_eq!(patches.search_index(6).unwrap(), SearchResult::NotFound(2)); + assert_eq!(patches.search_index(9).unwrap(), SearchResult::NotFound(3)); } #[test] @@ -1837,18 +1869,27 @@ mod test { assert!(patches.chunk_offsets.is_some()); // chunk 0: patches at 100, 200 - assert_eq!(patches.search_index(100), SearchResult::Found(0)); - assert_eq!(patches.search_index(200), SearchResult::Found(1)); + assert_eq!(patches.search_index(100).unwrap(), SearchResult::Found(0)); + assert_eq!(patches.search_index(200).unwrap(), SearchResult::Found(1)); // chunks 1, 2: no patches - assert_eq!(patches.search_index(1500), SearchResult::NotFound(2)); - assert_eq!(patches.search_index(2000), SearchResult::NotFound(2)); + assert_eq!( + patches.search_index(1500).unwrap(), + SearchResult::NotFound(2) + ); + assert_eq!( + patches.search_index(2000).unwrap(), + SearchResult::NotFound(2) + ); // chunk 3: patches at 3000, 3100 - assert_eq!(patches.search_index(3000), SearchResult::Found(2)); - assert_eq!(patches.search_index(3100), SearchResult::Found(3)); + assert_eq!(patches.search_index(3000).unwrap(), SearchResult::Found(2)); + assert_eq!(patches.search_index(3100).unwrap(), SearchResult::Found(3)); - assert_eq!(patches.search_index(1024), SearchResult::NotFound(2)); + assert_eq!( + patches.search_index(1024).unwrap(), + SearchResult::NotFound(2) + ); } #[test] @@ -1858,16 +1899,16 @@ mod test { let chunk_offsets = buffer![0u64, 2, 6].into_array(); let patches = Patches::new(3000, 0, indices, values, Some(chunk_offsets)); - let sliced = patches.slice(1000..2200).unwrap(); + let sliced = patches.slice(1000..2200).unwrap().unwrap(); assert!(sliced.chunk_offsets.is_some()); assert_eq!(sliced.offset(), 1000); - assert_eq!(sliced.search_index(200), SearchResult::Found(0)); - assert_eq!(sliced.search_index(500), SearchResult::Found(2)); - assert_eq!(sliced.search_index(1100), SearchResult::Found(4)); + assert_eq!(sliced.search_index(200).unwrap(), SearchResult::Found(0)); + assert_eq!(sliced.search_index(500).unwrap(), SearchResult::Found(2)); + assert_eq!(sliced.search_index(1100).unwrap(), SearchResult::Found(4)); - assert_eq!(sliced.search_index(250), SearchResult::NotFound(1)); + assert_eq!(sliced.search_index(250).unwrap(), SearchResult::NotFound(1)); } #[test] @@ -1877,16 +1918,16 @@ mod test { let chunk_offsets = buffer![0u64, 2, 6].into_array(); let patches = Patches::new(3000, 0, indices, values, Some(chunk_offsets)); - let sliced = patches.slice(1300..2200).unwrap(); + let sliced = patches.slice(1300..2200).unwrap().unwrap(); assert!(sliced.chunk_offsets.is_some()); assert_eq!(sliced.offset(), 1300); - assert_eq!(sliced.search_index(0), SearchResult::Found(0)); - assert_eq!(sliced.search_index(200), SearchResult::Found(1)); - assert_eq!(sliced.search_index(500), SearchResult::Found(2)); - assert_eq!(sliced.search_index(250), SearchResult::NotFound(2)); - assert_eq!(sliced.search_index(900), SearchResult::NotFound(4)); + assert_eq!(sliced.search_index(0).unwrap(), SearchResult::Found(0)); + assert_eq!(sliced.search_index(200).unwrap(), SearchResult::Found(1)); + assert_eq!(sliced.search_index(500).unwrap(), SearchResult::Found(2)); + assert_eq!(sliced.search_index(250).unwrap(), SearchResult::NotFound(2)); + assert_eq!(sliced.search_index(900).unwrap(), SearchResult::NotFound(4)); } #[test] @@ -1896,7 +1937,7 @@ mod test { let chunk_offsets = buffer![0u64, 2].into_array(); let patches = Patches::new(4000, 0, indices, values, Some(chunk_offsets)); - let sliced = patches.slice(1000..2000); + let sliced = patches.slice(1000..2000).unwrap(); assert!(sliced.is_none()); } @@ -1907,13 +1948,13 @@ mod test { let chunk_offsets = buffer![0u64, 1, 3].into_array(); let patches = Patches::new(3000, 0, indices, values, Some(chunk_offsets)); - let sliced = patches.slice(1100..1250).unwrap(); + let sliced = patches.slice(1100..1250).unwrap().unwrap(); assert_eq!(sliced.num_patches(), 1); assert_eq!(sliced.offset(), 1100); - assert_eq!(sliced.search_index(100), SearchResult::Found(0)); // 1200 - 1100 = 100 - assert_eq!(sliced.search_index(50), SearchResult::NotFound(0)); - assert_eq!(sliced.search_index(150), SearchResult::NotFound(1)); + assert_eq!(sliced.search_index(100).unwrap(), SearchResult::Found(0)); // 1200 - 1100 = 100 + assert_eq!(sliced.search_index(50).unwrap(), SearchResult::NotFound(0)); + assert_eq!(sliced.search_index(150).unwrap(), SearchResult::NotFound(1)); } #[test] @@ -1923,15 +1964,15 @@ mod test { let chunk_offsets = buffer![0u64, 2, 4].into_array(); let patches = Patches::new(3000, 0, indices, values, Some(chunk_offsets)); - let sliced = patches.slice(150..2150).unwrap(); + let sliced = patches.slice(150..2150).unwrap().unwrap(); assert_eq!(sliced.num_patches(), 4); assert_eq!(sliced.offset(), 150); - assert_eq!(sliced.search_index(50), SearchResult::Found(0)); // 200 - 150 = 50 - assert_eq!(sliced.search_index(950), SearchResult::Found(1)); // 1100 - 150 = 950 - assert_eq!(sliced.search_index(1050), SearchResult::Found(2)); // 1200 - 150 = 1050 - assert_eq!(sliced.search_index(1950), SearchResult::Found(3)); // 2100 - 150 = 1950 + assert_eq!(sliced.search_index(50).unwrap(), SearchResult::Found(0)); // 200 - 150 = 50 + assert_eq!(sliced.search_index(950).unwrap(), SearchResult::Found(1)); // 1100 - 150 = 950 + assert_eq!(sliced.search_index(1050).unwrap(), SearchResult::Found(2)); // 1200 - 150 = 1050 + assert_eq!(sliced.search_index(1950).unwrap(), SearchResult::Found(3)); // 2100 - 150 = 1950 } #[test] @@ -1941,14 +1982,20 @@ mod test { let chunk_offsets = buffer![0u64, 1, 4].into_array(); let patches = Patches::new(3000, 0, indices, values, Some(chunk_offsets)); - assert_eq!(patches.search_index(1023), SearchResult::Found(0)); - assert_eq!(patches.search_index(1024), SearchResult::Found(1)); - assert_eq!(patches.search_index(1025), SearchResult::Found(2)); - assert_eq!(patches.search_index(2047), SearchResult::Found(3)); - assert_eq!(patches.search_index(2048), SearchResult::Found(4)); + assert_eq!(patches.search_index(1023).unwrap(), SearchResult::Found(0)); + assert_eq!(patches.search_index(1024).unwrap(), SearchResult::Found(1)); + assert_eq!(patches.search_index(1025).unwrap(), SearchResult::Found(2)); + assert_eq!(patches.search_index(2047).unwrap(), SearchResult::Found(3)); + assert_eq!(patches.search_index(2048).unwrap(), SearchResult::Found(4)); - assert_eq!(patches.search_index(1022), SearchResult::NotFound(0)); - assert_eq!(patches.search_index(2046), SearchResult::NotFound(3)); + assert_eq!( + patches.search_index(1022).unwrap(), + SearchResult::NotFound(0) + ); + assert_eq!( + patches.search_index(2046).unwrap(), + SearchResult::NotFound(3) + ); } #[test] @@ -1959,16 +2006,16 @@ mod test { let patches = Patches::new(3000, 0, indices, values, Some(chunk_offsets)); // Slice at the very beginning - let sliced = patches.slice(0..10).unwrap(); + let sliced = patches.slice(0..10).unwrap().unwrap(); assert_eq!(sliced.num_patches(), 2); - assert_eq!(sliced.search_index(0), SearchResult::Found(0)); - assert_eq!(sliced.search_index(1), SearchResult::Found(1)); + assert_eq!(sliced.search_index(0).unwrap(), SearchResult::Found(0)); + assert_eq!(sliced.search_index(1).unwrap(), SearchResult::Found(1)); // Slice at the very end - let sliced = patches.slice(2040..3000).unwrap(); + let sliced = patches.slice(2040..3000).unwrap().unwrap(); assert_eq!(sliced.num_patches(), 2); // patches at 2047 and 2048 - assert_eq!(sliced.search_index(7), SearchResult::Found(0)); // 2047 - 2040 - assert_eq!(sliced.search_index(8), SearchResult::Found(1)); // 2048 - 2040 + assert_eq!(sliced.search_index(7).unwrap(), SearchResult::Found(0)); // 2047 - 2040 + assert_eq!(sliced.search_index(8).unwrap(), SearchResult::Found(1)); // 2048 - 2040 } #[test] @@ -1978,15 +2025,18 @@ mod test { let chunk_offsets = buffer![0u64].into_array(); let patches = Patches::new(1000, 0, indices, values, Some(chunk_offsets)); - let sliced1 = patches.slice(150..550).unwrap(); + let sliced1 = patches.slice(150..550).unwrap().unwrap(); assert_eq!(sliced1.num_patches(), 4); // 200, 300, 400, 500 - let sliced2 = sliced1.slice(100..250).unwrap(); + let sliced2 = sliced1.slice(100..250).unwrap().unwrap(); assert_eq!(sliced2.num_patches(), 1); // 300 assert_eq!(sliced2.offset(), 250); - assert_eq!(sliced2.search_index(50), SearchResult::Found(0)); // 300 - 250 - assert_eq!(sliced2.search_index(150), SearchResult::NotFound(1)); + assert_eq!(sliced2.search_index(50).unwrap(), SearchResult::Found(0)); // 300 - 250 + assert_eq!( + sliced2.search_index(150).unwrap(), + SearchResult::NotFound(1) + ); } #[test] @@ -1995,6 +2045,9 @@ mod test { let indices = buffer![1023u64].into_array(); let values = buffer![42i32].into_array(); let patches = Patches::new(1024, 0, indices, values, Some(chunk_offsets)); - assert_eq!(patches.search_index(2048), SearchResult::NotFound(1)); + assert_eq!( + patches.search_index(2048).unwrap(), + SearchResult::NotFound(1) + ); } } diff --git a/vortex-array/src/search_sorted.rs b/vortex-array/src/search_sorted.rs index d6b3fdde81f..e050cf8ad74 100644 --- a/vortex-array/src/search_sorted.rs +++ b/vortex-array/src/search_sorted.rs @@ -10,6 +10,7 @@ use std::fmt::Display; use std::fmt::Formatter; use std::hint; +use vortex_error::VortexResult; use vortex_scalar::Scalar; use crate::Array; @@ -98,22 +99,22 @@ impl Display for SearchResult { pub trait IndexOrd { /// PartialOrd of the value at index `idx` with `elem`. /// For example, if self\[idx\] > elem, return Some(Greater). - fn index_cmp(&self, idx: usize, elem: &V) -> Option; + fn index_cmp(&self, idx: usize, elem: &V) -> VortexResult>; - fn index_lt(&self, idx: usize, elem: &V) -> bool { - matches!(self.index_cmp(idx, elem), Some(Less)) + fn index_lt(&self, idx: usize, elem: &V) -> VortexResult { + Ok(matches!(self.index_cmp(idx, elem)?, Some(Less))) } - fn index_le(&self, idx: usize, elem: &V) -> bool { - matches!(self.index_cmp(idx, elem), Some(Less | Equal)) + fn index_le(&self, idx: usize, elem: &V) -> VortexResult { + Ok(matches!(self.index_cmp(idx, elem)?, Some(Less | Equal))) } - fn index_gt(&self, idx: usize, elem: &V) -> bool { - matches!(self.index_cmp(idx, elem), Some(Greater)) + fn index_gt(&self, idx: usize, elem: &V) -> VortexResult { + Ok(matches!(self.index_cmp(idx, elem)?, Some(Greater))) } - fn index_ge(&self, idx: usize, elem: &V) -> bool { - matches!(self.index_cmp(idx, elem), Some(Greater | Equal)) + fn index_ge(&self, idx: usize, elem: &V) -> VortexResult { + Ok(matches!(self.index_cmp(idx, elem)?, Some(Greater | Equal))) } /// Get the length of the underlying ordered collection @@ -129,43 +130,30 @@ pub trait IndexOrd { /// |left |array\[i-1\] < value <= array\[i\]| /// |right|array\[i-1\] <= value < array\[i\]| pub trait SearchSorted { - fn search_sorted_many>( - &self, - values: I, - side: SearchSortedSide, - ) -> impl Iterator - where - Self: IndexOrd, - { - values - .into_iter() - .map(move |value| self.search_sorted(&value, side)) - } - - fn search_sorted(&self, value: &T, side: SearchSortedSide) -> SearchResult + fn search_sorted(&self, value: &T, side: SearchSortedSide) -> VortexResult where Self: IndexOrd, { match side { SearchSortedSide::Left => self.search_sorted_by( - |idx| self.index_cmp(idx, value).unwrap_or(Less), + |idx| Ok(self.index_cmp(idx, value)?.unwrap_or(Less)), |idx| { - if self.index_lt(idx, value) { + Ok(if self.index_lt(idx, value)? { Less } else { Greater - } + }) }, side, ), SearchSortedSide::Right => self.search_sorted_by( - |idx| self.index_cmp(idx, value).unwrap_or(Less), + |idx| Ok(self.index_cmp(idx, value)?.unwrap_or(Less)), |idx| { - if self.index_le(idx, value) { + Ok(if self.index_le(idx, value)? { Less } else { Greater - } + }) }, side, ), @@ -174,12 +162,15 @@ pub trait SearchSorted { /// find function is used to find the element if it exists, if element exists side_find will be /// used to find desired index amongst equal values - fn search_sorted_by Ordering, N: FnMut(usize) -> Ordering>( + fn search_sorted_by< + F: FnMut(usize) -> VortexResult, + N: FnMut(usize) -> VortexResult, + >( &self, find: F, side_find: N, side: SearchSortedSide, - ) -> SearchResult; + ) -> VortexResult; } // Default implementation for types that implement IndexOrd. @@ -187,41 +178,44 @@ impl SearchSorted for S where S: IndexOrd + ?Sized, { - fn search_sorted_by Ordering, N: FnMut(usize) -> Ordering>( + fn search_sorted_by< + F: FnMut(usize) -> VortexResult, + N: FnMut(usize) -> VortexResult, + >( &self, find: F, side_find: N, side: SearchSortedSide, - ) -> SearchResult { - match search_sorted_side_idx(find, 0, self.index_len()) { + ) -> VortexResult { + match search_sorted_side_idx(find, 0, self.index_len())? { SearchResult::Found(found) => { let idx_search = match side { - SearchSortedSide::Left => search_sorted_side_idx(side_find, 0, found), + SearchSortedSide::Left => search_sorted_side_idx(side_find, 0, found)?, SearchSortedSide::Right => { - search_sorted_side_idx(side_find, found, self.index_len()) + search_sorted_side_idx(side_find, found, self.index_len())? } }; match idx_search { - SearchResult::NotFound(i) => SearchResult::Found(i), + SearchResult::NotFound(i) => Ok(SearchResult::Found(i)), _ => unreachable!( "searching amongst equal values should never return Found result" ), } } - s => s, + s => Ok(s), } } } // Code adapted from Rust standard library slice::binary_search_by -fn search_sorted_side_idx Ordering>( +fn search_sorted_side_idx VortexResult>( mut find: F, from: usize, to: usize, -) -> SearchResult { +) -> VortexResult { let mut size = to - from; if size == 0 { - return SearchResult::NotFound(0); + return Ok(SearchResult::NotFound(0)); } let mut base = from; @@ -236,7 +230,7 @@ fn search_sorted_side_idx Ordering>( // SAFETY: the call is made safe by the following inconstants: // - `mid >= 0`: by definition // - `mid < size`: `mid = size / 2 + size / 4 + size / 8 ...` - let cmp = find(mid); + let cmp = find(mid)?; // Binary search interacts poorly with branch prediction, so force // the compiler to use conditional moves if supported by the target @@ -254,24 +248,24 @@ fn search_sorted_side_idx Ordering>( } // SAFETY: base is always in [0, size) because base <= mid. - let cmp = find(base); + let cmp = find(base)?; if cmp == Equal { // SAFETY: same as the call to `find` above. unsafe { hint::assert_unchecked(base < to) }; - SearchResult::Found(base) + Ok(SearchResult::Found(base)) } else { let result = base + (cmp == Less) as usize; // SAFETY: same as the call to `find` above. // Note that this is `<=`, unlike the assert in the `Found` path. unsafe { hint::assert_unchecked(result <= to) }; - SearchResult::NotFound(result) + Ok(SearchResult::NotFound(result)) } } impl IndexOrd for dyn Array + '_ { - fn index_cmp(&self, idx: usize, elem: &Scalar) -> Option { - let scalar_a = self.scalar_at(idx); - scalar_a.partial_cmp(elem) + fn index_cmp(&self, idx: usize, elem: &Scalar) -> VortexResult> { + let scalar_a = self.scalar_at(idx)?; + Ok(scalar_a.partial_cmp(elem)) } fn index_len(&self) -> usize { @@ -280,9 +274,9 @@ impl IndexOrd for dyn Array + '_ { } impl IndexOrd for [T] { - fn index_cmp(&self, idx: usize, elem: &T) -> Option { + fn index_cmp(&self, idx: usize, elem: &T) -> VortexResult> { // SAFETY: Used in search_sorted_by same as the standard library. The search_sorted ensures idx is in bounds - unsafe { self.get_unchecked(idx) }.partial_cmp(elem) + Ok(unsafe { self.get_unchecked(idx) }.partial_cmp(elem)) } fn index_len(&self) -> usize { @@ -291,56 +285,64 @@ impl IndexOrd for [T] { } #[cfg(test)] -mod test { +mod tests { + use vortex_error::VortexResult; + use crate::search_sorted::SearchResult; use crate::search_sorted::SearchSorted; use crate::search_sorted::SearchSortedSide; #[test] - fn left_side_equal() { + fn left_side_equal() -> VortexResult<()> { let arr = [0, 1, 2, 2, 2, 2, 3, 4, 5, 6, 7, 8, 9]; - let res = arr.search_sorted(&2, SearchSortedSide::Left); + let res = arr.search_sorted(&2, SearchSortedSide::Left)?; assert_eq!(arr[res.to_index()], 2); assert_eq!(res, SearchResult::Found(2)); + Ok(()) } #[test] - fn right_side_equal() { + fn right_side_equal() -> VortexResult<()> { let arr = [0, 1, 2, 2, 2, 2, 3, 4, 5, 6, 7, 8, 9]; - let res = arr.search_sorted(&2, SearchSortedSide::Right); + let res = arr.search_sorted(&2, SearchSortedSide::Right)?; assert_eq!(arr[res.to_index() - 1], 2); assert_eq!(res, SearchResult::Found(6)); + Ok(()) } #[test] - fn left_side_equal_beginning() { + fn left_side_equal_beginning() -> VortexResult<()> { let arr = [0, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9]; - let res = arr.search_sorted(&0, SearchSortedSide::Left); + let res = arr.search_sorted(&0, SearchSortedSide::Left)?; assert_eq!(arr[res.to_index()], 0); assert_eq!(res, SearchResult::Found(0)); + Ok(()) } #[test] - fn right_side_equal_beginning() { + fn right_side_equal_beginning() -> VortexResult<()> { let arr = [0, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9]; - let res = arr.search_sorted(&0, SearchSortedSide::Right); + let res = arr.search_sorted(&0, SearchSortedSide::Right)?; assert_eq!(arr[res.to_index() - 1], 0); assert_eq!(res, SearchResult::Found(4)); + Ok(()) } #[test] - fn left_side_equal_end() { + fn left_side_equal_end() -> VortexResult<()> { let arr = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 9, 9, 9]; - let res = arr.search_sorted(&9, SearchSortedSide::Left); + let res = arr.search_sorted(&9, SearchSortedSide::Left)?; assert_eq!(arr[res.to_index()], 9); assert_eq!(res, SearchResult::Found(9)); + Ok(()) } #[test] - fn right_side_equal_end() { + fn right_side_equal_end() -> VortexResult<()> { let arr = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 9, 9, 9]; - let res = arr.search_sorted(&9, SearchSortedSide::Right); + let res = arr.search_sorted(&9, SearchSortedSide::Right)?; assert_eq!(arr[res.to_index() - 1], 9); assert_eq!(res, SearchResult::Found(13)); + Ok(()) } } diff --git a/vortex-array/src/stats/array.rs b/vortex-array/src/stats/array.rs index fa6ffb98375..9df19542f15 100644 --- a/vortex-array/src/stats/array.rs +++ b/vortex-array/src/stats/array.rs @@ -160,7 +160,7 @@ impl StatsSetRef<'_> { }) .transpose()? } - Stat::NullCount => Some(self.dyn_array_ref.invalid_count().into()), + Stat::NullCount => self.dyn_array_ref.invalid_count().ok().map(Into::into), Stat::IsConstant => { if self.dyn_array_ref.is_empty() { None diff --git a/vortex-array/src/test_harness.rs b/vortex-array/src/test_harness.rs index f45335b1cae..4160919ad48 100644 --- a/vortex-array/src/test_harness.rs +++ b/vortex-array/src/test_harness.rs @@ -32,7 +32,7 @@ where /// Outputs the indices of the true values in a BoolArray pub fn to_int_indices(indices_bits: BoolArray) -> VortexResult> { let buffer = indices_bits.bit_buffer(); - let mask = indices_bits.validity_mask(); + let mask = indices_bits.validity_mask()?; Ok(buffer .iter() .enumerate() diff --git a/vortex-array/src/validity.rs b/vortex-array/src/validity.rs index 5b224a7278e..7a359f4da7d 100644 --- a/vortex-array/src/validity.rs +++ b/vortex-array/src/validity.rs @@ -131,13 +131,12 @@ impl Validity { match self { Self::NonNullable | Self::AllValid => true, Self::AllInvalid => false, - Self::Array(a) => { - let scalar = a.scalar_at(index); - scalar - .as_bool() - .value() - .vortex_expect("Validity must be non-nullable") - } + Self::Array(a) => a + .scalar_at(index) + .vortex_expect("Validity array must support scalar_at") + .as_bool() + .value() + .vortex_expect("Validity must be non-nullable"), } } @@ -156,7 +155,7 @@ impl Validity { pub fn take(&self, indices: &dyn Array) -> VortexResult { match self { - Self::NonNullable => match indices.validity_mask().bit_buffer() { + Self::NonNullable => match indices.validity_mask()?.bit_buffer() { AllOr::All => { if indices.dtype().is_nullable() { Ok(Self::AllValid) @@ -167,7 +166,7 @@ impl Validity { AllOr::None => Ok(Self::AllInvalid), AllOr::Some(buf) => Ok(Validity::from(buf.clone())), }, - Self::AllValid => match indices.validity_mask().bit_buffer() { + Self::AllValid => match indices.validity_mask()?.bit_buffer() { AllOr::All => Ok(Self::AllValid), AllOr::None => Ok(Self::AllInvalid), AllOr::Some(buf) => Ok(Validity::from(buf.clone())), @@ -360,8 +359,11 @@ impl Validity { /// Create Validity by copying the given array's validity. #[inline] - pub fn copy_from_array(array: &dyn Array) -> Self { - Validity::from_mask(array.validity_mask(), array.dtype().nullability()) + pub fn copy_from_array(array: &dyn Array) -> VortexResult { + Ok(Validity::from_mask( + array.validity_mask()?, + array.dtype().nullability(), + )) } /// Create Validity from boolean array with given nullability of the array. diff --git a/vortex-array/src/variants.rs b/vortex-array/src/variants.rs index 3f733d918e7..d2edc0ccbd1 100644 --- a/vortex-array/src/variants.rs +++ b/vortex-array/src/variants.rs @@ -109,23 +109,28 @@ impl PrimitiveTyped<'_> { } /// Return the primitive value at the given index. - pub fn value(&self, idx: usize) -> Option { - self.0.is_valid(idx).then(|| self.value_unchecked(idx)) + pub fn value(&self, idx: usize) -> VortexResult> { + self.0 + .is_valid(idx)? + .then(|| self.value_unchecked(idx)) + .transpose() } /// Return the primitive value at the given index, ignoring nullability. - pub fn value_unchecked(&self, idx: usize) -> PValue { - self.0 - .scalar_at(idx) + pub fn value_unchecked(&self, idx: usize) -> VortexResult { + Ok(self + .0 + .scalar_at(idx)? .as_primitive() .pvalue() - .unwrap_or_else(|| PValue::zero(self.ptype())) + .unwrap_or_else(|| PValue::zero(self.ptype()))) } } impl IndexOrd> for PrimitiveTyped<'_> { - fn index_cmp(&self, idx: usize, elem: &Option) -> Option { - self.value(idx).partial_cmp(elem) + fn index_cmp(&self, idx: usize, elem: &Option) -> VortexResult> { + let value = self.value(idx)?; + Ok(value.partial_cmp(elem)) } fn index_len(&self) -> usize { @@ -135,9 +140,10 @@ impl IndexOrd> for PrimitiveTyped<'_> { // TODO(ngates): add generics to the `value` function and implement this over T. impl IndexOrd for PrimitiveTyped<'_> { - fn index_cmp(&self, idx: usize, elem: &PValue) -> Option { - assert!(self.0.all_valid()); - self.value_unchecked(idx).partial_cmp(elem) + fn index_cmp(&self, idx: usize, elem: &PValue) -> VortexResult> { + assert!(self.0.all_valid()?); + let value = self.value_unchecked(idx)?; + Ok(value.partial_cmp(elem)) } fn index_len(&self) -> usize { diff --git a/vortex-array/src/vtable/canonical.rs b/vortex-array/src/vtable/canonical.rs deleted file mode 100644 index d2d958b9c67..00000000000 --- a/vortex-array/src/vtable/canonical.rs +++ /dev/null @@ -1,6 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use crate::vtable::VTable; - -pub trait CanonicalVTable {} diff --git a/vortex-array/src/vtable/mod.rs b/vortex-array/src/vtable/mod.rs index 5382c29f295..12ae5c87c40 100644 --- a/vortex-array/src/vtable/mod.rs +++ b/vortex-array/src/vtable/mod.rs @@ -4,7 +4,6 @@ //! This module contains the VTable definitions for a Vortex encoding. mod array; -mod canonical; mod compute; mod dyn_; mod operations; @@ -16,7 +15,6 @@ use std::ops::Deref; use std::ops::Range; pub use array::*; -pub use canonical::*; pub use compute::*; pub use dyn_::*; pub use operations::*; diff --git a/vortex-array/src/vtable/operations.rs b/vortex-array/src/vtable/operations.rs index 951ea7105c6..2c0353eaa93 100644 --- a/vortex-array/src/vtable/operations.rs +++ b/vortex-array/src/vtable/operations.rs @@ -1,7 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_error::vortex_panic; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_scalar::Scalar; use crate::vtable::NotSupported; @@ -14,12 +15,12 @@ pub trait OperationsVTable { /// /// Bounds-checking has already been performed by the time this function is called, /// and the index is guaranteed to be non-null. - fn scalar_at(array: &V::Array, index: usize) -> Scalar; + fn scalar_at(array: &V::Array, index: usize) -> VortexResult; } impl OperationsVTable for NotSupported { - fn scalar_at(array: &V::Array, _index: usize) -> Scalar { - vortex_panic!( + fn scalar_at(array: &V::Array, _index: usize) -> VortexResult { + vortex_bail!( "Legacy scalar_at operation is not supported for {} arrays", array.encoding_id() ) diff --git a/vortex-array/src/vtable/validity.rs b/vortex-array/src/vtable/validity.rs index df1ae03b3d9..89e517365ba 100644 --- a/vortex-array/src/vtable/validity.rs +++ b/vortex-array/src/vtable/validity.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_panic; use vortex_mask::Mask; @@ -20,10 +19,8 @@ pub trait ValidityVTable { /// - The array DType is nullable. fn validity(array: &V::Array) -> VortexResult; - fn validity_mask(array: &V::Array) -> Mask { - Self::validity(array) - .vortex_expect("TODO: make this fallible") - .to_mask(array.len()) + fn validity_mask(array: &V::Array) -> VortexResult { + Ok(Self::validity(array)?.to_mask(array.len())) } } @@ -91,7 +88,7 @@ where V::validity_child(array).validity() } - fn validity_mask(array: &V::Array) -> Mask { + fn validity_mask(array: &V::Array) -> VortexResult { V::validity_child(array).validity_mask() } } diff --git a/vortex-btrblocks/src/float.rs b/vortex-btrblocks/src/float.rs index d3c371c8778..8a80503a2ce 100644 --- a/vortex-btrblocks/src/float.rs +++ b/vortex-btrblocks/src/float.rs @@ -173,13 +173,14 @@ impl Scheme for ConstantScheme { _allowed_cascading: usize, _excludes: &[FloatCode], ) -> VortexResult { - let scalar_idx = (0..stats.source().len()).position(|idx| stats.source().is_valid(idx)); + let scalar_idx = + (0..stats.source().len()).position(|idx| stats.source().is_valid(idx).unwrap_or(false)); match scalar_idx { Some(idx) => { - let scalar = stats.source().scalar_at(idx); + let scalar = stats.source().scalar_at(idx)?; let const_arr = ConstantArray::new(scalar, stats.src.len()).into_array(); - if !stats.source().all_valid() { + if !stats.source().all_valid()? { Ok(MaskedArray::try_new(const_arr, stats.src.validity().clone())?.into_array()) } else { Ok(const_arr) diff --git a/vortex-btrblocks/src/float/stats.rs b/vortex-btrblocks/src/float/stats.rs index af46197bfd3..e4bb660a58a 100644 --- a/vortex-btrblocks/src/float/stats.rs +++ b/vortex-btrblocks/src/float/stats.rs @@ -121,7 +121,7 @@ where } .into(), }; - } else if array.all_invalid() { + } else if array.all_invalid().vortex_expect("all_invalid") { return FloatStats { src: array.clone(), null_count: array.len().try_into().vortex_expect("null_count"), @@ -149,7 +149,7 @@ where HashSet::with_hasher(FxBuildHasher) }; - let validity = array.validity_mask(); + let validity = array.validity_mask().vortex_expect("validity_mask"); let mut runs = 1; let head_idx = validity diff --git a/vortex-btrblocks/src/integer.rs b/vortex-btrblocks/src/integer.rs index bcec54076fe..1f66d61df4a 100644 --- a/vortex-btrblocks/src/integer.rs +++ b/vortex-btrblocks/src/integer.rs @@ -225,13 +225,14 @@ impl Scheme for ConstantScheme { _allowed_cascading: usize, _excludes: &[IntCode], ) -> VortexResult { - let scalar_idx = (0..stats.source().len()).position(|idx| stats.source().is_valid(idx)); + let scalar_idx = + (0..stats.source().len()).position(|idx| stats.source().is_valid(idx).unwrap_or(false)); match scalar_idx { Some(idx) => { - let scalar = stats.source().scalar_at(idx); + let scalar = stats.source().scalar_at(idx)?; let const_arr = ConstantArray::new(scalar, stats.src.len()).into_array(); - if !stats.source().all_valid() { + if !stats.source().all_valid()? { Ok(MaskedArray::try_new(const_arr, stats.src.validity().clone())?.into_array()) } else { Ok(const_arr) diff --git a/vortex-btrblocks/src/integer/stats.rs b/vortex-btrblocks/src/integer/stats.rs index 0e3aa3d7b68..0075e692048 100644 --- a/vortex-btrblocks/src/integer/stats.rs +++ b/vortex-btrblocks/src/integer/stats.rs @@ -196,7 +196,7 @@ where } .into(), }; - } else if array.all_invalid() { + } else if array.all_invalid().vortex_expect("all_invalid") { return IntegerStats { src: array.clone(), null_count: array.len().try_into().vortex_expect("null_count"), @@ -214,7 +214,7 @@ where }; } - let validity = array.validity_mask(); + let validity = array.validity_mask().vortex_expect("validity_mask"); let null_count = validity.false_count(); let value_count = validity.true_count(); diff --git a/vortex-btrblocks/src/lib.rs b/vortex-btrblocks/src/lib.rs index 4f66e5fbad6..6a672c1c5c6 100644 --- a/vortex-btrblocks/src/lib.rs +++ b/vortex-btrblocks/src/lib.rs @@ -494,7 +494,7 @@ impl BtrBlocksCompressor { .unwrap_or_default() { return Ok(ConstantArray::new( - temporal_array.as_ref().scalar_at(0), + temporal_array.as_ref().scalar_at(0)?, ext_array.len(), ) .into_array()); diff --git a/vortex-btrblocks/src/patches.rs b/vortex-btrblocks/src/patches.rs index 57133a5ed88..9a195b82b7c 100644 --- a/vortex-btrblocks/src/patches.rs +++ b/vortex-btrblocks/src/patches.rs @@ -20,7 +20,7 @@ pub fn compress_patches(patches: &Patches) -> VortexResult { .compute_is_constant() .unwrap_or_default() { - ConstantArray::new(values.scalar_at(0), values.len()).into_array() + ConstantArray::new(values.scalar_at(0)?, values.len()).into_array() } else { values.clone() }; diff --git a/vortex-btrblocks/src/string.rs b/vortex-btrblocks/src/string.rs index 8f1d765056a..479727ac2fd 100644 --- a/vortex-btrblocks/src/string.rs +++ b/vortex-btrblocks/src/string.rs @@ -347,13 +347,14 @@ impl Scheme for ConstantScheme { _allowed_cascading: usize, _excludes: &[Self::CodeType], ) -> VortexResult { - let scalar_idx = (0..stats.source().len()).position(|idx| stats.source().is_valid(idx)); + let scalar_idx = + (0..stats.source().len()).position(|idx| stats.source().is_valid(idx).unwrap_or(false)); match scalar_idx { Some(idx) => { - let scalar = stats.source().scalar_at(idx); + let scalar = stats.source().scalar_at(idx)?; let const_arr = ConstantArray::new(scalar, stats.src.len()).into_array(); - if !stats.source().all_valid() { + if !stats.source().all_valid()? { Ok(MaskedArray::try_new(const_arr, stats.src.validity().clone())?.into_array()) } else { Ok(const_arr) diff --git a/vortex-duckdb/src/convert/vector.rs b/vortex-duckdb/src/convert/vector.rs index 7dbedc69813..51c2fbc0f30 100644 --- a/vortex-duckdb/src/convert/vector.rs +++ b/vortex-duckdb/src/convert/vector.rs @@ -503,7 +503,7 @@ mod tests { assert_eq!(values_slice, values); assert_eq!( - vortex_values.validity_mask(), + vortex_values.validity_mask().unwrap(), Mask::from_indices(3, vec![0, 2]) ); } @@ -609,7 +609,7 @@ mod tests { assert_eq!(vortex_slice, values); assert_eq!( - vortex_array.validity_mask(), + vortex_array.validity_mask().unwrap(), Mask::from_indices(3, vec![0, 2]) ); } @@ -775,7 +775,10 @@ mod tests { vortex_array.list_elements_at(0), PrimitiveArray::from_option_iter([Some(1i32), Some(2), Some(3), Some(4)]) ); - assert_eq!(vortex_array.validity_mask(), Mask::from_indices(2, vec![0])); + assert_eq!( + vortex_array.validity_mask().unwrap(), + Mask::from_indices(2, vec![0]) + ); } #[test] @@ -886,7 +889,7 @@ mod tests { assert_eq!(sizes.as_slice::()[1], 0); assert_eq!( - vortex_array.validity_mask(), + vortex_array.validity_mask().unwrap(), Mask::from_indices(3, vec![0, 2]) ); } diff --git a/vortex-duckdb/src/exporter/decimal.rs b/vortex-duckdb/src/exporter/decimal.rs index 5b9c624bc29..a286503a685 100644 --- a/vortex-duckdb/src/exporter/decimal.rs +++ b/vortex-duckdb/src/exporter/decimal.rs @@ -139,7 +139,7 @@ mod tests { pub(crate) fn new_zero_copy_exporter( array: &DecimalArray, ) -> VortexResult> { - let validity = array.validity_mask(); + let validity = array.validity_mask()?; let dest_values_type = precision_to_duckdb_storage_size(&array.decimal_dtype())?; assert_eq!(array.values_type(), dest_values_type); diff --git a/vortex-duckdb/src/exporter/dict.rs b/vortex-duckdb/src/exporter/dict.rs index be045c6eed9..9dc9934ae94 100644 --- a/vortex-duckdb/src/exporter/dict.rs +++ b/vortex-duckdb/src/exporter/dict.rs @@ -53,13 +53,13 @@ pub(crate) fn new_exporter_with_flatten( if let Some(constant) = values.as_opt::() { return constant::new_exporter_with_mask( ConstantArray::new(constant.scalar().clone(), array.codes().len()), - array.codes().validity_mask(), + array.codes().validity_mask()?, cache, ctx, ); } - let codes_mask = array.codes().validity_mask(); + let codes_mask = array.codes().validity_mask()?; match codes_mask { Mask::AllTrue(_) => {} diff --git a/vortex-duckdb/src/exporter/run_end.rs b/vortex-duckdb/src/exporter/run_end.rs index f37e5a1146f..32b102f949a 100644 --- a/vortex-duckdb/src/exporter/run_end.rs +++ b/vortex-duckdb/src/exporter/run_end.rs @@ -65,7 +65,7 @@ impl ColumnExporter for RunEndExporter { // Find the run that contains the start offset. let start_run_idx = ends_slice - .search_sorted(&offset, SearchSortedSide::Right) + .search_sorted(&offset, SearchSortedSide::Right)? .to_ends_index(ends_slice.len()); // Find the final run in case we can short-circuit and return a constant vector. @@ -73,13 +73,13 @@ impl ColumnExporter for RunEndExporter { .search_sorted( &offset.add(E::from_usize(len).vortex_expect("len out of bounds")), SearchSortedSide::Right, - ) + )? .to_ends_index(ends_slice.len()); if start_run_idx == end_run_idx { // NOTE(ngates): would be great if we could just export and set type == CONSTANT // self.values_exporter.export(start_run_idx, 1, vector, cache); - let constant = self.values.scalar_at(start_run_idx); + let constant = self.values.scalar_at(start_run_idx)?; let value = constant.try_to_duckdb_scalar()?; vector.reference_value(&value); return Ok(()); diff --git a/vortex-error/src/lib.rs b/vortex-error/src/lib.rs index 85c83782416..bf1079f5e66 100644 --- a/vortex-error/src/lib.rs +++ b/vortex-error/src/lib.rs @@ -310,7 +310,6 @@ where #[inline(always)] fn vortex_expect(self, msg: &'static str) -> Self::Output { self.map_err(|err| err.into()) - .inspect_err(|e| println!("got a big ERROR {e}")) .unwrap_or_else(|e| vortex_panic!(e.with_context(msg.to_string()))) } } diff --git a/vortex-ffi/src/array.rs b/vortex-ffi/src/array.rs index cf728388fbd..c2475a2060e 100644 --- a/vortex-ffi/src/array.rs +++ b/vortex-ffi/src/array.rs @@ -87,7 +87,9 @@ pub unsafe extern "C-unwind" fn vx_array_is_null( _error_out: *mut *mut vx_error, ) -> bool { let array = vx_array::as_ref(array); - array.is_invalid(index as usize) + array + .is_invalid(index as usize) + .vortex_expect("is_invalid failed") } // TODO(robert): Make this return usize and remove error @@ -97,7 +99,7 @@ pub unsafe extern "C-unwind" fn vx_array_null_count( error_out: *mut *mut vx_error, ) -> u32 { let array = vx_array::as_ref(array); - try_or_default(error_out, || Ok(array.invalid_count().try_into()?)) + try_or_default(error_out, || Ok(array.invalid_count()?.try_into()?)) } macro_rules! ffiarray_get_ptype { @@ -106,7 +108,7 @@ macro_rules! ffiarray_get_ptype { #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn [](array: *const vx_array, index: u32) -> $ptype { let array = vx_array::as_ref(array); - let value = array.scalar_at(index as usize); + let value = array.scalar_at(index as usize).vortex_expect("scalar_at failed"); value.as_primitive() .as_::<$ptype>() .vortex_expect("null value") @@ -115,7 +117,7 @@ macro_rules! ffiarray_get_ptype { #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn [](array: *const vx_array, index: u32) -> $ptype { let array = vx_array::as_ref(array); - let value = array.scalar_at(index as usize); + let value = array.scalar_at(index as usize).vortex_expect("scalar_at failed"); value.as_extension() .storage() .as_primitive() @@ -146,7 +148,9 @@ pub unsafe extern "C-unwind" fn vx_array_get_utf8( index: u32, ) -> *const vx_string { let array = vx_array::as_ref(array); - let value = array.scalar_at(index as usize); + let value = array + .scalar_at(index as usize) + .vortex_expect("scalar_at failed"); let utf8_scalar = value.as_utf8(); if let Some(buffer) = utf8_scalar.value() { vx_string::new(Arc::from(buffer.as_str())) @@ -163,7 +167,9 @@ pub unsafe extern "C-unwind" fn vx_array_get_binary( index: u32, ) -> *const vx_binary { let array = vx_array::as_ref(array); - let value = array.scalar_at(index as usize); + let value = array + .scalar_at(index as usize) + .vortex_expect("scalar_at failed"); let binary_scalar = value.as_binary(); if let Some(bytes) = binary_scalar.value() { vx_binary::new(Arc::from(bytes.as_bytes())) diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 2bcc0b73739..6d877aa4762 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -23,6 +23,7 @@ use vortex_array::arrays::DictVTable; use vortex_array::arrays::ListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; +use vortex_array::arrays::TemporalArray; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::assert_arrays_eq; @@ -52,10 +53,14 @@ use vortex_buffer::ByteBufferMut; use vortex_buffer::buffer; use vortex_dtype::DType; use vortex_dtype::DecimalDType; +use vortex_dtype::ExtDType; use vortex_dtype::Nullability; use vortex_dtype::PType; use vortex_dtype::PType::I32; use vortex_dtype::StructFields; +use vortex_dtype::datetime::TIMESTAMP_ID; +use vortex_dtype::datetime::TemporalMetadata; +use vortex_dtype::datetime::TimeUnit; use vortex_error::VortexResult; use vortex_io::session::RuntimeSession; use vortex_layout::session::LayoutSession; @@ -1098,10 +1103,10 @@ async fn test_repeated_projection() { assert_eq!( (0..actual.len()) - .map(|index| actual.scalar_at(index)) + .map(|index| actual.scalar_at(index).unwrap()) .collect_vec(), (0..expected.len()) - .map(|index| expected.scalar_at(index)) + .map(|index| expected.scalar_at(index).unwrap()) .collect_vec() ); } @@ -1252,12 +1257,12 @@ async fn write_nullable_nested_struct() -> VortexResult<()> { assert_eq!(result.len(), 3); assert_eq!(result.fields().len(), 1); - assert!(result.all_valid()); + assert!(result.all_valid()?); let nested_struct = result.field_by_name("struct")?.to_struct(); assert_eq!(nested_struct.dtype(), &nested_dtype); assert_eq!(nested_struct.len(), 3); - assert!(nested_struct.all_invalid()); + assert!(nested_struct.all_invalid()?); Ok(()) } @@ -1589,3 +1594,45 @@ async fn test_writer_with_statistics() -> VortexResult<()> { Ok(()) } + +#[tokio::test] +async fn main_test() -> Result<(), Box> { + // Write file with MILLISECONDS timestamps + let ts_array = PrimitiveArray::from_iter(vec![1704067200000i64, 1704153600000, 1704240000000]) + .into_array(); + let temporal = TemporalArray::new_timestamp(ts_array, TimeUnit::Milliseconds, None); + + let mut buf = ByteBufferMut::empty(); + SESSION + .write_options() + .write(&mut buf, temporal.into_array().to_array_stream()) + .await?; + + // Read with SECONDS filter scalar + let seconds_ext_dtype = Arc::new(ExtDType::new( + TIMESTAMP_ID.clone(), + Arc::new(DType::Primitive(PType::I64, Nullability::Nullable)), + Some(TemporalMetadata::Timestamp(TimeUnit::Seconds, None).into()), + )); + let filter_expr = gt( + root(), + lit(Scalar::extension( + seconds_ext_dtype, + Scalar::from(1704153600i64), + )), + ); + + let stream = SESSION + .open_options() + .open_buffer(buf)? + .scan()? + .with_filter(filter_expr) + .into_array_stream()?; + + let results = stream.try_collect::>().await; + + let err = results.err().unwrap(); + println!("Expected error: {}", err); + + Ok(()) +} diff --git a/vortex-jni/src/array.rs b/vortex-jni/src/array.rs index 3a48b1438a8..3c1faffea92 100644 --- a/vortex-jni/src/array.rs +++ b/vortex-jni/src/array.rs @@ -258,7 +258,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeArrayMethods_getNull( ) -> jboolean { let array_ref = unsafe { NativeArray::from_ptr(array_ptr) }; try_or_throw(&mut env, |_| { - let is_null = array_ref.inner.is_invalid(index as usize); + let is_null = array_ref.inner.is_invalid(index as usize)?; if is_null { Ok(JNI_TRUE) } else { Ok(JNI_FALSE) } }) } @@ -271,7 +271,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeArrayMethods_getNullCount( ) -> jint { let array_ref = unsafe { NativeArray::from_ptr(array_ptr) }; try_or_throw(&mut env, |_| { - let count = array_ref.inner.invalid_count(); + let count = array_ref.inner.invalid_count()?; Ok(jint::try_from(count).unwrap_or(-1)) }) } @@ -292,9 +292,9 @@ macro_rules! get_primitive { .inner .to_extension() .storage() - .scalar_at(index as usize) + .scalar_at(index as usize)? } else { - array_ref.inner.scalar_at(index as usize) + array_ref.inner.scalar_at(index as usize)? }; Ok(scalar_value @@ -331,9 +331,9 @@ pub extern "system" fn Java_dev_vortex_jni_NativeArrayMethods_getBigDecimal( .inner .to_extension() .storage() - .scalar_at(index as usize) + .scalar_at(index as usize)? } else { - array_ref.inner.scalar_at(index as usize) + array_ref.inner.scalar_at(index as usize)? }; let decimal_scalar = scalar_value.as_decimal(); @@ -395,7 +395,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeArrayMethods_getBool( ) -> jboolean { let array_ref = unsafe { NativeArray::from_ptr(array_ptr) }; try_or_throw(&mut env, |_| { - let value = array_ref.inner.scalar_at(index as usize); + let value = array_ref.inner.scalar_at(index as usize)?; match value.as_bool().value() { None => Ok(JNI_FALSE), Some(b) => { @@ -418,7 +418,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeArrayMethods_getUTF8<'local>( ) -> jstring { let array_ref = unsafe { NativeArray::from_ptr(array_ptr) }; try_or_throw(&mut env, |env| { - let value = array_ref.inner.scalar_at(index as usize); + let value = array_ref.inner.scalar_at(index as usize)?; match value.as_utf8().value() { None => Ok(JObject::null().into_raw()), Some(buf_str) => Ok(env.new_string(buf_str.as_str())?.into_raw()), @@ -467,7 +467,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeArrayMethods_getBinary<'local>( ) -> jbyteArray { let array_ref = unsafe { NativeArray::from_ptr(array_ptr) }; try_or_throw(&mut env, |env| { - let value = array_ref.inner.scalar_at(index as usize); + let value = array_ref.inner.scalar_at(index as usize)?; match value.as_binary().value() { None => Ok(JObject::null().into_raw()), Some(buf) => Ok(env.byte_array_from_slice(buf.as_slice())?.into_raw()), diff --git a/vortex-layout/src/layouts/dict/reader.rs b/vortex-layout/src/layouts/dict/reader.rs index 0c73dfb3125..df78e5bb10c 100644 --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs @@ -479,7 +479,7 @@ mod tests { .unwrap() .await .unwrap(); - let expected = array.validity_mask().into_array(); + let expected = array.validity_mask().unwrap().into_array(); assert_arrays_eq!( actual .to_canonical() diff --git a/vortex-layout/src/layouts/file_stats.rs b/vortex-layout/src/layouts/file_stats.rs index 38f0bba6441..7c923f97b6d 100644 --- a/vortex-layout/src/layouts/file_stats.rs +++ b/vortex-layout/src/layouts/file_stats.rs @@ -113,6 +113,7 @@ impl FileStatsAccumulator { .iter_mut() .map(|acc| { acc.as_stats_table() + .vortex_expect("as_stats_table should not fail") .map(|table| { table .to_stats_set(&self.stats) diff --git a/vortex-layout/src/layouts/flat/writer.rs b/vortex-layout/src/layouts/flat/writer.rs index f28134f78f8..d9b261f36cb 100644 --- a/vortex-layout/src/layouts/flat/writer.rs +++ b/vortex-layout/src/layouts/flat/writer.rs @@ -347,7 +347,7 @@ mod tests { .unwrap(); assert_eq!( - result.validity_mask().bit_buffer(), + result.validity_mask().unwrap().bit_buffer(), AllOr::Some(&validity_boolean_buffer) ); assert_eq!( diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index a14495334b1..34524249464 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -702,7 +702,10 @@ mod tests { ); // ...and the result is masked with the validity of the parent StructArray - assert_eq!(result.scalar_at(0), Scalar::null(result.dtype().clone()),); + assert_eq!( + result.scalar_at(0).unwrap(), + Scalar::null(result.dtype().clone()), + ); assert_nth_scalar!(result, 1, 2); assert_nth_scalar!(result, 2, 3); } @@ -729,19 +732,30 @@ mod tests { // Struct scalars holding the "c" field value scalars assert_eq!( - result.scalar_at(0).as_struct().field_by_idx(0).unwrap(), + result + .scalar_at(0) + .unwrap() + .as_struct() + .field_by_idx(0) + .unwrap(), Scalar::primitive(4, Nullability::Nullable) ); assert!( result .scalar_at(1) + .unwrap() .as_struct() .field_by_idx(0) .unwrap() .is_null(), ); assert_eq!( - result.scalar_at(2).as_struct().field_by_idx(0).unwrap(), + result + .scalar_at(2) + .unwrap() + .as_struct() + .field_by_idx(0) + .unwrap(), Scalar::primitive(6, Nullability::Nullable) ); } diff --git a/vortex-layout/src/layouts/struct_/writer.rs b/vortex-layout/src/layouts/struct_/writer.rs index 277908d7c16..585b5248c11 100644 --- a/vortex-layout/src/layouts/struct_/writer.rs +++ b/vortex-layout/src/layouts/struct_/writer.rs @@ -107,7 +107,7 @@ impl LayoutStrategy for StructStrategy { if is_nullable { columns.push(( sequence_pointer.advance(), - chunk.validity_mask().into_array(), + chunk.validity_mask()?.into_array(), )); } diff --git a/vortex-layout/src/layouts/table.rs b/vortex-layout/src/layouts/table.rs index 3ccb9c7b04a..561cceaedea 100644 --- a/vortex-layout/src/layouts/table.rs +++ b/vortex-layout/src/layouts/table.rs @@ -251,7 +251,7 @@ impl LayoutStrategy for TableStrategy { if is_nullable { columns.push(( sequence_pointer.advance(), - chunk.validity_mask().into_array(), + chunk.validity_mask()?.into_array(), )); } diff --git a/vortex-layout/src/layouts/zoned/builder.rs b/vortex-layout/src/layouts/zoned/builder.rs index 170b116b181..45462488c73 100644 --- a/vortex-layout/src/layouts/zoned/builder.rs +++ b/vortex-layout/src/layouts/zoned/builder.rs @@ -67,7 +67,7 @@ pub struct NamedArrays { } impl NamedArrays { - pub fn all_invalid(&self) -> bool { + pub fn all_invalid(&self) -> VortexResult { // by convention we assume that the first array is the one we care about for logical validity self.arrays[0].all_invalid() } diff --git a/vortex-layout/src/layouts/zoned/writer.rs b/vortex-layout/src/layouts/zoned/writer.rs index 7adedde2c5e..1800532b1ee 100644 --- a/vortex-layout/src/layouts/zoned/writer.rs +++ b/vortex-layout/src/layouts/zoned/writer.rs @@ -134,7 +134,7 @@ impl LayoutStrategy for ZonedStrategy { ) .await?; - let Some(stats_table) = stats_accumulator.lock().as_stats_table() else { + let Some(stats_table) = stats_accumulator.lock().as_stats_table()? else { // If we have no stats (e.g. the DType doesn't support them), then we just return the // child layout. return Ok(data_layout); diff --git a/vortex-layout/src/layouts/zoned/zone_map.rs b/vortex-layout/src/layouts/zoned/zone_map.rs index d1fbc4ac7ad..0d8a693b8a4 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -218,7 +218,7 @@ impl StatsAccumulator { /// /// Returns `None` if none of the requested statistics can be computed, for example they are /// not applicable to the column's data type. - pub fn as_stats_table(&mut self) -> Option { + pub fn as_stats_table(&mut self) -> VortexResult> { let mut names = Vec::new(); let mut fields = Vec::new(); let mut stats = Vec::new(); @@ -232,7 +232,7 @@ impl StatsAccumulator { let values = builder.finish(); // We drop any all-null stats columns - if values.all_invalid() { + if values.all_invalid()? { continue; } @@ -242,14 +242,14 @@ impl StatsAccumulator { } if names.is_empty() { - return None; + return Ok(None); } - Some(ZoneMap { + Ok(Some(ZoneMap { array: StructArray::try_new(names.into(), fields, self.length, Validity::NonNullable) .vortex_expect("Failed to create zone map"), stats: stats.into(), - }) + })) } } @@ -305,7 +305,10 @@ mod tests { .vortex_expect("push_chunk should succeed for test data"); acc.push_chunk(&builder2.finish()) .vortex_expect("push_chunk should succeed for test data"); - let stats_table = acc.as_stats_table().vortex_expect("Must have stats table"); + let stats_table = acc + .as_stats_table() + .unwrap() + .expect("Must have stats table"); assert_eq!( stats_table.array.names().as_ref(), &[ @@ -331,7 +334,10 @@ mod tests { let mut acc = StatsAccumulator::new(array.dtype(), &[Stat::Max, Stat::Min, Stat::Sum], 12); acc.push_chunk(&array) .vortex_expect("push_chunk should succeed for test array"); - let stats_table = acc.as_stats_table().vortex_expect("Must have stats table"); + let stats_table = acc + .as_stats_table() + .unwrap() + .expect("Must have stats table"); assert_eq!( stats_table.array.names().as_ref(), &[ diff --git a/vortex-python/src/arrays/mod.rs b/vortex-python/src/arrays/mod.rs index d33af824064..0f540b6b3ae 100644 --- a/vortex-python/src/arrays/mod.rs +++ b/vortex-python/src/arrays/mod.rs @@ -551,8 +551,7 @@ impl PyArray { /// >>> vx.array([10, 42, 999, 1992]).scalar_at(10) /// Traceback (most recent call last): /// ... - /// pyo3_runtime.PanicException: index 10 out of bounds - /// ... + /// ValueError: index 10 out of bounds from 0 to 4 /// ``` /// /// Unlike Python, negative indices are not supported: @@ -567,7 +566,7 @@ impl PyArray { fn scalar_at(slf: Bound, index: usize) -> PyResult> { let py = slf.py(); let slf = PyArrayRef::extract(slf.as_any().as_borrowed())?.into_inner(); - PyScalar::init(py, slf.scalar_at(index)) + PyScalar::init(py, slf.scalar_at(index)?) } /// Filter, permute, and/or repeat elements by their index. diff --git a/vortex-python/src/arrays/py/vtable.rs b/vortex-python/src/arrays/py/vtable.rs index 7469f244f68..8e65e50993e 100644 --- a/vortex-python/src/arrays/py/vtable.rs +++ b/vortex-python/src/arrays/py/vtable.rs @@ -142,7 +142,7 @@ impl BaseArrayVTable for PythonVTable { } impl OperationsVTable for PythonVTable { - fn scalar_at(_array: &PythonArray, _index: usize) -> Scalar { + fn scalar_at(_array: &PythonArray, _index: usize) -> VortexResult { todo!() } } @@ -152,7 +152,7 @@ impl ValidityVTable for PythonVTable { todo!() } - fn validity_mask(_array: &PythonArray) -> Mask { + fn validity_mask(_array: &PythonArray) -> VortexResult { todo!() } } diff --git a/vortex-python/src/scan.rs b/vortex-python/src/scan.rs index 53427372da2..6ea00bda112 100644 --- a/vortex-python/src/scan.rs +++ b/vortex-python/src/scan.rs @@ -67,7 +67,7 @@ impl PyRepeatedScan { if array.is_empty() { continue; } - let scalar = array.scalar_at(0); + let scalar = array.scalar_at(0)?; return PyScalar::init(slf.py(), scalar); } diff --git a/vortex-tui/src/browse/ui/layouts.rs b/vortex-tui/src/browse/ui/layouts.rs index 1e51eb7a804..6ed9355cd2c 100644 --- a/vortex-tui/src/browse/ui/layouts.rs +++ b/vortex-tui/src/browse/ui/layouts.rs @@ -169,11 +169,13 @@ fn render_array(app: &AppState<'_>, area: Rect, buf: &mut Buffer, is_stats_table // TODO: trim the number of displayed rows and allow paging through column stats. let rows = (0..array.len()).map(|chunk_id| { std::iter::once(Cell::from(Text::from(format!("{chunk_id}")))) - .chain( - field_arrays - .iter() - .map(|arr| Cell::from(Text::from(arr.scalar_at(chunk_id).to_string()))), - ) + .chain(field_arrays.iter().map(|arr| { + Cell::from(Text::from( + arr.scalar_at(chunk_id) + .vortex_expect("scalar_at failed") + .to_string(), + )) + })) .collect::() });