vortex_btrblocks/schemes/float/
alprd.rs1use vortex_alp::ALPRDArrayExt;
7use vortex_alp::ALPRDArrayOwnedExt;
8use vortex_alp::RDEncoder;
9use vortex_alp::RDEncoderExt;
10use vortex_array::ArrayId;
11use vortex_array::ArrayRef;
12use vortex_array::Canonical;
13use vortex_array::ExecutionCtx;
14use vortex_array::IntoArray;
15use vortex_array::VTable;
16use vortex_array::arrays::primitive::PrimitiveArrayExt;
17use vortex_array::dtype::PType;
18use vortex_compressor::scheme::CompressionEstimate;
19use vortex_compressor::scheme::DeferredEstimate;
20use vortex_compressor::scheme::EstimateVerdict;
21use vortex_error::VortexResult;
22use vortex_error::vortex_panic;
23
24use crate::ArrayAndStats;
25use crate::CascadingCompressor;
26use crate::CompressorContext;
27use crate::Scheme;
28use crate::compress_patches;
29
30#[derive(Debug, Copy, Clone, PartialEq, Eq)]
32pub struct ALPRDScheme;
33
34impl Scheme for ALPRDScheme {
35 fn scheme_name(&self) -> &'static str {
36 "vortex.float.alprd"
37 }
38
39 fn matches(&self, canonical: &Canonical) -> bool {
40 canonical.dtype().is_float()
41 }
42
43 fn produced_encodings(&self) -> Vec<ArrayId> {
44 vec![vortex_alp::ALPRD.id()]
45 }
46
47 fn expected_compression_ratio(
48 &self,
49 data: &ArrayAndStats,
50 _compress_ctx: CompressorContext,
51 _exec_ctx: &mut ExecutionCtx,
52 ) -> CompressionEstimate {
53 if data.array_as_primitive().ptype() == PType::F16 {
55 return CompressionEstimate::Verdict(EstimateVerdict::Skip);
56 }
57
58 CompressionEstimate::Deferred(DeferredEstimate::Sample)
59 }
60
61 fn compress(
62 &self,
63 _compressor: &CascadingCompressor,
64 data: &ArrayAndStats,
65 _compress_ctx: CompressorContext,
66 exec_ctx: &mut ExecutionCtx,
67 ) -> VortexResult<ArrayRef> {
68 let primitive_array = data.array_as_primitive();
69
70 let encoder = match primitive_array.ptype() {
71 PType::F32 => RDEncoder::new(primitive_array.as_slice::<f32>()),
72 PType::F64 => RDEncoder::new(primitive_array.as_slice::<f64>()),
73 ptype => vortex_panic!("cannot ALPRD compress ptype {ptype}"),
74 };
75
76 let alp_rd = encoder.encode(primitive_array);
77 let dtype = alp_rd.dtype().clone();
78 let right_bit_width = alp_rd.right_bit_width();
79 let mut parts = ALPRDArrayOwnedExt::into_data_parts(alp_rd);
80 parts.left_parts_patches = parts
81 .left_parts_patches
82 .map(|p| compress_patches(p, exec_ctx))
83 .transpose()?;
84
85 Ok(vortex_alp::ALPRD::try_new(
86 dtype,
87 parts.left_parts,
88 parts.left_parts_dictionary,
89 parts.right_parts,
90 right_bit_width,
91 parts.left_parts_patches,
92 )?
93 .into_array())
94 }
95}