Skip to main content

vortex_btrblocks/schemes/float/
alprd.rs

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