vortex_layout/layouts/chunked/
writer.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::sync::Arc;
5
6use async_stream::stream;
7use async_trait::async_trait;
8use futures::StreamExt;
9use futures::TryStreamExt;
10use futures::stream;
11use vortex_error::VortexExpect;
12use vortex_error::VortexResult;
13use vortex_io::session::RuntimeSessionExt;
14use vortex_session::VortexSession;
15
16use crate::LayoutRef;
17use crate::LayoutStrategy;
18use crate::LayoutWriterContext;
19use crate::children::OwnedLayoutChildren;
20use crate::layouts::chunked::ChunkedLayout;
21use crate::segments::SegmentSinkRef;
22use crate::sequence::SendableSequentialStream;
23use crate::sequence::SequencePointer;
24use crate::sequence::SequentialStreamAdapter;
25use crate::sequence::SequentialStreamExt as _;
26
27#[derive(Clone)]
28pub struct ChunkedLayoutStrategy {
29    /// The layout strategy for each chunk.
30    pub chunk_strategy: Arc<dyn LayoutStrategy>,
31}
32
33impl ChunkedLayoutStrategy {
34    pub fn new<S: LayoutStrategy>(chunk_strategy: S) -> Self {
35        Self {
36            chunk_strategy: Arc::new(chunk_strategy),
37        }
38    }
39}
40
41#[async_trait]
42impl LayoutStrategy for ChunkedLayoutStrategy {
43    async fn write_stream(
44        &self,
45        ctx: LayoutWriterContext,
46        segment_sink: SegmentSinkRef,
47        stream: SendableSequentialStream,
48        mut eof: SequencePointer,
49        session: &VortexSession,
50    ) -> VortexResult<LayoutRef> {
51        let dtype = stream.dtype().clone();
52        let dtype2 = dtype.clone();
53        let chunk_strategy = Arc::clone(&self.chunk_strategy);
54        let handle = session.handle();
55
56        // We spawn each child to allow parallelism when processing chunks.
57        let stream = stream! {
58            let mut stream = stream;
59            while let Some(chunk) = stream.next().await {
60                let chunk_eof = eof.split_off();
61
62                let chunk_strategy = Arc::clone(&chunk_strategy);
63                let ctx = ctx.clone();
64                let segment_sink = Arc::clone(&segment_sink);
65                let dtype = dtype2.clone();
66                let session = session.clone();
67
68                yield handle.spawn_nested(move |_| async move {
69                    chunk_strategy
70                        .write_stream(
71                            ctx,
72                            segment_sink,
73                            SequentialStreamAdapter::new(
74                                dtype,
75                                stream::iter([chunk]),
76                            )
77                            .sendable(),
78                            chunk_eof,
79                            &session,
80                        )
81                        .await
82                })
83            }
84        };
85
86        // Poll all of our children concurrently to accumulate their layouts.
87        let mut child_layouts: Vec<LayoutRef> = stream.buffered(usize::MAX).try_collect().await?;
88
89        if child_layouts.len() == 1 {
90            Ok(child_layouts.pop().vortex_expect("must have one child"))
91        } else {
92            let row_count = child_layouts.iter().map(|layout| layout.row_count()).sum();
93            Ok(ChunkedLayout::new(
94                row_count,
95                dtype,
96                OwnedLayoutChildren::layout_children(child_layouts),
97            )
98            .into_layout())
99        }
100    }
101}
vortex_layout/layouts/chunked/writer.rs

vortex_layout/layouts/chunked/
writer.rs
