polars_io/utils/
file.rs

1use std::io;
2use std::num::NonZeroUsize;
3use std::ops::{Deref, DerefMut};
4use std::sync::Arc;
5
6use polars_error::{PolarsResult, feature_gated, polars_err};
7use polars_utils::file::close_file;
8use polars_utils::io::create_file;
9use polars_utils::mmap::ensure_not_mapped;
10use polars_utils::pl_path::{PlRefPath, format_file_uri};
11
12use super::sync_on_close::SyncOnCloseType;
13use crate::cloud::CloudOptions;
14use crate::metrics::IOMetrics;
15use crate::resolve_homedir;
16
17// TODO document precise contract.
18pub trait WritableTrait: std::io::Write {
19    fn close(&mut self) -> std::io::Result<()>;
20    fn sync_all(&self) -> std::io::Result<()>;
21    fn sync_data(&self) -> std::io::Result<()>;
22}
23
24/// Holds a non-async writable file, abstracted over local files or cloud files.
25///
26/// This implements `DerefMut` to a trait object implementing [`std::io::Write`].
27///
28/// Also see: `Writable::try_into_async_writable` and `AsyncWritable`.
29#[allow(clippy::large_enum_variant)] // It will be boxed
30pub enum Writable {
31    /// An abstract implementation for writable.
32    ///
33    /// This is used to implement writing to in-memory and arbitrary file descriptors.
34    Dyn(Box<dyn WritableTrait + Send>),
35    Local(std::fs::File),
36    #[cfg(feature = "cloud")]
37    Cloud(crate::cloud::cloud_writer::CloudWriterIoTraitWrap),
38}
39
40impl Writable {
41    pub fn try_new(
42        path: PlRefPath,
43        #[cfg_attr(not(feature = "cloud"), expect(unused))] cloud_options: Option<&CloudOptions>,
44        #[cfg_attr(not(feature = "cloud"), expect(unused))] cloud_upload_chunk_size: Option<
45            NonZeroUsize,
46        >,
47        #[cfg_attr(not(feature = "cloud"), expect(unused))] cloud_upload_concurrency: usize,
48        io_metrics: Option<Arc<IOMetrics>>,
49    ) -> PolarsResult<Self> {
50        Ok(if path.has_scheme() {
51            feature_gated!("cloud", {
52                use polars_core::runtime::ASYNC;
53
54                use crate::cloud::cloud_writer::CloudWriterIoTraitWrap;
55
56                let writer = ASYNC.block_in_place_on(new_cloud_writer(
57                    path,
58                    cloud_options,
59                    cloud_upload_chunk_size,
60                    cloud_upload_concurrency.try_into().unwrap(),
61                    io_metrics,
62                ))?;
63
64                Self::Cloud(CloudWriterIoTraitWrap::from(writer))
65            })
66        } else if polars_config::config().force_async() {
67            feature_gated!("cloud", {
68                let path = resolve_homedir(path.as_std_path());
69                create_file(&path)?;
70                let path = std::fs::canonicalize(&path)?;
71
72                ensure_not_mapped(&path.metadata()?)?;
73
74                let path = path.to_str().ok_or_else(|| polars_err!(non_utf8_path))?;
75                let path = format_file_uri(path);
76
77                use polars_core::runtime::ASYNC;
78
79                use crate::cloud::cloud_writer::CloudWriterIoTraitWrap;
80
81                let writer = ASYNC.block_in_place_on(new_cloud_writer(
82                    path,
83                    cloud_options,
84                    cloud_upload_chunk_size,
85                    cloud_upload_concurrency.try_into().unwrap(),
86                    io_metrics,
87                ))?;
88
89                Self::Cloud(CloudWriterIoTraitWrap::from(writer))
90            })
91        } else {
92            let path = resolve_homedir(path.as_std_path());
93            create_file(&path)?;
94
95            Self::Local(polars_utils::io::open_file_write(&path)?)
96        })
97    }
98
99    /// If this writer holds a cloud writer, it will `mem::take(T)`. `T` is unmodified for other
100    /// writer types.
101    #[cfg(feature = "cloud")]
102    pub async fn write_all_owned<T>(&mut self, src: &mut T) -> io::Result<()>
103    where
104        T: AsRef<[u8]> + Default + Drop, // `Drop` is to exclude `&[u8]` slices.
105        bytes::Bytes: From<T>,
106    {
107        match self {
108            Self::Cloud(v) => {
109                v.write_all_owned(bytes::Bytes::from(std::mem::take(src)))
110                    .await
111            },
112            Self::Dyn(_) | Self::Local(_) => self.write_all(src.as_ref()),
113        }
114    }
115
116    pub fn as_buffered(&mut self) -> BufferedWritable<'_> {
117        match self {
118            Writable::Dyn(v) => BufferedWritable::BufWriter(std::io::BufWriter::new(v.as_mut())),
119            Writable::Local(v) => BufferedWritable::BufWriter(std::io::BufWriter::new(v)),
120            #[cfg(feature = "cloud")]
121            Writable::Cloud(v) => BufferedWritable::Direct(v as _),
122        }
123    }
124
125    pub fn sync_all(&self) -> io::Result<()> {
126        match self {
127            Self::Dyn(v) => v.sync_all(),
128            Self::Local(v) => v.sync_all(),
129            #[cfg(feature = "cloud")]
130            Self::Cloud(v) => v.sync_all(),
131        }
132    }
133
134    pub fn sync_data(&self) -> io::Result<()> {
135        match self {
136            Self::Dyn(v) => v.sync_data(),
137            Self::Local(v) => v.sync_data(),
138            #[cfg(feature = "cloud")]
139            Self::Cloud(v) => v.sync_data(),
140        }
141    }
142
143    pub fn close(self, sync: SyncOnCloseType) -> std::io::Result<()> {
144        match sync {
145            SyncOnCloseType::All => self.sync_all()?,
146            SyncOnCloseType::Data => self.sync_data()?,
147            SyncOnCloseType::None => {},
148        }
149
150        match self {
151            Self::Dyn(mut v) => v.close(),
152            Self::Local(v) => close_file(v),
153            #[cfg(feature = "cloud")]
154            Self::Cloud(mut v) => v.close(),
155        }
156    }
157}
158
159impl io::Write for Writable {
160    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
161        match self {
162            Self::Dyn(v) => v.write(buf),
163            Self::Local(v) => v.write(buf),
164            #[cfg(feature = "cloud")]
165            Self::Cloud(v) => v.write(buf),
166        }
167    }
168
169    fn flush(&mut self) -> io::Result<()> {
170        match self {
171            Self::Dyn(v) => v.flush(),
172            Self::Local(v) => v.flush(),
173            #[cfg(feature = "cloud")]
174            Self::Cloud(v) => v.flush(),
175        }
176    }
177}
178
179impl Deref for Writable {
180    type Target = dyn io::Write + Send;
181
182    fn deref(&self) -> &Self::Target {
183        match self {
184            Self::Dyn(v) => v,
185            Self::Local(v) => v,
186            #[cfg(feature = "cloud")]
187            Self::Cloud(v) => v,
188        }
189    }
190}
191
192impl DerefMut for Writable {
193    fn deref_mut(&mut self) -> &mut Self::Target {
194        match self {
195            Self::Dyn(v) => v,
196            Self::Local(v) => v,
197            #[cfg(feature = "cloud")]
198            Self::Cloud(v) => v,
199        }
200    }
201}
202
203/// Avoid BufWriter wrapping on writers that already have internal buffering.
204pub enum BufferedWritable<'a> {
205    BufWriter(std::io::BufWriter<&'a mut (dyn std::io::Write + Send)>),
206    Direct(&'a mut (dyn std::io::Write + Send)),
207}
208
209impl<'a> Deref for BufferedWritable<'a> {
210    type Target = dyn io::Write + Send + 'a;
211
212    fn deref(&self) -> &Self::Target {
213        match self {
214            Self::BufWriter(v) => v as _,
215            Self::Direct(v) => v,
216        }
217    }
218}
219
220impl DerefMut for BufferedWritable<'_> {
221    fn deref_mut(&mut self) -> &mut Self::Target {
222        match self {
223            Self::BufWriter(v) => v as _,
224            Self::Direct(v) => v,
225        }
226    }
227}
228
229#[cfg(feature = "cloud")]
230async fn new_cloud_writer(
231    path: PlRefPath,
232    cloud_options: Option<&CloudOptions>,
233    cloud_upload_chunk_size: Option<NonZeroUsize>,
234    cloud_upload_concurrency: NonZeroUsize,
235    io_metrics: Option<Arc<IOMetrics>>,
236) -> PolarsResult<crate::cloud::cloud_writer::CloudWriter> {
237    use crate::cloud::cloud_writer::CloudWriter;
238    use crate::cloud::object_path_from_str;
239
240    let (cloud_location, object_store) =
241        crate::cloud::build_object_store(path, cloud_options, false).await?;
242
243    let mut writer = CloudWriter::new(
244        object_store,
245        object_path_from_str(&cloud_location.prefix)?,
246        cloud_upload_chunk_size,
247        cloud_upload_concurrency,
248        io_metrics,
249    );
250
251    writer.start().await?;
252
253    Ok(writer)
254}
polars_io/utils/file.rs

polars_io/utils/
file.rs
