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 |handle| async move {
69                    let session = session.with_handle(handle);
70                    chunk_strategy
71                        .write_stream(
72                            ctx,
73                            segment_sink,
74                            SequentialStreamAdapter::new(
75                                dtype,
76                                stream::iter([chunk]),
77                            )
78                            .sendable(),
79                            chunk_eof,
80                            &session,
81                        )
82                        .await
83                })
84            }
85        };
86
87        // Poll all of our children concurrently to accumulate their layouts.
88        let mut child_layouts: Vec<LayoutRef> = stream.buffered(usize::MAX).try_collect().await?;
89
90        if child_layouts.len() == 1 {
91            Ok(child_layouts.pop().vortex_expect("must have one child"))
92        } else {
93            let row_count = child_layouts.iter().map(|layout| layout.row_count()).sum();
94            Ok(ChunkedLayout::new(
95                row_count,
96                dtype,
97                OwnedLayoutChildren::layout_children(child_layouts),
98            )
99            .into_layout())
100        }
101    }
102}
vortex_layout/layouts/chunked/writer.rs

vortex_layout/layouts/chunked/
writer.rs
