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::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::Patched;
17use vortex_array::arrays::patched::use_experimental_patches;
18use vortex_array::arrays::primitive::PrimitiveArrayExt;
19use vortex_array::dtype::PType;
20use vortex_compressor::scheme::CompressionEstimate;
21use vortex_compressor::scheme::DeferredEstimate;
22use vortex_compressor::scheme::EstimateVerdict;
23use vortex_error::VortexResult;
24
25use crate::ArrayAndStats;
26use crate::CascadingCompressor;
27use crate::CompressorContext;
28use crate::Scheme;
29use crate::SchemeExt;
30use crate::compress_patches;
31
32#[derive(Debug, Copy, Clone, PartialEq, Eq)]
34pub struct ALPScheme;
35
36impl Scheme for ALPScheme {
37 fn scheme_name(&self) -> &'static str {
38 "vortex.float.alp"
39 }
40
41 fn matches(&self, canonical: &Canonical) -> bool {
42 canonical.dtype().is_float()
43 }
44
45 fn produced_encodings(&self) -> Vec<ArrayId> {
46 let mut encodings = vec![ALP.id()];
47 if use_experimental_patches() {
48 encodings.push(Patched.id());
49 }
50 encodings
51 }
52
53 fn num_children(&self) -> usize {
55 1
56 }
57
58 fn expected_compression_ratio(
59 &self,
60 data: &ArrayAndStats,
61 compress_ctx: CompressorContext,
62 _exec_ctx: &mut ExecutionCtx,
63 ) -> CompressionEstimate {
64 if compress_ctx.finished_cascading() {
67 return CompressionEstimate::Verdict(EstimateVerdict::Skip);
68 }
69
70 if data.array_as_primitive().ptype() == PType::F16 {
72 return CompressionEstimate::Verdict(EstimateVerdict::Skip);
73 }
74
75 CompressionEstimate::Deferred(DeferredEstimate::Sample)
76 }
77
78 fn compress(
79 &self,
80 compressor: &CascadingCompressor,
81 data: &ArrayAndStats,
82 compress_ctx: CompressorContext,
83 exec_ctx: &mut ExecutionCtx,
84 ) -> VortexResult<ArrayRef> {
85 let alp_encoded = alp_encode(data.array_as_primitive(), None, exec_ctx)?;
86
87 let compressed_alp_ints = compressor.compress_child(
89 alp_encoded.encoded(),
90 &compress_ctx,
91 self.id(),
92 0,
93 exec_ctx,
94 )?;
95
96 let alp_stats = alp_encoded.as_array().statistics().to_owned();
97 let exponents = alp_encoded.exponents();
98
99 if use_experimental_patches() {
100 let patches = alp_encoded.patches();
101
102 let alp_array = ALP::new(compressed_alp_ints, exponents, None).into_array();
104
105 match patches {
106 None => Ok(alp_array),
107 Some(p) => Ok(Patched::from_array_and_patches(alp_array, &p, exec_ctx)?
108 .with_stats_set(alp_stats)
109 .into_array()),
110 }
111 } else {
112 let patches = alp_encoded
113 .patches()
114 .map(|p| compress_patches(p, exec_ctx))
115 .transpose()?;
116
117 Ok(ALP::new(compressed_alp_ints, exponents, patches).into_array())
118 }
119 }
120}