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