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