Skip to main content

vortex_layout/
strategy.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use async_trait::async_trait;
5use vortex_array::ArrayContext;
6use vortex_error::VortexResult;
7use vortex_session::VortexSession;
8
9use crate::LayoutRef;
10use crate::segments::SegmentSinkRef;
11use crate::sequence::SendableSequentialStream;
12use crate::sequence::SequencePointer;
13
14// [layout writer]
15/// Writes an ordered array stream into a layout tree and segment sink.
16///
17/// Layout strategies are writer-side extension points. Strategies may repartition, buffer,
18/// collect columns, compute statistics, compress arrays, or delegate to child strategies before
19/// finally emitting segments. They must preserve the logical row order represented by the
20/// [`SequencePointer`]s in the input stream.
21#[async_trait]
22pub trait LayoutStrategy: 'static + Send + Sync {
23    /// Asynchronously process an ordered stream of array chunks, emitting them into a sink and
24    /// returning the [`Layout`][crate::Layout] instance that can be parsed to retrieve the data
25    /// from rest.
26    ///
27    /// This trait uses the `#[async_trait]` attribute to denote that trait objects of this type
28    /// can be `Box`ed or `Arc`ed and shared around. Commonly, these strategies are composed to
29    /// form a operator of operations, each of which modifies the chunk stream in some way before
30    /// passing the data on to a downstream writer.
31    ///
32    /// # Sequencing and EOF
33    ///
34    /// The `stream` parameter is a stream of ordered array chunks, each of which is associated
35    /// with a sequence pointer that indicates its position in the overall array. By passing
36    /// around these pointers (essentially vector clocks), the writer can support concurrent
37    /// and parallel processing while maintaining a deterministic order of data in the file.
38    ///
39    /// The `eof` parameter is a guaranteed to be greater than all sequence pointers in the stream.
40    ///
41    /// Because child strategies can write to the end-of-file pointer, it is very important that
42    /// **all strategies must await all children concurrently**. Otherwise it is possible to
43    /// deadlock if one child is waiting to write to EOF while your strategy is preventing the
44    /// stream from progressing to completion.
45    ///
46    /// # Blocking operations
47    ///
48    /// This is an async trait method, which will return a `BoxFuture` that you can await from
49    /// any runtime. Implementations should avoid directly performing blocking work within the
50    /// `write_stream`, and should instead spawn it onto an appropriate runtime or threadpool
51    /// dedicated to such work.
52    ///
53    /// Such operations are common, and include things like compression and parsing large blobs
54    /// of data, or serializing very large messages to flatbuffers.
55    async fn write_stream(
56        &self,
57        ctx: ArrayContext,
58        segment_sink: SegmentSinkRef,
59        stream: SendableSequentialStream,
60        eof: SequencePointer,
61        session: &VortexSession,
62    ) -> VortexResult<LayoutRef>;
63
64    /// Returns the number of bytes currently buffered by this strategy and any child strategies.
65    ///
66    /// This method allows tracking of data that has been processed by the strategy but not yet
67    /// written to the underlying sink, providing more accurate estimates of final file size
68    /// during write operations.
69    fn buffered_bytes(&self) -> u64 {
70        0
71    }
72}
73
74#[async_trait]
75impl LayoutStrategy for std::sync::Arc<dyn LayoutStrategy> {
76    async fn write_stream(
77        &self,
78        ctx: ArrayContext,
79        segment_sink: SegmentSinkRef,
80        stream: SendableSequentialStream,
81        eof: SequencePointer,
82        session: &VortexSession,
83    ) -> VortexResult<LayoutRef> {
84        (**self)
85            .write_stream(ctx, segment_sink, stream, eof, session)
86            .await
87    }
88
89    fn buffered_bytes(&self) -> u64 {
90        (**self).buffered_bytes()
91    }
92}
93// [layout writer]