polars_io/cloud/cloud_writer/
writer.rs

1use std::num::NonZeroUsize;
2use std::sync::Arc;
3
4use bytes::Bytes;
5use polars_error::PolarsResult;
6
7use crate::cloud::PolarsObjectStore;
8use crate::cloud::cloud_writer::bufferer::BytesBufferer;
9use crate::cloud::cloud_writer::internal_writer::{InternalCloudWriter, InternalCloudWriterState};
10use crate::metrics::{IOMetrics, OptIOMetrics};
11
12pub struct CloudWriter {
13    writer: InternalCloudWriter,
14    bufferer: BytesBufferer,
15}
16
17impl CloudWriter {
18    pub fn new(
19        store: PolarsObjectStore,
20        path: object_store::path::Path,
21        upload_chunk_size: Option<NonZeroUsize>,
22        max_concurrency: NonZeroUsize,
23        io_metrics: Option<Arc<IOMetrics>>,
24    ) -> Self {
25        let bufferer = BytesBufferer::new(upload_chunk_size);
26
27        Self {
28            writer: InternalCloudWriter {
29                store,
30                path,
31                max_concurrency,
32                io_metrics: OptIOMetrics(io_metrics),
33                state: InternalCloudWriterState::NotStarted,
34            },
35            bufferer,
36        }
37    }
38
39    pub async fn write_all_owned(&mut self, mut bytes: Bytes) -> PolarsResult<()> {
40        while !bytes.is_empty() {
41            self.bufferer.push_owned(&mut bytes);
42
43            if let Some(payload) = self.bufferer.flush_full_chunk() {
44                self.writer.put(payload).await?;
45            }
46        }
47
48        Ok(())
49    }
50
51    pub(super) fn fill_buffer_from_slice(&mut self, bytes: &mut &[u8]) -> bool {
52        self.bufferer.push_slice(bytes);
53        self.bufferer.is_full()
54    }
55
56    pub(super) async fn flush_full_chunk(&mut self) -> PolarsResult<()> {
57        if let Some(payload) = self.bufferer.flush_full_chunk() {
58            self.writer.put(payload).await?;
59        }
60
61        Ok(())
62    }
63
64    pub async fn finish(&mut self) -> PolarsResult<()> {
65        let tail = self.bufferer.flush();
66        assert!(self.bufferer.is_empty());
67
68        self.writer.finish(tail).await
69    }
70}
polars_io/cloud/cloud_writer/writer.rs

polars_io/cloud/cloud_writer/
writer.rs
