polars_io/cloud/cloud_writer/
writer.rs

1use std::num::NonZeroUsize;
2use std::sync::Arc;
3
4use bytes::Bytes;
5use object_store::PutPayload;
6use polars_error::PolarsResult;
7
8use crate::cloud::PolarsObjectStore;
9use crate::cloud::cloud_writer::bufferer::BytesBufferer;
10use crate::cloud::cloud_writer::internal_writer::{InternalCloudWriter, InternalCloudWriterState};
11use crate::metrics::{IOMetrics, OptIOMetrics};
12
13pub struct CloudWriter {
14    writer: InternalCloudWriter,
15    bufferer: BytesBufferer,
16}
17
18impl CloudWriter {
19    pub fn new(
20        store: PolarsObjectStore,
21        path: object_store::path::Path,
22        upload_chunk_size: usize,
23        max_concurrency: NonZeroUsize,
24        io_metrics: Option<Arc<IOMetrics>>,
25    ) -> Self {
26        let bufferer = BytesBufferer::new(upload_chunk_size);
27
28        Self {
29            writer: InternalCloudWriter {
30                store,
31                path,
32                max_concurrency,
33                io_metrics: OptIOMetrics(io_metrics),
34                state: InternalCloudWriterState::NotStarted,
35            },
36            bufferer,
37        }
38    }
39
40    pub async fn start(&mut self) -> PolarsResult<()> {
41        self.writer.start().await
42    }
43
44    pub async fn write_all_owned(&mut self, mut bytes: Bytes) -> PolarsResult<()> {
45        while !bytes.is_empty() {
46            self.bufferer.push_owned(&mut bytes);
47
48            if let Some(payload) = self.bufferer.flush_full_chunk() {
49                self.writer.put(payload).await?;
50            }
51        }
52
53        Ok(())
54    }
55
56    /// Write multiple Bytes without buffering.
57    pub async fn write_multiple_owned_unbuffered<I, T>(&mut self, bytes: I) -> PolarsResult<()>
58    where
59        I: IntoIterator<Item = T>,
60        Bytes: From<T>,
61    {
62        let payload = PutPayload::from_iter(bytes.into_iter().map(Bytes::from));
63
64        if payload.iter().next().is_some() {
65            if let Some(payload) = self.bufferer.flush() {
66                self.writer.put(payload).await?;
67            }
68
69            self.writer.put(payload).await?;
70        }
71
72        Ok(())
73    }
74
75    pub(super) fn fill_buffer_from_slice(&mut self, bytes: &mut &[u8]) -> bool {
76        self.bufferer.push_slice(bytes);
77        self.bufferer.is_full()
78    }
79
80    pub(super) async fn flush_full_chunk(&mut self) -> PolarsResult<()> {
81        if let Some(payload) = self.bufferer.flush_full_chunk() {
82            self.writer.put(payload).await?;
83        }
84
85        Ok(())
86    }
87
88    pub(super) async fn flush(&mut self) -> PolarsResult<()> {
89        if let Some(payload) = self.bufferer.flush() {
90            self.writer.put(payload).await?;
91        }
92
93        assert!(self.bufferer.is_empty());
94
95        Ok(())
96    }
97
98    pub(super) fn has_buffered_bytes(&self) -> bool {
99        !self.bufferer.is_empty()
100    }
101
102    pub async fn finish(&mut self) -> PolarsResult<()> {
103        if let Some(payload) = self.bufferer.flush() {
104            self.writer.put(payload).await?;
105        }
106
107        assert!(self.bufferer.is_empty());
108
109        self.writer.finish().await
110    }
111}
polars_io/cloud/cloud_writer/writer.rs

polars_io/cloud/cloud_writer/
writer.rs
