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]