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::file::close_file;
10use polars_utils::io::create_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::io::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        match self {
186            Self::Dyn(v) => v.flush(),
187            Self::Local(v) => v.flush(),
188            #[cfg(feature = "cloud")]
189            Self::Cloud(v) => v.flush(),
190        }
191    }
192}
193
194impl Deref for Writable {
195    type Target = dyn io::Write + Send;
196
197    fn deref(&self) -> &Self::Target {
198        match self {
199            Self::Dyn(v) => v,
200            Self::Local(v) => v,
201            #[cfg(feature = "cloud")]
202            Self::Cloud(v) => v,
203        }
204    }
205}
206
207impl DerefMut for Writable {
208    fn deref_mut(&mut self) -> &mut Self::Target {
209        match self {
210            Self::Dyn(v) => v,
211            Self::Local(v) => v,
212            #[cfg(feature = "cloud")]
213            Self::Cloud(v) => v,
214        }
215    }
216}
217
218/// Avoid BufWriter wrapping on writers that already have internal buffering.
219pub enum BufferedWritable<'a> {
220    BufWriter(std::io::BufWriter<&'a mut (dyn std::io::Write + Send)>),
221    Direct(&'a mut (dyn std::io::Write + Send)),
222}
223
224impl<'a> Deref for BufferedWritable<'a> {
225    type Target = dyn io::Write + Send + 'a;
226
227    fn deref(&self) -> &Self::Target {
228        match self {
229            Self::BufWriter(v) => v as _,
230            Self::Direct(v) => v,
231        }
232    }
233}
234
235impl DerefMut for BufferedWritable<'_> {
236    fn deref_mut(&mut self) -> &mut Self::Target {
237        match self {
238            Self::BufWriter(v) => v as _,
239            Self::Direct(v) => v,
240        }
241    }
242}
243
244#[cfg(feature = "cloud")]
245async fn new_cloud_writer(
246    path: PlRefPath,
247    cloud_options: Option<&CloudOptions>,
248    cloud_upload_chunk_size: Option<NonZeroUsize>,
249    cloud_upload_concurrency: NonZeroUsize,
250    io_metrics: Option<Arc<IOMetrics>>,
251) -> PolarsResult<crate::cloud::cloud_writer::CloudWriter> {
252    use crate::cloud::cloud_writer::CloudWriter;
253    use crate::cloud::object_path_from_str;
254
255    let (cloud_location, object_store) =
256        crate::cloud::build_object_store(path, cloud_options, false).await?;
257
258    let mut writer = CloudWriter::new(
259        object_store,
260        object_path_from_str(&cloud_location.prefix)?,
261        cloud_upload_chunk_size,
262        cloud_upload_concurrency,
263        io_metrics,
264    );
265
266    writer.start().await?;
267
268    Ok(writer)
269}
270
271#[cfg(feature = "cloud")]
272mod async_writable {
273    use std::io;
274    use std::num::NonZeroUsize;
275    use std::ops::{Deref, DerefMut};
276    use std::pin::Pin;
277    use std::sync::Arc;
278    use std::task::{Context, Poll};
279
280    use bytes::Bytes;
281    use polars_error::{PolarsError, PolarsResult};
282    use polars_utils::file::close_file;
283    use polars_utils::pl_path::PlRefPath;
284    use tokio::io::AsyncWriteExt;
285    use tokio::task;
286
287    use super::{Writable, WritableTrait};
288    use crate::cloud::CloudOptions;
289    use crate::metrics::IOMetrics;
290    use crate::utils::sync_on_close::SyncOnCloseType;
291
292    /// Turn an abstract io::Write into an abstract tokio::io::AsyncWrite.
293    pub struct AsyncDynWritable(pub Box<dyn WritableTrait + Send>);
294
295    impl tokio::io::AsyncWrite for AsyncDynWritable {
296        fn poll_write(
297            self: Pin<&mut Self>,
298            _cx: &mut Context<'_>,
299            buf: &[u8],
300        ) -> Poll<io::Result<usize>> {
301            let result = task::block_in_place(|| self.get_mut().0.write(buf));
302            Poll::Ready(result)
303        }
304
305        fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
306            let result = task::block_in_place(|| self.get_mut().0.flush());
307            Poll::Ready(result)
308        }
309
310        fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
311            self.poll_flush(cx)
312        }
313    }
314
315    /// Holds an async writable file, abstracted over local files or cloud files.
316    ///
317    /// This implements `DerefMut` to a trait object implementing [`tokio::io::AsyncWrite`].
318    ///
319    /// Note: It is important that you do not call `shutdown()` on the deref'ed `AsyncWrite` object.
320    /// You should instead call the [`AsyncWritable::close`] at the end.
321    pub enum AsyncWritable {
322        Dyn(AsyncDynWritable),
323        Local(tokio::fs::File),
324        Cloud(crate::cloud::cloud_writer::CloudWriterIoTraitWrap),
325    }
326
327    impl AsyncWritable {
328        pub async fn try_new(
329            path: PlRefPath,
330            cloud_options: Option<&CloudOptions>,
331            cloud_upload_chunk_size: Option<NonZeroUsize>,
332            cloud_upload_concurrency: usize,
333            io_metrics: Option<Arc<IOMetrics>>,
334        ) -> PolarsResult<Self> {
335            // TODO: Native async impl
336            Writable::try_new(
337                path,
338                cloud_options,
339                cloud_upload_chunk_size,
340                cloud_upload_concurrency,
341                io_metrics,
342            )
343            .and_then(|x| x.try_into_async_writable())
344        }
345
346        /// If this writer holds a cloud writer, it will `mem::take(T)`. `T` is unmodified for other
347        /// writer types.
348        pub async fn write_all_owned<T>(&mut self, src: &mut T) -> io::Result<()>
349        where
350            T: AsRef<[u8]> + Default + Drop, // `Drop` is to exclude `&[u8]` slices.
351            Bytes: From<T>,
352        {
353            match self {
354                Self::Cloud(v) => v.write_all_owned(Bytes::from(std::mem::take(src))).await,
355                Self::Dyn(_) | Self::Local(_) => self.write_all(src.as_ref()).await,
356            }
357        }
358
359        pub async fn sync_all(&mut self) -> io::Result<()> {
360            match self {
361                Self::Dyn(v) => task::block_in_place(|| v.0.as_ref().sync_all()),
362                Self::Local(v) => v.sync_all().await,
363                Self::Cloud(_) => Ok(()),
364            }
365        }
366
367        pub async fn sync_data(&mut self) -> io::Result<()> {
368            match self {
369                Self::Dyn(v) => task::block_in_place(|| v.0.as_ref().sync_data()),
370                Self::Local(v) => v.sync_data().await,
371                Self::Cloud(_) => Ok(()),
372            }
373        }
374
375        pub async fn close(mut self, sync: SyncOnCloseType) -> PolarsResult<()> {
376            match sync {
377                SyncOnCloseType::All => self.sync_all().await?,
378                SyncOnCloseType::Data => self.sync_data().await?,
379                SyncOnCloseType::None => {},
380            }
381
382            match self {
383                Self::Dyn(mut v) => {
384                    v.shutdown().await.map_err(PolarsError::from)?;
385                    Ok(task::block_in_place(|| v.0.close())?)
386                },
387                Self::Local(v) => async {
388                    let f = v.into_std().await;
389                    close_file(f)
390                }
391                .await
392                .map_err(PolarsError::from),
393                Self::Cloud(mut v) => v.shutdown().await.map_err(PolarsError::from),
394            }
395        }
396    }
397
398    impl Deref for AsyncWritable {
399        type Target = dyn tokio::io::AsyncWrite + Send + Unpin;
400
401        fn deref(&self) -> &Self::Target {
402            match self {
403                Self::Dyn(v) => v,
404                Self::Local(v) => v,
405                Self::Cloud(v) => v,
406            }
407        }
408    }
409
410    impl DerefMut for AsyncWritable {
411        fn deref_mut(&mut self) -> &mut Self::Target {
412            match self {
413                Self::Dyn(v) => v,
414                Self::Local(v) => v,
415                Self::Cloud(v) => v,
416            }
417        }
418    }
419}
polars_io/utils/file.rs

polars_io/utils/
file.rs
