diff --git a/bench-vortex/src/bin/compress.rs b/bench-vortex/src/bin/compress.rs index 36b547c7c57..a1858672817 100644 --- a/bench-vortex/src/bin/compress.rs +++ b/bench-vortex/src/bin/compress.rs @@ -43,6 +43,7 @@ use vortex::array::IntoArray; use vortex::array::arrays::ChunkedArray; use vortex::array::arrays::ChunkedVTable; use vortex::array::builders::builder_with_capacity; +use vortex::error::VortexExpect; use vortex::utils::aliases::hash_map::HashMap; #[derive(Parser, Debug)] @@ -235,7 +236,9 @@ pub fn benchmark_compress( .iter() .map(|chunk| { let mut builder = builder_with_capacity(chunk.dtype(), chunk.len()); - chunk.append_to_builder(builder.as_mut()); + chunk + .append_to_builder(builder.as_mut()) + .vortex_expect("append_to_builder"); builder.finish() }), ) diff --git a/bench-vortex/src/bin/random_access.rs b/bench-vortex/src/bin/random_access.rs index aad95f579d4..2dc92b8da41 100644 --- a/bench-vortex/src/bin/random_access.rs +++ b/bench-vortex/src/bin/random_access.rs @@ -231,7 +231,7 @@ fn random_access( } fn validate_vortex_array(array: ArrayRef) { - let struct_ = array.to_struct(); + let struct_ = array.to_struct().vortex_expect("to_struct"); assert_eq!(struct_.len(), 6, "expected 6 rows"); let pu_location_id = struct_ .field_by_name("PULocationID") diff --git a/bench-vortex/src/datasets/tpch_l_comment.rs b/bench-vortex/src/datasets/tpch_l_comment.rs index 276ee8a13e7..47a3973b77c 100644 --- a/bench-vortex/src/datasets/tpch_l_comment.rs +++ b/bench-vortex/src/datasets/tpch_l_comment.rs @@ -55,7 +55,7 @@ impl Dataset for TPCHLCommentChunked { let file_chunks: Vec<_> = file .scan()? .with_projection(pack(vec![("l_comment", col("l_comment"))], NonNullable)) - .map(|a| Ok(a.to_canonical().into_array())) + .map(|a| Ok(a.to_canonical()?.into_array())) .into_array_stream()? .try_collect() .await?; @@ -78,7 +78,7 @@ impl Dataset for TPCHLCommentCanonical { let comments_canonical = TPCHLCommentChunked .to_vortex_array() .await? - .to_struct() + .to_struct()? .into_array(); Ok(ChunkedArray::from_iter([comments_canonical]).into_array()) } diff --git a/bench-vortex/src/random_access/take.rs b/bench-vortex/src/random_access/take.rs index b51913ea818..f9ca77bdf22 100644 --- a/bench-vortex/src/random_access/take.rs +++ b/bench-vortex/src/random_access/take.rs @@ -53,7 +53,7 @@ async fn take_vortex(reader: impl AsRef, indices: Buffer) -> anyhow:: .read_all() .await? // We canonicalize / decompress for equivalence to Arrow's `RecordBatch`es. - .to_canonical() + .to_canonical()? .into_array()) } diff --git a/encodings/alp/src/alp/array.rs b/encodings/alp/src/alp/array.rs index d2812f32d9f..bcb283bc8c8 100644 --- a/encodings/alp/src/alp/array.rs +++ b/encodings/alp/src/alp/array.rs @@ -420,8 +420,8 @@ impl BaseArrayVTable for ALPVTable { } impl CanonicalVTable for ALPVTable { - fn canonicalize(array: &ALPArray) -> Canonical { - Canonical::Primitive(decompress_into_array(array.clone())) + fn canonicalize(array: &ALPArray) -> VortexResult { + Ok(Canonical::Primitive(decompress_into_array(array.clone())?)) } } @@ -482,7 +482,7 @@ mod tests { let result_vector = encoded.to_array().execute(&SESSION).unwrap(); // Compare against the traditional array-based decompress path - let expected = decompress_into_array(encoded); + let expected = decompress_into_array(encoded).unwrap(); assert_eq!(result_vector.len(), size); @@ -506,7 +506,7 @@ mod tests { let result_vector = encoded.to_array().execute(&SESSION).unwrap(); // Compare against the traditional array-based decompress path - let expected = decompress_into_array(encoded); + let expected = decompress_into_array(encoded).unwrap(); assert_eq!(result_vector.len(), size); @@ -536,7 +536,7 @@ mod tests { let result_vector = encoded.to_array().execute(&SESSION).unwrap(); // Compare against the traditional array-based decompress path - let expected = decompress_into_array(encoded); + let expected = decompress_into_array(encoded).unwrap(); assert_eq!(result_vector.len(), size); @@ -564,7 +564,7 @@ mod tests { let result_vector = encoded.to_array().execute(&SESSION).unwrap(); // Compare against the traditional array-based decompress path - let expected = decompress_into_array(encoded); + let expected = decompress_into_array(encoded).unwrap(); assert_eq!(result_vector.len(), size); @@ -575,7 +575,7 @@ mod tests { for idx in 0..size { assert_eq!( result_primitive.validity().value(idx), - expected.validity().is_valid(idx) + expected.validity().is_valid(idx).unwrap() ); } } @@ -603,7 +603,7 @@ mod tests { let result_vector = encoded.to_array().execute(&SESSION).unwrap(); // Compare against the traditional array-based decompress path - let expected = decompress_into_array(encoded); + let expected = decompress_into_array(encoded).unwrap(); assert_eq!(result_vector.len(), size); @@ -614,7 +614,7 @@ mod tests { for idx in 0..size { assert_eq!( result_primitive.validity().value(idx), - expected.validity().is_valid(idx) + expected.validity().is_valid(idx).unwrap() ); } } diff --git a/encodings/alp/src/alp/compress.rs b/encodings/alp/src/alp/compress.rs index d0d5d0dc089..3b853ae6f27 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) = @@ -150,7 +150,7 @@ mod tests { assert_arrays_eq!(encoded.encoded(), expected_encoded); assert_eq!(encoded.exponents(), Exponents { e: 9, f: 6 }); - let decoded = decompress_into_array(encoded); + let decoded = decompress_into_array(encoded).unwrap(); assert_arrays_eq!(decoded, array); } @@ -163,9 +163,9 @@ mod tests { assert_arrays_eq!(encoded.encoded(), expected_encoded); assert_eq!(encoded.exponents(), Exponents { e: 9, f: 6 }); - let decoded = decompress_into_array(encoded); + let decoded = decompress_into_array(encoded).unwrap(); let expected = PrimitiveArray::from_option_iter(vec![None, Some(1.234f32), None]); - assert_arrays_eq!(decoded, expected); + assert_arrays_eq!(decoded.as_ref(), expected.as_ref()); } #[test] @@ -179,9 +179,9 @@ mod tests { assert_arrays_eq!(encoded.encoded(), expected_encoded); assert_eq!(encoded.exponents(), Exponents { e: 16, f: 13 }); - let decoded = decompress_into_array(encoded); + let decoded = decompress_into_array(encoded).unwrap(); let expected_decoded = PrimitiveArray::new(values, Validity::NonNullable); - assert_arrays_eq!(decoded, expected_decoded); + assert_arrays_eq!(decoded.as_ref(), expected_decoded.as_ref()); } #[test] @@ -196,7 +196,7 @@ mod tests { assert_arrays_eq!(encoded.encoded(), expected_encoded); assert_eq!(encoded.exponents(), Exponents { e: 16, f: 13 }); - let decoded = decompress_into_array(encoded); + let decoded = decompress_into_array(encoded).unwrap(); assert_arrays_eq!(decoded, array); } @@ -215,17 +215,17 @@ mod tests { assert_eq!(encoded.exponents(), Exponents { e: 16, f: 13 }); - assert_arrays_eq!(encoded, array); + assert_arrays_eq!(encoded.as_ref(), array.as_ref()); - let _decoded = decompress_into_array(encoded); + let _decoded = decompress_into_array(encoded).unwrap(); } #[test] fn roundtrips_close_fractional() { let original = PrimitiveArray::from_iter([195.26274f32, 195.27837, -48.815685]); let alp_arr = alp_encode(&original, None).unwrap(); - let decompressed = alp_arr.to_primitive(); - assert_arrays_eq!(decompressed, original); + let decompressed = alp_arr.to_primitive().unwrap(); + assert_arrays_eq!(decompressed.as_ref(), original.as_ref()); } #[test] @@ -233,7 +233,7 @@ mod tests { let original = PrimitiveArray::new(buffer![195.26274f64, PI, -48.815685], Validity::AllInvalid); let alp_arr = alp_encode(&original, None).unwrap(); - let decompressed = alp_arr.to_primitive(); + let decompressed = alp_arr.to_primitive().unwrap(); assert_eq!( // The second and third values become exceptions and are replaced @@ -241,7 +241,7 @@ mod tests { decompressed.as_slice::() ); - assert_arrays_eq!(decompressed, original); + assert_arrays_eq!(decompressed.as_ref(), original.as_ref()); } #[test] @@ -251,7 +251,7 @@ mod tests { Validity::NonNullable, ); let encoded = alp_encode(&original, None).unwrap(); - let decoded = encoded.to_primitive(); + let decoded = encoded.to_primitive().unwrap(); for idx in 0..original.len() { let decoded_val = decoded.as_slice::()[idx]; let original_val = original.as_slice::()[idx]; @@ -274,17 +274,23 @@ mod tests { let encoded = alp_encode(&array, None).unwrap(); let patches = encoded.patches().unwrap(); - let chunk_offsets = patches.chunk_offsets().clone().unwrap().to_primitive(); + let chunk_offsets = patches + .chunk_offsets() + .clone() + .unwrap() + .as_ref() + .to_primitive() + .unwrap(); let expected_offsets = PrimitiveArray::from_iter(vec![0u64, 1, 3]); - assert_arrays_eq!(chunk_offsets, expected_offsets); + assert_arrays_eq!(chunk_offsets.as_ref(), expected_offsets.as_ref()); - let patch_indices = patches.indices().to_primitive(); + let patch_indices = patches.indices().as_ref().to_primitive().unwrap(); let expected_indices = PrimitiveArray::from_iter(vec![1023u64, 1024, 1025]); - assert_arrays_eq!(patch_indices, expected_indices); + assert_arrays_eq!(patch_indices.as_ref(), expected_indices.as_ref()); - let patch_values = patches.values().to_primitive(); + let patch_values = patches.values().as_ref().to_primitive().unwrap(); let expected_values = PrimitiveArray::from_iter(vec![PI, E, PI]); - assert_arrays_eq!(patch_values, expected_values); + assert_arrays_eq!(patch_values.as_ref(), expected_values.as_ref()); } #[test] @@ -297,17 +303,23 @@ mod tests { let encoded = alp_encode(&array, None).unwrap(); let patches = encoded.patches().unwrap(); - let chunk_offsets = patches.chunk_offsets().clone().unwrap().to_primitive(); + let chunk_offsets = patches + .chunk_offsets() + .clone() + .unwrap() + .as_ref() + .to_primitive() + .unwrap(); let expected_offsets = PrimitiveArray::from_iter(vec![0u64, 1, 1]); - assert_arrays_eq!(chunk_offsets, expected_offsets); + assert_arrays_eq!(chunk_offsets.as_ref(), expected_offsets.as_ref()); - let patch_indices = patches.indices().to_primitive(); + let patch_indices = patches.indices().as_ref().to_primitive().unwrap(); let expected_indices = PrimitiveArray::from_iter(vec![0u64, 2048]); - assert_arrays_eq!(patch_indices, expected_indices); + assert_arrays_eq!(patch_indices.as_ref(), expected_indices.as_ref()); - let patch_values = patches.values().to_primitive(); + let patch_values = patches.values().as_ref().to_primitive().unwrap(); let expected_values = PrimitiveArray::from_iter(vec![PI, E]); - assert_arrays_eq!(patch_values, expected_values); + assert_arrays_eq!(patch_values.as_ref(), expected_values.as_ref()); } #[test] @@ -319,17 +331,23 @@ mod tests { let encoded = alp_encode(&array, None).unwrap(); let patches = encoded.patches().unwrap(); - let chunk_offsets = patches.chunk_offsets().clone().unwrap().to_primitive(); + let chunk_offsets = patches + .chunk_offsets() + .clone() + .unwrap() + .as_ref() + .to_primitive() + .unwrap(); let expected_offsets = PrimitiveArray::from_iter(vec![0u64, 1, 1]); - assert_arrays_eq!(chunk_offsets, expected_offsets); + assert_arrays_eq!(chunk_offsets.as_ref(), expected_offsets.as_ref()); - let patch_indices = patches.indices().to_primitive(); + let patch_indices = patches.indices().as_ref().to_primitive().unwrap(); let expected_indices = PrimitiveArray::from_iter(vec![0u64]); - assert_arrays_eq!(patch_indices, expected_indices); + assert_arrays_eq!(patch_indices.as_ref(), expected_indices.as_ref()); - let patch_values = patches.values().to_primitive(); + let patch_values = patches.values().as_ref().to_primitive().unwrap(); let expected_values = PrimitiveArray::from_iter(vec![PI]); - assert_arrays_eq!(patch_values, expected_values); + assert_arrays_eq!(patch_values.as_ref(), expected_values.as_ref()); } #[test] @@ -342,17 +360,23 @@ mod tests { let encoded = alp_encode(&array, None).unwrap(); let patches = encoded.patches().unwrap(); - let chunk_offsets = patches.chunk_offsets().clone().unwrap().to_primitive(); + let chunk_offsets = patches + .chunk_offsets() + .clone() + .unwrap() + .as_ref() + .to_primitive() + .unwrap(); let expected_offsets = PrimitiveArray::from_iter(vec![0u64]); - assert_arrays_eq!(chunk_offsets, expected_offsets); + assert_arrays_eq!(chunk_offsets.as_ref(), expected_offsets.as_ref()); - let patch_indices = patches.indices().to_primitive(); + let patch_indices = patches.indices().as_ref().to_primitive().unwrap(); let expected_indices = PrimitiveArray::from_iter(vec![0u64, 100]); - assert_arrays_eq!(patch_indices, expected_indices); + assert_arrays_eq!(patch_indices.as_ref(), expected_indices.as_ref()); - let patch_values = patches.values().to_primitive(); + let patch_values = patches.values().as_ref().to_primitive().unwrap(); let expected_values = PrimitiveArray::from_iter(vec![PI, E]); - assert_arrays_eq!(patch_values, expected_values); + assert_arrays_eq!(patch_values.as_ref(), expected_values.as_ref()); } #[test] @@ -363,9 +387,9 @@ mod tests { let encoded = alp_encode(&original, None).unwrap(); let sliced_alp = encoded.slice(512..1024); - let decoded = sliced_alp.to_primitive(); + let decoded = sliced_alp.to_primitive().unwrap(); - let expected_slice = original.slice(512..1024).to_primitive(); + let expected_slice = original.slice(512..1024).as_ref().to_primitive().unwrap(); assert_eq!(expected_slice.as_slice::(), decoded.as_slice::()); } @@ -376,9 +400,9 @@ mod tests { let encoded = alp_encode(&original, None).unwrap(); let sliced_alp = encoded.slice(512..1024); - let decoded = sliced_alp.to_primitive(); + let decoded = sliced_alp.to_primitive().unwrap(); - let expected_slice = original.slice(512..1024).to_primitive(); + let expected_slice = original.slice(512..1024).as_ref().to_primitive().unwrap(); assert_eq!(expected_slice.as_slice::(), decoded.as_slice::()); } @@ -393,9 +417,9 @@ mod tests { let encoded = alp_encode(&original, None).unwrap(); let sliced_alp = encoded.slice(512..1024); - let decoded = sliced_alp.to_primitive(); + let decoded = sliced_alp.to_primitive().unwrap(); - let expected_slice = original.slice(512..1024).to_primitive(); + let expected_slice = original.slice(512..1024).as_ref().to_primitive().unwrap(); assert_eq!(expected_slice.as_slice::(), decoded.as_slice::()); assert!(encoded.patches().is_some()); } @@ -414,9 +438,9 @@ mod tests { let encoded = alp_encode(&original, None).unwrap(); let sliced_alp = encoded.slice(1023..1025); - let decoded = sliced_alp.to_primitive(); + let decoded = sliced_alp.to_primitive().unwrap(); - let expected_slice = original.slice(1023..1025).to_primitive(); + let expected_slice = original.slice(1023..1025).as_ref().to_primitive().unwrap(); assert_eq!(expected_slice.as_slice::(), decoded.as_slice::()); assert!(encoded.patches().is_some()); } @@ -431,10 +455,10 @@ mod tests { let encoded = alp_encode(&original, None).unwrap(); let sliced_alp = encoded.slice(512..1024); - let decoded = sliced_alp.to_primitive(); + let decoded = sliced_alp.to_primitive().unwrap(); let expected_slice = original.slice(512..1024); - assert_arrays_eq!(decoded, expected_slice); + assert_arrays_eq!(decoded.as_ref(), expected_slice.as_ref()); } #[test] @@ -444,7 +468,7 @@ mod tests { let encoded = alp_encode(&array, None).unwrap(); assert!(encoded.patches().is_none()); - let decoded = decompress_into_array(encoded); + let decoded = decompress_into_array(encoded).unwrap(); assert_eq!(array.as_slice::(), decoded.as_slice::()); } @@ -455,7 +479,7 @@ mod tests { let encoded = alp_encode(&array, None).unwrap(); assert!(encoded.patches().is_none()); - let decoded = decompress_into_array(encoded); + let decoded = decompress_into_array(encoded).unwrap(); assert_eq!(array.as_slice::(), decoded.as_slice::()); } @@ -472,7 +496,7 @@ mod tests { let encoded = alp_encode(&array, None).unwrap(); assert!(encoded.patches().is_some()); - let decoded = decompress_into_array(encoded); + let decoded = decompress_into_array(encoded).unwrap(); assert_eq!(values.as_slice(), decoded.as_slice::()); } @@ -493,7 +517,7 @@ mod tests { let encoded = alp_encode(&array, None).unwrap(); assert!(encoded.patches().is_some()); - let decoded = decompress_into_array(encoded); + let decoded = decompress_into_array(encoded).unwrap(); for idx in 0..size { let decoded_val = decoded.as_slice::()[idx]; @@ -520,9 +544,9 @@ mod tests { let array = PrimitiveArray::from_option_iter(values); let encoded = alp_encode(&array, None).unwrap(); - let decoded = decompress_into_array(encoded); + let decoded = decompress_into_array(encoded).unwrap(); - assert_arrays_eq!(decoded, array); + assert_arrays_eq!(decoded.as_ref(), array.as_ref()); } #[test] @@ -540,8 +564,8 @@ mod tests { let array = PrimitiveArray::new(Buffer::from(values), validity); let encoded = alp_encode(&array, None).unwrap(); - let decoded = decompress_into_array(encoded); + let decoded = decompress_into_array(encoded).unwrap(); - assert_arrays_eq!(decoded, array); + assert_arrays_eq!(decoded.as_ref(), array.as_ref()); } } diff --git a/encodings/alp/src/alp/compute/between.rs b/encodings/alp/src/alp/compute/between.rs index d9f7aee4d59..9512419ea05 100644 --- a/encodings/alp/src/alp/compute/between.rs +++ b/encodings/alp/src/alp/compute/between.rs @@ -114,7 +114,9 @@ mod tests { fn between_test(arr: &ALPArray, lower: f32, upper: f32, options: &BetweenOptions) -> bool { let res = between_impl(arr, lower, upper, Nullability::Nullable, options) .unwrap() + .as_ref() .to_bool() + .unwrap() .bit_buffer() .iter() .collect_vec(); @@ -130,7 +132,12 @@ mod tests { let encoded = alp_encode(&array, None).unwrap(); assert!(encoded.patches().is_none()); assert_eq!( - encoded.encoded().to_primitive().as_slice::(), + encoded + .encoded() + .as_ref() + .to_primitive() + .unwrap() + .as_slice::(), vec![605; 1] ); diff --git a/encodings/alp/src/alp/compute/cast.rs b/encodings/alp/src/alp/compute/cast.rs index a653fc35840..36af421ca25 100644 --- a/encodings/alp/src/alp/compute/cast.rs +++ b/encodings/alp/src/alp/compute/cast.rs @@ -68,7 +68,7 @@ mod tests { let values = buffer![1.5f32, 2.5, 3.5, 4.5].into_array(); let alp = ALPVTable .as_vtable() - .encode(&values.to_canonical(), None) + .encode(&values.as_ref().to_canonical().unwrap(), None) .unwrap() .unwrap(); @@ -82,7 +82,7 @@ mod tests { &DType::Primitive(PType::F64, Nullability::NonNullable) ); - let decoded = casted.to_canonical().into_primitive(); + let decoded = casted.as_ref().to_canonical().unwrap().into_primitive(); let values = decoded.as_slice::(); assert_eq!(values.len(), 4); assert!((values[0] - 1.5).abs() < f64::EPSILON); @@ -94,7 +94,7 @@ mod tests { let values = buffer![1.0f32, 2.0, 3.0, 4.0].into_array(); let alp = ALPVTable .as_vtable() - .encode(&values.to_canonical(), None) + .encode(&values.as_ref().to_canonical().unwrap(), None) .unwrap() .unwrap(); @@ -108,8 +108,9 @@ mod tests { &DType::Primitive(PType::I32, Nullability::NonNullable) ); - let decoded = casted.to_canonical().into_primitive(); - assert_arrays_eq!(decoded, PrimitiveArray::from_iter([1i32, 2, 3, 4])); + let decoded = casted.as_ref().to_canonical().unwrap().into_primitive(); + let expected = PrimitiveArray::from_iter([1i32, 2, 3, 4]); + assert_arrays_eq!(decoded.as_ref(), expected.as_ref()); } #[rstest] @@ -121,7 +122,7 @@ mod tests { fn test_cast_alp_conformance(#[case] array: vortex_array::ArrayRef) { let alp = ALPVTable .as_vtable() - .encode(&array.to_canonical(), None) + .encode(&array.as_ref().to_canonical().unwrap(), None) .unwrap() .unwrap(); test_cast_conformance(alp.as_ref()); diff --git a/encodings/alp/src/alp/compute/compare.rs b/encodings/alp/src/alp/compute/compare.rs index 92a59aa9e42..9088b40c1e4 100644 --- a/encodings/alp/src/alp/compute/compare.rs +++ b/encodings/alp/src/alp/compute/compare.rs @@ -163,7 +163,7 @@ mod tests { { alp_scalar_compare(alp, value, operator) .unwrap() - .map(|a| a.to_bool().bit_buffer().iter().collect()) + .map(|a| a.as_ref().to_bool().unwrap().bit_buffer().iter().collect()) } #[test] @@ -172,14 +172,21 @@ mod tests { let encoded = alp_encode(&array, None).unwrap(); assert!(encoded.patches().is_none()); assert_eq!( - encoded.encoded().to_primitive().as_slice::(), + encoded + .encoded() + .as_ref() + .to_primitive() + .unwrap() + .as_slice::(), vec![1234; 1025] ); let r = alp_scalar_compare(&encoded, 1.3_f32, Operator::Eq) .unwrap() .unwrap() - .to_bool(); + .as_ref() + .to_bool() + .unwrap(); for v in r.bit_buffer().iter() { assert!(!v); @@ -188,7 +195,9 @@ mod tests { let r = alp_scalar_compare(&encoded, 1.234f32, Operator::Eq) .unwrap() .unwrap() - .to_bool(); + .as_ref() + .to_bool() + .unwrap(); for v in r.bit_buffer().iter() { assert!(v); @@ -201,7 +210,12 @@ mod tests { let encoded = alp_encode(&array, None).unwrap(); assert!(encoded.patches().is_none()); assert_eq!( - encoded.encoded().to_primitive().as_slice::(), + encoded + .encoded() + .as_ref() + .to_primitive() + .unwrap() + .as_slice::(), vec![1234; 1025] ); @@ -209,7 +223,9 @@ mod tests { let r_eq = alp_scalar_compare(&encoded, 1.234444_f32, Operator::Eq) .unwrap() .unwrap() - .to_bool(); + .as_ref() + .to_bool() + .unwrap(); assert!(r_eq.bit_buffer().iter().all(|v| !v)); @@ -217,7 +233,9 @@ mod tests { let r_neq = alp_scalar_compare(&encoded, 1.234444f32, Operator::NotEq) .unwrap() .unwrap() - .to_bool(); + .as_ref() + .to_bool() + .unwrap(); assert!(r_neq.bit_buffer().iter().all(|v| v)); } @@ -228,14 +246,21 @@ mod tests { let encoded = alp_encode(&array, None).unwrap(); assert!(encoded.patches().is_none()); assert_eq!( - encoded.encoded().to_primitive().as_slice::(), + encoded + .encoded() + .as_ref() + .to_primitive() + .unwrap() + .as_slice::(), vec![605; 10] ); let r_gte = alp_scalar_compare(&encoded, 0.06051_f32, Operator::Gte) .unwrap() .unwrap() - .to_bool(); + .as_ref() + .to_bool() + .unwrap(); // !(0.0605_f32 >= 0.06051_f32); assert!(r_gte.bit_buffer().iter().all(|v| !v)); @@ -243,7 +268,9 @@ mod tests { let r_gt = alp_scalar_compare(&encoded, 0.06051_f32, Operator::Gt) .unwrap() .unwrap() - .to_bool(); + .as_ref() + .to_bool() + .unwrap(); // (0.0605_f32 > 0.06051_f32); assert!(r_gt.bit_buffer().iter().all(|v| !v)); @@ -251,7 +278,9 @@ mod tests { let r_lte = alp_scalar_compare(&encoded, 0.06051_f32, Operator::Lte) .unwrap() .unwrap() - .to_bool(); + .as_ref() + .to_bool() + .unwrap(); // 0.0605_f32 <= 0.06051_f32; assert!(r_lte.bit_buffer().iter().all(|v| v)); @@ -259,7 +288,9 @@ mod tests { let r_lt = alp_scalar_compare(&encoded, 0.06051_f32, Operator::Lt) .unwrap() .unwrap() - .to_bool(); + .as_ref() + .to_bool() + .unwrap(); //0.0605_f32 < 0.06051_f32; assert!(r_lt.bit_buffer().iter().all(|v| v)); @@ -271,7 +302,12 @@ mod tests { let encoded = alp_encode(&array, None).unwrap(); assert!(encoded.patches().is_none()); assert_eq!( - encoded.encoded().to_primitive().as_slice::(), + encoded + .encoded() + .as_ref() + .to_primitive() + .unwrap() + .as_slice::(), vec![0; 10] ); @@ -324,7 +360,9 @@ mod tests { let r = compare(encoded.as_ref(), other.as_ref(), Operator::Eq) .unwrap() - .to_bool(); + .as_ref() + .to_bool() + .unwrap(); for v in r.bit_buffer().iter() { assert!(!v); diff --git a/encodings/alp/src/alp/compute/filter.rs b/encodings/alp/src/alp/compute/filter.rs index 891e943c37a..4558038f71e 100644 --- a/encodings/alp/src/alp/compute/filter.rs +++ b/encodings/alp/src/alp/compute/filter.rs @@ -59,7 +59,7 @@ mod test { fn test_filter_alp_conformance(#[case] array: ArrayRef) { let alp = ALPVTable .as_vtable() - .encode(&array.to_canonical(), None) + .encode(&array.as_ref().to_canonical().unwrap(), None) .unwrap() .unwrap(); test_filter_conformance(alp.as_ref()); diff --git a/encodings/alp/src/alp/compute/mask.rs b/encodings/alp/src/alp/compute/mask.rs index 364e5c09f09..96d36c50c43 100644 --- a/encodings/alp/src/alp/compute/mask.rs +++ b/encodings/alp/src/alp/compute/mask.rs @@ -57,7 +57,7 @@ mod test { fn test_mask_alp_conformance(#[case] array: vortex_array::ArrayRef) { let alp = ALPVTable .as_vtable() - .encode(&array.to_canonical(), None) + .encode(&array.as_ref().to_canonical().unwrap(), None) .unwrap() .unwrap(); test_mask_conformance(alp.as_ref()); diff --git a/encodings/alp/src/alp/compute/take.rs b/encodings/alp/src/alp/compute/take.rs index b24cbe7e455..cf4cb1a4dcf 100644 --- a/encodings/alp/src/alp/compute/take.rs +++ b/encodings/alp/src/alp/compute/take.rs @@ -54,7 +54,7 @@ mod test { fn test_take_alp_conformance(#[case] array: vortex_array::ArrayRef) { let alp = ALPVTable .as_vtable() - .encode(&array.to_canonical(), None) + .encode(&array.as_ref().to_canonical().unwrap(), None) .unwrap() .unwrap(); test_take_conformance(alp.as_ref()); diff --git a/encodings/alp/src/alp/decompress.rs b/encodings/alp/src/alp/decompress.rs index 84e33517b87..e8e2a5d76bc 100644 --- a/encodings/alp/src/alp/decompress.rs +++ b/encodings/alp/src/alp/decompress.rs @@ -30,7 +30,7 @@ use crate::match_each_alp_float_ptype; /// # Returns /// /// A `PrimitiveArray` containing the decompressed floating-point values with all patches applied. -pub fn decompress_into_array(array: ALPArray) -> PrimitiveArray { +pub fn decompress_into_array(array: ALPArray) -> VortexResult { let (encoded, exponents, patches, dtype) = array.into_parts(); if let Some(ref patches) = patches && let Some(chunk_offsets) = patches.chunk_offsets() @@ -39,7 +39,7 @@ pub fn decompress_into_array(array: ALPArray) -> PrimitiveArray { encoded, exponents, patches, - &chunk_offsets.as_ref().to_primitive(), + &chunk_offsets.as_ref().to_primitive()?, dtype, ) } else { @@ -102,15 +102,15 @@ fn decompress_chunked( patches: &Patches, patches_chunk_offsets: &PrimitiveArray, dtype: DType, -) -> PrimitiveArray { - let encoded = array.to_primitive(); +) -> VortexResult { + let encoded = array.to_primitive()?; let validity = encoded.validity().clone(); - let patches_indices = patches.indices().to_primitive(); - let patches_values = patches.values().to_primitive(); + let patches_indices = patches.indices().to_primitive()?; + let patches_values = patches.values().to_primitive()?; let ptype = dtype.as_ptype(); - let array_len = array.len(); + let array_len = encoded.len(); let patches_offset = patches.offset(); // We need to drop ALPArray here in case converting encoded buffer into @@ -118,7 +118,7 @@ fn decompress_chunked( // will hold a reference to the buffer we want to mutate. drop(array); - match_each_alp_float_ptype!(ptype, |T| { + Ok(match_each_alp_float_ptype!(ptype, |T| { let patches_values = patches_values.as_slice::(); let mut alp_buffer = encoded.into_buffer_mut(); match_each_unsigned_integer_ptype!(patches_chunk_offsets.ptype(), |C| { @@ -153,7 +153,7 @@ fn decompress_chunked( PrimitiveArray::new::(decoded_buffer.freeze(), validity) }) }) - }) + })) } /// Decompresses an ALP-encoded array without chunk offsets. @@ -166,8 +166,8 @@ fn decompress_unchunked( exponents: Exponents, patches: Option, dtype: DType, -) -> PrimitiveArray { - let encoded = array.to_primitive(); +) -> VortexResult { + let encoded = array.to_primitive()?; // We need to drop ALPArray here in case converting encoded buffer into // primitive didn't create a copy. In that case both alp_encoded and array @@ -184,9 +184,9 @@ fn decompress_unchunked( PrimitiveArray::new::(decoded_buffer.freeze(), validity) }); - if let Some(patches) = patches { + Ok(if let Some(patches) = patches { decoded.patch(&patches) } else { decoded - } + }) } diff --git a/encodings/alp/src/alp_rd/array.rs b/encodings/alp/src/alp_rd/array.rs index cb36e4f5e5d..bf3c18e37c2 100644 --- a/encodings/alp/src/alp_rd/array.rs +++ b/encodings/alp/src/alp_rd/array.rs @@ -243,7 +243,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. @@ -368,9 +368,9 @@ impl BaseArrayVTable for ALPRDVTable { } impl CanonicalVTable for ALPRDVTable { - fn canonicalize(array: &ALPRDArray) -> Canonical { - let left_parts = array.left_parts().to_primitive(); - let right_parts = array.right_parts().to_primitive(); + fn canonicalize(array: &ALPRDArray) -> VortexResult { + let left_parts = array.left_parts().to_primitive()?; + let right_parts = array.right_parts().to_primitive()?; // Decode the left_parts using our builtin dictionary. let left_parts_dict = array.left_parts_dictionary(); @@ -383,7 +383,7 @@ impl CanonicalVTable for ALPRDVTable { array.right_bit_width, right_parts.into_buffer_mut::(), array.left_parts_patches(), - ), + )?, Validity::copy_from_array(array.as_ref()), ) } else { @@ -394,12 +394,12 @@ impl CanonicalVTable for ALPRDVTable { array.right_bit_width, right_parts.into_buffer_mut::(), array.left_parts_patches(), - ), + )?, Validity::copy_from_array(array.as_ref()), ) }; - Canonical::Primitive(decoded_array) + Ok(Canonical::Primitive(decoded_array)) } } @@ -482,9 +482,10 @@ mod test { let rd_array = encoder.encode(&real_array); - let decoded = rd_array.to_primitive(); + let decoded = rd_array.to_primitive().unwrap(); + let expected = PrimitiveArray::from_option_iter(reals); - assert_arrays_eq!(decoded, PrimitiveArray::from_option_iter(reals)); + assert_arrays_eq!(decoded.as_ref(), expected.as_ref()); } #[cfg_attr(miri, ignore)] diff --git a/encodings/alp/src/alp_rd/compute/cast.rs b/encodings/alp/src/alp_rd/compute/cast.rs index 7ede0ef8a91..50e0165b95d 100644 --- a/encodings/alp/src/alp_rd/compute/cast.rs +++ b/encodings/alp/src/alp_rd/compute/cast.rs @@ -80,7 +80,7 @@ mod tests { &DType::Primitive(PType::F64, Nullability::NonNullable) ); - let decoded = casted.to_primitive(); + let decoded = casted.as_ref().to_primitive().unwrap(); let f64_values = decoded.as_slice::(); assert_eq!(f64_values.len(), 5); assert!((f64_values[0] - 1.0).abs() < f64::EPSILON); diff --git a/encodings/alp/src/alp_rd/compute/filter.rs b/encodings/alp/src/alp_rd/compute/filter.rs index be3a9109244..df9dc8fb641 100644 --- a/encodings/alp/src/alp_rd/compute/filter.rs +++ b/encodings/alp/src/alp_rd/compute/filter.rs @@ -64,8 +64,11 @@ mod test { // The first two values need no patching let filtered = filter(encoded.as_ref(), &Mask::from_iter([true, false, true])) .unwrap() - .to_primitive(); - assert_arrays_eq!(filtered, PrimitiveArray::from_iter([a, outlier])); + .as_ref() + .to_primitive() + .unwrap(); + let expected = PrimitiveArray::from_iter([a, outlier]); + assert_arrays_eq!(filtered.as_ref(), expected.as_ref()); } #[rstest] diff --git a/encodings/alp/src/alp_rd/compute/take.rs b/encodings/alp/src/alp_rd/compute/take.rs index 1b43c0c423f..86063882de8 100644 --- a/encodings/alp/src/alp_rd/compute/take.rs +++ b/encodings/alp/src/alp_rd/compute/take.rs @@ -86,9 +86,12 @@ mod test { let taken = take(encoded.as_ref(), buffer![0, 2].into_array().as_ref()) .unwrap() - .to_primitive(); + .as_ref() + .to_primitive() + .unwrap(); - assert_arrays_eq!(taken, PrimitiveArray::from_iter([a, outlier])); + let expected = PrimitiveArray::from_iter([a, outlier]); + assert_arrays_eq!(taken.as_ref(), expected.as_ref()); } #[rstest] @@ -112,12 +115,12 @@ mod test { PrimitiveArray::from_option_iter([Some(0), Some(2), None]).as_ref(), ) .unwrap() - .to_primitive(); + .as_ref() + .to_primitive() + .unwrap(); - assert_arrays_eq!( - taken, - PrimitiveArray::from_option_iter([Some(a), Some(outlier), None]) - ); + let expected = PrimitiveArray::from_option_iter([Some(a), Some(outlier), None]); + assert_arrays_eq!(taken.as_ref(), expected.as_ref()); } #[rstest] diff --git a/encodings/alp/src/alp_rd/mod.rs b/encodings/alp/src/alp_rd/mod.rs index 6cc781aa45f..eba8972f01f 100644 --- a/encodings/alp/src/alp_rd/mod.rs +++ b/encodings/alp/src/alp_rd/mod.rs @@ -31,6 +31,7 @@ use vortex_dtype::DType; use vortex_dtype::NativePType; use vortex_dtype::match_each_integer_ptype; use vortex_error::VortexExpect; +use vortex_error::VortexResult; use vortex_error::VortexUnwrap; use vortex_error::vortex_panic; use vortex_utils::aliases::hash_map::HashMap; @@ -289,7 +290,7 @@ pub fn alp_rd_decode( right_bit_width: u8, right_parts: BufferMut, left_parts_patches: Option<&Patches>, -) -> Buffer { +) -> VortexResult> { if left_parts.len() != right_parts.len() { vortex_panic!("alp_rd_decode: left_parts.len != right_parts.len"); } @@ -303,8 +304,8 @@ pub fn alp_rd_decode( // Apply any patches if let Some(patches) = left_parts_patches { - let indices = patches.indices().to_primitive(); - let patch_values = patches.values().to_primitive(); + let indices = patches.indices().to_primitive()?; + let patch_values = patches.values().to_primitive()?; match_each_integer_ptype!(indices.ptype(), |T| { indices .as_slice::() @@ -318,14 +319,14 @@ pub fn alp_rd_decode( // Shift the left-parts and add in the right-parts. let mut index = 0; - right_parts + Ok(right_parts .map_each_in_place(|right| { let left = values[index]; index += 1; let left = ::from_u16(left); T::from_bits((left << (right_bit_width as usize)) | right) }) - .freeze() + .freeze()) } /// Find the best "cut point" for a set of floating point values such that we can diff --git a/encodings/bytebool/src/array.rs b/encodings/bytebool/src/array.rs index 46968cb0977..1f582c51beb 100644 --- a/encodings/bytebool/src/array.rs +++ b/encodings/bytebool/src/array.rs @@ -187,10 +187,13 @@ impl BaseArrayVTable for ByteBoolVTable { } impl CanonicalVTable for ByteBoolVTable { - fn canonicalize(array: &ByteBoolArray) -> Canonical { + fn canonicalize(array: &ByteBoolArray) -> VortexResult { let boolean_buffer = BitBuffer::from(array.as_slice()); let validity = array.validity().clone(); - Canonical::Bool(BoolArray::from_bit_buffer(boolean_buffer, validity)) + Ok(Canonical::Bool(BoolArray::from_bit_buffer( + boolean_buffer, + validity, + ))) } } @@ -259,14 +262,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]; @@ -276,7 +279,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/bytebool/src/compute.rs b/encodings/bytebool/src/compute.rs index 011b9b23e3d..9c68e62f5df 100644 --- a/encodings/bytebool/src/compute.rs +++ b/encodings/bytebool/src/compute.rs @@ -49,7 +49,7 @@ register_kernel!(CastKernelAdapter(ByteBoolVTable).lift()); impl MaskKernel for ByteBoolVTable { fn mask(&self, array: &ByteBoolArray, mask: &Mask) -> VortexResult { - Ok(ByteBoolArray::new(array.buffer().clone(), array.validity().mask(mask)).into_array()) + Ok(ByteBoolArray::new(array.buffer().clone(), array.validity().mask(mask)?).into_array()) } } @@ -57,7 +57,7 @@ register_kernel!(MaskKernelAdapter(ByteBoolVTable).lift()); impl TakeKernel for ByteBoolVTable { fn take(&self, array: &ByteBoolArray, indices: &dyn Array) -> VortexResult { - let indices = indices.to_primitive(); + let indices = indices.to_primitive()?; let bools = array.as_slice(); // This handles combining validity from both source array and nullable indices diff --git a/encodings/datetime-parts/src/canonical.rs b/encodings/datetime-parts/src/canonical.rs index d829c386bf7..ba12808a058 100644 --- a/encodings/datetime-parts/src/canonical.rs +++ b/encodings/datetime-parts/src/canonical.rs @@ -17,21 +17,22 @@ use vortex_dtype::datetime::TemporalMetadata; use vortex_dtype::datetime::TimeUnit; use vortex_dtype::match_each_integer_ptype; use vortex_error::VortexExpect as _; +use vortex_error::VortexResult; use vortex_error::vortex_panic; use crate::DateTimePartsArray; use crate::DateTimePartsVTable; impl CanonicalVTable for DateTimePartsVTable { - fn canonicalize(array: &DateTimePartsArray) -> Canonical { - Canonical::Extension(decode_to_temporal(array).into()) + fn canonicalize(array: &DateTimePartsArray) -> VortexResult { + Ok(Canonical::Extension(decode_to_temporal(array)?.into())) } } /// Decode an [Array] into a [TemporalArray]. /// /// Enforces that the passed array is actually a [DateTimePartsArray] with proper metadata. -pub fn decode_to_temporal(array: &DateTimePartsArray) -> TemporalArray { +pub fn decode_to_temporal(array: &DateTimePartsArray) -> VortexResult { let DType::Extension(ext) = array.dtype().clone() else { vortex_panic!(ComputeError: "expected dtype to be DType::Extension variant") }; @@ -53,7 +54,7 @@ pub fn decode_to_temporal(array: &DateTimePartsArray) -> TemporalArray { &DType::Primitive(PType::I64, array.dtype().nullability()), ) .vortex_expect("must be able to cast days to i64") - .to_primitive(); + .to_primitive()?; // We start with the days component, which is always present. // And then add the seconds and subseconds components. @@ -73,7 +74,7 @@ pub fn decode_to_temporal(array: &DateTimePartsArray) -> TemporalArray { *v += seconds; } } else { - let seconds_buf = array.seconds().to_primitive(); + let seconds_buf = array.seconds().to_primitive()?; match_each_integer_ptype!(seconds_buf.ptype(), |S| { for (v, second) in values.iter_mut().zip(seconds_buf.as_slice::()) { let second: i64 = second.as_(); @@ -91,7 +92,7 @@ pub fn decode_to_temporal(array: &DateTimePartsArray) -> TemporalArray { *v += subseconds; } } else { - let subseconds_buf = array.subseconds().to_primitive(); + let subseconds_buf = array.subseconds().to_primitive()?; match_each_integer_ptype!(subseconds_buf.ptype(), |S| { for (v, subseconds) in values.iter_mut().zip(subseconds_buf.as_slice::()) { let subseconds: i64 = subseconds.as_(); @@ -100,12 +101,12 @@ pub fn decode_to_temporal(array: &DateTimePartsArray) -> TemporalArray { }); } - TemporalArray::new_timestamp( + Ok(TemporalArray::new_timestamp( 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), - ) + )) } #[cfg(test)] @@ -119,6 +120,7 @@ mod test { use vortex_array::vtable::ValidityHelper; use vortex_buffer::buffer; use vortex_dtype::datetime::TimeUnit; + use vortex_error::VortexResult; use crate::DateTimePartsArray; use crate::canonical::decode_to_temporal; @@ -128,7 +130,7 @@ mod test { #[case(Validity::AllValid)] #[case(Validity::AllInvalid)] #[case(Validity::from_iter([true, true, false, false, true, true]))] - fn test_decode_to_temporal(#[case] validity: Validity) { + fn test_decode_to_temporal(#[case] validity: Validity) -> VortexResult<()> { let milliseconds = PrimitiveArray::new( buffer![ 86_400i64, // element with only day component @@ -148,18 +150,20 @@ mod test { .unwrap(); assert_eq!( - date_times.validity_mask(), - validity.to_mask(date_times.len()) + date_times.validity_mask()?, + validity.to_mask(date_times.len())? ); let primitive_values = decode_to_temporal(&date_times) + .unwrap() .temporal_values() - .to_primitive(); + .to_primitive()?; assert_eq!( primitive_values.as_slice::(), milliseconds.as_slice::() ); assert_eq!(primitive_values.validity(), &validity); + Ok(()) } } diff --git a/encodings/datetime-parts/src/compress.rs b/encodings/datetime-parts/src/compress.rs index 441166f45e8..5df1212103a 100644 --- a/encodings/datetime-parts/src/compress.rs +++ b/encodings/datetime-parts/src/compress.rs @@ -28,14 +28,14 @@ pub struct TemporalParts { /// Splitting the components by granularity creates more small values, which enables better /// cascading compression. pub fn split_temporal(array: TemporalArray) -> VortexResult { - let temporal_values = array.temporal_values().to_primitive(); + let temporal_values = array.temporal_values().to_primitive()?; // After this operation, timestamps will be a PrimitiveArray let timestamps = cast( temporal_values.as_ref(), &DType::Primitive(PType::I64, temporal_values.dtype().nullability()), )? - .to_primitive(); + .to_primitive()?; let length = timestamps.len(); let mut days = BufferMut::with_capacity(length); @@ -110,8 +110,14 @@ mod tests { seconds, subseconds, } = split_temporal(temporal_array).unwrap(); - assert_eq!(days.to_primitive().validity(), &validity); - assert_eq!(seconds.to_primitive().validity(), &Validity::NonNullable); - assert_eq!(subseconds.to_primitive().validity(), &Validity::NonNullable); + assert_eq!(days.to_primitive().unwrap().validity(), &validity); + assert_eq!( + seconds.to_primitive().unwrap().validity(), + &Validity::NonNullable + ); + assert_eq!( + subseconds.to_primitive().unwrap().validity(), + &Validity::NonNullable + ); } } diff --git a/encodings/datetime-parts/src/compute/take.rs b/encodings/datetime-parts/src/compute/take.rs index c5e1f62f461..0e9da541f64 100644 --- a/encodings/datetime-parts/src/compute/take.rs +++ b/encodings/datetime-parts/src/compute/take.rs @@ -23,7 +23,7 @@ use crate::DateTimePartsVTable; impl TakeKernel for DateTimePartsVTable { fn take(&self, array: &DateTimePartsArray, indices: &dyn Array) -> VortexResult { // we go ahead and canonicalize here to avoid worst-case canonicalizing 3 separate times - let indices = indices.to_primitive(); + let indices = indices.to_primitive()?; let taken_days = take(array.days(), indices.as_ref())?; let taken_seconds = take(array.seconds(), indices.as_ref())?; @@ -75,7 +75,7 @@ impl TakeKernel for DateTimePartsVTable { .map(|s| s.into_inner()) .unwrap_or_else(|| Scalar::primitive(0i64, Nullability::NonNullable)) .cast(array.seconds().dtype())?; - let taken_seconds = fill_null(taken_seconds.as_ref(), &seconds_fill)?; + let taken_seconds = fill_null(&taken_seconds, &seconds_fill)?; let subseconds_fill = array .subseconds() @@ -84,7 +84,7 @@ impl TakeKernel for DateTimePartsVTable { .map(|s| s.into_inner()) .unwrap_or_else(|| Scalar::primitive(0i64, Nullability::NonNullable)) .cast(array.subseconds().dtype())?; - let taken_subseconds = fill_null(taken_subseconds.as_ref(), &subseconds_fill)?; + let taken_subseconds = fill_null(&taken_subseconds, &subseconds_fill)?; Ok( DateTimePartsArray::try_new(dtype, taken_days, taken_seconds, taken_subseconds)? diff --git a/encodings/datetime-parts/src/ops.rs b/encodings/datetime-parts/src/ops.rs index 994d2458751..3a591db3d29 100644 --- a/encodings/datetime-parts/src/ops.rs +++ b/encodings/datetime-parts/src/ops.rs @@ -44,7 +44,7 @@ impl OperationsVTable for DateTimePartsVTable { vortex_panic!(ComputeError: "must decode TemporalMetadata from extension metadata"); }; - if !array.is_valid(index) { + if !array.is_valid(index).vortex_expect("is_valid") { return Scalar::null(DType::Extension(ext)); } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs index d51a7d15ed8..0035807bb94 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs @@ -79,7 +79,7 @@ mod tests { ); // Verify the values are preserved - let decoded = casted.to_decimal(); + let decoded = casted.to_decimal().unwrap(); assert_eq!(decoded.len(), 4); } 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 e9c798d2425..3fac710e9f5 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -205,13 +205,13 @@ impl BaseArrayVTable for DecimalBytePartsVTable { } impl CanonicalVTable for DecimalBytePartsVTable { - fn canonicalize(array: &DecimalBytePartsArray) -> Canonical { + fn canonicalize(array: &DecimalBytePartsArray) -> VortexResult { // TODO(joe): support parts len != 1 - let prim = array.msp.to_primitive(); + let prim = array.msp.to_primitive()?; // Depending on the decimal type and the min/max of the primitive array we can choose // the correct buffer size - match_each_signed_integer_ptype!(prim.ptype(), |P| { + Ok(match_each_signed_integer_ptype!(prim.ptype(), |P| { // SAFETY: The primitive array's buffer is already validated with correct type. // The decimal dtype matches the array's dtype, and validity is preserved. Canonical::Decimal(unsafe { @@ -221,7 +221,7 @@ impl CanonicalVTable for DecimalBytePartsVTable { prim.validity().clone(), ) }) - }) + })) } } diff --git a/encodings/fastlanes/benches/bitpacking_decompress_selection.rs b/encodings/fastlanes/benches/bitpacking_decompress_selection.rs index e3544d673cb..dc756820e0f 100644 --- a/encodings/fastlanes/benches/bitpacking_decompress_selection.rs +++ b/encodings/fastlanes/benches/bitpacking_decompress_selection.rs @@ -36,7 +36,8 @@ fn decompress_bitpacking_early_filter(bencher: Bencher, fractio .map(|_| T::from(rng.random_range(0..100)).unwrap()) .collect::>() .into_array() - .to_primitive(); + .to_primitive() + .unwrap(); let array = bitpack_to_best_bit_width(&values).unwrap(); let mask = (0..100_000) .map(|_| rng.random_bool(fraction_kept)) @@ -56,7 +57,8 @@ fn decompress_bitpacking_late_filter(bencher: Bencher, fraction .map(|_| T::from(rng.random_range(0..100)).unwrap()) .collect::>() .into_array() - .to_primitive(); + .to_primitive() + .unwrap(); let array = bitpack_to_best_bit_width(&values).unwrap(); @@ -67,5 +69,5 @@ fn decompress_bitpacking_late_filter(bencher: Bencher, fraction bencher // Be sure to reconstruct the mask to avoid cached set_indices .with_inputs(|| (&array, Mask::from_buffer(mask.clone()))) - .bench_refs(|(array, mask)| filter(array.to_canonical().as_ref(), mask).unwrap()); + .bench_refs(|(array, mask)| filter(array.to_canonical().unwrap().as_ref(), mask).unwrap()); } diff --git a/encodings/fastlanes/benches/canonicalize_bench.rs b/encodings/fastlanes/benches/canonicalize_bench.rs index ba43f4906c6..a7c352b64f6 100644 --- a/encodings/fastlanes/benches/canonicalize_bench.rs +++ b/encodings/fastlanes/benches/canonicalize_bench.rs @@ -70,7 +70,9 @@ fn canonical_into_non_nullable( chunked.dtype().nullability(), chunk_len * chunk_count, ); - chunked.append_to_builder(&mut primitive_builder); + chunked + .append_to_builder(&mut primitive_builder) + .vortex_expect("append_to_builder"); primitive_builder.finish() }); } @@ -125,7 +127,9 @@ fn canonical_into_nullable( chunked.dtype().nullability(), chunk_len * chunk_count, ); - chunked.append_to_builder(&mut primitive_builder); + chunked + .append_to_builder(&mut primitive_builder) + .vortex_expect("append_to_builder"); primitive_builder.finish() }); } diff --git a/encodings/fastlanes/benches/compute_between.rs b/encodings/fastlanes/benches/compute_between.rs index 011936af5b2..cd4e5d33915 100644 --- a/encodings/fastlanes/benches/compute_between.rs +++ b/encodings/fastlanes/benches/compute_between.rs @@ -50,7 +50,7 @@ fn generate_alp_bit_pack_primitive_array( let alp = alp_encode(&a, None).vortex_expect(""); - let encoded = alp.encoded().to_primitive(); + let encoded = alp.encoded().to_primitive().vortex_expect(""); let bp = bitpack_to_best_bit_width(&encoded) .vortex_expect("") diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs index e9a1619c37a..bf361e1c71b 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs @@ -200,7 +200,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| { @@ -300,7 +300,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::() { @@ -396,7 +396,7 @@ pub mod test_harness { .collect::>(); let values = if fraction_null == 0.0 { - values.into_array().to_primitive() + values.into_array().to_primitive()? } else { let validity = Validity::from_iter((0..len).map(|_| !rng.random_bool(fraction_null))); PrimitiveArray::new(values, validity) @@ -446,6 +446,7 @@ mod test { (0..(1 << 4)).collect::>(), compressed .validity_mask() + .unwrap() .to_bit_buffer() .set_indices() .collect::>() @@ -471,20 +472,23 @@ mod test { .collect::>(); let chunked = ChunkedArray::from_iter(chunks).into_array(); - let into_ca = chunked.clone().to_primitive(); + let into_ca = chunked.clone().to_primitive().unwrap(); let mut primitive_builder = PrimitiveBuilder::::with_capacity(chunked.dtype().nullability(), 10 * 100); - chunked.clone().append_to_builder(&mut primitive_builder); + chunked + .clone() + .append_to_builder(&mut primitive_builder) + .unwrap(); let ca_into = primitive_builder.finish(); - assert_arrays_eq!(into_ca, ca_into.to_primitive()); + assert_arrays_eq!(into_ca, ca_into.to_primitive().unwrap()); let mut primitive_builder = PrimitiveBuilder::::with_capacity(chunked.dtype().nullability(), 10 * 100); primitive_builder.extend_from_array(&chunked); let ca_into = primitive_builder.finish(); - assert_arrays_eq!(into_ca, ca_into.to_primitive()); + assert_arrays_eq!(into_ca, ca_into.to_primitive().unwrap()); } #[test] @@ -501,7 +505,12 @@ mod test { let bitpacked = bitpack_encode(&array, 4, None).unwrap(); let patches = bitpacked.patches().unwrap(); - let chunk_offsets = patches.chunk_offsets().as_ref().unwrap().to_primitive(); + let chunk_offsets = patches + .chunk_offsets() + .as_ref() + .unwrap() + .to_primitive() + .unwrap(); // chunk 0 (0-1023): patches at 100, 200 -> starts at patch index 0 // chunk 1 (1024-2047): no patches -> points to patch index 2 @@ -524,7 +533,12 @@ mod test { let bitpacked = bitpack_encode(&array, 4, None).unwrap(); let patches = bitpacked.patches().unwrap(); - let chunk_offsets = patches.chunk_offsets().as_ref().unwrap().to_primitive(); + let chunk_offsets = patches + .chunk_offsets() + .as_ref() + .unwrap() + .to_primitive() + .unwrap(); assert_arrays_eq!(chunk_offsets, PrimitiveArray::from_iter([0u64, 2, 2])); } @@ -543,7 +557,12 @@ mod test { let bitpacked = bitpack_encode(&array, 4, None).unwrap(); let patches = bitpacked.patches().unwrap(); - let chunk_offsets = patches.chunk_offsets().as_ref().unwrap().to_primitive(); + let chunk_offsets = patches + .chunk_offsets() + .as_ref() + .unwrap() + .to_primitive() + .unwrap(); // chunk 0 (0-1023): patches at 100, 200 -> starts at patch index 0 // chunk 1 (1024-2047): patch at 1500 -> starts at patch index 2 @@ -567,7 +586,12 @@ mod test { let bitpacked = bitpack_encode(&array, 4, None).unwrap(); let patches = bitpacked.patches().unwrap(); - let chunk_offsets = patches.chunk_offsets().as_ref().unwrap().to_primitive(); + let chunk_offsets = patches + .chunk_offsets() + .as_ref() + .unwrap() + .to_primitive() + .unwrap(); // Single chunk starting at patch index 0. assert_arrays_eq!(chunk_offsets, PrimitiveArray::from_iter([0u64])); diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs index 7c6cb0eb0b3..28d8e3077ad 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs @@ -15,6 +15,7 @@ use vortex_dtype::NativePType; use vortex_dtype::match_each_integer_ptype; use vortex_dtype::match_each_unsigned_integer_ptype; use vortex_error::VortexExpect; +use vortex_error::VortexResult; use vortex_error::vortex_panic; use vortex_mask::Mask; use vortex_scalar::Scalar; @@ -25,14 +26,16 @@ use crate::BitPackedArray; use crate::unpack_iter::BitPacked; /// Unpacks a bit-packed array into a primitive vector. -pub fn unpack_to_primitive_vector(array: &BitPackedArray) -> PrimitiveVectorMut { - match_each_integer_ptype!(array.ptype(), |P| { unpack_to_pvector::

(array).into() }) +pub fn unpack_to_primitive_vector(array: &BitPackedArray) -> VortexResult { + match_each_integer_ptype!(array.ptype(), |P| { + Ok(unpack_to_pvector::

(array)?.into()) + }) } /// Unpacks a bit-packed array into a generic [`PVectorMut`]. -pub fn unpack_to_pvector(array: &BitPackedArray) -> PVectorMut

{ +pub fn unpack_to_pvector(array: &BitPackedArray) -> VortexResult> { if array.is_empty() { - return PVectorMut::with_capacity(0); + return Ok(PVectorMut::with_capacity(0)); } let len = array.len(); @@ -45,7 +48,7 @@ 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()?.into_mut(); debug_assert_eq!(validity.len(), len); // TODO(connor): Implement a fused version of patching instead. @@ -58,7 +61,7 @@ pub fn unpack_to_pvector(array: &BitPackedArray) -> PVectorMut

} // SAFETY: `elements` and `validity` have the same length. - unsafe { PVectorMut::new_unchecked(elements, validity) } + Ok(unsafe { PVectorMut::new_unchecked(elements, validity) }) } pub fn unpack_array(array: &BitPackedArray) -> PrimitiveArray { @@ -67,7 +70,8 @@ pub fn unpack_array(array: &BitPackedArray) -> PrimitiveArray { pub fn unpack_primitive_array(array: &BitPackedArray) -> PrimitiveArray { let mut builder = PrimitiveBuilder::with_capacity(array.dtype().nullability(), array.len()); - unpack_into_primitive_builder::(array, &mut builder); + unpack_into_primitive_builder::(array, &mut builder) + .vortex_expect("unpack_into_primitive_builder cannot fail"); assert_eq!(builder.len(), array.len()); builder.finish_into_primitive() } @@ -76,10 +80,10 @@ pub(crate) fn unpack_into_primitive_builder( array: &BitPackedArray, // TODO(ngates): do we want to use fastlanes alignment for this buffer? builder: &mut PrimitiveBuilder, -) { +) -> VortexResult<()> { // If the array is empty, then we don't need to add anything to the builder. if array.is_empty() { - return; + return Ok(()); } let mut uninit_range = builder.uninit_range(array.len()); @@ -87,7 +91,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. @@ -105,6 +109,7 @@ pub(crate) fn unpack_into_primitive_builder( unsafe { uninit_range.finish(); } + Ok(()) } pub fn apply_patches_to_uninit_range(dst: &mut UninitRange, patches: &Patches) { @@ -118,9 +123,17 @@ pub fn apply_patches_to_uninit_range_fn T>( ) { assert_eq!(patches.array_len(), dst.len()); - let indices = patches.indices().to_primitive(); - let values = patches.values().to_primitive(); - let validity = values.validity_mask(); + let indices = patches + .indices() + .to_primitive() + .vortex_expect("indices cannot fail"); + let values = patches + .values() + .to_primitive() + .vortex_expect("values cannot fail"); + let validity = values + .validity_mask() + .vortex_expect("validity_mask cannot fail"); let values = values.as_slice::(); match_each_unsigned_integer_ptype!(indices.ptype(), |P| { @@ -223,7 +236,7 @@ mod tests { fn compression_roundtrip(n: usize) { let values = PrimitiveArray::from_iter((0..n).map(|i| (i % 2047) as u16)); let compressed = BitPackedArray::encode(values.as_ref(), 11).unwrap(); - let decompressed = compressed.to_primitive(); + let decompressed = compressed.to_primitive().unwrap(); assert_arrays_eq!(decompressed, values); values @@ -251,7 +264,7 @@ mod tests { #[test] fn test_all_zeros() { - let zeros = buffer![0u16, 0, 0, 0].into_array().to_primitive(); + let zeros = buffer![0u16, 0, 0, 0].into_array().to_primitive().unwrap(); let bitpacked = bitpack_encode(&zeros, 0, None).unwrap(); let actual = unpack_array(&bitpacked); assert_arrays_eq!(actual, PrimitiveArray::from_iter([0u16, 0, 0, 0])); @@ -259,7 +272,7 @@ mod tests { #[test] fn test_simple_patches() { - let zeros = buffer![0u16, 1, 0, 1].into_array().to_primitive(); + let zeros = buffer![0u16, 1, 0, 1].into_array().to_primitive().unwrap(); let bitpacked = bitpack_encode(&zeros, 0, None).unwrap(); let actual = unpack_array(&bitpacked); assert_arrays_eq!(actual, PrimitiveArray::from_iter([0u16, 1, 0, 1])); @@ -267,7 +280,10 @@ mod tests { #[test] fn test_one_full_chunk() { - let zeros = BufferMut::from_iter(0u16..1024).into_array().to_primitive(); + let zeros = BufferMut::from_iter(0u16..1024) + .into_array() + .to_primitive() + .unwrap(); let bitpacked = bitpack_encode(&zeros, 10, None).unwrap(); let actual = unpack_array(&bitpacked); assert_arrays_eq!(actual, PrimitiveArray::from_iter(0u16..1024)); @@ -277,7 +293,8 @@ mod tests { fn test_three_full_chunks_with_patches() { let zeros = BufferMut::from_iter((5u16..1029).chain(5u16..1029).chain(5u16..1029)) .into_array() - .to_primitive(); + .to_primitive() + .unwrap(); let bitpacked = bitpack_encode(&zeros, 10, None).unwrap(); assert!(bitpacked.patches().is_some()); let actual = unpack_array(&bitpacked); @@ -289,7 +306,10 @@ mod tests { #[test] fn test_one_full_chunk_and_one_short_chunk_no_patch() { - let zeros = BufferMut::from_iter(0u16..1025).into_array().to_primitive(); + let zeros = BufferMut::from_iter(0u16..1025) + .into_array() + .to_primitive() + .unwrap(); let bitpacked = bitpack_encode(&zeros, 11, None).unwrap(); assert!(bitpacked.patches().is_none()); let actual = unpack_array(&bitpacked); @@ -300,7 +320,8 @@ mod tests { fn test_one_full_chunk_and_one_short_chunk_with_patches() { let zeros = BufferMut::from_iter(512u16..1537) .into_array() - .to_primitive(); + .to_primitive() + .unwrap(); let bitpacked = bitpack_encode(&zeros, 10, None).unwrap(); assert_eq!(bitpacked.len(), 1025); assert!(bitpacked.patches().is_some()); @@ -312,7 +333,8 @@ mod tests { fn test_offset_and_short_chunk_and_patches() { let zeros = BufferMut::from_iter(512u16..1537) .into_array() - .to_primitive(); + .to_primitive() + .unwrap(); let bitpacked = bitpack_encode(&zeros, 10, None).unwrap(); assert_eq!(bitpacked.len(), 1025); assert!(bitpacked.patches().is_some()); @@ -325,7 +347,8 @@ mod tests { fn test_offset_and_short_chunk_with_chunks_between_and_patches() { let zeros = BufferMut::from_iter(512u16..2741) .into_array() - .to_primitive(); + .to_primitive() + .unwrap(); let bitpacked = bitpack_encode(&zeros, 10, None).unwrap(); assert_eq!(bitpacked.len(), 2229); assert!(bitpacked.patches().is_some()); @@ -343,7 +366,7 @@ mod tests { let bitpacked = bitpack_encode(&empty, 0, None).unwrap(); let mut builder = PrimitiveBuilder::::new(Nullability::NonNullable); - unpack_into_primitive_builder(&bitpacked, &mut builder); + unpack_into_primitive_builder(&bitpacked, &mut builder).unwrap(); let result = builder.finish_into_primitive(); assert_eq!( @@ -366,7 +389,7 @@ mod tests { // Unpack into a new builder. let mut builder = PrimitiveBuilder::::with_capacity(Nullability::Nullable, 5); - unpack_into_primitive_builder(&bitpacked, &mut builder); + unpack_into_primitive_builder(&bitpacked, &mut builder).unwrap(); let result = builder.finish_into_primitive(); @@ -397,7 +420,7 @@ mod tests { // Unpack into a new builder. let mut builder = PrimitiveBuilder::::with_capacity(Nullability::NonNullable, 100); - unpack_into_primitive_builder(&bitpacked, &mut builder); + unpack_into_primitive_builder(&bitpacked, &mut builder).unwrap(); let result = builder.finish_into_primitive(); @@ -411,7 +434,7 @@ mod tests { // Test with u8 values. let u8_values = PrimitiveArray::from_iter([5u8, 10, 15, 20, 25]); let u8_bitpacked = bitpack_encode(&u8_values, 5, None).unwrap(); - let u8_vector = unpack_to_primitive_vector(&u8_bitpacked); + let u8_vector = unpack_to_primitive_vector(&u8_bitpacked).unwrap(); // Compare with existing unpack method. let expected = unpack_array(&u8_bitpacked); assert_eq!(u8_vector.len(), expected.len()); @@ -422,19 +445,19 @@ mod tests { // Test with u32 values - empty array. let u32_empty: PrimitiveArray = PrimitiveArray::from_iter(Vec::::new()); let u32_empty_bp = bitpack_encode(&u32_empty, 0, None).unwrap(); - let u32_empty_vec = unpack_to_primitive_vector(&u32_empty_bp); + let u32_empty_vec = unpack_to_primitive_vector(&u32_empty_bp).unwrap(); assert_eq!(u32_empty_vec.len(), 0); // Test with u16 values - exactly one chunk (1024 elements). let u16_values = PrimitiveArray::from_iter(0u16..1024); let u16_bitpacked = bitpack_encode(&u16_values, 10, None).unwrap(); - let u16_vector = unpack_to_primitive_vector(&u16_bitpacked); + let u16_vector = unpack_to_primitive_vector(&u16_bitpacked).unwrap(); assert_eq!(u16_vector.len(), 1024); // Test with i32 values - partial chunk (1025 elements). let i32_values = PrimitiveArray::from_iter((0i32..1025).map(|x| x % 512)); let i32_bitpacked = bitpack_encode(&i32_values, 9, None).unwrap(); - let i32_vector = unpack_to_primitive_vector(&i32_bitpacked); + let i32_vector = unpack_to_primitive_vector(&i32_bitpacked).unwrap(); assert_eq!(i32_vector.len(), 1025); // Verify consistency: unpack_to_primitive_vector and unpack_array should produce same values. @@ -458,7 +481,7 @@ mod tests { assert!(bitpacked.patches().is_some(), "Should have patches"); // Unpack to vector. - let vector = unpack_to_primitive_vector(&bitpacked); + let vector = unpack_to_primitive_vector(&bitpacked).unwrap(); // Verify length and that patches were applied. assert_eq!(vector.len(), values.len()); @@ -480,7 +503,7 @@ mod tests { let large_bitpacked = bitpack_encode(&large_array, 8, None).unwrap(); assert!(large_bitpacked.patches().is_some()); - let large_vector = unpack_to_primitive_vector(&large_bitpacked); + let large_vector = unpack_to_primitive_vector(&large_bitpacked).unwrap(); assert_eq!(large_vector.len(), 3072); } @@ -493,7 +516,7 @@ mod tests { let array = PrimitiveArray::new(values, validity); let bitpacked = bitpack_encode(&array, 9, None).unwrap(); - let vector = unpack_to_primitive_vector(&bitpacked); + let vector = unpack_to_primitive_vector(&bitpacked).unwrap(); // Verify length. assert_eq!(vector.len(), 7); @@ -507,7 +530,7 @@ mod tests { let patch_bitpacked = bitpack_encode(&patch_array, 5, None).unwrap(); assert!(patch_bitpacked.patches().is_some()); - let patch_vector = unpack_to_primitive_vector(&patch_bitpacked); + let patch_vector = unpack_to_primitive_vector(&patch_bitpacked).unwrap(); assert_eq!(patch_vector.len(), 7); // Test all nulls edge case. @@ -516,7 +539,7 @@ mod tests { Validity::from_iter([false, false, false, false]), ); let all_nulls_bp = bitpack_encode(&all_nulls, 0, None).unwrap(); - let all_nulls_vec = unpack_to_primitive_vector(&all_nulls_bp); + let all_nulls_vec = unpack_to_primitive_vector(&all_nulls_bp).unwrap(); assert_eq!(all_nulls_vec.len(), 4); } @@ -530,7 +553,7 @@ mod tests { let bitpacked = bitpack_encode(array, bit_width, None).unwrap(); // Method 1: Using the new unpack_to_primitive_vector. - let vector_result = unpack_to_primitive_vector(&bitpacked); + let vector_result = unpack_to_primitive_vector(&bitpacked).unwrap(); // Method 2: Using the old unpack_array. let unpacked_array = unpack_array(&bitpacked); @@ -591,7 +614,7 @@ mod tests { // Test all three methods on the sliced array. let sliced_bp = sliced.as_::(); - let vector_result = unpack_to_primitive_vector(sliced_bp); + let vector_result = unpack_to_primitive_vector(sliced_bp).unwrap(); let unpacked_array = unpack_array(sliced_bp); let executed = sliced.execute(&SESSION).unwrap(); @@ -624,13 +647,13 @@ mod tests { // Empty array. let empty: PrimitiveArray = PrimitiveArray::from_iter(Vec::::new()); let empty_bp = bitpack_encode(&empty, 0, None).unwrap(); - let empty_vec = unpack_to_primitive_vector(&empty_bp); + let empty_vec = unpack_to_primitive_vector(&empty_bp).unwrap(); assert_eq!(empty_vec.len(), 0); // All zeros (bit_width = 0). let zeros = PrimitiveArray::from_iter([0u32; 100]); let zeros_bp = bitpack_encode(&zeros, 0, None).unwrap(); - let zeros_vec = unpack_to_primitive_vector(&zeros_bp); + let zeros_vec = unpack_to_primitive_vector(&zeros_bp).unwrap(); assert_eq!(zeros_vec.len(), 100); // Verify consistency with unpack_array. let zeros_array = unpack_array(&zeros_bp); @@ -639,7 +662,7 @@ mod tests { // Maximum bit width for u16 (15 bits, since bitpacking requires bit_width < type bit width). let max_values = PrimitiveArray::from_iter([32767u16; 50]); // 2^15 - 1 let max_bp = bitpack_encode(&max_values, 15, None).unwrap(); - let max_vec = unpack_to_primitive_vector(&max_bp); + let max_vec = unpack_to_primitive_vector(&max_bp).unwrap(); assert_eq!(max_vec.len(), 50); // Exactly 3072 elements with patches across chunks. @@ -656,7 +679,7 @@ mod tests { let boundary_bp = bitpack_encode(&boundary_array, 7, None).unwrap(); assert!(boundary_bp.patches().is_some()); - let boundary_vec = unpack_to_primitive_vector(&boundary_bp); + let boundary_vec = unpack_to_primitive_vector(&boundary_bp).unwrap(); assert_eq!(boundary_vec.len(), 3072); // Verify consistency. let boundary_unpacked = unpack_array(&boundary_bp); @@ -665,7 +688,7 @@ mod tests { // Single element. let single = PrimitiveArray::from_iter([42u8]); let single_bp = bitpack_encode(&single, 6, None).unwrap(); - let single_vec = unpack_to_primitive_vector(&single_bp); + let single_vec = unpack_to_primitive_vector(&single_bp).unwrap(); assert_eq!(single_vec.len(), 1); } } diff --git a/encodings/fastlanes/src/bitpacking/array/mod.rs b/encodings/fastlanes/src/bitpacking/array/mod.rs index 9d7f2e46ec0..967d541a7a8 100644 --- a/encodings/fastlanes/src/bitpacking/array/mod.rs +++ b/encodings/fastlanes/src/bitpacking/array/mod.rs @@ -306,7 +306,7 @@ mod test { let uncompressed = PrimitiveArray::from_option_iter(values); let packed = BitPackedArray::encode(uncompressed.as_ref(), 1).unwrap(); let expected = &[1, 0, 1, 0, 1, 0, u64::MAX]; - let results = packed.to_primitive().as_slice::().to_vec(); + let results = packed.to_primitive().unwrap().as_slice::().to_vec(); assert_eq!(results, expected); } @@ -328,7 +328,10 @@ mod test { let packed_with_patches = BitPackedArray::encode(&parray, 9).unwrap(); assert!(packed_with_patches.patches().is_some()); assert_eq!( - packed_with_patches.to_primitive().as_slice::(), + packed_with_patches + .to_primitive() + .unwrap() + .as_slice::(), values.as_slice() ); } diff --git a/encodings/fastlanes/src/bitpacking/compute/between.rs b/encodings/fastlanes/src/bitpacking/compute/between.rs index 13f34c0ac21..a231f55d8b6 100644 --- a/encodings/fastlanes/src/bitpacking/compute/between.rs +++ b/encodings/fastlanes/src/bitpacking/compute/between.rs @@ -27,7 +27,7 @@ impl BetweenKernel for BitPackedVTable { }; between( - &array.clone().to_canonical().into_array(), + &array.clone().to_canonical()?.into_array(), lower, upper, options, diff --git a/encodings/fastlanes/src/bitpacking/compute/cast.rs b/encodings/fastlanes/src/bitpacking/compute/cast.rs index 5c240fbc1b6..52abd6e263d 100644 --- a/encodings/fastlanes/src/bitpacking/compute/cast.rs +++ b/encodings/fastlanes/src/bitpacking/compute/cast.rs @@ -86,7 +86,7 @@ mod tests { &DType::Primitive(PType::U32, Nullability::NonNullable) ); - let decoded = casted.to_primitive(); + let decoded = casted.to_primitive().unwrap(); assert_arrays_eq!( decoded.as_ref(), PrimitiveArray::from_iter([10u32, 20, 30, 40, 50, 60]) diff --git a/encodings/fastlanes/src/bitpacking/compute/filter.rs b/encodings/fastlanes/src/bitpacking/compute/filter.rs index a97a9793ab6..1910ee3b181 100644 --- a/encodings/fastlanes/src/bitpacking/compute/filter.rs +++ b/encodings/fastlanes/src/bitpacking/compute/filter.rs @@ -70,8 +70,8 @@ fn filter_primitive( // with a "Cascade Lake" CPU. }; if mask.density() >= full_decompression_threshold { - let decompressed_array = array.to_primitive(); - Ok(filter(decompressed_array.as_ref(), mask)?.to_primitive()) + let decompressed_array = array.to_primitive()?; + Ok(filter(decompressed_array.as_ref(), mask)?.to_primitive()?) } else { filter_primitive_no_decompression::(array, mask) } @@ -186,7 +186,10 @@ mod test { let mask = Mask::from_indices(bitpacked.len(), vec![0, 125, 2047, 2049, 2151, 2790]); - let primitive_result = filter(bitpacked.as_ref(), &mask).unwrap().to_primitive(); + let primitive_result = filter(bitpacked.as_ref(), &mask) + .unwrap() + .to_primitive() + .unwrap(); assert_arrays_eq!( primitive_result, PrimitiveArray::from_iter([0u8, 62, 31, 33, 9, 18]) @@ -202,7 +205,7 @@ mod test { let mask = Mask::from_indices(sliced.len(), vec![1919, 1921]); - let primitive_result = filter(&sliced, &mask).unwrap().to_primitive(); + let primitive_result = filter(&sliced, &mask).unwrap().to_primitive().unwrap(); assert_arrays_eq!(primitive_result, PrimitiveArray::from_iter([31u8, 33])); } @@ -216,7 +219,7 @@ mod test { ) .unwrap(); assert_arrays_eq!( - filtered.to_primitive(), + filtered.to_primitive().unwrap(), PrimitiveArray::from_iter((0..1024).map(|i| (i % 63) as u8)) ); } @@ -231,7 +234,8 @@ mod test { &Mask::from_indices(values.len(), (0..250).collect()), ) .unwrap() - .to_primitive(); + .to_primitive() + .unwrap(); assert_arrays_eq!( filtered, diff --git a/encodings/fastlanes/src/bitpacking/compute/is_constant.rs b/encodings/fastlanes/src/bitpacking/compute/is_constant.rs index 64856c4ad47..4c851713f1c 100644 --- a/encodings/fastlanes/src/bitpacking/compute/is_constant.rs +++ b/encodings/fastlanes/src/bitpacking/compute/is_constant.rs @@ -16,6 +16,7 @@ use vortex_array::register_kernel; use vortex_dtype::IntegerPType; use vortex_dtype::match_each_integer_ptype; use vortex_dtype::match_each_unsigned_integer_ptype; +use vortex_error::VortexError; use vortex_error::VortexResult; use crate::BitPackedArray; @@ -44,12 +45,15 @@ fn bitpacked_is_constant( array: &BitPackedArray, ) -> VortexResult { let mut bit_unpack_iterator = array.unpacked_chunks::(); - let patches = array.patches().map(|p| { - let values = p.values().to_primitive(); - let indices = p.indices().to_primitive(); - let offset = p.offset(); - (indices, values, offset) - }); + let patches = array + .patches() + .map(|p| { + let values = p.values().to_primitive()?; + let indices = p.indices().to_primitive()?; + let offset = p.offset(); + Ok::<_, VortexError>((indices, values, offset)) + }) + .transpose()?; let mut header_constant_value = None; let mut current_idx = 0; diff --git a/encodings/fastlanes/src/bitpacking/compute/take.rs b/encodings/fastlanes/src/bitpacking/compute/take.rs index e4bcf03e15a..b754b2d0bfd 100644 --- a/encodings/fastlanes/src/bitpacking/compute/take.rs +++ b/encodings/fastlanes/src/bitpacking/compute/take.rs @@ -40,7 +40,7 @@ impl TakeKernel for BitPackedVTable { fn take(&self, array: &BitPackedArray, indices: &dyn Array) -> VortexResult { // If the indices are large enough, it's faster to flatten and take the primitive array. if indices.len() * UNPACK_CHUNK_THRESHOLD > array.len() { - return take(array.to_primitive().as_ref(), indices); + return take(array.to_primitive()?.as_ref(), indices); } // NOTE: we use the unsigned PType because all values in the BitPackedArray must @@ -49,7 +49,7 @@ impl TakeKernel for BitPackedVTable { let validity = array.validity(); let taken_validity = validity.take(indices)?; - let indices = indices.to_primitive(); + let indices = indices.to_primitive()?; let taken = match_each_unsigned_integer_ptype!(ptype.to_unsigned(), |T| { match_each_integer_ptype!(indices.ptype(), |I| { take_primitive::(array, &indices, taken_validity)? @@ -170,7 +170,10 @@ mod test { let unpacked = PrimitiveArray::from_iter((0..4096).map(|i| (i % 63) as u8)); let bitpacked = BitPackedArray::encode(unpacked.as_ref(), 6).unwrap(); - let primitive_result = take(bitpacked.as_ref(), &indices).unwrap().to_primitive(); + let primitive_result = take(bitpacked.as_ref(), &indices) + .unwrap() + .to_primitive() + .unwrap(); assert_arrays_eq!( primitive_result, PrimitiveArray::from_iter([0u8, 62, 31, 33, 9, 18]) @@ -186,7 +189,8 @@ mod test { let primitive_result = take(bitpacked.as_ref(), indices.as_ref()) .unwrap() - .to_primitive(); + .to_primitive() + .unwrap(); assert_arrays_eq!(primitive_result, PrimitiveArray::from_iter([0u32, 2, 4, 6])); } @@ -199,7 +203,7 @@ mod test { let bitpacked = BitPackedArray::encode(unpacked.as_ref(), 6).unwrap(); let sliced = bitpacked.slice(128..2050); - let primitive_result = take(&sliced, &indices).unwrap().to_primitive(); + let primitive_result = take(&sliced, &indices).unwrap().to_primitive().unwrap(); assert_arrays_eq!(primitive_result, PrimitiveArray::from_iter([31u8, 33])); } @@ -260,12 +264,13 @@ mod test { PrimitiveArray::from_option_iter([Some(0u64), Some(1), None, Some(3)]).as_ref(), ) .unwrap() - .to_primitive(); + .to_primitive() + .unwrap(); assert_arrays_eq!( taken_primitive, PrimitiveArray::from_option_iter([Some(1i32), Some(2), None, Some(4)]) ); - assert_eq!(taken_primitive.invalid_count(), 1); + assert_eq!(taken_primitive.invalid_count().unwrap(), 1); } #[rstest] diff --git a/encodings/fastlanes/src/bitpacking/vtable/canonical.rs b/encodings/fastlanes/src/bitpacking/vtable/canonical.rs index bf457e8e107..da159a48126 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/canonical.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/canonical.rs @@ -6,6 +6,7 @@ use vortex_array::builders::ArrayBuilder; use vortex_array::vtable::CanonicalVTable; use vortex_dtype::match_each_integer_ptype; use vortex_error::VortexExpect; +use vortex_error::VortexResult; use crate::BitPackedArray; use crate::BitPackedVTable; @@ -13,11 +14,14 @@ use crate::bitpack_decompress::unpack_array; use crate::bitpack_decompress::unpack_into_primitive_builder; impl CanonicalVTable for BitPackedVTable { - fn canonicalize(array: &BitPackedArray) -> Canonical { - Canonical::Primitive(unpack_array(array)) + fn canonicalize(array: &BitPackedArray) -> VortexResult { + Ok(Canonical::Primitive(unpack_array(array))) } - fn append_to_builder(array: &BitPackedArray, builder: &mut dyn ArrayBuilder) { + fn append_to_builder( + array: &BitPackedArray, + builder: &mut dyn ArrayBuilder, + ) -> VortexResult<()> { match_each_integer_ptype!(array.ptype(), |T| { unpack_into_primitive_builder::( array, @@ -25,7 +29,8 @@ impl CanonicalVTable for BitPackedVTable { .as_any_mut() .downcast_mut() .vortex_expect("bit packed array must canonicalize into a primitive array"), - ) - }) + )? + }); + Ok(()) } } diff --git a/encodings/fastlanes/src/bitpacking/vtable/mod.rs b/encodings/fastlanes/src/bitpacking/vtable/mod.rs index e47ed45f044..cafba1f021d 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/mod.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/mod.rs @@ -173,7 +173,7 @@ impl VTable for BitPackedVTable { } fn batch_execute(array: &BitPackedArray, _ctx: &mut ExecutionCtx) -> VortexResult { - Ok(unpack_to_primitive_vector(array).freeze().into()) + Ok(unpack_to_primitive_vector(array)?.freeze().into()) } } diff --git a/encodings/fastlanes/src/delta/array/delta_compress.rs b/encodings/fastlanes/src/delta/array/delta_compress.rs index d8e7955f8f3..55308d60e74 100644 --- a/encodings/fastlanes/src/delta/array/delta_compress.rs +++ b/encodings/fastlanes/src/delta/array/delta_compress.rs @@ -119,7 +119,7 @@ mod tests { fn do_roundtrip_test(input: PrimitiveArray) { let delta = DeltaArray::try_from_primitive_array(&input).unwrap(); assert_eq!(delta.len(), input.len()); - let decompressed = delta_decompress(&delta); + let decompressed = delta_decompress(&delta).unwrap(); let decompressed_slice = decompressed.as_slice::(); assert_eq!(decompressed_slice.len(), input.len()); for (actual, expected) in decompressed_slice.iter().zip(input.as_slice()) { diff --git a/encodings/fastlanes/src/delta/array/delta_decompress.rs b/encodings/fastlanes/src/delta/array/delta_decompress.rs index c68df2e5a0b..9d8f6710987 100644 --- a/encodings/fastlanes/src/delta/array/delta_decompress.rs +++ b/encodings/fastlanes/src/delta/array/delta_decompress.rs @@ -15,18 +15,19 @@ use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_dtype::NativePType; use vortex_dtype::match_each_unsigned_integer_ptype; +use vortex_error::VortexResult; use crate::DeltaArray; -pub fn delta_decompress(array: &DeltaArray) -> PrimitiveArray { - let bases = array.bases().to_primitive(); - let deltas = array.deltas().to_primitive(); +pub fn delta_decompress(array: &DeltaArray) -> VortexResult { + let bases = array.bases().to_primitive()?; + let deltas = array.deltas().to_primitive()?; let decoded = match_each_unsigned_integer_ptype!(deltas.ptype(), |T| { const LANES: usize = T::LANES; PrimitiveArray::new( decompress_primitive::(bases.as_slice(), deltas.as_slice()), - Validity::from_mask(array.deltas().validity_mask(), array.dtype().nullability()), + Validity::from_mask(array.deltas().validity_mask()?, array.dtype().nullability()), ) }); diff --git a/encodings/fastlanes/src/delta/compute/cast.rs b/encodings/fastlanes/src/delta/compute/cast.rs index 443e5be8c06..70b7a7d8f4a 100644 --- a/encodings/fastlanes/src/delta/compute/cast.rs +++ b/encodings/fastlanes/src/delta/compute/cast.rs @@ -81,7 +81,7 @@ mod tests { ); // Verify by decoding - let decoded = casted.to_primitive(); + let decoded = casted.to_primitive().unwrap(); assert_arrays_eq!(decoded, PrimitiveArray::from_iter([10u32, 20, 30, 40, 50])); } diff --git a/encodings/fastlanes/src/delta/vtable/canonical.rs b/encodings/fastlanes/src/delta/vtable/canonical.rs index 8707f106c81..dcc3136a030 100644 --- a/encodings/fastlanes/src/delta/vtable/canonical.rs +++ b/encodings/fastlanes/src/delta/vtable/canonical.rs @@ -3,13 +3,14 @@ use vortex_array::Canonical; use vortex_array::vtable::CanonicalVTable; +use vortex_error::VortexResult; use super::DeltaVTable; use crate::DeltaArray; use crate::delta::array::delta_decompress::delta_decompress; impl CanonicalVTable for DeltaVTable { - fn canonicalize(array: &DeltaArray) -> Canonical { - Canonical::Primitive(delta_decompress(array)) + fn canonicalize(array: &DeltaArray) -> VortexResult { + Ok(Canonical::Primitive(delta_decompress(array)?)) } } diff --git a/encodings/fastlanes/src/delta/vtable/operations.rs b/encodings/fastlanes/src/delta/vtable/operations.rs index d8a8014fff4..6a6ef15e3ce 100644 --- a/encodings/fastlanes/src/delta/vtable/operations.rs +++ b/encodings/fastlanes/src/delta/vtable/operations.rs @@ -7,12 +7,13 @@ use std::ops::Range; use vortex_array::Array; use vortex_array::ArrayRef; use vortex_array::IntoArray; -use vortex_array::ToCanonical; use vortex_array::vtable::OperationsVTable; +use vortex_error::VortexExpect; use vortex_scalar::Scalar; use super::DeltaVTable; use crate::DeltaArray; +use crate::delta::array::delta_decompress::delta_decompress; impl OperationsVTable for DeltaVTable { fn slice(array: &DeltaArray, range: Range) -> ArrayRef { @@ -42,7 +43,11 @@ impl OperationsVTable for DeltaVTable { } fn scalar_at(array: &DeltaArray, index: usize) -> Scalar { - let decompressed = array.slice(index..index + 1).to_primitive(); + let sliced = Self::slice(array, index..index + 1); + let delta_array = sliced + .as_opt::() + .vortex_expect("sliced must be DeltaArray"); + let decompressed = delta_decompress(delta_array).vortex_expect("delta_decompress failed"); decompressed.scalar_at(0) } } diff --git a/encodings/fastlanes/src/for/array/for_compress.rs b/encodings/fastlanes/src/for/array/for_compress.rs index 9af7b7941e6..47f880a73af 100644 --- a/encodings/fastlanes/src/for/array/for_compress.rs +++ b/encodings/fastlanes/src/for/array/for_compress.rs @@ -67,7 +67,7 @@ mod test { let compressed = FoRArray::encode(array.clone()).unwrap(); assert_eq!(i32::try_from(compressed.reference_scalar()).unwrap(), 1); - let decompressed = compressed.to_primitive(); + let decompressed = compressed.to_primitive().unwrap(); assert_arrays_eq!(decompressed, array); } @@ -107,7 +107,7 @@ mod test { // Create a range offset by a million. let array = PrimitiveArray::from_iter((0u32..100_000).step_by(1024).map(|v| v + 1_000_000)); let compressed = FoRArray::encode(array.clone()).unwrap(); - let decompressed = compressed.to_primitive(); + let decompressed = compressed.to_primitive().unwrap(); assert_arrays_eq!(decompressed, array); } @@ -118,7 +118,7 @@ mod test { let array = PrimitiveArray::from_iter((0u32..1024).map(|x| x % 7)); let bp = BitPackedArray::encode(array.as_ref(), 3).unwrap(); let compressed = FoRArray::try_new(bp.into_array(), 10u32.into()).unwrap(); - let decompressed = compressed.to_primitive(); + let decompressed = compressed.to_primitive().unwrap(); assert_arrays_eq!(decompressed, expect); } @@ -129,7 +129,7 @@ mod test { let array = PrimitiveArray::from_iter((0u32..1024).map(|x| x % 7)); let bp = BitPackedArray::encode(array.as_ref(), 2).unwrap(); let compressed = FoRArray::try_new(bp.clone().into_array(), 10u32.into()).unwrap(); - let decompressed = fused_decompress::(&compressed, &bp); + let decompressed = fused_decompress::(&compressed, &bp).unwrap(); assert_arrays_eq!(decompressed, expect); } @@ -149,12 +149,13 @@ mod test { let encoded = compressed .encoded() .to_primitive() + .unwrap() .reinterpret_cast(PType::U8); let unsigned: Vec = (0..=u8::MAX).collect_vec(); let expected_unsigned = PrimitiveArray::from_iter(unsigned); assert_arrays_eq!(encoded, expected_unsigned); - let decompressed = decompress(&compressed); + let decompressed = decompress(&compressed).unwrap(); array .as_slice::() .iter() diff --git a/encodings/fastlanes/src/for/array/for_decompress.rs b/encodings/fastlanes/src/for/array/for_decompress.rs index c1b32d59668..eb4bdc2fa84 100644 --- a/encodings/fastlanes/src/for/array/for_decompress.rs +++ b/encodings/fastlanes/src/for/array/for_decompress.rs @@ -16,6 +16,7 @@ use vortex_dtype::UnsignedPType; use vortex_dtype::match_each_integer_ptype; use vortex_dtype::match_each_unsigned_integer_ptype; use vortex_error::VortexExpect; +use vortex_error::VortexResult; use crate::BitPackedArray; use crate::BitPackedVTable; @@ -44,7 +45,7 @@ impl + FoR> UnpackStrategy for FoRStrategy } } -pub fn decompress(array: &FoRArray) -> PrimitiveArray { +pub fn decompress(array: &FoRArray) -> VortexResult { let ptype = array.ptype(); // Try to do fused unpack. @@ -57,10 +58,10 @@ pub fn decompress(array: &FoRArray) -> PrimitiveArray { } // TODO(ngates): Do we need this to be into_encoded() somehow? - let encoded = array.encoded().to_primitive(); + let encoded = array.encoded().to_primitive()?; let validity = encoded.validity().clone(); - match_each_integer_ptype!(ptype, |T| { + Ok(match_each_integer_ptype!(ptype, |T| { let min = array .reference_scalar() .as_primitive() @@ -74,7 +75,7 @@ pub fn decompress(array: &FoRArray) -> PrimitiveArray { validity, ) } - }) + })) } pub(crate) fn fused_decompress< @@ -82,7 +83,7 @@ pub(crate) fn fused_decompress< >( for_: &FoRArray, bp: &BitPackedArray, -) -> PrimitiveArray { +) -> VortexResult { let ref_ = for_ .reference_scalar() .as_primitive() @@ -107,7 +108,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. @@ -128,7 +129,7 @@ pub(crate) fn fused_decompress< uninit_range.finish(); } - builder.finish_into_primitive() + Ok(builder.finish_into_primitive()) } fn decompress_primitive( diff --git a/encodings/fastlanes/src/for/compute/cast.rs b/encodings/fastlanes/src/for/compute/cast.rs index 39103b5776b..81b7e88bfef 100644 --- a/encodings/fastlanes/src/for/compute/cast.rs +++ b/encodings/fastlanes/src/for/compute/cast.rs @@ -68,7 +68,7 @@ mod tests { ); // Verify the values after decoding - let decoded = casted.to_primitive(); + let decoded = casted.to_primitive().unwrap(); assert_arrays_eq!( decoded, PrimitiveArray::from_iter([100i64, 110, 120, 130, 140]) diff --git a/encodings/fastlanes/src/for/compute/compare.rs b/encodings/fastlanes/src/for/compute/compare.rs index 3363ac3b65e..c14c1cc87d2 100644 --- a/encodings/fastlanes/src/for/compute/compare.rs +++ b/encodings/fastlanes/src/for/compute/compare.rs @@ -213,7 +213,7 @@ mod tests { result: VortexResult>, expected: T, ) { - let result = result.unwrap().unwrap().to_bool(); + let result = result.unwrap().unwrap().to_bool().unwrap(); assert_eq!(result.bit_buffer(), &BitBuffer::from_iter(expected)); } } diff --git a/encodings/fastlanes/src/for/compute/is_sorted.rs b/encodings/fastlanes/src/for/compute/is_sorted.rs index ed5c476c561..d65eff07f63 100644 --- a/encodings/fastlanes/src/for/compute/is_sorted.rs +++ b/encodings/fastlanes/src/for/compute/is_sorted.rs @@ -72,7 +72,7 @@ use crate::FoRVTable; /// unwrapped values. impl IsSortedKernel for FoRVTable { fn is_sorted(&self, array: &FoRArray) -> VortexResult> { - let encoded = array.encoded().to_primitive(); + let encoded = array.encoded().to_primitive()?; is_sorted( &encoded .reinterpret_cast(encoded.ptype().to_unsigned()) @@ -81,7 +81,7 @@ impl IsSortedKernel for FoRVTable { } fn is_strict_sorted(&self, array: &FoRArray) -> VortexResult> { - let encoded = array.encoded().to_primitive(); + let encoded = array.encoded().to_primitive()?; is_strict_sorted( &encoded .reinterpret_cast(encoded.ptype().to_unsigned()) diff --git a/encodings/fastlanes/src/for/vtable/canonical.rs b/encodings/fastlanes/src/for/vtable/canonical.rs index 7ab5e372a4b..fdb401d6550 100644 --- a/encodings/fastlanes/src/for/vtable/canonical.rs +++ b/encodings/fastlanes/src/for/vtable/canonical.rs @@ -3,13 +3,14 @@ use vortex_array::Canonical; use vortex_array::vtable::CanonicalVTable; +use vortex_error::VortexResult; use super::FoRVTable; use crate::FoRArray; use crate::r#for::array::for_decompress::decompress; impl CanonicalVTable for FoRVTable { - fn canonicalize(array: &FoRArray) -> Canonical { - Canonical::Primitive(decompress(array)) + fn canonicalize(array: &FoRArray) -> VortexResult { + Ok(Canonical::Primitive(decompress(array)?)) } } diff --git a/encodings/fastlanes/src/rle/array/mod.rs b/encodings/fastlanes/src/rle/array/mod.rs index 9e6bf378477..a964bcc3338 100644 --- a/encodings/fastlanes/src/rle/array/mod.rs +++ b/encodings/fastlanes/src/rle/array/mod.rs @@ -278,9 +278,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] @@ -313,10 +313,10 @@ mod tests { .unwrap(); let valid_slice = rle_array.slice(0..3); - 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] @@ -349,10 +349,10 @@ mod tests { .unwrap(); let invalid_slice = rle_array.slice(2..5); - 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] @@ -385,9 +385,11 @@ 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); + let expected_mask = Validity::from_iter([false, true, false]) + .to_mask(3) + .unwrap(); assert_eq!(validity_mask.len(), expected_mask.len()); assert_eq!(validity_mask, expected_mask); } @@ -430,7 +432,7 @@ mod tests { let rle_array = RLEArray::encode(&primitive).unwrap(); assert_eq!(rle_array.len(), 2048); - let original_data = rle_array.to_primitive(); + let original_data = rle_array.to_primitive().unwrap(); let original_values = original_data.as_slice::(); let ctx = ArrayContext::empty().with(RLEVTable.as_vtable()); @@ -454,7 +456,7 @@ mod tests { ) .unwrap(); - let decoded_data = decoded.to_primitive(); + let decoded_data = decoded.to_primitive().unwrap(); let decoded_values = decoded_data.as_slice::(); assert_eq!(original_values, decoded_values); @@ -481,8 +483,8 @@ mod tests { let parts = ArrayParts::try_from(concat).unwrap(); let decoded = parts.decode(&ctx, sliced.dtype(), sliced.len()).unwrap(); - let original_data = sliced.to_primitive(); - let decoded_data = decoded.to_primitive(); + let original_data = sliced.to_primitive().unwrap(); + let decoded_data = decoded.to_primitive().unwrap(); let original_values = original_data.as_slice::(); let decoded_values = decoded_data.as_slice::(); diff --git a/encodings/fastlanes/src/rle/array/rle_compress.rs b/encodings/fastlanes/src/rle/array/rle_compress.rs index 3f3cc6f0871..83dc3f33d18 100644 --- a/encodings/fastlanes/src/rle/array/rle_compress.rs +++ b/encodings/fastlanes/src/rle/array/rle_compress.rs @@ -94,7 +94,7 @@ where RLEArray::try_new( values_buf.into_array(), - PrimitiveArray::new(indices_buf.freeze(), padded_validity(array)).into_array(), + PrimitiveArray::new(indices_buf.freeze(), padded_validity(array)?).into_array(), values_idx_offsets.into_array(), 0, array.len(), @@ -102,8 +102,8 @@ where } /// Returns validity padded to the next 1024 chunk for a given array. -fn padded_validity(array: &PrimitiveArray) -> Validity { - match array.validity() { +fn padded_validity(array: &PrimitiveArray) -> VortexResult { + Ok(match array.validity() { Validity::NonNullable => Validity::NonNullable, Validity::AllValid => Validity::AllValid, Validity::AllInvalid => Validity::AllInvalid, @@ -112,18 +112,18 @@ fn padded_validity(array: &PrimitiveArray) -> Validity { let padded_len = len.next_multiple_of(FL_CHUNK_SIZE); if len == padded_len { - return Validity::Array(validity_array.clone()); + return Ok(Validity::Array(validity_array.clone())); } let mut builder = BitBufferMut::with_capacity(padded_len); - let bool_array = validity_array.to_bool(); + let bool_array = validity_array.to_bool()?; builder.append_buffer(bool_array.bit_buffer()); builder.append_n(false, padded_len - len); Validity::from(builder.freeze()) } - } + }) } #[cfg(test)] @@ -142,24 +142,24 @@ mod tests { // u8 let values_u8: Buffer = [1, 1, 2, 2, 3, 3].iter().copied().collect(); let array_u8 = values_u8.into_array(); - let encoded_u8 = RLEArray::encode(&array_u8.to_primitive()).unwrap(); - let decoded_u8 = encoded_u8.to_primitive(); + let encoded_u8 = RLEArray::encode(&array_u8.to_primitive().unwrap()).unwrap(); + let decoded_u8 = encoded_u8.to_primitive().unwrap(); let expected_u8 = PrimitiveArray::from_iter(vec![1u8, 1, 2, 2, 3, 3]); assert_arrays_eq!(decoded_u8, expected_u8); // u16 let values_u16: Buffer = [100, 100, 200, 200].iter().copied().collect(); let array_u16 = values_u16.into_array(); - let encoded_u16 = RLEArray::encode(&array_u16.to_primitive()).unwrap(); - let decoded_u16 = encoded_u16.to_primitive(); + let encoded_u16 = RLEArray::encode(&array_u16.to_primitive().unwrap()).unwrap(); + let decoded_u16 = encoded_u16.to_primitive().unwrap(); let expected_u16 = PrimitiveArray::from_iter(vec![100u16, 100, 200, 200]); assert_arrays_eq!(decoded_u16, expected_u16); // u64 let values_u64: Buffer = [1000, 1000, 2000].iter().copied().collect(); let array_u64 = values_u64.into_array(); - let encoded_u64 = RLEArray::encode(&array_u64.to_primitive()).unwrap(); - let decoded_u64 = encoded_u64.to_primitive(); + let encoded_u64 = RLEArray::encode(&array_u64.to_primitive().unwrap()).unwrap(); + let decoded_u64 = encoded_u64.to_primitive().unwrap(); let expected_u64 = PrimitiveArray::from_iter(vec![1000u64, 1000, 2000]); assert_arrays_eq!(decoded_u64, expected_u64); } @@ -168,7 +168,7 @@ mod tests { fn test_length() { let values: Buffer = [1, 1, 2, 2, 2, 3].iter().copied().collect(); let array = values.into_array(); - let encoded = RLEArray::encode(&array.to_primitive()).unwrap(); + let encoded = RLEArray::encode(&array.to_primitive().unwrap()).unwrap(); assert_eq!(encoded.len(), 6); } @@ -176,7 +176,7 @@ mod tests { fn test_empty_length() { let values: Buffer = Buffer::empty(); let array = values.into_array(); - let encoded = RLEArray::encode(&array.to_primitive()).unwrap(); + let encoded = RLEArray::encode(&array.to_primitive().unwrap()).unwrap(); assert_eq!(encoded.len(), 0); assert_eq!(encoded.values().len(), 0); @@ -187,10 +187,10 @@ mod tests { let values: Buffer = vec![42; 2000].into_iter().collect(); let array = values.into_array(); - let encoded = RLEArray::encode(&array.to_primitive()).unwrap(); + let encoded = RLEArray::encode(&array.to_primitive().unwrap()).unwrap(); assert_eq!(encoded.values().len(), 2); // 2 chunks, each storing value 42 - let decoded = encoded.to_primitive(); // Verify round-trip + let decoded = encoded.to_primitive().unwrap(); // Verify round-trip let expected = PrimitiveArray::from_iter(vec![42u16; 2000]); assert_arrays_eq!(decoded, expected); } @@ -200,10 +200,10 @@ mod tests { let values: Buffer = (0u8..=255).collect(); let array = values.into_array(); - let encoded = RLEArray::encode(&array.to_primitive()).unwrap(); + let encoded = RLEArray::encode(&array.to_primitive().unwrap()).unwrap(); assert_eq!(encoded.values().len(), 256); - let decoded = encoded.to_primitive(); // Verify round-trip + let decoded = encoded.to_primitive().unwrap(); // Verify round-trip let expected = PrimitiveArray::from_iter((0u8..=255).collect::>()); assert_arrays_eq!(decoded, expected); } @@ -214,8 +214,8 @@ mod tests { let values: Buffer = (0..1500).map(|i| (i / 100) as u32).collect(); let array = values.into_array(); - let encoded = RLEArray::encode(&array.to_primitive()).unwrap(); - let decoded = encoded.to_primitive(); + let encoded = RLEArray::encode(&array.to_primitive().unwrap()).unwrap(); + let decoded = encoded.to_primitive().unwrap(); assert_eq!(encoded.len(), 1500); assert_arrays_eq!(decoded, array); @@ -229,8 +229,8 @@ mod tests { let values: Buffer = (0..2048).map(|i| (i / 100) as u32).collect(); let array = values.into_array(); - let encoded = RLEArray::encode(&array.to_primitive()).unwrap(); - let decoded = encoded.to_primitive(); + let encoded = RLEArray::encode(&array.to_primitive().unwrap()).unwrap(); + let decoded = encoded.to_primitive().unwrap(); assert_eq!(encoded.len(), 2048); assert_arrays_eq!(decoded, array); @@ -250,9 +250,9 @@ mod tests { #[case::f32((-2000..2000).map(|i| i as f32).collect::>())] #[case::f64((-2000..2000).map(|i| i as f64).collect::>())] fn test_roundtrip_primitive_types(#[case] values: Buffer) { - let primitive = values.clone().into_array().to_primitive(); + let primitive = values.clone().into_array().to_primitive().unwrap(); let result = RLEArray::encode(&primitive).unwrap(); - let decoded = result.to_primitive(); + let decoded = result.to_primitive().unwrap(); let expected = PrimitiveArray::new(values, primitive.validity().clone()); assert_arrays_eq!(decoded, expected); } diff --git a/encodings/fastlanes/src/rle/array/rle_decompress.rs b/encodings/fastlanes/src/rle/array/rle_decompress.rs index d421a61ca11..a0096e12bc7 100644 --- a/encodings/fastlanes/src/rle/array/rle_decompress.rs +++ b/encodings/fastlanes/src/rle/array/rle_decompress.rs @@ -13,6 +13,7 @@ use vortex_buffer::BufferMut; use vortex_dtype::NativePType; use vortex_dtype::match_each_native_ptype; use vortex_dtype::match_each_unsigned_integer_ptype; +use vortex_error::VortexExpect; use vortex_error::vortex_panic; use crate::FL_CHUNK_SIZE; @@ -26,6 +27,7 @@ use crate::RLEArray; pub fn rle_decompress(array: &RLEArray) -> PrimitiveArray { match_each_native_ptype!(array.values().dtype().as_ptype(), |V| { match_each_unsigned_integer_ptype!(array.values_idx_offsets().dtype().as_ptype(), |O| { + use vortex_dtype::PType; // RLE indices are always u16 (or u8 if downcasted). match array.indices().dtype().as_ptype() { PType::U8 => rle_decode_typed::(array), @@ -46,10 +48,16 @@ where I: NativePType + Into, O: NativePType + AsPrimitive, { - let values = array.values().to_primitive(); + let values = array + .values() + .to_primitive() + .vortex_expect("values cannot fail"); let values = values.as_slice::(); - let indices = array.indices().to_primitive(); + let indices = array + .indices() + .to_primitive() + .vortex_expect("indices cannot fail"); let indices = indices.as_slice::(); assert_eq!(indices.len() % FL_CHUNK_SIZE, 0); @@ -60,7 +68,10 @@ where let mut buffer = BufferMut::::with_capacity(num_chunks * FL_CHUNK_SIZE); let buffer_uninit = buffer.spare_capacity_mut(); - let values_idx_offsets = array.values_idx_offsets().to_primitive(); + let values_idx_offsets = array + .values_idx_offsets() + .to_primitive() + .vortex_expect("values_idx_offsets cannot fail"); let values_idx_offsets = values_idx_offsets.as_slice::(); for chunk_idx in 0..num_chunks { diff --git a/encodings/fastlanes/src/rle/vtable/canonical.rs b/encodings/fastlanes/src/rle/vtable/canonical.rs index 1ec4520d840..2fd62ed0f42 100644 --- a/encodings/fastlanes/src/rle/vtable/canonical.rs +++ b/encodings/fastlanes/src/rle/vtable/canonical.rs @@ -3,13 +3,14 @@ use vortex_array::Canonical; use vortex_array::vtable::CanonicalVTable; +use vortex_error::VortexResult; use super::RLEVTable; use crate::RLEArray; use crate::rle::array::rle_decompress::rle_decompress; impl CanonicalVTable for RLEVTable { - fn canonicalize(array: &RLEArray) -> Canonical { - Canonical::Primitive(rle_decompress(array)) + fn canonicalize(array: &RLEArray) -> VortexResult { + Ok(Canonical::Primitive(rle_decompress(array))) } } diff --git a/encodings/fastlanes/src/rle/vtable/operations.rs b/encodings/fastlanes/src/rle/vtable/operations.rs index 1ea116887d7..ca4a884bffb 100644 --- a/encodings/fastlanes/src/rle/vtable/operations.rs +++ b/encodings/fastlanes/src/rle/vtable/operations.rs @@ -202,7 +202,7 @@ mod tests { let expected: Vec = (0..3000).map(|i| (i / 50) as u16).collect(); let array = values.into_array(); - let encoded = RLEArray::encode(&array.to_primitive()).unwrap(); + let encoded = RLEArray::encode(&array.to_primitive().unwrap()).unwrap(); // Access scalars from multiple chunks. for &idx in &[1023, 1024, 1025, 2047, 2048, 2049] { @@ -283,7 +283,7 @@ mod tests { #[test] fn test_slice_decode_with_nulls() { let array = fixture::rle_array_with_nulls(); - let sliced = array.slice(1..4).to_array().to_primitive(); // [null, 20, 20] + let sliced = array.slice(1..4).to_array().to_primitive().unwrap(); // [null, 20, 20] let expected = PrimitiveArray::from_option_iter([Option::::None, Some(20), Some(20)]); assert_arrays_eq!(sliced.to_array(), expected.to_array()); @@ -303,7 +303,7 @@ mod tests { let expected: Vec = (0..2100).map(|i| (i / 100) as u32).collect(); let array = values.into_array(); - let encoded = RLEArray::encode(&array.to_primitive()).unwrap(); + let encoded = RLEArray::encode(&array.to_primitive().unwrap()).unwrap(); // Slice across first and second chunk. let slice = encoded.slice(500..1500); diff --git a/encodings/fsst/benches/chunked_dict_fsst_builder.rs b/encodings/fsst/benches/chunked_dict_fsst_builder.rs index 509932da372..697d908689c 100644 --- a/encodings/fsst/benches/chunked_dict_fsst_builder.rs +++ b/encodings/fsst/benches/chunked_dict_fsst_builder.rs @@ -9,6 +9,7 @@ use vortex_array::arrays::ChunkedArray; use vortex_array::builders::builder_with_capacity; use vortex_array::compute::warm_up_vtables; use vortex_dtype::NativePType; +use vortex_error::VortexExpect; use vortex_fsst::test_utils::gen_dict_fsst_test_data; fn main() { @@ -45,7 +46,9 @@ fn chunked_dict_fsst_canonical_into( bencher.with_inputs(|| &chunk).bench_refs(|chunk| { let mut builder = builder_with_capacity(chunk.dtype(), len * chunk_count); - chunk.append_to_builder(builder.as_mut()); + chunk + .append_to_builder(builder.as_mut()) + .vortex_expect("append_to_builder"); builder.finish() }) } diff --git a/encodings/fsst/benches/fsst_compress.rs b/encodings/fsst/benches/fsst_compress.rs index 1de3dae93f9..f5e19c8fcf4 100644 --- a/encodings/fsst/benches/fsst_compress.rs +++ b/encodings/fsst/benches/fsst_compress.rs @@ -1,8 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -#![allow(clippy::unwrap_used)] - use divan::Bencher; use rand::Rng; use rand::SeedableRng; @@ -18,6 +16,7 @@ use vortex_array::compute::compare; use vortex_array::compute::warm_up_vtables; use vortex_dtype::DType; use vortex_dtype::Nullability; +use vortex_error::VortexExpect; use vortex_fsst::fsst_compress; use vortex_fsst::fsst_train_compressor; use vortex_scalar::Scalar; @@ -81,7 +80,7 @@ fn pushdown_compare(bencher: Bencher, (string_count, avg_len, unique_chars): (us bencher .with_inputs(|| (&fsst_array, &constant)) .bench_refs(|(fsst_array, constant)| { - compare(fsst_array.as_ref(), constant.as_ref(), Operator::Eq).unwrap(); + compare(fsst_array.as_ref(), constant.as_ref(), Operator::Eq).vortex_expect("compare"); }) } @@ -99,11 +98,14 @@ fn canonicalize_compare( .with_inputs(|| (&fsst_array, &constant)) .bench_refs(|(fsst_array, constant)| { compare( - fsst_array.to_canonical().as_ref(), + fsst_array + .to_canonical() + .vortex_expect("to_canonical") + .as_ref(), constant.as_ref(), Operator::Eq, ) - .unwrap() + .vortex_expect("compare") }); } @@ -130,7 +132,9 @@ fn chunked_canonicalize_into( bencher.with_inputs(|| &array).bench_refs(|array| { let mut builder = VarBinViewBuilder::with_capacity(DType::Binary(Nullability::NonNullable), array.len()); - array.append_to_builder(&mut builder); + array + .append_to_builder(&mut builder) + .vortex_expect("append_to_builder"); builder.finish() }); } diff --git a/encodings/fsst/src/canonical.rs b/encodings/fsst/src/canonical.rs index 492ba3e6e05..6d6d9040501 100644 --- a/encodings/fsst/src/canonical.rs +++ b/encodings/fsst/src/canonical.rs @@ -16,37 +16,40 @@ use vortex_buffer::BufferMut; use vortex_buffer::ByteBuffer; use vortex_buffer::ByteBufferMut; use vortex_dtype::match_each_integer_ptype; +use vortex_error::VortexResult; use vortex_vector::binaryview::BinaryView; use crate::FSSTArray; use crate::FSSTVTable; impl CanonicalVTable for FSSTVTable { - fn canonicalize(array: &FSSTArray) -> Canonical { - let (buffer, views) = fsst_decode_views(array, 0); + fn canonicalize(array: &FSSTArray) -> VortexResult { + let (buffer, views) = fsst_decode_views(array, 0)?; // SAFETY: FSST already validates the bytes for binary/UTF-8. We build views directly on // top of them, so the view pointers will all be valid. - unsafe { + Ok(unsafe { Canonical::VarBinView(VarBinViewArray::new_unchecked( views, Arc::new([buffer]), array.dtype().clone(), array.codes().validity().clone(), )) - } + }) } - fn append_to_builder(array: &FSSTArray, builder: &mut dyn ArrayBuilder) { + fn append_to_builder(array: &FSSTArray, builder: &mut dyn ArrayBuilder) -> VortexResult<()> { let Some(builder) = builder.as_any_mut().downcast_mut::() else { - return builder.extend_from_array(&array.to_canonical().into_array()); + builder.extend_from_array(&array.to_canonical()?.into_array()); + return Ok(()); }; // Decompress the whole block of data into a new buffer, and create some views // from it instead. - let (buffer, views) = fsst_decode_views(array, builder.completed_block_count()); + let (buffer, views) = fsst_decode_views(array, builder.completed_block_count())?; - builder.push_buffer_and_adjusted_views(&[buffer], &views, array.validity_mask()); + builder.push_buffer_and_adjusted_views(&[buffer], &views, array.validity_mask()?); + Ok(()) } } @@ -54,7 +57,10 @@ impl CanonicalVTable for FSSTVTable { clippy::cast_possible_truncation, reason = "truncation is intentional for buffer index" )] -fn fsst_decode_views(fsst_array: &FSSTArray, buf_index: u32) -> (ByteBuffer, Buffer) { +fn fsst_decode_views( + fsst_array: &FSSTArray, + buf_index: u32, +) -> VortexResult<(ByteBuffer, Buffer)> { // FSSTArray has two child arrays: // 1. A VarBinArray, which holds the string heap of the compressed codes. // 2. An uncompressed_lengths primitive array, storing the length of each original @@ -64,7 +70,7 @@ fn fsst_decode_views(fsst_array: &FSSTArray, buf_index: u32) -> (ByteBuffer, Buf // necessary for a VarBinViewArray and construct the canonical array. let bytes = fsst_array.codes().sliced_bytes(); - let uncompressed_lens_array = fsst_array.uncompressed_lengths().to_primitive(); + let uncompressed_lens_array = fsst_array.uncompressed_lengths().to_primitive()?; // Decompress the full dataset. #[allow(clippy::cast_possible_truncation)] @@ -101,7 +107,7 @@ fn fsst_decode_views(fsst_array: &FSSTArray, buf_index: u32) -> (ByteBuffer, Buf } }); - (uncompressed_bytes.freeze(), views.freeze()) + Ok((uncompressed_bytes.freeze(), views.freeze())) } #[cfg(test)] @@ -178,7 +184,7 @@ mod tests { let mut builder = VarBinViewBuilder::with_capacity(chunked_arr.dtype().clone(), chunked_arr.len()); - chunked_arr.append_to_builder(&mut builder); + chunked_arr.append_to_builder(&mut builder).unwrap(); { let arr = builder.finish_into_canonical().into_varbinview(); @@ -188,7 +194,7 @@ mod tests { }; { - let arr2 = chunked_arr.to_varbinview(); + let arr2 = chunked_arr.to_varbinview().unwrap(); let res2 = arr2.with_iterator(|iter| iter.map(|b| b.map(|v| v.to_vec())).collect::>()); assert_eq!(data, res2) diff --git a/encodings/fsst/src/compute/compare.rs b/encodings/fsst/src/compute/compare.rs index a1e8310c58f..7408e8b4e9f 100644 --- a/encodings/fsst/src/compute/compare.rs +++ b/encodings/fsst/src/compute/compare.rs @@ -67,7 +67,7 @@ fn compare_fsst_constant( // No value is lt "" Operator::Lt => BitBuffer::new_unset(left.len()), _ => { - let uncompressed_lengths = left.uncompressed_lengths().to_primitive(); + let uncompressed_lengths = left.uncompressed_lengths().to_primitive()?; match_each_integer_ptype!(uncompressed_lengths.ptype(), |P| { compare_lengths_to_empty( uncompressed_lengths.as_slice::

().iter().copied(), @@ -156,7 +156,8 @@ mod tests { // Ensure fastpath for Eq exists, and returns correct answer let equals = compare(lhs.as_ref(), rhs.as_ref(), Operator::Eq) .unwrap() - .to_bool(); + .to_bool() + .unwrap(); assert_eq!(equals.dtype(), &DType::Bool(Nullability::Nullable)); @@ -168,7 +169,8 @@ mod tests { // Ensure fastpath for Eq exists, and returns correct answer let not_equals = compare(lhs.as_ref(), rhs.as_ref(), Operator::NotEq) .unwrap() - .to_bool(); + .to_bool() + .unwrap(); assert_eq!(not_equals.dtype(), &DType::Bool(Nullability::Nullable)); assert_eq!( diff --git a/encodings/fsst/src/tests.rs b/encodings/fsst/src/tests.rs index 95700b875b5..9022ad96997 100644 --- a/encodings/fsst/src/tests.rs +++ b/encodings/fsst/src/tests.rs @@ -12,6 +12,7 @@ use vortex_array::compute::take; use vortex_buffer::buffer; use vortex_dtype::DType; use vortex_dtype::Nullability; +use vortex_error::VortexResult; use vortex_mask::Mask; use crate::FSSTVTable; @@ -39,7 +40,7 @@ pub(crate) fn build_fsst_array() -> ArrayRef { } #[test] -fn test_fsst_array_ops() { +fn test_fsst_array_ops() -> VortexResult<()> { // first test the scalar_at values let fsst_array = build_fsst_array(); assert_nth_scalar!( @@ -75,7 +76,7 @@ fn test_fsst_array_ops() { // test take let indices = buffer![0, 2].into_array(); - let fsst_taken = take(&fsst_array, &indices).unwrap(); + let fsst_taken = take(&fsst_array, &indices)?; assert!(fsst_taken.is::()); assert_eq!(fsst_taken.len(), 2); assert_nth_scalar!( @@ -92,7 +93,7 @@ fn test_fsst_array_ops() { // test filter let mask = Mask::from_iter([false, true, true]); - let fsst_filtered = filter(&fsst_array, &mask).unwrap(); + let fsst_filtered = filter(&fsst_array, &mask)?; assert!(fsst_filtered.is::()); assert_eq!(fsst_filtered.len(), 2); assert_nth_scalar!( @@ -102,7 +103,9 @@ fn test_fsst_array_ops() { ); // test to_canonical - let canonical_array = fsst_array.to_varbinview().into_array(); + let canonical_array = fsst_array.to_varbinview()?.into_array(); assert_arrays_eq!(fsst_array.to_array(), canonical_array); + + Ok(()) } diff --git a/encodings/pco/benches/pco.rs b/encodings/pco/benches/pco.rs index aa1aa88fade..d0ec478bfff 100644 --- a/encodings/pco/benches/pco.rs +++ b/encodings/pco/benches/pco.rs @@ -46,7 +46,8 @@ pub fn pco_canonical(bencher: Bencher, (size, selectivity): (usize, f64)) { .map(|i| (i % 10000) as i32) .collect::>() .into_array() - .to_primitive(); + .to_primitive() + .unwrap(); let pco_array = PcoArray::from_primitive(&values, 3, 0).unwrap(); let mask = (0..size) @@ -56,5 +57,7 @@ pub fn pco_canonical(bencher: Bencher, (size, selectivity): (usize, f64)) { bencher // Be sure to reconstruct the mask to avoid cached set_indices .with_inputs(|| (&pco_array, Mask::from_buffer(mask.clone()))) - .bench_refs(|(pco_array, mask)| filter(pco_array.to_canonical().as_ref(), mask).unwrap()); + .bench_refs(|(pco_array, mask)| { + filter(pco_array.to_canonical().unwrap().as_ref(), mask).unwrap() + }); } diff --git a/encodings/pco/src/array.rs b/encodings/pco/src/array.rs index 1f8fec7ac09..767be40843e 100644 --- a/encodings/pco/src/array.rs +++ b/encodings/pco/src/array.rs @@ -184,8 +184,8 @@ pub(crate) fn number_type_from_dtype(dtype: &DType) -> NumberType { } fn collect_valid(parray: &PrimitiveArray) -> VortexResult { - let mask = parray.validity_mask(); - Ok(filter(&parray.to_array(), &mask)?.to_primitive()) + let mask = parray.validity_mask()?; + filter(&parray.to_array(), &mask)?.to_primitive() } pub(crate) fn vortex_err_from_pco(err: PcoError) -> VortexError { @@ -350,6 +350,7 @@ impl PcoArray { let slice_value_indices = self .unsliced_validity .to_mask(self.unsliced_n_rows) + .vortex_unwrap() .valid_counts_for_indices(&[self.slice_start, self.slice_stop]); let slice_value_start = slice_value_indices[0]; let slice_value_stop = slice_value_indices[1]; @@ -505,7 +506,7 @@ impl BaseArrayVTable for PcoVTable { } impl CanonicalVTable for PcoVTable { - fn canonicalize(array: &PcoArray) -> Canonical { + fn canonicalize(array: &PcoArray) -> VortexResult { array.decompress().to_canonical() } } @@ -566,7 +567,7 @@ mod tests { Validity::from_iter([false, true, true, true, true, false]), ); let pco = PcoArray::from_primitive(&values, 0, 128).unwrap(); - let decoded = pco.to_primitive(); + let decoded = pco.to_primitive().unwrap(); assert_arrays_eq!( decoded, PrimitiveArray::from_option_iter([ @@ -585,6 +586,6 @@ mod tests { PrimitiveArray::from_option_iter([Some(20u32), Some(30), Some(40), Some(50)]) .into_array(); assert_arrays_eq!(sliced, expected); - assert_arrays_eq!(sliced.to_canonical().into_array(), expected); + assert_arrays_eq!(sliced.to_canonical().unwrap().into_array(), expected); } } diff --git a/encodings/pco/src/compute/cast.rs b/encodings/pco/src/compute/cast.rs index 35536624bf8..4c225593138 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. @@ -86,7 +86,7 @@ mod tests { &DType::Primitive(PType::F64, Nullability::NonNullable) ); - let decoded = casted.to_primitive(); + let decoded = casted.to_primitive().unwrap(); let f64_values = decoded.as_slice::(); assert_eq!(f64_values.len(), 5); assert!((f64_values[0] - 1.0).abs() < f64::EPSILON); @@ -130,7 +130,7 @@ mod tests { &DType::Primitive(PType::U32, Nullability::NonNullable) ); // Verify the values are correct - let decoded = casted.to_primitive(); + let decoded = casted.to_primitive().unwrap(); let u32_values = decoded.as_slice::(); assert_eq!(u32_values, &[20, 30, 40, 50]); } @@ -156,7 +156,7 @@ mod tests { casted.dtype(), &DType::Primitive(PType::U32, Nullability::NonNullable) ); - let decoded = casted.to_primitive(); + let decoded = casted.to_primitive().unwrap(); let expected = PrimitiveArray::from_iter([20u32, 30, 40, 50]); assert_arrays_eq!(decoded, expected); } diff --git a/encodings/pco/src/test.rs b/encodings/pco/src/test.rs index 9c4e96090c5..a855ebda5e0 100644 --- a/encodings/pco/src/test.rs +++ b/encodings/pco/src/test.rs @@ -20,6 +20,7 @@ use vortex_buffer::BufferMut; use vortex_dtype::DType; use vortex_dtype::Nullability; use vortex_dtype::PType; +use vortex_error::VortexResult; use vortex_mask::Mask; use crate::PcoArray; @@ -32,7 +33,7 @@ macro_rules! assert_nth_scalar { } #[test] -fn test_compress_decompress() { +fn test_compress_decompress() -> VortexResult<()> { let data: Vec = (0..200).collect(); let array = PrimitiveArray::from_iter(data.clone()); let compressed = PcoArray::from_primitive(&array, 3, 0).unwrap(); @@ -48,15 +49,17 @@ fn test_compress_decompress() { for i in 0_i32..5 { assert_nth_scalar!(slice, i as usize, 100 + i); } - let primitive = slice.to_primitive(); + let primitive = slice.to_primitive()?; assert_arrays_eq!( primitive, PrimitiveArray::from_iter([100, 101, 102, 103, 104]) ); let slice = compressed.slice(200..200); - let primitive = slice.to_primitive(); + let primitive = slice.to_primitive()?; assert_arrays_eq!(primitive, PrimitiveArray::from_iter(Vec::::new())); + + Ok(()) } #[test] @@ -81,7 +84,7 @@ fn test_empty() { } #[test] -fn test_validity_and_multiple_chunks_and_pages() { +fn test_validity_and_multiple_chunks_and_pages() -> VortexResult<()> { let data: Vec = (0..200).collect(); let mut validity: Vec = vec![true; 200]; validity[7..15].fill(false); @@ -115,15 +118,17 @@ fn test_validity_and_multiple_chunks_and_pages() { let slice = compressed.slice(100..103); assert_nth_scalar!(slice, 0, 100); assert_nth_scalar!(slice, 2, 102); - let primitive = slice.to_primitive(); + let primitive = slice.to_primitive()?; assert_eq!( primitive.validity(), &Validity::Array(BoolArray::from_iter(vec![true, false, true]).to_array()) ); + + Ok(()) } #[test] -fn test_validity_vtable() { +fn test_validity_vtable() -> VortexResult<()> { let data: Vec = (0..5).collect(); let mask_bools = vec![false, true, true, false, true]; let array = PrimitiveArray::new( @@ -131,11 +136,13 @@ 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.validity_mask()?, Mask::from_iter(mask_bools)); assert_eq!( - compressed.slice(1..4).validity_mask(), + compressed.slice(1..4).validity_mask()?, Mask::from_iter(vec![true, true, false]) ); + + Ok(()) } #[test] diff --git a/encodings/runend/benches/run_end_compress.rs b/encodings/runend/benches/run_end_compress.rs index aeb0d9b9c73..aac12122df6 100644 --- a/encodings/runend/benches/run_end_compress.rs +++ b/encodings/runend/benches/run_end_compress.rs @@ -90,8 +90,8 @@ fn take_indices(bencher: Bencher, (length, run_step): (usize, usize)) { ); let source_array = PrimitiveArray::from_iter(0..(length as i32)).into_array(); - let (ends, values) = runend_encode(&values); - let runend_array = RunEndArray::try_new(ends.into_array(), values) + let (ends, values) = runend_encode(&values).unwrap(); + let runend_array = RunEndArray::try_new(ends.into_array(), values.into_array()) .unwrap() .to_array(); diff --git a/encodings/runend/src/array.rs b/encodings/runend/src/array.rs index 13f3b097d24..2472c4f2d3e 100644 --- a/encodings/runend/src/array.rs +++ b/encodings/runend/src/array.rs @@ -320,7 +320,7 @@ impl RunEndArray { /// Run the array through run-end encoding. pub fn encode(array: ArrayRef) -> VortexResult { if let Some(parray) = array.as_opt::() { - let (ends, values) = runend_encode(parray); + let (ends, values) = runend_encode(parray)?; // SAFETY: runend_encode handles this unsafe { Ok(Self::new_unchecked( @@ -391,21 +391,21 @@ impl BaseArrayVTable for RunEndVTable { } impl ValidityVTable for RunEndVTable { - fn is_valid(array: &RunEndArray, index: usize) -> bool { + fn is_valid(array: &RunEndArray, index: usize) -> VortexResult { let physical_idx = array.find_physical_index(index); array.values().is_valid(physical_idx) } - fn all_valid(array: &RunEndArray) -> bool { + fn all_valid(array: &RunEndArray) -> VortexResult { array.values().all_valid() } - fn all_invalid(array: &RunEndArray) -> bool { + fn all_invalid(array: &RunEndArray) -> VortexResult { array.values().all_invalid() } - 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) => { @@ -420,36 +420,36 @@ impl ValidityVTable for RunEndVTable { ) .into_array() }; - Mask::from_buffer(ree_validity.to_bool().bit_buffer().clone()) + Mask::from_buffer(ree_validity.to_bool()?.bit_buffer().clone()) } - } + }) } } impl CanonicalVTable for RunEndVTable { - fn canonicalize(array: &RunEndArray) -> Canonical { - let pends = array.ends().to_primitive(); - match array.dtype() { + fn canonicalize(array: &RunEndArray) -> VortexResult { + let pends = array.ends().to_primitive()?; + Ok(match array.dtype() { DType::Bool(_) => { - let bools = array.values().to_bool(); + let bools = array.values().to_bool()?; Canonical::Bool(runend_decode_bools( pends, bools, array.offset(), array.len(), - )) + )?) } DType::Primitive(..) => { - let pvalues = array.values().to_primitive(); + let pvalues = array.values().to_primitive()?; Canonical::Primitive(runend_decode_primitive( pends, pvalues, array.offset(), array.len(), - )) + )?) } _ => vortex_panic!("Only Primitive and Bool values are supported"), - } + }) } } diff --git a/encodings/runend/src/compress.rs b/encodings/runend/src/compress.rs index 67f7300ba46..b2cf2995b41 100644 --- a/encodings/runend/src/compress.rs +++ b/encodings/runend/src/compress.rs @@ -20,24 +20,25 @@ use vortex_dtype::Nullability; use vortex_dtype::match_each_native_ptype; use vortex_dtype::match_each_unsigned_integer_ptype; use vortex_error::VortexExpect; +use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_scalar::Scalar; use crate::iter::trimmed_ends_iter; /// Run-end encode a `PrimitiveArray`, returning a tuple of `(ends, values)`. -pub fn runend_encode(array: &PrimitiveArray) -> (PrimitiveArray, ArrayRef) { +pub fn runend_encode(array: &PrimitiveArray) -> VortexResult<(PrimitiveArray, ArrayRef)> { let validity = match array.validity() { Validity::NonNullable => None, Validity::AllValid => None, Validity::AllInvalid => { // We can trivially return an all-null REE array - return ( + return Ok(( PrimitiveArray::new(buffer![array.len() as u64], Validity::NonNullable), ConstantArray::new(Scalar::null(array.dtype().clone()), 1).into_array(), - ); + )); } - Validity::Array(a) => Some(a.to_bool().bit_buffer().clone()), + Validity::Array(a) => Some(a.to_bool()?.bit_buffer().clone()), }; let (ends, values) = match validity { @@ -65,9 +66,9 @@ pub fn runend_encode(array: &PrimitiveArray) -> (PrimitiveArray, ArrayRef) { let ends = ends .narrow() .vortex_expect("Ends must succeed downcasting") - .to_primitive(); + .to_primitive()?; - (ends, values) + Ok((ends, values)) } fn runend_encode_primitive(elements: &[T]) -> (Buffer, Buffer) { @@ -162,18 +163,18 @@ pub fn runend_decode_primitive( values: PrimitiveArray, offset: usize, length: usize, -) -> PrimitiveArray { - match_each_native_ptype!(values.ptype(), |P| { +) -> VortexResult { + Ok(match_each_native_ptype!(values.ptype(), |P| { match_each_unsigned_integer_ptype!(ends.ptype(), |E| { runend_decode_typed_primitive( trimmed_ends_iter(ends.as_slice::(), offset, length), values.as_slice::

(), - values.validity_mask(), + values.validity_mask()?, values.dtype().nullability(), length, ) }) - }) + })) } pub fn runend_decode_bools( @@ -181,16 +182,16 @@ pub fn runend_decode_bools( values: BoolArray, offset: usize, length: usize, -) -> BoolArray { - match_each_unsigned_integer_ptype!(ends.ptype(), |E| { +) -> VortexResult { + Ok(match_each_unsigned_integer_ptype!(ends.ptype(), |E| { runend_decode_typed_bool( trimmed_ends_iter(ends.as_slice::(), offset, length), values.bit_buffer(), - values.validity_mask(), + values.validity_mask()?, values.dtype().nullability(), length, ) - }) + })) } pub fn runend_decode_typed_primitive( @@ -300,8 +301,8 @@ mod test { #[test] fn encode() { let arr = PrimitiveArray::from_iter([1i32, 1, 2, 2, 2, 3, 3, 3, 3, 3]); - let (ends, values) = runend_encode(&arr); - let values = values.to_primitive(); + let (ends, values) = runend_encode(&arr).unwrap(); + let values = values.to_primitive().unwrap(); let expected_ends = PrimitiveArray::from_iter(vec![2u8, 5, 10]); assert_arrays_eq!(ends, expected_ends); @@ -317,8 +318,8 @@ mod test { true, true, false, false, true, true, true, true, false, false, ])), ); - let (ends, values) = runend_encode(&arr); - let values = values.to_primitive(); + let (ends, values) = runend_encode(&arr).unwrap(); + let values = values.to_primitive().unwrap(); let expected_ends = PrimitiveArray::from_iter(vec![2u8, 4, 5, 8, 10]); assert_arrays_eq!(ends, expected_ends); @@ -333,8 +334,8 @@ mod test { buffer![0, 0, 0, 0, 0], Validity::from(BitBuffer::new_unset(5)), ); - let (ends, values) = runend_encode(&arr); - let values = values.to_primitive(); + let (ends, values) = runend_encode(&arr).unwrap(); + let values = values.to_primitive().unwrap(); let expected_ends = PrimitiveArray::from_iter(vec![5u64]); assert_arrays_eq!(ends, expected_ends); @@ -346,7 +347,7 @@ mod test { fn decode() { let ends = PrimitiveArray::from_iter([2u32, 5, 10]); let values = PrimitiveArray::from_iter([1i32, 2, 3]); - let decoded = runend_decode_primitive(ends, values, 0, 10); + let decoded = runend_decode_primitive(ends, values, 0, 10).unwrap(); let expected = PrimitiveArray::from_iter(vec![1i32, 1, 2, 2, 2, 3, 3, 3, 3, 3]); assert_arrays_eq!(decoded, expected); diff --git a/encodings/runend/src/compute/cast.rs b/encodings/runend/src/compute/cast.rs index 5663e42504a..91580bad116 100644 --- a/encodings/runend/src/compute/cast.rs +++ b/encodings/runend/src/compute/cast.rs @@ -72,7 +72,7 @@ mod tests { ); // Verify by decoding to canonical form - let decoded = casted.to_primitive(); + let decoded = casted.to_primitive().unwrap(); // RunEnd encoding should expand to [100, 100, 100, 200, 200, 100, 100, 100, 300, 300] assert_eq!(decoded.len(), 10); assert_eq!( @@ -133,7 +133,7 @@ mod tests { let sliced = runend.slice(3..8); // Verify the slice is correct before casting - let sliced_decoded = sliced.to_primitive(); + let sliced_decoded = sliced.to_primitive().unwrap(); assert_eq!(sliced_decoded.len(), 5); assert_arrays_eq!( sliced_decoded, @@ -148,7 +148,7 @@ mod tests { .unwrap(); // Verify the cast preserved the offset - let casted_decoded = casted.to_primitive(); + let casted_decoded = casted.to_primitive().unwrap(); assert_eq!(casted_decoded.len(), 5); assert_arrays_eq!( casted_decoded, diff --git a/encodings/runend/src/compute/compare.rs b/encodings/runend/src/compute/compare.rs index 3b10b7180db..eabc1d1e40d 100644 --- a/encodings/runend/src/compute/compare.rs +++ b/encodings/runend/src/compute/compare.rs @@ -31,10 +31,10 @@ impl CompareKernel for RunEndVTable { ConstantArray::new(const_scalar, lhs.values().len()).as_ref(), operator, ) - .map(|values| { + .and_then(|values| { runend_decode_bools( - lhs.ends().to_primitive(), - values.to_bool(), + lhs.ends().to_primitive()?, + values.to_bool()?, lhs.offset(), lhs.len(), ) @@ -78,7 +78,7 @@ mod test { Operator::Eq, ) .unwrap(); - let res_canon = res.to_bool(); + let res_canon = res.to_bool().unwrap(); assert_eq!( res_canon.bit_buffer(), &BitBuffer::from(vec![ diff --git a/encodings/runend/src/compute/filter.rs b/encodings/runend/src/compute/filter.rs index 3bc373d4fdf..0d679dd18c1 100644 --- a/encodings/runend/src/compute/filter.rs +++ b/encodings/runend/src/compute/filter.rs @@ -45,7 +45,7 @@ impl FilterKernel for RunEndVTable { } else { // This strategy ends up being close to fixed cost based on the number of runs, // rather than the number of indices. - let primitive_run_ends = array.ends().to_primitive(); + let primitive_run_ends = array.ends().to_primitive()?; let (run_ends, values_mask) = match_each_unsigned_integer_ptype!(primitive_run_ends.ptype(), |P| { filter_run_end_primitive( @@ -77,7 +77,7 @@ register_kernel!(FilterKernelAdapter(RunEndVTable).lift()); // We expose this function to our benchmarks. pub fn filter_run_end(array: &RunEndArray, mask: &Mask) -> VortexResult { - let primitive_run_ends = array.ends().to_primitive(); + let primitive_run_ends = array.ends().to_primitive()?; let (run_ends, values_mask) = match_each_unsigned_integer_ptype!(primitive_run_ends.ptype(), |P| { filter_run_end_primitive( @@ -171,11 +171,19 @@ mod tests { let filtered_run_end = filtered.as_::(); assert_eq!( - filtered_run_end.ends().to_primitive().as_slice::(), + filtered_run_end + .ends() + .to_primitive() + .unwrap() + .as_slice::(), [2, 4] ); assert_eq!( - filtered_run_end.values().to_primitive().as_slice::(), + filtered_run_end + .values() + .to_primitive() + .unwrap() + .as_slice::(), [1, 5] ); } @@ -191,11 +199,19 @@ mod tests { let filtered_run_end = filtered.as_::(); assert_eq!( - filtered_run_end.ends().to_primitive().as_slice::(), + filtered_run_end + .ends() + .to_primitive() + .unwrap() + .as_slice::(), [1, 2, 3] ); assert_eq!( - filtered_run_end.values().to_primitive().as_slice::(), + filtered_run_end + .values() + .to_primitive() + .unwrap() + .as_slice::(), [1, 4, 2] ); } diff --git a/encodings/runend/src/compute/is_sorted.rs b/encodings/runend/src/compute/is_sorted.rs index b49a5b96e83..89033620dbd 100644 --- a/encodings/runend/src/compute/is_sorted.rs +++ b/encodings/runend/src/compute/is_sorted.rs @@ -17,7 +17,7 @@ impl IsSortedKernel for RunEndVTable { } fn is_strict_sorted(&self, array: &RunEndArray) -> VortexResult> { - is_strict_sorted(array.to_canonical().as_ref()) + is_strict_sorted(array.to_canonical()?.as_ref()) } } diff --git a/encodings/runend/src/compute/take.rs b/encodings/runend/src/compute/take.rs index 75ea164f3f7..cf3d810155b 100644 --- a/encodings/runend/src/compute/take.rs +++ b/encodings/runend/src/compute/take.rs @@ -30,7 +30,7 @@ impl TakeKernel for RunEndVTable { reason = "index cast to usize inside macro" )] fn take(&self, array: &RunEndArray, indices: &dyn Array) -> VortexResult { - let primitive_indices = indices.to_primitive(); + let primitive_indices = indices.to_primitive()?; let checked_indices = match_each_integer_ptype!(primitive_indices.ptype(), |P| { primitive_indices @@ -59,7 +59,7 @@ pub fn take_indices_unchecked>( indices: &[T], validity: &Validity, ) -> VortexResult { - let ends = array.ends().to_primitive(); + let ends = array.ends().to_primitive()?; let ends_len = ends.len(); // TODO(joe): use the validity mask to skip search sorted. diff --git a/encodings/runend/src/lib.rs b/encodings/runend/src/lib.rs index c5e51d39c6a..58913be532d 100644 --- a/encodings/runend/src/lib.rs +++ b/encodings/runend/src/lib.rs @@ -36,7 +36,7 @@ impl EncodeVTable for RunEndVTable { _like: Option<&RunEndArray>, ) -> VortexResult> { let parray = canonical.clone().into_primitive(); - let (ends, values) = runend_encode(&parray); + let (ends, values) = runend_encode(&parray)?; // SAFETY: runend_decode implementation must return valid RunEndArray // components. unsafe { diff --git a/encodings/sequence/src/array.rs b/encodings/sequence/src/array.rs index 8be4483b419..8aa3557e395 100644 --- a/encodings/sequence/src/array.rs +++ b/encodings/sequence/src/array.rs @@ -322,7 +322,7 @@ impl BaseArrayVTable for SequenceVTable { } impl CanonicalVTable for SequenceVTable { - fn canonicalize(array: &SequenceArray) -> Canonical { + fn canonicalize(array: &SequenceArray) -> VortexResult { let prim = match_each_native_ptype!(array.ptype(), |P| { let base = array.base().cast::

(); let multiplier = array.multiplier().cast::

(); @@ -333,7 +333,7 @@ impl CanonicalVTable for SequenceVTable { PrimitiveArray::new(values, array.dtype.nullability().into()) }); - Canonical::Primitive(prim) + Ok(Canonical::Primitive(prim)) } } @@ -358,20 +358,20 @@ impl OperationsVTable for SequenceVTable { } impl ValidityVTable for SequenceVTable { - fn is_valid(_array: &SequenceArray, _index: usize) -> bool { - true + fn is_valid(_array: &SequenceArray, _index: usize) -> VortexResult { + Ok(true) } - fn all_valid(_array: &SequenceArray) -> bool { - true + fn all_valid(_array: &SequenceArray) -> VortexResult { + Ok(true) } - fn all_invalid(_array: &SequenceArray) -> bool { - false + fn all_invalid(_array: &SequenceArray) -> VortexResult { + Ok(false) } - fn validity_mask(array: &SequenceArray) -> Mask { - Mask::AllTrue(array.len()) + fn validity_mask(array: &SequenceArray) -> VortexResult { + Ok(Mask::AllTrue(array.len())) } } @@ -414,7 +414,7 @@ mod tests { let canon = PrimitiveArray::from_iter((0..4).map(|i| 2i64 + i * 3)); assert_eq!( - arr.to_primitive().as_slice::(), + arr.to_primitive().unwrap().as_slice::(), canon.as_slice::() ) } @@ -428,7 +428,7 @@ mod tests { let canon = PrimitiveArray::from_iter((2..3).map(|i| 2i64 + i * 3)); assert_eq!( - arr.to_primitive().as_slice::(), + arr.to_primitive().unwrap().as_slice::(), canon.as_slice::() ) } diff --git a/encodings/sequence/src/compress.rs b/encodings/sequence/src/compress.rs index c24838b2cc3..b472b95d807 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 indices = indices.to_primitive(); + let mask = indices.validity_mask()?; + let indices = indices.to_primitive()?; let result_nullability = array.dtype().nullability() | indices.dtype().nullability(); match_each_integer_ptype!(indices.ptype(), |T| { diff --git a/encodings/sparse/src/canonical.rs b/encodings/sparse/src/canonical.rs index 3afb5de625d..d5e59453a09 100644 --- a/encodings/sparse/src/canonical.rs +++ b/encodings/sparse/src/canonical.rs @@ -43,6 +43,7 @@ use vortex_dtype::match_each_integer_ptype; use vortex_dtype::match_each_native_ptype; use vortex_error::VortexError; use vortex_error::VortexExpect; +use vortex_error::VortexResult; use vortex_error::vortex_panic; use vortex_scalar::DecimalScalar; use vortex_scalar::ListScalar; @@ -54,7 +55,7 @@ use crate::SparseArray; use crate::SparseVTable; impl CanonicalVTable for SparseVTable { - fn canonicalize(array: &SparseArray) -> Canonical { + fn canonicalize(array: &SparseArray) -> VortexResult { if array.patches().num_patches() == 0 { return ConstantArray::new(array.fill_scalar().clone(), array.len()).to_canonical(); } @@ -62,7 +63,7 @@ impl CanonicalVTable for SparseVTable { match array.dtype() { DType::Null => { assert!(array.fill_scalar().is_null()); - Canonical::Null(NullArray::new(array.len())) + Ok(Canonical::Null(NullArray::new(array.len()))) } DType::Bool(..) => { let resolved_patches = array.resolved_patches(); @@ -123,7 +124,7 @@ fn canonicalize_sparse_lists( array: &SparseArray, values_dtype: Arc, nullability: Nullability, -) -> Canonical { +) -> VortexResult { // TODO(connor): We should move this to `vortex-dtype` so that we can use this elsewhere. macro_rules! match_smallest_offset_type { ($n_elements:expr, | $offset_type:ident | $body:block) => {{ @@ -147,16 +148,16 @@ fn canonicalize_sparse_lists( let resolved_patches = array.resolved_patches(); - let indices = resolved_patches.indices().to_primitive(); - let values = resolved_patches.values().to_listview(); + let indices = resolved_patches.indices().to_primitive()?; + let values = resolved_patches.values().to_listview()?; let fill_value = array.fill_scalar().as_list(); 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()?, nullability); - match_each_integer_ptype!(indices.ptype(), |I| { + Ok(match_each_integer_ptype!(indices.ptype(), |I| { match_smallest_offset_type!(total_canonical_values, |O| { canonicalize_sparse_lists_inner::( indices.as_slice(), @@ -168,7 +169,7 @@ fn canonicalize_sparse_lists( validity, ) }) - }) + })) } fn canonicalize_sparse_lists_inner( @@ -218,13 +219,16 @@ fn canonicalize_sparse_lists_inner( } /// Canonicalize a sparse [`FixedSizeListArray`] by expanding it into a dense representation. -fn canonicalize_sparse_fixed_size_list(array: &SparseArray, nullability: Nullability) -> Canonical { +fn canonicalize_sparse_fixed_size_list( + array: &SparseArray, + nullability: Nullability, +) -> VortexResult { let resolved_patches = array.resolved_patches(); - let indices = resolved_patches.indices().to_primitive(); - let values = resolved_patches.values().to_fixed_size_list(); + let indices = resolved_patches.indices().to_primitive()?; + 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()?, nullability); match_each_integer_ptype!(indices.ptype(), |I| { canonicalize_sparse_fixed_size_list_inner::( @@ -249,7 +253,7 @@ fn canonicalize_sparse_fixed_size_list_inner( fill_value: ListScalar, array_len: usize, validity: Validity, -) -> Canonical { +) -> VortexResult { let list_size = values.list_size(); let element_dtype = values.elements().dtype(); let total_elements = array_len * list_size as usize; @@ -271,7 +275,7 @@ fn canonicalize_sparse_fixed_size_list_inner( ); // Append the patch value, handling null patches by appending defaults. - if values.validity().is_valid(patch_idx) { + if values.validity().is_valid(patch_idx)? { let patch_list = values.fixed_size_list_elements_at(patch_idx); for i in 0..list_size as usize { builder @@ -296,9 +300,9 @@ fn canonicalize_sparse_fixed_size_list_inner( let elements = builder.finish(); // SAFETY: elements.len() == array_len * list_size, validity length matches array_len. - Canonical::FixedSizeList(unsafe { + Ok(Canonical::FixedSizeList(unsafe { FixedSizeListArray::new_unchecked(elements, list_size, validity, array_len) - }) + })) } /// Append `count` copies of a fixed-size list to the builder. @@ -324,7 +328,7 @@ fn append_n_lists( } } -fn canonicalize_sparse_bools(patches: &Patches, fill_value: &Scalar) -> Canonical { +fn canonicalize_sparse_bools(patches: &Patches, fill_value: &Scalar) -> VortexResult { let (fill_bool, validity) = if fill_value.is_null() { (false, Validity::AllInvalid) } else { @@ -343,7 +347,7 @@ fn canonicalize_sparse_bools(patches: &Patches, fill_value: &Scalar) -> Canonica let bools = BoolArray::from_bit_buffer(BitBuffer::full(fill_bool, patches.array_len()), validity); - Canonical::Bool(bools.patch(patches)) + Ok(Canonical::Bool(bools.patch(patches))) } fn canonicalize_sparse_primitives< @@ -351,7 +355,7 @@ fn canonicalize_sparse_primitives< >( patches: &Patches, fill_value: &Scalar, -) -> Canonical { +) -> VortexResult { let (primitive_fill, validity) = if fill_value.is_null() { (T::default(), Validity::AllInvalid) } else { @@ -369,7 +373,7 @@ fn canonicalize_sparse_primitives< let parray = PrimitiveArray::new(buffer![primitive_fill; patches.array_len()], validity); - Canonical::Primitive(parray.patch(patches)) + Ok(Canonical::Primitive(parray.patch(patches))) } fn canonicalize_sparse_struct( @@ -379,7 +383,7 @@ fn canonicalize_sparse_struct( // Resolution is unnecessary b/c we're just pushing the patches into the fields. unresolved_patches: &Patches, len: usize, -) -> Canonical { +) -> VortexResult { let (fill_values, top_level_fill_validity) = match fill_struct.fields() { Some(fill_values) => (fill_values.collect::>(), Validity::AllValid), None => ( @@ -390,7 +394,7 @@ fn canonicalize_sparse_struct( Validity::AllInvalid, ), }; - let patch_values_as_struct = unresolved_patches.values().to_struct(); + let patch_values_as_struct = unresolved_patches.values().to_struct()?; let columns_patch_values = patch_values_as_struct.fields(); let names = patch_values_as_struct.names(); let validity = if dtype.is_nullable() { @@ -399,17 +403,17 @@ fn canonicalize_sparse_struct( unresolved_patches.offset(), unresolved_patches.indices(), &Validity::from_mask( - unresolved_patches.values().validity_mask(), + unresolved_patches.values().validity_mask()?, Nullability::Nullable, ), - ) + )? } else { top_level_fill_validity .into_non_nullable(len) .unwrap_or_else(|| vortex_panic!("fill validity should match sparse array nullability")) }; - StructArray::try_from_iter_with_validity( + Ok(StructArray::try_from_iter_with_validity( names.iter().zip_eq( columns_patch_values .iter() @@ -428,7 +432,7 @@ fn canonicalize_sparse_struct( validity, ) .map(Canonical::Struct) - .vortex_expect("Creating struct array") + .vortex_expect("Creating struct array")) } fn canonicalize_sparse_decimal( @@ -437,7 +441,7 @@ fn canonicalize_sparse_decimal( fill_value: DecimalScalar, patches: &Patches, len: usize, -) -> Canonical { +) -> VortexResult { let mut builder = DecimalBuilder::with_capacity::(len, decimal_dtype, nullability); match fill_value.decimal_value() { Some(fill_value) => { @@ -454,24 +458,24 @@ fn canonicalize_sparse_decimal( } let filled_array = builder.finish_into_decimal(); let array = filled_array.patch(patches); - Canonical::Decimal(array) + Ok(Canonical::Decimal(array)) } fn canonicalize_varbin( array: &SparseArray, dtype: DType, fill_value: Option, -) -> Canonical { +) -> VortexResult { 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 indices = patches.indices().to_primitive()?; + let values = patches.values().to_varbinview()?; + let validity = Validity::from_mask(array.validity_mask()?, dtype.nullability()); let len = array.len(); - match_each_integer_ptype!(indices.ptype(), |I| { + Ok(match_each_integer_ptype!(indices.ptype(), |I| { let indices = indices.buffer::(); canonicalize_varbin_inner::(fill_value, indices, values, dtype, validity, len) - }) + })) } fn canonicalize_varbin_inner( @@ -558,7 +562,7 @@ mod test { let values = BoolArray::from_iter([Some(true), None, Some(false)]).into_array(); let sparse_bools = SparseArray::try_new(indices, values, 10, Scalar::from(fill_value)).unwrap(); - let actual = sparse_bools.to_bool(); + let actual = sparse_bools.to_bool().unwrap(); let expected = BoolArray::from_iter([ Some(true), @@ -587,7 +591,7 @@ mod test { SparseArray::try_new(indices, values, 10, Scalar::from(fill_value)).unwrap(); assert_eq!(*sparse_ints.dtype(), DType::Primitive(PType::I32, Nullable)); - let flat_ints = sparse_ints.to_primitive(); + let flat_ints = sparse_ints.to_primitive().unwrap(); let expected = PrimitiveArray::from_option_iter([ Some(0i32), None, @@ -670,7 +674,7 @@ mod test { .unwrap() .to_array(); - let actual = sparse_struct.to_struct(); + let actual = sparse_struct.to_struct().unwrap(); assert_arrays_eq!(actual, expected); } @@ -737,7 +741,7 @@ mod test { .unwrap() .to_array(); - let actual = sparse_struct.to_struct(); + let actual = sparse_struct.to_struct().unwrap(); assert_arrays_eq!(actual, expected); } @@ -767,6 +771,7 @@ mod test { let actual = sparse_struct .to_decimal() + .unwrap() .to_array() .into_arrow_preferred() .unwrap(); @@ -796,7 +801,7 @@ mod test { ) .unwrap(); - let actual = array.to_varbinview().into_array(); + let actual = array.to_varbinview().unwrap().into_array(); let expected = >::from_iter([ Some("hello"), Some("123"), @@ -837,7 +842,7 @@ mod test { ) .unwrap(); - let actual = array.to_varbinview().into_array(); + let actual = array.to_varbinview().unwrap().into_array(); let expected = >::from_iter([ Some("hello"), None, @@ -871,7 +876,7 @@ mod test { ) .unwrap(); - let actual = array.to_varbinview().into_array(); + let actual = array.to_varbinview().unwrap().into_array(); let expected = VarBinViewArray::from_iter_str([ "hello", "123", "123", "goodbye", "hello", "bonjour", "123", "123", "你好", ]) @@ -901,7 +906,7 @@ mod test { ) .unwrap(); - let actual = array.to_varbinview().into_array(); + let actual = array.to_varbinview().unwrap().into_array(); let expected = >::from_iter([ Some("hello"), None, @@ -942,7 +947,7 @@ mod test { ) .unwrap(); - let actual = array.to_varbinview().into_array(); + let actual = array.to_varbinview().unwrap().into_array(); let expected = VarBinViewArray::from_iter_nullable_bin([ Some(b"hello" as &[u8]), Some(b"123"), @@ -985,8 +990,8 @@ mod test { .unwrap() .into_array(); - let actual = sparse.to_canonical().into_array(); - let result_listview = actual.to_listview(); + let actual = sparse.to_canonical().unwrap().into_array(); + let result_listview = actual.to_listview().unwrap(); // Check the structure assert_eq!(result_listview.len(), 6); @@ -1000,7 +1005,7 @@ mod test { assert_eq!(result_listview.size_at(5), 1); // [2] // Verify actual values - let elements_array = result_listview.elements().to_primitive(); + let elements_array = result_listview.elements().to_primitive().unwrap(); let elements_slice = elements_array.as_slice::(); let list0_offset = result_listview.offset_at(0); @@ -1037,8 +1042,8 @@ mod test { .unwrap() .into_array(); - let actual = sparse.to_canonical().into_array(); - let result_listview = actual.to_listview(); + let actual = sparse.to_canonical().unwrap().into_array(); + let result_listview = actual.to_listview().unwrap(); // Check the structure assert_eq!(result_listview.len(), 6); @@ -1052,7 +1057,7 @@ mod test { assert_eq!(result_listview.size_at(5), 1); // [2] - extra element beyond original slice // Verify actual values - let elements_array = result_listview.elements().to_primitive(); + let elements_array = result_listview.elements().to_primitive().unwrap(); let elements_slice = elements_array.as_slice::(); let list0_offset = result_listview.offset_at(0); @@ -1080,8 +1085,8 @@ mod test { .unwrap() .into_array(); - let actual = sparse.to_canonical().into_array(); - let result_listview = actual.to_listview(); + let actual = sparse.to_canonical().unwrap().into_array(); + let result_listview = actual.to_listview().unwrap(); // Check the structure assert_eq!(result_listview.len(), 6); @@ -1095,7 +1100,7 @@ mod test { assert_eq!(result_listview.size_at(5), 1); // [2] from sparse // Verify actual values - let elements_array = result_listview.elements().to_primitive(); + let elements_array = result_listview.elements().to_primitive().unwrap(); let elements_slice = elements_array.as_slice::(); // List 0: [1] @@ -1152,7 +1157,7 @@ mod test { ) .unwrap(); - let actual = array.to_varbinview().into_array(); + let actual = array.to_varbinview().unwrap().into_array(); let expected = VarBinViewArray::from_iter_nullable_bin([ Some(b"hello" as &[u8]), None, @@ -1190,7 +1195,7 @@ mod test { .unwrap() .into_array(); - let actual = sparse.to_canonical().into_array(); + let actual = sparse.to_canonical().unwrap().into_array(); // Expected: [1,2,3], null, [4,5,6], [7,8,9], null. let expected_elements = @@ -1227,7 +1232,7 @@ mod test { .unwrap() .into_array(); - let actual = sparse.to_canonical().into_array(); + let actual = sparse.to_canonical().unwrap().into_array(); // Expected: [1,2], [99,88], [3,4], [99,88], [5,6], [99,88]. let expected_elements = buffer![1i32, 2, 99, 88, 3, 4, 99, 88, 5, 6, 99, 88].into_array(); @@ -1264,7 +1269,7 @@ mod test { .unwrap() .into_array(); - let actual = sparse.to_canonical().into_array(); + let actual = sparse.to_canonical().unwrap().into_array(); // Expected validity: [true, true, true, false, true, true]. // Expected elements: [7,8], [10,20], [7,8], [30,40], [50,60], [7,8]. @@ -1311,7 +1316,7 @@ mod test { .unwrap() .into_array(); - let actual = sparse.to_canonical().into_array(); + let actual = sparse.to_canonical().unwrap().into_array(); // Build expected: 97 copies of [99,99] with patches at positions 5, 50, 95. let mut expected_elements_vec = Vec::with_capacity(200); @@ -1367,7 +1372,7 @@ mod test { .unwrap() .into_array(); - let actual = sparse.to_canonical().into_array(); + let actual = sparse.to_canonical().unwrap().into_array(); // Expected: just [42, 43]. let expected_elements = buffer![42i32, 43].into_array(); @@ -1392,7 +1397,7 @@ mod test { .unwrap() .into_array(); - let actual = sparse.to_canonical().into_array(); + let actual = sparse.to_canonical().unwrap().into_array(); let mut expected_elements = buffer_mut![1, 2, 1, 2]; expected_elements.extend(buffer![42i32; 252]); let expected = ListArray::try_new( @@ -1404,7 +1409,7 @@ mod test { .into_array(); assert_eq!( - actual.to_listview().offsets().dtype(), + actual.to_listview().unwrap().offsets().dtype(), &DType::Primitive(PType::U16, NonNullable) ); assert_arrays_eq!(&actual, &expected); @@ -1458,8 +1463,8 @@ mod test { .unwrap(); // Convert to canonical form - this triggers the function we're testing - let canonical = sparse.to_canonical().into_array(); - let result_listview = canonical.to_listview(); + let canonical = sparse.to_canonical().unwrap().into_array(); + let result_listview = canonical.to_listview().unwrap(); // Verify the structure assert_eq!(result_listview.len(), 10); @@ -1471,7 +1476,7 @@ mod test { if size == 0 { vec![] // null/empty list } else { - let elements = result_listview.elements().to_primitive(); + let elements = result_listview.elements().to_primitive().unwrap(); let slice = elements.as_slice::(); slice[offset..offset + size].to_vec() } @@ -1536,8 +1541,8 @@ mod test { let sparse = SparseArray::try_new(indices, values, 5, Scalar::null(sliced.dtype().clone())).unwrap(); - let canonical = sparse.to_canonical().into_array(); - let result_listview = canonical.to_listview(); + let canonical = sparse.to_canonical().unwrap().into_array(); + let result_listview = canonical.to_listview().unwrap(); assert_eq!(result_listview.len(), 5); @@ -1548,7 +1553,7 @@ mod test { if size == 0 { vec![] // null/empty list } else { - let elements = result_listview.elements().to_primitive(); + let elements = result_listview.elements().to_primitive().unwrap(); let slice = elements.as_slice::(); slice[offset..offset + size].to_vec() } diff --git a/encodings/sparse/src/compute/cast.rs b/encodings/sparse/src/compute/cast.rs index 341898a15e7..3c087f5ad62 100644 --- a/encodings/sparse/src/compute/cast.rs +++ b/encodings/sparse/src/compute/cast.rs @@ -66,7 +66,7 @@ mod tests { &DType::Primitive(PType::I64, Nullability::NonNullable) ); - let decoded = casted.to_primitive(); + let decoded = casted.to_primitive().unwrap(); let values = decoded.as_slice::(); assert_eq!(values[2], 100); assert_eq!(values[5], 200); diff --git a/encodings/sparse/src/compute/take.rs b/encodings/sparse/src/compute/take.rs index 9a9a553619b..3ee4d2c38e2 100644 --- a/encodings/sparse/src/compute/take.rs +++ b/encodings/sparse/src/compute/take.rs @@ -122,11 +122,21 @@ mod test { let taken = taken_arr.as_::(); assert_eq!( - taken.patches().indices().to_primitive().as_slice::(), + taken + .patches() + .indices() + .to_primitive() + .unwrap() + .as_slice::(), [1] ); assert_eq!( - taken.patches().values().to_primitive().as_slice::(), + taken + .patches() + .values() + .to_primitive() + .unwrap() + .as_slice::(), [0.47f64] ); assert_eq!(taken.len(), 2); diff --git a/encodings/sparse/src/lib.rs b/encodings/sparse/src/lib.rs index ea2aea04739..b0324a5ec88 100644 --- a/encodings/sparse/src/lib.rs +++ b/encodings/sparse/src/lib.rs @@ -257,7 +257,7 @@ impl SparseArray { fill_value.dtype() ) } - let mask = array.validity_mask(); + let mask = array.validity_mask()?; if mask.all_false() { // Array is constant NULL @@ -300,7 +300,7 @@ impl SparseArray { } else { // TODO(robert): Support other dtypes, only thing missing is getting most common value out of the array let (top_pvalue, _) = array - .to_primitive() + .to_primitive()? .top_value()? .vortex_expect("Non empty or all null array"); @@ -313,7 +313,7 @@ impl SparseArray { &compare(array, &fill_array, Operator::NotEq)?, &Scalar::bool(true, Nullability::NonNullable), )? - .to_bool() + .to_bool()? .bit_buffer() .clone(), ); @@ -361,43 +361,43 @@ impl BaseArrayVTable for SparseVTable { } impl ValidityVTable for SparseVTable { - fn is_valid(array: &SparseArray, index: usize) -> bool { - match array.patches().get_patched(index) { + fn is_valid(array: &SparseArray, index: usize) -> VortexResult { + Ok(match array.patches().get_patched(index) { None => array.fill_scalar().is_valid(), Some(patch_value) => patch_value.is_valid(), - } + }) } - fn all_valid(array: &SparseArray) -> bool { + fn all_valid(array: &SparseArray) -> VortexResult { if array.fill_scalar().is_null() { // We need _all_ values to be patched, and all patches to be valid - return array.patches().values().len() == array.len() - && array.patches().values().all_valid(); + return Ok(array.patches().values().len() == array.len() + && array.patches().values().all_valid()?); } array.patches().values().all_valid() } - fn all_invalid(array: &SparseArray) -> bool { + fn all_invalid(array: &SparseArray) -> VortexResult { if !array.fill_scalar().is_null() { // We need _all_ values to be patched, and all patches to be invalid - return array.patches().values().len() == array.len() - && array.patches().values().all_invalid(); + return Ok(array.patches().values().len() == array.len() + && array.patches().values().all_invalid()?); } array.patches().values().all_invalid() } - 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 { @@ -406,7 +406,7 @@ impl ValidityVTable for SparseVTable { BitBufferMut::new_unset(len) }; - let indices = array.patches().indices().to_primitive(); + let indices = array.patches().indices().to_primitive()?; let index_offset = array.patches().offset(); match_each_integer_ptype!(indices.ptype(), |I| { @@ -414,7 +414,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())) } } @@ -558,7 +558,7 @@ mod 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]) ); } @@ -579,7 +579,7 @@ mod test { .slice(2..7); assert_eq!( - sliced.validity_mask(), + sliced.validity_mask().unwrap(), Mask::from_iter(vec![false, true, true, false, true]) ); } @@ -597,7 +597,12 @@ mod 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 ] @@ -607,7 +612,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] @@ -640,9 +645,9 @@ mod test { None, ) .vortex_unwrap(); - let canonical = sparse.to_primitive(); + let canonical = sparse.to_primitive().unwrap(); 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, ]) @@ -659,7 +664,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 b1184bad4ed..ca2b7d4502d 100644 --- a/encodings/sparse/src/ops.rs +++ b/encodings/sparse/src/ops.rs @@ -59,7 +59,7 @@ mod tests { let mut expected = vec![999u64; 1000]; expected[0] = 0; - let values = sliced.to_primitive(); + let values = sliced.to_primitive().unwrap(); assert_arrays_eq!(values, PrimitiveArray::from_iter(expected)); } } diff --git a/encodings/zigzag/src/array.rs b/encodings/zigzag/src/array.rs index 664f18475c3..7f182eb8b85 100644 --- a/encodings/zigzag/src/array.rs +++ b/encodings/zigzag/src/array.rs @@ -164,8 +164,10 @@ impl BaseArrayVTable for ZigZagVTable { } impl CanonicalVTable for ZigZagVTable { - fn canonicalize(array: &ZigZagArray) -> Canonical { - Canonical::Primitive(zigzag_decode(array.encoded().to_primitive())) + fn canonicalize(array: &ZigZagArray) -> VortexResult { + Ok(Canonical::Primitive(zigzag_decode( + array.encoded().to_primitive()?, + ))) } } @@ -239,7 +241,7 @@ mod test { #[test] fn test_compute_statistics() { let array = buffer![1i32, -5i32, 2, 3, 4, 5, 6, 7, 8, 9, 10].into_array(); - let canonical = array.to_canonical(); + let canonical = array.to_canonical().unwrap(); let zigzag = ZigZagVTable .as_vtable() .encode(&canonical, None) diff --git a/encodings/zigzag/src/compress.rs b/encodings/zigzag/src/compress.rs index d3a30a557ea..8f1aab9d45f 100644 --- a/encodings/zigzag/src/compress.rs +++ b/encodings/zigzag/src/compress.rs @@ -81,7 +81,7 @@ mod test { let compressed = zigzag_encode(PrimitiveArray::from_iter(-100_i8..100)).unwrap(); assert!(compressed.is::()); assert_eq!( - compressed.to_primitive().as_slice::(), + compressed.to_primitive().unwrap().as_slice::(), (-100_i8..100).collect::>() ); } @@ -90,7 +90,7 @@ mod test { let compressed = zigzag_encode(PrimitiveArray::from_iter(-100_i16..100)).unwrap(); assert!(compressed.is::()); assert_eq!( - compressed.to_primitive().as_slice::(), + compressed.to_primitive().unwrap().as_slice::(), (-100_i16..100).collect::>() ); } @@ -99,7 +99,7 @@ mod test { let compressed = zigzag_encode(PrimitiveArray::from_iter(-100_i32..100)).unwrap(); assert!(compressed.is::()); assert_eq!( - compressed.to_primitive().as_slice::(), + compressed.to_primitive().unwrap().as_slice::(), (-100_i32..100).collect::>() ); } @@ -108,7 +108,7 @@ mod test { let compressed = zigzag_encode(PrimitiveArray::from_iter(-100_i64..100)).unwrap(); assert!(compressed.is::()); assert_eq!( - compressed.to_primitive().as_slice::(), + compressed.to_primitive().unwrap().as_slice::(), (-100_i64..100).collect::>() ); } diff --git a/encodings/zigzag/src/compute/cast.rs b/encodings/zigzag/src/compute/cast.rs index 15c6040507f..6e06d00215f 100644 --- a/encodings/zigzag/src/compute/cast.rs +++ b/encodings/zigzag/src/compute/cast.rs @@ -67,7 +67,7 @@ mod tests { "Cast should preserve ZigZag encoding" ); - let decoded = casted.to_primitive(); + let decoded = casted.to_primitive().unwrap(); assert_arrays_eq!(decoded, PrimitiveArray::from_iter([-100i64, -1, 0, 1, 100])); } @@ -88,7 +88,7 @@ mod tests { "Should remain ZigZag encoded" ); - let decoded = casted.to_primitive(); + let decoded = casted.to_primitive().unwrap(); assert_arrays_eq!( decoded, PrimitiveArray::from_iter([100i16, -50, 0, 25, -100]) @@ -109,7 +109,7 @@ mod tests { "Should remain ZigZag encoded" ); - let decoded64 = casted64.to_primitive(); + let decoded64 = casted64.to_primitive().unwrap(); assert_arrays_eq!( decoded64, PrimitiveArray::from_iter([1000i64, -500, 0, 250, -1000]) diff --git a/encodings/zigzag/src/compute/mod.rs b/encodings/zigzag/src/compute/mod.rs index 56f2c81f3b4..e34e40becb8 100644 --- a/encodings/zigzag/src/compute/mod.rs +++ b/encodings/zigzag/src/compute/mod.rs @@ -98,7 +98,9 @@ mod tests { let zigzag = ZigZagVTable .as_vtable() .encode( - &PrimitiveArray::new(buffer![-189, -160, 1], Validity::AllValid).to_canonical(), + &PrimitiveArray::new(buffer![-189, -160, 1], Validity::AllValid) + .to_canonical() + .unwrap(), None, ) .unwrap() @@ -113,18 +115,22 @@ mod tests { fn take_zigzag() { let zigzag = ZigZagVTable .as_vtable() - .encode(&buffer![-189, -160, 1].into_array().to_canonical(), None) + .encode( + &buffer![-189, -160, 1].into_array().to_canonical().unwrap(), + None, + ) .unwrap() .unwrap(); let indices = buffer![0, 2].into_array(); - let actual = take(&zigzag, &indices).unwrap().to_primitive(); + let actual = take(&zigzag, &indices).unwrap().to_primitive().unwrap(); let expected = ZigZagVTable .as_vtable() - .encode(&buffer![-189, 1].into_array().to_canonical(), None) + .encode(&buffer![-189, 1].into_array().to_canonical().unwrap(), None) .unwrap() .unwrap() - .to_primitive(); + .to_primitive() + .unwrap(); assert_arrays_eq!(actual, expected); } @@ -132,17 +138,24 @@ mod tests { fn filter_zigzag() { let zigzag = ZigZagVTable .as_vtable() - .encode(&buffer![-189, -160, 1].into_array().to_canonical(), None) + .encode( + &buffer![-189, -160, 1].into_array().to_canonical().unwrap(), + None, + ) .unwrap() .unwrap(); let filter_mask = BitBuffer::from(vec![true, false, true]).into(); - let actual = filter(&zigzag, &filter_mask).unwrap().to_primitive(); + let actual = filter(&zigzag, &filter_mask) + .unwrap() + .to_primitive() + .unwrap(); let expected = ZigZagVTable .as_vtable() - .encode(&buffer![-189, 1].into_array().to_canonical(), None) + .encode(&buffer![-189, 1].into_array().to_canonical().unwrap(), None) .unwrap() .unwrap() - .to_primitive(); + .to_primitive() + .unwrap(); assert_arrays_eq!(actual, expected); } @@ -156,7 +169,8 @@ mod tests { .encode( &buffer![-189i32, -160, 1, 42, -73] .into_array() - .to_canonical(), + .to_canonical() + .unwrap(), None, ) .unwrap() @@ -169,7 +183,8 @@ mod tests { .encode( &buffer![1000i64, -2000, 3000, -4000, 5000] .into_array() - .to_canonical(), + .to_canonical() + .unwrap(), None, ) .unwrap() @@ -181,7 +196,7 @@ mod tests { PrimitiveArray::from_option_iter([Some(-10i16), None, Some(20), Some(-30), None]); let zigzag = ZigZagVTable .as_vtable() - .encode(&array.to_canonical(), None) + .encode(&array.to_canonical().unwrap(), None) .unwrap() .unwrap(); test_filter_conformance(zigzag.as_ref()); @@ -197,7 +212,8 @@ mod tests { .encode( &buffer![-100i32, 200, -300, 400, -500] .into_array() - .to_canonical(), + .to_canonical() + .unwrap(), None, ) .unwrap() @@ -208,7 +224,10 @@ mod tests { let zigzag = ZigZagVTable .as_vtable() .encode( - &buffer![-127i8, 0, 127, -1, 1].into_array().to_canonical(), + &buffer![-127i8, 0, 127, -1, 1] + .into_array() + .to_canonical() + .unwrap(), None, ) .unwrap() @@ -226,7 +245,7 @@ mod tests { let zigzag = ZigZagVTable .as_vtable() - .encode(&array.to_canonical(), None) + .encode(&array.to_canonical().unwrap(), None) .unwrap() .unwrap(); test_take_conformance(zigzag.as_ref()); diff --git a/encodings/zstd/src/array.rs b/encodings/zstd/src/array.rs index aba5f711c02..b8ecfee4779 100644 --- a/encodings/zstd/src/array.rs +++ b/encodings/zstd/src/array.rs @@ -197,12 +197,12 @@ fn choose_max_dict_size(uncompressed_size: usize) -> usize { } fn collect_valid_primitive(parray: &PrimitiveArray) -> VortexResult { - let mask = parray.validity_mask(); - Ok(filter(&parray.to_array(), &mask)?.to_primitive()) + let mask = parray.validity_mask()?; + 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()), _ => { @@ -457,7 +457,7 @@ impl ZstdArray { } pub fn from_array(array: ArrayRef, level: i32, values_per_frame: usize) -> VortexResult { - Self::from_canonical(&array.to_canonical(), level, values_per_frame)? + Self::from_canonical(&array.to_canonical()?, level, values_per_frame)? .ok_or_else(|| vortex_err!("Zstd can only encode Primitive and VarBinView arrays")) } @@ -469,14 +469,14 @@ impl ZstdArray { } } - pub fn decompress(&self) -> ArrayRef { + pub fn decompress(&self) -> VortexResult { // To start, we figure out which frames we need to decompress, and with // what row offset into the first such frame. let byte_width = self.byte_width(); let slice_n_rows = self.slice_stop - self.slice_start; let slice_value_indices = self .unsliced_validity - .to_mask(self.unsliced_n_rows) + .to_mask(self.unsliced_n_rows)? .valid_counts_for_indices(&[self.slice_start, self.slice_stop]); let slice_value_idx_start = slice_value_indices[0]; @@ -558,7 +558,7 @@ impl ZstdArray { // implied by the DType. if !self.dtype().is_nullable() && slice_validity != Validity::NonNullable { assert!( - slice_validity.all_valid(slice_n_rows), + slice_validity.all_valid(slice_n_rows)?, "ZSTD array expects to be non-nullable but there are nulls after decompression" ); @@ -570,7 +570,7 @@ impl ZstdArray { // END OF IMPORTANT BLOCK // - match &self.dtype { + Ok(match &self.dtype { DType::Primitive(..) => { let slice_values_buffer = decompressed.slice( (slice_value_idx_start - n_skipped_values) * byte_width @@ -586,7 +586,7 @@ impl ZstdArray { primitive.into_array() } DType::Binary(_) | DType::Utf8(_) => { - match slice_validity.to_mask(slice_n_rows).indices() { + match slice_validity.to_mask(slice_n_rows)?.indices() { AllOr::All => { // the decompressed buffer is a bunch of interleaved u32 lengths // and strings of those lengths, we need to reconstruct the @@ -639,7 +639,7 @@ impl ZstdArray { } } _ => vortex_panic!("Unsupported dtype for Zstd array: {}", self.dtype), - } + }) } pub(crate) fn _slice(&self, start: usize, stop: usize) -> ZstdArray { @@ -749,8 +749,8 @@ impl BaseArrayVTable for ZstdVTable { } impl CanonicalVTable for ZstdVTable { - fn canonicalize(array: &ZstdArray) -> Canonical { - array.decompress().to_canonical() + fn canonicalize(array: &ZstdArray) -> VortexResult { + array.decompress()?.to_canonical() } } @@ -760,7 +760,11 @@ impl OperationsVTable for ZstdVTable { } fn scalar_at(array: &ZstdArray, index: usize) -> Scalar { - array._slice(index, index + 1).decompress().scalar_at(0) + array + ._slice(index, index + 1) + .decompress() + .vortex_expect("decompress failed") + .scalar_at(0) } } diff --git a/encodings/zstd/src/compute/cast.rs b/encodings/zstd/src/compute/cast.rs index b51b6658062..a3223cdcc1d 100644 --- a/encodings/zstd/src/compute/cast.rs +++ b/encodings/zstd/src/compute/cast.rs @@ -46,7 +46,7 @@ impl CastKernel for ZstdVTable { let has_nulls = !array .unsliced_validity .slice(array.slice_start()..array.slice_stop()) - .all_valid(sliced_len); + .all_valid(sliced_len)?; // We don't attempt to handle casting when there are nulls. if has_nulls { @@ -106,7 +106,7 @@ mod tests { &DType::Primitive(PType::I64, Nullability::NonNullable) ); - let decoded = casted.to_primitive(); + let decoded = casted.to_primitive().unwrap(); assert_arrays_eq!(decoded, PrimitiveArray::from_iter([1i64, 2, 3, 4, 5])); } @@ -147,7 +147,7 @@ mod tests { &DType::Primitive(PType::U32, Nullability::NonNullable) ); // Verify the values are correct - let decoded = casted.to_primitive(); + let decoded = casted.to_primitive().unwrap(); let u32_values = decoded.as_slice::(); assert_eq!(u32_values, &[20, 30, 40, 50]); } @@ -173,7 +173,7 @@ mod tests { casted.dtype(), &DType::Primitive(PType::U32, Nullability::NonNullable) ); - let decoded = casted.to_primitive(); + let decoded = casted.to_primitive().unwrap(); let expected = PrimitiveArray::from_iter([20u32, 30, 40, 50]); assert_arrays_eq!(decoded, expected); } diff --git a/encodings/zstd/src/test.rs b/encodings/zstd/src/test.rs index dbdd5f21a51..9962290f24a 100644 --- a/encodings/zstd/src/test.rs +++ b/encodings/zstd/src/test.rs @@ -12,6 +12,7 @@ use vortex_array::vtable::ValidityHelper; use vortex_buffer::Buffer; use vortex_dtype::DType; use vortex_dtype::Nullability; +use vortex_error::VortexResult; use vortex_mask::Mask; use crate::ZstdArray; @@ -23,7 +24,7 @@ macro_rules! assert_nth_scalar { } #[test] -fn test_zstd_compress_decompress() { +fn test_zstd_compress_decompress() -> VortexResult<()> { let data: Vec = (0..200).collect(); let array = PrimitiveArray::from_iter(data.clone()); @@ -33,7 +34,7 @@ fn test_zstd_compress_decompress() { assert!(compressed.dictionary.is_none()); // check full decompression works - let decompressed = compressed.decompress().to_primitive(); + let decompressed = compressed.decompress()?.to_primitive()?; assert_arrays_eq!(decompressed, PrimitiveArray::from_iter(data)); // check slicing works @@ -41,19 +42,21 @@ fn test_zstd_compress_decompress() { for i in 0_i32..5 { assert_nth_scalar!(slice, i as usize, 100 + i); } - let primitive = slice.to_primitive(); + let primitive = slice.to_primitive()?; assert_arrays_eq!( primitive, PrimitiveArray::from_iter([100, 101, 102, 103, 104]) ); let slice = compressed.slice(200..200); - let primitive = slice.to_primitive(); + let primitive = slice.to_primitive()?; assert_arrays_eq!(primitive, PrimitiveArray::from_iter(Vec::::new())); + + Ok(()) } #[test] -fn test_zstd_empty() { +fn test_zstd_empty() -> VortexResult<()> { let data: Vec = vec![]; let array = PrimitiveArray::new( data.iter().cloned().collect::>(), @@ -62,12 +65,14 @@ fn test_zstd_empty() { let compressed = ZstdArray::from_primitive(&array, 3, 100).unwrap(); - let primitive = compressed.to_primitive(); + let primitive = compressed.to_primitive()?; assert_arrays_eq!(primitive, PrimitiveArray::from_iter(data)); + + Ok(()) } #[test] -fn test_zstd_with_validity_and_multi_frame() { +fn test_zstd_with_validity_and_multi_frame() -> VortexResult<()> { let data: Vec = (0..200).collect(); let mut validity: Vec = vec![false; 200]; validity[3] = true; @@ -84,7 +89,7 @@ fn test_zstd_with_validity_and_multi_frame() { assert_nth_scalar!(compressed, 10, None::); assert_nth_scalar!(compressed, 177, 177); - let decompressed = compressed.decompress().to_primitive(); + let decompressed = compressed.decompress()?.to_primitive()?; let decompressed_values = decompressed.as_slice::(); assert_eq!(decompressed_values[3], 3); assert_eq!(decompressed_values[177], 177); @@ -92,7 +97,7 @@ fn test_zstd_with_validity_and_multi_frame() { // check slicing works let slice = compressed.slice(176..179); - let primitive = slice.to_primitive(); + let primitive = slice.to_primitive()?; assert_eq!( TryInto::::try_into(primitive.scalar_at(1).as_ref()) .ok() @@ -103,10 +108,12 @@ fn test_zstd_with_validity_and_multi_frame() { primitive.validity(), &Validity::Array(BoolArray::from_iter(vec![false, true, false]).to_array()) ); + + Ok(()) } #[test] -fn test_zstd_with_dict() { +fn test_zstd_with_dict() -> VortexResult<()> { let data: Vec = (0..200).collect(); let array = PrimitiveArray::new( data.iter().cloned().collect::>(), @@ -118,29 +125,33 @@ fn test_zstd_with_dict() { assert_nth_scalar!(compressed, 0, 0); assert_nth_scalar!(compressed, 199, 199); - let decompressed = compressed.decompress().to_primitive(); + let decompressed = compressed.decompress()?.to_primitive()?; assert_arrays_eq!(decompressed, PrimitiveArray::from_iter(data)); assert_eq!(decompressed.validity(), array.validity()); // check slicing works let slice = compressed.slice(176..179); - let primitive = slice.to_primitive(); + let primitive = slice.to_primitive()?; assert_arrays_eq!(primitive, PrimitiveArray::from_iter([176, 177, 178])); + + Ok(()) } #[test] -fn test_validity_vtable() { +fn test_validity_vtable() -> VortexResult<()> { let mask_bools = vec![false, true, true, false, true]; let array = PrimitiveArray::new( (0..5).collect::>(), 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.validity_mask()?, Mask::from_iter(mask_bools)); assert_eq!( - compressed.slice(1..4).validity_mask(), + compressed.slice(1..4).validity_mask()?, Mask::from_iter(vec![true, true, false]) ); + + Ok(()) } #[test] @@ -169,7 +180,7 @@ fn test_zstd_var_bin_view() { } #[test] -fn test_zstd_decompress_var_bin_view() { +fn test_zstd_decompress_var_bin_view() -> VortexResult<()> { let data: [Option<&'static [u8]>; 5] = [ Some(b"foo"), Some(b"bar"), @@ -186,12 +197,14 @@ fn test_zstd_decompress_var_bin_view() { assert_nth_scalar!(compressed, 2, None::); assert_nth_scalar!(compressed, 3, "Lorem ipsum dolor sit amet"); assert_nth_scalar!(compressed, 4, "baz"); - let decompressed = compressed.decompress().to_varbinview(); + let decompressed = compressed.decompress()?.to_varbinview()?; assert_nth_scalar!(decompressed, 0, "foo"); assert_nth_scalar!(decompressed, 1, "bar"); assert_nth_scalar!(decompressed, 2, None::); assert_nth_scalar!(decompressed, 3, "Lorem ipsum dolor sit amet"); assert_nth_scalar!(decompressed, 4, "baz"); + + Ok(()) } #[test] diff --git a/fuzz/fuzz_targets/array_ops.rs b/fuzz/fuzz_targets/array_ops.rs index 401a4aab1b5..27adab9f334 100644 --- a/fuzz/fuzz_targets/array_ops.rs +++ b/fuzz/fuzz_targets/array_ops.rs @@ -44,10 +44,10 @@ fuzz_target!(|fuzz_action: FuzzArrayAction| -> Corpus { Action::Compress(strategy) => { current_array = match strategy { CompressorStrategy::Default => BtrBlocksCompressor::default() - .compress(current_array.to_canonical().as_ref()) + .compress(current_array.to_canonical().vortex_unwrap().as_ref()) .vortex_unwrap(), CompressorStrategy::Compact => CompactCompressor::default() - .compress(current_array.to_canonical().as_ref()) + .compress(current_array.to_canonical().vortex_unwrap().as_ref()) .vortex_unwrap(), }; assert_array_eq(&expected.array(), ¤t_array, i).unwrap(); diff --git a/fuzz/fuzz_targets/file_io.rs b/fuzz/fuzz_targets/file_io.rs index df4cc45a2c4..e6d767110e6 100644 --- a/fuzz/fuzz_targets/file_io.rs +++ b/fuzz/fuzz_targets/file_io.rs @@ -53,7 +53,10 @@ fuzz_target!(|fuzz: FuzzFileAction| -> Corpus { .unwrap_or_else(|| lit(true)) .evaluate(&array_data) .vortex_unwrap(); - let mask = bool_mask.to_bool().to_mask_fill_null_false(); + let mask = bool_mask + .to_bool() + .vortex_unwrap() + .to_mask_fill_null_false(); let filtered = filter(&array_data, &mask).vortex_unwrap(); projection_expr .clone() @@ -114,9 +117,11 @@ fuzz_target!(|fuzz: FuzzFileAction| -> Corpus { let bool_result = compare(&expected_array, &output_array, Operator::Eq) .vortex_unwrap() - .to_bool(); + .to_bool() + .vortex_unwrap(); 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_unwrap() || expected_array.all_valid().vortex_unwrap()) { vortex_panic!( "Failed to match original array {}with{}", diff --git a/fuzz/src/array/cast.rs b/fuzz/src/array/cast.rs index 906a2e33322..6a12deb0e7e 100644 --- a/fuzz/src/array/cast.rs +++ b/fuzz/src/array/cast.rs @@ -35,12 +35,12 @@ pub fn cast_canonical_array(array: &ArrayRef, target: &DType) -> VortexResult() .iter() .map(|v| *v as Out) .collect::>(), - Validity::from_mask(array.validity_mask(), target.nullability()), + Validity::from_mask(array.validity_mask()?, target.nullability()), ) .to_array() }) @@ -61,12 +61,12 @@ pub fn cast_canonical_array(array: &ArrayRef, target: &DType) -> VortexResult Ok(Some( PrimitiveArray::new( array - .to_primitive() + .to_primitive()? .as_slice::() .iter() .map(|v| *v as f64) .collect::>(), - Validity::from_mask(array.validity_mask(), target.nullability()), + Validity::from_mask(array.validity_mask()?, target.nullability()), ) .to_array(), )), @@ -76,12 +76,12 @@ pub fn cast_canonical_array(array: &ArrayRef, target: &DType) -> VortexResult() .iter() .map(|v| *v as f32) .collect::>(), - 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..7e538fd5c82 100644 --- a/fuzz/src/array/compare.rs +++ b/fuzz/src/array/compare.rs @@ -17,18 +17,26 @@ use vortex_dtype::Nullability; use vortex_dtype::match_each_decimal_value_type; use vortex_dtype::match_each_native_ptype; use vortex_error::VortexExpect; +use vortex_error::VortexResult; use vortex_error::vortex_panic; use vortex_scalar::Scalar; -pub fn compare_canonical_array(array: &dyn Array, value: &Scalar, operator: Operator) -> ArrayRef { +pub fn compare_canonical_array( + array: &dyn Array, + value: &Scalar, + operator: Operator, +) -> VortexResult { if value.is_null() { - return BoolArray::from_bit_buffer(BitBuffer::new_unset(array.len()), Validity::AllInvalid) - .into_array(); + return Ok(BoolArray::from_bit_buffer( + BitBuffer::new_unset(array.len()), + Validity::AllInvalid, + ) + .into_array()); } let result_nullability = array.dtype().nullability() | value.dtype().nullability(); - match array.dtype() { + Ok(match array.dtype() { DType::Bool(_) => { let bool = value .as_bool() @@ -36,10 +44,10 @@ pub fn compare_canonical_array(array: &dyn Array, value: &Scalar, operator: Oper .vortex_expect("nulls handled before"); compare_to( array - .to_bool() + .to_bool()? .bit_buffer() .iter() - .zip(array.validity_mask().to_bit_buffer().iter()) + .zip(array.validity_mask()?.to_bit_buffer().iter()) .map(|(b, v)| v.then_some(b)), bool, operator, @@ -48,7 +56,7 @@ pub fn compare_canonical_array(array: &dyn Array, value: &Scalar, operator: Oper } DType::Primitive(p, _) => { let primitive = value.as_primitive(); - let primitive_array = array.to_primitive(); + let primitive_array = array.to_primitive()?; match_each_native_ptype!(p, |P| { let pval = primitive .typed_value::

() @@ -58,7 +66,7 @@ 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()?.to_bit_buffer().iter()) .map(|(b, v)| v.then_some(NativeValue(b))), NativeValue(pval), operator, @@ -68,7 +76,7 @@ pub fn compare_canonical_array(array: &dyn Array, value: &Scalar, operator: Oper } DType::Decimal(..) => { let decimal = value.as_decimal(); - let decimal_array = array.to_decimal(); + let decimal_array = array.to_decimal()?; match_each_decimal_value_type!(decimal_array.values_type(), |D| { let dval = decimal .decimal_value() @@ -80,7 +88,7 @@ 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()?.to_bit_buffer().iter()) .map(|(b, v)| v.then_some(b)), dval, operator, @@ -88,7 +96,7 @@ pub fn compare_canonical_array(array: &dyn Array, value: &Scalar, operator: Oper ) }) } - DType::Utf8(_) => array.to_varbinview().with_iterator(|iter| { + DType::Utf8(_) => array.to_varbinview()?.with_iterator(|iter| { let utf8_value = value .as_utf8() .value() @@ -100,7 +108,7 @@ pub fn compare_canonical_array(array: &dyn Array, value: &Scalar, operator: Oper result_nullability, ) }), - DType::Binary(_) => array.to_varbinview().with_iterator(|iter| { + DType::Binary(_) => array.to_varbinview()?.with_iterator(|iter| { let binary_value = value .as_binary() .value() @@ -126,7 +134,7 @@ pub fn compare_canonical_array(array: &dyn Array, value: &Scalar, operator: Oper d @ (DType::Null | DType::Extension(_)) => { unreachable!("DType {d} not supported for fuzzing") } - } + }) } fn compare_to( diff --git a/fuzz/src/array/fill_null.rs b/fuzz/src/array/fill_null.rs index 0235174190b..9153e256bf4 100644 --- a/fuzz/src/array/fill_null.rs +++ b/fuzz/src/array/fill_null.rs @@ -65,7 +65,7 @@ fn fill_bool_array( } Validity::AllInvalid => ConstantArray::new(fill_value.clone(), array.len()).into_array(), Validity::Array(validity_array) => { - let validity_bool_array = validity_array.to_bool(); + let validity_bool_array = validity_array.to_bool().vortex_unwrap(); let validity_bits = validity_bool_array.bit_buffer(); let data_bits = array.bit_buffer(); @@ -100,7 +100,7 @@ fn fill_primitive_array( ConstantArray::new(fill_value.clone(), array.len()).into_array() } Validity::Array(validity_array) => { - let validity_bool_array = validity_array.to_bool(); + let validity_bool_array = validity_array.to_bool().vortex_unwrap(); let validity_bits = validity_bool_array.bit_buffer(); let data_slice = array.as_slice::(); @@ -143,7 +143,7 @@ fn fill_decimal_array( ConstantArray::new(fill_value.clone(), array.len()).into_array() } Validity::Array(validity_array) => { - let validity_bool_array = validity_array.to_bool(); + let validity_bool_array = validity_array.to_bool().vortex_unwrap(); let validity_bits = validity_bool_array.bit_buffer(); let data_buffer = array.buffer::(); @@ -173,7 +173,7 @@ fn fill_varbinview_array( Validity::NonNullable | Validity::AllValid => array.clone().into_array(), Validity::AllInvalid => ConstantArray::new(fill_value.clone(), array.len()).into_array(), Validity::Array(validity_array) => { - let validity_bool_array = validity_array.to_bool(); + let validity_bool_array = validity_array.to_bool().vortex_unwrap(); let validity_bits = validity_bool_array.bit_buffer(); match array.dtype() { @@ -200,8 +200,8 @@ fn fill_varbinview_array( let result = VarBinViewArray::from_iter_str(string_refs).into_array(); if result_nullability == Nullability::Nullable { VarBinViewArray::new( - result.to_varbinview().views().clone(), - result.to_varbinview().buffers().clone(), + result.to_varbinview().vortex_unwrap().views().clone(), + result.to_varbinview().vortex_unwrap().buffers().clone(), result.dtype().as_nullable(), result_nullability.into(), ) @@ -233,8 +233,8 @@ fn fill_varbinview_array( let result = VarBinViewArray::from_iter_bin(binary_refs).into_array(); if result_nullability == Nullability::Nullable { VarBinViewArray::new( - result.to_varbinview().views().clone(), - result.to_varbinview().buffers().clone(), + result.to_varbinview().vortex_unwrap().views().clone(), + result.to_varbinview().vortex_unwrap().buffers().clone(), result.dtype().as_nullable(), result_nullability.into(), ) @@ -275,7 +275,7 @@ mod tests { let array = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, Some(5)]); let fill_value = Scalar::from(42i32); - let result = fill_null_canonical_array(array.to_canonical(), &fill_value).unwrap(); + let result = fill_null_canonical_array(array.to_canonical().unwrap(), &fill_value).unwrap(); let expected = PrimitiveArray::from_iter([1i32, 42, 3, 42, 5]); assert_arrays_eq!(expected, result); @@ -288,7 +288,7 @@ mod tests { let array = BoolArray::from_bit_buffer(data_buffer, Validity::from(validity_buffer)); let fill_value = Scalar::from(true); - let result = fill_null_canonical_array(array.to_canonical(), &fill_value).unwrap(); + let result = fill_null_canonical_array(array.to_canonical().unwrap(), &fill_value).unwrap(); let expected = BoolArray::from(BitBuffer::from(vec![true, true, false, true])); assert_arrays_eq!(expected, result); @@ -302,7 +302,7 @@ mod tests { ); let fill_value = Scalar::utf8("default", Nullability::NonNullable); - let result = fill_null_canonical_array(array.to_canonical(), &fill_value).unwrap(); + let result = fill_null_canonical_array(array.to_canonical().unwrap(), &fill_value).unwrap(); let expected = VarBinViewArray::from_iter_str(["hello", "default", "world"]); assert_arrays_eq!(expected, result); @@ -313,7 +313,7 @@ mod tests { let array = PrimitiveArray::from_option_iter([None::, None, None]); let fill_value = Scalar::from(100i32); - let result = fill_null_canonical_array(array.to_canonical(), &fill_value).unwrap(); + let result = fill_null_canonical_array(array.to_canonical().unwrap(), &fill_value).unwrap(); let expected = PrimitiveArray::from_iter([100i32, 100, 100]); assert_arrays_eq!(expected, result); @@ -324,7 +324,7 @@ mod tests { let array = PrimitiveArray::from_iter([1i32, 2, 3]); let fill_value = Scalar::from(42i32); - let result = fill_null_canonical_array(array.to_canonical(), &fill_value).unwrap(); + let result = fill_null_canonical_array(array.to_canonical().unwrap(), &fill_value).unwrap(); let expected = PrimitiveArray::from_iter([1i32, 2, 3]); assert_arrays_eq!(expected, result); @@ -336,7 +336,7 @@ mod tests { let array = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]); let fill_value = Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)); - let result = fill_null_canonical_array(array.to_canonical(), &fill_value); + let result = fill_null_canonical_array(array.to_canonical().unwrap(), &fill_value); assert!(result.is_err()); assert!( @@ -359,7 +359,7 @@ mod tests { Nullability::NonNullable, ); - let result = fill_null_canonical_array(array.to_canonical(), &fill_value).unwrap(); + let result = fill_null_canonical_array(array.to_canonical().unwrap(), &fill_value).unwrap(); let expected = DecimalArray::from_iter( [100i32, 999i32, 300i32, 999i32, 500i32], @@ -380,7 +380,7 @@ mod tests { Nullability::NonNullable, ); - let result = fill_null_canonical_array(array.to_canonical(), &fill_value).unwrap(); + let result = fill_null_canonical_array(array.to_canonical().unwrap(), &fill_value).unwrap(); let expected = DecimalArray::from_iter([1000i64, 9999i64, 3000i64], DecimalDType::new(15, 3)); @@ -399,7 +399,7 @@ mod tests { Nullability::NonNullable, ); - let result = fill_null_canonical_array(array.to_canonical(), &fill_value).unwrap(); + let result = fill_null_canonical_array(array.to_canonical().unwrap(), &fill_value).unwrap(); let expected = DecimalArray::from_iter( [10000i128, 99999i128, 30000i128, 99999i128], @@ -418,7 +418,7 @@ mod tests { Nullability::NonNullable, ); - let result = fill_null_canonical_array(array.to_canonical(), &fill_value).unwrap(); + let result = fill_null_canonical_array(array.to_canonical().unwrap(), &fill_value).unwrap(); let expected = cast( &DecimalArray::from_option_iter( @@ -444,7 +444,7 @@ mod tests { Nullability::NonNullable, ); - let result = fill_null_canonical_array(array.to_canonical(), &fill_value).unwrap(); + let result = fill_null_canonical_array(array.to_canonical().unwrap(), &fill_value).unwrap(); let expected = cast( &DecimalArray::from_option_iter( diff --git a/fuzz/src/array/filter.rs b/fuzz/src/array/filter.rs index 4d6fe1c0e03..5b495558078 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() @@ -37,7 +37,7 @@ pub fn filter_canonical_array(array: &dyn Array, filter: &[bool]) -> VortexResul match array.dtype() { DType::Bool(_) => { - let bool_array = array.to_bool(); + let bool_array = array.to_bool()?; Ok(BoolArray::from_bit_buffer( BitBuffer::from_iter( filter @@ -51,7 +51,7 @@ pub fn filter_canonical_array(array: &dyn Array, filter: &[bool]) -> VortexResul .into_array()) } DType::Primitive(p, _) => match_each_native_ptype!(p, |P| { - let primitive_array = array.to_primitive(); + let primitive_array = array.to_primitive()?; Ok(PrimitiveArray::new( filter .iter() @@ -64,7 +64,7 @@ pub fn filter_canonical_array(array: &dyn Array, filter: &[bool]) -> VortexResul .into_array()) }), DType::Decimal(d, _) => { - let decimal_array = array.to_decimal(); + let decimal_array = array.to_decimal()?; match_each_decimal_value_type!(decimal_array.values_type(), |D| { let buf = decimal_array.buffer::(); Ok(DecimalArray::new( @@ -81,7 +81,7 @@ pub fn filter_canonical_array(array: &dyn Array, filter: &[bool]) -> VortexResul }) } DType::Utf8(_) | DType::Binary(_) => { - let utf8 = array.to_varbinview(); + let utf8 = array.to_varbinview()?; let values = utf8.with_iterator(|iter| { iter.zip(filter.iter()) .filter(|(_, f)| **f) @@ -91,7 +91,7 @@ pub fn filter_canonical_array(array: &dyn Array, filter: &[bool]) -> VortexResul Ok(VarBinViewArray::from_iter(values, array.dtype().clone()).into_array()) } DType::Struct(..) => { - let struct_array = array.to_struct(); + let struct_array = array.to_struct()?; let filtered_children = struct_array .fields() .iter() diff --git a/fuzz/src/array/mask.rs b/fuzz/src/array/mask.rs index d604f299a66..fbd2850a7fb 100644 --- a/fuzz/src/array/mask.rs +++ b/fuzz/src/array/mask.rs @@ -31,11 +31,11 @@ pub fn mask_canonical_array(canonical: Canonical, mask: &Mask) -> VortexResult { - let new_validity = array.validity().mask(mask); + let new_validity = array.validity().mask(mask).vortex_unwrap(); BoolArray::from_bit_buffer(array.bit_buffer().clone(), new_validity).into_array() } Canonical::Primitive(array) => { - let new_validity = array.validity().mask(mask); + let new_validity = array.validity().mask(mask).vortex_unwrap(); PrimitiveArray::from_byte_buffer( array.byte_buffer().clone(), array.ptype(), @@ -44,14 +44,14 @@ pub fn mask_canonical_array(canonical: Canonical, mask: &Mask) -> VortexResult { - let new_validity = array.validity().mask(mask); + let new_validity = array.validity().mask(mask).vortex_unwrap(); match_each_decimal_value_type!(array.values_type(), |D| { DecimalArray::new(array.buffer::(), array.decimal_dtype(), new_validity) .into_array() }) } Canonical::VarBinView(array) => { - let new_validity = array.validity().mask(mask); + let new_validity = array.validity().mask(mask).vortex_unwrap(); VarBinViewArray::new( array.views().clone(), array.buffers().clone(), @@ -61,7 +61,7 @@ pub fn mask_canonical_array(canonical: Canonical, mask: &Mask) -> VortexResult { - let new_validity = array.validity().mask(mask); + let new_validity = array.validity().mask(mask).vortex_unwrap(); // SAFETY: Since we are only masking the validity and everything else comes from an // already valid `ListViewArray`, all of the invariants are still upheld. @@ -77,7 +77,7 @@ pub fn mask_canonical_array(canonical: Canonical, mask: &Mask) -> VortexResult { - let new_validity = array.validity().mask(mask); + let new_validity = array.validity().mask(mask).vortex_unwrap(); FixedSizeListArray::new( array.elements().clone(), array.list_size(), @@ -87,7 +87,7 @@ pub fn mask_canonical_array(canonical: Canonical, mask: &Mask) -> VortexResult { - let new_validity = array.validity().mask(mask); + let new_validity = array.validity().mask(mask).vortex_unwrap(); StructArray::try_new_with_dtype( array.fields().clone(), array.struct_fields().clone(), @@ -100,7 +100,7 @@ pub fn mask_canonical_array(canonical: Canonical, mask: &Mask) -> VortexResult { // Recursively mask the storage array let masked_storage = - mask_canonical_array(array.storage().to_canonical(), mask).vortex_unwrap(); + mask_canonical_array(array.storage().to_canonical()?, mask).vortex_unwrap(); if masked_storage.dtype().nullability() == array.ext_dtype().storage_dtype().nullability() @@ -144,12 +144,12 @@ mod tests { let array = NullArray::new(5); let mask = Mask::from_iter([true, false, true, false, true]); - let result = mask_canonical_array(array.to_canonical(), &mask).unwrap(); + let result = mask_canonical_array(array.to_canonical().unwrap(), &mask).unwrap(); 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()); } } @@ -158,13 +158,13 @@ mod tests { let array = BoolArray::from_iter([true, false, true, false, true]); let mask = Mask::from_iter([true, false, false, true, false]); - let result = mask_canonical_array(array.to_canonical(), &mask).unwrap(); + let result = mask_canonical_array(array.to_canonical().unwrap(), &mask).unwrap(); assert_eq!(result.len(), 5); - assert!(!result.is_valid(0)); + assert!(!result.is_valid(0).unwrap()); assert_eq!(result.scalar_at(1), Scalar::from(Some(false))); assert_eq!(result.scalar_at(2), Scalar::from(Some(true))); - assert!(!result.is_valid(3)); + assert!(!result.is_valid(3).unwrap()); assert_eq!(result.scalar_at(4), Scalar::from(Some(true))); } @@ -173,13 +173,13 @@ mod tests { let array = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]); let mask = Mask::from_iter([false, true, false, true, false]); - let result = mask_canonical_array(array.to_canonical(), &mask).unwrap(); + let result = mask_canonical_array(array.to_canonical().unwrap(), &mask).unwrap(); assert_eq!(result.len(), 5); assert_eq!(result.scalar_at(0), Scalar::from(Some(1))); - assert!(!result.is_valid(1)); + assert!(!result.is_valid(1).unwrap()); assert_eq!(result.scalar_at(2), Scalar::from(Some(3))); - assert!(!result.is_valid(3)); + assert!(!result.is_valid(3).unwrap()); assert_eq!(result.scalar_at(4), Scalar::from(Some(5))); } @@ -188,14 +188,14 @@ mod tests { let array = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4), None]); let mask = Mask::from_iter([true, false, false, true, false]); - let result = mask_canonical_array(array.to_canonical(), &mask).unwrap(); + let result = mask_canonical_array(array.to_canonical().unwrap(), &mask).unwrap(); assert_eq!(result.len(), 5); - assert!(!result.is_valid(0)); - assert!(!result.is_valid(1)); // was already null + assert!(!result.is_valid(0).unwrap()); + assert!(!result.is_valid(1).unwrap()); // was already null assert_eq!(result.scalar_at(2), Scalar::from(Some(3))); - assert!(!result.is_valid(3)); - assert!(!result.is_valid(4)); // was already null + assert!(!result.is_valid(3).unwrap()); + assert!(!result.is_valid(4).unwrap()); // was already null } #[test] @@ -206,14 +206,14 @@ mod tests { ); let mask = Mask::from_iter([false, false, true, false, false]); - let result = mask_canonical_array(array.to_canonical(), &mask).unwrap(); + let result = mask_canonical_array(array.to_canonical().unwrap(), &mask).unwrap(); assert_eq!(result.len(), 5); - 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] @@ -221,20 +221,20 @@ mod tests { let array = VarBinViewArray::from_iter_str(["one", "two", "three", "four", "five"]); let mask = Mask::from_iter([true, false, true, false, true]); - let result = mask_canonical_array(array.to_canonical(), &mask).unwrap(); + let result = mask_canonical_array(array.to_canonical().unwrap(), &mask).unwrap(); assert_eq!(result.len(), 5); - assert!(!result.is_valid(0)); + assert!(!result.is_valid(0).unwrap()); assert_eq!( result.scalar_at(1), Scalar::utf8("two", Nullability::Nullable) ); - assert!(!result.is_valid(2)); + assert!(!result.is_valid(2).unwrap()); assert_eq!( result.scalar_at(3), Scalar::utf8("four", Nullability::Nullable) ); - assert!(!result.is_valid(4)); + assert!(!result.is_valid(4).unwrap()); } #[test] @@ -249,12 +249,12 @@ mod tests { let mask = Mask::from_iter([false, true, false]); - let result = mask_canonical_array(array.to_canonical(), &mask).unwrap(); + 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] @@ -265,12 +265,12 @@ mod tests { let mask = Mask::from_iter([true, false, true]); - let result = mask_canonical_array(array.to_canonical(), &mask).unwrap(); + 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] @@ -289,12 +289,12 @@ mod tests { let mask = Mask::from_iter([false, true, false]); - let result = mask_canonical_array(array.to_canonical(), &mask).unwrap(); + 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] @@ -302,12 +302,12 @@ mod tests { let array = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]); let mask = Mask::AllTrue(5); - let result = mask_canonical_array(array.to_canonical(), &mask).unwrap(); + let result = mask_canonical_array(array.to_canonical().unwrap(), &mask).unwrap(); assert_eq!(result.len(), 5); // All values should be masked out (null) for i in 0..5 { - assert!(!result.is_valid(i)); + assert!(!result.is_valid(i).unwrap()); } } @@ -316,12 +316,12 @@ mod tests { let array = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]); let mask = Mask::AllFalse(5); - let result = mask_canonical_array(array.to_canonical(), &mask).unwrap(); + let result = mask_canonical_array(array.to_canonical().unwrap(), &mask).unwrap(); assert_eq!(result.len(), 5); // No values should be masked out for i in 0..5 { - assert!(result.is_valid(i)); + assert!(result.is_valid(i).unwrap()); #[allow(clippy::cast_possible_truncation)] let expected = (i + 1) as i32; assert_eq!(result.scalar_at(i), Scalar::from(Some(expected))); @@ -333,7 +333,7 @@ mod tests { let array = PrimitiveArray::from_iter(Vec::::new()); let mask = Mask::AllFalse(0); - let result = mask_canonical_array(array.to_canonical(), &mask).unwrap(); + let result = mask_canonical_array(array.to_canonical().unwrap(), &mask).unwrap(); assert_eq!(result.len(), 0); } diff --git a/fuzz/src/array/mod.rs b/fuzz/src/array/mod.rs index 2d9e0aab8fe..96b3af61573 100644 --- a/fuzz/src/array/mod.rs +++ b/fuzz/src/array/mod.rs @@ -266,7 +266,8 @@ impl<'a> Arbitrary<'a> for FuzzArrayAction { }; let op = u.arbitrary()?; - current_array = compare_canonical_array(¤t_array, &scalar, op); + current_array = + compare_canonical_array(¤t_array, &scalar, op).vortex_unwrap(); ( Action::Compare(scalar, op), ExpectedValue::Array(current_array.to_array()), @@ -294,13 +295,15 @@ impl<'a> Arbitrary<'a> for FuzzArrayAction { // Sum - returns a scalar, does NOT update current_array (terminal operation) let sum_result = - sum_canonical_array(current_array.to_canonical()).vortex_unwrap(); + sum_canonical_array(current_array.to_canonical().vortex_unwrap()) + .vortex_unwrap(); (Action::Sum, ExpectedValue::Scalar(sum_result)) } ActionType::MinMax => { // MinMax - returns a scalar, does NOT update current_array (terminal operation) let min_max_result = - min_max_canonical_array(current_array.to_canonical()).vortex_unwrap(); + min_max_canonical_array(current_array.to_canonical().vortex_unwrap()) + .vortex_unwrap(); (Action::MinMax, ExpectedValue::MinMax(min_max_result)) } ActionType::FillNull => { @@ -324,9 +327,11 @@ impl<'a> Arbitrary<'a> for FuzzArrayAction { } // Compute expected result on canonical form - let expected_result = - fill_null_canonical_array(current_array.to_canonical(), &fill_value) - .vortex_unwrap(); + let expected_result = fill_null_canonical_array( + current_array.to_canonical().vortex_unwrap(), + &fill_value, + ) + .vortex_unwrap(); // Update current_array to the result for chaining current_array = expected_result.clone(); ( @@ -342,7 +347,7 @@ impl<'a> Arbitrary<'a> for FuzzArrayAction { // Compute expected result on canonical form let expected_result = mask_canonical_array( - current_array.to_canonical(), + current_array.to_canonical().vortex_unwrap(), &Mask::from_iter(mask.iter().copied()), ) .vortex_unwrap(); @@ -372,8 +377,11 @@ impl<'a> Arbitrary<'a> for FuzzArrayAction { let expected_scalars: Vec = indices_vec .iter() .map(|&idx| { - scalar_at_canonical_array(current_array.to_canonical(), idx) - .vortex_unwrap() + scalar_at_canonical_array( + current_array.to_canonical().vortex_unwrap(), + idx, + ) + .vortex_unwrap() }) .collect(); diff --git a/fuzz/src/array/scalar_at.rs b/fuzz/src/array/scalar_at.rs index e2bcce4b35e..66fa49a6e41 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 { @@ -44,7 +44,10 @@ pub fn scalar_at_canonical_array(canonical: Canonical, index: usize) -> VortexRe Canonical::List(array) => { let list = array.list_elements_at(index); let children: Vec = (0..list.len()) - .map(|i| scalar_at_canonical_array(list.to_canonical(), i).vortex_unwrap()) + .map(|i| { + scalar_at_canonical_array(list.to_canonical().vortex_unwrap(), i) + .vortex_unwrap() + }) .collect(); Scalar::list( Arc::new(list.dtype().clone()), @@ -55,7 +58,10 @@ pub fn scalar_at_canonical_array(canonical: Canonical, index: usize) -> VortexRe Canonical::FixedSizeList(array) => { let list = array.fixed_size_list_elements_at(index); let children: Vec = (0..list.len()) - .map(|i| scalar_at_canonical_array(list.to_canonical(), i).vortex_unwrap()) + .map(|i| { + scalar_at_canonical_array(list.to_canonical().vortex_unwrap(), i) + .vortex_unwrap() + }) .collect(); Scalar::fixed_size_list(list.dtype().clone(), children, array.dtype().nullability()) } @@ -63,12 +69,15 @@ pub fn scalar_at_canonical_array(canonical: Canonical, index: usize) -> VortexRe let field_scalars: Vec = array .fields() .iter() - .map(|field| scalar_at_canonical_array(field.to_canonical(), index).vortex_unwrap()) + .map(|field| { + scalar_at_canonical_array(field.to_canonical().vortex_unwrap(), index) + .vortex_unwrap() + }) .collect(); Scalar::struct_(array.dtype().clone(), field_scalars) } Canonical::Extension(array) => { - let storage_scalar = scalar_at_canonical_array(array.storage().to_canonical(), index)?; + let storage_scalar = scalar_at_canonical_array(array.storage().to_canonical()?, index)?; Scalar::extension(array.ext_dtype().clone(), storage_scalar) } }) diff --git a/fuzz/src/array/search_sorted.rs b/fuzz/src/array/search_sorted.rs index 0f0e7349e99..0b246fbca44 100644 --- a/fuzz/src/array/search_sorted.rs +++ b/fuzz/src/array/search_sorted.rs @@ -63,8 +63,8 @@ pub fn search_sorted_canonical_array( ) -> VortexResult { match array.dtype() { DType::Bool(_) => { - let bool_array = array.to_bool(); - let validity = bool_array.validity_mask().to_bit_buffer(); + let bool_array = array.to_bool()?; + let validity = bool_array.validity_mask()?.to_bit_buffer(); let opt_values = bool_array .bit_buffer() .iter() @@ -75,8 +75,8 @@ pub fn search_sorted_canonical_array( Ok(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 primitive_array = array.to_primitive()?; + let validity = primitive_array.validity_mask()?.to_bit_buffer(); match_each_native_ptype!(p, |P| { let opt_values = primitive_array .as_slice::

() @@ -90,8 +90,8 @@ pub fn search_sorted_canonical_array( }) } DType::Decimal(d, _) => { - let decimal_array = array.to_decimal(); - let validity = decimal_array.validity_mask().to_bit_buffer(); + let decimal_array = array.to_decimal()?; + 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 @@ -115,7 +115,7 @@ pub fn search_sorted_canonical_array( }) } DType::Utf8(_) | DType::Binary(_) => { - let utf8 = array.to_varbinview(); + let utf8 = array.to_varbinview()?; let opt_values = utf8.with_iterator(|iter| iter.map(|v| v.map(|u| u.to_vec())).collect::>()); let to_find = if matches!(array.dtype(), DType::Utf8(_)) { diff --git a/fuzz/src/array/slice.rs b/fuzz/src/array/slice.rs index eaa22a5dcf1..df74c0706c4 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 @@ -34,12 +34,12 @@ pub fn slice_canonical_array( match array.dtype() { DType::Bool(_) => { - let bool_array = array.to_bool(); + let bool_array = array.to_bool()?; let sliced_bools = bool_array.bit_buffer().slice(start..stop); Ok(BoolArray::from_bit_buffer(sliced_bools, validity).into_array()) } DType::Primitive(p, _) => { - let primitive_array = array.to_primitive(); + let primitive_array = array.to_primitive()?; match_each_native_ptype!(p, |P| { Ok( PrimitiveArray::new(primitive_array.buffer::

().slice(start..stop), validity) @@ -48,7 +48,7 @@ pub fn slice_canonical_array( }) } DType::Utf8(_) | DType::Binary(_) => { - let utf8 = array.to_varbinview(); + let utf8 = array.to_varbinview()?; let values = utf8.with_iterator(|iter| iter.map(|v| v.map(|u| u.to_vec())).collect::>()); Ok(VarBinViewArray::from_iter( @@ -58,7 +58,7 @@ pub fn slice_canonical_array( .into_array()) } DType::Struct(..) => { - let struct_array = array.to_struct(); + let struct_array = array.to_struct()?; let sliced_children = struct_array .fields() .iter() @@ -73,7 +73,7 @@ pub fn slice_canonical_array( .map(|a| a.into_array()) } DType::List(..) => { - let list_array = array.to_listview(); + let list_array = array.to_listview()?; let offsets = slice_canonical_array(list_array.offsets(), start, stop)?; let sizes = slice_canonical_array(list_array.sizes(), start, stop)?; @@ -91,7 +91,7 @@ pub fn slice_canonical_array( .into_array()) } DType::FixedSizeList(..) => { - let fsl_array = array.to_fixed_size_list(); + let fsl_array = array.to_fixed_size_list()?; let list_size = fsl_array.list_size() as usize; let elements = slice_canonical_array(fsl_array.elements(), start * list_size, stop * list_size)?; @@ -101,7 +101,7 @@ pub fn slice_canonical_array( .map(|a| a.into_array()) } DType::Decimal(decimal_dtype, _) => { - let decimal_array = array.to_decimal(); + let decimal_array = array.to_decimal()?; Ok( match_each_decimal_value_type!(decimal_array.values_type(), |D| { DecimalArray::new( diff --git a/fuzz/src/array/sort.rs b/fuzz/src/array/sort.rs index b1ffaf2f28e..c4ac38d6316 100644 --- a/fuzz/src/array/sort.rs +++ b/fuzz/src/array/sort.rs @@ -24,24 +24,24 @@ use crate::array::take_canonical_array_non_nullable_indices; pub fn sort_canonical_array(array: &dyn Array) -> VortexResult { match array.dtype() { DType::Bool(_) => { - let bool_array = array.to_bool(); + let bool_array = array.to_bool()?; 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(); Ok(BoolArray::from_iter(opt_values).into_array()) } DType::Primitive(p, _) => { - let primitive_array = array.to_primitive(); + let primitive_array = array.to_primitive()?; match_each_native_ptype!(p, |P| { let mut opt_values = primitive_array .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); @@ -49,14 +49,14 @@ pub fn sort_canonical_array(array: &dyn Array) -> VortexResult { }) } DType::Decimal(d, _) => { - let decimal_array = array.to_decimal(); + let decimal_array = array.to_decimal()?; match_each_decimal_value_type!(decimal_array.values_type(), |D| { let buf = decimal_array.buffer::(); let mut opt_values = buf .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(); @@ -64,7 +64,7 @@ pub fn sort_canonical_array(array: &dyn Array) -> VortexResult { }) } DType::Utf8(_) | DType::Binary(_) => { - let utf8 = array.to_varbinview(); + let utf8 = array.to_varbinview()?; let mut opt_values = utf8.with_iterator(|iter| iter.map(|v| v.map(|u| u.to_vec())).collect::>()); opt_values.sort(); diff --git a/fuzz/src/array/take.rs b/fuzz/src/array/take.rs index 6d4a3b7660b..2763a98a7fc 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 @@ -61,7 +61,7 @@ pub fn take_canonical_array( match array.dtype() { DType::Bool(_) => { - let bool_array = array.to_bool(); + let bool_array = array.to_bool()?; let vec_values = bool_array.bit_buffer().iter().collect::>(); Ok(BoolArray::from_bit_buffer( indices_slice_non_opt @@ -73,7 +73,7 @@ pub fn take_canonical_array( .into_array()) } DType::Primitive(p, _) => { - let primitive_array = array.to_primitive(); + let primitive_array = array.to_primitive()?; match_each_native_ptype!(p, |P| { Ok(take_primitive::

( primitive_array, @@ -83,7 +83,7 @@ pub fn take_canonical_array( }) } DType::Decimal(d, _) => { - let decimal_array = array.to_decimal(); + let decimal_array = array.to_decimal()?; match_each_decimal_value_type!(decimal_array.values_type(), |D| { Ok(take_decimal::( @@ -95,7 +95,7 @@ pub fn take_canonical_array( }) } DType::Utf8(_) | DType::Binary(_) => { - let utf8 = array.to_varbinview(); + let utf8 = array.to_varbinview()?; let values = utf8.with_iterator(|iter| iter.map(|v| v.map(|u| u.to_vec())).collect::>()); Ok(VarBinViewArray::from_iter( @@ -107,7 +107,7 @@ pub fn take_canonical_array( .into_array()) } DType::Struct(..) => { - let struct_array = array.to_struct(); + let struct_array = array.to_struct()?; let taken_children = struct_array .fields() .iter() diff --git a/vortex-array/benches/chunk_array_builder.rs b/vortex-array/benches/chunk_array_builder.rs index 98bf666e4d1..870fe89e5aa 100644 --- a/vortex-array/benches/chunk_array_builder.rs +++ b/vortex-array/benches/chunk_array_builder.rs @@ -14,6 +14,7 @@ use vortex_array::builders::ArrayBuilder; use vortex_array::builders::VarBinViewBuilder; use vortex_array::builders::builder_with_capacity; use vortex_dtype::DType; +use vortex_error::VortexExpect; fn main() { divan::main(); @@ -32,7 +33,9 @@ fn chunked_bool_canonical_into(bencher: Bencher, (len, chunk_count): (usize, usi bencher.with_inputs(|| &chunk).bench_refs(|chunk| { let mut builder = builder_with_capacity(chunk.dtype(), len * chunk_count); - chunk.append_to_builder(builder.as_mut()); + chunk + .append_to_builder(builder.as_mut()) + .vortex_expect("append_to_builder"); builder.finish() }) } @@ -43,7 +46,9 @@ fn chunked_opt_bool_canonical_into(bencher: Bencher, (len, chunk_count): (usize, bencher.with_inputs(|| &chunk).bench_refs(|chunk| { let mut builder = builder_with_capacity(chunk.dtype(), len * chunk_count); - chunk.append_to_builder(builder.as_mut()); + chunk + .append_to_builder(builder.as_mut()) + .vortex_expect("append_to_builder"); builder.finish() }) } @@ -75,7 +80,9 @@ fn chunked_varbinview_canonical_into(bencher: Bencher, (len, chunk_count): (usiz DType::Utf8(chunk.dtype().nullability()), len * chunk_count, ); - chunk.append_to_builder(&mut builder); + chunk + .append_to_builder(&mut builder) + .vortex_expect("append_to_builder"); builder.finish() }) } @@ -98,7 +105,9 @@ fn chunked_varbinview_opt_canonical_into(bencher: Bencher, (len, chunk_count): ( DType::Utf8(chunk.dtype().nullability()), len * chunk_count, ); - chunk.append_to_builder(&mut builder); + chunk + .append_to_builder(&mut builder) + .vortex_expect("append_to_builder"); builder.finish() }) } diff --git a/vortex-array/benches/chunked_dict_builder.rs b/vortex-array/benches/chunked_dict_builder.rs index 1c549ae72ba..a9cca19631f 100644 --- a/vortex-array/benches/chunked_dict_builder.rs +++ b/vortex-array/benches/chunked_dict_builder.rs @@ -9,6 +9,7 @@ use vortex_array::arrays::dict_test::gen_dict_primitive_chunks; use vortex_array::builders::builder_with_capacity; use vortex_array::compute::warm_up_vtables; use vortex_dtype::NativePType; +use vortex_error::VortexExpect; fn main() { warm_up_vtables(); @@ -35,7 +36,9 @@ fn chunked_dict_primitive_canonical_into( bencher.with_inputs(|| &chunk).bench_refs(|chunk| { let mut builder = builder_with_capacity(chunk.dtype(), len * chunk_count); - chunk.append_to_builder(builder.as_mut()); + chunk + .append_to_builder(builder.as_mut()) + .vortex_expect("append_to_builder"); builder.finish() }) } diff --git a/vortex-array/benches/take_patches.rs b/vortex-array/benches/take_patches.rs index aeefba5008b..7ed0ecafcdf 100644 --- a/vortex-array/benches/take_patches.rs +++ b/vortex-array/benches/take_patches.rs @@ -50,7 +50,9 @@ fn take_search(bencher: Bencher, (patches_sparsity, index_multiple): (f64, f64)) bencher .with_inputs(|| (&patches, &indices)) - .bench_refs(|(patches, indices)| patches.take_search(indices.to_primitive(), false)); + .bench_refs(|(patches, indices)| { + patches.take_search(indices.to_primitive().unwrap(), false) + }); } #[divan::bench(args = BENCH_ARGS)] @@ -65,7 +67,9 @@ fn take_search_chunked(bencher: Bencher, (patches_sparsity, index_multiple): (f6 bencher .with_inputs(|| (&patches, &indices)) - .bench_refs(|(patches, indices)| patches.take_search(indices.to_primitive(), false)); + .bench_refs(|(patches, indices)| { + patches.take_search(indices.to_primitive().unwrap(), false) + }); } #[divan::bench(args = BENCH_ARGS)] @@ -80,7 +84,7 @@ fn take_map(bencher: Bencher, (patches_sparsity, index_multiple): (f64, f64)) { bencher .with_inputs(|| (&patches, &indices)) - .bench_refs(|(patches, indices)| patches.take_map(indices.to_primitive(), false)); + .bench_refs(|(patches, indices)| patches.take_map(indices.to_primitive().unwrap(), false)); } fn fixture(len: usize, sparsity: f64, rng: &mut StdRng) -> Patches { diff --git a/vortex-array/benches/take_strings.rs b/vortex-array/benches/take_strings.rs index 3c82ab1dd4d..30c2068a71d 100644 --- a/vortex-array/benches/take_strings.rs +++ b/vortex-array/benches/take_strings.rs @@ -32,7 +32,7 @@ fn varbin(bencher: Bencher) { #[divan::bench] fn varbinview(bencher: Bencher) { - let array = fixture(65_535).to_varbinview(); + let array = fixture(65_535).to_varbinview().unwrap(); let indices = indices(1024, 65_535); bencher @@ -52,7 +52,7 @@ fn varbin_non_null(bencher: Bencher) { #[divan::bench] fn varbinview_non_null(bencher: Bencher) { - let array = non_null_fixture(65_535).to_varbinview(); + let array = non_null_fixture(65_535).to_varbinview().unwrap(); let indices = indices(1024, 65_535); bencher diff --git a/vortex-array/benches/varbinview_compact.rs b/vortex-array/benches/varbinview_compact.rs index e6d5cbe636a..eb1d6deeb20 100644 --- a/vortex-array/benches/varbinview_compact.rs +++ b/vortex-array/benches/varbinview_compact.rs @@ -43,7 +43,7 @@ fn compact_impl(bencher: Bencher, (output_size, utilization_pct): (usize, usize) let base_array = build_varbinview_fixture(base_size); let indices = random_indices(output_size, base_size); let taken = take(base_array.as_ref(), &indices).vortex_unwrap(); - let array = taken.to_varbinview(); + let array = taken.to_varbinview().vortex_unwrap(); bencher .with_inputs(|| &array) @@ -54,7 +54,7 @@ fn compact_sliced_impl(bencher: Bencher, (output_size, utilization_pct): (usize, let base_size = (output_size * 100) / utilization_pct; let base_array = build_varbinview_fixture(base_size); let sliced = base_array.as_ref().slice(0..output_size); - let array = sliced.to_varbinview(); + let array = sliced.to_varbinview().vortex_unwrap(); bencher .with_inputs(|| &array) diff --git a/vortex-array/src/array/mod.rs b/vortex-array/src/array/mod.rs index ce74762b840..a988826385c 100644 --- a/vortex-array/src/array/mod.rs +++ b/vortex-array/src/array/mod.rs @@ -19,7 +19,6 @@ use vortex_error::VortexExpect; 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; use vortex_session::VortexSession; @@ -131,37 +130,37 @@ pub trait Array: } /// 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`. - 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`. - 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 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) -> Canonical; + fn to_canonical(&self) -> VortexResult; /// Writes the array into the canonical builder. /// /// The [`DType`] of the builder must match that of the array. - fn append_to_builder(&self, builder: &mut dyn ArrayBuilder); + fn append_to_builder(&self, builder: &mut dyn ArrayBuilder) -> VortexResult<()>; /// Returns the statistics of the array. // TODO(ngates): change how this works. It's weird. @@ -235,45 +234,45 @@ impl Array for Arc { } #[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() } #[inline] - fn validity_mask(&self) -> Mask { + fn validity_mask(&self) -> VortexResult { self.as_ref().validity_mask() } - fn to_canonical(&self) -> Canonical { + fn to_canonical(&self) -> VortexResult { self.as_ref().to_canonical() } - fn append_to_builder(&self, builder: &mut dyn ArrayBuilder) { + fn append_to_builder(&self, builder: &mut dyn ArrayBuilder) -> VortexResult<()> { self.as_ref().append_to_builder(builder) } @@ -508,7 +507,10 @@ impl Array for ArrayAdapter { fn scalar_at(&self, index: usize) -> Scalar { assert!(index < self.len(), "index {index} out of bounds"); - if self.is_invalid(index) { + if self + .is_invalid(index) + .vortex_expect("is_invalid should not fail") + { return Scalar::null(self.dtype().clone()); } let scalar = >::scalar_at(&self.0, index); @@ -516,65 +518,65 @@ impl Array for ArrayAdapter { scalar } - fn is_valid(&self, index: usize) -> bool { + fn is_valid(&self, index: usize) -> VortexResult { if index >= self.len() { - vortex_panic!(OutOfBounds: index, 0, self.len()); + vortex_bail!(OutOfBounds: index, 0, self.len()); } >::is_valid(&self.0, index) } - 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 { + fn all_valid(&self) -> VortexResult { >::all_valid(&self.0) } - fn all_invalid(&self) -> bool { + fn all_invalid(&self) -> VortexResult { >::all_invalid(&self.0) } - 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 = >::valid_count(&self.0); + let count = >::valid_count(&self.0)?; assert!(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 { + fn invalid_count(&self) -> VortexResult { if let Some(Precision::Exact(invalid_count)) = self.statistics().get_as::(Stat::NullCount) { - return invalid_count; + return Ok(invalid_count); } - let count = >::invalid_count(&self.0); + let count = >::invalid_count(&self.0)?; assert!(count <= self.len(), "Invalid count exceeds array length"); self.statistics() .set(Stat::NullCount, Precision::exact(count)); - count + Ok(count) } - fn validity_mask(&self) -> Mask { - let mask = >::validity_mask(&self.0); + fn validity_mask(&self) -> VortexResult { + let mask = >::validity_mask(&self.0)?; assert_eq!(mask.len(), self.len(), "Validity mask length mismatch"); - mask + Ok(mask) } - fn to_canonical(&self) -> Canonical { - let canonical = >::canonicalize(&self.0); + fn to_canonical(&self) -> VortexResult { + let canonical = >::canonicalize(&self.0)?; assert_eq!( self.len(), canonical.as_ref().len(), @@ -595,12 +597,12 @@ impl Array for ArrayAdapter { .as_ref() .statistics() .inherit_from(self.statistics()); - canonical + Ok(canonical) } - fn append_to_builder(&self, builder: &mut dyn ArrayBuilder) { + fn append_to_builder(&self, builder: &mut dyn ArrayBuilder) -> VortexResult<()> { if builder.dtype() != self.dtype() { - vortex_panic!( + vortex_bail!( "Builder dtype mismatch: expected {}, got {}", self.dtype(), builder.dtype(), @@ -608,13 +610,14 @@ impl Array for ArrayAdapter { } let len = builder.len(); - >::append_to_builder(&self.0, builder); + >::append_to_builder(&self.0, builder)?; assert_eq!( len + self.len(), builder.len(), "Builder length mismatch after writing array for encoding {}", self.encoding_id(), ); + Ok(()) } fn statistics(&self) -> StatsSetRef<'_> { diff --git a/vortex-array/src/arrays/arbitrary.rs b/vortex-array/src/arrays/arbitrary.rs index 22e0d95e915..f78161fc790 100644 --- a/vortex-array/src/arrays/arbitrary.rs +++ b/vortex-array/src/arrays/arbitrary.rs @@ -108,6 +108,7 @@ fn random_array_chunk( PType::I64 => random_primitive::(u, *n, chunk_len), PType::F16 => Ok(random_primitive::(u, *n, chunk_len)? .to_primitive() + .vortex_expect("to_primitive") .reinterpret_cast(PType::F16) .into_array()), PType::F32 => random_primitive::(u, *n, chunk_len), diff --git a/vortex-array/src/arrays/bool/array.rs b/vortex-array/src/arrays/bool/array.rs index ba9830d32cb..df16d4cc134 100644 --- a/vortex-array/src/arrays/bool/array.rs +++ b/vortex-array/src/arrays/bool/array.rs @@ -198,6 +198,7 @@ impl BoolArray { pub fn maybe_to_mask(&self) -> Option { self.all_valid() + .ok()? .then(|| Mask::from_buffer(self.bit_buffer().clone())) } @@ -211,7 +212,7 @@ impl BoolArray { } } // Extract a boolean buffer, treating null values to false - let buffer = match self.validity_mask() { + let buffer = match self.validity_mask().vortex_expect("validity_mask") { Mask::AllTrue(_) => self.bit_buffer().clone(), Mask::AllFalse(_) => return Mask::new_false(self.len()), Mask::Values(validity) => validity.bit_buffer() & self.bit_buffer(), @@ -319,7 +320,7 @@ mod tests { let arr = BoolArray::from(BitBuffer::new_set(12)); let sliced = arr.slice(4..12); assert_eq!(sliced.len(), 8); - let values = sliced.to_bool().into_bit_buffer().into_mut(); + let values = sliced.to_bool().unwrap().into_bit_buffer().into_mut(); assert_eq!(values.len(), 8); assert_eq!(values.as_slice(), &[255, 255]); @@ -330,7 +331,7 @@ mod tests { }; let sliced = arr.slice(4..12); let sliced_len = sliced.len(); - let values = sliced.to_bool().into_bit_buffer().into_mut(); + let values = sliced.to_bool().unwrap().into_bit_buffer().into_mut(); assert_eq!(values.as_slice(), &[254, 15]); // patch the underlying array @@ -343,12 +344,12 @@ mod tests { ); let arr = arr.patch(&patches); let arr_len = arr.len(); - let values = arr.to_bool().into_bit_buffer().into_mut(); + let values = arr.to_bool().unwrap().into_bit_buffer().into_mut(); assert_eq!(values.len(), arr_len); assert_eq!(values.as_slice(), &[238, 15]); // the slice should be unchanged - let values = sliced.to_bool().into_bit_buffer().into_mut(); + let values = sliced.to_bool().unwrap().into_bit_buffer().into_mut(); assert_eq!(values.len(), sliced_len); assert_eq!(values.as_slice(), &[254, 15]); // unchanged } @@ -358,7 +359,7 @@ mod tests { let arr = BoolArray::from(BitBuffer::new_set(16)); let sliced = arr.slice(4..12); let sliced_len = sliced.len(); - let values = sliced.to_bool().into_bit_buffer().into_mut(); + let values = sliced.to_bool().unwrap().into_bit_buffer().into_mut(); assert_eq!(values.len(), sliced_len); assert_eq!(values.as_slice(), &[255, 255]); } @@ -386,7 +387,7 @@ mod tests { fn patch_sliced_bools_offset() { let arr = BoolArray::from(BitBuffer::new_set(15)); let sliced = arr.slice(4..15); - let values = sliced.to_bool().into_bit_buffer().into_mut(); + let values = sliced.to_bool().unwrap().into_bit_buffer().into_mut(); assert_eq!(values.as_slice(), &[255, 255]); } } diff --git a/vortex-array/src/arrays/bool/compute/fill_null.rs b/vortex-array/src/arrays/bool/compute/fill_null.rs index 79a5fab1a03..66afb622215 100644 --- a/vortex-array/src/arrays/bool/compute/fill_null.rs +++ b/vortex-array/src/arrays/bool/compute/fill_null.rs @@ -25,10 +25,11 @@ impl FillNullKernel for BoolVTable { Ok(match array.validity() { Validity::Array(v) => { + let v_bool = v.to_bool()?; let bool_buffer = if fill { - array.bit_buffer() | &!v.to_bool().bit_buffer() + array.bit_buffer() | &!v_bool.bit_buffer() } else { - array.bit_buffer() & v.to_bool().bit_buffer() + array.bit_buffer() & v_bool.bit_buffer() }; BoolArray::from_bit_buffer(bool_buffer, fill_value.dtype().nullability().into()) .into_array() @@ -63,7 +64,8 @@ mod tests { ); let non_null_array = fill_null(bool_array.as_ref(), &fill_value.into()) .unwrap() - .to_bool(); + .to_bool() + .unwrap(); assert_eq!(non_null_array.bit_buffer(), &expected); assert_eq!( non_null_array.dtype(), diff --git a/vortex-array/src/arrays/bool/compute/filter.rs b/vortex-array/src/arrays/bool/compute/filter.rs index 8117236a507..d8de9ead23d 100644 --- a/vortex-array/src/arrays/bool/compute/filter.rs +++ b/vortex-array/src/arrays/bool/compute/filter.rs @@ -95,7 +95,7 @@ mod test { let arr = BoolArray::from_iter([true, true, false]); let mask = Mask::from_iter([true, false, true]); - let filtered = filter(arr.as_ref(), &mask).unwrap().to_bool(); + let filtered = filter(arr.as_ref(), &mask).unwrap().to_bool().unwrap(); assert_eq!(2, filtered.len()); assert_eq!( 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/mask.rs b/vortex-array/src/arrays/bool/compute/mask.rs index 1798c644eed..0bc92641249 100644 --- a/vortex-array/src/arrays/bool/compute/mask.rs +++ b/vortex-array/src/arrays/bool/compute/mask.rs @@ -16,7 +16,7 @@ use crate::vtable::ValidityHelper; impl MaskKernel for BoolVTable { fn mask(&self, array: &BoolArray, mask: &Mask) -> VortexResult { Ok( - BoolArray::from_bit_buffer(array.bit_buffer().clone(), array.validity().mask(mask)) + BoolArray::from_bit_buffer(array.bit_buffer().clone(), array.validity().mask(mask)?) .into_array(), ) } diff --git a/vortex-array/src/arrays/bool/compute/min_max.rs b/vortex-array/src/arrays/bool/compute/min_max.rs index d42b5796769..de6f78b0d26 100644 --- a/vortex-array/src/arrays/bool/compute/min_max.rs +++ b/vortex-array/src/arrays/bool/compute/min_max.rs @@ -18,7 +18,7 @@ use crate::register_kernel; impl MinMaxKernel for BoolVTable { fn min_max(&self, array: &BoolArray) -> VortexResult> { - let mask = array.validity_mask(); + let mask = array.validity_mask()?; let true_non_null = match mask { Mask::AllTrue(_) => array.bit_buffer().clone(), Mask::AllFalse(_) => return Ok(None), 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..e6c93db1aa8 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( @@ -36,7 +36,7 @@ impl TakeKernel for BoolVTable { } Mask::Values(_) => fill_null(indices, &Scalar::from(0).cast(indices.dtype())?)?, }; - let indices_nulls_zeroed = indices_nulls_zeroed.to_primitive(); + let indices_nulls_zeroed = indices_nulls_zeroed.to_primitive()?; let buffer = match_each_integer_ptype!(indices_nulls_zeroed.ptype(), |I| { take_valid_indices(array.bit_buffer(), indices_nulls_zeroed.as_slice::()) }); @@ -101,7 +101,8 @@ mod test { let b = take(reference.as_ref(), buffer![0, 3, 4].into_array().as_ref()) .unwrap() - .to_bool(); + .to_bool() + .unwrap(); assert_eq!( b.bit_buffer(), BoolArray::from_iter([Some(false), None, Some(false)]).bit_buffer() diff --git a/vortex-array/src/arrays/bool/patch.rs b/vortex-array/src/arrays/bool/patch.rs index c2826e958f1..99cd8fd3452 100644 --- a/vortex-array/src/arrays/bool/patch.rs +++ b/vortex-array/src/arrays/bool/patch.rs @@ -3,6 +3,7 @@ use itertools::Itertools; use vortex_dtype::match_each_unsigned_integer_ptype; +use vortex_error::VortexExpect; use crate::ToCanonical; use crate::arrays::BoolArray; @@ -13,13 +14,17 @@ impl BoolArray { pub fn patch(self, patches: &Patches) -> Self { let len = self.len(); let offset = patches.offset(); - let indices = patches.indices().to_primitive(); - let values = patches.values().to_bool(); + let indices = patches + .indices() + .to_primitive() + .vortex_expect("to_primitive"); + let values = patches.values().to_bool().vortex_expect("to_bool"); - let patched_validity = - self.validity() - .clone() - .patch(len, offset, indices.as_ref(), values.validity()); + let patched_validity = self + .validity() + .clone() + .patch(len, offset, indices.as_ref(), values.validity()) + .vortex_expect("patch validity"); let mut own_values = self.into_bit_buffer().into_mut(); match_each_unsigned_integer_ptype!(indices.ptype(), |I| { @@ -40,6 +45,7 @@ impl BoolArray { #[cfg(test)] mod tests { use vortex_buffer::BitBuffer; + use vortex_error::VortexExpect; use crate::ToCanonical; use crate::arrays::BoolArray; @@ -48,7 +54,11 @@ mod tests { fn patch_sliced_bools() { let arr = BoolArray::from(BitBuffer::new_set(12)); let sliced = arr.slice(4..12); - let values = sliced.to_bool().into_bit_buffer().into_mut(); + let values = sliced + .to_bool() + .vortex_expect("to_bool") + .into_bit_buffer() + .into_mut(); assert_eq!(values.len(), 8); assert_eq!(values.as_slice(), &[255, 255]); } @@ -57,7 +67,11 @@ mod tests { fn patch_sliced_bools_offset() { let arr = BoolArray::from(BitBuffer::new_set(15)); let sliced = arr.slice(4..15); - let values = sliced.to_bool().into_bit_buffer().into_mut(); + let values = sliced + .to_bool() + .vortex_expect("to_bool") + .into_bit_buffer() + .into_mut(); assert_eq!(values.len(), 11); assert_eq!(values.as_slice(), &[255, 255]); } diff --git a/vortex-array/src/arrays/bool/test_harness.rs b/vortex-array/src/arrays/bool/test_harness.rs index f1df7f8f2d2..67cc07b490e 100644 --- a/vortex-array/src/arrays/bool/test_harness.rs +++ b/vortex-array/src/arrays/bool/test_harness.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +#![allow(clippy::unwrap_used)] + use vortex_error::vortex_panic; use crate::arrays::BoolArray; @@ -8,6 +10,7 @@ use crate::arrays::BoolArray; impl BoolArray { pub fn opt_bool_vec(&self) -> Vec> { self.validity_mask() + .unwrap() .to_bit_buffer() .iter() .zip(self.bit_buffer().iter()) @@ -17,6 +20,7 @@ impl BoolArray { pub fn bool_vec(&self) -> Vec { self.validity_mask() + .unwrap() .to_bit_buffer() .iter() .zip(self.bit_buffer().iter()) diff --git a/vortex-array/src/arrays/bool/vtable/canonical.rs b/vortex-array/src/arrays/bool/vtable/canonical.rs index 05737010231..553bdfee52c 100644 --- a/vortex-array/src/arrays/bool/vtable/canonical.rs +++ b/vortex-array/src/arrays/bool/vtable/canonical.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_error::VortexResult; + use crate::Canonical; use crate::arrays::BoolArray; use crate::arrays::BoolVTable; @@ -8,11 +10,12 @@ use crate::builders::ArrayBuilder; use crate::vtable::CanonicalVTable; impl CanonicalVTable for BoolVTable { - fn canonicalize(array: &BoolArray) -> Canonical { - Canonical::Bool(array.clone()) + fn canonicalize(array: &BoolArray) -> VortexResult { + Ok(Canonical::Bool(array.clone())) } - fn append_to_builder(array: &BoolArray, builder: &mut dyn ArrayBuilder) { - builder.extend_from_array(array.as_ref()) + fn append_to_builder(array: &BoolArray, builder: &mut dyn ArrayBuilder) -> VortexResult<()> { + builder.extend_from_array(array.as_ref()); + Ok(()) } } diff --git a/vortex-array/src/arrays/bool/vtable/mod.rs b/vortex-array/src/arrays/bool/vtable/mod.rs index eb049e9fd58..01370a64a84 100644 --- a/vortex-array/src/arrays/bool/vtable/mod.rs +++ b/vortex-array/src/arrays/bool/vtable/mod.rs @@ -107,7 +107,7 @@ impl VTable for BoolVTable { } fn batch_execute(array: &Self::Array, _ctx: &mut ExecutionCtx) -> VortexResult { - Ok(BoolVector::new(array.bit_buffer().clone(), array.validity_mask()).into()) + Ok(BoolVector::new(array.bit_buffer().clone(), array.validity_mask()?).into()) } } diff --git a/vortex-array/src/arrays/bool/vtable/operations.rs b/vortex-array/src/arrays/bool/vtable/operations.rs index 0788a80db75..daef671bbb9 100644 --- a/vortex-array/src/arrays/bool/vtable/operations.rs +++ b/vortex-array/src/arrays/bool/vtable/operations.rs @@ -30,6 +30,8 @@ impl OperationsVTable for BoolVTable { mod tests { use std::iter; + use vortex_error::VortexExpect; + use super::*; use crate::ToCanonical; use crate::assert_arrays_eq; @@ -37,7 +39,7 @@ mod tests { #[test] fn test_slice_hundred_elements() { let arr = BoolArray::from_iter(iter::repeat_n(Some(true), 100)); - let sliced_arr = arr.slice(8..16).to_bool(); + let sliced_arr = arr.slice(8..16).to_bool().vortex_expect("to_bool"); assert_eq!(sliced_arr.len(), 8); assert_eq!(sliced_arr.bit_buffer().len(), 8); assert_eq!(sliced_arr.bit_buffer().offset(), 0); @@ -46,7 +48,7 @@ mod tests { #[test] fn test_slice() { let arr = BoolArray::from_iter([Some(true), Some(true), None, Some(false), None]); - let sliced_arr = arr.slice(1..4).to_bool(); + let sliced_arr = arr.slice(1..4).to_bool().vortex_expect("to_bool"); assert_arrays_eq!( sliced_arr, diff --git a/vortex-array/src/arrays/bool/vtable/operator.rs b/vortex-array/src/arrays/bool/vtable/operator.rs index b3613e27fab..ded3d0c81fe 100644 --- a/vortex-array/src/arrays/bool/vtable/operator.rs +++ b/vortex-array/src/arrays/bool/vtable/operator.rs @@ -40,7 +40,7 @@ impl ArrayParentReduceRule, Exact> for BoolMaske Ok(Some( BoolArray::from_bit_buffer( array.bit_buffer().clone(), - array.validity().clone().and(parent.validity().clone()), + array.validity().clone().and(parent.validity().clone())?, ) .into_array(), )) diff --git a/vortex-array/src/arrays/chunked/array.rs b/vortex-array/src/arrays/chunked/array.rs index 07f68ebe3be..92685c888db 100644 --- a/vortex-array/src/arrays/chunked/array.rs +++ b/vortex-array/src/arrays/chunked/array.rs @@ -174,7 +174,7 @@ impl ChunkedArray { // All chunks are guaranteed to be valid arrays matching self.dtype(). unsafe { ChunkedArray::new_unchecked(chunks_to_combine, self.dtype().clone()) - .to_canonical() + .to_canonical()? .into_array() }, ); @@ -198,7 +198,7 @@ impl ChunkedArray { // SAFETY: chunks_to_combine contains valid chunks of the same dtype as self. // All chunks are guaranteed to be valid arrays matching self.dtype(). ChunkedArray::new_unchecked(chunks_to_combine, self.dtype().clone()) - .to_canonical() + .to_canonical()? .into_array() }); } @@ -259,11 +259,23 @@ mod test { let chunked = array.as_::(); let chunks_out = chunked.chunks(); - let results = chunks_out[0].to_primitive().as_slice::().to_vec(); + let results = chunks_out[0] + .to_primitive() + .unwrap() + .as_slice::() + .to_vec(); assert_eq!(results, &[0u64, 1, 2]); - let results = chunks_out[1].to_primitive().as_slice::().to_vec(); + let results = chunks_out[1] + .to_primitive() + .unwrap() + .as_slice::() + .to_vec(); assert_eq!(results, &[3u64, 4, 5]); - let results = chunks_out[2].to_primitive().as_slice::().to_vec(); + let results = chunks_out[2] + .to_primitive() + .unwrap() + .as_slice::() + .to_vec(); assert_eq!(results, &[6u64, 7, 8]); } @@ -347,8 +359,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(()) } @@ -367,8 +379,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(()) } @@ -387,8 +399,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/cast.rs b/vortex-array/src/arrays/chunked/compute/cast.rs index 25101d7072a..5b93851711a 100644 --- a/vortex-array/src/arrays/chunked/compute/cast.rs +++ b/vortex-array/src/arrays/chunked/compute/cast.rs @@ -73,6 +73,7 @@ mod test { ) .unwrap() .to_primitive() + .unwrap() .as_slice::(), &[0u64, 1, 2, 3], ); diff --git a/vortex-array/src/arrays/chunked/compute/take.rs b/vortex-array/src/arrays/chunked/compute/take.rs index 71c0e3a4458..8feeba91d65 100644 --- a/vortex-array/src/arrays/chunked/compute/take.rs +++ b/vortex-array/src/arrays/chunked/compute/take.rs @@ -26,11 +26,11 @@ impl TakeKernel for ChunkedVTable { indices, &DType::Primitive(PType::U64, indices.dtype().nullability()), )? - .to_primitive(); + .to_primitive()?; // 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(); @@ -112,7 +112,10 @@ mod test { assert_eq!(arr.len(), 9); let indices = buffer![0u64, 0, 6, 4].into_array(); - let result = take(arr.as_ref(), indices.as_ref()).unwrap().to_primitive(); + let result = take(arr.as_ref(), indices.as_ref()) + .unwrap() + .to_primitive() + .unwrap(); assert_eq!(result.as_slice::(), &[1, 1, 1, 2]); } @@ -152,7 +155,10 @@ mod test { assert_eq!(arr.len(), 9); let indices = PrimitiveArray::empty::(Nullability::NonNullable); - let result = take(arr.as_ref(), indices.as_ref()).unwrap().to_primitive(); + let result = take(arr.as_ref(), indices.as_ref()) + .unwrap() + .to_primitive() + .unwrap(); assert!(result.is_empty()); assert_eq!(result.dtype(), arr.dtype()); diff --git a/vortex-array/src/arrays/chunked/compute/zip.rs b/vortex-array/src/arrays/chunked/compute/zip.rs index db2e95997c3..fc93ff93fc7 100644 --- a/vortex-array/src/arrays/chunked/compute/zip.rs +++ b/vortex-array/src/arrays/chunked/compute/zip.rs @@ -120,7 +120,7 @@ mod tests { assert_eq!(zipped.nchunks(), 4); let mut values: Vec = Vec::new(); for chunk in zipped.chunks() { - let primitive = chunk.to_primitive(); + let primitive = chunk.to_primitive().unwrap(); values.extend_from_slice(primitive.as_slice::()); } assert_eq!(values, vec![1, 11, 3, 13, 5]); diff --git a/vortex-array/src/arrays/chunked/tests.rs b/vortex-array/src/arrays/chunked/tests.rs index aa7dcba56ed..e49d3fadff1 100644 --- a/vortex-array/src/arrays/chunked/tests.rs +++ b/vortex-array/src/arrays/chunked/tests.rs @@ -39,10 +39,10 @@ fn assert_equal_slices(arr: &dyn Array, slice: &[T]) { if let Some(arr) = arr.as_opt::() { arr.chunks() .iter() - .map(|a| a.to_primitive()) + .map(|a| a.to_primitive().unwrap()) .for_each(|a| values.extend_from_slice(a.as_slice::())); } else { - values.extend_from_slice(arr.to_primitive().as_slice::()); + values.extend_from_slice(arr.to_primitive().unwrap().as_slice::()); } assert_eq!(values, slice); } @@ -158,9 +158,9 @@ pub fn pack_nested_structs() { ) .unwrap() .into_array(); - let canonical_struct = chunked.to_struct(); - let canonical_varbin = canonical_struct.fields()[0].to_varbinview(); - let original_varbin = struct_array.fields()[0].to_varbinview(); + let canonical_struct = chunked.to_struct().unwrap(); + let canonical_varbin = canonical_struct.fields()[0].to_varbinview().unwrap(); + let original_varbin = struct_array.fields()[0].to_varbinview().unwrap(); let orig_values = original_varbin.with_iterator(|it| it.map(|a| a.map(|v| v.to_vec())).collect::>()); let canon_values = @@ -192,7 +192,7 @@ pub fn pack_nested_lists() { ), ); - let canon_values = chunked_list.unwrap().to_listview(); + let canon_values = chunked_list.unwrap().to_listview().unwrap(); assert_eq!(l1.scalar_at(0), canon_values.scalar_at(0)); assert_eq!(l2.scalar_at(0), canon_values.scalar_at(1)); diff --git a/vortex-array/src/arrays/chunked/vtable/canonical.rs b/vortex-array/src/arrays/chunked/vtable/canonical.rs index d9e6968139b..13895e688c3 100644 --- a/vortex-array/src/arrays/chunked/vtable/canonical.rs +++ b/vortex-array/src/arrays/chunked/vtable/canonical.rs @@ -7,6 +7,7 @@ use vortex_dtype::Nullability; use vortex_dtype::PType; use vortex_dtype::StructFields; use vortex_error::VortexExpect; +use vortex_error::VortexResult; use crate::Array; use crate::ArrayRef; @@ -26,9 +27,9 @@ use crate::validity::Validity; use crate::vtable::CanonicalVTable; impl CanonicalVTable for ChunkedVTable { - fn canonicalize(array: &ChunkedArray) -> Canonical { + fn canonicalize(array: &ChunkedArray) -> VortexResult { if array.nchunks() == 0 { - return Canonical::empty(array.dtype()); + return Ok(Canonical::empty(array.dtype())); } if array.nchunks() == 1 { return array.chunks()[0].to_canonical(); @@ -40,26 +41,27 @@ impl CanonicalVTable for ChunkedVTable { array.chunks(), Validity::copy_from_array(array.as_ref()), struct_dtype, - ); - Canonical::Struct(struct_array) + )?; + Ok(Canonical::Struct(struct_array)) } - DType::List(elem_dtype, _) => Canonical::List(swizzle_list_chunks( + DType::List(elem_dtype, _) => Ok(Canonical::List(swizzle_list_chunks( array.chunks(), Validity::copy_from_array(array.as_ref()), elem_dtype, - )), + )?)), _ => { let mut builder = builder_with_capacity(array.dtype(), array.len()); - array.append_to_builder(builder.as_mut()); - builder.finish_into_canonical() + array.append_to_builder(builder.as_mut())?; + Ok(builder.finish_into_canonical()) } } } - fn append_to_builder(array: &ChunkedArray, builder: &mut dyn ArrayBuilder) { + fn append_to_builder(array: &ChunkedArray, builder: &mut dyn ArrayBuilder) -> VortexResult<()> { for chunk in array.chunks() { - chunk.append_to_builder(builder); + chunk.append_to_builder(builder)?; } + Ok(()) } } @@ -71,7 +73,7 @@ fn pack_struct_chunks( chunks: &[ArrayRef], validity: Validity, struct_dtype: &StructFields, -) -> StructArray { +) -> VortexResult { let len = chunks.iter().map(|chunk| chunk.len()).sum(); let mut field_arrays = Vec::new(); @@ -79,13 +81,14 @@ fn pack_struct_chunks( let field_chunks = chunks .iter() .map(|c| { - c.to_struct() - .fields() - .get(field_idx) - .vortex_expect("Invalid field index") - .to_array() + c.to_struct().map(|s| { + s.fields() + .get(field_idx) + .vortex_expect("Invalid field index") + .to_array() + }) }) - .collect::>(); + .collect::>>()?; // SAFETY: field_chunks are extracted from valid StructArrays with matching dtypes. // Each chunk's field array is guaranteed to be valid for field_dtype. @@ -95,7 +98,7 @@ fn pack_struct_chunks( // SAFETY: field_arrays are built from corresponding chunks of same length, dtypes match by // construction. - unsafe { StructArray::new_unchecked(field_arrays, struct_dtype.clone(), len, validity) } + Ok(unsafe { StructArray::new_unchecked(field_arrays, struct_dtype.clone(), len, validity) }) } /// Packs [`ListViewArray`]s together into a chunked `ListViewArray`. @@ -107,7 +110,7 @@ fn swizzle_list_chunks( chunks: &[ArrayRef], validity: Validity, elem_dtype: &DType, -) -> ListViewArray { +) -> VortexResult { let len: usize = chunks.iter().map(|c| c.len()).sum(); assert_eq!( @@ -133,7 +136,7 @@ fn swizzle_list_chunks( let mut sizes = BufferMut::::with_capacity(len); for chunk in chunks { - let chunk_array = chunk.to_listview(); + let chunk_array = chunk.to_listview()?; // By rebuilding as zero-copy to `List` and trimming all elements (to prevent gaps), we make // the final output `ListView` also zero-copyable to `List`. let chunk_array = chunk_array.rebuild(ListViewRebuildMode::MakeExact); @@ -145,16 +148,14 @@ fn swizzle_list_chunks( let offsets_arr = cast( chunk_array.offsets(), &DType::Primitive(PType::U64, Nullability::NonNullable), - ) - .vortex_expect("Must be able to fit array offsets in u64") - .to_primitive(); + )? + .to_primitive()?; let sizes_arr = cast( chunk_array.sizes(), &DType::Primitive(PType::U64, Nullability::NonNullable), - ) - .vortex_expect("Must be able to fit array offsets in u64") - .to_primitive(); + )? + .to_primitive()?; let offsets_slice = offsets_arr.as_slice::(); let sizes_slice = sizes_arr.as_slice::(); @@ -181,10 +182,10 @@ fn swizzle_list_chunks( // - Validity came from the outer chunked array so it must have the same length // - Since we made sure that all chunks were zero-copyable to a list above, we know that the // final concatenated output is also zero-copyable to a list. - unsafe { + Ok(unsafe { ListViewArray::new_unchecked(chunked_elements, offsets, sizes, validity) .with_zero_copy_to_list(true) - } + }) } #[cfg(test)] @@ -226,9 +227,9 @@ mod tests { ) .unwrap() .into_array(); - let canonical_struct = chunked.to_struct(); - let canonical_varbin = canonical_struct.fields()[0].to_varbinview(); - let original_varbin = struct_array.fields()[0].to_varbinview(); + let canonical_struct = chunked.to_struct().unwrap(); + let canonical_varbin = canonical_struct.fields()[0].to_varbinview().unwrap(); + let original_varbin = struct_array.fields()[0].to_varbinview().unwrap(); let orig_values = original_varbin .with_iterator(|it| it.map(|a| a.map(|v| v.to_vec())).collect::>()); let canon_values = canonical_varbin @@ -257,7 +258,7 @@ mod tests { List(Arc::new(Primitive(I32, NonNullable)), NonNullable), ); - let canon_values = chunked_list.unwrap().to_listview(); + let canon_values = chunked_list.unwrap().to_listview().unwrap(); assert_eq!(l1.scalar_at(0), canon_values.scalar_at(0)); assert_eq!(l2.scalar_at(0), canon_values.scalar_at(1)); diff --git a/vortex-array/src/arrays/chunked/vtable/compute.rs b/vortex-array/src/arrays/chunked/vtable/compute.rs index d7360d462be..9ea4364986f 100644 --- a/vortex-array/src/arrays/chunked/vtable/compute.rs +++ b/vortex-array/src/arrays/chunked/vtable/compute.rs @@ -108,7 +108,8 @@ mod tests { let result = boolean(chunked1.as_ref(), chunked2.as_ref(), BooleanOperator::Or) .unwrap() - .to_bool(); + .to_bool() + .unwrap(); assert_eq!( result.bit_buffer(), &BitBuffer::from_iter([true, true, false, false, true]) diff --git a/vortex-array/src/arrays/chunked/vtable/mod.rs b/vortex-array/src/arrays/chunked/vtable/mod.rs index bb0491e54db..988536c9d13 100644 --- a/vortex-array/src/arrays/chunked/vtable/mod.rs +++ b/vortex-array/src/arrays/chunked/vtable/mod.rs @@ -91,7 +91,7 @@ impl VTable for ChunkedVTable { // 1 extra offset for the end of the last chunk nchunks + 1, )? - .to_primitive(); + .to_primitive()?; let chunk_offsets_buf = chunk_offsets_array.buffer::(); diff --git a/vortex-array/src/arrays/chunked/vtable/operations.rs b/vortex-array/src/arrays/chunked/vtable/operations.rs index 0d72993bcbd..3f66c901d32 100644 --- a/vortex-array/src/arrays/chunked/vtable/operations.rs +++ b/vortex-array/src/arrays/chunked/vtable/operations.rs @@ -93,10 +93,10 @@ mod tests { if let Some(arr) = arr.as_opt::() { arr.chunks() .iter() - .map(|a| a.to_primitive()) + .map(|a| a.to_primitive().unwrap()) .for_each(|a| values.extend_from_slice(a.as_slice::())); } else { - values.extend_from_slice(arr.to_primitive().as_slice::()); + values.extend_from_slice(arr.to_primitive().unwrap().as_slice::()); } assert_eq!(values, slice); } diff --git a/vortex-array/src/arrays/chunked/vtable/validity.rs b/vortex-array/src/arrays/chunked/vtable/validity.rs index 12cdd857601..4d150ab26bf 100644 --- a/vortex-array/src/arrays/chunked/vtable/validity.rs +++ b/vortex-array/src/arrays/chunked/vtable/validity.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_error::VortexResult; use vortex_mask::Mask; use crate::Array; @@ -9,39 +10,43 @@ use crate::arrays::ChunkedVTable; use crate::vtable::ValidityVTable; impl ValidityVTable for ChunkedVTable { - fn is_valid(array: &ChunkedArray, index: usize) -> bool { + fn is_valid(array: &ChunkedArray, index: usize) -> VortexResult { if !array.dtype.is_nullable() { - return true; + return Ok(true); } let (chunk, offset_in_chunk) = array.find_chunk_idx(index); array.chunk(chunk).is_valid(offset_in_chunk) } - fn all_valid(array: &ChunkedArray) -> bool { + fn all_valid(array: &ChunkedArray) -> VortexResult { if !array.dtype().is_nullable() { - return true; + return Ok(true); } for chunk in array.non_empty_chunks() { - if !chunk.all_valid() { - return false; + if !chunk.all_valid()? { + return Ok(false); } } - true + Ok(true) } - fn all_invalid(array: &ChunkedArray) -> bool { + fn all_invalid(array: &ChunkedArray) -> VortexResult { if !array.dtype().is_nullable() { - return false; + return Ok(false); } for chunk in array.non_empty_chunks() { - if !chunk.all_invalid() { - return false; + if !chunk.all_invalid()? { + return Ok(false); } } - true + Ok(true) } - fn validity_mask(array: &ChunkedArray) -> Mask { - array.chunks().iter().map(|a| a.validity_mask()).collect() + fn validity_mask(array: &ChunkedArray) -> VortexResult { + array + .chunks() + .iter() + .map(|a| a.validity_mask()) + .collect::>() } } diff --git a/vortex-array/src/arrays/constant/compute/boolean.rs b/vortex-array/src/arrays/constant/compute/boolean.rs index de72cc6022b..2ba8fd05b3e 100644 --- a/vortex-array/src/arrays/constant/compute/boolean.rs +++ b/vortex-array/src/arrays/constant/compute/boolean.rs @@ -89,6 +89,7 @@ fn kleene_or(left: Option, right: Option) -> Option { #[cfg(test)] mod test { use rstest::rstest; + use vortex_error::VortexResult; use crate::Array; use crate::ArrayRef; @@ -104,8 +105,8 @@ mod test { )] #[case(BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)].into_iter()).into_array(), ConstantArray::new(true, 4).into_array() )] - fn test_or(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) { - let r = or(&lhs, &rhs).unwrap().to_bool().into_array(); + fn test_or(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) -> VortexResult<()> { + 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(); @@ -116,6 +117,7 @@ mod test { assert!(v1.unwrap()); assert!(v2.unwrap()); assert!(v3.unwrap()); + Ok(()) } #[rstest] @@ -123,8 +125,8 @@ mod test { )] #[case(BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)].into_iter()).into_array(), ConstantArray::new(true, 4).into_array())] - fn test_and(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) { - let r = and(&lhs, &rhs).unwrap().to_bool().into_array(); + fn test_and(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) -> VortexResult<()> { + 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(); @@ -135,5 +137,6 @@ mod test { assert!(!v1.unwrap()); assert!(v2.unwrap()); assert!(!v3.unwrap()); + Ok(()) } } diff --git a/vortex-array/src/arrays/constant/compute/take.rs b/vortex-array/src/arrays/constant/compute/take.rs index 4dd83f164d8..69c5dbe6da7 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 @@ -58,6 +58,7 @@ mod tests { use rstest::rstest; use vortex_buffer::buffer; use vortex_dtype::Nullability; + use vortex_error::VortexResult; use vortex_mask::AllOr; use vortex_scalar::Scalar; @@ -71,7 +72,7 @@ mod tests { use crate::validity::Validity; #[test] - fn take_nullable_indices() { + fn take_nullable_indices() -> VortexResult<()> { let array = ConstantArray::new(42, 10).to_array(); let taken = take( &array, @@ -87,12 +88,13 @@ mod tests { &array.dtype().with_nullability(Nullability::Nullable), taken.dtype() ); - assert_eq!(taken.to_primitive().as_slice::(), &[42, 42, 42]); - assert_eq!(taken.validity_mask().indices(), AllOr::Some(valid_indices)); + assert_eq!(taken.to_primitive()?.as_slice::(), &[42, 42, 42]); + assert_eq!(taken.validity_mask()?.indices(), AllOr::Some(valid_indices)); + Ok(()) } #[test] - fn take_all_valid_indices() { + fn take_all_valid_indices() -> VortexResult<()> { let array = ConstantArray::new(42, 10).to_array(); let taken = take( &array, @@ -103,8 +105,9 @@ mod tests { &array.dtype().with_nullability(Nullability::Nullable), taken.dtype() ); - assert_eq!(taken.to_primitive().as_slice::(), &[42, 42, 42]); - assert_eq!(taken.validity_mask().indices(), AllOr::All); + assert_eq!(taken.to_primitive()?.as_slice::(), &[42, 42, 42]); + assert_eq!(taken.validity_mask()?.indices(), AllOr::All); + Ok(()) } #[rstest] diff --git a/vortex-array/src/arrays/constant/vtable/canonical.rs b/vortex-array/src/arrays/constant/vtable/canonical.rs index dcdf18c4ed3..e1e9b7b3549 100644 --- a/vortex-array/src/arrays/constant/vtable/canonical.rs +++ b/vortex-array/src/arrays/constant/vtable/canonical.rs @@ -13,6 +13,7 @@ use vortex_dtype::match_each_decimal_value; use vortex_dtype::match_each_decimal_value_type; use vortex_dtype::match_each_native_ptype; use vortex_error::VortexExpect; +use vortex_error::VortexResult; use vortex_scalar::BinaryScalar; use vortex_scalar::BoolScalar; use vortex_scalar::DecimalValue; @@ -41,7 +42,7 @@ use crate::validity::Validity; use crate::vtable::CanonicalVTable; impl CanonicalVTable for ConstantVTable { - fn canonicalize(array: &ConstantArray) -> Canonical { + fn canonicalize(array: &ConstantArray) -> VortexResult { let scalar = array.scalar(); let validity = match array.dtype().nullability() { @@ -53,8 +54,8 @@ impl CanonicalVTable for ConstantVTable { }; match array.dtype() { - DType::Null => Canonical::Null(NullArray::new(array.len())), - DType::Bool(..) => Canonical::Bool(BoolArray::from_bit_buffer( + DType::Null => Ok(Canonical::Null(NullArray::new(array.len()))), + DType::Bool(..) => Ok(Canonical::Bool(BoolArray::from_bit_buffer( if BoolScalar::try_from(scalar) .vortex_expect("must be bool") .value() @@ -65,23 +66,21 @@ impl CanonicalVTable for ConstantVTable { BitBuffer::new_unset(array.len()) }, validity, - )), - DType::Primitive(ptype, ..) => { - match_each_native_ptype!(ptype, |P| { - Canonical::Primitive(PrimitiveArray::new( - if scalar.is_valid() { - Buffer::full( - P::try_from(scalar) - .vortex_expect("Couldn't unwrap scalar to primitive"), - array.len(), - ) - } else { - Buffer::zeroed(array.len()) - }, - validity, - )) - }) - } + ))), + DType::Primitive(ptype, ..) => Ok(match_each_native_ptype!(ptype, |P| { + Canonical::Primitive(PrimitiveArray::new( + if scalar.is_valid() { + Buffer::full( + P::try_from(scalar) + .vortex_expect("Couldn't unwrap scalar to primitive"), + array.len(), + ) + } else { + Buffer::zeroed(array.len()) + }, + validity, + )) + })), DType::Decimal(decimal_type, ..) => { let size = DecimalType::smallest_decimal_value_type(decimal_type); let decimal = scalar.as_decimal(); @@ -96,7 +95,7 @@ impl CanonicalVTable for ConstantVTable { ) } }); - return Canonical::Decimal(all_null); + return Ok(Canonical::Decimal(all_null)); }; let decimal_array = match_each_decimal_value!(value, |value| { @@ -109,29 +108,29 @@ impl CanonicalVTable for ConstantVTable { ) } }); - Canonical::Decimal(decimal_array) + Ok(Canonical::Decimal(decimal_array)) } DType::Utf8(_) => { let value = Utf8Scalar::try_from(scalar) .vortex_expect("Must be a utf8 scalar") .value(); let const_value = value.as_ref().map(|v| v.as_bytes()); - Canonical::VarBinView(constant_canonical_byte_view( + Ok(Canonical::VarBinView(constant_canonical_byte_view( const_value, array.dtype(), array.len(), - )) + ))) } DType::Binary(_) => { let value = BinaryScalar::try_from(scalar) .vortex_expect("must be a binary scalar") .value(); let const_value = value.as_ref().map(|v| v.as_slice()); - Canonical::VarBinView(constant_canonical_byte_view( + Ok(Canonical::VarBinView(constant_canonical_byte_view( const_value, array.dtype(), array.len(), - )) + ))) } DType::Struct(struct_dtype, _) => { let value = StructScalar::try_from(scalar).vortex_expect("must be struct"); @@ -141,7 +140,11 @@ impl CanonicalVTable for ConstantVTable { .map(|s| ConstantArray::new(s, array.len()).into_array()) .collect(), None => { - assert!(validity.all_invalid(array.len())); + assert!( + validity + .all_invalid(array.len()) + .vortex_expect("all_invalid") + ); struct_dtype .fields() .map(|dt| { @@ -153,20 +156,25 @@ impl CanonicalVTable for ConstantVTable { }; // SAFETY: Fields are constructed from the same struct scalar, all have same // length, dtypes match by construction. - Canonical::Struct(unsafe { + Ok(Canonical::Struct(unsafe { StructArray::new_unchecked(fields, struct_dtype.clone(), array.len(), validity) - }) + })) } - DType::List(..) => Canonical::List(constant_canonical_list_array(scalar, array.len())), + DType::List(..) => Ok(Canonical::List(constant_canonical_list_array( + scalar, + array.len(), + ))), DType::FixedSizeList(element_dtype, list_size, _) => { let value = ListScalar::try_from(scalar).vortex_expect("must be list"); - Canonical::FixedSizeList(constant_canonical_fixed_size_list_array( - value.elements(), - element_dtype, - *list_size, - value.dtype().nullability(), - array.len(), + Ok(Canonical::FixedSizeList( + constant_canonical_fixed_size_list_array( + value.elements(), + element_dtype, + *list_size, + value.dtype().nullability(), + array.len(), + ), )) } DType::Extension(ext_dtype) => { @@ -174,7 +182,10 @@ impl CanonicalVTable for ConstantVTable { let storage_scalar = s.storage(); let storage_self = ConstantArray::new(storage_scalar, array.len()).into_array(); - Canonical::Extension(ExtensionArray::new(ext_dtype.clone(), storage_self)) + Ok(Canonical::Extension(ExtensionArray::new( + ext_dtype.clone(), + storage_self, + ))) } } } @@ -329,6 +340,7 @@ mod tests { use vortex_dtype::Nullability; use vortex_dtype::PType; use vortex_dtype::half::f16; + use vortex_error::VortexResult; use vortex_scalar::Scalar; use crate::Array; @@ -342,36 +354,38 @@ mod tests { use crate::vtable::ValidityHelper; #[test] - fn test_canonicalize_null() { + fn test_canonicalize_null() -> VortexResult<()> { let const_null = ConstantArray::new(Scalar::null(DType::Null), 42); - let actual = const_null.to_null(); + let actual = const_null.to_null()?; assert_eq!(actual.len(), 42); assert_eq!(actual.scalar_at(33), Scalar::null(DType::Null)); + Ok(()) } #[test] - fn test_canonicalize_const_str() { + fn test_canonicalize_const_str() -> VortexResult<()> { let const_array = ConstantArray::new("four".to_string(), 4); // Check all values correct. - let canonical = const_array.to_varbinview(); + let canonical = const_array.to_varbinview()?; assert_eq!(canonical.len(), 4); for i in 0..=3 { assert_eq!(canonical.scalar_at(i), "four".into()); } + Ok(()) } #[test] - fn test_canonicalize_propagates_stats() { + fn test_canonicalize_propagates_stats() -> VortexResult<()> { let scalar = Scalar::bool(true, Nullability::NonNullable); let const_array = ConstantArray::new(scalar, 4).into_array(); let stats = const_array .statistics() .compute_all(&all::().collect_vec()) .unwrap(); - let canonical = const_array.to_canonical(); + let canonical = const_array.to_canonical()?; let canonical_stats = canonical.as_ref().statistics(); let stats_ref = stats.as_typed_ref(canonical.as_ref().dtype()); @@ -386,86 +400,91 @@ mod tests { "stat mismatch {stat}" ); } + Ok(()) } #[test] - fn test_canonicalize_scalar_values() { + fn test_canonicalize_scalar_values() -> VortexResult<()> { let f16_value = f16::from_f32(5.722046e-6); let f16_scalar = Scalar::primitive(f16_value, Nullability::NonNullable); // Create a ConstantArray with the f16 scalar let const_array = ConstantArray::new(f16_scalar.clone(), 1).into_array(); - let canonical_const = const_array.to_primitive(); + let canonical_const = const_array.to_primitive()?; // Verify the scalar value is preserved through canonicalization assert_eq!(canonical_const.scalar_at(0), f16_scalar); + Ok(()) } #[test] - fn test_canonicalize_lists() { + fn test_canonicalize_lists() -> VortexResult<()> { let list_scalar = Scalar::list( Arc::new(DType::Primitive(PType::U64, Nullability::NonNullable)), vec![1u64.into(), 2u64.into()], Nullability::NonNullable, ); let const_array = ConstantArray::new(list_scalar, 2).into_array(); - let canonical_const = const_array.to_listview(); + let canonical_const = const_array.to_listview()?; let list_array = canonical_const.rebuild(ListViewRebuildMode::MakeZeroCopyToList); assert_eq!( - list_array.elements().to_primitive().as_slice::(), + list_array.elements().to_primitive()?.as_slice::(), [1u64, 2, 1, 2] ); assert_eq!( - list_array.offsets().to_primitive().as_slice::(), + list_array.offsets().to_primitive()?.as_slice::(), [0u64, 2] ); assert_eq!( - list_array.sizes().to_primitive().as_slice::(), + list_array.sizes().to_primitive()?.as_slice::(), [2u64, 2] ); + Ok(()) } #[test] - fn test_canonicalize_empty_list() { + fn test_canonicalize_empty_list() -> VortexResult<()> { let list_scalar = Scalar::list( Arc::new(DType::Primitive(PType::U64, Nullability::NonNullable)), vec![], Nullability::NonNullable, ); let const_array = ConstantArray::new(list_scalar, 2).into_array(); - let canonical_const = const_array.to_listview(); - assert!(canonical_const.elements().to_primitive().is_empty()); + let canonical_const = const_array.to_listview()?; + assert!(canonical_const.elements().to_primitive()?.is_empty()); assert_eq!( - canonical_const.offsets().to_primitive().as_slice::(), + canonical_const.offsets().to_primitive()?.as_slice::(), [0u64, 0] ); assert_eq!( - canonical_const.sizes().to_primitive().as_slice::(), + canonical_const.sizes().to_primitive()?.as_slice::(), [0u64, 0] ); + Ok(()) } #[test] - fn test_canonicalize_null_list() { + fn test_canonicalize_null_list() -> VortexResult<()> { let list_scalar = Scalar::null(DType::List( Arc::new(DType::Primitive(PType::U64, Nullability::NonNullable)), Nullability::Nullable, )); let const_array = ConstantArray::new(list_scalar, 2).into_array(); - let canonical_const = const_array.to_listview(); - assert!(canonical_const.elements().to_primitive().is_empty()); + let canonical_const = const_array.to_listview()?; + assert!(canonical_const.elements().to_primitive()?.is_empty()); assert_eq!( - canonical_const.offsets().to_primitive().as_slice::(), + canonical_const.offsets().to_primitive()?.as_slice::(), [0u64, 0] ); assert_eq!( - canonical_const.sizes().to_primitive().as_slice::(), + canonical_const.sizes().to_primitive()?.as_slice::(), [0u64, 0] ); + Ok(()) } #[test] - fn test_canonicalize_nullable_struct() { + fn test_canonicalize_nullable_struct() -> VortexResult<()> { let array = ConstantArray::new( Scalar::null(DType::struct_( [( @@ -477,9 +496,9 @@ mod tests { 3, ); - let struct_array = array.to_struct(); + 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()?, 0); let field = struct_array.field_by_name("non_null_field").unwrap(); @@ -487,10 +506,11 @@ mod tests { field.dtype(), &DType::Primitive(PType::I8, Nullability::NonNullable) ); + Ok(()) } #[test] - fn test_canonicalize_fixed_size_list_non_null() { + fn test_canonicalize_fixed_size_list_non_null() -> VortexResult<()> { // Test with a non-null fixed-size list constant. let fsl_scalar = Scalar::fixed_size_list( Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), @@ -503,7 +523,7 @@ mod tests { ); let const_array = ConstantArray::new(fsl_scalar, 4).into_array(); - let canonical = const_array.to_fixed_size_list(); + let canonical = const_array.to_fixed_size_list()?; assert_eq!(canonical.len(), 4); assert_eq!(canonical.list_size(), 3); @@ -512,13 +532,14 @@ mod tests { // Check that each list is [10, 20, 30]. for i in 0..4 { let list = canonical.fixed_size_list_elements_at(i); - let list_primitive = list.to_primitive(); + let list_primitive = list.to_primitive()?; assert_eq!(list_primitive.as_slice::(), [10, 20, 30]); } + Ok(()) } #[test] - fn test_canonicalize_fixed_size_list_nullable() { + fn test_canonicalize_fixed_size_list_nullable() -> VortexResult<()> { // Test with a nullable but non-null fixed-size list constant. let fsl_scalar = Scalar::fixed_size_list( Arc::new(DType::Primitive(PType::F64, Nullability::NonNullable)), @@ -530,19 +551,20 @@ mod tests { ); let const_array = ConstantArray::new(fsl_scalar, 3).into_array(); - let canonical = const_array.to_fixed_size_list(); + let canonical = const_array.to_fixed_size_list()?; assert_eq!(canonical.len(), 3); assert_eq!(canonical.list_size(), 2); assert_eq!(canonical.validity(), &Validity::AllValid); // Check elements. - let elements = canonical.elements().to_primitive(); + let elements = canonical.elements().to_primitive()?; assert_eq!(elements.as_slice::(), [1.5, 2.5, 1.5, 2.5, 1.5, 2.5]); + Ok(()) } #[test] - fn test_canonicalize_fixed_size_list_null() { + fn test_canonicalize_fixed_size_list_null() -> VortexResult<()> { // Test with a null fixed-size list constant. let fsl_scalar = Scalar::null(DType::FixedSizeList( Arc::new(DType::Primitive(PType::U64, Nullability::NonNullable)), @@ -551,20 +573,21 @@ mod tests { )); let const_array = ConstantArray::new(fsl_scalar, 5).into_array(); - let canonical = const_array.to_fixed_size_list(); + let canonical = const_array.to_fixed_size_list()?; assert_eq!(canonical.len(), 5); assert_eq!(canonical.list_size(), 4); assert_eq!(canonical.validity(), &Validity::AllInvalid); // Elements should be defaults (zeros). - let elements = canonical.elements().to_primitive(); + let elements = canonical.elements().to_primitive()?; assert_eq!(elements.len(), 20); // 5 lists * 4 elements each assert!(elements.as_slice::().iter().all(|&x| x == 0)); + Ok(()) } #[test] - fn test_canonicalize_fixed_size_list_empty() { + fn test_canonicalize_fixed_size_list_empty() -> VortexResult<()> { // Test with size-0 lists (edge case). let fsl_scalar = Scalar::fixed_size_list( Arc::new(DType::Primitive(PType::I8, Nullability::NonNullable)), @@ -573,7 +596,7 @@ mod tests { ); let const_array = ConstantArray::new(fsl_scalar, 10).into_array(); - let canonical = const_array.to_fixed_size_list(); + let canonical = const_array.to_fixed_size_list()?; assert_eq!(canonical.len(), 10); assert_eq!(canonical.list_size(), 0); @@ -581,10 +604,11 @@ mod tests { // Elements array should be empty. assert!(canonical.elements().is_empty()); + Ok(()) } #[test] - fn test_canonicalize_fixed_size_list_nested() { + fn test_canonicalize_fixed_size_list_nested() -> VortexResult<()> { // Test with nested data types (list of strings). let fsl_scalar = Scalar::fixed_size_list( Arc::new(DType::Utf8(Nullability::NonNullable)), @@ -593,21 +617,22 @@ mod tests { ); let const_array = ConstantArray::new(fsl_scalar, 2).into_array(); - let canonical = const_array.to_fixed_size_list(); + let canonical = const_array.to_fixed_size_list()?; assert_eq!(canonical.len(), 2); assert_eq!(canonical.list_size(), 2); // Check elements are repeated correctly. - let elements = canonical.elements().to_varbinview(); + 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()); + Ok(()) } #[test] - fn test_canonicalize_fixed_size_list_single_element() { + fn test_canonicalize_fixed_size_list_single_element() -> VortexResult<()> { // Test with a single-element list. let fsl_scalar = Scalar::fixed_size_list( Arc::new(DType::Primitive(PType::I16, Nullability::NonNullable)), @@ -616,17 +641,18 @@ mod tests { ); let const_array = ConstantArray::new(fsl_scalar, 1).into_array(); - let canonical = const_array.to_fixed_size_list(); + let canonical = const_array.to_fixed_size_list()?; assert_eq!(canonical.len(), 1); assert_eq!(canonical.list_size(), 1); - let elements = canonical.elements().to_primitive(); + let elements = canonical.elements().to_primitive()?; assert_eq!(elements.as_slice::(), [42]); + Ok(()) } #[test] - fn test_canonicalize_fixed_size_list_with_null_elements() { + fn test_canonicalize_fixed_size_list_with_null_elements() -> VortexResult<()> { // Test FSL with nullable element type where some elements are null. let fsl_scalar = Scalar::fixed_size_list( Arc::new(DType::Primitive(PType::I32, Nullability::Nullable)), @@ -639,32 +665,33 @@ mod tests { ); let const_array = ConstantArray::new(fsl_scalar, 3).into_array(); - let canonical = const_array.to_fixed_size_list(); + let canonical = const_array.to_fixed_size_list()?; assert_eq!(canonical.len(), 3); assert_eq!(canonical.list_size(), 3); assert_eq!(canonical.validity(), &Validity::NonNullable); // Check elements including nulls. - let elements = canonical.elements().to_primitive(); + let elements = canonical.elements().to_primitive()?; assert_eq!(elements.as_slice::()[0], 100); assert_eq!(elements.as_slice::()[1], 0); // null becomes 0 assert_eq!(elements.as_slice::()[2], 200); // Check element validity. let element_validity = elements.validity(); - assert!(element_validity.is_valid(0)); - assert!(!element_validity.is_valid(1)); - assert!(element_validity.is_valid(2)); + assert!(element_validity.is_valid(0).unwrap()); + assert!(!element_validity.is_valid(1).unwrap()); + assert!(element_validity.is_valid(2).unwrap()); // Pattern should repeat. - assert!(element_validity.is_valid(3)); - assert!(!element_validity.is_valid(4)); - assert!(element_validity.is_valid(5)); + assert!(element_validity.is_valid(3).unwrap()); + assert!(!element_validity.is_valid(4).unwrap()); + assert!(element_validity.is_valid(5).unwrap()); + Ok(()) } #[test] - fn test_canonicalize_fixed_size_list_large() { + fn test_canonicalize_fixed_size_list_large() -> VortexResult<()> { // Test with a large constant array. let fsl_scalar = Scalar::fixed_size_list( Arc::new(DType::Primitive(PType::U8, Nullability::NonNullable)), @@ -679,12 +706,12 @@ mod tests { ); let const_array = ConstantArray::new(fsl_scalar, 1000).into_array(); - let canonical = const_array.to_fixed_size_list(); + let canonical = const_array.to_fixed_size_list()?; assert_eq!(canonical.len(), 1000); assert_eq!(canonical.list_size(), 5); - let elements = canonical.elements().to_primitive(); + let elements = canonical.elements().to_primitive()?; assert_eq!(elements.len(), 5000); // Check pattern repeats correctly. @@ -696,5 +723,6 @@ mod tests { assert_eq!(elements.as_slice::()[base + 3], 4); assert_eq!(elements.as_slice::()[base + 4], 5); } + Ok(()) } } diff --git a/vortex-array/src/arrays/constant/vtable/validity.rs b/vortex-array/src/arrays/constant/vtable/validity.rs index 8c744e3b68a..4576fac7a5f 100644 --- a/vortex-array/src/arrays/constant/vtable/validity.rs +++ b/vortex-array/src/arrays/constant/vtable/validity.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_error::VortexResult; use vortex_mask::Mask; use crate::arrays::ConstantArray; @@ -8,22 +9,22 @@ use crate::arrays::ConstantVTable; use crate::vtable::ValidityVTable; impl ValidityVTable for ConstantVTable { - fn is_valid(array: &ConstantArray, _index: usize) -> bool { - !array.scalar().is_null() + fn is_valid(array: &ConstantArray, _index: usize) -> VortexResult { + Ok(!array.scalar().is_null()) } - fn all_valid(array: &ConstantArray) -> bool { - !array.scalar().is_null() + fn all_valid(array: &ConstantArray) -> VortexResult { + Ok(!array.scalar().is_null()) } - fn all_invalid(array: &ConstantArray) -> bool { - array.scalar().is_null() + fn all_invalid(array: &ConstantArray) -> VortexResult { + Ok(array.scalar().is_null()) } - 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 6b710ad52ca..69b54ef15d7 100644 --- a/vortex-array/src/arrays/datetime/test.rs +++ b/vortex-array/src/arrays/datetime/test.rs @@ -5,6 +5,7 @@ use rstest::rstest; use vortex_buffer::buffer; use vortex_dtype::datetime::TemporalMetadata; use vortex_dtype::datetime::TimeUnit; +use vortex_error::VortexResult; use crate::IntoArray; use crate::ToCanonical; @@ -18,7 +19,19 @@ macro_rules! test_temporal_roundtrip { ($prim:ty, $constructor:expr, $unit:expr) => {{ let array = buffer![100 as $prim].into_array(); let temporal: TemporalArray = $constructor(array, $unit); - let prims = temporal.temporal_values().to_primitive(); + let prims = temporal.temporal_values().to_primitive()?; + + assert_eq!(prims.as_slice::<$prim>(), vec![100 as $prim].as_slice(),); + assert_eq!(temporal.temporal_metadata().time_unit(), $unit); + Ok(()) + }}; +} + +macro_rules! test_temporal_roundtrip_panic { + ($prim:ty, $constructor:expr, $unit:expr) => {{ + let array = buffer![100 as $prim].into_array(); + let temporal: TemporalArray = $constructor(array, $unit); + let prims = temporal.temporal_values().to_primitive().unwrap(); assert_eq!(prims.as_slice::<$prim>(), vec![100 as $prim].as_slice(),); assert_eq!(temporal.temporal_metadata().time_unit(), $unit); @@ -28,8 +41,8 @@ macro_rules! test_temporal_roundtrip { macro_rules! test_success_case { ($name:ident, $prim:ty, $constructor:expr, $unit:expr) => { #[test] - fn $name() { - test_temporal_roundtrip!($prim, $constructor, $unit); + fn $name() -> VortexResult<()> { + test_temporal_roundtrip!($prim, $constructor, $unit) } }; } @@ -39,7 +52,7 @@ macro_rules! test_fail_case { #[test] #[should_panic] fn $name() { - test_temporal_roundtrip!($prim, $constructor, $unit) + test_temporal_roundtrip_panic!($prim, $constructor, $unit) } }; } @@ -132,7 +145,7 @@ test_fail_case!( // We test Timestamp explicitly to avoid the macro getting too complex. #[test] -fn test_timestamp() { +fn test_timestamp() -> VortexResult<()> { let ts = buffer![100i64].into_array(); let ts_array = ts.into_array(); @@ -146,7 +159,7 @@ fn test_timestamp() { let temporal_array = TemporalArray::new_timestamp(ts_array.to_array(), unit, tz.clone()); - let values = temporal_array.temporal_values().to_primitive(); + let values = temporal_array.temporal_values().to_primitive()?; assert_eq!(values.as_slice::(), vec![100i64].as_slice()); assert_eq!( temporal_array.temporal_metadata(), @@ -154,6 +167,7 @@ fn test_timestamp() { ); } } + Ok(()) } #[test] @@ -170,7 +184,7 @@ fn test_timestamp_fails_i32() { #[case(Validity::AllValid)] #[case(Validity::AllInvalid)] #[case(Validity::from_iter([true, false, true]))] -fn test_validity_preservation(#[case] validity: Validity) { +fn test_validity_preservation(#[case] validity: Validity) -> VortexResult<()> { let milliseconds = PrimitiveArray::new( buffer![ 86_400i64, // element with only day component @@ -186,7 +200,8 @@ fn test_validity_preservation(#[case] validity: Validity) { Some("UTC".to_string()), ); assert_eq!( - temporal_array.temporal_values().to_primitive().validity(), + temporal_array.temporal_values().to_primitive()?.validity(), &validity ); + Ok(()) } diff --git a/vortex-array/src/arrays/decimal/array.rs b/vortex-array/src/arrays/decimal/array.rs index 29671af0808..411e979f29d 100644 --- a/vortex-array/src/arrays/decimal/array.rs +++ b/vortex-array/src/arrays/decimal/array.rs @@ -259,15 +259,25 @@ impl DecimalArray { )] pub fn patch(self, patches: &Patches) -> Self { let offset = patches.offset(); - let patch_indices = patches.indices().to_primitive(); - let patch_values = patches.values().to_decimal(); + let patch_indices = patches + .indices() + .to_primitive() + .vortex_expect("patch indices must be primitive"); + let patch_values = patches + .values() + .to_decimal() + .vortex_expect("patch values must be decimal"); - let patched_validity = self.validity().clone().patch( - self.len(), - offset, - patch_indices.as_ref(), - patch_values.validity(), - ); + let patched_validity = self + .validity() + .clone() + .patch( + self.len(), + offset, + patch_indices.as_ref(), + patch_values.validity(), + ) + .vortex_expect("failed to patch validity"); assert_eq!(self.decimal_dtype(), patch_values.decimal_dtype()); match_each_integer_ptype!(patch_indices.ptype(), |I| { diff --git a/vortex-array/src/arrays/decimal/compute/between.rs b/vortex-array/src/arrays/decimal/compute/between.rs index e732be5765c..dcae5fb0eaf 100644 --- a/vortex-array/src/arrays/decimal/compute/between.rs +++ b/vortex-array/src/arrays/decimal/compute/between.rs @@ -185,6 +185,6 @@ mod tests { } fn bool_to_vec(array: &dyn Array) -> Vec { - array.to_bool().bit_buffer().iter().collect() + array.to_bool().unwrap().bit_buffer().iter().collect() } } diff --git a/vortex-array/src/arrays/decimal/compute/cast.rs b/vortex-array/src/arrays/decimal/compute/cast.rs index 040a64ee05a..a93bf017e47 100644 --- a/vortex-array/src/arrays/decimal/compute/cast.rs +++ b/vortex-array/src/arrays/decimal/compute/cast.rs @@ -92,7 +92,10 @@ mod tests { // Cast to nullable let nullable_dtype = DType::Decimal(decimal_dtype, Nullability::Nullable); - let casted = cast(array.as_ref(), &nullable_dtype).unwrap().to_decimal(); + let casted = cast(array.as_ref(), &nullable_dtype) + .unwrap() + .to_decimal() + .unwrap(); assert_eq!(casted.dtype(), &nullable_dtype); assert_eq!(casted.validity(), &Validity::AllValid); @@ -110,7 +113,8 @@ mod tests { let non_nullable_dtype = DType::Decimal(decimal_dtype, Nullability::NonNullable); let casted = cast(array.as_ref(), &non_nullable_dtype) .unwrap() - .to_decimal(); + .to_decimal() + .unwrap(); assert_eq!(casted.dtype(), &non_nullable_dtype); assert_eq!(casted.validity(), &Validity::NonNullable); diff --git a/vortex-array/src/arrays/decimal/compute/fill_null.rs b/vortex-array/src/arrays/decimal/compute/fill_null.rs index 0f17b9aa5c5..bb701b67d66 100644 --- a/vortex-array/src/arrays/decimal/compute/fill_null.rs +++ b/vortex-array/src/arrays/decimal/compute/fill_null.rs @@ -25,7 +25,7 @@ impl FillNullKernel for DecimalVTable { Ok(match array.validity() { Validity::Array(is_valid) => { - let is_invalid = is_valid.to_bool().bit_buffer().not(); + let is_invalid = is_valid.to_bool()?.bit_buffer().not(); match_each_decimal_value_type!(array.values_type(), |T| { let mut buffer = array.buffer::().into_mut(); let fill_value = fill_value @@ -77,7 +77,8 @@ mod tests { ), ) .unwrap() - .to_decimal(); + .to_decimal() + .unwrap(); assert_arrays_eq!( p, DecimalArray::from_iter([4200, 800, 4200, 1000, 4200], decimal_dtype) @@ -86,7 +87,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] @@ -107,7 +108,8 @@ mod tests { ), ) .unwrap() - .to_decimal(); + .to_decimal() + .unwrap(); assert_arrays_eq!( p, DecimalArray::from_iter([25500, 25500, 25500, 25500, 25500], decimal_dtype) @@ -132,7 +134,8 @@ mod tests { ), ) .unwrap() - .to_decimal(); + .to_decimal() + .unwrap(); assert_arrays_eq!( p, DecimalArray::from_iter([800i128, 1000, 1200, 1400, 1600], decimal_dtype) 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| { @@ -52,8 +53,8 @@ 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 primitive_codes = array.codes().to_primitive()?; + 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)] @@ -63,6 +64,6 @@ impl ValidityVTable for DictVTable { }); Mask::from_buffer(is_valid_buffer) } - } + }) } } diff --git a/vortex-array/src/arrays/extension/vtable/canonical.rs b/vortex-array/src/arrays/extension/vtable/canonical.rs index f9f24922ce3..2d964e8bab6 100644 --- a/vortex-array/src/arrays/extension/vtable/canonical.rs +++ b/vortex-array/src/arrays/extension/vtable/canonical.rs @@ -1,13 +1,15 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_error::VortexResult; + use crate::Canonical; use crate::arrays::extension::ExtensionArray; use crate::arrays::extension::ExtensionVTable; use crate::vtable::CanonicalVTable; impl CanonicalVTable for ExtensionVTable { - fn canonicalize(array: &ExtensionArray) -> Canonical { - Canonical::Extension(array.clone()) + fn canonicalize(array: &ExtensionArray) -> VortexResult { + Ok(Canonical::Extension(array.clone())) } } diff --git a/vortex-array/src/arrays/fixed_size_list/array.rs b/vortex-array/src/arrays/fixed_size_list/array.rs index 814c4c20739..3437d70a0ea 100644 --- a/vortex-array/src/arrays/fixed_size_list/array.rs +++ b/vortex-array/src/arrays/fixed_size_list/array.rs @@ -223,7 +223,7 @@ impl FixedSizeListArray { "index out of bounds: the len is {} but the index is {index}", self.len ); - debug_assert!(self.validity.is_valid(index)); + debug_assert!(self.validity.is_valid(index).unwrap_or(false)); let start = self.list_size as usize * index; let end = self.list_size as usize * (index + 1); diff --git a/vortex-array/src/arrays/fixed_size_list/compute/mask.rs b/vortex-array/src/arrays/fixed_size_list/compute/mask.rs index 26232addc37..2b27db722cf 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/mask.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/mask.rs @@ -24,7 +24,7 @@ impl MaskKernel for FixedSizeListVTable { FixedSizeListArray::new_unchecked( array.elements().clone(), array.list_size(), - array.validity().mask(mask), + array.validity().mask(mask)?, array.len(), ) } 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..aac19509412 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/take.rs @@ -47,7 +47,7 @@ fn take_with_indices( ) -> VortexResult { let list_size = array.list_size() as usize; - let indices_array = indices.to_primitive(); + let indices_array = indices.to_primitive()?; // Make sure to handle degenerate case where lists have size 0 (these can take fast paths). if list_size == 0 { @@ -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/nested.rs b/vortex-array/src/arrays/fixed_size_list/tests/nested.rs index 9a2180d10c2..dd2526bd357 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/nested.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/nested.rs @@ -82,6 +82,7 @@ fn test_fsl_of_fsl_basic() { // Check first inner list [1,2]. let inner_list_0 = first_outer_list .to_fixed_size_list() + .unwrap() .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()); @@ -89,6 +90,7 @@ fn test_fsl_of_fsl_basic() { // Check second inner list [3,4]. let inner_list_1 = first_outer_list .to_fixed_size_list() + .unwrap() .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()); @@ -96,6 +98,7 @@ fn test_fsl_of_fsl_basic() { // Check third inner list [5,6]. let inner_list_2 = first_outer_list .to_fixed_size_list() + .unwrap() .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()); @@ -106,6 +109,7 @@ fn test_fsl_of_fsl_basic() { // Check first inner list [7,8]. let inner_list_0 = second_outer_list .to_fixed_size_list() + .unwrap() .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()); @@ -113,6 +117,7 @@ fn test_fsl_of_fsl_basic() { // Check second inner list [9,10]. let inner_list_1 = second_outer_list .to_fixed_size_list() + .unwrap() .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()); @@ -120,6 +125,7 @@ fn test_fsl_of_fsl_basic() { // Check third inner list [11,12]. let inner_list_2 = second_outer_list .to_fixed_size_list() + .unwrap() .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()); @@ -218,26 +224,40 @@ fn test_deeply_nested_fsl() { let top_level = level3.fixed_size_list_elements_at(0); let level2_0 = top_level .to_fixed_size_list() + .unwrap() .fixed_size_list_elements_at(0); let level2_1 = top_level .to_fixed_size_list() + .unwrap() .fixed_size_list_elements_at(1); // First level-2 list: [[1,2],[3,4]]. - let level1_0_0 = level2_0.to_fixed_size_list().fixed_size_list_elements_at(0); + let level1_0_0 = level2_0 + .to_fixed_size_list() + .unwrap() + .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()); - let level1_0_1 = level2_0.to_fixed_size_list().fixed_size_list_elements_at(1); + let level1_0_1 = level2_0 + .to_fixed_size_list() + .unwrap() + .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()); // Second level-2 list: [[5,6],[7,8]]. - let level1_1_0 = level2_1.to_fixed_size_list().fixed_size_list_elements_at(0); + let level1_1_0 = level2_1 + .to_fixed_size_list() + .unwrap() + .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()); - let level1_1_1 = level2_1.to_fixed_size_list().fixed_size_list_elements_at(1); + let level1_1_1 = level2_1 + .to_fixed_size_list() + .unwrap() + .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()); } diff --git a/vortex-array/src/arrays/fixed_size_list/vtable/canonical.rs b/vortex-array/src/arrays/fixed_size_list/vtable/canonical.rs index 4f6d4d86039..dac77805e2e 100644 --- a/vortex-array/src/arrays/fixed_size_list/vtable/canonical.rs +++ b/vortex-array/src/arrays/fixed_size_list/vtable/canonical.rs @@ -1,13 +1,15 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_error::VortexResult; + use crate::Canonical; use crate::arrays::FixedSizeListArray; use crate::arrays::FixedSizeListVTable; use crate::vtable::CanonicalVTable; impl CanonicalVTable for FixedSizeListVTable { - fn canonicalize(array: &FixedSizeListArray) -> Canonical { - Canonical::FixedSizeList(array.clone()) + fn canonicalize(array: &FixedSizeListArray) -> VortexResult { + Ok(Canonical::FixedSizeList(array.clone())) } } diff --git a/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs b/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs index 0678275db17..b05a9482777 100644 --- a/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs +++ b/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs @@ -113,7 +113,7 @@ impl VTable for FixedSizeListVTable { FixedSizeListVector::new_unchecked( Arc::new(array.elements().batch_execute(ctx)?), array.list_size(), - array.validity_mask(), + array.validity_mask()?, ) } .into()) diff --git a/vortex-array/src/arrays/list/array.rs b/vortex-array/src/arrays/list/array.rs index f7d3c13abcd..e57075c35ef 100644 --- a/vortex-array/src/arrays/list/array.rs +++ b/vortex-array/src/arrays/list/array.rs @@ -290,7 +290,7 @@ impl ListArray { pub fn reset_offsets(&self, recurse: bool) -> VortexResult { let mut elements = self.sliced_elements(); if recurse && elements.is_canonical() { - elements = elements.to_canonical().compact()?.into_array(); + elements = elements.to_canonical()?.compact()?.into_array(); } else if recurse && let Some(child_list_array) = elements.as_opt::() { elements = child_list_array.reset_offsets(recurse)?.into_array(); } diff --git a/vortex-array/src/arrays/list/compute/filter.rs b/vortex-array/src/arrays/list/compute/filter.rs index aafd45cbda3..84c5a6782f6 100644 --- a/vortex-array/src/arrays/list/compute/filter.rs +++ b/vortex-array/src/arrays/list/compute/filter.rs @@ -34,7 +34,7 @@ const MASK_EXPANSION_DENSITY_THRESHOLD: f64 = 0.05; impl FilterKernel for ListVTable { fn filter(&self, array: &ListArray, selection_mask: &Mask) -> VortexResult { let elements = array.elements(); - let offsets = array.offsets().to_primitive(); + let offsets = array.offsets().to_primitive()?; let new_validity = array.validity().filter(selection_mask)?; debug_assert!( diff --git a/vortex-array/src/arrays/list/compute/mask.rs b/vortex-array/src/arrays/list/compute/mask.rs index 877386133c1..c923860d754 100644 --- a/vortex-array/src/arrays/list/compute/mask.rs +++ b/vortex-array/src/arrays/list/compute/mask.rs @@ -18,7 +18,7 @@ impl MaskKernel for ListVTable { ListArray::try_new( array.elements().clone(), array.offsets().clone(), - array.validity().mask(mask), + array.validity().mask(mask)?, ) .map(|a| a.into_array()) } diff --git a/vortex-array/src/arrays/list/compute/take.rs b/vortex-array/src/arrays/list/compute/take.rs index 4cf4cbb15cf..dd8a1b124d2 100644 --- a/vortex-array/src/arrays/list/compute/take.rs +++ b/vortex-array/src/arrays/list/compute/take.rs @@ -35,8 +35,8 @@ use crate::vtable::ValidityHelper; /// non-contiguous indices would violate this requirement. impl TakeKernel for ListVTable { fn take(&self, array: &ListArray, indices: &dyn Array) -> VortexResult { - let indices = indices.to_primitive(); - let offsets = array.offsets().to_primitive(); + let indices = indices.to_primitive()?; + let offsets = array.offsets().to_primitive()?; match_each_integer_ptype!(offsets.dtype().as_ptype(), |O| { match_each_integer_ptype!(indices.ptype(), |I| { @@ -44,8 +44,8 @@ impl TakeKernel for ListVTable { array, offsets.as_slice::(), &indices, - array.validity_mask(), - indices.validity_mask(), + array.validity_mask()?, + indices.validity_mask()?, ) }) }) @@ -116,7 +116,7 @@ fn _take( indices_array .validity() .clone() - .and(array.validity().clone()), + .and(array.validity().clone())?, )? .to_array()) } @@ -234,13 +234,13 @@ mod test { ) ); - let result = result.to_listview(); + let result = result.to_listview().unwrap(); assert_eq!(result.len(), 4); 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), Scalar::list( @@ -250,9 +250,9 @@ 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), Scalar::list( @@ -262,7 +262,7 @@ mod test { ) ); - assert!(result.is_valid(3)); + assert!(result.is_valid(3).unwrap()); assert_eq!( result.scalar_at(3), Scalar::list(element_dtype, vec![], Nullability::Nullable) @@ -314,13 +314,13 @@ mod test { ) ); - let result = result.to_listview(); + let result = result.to_listview().unwrap(); assert_eq!(result.len(), 3); 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), Scalar::list( @@ -330,7 +330,7 @@ mod test { ) ); - assert!(result.is_valid(1)); + assert!(result.is_valid(1).unwrap()); assert_eq!( result.scalar_at(1), Scalar::list( @@ -340,7 +340,7 @@ mod test { ) ); - assert!(result.is_valid(2)); + assert!(result.is_valid(2).unwrap()); assert_eq!( result.scalar_at(2), Scalar::list(element_dtype, vec![], Nullability::NonNullable) diff --git a/vortex-array/src/arrays/list/vtable/canonical.rs b/vortex-array/src/arrays/list/vtable/canonical.rs index 8b757616b26..03d6110cba5 100644 --- a/vortex-array/src/arrays/list/vtable/canonical.rs +++ b/vortex-array/src/arrays/list/vtable/canonical.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_error::VortexResult; + use crate::Canonical; use crate::arrays::ListArray; use crate::arrays::ListVTable; @@ -8,7 +10,7 @@ use crate::arrays::list_view_from_list; use crate::vtable::CanonicalVTable; impl CanonicalVTable for ListVTable { - fn canonicalize(array: &ListArray) -> Canonical { - Canonical::List(list_view_from_list(array.clone())) + fn canonicalize(array: &ListArray) -> VortexResult { + Ok(Canonical::List(list_view_from_list(array.clone()))) } } diff --git a/vortex-array/src/arrays/listview/array.rs b/vortex-array/src/arrays/listview/array.rs index 599d09ae2e9..c17d5b6530e 100644 --- a/vortex-array/src/arrays/listview/array.rs +++ b/vortex-array/src/arrays/listview/array.rs @@ -242,8 +242,8 @@ impl ListViewArray { ); } - let offsets_primitive = offsets.to_primitive(); - let sizes_primitive = sizes.to_primitive(); + let offsets_primitive = offsets.to_primitive()?; + let sizes_primitive = sizes.to_primitive()?; // Validate the `offsets` and `sizes` arrays. match_each_integer_ptype!(offset_ptype, |O| { @@ -284,8 +284,12 @@ impl ListViewArray { if cfg!(debug_assertions) && is_zctl { validate_zctl( &self.elements, - self.offsets.to_primitive(), - self.sizes.to_primitive(), + self.offsets + .to_primitive() + .vortex_expect("Failed to get primitive offsets"), + self.sizes + .to_primitive() + .vortex_expect("Failed to get primitive sizes"), ) .vortex_expect("Failed to validate zero-copy to list flag"); } @@ -311,12 +315,14 @@ impl ListViewArray { /// [`ListArray`]: crate::arrays::ListArray /// [`with_zero_copy_to_list`]: Self::with_zero_copy_to_list pub fn verify_is_zero_copy_to_list(&self) -> bool { - validate_zctl( - &self.elements, - self.offsets.to_primitive(), - self.sizes.to_primitive(), - ) - .is_ok() + let Ok(offsets_primitive) = self.offsets.to_primitive() else { + return false; + }; + let Ok(sizes_primitive) = self.sizes.to_primitive() else { + return false; + }; + + validate_zctl(&self.elements, offsets_primitive, sizes_primitive).is_ok() } /// Returns the offset at the given index. diff --git a/vortex-array/src/arrays/listview/compute/mask.rs b/vortex-array/src/arrays/listview/compute/mask.rs index e34ec88527f..e50a9b666ca 100644 --- a/vortex-array/src/arrays/listview/compute/mask.rs +++ b/vortex-array/src/arrays/listview/compute/mask.rs @@ -22,7 +22,7 @@ impl MaskKernel for ListViewVTable { array.elements().clone(), array.offsets().clone(), array.sizes().clone(), - array.validity().mask(mask), + array.validity().mask(mask)?, ) .with_zero_copy_to_list(array.is_zero_copy_to_list()) } diff --git a/vortex-array/src/arrays/listview/conversion.rs b/vortex-array/src/arrays/listview/conversion.rs index fff033075da..25eef063c7d 100644 --- a/vortex-array/src/arrays/listview/conversion.rs +++ b/vortex-array/src/arrays/listview/conversion.rs @@ -74,7 +74,10 @@ fn build_sizes_from_offsets(list: &ListArray) -> ArrayRef { // Create `UninitRange` for direct memory access. let mut sizes_range = sizes_builder.uninit_range(len); - let offsets = list.offsets().to_primitive(); + let offsets = list + .offsets() + .to_primitive() + .vortex_expect("list offsets must be primitive"); let offsets_slice = offsets.as_slice::(); debug_assert_eq!(len + 1, offsets_slice.len()); debug_assert!(offsets_slice.is_sorted()); @@ -144,7 +147,10 @@ unsafe fn build_list_offsets_from_list_view( // Create uninit range for direct memory access. let mut offsets_range = offsets_builder.uninit_range(len + 1); - let offsets = list_view.offsets().to_primitive(); + let offsets = list_view + .offsets() + .to_primitive() + .vortex_expect("listview offsets must be primitive"); let offsets_slice = offsets.as_slice::(); debug_assert!(offsets_slice.is_sorted()); @@ -182,7 +188,9 @@ pub fn recursive_list_from_list_view(array: ArrayRef) -> ArrayRef { return array; } - let canonical = array.to_canonical(); + let canonical = array + .to_canonical() + .vortex_expect("canonicalization failed"); match canonical { Canonical::List(listview) => { diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index 2904fec1fcd..3057710b116 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -114,9 +114,15 @@ impl ListViewArray { .as_list_element_opt() .vortex_expect("somehow had a canonical list that was not a list"); - let offsets_canonical = self.offsets().to_primitive(); + let offsets_canonical = self + .offsets() + .to_primitive() + .vortex_expect("offsets must be primitive"); let offsets_slice = offsets_canonical.as_slice::(); - let sizes_canonical = self.sizes().to_primitive(); + let sizes_canonical = self + .sizes() + .to_primitive() + .vortex_expect("sizes must be primitive"); let sizes_slice = sizes_canonical.as_slice::(); let len = offsets_slice.len(); @@ -129,7 +135,11 @@ impl ListViewArray { let mut new_sizes = BufferMut::::with_capacity(len); // Canonicalize the elements up front as we will be slicing the elements quite a lot. - let elements_canonical = self.elements().to_canonical().into_array(); + let elements_canonical = self + .elements() + .to_canonical() + .vortex_expect("elements must be canonical") + .into_array(); // Note that we do not know what the exact capacity should be of the new elements since // there could be overlaps in the existing `ListViewArray`. @@ -138,7 +148,7 @@ impl ListViewArray { let mut n_elements = NewOffset::zero(); for index in 0..len { - if !self.is_valid(index) { + if !self.is_valid(index).vortex_expect("validity check failed") { // 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); @@ -293,10 +303,10 @@ mod tests { assert_eq!(flattened.size_at(1), 2); // Verify the data is correct - let list0 = flattened.list_elements_at(0).to_primitive(); + let list0 = flattened.list_elements_at(0).to_primitive().unwrap(); assert_eq!(list0.as_slice::(), &[1, 2, 3]); - let list1 = flattened.list_elements_at(1).to_primitive(); + let list1 = flattened.list_elements_at(1).to_primitive().unwrap(); assert_eq!(list1.as_slice::(), &[2, 3]); } @@ -322,15 +332,15 @@ mod tests { // Verify nullability is preserved assert_eq!(flattened.dtype().nullability(), Nullability::Nullable); - assert!(flattened.validity().is_valid(0)); - assert!(!flattened.validity().is_valid(1)); - assert!(flattened.validity().is_valid(2)); + assert!(flattened.validity().is_valid(0).unwrap()); + assert!(!flattened.validity().is_valid(1).unwrap()); + assert!(flattened.validity().is_valid(2).unwrap()); // Verify valid lists contain correct data - let list0 = flattened.list_elements_at(0).to_primitive(); + let list0 = flattened.list_elements_at(0).to_primitive().unwrap(); assert_eq!(list0.as_slice::(), &[1, 2]); - let list2 = flattened.list_elements_at(2).to_primitive(); + let list2 = flattened.list_elements_at(2).to_primitive().unwrap(); assert_eq!(list2.as_slice::(), &[3]); } @@ -362,14 +372,14 @@ mod tests { assert_eq!(trimmed.size_at(1), 2); // Verify the data is correct. - let list0 = trimmed.list_elements_at(0).to_primitive(); + let list0 = trimmed.list_elements_at(0).to_primitive().unwrap(); assert_eq!(list0.as_slice::(), &[1, 2]); - let list1 = trimmed.list_elements_at(1).to_primitive(); + let list1 = trimmed.list_elements_at(1).to_primitive().unwrap(); assert_eq!(list1.as_slice::(), &[3, 4]); // Note that element at index 2 (97) is preserved as a gap. - let all_elements = trimmed.elements().to_primitive(); + let all_elements = trimmed.elements().to_primitive().unwrap(); assert_eq!(all_elements.scalar_at(2), 97i32.into()); } @@ -409,16 +419,16 @@ 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 - let list0 = exact.list_elements_at(0).to_primitive(); + let list0 = exact.list_elements_at(0).to_primitive().unwrap(); assert_eq!(list0.as_slice::(), &[1, 2]); - let list1 = exact.list_elements_at(1).to_primitive(); + let list1 = exact.list_elements_at(1).to_primitive().unwrap(); assert_eq!(list1.as_slice::(), &[3, 4]); } } diff --git a/vortex-array/src/arrays/listview/tests/basic.rs b/vortex-array/src/arrays/listview/tests/basic.rs index 3228511d16f..f22f1eb6e42 100644 --- a/vortex-array/src/arrays/listview/tests/basic.rs +++ b/vortex-array/src/arrays/listview/tests/basic.rs @@ -8,6 +8,7 @@ use vortex_buffer::buffer; use vortex_dtype::DType; use vortex_dtype::Nullability; use vortex_dtype::PType; +use vortex_error::VortexResult; use vortex_scalar::Scalar; use crate::Array; @@ -117,7 +118,7 @@ fn test_empty_listview() { } #[test] -fn test_from_list_array() { +fn test_from_list_array() -> VortexResult<()> { // Test conversion from ListArray to ListViewArray. // Logical lists: [[1,2], null, [5,6,7]] let offsets = buffer![0i64, 2, 4, 7].into_array(); @@ -135,15 +136,17 @@ fn test_from_list_array() { assert_eq!(first.scalar_at(1), 2i32.into()); // 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)?); + assert!(list_view.is_invalid(1)?); + assert!(list_view.is_valid(2)?); // Check third list. let third = list_view.list_elements_at(2); assert_eq!(third.scalar_at(0), 5i32.into()); assert_eq!(third.scalar_at(1), 6i32.into()); assert_eq!(third.scalar_at(2), 7i32.into()); + + Ok(()) } // Parameterized tests for ConstantArray scenarios. diff --git a/vortex-array/src/arrays/listview/tests/filter.rs b/vortex-array/src/arrays/listview/tests/filter.rs index ce3be49dce0..10e6e867379 100644 --- a/vortex-array/src/arrays/listview/tests/filter.rs +++ b/vortex-array/src/arrays/listview/tests/filter.rs @@ -3,6 +3,7 @@ use rstest::rstest; use vortex_buffer::buffer; +use vortex_error::VortexResult; use vortex_mask::Mask; use super::common::create_basic_listview; @@ -32,7 +33,7 @@ fn test_filter_listview_conformance(#[case] listview: ListViewArray) { #[ignore = "TODO(connor)[ListView]: Don't rebuild ListView after every `filter`"] #[test] -fn test_filter_preserves_unreferenced_elements() { +fn test_filter_preserves_unreferenced_elements() -> VortexResult<()> { // ListView-specific: Test that filter preserves the entire elements array. // // Logical list: [[5,6,7], [2,3], [8,9], [0,1], [1,2,3,4]] @@ -47,12 +48,12 @@ fn test_filter_preserves_unreferenced_elements() { // Filter to keep only 2 lists. let mask = Mask::from_iter([true, false, false, true, false]); let result = filter(&listview, &mask).unwrap(); - let result_list = result.to_listview(); + let result_list = result.to_listview()?; assert_eq!(result_list.len(), 2, "Wrong number of filtered lists"); // Verify the entire elements array is preserved. - let result_elements = result_list.elements().to_primitive(); + let result_elements = result_list.elements().to_primitive()?; assert_eq!( result_elements.as_slice::(), &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], @@ -62,11 +63,13 @@ fn test_filter_preserves_unreferenced_elements() { // Verify offsets are unchanged. assert_eq!(result_list.offset_at(0), 5, "Wrong offset at index 0"); assert_eq!(result_list.offset_at(1), 0, "Wrong offset at index 3"); + + Ok(()) } #[ignore = "TODO(connor)[ListView]: Don't rebuild ListView after every `filter`"] #[test] -fn test_filter_with_gaps() { +fn test_filter_with_gaps() -> VortexResult<()> { // ListView-specific: Test filtering with gaps in elements array. // // Logical list: [[1,2,3], [7,8,9], [11,12], [2,3], [8,9]] @@ -81,12 +84,12 @@ fn test_filter_with_gaps() { // Filter to keep lists with gaps and overlaps. let mask = Mask::from_iter([false, true, true, true, false]); let result = filter(&listview, &mask).unwrap(); - let result_list = result.to_listview(); + let result_list = result.to_listview()?; assert_eq!(result_list.len(), 3, "Wrong filter result length"); // Verify the entire elements array is preserved including gaps. - let result_elements = result_list.elements().to_primitive(); + let result_elements = result_list.elements().to_primitive()?; assert_eq!( result_elements.as_slice::(), &[1, 2, 3, 999, 999, 999, 7, 8, 9, 999, 11, 12] @@ -102,11 +105,13 @@ fn test_filter_with_gaps() { assert_eq!(list0.scalar_at(0).as_primitive().as_::().unwrap(), 7); assert_eq!(list0.scalar_at(1).as_primitive().as_::().unwrap(), 8); assert_eq!(list0.scalar_at(2).as_primitive().as_::().unwrap(), 9); + + Ok(()) } #[ignore = "TODO(connor)[ListView]: Don't rebuild ListView after every `filter`"] #[test] -fn test_filter_constant_arrays() { +fn test_filter_constant_arrays() -> VortexResult<()> { // ListView-specific: Test filter with ConstantArray for offsets/sizes. let elements = buffer![100i32, 200, 300, 400, 500, 600, 700, 800].into_array(); @@ -125,7 +130,7 @@ fn test_filter_constant_arrays() { let mask1 = Mask::from_iter([true, false, true, false]); let result1 = filter(&const_offset_list, &mask1).unwrap(); - let result1_list = result1.to_listview(); + let result1_list = result1.to_listview()?; assert_eq!(result1_list.len(), 2); assert_eq!(result1_list.offset_at(0), 2); // Both offsets are 2 @@ -148,18 +153,20 @@ fn test_filter_constant_arrays() { let mask2 = Mask::from_iter([true, false, true]); let result2 = filter(&both_const_list, &mask2).unwrap(); - let result2_list = result2.to_listview(); + let result2_list = result2.to_listview()?; assert_eq!(result2_list.len(), 2); assert_eq!(result2_list.offset_at(0), 1); assert_eq!(result2_list.offset_at(1), 1); assert_eq!(result2_list.size_at(0), 3); assert_eq!(result2_list.size_at(1), 3); + + Ok(()) } #[ignore = "TODO(connor)[ListView]: Don't rebuild ListView after every `filter`"] #[test] -fn test_filter_extreme_offsets() { +fn test_filter_extreme_offsets() -> VortexResult<()> { // ListView-specific: Test with very large offsets. let elements = PrimitiveArray::from_iter(0i32..10000).into_array(); @@ -174,7 +181,7 @@ fn test_filter_extreme_offsets() { // Filter to keep only 2 lists, demonstrating we keep all 10000 elements. let mask = Mask::from_iter([false, true, false, false, true]); let result = filter(&listview, &mask).unwrap(); - let result_list = result.to_listview(); + let result_list = result.to_listview()?; assert_eq!(result_list.len(), 2); @@ -199,10 +206,12 @@ fn test_filter_extreme_offsets() { // Test sparse selection from large dataset. let sparse_mask = Mask::from_iter((0..5).map(|i| i == 0 || i == 4)); let sparse_result = filter(&listview, &sparse_mask).unwrap(); - let sparse_list = sparse_result.to_listview(); + let sparse_list = sparse_result.to_listview()?; assert_eq!(sparse_list.len(), 2); assert_eq!(sparse_list.offset_at(0), 0); // First list assert_eq!(sparse_list.offset_at(1), 7500); // Last list assert_eq!(sparse_list.elements().len(), 10000); // Still keeps all elements + + Ok(()) } diff --git a/vortex-array/src/arrays/listview/tests/nullability.rs b/vortex-array/src/arrays/listview/tests/nullability.rs index f4b947f4ca6..fac10c178c4 100644 --- a/vortex-array/src/arrays/listview/tests/nullability.rs +++ b/vortex-array/src/arrays/listview/tests/nullability.rs @@ -8,6 +8,7 @@ use vortex_buffer::buffer; use vortex_dtype::DType; use vortex_dtype::Nullability; use vortex_dtype::PType; +use vortex_error::VortexResult; use vortex_scalar::Scalar; use crate::IntoArray; @@ -17,7 +18,7 @@ use crate::arrays::PrimitiveArray; use crate::validity::Validity; #[test] -fn test_nullable_listview_comprehensive() { +fn test_nullable_listview_comprehensive() -> VortexResult<()> { // Comprehensive test for nullable ListView including scalar_at with nulls. // Logical lists: [[1,2], null, [5,6]] let elements = buffer![1i32, 2, 3, 4, 5, 6].into_array(); @@ -33,9 +34,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)?); + assert!(listview.is_invalid(1)?); + assert!(listview.is_valid(2)?); // Check dtype reflects nullability. assert!(matches!( @@ -74,6 +75,8 @@ fn test_nullable_listview_comprehensive() { 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()); + + Ok(()) } // Parameterized tests for different null patterns. @@ -81,7 +84,10 @@ fn test_nullable_listview_comprehensive() { #[case::all_nulls(Validity::AllInvalid, vec![false, false, false])] #[case::all_valid(Validity::AllValid, vec![true, true, true])] #[case::mixed(Validity::from_iter([false, true, false]), vec![false, true, false])] -fn test_nullable_patterns(#[case] validity: Validity, #[case] expected_validity: Vec) { +fn test_nullable_patterns( + #[case] validity: Validity, + #[case] expected_validity: Vec, +) -> VortexResult<()> { // Logical lists: [[1,2], [3,4], [5,6]] with varying validity let elements = buffer![1i32, 2, 3, 4, 5, 6].into_array(); let offsets = buffer![0i32, 2, 4].into_array(); @@ -90,8 +96,10 @@ 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)?, expected); } + + Ok(()) } #[test] diff --git a/vortex-array/src/arrays/listview/tests/operations.rs b/vortex-array/src/arrays/listview/tests/operations.rs index d0ccd497884..1a859aca4cf 100644 --- a/vortex-array/src/arrays/listview/tests/operations.rs +++ b/vortex-array/src/arrays/listview/tests/operations.rs @@ -8,6 +8,7 @@ use vortex_buffer::buffer; use vortex_dtype::DType; use vortex_dtype::Nullability; use vortex_dtype::PType; +use vortex_error::VortexResult; use vortex_mask::Mask; use super::common::create_basic_listview; @@ -121,7 +122,7 @@ fn test_slice_out_of_order() { } #[test] -fn test_slice_with_nulls() { +fn test_slice_with_nulls() -> VortexResult<()> { // Test slicing with nullable ListView. // Logical lists: [[1,2], null, [5,6], null] let elements = buffer![1i32, 2, 3, 4, 5, 6, 7, 8].into_array(); @@ -141,14 +142,16 @@ 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)?); // Original index 1 was null. + assert!(sliced_list.is_valid(1)?); // Original index 2 was valid. // Verify offsets and sizes are preserved. assert_eq!(sliced_list.offset_at(0), 2); assert_eq!(sliced_list.size_at(0), 2); assert_eq!(sliced_list.offset_at(1), 4); assert_eq!(sliced_list.size_at(1), 2); + + Ok(()) } // Parameterized edge case tests. @@ -194,7 +197,7 @@ fn test_slice_edge_cases( #[case::i32_to_i64(PType::I32, PType::I64)] #[case::f32_to_f64(PType::F32, PType::F64)] #[case::u8_to_u16(PType::U8, PType::U16)] -fn test_cast_numeric_types(#[case] from_ptype: PType, #[case] to_ptype: PType) { +fn test_cast_numeric_types(#[case] from_ptype: PType, #[case] to_ptype: PType) -> VortexResult<()> { let elements = match from_ptype { PType::I32 => buffer![1i32, 2, 3, 4, 5, 6].into_array(), PType::F32 => buffer![1.0f32, 2.0, 3.0, 4.0].into_array(), @@ -229,7 +232,7 @@ fn test_cast_numeric_types(#[case] from_ptype: PType, #[case] to_ptype: PType) { let result = cast(&listview, &target_dtype).unwrap(); assert_eq!(result.dtype(), &target_dtype); - let result_list = result.to_listview(); + let result_list = result.to_listview()?; assert!( result_list.len() == 3 || result_list.len() == 2, "Expected 2 or 3 lists" @@ -241,10 +244,12 @@ fn test_cast_numeric_types(#[case] from_ptype: PType, #[case] to_ptype: PType) { elements.dtype(), &DType::Primitive(to_ptype, Nullability::NonNullable) ); + + Ok(()) } #[test] -fn test_cast_with_nulls() { +fn test_cast_with_nulls() -> VortexResult<()> { // Logical lists: [[10,20], null] let elements = buffer![10i32, 20, 30, 40].into_array(); let offsets = buffer![0u32, 2].into_array(); @@ -265,15 +270,20 @@ fn test_cast_with_nulls() { let result = cast(&listview, &target_dtype).unwrap(); assert_eq!(result.dtype(), &target_dtype); - let result_list = result.to_listview(); - assert!(result_list.is_valid(0)); - assert!(result_list.is_invalid(1)); + let result_list = result.to_listview()?; + assert!(result_list.is_valid(0)?); + assert!(result_list.is_invalid(1)?); + + Ok(()) } #[rstest] #[case::empty_lists(vec![0, 1, 0, 1], 4)] #[case::overlapping(vec![3, 3, 5], 3)] -fn test_cast_special_patterns(#[case] expected_sizes: Vec, #[case] list_count: usize) { +fn test_cast_special_patterns( + #[case] expected_sizes: Vec, + #[case] list_count: usize, +) -> VortexResult<()> { let is_empty_case = list_count == 4; let (elements, offsets, sizes) = if is_empty_case { @@ -307,17 +317,19 @@ fn test_cast_special_patterns(#[case] expected_sizes: Vec, #[case] list_c }; let result = cast(&listview, &target_dtype).unwrap(); - let result_list = result.to_listview(); + let result_list = result.to_listview()?; assert_eq!(result_list.len(), list_count); for (i, expected_size) in expected_sizes.iter().enumerate() { assert_eq!(result_list.size_at(i), *expected_size); } + + Ok(()) } #[test] -fn test_cast_large_dataset() { +fn test_cast_large_dataset() -> VortexResult<()> { // Test with larger data. // Logical lists: [[0..4], [4..8], [8..12], ..., [76..80]] (20 lists of size 4) let elements = buffer![0u16..100].into_array(); @@ -339,12 +351,14 @@ fn test_cast_large_dataset() { ); let result = cast(&listview, &target_dtype).unwrap(); - let result_list = result.to_listview(); + let result_list = result.to_listview()?; assert_eq!(result_list.len(), 20); for i in 0..20 { assert_eq!(result_list.size_at(i), 4); } + + Ok(()) } //////////////////////////////////////////////////////////////////////////////////////////////////// @@ -489,7 +503,7 @@ fn test_mask_listview_conformance(#[case] listview: ListViewArray) { } #[test] -fn test_mask_preserves_structure() { +fn test_mask_preserves_structure() -> VortexResult<()> { // ListView-specific: Verify mask preserves offsets and sizes. // Logical lists: [[1,2], [3,4], [5,6], [7,8]] let elements = buffer![1i32, 2, 3, 4, 5, 6, 7, 8].into_array(); @@ -507,13 +521,13 @@ fn test_mask_preserves_structure() { let result = mask(&listview, &selection).unwrap(); assert_eq!(result.len(), 4); // Length is preserved. - let result_list = result.to_listview(); + 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)?); // Masked. + assert!(result_list.is_valid(1)?); // Not masked. + assert!(!result_list.is_valid(2)?); // Masked. + assert!(!result_list.is_valid(3)?); // Masked. // Offsets and sizes are preserved. assert_eq!(result_list.offset_at(0), 0); @@ -524,10 +538,12 @@ fn test_mask_preserves_structure() { assert_eq!(result_list.size_at(2), 2); assert_eq!(result_list.offset_at(3), 6); assert_eq!(result_list.size_at(3), 2); + + Ok(()) } #[test] -fn test_mask_with_existing_nulls() { +fn test_mask_with_existing_nulls() -> VortexResult<()> { // ListView-specific: Test interaction between existing nulls and mask. // Logical lists: [[10,20], null, [50,60]] let elements = buffer![10i32, 20, 30, 40, 50, 60].into_array(); @@ -544,16 +560,18 @@ fn test_mask_with_existing_nulls() { // Mask additional elements. let selection = Mask::from_iter([false, true, true]); let result = mask(&listview, &selection).unwrap(); - let result_list = result.to_listview(); + 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)?); // 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. + + Ok(()) } #[test] -fn test_mask_with_gaps() { +fn test_mask_with_gaps() -> VortexResult<()> { // ListView-specific: Mask with gaps in elements. // Logical lists: [[1,2], [5,6], [9,10]] (999 values are gaps) let elements = buffer![1i32, 2, 999, 999, 5, 6, 999, 999, 9, 10].into_array(); @@ -564,20 +582,22 @@ fn test_mask_with_gaps() { let selection = Mask::from_iter([true, false, false]); let result = mask(&listview, &selection).unwrap(); - let result_list = result.to_listview(); + 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)?); // Masked + assert!(result_list.is_valid(1)?); // Not masked + assert!(result_list.is_valid(2)?); // Not masked // Offsets and sizes still preserved assert_eq!(result_list.offset_at(1), 4); assert_eq!(result_list.size_at(1), 2); + + Ok(()) } #[test] -fn test_mask_constant_arrays() { +fn test_mask_constant_arrays() -> VortexResult<()> { // ListView-specific: Test mask with ConstantArray offsets/sizes. // Logical lists: [[200,300], [200,300], [200,300]] let elements = buffer![100i32, 200, 300, 400, 500, 600].into_array(); @@ -596,12 +616,12 @@ fn test_mask_constant_arrays() { let selection = Mask::from_iter([false, true, false]); let result = mask(&const_list, &selection).unwrap(); - let result_list = result.to_listview(); + 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)?); + assert!(!result_list.is_valid(1)?); // Masked + assert!(result_list.is_valid(2)?); // All offsets and sizes remain constant assert_eq!(result_list.offset_at(0), 1); @@ -610,4 +630,6 @@ fn test_mask_constant_arrays() { assert_eq!(result_list.size_at(0), 2); assert_eq!(result_list.size_at(1), 2); assert_eq!(result_list.size_at(2), 2); + + Ok(()) } diff --git a/vortex-array/src/arrays/listview/tests/take.rs b/vortex-array/src/arrays/listview/tests/take.rs index f0580423985..ae06e162016 100644 --- a/vortex-array/src/arrays/listview/tests/take.rs +++ b/vortex-array/src/arrays/listview/tests/take.rs @@ -3,6 +3,7 @@ use rstest::rstest; use vortex_buffer::buffer; +use vortex_error::VortexResult; use super::common::create_basic_listview; use super::common::create_empty_lists_listview; @@ -34,7 +35,7 @@ fn test_take_listview_conformance(#[case] listview: ListViewArray) { #[ignore = "TODO(connor)[ListView]: Don't rebuild ListView after every `take`"] #[test] -fn test_take_preserves_unreferenced_elements() { +fn test_take_preserves_unreferenced_elements() -> VortexResult<()> { // ListView-specific: Test that take preserves the entire elements array // even when taking only a subset of lists. let elements = buffer![0i32, 1, 2, 3, 4, 5, 6, 7, 8, 9].into_array(); @@ -47,12 +48,12 @@ fn test_take_preserves_unreferenced_elements() { // Take only 2 lists. let indices = buffer![1u32, 3].into_array(); let result = take(&listview, &indices).unwrap(); - let result_list = result.to_listview(); + let result_list = result.to_listview()?; assert_eq!(result_list.len(), 2); // Verify the entire elements array is preserved. - let result_elements = result_list.elements().to_primitive(); + let result_elements = result_list.elements().to_primitive()?; assert_eq!( result_elements.as_slice::(), &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9] @@ -61,11 +62,13 @@ fn test_take_preserves_unreferenced_elements() { // Verify offsets are preserved. assert_eq!(result_list.offset_at(0), 2); // List 1 assert_eq!(result_list.offset_at(1), 0); // List 3 + + Ok(()) } #[ignore = "TODO(connor)[ListView]: Don't rebuild ListView after every `take`"] #[test] -fn test_take_with_gaps() { +fn test_take_with_gaps() -> VortexResult<()> { // ListView-specific: Test with gaps in elements array. // Elements with gaps (999 values are "gaps" between used ranges). let elements = buffer![1i32, 2, 3, 999, 999, 999, 7, 8, 9, 999, 11, 12].into_array(); @@ -77,10 +80,10 @@ fn test_take_with_gaps() { let indices = buffer![1u32, 3, 4, 2].into_array(); let result = take(&listview, &indices).unwrap(); - let result_list = result.to_listview(); + let result_list = result.to_listview()?; // Verify the entire elements array is preserved including gaps. - let result_elements = result_list.elements().to_primitive(); + let result_elements = result_list.elements().to_primitive()?; assert_eq!( result_elements.as_slice::(), &[1, 2, 3, 999, 999, 999, 7, 8, 9, 999, 11, 12] @@ -91,11 +94,13 @@ fn test_take_with_gaps() { assert_eq!(list0.scalar_at(0).as_primitive().as_::().unwrap(), 7); assert_eq!(list0.scalar_at(1).as_primitive().as_::().unwrap(), 8); assert_eq!(list0.scalar_at(2).as_primitive().as_::().unwrap(), 9); + + Ok(()) } #[ignore = "TODO(connor)[ListView]: Don't rebuild ListView after every `take`"] #[test] -fn test_take_constant_arrays() { +fn test_take_constant_arrays() -> VortexResult<()> { // ListView-specific: Test with ConstantArray for offsets/sizes. let elements = buffer![100i32, 200, 300, 400, 500, 600, 700, 800].into_array(); @@ -113,7 +118,7 @@ fn test_take_constant_arrays() { let indices = buffer![3u32, 0, 2].into_array(); let result = take(&const_offset_list, &indices).unwrap(); - let result_list = result.to_listview(); + let result_list = result.to_listview()?; assert_eq!(result_list.len(), 3); assert_eq!(result_list.offset_at(0), 2); // All offsets are 2 @@ -137,18 +142,20 @@ fn test_take_constant_arrays() { let indices2 = buffer![2u32, 0].into_array(); let result2 = take(&both_const_list, &indices2).unwrap(); - let result2_list = result2.to_listview(); + let result2_list = result2.to_listview()?; assert_eq!(result2_list.len(), 2); assert_eq!(result2_list.offset_at(0), 1); assert_eq!(result2_list.offset_at(1), 1); assert_eq!(result2_list.size_at(0), 3); assert_eq!(result2_list.size_at(1), 3); + + Ok(()) } #[ignore = "TODO(connor)[ListView]: Don't rebuild ListView after every `take`"] #[test] -fn test_take_extreme_offsets() { +fn test_take_extreme_offsets() -> VortexResult<()> { // ListView-specific: Test with very large offsets to demonstrate // that we keep unreferenced elements. let elements = PrimitiveArray::from_iter(0i32..10000).into_array(); @@ -163,7 +170,7 @@ fn test_take_extreme_offsets() { // Take only 2 lists, demonstrating we keep all 10000 elements. let indices = buffer![1u32, 4].into_array(); let result = take(&listview, &indices).unwrap(); - let result_list = result.to_listview(); + let result_list = result.to_listview()?; assert_eq!(result_list.len(), 2); @@ -184,4 +191,6 @@ fn test_take_extreme_offsets() { list0.scalar_at(1).as_primitive().as_::().unwrap(), 5000 ); + + Ok(()) } diff --git a/vortex-array/src/arrays/listview/vtable/canonical.rs b/vortex-array/src/arrays/listview/vtable/canonical.rs index 78cdda5b8bb..ccea3126ccf 100644 --- a/vortex-array/src/arrays/listview/vtable/canonical.rs +++ b/vortex-array/src/arrays/listview/vtable/canonical.rs @@ -1,13 +1,15 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_error::VortexResult; + use crate::Canonical; use crate::arrays::ListViewArray; use crate::arrays::ListViewVTable; use crate::vtable::CanonicalVTable; impl CanonicalVTable for ListViewVTable { - fn canonicalize(array: &ListViewArray) -> Canonical { - Canonical::List(array.clone()) + fn canonicalize(array: &ListViewArray) -> VortexResult { + Ok(Canonical::List(array.clone())) } } diff --git a/vortex-array/src/arrays/listview/vtable/mod.rs b/vortex-array/src/arrays/listview/vtable/mod.rs index e1c62a109c3..95998e4f342 100644 --- a/vortex-array/src/arrays/listview/vtable/mod.rs +++ b/vortex-array/src/arrays/listview/vtable/mod.rs @@ -146,7 +146,7 @@ impl VTable for ListViewVTable { Arc::new(array.elements().batch_execute(ctx)?), array.offsets().batch_execute(ctx)?.into_primitive(), array.sizes().batch_execute(ctx)?.into_primitive(), - array.validity_mask(), + array.validity_mask()?, ) } .into()) diff --git a/vortex-array/src/arrays/masked/array.rs b/vortex-array/src/arrays/masked/array.rs index 08d8314e7b2..a375f85f3a0 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"); } @@ -52,7 +52,7 @@ impl MaskedArray { pub(crate) fn masked_child(&self) -> VortexResult { // Invert the validity mask - we want to set values to null where validity is false. - let inverted_mask = !self.validity.to_mask(self.len()); + let inverted_mask = !self.validity.to_mask(self.len())?; mask(&self.child, &inverted_mask) } } diff --git a/vortex-array/src/arrays/masked/compute/compare.rs b/vortex-array/src/arrays/masked/compute/compare.rs index 3e73cd7926f..d4d455dec4f 100644 --- a/vortex-array/src/arrays/masked/compute/compare.rs +++ b/vortex-array/src/arrays/masked/compute/compare.rs @@ -28,8 +28,8 @@ impl CompareKernel for MaskedVTable { let compare_result = compare(&lhs.child, rhs, operator)?; // Get the boolean buffer from the comparison result - let bool_array = compare_result.to_bool(); - let combined_validity = bool_array.validity().clone().and(lhs.validity().clone()); + let bool_array = compare_result.to_bool()?; + let combined_validity = bool_array.validity().clone().and(lhs.validity().clone())?; // Return a plain BoolArray with the combined validity Ok(Some( @@ -70,7 +70,7 @@ mod tests { Operator::Eq, ) .unwrap(); - let res = res.to_bool(); + let res = res.to_bool().unwrap(); assert_eq!( res.bit_buffer().iter().collect::>(), vec![false, true, false] @@ -91,7 +91,7 @@ mod tests { Operator::Gt, ) .unwrap(); - let res = res.to_bool(); + let res = res.to_bool().unwrap(); assert_eq!( res.bit_buffer().iter().collect::>(), vec![false, false, true] @@ -113,13 +113,16 @@ mod tests { Operator::Eq, ) .unwrap(); - let res = res.to_bool(); + let res = res.to_bool().unwrap(); assert_eq!( res.bit_buffer().iter().collect::>(), 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] @@ -135,13 +138,16 @@ mod tests { let rhs = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]); let res = compare(masked.as_ref(), rhs.as_ref(), Operator::Eq).unwrap(); - let res = res.to_bool(); + let res = res.to_bool().unwrap(); assert_eq!( res.bit_buffer().iter().collect::>(), vec![true, false, true] ); 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/mask.rs b/vortex-array/src/arrays/masked/compute/mask.rs index a8d3c3eacd2..b679e9dc89a 100644 --- a/vortex-array/src/arrays/masked/compute/mask.rs +++ b/vortex-array/src/arrays/masked/compute/mask.rs @@ -17,7 +17,7 @@ impl MaskKernel for MaskedVTable { fn mask(&self, array: &MaskedArray, mask_arg: &MaskType) -> VortexResult { // Combine the mask with the existing validity // The child remains unchanged (no nulls), only validity is updated - let combined_validity = array.validity().mask(mask_arg); + let combined_validity = array.validity().mask(mask_arg)?; Ok(MaskedArray::try_new(array.child.clone(), combined_validity)?.into_array()) } 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 311efb254ce..5119d8e0ac6 100644 --- a/vortex-array/src/arrays/masked/tests.rs +++ b/vortex-array/src/arrays/masked/tests.rs @@ -4,6 +4,7 @@ use rstest::rstest; use vortex_dtype::DType; use vortex_dtype::Nullability; +use vortex_error::VortexResult; use super::*; use crate::Array; @@ -37,52 +38,55 @@ fn test_dtype_nullability_with_nullable_child() { } #[test] -fn test_canonical_dtype_matches_array_dtype() { +fn test_canonical_dtype_matches_array_dtype() -> VortexResult<()> { // The canonical form should have the same nullability as the array's dtype. let child = PrimitiveArray::from_iter([1i32, 2, 3]).into_array(); let array = MaskedArray::try_new(child, Validity::AllValid).unwrap(); - let canonical = array.to_canonical(); + let canonical = array.to_canonical()?; assert_eq!(canonical.as_ref().dtype(), array.dtype()); + Ok(()) } #[test] -fn test_masked_child_with_validity() { +fn test_masked_child_with_validity() -> VortexResult<()> { // When validity has nulls, masked_child should apply inverted mask. let child = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]).into_array(); let array = MaskedArray::try_new(child, Validity::from_iter([true, false, true, false, true])).unwrap(); let masked = array.masked_child().unwrap(); - let prim = masked.to_primitive(); + let prim = masked.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()?, 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!( array.as_ref().display_values().to_string(), masked.display_values().to_string() ); + Ok(()) } #[test] -fn test_masked_child_all_valid() { +fn test_masked_child_all_valid() -> VortexResult<()> { // When validity is AllValid, masked_child should invert to AllInvalid. let child = PrimitiveArray::from_iter([10i32, 20, 30]).into_array(); let array = MaskedArray::try_new(child, Validity::AllValid).unwrap(); let masked = array.masked_child().unwrap(); assert_eq!(masked.len(), 3); - assert_eq!(masked.valid_count(), 3); + assert_eq!(masked.valid_count()?, 3); assert_eq!( array.as_ref().display_values().to_string(), masked.display_values().to_string() ); + Ok(()) } #[rstest] @@ -90,7 +94,7 @@ fn test_masked_child_all_valid() { #[case(Validity::from_iter([true, true, true]))] #[case(Validity::from_iter([false, false, false]))] #[case(Validity::from_iter([true, false, true, false]))] -fn test_masked_child_preserves_length(#[case] validity: Validity) { +fn test_masked_child_preserves_length(#[case] validity: Validity) -> VortexResult<()> { let len = match &validity { Validity::Array(arr) => arr.len(), _ => 3, @@ -102,9 +106,10 @@ fn test_masked_child_preserves_length(#[case] validity: Validity) { let masked = array.masked_child().unwrap(); assert_eq!(masked.len(), len); - assert_eq!(masked.validity_mask(), validity.to_mask(len)); + assert_eq!(masked.validity_mask()?, validity.to_mask(len)?); assert_eq!( array.as_ref().display_values().to_string(), masked.display_values().to_string() ); + Ok(()) } diff --git a/vortex-array/src/arrays/masked/vtable/canonical.rs b/vortex-array/src/arrays/masked/vtable/canonical.rs index 4492eb99dce..611fbb999b5 100644 --- a/vortex-array/src/arrays/masked/vtable/canonical.rs +++ b/vortex-array/src/arrays/masked/vtable/canonical.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 crate::Array; use crate::Canonical; @@ -12,20 +12,16 @@ use crate::compute::mask; use crate::vtable::CanonicalVTable; impl CanonicalVTable for MaskedVTable { - fn canonicalize(array: &MaskedArray) -> Canonical { + fn canonicalize(array: &MaskedArray) -> VortexResult { if array.child.is::() { // To allow constant array to produce masked array from mask call, we have to unwrap constant here and canonicalize it first mask( - array.child.to_canonical().as_ref(), - &!array.validity.to_mask(array.len()), - ) - .vortex_expect("constant masked to canonical") + array.child.to_canonical()?.as_ref(), + &!array.validity.to_mask(array.len())?, + )? .to_canonical() } else { - array - .masked_child() - .vortex_expect("masked child of a masked array") - .to_canonical() + array.masked_child()?.to_canonical() } } } @@ -60,7 +56,7 @@ mod tests { #[case] array: MaskedArray, #[case] expected_nullability: Nullability, ) { - let canonical = array.to_canonical(); + let canonical = array.to_canonical().unwrap(); assert_eq!( canonical.as_ref().dtype().nullability(), expected_nullability @@ -76,16 +72,16 @@ mod tests { ) .unwrap(); - let canonical = array.to_canonical(); - let prim = canonical.as_ref().to_primitive(); + let canonical = array.to_canonical().unwrap(); + let prim = canonical.as_ref().to_primitive().unwrap(); // 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()); } #[test] @@ -96,8 +92,8 @@ mod tests { ) .unwrap(); - let canonical = array.to_canonical(); - assert_eq!(canonical.as_ref().valid_count(), 3); + let canonical = array.to_canonical().unwrap(); + 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/mod.rs b/vortex-array/src/arrays/masked/vtable/mod.rs index f35cfe042f4..a0acc40d1c7 100644 --- a/vortex-array/src/arrays/masked/vtable/mod.rs +++ b/vortex-array/src/arrays/masked/vtable/mod.rs @@ -107,7 +107,7 @@ impl VTable for MaskedVTable { fn batch_execute(array: &Self::Array, ctx: &mut ExecutionCtx) -> VortexResult { let vector = array.child().batch_execute(ctx)?; - Ok(MaskValidity::mask_validity(vector, &array.validity_mask())) + Ok(MaskValidity::mask_validity(vector, &array.validity_mask()?)) } } diff --git a/vortex-array/src/arrays/null/compute/mod.rs b/vortex-array/src/arrays/null/compute/mod.rs index 237a558738b..b95544698ac 100644 --- a/vortex-array/src/arrays/null/compute/mod.rs +++ b/vortex-array/src/arrays/null/compute/mod.rs @@ -26,10 +26,10 @@ mod test { #[test] fn test_slice_nulls() { let nulls = NullArray::new(10); - let sliced = nulls.slice(0..4).to_null(); + let sliced = nulls.slice(0..4).to_null().unwrap(); assert_eq!(sliced.len(), 4); - assert!(matches!(sliced.validity_mask(), Mask::AllFalse(4))); + assert!(matches!(sliced.validity_mask().unwrap(), Mask::AllFalse(4))); } #[test] @@ -37,10 +37,11 @@ mod test { let nulls = NullArray::new(10); let taken = take(nulls.as_ref(), &buffer![0u64, 2, 4, 6, 8].into_array()) .unwrap() - .to_null(); + .to_null() + .unwrap(); assert_eq!(taken.len(), 5); - assert!(matches!(taken.validity_mask(), Mask::AllFalse(5))); + assert!(matches!(taken.validity_mask().unwrap(), Mask::AllFalse(5))); } #[test] diff --git a/vortex-array/src/arrays/null/compute/take.rs b/vortex-array/src/arrays/null/compute/take.rs index 4d3d595bd7d..53ce41ca298 100644 --- a/vortex-array/src/arrays/null/compute/take.rs +++ b/vortex-array/src/arrays/null/compute/take.rs @@ -18,7 +18,7 @@ use crate::register_kernel; impl TakeKernel for NullVTable { #[allow(clippy::cast_possible_truncation)] fn take(&self, array: &NullArray, indices: &dyn Array) -> VortexResult { - let indices = indices.to_primitive(); + let indices = indices.to_primitive()?; // Enforce all indices are valid match_each_integer_ptype!(indices.ptype(), |T| { diff --git a/vortex-array/src/arrays/null/mod.rs b/vortex-array/src/arrays/null/mod.rs index 3aeb4508d87..f5bf0432c3c 100644 --- a/vortex-array/src/arrays/null/mod.rs +++ b/vortex-array/src/arrays/null/mod.rs @@ -159,8 +159,8 @@ impl VisitorVTable for NullVTable { } impl CanonicalVTable for NullVTable { - fn canonicalize(array: &NullArray) -> Canonical { - Canonical::Null(array.clone()) + fn canonicalize(array: &NullArray) -> VortexResult { + Ok(Canonical::Null(array.clone())) } } @@ -175,19 +175,19 @@ impl OperationsVTable for NullVTable { } impl ValidityVTable for NullVTable { - fn is_valid(_array: &NullArray, _index: usize) -> bool { - false + fn is_valid(_array: &NullArray, _index: usize) -> VortexResult { + Ok(false) } - fn all_valid(array: &NullArray) -> bool { - array.is_empty() + fn all_valid(array: &NullArray) -> VortexResult { + Ok(array.is_empty()) } - fn all_invalid(array: &NullArray) -> bool { - !array.is_empty() + fn all_invalid(array: &NullArray) -> VortexResult { + Ok(!array.is_empty()) } - 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/accessor.rs b/vortex-array/src/arrays/primitive/array/accessor.rs index d3f7dde0161..bc5bf91da0b 100644 --- a/vortex-array/src/arrays/primitive/array/accessor.rs +++ b/vortex-array/src/arrays/primitive/array/accessor.rs @@ -4,6 +4,7 @@ use std::iter; use vortex_dtype::NativePType; +use vortex_error::VortexExpect; use crate::ToCanonical; use crate::accessor::ArrayAccessor; @@ -23,7 +24,7 @@ impl ArrayAccessor for PrimitiveArray { } Validity::AllInvalid => f(&mut iter::repeat_n(None, self.len())), Validity::Array(v) => { - let validity = v.to_bool(); + let validity = v.to_bool().vortex_expect("to_bool"); let mut iter = self .as_slice::() .iter() diff --git a/vortex-array/src/arrays/primitive/array/cast.rs b/vortex-array/src/arrays/primitive/array/cast.rs index c559f0e7a8d..6e5c3afd605 100644 --- a/vortex-array/src/arrays/primitive/array/cast.rs +++ b/vortex-array/src/arrays/primitive/array/cast.rs @@ -72,52 +72,52 @@ impl PrimitiveArray { if min < 0 || max < 0 { // Signed if min >= i8::MIN as i64 && max <= i8::MAX as i64 { - return Ok(cast( + return cast( self.as_ref(), &DType::Primitive(PType::I8, self.dtype().nullability()), )? - .to_primitive()); + .to_primitive(); } if min >= i16::MIN as i64 && max <= i16::MAX as i64 { - return Ok(cast( + return cast( self.as_ref(), &DType::Primitive(PType::I16, self.dtype().nullability()), )? - .to_primitive()); + .to_primitive(); } if min >= i32::MIN as i64 && max <= i32::MAX as i64 { - return Ok(cast( + return cast( self.as_ref(), &DType::Primitive(PType::I32, self.dtype().nullability()), )? - .to_primitive()); + .to_primitive(); } } else { // Unsigned if max <= u8::MAX as i64 { - return Ok(cast( + return cast( self.as_ref(), &DType::Primitive(PType::U8, self.dtype().nullability()), )? - .to_primitive()); + .to_primitive(); } if max <= u16::MAX as i64 { - return Ok(cast( + return cast( self.as_ref(), &DType::Primitive(PType::U16, self.dtype().nullability()), )? - .to_primitive()); + .to_primitive(); } if max <= u32::MAX as i64 { - return Ok(cast( + return cast( self.as_ref(), &DType::Primitive(PType::U32, self.dtype().nullability()), )? - .to_primitive()); + .to_primitive(); } } diff --git a/vortex-array/src/arrays/primitive/array/conversion.rs b/vortex-array/src/arrays/primitive/array/conversion.rs index 50db4459edc..698f3b3fe96 100644 --- a/vortex-array/src/arrays/primitive/array/conversion.rs +++ b/vortex-array/src/arrays/primitive/array/conversion.rs @@ -176,11 +176,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 2443fa46abf..e423c1ed387 100644 --- a/vortex-array/src/arrays/primitive/array/mod.rs +++ b/vortex-array/src/arrays/primitive/array/mod.rs @@ -176,7 +176,9 @@ impl PrimitiveArray { Validity::AllValid | Validity::NonNullable => valid_elems_buffer.aligned(alignment), Validity::AllInvalid => ByteBuffer::zeroed_aligned(n_rows * byte_width, alignment), Validity::Array(is_valid) => { - let bool_array = is_valid.to_bool(); + let bool_array = is_valid + .to_bool() + .vortex_expect("Failed to convert to bool"); let bool_buffer = bool_array.bit_buffer(); let mut bytes = ByteBufferMut::zeroed_aligned(n_rows * byte_width, alignment); for (i, valid_i) in bool_buffer.set_indices().enumerate() { @@ -232,7 +234,7 @@ impl PrimitiveArray { BufferMut::::from_iter(buf_iter.zip(iter::repeat(false)).map(f)) } Validity::Array(val) => { - let val = val.to_bool(); + let val = val.to_bool()?; BufferMut::::from_iter(buf_iter.zip(val.bit_buffer()).map(f)) } }; diff --git a/vortex-array/src/arrays/primitive/array/patch.rs b/vortex-array/src/arrays/primitive/array/patch.rs index b9524ed8c76..102cc603a2d 100644 --- a/vortex-array/src/arrays/primitive/array/patch.rs +++ b/vortex-array/src/arrays/primitive/array/patch.rs @@ -6,6 +6,7 @@ use vortex_dtype::NativePType; use vortex_dtype::UnsignedPType; use vortex_dtype::match_each_integer_ptype; use vortex_dtype::match_each_native_ptype; +use vortex_error::VortexExpect; use crate::ToCanonical; use crate::arrays::PrimitiveArray; @@ -15,15 +16,25 @@ use crate::vtable::ValidityHelper; impl PrimitiveArray { pub fn patch(self, patches: &Patches) -> Self { - let patch_indices = patches.indices().to_primitive(); - let patch_values = patches.values().to_primitive(); + let patch_indices = patches + .indices() + .to_primitive() + .vortex_expect("patch indices must be primitive"); + let patch_values = patches + .values() + .to_primitive() + .vortex_expect("patch values must be primitive"); - let patched_validity = self.validity().clone().patch( - self.len(), - patches.offset(), - patch_indices.as_ref(), - patch_values.validity(), - ); + let patched_validity = self + .validity() + .clone() + .patch( + self.len(), + patches.offset(), + patch_indices.as_ref(), + patch_values.validity(), + ) + .vortex_expect("patch validity"); match_each_integer_ptype!(patch_indices.ptype(), |I| { match_each_native_ptype!(self.ptype(), |T| { self.patch_typed::( @@ -119,6 +130,6 @@ mod tests { fn patch_sliced() { let input = PrimitiveArray::new(buffer![2u32; 10], Validity::AllValid); let sliced = input.slice(2..8); - assert_eq!(sliced.to_primitive().as_slice::(), &[2u32; 6]); + assert_eq!(sliced.to_primitive().unwrap().as_slice::(), &[2u32; 6]); } } 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 68927d4060f..abd1ad077ac 100644 --- a/vortex-array/src/arrays/primitive/compute/cast.rs +++ b/vortex-array/src/arrays/primitive/compute/cast.rs @@ -45,7 +45,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| { @@ -117,7 +117,10 @@ mod test { let arr = buffer![0u32, 10, 200].into_array(); // cast from u32 to u8 - let p = cast(&arr, PType::U8.into()).unwrap().to_primitive(); + let p = cast(&arr, PType::U8.into()) + .unwrap() + .to_primitive() + .unwrap(); assert_eq!(p.as_slice::(), vec![0u8, 10, 200]); assert_eq!(p.validity(), &Validity::NonNullable); @@ -127,7 +130,8 @@ mod test { &DType::Primitive(PType::U8, Nullability::Nullable), ) .unwrap() - .to_primitive(); + .to_primitive() + .unwrap(); assert_eq!(p.as_slice::(), vec![0u8, 10, 200]); assert_eq!(p.validity(), &Validity::AllValid); @@ -137,7 +141,8 @@ mod test { &DType::Primitive(PType::U8, Nullability::NonNullable), ) .unwrap() - .to_primitive(); + .to_primitive() + .unwrap(); assert_eq!(p.as_slice::(), vec![0u8, 10, 200]); assert_eq!(p.validity(), &Validity::NonNullable); @@ -147,7 +152,8 @@ mod test { &DType::Primitive(PType::U32, Nullability::Nullable), ) .unwrap() - .to_primitive(); + .to_primitive() + .unwrap(); assert_eq!(p.as_slice::(), vec![0u32, 10, 200]); assert_eq!(p.validity(), &Validity::AllValid); @@ -157,7 +163,8 @@ mod test { &DType::Primitive(PType::U8, Nullability::NonNullable), ) .unwrap() - .to_primitive(); + .to_primitive() + .unwrap(); assert_eq!(p.as_slice::(), vec![0u8, 10, 200]); assert_eq!(p.validity(), &Validity::NonNullable); } @@ -165,7 +172,10 @@ mod test { #[test] fn cast_u32_f32() { let arr = buffer![0u32, 10, 200].into_array(); - let u8arr = cast(&arr, PType::F32.into()).unwrap().to_primitive(); + let u8arr = cast(&arr, PType::F32.into()) + .unwrap() + .to_primitive() + .unwrap(); assert_eq!(u8arr.as_slice::(), vec![0.0f32, 10., 200.]); } @@ -203,10 +213,11 @@ mod test { &DType::Primitive(PType::U32, Nullability::Nullable), ) .unwrap() - .to_primitive(); + .to_primitive() + .unwrap(); assert_eq!(p.as_slice::(), vec![0, 0, 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 3036c69c153..efc5badc069 100644 --- a/vortex-array/src/arrays/primitive/compute/fill_null.rs +++ b/vortex-array/src/arrays/primitive/compute/fill_null.rs @@ -25,7 +25,7 @@ impl FillNullKernel for PrimitiveVTable { Ok(match array.validity() { Validity::Array(is_valid) => { - let is_invalid = is_valid.to_bool().bit_buffer().not(); + let is_invalid = is_valid.to_bool()?.bit_buffer().not(); match_each_native_ptype!(array.ptype(), |T| { let mut buffer = array.buffer::().into_mut(); let fill_value = fill_value @@ -62,9 +62,10 @@ mod test { let arr = PrimitiveArray::from_option_iter([None, Some(8u8), None, Some(10), None]); let p = fill_null(arr.as_ref(), &Scalar::from(42u8)) .unwrap() - .to_primitive(); + .to_primitive() + .unwrap(); assert_eq!(p.as_slice::(), vec![42, 8, 42, 10, 42]); - assert!(p.validity_mask().all_true()); + assert!(p.validity_mask().unwrap().all_true()); } #[test] @@ -73,9 +74,10 @@ mod test { let p = fill_null(arr.as_ref(), &Scalar::from(255u8)) .unwrap() - .to_primitive(); + .to_primitive() + .unwrap(); assert_eq!(p.as_slice::(), vec![255, 255, 255, 255, 255]); - assert!(p.validity_mask().all_true()); + assert!(p.validity_mask().unwrap().all_true()); } #[test] @@ -86,9 +88,10 @@ mod test { ); let p = fill_null(arr.as_ref(), &Scalar::from(255u8)) .unwrap() - .to_primitive(); + .to_primitive() + .unwrap(); assert_eq!(p.as_slice::(), vec![8, 10, 12, 14, 16]); - assert!(p.validity_mask().all_true()); + assert!(p.validity_mask().unwrap().all_true()); } #[test] @@ -96,8 +99,9 @@ mod test { let arr = buffer![8u8, 10, 12, 14, 16].into_array(); let p = fill_null(&arr, &Scalar::from(255u8)) .unwrap() - .to_primitive(); + .to_primitive() + .unwrap(); assert_eq!(p.as_slice::(), vec![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/filter.rs b/vortex-array/src/arrays/primitive/compute/filter.rs index 1ca069a0df7..ee40a2b8b1f 100644 --- a/vortex-array/src/arrays/primitive/compute/filter.rs +++ b/vortex-array/src/arrays/primitive/compute/filter.rs @@ -97,7 +97,8 @@ mod test { let filtered = filter(arr.as_ref(), &Mask::from_iter(mask)) .unwrap() - .to_primitive(); + .to_primitive() + .unwrap(); assert_eq!( filtered.len(), mask.iter().filter(|x| **x).collect_vec().len() diff --git a/vortex-array/src/arrays/primitive/compute/is_sorted.rs b/vortex-array/src/arrays/primitive/compute/is_sorted.rs index 82c8d804547..0d24bf1a56a 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/mask.rs b/vortex-array/src/arrays/primitive/compute/mask.rs index 8a701ce72a8..eae8f488788 100644 --- a/vortex-array/src/arrays/primitive/compute/mask.rs +++ b/vortex-array/src/arrays/primitive/compute/mask.rs @@ -15,7 +15,7 @@ use crate::vtable::ValidityHelper; impl MaskKernel for PrimitiveVTable { fn mask(&self, array: &PrimitiveArray, mask: &Mask) -> VortexResult { - let validity = array.validity().mask(mask); + let validity = array.validity().mask(mask)?; Ok( PrimitiveArray::from_byte_buffer(array.byte_buffer().clone(), array.ptype(), validity) .into_array(), 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..6fff7cd6d21 100644 --- a/vortex-array/src/arrays/primitive/compute/take/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/take/mod.rs @@ -88,10 +88,10 @@ impl TakeKernel for PrimitiveVTable { }; let unsigned_indices = if ptype.is_unsigned_int() { - indices.to_primitive() + indices.to_primitive()? } else { // This will fail if all values cannot be converted to unsigned - cast(indices, &DType::Primitive(ptype.to_unsigned(), *null))?.to_primitive() + cast(indices, &DType::Primitive(ptype.to_unsigned(), *null))?.to_primitive()? }; let validity = array.validity().take(unsigned_indices.as_ref())?; diff --git a/vortex-array/src/arrays/primitive/vtable/canonical.rs b/vortex-array/src/arrays/primitive/vtable/canonical.rs index 4c657bc8a12..5ac34804e9f 100644 --- a/vortex-array/src/arrays/primitive/vtable/canonical.rs +++ b/vortex-array/src/arrays/primitive/vtable/canonical.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_error::VortexResult; + use crate::Canonical; use crate::arrays::PrimitiveArray; use crate::arrays::PrimitiveVTable; @@ -8,11 +10,15 @@ use crate::builders::ArrayBuilder; use crate::vtable::CanonicalVTable; impl CanonicalVTable for PrimitiveVTable { - fn canonicalize(array: &PrimitiveArray) -> Canonical { - Canonical::Primitive(array.clone()) + fn canonicalize(array: &PrimitiveArray) -> VortexResult { + Ok(Canonical::Primitive(array.clone())) } - fn append_to_builder(array: &PrimitiveArray, builder: &mut dyn ArrayBuilder) { - builder.extend_from_array(array.as_ref()) + fn append_to_builder( + array: &PrimitiveArray, + builder: &mut dyn ArrayBuilder, + ) -> VortexResult<()> { + builder.extend_from_array(array.as_ref()); + Ok(()) } } diff --git a/vortex-array/src/arrays/primitive/vtable/mod.rs b/vortex-array/src/arrays/primitive/vtable/mod.rs index 986bc159007..f010898281b 100644 --- a/vortex-array/src/arrays/primitive/vtable/mod.rs +++ b/vortex-array/src/arrays/primitive/vtable/mod.rs @@ -118,7 +118,7 @@ impl VTable for PrimitiveVTable { fn batch_execute(array: &Self::Array, _ctx: &mut ExecutionCtx) -> VortexResult { Ok(match_each_native_ptype!(array.ptype(), |T| { - PVector::new(array.buffer::(), array.validity_mask()).into() + PVector::new(array.buffer::(), array.validity_mask()?).into() })) } } diff --git a/vortex-array/src/arrays/primitive/vtable/operator.rs b/vortex-array/src/arrays/primitive/vtable/operator.rs index 818a4991a91..fea5d0f60fa 100644 --- a/vortex-array/src/arrays/primitive/vtable/operator.rs +++ b/vortex-array/src/arrays/primitive/vtable/operator.rs @@ -47,7 +47,7 @@ impl ArrayParentReduceRule, Exact> unsafe { PrimitiveArray::new_unchecked( Buffer::::from_byte_buffer(array.byte_buffer().clone()), - array.validity().clone().and(parent.validity().clone()), + array.validity().clone().and(parent.validity().clone())?, ) } .into_array() diff --git a/vortex-array/src/arrays/scalar_fn/vtable/canonical.rs b/vortex-array/src/arrays/scalar_fn/vtable/canonical.rs index f1222998ab9..b9d7a8ac532 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/canonical.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/canonical.rs @@ -2,7 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use itertools::Itertools; -use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_err; use vortex_vector::Datum; use crate::Array; @@ -15,27 +16,22 @@ use crate::vectors::VectorIntoArray; use crate::vtable::CanonicalVTable; impl CanonicalVTable for ScalarFnVTable { - fn canonicalize(array: &ScalarFnArray) -> Canonical { + fn canonicalize(array: &ScalarFnArray) -> VortexResult { let child_dtypes: Vec<_> = array.children.iter().map(|c| c.dtype().clone()).collect(); let child_datums: Vec<_> = array .children() .iter() // TODO(ngates): we could make all execution operate over datums .map(|child| child.execute(&SCALAR_FN_SESSION).map(Datum::Vector)) - .try_collect() - // FIXME(ngates): canonicalizing really ought to be fallible - .vortex_expect( - "Failed to execute child array during canonicalization of ScalarFnArray", - ); + .try_collect()?; let ctx = ExecutionArgs::new(array.len, array.dtype.clone(), child_dtypes, child_datums); let result_vector = array .scalar_fn - .execute(&ctx) - .vortex_expect("Canonicalize should be fallible") + .execute(&ctx)? .into_vector() - .vortex_expect("Canonicalize should return a vector"); + .ok_or_else(|| vortex_err!("Canonicalize should return a vector"))?; result_vector.into_array(&array.dtype).to_canonical() } diff --git a/vortex-array/src/arrays/scalar_fn/vtable/validity.rs b/vortex-array/src/arrays/scalar_fn/vtable/validity.rs index eb3fada3d50..1f136d8075d 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/validity.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/validity.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_mask::Mask; use crate::Array; @@ -12,40 +12,48 @@ use crate::expr::functions::NullHandling; use crate::vtable::ValidityVTable; impl ValidityVTable for ScalarFnVTable { - fn is_valid(array: &ScalarFnArray, index: usize) -> bool { - array.scalar_at(index).is_valid() + fn is_valid(array: &ScalarFnArray, index: usize) -> VortexResult { + Ok(array.scalar_at(index).is_valid()) } - fn all_valid(array: &ScalarFnArray) -> bool { + fn all_valid(array: &ScalarFnArray) -> VortexResult { match array.scalar_fn.signature().null_handling() { NullHandling::Propagate | NullHandling::AbsorbsNull => { // Requires all children to guarantee all_valid - array.children().iter().all(|child| child.all_valid()) + for child in array.children().iter() { + if !child.all_valid()? { + return Ok(false); + } + } + Ok(true) } NullHandling::Custom => { // We cannot guarantee that the array is all valid without evaluating the function - false + Ok(false) } } } - fn all_invalid(array: &ScalarFnArray) -> bool { + fn all_invalid(array: &ScalarFnArray) -> VortexResult { match array.scalar_fn.signature().null_handling() { NullHandling::Propagate => { // All null if any child is all null - array.children().iter().any(|child| child.all_invalid()) + for child in array.children().iter() { + if child.all_invalid()? { + return Ok(true); + } + } + Ok(false) } NullHandling::AbsorbsNull | NullHandling::Custom => { // We cannot guarantee that the array is all valid without evaluating the function - false + Ok(false) } } } - fn validity_mask(array: &ScalarFnArray) -> Mask { - let vector = array - .execute(&SCALAR_FN_SESSION) - .vortex_expect("Validity mask computation should be fallible"); - Mask::from_buffer(vector.into_bool().into_bits()) + fn validity_mask(array: &ScalarFnArray) -> VortexResult { + let vector = array.execute(&SCALAR_FN_SESSION)?; + Ok(Mask::from_buffer(vector.into_bool().into_bits())) } } diff --git a/vortex-array/src/arrays/struct_/compute/mask.rs b/vortex-array/src/arrays/struct_/compute/mask.rs index fa8b8a2183b..2d5a5519e06 100644 --- a/vortex-array/src/arrays/struct_/compute/mask.rs +++ b/vortex-array/src/arrays/struct_/compute/mask.rs @@ -15,7 +15,7 @@ use crate::vtable::ValidityHelper; impl MaskKernel for StructVTable { fn mask(&self, array: &StructArray, filter_mask: &Mask) -> VortexResult { - let validity = array.validity().mask(filter_mask); + let validity = array.validity().mask(filter_mask)?; StructArray::try_new_with_dtype( array.fields().clone(), diff --git a/vortex-array/src/arrays/struct_/compute/zip.rs b/vortex-array/src/arrays/struct_/compute/zip.rs index db419c5e923..53bd9fdf7a3 100644 --- a/vortex-array/src/arrays/struct_/compute/zip.rs +++ b/vortex-array/src/arrays/struct_/compute/zip.rs @@ -48,8 +48,8 @@ impl ZipKernel for StructVTable { (&Validity::AllInvalid, &Validity::AllInvalid) => Validity::AllInvalid, (v1, v2) => { - let v1m = v1.to_mask(if_true.len()); - let v2m = v2.to_mask(if_false.len()); + let v1m = v1.to_mask(if_true.len())?; + let v2m = v2.to_mask(if_false.len())?; let combined = (v1m.bitand(mask)).bitor(&v2m.bitand(&mask.not())); Validity::from_mask( diff --git a/vortex-array/src/arrays/struct_/tests.rs b/vortex-array/src/arrays/struct_/tests.rs index d768adff136..b7bccad4e33 100644 --- a/vortex-array/src/arrays/struct_/tests.rs +++ b/vortex-array/src/arrays/struct_/tests.rs @@ -19,7 +19,7 @@ use crate::arrays::varbin::VarBinArray; use crate::validity::Validity; #[test] -fn test_project() { +fn test_project() -> vortex_error::VortexResult<()> { let xs = PrimitiveArray::new(buffer![0i64, 1, 2, 3, 4], Validity::NonNullable); let ys = VarBinArray::from_vec( vec!["a", "b", "c", "d", "e"], @@ -47,16 +47,17 @@ fn test_project() { let bools = &struct_b.fields[0]; assert_eq!( - bools.to_bool().bit_buffer().iter().collect::>(), + bools.to_bool()?.bit_buffer().iter().collect::>(), vec![true, true, true, false, false] ); let prims = &struct_b.fields[1]; - assert_eq!(prims.to_primitive().as_slice::(), [0i64, 1, 2, 3, 4]); + assert_eq!(prims.to_primitive()?.as_slice::(), [0i64, 1, 2, 3, 4]); + Ok(()) } #[test] -fn test_remove_column() { +fn test_remove_column() -> vortex_error::VortexResult<()> { let xs = PrimitiveArray::new(buffer![0i64, 1, 2, 3, 4], Validity::NonNullable); let ys = PrimitiveArray::new(buffer![4u64, 5, 6, 7, 8], Validity::NonNullable); @@ -73,7 +74,10 @@ fn test_remove_column() { removed.dtype(), &DType::Primitive(PType::I64, Nullability::NonNullable) ); - assert_eq!(removed.to_primitive().as_slice::(), [0i64, 1, 2, 3, 4]); + assert_eq!( + removed.to_primitive()?.as_slice::(), + [0i64, 1, 2, 3, 4] + ); assert_eq!(struct_a.names(), &["ys"]); assert_eq!(struct_a.fields.len(), 1); @@ -83,7 +87,7 @@ fn test_remove_column() { &DType::Primitive(PType::U64, Nullability::NonNullable) ); assert_eq!( - struct_a.fields[0].to_primitive().as_slice::(), + struct_a.fields[0].to_primitive()?.as_slice::(), [4u64, 5, 6, 7, 8] ); @@ -93,10 +97,11 @@ fn test_remove_column() { "Expected None when removing non-existent column" ); assert_eq!(struct_a.names(), &["ys"]); + Ok(()) } #[test] -fn test_duplicate_field_names() { +fn test_duplicate_field_names() -> vortex_error::VortexResult<()> { // Test that StructArray allows duplicate field names and returns the first match let field1 = buffer![1i32, 2, 3].into_array(); let field2 = buffer![10i32, 20, 30].into_array(); @@ -114,27 +119,28 @@ fn test_duplicate_field_names() { // field_by_name should return the first field with the matching name let first_value_field = struct_array.field_by_name("value").unwrap(); assert_eq!( - first_value_field.to_primitive().as_slice::(), + first_value_field.to_primitive()?.as_slice::(), [1i32, 2, 3] // This is field1, not field3 ); // Verify field_by_name_opt also returns the first match let opt_field = struct_array.field_by_name_opt("value").unwrap(); assert_eq!( - opt_field.to_primitive().as_slice::(), + opt_field.to_primitive()?.as_slice::(), [1i32, 2, 3] // First "value" field ); // Verify the third field (second "value") can be accessed by index let third_field = &struct_array.fields()[2]; assert_eq!( - third_field.to_primitive().as_slice::(), + third_field.to_primitive()?.as_slice::(), [100i32, 200, 300] ); + Ok(()) } #[test] -fn test_uncompressed_size_in_bytes() { +fn test_uncompressed_size_in_bytes() -> vortex_error::VortexResult<()> { let struct_array = StructArray::new( FieldNames::from(["integers"]), vec![ConstantArray::new(5, 1000).into_array()], @@ -142,11 +148,12 @@ fn test_uncompressed_size_in_bytes() { Validity::NonNullable, ); - let canonical_size = struct_array.to_canonical().into_array().nbytes(); + let canonical_size = struct_array.to_canonical()?.into_array().nbytes(); let uncompressed_size = struct_array .statistics() .compute_uncompressed_size_in_bytes(); assert_eq!(canonical_size, 2); assert_eq!(uncompressed_size, Some(4000)); + Ok(()) } diff --git a/vortex-array/src/arrays/struct_/vtable/canonical.rs b/vortex-array/src/arrays/struct_/vtable/canonical.rs index 8ac579b90b1..85c4d95b680 100644 --- a/vortex-array/src/arrays/struct_/vtable/canonical.rs +++ b/vortex-array/src/arrays/struct_/vtable/canonical.rs @@ -1,13 +1,15 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_error::VortexResult; + use crate::Canonical; use crate::arrays::struct_::StructArray; use crate::arrays::struct_::StructVTable; use crate::vtable::CanonicalVTable; impl CanonicalVTable for StructVTable { - fn canonicalize(array: &StructArray) -> Canonical { - Canonical::Struct(array.clone()) + fn canonicalize(array: &StructArray) -> VortexResult { + Ok(Canonical::Struct(array.clone())) } } diff --git a/vortex-array/src/arrays/struct_/vtable/mod.rs b/vortex-array/src/arrays/struct_/vtable/mod.rs index a621793985e..adc2a5bfc76 100644 --- a/vortex-array/src/arrays/struct_/vtable/mod.rs +++ b/vortex-array/src/arrays/struct_/vtable/mod.rs @@ -113,7 +113,7 @@ impl VTable for StructVTable { .map(|field| field.batch_execute(ctx)) .try_collect()?; // SAFETY: we know that all field lengths match the struct array length, and the validity - Ok(unsafe { StructVector::new_unchecked(Arc::new(fields), array.validity_mask()) }.into()) + Ok(unsafe { StructVector::new_unchecked(Arc::new(fields), array.validity_mask()?) }.into()) } } diff --git a/vortex-array/src/arrays/varbin/accessor.rs b/vortex-array/src/arrays/varbin/accessor.rs index cbe478d1fa9..7e15a0559ed 100644 --- a/vortex-array/src/arrays/varbin/accessor.rs +++ b/vortex-array/src/arrays/varbin/accessor.rs @@ -4,6 +4,7 @@ use std::iter; use vortex_dtype::match_each_integer_ptype; +use vortex_error::VortexExpect; use crate::ToCanonical; use crate::accessor::ArrayAccessor; @@ -16,7 +17,7 @@ impl ArrayAccessor<[u8]> for VarBinArray { where F: for<'a> FnOnce(&mut dyn Iterator>) -> R, { - let offsets = self.offsets().to_primitive(); + let offsets = self.offsets().to_primitive().vortex_expect("to_primitive"); let validity = self.validity(); let bytes = self.bytes(); @@ -35,7 +36,7 @@ impl ArrayAccessor<[u8]> for VarBinArray { } Validity::AllInvalid => f(&mut iter::repeat_n(None, self.len())), Validity::Array(v) => { - let validity = v.to_bool(); + let validity = v.to_bool().vortex_expect("to_bool"); let mut iter = offsets .windows(2) .zip(validity.bit_buffer()) diff --git a/vortex-array/src/arrays/varbin/array.rs b/vortex-array/src/arrays/varbin/array.rs index 068d11fb6c3..eab8d0390fd 100644 --- a/vortex-array/src/arrays/varbin/array.rs +++ b/vortex-array/src/arrays/varbin/array.rs @@ -170,7 +170,7 @@ impl VarBinArray { // Validate UTF-8 for Utf8 dtype if matches!(dtype, DType::Utf8(_)) { - let primitive_offsets = offsets.to_primitive(); + let primitive_offsets = offsets.to_primitive()?; match_each_integer_ptype!(primitive_offsets.dtype().as_ptype(), |O| { let offsets_slice = primitive_offsets.as_slice::(); for (i, (start, end)) in offsets_slice @@ -178,7 +178,7 @@ impl VarBinArray { .map(|o| (o[0].as_(), o[1].as_())) .enumerate() { - if validity.is_null(i) { + if validity.is_null(i)? { continue; } diff --git a/vortex-array/src/arrays/varbin/compute/compare.rs b/vortex-array/src/arrays/varbin/compute/compare.rs index 09e90f29035..710aeea82c1 100644 --- a/vortex-array/src/arrays/varbin/compute/compare.rs +++ b/vortex-array/src/arrays/varbin/compute/compare.rs @@ -61,7 +61,7 @@ impl CompareKernel for VarBinVTable { Operator::Gte => BitBuffer::new_set(len), // Every possible value is >= "" Operator::Lt => BitBuffer::new_unset(len), // No value is < "" Operator::Eq | Operator::NotEq | Operator::Gt | Operator::Lte => { - let lhs_offsets = lhs.offsets().to_primitive(); + let lhs_offsets = lhs.offsets().to_primitive()?; match_each_integer_ptype!(lhs_offsets.ptype(), |P| { compare_offsets_to_empty::

(lhs_offsets, operator) }) @@ -114,7 +114,7 @@ impl CompareKernel for VarBinVTable { // NOTE: If the rhs is not a VarBin array it will be canonicalized to a VarBinView // Arrow doesn't support comparing VarBin to VarBinView arrays, so we convert ourselves // to VarBinView and re-invoke. - return Ok(Some(compare(lhs.to_varbinview().as_ref(), rhs, operator)?)); + return Ok(Some(compare(lhs.to_varbinview()?.as_ref(), rhs, operator)?)); } else { Ok(None) } @@ -166,10 +166,11 @@ mod test { Operator::Eq, ) .unwrap() - .to_bool(); + .to_bool() + .unwrap(); assert_eq!( - &result.validity_mask().to_bit_buffer(), + &result.validity_mask().unwrap().to_bit_buffer(), &BitBuffer::from_iter([true, false, true]) ); assert_eq!( @@ -190,10 +191,11 @@ mod test { ); let result = compare(array.as_ref(), vbv.as_ref(), Operator::Eq) .unwrap() - .to_bool(); + .to_bool() + .unwrap(); 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..ecf2461638f 100644 --- a/vortex-array/src/arrays/varbin/compute/filter.rs +++ b/vortex-array/src/arrays/varbin/compute/filter.rs @@ -52,14 +52,14 @@ fn filter_select_var_bin_by_slice( mask_slices: &[(usize, usize)], selection_count: usize, ) -> VortexResult { - let offsets = values.offsets().to_primitive(); + let offsets = values.offsets().to_primitive()?; match_each_integer_ptype!(offsets.ptype(), |O| { filter_select_var_bin_by_slice_primitive_offset( values.dtype().clone(), offsets.as_slice::(), values.bytes().as_slice(), mask_slices, - values.validity_mask(), + values.validity_mask()?, selection_count, ) }) @@ -150,7 +150,7 @@ fn filter_select_var_bin_by_index( mask_indices: &[usize], selection_count: usize, ) -> VortexResult { - let offsets = values.offsets().to_primitive(); + let offsets = values.offsets().to_primitive()?; match_each_integer_ptype!(offsets.ptype(), |O| { filter_select_var_bin_by_index_primitive_offset( values.dtype().clone(), @@ -174,7 +174,7 @@ fn filter_select_var_bin_by_index_primitive_offset( ) -> VortexResult { let mut builder = VarBinBuilder::::with_capacity(selection_count); for idx in mask_indices.iter().copied() { - if validity.is_valid(idx) { + if validity.is_valid(idx)? { let (start, end) = ( offsets[idx].to_usize().ok_or_else(|| { vortex_err!("Failed to convert offset to usize: {}", offsets[idx]) diff --git a/vortex-array/src/arrays/varbin/compute/mask.rs b/vortex-array/src/arrays/varbin/compute/mask.rs index 5119ebbeb82..5e6cb21d0f9 100644 --- a/vortex-array/src/arrays/varbin/compute/mask.rs +++ b/vortex-array/src/arrays/varbin/compute/mask.rs @@ -19,7 +19,7 @@ impl MaskKernel for VarBinVTable { array.offsets().clone(), array.bytes().clone(), array.dtype().as_nullable(), - array.validity().mask(mask), + array.validity().mask(mask)?, )? .into_array()) } diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index caa32c451a0..2e10b90073d 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -26,9 +26,9 @@ use crate::validity::Validity; impl TakeKernel for VarBinVTable { fn take(&self, array: &VarBinArray, indices: &dyn Array) -> VortexResult { - let offsets = array.offsets().to_primitive(); + let offsets = array.offsets().to_primitive()?; let data = array.bytes(); - let indices = indices.to_primitive(); + let indices = indices.to_primitive()?; let dtype = array .dtype() .clone() @@ -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 e9ef012e568..b3c655417ae 100644 --- a/vortex-array/src/arrays/varbin/vtable/canonical.rs +++ b/vortex-array/src/arrays/varbin/vtable/canonical.rs @@ -8,7 +8,7 @@ use arrow_array::StringViewArray; use arrow_array::cast::AsArray; use arrow_schema::DataType; use vortex_dtype::DType; -use vortex_error::VortexExpect; +use vortex_error::VortexResult; use crate::ArrayRef; use crate::Canonical; @@ -20,14 +20,11 @@ use crate::arrow::IntoArrowArray; use crate::vtable::CanonicalVTable; impl CanonicalVTable for VarBinVTable { - fn canonicalize(array: &VarBinArray) -> Canonical { + fn canonicalize(array: &VarBinArray) -> VortexResult { let dtype = array.dtype().clone(); let nullable = dtype.is_nullable(); - let array_ref = array - .to_array() - .into_arrow_preferred() - .vortex_expect("VarBinArray must be convertible to arrow array"); + let array_ref = array.to_array().into_arrow_preferred()?; let array = match (&dtype, array_ref.data_type()) { (DType::Utf8(_), DataType::Utf8) => { @@ -51,7 +48,9 @@ impl CanonicalVTable for VarBinVTable { } _ => unreachable!("VarBinArray must have Utf8 or Binary dtype, instead got: {dtype}",), }; - Canonical::VarBinView(ArrayRef::from_arrow(array.as_ref(), nullable).to_varbinview()) + Ok(Canonical::VarBinView( + ArrayRef::from_arrow(array.as_ref(), nullable).to_varbinview()?, + )) } } @@ -77,11 +76,11 @@ mod tests { varbin.append_value("1234567890123".as_bytes()); let varbin = varbin.finish(dtype.clone()); - let canonical = varbin.to_varbinview(); + let canonical = varbin.to_varbinview().unwrap(); assert_eq!(canonical.dtype(), &dtype); - assert!(!canonical.is_valid(0)); - assert!(!canonical.is_valid(1)); + assert!(!canonical.is_valid(0).unwrap()); + assert!(!canonical.is_valid(1).unwrap()); // First value is inlined (12 bytes) assert!(canonical.views()[2].is_inlined()); diff --git a/vortex-array/src/arrays/varbinview/accessor.rs b/vortex-array/src/arrays/varbinview/accessor.rs index 4dbe8c10887..ad7e62d693e 100644 --- a/vortex-array/src/arrays/varbinview/accessor.rs +++ b/vortex-array/src/arrays/varbinview/accessor.rs @@ -3,6 +3,8 @@ use std::iter; +use vortex_error::VortexExpect; + use crate::ToCanonical; use crate::accessor::ArrayAccessor; use crate::arrays::varbinview::VarBinViewArray; @@ -35,7 +37,7 @@ impl ArrayAccessor<[u8]> for VarBinViewArray { } Validity::AllInvalid => f(&mut iter::repeat_n(None, views.len())), Validity::Array(v) => { - let validity = v.to_bool(); + let validity = v.to_bool().vortex_expect("to_bool"); let mut iter = views .iter() .zip(validity.bit_buffer()) diff --git a/vortex-array/src/arrays/varbinview/array.rs b/vortex-array/src/arrays/varbinview/array.rs index 1c7aacc45d0..8a4ee16a18e 100644 --- a/vortex-array/src/arrays/varbinview/array.rs +++ b/vortex-array/src/arrays/varbinview/array.rs @@ -210,7 +210,7 @@ impl VarBinViewArray { F: Fn(&[u8]) -> bool, { for (idx, &view) in views.iter().enumerate() { - if validity.is_null(idx) { + if validity.is_null(idx)? { continue; } diff --git a/vortex-array/src/arrays/varbinview/compact.rs b/vortex-array/src/arrays/varbinview/compact.rs index 9df4dea11a6..6bebc78b236 100644 --- a/vortex-array/src/arrays/varbinview/compact.rs +++ b/vortex-array/src/arrays/varbinview/compact.rs @@ -65,7 +65,7 @@ impl VarBinViewArray { .iter() .enumerate() .map(|(idx, &view)| { - if !self.is_valid(idx) || view.is_inlined() { + if !self.is_valid(idx).unwrap_or(false) || view.is_inlined() { 0u64 } else { view.len() as u64 @@ -90,7 +90,7 @@ impl VarBinViewArray { } for (idx, &view) in self.views().iter().enumerate() { - if !self.is_valid(idx) || view.is_inlined() { + if !self.is_valid(idx).unwrap_or(false) || view.is_inlined() { continue; } let view = view.as_view(); diff --git a/vortex-array/src/arrays/varbinview/compute/mask.rs b/vortex-array/src/arrays/varbinview/compute/mask.rs index 884419c7578..b634bc34ff6 100644 --- a/vortex-array/src/arrays/varbinview/compute/mask.rs +++ b/vortex-array/src/arrays/varbinview/compute/mask.rs @@ -21,7 +21,7 @@ impl MaskKernel for VarBinViewVTable { array.views().clone(), array.buffers().clone(), array.dtype().as_nullable(), - array.validity().mask(mask), + array.validity().mask(mask)?, ) .into_array()) } diff --git a/vortex-array/src/arrays/varbinview/compute/mod.rs b/vortex-array/src/arrays/varbinview/compute/mod.rs index ce194b684df..b729299c27d 100644 --- a/vortex-array/src/arrays/varbinview/compute/mod.rs +++ b/vortex-array/src/arrays/varbinview/compute/mod.rs @@ -36,7 +36,7 @@ mod tests { assert!(taken.dtype().is_nullable()); assert_eq!( - taken.to_varbinview().with_iterator(|it| it + taken.to_varbinview().unwrap().with_iterator(|it| it .map(|v| v.map(|b| unsafe { String::from_utf8_unchecked(b.to_vec()) })) .collect::>()), [Some("one".to_string()), Some("four".to_string())] diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index e400dac0bde..3a4963ac3ad 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/vortex-array/src/arrays/varbinview/compute/take.rs @@ -28,7 +28,7 @@ impl TakeKernel for VarBinViewVTable { // This is valid since all elements (of all arrays) even null values must be inside // min-max valid range. let validity = array.validity().take(indices)?; - let indices = indices.to_primitive(); + let indices = indices.to_primitive()?; let views_buffer = match_each_integer_ptype!(indices.ptype(), |I| { // This is valid since all elements even null values are inside the min-max valid range. @@ -92,7 +92,7 @@ mod tests { assert!(taken.dtype().is_nullable()); assert_eq!( - taken.to_varbinview().with_iterator(|it| it + taken.to_varbinview().unwrap().with_iterator(|it| it .map(|v| v.map(|b| unsafe { String::from_utf8_unchecked(b.to_vec()) })) .collect::>()), [Some("one".to_string()), Some("four".to_string())] @@ -111,7 +111,7 @@ mod tests { assert!(taken.dtype().is_nullable()); assert_eq!( - taken.to_varbinview().with_iterator(|it| it + taken.to_varbinview().unwrap().with_iterator(|it| it .map(|v| v.map(|b| unsafe { String::from_utf8_unchecked(b.to_vec()) })) .collect::>()), [Some("two".to_string()), None] diff --git a/vortex-array/src/arrays/varbinview/compute/zip.rs b/vortex-array/src/arrays/varbinview/compute/zip.rs index c1ae7f4f6e2..a3422156460 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( @@ -244,7 +244,10 @@ mod tests { let mask = Mask::from_iter([true, false, true, false, false, true]); - let zipped = zip(a.as_ref(), b.as_ref(), &mask).unwrap().to_varbinview(); + let zipped = zip(a.as_ref(), b.as_ref(), &mask) + .unwrap() + .to_varbinview() + .unwrap(); let values = zipped.with_iterator(|it| { it.map(|v| v.map(|bytes| String::from_utf8(bytes.to_vec()).unwrap())) diff --git a/vortex-array/src/arrays/varbinview/tests.rs b/vortex-array/src/arrays/varbinview/tests.rs index f33beba8a16..bb4b76fb3e5 100644 --- a/vortex-array/src/arrays/varbinview/tests.rs +++ b/vortex-array/src/arrays/varbinview/tests.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 vortex_vector::binaryview::BinaryView; @@ -32,11 +33,12 @@ pub fn slice_array() { } #[test] -pub fn flatten_array() { +pub fn flatten_array() -> VortexResult<()> { let binary_arr = VarBinViewArray::from_iter_str(["string1", "string2"]); - let var_bin = binary_arr.to_varbinview(); + let var_bin = binary_arr.to_varbinview()?; assert_eq!(var_bin.scalar_at(0), Scalar::from("string1")); assert_eq!(var_bin.scalar_at(1), Scalar::from("string2")); + Ok(()) } #[test] diff --git a/vortex-array/src/arrays/varbinview/vtable/canonical.rs b/vortex-array/src/arrays/varbinview/vtable/canonical.rs index e4bbacc5581..da4b4f6dd6a 100644 --- a/vortex-array/src/arrays/varbinview/vtable/canonical.rs +++ b/vortex-array/src/arrays/varbinview/vtable/canonical.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_error::VortexResult; + use crate::Canonical; use crate::arrays::VarBinViewVTable; use crate::arrays::varbinview::VarBinViewArray; @@ -8,11 +10,15 @@ use crate::builders::ArrayBuilder; use crate::vtable::CanonicalVTable; impl CanonicalVTable for VarBinViewVTable { - fn canonicalize(array: &VarBinViewArray) -> Canonical { - Canonical::VarBinView(array.clone()) + fn canonicalize(array: &VarBinViewArray) -> VortexResult { + Ok(Canonical::VarBinView(array.clone())) } - fn append_to_builder(array: &VarBinViewArray, builder: &mut dyn ArrayBuilder) { - builder.extend_from_array(array.as_ref()) + fn append_to_builder( + array: &VarBinViewArray, + builder: &mut dyn ArrayBuilder, + ) -> VortexResult<()> { + builder.extend_from_array(array.as_ref()); + Ok(()) } } diff --git a/vortex-array/src/arrays/varbinview/vtable/mod.rs b/vortex-array/src/arrays/varbinview/vtable/mod.rs index 677cd767ccc..77283d4aa1a 100644 --- a/vortex-array/src/arrays/varbinview/vtable/mod.rs +++ b/vortex-array/src/arrays/varbinview/vtable/mod.rs @@ -110,7 +110,7 @@ impl VTable for VarBinViewVTable { StringVector::new_unchecked( array.views().clone(), Arc::new(array.buffers().to_vec().into_boxed_slice()), - array.validity_mask(), + array.validity_mask()?, ) } .into(), @@ -118,7 +118,7 @@ impl VTable for VarBinViewVTable { BinaryVector::new_unchecked( array.views().clone(), Arc::new(array.buffers().to_vec().into_boxed_slice()), - array.validity_mask(), + array.validity_mask()?, ) } .into(), diff --git a/vortex-array/src/arrow/array.rs b/vortex-array/src/arrow/array.rs index bd31e6bcb8c..8cd3c7ce8bd 100644 --- a/vortex-array/src/arrow/array.rs +++ b/vortex-array/src/arrow/array.rs @@ -139,7 +139,7 @@ impl BaseArrayVTable for ArrowVTable { } impl CanonicalVTable for ArrowVTable { - fn canonicalize(array: &ArrowArray) -> Canonical { + fn canonicalize(array: &ArrowArray) -> VortexResult { ArrayRef::from_arrow(array.inner.as_ref(), array.dtype.is_nullable()).to_canonical() } } @@ -161,24 +161,24 @@ impl OperationsVTable for ArrowVTable { } impl ValidityVTable for ArrowVTable { - fn is_valid(array: &ArrowArray, index: usize) -> bool { - array.inner.is_valid(index) + fn is_valid(array: &ArrowArray, index: usize) -> VortexResult { + Ok(array.inner.is_valid(index)) } - fn all_valid(array: &ArrowArray) -> bool { - array.inner.logical_null_count() == 0 + fn all_valid(array: &ArrowArray) -> VortexResult { + Ok(array.inner.logical_null_count() == 0) } - fn all_invalid(array: &ArrowArray) -> bool { - array.inner.logical_null_count() == array.inner.len() + fn all_invalid(array: &ArrowArray) -> VortexResult { + Ok(array.inner.logical_null_count() == array.inner.len()) } - 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/compute/to_arrow/canonical.rs b/vortex-array/src/arrow/compute/to_arrow/canonical.rs index 3e0223cae84..dadb31a4b6d 100644 --- a/vortex-array/src/arrow/compute/to_arrow/canonical.rs +++ b/vortex-array/src/arrow/compute/to_arrow/canonical.rs @@ -111,7 +111,7 @@ impl Kernel for ToArrowCanonical { // preferred Arrow type if the array has child arrays (struct, list, and fixed-size list). let to_preferred = arrow_type_opt.is_none(); - let arrow_array = match (array.to_canonical(), &arrow_type) { + let arrow_array = match (array.to_canonical()?, &arrow_type) { (Canonical::Null(array), DataType::Null) => to_arrow_null(array), (Canonical::Bool(array), DataType::Boolean) => to_arrow_bool(array), (Canonical::Primitive(array), DataType::Int8) if matches!(array.ptype(), PType::I8) => { @@ -294,12 +294,12 @@ fn to_arrow_null(array: NullArray) -> VortexResult { fn to_arrow_bool(array: BoolArray) -> VortexResult { Ok(Arc::new(ArrowBoolArray::new( array.bit_buffer().clone().into(), - to_null_buffer(array.validity_mask()), + to_null_buffer(array.validity_mask()?), ))) } fn to_arrow_primitive(array: PrimitiveArray) -> VortexResult { - let null_buffer = to_null_buffer(array.validity_mask()); + let null_buffer = to_null_buffer(array.validity_mask()?); let len = array.len(); let buffer = array.into_byte_buffer().into_arrow_buffer(); Ok(Arc::new(ArrowPrimitiveArray::::new( @@ -309,7 +309,7 @@ fn to_arrow_primitive(array: PrimitiveArray) -> VortexRes } 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_())) @@ -353,7 +353,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_())) @@ -392,7 +392,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_())) @@ -426,7 +426,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_())) @@ -477,7 +477,7 @@ fn to_arrow_struct( // We check that the Vortex array nullability is compatible with the field // nullability. In other words, make sure we don't return any nulls for a // non-nullable field. - if arr.dtype().is_nullable() && !field.is_nullable() && !arr.all_valid() { + if arr.dtype().is_nullable() && !field.is_nullable() && !arr.all_valid()? { vortex_bail!( "Field {} is non-nullable but has nulls {}", field, @@ -494,7 +494,7 @@ fn to_arrow_struct( }) .collect::>>()?; - let nulls = to_null_buffer(array.validity_mask()); + let nulls = to_null_buffer(array.validity_mask()?); if field_arrays.is_empty() { return Ok(Arc::new(ArrowStructArray::new_empty_fields( @@ -575,10 +575,10 @@ fn to_arrow_list( // Convert the child `offsets` and `validity` array to Arrow. let offsets = list_array .offsets() - .to_primitive() + .to_primitive()? .buffer::() .into_arrow_offset_buffer(); - let nulls = to_null_buffer(list_array.validity_mask()); + let nulls = to_null_buffer(list_array.validity_mask()?); Ok(Arc::new(GenericListArray::new( element_field, @@ -597,15 +597,15 @@ fn to_arrow_listview( let offsets_dtype = DType::Primitive(O::PTYPE, array.dtype().nullability()); let offsets = cast(array.offsets(), &offsets_dtype) .map_err(|err| err.with_context(format!("Failed to cast offsets to {offsets_dtype}")))? - .to_primitive(); + .to_primitive()?; let sizes = cast(array.sizes(), &offsets_dtype) .map_err(|err| err.with_context(format!("Failed to cast sizes to {offsets_dtype}")))? - .to_primitive(); + .to_primitive()?; // Convert `offsets`, `sizes`, and `validity` to Arrow buffers. let arrow_offsets = offsets.buffer::().into_arrow_scalar_buffer(); let arrow_sizes = sizes.buffer::().into_arrow_scalar_buffer(); - let nulls = to_null_buffer(array.validity_mask()); + let nulls = to_null_buffer(array.validity_mask()?); // Convert the child `elements` array to Arrow. let (elements, element_field) = { @@ -665,7 +665,7 @@ fn to_arrow_fixed_size_list( )); (values, element_field) }; - let nulls = to_null_buffer(array.validity_mask()); + let nulls = to_null_buffer(array.validity_mask()?); // TODO(connor): Revert this once the issue below is resolved. // Ok(Arc::new(ArrowFixedSizeListArray::new( @@ -699,7 +699,7 @@ fn to_arrow_varbinview(array: VarBinViewArray) -> VortexResult< .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. Ok(Arc::new(unsafe { diff --git a/vortex-array/src/arrow/compute/to_arrow/list.rs b/vortex-array/src/arrow/compute/to_arrow/list.rs index 3046a6c5253..f8b1f4afd9b 100644 --- a/vortex-array/src/arrow/compute/to_arrow/list.rs +++ b/vortex-array/src/arrow/compute/to_arrow/list.rs @@ -65,11 +65,11 @@ fn list_array_to_arrow_list( let offsets_dtype = DType::Primitive(O::PTYPE, array.dtype().nullability()); let offsets = cast(array.offsets(), &offsets_dtype) .map_err(|err| err.with_context(format!("Failed to cast offsets to {offsets_dtype}")))? - .to_primitive(); + .to_primitive()?; // Convert `offsets` and `validity` to Arrow buffers. let arrow_offsets = offsets.buffer::().into_arrow_offset_buffer(); - let nulls = to_null_buffer(array.validity_mask()); + let nulls = to_null_buffer(array.validity_mask()?); // Convert the child `elements` array to Arrow. let (elements, element_field) = { diff --git a/vortex-array/src/arrow/compute/to_arrow/mod.rs b/vortex-array/src/arrow/compute/to_arrow/mod.rs index 11c4d6c6533..0abb3c26e9e 100644 --- a/vortex-array/src/arrow/compute/to_arrow/mod.rs +++ b/vortex-array/src/arrow/compute/to_arrow/mod.rs @@ -127,7 +127,7 @@ impl ComputeFnVTable for ToArrow { // Fall back to canonicalizing and then converting. if !array.is_canonical() { - let canonical_array = array.to_canonical(); + let canonical_array = array.to_canonical()?; let arrow_array = to_arrow_opts( canonical_array.as_ref(), &ToArrowOptions { diff --git a/vortex-array/src/arrow/compute/to_arrow/temporal.rs b/vortex-array/src/arrow/compute/to_arrow/temporal.rs index 9a58e2016c7..9b3cc2efa68 100644 --- a/vortex-array/src/arrow/compute/to_arrow/temporal.rs +++ b/vortex-array/src/arrow/compute/to_arrow/temporal.rs @@ -128,10 +128,10 @@ where { let values_dtype = DType::Primitive(T::Native::PTYPE, array.dtype().nullability()); let values = cast(array.temporal_values(), &values_dtype)? - .to_primitive() + .to_primitive()? .into_buffer() .into_arrow_scalar_buffer(); - let nulls = to_null_buffer(array.temporal_values().validity_mask()); + let nulls = to_null_buffer(array.temporal_values().validity_mask()?); Ok(ArrowPrimitiveArray::::new(values, nulls)) } diff --git a/vortex-array/src/arrow/compute/to_arrow/varbin.rs b/vortex-array/src/arrow/compute/to_arrow/varbin.rs index 189a1151260..f940b9d6d73 100644 --- a/vortex-array/src/arrow/compute/to_arrow/varbin.rs +++ b/vortex-array/src/arrow/compute/to_arrow/varbin.rs @@ -82,9 +82,9 @@ fn to_arrow(array: &VarBinArray) -> VortexRes array.offsets(), &DType::Primitive(O::PTYPE, Nullability::NonNullable), )? - .to_primitive(); + .to_primitive()?; - let nulls = to_null_buffer(array.validity_mask()); + let nulls = to_null_buffer(array.validity_mask()?); let data = array.bytes().clone(); // Match on the `DType`. diff --git a/vortex-array/src/arrow/record_batch.rs b/vortex-array/src/arrow/record_batch.rs index 3b7ca6a8314..1583de7cbfc 100644 --- a/vortex-array/src/arrow/record_batch.rs +++ b/vortex-array/src/arrow/record_batch.rs @@ -20,12 +20,12 @@ impl TryFrom<&dyn Array> for RecordBatch { type Error = VortexError; fn try_from(value: &dyn Array) -> VortexResult { - let Canonical::Struct(struct_array) = value.to_canonical() else { + let Canonical::Struct(struct_array) = value.to_canonical()? else { vortex_bail!("RecordBatch can only be constructed from ") }; 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 85c23caa8f0..143fe95388e 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; @@ -114,10 +115,11 @@ impl ArrayBuilder for BoolBuilder { } unsafe fn extend_from_array_unchecked(&mut self, array: &dyn Array) { - let bool_array = array.to_bool(); + let bool_array = array.to_bool().vortex_expect("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")); } fn reserve_exact(&mut self, additional: usize) { @@ -146,6 +148,7 @@ mod tests { use rand::prelude::StdRng; use vortex_dtype::DType; use vortex_dtype::Nullability; + use vortex_error::VortexResult; use vortex_scalar::Scalar; use crate::ArrayRef; @@ -178,19 +181,20 @@ mod tests { } #[test] - fn tests() { + fn tests() -> VortexResult<()> { let len = 1000; let chunk_count = 10; let chunk = make_opt_bool_chunks(len, chunk_count); let mut builder = builder_with_capacity(chunk.dtype(), len * chunk_count); - chunk.clone().append_to_builder(builder.as_mut()); + chunk.clone().append_to_builder(builder.as_mut())?; - let canon_into = builder.finish().to_bool(); - let into_canon = chunk.to_bool(); + let canon_into = builder.finish().to_bool()?; + let into_canon = chunk.to_bool()?; assert_eq!(canon_into.validity(), into_canon.validity()); assert_eq!(canon_into.bit_buffer(), into_canon.bit_buffer()); + Ok(()) } #[test] diff --git a/vortex-array/src/builders/decimal.rs b/vortex-array/src/builders/decimal.rs index 96be211e64d..90dcef90333 100644 --- a/vortex-array/src/builders/decimal.rs +++ b/vortex-array/src/builders/decimal.rs @@ -189,7 +189,7 @@ impl ArrayBuilder for DecimalBuilder { } unsafe fn extend_from_array_unchecked(&mut self, array: &dyn Array) { - let decimal_array = array.to_decimal(); + let decimal_array = array.to_decimal().vortex_expect("to_decimal"); match_each_decimal_value_type!(decimal_array.values_type(), |D| { // Extends the values buffer from another buffer of type D where D can be coerced to the @@ -199,7 +199,7 @@ impl ArrayBuilder for DecimalBuilder { }); self.nulls - .append_validity_mask(decimal_array.validity_mask()); + .append_validity_mask(decimal_array.validity_mask().vortex_expect("validity_mask")); } fn reserve_exact(&mut self, additional: usize) { diff --git a/vortex-array/src/builders/dict/bytes.rs b/vortex-array/src/builders/dict/bytes.rs index ac769037e26..9873dc14599 100644 --- a/vortex-array/src/builders/dict/bytes.rs +++ b/vortex-array/src/builders/dict/bytes.rs @@ -174,7 +174,7 @@ impl DictEncoder for BytesDictBuilder { } else { // NOTE(aduffy): it is very rare that this path would be taken, only e.g. // if we're performing dictionary encoding downstream of some other compression. - self.encode_bytes(&array.to_varbinview(), len) + self.encode_bytes(&array.to_varbinview().vortex_unwrap(), len) } } @@ -211,17 +211,20 @@ mod test { let arr = VarBinArray::from(vec!["hello", "world", "hello", "again", "world"]); let dict = dict_encode(arr.as_ref()).unwrap(); assert_eq!( - dict.codes().to_primitive().as_slice::(), + dict.codes().to_primitive().unwrap().as_slice::(), &[0, 1, 0, 2, 1] ); - dict.values().to_varbinview().with_iterator(|iter| { - assert_eq!( - iter.flatten() - .map(|b| unsafe { str::from_utf8_unchecked(b) }) - .collect::>(), - vec!["hello", "world", "again"] - ); - }); + dict.values() + .to_varbinview() + .unwrap() + .with_iterator(|iter| { + assert_eq!( + iter.flatten() + .map(|b| unsafe { str::from_utf8_unchecked(b) }) + .collect::>(), + vec!["hello", "world", "again"] + ); + }); } #[test] @@ -240,32 +243,38 @@ mod test { .collect(); let dict = dict_encode(arr.as_ref()).unwrap(); assert_eq!( - dict.codes().to_primitive().as_slice::(), + dict.codes().to_primitive().unwrap().as_slice::(), &[0, 1, 2, 0, 1, 3, 2, 1] ); - dict.values().to_varbinview().with_iterator(|iter| { - assert_eq!( - iter.map(|b| b.map(|v| unsafe { str::from_utf8_unchecked(v) })) - .collect::>(), - vec![Some("hello"), None, Some("world"), Some("again")] - ); - }); + dict.values() + .to_varbinview() + .unwrap() + .with_iterator(|iter| { + assert_eq!( + iter.map(|b| b.map(|v| unsafe { str::from_utf8_unchecked(v) })) + .collect::>(), + vec![Some("hello"), None, Some("world"), Some("again")] + ); + }); } #[test] fn repeated_values() { let arr = VarBinArray::from(vec!["a", "a", "b", "b", "a", "b", "a", "b"]); let dict = dict_encode(arr.as_ref()).unwrap(); - dict.values().to_varbinview().with_iterator(|iter| { - assert_eq!( - iter.flatten() - .map(|b| unsafe { str::from_utf8_unchecked(b) }) - .collect::>(), - vec!["a", "b"] - ); - }); + dict.values() + .to_varbinview() + .unwrap() + .with_iterator(|iter| { + assert_eq!( + iter.flatten() + .map(|b| unsafe { str::from_utf8_unchecked(b) }) + .collect::>(), + vec!["a", "b"] + ); + }); assert_eq!( - dict.codes().to_primitive().as_slice::(), + dict.codes().to_primitive().unwrap().as_slice::(), &[0, 0, 1, 1, 0, 1, 0, 1] ); } diff --git a/vortex-array/src/builders/dict/mod.rs b/vortex-array/src/builders/dict/mod.rs index 8d4c2306bb7..7ec0f01604b 100644 --- a/vortex-array/src/builders/dict/mod.rs +++ b/vortex-array/src/builders/dict/mod.rs @@ -59,7 +59,7 @@ pub fn dict_encode_with_constraints( constraints: &DictConstraints, ) -> VortexResult { let mut encoder = dict_encoder(array, constraints); - let codes = encoder.encode(array).to_primitive().narrow()?; + let codes = encoder.encode(array).to_primitive()?.narrow()?; // SAFETY: The encoding process will produce a value set of codes and values // All values in the dictionary are guaranteed to be referenced by at least one code // since we build the dictionary from the codes we observe during encoding diff --git a/vortex-array/src/builders/dict/primitive.rs b/vortex-array/src/builders/dict/primitive.rs index 374c5f2b6c9..4fe538f07ad 100644 --- a/vortex-array/src/builders/dict/primitive.rs +++ b/vortex-array/src/builders/dict/primitive.rs @@ -10,6 +10,7 @@ use vortex_buffer::BufferMut; use vortex_dtype::NativePType; use vortex_dtype::Nullability; use vortex_dtype::UnsignedPType; +use vortex_error::VortexExpect; use vortex_error::vortex_panic; use vortex_utils::aliases::hash_map::Entry; use vortex_utils::aliases::hash_map::HashMap; @@ -126,14 +127,17 @@ where fn encode(&mut self, array: &dyn Array) -> ArrayRef { let mut codes = BufferMut::::with_capacity(array.len()); - array.to_primitive().with_iterator(|it| { - for value in it { - let Some(code) = self.encode_value(value.copied()) else { - break; - }; - unsafe { codes.push_unchecked(code) } - } - }); + array + .to_primitive() + .vortex_expect("dict encode expects primitive array") + .with_iterator(|it| { + for value in it { + let Some(code) = self.encode_value(value.copied()) else { + break; + }; + unsafe { codes.push_unchecked(code) } + } + }); PrimitiveArray::new(codes, Validity::NonNullable).into_array() } diff --git a/vortex-array/src/builders/extension.rs b/vortex-array/src/builders/extension.rs index a7a207dc177..3dbfc4ac18f 100644 --- a/vortex-array/src/builders/extension.rs +++ b/vortex-array/src/builders/extension.rs @@ -6,6 +6,7 @@ use std::sync::Arc; use vortex_dtype::DType; use vortex_dtype::ExtDType; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_mask::Mask; @@ -101,7 +102,7 @@ impl ArrayBuilder for ExtensionBuilder { } unsafe fn extend_from_array_unchecked(&mut self, array: &dyn Array) { - let ext_array = array.to_extension(); + let ext_array = array.to_extension().vortex_expect("to_extension"); self.storage.extend_from_array(ext_array.storage()) } diff --git a/vortex-array/src/builders/fixed_size_list.rs b/vortex-array/src/builders/fixed_size_list.rs index dfae6a1ffec..7089a3bd09a 100644 --- a/vortex-array/src/builders/fixed_size_list.rs +++ b/vortex-array/src/builders/fixed_size_list.rs @@ -208,13 +208,16 @@ impl ArrayBuilder for FixedSizeListBuilder { /// This will increase the capacity if extending with this `array` would go past the original /// capacity. unsafe fn extend_from_array_unchecked(&mut self, array: &dyn Array) { - let fsl = array.to_fixed_size_list(); + let fsl = array + .to_fixed_size_list() + .vortex_expect("to_fixed_size_list"); if fsl.is_empty() { return; } 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")); } fn reserve_exact(&mut self, additional: usize) { @@ -297,7 +300,7 @@ mod tests { let fsl = builder.finish(); assert_eq!(fsl.len(), 2); - let fsl_array = fsl.to_fixed_size_list(); + let fsl_array = fsl.to_fixed_size_list().unwrap(); assert_eq!(fsl_array.elements().len(), 6); assert_eq!(fsl_array.list_size(), 3); } @@ -318,7 +321,7 @@ mod tests { let fsl = builder.finish(); assert_eq!(fsl.len(), 100); - let fsl_array = fsl.to_fixed_size_list(); + let fsl_array = fsl.to_fixed_size_list().unwrap(); assert_eq!(fsl_array.list_size(), 0); // The elements array should be empty since list_size is 0. assert_eq!(fsl_array.elements().len(), 0); @@ -346,7 +349,7 @@ mod tests { let fsl = builder.finish(); assert_eq!(fsl.len(), 100); - let fsl_array = fsl.to_fixed_size_list(); + let fsl_array = fsl.to_fixed_size_list().unwrap(); assert_eq!(fsl_array.list_size(), 0); assert_eq!(fsl_array.elements().len(), 0); } @@ -374,7 +377,7 @@ mod tests { let fsl = builder.finish(); assert_eq!(fsl.len(), 5); - let fsl_array = fsl.to_fixed_size_list(); + let fsl_array = fsl.to_fixed_size_list().unwrap(); assert_eq!(fsl_array.elements().len(), 10); } @@ -387,7 +390,7 @@ mod tests { let fsl = builder.finish(); assert_eq!(fsl.len(), 0); - let fsl_array = fsl.to_fixed_size_list(); + let fsl_array = fsl.to_fixed_size_list().unwrap(); assert_eq!(fsl_array.list_size(), 100000000); assert_eq!(fsl_array.elements().len(), 0); } @@ -415,10 +418,10 @@ mod tests { let fsl = builder.finish(); assert_eq!(fsl.len(), 3); - let fsl_array = fsl.to_fixed_size_list(); - assert!(fsl_array.validity().is_valid(0)); - assert!(!fsl_array.validity().is_valid(1)); - assert!(fsl_array.validity().is_valid(2)); + let fsl_array = fsl.to_fixed_size_list().unwrap(); + assert!(fsl_array.validity().is_valid(0).unwrap()); + assert!(!fsl_array.validity().is_valid(1).unwrap()); + assert!(fsl_array.validity().is_valid(2).unwrap()); } #[test] @@ -459,7 +462,7 @@ mod tests { let fsl = builder.finish(); assert_eq!(fsl.len(), 2); - let fsl_array = fsl.to_fixed_size_list(); + let fsl_array = fsl.to_fixed_size_list().unwrap(); assert_eq!(fsl_array.elements().len(), 6); } @@ -473,12 +476,12 @@ mod tests { let fsl = builder.finish(); assert_eq!(fsl.len(), 5); - let fsl_array = fsl.to_fixed_size_list(); + let fsl_array = fsl.to_fixed_size_list().unwrap(); assert_eq!(fsl_array.list_size(), 3); assert_eq!(fsl_array.elements().len(), 15); // Check that all elements are zeros. - let elements_array = fsl_array.elements().to_primitive(); + let elements_array = fsl_array.elements().to_primitive().unwrap(); let elements = elements_array.as_slice::(); assert!(elements.iter().all(|&x| x == 0)); } @@ -496,12 +499,12 @@ mod tests { let fsl = builder.finish(); assert_eq!(fsl.len(), 3); - let fsl_array = fsl.to_fixed_size_list(); + let fsl_array = fsl.to_fixed_size_list().unwrap(); assert_eq!(fsl_array.list_size(), 2); // Check that all lists are null. for i in 0..3 { - assert!(!fsl_array.validity().is_valid(i)); + assert!(!fsl_array.validity().is_valid(i).unwrap()); } } @@ -520,11 +523,11 @@ mod tests { let fsl = builder.finish(); assert_eq!(fsl.len(), 1); - let fsl_array = fsl.to_fixed_size_list(); + let fsl_array = fsl.to_fixed_size_list().unwrap(); assert_eq!(fsl_array.list_size(), 2); // Check that all lists are null. - assert!(!fsl_array.validity().is_valid(0)); + assert!(!fsl_array.validity().is_valid(0).unwrap()); } #[test] @@ -539,7 +542,7 @@ mod tests { let fsl = builder.finish(); assert_eq!(fsl.len(), 1000); - let fsl_array = fsl.to_fixed_size_list(); + let fsl_array = fsl.to_fixed_size_list().unwrap(); assert_eq!(fsl_array.list_size(), 0); assert_eq!(fsl_array.elements().len(), 0); } @@ -589,16 +592,16 @@ mod tests { let fsl = builder.finish(); assert_eq!(fsl.len(), 6); - let fsl_array = fsl.to_fixed_size_list(); + let fsl_array = fsl.to_fixed_size_list().unwrap(); assert_eq!(fsl_array.elements().len(), 12); // Check validity pattern is repeated. - assert!(fsl_array.validity().is_valid(0)); - assert!(!fsl_array.validity().is_valid(1)); - assert!(fsl_array.validity().is_valid(2)); - assert!(fsl_array.validity().is_valid(3)); - assert!(!fsl_array.validity().is_valid(4)); - assert!(fsl_array.validity().is_valid(5)); + assert!(fsl_array.validity().is_valid(0).unwrap()); + assert!(!fsl_array.validity().is_valid(1).unwrap()); + assert!(fsl_array.validity().is_valid(2).unwrap()); + assert!(fsl_array.validity().is_valid(3).unwrap()); + assert!(!fsl_array.validity().is_valid(4).unwrap()); + assert!(fsl_array.validity().is_valid(5).unwrap()); } #[test] @@ -628,16 +631,16 @@ mod tests { let fsl = builder.finish(); assert_eq!(fsl.len(), 5); - let fsl_array = fsl.to_fixed_size_list(); + let fsl_array = fsl.to_fixed_size_list().unwrap(); assert_eq!(fsl_array.list_size(), 0); assert_eq!(fsl_array.elements().len(), 0); // Check validity pattern. - assert!(fsl_array.validity().is_valid(0)); - assert!(!fsl_array.validity().is_valid(1)); - assert!(fsl_array.validity().is_valid(2)); - assert!(!fsl_array.validity().is_valid(3)); - assert!(fsl_array.validity().is_valid(4)); + assert!(fsl_array.validity().is_valid(0).unwrap()); + assert!(!fsl_array.validity().is_valid(1).unwrap()); + assert!(fsl_array.validity().is_valid(2).unwrap()); + assert!(!fsl_array.validity().is_valid(3).unwrap()); + assert!(fsl_array.validity().is_valid(4).unwrap()); } #[test] @@ -709,16 +712,16 @@ mod tests { let fsl = builder.finish(); assert_eq!(fsl.len(), 6); - let fsl_array = fsl.to_fixed_size_list(); + let fsl_array = fsl.to_fixed_size_list().unwrap(); assert_eq!(fsl_array.elements().len(), 12); // Check validity. - assert!(fsl_array.validity().is_valid(0)); // append_value - assert!(!fsl_array.validity().is_valid(1)); // append_null - assert!(fsl_array.validity().is_valid(2)); // append_zeros - assert!(fsl_array.validity().is_valid(3)); // append_zeros - assert!(!fsl_array.validity().is_valid(4)); // append_nulls - assert!(fsl_array.validity().is_valid(5)); // extend_from_array + assert!(fsl_array.validity().is_valid(0).unwrap()); // append_value + assert!(!fsl_array.validity().is_valid(1).unwrap()); // append_null + assert!(fsl_array.validity().is_valid(2).unwrap()); // append_zeros + assert!(fsl_array.validity().is_valid(3).unwrap()); // append_zeros + assert!(!fsl_array.validity().is_valid(4).unwrap()); // append_nulls + assert!(fsl_array.validity().is_valid(5).unwrap()); // extend_from_array } #[test] @@ -761,9 +764,9 @@ mod tests { } // Check validity - first two should be valid, third should be null. - assert!(array.validity().is_valid(0)); - assert!(array.validity().is_valid(1)); - assert!(!array.validity().is_valid(2)); + assert!(array.validity().is_valid(0).unwrap()); + assert!(array.validity().is_valid(1).unwrap()); + assert!(!array.validity().is_valid(2).unwrap()); // Test wrong dtype error. let mut builder = FixedSizeListBuilder::with_capacity(dtype, 2, NonNullable, 10); diff --git a/vortex-array/src/builders/list.rs b/vortex-array/src/builders/list.rs index ff3fb43563c..86009a7fcd3 100644 --- a/vortex-array/src/builders/list.rs +++ b/vortex-array/src/builders/list.rs @@ -193,18 +193,30 @@ impl ArrayBuilder for ListBuilder { } unsafe fn extend_from_array_unchecked(&mut self, array: &dyn Array) { - let list = array.to_listview(); + let list = array + .to_listview() + .vortex_expect("failed to convert to listview"); if list.is_empty() { return; } // Append validity information. - self.nulls.append_validity_mask(array.validity_mask()); + self.nulls.append_validity_mask( + array + .validity_mask() + .vortex_expect("failed to get validity mask"), + ); // Note that `ListViewArray` has `n` offsets and sizes, not `n+1` offsets like `ListArray`. let elements = list.elements(); - let offsets = list.offsets().to_primitive(); - let sizes = list.sizes().to_primitive(); + let offsets = list + .offsets() + .to_primitive() + .vortex_expect("failed to convert offsets to primitive"); + let sizes = list + .sizes() + .to_primitive() + .vortex_expect("failed to convert sizes to primitive"); fn extend_inner( builder: &mut ListBuilder, @@ -339,7 +351,7 @@ mod tests { let list = builder.finish(); assert_eq!(list.len(), 2); - let list_array = list.to_listview(); + let list_array = list.to_listview().unwrap(); assert_eq!(list_array.list_elements_at(0).len(), 3); assert_eq!(list_array.list_elements_at(1).len(), 3); @@ -391,7 +403,7 @@ mod tests { let list = builder.finish(); assert_eq!(list.len(), 3); - let list_array = list.to_listview(); + let list_array = list.to_listview().unwrap(); assert_eq!(list_array.list_elements_at(0).len(), 3); assert_eq!(list_array.list_elements_at(1).len(), 0); @@ -427,18 +439,23 @@ mod tests { Arc::new(DType::Primitive(I32, NonNullable)), ) .unwrap() - .to_listview(); + .to_listview() + .unwrap(); let actual = builder.finish_into_canonical().into_listview(); assert_eq!( - actual.elements().to_primitive().as_slice::(), - expected.elements().to_primitive().as_slice::() + actual.elements().to_primitive().unwrap().as_slice::(), + expected + .elements() + .to_primitive() + .unwrap() + .as_slice::() ); assert_eq!( - actual.offsets().to_primitive().as_slice::(), - expected.offsets().to_primitive().as_slice::() + actual.offsets().to_primitive().unwrap().as_slice::(), + expected.offsets().to_primitive().unwrap().as_slice::() ); assert_eq!(actual.validity(), expected.validity()) @@ -481,7 +498,7 @@ mod tests { DType::List(Arc::new(DType::Primitive(I32, NonNullable)), NonNullable), ); - let canon_values = chunked_list.unwrap().to_listview(); + let canon_values = chunked_list.unwrap().to_listview().unwrap(); assert_eq!( one_trailing_unused_element.scalar_at(0), @@ -538,9 +555,9 @@ mod tests { assert!(list2.is_null()); // This should be null. // Check validity. - assert!(array.validity().is_valid(0)); - assert!(array.validity().is_valid(1)); - assert!(!array.validity().is_valid(2)); + assert!(array.validity().is_valid(0).unwrap()); + assert!(array.validity().is_valid(1).unwrap()); + assert!(!array.validity().is_valid(2).unwrap()); // Test wrong dtype error. let mut builder = ListBuilder::::with_capacity(dtype, NonNullable, 20, 10); diff --git a/vortex-array/src/builders/listview.rs b/vortex-array/src/builders/listview.rs index d8a2db80af8..41c76bea7db 100644 --- a/vortex-array/src/builders/listview.rs +++ b/vortex-array/src/builders/listview.rs @@ -254,7 +254,9 @@ impl ArrayBuilder for ListViewBuilder { } unsafe fn extend_from_array_unchecked(&mut self, array: &dyn Array) { - let listview = array.to_listview(); + let listview = array + .to_listview() + .vortex_expect("failed to convert to listview"); if listview.is_empty() { return; } @@ -277,7 +279,11 @@ impl ArrayBuilder for ListViewBuilder { let listview = listview.rebuild(ListViewRebuildMode::MakeExact); 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("failed to get validity mask"), + ); // Bulk append the new elements (which should have no gaps or overlaps). let old_elements_len = self.elements_builder.len(); @@ -303,7 +309,10 @@ impl ArrayBuilder for ListViewBuilder { let uninit_range = self.offsets_builder.uninit_range(extend_length); // This should be cheap because we didn't compress after rebuilding. - let new_offsets = listview.offsets().to_primitive(); + let new_offsets = listview + .offsets() + .to_primitive() + .vortex_expect("failed to convert offsets to primitive"); match_each_integer_ptype!(new_offsets.ptype(), |A| { adjust_and_extend_offsets::( @@ -449,7 +458,7 @@ mod tests { assert_eq!(empty_list.len(), 0); // Check null list. - assert!(!listview.validity().is_valid(2)); + assert!(!listview.validity().is_valid(2).unwrap()); // Check last list: [4, 5]. let last_list = listview.list_elements_at(3); @@ -548,8 +557,8 @@ mod tests { assert_eq!(listview.list_elements_at(1).len(), 0); // Next two are nulls. - assert!(!listview.validity().is_valid(2)); - assert!(!listview.validity().is_valid(3)); + assert!(!listview.validity().is_valid(2).unwrap()); + assert!(!listview.validity().is_valid(3).unwrap()); // Last is the regular list: [10, 20]. let last_list = listview.list_elements_at(4); @@ -604,7 +613,7 @@ mod tests { assert_eq!(second.scalar_at(2), 3i32.into()); // Third list: null (from source). - assert!(!listview.validity().is_valid(2)); + assert!(!listview.validity().is_valid(2).unwrap()); // Fourth list: [4, 5] (from source). let fourth = listview.list_elements_at(3); diff --git a/vortex-array/src/builders/mod.rs b/vortex-array/src/builders/mod.rs index 072fba3dd35..a3f95c04aad 100644 --- a/vortex-array/src/builders/mod.rs +++ b/vortex-array/src/builders/mod.rs @@ -33,6 +33,7 @@ use std::any::Any; use vortex_dtype::DType; use vortex_dtype::match_each_decimal_value_type; use vortex_dtype::match_each_native_ptype; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_panic; use vortex_mask::Mask; @@ -212,7 +213,7 @@ pub trait ArrayBuilder: Send { /// then converts it to canonical form. Specific builders can override this with optimized /// implementations that avoid the intermediate [`Array`] creation. fn finish_into_canonical(&mut self) -> Canonical { - self.finish().to_canonical() + self.finish().to_canonical().vortex_expect("to_canonical") } } diff --git a/vortex-array/src/builders/primitive.rs b/vortex-array/src/builders/primitive.rs index 88eed215e1a..2517ab42c1a 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; @@ -165,7 +166,7 @@ impl ArrayBuilder for PrimitiveBuilder { } unsafe fn extend_from_array_unchecked(&mut self, array: &dyn Array) { - let array = array.to_primitive(); + let array = array.to_primitive().vortex_expect("to_primitive"); // This should be checked in `extend_from_array` but we can check it again. debug_assert_eq!( @@ -175,7 +176,8 @@ 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")); } fn reserve_exact(&mut self, additional: usize) { @@ -605,9 +607,9 @@ mod tests { // Check validity - first two should be valid, third should be null. use crate::vtable::ValidityHelper; - assert!(array.validity().is_valid(0)); - assert!(array.validity().is_valid(1)); - assert!(!array.validity().is_valid(2)); + assert!(array.validity().is_valid(0).unwrap()); + assert!(array.validity().is_valid(1).unwrap()); + assert!(!array.validity().is_valid(2).unwrap()); // Test wrong dtype error. let mut builder = PrimitiveBuilder::::with_capacity(Nullability::NonNullable, 10); diff --git a/vortex-array/src/builders/struct_.rs b/vortex-array/src/builders/struct_.rs index 7ca3a6c80c3..83a8b77ff96 100644 --- a/vortex-array/src/builders/struct_.rs +++ b/vortex-array/src/builders/struct_.rs @@ -167,7 +167,7 @@ impl ArrayBuilder for StructBuilder { } unsafe fn extend_from_array_unchecked(&mut self, array: &dyn Array) { - let array = array.to_struct(); + let array = array.to_struct().vortex_expect("to_struct"); for (a, builder) in array .fields() @@ -175,10 +175,12 @@ impl ArrayBuilder for StructBuilder { .cloned() .zip_eq(self.builders.iter_mut()) { - a.append_to_builder(builder.as_mut()); + a.append_to_builder(builder.as_mut()) + .vortex_expect("append_to_builder"); } - self.nulls.append_validity_mask(array.validity_mask()); + self.nulls + .append_validity_mask(array.validity_mask().vortex_expect("validity_mask")); } fn reserve_exact(&mut self, capacity: usize) { diff --git a/vortex-array/src/builders/tests.rs b/vortex-array/src/builders/tests.rs index 6c859041e37..c54b19cbcc3 100644 --- a/vortex-array/src/builders/tests.rs +++ b/vortex-array/src/builders/tests.rs @@ -243,7 +243,7 @@ where // Get canonical arrays using both methods. let canonical_direct = builder1.finish_into_canonical(); - let canonical_indirect = builder2.finish().to_canonical(); + let canonical_indirect = builder2.finish().to_canonical().unwrap(); // Convert both to arrays for comparison. let array_direct = canonical_direct.into_array(); diff --git a/vortex-array/src/builders/varbinview.rs b/vortex-array/src/builders/varbinview.rs index f1362ab4956..21854cb4bc4 100644 --- a/vortex-array/src/builders/varbinview.rs +++ b/vortex-array/src/builders/varbinview.rs @@ -275,10 +275,16 @@ impl ArrayBuilder for VarBinViewBuilder { } unsafe fn extend_from_array_unchecked(&mut self, array: &dyn Array) { - let array = array.to_varbinview(); + let array = array + .to_varbinview() + .vortex_expect("failed to convert to varbinview"); self.flush_in_progress(); - self.push_only_validity_mask(array.validity_mask()); + self.push_only_validity_mask( + array + .validity_mask() + .vortex_expect("failed to get validity mask"), + ); let view_adjustment = self.completed @@ -296,7 +302,10 @@ impl ArrayBuilder for VarBinViewBuilder { ), ViewAdjustment::Rewriting(adjustment) => { for (idx, &view) in array.views().iter().enumerate() { - let new_view = if !array.is_valid(idx) { + let new_view = if !array + .is_valid(idx) + .vortex_expect("failed to check validity") + { BinaryView::empty_view() } else if view.is_inlined() { view @@ -776,6 +785,7 @@ impl RewritingViewAdjustment { mod tests { use vortex_dtype::DType; use vortex_dtype::Nullability; + use vortex_error::VortexResult; use crate::IntoArray; use crate::arrays::VarBinViewArray; @@ -839,7 +849,7 @@ mod tests { } #[test] - fn test_buffer_deduplication() { + fn test_buffer_deduplication() -> VortexResult<()> { let array = { let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 10); @@ -852,11 +862,11 @@ mod tests { let mut builder = VarBinViewBuilder::with_buffer_deduplication(DType::Utf8(Nullability::Nullable), 10); - array.append_to_builder(&mut builder); + array.append_to_builder(&mut builder)?; assert_eq!(builder.completed_block_count(), 1); - array.slice(1..2).append_to_builder(&mut builder); - array.slice(0..1).append_to_builder(&mut builder); + array.slice(1..2).append_to_builder(&mut builder)?; + array.slice(0..1).append_to_builder(&mut builder)?; assert_eq!(builder.completed_block_count(), 1); let array2 = { @@ -866,12 +876,13 @@ mod tests { builder.finish_into_varbinview() }; - array2.append_to_builder(&mut builder); + array2.append_to_builder(&mut builder)?; assert_eq!(builder.completed_block_count(), 2); - array.slice(0..1).append_to_builder(&mut builder); - array2.slice(0..1).append_to_builder(&mut builder); + array.slice(0..1).append_to_builder(&mut builder)?; + array2.slice(0..1).append_to_builder(&mut builder)?; assert_eq!(builder.completed_block_count(), 2); + Ok(()) } #[test] diff --git a/vortex-array/src/canonical.rs b/vortex-array/src/canonical.rs index 34fa4b440e2..66c06a7e72f 100644 --- a/vortex-array/src/canonical.rs +++ b/vortex-array/src/canonical.rs @@ -308,74 +308,74 @@ impl IntoArray for Canonical { /// This trait has a blanket implementation for all types implementing [ToCanonical]. pub trait ToCanonical { /// Canonicalize into a [`NullArray`] if the target is [`Null`](DType::Null) typed. - fn to_null(&self) -> NullArray; + fn to_null(&self) -> VortexResult; /// Canonicalize into a [`BoolArray`] if the target is [`Bool`](DType::Bool) typed. - fn to_bool(&self) -> BoolArray; + fn to_bool(&self) -> VortexResult; /// Canonicalize into a [`PrimitiveArray`] if the target is [`Primitive`](DType::Primitive) /// typed. - fn to_primitive(&self) -> PrimitiveArray; + fn to_primitive(&self) -> VortexResult; /// Canonicalize into a [`DecimalArray`] if the target is [`Decimal`](DType::Decimal) /// typed. - fn to_decimal(&self) -> DecimalArray; + fn to_decimal(&self) -> VortexResult; /// Canonicalize into a [`StructArray`] if the target is [`Struct`](DType::Struct) typed. - fn to_struct(&self) -> StructArray; + fn to_struct(&self) -> VortexResult; /// Canonicalize into a [`ListViewArray`] if the target is [`List`](DType::List) typed. - fn to_listview(&self) -> ListViewArray; + fn to_listview(&self) -> VortexResult; /// Canonicalize into a [`FixedSizeListArray`] if the target is [`List`](DType::FixedSizeList) /// typed. - fn to_fixed_size_list(&self) -> FixedSizeListArray; + fn to_fixed_size_list(&self) -> VortexResult; /// Canonicalize into a [`VarBinViewArray`] if the target is [`Utf8`](DType::Utf8) /// or [`Binary`](DType::Binary) typed. - fn to_varbinview(&self) -> VarBinViewArray; + fn to_varbinview(&self) -> VortexResult; /// Canonicalize into an [`ExtensionArray`] if the array is [`Extension`](DType::Extension) /// typed. - fn to_extension(&self) -> ExtensionArray; + fn to_extension(&self) -> VortexResult; } // Blanket impl for all Array encodings. impl ToCanonical for A { - fn to_null(&self) -> NullArray { - self.to_canonical().into_null() + fn to_null(&self) -> VortexResult { + Ok(self.to_canonical()?.into_null()) } - fn to_bool(&self) -> BoolArray { - self.to_canonical().into_bool() + fn to_bool(&self) -> VortexResult { + Ok(self.to_canonical()?.into_bool()) } - fn to_primitive(&self) -> PrimitiveArray { - self.to_canonical().into_primitive() + fn to_primitive(&self) -> VortexResult { + Ok(self.to_canonical()?.into_primitive()) } - fn to_decimal(&self) -> DecimalArray { - self.to_canonical().into_decimal() + fn to_decimal(&self) -> VortexResult { + Ok(self.to_canonical()?.into_decimal()) } - fn to_struct(&self) -> StructArray { - self.to_canonical().into_struct() + fn to_struct(&self) -> VortexResult { + Ok(self.to_canonical()?.into_struct()) } - fn to_listview(&self) -> ListViewArray { - self.to_canonical().into_listview() + fn to_listview(&self) -> VortexResult { + Ok(self.to_canonical()?.into_listview()) } - fn to_fixed_size_list(&self) -> FixedSizeListArray { - self.to_canonical().into_fixed_size_list() + fn to_fixed_size_list(&self) -> VortexResult { + Ok(self.to_canonical()?.into_fixed_size_list()) } - fn to_varbinview(&self) -> VarBinViewArray { - self.to_canonical().into_varbinview() + fn to_varbinview(&self) -> VortexResult { + Ok(self.to_canonical()?.into_varbinview()) } - fn to_extension(&self) -> ExtensionArray { - self.to_canonical().into_extension() + fn to_extension(&self) -> VortexResult { + Ok(self.to_canonical()?.into_extension()) } } diff --git a/vortex-array/src/compute/between.rs b/vortex-array/src/compute/between.rs index ee0e8409ee6..12af93d4eb3 100644 --- a/vortex-array/src/compute/between.rs +++ b/vortex-array/src/compute/between.rs @@ -121,7 +121,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()) { @@ -316,7 +316,8 @@ mod tests { }, ) .unwrap() - .to_bool(); + .to_bool() + .unwrap(); let indices = to_int_indices(matches).unwrap(); assert_eq!(indices, expected); @@ -343,7 +344,8 @@ mod tests { }, ) .unwrap() - .to_bool(); + .to_bool() + .unwrap(); let indices = to_int_indices(matches).unwrap(); assert!(indices.is_empty()); @@ -360,7 +362,8 @@ mod tests { }, ) .unwrap() - .to_bool(); + .to_bool() + .unwrap(); let indices = to_int_indices(matches).unwrap(); assert_eq!(indices, vec![0, 1, 3]); @@ -377,7 +380,8 @@ mod tests { }, ) .unwrap() - .to_bool(); + .to_bool() + .unwrap(); let indices = to_int_indices(matches).unwrap(); assert_eq!(indices, vec![0, 1, 2, 3, 4]); } diff --git a/vortex-array/src/compute/boolean.rs b/vortex-array/src/compute/boolean.rs index 9791e0b18fb..257ee7d8713 100644 --- a/vortex-array/src/compute/boolean.rs +++ b/vortex-array/src/compute/boolean.rs @@ -336,7 +336,7 @@ mod tests { fn test_or(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) { let r = or(&lhs, &rhs).unwrap(); - let r = r.to_bool().into_array(); + let r = r.to_bool().unwrap().into_array(); let v0 = r.scalar_at(0).as_bool().value(); let v1 = r.scalar_at(1).as_bool().value(); @@ -356,7 +356,7 @@ mod tests { #[case(BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)].into_iter()).into_array(), BoolArray::from_iter([Some(true), Some(true), Some(false), Some(false)].into_iter()).into_array())] fn test_and(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) { - let r = and(&lhs, &rhs).unwrap().to_bool().into_array(); + let r = and(&lhs, &rhs).unwrap().to_bool().unwrap().into_array(); let v0 = r.scalar_at(0).as_bool().value(); let v1 = r.scalar_at(1).as_bool().value(); diff --git a/vortex-array/src/compute/cast.rs b/vortex-array/src/compute/cast.rs index 88ca78917ab..97891144e9d 100644 --- a/vortex-array/src/compute/cast.rs +++ b/vortex-array/src/compute/cast.rs @@ -85,7 +85,7 @@ impl ComputeFnVTable for Cast { ); } - Ok(cast(array.to_canonical().as_ref(), dtype)?.into()) + Ok(cast(array.to_canonical()?.as_ref(), dtype)?.into()) } fn return_dtype(&self, args: &InvocationArgs) -> VortexResult { diff --git a/vortex-array/src/compute/compare.rs b/vortex-array/src/compute/compare.rs index 7f8b67b5634..1b5a83743d2 100644 --- a/vortex-array/src/compute/compare.rs +++ b/vortex-array/src/compute/compare.rs @@ -422,13 +422,15 @@ mod tests { let matches = compare(arr.as_ref(), arr.as_ref(), Operator::Eq) .unwrap() - .to_bool(); + .to_bool() + .unwrap(); assert_eq!(to_int_indices(matches).unwrap(), [1u64, 2, 3, 4]); let matches = compare(arr.as_ref(), arr.as_ref(), Operator::NotEq) .unwrap() - .to_bool(); + .to_bool() + .unwrap(); let empty: [u64; 0] = []; assert_eq!(to_int_indices(matches).unwrap(), empty); @@ -439,22 +441,26 @@ mod tests { let matches = compare(arr.as_ref(), other.as_ref(), Operator::Lte) .unwrap() - .to_bool(); + .to_bool() + .unwrap(); assert_eq!(to_int_indices(matches).unwrap(), [2u64, 3, 4]); let matches = compare(arr.as_ref(), other.as_ref(), Operator::Lt) .unwrap() - .to_bool(); + .to_bool() + .unwrap(); assert_eq!(to_int_indices(matches).unwrap(), [4u64]); let matches = compare(other.as_ref(), arr.as_ref(), Operator::Gte) .unwrap() - .to_bool(); + .to_bool() + .unwrap(); assert_eq!(to_int_indices(matches).unwrap(), [2u64, 3, 4]); let matches = compare(other.as_ref(), arr.as_ref(), Operator::Gt) .unwrap() - .to_bool(); + .to_bool() + .unwrap(); assert_eq!(to_int_indices(matches).unwrap(), [4u64]); } @@ -495,7 +501,7 @@ mod tests { #[case(VarBinViewArray::from_iter_bin(["a".as_bytes(), "b".as_bytes()]).into_array(), VarBinArray::from(vec!["a".as_bytes(), "b".as_bytes()]).into_array())] fn arrow_compare_different_encodings(#[case] left: ArrayRef, #[case] right: ArrayRef) { let res = compare(&left, &right, Operator::Eq).unwrap(); - assert_eq!(res.to_bool().bit_buffer().true_count(), left.len()); + assert_eq!(res.to_bool().unwrap().bit_buffer().true_count(), left.len()); } #[ignore = "Arrow's ListView cannot be compared"] @@ -522,21 +528,21 @@ mod tests { // Test equality - first two lists should be equal, third should be different let result = compare(list1.as_ref(), list2.as_ref(), Operator::Eq).unwrap(); - let bool_result = result.to_bool(); + let bool_result = result.to_bool().unwrap(); assert!(bool_result.bit_buffer().value(0)); // [1,2] == [1,2] assert!(bool_result.bit_buffer().value(1)); // [3,4] == [3,4] assert!(!bool_result.bit_buffer().value(2)); // [5,6] != [7,8] // Test inequality let result = compare(list1.as_ref(), list2.as_ref(), Operator::NotEq).unwrap(); - let bool_result = result.to_bool(); + let bool_result = result.to_bool().unwrap(); assert!(!bool_result.bit_buffer().value(0)); assert!(!bool_result.bit_buffer().value(1)); assert!(bool_result.bit_buffer().value(2)); // Test less than let result = compare(list1.as_ref(), list2.as_ref(), Operator::Lt).unwrap(); - let bool_result = result.to_bool(); + let bool_result = result.to_bool().unwrap(); assert!(!bool_result.bit_buffer().value(0)); // [1,2] < [1,2] = false assert!(!bool_result.bit_buffer().value(1)); // [3,4] < [3,4] = false assert!(bool_result.bit_buffer().value(2)); // [5,6] < [7,8] = true @@ -570,7 +576,7 @@ mod tests { // Compare list with constant - all should be compared to [3,4] let result = compare(list.as_ref(), constant.as_ref(), Operator::Eq).unwrap(); - let bool_result = result.to_bool(); + let bool_result = result.to_bool().unwrap(); assert!(!bool_result.bit_buffer().value(0)); // [1,2] != [3,4] assert!(bool_result.bit_buffer().value(1)); // [3,4] == [3,4] assert!(!bool_result.bit_buffer().value(2)); // [5,6] != [3,4] @@ -599,14 +605,14 @@ mod tests { // Test equality let result = compare(struct1.as_ref(), struct2.as_ref(), Operator::Eq).unwrap(); - let bool_result = result.to_bool(); + let bool_result = result.to_bool().unwrap(); assert!(bool_result.bit_buffer().value(0)); // {true, 1} == {true, 1} assert!(bool_result.bit_buffer().value(1)); // {false, 2} == {false, 2} assert!(!bool_result.bit_buffer().value(2)); // {true, 3} != {false, 4} // Test greater than let result = compare(struct1.as_ref(), struct2.as_ref(), Operator::Gt).unwrap(); - let bool_result = result.to_bool(); + let bool_result = result.to_bool().unwrap(); assert!(!bool_result.bit_buffer().value(0)); // {true, 1} > {true, 1} = false assert!(!bool_result.bit_buffer().value(1)); // {false, 2} > {false, 2} = false assert!(bool_result.bit_buffer().value(2)); // {true, 3} > {false, 4} = true (bool field takes precedence) @@ -631,7 +637,7 @@ mod tests { .unwrap(); let result = compare(empty1.as_ref(), empty2.as_ref(), Operator::Eq).unwrap(); - let result = result.to_bool(); + let result = result.to_bool().unwrap(); for idx in 0..5 { assert!(result.bit_buffer().value(idx)); diff --git a/vortex-array/src/compute/conformance/binary_numeric.rs b/vortex-array/src/compute/conformance/binary_numeric.rs index 78e36147acc..15e06416e8e 100644 --- a/vortex-array/src/compute/conformance/binary_numeric.rs +++ b/vortex-array/src/compute/conformance/binary_numeric.rs @@ -88,7 +88,7 @@ fn test_standard_binary_numeric(array: ArrayRef) where Scalar: From, { - let canonicalized_array = array.to_primitive(); + let canonicalized_array = array.to_primitive().vortex_unwrap(); let original_values = to_vec_of_scalar(&canonicalized_array.into_array()); let one = T::from(1) @@ -322,7 +322,7 @@ where T: NativePType + Num + Copy + std::fmt::Debug, Scalar: From, { - let canonicalized_array = array.to_primitive(); + let canonicalized_array = array.to_primitive().vortex_unwrap(); let original_values = to_vec_of_scalar(&canonicalized_array.into_array()); let scalar = Scalar::from(scalar_value) diff --git a/vortex-array/src/compute/conformance/cast.rs b/vortex-array/src/compute/conformance/cast.rs index 691dc22ae4c..bbebad61d71 100644 --- a/vortex-array/src/compute/conformance/cast.rs +++ b/vortex-array/src/compute/conformance/cast.rs @@ -100,7 +100,7 @@ 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().unwrap() == 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()); @@ -223,7 +223,10 @@ 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_unwrap(), + casted.validity_mask().vortex_unwrap() + ); for i in 0..array.len().min(10) { let original = array.scalar_at(i); let casted = casted.scalar_at(i); diff --git a/vortex-array/src/compute/conformance/consistency.rs b/vortex-array/src/compute/conformance/consistency.rs index df9797849d1..9ce04d87fef 100644 --- a/vortex-array/src/compute/conformance/consistency.rs +++ b/vortex-array/src/compute/conformance/consistency.rs @@ -821,17 +821,17 @@ fn test_slice_aggregate_consistency(array: &dyn Array) { // Get sliced array and canonical slice let sliced = array.slice(start..end); - let canonical = array.to_canonical(); + let canonical = array.to_canonical().vortex_unwrap(); let canonical_sliced = canonical.as_ref().slice(start..end); // Test null count through invalid_count assert_eq!( - sliced.invalid_count(), - canonical_sliced.invalid_count(), + sliced.invalid_count().unwrap(), + canonical_sliced.invalid_count().unwrap(), "null_count on sliced array should match canonical. \ Sliced: {}, Canonical: {}", - sliced.invalid_count(), - canonical_sliced.invalid_count() + sliced.invalid_count().unwrap(), + canonical_sliced.invalid_count().unwrap() ); // Test sum for numeric types @@ -910,7 +910,7 @@ fn test_cast_slice_consistency(array: &dyn Array) { let end = 7.min(len - 2).max(start + 1); // Ensure we have at least 1 element // Get canonical form of the original array - let canonical = array.to_canonical(); + let canonical = array.to_canonical().vortex_unwrap(); // Choose appropriate target dtype based on the array's type let target_dtypes = match array.dtype() { diff --git a/vortex-array/src/compute/conformance/mask.rs b/vortex-array/src/compute/conformance/mask.rs index 719b0e0b406..62ab26e6f48 100644 --- a/vortex-array/src/compute/conformance/mask.rs +++ b/vortex-array/src/compute/conformance/mask.rs @@ -45,7 +45,7 @@ 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).unwrap()); } else { assert_eq!(masked.scalar_at(i), array.scalar_at(i).into_nullable()); } @@ -78,7 +78,7 @@ 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).unwrap()); } } @@ -93,7 +93,7 @@ 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).unwrap()); } else { assert_eq!(masked.scalar_at(i), array.scalar_at(i).into_nullable()); } @@ -115,13 +115,13 @@ 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).unwrap()).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).unwrap()) .count(); assert_eq!(valid_count, len - expected_invalid_count); @@ -137,7 +137,7 @@ fn test_single_element_mask(array: &dyn Array) { let mask_array = Mask::from_iter(pattern); let masked = mask(array, &mask_array).vortex_unwrap(); - assert!(!masked.is_valid(0)); + assert!(!masked.is_valid(0).unwrap()); for i in 1..len { assert_eq!(masked.scalar_at(i), array.scalar_at(i).into_nullable()); @@ -161,7 +161,7 @@ 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).unwrap()); } else { assert_eq!( double_masked.scalar_at(i), @@ -192,7 +192,7 @@ 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).unwrap()); } else { assert_eq!(masked.scalar_at(i), array.scalar_at(i).into_nullable()); } diff --git a/vortex-array/src/compute/conformance/mod.rs b/vortex-array/src/compute/conformance/mod.rs index 68d80a0998e..7027d379d79 100644 --- a/vortex-array/src/compute/conformance/mod.rs +++ b/vortex-array/src/compute/conformance/mod.rs @@ -1,6 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +// Conformance tests are allowed to use unwrap for brevity +#![allow(clippy::unwrap_used)] +#![allow(clippy::unwrap_in_result)] + pub mod binary_numeric; pub mod cast; pub mod consistency; diff --git a/vortex-array/src/compute/conformance/take.rs b/vortex-array/src/compute/conformance/take.rs index 453e65f7671..62f728b57aa 100644 --- a/vortex-array/src/compute/conformance/take.rs +++ b/vortex-array/src/compute/conformance/take.rs @@ -61,7 +61,10 @@ fn test_take_all(array: &dyn Array) { assert_eq!(result.dtype(), array.dtype()); // Verify elements match - match (&array.to_canonical(), &result.to_canonical()) { + match ( + &array.to_canonical().vortex_unwrap(), + &result.to_canonical().vortex_unwrap(), + ) { (Canonical::Primitive(orig_prim), Canonical::Primitive(result_prim)) => { assert_eq!(orig_prim.byte_buffer(), result_prim.byte_buffer()); } diff --git a/vortex-array/src/compute/fill_null.rs b/vortex-array/src/compute/fill_null.rs index 0bbb74b07c3..4e0f15936b7 100644 --- a/vortex-array/src/compute/fill_null.rs +++ b/vortex-array/src/compute/fill_null.rs @@ -96,11 +96,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()); @@ -121,7 +121,7 @@ impl ComputeFnVTable for FillNull { log::debug!("FillNullFn not implemented for {}", array.encoding_id()); if !array.is_canonical() { - let canonical_arr = array.to_canonical().into_array(); + let canonical_arr = array.to_canonical()?.into_array(); return Ok(fill_null(canonical_arr.as_ref(), fill_value)?.into()); } diff --git a/vortex-array/src/compute/filter.rs b/vortex-array/src/compute/filter.rs index 4f10dbe88cb..2028efb8434 100644 --- a/vortex-array/src/compute/filter.rs +++ b/vortex-array/src/compute/filter.rs @@ -106,7 +106,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() @@ -135,7 +135,7 @@ impl ComputeFnVTable for Filter { log::debug!("No filter implementation found for {}", array.encoding_id(),); if !array.is_canonical() { - let canonical = array.to_canonical().into_array(); + let canonical = array.to_canonical()?.into_array(); return filter(&canonical, mask).map(Into::into); }; @@ -237,7 +237,7 @@ impl dyn Array + '_ { // Convert nulls to false first in case this can be done cheaply by the encoding. let array = fill_null(self, &Scalar::bool(false, self.dtype().nullability()))?; - Ok(array.to_bool().to_mask_fill_null_false()) + Ok(array.to_bool()?.to_mask_fill_null_false()) } } @@ -273,7 +273,7 @@ mod test { let filtered = filter(&items, &mask).unwrap(); assert_eq!( - filtered.to_primitive().as_slice::(), + filtered.to_primitive().unwrap().as_slice::(), &[0i32, 1i32, 2i32] ); } diff --git a/vortex-array/src/compute/invert.rs b/vortex-array/src/compute/invert.rs index 44a755ccf2e..ea5185c62b6 100644 --- a/vortex-array/src/compute/invert.rs +++ b/vortex-array/src/compute/invert.rs @@ -72,7 +72,7 @@ impl ComputeFnVTable for Invert { if array.is_canonical() { vortex_panic!("Canonical bool array does not implement invert"); } - Ok(invert(&array.to_bool().into_array())?.into()) + Ok(invert(&array.to_bool()?.into_array())?.into()) } fn return_dtype(&self, args: &InvocationArgs) -> VortexResult { diff --git a/vortex-array/src/compute/is_constant.rs b/vortex-array/src/compute/is_constant.rs index 7dd31e286f0..7cf03dff6cc 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 { @@ -191,7 +191,7 @@ fn is_constant_impl( ); if options.cost == Cost::Canonicalize && !array.is_canonical() { - let array = array.to_canonical(); + let array = array.to_canonical()?; let is_constant = is_constant_opts(array.as_ref(), options)?; return Ok(is_constant); } diff --git a/vortex-array/src/compute/is_sorted.rs b/vortex-array/src/compute/is_sorted.rs index 3c29d4b58d1..c96595c4edb 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)); } } @@ -290,7 +290,7 @@ fn is_sorted_impl( ); // Recurse to canonical implementation - let array = array.to_canonical(); + let array = array.to_canonical()?; return if strict { is_strict_sorted(array.as_ref()) diff --git a/vortex-array/src/compute/list_contains.rs b/vortex-array/src/compute/list_contains.rs index 95c01bddc76..804c6d17d38 100644 --- a/vortex-array/src/compute/list_contains.rs +++ b/vortex-array/src/compute/list_contains.rs @@ -99,7 +99,7 @@ pub(crate) fn warm_up_vtable() -> usize { /// list_array.len()).as_ref() /// ).unwrap(); /// -/// assert_eq!(matches.to_bool().bit_buffer(), &bitbuffer![false, true, false]); +/// assert_eq!(matches.to_bool().unwrap().bit_buffer(), &bitbuffer![false, true, false]); /// ``` pub fn list_contains(array: &dyn Array, value: &dyn Array) -> VortexResult { LIST_CONTAINS_FN @@ -135,7 +135,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( return Ok(ConstantArray::new(contains.scalar_at(0), array.len()).into_array()); } - let list_array = array.to_listview(); + let list_array = array.to_listview()?; let elems = list_array.elements(); if elems.is_empty() { @@ -270,7 +270,7 @@ fn list_contains_scalar( let rhs = ConstantArray::new(value.clone(), elems.len()); let matching_elements = compute::compare(elems, rhs.as_ref(), Operator::Eq)?; - let matches = matching_elements.to_bool(); + let matches = matching_elements.to_bool()?; // Fast path: no elements match. if let Some(pred) = matches.as_constant() { @@ -302,8 +302,8 @@ fn list_contains_scalar( } // Get the offsets and sizes as primitive arrays. - let offsets = list_array.offsets().to_primitive(); - let sizes = list_array.sizes().to_primitive(); + let offsets = list_array.offsets().to_primitive()?; + let sizes = list_array.sizes().to_primitive()?; // Process based on the offset and size types. let list_matches = match_each_integer_ptype!(offsets.ptype(), |O| { @@ -400,7 +400,7 @@ fn list_is_not_empty( .into_array()); } - let sizes = list_array.sizes().to_primitive(); + let sizes = list_array.sizes().to_primitive()?; let buffer = match_each_integer_ptype!(sizes.ptype(), |S| { BitBuffer::from_iter(sizes.as_slice::().iter().map(|&size| size != S::zero())) }); @@ -544,7 +544,7 @@ mod tests { }; let elem = ConstantArray::new(scalar, list_array.len()); let result = list_contains(&list_array, elem.as_ref()).expect("list_contains failed"); - let bool_result = result.to_bool(); + let bool_result = result.to_bool().unwrap(); assert_eq!(bool_result.opt_bool_vec(), expected.opt_bool_vec()); assert_eq!(bool_result.validity(), expected.validity()); } @@ -567,7 +567,10 @@ mod tests { ) .unwrap(); assert!(contains.is::(), "Expected constant result"); - assert_eq!(contains.to_bool().bit_buffer(), &bitbuffer![true, true],); + assert_eq!( + contains.to_bool().unwrap().bit_buffer(), + &bitbuffer![true, true], + ); } #[test] @@ -589,7 +592,10 @@ mod tests { assert!(contains.is::(), "Expected constant result"); assert_eq!(contains.len(), 5); - assert_eq!(contains.to_bool().validity(), &Validity::AllInvalid); + assert_eq!( + contains.to_bool().unwrap().validity(), + &Validity::AllInvalid + ); } #[test] @@ -608,7 +614,7 @@ mod tests { assert_eq!(contains.len(), 7); assert_eq!( - contains.to_bool().opt_bool_vec(), + contains.to_bool().unwrap().opt_bool_vec(), vec![ Some(false), Some(true), @@ -645,7 +651,7 @@ mod tests { // All lists are empty, so all should return false assert_eq!(result.len(), 4); assert_eq!( - result.to_bool().bool_vec(), + result.to_bool().unwrap().bool_vec(), vec![false, false, false, false] ); } @@ -676,7 +682,7 @@ mod tests { // Searching for null in lists with null elements should return null assert_eq!(result.len(), 3); - assert_eq!(result.to_bool().validity(), &Validity::AllInvalid); + assert_eq!(result.to_bool().unwrap().validity(), &Validity::AllInvalid); // Test searching for a non-null value let non_null_search = ConstantArray::new(Scalar::from(42i32), list_array.len()); @@ -684,7 +690,10 @@ mod tests { // All comparisons result in null, but search is not null, so should return false assert_eq!(result2.len(), 3); - assert_eq!(result2.to_bool().bool_vec(), vec![false, false, false]); + assert_eq!( + result2.to_bool().unwrap().bool_vec(), + vec![false, false, false] + ); } #[test] @@ -710,7 +719,7 @@ mod tests { assert_eq!(result.len(), 4); assert_eq!( - result.to_bool().bool_vec(), + result.to_bool().unwrap().bool_vec(), vec![false, true, false, false] // Value 2 is only in list 1 ); @@ -719,7 +728,7 @@ mod tests { let result5 = list_contains(list_array.as_ref(), search5.as_ref()).unwrap(); assert_eq!( - result5.to_bool().bool_vec(), + result5.to_bool().unwrap().bool_vec(), vec![false, false, true, false] // Value 5 is only in list 2 ); } @@ -742,14 +751,17 @@ mod tests { let result = list_contains(list_array.as_ref(), search.as_ref()).unwrap(); assert_eq!(result.len(), 4); - assert_eq!(result.to_bool().bool_vec(), vec![false, false, false, true]); + assert_eq!( + result.to_bool().unwrap().bool_vec(), + vec![false, false, false, true] + ); // Search for value 0 which should only be in the first list let search_zero = ConstantArray::new(Scalar::from(0i32), list_array.len()); let result_zero = list_contains(list_array.as_ref(), search_zero.as_ref()).unwrap(); assert_eq!( - result_zero.to_bool().bool_vec(), + result_zero.to_bool().unwrap().bool_vec(), vec![true, false, false, false] ); } diff --git a/vortex-array/src/compute/mask.rs b/vortex-array/src/compute/mask.rs index 86d59f5f8e4..92a824651b4 100644 --- a/vortex-array/src/compute/mask.rs +++ b/vortex-array/src/compute/mask.rs @@ -58,11 +58,11 @@ pub(crate) fn warm_up_vtable() -> usize { /// /// let masked = mask(array.as_ref(), &mask_array).unwrap(); /// assert_eq!(masked.len(), 5); -/// assert!(!masked.is_valid(0)); -/// assert!(!masked.is_valid(1)); +/// assert!(!masked.is_valid(0).unwrap()); +/// assert!(!masked.is_valid(1).unwrap()); /// assert_eq!(masked.scalar_at(2), Scalar::from(Some(1))); -/// assert!(!masked.is_valid(3)); -/// assert!(!masked.is_valid(4)); +/// assert!(!masked.is_valid(3).unwrap()); +/// assert!(!masked.is_valid(4).unwrap()); /// ``` /// pub fn mask(array: &dyn Array, mask: &Mask) -> VortexResult { @@ -127,7 +127,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 60a26e3ec4b..71139ac4808 100644 --- a/vortex-array/src/compute/min_max.rs +++ b/vortex-array/src/compute/min_max.rs @@ -143,7 +143,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); } @@ -182,7 +182,7 @@ fn min_max_impl( } if !array.is_canonical() { - let array = array.to_canonical(); + let array = array.to_canonical()?; return min_max(array.as_ref()); } diff --git a/vortex-array/src/compute/nan_count.rs b/vortex-array/src/compute/nan_count.rs index 8c31941ede8..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); } @@ -151,7 +151,7 @@ fn nan_count_impl(array: &dyn Array, kernels: &[ArcRef]) -> VortexRe } if !array.is_canonical() { - let canonical = array.to_canonical(); + let canonical = array.to_canonical()?; return nan_count(canonical.as_ref()); } diff --git a/vortex-array/src/compute/numeric.rs b/vortex-array/src/compute/numeric.rs index ccba7e8cc02..e974bf87290 100644 --- a/vortex-array/src/compute/numeric.rs +++ b/vortex-array/src/compute/numeric.rs @@ -296,6 +296,7 @@ mod test { let results = sub_scalar(&values, 1u16.into()) .unwrap() .to_primitive() + .unwrap() .as_slice::() .to_vec(); assert_eq!(results, &[0u16, 1, 2]); @@ -307,6 +308,7 @@ mod test { let results = sub_scalar(&values, (-1i64).into()) .unwrap() .to_primitive() + .unwrap() .as_slice::() .to_vec(); assert_eq!(results, &[2i64, 3, 4]); @@ -317,7 +319,8 @@ mod test { let values = PrimitiveArray::from_option_iter([Some(1u16), Some(2), None, Some(3)]); let result = sub_scalar(values.as_ref(), Some(1u16).into()) .unwrap() - .to_primitive(); + .to_primitive() + .unwrap(); let actual = (0..result.len()) .map(|index| result.scalar_at(index)) @@ -340,6 +343,7 @@ mod test { let results = sub_scalar(&values, to_subtract.into()) .unwrap() .to_primitive() + .unwrap() .as_slice::() .to_vec(); assert_eq!(results, &[2.0f64, 3.0, 4.0]); diff --git a/vortex-array/src/compute/sum.rs b/vortex-array/src/compute/sum.rs index b5f8b263dba..5ae44047560 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()); } @@ -264,7 +264,7 @@ pub fn sum_impl( array.encoding_id() ); } - sum_with_accumulator(array.to_canonical().as_ref(), accumulator) + sum_with_accumulator(array.to_canonical()?.as_ref(), accumulator) } #[cfg(test)] diff --git a/vortex-array/src/compute/take.rs b/vortex-array/src/compute/take.rs index 1db94d382a4..8ddc0e3b75a 100644 --- a/vortex-array/src/compute/take.rs +++ b/vortex-array/src/compute/take.rs @@ -80,7 +80,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(), @@ -131,7 +131,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()? { 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 @@ -150,7 +150,8 @@ fn propagate_take_stats( st.combine_sets( &(unsafe { StatsSet::new_unchecked(inexact_min_max) }).as_typed_ref(source.dtype()), ) - }) + })?; + Ok(()) } fn take_impl( @@ -186,7 +187,7 @@ fn take_impl( // Otherwise, canonicalize and try again. if !array.is_canonical() { log::debug!("No take implementation found for {}", array.encoding_id()); - let canonical = array.to_canonical(); + let canonical = array.to_canonical()?; return take(canonical.as_ref(), indices); } diff --git a/vortex-array/src/compute/zip.rs b/vortex-array/src/compute/zip.rs index 86b585eda89..8fff7df64fa 100644 --- a/vortex-array/src/compute/zip.rs +++ b/vortex-array/src/compute/zip.rs @@ -87,16 +87,16 @@ impl ComputeFnVTable for Zip { if !if_true.is_canonical() || !if_false.is_canonical() { return zip( - if_true.to_canonical().as_ref(), - if_false.to_canonical().as_ref(), + if_true.to_canonical()?.as_ref(), + if_false.to_canonical()?.as_ref(), mask, ) .map(Into::into); } Ok(zip_impl( - if_true.to_canonical().as_ref(), - if_false.to_canonical().as_ref(), + if_true.to_canonical()?.as_ref(), + if_false.to_canonical()?.as_ref(), mask, )? .into()) @@ -272,8 +272,8 @@ mod tests { let expected = buffer![10, 2, 3, 40, 5].into_array(); assert_eq!( - result.to_primitive().as_slice::(), - expected.to_primitive().as_slice::() + result.to_primitive().unwrap().as_slice::(), + expected.to_primitive().unwrap().as_slice::() ); } @@ -287,8 +287,8 @@ mod tests { let result = zip(&if_true, &if_false, &mask).unwrap(); assert_eq!( - result.to_primitive().as_slice::(), - if_true.to_primitive().as_slice::() + result.to_primitive().unwrap().as_slice::(), + if_true.to_primitive().unwrap().as_slice::() ); // result must be nullable even if_true was not diff --git a/vortex-array/src/display/mod.rs b/vortex-array/src/display/mod.rs index 469439f1caa..d566b3eb696 100644 --- a/vortex-array/src/display/mod.rs +++ b/vortex-array/src/display/mod.rs @@ -265,11 +265,13 @@ impl dyn Array + '_ { return write!(f, "{table}"); }; - let struct_ = self.to_struct(); + use vortex_error::VortexExpect as _; + + let struct_ = self.to_struct().vortex_expect("to_struct"); 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).vortex_expect("is_valid") { let null_row = vec!["null".to_string(); sf.names().len()]; builder.push_record(null_row); } else { @@ -291,7 +293,7 @@ impl dyn Array + '_ { } for row_idx in 0..self.len() { - if !self.is_valid(row_idx) { + if !self.is_valid(row_idx).vortex_expect("is_valid") { table.modify( (1 + row_idx, 0), tabled::settings::Span::column(sf.names().len() as isize), diff --git a/vortex-array/src/expr/exprs/binary.rs b/vortex-array/src/expr/exprs/binary.rs index 59fee0b2b20..7554888cf8d 100644 --- a/vortex-array/src/expr/exprs/binary.rs +++ b/vortex-array/src/expr/exprs/binary.rs @@ -304,7 +304,7 @@ impl ExpressionView<'_, Binary> { /// let result = eq(root(), lit(3)).evaluate(&xs.to_array()).unwrap(); /// /// assert_eq!( -/// result.to_bool().bit_buffer(), +/// result.to_bool().unwrap().bit_buffer(), /// BoolArray::from_iter(vec![false, false, true]).bit_buffer(), /// ); /// ``` @@ -328,7 +328,7 @@ pub fn eq(lhs: Expression, rhs: Expression) -> Expression { /// let result = not_eq(root(), lit(3)).evaluate(&xs.to_array()).unwrap(); /// /// assert_eq!( -/// result.to_bool().bit_buffer(), +/// result.to_bool().unwrap().bit_buffer(), /// BoolArray::from_iter(vec![true, true, false]).bit_buffer(), /// ); /// ``` @@ -352,7 +352,7 @@ pub fn not_eq(lhs: Expression, rhs: Expression) -> Expression { /// let result = gt_eq(root(), lit(3)).evaluate(&xs.to_array()).unwrap(); /// /// assert_eq!( -/// result.to_bool().bit_buffer(), +/// result.to_bool().unwrap().bit_buffer(), /// BoolArray::from_iter(vec![false, false, true]).bit_buffer(), /// ); /// ``` @@ -376,7 +376,7 @@ pub fn gt_eq(lhs: Expression, rhs: Expression) -> Expression { /// let result = gt(root(), lit(2)).evaluate(&xs.to_array()).unwrap(); /// /// assert_eq!( -/// result.to_bool().bit_buffer(), +/// result.to_bool().unwrap().bit_buffer(), /// BoolArray::from_iter(vec![false, false, true]).bit_buffer(), /// ); /// ``` @@ -400,7 +400,7 @@ pub fn gt(lhs: Expression, rhs: Expression) -> Expression { /// let result = lt_eq(root(), lit(2)).evaluate(&xs.to_array()).unwrap(); /// /// assert_eq!( -/// result.to_bool().bit_buffer(), +/// result.to_bool().unwrap().bit_buffer(), /// BoolArray::from_iter(vec![true, true, false]).bit_buffer(), /// ); /// ``` @@ -424,7 +424,7 @@ pub fn lt_eq(lhs: Expression, rhs: Expression) -> Expression { /// let result = lt(root(), lit(3)).evaluate(&xs.to_array()).unwrap(); /// /// assert_eq!( -/// result.to_bool().bit_buffer(), +/// result.to_bool().unwrap().bit_buffer(), /// BoolArray::from_iter(vec![true, true, false]).bit_buffer(), /// ); /// ``` @@ -446,7 +446,7 @@ pub fn lt(lhs: Expression, rhs: Expression) -> Expression { /// let result = or(root(), lit(false)).evaluate(&xs.to_array()).unwrap(); /// /// assert_eq!( -/// result.to_bool().bit_buffer(), +/// result.to_bool().unwrap().bit_buffer(), /// BoolArray::from_iter(vec![true, false, true]).bit_buffer(), /// ); /// ``` @@ -480,7 +480,7 @@ where /// let result = and(root(), lit(true)).evaluate(&xs.to_array()).unwrap(); /// /// assert_eq!( -/// result.to_bool().bit_buffer(), +/// result.to_bool().unwrap().bit_buffer(), /// BoolArray::from_iter(vec![true, false, true]).bit_buffer(), /// ); /// ``` diff --git a/vortex-array/src/expr/exprs/get_item/mod.rs b/vortex-array/src/expr/exprs/get_item/mod.rs index d0f8bb74ff8..c8efd483a25 100644 --- a/vortex-array/src/expr/exprs/get_item/mod.rs +++ b/vortex-array/src/expr/exprs/get_item/mod.rs @@ -104,12 +104,12 @@ impl VTable for GetItem { } fn evaluate(&self, expr: &ExpressionView, scope: &ArrayRef) -> VortexResult { - let input = expr.children()[0].evaluate(scope)?.to_struct(); + let input = expr.children()[0].evaluate(scope)?.to_struct()?; let field = input.field_by_name(expr.data()).cloned()?; match input.dtype().nullability() { Nullability::NonNullable => Ok(field), - Nullability::Nullable => mask(&field, &input.validity_mask().not()), + Nullability::Nullable => mask(&field, &input.validity_mask()?.not()), } } diff --git a/vortex-array/src/expr/exprs/is_null.rs b/vortex-array/src/expr/exprs/is_null.rs index 6e95351f6eb..1631634e933 100644 --- a/vortex-array/src/expr/exprs/is_null.rs +++ b/vortex-array/src/expr/exprs/is_null.rs @@ -78,7 +78,7 @@ impl VTable for IsNull { fn evaluate(&self, expr: &ExpressionView, scope: &ArrayRef) -> VortexResult { let array = expr.child(0).evaluate(scope)?; - match array.validity_mask() { + match array.validity_mask()? { Mask::AllTrue(len) => Ok(ConstantArray::new(false, len).into_array()), Mask::AllFalse(len) => Ok(ConstantArray::new(true, len).into_array()), Mask::Values(mask) => Ok(BoolArray::from(mask.bit_buffer().not()).into_array()), diff --git a/vortex-array/src/expr/exprs/like.rs b/vortex-array/src/expr/exprs/like.rs index 65f6d823dce..4461c715c22 100644 --- a/vortex-array/src/expr/exprs/like.rs +++ b/vortex-array/src/expr/exprs/like.rs @@ -171,6 +171,7 @@ mod tests { .evaluate(&bools.to_array()) .unwrap() .to_bool() + .unwrap() .bit_buffer() .iter() .collect::>(), diff --git a/vortex-array/src/expr/exprs/list_contains.rs b/vortex-array/src/expr/exprs/list_contains.rs index 8a4664b8e2a..05134a1a50b 100644 --- a/vortex-array/src/expr/exprs/list_contains.rs +++ b/vortex-array/src/expr/exprs/list_contains.rs @@ -277,7 +277,7 @@ mod tests { let item = expr.evaluate(&arr).unwrap(); assert_eq!(item.scalar_at(0), Scalar::bool(true, Nullability::Nullable)); - assert!(!item.is_valid(1)); + assert!(!item.is_valid(1).unwrap()); } #[test] diff --git a/vortex-array/src/expr/exprs/merge/mod.rs b/vortex-array/src/expr/exprs/merge/mod.rs index 3d2a2ec1ee5..6c6321c1bd1 100644 --- a/vortex-array/src/expr/exprs/merge/mod.rs +++ b/vortex-array/src/expr/exprs/merge/mod.rs @@ -137,7 +137,7 @@ impl VTable for Merge { if !array.dtype().is_struct() { vortex_bail!("merge expects struct input"); } - let array = array.to_struct(); + let array = array.to_struct()?; for (field_name, array) in array.names().iter().zip_eq(array.fields().iter().cloned()) { // Update or insert field. @@ -234,11 +234,11 @@ mod tests { vortex_bail!("empty field path"); }; - let mut array = array.to_struct().field_by_name(field)?.clone(); + let mut array = array.to_struct().unwrap().field_by_name(field)?.clone(); for field in field_path { - array = array.to_struct().field_by_name(field)?.clone(); + array = array.to_struct().unwrap().field_by_name(field)?.clone(); } - Ok(array.to_primitive()) + Ok(array.to_primitive().unwrap()) } #[test] @@ -418,13 +418,18 @@ mod tests { ]) .unwrap() .into_array(); - let actual_array = expr.evaluate(&test_array.clone()).unwrap().to_struct(); + let actual_array = expr + .evaluate(&test_array.clone()) + .unwrap() + .to_struct() + .unwrap(); assert_eq!( actual_array .field_by_name("a") .unwrap() .to_struct() + .unwrap() .names() .iter() .map(|name| name.as_ref()) @@ -459,7 +464,11 @@ mod tests { ]) .unwrap() .into_array(); - let actual_array = expr.evaluate(&test_array.clone()).unwrap().to_struct(); + let actual_array = expr + .evaluate(&test_array.clone()) + .unwrap() + .to_struct() + .unwrap(); assert_eq!(actual_array.names(), ["a", "c", "b", "d"]); } diff --git a/vortex-array/src/expr/exprs/not.rs b/vortex-array/src/expr/exprs/not.rs index 2a38bbf7677..919d2a0e068 100644 --- a/vortex-array/src/expr/exprs/not.rs +++ b/vortex-array/src/expr/exprs/not.rs @@ -125,6 +125,7 @@ mod tests { .evaluate(&bools.to_array()) .unwrap() .to_bool() + .unwrap() .bit_buffer() .iter() .collect::>(), diff --git a/vortex-array/src/expr/exprs/pack.rs b/vortex-array/src/expr/exprs/pack.rs index a03136ca6f4..ae52e360aba 100644 --- a/vortex-array/src/expr/exprs/pack.rs +++ b/vortex-array/src/expr/exprs/pack.rs @@ -231,11 +231,11 @@ mod tests { vortex_bail!("empty field path"); }; - let mut array = array.to_struct().field_by_name(field)?.clone(); + let mut array = array.to_struct()?.field_by_name(field)?.clone(); for field in field_path { - array = array.to_struct().field_by_name(field)?.clone(); + array = array.to_struct()?.field_by_name(field)?.clone(); } - Ok(array.to_primitive()) + Ok(array.to_primitive().unwrap()) } #[test] @@ -251,7 +251,10 @@ mod tests { let test_array = test_array(); let actual_array = expr.evaluate(&test_array.clone()).unwrap(); assert_eq!(actual_array.len(), test_array.len()); - assert_eq!(actual_array.to_struct().struct_fields().nfields(), 0); + assert_eq!( + actual_array.to_struct().unwrap().struct_fields().nfields(), + 0 + ); } #[test] @@ -264,7 +267,7 @@ mod tests { [col("a"), col("b"), col("a")], ); - let actual_array = expr.evaluate(&test_array()).unwrap().to_struct(); + let actual_array = expr.evaluate(&test_array()).unwrap().to_struct().unwrap(); assert_eq!(actual_array.names(), ["one", "two", "three"]); assert_eq!(actual_array.validity(), &Validity::NonNullable); @@ -309,7 +312,7 @@ mod tests { ], ); - let actual_array = expr.evaluate(&test_array()).unwrap().to_struct(); + let actual_array = expr.evaluate(&test_array()).unwrap().to_struct().unwrap(); assert_eq!(actual_array.names(), ["one", "two", "three"]); @@ -349,7 +352,7 @@ mod tests { [col("a"), col("b"), col("a")], ); - let actual_array = expr.evaluate(&test_array()).unwrap().to_struct(); + let actual_array = expr.evaluate(&test_array()).unwrap().to_struct().unwrap(); assert_eq!(actual_array.names(), ["one", "two", "three"]); assert_eq!(actual_array.validity(), &Validity::AllValid); diff --git a/vortex-array/src/expr/exprs/select/mod.rs b/vortex-array/src/expr/exprs/select/mod.rs index 8b27fb73d3a..d5f44237b90 100644 --- a/vortex-array/src/expr/exprs/select/mod.rs +++ b/vortex-array/src/expr/exprs/select/mod.rs @@ -145,7 +145,7 @@ impl VTable for Select { } fn evaluate(&self, expr: &ExpressionView, scope: &ArrayRef) -> VortexResult { - let batch = expr.child().evaluate(scope)?.to_struct(); + let batch = expr.child().evaluate(scope)?.to_struct()?; Ok(match expr.data() { FieldSelection::Include(f) => batch.project(f.as_ref()), FieldSelection::Exclude(names) => { @@ -353,7 +353,11 @@ mod tests { pub fn include_columns() { let st = test_array(); let select = select(vec![FieldName::from("a")], root()); - let selected = select.evaluate(&st.to_array()).unwrap().to_struct(); + let selected = select + .evaluate(&st.to_array()) + .unwrap() + .to_struct() + .unwrap(); let selected_names = selected.names().clone(); assert_eq!(selected_names.as_ref(), &["a"]); } @@ -362,7 +366,11 @@ mod tests { pub fn exclude_columns() { let st = test_array(); let select = select_exclude(vec![FieldName::from("a")], root()); - let selected = select.evaluate(&st.to_array()).unwrap().to_struct(); + let selected = select + .evaluate(&st.to_array()) + .unwrap() + .to_struct() + .unwrap(); let selected_names = selected.names().clone(); assert_eq!(selected_names.as_ref(), &["b"]); } diff --git a/vortex-array/src/patches.rs b/vortex-array/src/patches.rs index 1c9e35e7e7f..3acf12c355b 100644 --- a/vortex-array/src/patches.rs +++ b/vortex-array/src/patches.rs @@ -379,7 +379,9 @@ impl Patches { /// with the insertion point if not found. fn search_index_binary_search(indices: &dyn Array, needle: usize) -> SearchResult { if indices.is_canonical() { - let primitive = indices.to_primitive(); + let Ok(primitive) = indices.to_primitive() else { + return SearchResult::NotFound(0); + }; match_each_integer_ptype!(primitive.ptype(), |T| { let Ok(needle) = T::try_from(needle) else { // If the needle is not of type T, then it cannot possibly be in this array. @@ -566,7 +568,7 @@ impl Patches { AllOr::All => Ok(Some(self.clone())), AllOr::None => Ok(None), AllOr::Some(mask_indices) => { - let flat_indices = self.indices().to_primitive(); + let flat_indices = self.indices().to_primitive()?; match_each_unsigned_integer_ptype!(flat_indices.ptype(), |I| { filter_patches_with_mask( flat_indices.as_slice::(), @@ -595,7 +597,7 @@ impl Patches { AllOr::All => return Ok(None), AllOr::None => return Ok(Some(self.clone())), AllOr::Some(masked) => { - let patch_indices = self.indices().to_primitive(); + let patch_indices = self.indices().to_primitive()?; match_each_unsigned_integer_ptype!(patch_indices.ptype(), |P| { let patch_indices = patch_indices.as_slice::

(); Mask::from_buffer(BitBuffer::collect_bool(patch_indices.len(), |i| { @@ -680,7 +682,7 @@ impl Patches { return Ok(None); } - let take_indices = take_indices.to_primitive(); + let take_indices = take_indices.to_primitive()?; if self.is_map_faster_than_search(&take_indices) { self.take_map(take_indices, true) } else { @@ -696,7 +698,7 @@ impl Patches { return Ok(None); } - let take_indices = take_indices.to_primitive(); + let take_indices = take_indices.to_primitive()?; if self.is_map_faster_than_search(&take_indices) { self.take_map(take_indices, false) } else { @@ -714,8 +716,12 @@ impl Patches { include_nulls: bool, ) -> VortexResult> { let take_indices_validity = take_indices.validity(); - let patch_indices = self.indices.to_primitive(); - let chunk_offsets = self.chunk_offsets().as_ref().map(|co| co.to_primitive()); + let patch_indices = self.indices.to_primitive()?; + let chunk_offsets = self + .chunk_offsets() + .as_ref() + .map(|co| co.to_primitive()) + .transpose()?; let (values_indices, new_indices): (BufferMut, BufferMut) = match_each_unsigned_integer_ptype!(patch_indices.ptype(), |PatchT| { @@ -729,7 +735,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( @@ -744,7 +750,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| { let Some(offset) = ::from(self.offset) else { @@ -786,7 +792,7 @@ impl Patches { take_indices: PrimitiveArray, include_nulls: bool, ) -> VortexResult> { - let indices = self.indices.to_primitive(); + let indices = self.indices.to_primitive()?; let new_length = take_indices.len(); let Some((new_sparse_indices, value_indices)) = @@ -830,8 +836,14 @@ impl Patches { /// - All patch indices after offset adjustment must be valid indices into the buffer. /// - The buffer and validity mask must have the same length. pub unsafe fn apply_to_buffer(&self, buffer: &mut [P], validity: &mut MaskMut) { - let patch_indices = self.indices.to_primitive(); - let patch_values = self.values.to_primitive(); + let patch_indices = self + .indices + .to_primitive() + .vortex_expect("patch indices must be primitive"); + let patch_values = self + .values + .to_primitive() + .vortex_expect("patch values must be primitive"); let patches_validity = patch_values.validity(); let patch_values_slice = patch_values.as_slice::

(); @@ -932,7 +944,7 @@ unsafe fn apply_patches_to_buffer_inner( } Validity::Array(array) => { // Some patch values may be null, check each one. - let bool_array = array.to_bool(); + let bool_array = array.to_bool().vortex_expect("validity array must be bool"); let mask = bool_array.bit_buffer(); for (patch_idx, (&i, &value)) in patch_indices.iter().zip_eq(patch_values).enumerate() { let index = i.as_() - patch_offset; @@ -984,7 +996,8 @@ where iter.enumerate() .filter_map(|(idx_in_take, ti)| { // If we have to take nulls the take index doesn't matter, make it 0 for consistency - if include_nulls && take_indices_validity.is_null(idx_in_take) { + if include_nulls && take_indices_validity.is_null(idx_in_take).unwrap_or(false) + { Some((idx_in_take as u64, 0)) } else if ti < min_index || ti > max_index { None @@ -1163,8 +1176,8 @@ mod test { .unwrap() .unwrap(); - let indices = filtered.indices().to_primitive(); - let values = filtered.values().to_primitive(); + let indices = filtered.indices().to_primitive().unwrap(); + let values = filtered.values().to_primitive().unwrap(); assert_eq!(indices.as_slice::(), &[0, 1]); assert_eq!(values.as_slice::(), &[100, 200]); } @@ -1186,13 +1199,13 @@ mod test { ) .unwrap() .unwrap(); - let primitive_values = taken.values().to_primitive(); - let primitive_indices = taken.indices().to_primitive(); + let primitive_values = taken.values().to_primitive().unwrap(); + let primitive_indices = taken.indices().to_primitive().unwrap(); assert_eq!(taken.array_len(), 2); assert_eq!(primitive_values.as_slice::(), [44]); assert_eq!(primitive_indices.as_slice::(), [0]); assert_eq!( - primitive_values.validity_mask(), + primitive_values.validity_mask().unwrap(), Mask::from_iter(vec![true]) ); } @@ -1215,14 +1228,14 @@ mod test { .unwrap() .unwrap(); - let primitive_values = taken.values().to_primitive(); - let primitive_indices = taken.indices().to_primitive(); + let primitive_values = taken.values().to_primitive().unwrap(); + let primitive_indices = taken.indices().to_primitive().unwrap(); assert_eq!(taken.array_len(), 2); assert_eq!(primitive_values.as_slice::(), [44, 33]); assert_eq!(primitive_indices.as_slice::(), [0, 1]); assert_eq!( - primitive_values.validity_mask(), + primitive_values.validity_mask().unwrap(), Mask::from_iter([true, false]) ); } @@ -1245,7 +1258,7 @@ mod test { .unwrap() .unwrap(); - let primitive_values = taken.values().to_primitive(); + let primitive_values = taken.values().to_primitive().unwrap(); assert_eq!(taken.array_len(), 3); assert_eq!(primitive_values.as_slice::(), [20, 30]); } @@ -1288,7 +1301,7 @@ mod test { .unwrap() .unwrap(); - let primitive_values = taken.values().to_primitive(); + let primitive_values = taken.values().to_primitive().unwrap(); assert_eq!(taken.array_len(), 4); assert_eq!(primitive_values.as_slice::(), [200, 300]); } @@ -1323,7 +1336,7 @@ mod test { let sliced = patches.slice(15..100).unwrap(); assert_eq!(sliced.array_len(), 100 - 15); - let primitive = sliced.values().to_primitive(); + let primitive = sliced.values().to_primitive().unwrap(); assert_eq!(primitive.as_slice::(), &[13531]); } @@ -1337,12 +1350,12 @@ mod test { let sliced = patches.slice(15..100).unwrap(); assert_eq!(sliced.array_len(), 100 - 15); - let primitive = sliced.values().to_primitive(); + let primitive = sliced.values().to_primitive().unwrap(); assert_eq!(primitive.as_slice::(), &[13531]); let doubly_sliced = sliced.slice(35..36).unwrap(); - let primitive_doubly_sliced = doubly_sliced.values().to_primitive(); + let primitive_doubly_sliced = doubly_sliced.values().to_primitive().unwrap(); assert_eq!(primitive_doubly_sliced.as_slice::(), &[13531]); } @@ -1376,14 +1389,14 @@ mod test { let masked = patches.mask(&mask).unwrap().unwrap(); // No patch values should be masked - let masked_values = masked.values().to_primitive(); + let masked_values = masked.values().to_primitive().unwrap(); assert_eq!(masked_values.as_slice::(), &[100, 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 - let indices = masked.indices().to_primitive(); + let indices = masked.indices().to_primitive().unwrap(); assert_eq!(indices.as_slice::(), &[2, 5, 8]); } @@ -1404,12 +1417,12 @@ mod test { let masked = patches.mask(&mask).unwrap().unwrap(); // Only the patch at index 5 should remain - let masked_values = masked.values().to_primitive(); + let masked_values = masked.values().to_primitive().unwrap(); assert_eq!(masked_values.len(), 1); assert_eq!(masked_values.as_slice::(), &[200]); // Only index 5 should remain - let indices = masked.indices().to_primitive(); + let indices = masked.indices().to_primitive().unwrap(); assert_eq!(indices.as_slice::(), &[5]); } @@ -1431,9 +1444,9 @@ mod test { let masked = patches.mask(&mask).unwrap().unwrap(); assert_eq!(masked.array_len(), 10); assert_eq!(masked.offset(), 5); - let indices = masked.indices().to_primitive(); + let indices = masked.indices().to_primitive().unwrap(); assert_eq!(indices.as_slice::(), &[10, 13]); - let masked_values = masked.values().to_primitive(); + let masked_values = masked.values().to_primitive().unwrap(); assert_eq!(masked_values.as_slice::(), &[200, 300]); } @@ -1454,14 +1467,14 @@ mod test { let masked = patches.mask(&mask).unwrap().unwrap(); // Patches at indices 5 and 8 should remain - let indices = masked.indices().to_primitive(); + let indices = masked.indices().to_primitive().unwrap(); assert_eq!(indices.as_slice::(), &[5, 8]); // Values should be the null and 300 - let masked_values = masked.values().to_primitive(); + let masked_values = masked.values().to_primitive().unwrap(); 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!(!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(), 300i32); } @@ -1479,8 +1492,8 @@ mod test { let mask = Mask::from_indices(10, (0..10).collect()); let filtered = patches.filter(&mask).unwrap().unwrap(); - let indices = filtered.indices().to_primitive(); - let values = filtered.values().to_primitive(); + let indices = filtered.indices().to_primitive().unwrap(); + let values = filtered.values().to_primitive().unwrap(); assert_eq!(indices.as_slice::(), &[2, 5, 8]); assert_eq!(values.as_slice::(), &[100, 200, 300]); } @@ -1515,8 +1528,8 @@ mod test { let mask = Mask::from_indices(10, vec![2, 5, 9]); let filtered = patches.filter(&mask).unwrap().unwrap(); - let indices = filtered.indices().to_primitive(); - let values = filtered.values().to_primitive(); + let indices = filtered.indices().to_primitive().unwrap(); + let values = filtered.values().to_primitive().unwrap(); assert_eq!(indices.as_slice::(), &[0, 1]); // Adjusted indices assert_eq!(values.as_slice::(), &[100, 200]); } @@ -1533,8 +1546,8 @@ mod test { let sliced = patches.slice(0..10).unwrap(); - let indices = sliced.indices().to_primitive(); - let values = sliced.values().to_primitive(); + let indices = sliced.indices().to_primitive().unwrap(); + let values = sliced.values().to_primitive().unwrap(); assert_eq!(indices.as_slice::(), &[2, 5, 8]); assert_eq!(values.as_slice::(), &[100, 200, 300]); } @@ -1552,8 +1565,8 @@ mod test { // Slice from 3 to 8 (includes patch at 5) let sliced = patches.slice(3..8).unwrap(); - let indices = sliced.indices().to_primitive(); - let values = sliced.values().to_primitive(); + let indices = sliced.indices().to_primitive().unwrap(); + let values = sliced.values().to_primitive().unwrap(); assert_eq!(indices.as_slice::(), &[5]); // Index stays the same assert_eq!(values.as_slice::(), &[200]); assert_eq!(sliced.array_len(), 5); // 8 - 3 = 5 @@ -1588,8 +1601,8 @@ mod test { // Slice from 3 to 8 (includes patch at actual index 5) let sliced = patches.slice(3..8).unwrap(); - let indices = sliced.indices().to_primitive(); - let values = sliced.values().to_primitive(); + let indices = sliced.indices().to_primitive().unwrap(); + let values = sliced.values().to_primitive().unwrap(); assert_eq!(indices.as_slice::(), &[10]); // Index stays the same (offset + 5 = 10) assert_eq!(values.as_slice::(), &[200]); assert_eq!(sliced.offset(), 8); // New offset = 5 + 3 @@ -1605,7 +1618,7 @@ mod test { None, ); - let values = patches.values().to_primitive(); + let values = patches.values().to_primitive().unwrap(); 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); @@ -1664,9 +1677,9 @@ mod test { let masked = patches.mask(&mask).unwrap(); assert!(masked.is_some()); let masked = masked.unwrap(); - let indices = masked.indices().to_primitive(); + let indices = masked.indices().to_primitive().unwrap(); assert_eq!(indices.as_slice::(), &[9]); - let values = masked.values().to_primitive(); + let values = masked.values().to_primitive().unwrap(); assert_eq!(values.as_slice::(), &[200]); } @@ -1706,9 +1719,9 @@ mod test { ]); let masked = patches.mask(&mask).unwrap().unwrap(); - let indices = masked.indices().to_primitive(); + let indices = masked.indices().to_primitive().unwrap(); assert_eq!(indices.as_slice::(), &[2, 5, 8]); - let values = masked.values().to_primitive(); + let values = masked.values().to_primitive().unwrap(); assert_eq!(values.as_slice::(), &[100, 200, 300]); } @@ -1731,7 +1744,7 @@ mod test { // Mask that keeps the single patch let mask = Mask::from_iter([true, false, false, true, false]); let masked = patches.mask(&mask).unwrap().unwrap(); - let indices = masked.indices().to_primitive(); + let indices = masked.indices().to_primitive().unwrap(); assert_eq!(indices.as_slice::(), &[2]); } @@ -1752,9 +1765,9 @@ mod test { ]); let masked = patches.mask(&mask).unwrap().unwrap(); - let indices = masked.indices().to_primitive(); + let indices = masked.indices().to_primitive().unwrap(); assert_eq!(indices.as_slice::(), &[3, 6]); - let values = masked.values().to_primitive(); + let values = masked.values().to_primitive().unwrap(); assert_eq!(values.as_slice::(), &[100, 400]); } @@ -1776,9 +1789,9 @@ mod test { ]); let masked = patches.mask(&mask).unwrap().unwrap(); - let indices = masked.indices().to_primitive(); + let indices = masked.indices().to_primitive().unwrap(); assert_eq!(indices.as_slice::(), &[16, 19]); - let values = masked.values().to_primitive(); + let values = masked.values().to_primitive().unwrap(); assert_eq!(values.as_slice::(), &[100, 300]); } diff --git a/vortex-array/src/stats/array.rs b/vortex-array/src/stats/array.rs index 172279a1b48..d3195131891 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 => Some(self.dyn_array_ref.invalid_count()?.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 70714dcf241..04cdd715ac9 100644 --- a/vortex-array/src/validity.rs +++ b/vortex-array/src/validity.rs @@ -89,52 +89,52 @@ impl Validity { } #[inline] - pub fn all_valid(&self, len: usize) -> bool { + pub fn all_valid(&self, len: usize) -> VortexResult { match self { - _ if len == 0 => true, - Validity::NonNullable | Validity::AllValid => true, - Validity::AllInvalid => false, + _ if len == 0 => Ok(true), + Validity::NonNullable | Validity::AllValid => Ok(true), + Validity::AllInvalid => Ok(false), Validity::Array(array) => { - usize::try_from(&sum(array).vortex_expect("must have sum for bool array")) - .vortex_expect("sum must be a usize") - == array.len() + let s = sum(array)?; + let count = usize::try_from(&s)?; + Ok(count == array.len()) } } } #[inline] - pub fn all_invalid(&self, len: usize) -> bool { + pub fn all_invalid(&self, len: usize) -> VortexResult { match self { - _ if len == 0 => true, - Validity::NonNullable | Validity::AllValid => false, - Validity::AllInvalid => true, + _ if len == 0 => Ok(true), + Validity::NonNullable | Validity::AllValid => Ok(false), + Validity::AllInvalid => Ok(true), Validity::Array(array) => { - usize::try_from(&sum(array).vortex_expect("must have sum for bool array")) - .vortex_expect("sum must be a usize") - == 0 + let s = sum(array)?; + let count = usize::try_from(&s)?; + Ok(count == 0) } } } /// Returns whether the `index` item is valid. #[inline] - pub fn is_valid(&self, index: usize) -> bool { + pub fn is_valid(&self, index: usize) -> VortexResult { match self { - Self::NonNullable | Self::AllValid => true, - Self::AllInvalid => false, + Self::NonNullable | Self::AllValid => Ok(true), + Self::AllInvalid => Ok(false), Self::Array(a) => { let scalar = a.scalar_at(index); - scalar + Ok(scalar .as_bool() .value() - .vortex_expect("Validity must be non-nullable") + .vortex_expect("Validity must be non-nullable")) } } } #[inline] - pub fn is_null(&self, index: usize) -> bool { - !self.is_valid(index) + pub fn is_null(&self, index: usize) -> VortexResult { + Ok(!self.is_valid(index)?) } #[inline] @@ -147,7 +147,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) @@ -158,7 +158,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())), @@ -191,8 +191,8 @@ impl Validity { /// /// The result is always nullable. The result has the same length as self. #[inline] - pub fn mask(&self, mask: &Mask) -> Self { - match mask.bit_buffer() { + pub fn mask(&self, mask: &Mask) -> VortexResult { + Ok(match mask.bit_buffer() { AllOr::All => Validity::AllInvalid, AllOr::None => self.clone().into_nullable(), AllOr::Some(make_invalid) => match self { @@ -201,18 +201,18 @@ impl Validity { } Validity::AllInvalid => Validity::AllInvalid, Validity::Array(is_valid) => { - let is_valid = is_valid.to_bool(); + let is_valid = is_valid.to_bool()?; Validity::from(is_valid.bit_buffer() & !make_invalid) } }, - } + }) } #[inline] - pub fn to_mask(&self, length: usize) -> Mask { + pub fn to_mask(&self, length: usize) -> VortexResult { match self { - Self::NonNullable | Self::AllValid => Mask::AllTrue(length), - Self::AllInvalid => Mask::AllFalse(length), + Self::NonNullable | Self::AllValid => Ok(Mask::AllTrue(length)), + Self::AllInvalid => Ok(Mask::AllFalse(length)), Self::Array(is_valid) => { assert_eq!( is_valid.len(), @@ -221,15 +221,15 @@ impl Validity { is_valid.len(), length, ); - is_valid.to_bool().to_mask() + Ok(is_valid.to_bool()?.to_mask()) } } } /// Logically & two Validity values of the same length #[inline] - pub fn and(self, rhs: Validity) -> Validity { - match (self, rhs) { + pub fn and(self, rhs: Validity) -> VortexResult { + Ok(match (self, rhs) { // Should be pretty clear (Validity::NonNullable, Validity::NonNullable) => Validity::NonNullable, // Any `AllInvalid` makes the output all invalid values @@ -245,15 +245,15 @@ impl Validity { | (Validity::AllValid, Validity::AllValid) => Validity::AllValid, // Here we actually have to do some work (Validity::Array(lhs), Validity::Array(rhs)) => { - let lhs = lhs.to_bool(); - let rhs = rhs.to_bool(); + let lhs = lhs.to_bool()?; + let rhs = rhs.to_bool()?; let lhs = lhs.bit_buffer(); let rhs = rhs.bit_buffer(); Validity::from(lhs.bitand(rhs)) } - } + }) } pub fn patch( @@ -262,17 +262,17 @@ impl Validity { indices_offset: usize, indices: &dyn Array, patches: &Validity, - ) -> Self { + ) -> VortexResult { match (&self, patches) { - (Validity::NonNullable, Validity::NonNullable) => return Validity::NonNullable, + (Validity::NonNullable, Validity::NonNullable) => return Ok(Validity::NonNullable), (Validity::NonNullable, _) => { vortex_panic!("Can't patch a non-nullable validity with nullable validity") } (_, Validity::NonNullable) => { vortex_panic!("Can't patch a nullable validity with non-nullable validity") } - (Validity::AllValid, Validity::AllValid) => return Validity::AllValid, - (Validity::AllInvalid, Validity::AllInvalid) => return Validity::AllInvalid, + (Validity::AllValid, Validity::AllValid) => return Ok(Validity::AllValid), + (Validity::AllInvalid, Validity::AllInvalid) => return Ok(Validity::AllInvalid), _ => {} }; @@ -286,14 +286,14 @@ impl Validity { Validity::NonNullable => BoolArray::from(BitBuffer::new_set(len)), Validity::AllValid => BoolArray::from(BitBuffer::new_set(len)), Validity::AllInvalid => BoolArray::from(BitBuffer::new_unset(len)), - Validity::Array(a) => a.to_bool(), + Validity::Array(a) => a.to_bool()?, }; let patch_values = match patches { Validity::NonNullable => BoolArray::from(BitBuffer::new_set(indices.len())), Validity::AllValid => BoolArray::from(BitBuffer::new_set(indices.len())), Validity::AllInvalid => BoolArray::from(BitBuffer::new_unset(indices.len())), - Validity::Array(a) => a.to_bool(), + Validity::Array(a) => a.to_bool()?, }; let patches = Patches::new( @@ -305,7 +305,10 @@ impl Validity { None, ); - Self::from_array(source.patch(&patches).into_array(), own_nullability) + Ok(Self::from_array( + source.patch(&patches).into_array(), + own_nullability, + )) } /// Convert into a nullable variant @@ -352,7 +355,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()) + use vortex_error::VortexExpect as _; + Validity::from_mask( + array.validity_mask().vortex_expect("validity_mask failed"), + array.dtype().nullability(), + ) } /// Create Validity from boolean array with given nullability of the array. @@ -397,8 +404,12 @@ impl PartialEq for Validity { (Self::AllValid, Self::AllValid) => true, (Self::AllInvalid, Self::AllInvalid) => true, (Self::Array(a), Self::Array(b)) => { - let a = a.to_bool(); - let b = b.to_bool(); + let a = a + .to_bool() + .vortex_expect("Failed to convert validity to bool"); + let b = b + .to_bool() + .vortex_expect("Failed to convert validity to bool"); a.bit_buffer() == b.bit_buffer() } _ => false, @@ -529,13 +540,18 @@ mod tests { ) { let indices = PrimitiveArray::new(Buffer::copy_from(positions), Validity::NonNullable).into_array(); - assert_eq!(validity.patch(len, 0, &indices, &patches), expected); + assert_eq!( + validity.patch(len, 0, &indices, &patches).unwrap(), + expected + ); } #[test] #[should_panic] fn out_of_bounds_patch() { - Validity::NonNullable.patch(2, 0, &buffer![4].into_array(), &Validity::AllInvalid); + Validity::NonNullable + .patch(2, 0, &buffer![4].into_array(), &Validity::AllInvalid) + .unwrap(); } #[test] @@ -573,7 +589,7 @@ mod tests { fn mask_non_nullable() { assert_eq!( Validity::AllValid, - Validity::NonNullable.mask(&Mask::AllFalse(2)) + Validity::NonNullable.mask(&Mask::AllFalse(2)).unwrap() ) } } diff --git a/vortex-array/src/variants.rs b/vortex-array/src/variants.rs index 3f733d918e7..962ceeda5c5 100644 --- a/vortex-array/src/variants.rs +++ b/vortex-array/src/variants.rs @@ -110,7 +110,11 @@ 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)) + use vortex_error::VortexExpect as _; + self.0 + .is_valid(idx) + .vortex_expect("is_valid") + .then(|| self.value_unchecked(idx)) } /// Return the primitive value at the given index, ignoring nullability. @@ -136,7 +140,8 @@ 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()); + use vortex_error::VortexExpect as _; + assert!(self.0.all_valid().vortex_expect("all_valid")); self.value_unchecked(idx).partial_cmp(elem) } diff --git a/vortex-array/src/vtable/canonical.rs b/vortex-array/src/vtable/canonical.rs index 94910c35144..c9f42956cc6 100644 --- a/vortex-array/src/vtable/canonical.rs +++ b/vortex-array/src/vtable/canonical.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 crate::Canonical; use crate::builders::ArrayBuilder; @@ -15,21 +16,22 @@ pub trait CanonicalVTable { /// - The length is equal to that of the input array. /// - The [`vortex_dtype::DType`] is equal to that of the input array. // TODO(ngates): rename to `decode` - fn canonicalize(array: &V::Array) -> Canonical; + fn canonicalize(array: &V::Array) -> VortexResult; /// Writes the array into a canonical builder. /// /// ## Post-conditions /// - The length of the builder is incremented by the length of the input array. - fn append_to_builder(array: &V::Array, builder: &mut dyn ArrayBuilder) { - let canonical = Self::canonicalize(array); - builder.extend_from_array(canonical.as_ref()) + fn append_to_builder(array: &V::Array, builder: &mut dyn ArrayBuilder) -> VortexResult<()> { + let canonical = Self::canonicalize(array)?; + builder.extend_from_array(canonical.as_ref()); + Ok(()) } } impl CanonicalVTable for NotSupported { - fn canonicalize(array: &V::Array) -> Canonical { - vortex_panic!( + fn canonicalize(array: &V::Array) -> VortexResult { + vortex_bail!( "Legacy canonicalize 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 14211a7946d..6c797da2209 100644 --- a/vortex-array/src/vtable/validity.rs +++ b/vortex-array/src/vtable/validity.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_mask::Mask; use crate::Array; @@ -11,55 +12,55 @@ use crate::vtable::NotSupported; use crate::vtable::VTable; pub trait ValidityVTable { - fn is_valid(array: &V::Array, index: usize) -> bool; + fn is_valid(array: &V::Array, index: usize) -> VortexResult; - fn all_valid(array: &V::Array) -> bool; + fn all_valid(array: &V::Array) -> VortexResult; - fn all_invalid(array: &V::Array) -> bool; + fn all_invalid(array: &V::Array) -> VortexResult; /// Returns the number of valid elements in the array. /// /// ## Post-conditions /// - The count is less than or equal to the length of the array. - fn valid_count(array: &V::Array) -> usize { - Self::validity_mask(array).true_count() + fn valid_count(array: &V::Array) -> VortexResult { + Ok(Self::validity_mask(array)?.true_count()) } /// Returns the number of invalid elements in the array. /// /// ## Post-conditions /// - The count is less than or equal to the length of the array. - fn invalid_count(array: &V::Array) -> usize { - Self::validity_mask(array).false_count() + fn invalid_count(array: &V::Array) -> VortexResult { + Ok(Self::validity_mask(array)?.false_count()) } - fn validity_mask(array: &V::Array) -> Mask; + fn validity_mask(array: &V::Array) -> VortexResult; } impl ValidityVTable for NotSupported { - fn is_valid(array: &V::Array, _index: usize) -> bool { - vortex_panic!( + fn is_valid(array: &V::Array, _index: usize) -> VortexResult { + vortex_bail!( "Legacy is_valid is not supported for {} arrays", array.encoding_id() ) } - fn all_valid(array: &V::Array) -> bool { - vortex_panic!( + fn all_valid(array: &V::Array) -> VortexResult { + vortex_bail!( "Legacy all_valid is not supported for {} arrays", array.encoding_id() ) } - fn all_invalid(array: &V::Array) -> bool { - vortex_panic!( + fn all_invalid(array: &V::Array) -> VortexResult { + vortex_bail!( "Legacy all_invalid is not supported for {} arrays", array.encoding_id() ) } - fn validity_mask(array: &V::Array) -> Mask { - vortex_panic!( + fn validity_mask(array: &V::Array) -> VortexResult { + vortex_bail!( "Legacy validity_mask is not supported for {} arrays", array.encoding_id() ) @@ -78,19 +79,19 @@ impl ValidityVTable for ValidityVTableFromValidityHelper where V::Array: ValidityHelper, { - fn is_valid(array: &V::Array, index: usize) -> bool { + fn is_valid(array: &V::Array, index: usize) -> VortexResult { array.validity().is_valid(index) } - fn all_valid(array: &V::Array) -> bool { + fn all_valid(array: &V::Array) -> VortexResult { array.validity().all_valid(array.len()) } - fn all_invalid(array: &V::Array) -> bool { + fn all_invalid(array: &V::Array) -> VortexResult { array.validity().all_invalid(array.len()) } - fn validity_mask(array: &V::Array) -> Mask { + fn validity_mask(array: &V::Array) -> VortexResult { array.validity().to_mask(array.len()) } } @@ -112,20 +113,20 @@ impl ValidityVTable for ValidityVTableFromValiditySliceHelper where V::Array: ValiditySliceHelper, { - fn is_valid(array: &V::Array, index: usize) -> bool { + fn is_valid(array: &V::Array, index: usize) -> VortexResult { let (unsliced_validity, start, _) = array.unsliced_validity_and_slice(); unsliced_validity.is_valid(start + index) } - fn all_valid(array: &V::Array) -> bool { + fn all_valid(array: &V::Array) -> VortexResult { array.sliced_validity().all_valid(array.len()) } - fn all_invalid(array: &V::Array) -> bool { + fn all_invalid(array: &V::Array) -> VortexResult { array.sliced_validity().all_invalid(array.len()) } - fn validity_mask(array: &V::Array) -> Mask { + fn validity_mask(array: &V::Array) -> VortexResult { array.sliced_validity().to_mask(array.len()) } } @@ -142,19 +143,19 @@ impl ValidityVTable for ValidityVTableFromChild where V: ValidityChild, { - fn is_valid(array: &V::Array, index: usize) -> bool { + fn is_valid(array: &V::Array, index: usize) -> VortexResult { V::validity_child(array).is_valid(index) } - fn all_valid(array: &V::Array) -> bool { + fn all_valid(array: &V::Array) -> VortexResult { V::validity_child(array).all_valid() } - fn all_invalid(array: &V::Array) -> bool { + fn all_invalid(array: &V::Array) -> VortexResult { V::validity_child(array).all_invalid() } - fn validity_mask(array: &V::Array) -> Mask { + fn validity_mask(array: &V::Array) -> VortexResult { V::validity_child(array).validity_mask() } } @@ -176,20 +177,20 @@ impl ValidityVTable for ValidityVTableFromChildSliceHelper where V::Array: ValidityChildSliceHelper, { - fn is_valid(array: &V::Array, index: usize) -> bool { + fn is_valid(array: &V::Array, index: usize) -> VortexResult { let (unsliced_validity, start, _) = array.unsliced_child_and_slice(); unsliced_validity.is_valid(start + index) } - fn all_valid(array: &V::Array) -> bool { + fn all_valid(array: &V::Array) -> VortexResult { array.sliced_child_array().all_valid() } - fn all_invalid(array: &V::Array) -> bool { + fn all_invalid(array: &V::Array) -> VortexResult { array.sliced_child_array().all_invalid() } - fn validity_mask(array: &V::Array) -> Mask { + fn validity_mask(array: &V::Array) -> VortexResult { array.sliced_child_array().validity_mask() } } diff --git a/vortex-btrblocks/benches/compress.rs b/vortex-btrblocks/benches/compress.rs index 0c51c0efdec..0727baaef84 100644 --- a/vortex-btrblocks/benches/compress.rs +++ b/vortex-btrblocks/benches/compress.rs @@ -41,7 +41,7 @@ mod benchmarks { #[divan::bench] fn btrblocks(bencher: Bencher) { - let array = make_clickbench_window_name().to_primitive(); + let array = make_clickbench_window_name().to_primitive().unwrap(); bencher .with_inputs(|| &array) .input_counter(|array| ItemsCount::new(array.len())) diff --git a/vortex-btrblocks/src/float.rs b/vortex-btrblocks/src/float.rs index d19996ef002..ae217140fec 100644 --- a/vortex-btrblocks/src/float.rs +++ b/vortex-btrblocks/src/float.rs @@ -19,6 +19,7 @@ use vortex_array::vtable::ValidityHelper; use vortex_dtype::PType; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::VortexUnwrap; use vortex_error::vortex_panic; use vortex_scalar::Scalar; use vortex_sparse::SparseArray; @@ -173,13 +174,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).vortex_unwrap()); match scalar_idx { Some(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().vortex_unwrap() { Ok(MaskedArray::try_new(const_arr, stats.src.validity().clone())?.into_array()) } else { Ok(const_arr) @@ -241,10 +243,10 @@ impl Scheme for ALPScheme { ) -> VortexResult { let alp_encoded = ALPVTable .as_vtable() - .encode(&stats.source().to_canonical(), None)? + .encode(&stats.source().to_canonical().vortex_unwrap(), None)? .vortex_expect("Input is a supported floating point array"); let alp = alp_encoded.as_::(); - let alp_ints = alp.encoded().to_primitive(); + let alp_ints = alp.encoded().to_primitive().vortex_unwrap(); // Compress the ALP ints. // Patches are not compressed. They should be infrequent, and if they are not then we want @@ -364,7 +366,7 @@ impl Scheme for DictScheme { // Only compress the codes. let codes_stats = IntegerStats::generate_opts( - &dict_array.codes().to_primitive().narrow()?, + &dict_array.codes().to_primitive().vortex_unwrap().narrow()?, GenerateStatsOptions { count_distinct_values: false, }, @@ -383,7 +385,7 @@ impl Scheme for DictScheme { )?; let compressed_values = FloatCompressor::compress( - &dict_array.values().to_primitive(), + &dict_array.values().to_primitive().vortex_unwrap(), is_sample, allowed_cascading - 1, &[DICT_SCHEME], @@ -452,7 +454,12 @@ impl Scheme for NullDominated { // Don't attempt to compress the non-null values - let indices = sparse.patches().indices().to_primitive().narrow()?; + let indices = sparse + .patches() + .indices() + .to_primitive() + .vortex_unwrap() + .narrow()?; let compressed_indices = IntCompressor::compress_no_dict( &indices, is_sample, @@ -522,7 +529,7 @@ mod tests { values[i] = (i % 50) as f32; } - let floats = values.into_array().to_primitive(); + let floats = values.into_array().to_primitive().unwrap(); let compressed = FloatCompressor::compress(&floats, false, MAX_CASCADE, &[]).unwrap(); println!("compressed: {}", compressed.display_tree()) } diff --git a/vortex-btrblocks/src/float/stats.rs b/vortex-btrblocks/src/float/stats.rs index 6c6e0b2d8be..6379219d359 100644 --- a/vortex-btrblocks/src/float/stats.rs +++ b/vortex-btrblocks/src/float/stats.rs @@ -81,7 +81,9 @@ impl CompressorStats for FloatStats { } fn sample_opts(&self, sample_size: u32, sample_count: u32, opts: GenerateStatsOptions) -> Self { - let sampled = sample(self.src.as_ref(), sample_size, sample_count).to_primitive(); + let sampled = sample(self.src.as_ref(), sample_size, sample_count) + .to_primitive() + .vortex_unwrap(); Self::generate_opts(&sampled, opts) } @@ -122,7 +124,7 @@ where } .into(), }; - } else if array.all_invalid() { + } else if array.all_invalid().vortex_unwrap() { return FloatStats { src: array.clone(), null_count: array.len().try_into().vortex_expect("null_count"), @@ -150,7 +152,7 @@ where HashSet::with_hasher(FxBuildHasher) }; - let validity = array.validity_mask(); + let validity = array.validity_mask().vortex_unwrap(); let mut runs = 1; let head_idx = validity @@ -232,7 +234,7 @@ mod tests { #[test] fn test_float_stats() { let floats = buffer![0.0f32, 1.0f32, 2.0f32].into_array(); - let floats = floats.to_primitive(); + let floats = floats.to_primitive().unwrap(); let stats = FloatStats::generate(&floats); diff --git a/vortex-btrblocks/src/integer.rs b/vortex-btrblocks/src/integer.rs index 2a953485757..bd7afcafd31 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).vortex_unwrap()); match scalar_idx { Some(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().vortex_unwrap() { Ok(MaskedArray::try_new(const_arr, stats.src.validity().clone())?.into_array()) } else { Ok(const_arr) @@ -301,7 +302,7 @@ impl Scheme for FORScheme { excludes: &[IntCode], ) -> VortexResult { let for_array = FoRArray::encode(stats.src.clone())?; - let biased = for_array.encoded().to_primitive(); + let biased = for_array.encoded().to_primitive().vortex_unwrap(); let biased_stats = IntegerStats::generate_opts( &biased, GenerateStatsOptions { @@ -368,7 +369,7 @@ impl Scheme for ZigZagScheme { ) -> VortexResult { // Zigzag encode the values, then recursively compress the inner values. let zag = zigzag_encode(stats.src.clone())?; - let encoded = zag.encoded().to_primitive(); + let encoded = zag.encoded().to_primitive().vortex_unwrap(); // ZigZag should be after Dict, RunEnd or Sparse. // We should only do these "container" style compressors once. @@ -530,13 +531,18 @@ impl Scheme for SparseScheme { new_excludes.extend_from_slice(excludes); let compressed_values = IntCompressor::compress_no_dict( - &sparse.patches().values().to_primitive(), + &sparse.patches().values().to_primitive().vortex_unwrap(), is_sample, allowed_cascading - 1, &new_excludes, )?; - let indices = sparse.patches().indices().to_primitive().narrow()?; + let indices = sparse + .patches() + .indices() + .to_primitive() + .vortex_unwrap() + .narrow()?; let compressed_indices = IntCompressor::compress_no_dict( &indices, @@ -626,7 +632,7 @@ impl Scheme for DictScheme { new_excludes.extend_from_slice(excludes); let compressed_codes = IntCompressor::compress_no_dict( - &dict.codes().to_primitive().narrow()?, + &dict.codes().to_primitive().vortex_unwrap().narrow()?, is_sample, allowed_cascading - 1, &new_excludes, @@ -687,13 +693,13 @@ impl Scheme for RunEndScheme { assert!(allowed_cascading > 0); // run-end encode the ends - let (ends, values) = runend_encode(&stats.src); + let (ends, values) = runend_encode(&stats.src)?; let mut new_excludes = vec![RunEndScheme.code(), DictScheme.code()]; new_excludes.extend_from_slice(excludes); let ends_stats = IntegerStats::generate_opts( - &ends.to_primitive(), + &ends.to_primitive().vortex_unwrap(), GenerateStatsOptions { count_distinct_values: false, }, @@ -708,7 +714,7 @@ impl Scheme for RunEndScheme { ends_scheme.compress(&ends_stats, is_sample, allowed_cascading - 1, &new_excludes)?; let compressed_values = IntCompressor::compress_no_dict( - &values.to_primitive(), + &values.to_primitive().vortex_unwrap(), is_sample, allowed_cascading - 1, &new_excludes, @@ -831,7 +837,7 @@ mod tests { } } - let primitive = codes.freeze().into_array().to_primitive(); + let primitive = codes.freeze().into_array().to_primitive().unwrap(); let compressed = IntCompressor::compress(&primitive, false, 3, &[]).unwrap(); assert!(compressed.is::()); } @@ -857,7 +863,7 @@ mod tests { values[random] = 5 * (rng.next_u64() % 100) as i32; } - let array = values.freeze().into_array().to_primitive(); + let array = values.freeze().into_array().to_primitive().unwrap(); let compressed = IntCompressor::compress(&array, false, 3, &[]).unwrap(); log::info!("WindowName compressed: {}", compressed.display_tree()); } diff --git a/vortex-btrblocks/src/integer/stats.rs b/vortex-btrblocks/src/integer/stats.rs index 40b7fc38e8c..a3a2f120550 100644 --- a/vortex-btrblocks/src/integer/stats.rs +++ b/vortex-btrblocks/src/integer/stats.rs @@ -154,7 +154,9 @@ impl CompressorStats for IntegerStats { } fn sample_opts(&self, sample_size: u32, sample_count: u32, opts: GenerateStatsOptions) -> Self { - let sampled = sample(self.src.as_ref(), sample_size, sample_count).to_primitive(); + let sampled = sample(self.src.as_ref(), sample_size, sample_count) + .to_primitive() + .vortex_unwrap(); Self::generate_opts(&sampled, opts) } @@ -197,7 +199,7 @@ where } .into(), }; - } else if array.all_invalid() { + } else if array.all_invalid().vortex_unwrap() { return IntegerStats { src: array.clone(), null_count: array.len().try_into().vortex_expect("null_count"), @@ -215,7 +217,7 @@ where }; } - let validity = array.validity_mask(); + let validity = array.validity_mask().vortex_unwrap(); 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 e089b7adc9d..7b5363ddf76 100644 --- a/vortex-btrblocks/src/lib.rs +++ b/vortex-btrblocks/src/lib.rs @@ -368,7 +368,7 @@ impl BtrBlocksCompressor { /// First canonicalizes and compacts the array, then applies optimal compression schemes. pub fn compress(&self, array: &dyn Array) -> VortexResult { // Canonicalize the array - let canonical = array.to_canonical(); + let canonical = array.to_canonical()?; // Compact it, removing any wasted space before we attempt to compress it let compact = canonical.compact()?; @@ -428,7 +428,11 @@ impl BtrBlocksCompressor { // widths. let compressed_offsets = IntCompressor::compress_no_dict( - &list_array.offsets().to_primitive().narrow()?, + &list_array + .offsets() + .to_primitive() + .vortex_unwrap() + .narrow()?, false, MAX_CASCADE, &[], diff --git a/vortex-btrblocks/src/patches.rs b/vortex-btrblocks/src/patches.rs index 57133a5ed88..b0bf0f65832 100644 --- a/vortex-btrblocks/src/patches.rs +++ b/vortex-btrblocks/src/patches.rs @@ -7,11 +7,17 @@ use vortex_array::ToCanonical; use vortex_array::arrays::ConstantArray; use vortex_array::patches::Patches; use vortex_error::VortexResult; +use vortex_error::VortexUnwrap; /// Compresses the given patches by downscaling integers and checking for constant values. pub fn compress_patches(patches: &Patches) -> VortexResult { // Downscale the patch indices. - let indices = patches.indices().to_primitive().narrow()?.into_array(); + let indices = patches + .indices() + .to_primitive() + .vortex_unwrap() + .narrow()? + .into_array(); // Check if the values are constant. let values = patches.values(); diff --git a/vortex-btrblocks/src/rle.rs b/vortex-btrblocks/src/rle.rs index 507e8cf8c0e..b88eaa3b0de 100644 --- a/vortex-btrblocks/src/rle.rs +++ b/vortex-btrblocks/src/rle.rs @@ -9,6 +9,7 @@ use vortex_array::IntoArray; use vortex_array::ToCanonical; use vortex_array::arrays::PrimitiveArray; use vortex_error::VortexResult; +use vortex_error::VortexUnwrap; use vortex_fastlanes::RLEArray; use crate::CompressorStats; @@ -110,7 +111,7 @@ where new_excludes.extend_from_slice(excludes); let compressed_values = (self.compress_values_fn)( - &rle_array.values().to_primitive(), + &rle_array.values().to_primitive().vortex_unwrap(), is_sample, allowed_cascading - 1, &new_excludes, @@ -130,14 +131,22 @@ where // #[cfg(not(feature = "unstable_encodings"))] let compressed_indices = IntCompressor::compress_no_dict( - &rle_array.indices().to_primitive().narrow()?, + &rle_array + .indices() + .to_primitive() + .vortex_unwrap() + .narrow()?, is_sample, allowed_cascading - 1, &[], )?; let compressed_offsets = IntCompressor::compress_no_dict( - &rle_array.values_idx_offsets().to_primitive().narrow()?, + &rle_array + .values_idx_offsets() + .to_primitive() + .vortex_unwrap() + .narrow()?, is_sample, allowed_cascading - 1, &[], diff --git a/vortex-btrblocks/src/string.rs b/vortex-btrblocks/src/string.rs index 3d862fb9453..80ae3d297b1 100644 --- a/vortex-btrblocks/src/string.rs +++ b/vortex-btrblocks/src/string.rs @@ -14,6 +14,7 @@ use vortex_array::builders::dict::dict_encode; use vortex_array::vtable::ValidityHelper; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::VortexUnwrap; use vortex_fsst::FSSTArray; use vortex_fsst::fsst_compress; use vortex_fsst::fsst_train_compressor; @@ -90,7 +91,9 @@ impl CompressorStats for StringStats { } fn sample_opts(&self, sample_size: u32, sample_count: u32, opts: GenerateStatsOptions) -> Self { - let sampled = sample(self.src.as_ref(), sample_size, sample_count).to_varbinview(); + let sampled = sample(self.src.as_ref(), sample_size, sample_count) + .to_varbinview() + .vortex_unwrap(); Self::generate_opts(&sampled, opts) } @@ -231,7 +234,7 @@ impl Scheme for DictScheme { // Find best compressor for codes and values separately let compressed_codes = IntCompressor::compress( - &dict.codes().to_primitive(), + &dict.codes().to_primitive().vortex_unwrap(), is_sample, allowed_cascading - 1, &[integer::DictScheme.code(), integer::SequenceScheme.code()], @@ -240,7 +243,7 @@ impl Scheme for DictScheme { // Attempt to compress the values with non-Dict compression. // Currently this will only be FSST. let compressed_values = StringCompressor::compress( - &dict.values().to_varbinview(), + &dict.values().to_varbinview().vortex_unwrap(), is_sample, allowed_cascading - 1, &[DictScheme.code()], @@ -276,14 +279,23 @@ impl Scheme for FSSTScheme { let fsst = fsst_compress(&stats.src, &compressor); let compressed_original_lengths = IntCompressor::compress( - &fsst.uncompressed_lengths().to_primitive().narrow()?, + &fsst + .uncompressed_lengths() + .to_primitive() + .vortex_unwrap() + .narrow()?, is_sample, allowed_cascading, &[], )?; let compressed_codes_offsets = IntCompressor::compress( - &fsst.codes().offsets().to_primitive().narrow()?, + &fsst + .codes() + .offsets() + .to_primitive() + .vortex_unwrap() + .narrow()?, is_sample, allowed_cascading, &[], @@ -345,13 +357,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).vortex_unwrap()); match scalar_idx { Some(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().vortex_unwrap() { Ok(MaskedArray::try_new(const_arr, stats.src.validity().clone())?.into_array()) } else { Ok(const_arr) @@ -417,7 +430,12 @@ impl Scheme for NullDominated { let new_excludes = vec![integer::SparseScheme.code()]; // Don't attempt to compress the non-null values - let indices = sparse.patches().indices().to_primitive().narrow()?; + let indices = sparse + .patches() + .indices() + .to_primitive() + .vortex_unwrap() + .narrow()?; let compressed_indices = IntCompressor::compress_no_dict( &indices, is_sample, diff --git a/vortex-btrblocks/src/temporal.rs b/vortex-btrblocks/src/temporal.rs index 67468c1b311..a64e726536c 100644 --- a/vortex-btrblocks/src/temporal.rs +++ b/vortex-btrblocks/src/temporal.rs @@ -11,6 +11,7 @@ use vortex_datetime_parts::DateTimePartsArray; use vortex_datetime_parts::TemporalParts; use vortex_datetime_parts::split_temporal; use vortex_error::VortexResult; +use vortex_error::VortexUnwrap; use crate::Compressor; use crate::MAX_CASCADE; @@ -25,16 +26,20 @@ pub fn compress_temporal(array: TemporalArray) -> VortexResult { subseconds, } = split_temporal(array)?; - let days = - IntCompressor::compress(&days.to_primitive().narrow()?, false, MAX_CASCADE - 1, &[])?; + let days = IntCompressor::compress( + &days.to_primitive().vortex_unwrap().narrow()?, + false, + MAX_CASCADE - 1, + &[], + )?; let seconds = IntCompressor::compress( - &seconds.to_primitive().narrow()?, + &seconds.to_primitive().vortex_unwrap().narrow()?, false, MAX_CASCADE - 1, &[], )?; let subseconds = IntCompressor::compress( - &subseconds.to_primitive().narrow()?, + &subseconds.to_primitive().vortex_unwrap().narrow()?, false, MAX_CASCADE - 1, &[], diff --git a/vortex-duckdb/src/convert/vector.rs b/vortex-duckdb/src/convert/vector.rs index 431a285b593..d29ae9185ca 100644 --- a/vortex-duckdb/src/convert/vector.rs +++ b/vortex-duckdb/src/convert/vector.rs @@ -338,7 +338,7 @@ mod tests { // Test conversion let result = flat_vector_to_vortex(&mut vector, len).unwrap(); let vortex_array = TemporalArray::try_from(result).unwrap(); - let vortex_values = vortex_array.temporal_values().to_primitive(); + let vortex_values = vortex_array.temporal_values().to_primitive().unwrap(); let values_slice = vortex_values.as_slice::(); assert_eq!(values_slice, values); @@ -361,7 +361,7 @@ mod tests { // Test conversion let result = flat_vector_to_vortex(&mut vector, len).unwrap(); let vortex_array = TemporalArray::try_from(result).unwrap(); - let vortex_values = vortex_array.temporal_values().to_primitive(); + let vortex_values = vortex_array.temporal_values().to_primitive().unwrap(); let values_slice = vortex_values.as_slice::(); assert_eq!(values_slice, values); @@ -384,7 +384,7 @@ mod tests { // Test conversion let result = flat_vector_to_vortex(&mut vector, len).unwrap(); let vortex_array = TemporalArray::try_from(result).unwrap(); - let vortex_values = vortex_array.temporal_values().to_primitive(); + let vortex_values = vortex_array.temporal_values().to_primitive().unwrap(); let values_slice = vortex_values.as_slice::(); assert_eq!(values_slice, values); @@ -412,12 +412,12 @@ mod tests { // Test conversion let result = flat_vector_to_vortex(&mut vector, len).unwrap(); let vortex_array = TemporalArray::try_from(result).unwrap(); - let vortex_values = vortex_array.temporal_values().to_primitive(); + let vortex_values = vortex_array.temporal_values().to_primitive().unwrap(); let values_slice = vortex_values.as_slice::(); assert_eq!(values_slice, values); assert_eq!( - vortex_values.validity_mask(), + vortex_values.validity_mask().unwrap(), Mask::from_indices(3, vec![0, 2]) ); } @@ -446,7 +446,7 @@ mod tests { // Test conversion let result = flat_vector_to_vortex(&mut vector, len).unwrap(); let vortex_array = TemporalArray::try_from(result).unwrap(); - let vortex_values = vortex_array.temporal_values().to_primitive(); + let vortex_values = vortex_array.temporal_values().to_primitive().unwrap(); let values_slice = vortex_values.as_slice::(); assert_eq!(values_slice, values); @@ -469,7 +469,7 @@ mod tests { // Test conversion let result = flat_vector_to_vortex(&mut vector, len).unwrap(); let vortex_array = TemporalArray::try_from(result).unwrap(); - let vortex_values = vortex_array.temporal_values().to_primitive(); + let vortex_values = vortex_array.temporal_values().to_primitive().unwrap(); let values_slice = vortex_values.as_slice::(); assert_eq!(values_slice, values); @@ -491,7 +491,7 @@ mod tests { // Test conversion let result = flat_vector_to_vortex(&mut vector, len).unwrap(); - let vortex_array = result.to_bool(); + let vortex_array = result.to_bool().unwrap(); assert_eq!(vortex_array.len(), len); assert_eq!(vortex_array.bit_buffer().iter().collect::>(), values); @@ -518,12 +518,12 @@ mod tests { // Test conversion let result = flat_vector_to_vortex(&mut vector, len).unwrap(); - let vortex_array = result.to_primitive(); + let vortex_array = result.to_primitive().unwrap(); let vortex_slice = vortex_array.as_slice::(); assert_eq!(vortex_slice, values); assert_eq!( - vortex_array.validity_mask(), + vortex_array.validity_mask().unwrap(), Mask::from_indices(3, vec![0, 2]) ); } @@ -552,13 +552,14 @@ mod tests { // Test conversion let result = flat_vector_to_vortex(&mut vector, len).unwrap(); - let vortex_array = result.to_listview(); + let vortex_array = result.to_listview().unwrap(); assert_eq!(vortex_array.len(), len); assert_eq!( vortex_array .list_elements_at(0) .to_primitive() + .unwrap() .as_slice::(), &[1, 2, 3, 4] ); @@ -583,13 +584,14 @@ mod tests { // Test conversion let result = flat_vector_to_vortex(&mut vector, len).unwrap(); - let vortex_array = result.to_fixed_size_list(); + let vortex_array = result.to_fixed_size_list().unwrap(); assert_eq!(vortex_array.len(), len); assert_eq!( vortex_array .fixed_size_list_elements_at(0) .to_primitive() + .unwrap() .as_slice::(), &[1, 2, 3, 4] ); @@ -603,7 +605,7 @@ mod tests { // Test conversion let result = flat_vector_to_vortex(&mut vector, len).unwrap(); - let vortex_array = result.to_struct(); + let vortex_array = result.to_struct().unwrap(); assert_eq!(vortex_array.len(), len); assert_eq!(vortex_array.fields().len(), 0); @@ -638,16 +640,22 @@ mod tests { // Test conversion let result = flat_vector_to_vortex(&mut vector, len).unwrap(); - let vortex_array = result.to_struct(); + let vortex_array = result.to_struct().unwrap(); assert_eq!(vortex_array.len(), len); assert_eq!(vortex_array.fields().len(), 2); assert_eq!( - vortex_array.fields()[0].to_primitive().as_slice::(), + vortex_array.fields()[0] + .to_primitive() + .unwrap() + .as_slice::(), &[1, 2, 3, 4] ); assert_eq!( - vortex_array.fields()[1].to_primitive().as_slice::(), + vortex_array.fields()[1] + .to_primitive() + .unwrap() + .as_slice::(), &[5, 6, 7, 8] ); } diff --git a/vortex-duckdb/src/duckdb/vector.rs b/vortex-duckdb/src/duckdb/vector.rs index 8343b4eef3e..f75c9858bd2 100644 --- a/vortex-duckdb/src/duckdb/vector.rs +++ b/vortex-duckdb/src/duckdb/vector.rs @@ -371,7 +371,7 @@ mod tests { // Check that the right positions are null assert_eq!( - validity.to_mask(len), + validity.to_mask(len).unwrap(), Mask::from_indices(len, vec![0, 2, 4, 5, 6, 8, 9]) ); } @@ -389,7 +389,7 @@ mod tests { let validity = vector.validity_ref(len); let validity = validity.to_validity(); - assert!(validity.is_null(0)); + assert!(validity.is_null(0).unwrap()); } #[test] @@ -405,7 +405,7 @@ mod tests { let validity = vector.validity_ref(len); let validity = validity.to_validity(); - assert!(validity.is_valid(0)); + assert!(validity.is_valid(0).unwrap()); } #[test] diff --git a/vortex-duckdb/src/exporter/bool.rs b/vortex-duckdb/src/exporter/bool.rs index 6934eb1f7c9..4fbf020c5e9 100644 --- a/vortex-duckdb/src/exporter/bool.rs +++ b/vortex-duckdb/src/exporter/bool.rs @@ -16,7 +16,7 @@ struct BoolExporter { } pub(crate) fn new_exporter(array: &BoolArray) -> VortexResult> { - let validity_mask = array.validity_mask(); + let validity_mask = array.validity_mask()?; if validity_mask.all_false() { return Ok(all_invalid::new_exporter( array.len(), diff --git a/vortex-duckdb/src/exporter/constant.rs b/vortex-duckdb/src/exporter/constant.rs index 4ae90b76bb9..df7aca92c1c 100644 --- a/vortex-duckdb/src/exporter/constant.rs +++ b/vortex-duckdb/src/exporter/constant.rs @@ -31,7 +31,7 @@ pub fn new_exporter_with_mask( // TODO(joe): we can splat the constant in a specific exporter and save a copy. return Ok(validity::new_exporter( mask, - new_array_exporter(array.to_canonical().as_ref(), cache)?, + new_array_exporter(array.to_canonical()?.as_ref(), cache)?, )); } diff --git a/vortex-duckdb/src/exporter/decimal.rs b/vortex-duckdb/src/exporter/decimal.rs index 7f0b179ae14..00b0c9aef8b 100644 --- a/vortex-duckdb/src/exporter/decimal.rs +++ b/vortex-duckdb/src/exporter/decimal.rs @@ -34,7 +34,7 @@ struct DecimalZeroCopyExporter { } pub(crate) fn new_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())?; if array.values_type() == dest_values_type { @@ -128,7 +128,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 af4fde58130..61b0be9723a 100644 --- a/vortex-duckdb/src/exporter/dict.rs +++ b/vortex-duckdb/src/exporter/dict.rs @@ -50,12 +50,12 @@ 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, ); } - let codes_mask = array.codes().validity_mask(); + let codes_mask = array.codes().validity_mask()?; match codes_mask { Mask::AllTrue(_) => {} @@ -68,7 +68,7 @@ pub(crate) fn new_exporter_with_flatten( } let values_key = Arc::as_ptr(values).addr(); - let codes = array.codes().to_primitive(); + let codes = array.codes().to_primitive()?; let exporter_values = if flatten { let canonical = cache @@ -78,7 +78,7 @@ pub(crate) fn new_exporter_with_flatten( let canonical = match canonical { Some(c) => c, None => { - let canonical = values.to_canonical(); + let canonical = values.to_canonical()?; cache .canonical_cache .insert(values_key, (values.clone(), canonical.clone())); @@ -229,10 +229,13 @@ mod tests { let mut flat_chunk = DataChunk::new([LogicalType::new(cpp::duckdb_type::DUCKDB_TYPE_INTEGER)]); - new_array_exporter(arr.to_canonical().as_ref(), &ConversionCache::default()) - .unwrap() - .export(0, 3, &mut flat_chunk.get_vector(0)) - .unwrap(); + new_array_exporter( + arr.to_canonical().unwrap().as_ref(), + &ConversionCache::default(), + ) + .unwrap() + .export(0, 3, &mut flat_chunk.get_vector(0)) + .unwrap(); flat_chunk.set_len(3); assert_eq!( diff --git a/vortex-duckdb/src/exporter/fixed_size_list.rs b/vortex-duckdb/src/exporter/fixed_size_list.rs index 5f8c57375c5..32452c87cb6 100644 --- a/vortex-duckdb/src/exporter/fixed_size_list.rs +++ b/vortex-duckdb/src/exporter/fixed_size_list.rs @@ -39,12 +39,12 @@ pub(crate) fn new_exporter( let ltype: LogicalType = array.dtype().try_into()?; - if let Mask::AllFalse(len) = array.validity_mask() { + if let Mask::AllFalse(len) = array.validity_mask()? { return Ok(all_invalid::new_exporter(len, <ype)); } Ok(Box::new(FixedSizeListExporter { - validity: array.validity_mask(), + validity: array.validity_mask()?, elements_exporter, list_size: array.list_size(), })) diff --git a/vortex-duckdb/src/exporter/list.rs b/vortex-duckdb/src/exporter/list.rs index 2fa0e610595..7cadd986db0 100644 --- a/vortex-duckdb/src/exporter/list.rs +++ b/vortex-duckdb/src/exporter/list.rs @@ -72,13 +72,13 @@ pub(crate) fn new_exporter( } }; - let offsets = array.offsets().to_primitive(); - let sizes = array.sizes().to_primitive(); + let offsets = array.offsets().to_primitive()?; + let sizes = array.sizes().to_primitive()?; let boxed = match_each_integer_ptype!(offsets.ptype(), |O| { match_each_integer_ptype!(sizes.ptype(), |S| { Box::new(ListExporter { - validity: array.validity_mask(), + validity: array.validity_mask()?, duckdb_elements: shared_elements, offsets, sizes, diff --git a/vortex-duckdb/src/exporter/mod.rs b/vortex-duckdb/src/exporter/mod.rs index 0cd8f24ebe5..1cf3b821783 100644 --- a/vortex-duckdb/src/exporter/mod.rs +++ b/vortex-duckdb/src/exporter/mod.rs @@ -68,7 +68,7 @@ impl ArrayIteratorExporter { if self.array_exporter.is_none() { if let Some(array) = self.iter.next() { // Create a new array exporter for the current array. - let array = array?.to_struct(); + let array = array?.to_struct()?; self.array_exporter = Some(ArrayExporter::try_new(&array, &self.cache)?); } else { // No more arrays to export. @@ -184,7 +184,7 @@ fn new_array_exporter_with_flatten( } // Otherwise, we fall back to canonical - match array.to_canonical() { + match array.to_canonical()? { Canonical::Null(_) => Ok(all_invalid::new_exporter( array.len(), &LogicalType::new(DUCKDB_TYPE::DUCKDB_TYPE_SQLNULL), diff --git a/vortex-duckdb/src/exporter/primitive.rs b/vortex-duckdb/src/exporter/primitive.rs index 9a502e0085d..661d1629459 100644 --- a/vortex-duckdb/src/exporter/primitive.rs +++ b/vortex-duckdb/src/exporter/primitive.rs @@ -26,7 +26,7 @@ pub fn new_exporter(array: &PrimitiveArray) -> VortexResult }); Ok(if array.dtype().is_nullable() { - validity::new_exporter(array.validity_mask(), prim) + validity::new_exporter(array.validity_mask()?, prim) } else { prim }) diff --git a/vortex-duckdb/src/exporter/run_end.rs b/vortex-duckdb/src/exporter/run_end.rs index c781abc6eae..157747e3daa 100644 --- a/vortex-duckdb/src/exporter/run_end.rs +++ b/vortex-duckdb/src/exporter/run_end.rs @@ -35,7 +35,7 @@ pub(crate) fn new_exporter( array: &RunEndArray, cache: &ConversionCache, ) -> VortexResult> { - let ends = array.ends().to_primitive(); + let ends = array.ends().to_primitive()?; let values = array.values().clone(); let values_exporter = new_array_exporter(array.values(), cache)?; diff --git a/vortex-duckdb/src/exporter/struct_.rs b/vortex-duckdb/src/exporter/struct_.rs index 5484dd16dee..9e7b51f78a0 100644 --- a/vortex-duckdb/src/exporter/struct_.rs +++ b/vortex-duckdb/src/exporter/struct_.rs @@ -19,7 +19,7 @@ pub(crate) fn new_exporter( array: &StructArray, cache: &ConversionCache, ) -> VortexResult> { - let validity = array.validity_mask(); + let validity = array.validity_mask()?; // DuckDB requires that the validity of the child be a subset of the parent struct so we mask out children with // parents nullability let validity_for_mask = array.dtype().is_nullable().then(|| !&validity); diff --git a/vortex-duckdb/src/exporter/temporal.rs b/vortex-duckdb/src/exporter/temporal.rs index 7345c564b11..968040cdc4c 100644 --- a/vortex-duckdb/src/exporter/temporal.rs +++ b/vortex-duckdb/src/exporter/temporal.rs @@ -16,7 +16,7 @@ struct TemporalExporter { pub(crate) fn new_exporter(array: &TemporalArray) -> VortexResult> { Ok(Box::new(TemporalExporter { storage_type_exporter: primitive::new_exporter( - &array.temporal_values().clone().to_primitive(), + &array.temporal_values().clone().to_primitive()?, )?, })) } diff --git a/vortex-duckdb/src/exporter/varbinview.rs b/vortex-duckdb/src/exporter/varbinview.rs index 45b3a449b4f..c8dad7a9b8b 100644 --- a/vortex-duckdb/src/exporter/varbinview.rs +++ b/vortex-duckdb/src/exporter/varbinview.rs @@ -25,7 +25,7 @@ struct VarBinViewExporter { } pub(crate) fn new_exporter(array: &VarBinViewArray) -> VortexResult> { - let validity = array.validity_mask(); + let validity = array.validity_mask()?; if validity.all_false() { return Ok(all_invalid::new_exporter( array.len(), @@ -42,7 +42,7 @@ pub(crate) fn new_exporter(array: &VarBinViewArray) -> VortexResult bool { let array = vx_array::as_ref(array); - array.is_invalid(index as usize) + array.is_invalid(index as usize).unwrap_or(true) } // TODO(robert): Make this return usize and remove error @@ -97,7 +97,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 { diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index de616355e9d..ace9679e6dc 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -301,7 +301,7 @@ async fn test_read_projection() { ) ); - let actual = array.to_struct().fields()[0].clone(); + let actual = array.to_struct().unwrap().fields()[0].clone(); let expected = VarBinArray::from(strings_expected.to_vec()).into_array(); assert_arrays_eq!(actual.as_ref(), expected.as_ref()); @@ -323,7 +323,7 @@ async fn test_read_projection() { ) ); - let actual = array.to_struct().fields()[0].clone(); + let actual = array.to_struct().unwrap().fields()[0].clone(); let expected = Buffer::copy_from(numbers_expected).into_array(); assert_arrays_eq!(actual.as_ref(), expected.as_ref()); } @@ -369,9 +369,11 @@ async fn unequal_batches() { let numbers = array .to_struct() + .unwrap() .field_by_name("numbers") .unwrap() - .to_primitive(); + .to_primitive() + .unwrap(); assert_eq!(numbers.ptype(), PType::U32); } assert_eq!(item_count, 10); @@ -535,13 +537,13 @@ async fn filter_string() { .unwrap(); assert_eq!(result.len(), 1); - let names_actual = result[0].to_struct().fields()[0].clone(); + let names_actual = result[0].to_struct().unwrap().fields()[0].clone(); let names_expected = VarBinArray::from_iter(vec![Some("Joseph")], DType::Utf8(Nullability::Nullable)) .into_array(); assert_arrays_eq!(names_actual.as_ref(), names_expected.as_ref()); - let ages_actual = result[0].to_struct().fields()[1].clone(); + let ages_actual = result[0].to_struct().unwrap().fields()[1].clone(); let ages_expected = PrimitiveArray::from_option_iter([Some(25i32)]).into_array(); assert_arrays_eq!(ages_actual.as_ref(), ages_expected.as_ref()); } @@ -590,17 +592,18 @@ async fn filter_or() { .unwrap(); assert_eq!(result.len(), 1); - let names = result[0].to_struct().fields()[0].clone(); + let names = result[0].to_struct().unwrap().fields()[0].clone(); assert_eq!( - names.to_varbinview().with_iterator(|iter| iter + names.to_varbinview().unwrap().with_iterator(|iter| iter .flatten() .map(|s| unsafe { String::from_utf8_unchecked(s.to_vec()) }) .collect::>()), vec!["Joseph".to_string(), "Angela".to_string()] ); - let ages = result[0].to_struct().fields()[1].clone(); + let ages = result[0].to_struct().unwrap().fields()[1].clone(); assert_eq!( ages.to_primitive() + .unwrap() .with_iterator(|iter| iter.map(|x| x.cloned()).collect::>()), vec![Some(25), None] ); @@ -647,7 +650,7 @@ async fn filter_and() { .unwrap(); assert_eq!(result.len(), 1); - let names_actual = result[0].to_struct().fields()[0].clone(); + let names_actual = result[0].to_struct().unwrap().fields()[0].clone(); let names_expected = VarBinArray::from_iter( vec![Some("Joseph"), None], DType::Utf8(Nullability::Nullable), @@ -655,7 +658,7 @@ async fn filter_and() { .into_array(); assert_arrays_eq!(names_actual.as_ref(), names_expected.as_ref()); - let ages_actual = result[0].to_struct().fields()[1].clone(); + let ages_actual = result[0].to_struct().unwrap().fields()[1].clone(); let ages_expected = PrimitiveArray::from_option_iter([Some(25i32), Some(31i32)]).into_array(); assert_arrays_eq!(ages_actual.as_ref(), ages_expected.as_ref()); } @@ -696,7 +699,8 @@ async fn test_with_indices_simple() { .read_all() .await .unwrap() - .to_struct(); + .to_struct() + .unwrap(); assert_eq!(actual_kept_array.len(), 0); @@ -712,8 +716,9 @@ async fn test_with_indices_simple() { .read_all() .await .unwrap() - .to_struct(); - let actual_kept_numbers_array = actual_kept_array.fields()[0].to_primitive(); + .to_struct() + .unwrap(); + let actual_kept_numbers_array = actual_kept_array.fields()[0].to_primitive().unwrap(); let expected_kept_numbers: Vec = kept_indices .iter() @@ -732,7 +737,8 @@ async fn test_with_indices_simple() { .read_all() .await .unwrap() - .to_struct(); + .to_struct() + .unwrap(); let actual_numbers_array = actual_array.fields()[0].clone(); let expected_array = Buffer::copy_from(&expected_numbers).into_array(); assert_arrays_eq!(actual_numbers_array.as_ref(), expected_array.as_ref()); @@ -776,7 +782,9 @@ async fn test_with_indices_on_two_columns() { .await .unwrap() .to_struct() - .to_struct(); + .unwrap() + .to_struct() + .unwrap(); let strings_actual = array.fields()[0].clone(); let strings_expected_vec: Vec<&str> = kept_indices @@ -831,7 +839,8 @@ async fn test_with_indices_and_with_row_filter_simple() { .read_all() .await .unwrap() - .to_struct(); + .to_struct() + .unwrap(); assert_eq!(actual_kept_array.len(), 0); @@ -848,9 +857,10 @@ async fn test_with_indices_and_with_row_filter_simple() { .read_all() .await .unwrap() - .to_struct(); + .to_struct() + .unwrap(); - let actual_kept_numbers_array = actual_kept_array.fields()[0].to_primitive(); + let actual_kept_numbers_array = actual_kept_array.fields()[0].to_primitive().unwrap(); let expected_kept_numbers: Buffer = kept_indices .iter() @@ -871,9 +881,10 @@ async fn test_with_indices_and_with_row_filter_simple() { .read_all() .await .unwrap() - .to_struct(); + .to_struct() + .unwrap(); - let actual_numbers_array = actual_array.fields()[0].to_primitive(); + let actual_numbers_array = actual_array.fields()[0].to_primitive().unwrap(); let actual_numbers = actual_numbers_array.as_slice::(); assert_eq!( @@ -933,7 +944,8 @@ async fn filter_string_chunked() { .read_all() .await .unwrap() - .to_struct(); + .to_struct() + .unwrap(); assert_eq!(actual_array.len(), 1); let names_actual = actual_array.fields()[0].clone(); @@ -1024,7 +1036,8 @@ async fn test_pruning_with_or() { .read_all() .await .unwrap() - .to_struct(); + .to_struct() + .unwrap(); assert_eq!(actual_array.len(), 10); let letters_actual = actual_array.fields()[0].clone(); @@ -1094,7 +1107,8 @@ async fn test_repeated_projection() { .read_all() .await .unwrap() - .to_struct(); + .to_struct() + .unwrap(); assert_eq!( (0..actual.len()) @@ -1248,16 +1262,16 @@ async fn write_nullable_nested_struct() -> VortexResult<()> { )? .into_array(); - let result = round_trip(&array, Ok).await?.to_struct(); + let result = round_trip(&array, Ok).await?.to_struct()?; 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(); + 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(()) } @@ -1387,7 +1401,7 @@ async fn test_writer_multiple_pushes() -> VortexResult<()> { let result = file.scan()?.into_array_stream()?.read_all().await?; assert_eq!(result.len(), 9); - let numbers = result.to_struct().field_by_name("numbers")?.clone(); + let numbers = result.to_struct()?.field_by_name("numbers")?.clone(); let expected = buffer![1u32, 2, 3, 4, 5, 6, 7, 8, 9].into_array(); assert_arrays_eq!(numbers.as_ref(), expected.as_ref()); @@ -1418,7 +1432,7 @@ async fn test_writer_push_stream() -> VortexResult<()> { let result = file.scan()?.into_array_stream()?.read_all().await?; assert_eq!(result.len(), 6); - let numbers = result.to_struct().field_by_name("numbers")?.clone(); + let numbers = result.to_struct()?.field_by_name("numbers")?.clone(); let expected = buffer![1u32, 2, 3, 4, 5, 6].into_array(); assert_arrays_eq!(numbers.as_ref(), expected.as_ref()); @@ -1479,7 +1493,7 @@ async fn test_writer_empty_chunks() -> VortexResult<()> { let result = file.scan()?.into_array_stream()?.read_all().await?; assert_eq!(result.len(), 2); - let numbers = result.to_struct().field_by_name("numbers")?.clone(); + let numbers = result.to_struct()?.field_by_name("numbers")?.clone(); let expected = buffer![1u32, 2].into_array(); assert_arrays_eq!(numbers.as_ref(), expected.as_ref()); @@ -1514,7 +1528,7 @@ async fn test_writer_mixed_push_and_stream() -> VortexResult<()> { let result = file.scan()?.into_array_stream()?.read_all().await?; assert_eq!(result.len(), 6); - let numbers = result.to_struct().field_by_name("numbers")?.clone(); + let numbers = result.to_struct()?.field_by_name("numbers")?.clone(); let expected = buffer![1u32, 2, 3, 4, 5, 6].into_array(); assert_arrays_eq!(numbers.as_ref(), expected.as_ref()); @@ -1553,8 +1567,8 @@ async fn test_writer_with_complex_types() -> VortexResult<()> { assert_eq!(result.len(), 3); assert_eq!(result.dtype(), &dtype); - let strings_field = result.to_struct().field_by_name("strings").cloned()?; - let strings = strings_field.to_varbinview().with_iterator(|iter| { + let strings_field = result.to_struct()?.field_by_name("strings").cloned()?; + let strings = strings_field.to_varbinview()?.with_iterator(|iter| { iter.map(|s| s.map(|st| unsafe { String::from_utf8_unchecked(st.to_vec()) })) .collect::>() }); diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 50c109438ea..41435233dde 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -197,7 +197,7 @@ impl VortexWriteOptions { if self.file_statistics.is_empty() { None } else { - Some(FileStatistics(file_stats.stats_sets().into())) + Some(FileStatistics(file_stats.stats_sets()?.into())) }, ctx, ); diff --git a/vortex-ipc/src/iterator.rs b/vortex-ipc/src/iterator.rs index a9ec48c423c..7ca71c86071 100644 --- a/vortex-ipc/src/iterator.rs +++ b/vortex-ipc/src/iterator.rs @@ -186,9 +186,9 @@ mod test { SyncIPCReader::try_new(Cursor::new(ipc_buffer), session.registry().clone()).unwrap(); assert_eq!(reader.dtype(), array.dtype()); - let result = reader.read_all().unwrap().to_primitive(); + let result = reader.read_all().unwrap().to_primitive().unwrap(); assert_eq!( - array.to_primitive().as_slice::(), + array.to_primitive().unwrap().as_slice::(), result.as_slice::() ); } diff --git a/vortex-ipc/src/stream.rs b/vortex-ipc/src/stream.rs index 26789141133..af302e2eaec 100644 --- a/vortex-ipc/src/stream.rs +++ b/vortex-ipc/src/stream.rs @@ -226,9 +226,9 @@ mod test { .unwrap(); assert_eq!(reader.dtype(), array.dtype()); - let result = reader.read_all().await.unwrap().to_primitive(); + let result = reader.read_all().await.unwrap().to_primitive().unwrap(); assert_eq!( - array.to_primitive().as_slice::(), + array.to_primitive().unwrap().as_slice::(), result.as_slice::() ); } diff --git a/vortex-jni/src/array.rs b/vortex-jni/src/array.rs index 3a48b1438a8..716607c7640 100644 --- a/vortex-jni/src/array.rs +++ b/vortex-jni/src/array.rs @@ -224,7 +224,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeArrayMethods_getField( try_or_throw(&mut env, |_| { let field = array_ref .inner - .to_struct() + .to_struct()? .fields() .get(index as usize) .cloned() @@ -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)) }) } @@ -290,7 +290,7 @@ macro_rules! get_primitive { let scalar_value = if array_ref.is_extension { array_ref .inner - .to_extension() + .to_extension()? .storage() .scalar_at(index as usize) } else { @@ -329,7 +329,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeArrayMethods_getBigDecimal( let scalar_value = if array_ref.is_extension { array_ref .inner - .to_extension() + .to_extension()? .storage() .scalar_at(index as usize) } else { diff --git a/vortex-layout/src/layouts/compact.rs b/vortex-layout/src/layouts/compact.rs index 5aef1b443f4..7d45bbeb722 100644 --- a/vortex-layout/src/layouts/compact.rs +++ b/vortex-layout/src/layouts/compact.rs @@ -68,7 +68,7 @@ impl CompactCompressor { } pub fn compress(&self, array: &dyn Array) -> VortexResult { - self.compress_canonical(array.to_canonical()) + self.compress_canonical(array.to_canonical()?) } /// Compress a single array using the compact strategy @@ -144,9 +144,9 @@ impl CompactCompressor { // Note that since the type of our offsets and sizes is not encoded in our `DType`, // we can narrow the widths. let compressed_offsets = - self.compress(&listview.offsets().to_primitive().narrow()?.into_array())?; + self.compress(&listview.offsets().to_primitive()?.narrow()?.into_array())?; let compressed_sizes = - self.compress(&listview.sizes().to_primitive().narrow()?.into_array())?; + self.compress(&listview.sizes().to_primitive()?.narrow()?.into_array())?; // SAFETY: Since compression does not change the logical values of arrays, this is // effectively the same array but represented differently, so all invariants that @@ -247,9 +247,9 @@ mod tests { let compressed = compressor.compress(struct_array.as_ref()).unwrap(); // Verify we can decompress back to original - let decompressed = compressed.to_canonical().into_array(); + let decompressed = compressed.to_canonical().unwrap().into_array(); assert_eq!(decompressed.len(), n_rows); - let decompressed_struct = decompressed.to_struct(); + let decompressed_struct = decompressed.to_struct().unwrap(); // Verify each field can be accessed and has correct data for (i, name) in decompressed_struct.names().iter().enumerate() { diff --git a/vortex-layout/src/layouts/dict/reader.rs b/vortex-layout/src/layouts/dict/reader.rs index 02487311a90..9571a596cbd 100644 --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs @@ -94,8 +94,8 @@ impl DictReader { MaskFuture::new_true(values_len), ) .vortex_expect("must construct dict values array evaluation") + .and_then(|arr| async move { Ok(arr.to_canonical()?.into_array()) }) .map_err(Arc::new) - .map_ok(|arr| arr.to_canonical().into_array()) .boxed() .shared() }) @@ -174,7 +174,7 @@ impl LayoutReader for DictReader { let mask = mask.await?; // Short-circuit when the values are all true/false. - if values.all_valid() + if values.all_valid()? && let Some(MinMaxResult { min, max }) = min_max(&values)? { #[expect(clippy::bool_comparison, reason = "easy to follow")] @@ -185,7 +185,7 @@ impl LayoutReader for DictReader { #[expect(clippy::bool_comparison, reason = "easy to follow")] if min.as_bool().value().vortex_expect("not null") == true { // All values are true, but we still need to respect codes validity - return Ok(mask.bitand(&codes.validity_mask())); + return Ok(mask.bitand(&codes.validity_mask()?)); } } @@ -477,7 +477,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, expected); }) } diff --git a/vortex-layout/src/layouts/file_stats.rs b/vortex-layout/src/layouts/file_stats.rs index 38f0bba6441..39a77ec8ef4 100644 --- a/vortex-layout/src/layouts/file_stats.rs +++ b/vortex-layout/src/layouts/file_stats.rs @@ -92,12 +92,12 @@ impl FileStatsAccumulator { ) -> VortexResult<(SequenceId, ArrayRef)> { let (sequence_id, chunk) = chunk?; if chunk.dtype().is_struct() { - let chunk = chunk.to_struct(); + let chunk_struct = chunk.to_struct()?; for (acc, field) in self .accumulators .lock() .iter_mut() - .zip_eq(chunk.fields().iter()) + .zip_eq(chunk_struct.fields().iter()) { acc.push_chunk(field)?; } @@ -107,18 +107,21 @@ impl FileStatsAccumulator { Ok((sequence_id, chunk)) } - pub fn stats_sets(&self) -> Vec { + pub fn stats_sets(&self) -> VortexResult> { self.accumulators .lock() .iter_mut() .map(|acc| { - acc.as_stats_table() - .map(|table| { - table - .to_stats_set(&self.stats) - .vortex_expect("shouldn't fail to convert table we just created") - }) - .unwrap_or_default() + acc.as_stats_table().and_then(|table_opt| { + table_opt + .map(|table| { + table + .to_stats_set(&self.stats) + .vortex_expect("shouldn't fail to convert table we just created") + }) + .map(Ok) + .unwrap_or_else(|| Ok(StatsSet::default())) + }) }) .collect() } diff --git a/vortex-layout/src/layouts/flat/reader.rs b/vortex-layout/src/layouts/flat/reader.rs index 568d2ee0907..082ad3bcf64 100644 --- a/vortex-layout/src/layouts/flat/reader.rs +++ b/vortex-layout/src/layouts/flat/reader.rs @@ -273,10 +273,11 @@ mod test { .unwrap() .await .unwrap() - .to_primitive(); + .to_primitive() + .unwrap(); assert_eq!( - array.to_primitive().as_slice::(), + array.to_primitive().unwrap().as_slice::(), result.as_slice::() ); }) @@ -313,7 +314,8 @@ mod test { .unwrap() .await .unwrap() - .to_bool(); + .to_bool() + .unwrap(); assert_eq!( &BitBuffer::from_iter([false, false, false, true, true]), diff --git a/vortex-layout/src/layouts/flat/writer.rs b/vortex-layout/src/layouts/flat/writer.rs index e31ed09a4b9..78f90a4ff8c 100644 --- a/vortex-layout/src/layouts/flat/writer.rs +++ b/vortex-layout/src/layouts/flat/writer.rs @@ -348,24 +348,28 @@ mod tests { .unwrap(); assert_eq!( - result.validity_mask().bit_buffer(), + result.validity_mask().unwrap().bit_buffer(), AllOr::Some(&validity_boolean_buffer) ); assert_eq!( result .to_struct() + .unwrap() .field_by_name("a") .unwrap() .to_primitive() + .unwrap() .as_slice::(), &[1, 2] ); assert_eq!( result .to_struct() + .unwrap() .field_by_name("b") .unwrap() .to_primitive() + .unwrap() .as_slice::(), &[3, 4] ); diff --git a/vortex-layout/src/layouts/repartition.rs b/vortex-layout/src/layouts/repartition.rs index f53c1822fe5..399593c51b5 100644 --- a/vortex-layout/src/layouts/repartition.rs +++ b/vortex-layout/src/layouts/repartition.rs @@ -71,7 +71,7 @@ impl LayoutStrategy for RepartitionStrategy { dtype.clone(), stream.map(|chunk| { let (sequence_id, chunk) = chunk?; - VortexResult::Ok((sequence_id, chunk.to_canonical().into_array())) + VortexResult::Ok((sequence_id, chunk.to_canonical()?.into_array())) }), ) .sendable() @@ -104,7 +104,7 @@ impl LayoutStrategy for RepartitionStrategy { if !chunked.is_empty() { yield ( sequence_pointer.advance(), - chunked.to_canonical().into_array(), + chunked.to_canonical()?.into_array(), ) } } @@ -117,7 +117,7 @@ impl LayoutStrategy for RepartitionStrategy { if !to_flush.is_empty() { yield ( sequence_pointer.advance(), - to_flush.to_canonical().into_array(), + to_flush.to_canonical()?.into_array(), ) } } diff --git a/vortex-layout/src/layouts/row_idx/mod.rs b/vortex-layout/src/layouts/row_idx/mod.rs index 1160c91fc04..c33a2b625e3 100644 --- a/vortex-layout/src/layouts/row_idx/mod.rs +++ b/vortex-layout/src/layouts/row_idx/mod.rs @@ -339,7 +339,8 @@ mod tests { .unwrap() .await .unwrap() - .to_bool(); + .to_bool() + .unwrap(); assert_eq!( &BitBuffer::from_iter([false, false, true, false, false]), @@ -380,7 +381,8 @@ mod tests { .unwrap() .await .unwrap() - .to_bool(); + .to_bool() + .unwrap(); assert_eq!( &BitBuffer::from_iter([false, false, false, false, true]), @@ -425,7 +427,8 @@ mod tests { .unwrap() .await .unwrap() - .to_bool(); + .to_bool() + .unwrap(); assert_eq!( vec![true, false, true, false, true], diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index 342c78832d4..fdf2b325f3d 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -350,11 +350,11 @@ impl LayoutReader for StructReader { Ok(Box::pin(async move { if let Some(validity_fut) = validity_fut { let (array, validity) = try_join!(projected, validity_fut)?; - let mask = Mask::from_buffer(validity.to_bool().bit_buffer().not()); + let mask = Mask::from_buffer(validity.to_bool()?.bit_buffer().not()); // If root expression was a pack, then we apply the validity to each child field if is_pack_merge { - let struct_array = array.to_struct(); + let struct_array = array.to_struct()?; let masked_fields: Vec = struct_array .fields() .iter() @@ -608,7 +608,12 @@ mod tests { .unwrap(); assert_eq!( vec![true, false, false], - result.to_bool().bit_buffer().iter().collect::>() + result + .to_bool() + .unwrap() + .bit_buffer() + .iter() + .collect::>() ); } @@ -633,7 +638,12 @@ mod tests { assert_eq!( vec![true, false], - result.to_bool().bit_buffer().iter().collect::>() + result + .to_bool() + .unwrap() + .bit_buffer() + .iter() + .collect::>() ); } @@ -663,9 +673,11 @@ mod tests { assert_eq!( result .to_struct() + .unwrap() .field_by_name("a") .unwrap() .to_primitive() + .unwrap() .as_slice::(), [7, 2].as_slice() ); @@ -673,9 +685,11 @@ mod tests { assert_eq!( result .to_struct() + .unwrap() .field_by_name("b") .unwrap() .to_primitive() + .unwrap() .as_slice::(), [4, 5].as_slice() ); diff --git a/vortex-layout/src/layouts/struct_/writer.rs b/vortex-layout/src/layouts/struct_/writer.rs index b657bfacbf7..8ec333d8a62 100644 --- a/vortex-layout/src/layouts/struct_/writer.rs +++ b/vortex-layout/src/layouts/struct_/writer.rs @@ -94,12 +94,12 @@ impl LayoutStrategy for StructStrategy { let columns_vec_stream = stream.map(move |chunk| { let (sequence_id, chunk) = chunk?; let mut sequence_pointer = sequence_id.descend(); - let struct_chunk = chunk.to_struct(); + let struct_chunk = chunk.to_struct()?; let mut columns: Vec<(SequenceId, ArrayRef)> = Vec::new(); 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/reader.rs b/vortex-layout/src/layouts/zoned/reader.rs index f084950c966..8b073d486da 100644 --- a/vortex-layout/src/layouts/zoned/reader.rs +++ b/vortex-layout/src/layouts/zoned/reader.rs @@ -135,7 +135,7 @@ impl ZonedReader { .vortex_expect("Failed construct zone map evaluation"); async move { - let zones_array = zones_eval.await?.to_struct(); + let zones_array = zones_eval.await?.to_struct()?; // SAFETY: This is only fine to call because we perform validation above Ok(unsafe { ZoneMap::new_unchecked(zones_array, present_stats) }) } 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 519edc4e8be..029530c7816 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -214,7 +214,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(); @@ -228,7 +228,7 @@ impl StatsAccumulator { let values = builder.finish(); // We drop any all-null stats columns - if values.all_invalid() { + if values.all_invalid()? { continue; } @@ -238,14 +238,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(), - }) + })) } } @@ -299,7 +299,10 @@ mod tests { StatsAccumulator::new(builder.dtype(), &[Stat::Max, Stat::Min, Stat::Sum], 12); acc.push_chunk(&builder.finish()).vortex_unwrap(); acc.push_chunk(&builder2.finish()).vortex_unwrap(); - let stats_table = acc.as_stats_table().vortex_expect("Must have stats table"); + let stats_table = acc + .as_stats_table() + .vortex_expect("Must have stats table") + .unwrap(); assert_eq!( stats_table.array.names().as_ref(), &[ @@ -310,11 +313,17 @@ mod tests { ] ); assert_eq!( - stats_table.array.fields()[1].to_bool().bit_buffer(), + stats_table.array.fields()[1] + .to_bool() + .unwrap() + .bit_buffer(), &BitBuffer::from(vec![false, true]) ); assert_eq!( - stats_table.array.fields()[3].to_bool().bit_buffer(), + stats_table.array.fields()[3] + .to_bool() + .unwrap() + .bit_buffer(), &BitBuffer::from(vec![true, false]) ); } @@ -324,7 +333,10 @@ mod tests { let array = buffer![0, 1, 2].into_array(); let mut acc = StatsAccumulator::new(array.dtype(), &[Stat::Max, Stat::Min, Stat::Sum], 12); acc.push_chunk(&array).vortex_unwrap(); - let stats_table = acc.as_stats_table().vortex_expect("Must have stats table"); + let stats_table = acc + .as_stats_table() + .vortex_expect("Must have stats table") + .unwrap(); assert_eq!( stats_table.array.names().as_ref(), &[ @@ -336,11 +348,17 @@ mod tests { ] ); assert_eq!( - stats_table.array.fields()[1].to_bool().bit_buffer(), + stats_table.array.fields()[1] + .to_bool() + .unwrap() + .bit_buffer(), &BitBuffer::from(vec![false]) ); assert_eq!( - stats_table.array.fields()[3].to_bool().bit_buffer(), + stats_table.array.fields()[3] + .to_bool() + .unwrap() + .bit_buffer(), &BitBuffer::from(vec![false]) ); } diff --git a/vortex-python/src/arrays/compressed.rs b/vortex-python/src/arrays/compressed.rs index c294d7ee5c0..014a0a4e3fd 100644 --- a/vortex-python/src/arrays/compressed.rs +++ b/vortex-python/src/arrays/compressed.rs @@ -89,7 +89,7 @@ impl PyZigZagArray { #[staticmethod] pub fn encode(array: PyArrayRef) -> PyResult { Ok(PyVortex( - zigzag_encode(array.inner().clone().to_primitive())?.into_array(), + zigzag_encode(array.inner().clone().to_primitive()?)?.into_array(), )) } } diff --git a/vortex-python/src/arrays/mod.rs b/vortex-python/src/arrays/mod.rs index e9443c16aec..a843699138b 100644 --- a/vortex-python/src/arrays/mod.rs +++ b/vortex-python/src/arrays/mod.rs @@ -487,7 +487,7 @@ impl PyArray { /// ``` fn filter(slf: Bound, mask: PyArrayRef) -> PyResult { let slf = PyArrayRef::extract(slf.as_any().as_borrowed())?.into_inner(); - let mask = (&*mask as &dyn Array).to_bool().to_mask_fill_null_false(); + let mask = (&*mask as &dyn Array).to_bool()?.to_mask_fill_null_false(); let inner = vortex::compute::filter(&*slf, &mask)?; Ok(PyArrayRef::from(inner)) } diff --git a/vortex-python/src/arrays/py/vtable.rs b/vortex-python/src/arrays/py/vtable.rs index fbb722bfc42..1380b238bbc 100644 --- a/vortex-python/src/arrays/py/vtable.rs +++ b/vortex-python/src/arrays/py/vtable.rs @@ -168,7 +168,7 @@ impl BaseArrayVTable for PythonVTable { } impl CanonicalVTable for PythonVTable { - fn canonicalize(_array: &PythonArray) -> Canonical { + fn canonicalize(_array: &PythonArray) -> VortexResult { todo!() } } @@ -184,19 +184,19 @@ impl OperationsVTable for PythonVTable { } impl ValidityVTable for PythonVTable { - fn is_valid(_array: &PythonArray, _index: usize) -> bool { + fn is_valid(_array: &PythonArray, _index: usize) -> VortexResult { todo!() } - fn all_valid(_array: &PythonArray) -> bool { + fn all_valid(_array: &PythonArray) -> VortexResult { todo!() } - fn all_invalid(_array: &PythonArray) -> bool { + fn all_invalid(_array: &PythonArray) -> VortexResult { todo!() } - fn validity_mask(_array: &PythonArray) -> Mask { + fn validity_mask(_array: &PythonArray) -> VortexResult { todo!() } } diff --git a/vortex-python/src/dataset.rs b/vortex-python/src/dataset.rs index 9d83de6cc30..0ea1848b87e 100644 --- a/vortex-python/src/dataset.rs +++ b/vortex-python/src/dataset.rs @@ -59,7 +59,7 @@ pub fn read_array_from_reader( } if let Some(indices) = indices { - let indices = indices.to_primitive().into_buffer(); + let indices = indices.to_primitive()?.into_buffer(); scan = scan.with_row_indices(indices); } diff --git a/vortex-python/src/file.rs b/vortex-python/src/file.rs index a640a1dd9f1..b7630820053 100644 --- a/vortex-python/src/file.rs +++ b/vortex-python/src/file.rs @@ -183,7 +183,7 @@ impl PyVortexFile { if let Some(indices) = indices { let indices = cast(indices.as_ref(), &DType::Primitive(PType::U64, NonNullable))? - .to_primitive() + .to_primitive()? .into_buffer::(); builder = builder.with_row_indices(indices); } diff --git a/vortex-tui/src/browse/ui/layouts.rs b/vortex-tui/src/browse/ui/layouts.rs index 110d3eb70ac..5e3f4ff662a 100644 --- a/vortex-tui/src/browse/ui/layouts.rs +++ b/vortex-tui/src/browse/ui/layouts.rs @@ -153,7 +153,7 @@ fn render_array(app: &AppState<'_>, area: Rect, buf: &mut Buffer, is_stats_table if is_stats_table { // Render the stats table horizontally - let struct_array = array.to_struct(); + let struct_array = array.to_struct().vortex_expect("to_struct"); // add 1 for the chunk column let field_count = struct_array.struct_fields().nfields() + 1; let header = std::iter::once("chunk") diff --git a/vortex/benches/common_encoding_tree_throughput.rs b/vortex/benches/common_encoding_tree_throughput.rs index 53df57d49bb..7863326450c 100644 --- a/vortex/benches/common_encoding_tree_throughput.rs +++ b/vortex/benches/common_encoding_tree_throughput.rs @@ -70,10 +70,12 @@ mod setup { PrimitiveArray::from_iter((0..NUM_VALUES).map(|_| rng.random_range(42u32..256))); let int_array = cast(uint_array.as_ref(), PType::I32.into()) .unwrap() - .to_primitive(); + .to_primitive() + .unwrap(); let float_array = cast(uint_array.as_ref(), PType::F64.into()) .unwrap() - .to_primitive(); + .to_primitive() + .unwrap(); (uint_array, int_array, float_array) } @@ -94,7 +96,7 @@ mod setup { let alp_compressed = alp_encode(&float_array, None).unwrap(); // Manually construct ALP <- FoR <- BitPacked tree - let for_array = FoRArray::encode(alp_compressed.encoded().to_primitive()).unwrap(); + let for_array = FoRArray::encode(alp_compressed.encoded().to_primitive().unwrap()).unwrap(); let inner = for_array.encoded(); let bp = BitPackedArray::encode(inner, 8).unwrap(); let for_with_bp = @@ -165,7 +167,7 @@ mod setup { let runend = RunEndArray::encode(prim_array.into_array()).unwrap(); // Compress the ends with FoR <- BitPacked - let ends_prim = runend.ends().to_primitive(); + let ends_prim = runend.ends().to_primitive().unwrap(); let ends_for = FoRArray::encode(ends_prim).unwrap(); let ends_inner = ends_for.encoded(); let ends_bp = BitPackedArray::encode(ends_inner, 8).unwrap(); @@ -175,7 +177,7 @@ mod setup { .into_array(); // Compress the values with BitPacked - let values_prim = runend.values().to_primitive(); + let values_prim = runend.values().to_primitive().unwrap(); let compressed_values = BitPackedArray::encode(values_prim.as_ref(), 8) .unwrap() .into_array(); @@ -239,7 +241,7 @@ mod setup { // Compress the VarBin offsets with BitPacked let codes = fsst.codes(); - let offsets_prim = codes.offsets().to_primitive(); + let offsets_prim = codes.offsets().to_primitive().unwrap(); let offsets_bp = BitPackedArray::encode(offsets_prim.as_ref(), 20).unwrap(); // Rebuild VarBin with compressed offsets @@ -290,7 +292,7 @@ mod setup { let parts = split_temporal(temporal_array.clone()).unwrap(); // Compress days with FoR <- BitPacked - let days_prim = parts.days.to_primitive(); + let days_prim = parts.days.to_primitive().unwrap(); let days_for = FoRArray::encode(days_prim).unwrap(); let days_inner = days_for.encoded(); let days_bp = BitPackedArray::encode(days_inner, 16).unwrap(); @@ -300,7 +302,7 @@ mod setup { .into_array(); // Compress seconds with FoR <- BitPacked - let seconds_prim = parts.seconds.to_primitive(); + let seconds_prim = parts.seconds.to_primitive().unwrap(); let seconds_for = FoRArray::encode(seconds_prim).unwrap(); let seconds_inner = seconds_for.encoded(); let seconds_bp = BitPackedArray::encode(seconds_inner, 17).unwrap(); @@ -312,7 +314,7 @@ mod setup { .into_array(); // Compress subseconds with FoR <- BitPacked - let subseconds_prim = parts.subseconds.to_primitive(); + let subseconds_prim = parts.subseconds.to_primitive().unwrap(); let subseconds_for = FoRArray::encode(subseconds_prim).unwrap(); let subseconds_inner = subseconds_for.encoded(); let subseconds_bp = BitPackedArray::encode(subseconds_inner, 20).unwrap(); diff --git a/vortex/benches/single_encoding_throughput.rs b/vortex/benches/single_encoding_throughput.rs index daac0254559..36e56fd8209 100644 --- a/vortex/benches/single_encoding_throughput.rs +++ b/vortex/benches/single_encoding_throughput.rs @@ -58,10 +58,12 @@ fn setup_primitive_arrays() -> (PrimitiveArray, PrimitiveArray, PrimitiveArray) PrimitiveArray::from_iter((0..NUM_VALUES).map(|_| rng.random_range(42u32..256))); let int_array = cast(uint_array.as_ref(), PType::I32.into()) .unwrap() - .to_primitive(); + .to_primitive() + .unwrap(); let float_array = cast(uint_array.as_ref(), PType::F64.into()) .unwrap() - .to_primitive(); + .to_primitive() + .unwrap(); (uint_array, int_array, float_array) } diff --git a/vortex/src/lib.rs b/vortex/src/lib.rs index c2107a81dfd..65445c0057f 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -229,7 +229,7 @@ mod test { .await?; assert_eq!(recovered_array.len(), array.len()); - let recovered_primitive = recovered_array.to_primitive(); + let recovered_primitive = recovered_array.to_primitive().unwrap(); assert_eq!(recovered_primitive.validity(), array.validity()); assert_eq!(recovered_primitive.buffer::(), array.buffer::());