polars_io/utils/
file.rs

1use std::io;
2use std::num::NonZeroUsize;
3use std::ops::{Deref, DerefMut};
4use std::sync::Arc;
5
6#[cfg(feature = "cloud")]
7pub use async_writable::{AsyncDynWritable, AsyncWritable};
8use polars_error::{PolarsResult, feature_gated, polars_err};
9use polars_utils::create_file;
10use polars_utils::file::close_file;
11use polars_utils::mmap::ensure_not_mapped;
12use polars_utils::pl_path::{PlRefPath, format_file_uri};
13
14use super::sync_on_close::SyncOnCloseType;
15use crate::cloud::CloudOptions;
16use crate::metrics::IOMetrics;
17use crate::resolve_homedir;
18
19// TODO document precise contract.
20pub trait WritableTrait: std::io::Write {
21    fn close(&mut self) -> std::io::Result<()>;
22    fn sync_all(&self) -> std::io::Result<()>;
23    fn sync_data(&self) -> std::io::Result<()>;
24}
25
26/// Holds a non-async writable file, abstracted over local files or cloud files.
27///
28/// This implements `DerefMut` to a trait object implementing [`std::io::Write`].
29///
30/// Also see: `Writable::try_into_async_writable` and `AsyncWritable`.
31#[allow(clippy::large_enum_variant)] // It will be boxed
32pub enum Writable {
33    /// An abstract implementation for writable.
34    ///
35    /// This is used to implement writing to in-memory and arbitrary file descriptors.
36    Dyn(Box<dyn WritableTrait + Send>),
37    Local(std::fs::File),
38    #[cfg(feature = "cloud")]
39    Cloud(crate::cloud::cloud_writer::CloudWriterIoTraitWrap),
40}
41
42impl Writable {
43    pub fn try_new(
44        path: PlRefPath,
45        #[cfg_attr(not(feature = "cloud"), expect(unused))] cloud_options: Option<&CloudOptions>,
46        #[cfg_attr(not(feature = "cloud"), expect(unused))] cloud_upload_chunk_size: Option<
47            NonZeroUsize,
48        >,
49        #[cfg_attr(not(feature = "cloud"), expect(unused))] cloud_upload_concurrency: usize,
50        io_metrics: Option<Arc<IOMetrics>>,
51    ) -> PolarsResult<Self> {
52        Ok(if path.has_scheme() {
53            feature_gated!("cloud", {
54                use polars_core::runtime::ASYNC;
55
56                use crate::cloud::cloud_writer::CloudWriterIoTraitWrap;
57
58                let writer = ASYNC.block_in_place_on(new_cloud_writer(
59                    path,
60                    cloud_options,
61                    cloud_upload_chunk_size,
62                    cloud_upload_concurrency.try_into().unwrap(),
63                    io_metrics,
64                ))?;
65
66                Self::Cloud(CloudWriterIoTraitWrap::from(writer))
67            })
68        } else if polars_config::config().force_async() {
69            feature_gated!("cloud", {
70                let path = resolve_homedir(path.as_std_path());
71                create_file(&path)?;
72                let path = std::fs::canonicalize(&path)?;
73
74                ensure_not_mapped(&path.metadata()?)?;
75
76                let path = path.to_str().ok_or_else(|| polars_err!(non_utf8_path))?;
77                let path = format_file_uri(path);
78
79                use polars_core::runtime::ASYNC;
80
81                use crate::cloud::cloud_writer::CloudWriterIoTraitWrap;
82
83                let writer = ASYNC.block_in_place_on(new_cloud_writer(
84                    path,
85                    cloud_options,
86                    cloud_upload_chunk_size,
87                    cloud_upload_concurrency.try_into().unwrap(),
88                    io_metrics,
89                ))?;
90
91                Self::Cloud(CloudWriterIoTraitWrap::from(writer))
92            })
93        } else {
94            let path = resolve_homedir(path.as_std_path());
95            create_file(&path)?;
96
97            Self::Local(polars_utils::open_file_write(&path)?)
98        })
99    }
100
101    /// If this writer holds a cloud writer, it will `mem::take(T)`. `T` is unmodified for other
102    /// writer types.
103    #[cfg(feature = "cloud")]
104    pub async fn write_all_owned<T>(&mut self, src: &mut T) -> io::Result<()>
105    where
106        T: AsRef<[u8]> + Default + Drop, // `Drop` is to exclude `&[u8]` slices.
107        bytes::Bytes: From<T>,
108    {
109        match self {
110            Self::Cloud(v) => {
111                v.write_all_owned(bytes::Bytes::from(std::mem::take(src)))
112                    .await
113            },
114            Self::Dyn(_) | Self::Local(_) => self.write_all(src.as_ref()),
115        }
116    }
117
118    /// This returns `Result<>` - if a write was performed before calling this,
119    /// `CloudWriter` can be in an Err(_) state.
120    #[cfg(feature = "cloud")]
121    pub fn try_into_async_writable(self) -> PolarsResult<AsyncWritable> {
122        use self::async_writable::AsyncDynWritable;
123
124        match self {
125            Self::Dyn(v) => Ok(AsyncWritable::Dyn(AsyncDynWritable(v))),
126            Self::Local(v) => Ok(AsyncWritable::Local(tokio::fs::File::from_std(v))),
127            Self::Cloud(v) => Ok(AsyncWritable::Cloud(v)),
128        }
129    }
130
131    pub fn as_buffered(&mut self) -> BufferedWritable<'_> {
132        match self {
133            Writable::Dyn(v) => BufferedWritable::BufWriter(std::io::BufWriter::new(v.as_mut())),
134            Writable::Local(v) => BufferedWritable::BufWriter(std::io::BufWriter::new(v)),
135            #[cfg(feature = "cloud")]
136            Writable::Cloud(v) => BufferedWritable::Direct(v as _),
137        }
138    }
139
140    pub fn sync_all(&self) -> io::Result<()> {
141        match self {
142            Self::Dyn(v) => v.sync_all(),
143            Self::Local(v) => v.sync_all(),
144            #[cfg(feature = "cloud")]
145            Self::Cloud(v) => v.sync_all(),
146        }
147    }
148
149    pub fn sync_data(&self) -> io::Result<()> {
150        match self {
151            Self::Dyn(v) => v.sync_data(),
152            Self::Local(v) => v.sync_data(),
153            #[cfg(feature = "cloud")]
154            Self::Cloud(v) => v.sync_data(),
155        }
156    }
157
158    pub fn close(self, sync: SyncOnCloseType) -> std::io::Result<()> {
159        match sync {
160            SyncOnCloseType::All => self.sync_all()?,
161            SyncOnCloseType::Data => self.sync_data()?,
162            SyncOnCloseType::None => {},
163        }
164
165        match self {
166            Self::Dyn(mut v) => v.close(),
167            Self::Local(v) => close_file(v),
168            #[cfg(feature = "cloud")]
169            Self::Cloud(mut v) => v.close(),
170        }
171    }
172}
173
174impl io::Write for Writable {
175    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
176        match self {
177            Self::Dyn(v) => v.write(buf),
178            Self::Local(v) => v.write(buf),
179            #[cfg(feature = "cloud")]
180            Self::Cloud(v) => v.write(buf),
181        }
182    }
183
184    fn flush(&mut self) -> io::Result<()> {
185        self.sync_all()
186    }
187}
188
189impl Deref for Writable {
190    type Target = dyn io::Write + Send;
191
192    fn deref(&self) -> &Self::Target {
193        match self {
194            Self::Dyn(v) => v,
195            Self::Local(v) => v,
196            #[cfg(feature = "cloud")]
197            Self::Cloud(v) => v,
198        }
199    }
200}
201
202impl DerefMut for Writable {
203    fn deref_mut(&mut self) -> &mut Self::Target {
204        match self {
205            Self::Dyn(v) => v,
206            Self::Local(v) => v,
207            #[cfg(feature = "cloud")]
208            Self::Cloud(v) => v,
209        }
210    }
211}
212
213/// Avoid BufWriter wrapping on writers that already have internal buffering.
214pub enum BufferedWritable<'a> {
215    BufWriter(std::io::BufWriter<&'a mut (dyn std::io::Write + Send)>),
216    Direct(&'a mut (dyn std::io::Write + Send)),
217}
218
219impl<'a> Deref for BufferedWritable<'a> {
220    type Target = dyn io::Write + Send + 'a;
221
222    fn deref(&self) -> &Self::Target {
223        match self {
224            Self::BufWriter(v) => v as _,
225            Self::Direct(v) => v,
226        }
227    }
228}
229
230impl DerefMut for BufferedWritable<'_> {
231    fn deref_mut(&mut self) -> &mut Self::Target {
232        match self {
233            Self::BufWriter(v) => v as _,
234            Self::Direct(v) => v,
235        }
236    }
237}
238
239#[cfg(feature = "cloud")]
240async fn new_cloud_writer(
241    path: PlRefPath,
242    cloud_options: Option<&CloudOptions>,
243    cloud_upload_chunk_size: Option<NonZeroUsize>,
244    cloud_upload_concurrency: NonZeroUsize,
245    io_metrics: Option<Arc<IOMetrics>>,
246) -> PolarsResult<crate::cloud::cloud_writer::CloudWriter> {
247    use crate::cloud::cloud_writer::CloudWriter;
248    use crate::cloud::object_path_from_str;
249
250    let (cloud_location, object_store) =
251        crate::cloud::build_object_store(path, cloud_options, false).await?;
252
253    let mut writer = CloudWriter::new(
254        object_store,
255        object_path_from_str(&cloud_location.prefix)?,
256        cloud_upload_chunk_size,
257        cloud_upload_concurrency,
258        io_metrics,
259    );
260
261    writer.start().await?;
262
263    Ok(writer)
264}
265
266#[cfg(feature = "cloud")]
267mod async_writable {
268    use std::io;
269    use std::num::NonZeroUsize;
270    use std::ops::{Deref, DerefMut};
271    use std::pin::Pin;
272    use std::sync::Arc;
273    use std::task::{Context, Poll};
274
275    use bytes::Bytes;
276    use polars_error::{PolarsError, PolarsResult};
277    use polars_utils::file::close_file;
278    use polars_utils::pl_path::PlRefPath;
279    use tokio::io::AsyncWriteExt;
280    use tokio::task;
281
282    use super::{Writable, WritableTrait};
283    use crate::cloud::CloudOptions;
284    use crate::metrics::IOMetrics;
285    use crate::utils::sync_on_close::SyncOnCloseType;
286
287    /// Turn an abstract io::Write into an abstract tokio::io::AsyncWrite.
288    pub struct AsyncDynWritable(pub Box<dyn WritableTrait + Send>);
289
290    impl tokio::io::AsyncWrite for AsyncDynWritable {
291        fn poll_write(
292            self: Pin<&mut Self>,
293            _cx: &mut Context<'_>,
294            buf: &[u8],
295        ) -> Poll<io::Result<usize>> {
296            let result = task::block_in_place(|| self.get_mut().0.write(buf));
297            Poll::Ready(result)
298        }
299
300        fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
301            let result = task::block_in_place(|| self.get_mut().0.flush());
302            Poll::Ready(result)
303        }
304
305        fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
306            self.poll_flush(cx)
307        }
308    }
309
310    /// Holds an async writable file, abstracted over local files or cloud files.
311    ///
312    /// This implements `DerefMut` to a trait object implementing [`tokio::io::AsyncWrite`].
313    ///
314    /// Note: It is important that you do not call `shutdown()` on the deref'ed `AsyncWrite` object.
315    /// You should instead call the [`AsyncWritable::close`] at the end.
316    pub enum AsyncWritable {
317        Dyn(AsyncDynWritable),
318        Local(tokio::fs::File),
319        Cloud(crate::cloud::cloud_writer::CloudWriterIoTraitWrap),
320    }
321
322    impl AsyncWritable {
323        pub async fn try_new(
324            path: PlRefPath,
325            cloud_options: Option<&CloudOptions>,
326            cloud_upload_chunk_size: Option<NonZeroUsize>,
327            cloud_upload_concurrency: usize,
328            io_metrics: Option<Arc<IOMetrics>>,
329        ) -> PolarsResult<Self> {
330            // TODO: Native async impl
331            Writable::try_new(
332                path,
333                cloud_options,
334                cloud_upload_chunk_size,
335                cloud_upload_concurrency,
336                io_metrics,
337            )
338            .and_then(|x| x.try_into_async_writable())
339        }
340
341        /// If this writer holds a cloud writer, it will `mem::take(T)`. `T` is unmodified for other
342        /// writer types.
343        pub async fn write_all_owned<T>(&mut self, src: &mut T) -> io::Result<()>
344        where
345            T: AsRef<[u8]> + Default + Drop, // `Drop` is to exclude `&[u8]` slices.
346            Bytes: From<T>,
347        {
348            match self {
349                Self::Cloud(v) => v.write_all_owned(Bytes::from(std::mem::take(src))).await,
350                Self::Dyn(_) | Self::Local(_) => self.write_all(src.as_ref()).await,
351            }
352        }
353
354        pub async fn sync_all(&mut self) -> io::Result<()> {
355            match self {
356                Self::Dyn(v) => task::block_in_place(|| v.0.as_ref().sync_all()),
357                Self::Local(v) => v.sync_all().await,
358                Self::Cloud(_) => Ok(()),
359            }
360        }
361
362        pub async fn sync_data(&mut self) -> io::Result<()> {
363            match self {
364                Self::Dyn(v) => task::block_in_place(|| v.0.as_ref().sync_data()),
365                Self::Local(v) => v.sync_data().await,
366                Self::Cloud(_) => Ok(()),
367            }
368        }
369
370        pub async fn close(mut self, sync: SyncOnCloseType) -> PolarsResult<()> {
371            match sync {
372                SyncOnCloseType::All => self.sync_all().await?,
373                SyncOnCloseType::Data => self.sync_data().await?,
374                SyncOnCloseType::None => {},
375            }
376
377            match self {
378                Self::Dyn(mut v) => {
379                    v.shutdown().await.map_err(PolarsError::from)?;
380                    Ok(task::block_in_place(|| v.0.close())?)
381                },
382                Self::Local(v) => async {
383                    let f = v.into_std().await;
384                    close_file(f)
385                }
386                .await
387                .map_err(PolarsError::from),
388                Self::Cloud(mut v) => v.shutdown().await.map_err(PolarsError::from),
389            }
390        }
391    }
392
393    impl Deref for AsyncWritable {
394        type Target = dyn tokio::io::AsyncWrite + Send + Unpin;
395
396        fn deref(&self) -> &Self::Target {
397            match self {
398                Self::Dyn(v) => v,
399                Self::Local(v) => v,
400                Self::Cloud(v) => v,
401            }
402        }
403    }
404
405    impl DerefMut for AsyncWritable {
406        fn deref_mut(&mut self) -> &mut Self::Target {
407            match self {
408                Self::Dyn(v) => v,
409                Self::Local(v) => v,
410                Self::Cloud(v) => v,
411            }
412        }
413    }
414}
polars_io/utils/file.rs

polars_io/utils/
file.rs
