vortex_btrblocks/schemes/float/
alp.rs1use vortex_alp::ALP;
7use vortex_alp::ALPArrayExt;
8use vortex_alp::ALPArraySlotsExt;
9use vortex_alp::alp_encode;
10use vortex_array::ArrayRef;
11use vortex_array::Canonical;
12use vortex_array::ExecutionCtx;
13use vortex_array::IntoArray;
14use vortex_array::arrays::Patched;
15use vortex_array::arrays::patched::use_experimental_patches;
16use vortex_array::arrays::primitive::PrimitiveArrayExt;
17use vortex_array::dtype::PType;
18use vortex_compressor::estimate::CompressionEstimate;
19use vortex_compressor::estimate::DeferredEstimate;
20use vortex_compressor::estimate::EstimateVerdict;
21use vortex_error::VortexResult;
22
23use crate::ArrayAndStats;
24use crate::CascadingCompressor;
25use crate::CompressorContext;
26use crate::Scheme;
27use crate::SchemeExt;
28use crate::compress_patches;
29
30#[derive(Debug, Copy, Clone, PartialEq, Eq)]
32pub struct ALPScheme;
33
34impl Scheme for ALPScheme {
35 fn scheme_name(&self) -> &'static str {
36 "vortex.float.alp"
37 }
38
39 fn matches(&self, canonical: &Canonical) -> bool {
40 canonical.dtype().is_float()
41 }
42
43 fn num_children(&self) -> usize {
45 1
46 }
47
48 fn expected_compression_ratio(
49 &self,
50 data: &ArrayAndStats,
51 compress_ctx: CompressorContext,
52 _exec_ctx: &mut ExecutionCtx,
53 ) -> CompressionEstimate {
54 if compress_ctx.finished_cascading() {
57 return CompressionEstimate::Verdict(EstimateVerdict::Skip);
58 }
59
60 if data.array_as_primitive().ptype() == PType::F16 {
62 return CompressionEstimate::Verdict(EstimateVerdict::Skip);
63 }
64
65 CompressionEstimate::Deferred(DeferredEstimate::Sample)
66 }
67
68 fn compress(
69 &self,
70 compressor: &CascadingCompressor,
71 data: &ArrayAndStats,
72 compress_ctx: CompressorContext,
73 exec_ctx: &mut ExecutionCtx,
74 ) -> VortexResult<ArrayRef> {
75 let alp_encoded = alp_encode(data.array_as_primitive(), None, exec_ctx)?;
76
77 let compressed_alp_ints = compressor.compress_child(
79 alp_encoded.encoded(),
80 &compress_ctx,
81 self.id(),
82 0,
83 exec_ctx,
84 )?;
85
86 let alp_stats = alp_encoded.as_array().statistics().to_owned();
87 let exponents = alp_encoded.exponents();
88
89 if use_experimental_patches() {
90 let patches = alp_encoded.patches();
91
92 let alp_array = ALP::new(compressed_alp_ints, exponents, None).into_array();
94
95 match patches {
96 None => Ok(alp_array),
97 Some(p) => Ok(Patched::from_array_and_patches(alp_array, &p, exec_ctx)?
98 .with_stats_set(alp_stats)
99 .into_array()),
100 }
101 } else {
102 let patches = alp_encoded
103 .patches()
104 .map(|p| compress_patches(p, exec_ctx))
105 .transpose()?;
106
107 Ok(ALP::new(compressed_alp_ints, exponents, patches).into_array())
108 }
109 }
110}