feat: teach PrimitiveArrayTrait iterate_primitive_array! and RowMask …

…from_index_array (#1241)

This touches a dumb number of files and regardless of how the
`with_indices` implementation shakes out, I think we'll need it.

encodings/alp/src/alp/array.rs

@@ -105,13 +105,6 @@ impl ALPArray {
         })
     }
 
-    #[inline]
-    pub fn ptype(&self) -> PType {
-        self.dtype()
-            .try_into()
-            .vortex_expect("Failed to convert DType to PType")
-    }
-
     #[inline]
     fn encoded_dtype(&self) -> DType {
         match self.dtype() {

encodings/alp/src/alp/compress.rs

@@ -1,5 +1,6 @@
 use vortex_array::array::{PrimitiveArray, Sparse, SparseArray};
 use vortex_array::validity::Validity;
+use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::{Array, ArrayDType, ArrayDef, IntoArray, IntoArrayVariant};
 use vortex_dtype::{NativePType, PType};
 use vortex_error::{vortex_bail, VortexExpect as _, VortexResult};

encodings/alp/src/alp/compute.rs

@@ -5,6 +5,7 @@ use vortex_array::compute::{
 };
 use vortex_array::stats::{ArrayStatistics, Stat};
 use vortex_array::validity::Validity;
+use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::{Array, ArrayDType, IntoArray};
 use vortex_error::{VortexExpect, VortexResult};
 use vortex_scalar::{PValue, Scalar};

encodings/alp/src/alp_rd/mod.rs

@@ -1,5 +1,6 @@
 pub use array::*;
 use vortex_array::validity::Validity;
+use vortex_array::variants::PrimitiveArrayTrait;
 
 mod array;
 mod compute;

encodings/bytebool/src/compute.rs

@@ -2,6 +2,7 @@ use num_traits::AsPrimitive;
 use vortex_array::compute::unary::{FillForwardFn, ScalarAtFn};
 use vortex_array::compute::{ArrayCompute, SliceFn, TakeFn};
 use vortex_array::validity::{ArrayValidity, Validity};
+use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::{Array, ArrayDType, IntoArray};
 use vortex_dtype::{match_each_integer_ptype, Nullability};
 use vortex_error::{vortex_err, VortexResult};

encodings/dict/src/array.rs

@@ -9,6 +9,7 @@ use vortex_array::compute::unary::scalar_at;
 use vortex_array::encoding::ids;
 use vortex_array::stats::StatsSet;
 use vortex_array::validity::{ArrayValidity, LogicalValidity};
+use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::{
     impl_encoding, Array, ArrayDType, ArrayTrait, Canonical, IntoArray, IntoArrayVariant,
     IntoCanonical,

encodings/dict/src/compress.rs

@@ -9,6 +9,7 @@ use vortex_array::array::{
     ConstantArray, PrimitiveArray, SparseArray, VarBinArray, VarBinViewArray,
 };
 use vortex_array::validity::Validity;
+use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::{ArrayDType, IntoArray, IntoCanonical};
 use vortex_dtype::{match_each_native_ptype, DType, NativePType, ToBytes};
 use vortex_error::{VortexExpect as _, VortexUnwrap};

encodings/fastlanes/src/bitpacking/compress.rs

@@ -3,6 +3,7 @@ use fastlanes::BitPacking;
 use vortex_array::array::{PrimitiveArray, Sparse, SparseArray};
 use vortex_array::stats::ArrayStatistics;
 use vortex_array::validity::{ArrayValidity, Validity};
+use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::{Array, ArrayDType, ArrayDef, IntoArray, IntoArrayVariant};
 use vortex_buffer::Buffer;
 use vortex_dtype::{

encodings/fastlanes/src/bitpacking/compute/search_sorted.rs

@@ -10,6 +10,7 @@ use vortex_array::compute::{
     search_sorted_u64, IndexOrd, Len, SearchResult, SearchSorted, SearchSortedFn, SearchSortedSide,
 };
 use vortex_array::validity::Validity;
+use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::ArrayDType;
 use vortex_dtype::{match_each_unsigned_integer_ptype, NativePType};
 use vortex_error::{VortexError, VortexExpect as _, VortexResult};

encodings/fastlanes/src/bitpacking/compute/slice.rs

@@ -2,6 +2,7 @@ use std::cmp::max;
 
 use vortex_array::array::SparseArray;
 use vortex_array::compute::{slice, SliceFn};
+use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::{Array, IntoArray};
 use vortex_error::{VortexExpect, VortexResult};
 

encodings/fastlanes/src/bitpacking/compute/take.rs

@@ -4,6 +4,7 @@ use fastlanes::BitPacking;
 use itertools::Itertools;
 use vortex_array::array::{PrimitiveArray, SparseArray};
 use vortex_array::compute::{slice, take, TakeFn};
+use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::{Array, ArrayDType, IntoArray, IntoArrayVariant, IntoCanonical};
 use vortex_dtype::{
     match_each_integer_ptype, match_each_unsigned_integer_ptype, NativePType, PType,

encodings/fastlanes/src/bitpacking/mod.rs

@@ -14,9 +14,7 @@ use vortex_array::{
 };
 use vortex_buffer::Buffer;
 use vortex_dtype::{DType, NativePType, Nullability, PType};
-use vortex_error::{
-    vortex_bail, vortex_err, vortex_panic, VortexError, VortexExpect as _, VortexResult,
-};
+use vortex_error::{vortex_bail, vortex_err, VortexExpect as _, VortexResult};
 
 mod compress;
 mod compute;
@@ -197,17 +195,6 @@ impl BitPackedArray {
         }
     }
 
-    #[inline]
-    pub fn ptype(&self) -> PType {
-        self.dtype().try_into().unwrap_or_else(|err: VortexError| {
-            vortex_panic!(
-                err,
-                "Failed to convert BitpackedArray DType {} to PType",
-                self.dtype()
-            )
-        })
-    }
-
     #[inline]
     pub fn max_packed_value(&self) -> usize {
         (1 << self.bit_width()) - 1

encodings/fastlanes/src/delta/compress.rs

@@ -5,6 +5,7 @@ use vortex_array::array::PrimitiveArray;
 use vortex_array::compute::unary::fill_forward;
 use vortex_array::compute::SliceFn;
 use vortex_array::validity::Validity;
+use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::IntoArrayVariant;
 use vortex_dtype::{match_each_unsigned_integer_ptype, NativePType, Nullability};
 use vortex_error::VortexResult;

encodings/fastlanes/src/for/compress.rs

@@ -3,6 +3,7 @@ use num_traits::{PrimInt, WrappingAdd, WrappingSub};
 use vortex_array::array::{ConstantArray, PrimitiveArray, SparseArray};
 use vortex_array::stats::{trailing_zeros, ArrayStatistics, Stat};
 use vortex_array::validity::LogicalValidity;
+use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::{Array, ArrayDType, IntoArray, IntoArrayVariant};
 use vortex_dtype::{match_each_integer_ptype, NativePType};
 use vortex_error::{vortex_err, VortexResult};

encodings/fastlanes/src/for/compute.rs

@@ -6,6 +6,7 @@ use vortex_array::compute::{
     filter, search_sorted, slice, take, ArrayCompute, FilterFn, SearchResult, SearchSortedFn,
     SearchSortedSide, SliceFn, TakeFn,
 };
+use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::{Array, ArrayDType, IntoArray};
 use vortex_dtype::{match_each_integer_ptype, NativePType};
 use vortex_error::{VortexError, VortexExpect as _, VortexResult, VortexUnwrap as _};

encodings/fastlanes/src/for/mod.rs

@@ -8,8 +8,8 @@ use vortex_array::stats::{ArrayStatisticsCompute, StatsSet};
 use vortex_array::validity::{ArrayValidity, LogicalValidity};
 use vortex_array::variants::{ArrayVariants, PrimitiveArrayTrait};
 use vortex_array::{impl_encoding, Array, ArrayDType, ArrayTrait, Canonical, IntoCanonical};
-use vortex_dtype::{DType, PType};
-use vortex_error::{vortex_bail, vortex_panic, VortexExpect as _, VortexResult};
+use vortex_dtype::DType;
+use vortex_error::{vortex_bail, VortexExpect as _, VortexResult};
 use vortex_scalar::{Scalar, ScalarValue};
 
 mod compress;
@@ -79,17 +79,6 @@ impl FoRArray {
     pub fn shift(&self) -> u8 {
         self.metadata().shift
     }
-
-    #[inline]
-    pub fn ptype(&self) -> PType {
-        self.dtype().try_into().unwrap_or_else(|err| {
-            vortex_panic!(
-                err,
-                "Failed to convert FoRArray DType {} to PType",
-                self.dtype()
-            )
-        })
-    }
 }
 
 impl ArrayValidity for FoRArray {

encodings/roaring/src/integer/compress.rs

@@ -1,6 +1,7 @@
 use croaring::Bitmap;
 use num_traits::NumCast;
 use vortex_array::array::PrimitiveArray;
+use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_dtype::{NativePType, PType};
 use vortex_error::{vortex_bail, vortex_err, VortexResult};
 

encodings/roaring/src/integer/compute.rs

@@ -54,7 +54,7 @@ impl SliceFn for RoaringIntArray {
         };
 
         bitmap.and_inplace(&Bitmap::from_range(start..=stop_inclusive));
-        Self::try_new(bitmap, self.ptype()).map(IntoArray::into_array)
+        Self::try_new(bitmap, self.cached_ptype()).map(IntoArray::into_array)
     }
 }
 

encodings/roaring/src/integer/mod.rs

@@ -83,7 +83,7 @@ impl RoaringIntArray {
         )
     }
 
-    pub fn ptype(&self) -> PType {
+    pub fn cached_ptype(&self) -> PType {
         self.metadata().ptype
     }
 

encodings/runend-bool/src/compress.rs

@@ -5,6 +5,7 @@ use arrow_buffer::BooleanBufferBuilder;
 use num_traits::{AsPrimitive, FromPrimitive};
 use vortex_array::array::{BoolArray, PrimitiveArray};
 use vortex_array::validity::Validity;
+use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_dtype::{match_each_integer_ptype, NativePType};
 use vortex_error::{vortex_panic, VortexExpect as _, VortexResult};
 

encodings/runend-bool/src/compute.rs

@@ -1,6 +1,7 @@
 use vortex_array::array::BoolArray;
 use vortex_array::compute::unary::ScalarAtFn;
 use vortex_array::compute::{slice, ArrayCompute, SliceFn, TakeFn};
+use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::{Array, IntoArray, IntoArrayVariant, ToArray};
 use vortex_dtype::match_each_integer_ptype;
 use vortex_error::{vortex_bail, VortexExpect as _, VortexResult};

encodings/runend/src/compress.rs

@@ -6,6 +6,7 @@ use vortex_array::array::PrimitiveArray;
 use vortex_array::compute::unary::scalar_at;
 use vortex_array::stats::{ArrayStatistics, Stat};
 use vortex_array::validity::Validity;
+use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::ArrayDType;
 use vortex_dtype::{match_each_integer_ptype, match_each_native_ptype, NativePType, Nullability};
 use vortex_error::{vortex_panic, VortexResult, VortexUnwrap as _};

encodings/runend/src/compute.rs

@@ -2,6 +2,7 @@ use vortex_array::array::{ConstantArray, PrimitiveArray, SparseArray};
 use vortex_array::compute::unary::{scalar_at, scalar_at_unchecked, ScalarAtFn};
 use vortex_array::compute::{filter, slice, take, ArrayCompute, SliceFn, TakeFn};
 use vortex_array::validity::Validity;
+use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::{Array, ArrayDType, IntoArray, IntoArrayVariant};
 use vortex_dtype::match_each_integer_ptype;
 use vortex_error::{VortexExpect as _, VortexResult};

encodings/zigzag/src/array.rs

@@ -12,9 +12,7 @@ use vortex_array::{
     IntoCanonical,
 };
 use vortex_dtype::{DType, PType};
-use vortex_error::{
-    vortex_bail, vortex_err, vortex_panic, VortexExpect as _, VortexResult, VortexUnwrap as _,
-};
+use vortex_error::{vortex_bail, vortex_err, vortex_panic, VortexExpect as _, VortexResult};
 
 use crate::compress::zigzag_encode;
 use crate::zigzag_decode;
@@ -62,10 +60,6 @@ impl ZigZagArray {
             .child(0, &encoded, self.len())
             .vortex_expect("ZigZagArray is missing its encoded child array")
     }
-
-    pub fn ptype(&self) -> PType {
-        PType::try_from(self.dtype()).vortex_unwrap()
-    }
 }
 
 impl ArrayTrait for ZigZagArray {}

encodings/zigzag/src/compress.rs

@@ -1,5 +1,6 @@
 use vortex_array::array::PrimitiveArray;
 use vortex_array::validity::Validity;
+use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::IntoArray;
 use vortex_dtype::{NativePType, PType};
 use vortex_error::{vortex_bail, VortexResult};

encodings/zigzag/src/compute.rs

@@ -1,5 +1,6 @@
 use vortex_array::compute::unary::{scalar_at_unchecked, ScalarAtFn};
 use vortex_array::compute::{slice, ArrayCompute, SliceFn};
+use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::{Array, ArrayDType, IntoArray};
 use vortex_dtype::match_each_unsigned_integer_ptype;
 use vortex_error::{vortex_err, VortexResult, VortexUnwrap as _};

vortex-array/src/array/bool/compute/take.rs

@@ -5,6 +5,7 @@ use vortex_error::VortexResult;
 
 use crate::array::BoolArray;
 use crate::compute::TakeFn;
+use crate::variants::PrimitiveArrayTrait;
 use crate::{Array, IntoArray, IntoArrayVariant};
 
 impl TakeFn for BoolArray {

vortex-array/src/array/null/compute.rs

@@ -5,6 +5,7 @@ use vortex_scalar::Scalar;
 use crate::array::null::NullArray;
 use crate::compute::unary::ScalarAtFn;
 use crate::compute::{ArrayCompute, SliceFn, TakeFn};
+use crate::variants::PrimitiveArrayTrait;
 use crate::{Array, IntoArray, IntoArrayVariant};
 
 impl ArrayCompute for NullArray {

vortex-array/src/array/primitive/compute/cast.rs

@@ -4,6 +4,7 @@ use vortex_error::{vortex_bail, vortex_err, VortexResult};
 use crate::array::primitive::PrimitiveArray;
 use crate::compute::unary::CastFn;
 use crate::validity::Validity;
+use crate::variants::PrimitiveArrayTrait;
 use crate::{Array, ArrayDType, IntoArray};
 
 impl CastFn for PrimitiveArray {

vortex-array/src/array/primitive/compute/compare.rs

@@ -7,6 +7,7 @@ use vortex_scalar::PrimitiveScalar;
 use crate::array::primitive::PrimitiveArray;
 use crate::array::{BoolArray, ConstantArray};
 use crate::compute::{MaybeCompareFn, Operator};
+use crate::variants::PrimitiveArrayTrait;
 use crate::{Array, ArrayDType, IntoArray};
 
 impl MaybeCompareFn for PrimitiveArray {

vortex-array/src/array/primitive/compute/fill.rs

@@ -4,6 +4,7 @@ use vortex_error::{vortex_err, VortexResult};
 use crate::array::primitive::PrimitiveArray;
 use crate::compute::unary::FillForwardFn;
 use crate::validity::{ArrayValidity, Validity};
+use crate::variants::PrimitiveArrayTrait;
 use crate::{Array, ArrayDType, IntoArray};
 
 impl FillForwardFn for PrimitiveArray {

vortex-array/src/array/primitive/compute/filter.rs

@@ -3,7 +3,7 @@ use vortex_error::{vortex_err, VortexResult};
 
 use crate::array::primitive::PrimitiveArray;
 use crate::compute::FilterFn;
-use crate::variants::BoolArrayTrait;
+use crate::variants::{BoolArrayTrait, PrimitiveArrayTrait};
 use crate::{Array, IntoArray};
 
 impl FilterFn for PrimitiveArray {

vortex-array/src/array/primitive/compute/scalar_at.rs

@@ -4,6 +4,7 @@ use vortex_scalar::Scalar;
 
 use crate::array::primitive::PrimitiveArray;
 use crate::compute::unary::ScalarAtFn;
+use crate::variants::PrimitiveArrayTrait;
 use crate::ArrayDType;
 
 impl ScalarAtFn for PrimitiveArray {

vortex-array/src/array/primitive/compute/search_sorted.rs

@@ -8,6 +8,7 @@ use vortex_scalar::Scalar;
 use crate::array::primitive::PrimitiveArray;
 use crate::compute::{IndexOrd, Len, SearchResult, SearchSorted, SearchSortedFn, SearchSortedSide};
 use crate::validity::Validity;
+use crate::variants::PrimitiveArrayTrait;
 use crate::ArrayDType;
 
 impl SearchSortedFn for PrimitiveArray {

vortex-array/src/array/primitive/compute/slice.rs

@@ -2,6 +2,7 @@ use vortex_error::VortexResult;
 
 use crate::array::primitive::PrimitiveArray;
 use crate::compute::SliceFn;
+use crate::variants::PrimitiveArrayTrait;
 use crate::{Array, IntoArray};
 
 impl SliceFn for PrimitiveArray {

vortex-array/src/array/primitive/compute/subtract_scalar.rs

@@ -8,6 +8,7 @@ use crate::array::constant::ConstantArray;
 use crate::array::primitive::PrimitiveArray;
 use crate::compute::unary::SubtractScalarFn;
 use crate::validity::ArrayValidity;
+use crate::variants::PrimitiveArrayTrait;
 use crate::{Array, ArrayDType, IntoArray};
 
 impl SubtractScalarFn for PrimitiveArray {

vortex-array/src/array/primitive/compute/take.rs

@@ -4,6 +4,7 @@ use vortex_error::{vortex_panic, VortexResult};
 
 use crate::array::primitive::PrimitiveArray;
 use crate::compute::TakeFn;
+use crate::variants::PrimitiveArrayTrait;
 use crate::{Array, IntoArray, IntoArrayVariant};
 
 impl TakeFn for PrimitiveArray {

vortex-array/src/array/primitive/mod.rs

@@ -10,7 +10,7 @@ use num_traits::AsPrimitive;
 use serde::{Deserialize, Serialize};
 use vortex_buffer::Buffer;
 use vortex_dtype::{match_each_native_ptype, DType, NativePType, PType};
-use vortex_error::{vortex_bail, vortex_panic, VortexError, VortexExpect as _, VortexResult};
+use vortex_error::{vortex_bail, VortexExpect as _, VortexResult};
 
 use crate::array::visitor::{AcceptArrayVisitor, ArrayVisitor};
 use crate::elementwise::{dyn_cast_array_iter, BinaryFn, UnaryFn};
@@ -100,13 +100,6 @@ impl PrimitiveArray {
         })
     }
 
-    pub fn ptype(&self) -> PType {
-        // TODO(ngates): we can't really cache this anywhere?
-        self.dtype().try_into().unwrap_or_else(|err: VortexError| {
-            vortex_panic!(err, "Failed to convert dtype {} to ptype", self.dtype())
-        })
-    }
-
     pub fn buffer(&self) -> &Buffer {
         self.as_ref()
             .buffer()

vortex-array/src/array/primitive/stats.rs

@@ -12,6 +12,7 @@ use crate::aliases::hash_map::HashMap;
 use crate::array::primitive::PrimitiveArray;
 use crate::stats::{ArrayStatisticsCompute, Stat, StatsSet};
 use crate::validity::{ArrayValidity, LogicalValidity};
+use crate::variants::PrimitiveArrayTrait;
 use crate::{ArrayDType, IntoArrayVariant};
 
 trait PStatsType: NativePType + Into<Scalar> + BitWidth {}

vortex-array/src/array/sparse/canonical.rs

@@ -7,6 +7,7 @@ use crate::array::primitive::PrimitiveArray;
 use crate::array::sparse::SparseArray;
 use crate::array::BoolArray;
 use crate::validity::Validity;
+use crate::variants::PrimitiveArrayTrait;
 use crate::{ArrayDType, Canonical, IntoArrayVariant, IntoCanonical};
 
 impl IntoCanonical for SparseArray {