From d2904cb58ddb4629f5cfe2330517491ec4276188 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Tue, 19 Nov 2024 18:28:16 +0000 Subject: [PATCH] Remove as_primitive (#1376) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit It hides a hidden clone as well as panicking 👍 --- encodings/alp/src/alp/compress.rs | 18 +++++++++++--- encodings/alp/src/alp/compute.rs | 6 ++++- encodings/bytebool/src/compute.rs | 4 ++-- encodings/datetime-parts/src/compress.rs | 2 +- encodings/datetime-parts/src/compute.rs | 12 +++++++--- encodings/fsst/src/canonical.rs | 6 +++-- encodings/runend/src/compress.rs | 6 ++--- .../src/array/primitive/compute/cast.rs | 24 +++++++++++++------ .../src/array/primitive/compute/fill.rs | 8 +++---- vortex-array/src/array/primitive/mod.rs | 6 ----- vortex-file/src/lib.rs | 7 +++--- .../src/compressors/alp.rs | 4 ++-- .../src/compressors/bitpacked.rs | 4 ++-- .../src/compressors/runend.rs | 4 ++-- 14 files changed, 70 insertions(+), 41 deletions(-) diff --git a/encodings/alp/src/alp/compress.rs b/encodings/alp/src/alp/compress.rs index d0c6dd3092..3e467ee923 100644 --- a/encodings/alp/src/alp/compress.rs +++ b/encodings/alp/src/alp/compress.rs @@ -124,7 +124,11 @@ mod tests { let encoded = alp_encode(&array).unwrap(); assert!(encoded.patches().is_none()); assert_eq!( - encoded.encoded().as_primitive().maybe_null_slice::(), + encoded + .encoded() + .into_primitive() + .unwrap() + .maybe_null_slice::(), vec![1234; 1025] ); assert_eq!(encoded.exponents(), Exponents { e: 9, f: 6 }); @@ -142,7 +146,11 @@ mod tests { let encoded = alp_encode(&array).unwrap(); assert!(encoded.patches().is_none()); assert_eq!( - encoded.encoded().as_primitive().maybe_null_slice::(), + encoded + .encoded() + .into_primitive() + .unwrap() + .maybe_null_slice::(), vec![0, 1234, 0] ); assert_eq!(encoded.exponents(), Exponents { e: 9, f: 6 }); @@ -160,7 +168,11 @@ mod tests { let encoded = alp_encode(&array).unwrap(); assert!(encoded.patches().is_some()); assert_eq!( - encoded.encoded().as_primitive().maybe_null_slice::(), + encoded + .encoded() + .into_primitive() + .unwrap() + .maybe_null_slice::(), vec![1234i64, 2718, 1234, 4000] // fill forward ); assert_eq!(encoded.exponents(), Exponents { e: 16, f: 13 }); diff --git a/encodings/alp/src/alp/compute.rs b/encodings/alp/src/alp/compute.rs index 2d605b28b1..c7a17063a6 100644 --- a/encodings/alp/src/alp/compute.rs +++ b/encodings/alp/src/alp/compute.rs @@ -172,7 +172,11 @@ mod tests { let encoded = alp_encode(&array).unwrap(); assert!(encoded.patches().is_none()); assert_eq!( - encoded.encoded().as_primitive().maybe_null_slice::(), + encoded + .encoded() + .into_primitive() + .unwrap() + .maybe_null_slice::(), vec![1234; 1025] ); diff --git a/encodings/bytebool/src/compute.rs b/encodings/bytebool/src/compute.rs index 904aaf23aa..50e402bb6e 100644 --- a/encodings/bytebool/src/compute.rs +++ b/encodings/bytebool/src/compute.rs @@ -3,7 +3,7 @@ use vortex_array::compute::unary::{FillForwardFn, ScalarAtFn}; use vortex_array::compute::{ArrayCompute, SliceFn, TakeFn, TakeOptions}; use vortex_array::validity::{ArrayValidity, Validity}; use vortex_array::variants::PrimitiveArrayTrait; -use vortex_array::{ArrayDType, ArrayData, IntoArrayData}; +use vortex_array::{ArrayDType, ArrayData, IntoArrayData, IntoArrayVariant}; use vortex_dtype::{match_each_integer_ptype, Nullability}; use vortex_error::{vortex_err, VortexResult}; use vortex_scalar::Scalar; @@ -51,7 +51,7 @@ impl SliceFn for ByteBoolArray { impl TakeFn for ByteBoolArray { fn take(&self, indices: &ArrayData, _options: TakeOptions) -> VortexResult { let validity = self.validity(); - let indices = indices.clone().as_primitive(); + let indices = indices.clone().into_primitive()?; let bools = self.maybe_null_slice(); let arr = match validity { diff --git a/encodings/datetime-parts/src/compress.rs b/encodings/datetime-parts/src/compress.rs index 8519eb276f..72e6f01c51 100644 --- a/encodings/datetime-parts/src/compress.rs +++ b/encodings/datetime-parts/src/compress.rs @@ -24,7 +24,7 @@ pub fn split_temporal(array: TemporalArray) -> VortexResult { &temporal_values, &DType::Primitive(PType::I64, temporal_values.dtype().nullability()), )? - .as_primitive(); + .into_primitive()?; let divisor = match array.temporal_metadata().time_unit() { TimeUnit::Ns => 1_000_000_000, diff --git a/encodings/datetime-parts/src/compute.rs b/encodings/datetime-parts/src/compute.rs index 0e8b3c5f34..115d5853a2 100644 --- a/encodings/datetime-parts/src/compute.rs +++ b/encodings/datetime-parts/src/compute.rs @@ -176,9 +176,15 @@ mod test { seconds, subseconds, } = split_temporal(temporal_array.clone()).unwrap(); - assert_eq!(days.as_primitive().validity(), validity); - assert_eq!(seconds.as_primitive().validity(), Validity::NonNullable); - assert_eq!(subseconds.as_primitive().validity(), Validity::NonNullable); + assert_eq!(days.clone().into_primitive().unwrap().validity(), validity); + assert_eq!( + seconds.clone().into_primitive().unwrap().validity(), + Validity::NonNullable + ); + assert_eq!( + subseconds.clone().into_primitive().unwrap().validity(), + Validity::NonNullable + ); assert_eq!(validity, raw_millis.validity()); let date_times = DateTimePartsArray::try_new( diff --git a/encodings/fsst/src/canonical.rs b/encodings/fsst/src/canonical.rs index d86fdda17a..425fa33922 100644 --- a/encodings/fsst/src/canonical.rs +++ b/encodings/fsst/src/canonical.rs @@ -1,7 +1,9 @@ use arrow_array::builder::make_view; use arrow_buffer::Buffer; use vortex_array::array::{PrimitiveArray, VarBinArray, VarBinViewArray}; -use vortex_array::{ArrayDType, ArrayData, Canonical, IntoArrayData, IntoCanonical}; +use vortex_array::{ + ArrayDType, ArrayData, Canonical, IntoArrayData, IntoArrayVariant, IntoCanonical, +}; use vortex_error::VortexResult; use crate::FSSTArray; @@ -21,7 +23,7 @@ impl IntoCanonical for FSSTArray { let compressed_bytes = VarBinArray::try_from(self.codes())? .sliced_bytes()? - .as_primitive(); + .into_primitive()?; // Bulk-decompress the entire array. let uncompressed_bytes = diff --git a/encodings/runend/src/compress.rs b/encodings/runend/src/compress.rs index d6e0fa389e..363435fac3 100644 --- a/encodings/runend/src/compress.rs +++ b/encodings/runend/src/compress.rs @@ -148,7 +148,7 @@ pub fn runend_decode_typed_bool + FromPrimit mod test { use vortex_array::array::PrimitiveArray; use vortex_array::validity::{ArrayValidity, Validity}; - use vortex_array::IntoArrayData; + use vortex_array::{IntoArrayData, IntoArrayVariant}; use crate::compress::{runend_decode_primitive, runend_encode}; use crate::RunEndArray; @@ -190,8 +190,8 @@ mod test { .unwrap(); let decoded = runend_decode_primitive( - arr.ends().as_primitive(), - arr.values().as_primitive(), + arr.ends().into_primitive().unwrap(), + arr.values().into_primitive().unwrap(), arr.validity(), 0, arr.len(), diff --git a/vortex-array/src/array/primitive/compute/cast.rs b/vortex-array/src/array/primitive/compute/cast.rs index e77b2eb15f..f8c55b01c1 100644 --- a/vortex-array/src/array/primitive/compute/cast.rs +++ b/vortex-array/src/array/primitive/compute/cast.rs @@ -69,42 +69,49 @@ mod test { use crate::array::PrimitiveArray; use crate::compute::unary::try_cast; use crate::validity::Validity; - use crate::IntoArrayData; + use crate::{IntoArrayData, IntoArrayVariant}; #[test] fn cast_u32_u8() { let arr = vec![0u32, 10, 200].into_array(); // cast from u32 to u8 - let p = try_cast(&arr, PType::U8.into()).unwrap().as_primitive(); + let p = try_cast(&arr, PType::U8.into()) + .unwrap() + .into_primitive() + .unwrap(); assert_eq!(p.maybe_null_slice::(), vec![0u8, 10, 200]); assert_eq!(p.validity(), Validity::NonNullable); // to nullable let p = try_cast(&p, &DType::Primitive(PType::U8, Nullability::Nullable)) .unwrap() - .as_primitive(); + .into_primitive() + .unwrap(); assert_eq!(p.maybe_null_slice::(), vec![0u8, 10, 200]); assert_eq!(p.validity(), Validity::AllValid); // back to non-nullable let p = try_cast(&p, &DType::Primitive(PType::U8, Nullability::NonNullable)) .unwrap() - .as_primitive(); + .into_primitive() + .unwrap(); assert_eq!(p.maybe_null_slice::(), vec![0u8, 10, 200]); assert_eq!(p.validity(), Validity::NonNullable); // to nullable u32 let p = try_cast(&p, &DType::Primitive(PType::U32, Nullability::Nullable)) .unwrap() - .as_primitive(); + .into_primitive() + .unwrap(); assert_eq!(p.maybe_null_slice::(), vec![0u32, 10, 200]); assert_eq!(p.validity(), Validity::AllValid); // to non-nullable u8 let p = try_cast(&p, &DType::Primitive(PType::U8, Nullability::NonNullable)) .unwrap() - .as_primitive(); + .into_primitive() + .unwrap(); assert_eq!(p.maybe_null_slice::(), vec![0u8, 10, 200]); assert_eq!(p.validity(), Validity::NonNullable); } @@ -112,7 +119,10 @@ mod test { #[test] fn cast_u32_f32() { let arr = vec![0u32, 10, 200].into_array(); - let u8arr = try_cast(&arr, PType::F32.into()).unwrap().as_primitive(); + let u8arr = try_cast(&arr, PType::F32.into()) + .unwrap() + .into_primitive() + .unwrap(); assert_eq!(u8arr.maybe_null_slice::(), vec![0.0f32, 10., 200.]); } diff --git a/vortex-array/src/array/primitive/compute/fill.rs b/vortex-array/src/array/primitive/compute/fill.rs index a8555cbfd9..0f8709949b 100644 --- a/vortex-array/src/array/primitive/compute/fill.rs +++ b/vortex-array/src/array/primitive/compute/fill.rs @@ -52,13 +52,13 @@ mod test { use crate::array::BoolArray; use crate::compute::unary::fill_forward; use crate::validity::{ArrayValidity, Validity}; - use crate::IntoArrayData; + use crate::{IntoArrayData, IntoArrayVariant}; #[test] fn leading_none() { let arr = PrimitiveArray::from_nullable_vec(vec![None, Some(8u8), None, Some(10), None]) .into_array(); - let p = fill_forward(&arr).unwrap().as_primitive(); + let p = fill_forward(&arr).unwrap().into_primitive().unwrap(); assert_eq!(p.maybe_null_slice::(), vec![0, 8, 8, 10, 10]); assert!(p.logical_validity().all_valid()); } @@ -69,7 +69,7 @@ mod test { PrimitiveArray::from_nullable_vec(vec![Option::::None, None, None, None, None]) .into_array(); - let p = fill_forward(&arr).unwrap().as_primitive(); + let p = fill_forward(&arr).unwrap().into_primitive().unwrap(); assert_eq!(p.maybe_null_slice::(), vec![0, 0, 0, 0, 0]); assert!(p.logical_validity().all_valid()); } @@ -81,7 +81,7 @@ mod test { Validity::Array(BoolArray::from_iter([true, true, true, true, true]).into_array()), ) .into_array(); - let p = fill_forward(&arr).unwrap().as_primitive(); + let p = fill_forward(&arr).unwrap().into_primitive().unwrap(); assert_eq!(p.maybe_null_slice::(), vec![8, 10, 12, 14, 16]); assert!(p.logical_validity().all_valid()); } diff --git a/vortex-array/src/array/primitive/mod.rs b/vortex-array/src/array/primitive/mod.rs index 8fb51811dd..92f05561ae 100644 --- a/vortex-array/src/array/primitive/mod.rs +++ b/vortex-array/src/array/primitive/mod.rs @@ -334,12 +334,6 @@ impl AcceptArrayVisitor for PrimitiveArray { } } -impl ArrayData { - pub fn as_primitive(&self) -> PrimitiveArray { - PrimitiveArray::try_from(self).vortex_expect("Expected primitive array") - } -} - // This is an arbitrary value, tried a few seems like this is a better value than smaller ones, // I assume there's some hardware dependency here but this seems to be good enough const CHUNK_SIZE: usize = 1024; diff --git a/vortex-file/src/lib.rs b/vortex-file/src/lib.rs index 24ed58bf7d..80e14a2540 100644 --- a/vortex-file/src/lib.rs +++ b/vortex-file/src/lib.rs @@ -213,7 +213,7 @@ mod test { use vortex_array::array::{ChunkedArray, PrimitiveArray, PrimitiveEncoding}; use vortex_array::encoding::ArrayEncoding; use vortex_array::stream::ArrayStreamExt; - use vortex_array::{ArrayDType, Context, IntoArrayData}; + use vortex_array::{ArrayDType, Context, IntoArrayData, IntoArrayVariant}; use vortex_buffer::Buffer; use vortex_error::VortexResult; use vortex_io::VortexBufReader; @@ -286,13 +286,14 @@ mod test { assert_eq!(next.encoding().id(), PrimitiveEncoding.id()); assert_eq!( - next.as_primitive().maybe_null_slice::(), + next.into_primitive().unwrap().maybe_null_slice::(), vec![2999989, 2999988, 2999987, 2999986, 2899999, 0, 0] ); assert_eq!( block_on(async { take_iter.try_next().await })? .expect("Expected a chunk") - .as_primitive() + .into_primitive() + .unwrap() .maybe_null_slice::(), vec![5999999] ); diff --git a/vortex-sampling-compressor/src/compressors/alp.rs b/vortex-sampling-compressor/src/compressors/alp.rs index d57e4eade7..046ecd9afa 100644 --- a/vortex-sampling-compressor/src/compressors/alp.rs +++ b/vortex-sampling-compressor/src/compressors/alp.rs @@ -6,7 +6,7 @@ use vortex_array::array::PrimitiveArray; use vortex_array::encoding::EncodingRef; use vortex_array::stats::ArrayStatistics as _; use vortex_array::variants::PrimitiveArrayTrait; -use vortex_array::{ArrayData, ArrayDef, IntoArrayData}; +use vortex_array::{ArrayData, ArrayDef, IntoArrayData, IntoArrayVariant}; use vortex_dtype::PType; use vortex_error::VortexResult; @@ -45,7 +45,7 @@ impl EncodingCompressor for ALPCompressor { ctx: SamplingCompressor<'a>, ) -> VortexResult> { // TODO(robert): Fill forward nulls? - let parray = array.as_primitive(); + let parray = array.clone().into_primitive()?; let (exponents, encoded, patches) = match_each_alp_float_ptype!( parray.ptype(), |$T| { diff --git a/vortex-sampling-compressor/src/compressors/bitpacked.rs b/vortex-sampling-compressor/src/compressors/bitpacked.rs index 7d8b6007ee..4fd0700445 100644 --- a/vortex-sampling-compressor/src/compressors/bitpacked.rs +++ b/vortex-sampling-compressor/src/compressors/bitpacked.rs @@ -3,7 +3,7 @@ use vortex_array::array::PrimitiveArray; use vortex_array::encoding::EncodingRef; use vortex_array::stats::ArrayStatistics; use vortex_array::variants::PrimitiveArrayTrait; -use vortex_array::{ArrayData, IntoArrayData}; +use vortex_array::{ArrayData, IntoArrayData, IntoArrayVariant}; use vortex_error::{vortex_err, vortex_panic, VortexResult}; use vortex_fastlanes::{ bitpack, count_exceptions, find_best_bit_width, find_min_patchless_bit_width, gather_patches, @@ -77,7 +77,7 @@ impl EncodingCompressor for BitPackedCompressor { like: Option>, ctx: SamplingCompressor<'a>, ) -> VortexResult> { - let parray = array.as_primitive(); + let parray = array.clone().into_primitive()?; let bit_width_freq = parray .statistics() .compute_bit_width_freq() diff --git a/vortex-sampling-compressor/src/compressors/runend.rs b/vortex-sampling-compressor/src/compressors/runend.rs index cbb3d2c01b..689cea4d14 100644 --- a/vortex-sampling-compressor/src/compressors/runend.rs +++ b/vortex-sampling-compressor/src/compressors/runend.rs @@ -2,7 +2,7 @@ use vortex_array::aliases::hash_set::HashSet; use vortex_array::array::Primitive; use vortex_array::encoding::EncodingRef; use vortex_array::stats::ArrayStatistics; -use vortex_array::{ArrayData, ArrayDef, IntoArrayData}; +use vortex_array::{ArrayData, ArrayDef, IntoArrayData, IntoArrayVariant}; use vortex_error::VortexResult; use vortex_runend::compress::runend_encode; use vortex_runend::{RunEnd, RunEndArray, RunEndEncoding}; @@ -49,7 +49,7 @@ impl EncodingCompressor for RunEndCompressor { like: Option>, ctx: SamplingCompressor<'a>, ) -> VortexResult> { - let primitive_array = array.as_primitive(); + let primitive_array = array.clone().into_primitive()?; let (ends, values) = runend_encode(&primitive_array); let compressed_ends = ctx