vortex_sampling_compressor/compressors/
alp_rd.rs

1use std::any::Any;
2use std::sync::Arc;
3
4use vortex_alp::{match_each_alp_float_ptype, ALPRDEncoding, RDEncoder as ALPRDEncoder};
5use vortex_array::aliases::hash_set::HashSet;
6use vortex_array::array::PrimitiveArray;
7use vortex_array::variants::PrimitiveArrayTrait;
8use vortex_array::{Array, Encoding, EncodingId, IntoArray, IntoArrayVariant};
9use vortex_dtype::PType;
10use vortex_error::{vortex_bail, VortexResult};
11use vortex_fastlanes::BitPackedEncoding;
12
13use crate::compressors::{CompressedArray, CompressionTree, EncoderMetadata, EncodingCompressor};
14use crate::{constants, SamplingCompressor};
15
16#[derive(Debug)]
17pub struct ALPRDCompressor;
18
19impl EncoderMetadata for ALPRDEncoder {
20    fn as_any(&self) -> &dyn Any {
21        self
22    }
23}
24
25impl EncodingCompressor for ALPRDCompressor {
26    fn id(&self) -> &str {
27        ALPRDEncoding::ID.as_ref()
28    }
29
30    fn cost(&self) -> u8 {
31        constants::ALP_RD_COST
32    }
33
34    fn can_compress(&self, array: &Array) -> Option<&dyn EncodingCompressor> {
35        // Only support primitive arrays
36        let parray = PrimitiveArray::maybe_from(array)?;
37
38        // Only supports f32 and f64
39        if !matches!(parray.ptype(), PType::F32 | PType::F64) {
40            return None;
41        }
42
43        Some(self)
44    }
45
46    fn compress<'a>(
47        &'a self,
48        array: &Array,
49        like: Option<CompressionTree<'a>>,
50        _ctx: SamplingCompressor<'a>,
51    ) -> VortexResult<CompressedArray<'a>> {
52        let primitive = array.clone().into_primitive()?;
53
54        // Train a new compressor or reuse an existing compressor.
55        let encoder = like
56            .clone()
57            .and_then(|mut tree| tree.metadata())
58            .map(VortexResult::Ok)
59            .unwrap_or_else(|| Ok(Arc::new(alp_rd_new_encoder(&primitive))))?;
60
61        let Some(alp_rd_encoder) = encoder.as_any().downcast_ref::<ALPRDEncoder>() else {
62            vortex_bail!("Could not downcast metadata as ALPRDEncoder");
63        };
64
65        let encoded = alp_rd_encoder.encode(&primitive).into_array();
66        Ok(CompressedArray::compressed(
67            encoded,
68            Some(CompressionTree::new_with_metadata(self, vec![], encoder)),
69            array,
70        ))
71    }
72
73    fn used_encodings(&self) -> HashSet<EncodingId> {
74        HashSet::from([ALPRDEncoding::ID, BitPackedEncoding::ID])
75    }
76}
77
78/// Create a new `ALPRDEncoder` from the given array of samples.
79fn alp_rd_new_encoder(array: &PrimitiveArray) -> ALPRDEncoder {
80    match_each_alp_float_ptype!(array.ptype(), |$P| {
81        ALPRDEncoder::new(array.as_slice::<$P>())
82    })
83}
vortex_sampling_compressor/compressors/alp_rd.rs

vortex_sampling_compressor/compressors/
alp_rd.rs
