vortex_layout/layouts/
compressed.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::sync::Arc;
5
6use async_trait::async_trait;
7use futures::StreamExt as _;
8use vortex_array::ArrayRef;
9use vortex_array::ExecutionCtx;
10use vortex_array::VortexSessionExecute;
11use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes;
12use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeOpts;
13use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes;
14use vortex_array::expr::stats::Stat;
15use vortex_btrblocks::BtrBlocksCompressor;
16use vortex_error::VortexResult;
17use vortex_io::session::RuntimeSessionExt;
18use vortex_session::VortexSession;
19use vortex_utils::parallelism::get_available_parallelism;
20
21use crate::LayoutRef;
22use crate::LayoutStrategy;
23use crate::LayoutWriterContext;
24use crate::segments::SegmentSinkRef;
25use crate::sequence::SendableSequentialStream;
26use crate::sequence::SequencePointer;
27use crate::sequence::SequentialStreamAdapter;
28use crate::sequence::SequentialStreamExt;
29
30/// A boxed compressor function from arrays into compressed arrays.
31///
32/// API consumers are free to implement this trait to provide new plugin compressors.
33pub trait CompressorPlugin: Send + Sync + 'static {
34    fn compress_chunk(&self, chunk: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<ArrayRef>;
35}
36
37impl CompressorPlugin for Arc<dyn CompressorPlugin> {
38    fn compress_chunk(&self, chunk: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<ArrayRef> {
39        self.as_ref().compress_chunk(chunk, ctx)
40    }
41}
42
43impl<F> CompressorPlugin for F
44where
45    F: Fn(&ArrayRef, &mut ExecutionCtx) -> VortexResult<ArrayRef> + Send + Sync + 'static,
46{
47    fn compress_chunk(&self, chunk: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<ArrayRef> {
48        self(chunk, ctx)
49    }
50}
51
52impl CompressorPlugin for BtrBlocksCompressor {
53    fn compress_chunk(&self, chunk: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<ArrayRef> {
54        self.compress(chunk, ctx)
55    }
56}
57
58/// A layout writer that compresses chunks.
59#[derive(Clone)]
60pub struct CompressingStrategy {
61    child: Arc<dyn LayoutStrategy>,
62    compressor: Arc<dyn CompressorPlugin>,
63    stats: Arc<[Stat]>,
64    concurrency: usize,
65}
66
67impl CompressingStrategy {
68    /// Create a new compressing layout strategy with the given child strategy and compressor.
69    pub fn new<S: LayoutStrategy, C: CompressorPlugin>(child: S, compressor: C) -> Self {
70        Self {
71            child: Arc::new(child),
72            compressor: Arc::new(compressor),
73            stats: Stat::all().collect(),
74            concurrency: get_available_parallelism().unwrap_or(1),
75        }
76    }
77
78    pub fn with_concurrency(mut self, concurrency: usize) -> Self {
79        self.concurrency = concurrency;
80        self
81    }
82
83    /// Override the set of statistics computed on each chunk before compression.
84    /// Defaults to `Stat::all()`.
85    pub fn with_stats(mut self, stats: &[Stat]) -> Self {
86        self.stats = stats.into();
87        self
88    }
89}
90
91#[async_trait]
92impl LayoutStrategy for CompressingStrategy {
93    async fn write_stream(
94        &self,
95        ctx: LayoutWriterContext,
96        segment_sink: SegmentSinkRef,
97        stream: SendableSequentialStream,
98        eof: SequencePointer,
99        session: &VortexSession,
100    ) -> VortexResult<LayoutRef> {
101        let dtype = stream.dtype().clone();
102        let compressor = Arc::clone(&self.compressor);
103        let stats = Arc::clone(&self.stats);
104        let session = session.clone();
105        let compute_session = session.clone();
106
107        let handle = session.handle();
108        let stream = stream
109            .map(move |chunk| {
110                let compressor = Arc::clone(&compressor);
111                let stats = Arc::clone(&stats);
112                let session = compute_session.clone();
113                handle.spawn_cpu(move || {
114                    let (sequence_id, chunk) = chunk?;
115                    let mut ctx = session.create_execution_ctx();
116                    // Compute the stats for the chunk prior to compression. The exact
117                    // uncompressed size is measured once here and cached for the compressor.
118                    chunk.statistics().compute_all(&stats, &mut ctx)?;
119                    if UncompressedSizeInBytes::supports_dtype(chunk.dtype()) {
120                        uncompressed_size_in_bytes(
121                            &chunk,
122                            UncompressedSizeOpts::exact(),
123                            &mut ctx,
124                        )?;
125                    }
126                    Ok((sequence_id, compressor.compress_chunk(&chunk, &mut ctx)?))
127                })
128            })
129            .buffered(self.concurrency);
130
131        self.child
132            .write_stream(
133                ctx,
134                segment_sink,
135                SequentialStreamAdapter::new(dtype, stream).sendable(),
136                eof,
137                &session,
138            )
139            .await
140    }
141}
vortex_layout/layouts/compressed.rs

vortex_layout/layouts/
compressed.rs
