Skip to main content

vortex_btrblocks/schemes/integer/
pco.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Pco (pcodec) integer compression.
5
6use vortex_array::ArrayRef;
7use vortex_array::Canonical;
8use vortex_array::ExecutionCtx;
9use vortex_array::IntoArray;
10use vortex_compressor::estimate::CompressionEstimate;
11use vortex_compressor::estimate::DeferredEstimate;
12use vortex_compressor::estimate::EstimateVerdict;
13use vortex_error::VortexResult;
14
15use crate::ArrayAndStats;
16use crate::CascadingCompressor;
17use crate::CompressorContext;
18use crate::Scheme;
19
20/// Pco (pcodec) compression for integers.
21#[derive(Debug, Copy, Clone, PartialEq, Eq)]
22pub struct PcoScheme;
23
24impl Scheme for PcoScheme {
25    fn scheme_name(&self) -> &'static str {
26        "vortex.int.pco"
27    }
28
29    fn matches(&self, canonical: &Canonical) -> bool {
30        canonical.dtype().is_int()
31    }
32
33    fn expected_compression_ratio(
34        &self,
35        data: &ArrayAndStats,
36        _compress_ctx: CompressorContext,
37        _exec_ctx: &mut ExecutionCtx,
38    ) -> CompressionEstimate {
39        use vortex_array::dtype::PType;
40
41        // Pco does not support I8 or U8.
42        if matches!(data.array_as_primitive().ptype(), PType::I8 | PType::U8) {
43            return CompressionEstimate::Verdict(EstimateVerdict::Skip);
44        }
45
46        CompressionEstimate::Deferred(DeferredEstimate::Sample)
47    }
48
49    fn compress(
50        &self,
51        _compressor: &CascadingCompressor,
52        data: &ArrayAndStats,
53        _compress_ctx: CompressorContext,
54        exec_ctx: &mut ExecutionCtx,
55    ) -> VortexResult<ArrayRef> {
56        Ok(vortex_pco::Pco::from_primitive(
57            data.array_as_primitive(),
58            pco::DEFAULT_COMPRESSION_LEVEL,
59            8192,
60            exec_ctx,
61        )?
62        .into_array())
63    }
64}