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