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::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/// ALP (Adaptive Lossless floating-Point) encoding.
33#[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    /// Children: encoded_ints=0.
54    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        // ALP encodes floats as integers. Without integer compression afterward, the encoded ints
65        // are the same size.
66        if compress_ctx.finished_cascading() {
67            return CompressionEstimate::Verdict(EstimateVerdict::Skip);
68        }
69
70        // We don't support ALP for f16.
71        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        // Compress the ALP ints.
88        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            // Create ALP array without interior patches.
103            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}