Skip to main content

vortex_btrblocks/schemes/float/
alp.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! ALP (Adaptive Lossless floating-Point) encoding.
5
6use 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/// ALP (Adaptive Lossless floating-Point) encoding.
31#[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    /// Children: encoded_ints=0.
44    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        // ALP encodes floats as integers. Without integer compression afterward, the encoded ints
55        // are the same size.
56        if compress_ctx.finished_cascading() {
57            return CompressionEstimate::Verdict(EstimateVerdict::Skip);
58        }
59
60        // We don't support ALP for f16.
61        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        // Compress the ALP ints.
78        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            // Create ALP array without interior patches.
93            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}