vortex_array/
compress.rs

1// TODO(ngates): make this a function on a PrimitiveArray
2use vortex_dtype::{DType, PType};
3use vortex_error::{VortexExpect, VortexResult};
4use vortex_scalar::Scalar;
5
6use crate::arrays::{ConstantArray, PrimitiveArray, PrimitiveEncoding};
7use crate::compute::{min_max, try_cast};
8use crate::vtable::EncodingVTable;
9use crate::{Array, ArrayExt, ArrayRef, ToCanonical};
10
11/// Downscale a primitive array to the narrowest PType that fits all the values.
12pub fn downscale_integer_array(array: ArrayRef) -> VortexResult<ArrayRef> {
13    if !array.is_encoding(PrimitiveEncoding.id()) {
14        // This can happen if e.g. the array is ConstantArray.
15        return Ok(array);
16    }
17    if array.is_empty() {
18        return Ok(array);
19    }
20    let array = array
21        .as_opt::<PrimitiveArray>()
22        .vortex_expect("Checked earlier");
23
24    let Some(min_max) = min_max(array)? else {
25        // This array but be all nulls.
26        return Ok(
27            ConstantArray::new(Scalar::null(array.dtype().clone()), array.len()).into_array(),
28        );
29    };
30
31    // If we can't cast to i64, then leave the array as its original type.
32    // It's too big to downcast anyway.
33    let Ok(min) = i64::try_from(min_max.min.value()) else {
34        return Ok(array.to_array());
35    };
36    let Ok(max) = i64::try_from(min_max.max.value()) else {
37        return Ok(array.to_array());
38    };
39
40    downscale_primitive_integer_array(array.clone(), min, max).map(|a| a.into_array())
41}
42
43/// Downscale a primitive array to the narrowest PType that fits all the values.
44fn downscale_primitive_integer_array(
45    array: PrimitiveArray,
46    min: i64,
47    max: i64,
48) -> VortexResult<PrimitiveArray> {
49    if min < 0 || max < 0 {
50        // Signed
51        if min >= i8::MIN as i64 && max <= i8::MAX as i64 {
52            return try_cast(
53                &array,
54                &DType::Primitive(PType::I8, array.dtype().nullability()),
55            )?
56            .to_primitive();
57        }
58
59        if min >= i16::MIN as i64 && max <= i16::MAX as i64 {
60            return try_cast(
61                &array,
62                &DType::Primitive(PType::I16, array.dtype().nullability()),
63            )?
64            .to_primitive();
65        }
66
67        if min >= i32::MIN as i64 && max <= i32::MAX as i64 {
68            return try_cast(
69                &array,
70                &DType::Primitive(PType::I32, array.dtype().nullability()),
71            )?
72            .to_primitive();
73        }
74    } else {
75        // Unsigned
76        if max <= u8::MAX as i64 {
77            return try_cast(
78                &array,
79                &DType::Primitive(PType::U8, array.dtype().nullability()),
80            )?
81            .to_primitive();
82        }
83
84        if max <= u16::MAX as i64 {
85            return try_cast(
86                &array,
87                &DType::Primitive(PType::U16, array.dtype().nullability()),
88            )?
89            .to_primitive();
90        }
91
92        if max <= u32::MAX as i64 {
93            return try_cast(
94                &array,
95                &DType::Primitive(PType::U32, array.dtype().nullability()),
96            )?
97            .to_primitive();
98        }
99    }
100
101    Ok(array)
102}
vortex_array/compress.rs

vortex_array/
compress.rs
