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