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