feat(compress): retain Narrow integer children through compression

encodings/fastlanes/src/rle/array/mod.rs

@@ -124,8 +124,8 @@ mod tests {
     use vortex_array::Canonical;
     use vortex_array::IntoArray;
     use vortex_array::VortexSessionExecute;
+    use vortex_array::arrays::NarrowArray;
     use vortex_array::arrays::PrimitiveArray;
-    use vortex_array::arrays::primitive::PrimitiveArrayExt;
     use vortex_array::assert_arrays_eq;
     use vortex_array::dtype::DType;
     use vortex_array::dtype::Nullability;
@@ -457,11 +457,10 @@ mod tests {
 
         // Simulate cascading compression: narrow u16->u8 then re-encode with RLE,
         // matching the path taken by the BtrBlocks compressor.
-        let indices_prim = rle
-            .indices()
-            .clone()
-            .execute::<PrimitiveArray>(&mut ctx)?
-            .narrow(&mut ctx)?;
+        let indices_prim = NarrowArray::encode_values(
+            rle.indices().clone().execute::<PrimitiveArray>(&mut ctx)?,
+            &mut ctx,
+        )?;
         let re_encoded = RLEData::encode(indices_prim.as_view(), &mut ctx)?;
 
         // Reconstruct the outer RLE with re-encoded indices.

encodings/runend/src/compress.rs

@@ -9,6 +9,7 @@ use vortex_array::IntoArray;
 use vortex_array::arrays::BoolArray;
 use vortex_array::arrays::ConstantArray;
 use vortex_array::arrays::DecimalArray;
+use vortex_array::arrays::NarrowArray;
 use vortex_array::arrays::Primitive;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::arrays::VarBinViewArray;
@@ -87,9 +88,7 @@ pub fn runend_encode(
         }
     };
 
-    let ends = ends
-        .narrow(ctx)
-        .vortex_expect("Ends must succeed downcasting");
+    let ends = NarrowArray::encode_values(ends, ctx).vortex_expect("Ends must succeed downcasting");
 
     ends.statistics()
         .set(Stat::IsStrictSorted, Precision::Exact(true.into()));

vortex-array/src/arrays/narrow/encoding.rs

@@ -33,32 +33,62 @@ impl NarrowArray {
             "Narrow requires integer values, got {ptype}"
         );
 
-        if ptype.byte_width() == 1 {
-            return Ok(array.into_array());
+        let dtype = array.dtype().clone();
+        let values = Self::encode_values(array, ctx)?;
+        if values.dtype() == &dtype {
+            return Ok(values.into_array());
         }
 
-        let bounds = min_max(array.as_ref(), ctx, NumericalAggregateOpts::default())?;
-        for candidate in integer_types(ptype) {
-            if candidate.byte_width() >= ptype.byte_width() {
-                break;
-            }
+        Ok(Self::try_new(values.into_array(), dtype)?.into_array())
+    }
 
-            let storage_dtype = DType::Primitive(candidate, array.dtype().nullability());
-            if bounds.as_ref().is_some_and(|bounds| {
-                bounds.min.cast(&storage_dtype).is_err() || bounds.max.cast(&storage_dtype).is_err()
-            }) {
-                continue;
-            }
+    /// Selects the stored primitive width for internal buffers whose dtype is chosen at construction.
+    ///
+    /// This changes the returned dtype while preserving signedness and validity. For arrays with
+    /// an existing logical dtype contract, use [`Self::encode`] to retain that dtype. Floating-point
+    /// arrays are returned unchanged. Empty and all-null integers use the smallest width of the
+    /// same signedness.
+    pub fn encode_values(
+        array: PrimitiveArray,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<PrimitiveArray> {
+        let ptype = array.ptype();
+        if !ptype.is_int() || ptype.byte_width() == 1 {
+            return Ok(array);
+        }
 
-            let values = array
-                .as_ref()
-                .cast(storage_dtype)?
-                .execute::<PrimitiveArray>(ctx)?;
+        let bounds = min_max(array.as_ref(), ctx, NumericalAggregateOpts::default())?;
+        let storage_type = if let Some(bounds) = bounds {
+            let (min_type, max_type) = if ptype.is_signed_int() {
+                (
+                    PType::min_signed_ptype_for_value(i64::try_from(&bounds.min)?),
+                    PType::min_signed_ptype_for_value(i64::try_from(&bounds.max)?),
+                )
+            } else {
+                (
+                    PType::min_unsigned_ptype_for_value(u64::try_from(&bounds.min)?),
+                    PType::min_unsigned_ptype_for_value(u64::try_from(&bounds.max)?),
+                )
+            };
+            if min_type.byte_width() >= max_type.byte_width() {
+                min_type
+            } else {
+                max_type
+            }
+        } else if ptype.is_signed_int() {
+            PType::I8
+        } else {
+            PType::U8
+        };
 
-            return Ok(Self::try_new(values.into_array(), array.dtype().clone())?.into_array());
+        if storage_type.byte_width() >= ptype.byte_width() {
+            return Ok(array);
         }
 
-        Ok(array.into_array())
+        array
+            .as_ref()
+            .cast(DType::Primitive(storage_type, array.dtype().nullability()))?
+            .execute(ctx)
     }
 }
 

vortex-array/src/arrays/narrow/tests/mod.rs

@@ -467,3 +467,35 @@ fn test_reject_malformed_metadata(#[case] metadata: Vec<u8>) {
             .is_err()
     );
 }
+
+#[test]
+fn test_encode_values_changes_only_internal_dtype() -> VortexResult<()> {
+    let mut ctx = array_session().create_execution_ctx();
+    let input = PrimitiveArray::from_option_iter([Some(0i64), None, Some(127)]);
+    let values = NarrowArray::encode_values(input.clone(), &mut ctx)?;
+    let wrapped = NarrowArray::encode(input.clone(), &mut ctx)?;
+    assert_eq!(
+        values.dtype(),
+        &DType::Primitive(PType::I8, Nullability::Nullable)
+    );
+    assert_eq!(wrapped.dtype(), input.dtype());
+    assert_arrays_eq!(
+        values,
+        PrimitiveArray::from_option_iter([Some(0i8), None, Some(127)]),
+        &mut ctx
+    );
+    assert_arrays_eq!(wrapped, input, &mut ctx);
+
+    Ok(())
+}
+
+#[test]
+fn test_encode_values_preserves_floats() -> VortexResult<()> {
+    let mut ctx = array_session().create_execution_ctx();
+    let input = PrimitiveArray::from_iter([1.0f32, 2.0, 3.0]);
+    let values = NarrowArray::encode_values(input.clone(), &mut ctx)?;
+    assert_eq!(values.dtype(), input.dtype());
+    assert_arrays_eq!(values, input, &mut ctx);
+
+    Ok(())
+}

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

@@ -34,122 +34,3 @@ impl PrimitiveData {
         unsafe { std::slice::from_raw_parts(raw_slice.cast(), byte_buffer.len() / size_of::<T>()) }
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use std::sync::LazyLock;
-
-    use rstest::rstest;
-    use vortex_buffer::Buffer;
-    use vortex_buffer::buffer;
-    use vortex_session::VortexSession;
-
-    use crate::VortexSessionExecute;
-    use crate::arrays::PrimitiveArray;
-    use crate::arrays::primitive::PrimitiveArrayExt;
-    use crate::dtype::DType;
-    use crate::dtype::Nullability;
-    use crate::dtype::PType;
-    use crate::validity::Validity;
-
-    static SESSION: LazyLock<VortexSession> = LazyLock::new(crate::array_session);
-
-    #[test]
-    fn test_downcast_all_invalid() {
-        let array = PrimitiveArray::new(
-            buffer![0_u32, 0, 0, 0, 0, 0, 0, 0, 0, 0],
-            Validity::AllInvalid,
-        );
-
-        let result = array.narrow(&mut SESSION.create_execution_ctx()).unwrap();
-        assert_eq!(
-            result.dtype(),
-            &DType::Primitive(PType::U8, Nullability::Nullable)
-        );
-        assert!(matches!(result.validity(), Ok(Validity::AllInvalid)));
-    }
-
-    #[rstest]
-    #[case(vec![0_i64, 127], PType::U8)]
-    #[case(vec![-128_i64, 127], PType::I8)]
-    #[case(vec![-129_i64, 127], PType::I16)]
-    #[case(vec![-128_i64, 128], PType::I16)]
-    #[case(vec![-32768_i64, 32767], PType::I16)]
-    #[case(vec![-32769_i64, 32767], PType::I32)]
-    #[case(vec![-32768_i64, 32768], PType::I32)]
-    #[case(vec![i32::MIN as i64, i32::MAX as i64], PType::I32)]
-    fn test_downcast_signed(#[case] values: Vec<i64>, #[case] expected_ptype: PType) {
-        let array = PrimitiveArray::from_iter(values);
-        let result = array.narrow(&mut SESSION.create_execution_ctx()).unwrap();
-        assert_eq!(result.ptype(), expected_ptype);
-    }
-
-    #[rstest]
-    #[case(vec![0_u64, 255], PType::U8)]
-    #[case(vec![0_u64, 256], PType::U16)]
-    #[case(vec![0_u64, 65535], PType::U16)]
-    #[case(vec![0_u64, 65536], PType::U32)]
-    #[case(vec![0_u64, u32::MAX as u64], PType::U32)]
-    fn test_downcast_unsigned(#[case] values: Vec<u64>, #[case] expected_ptype: PType) {
-        let array = PrimitiveArray::from_iter(values);
-        let result = array.narrow(&mut SESSION.create_execution_ctx()).unwrap();
-        assert_eq!(result.ptype(), expected_ptype);
-    }
-
-    #[test]
-    fn test_downcast_keeps_original_if_too_large() {
-        let array = PrimitiveArray::from_iter(vec![0_u64, u64::MAX]);
-        let result = array.narrow(&mut SESSION.create_execution_ctx()).unwrap();
-        assert_eq!(result.ptype(), PType::U64);
-    }
-
-    #[test]
-    fn test_downcast_preserves_nullability() {
-        let array = PrimitiveArray::from_option_iter([Some(0_i32), None, Some(127)]);
-        let result = array.narrow(&mut SESSION.create_execution_ctx()).unwrap();
-        assert_eq!(
-            result.dtype(),
-            &DType::Primitive(PType::U8, Nullability::Nullable)
-        );
-        // Check that validity is preserved (the array should still have nullable values)
-        assert!(matches!(result.validity(), Ok(Validity::Array(_))));
-    }
-
-    #[test]
-    fn test_downcast_preserves_values() {
-        let values = vec![-100_i16, 0, 100];
-        let array = PrimitiveArray::from_iter(values);
-        let result = array.narrow(&mut SESSION.create_execution_ctx()).unwrap();
-
-        assert_eq!(result.ptype(), PType::I8);
-        // Check that the values were properly downscaled
-        let downscaled_values: Vec<i8> = result.as_slice::<i8>().to_vec();
-        assert_eq!(downscaled_values, vec![-100_i8, 0, 100]);
-    }
-
-    #[test]
-    fn test_downcast_with_mixed_signs_chooses_signed() {
-        let array = PrimitiveArray::from_iter(vec![-1_i32, 200]);
-        let result = array.narrow(&mut SESSION.create_execution_ctx()).unwrap();
-        assert_eq!(result.ptype(), PType::I16);
-    }
-
-    #[test]
-    fn test_downcast_floats() {
-        let array = PrimitiveArray::from_iter(vec![1.0_f32, 2.0, 3.0]);
-        let result = array.narrow(&mut SESSION.create_execution_ctx()).unwrap();
-        // Floats should remain unchanged since they can't be downscaled to integers
-        assert_eq!(result.ptype(), PType::F32);
-    }
-
-    #[test]
-    fn test_downcast_empty_array() {
-        let array = PrimitiveArray::new(Buffer::<i32>::empty(), Validity::AllInvalid);
-        let result = array.narrow(&mut SESSION.create_execution_ctx()).unwrap();
-        let array2 = PrimitiveArray::new(Buffer::<i64>::empty(), Validity::NonNullable);
-        let result2 = array2.narrow(&mut SESSION.create_execution_ctx()).unwrap();
-        // Empty arrays should not have their validity changed
-        assert!(matches!(result.validity(), Ok(Validity::AllInvalid)));
-        assert!(matches!(result2.validity(), Ok(Validity::NonNullable)));
-    }
-}

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

@@ -39,14 +39,11 @@ pub use patch::chunk_range;
 pub use patch::patch_chunk;
 
 use crate::ArrayRef;
-use crate::aggregate_fn::NumericalAggregateOpts;
-use crate::aggregate_fn::fns::min_max::min_max;
 use crate::array::child_to_validity;
 use crate::array::validity_to_child;
 use crate::array_slots;
 use crate::arrays::bool::BoolArrayExt;
 use crate::buffer::BufferHandle;
-use crate::builtins::ArrayBuiltins;
 
 #[array_slots(Primitive)]
 pub struct PrimitiveSlots {
@@ -152,93 +149,6 @@ pub trait PrimitiveArrayExt: TypedArrayRef<Primitive> {
             PrimitiveArrayExt::validity(self),
         )
     }
-
-    /// Narrow the array to the smallest possible integer type that can represent all values.
-    fn narrow(&self, ctx: &mut ExecutionCtx) -> VortexResult<PrimitiveArray> {
-        if !self.ptype().is_int() {
-            return Ok(self.to_owned());
-        }
-
-        let Some(min_max) = min_max(self.as_ref(), ctx, NumericalAggregateOpts::default())? else {
-            return Ok(PrimitiveArray::new(
-                Buffer::<u8>::zeroed(self.len()),
-                PrimitiveArrayExt::validity(self),
-            ));
-        };
-
-        // If we can't cast to i64, then leave the array as its original type.
-        // It's too big to downcast anyway.
-        let Ok(min) = min_max
-            .min
-            .cast(&PType::I64.into())
-            .and_then(|s| i64::try_from(&s))
-        else {
-            return Ok(self.to_owned());
-        };
-        let Ok(max) = min_max
-            .max
-            .cast(&PType::I64.into())
-            .and_then(|s| i64::try_from(&s))
-        else {
-            return Ok(self.to_owned());
-        };
-
-        let nullability = self.as_ref().dtype().nullability();
-
-        if min < 0 || max < 0 {
-            // Signed
-            if min >= i8::MIN as i64 && max <= i8::MAX as i64 {
-                let result = self
-                    .as_ref()
-                    .cast(DType::Primitive(PType::I8, nullability))?
-                    .execute::<PrimitiveArray>(ctx)?;
-                return Ok(result);
-            }
-
-            if min >= i16::MIN as i64 && max <= i16::MAX as i64 {
-                let result = self
-                    .as_ref()
-                    .cast(DType::Primitive(PType::I16, nullability))?
-                    .execute::<PrimitiveArray>(ctx)?;
-                return Ok(result);
-            }
-
-            if min >= i32::MIN as i64 && max <= i32::MAX as i64 {
-                let result = self
-                    .as_ref()
-                    .cast(DType::Primitive(PType::I32, nullability))?
-                    .execute::<PrimitiveArray>(ctx)?;
-                return Ok(result);
-            }
-        } else {
-            // Unsigned
-            if max <= u8::MAX as i64 {
-                let result = self
-                    .as_ref()
-                    .cast(DType::Primitive(PType::U8, nullability))?
-                    .execute::<PrimitiveArray>(ctx)?;
-                return Ok(result);
-            }
-
-            if max <= u16::MAX as i64 {
-                let result = self
-                    .as_ref()
-                    .cast(DType::Primitive(PType::U16, nullability))?
-                    .execute::<PrimitiveArray>(ctx)?;
-                return Ok(result);
-            }
-
-            if max <= u32::MAX as i64 {
-                let result = self
-                    .as_ref()
-                    .cast(DType::Primitive(PType::U32, nullability))?
-                    .execute::<PrimitiveArray>(ctx)?;
-                return Ok(result);
-            }
-        }
-
-        Ok(self.to_owned())
-    }
 }
 impl<T: TypedArrayRef<Primitive>> PrimitiveArrayExt for T {}