refactor(rust): Take sync parameter in Writeable::close() (#25475)

Author: nameexhaustion

crates/polars-io/src/utils/file.rs

@@ -128,20 +128,31 @@ impl Writeable {
         }
     }
 
-    pub fn sync_on_close(&mut self, sync_on_close: SyncOnCloseType) -> std::io::Result<()> {
+    pub fn sync_all(&self) -> io::Result<()> {
         match self {
-            Writeable::Dyn(d) => {
-                crate::utils::sync_on_close::sync_on_close_dyn(sync_on_close, d.as_mut())
-            },
-            Writeable::Local(file) => {
-                crate::utils::sync_on_close::sync_on_close(sync_on_close, file)
-            },
+            Self::Dyn(v) => v.sync_all(),
+            Self::Local(v) => v.sync_all(),
             #[cfg(feature = "cloud")]
-            Writeable::Cloud(_) => Ok(()),
+            Self::Cloud(v) => v.sync_all(),
         }
     }
 
-    pub fn close(self) -> std::io::Result<()> {
+    pub fn sync_data(&self) -> io::Result<()> {
+        match self {
+            Self::Dyn(v) => v.sync_data(),
+            Self::Local(v) => v.sync_data(),
+            #[cfg(feature = "cloud")]
+            Self::Cloud(v) => v.sync_data(),
+        }
+    }
+
+    pub fn close(self, sync: SyncOnCloseType) -> std::io::Result<()> {
+        match sync {
+            SyncOnCloseType::All => self.sync_all()?,
+            SyncOnCloseType::Data => self.sync_data()?,
+            SyncOnCloseType::None => {},
+        }
+
         match self {
             Self::Dyn(mut v) => v.close(),
             Self::Local(v) => close_file(v),
@@ -236,22 +247,29 @@ mod async_writeable {
             Writeable::try_new(path, cloud_options).and_then(|x| x.try_into_async_writeable())
         }
 
-        pub async fn sync_on_close(
-            &mut self,
-            sync_on_close: SyncOnCloseType,
-        ) -> std::io::Result<()> {
+        pub async fn sync_all(&mut self) -> io::Result<()> {
             match self {
-                Self::Dyn(d) => task::block_in_place(|| {
-                    crate::utils::sync_on_close::sync_on_close_dyn(sync_on_close, d.0.as_mut())
-                }),
-                Self::Local(file) => {
-                    crate::utils::sync_on_close::tokio_sync_on_close(sync_on_close, file).await
-                },
+                Self::Dyn(v) => task::block_in_place(|| v.0.as_ref().sync_all()),
+                Self::Local(v) => v.sync_all().await,
+                Self::Cloud(_) => Ok(()),
+            }
+        }
+
+        pub async fn sync_data(&mut self) -> io::Result<()> {
+            match self {
+                Self::Dyn(v) => task::block_in_place(|| v.0.as_ref().sync_data()),
+                Self::Local(v) => v.sync_data().await,
                 Self::Cloud(_) => Ok(()),
             }
         }
 
-        pub async fn close(self) -> PolarsResult<()> {
+        pub async fn close(mut self, sync: SyncOnCloseType) -> PolarsResult<()> {
+            match sync {
+                SyncOnCloseType::All => self.sync_all().await?,
+                SyncOnCloseType::Data => self.sync_data().await?,
+                SyncOnCloseType::None => {},
+            }
+
             match self {
                 Self::Dyn(mut v) => {
                     v.shutdown().await.map_err(PolarsError::from)?;

crates/polars-io/src/utils/sync_on_close.rs

@@ -1,7 +1,3 @@
-use std::{fs, io};
-
-use crate::utils::file::WriteableTrait;
-
 #[derive(Clone, Copy, PartialEq, Eq, Debug, Default, Hash)]
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
 #[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
@@ -15,34 +11,3 @@ pub enum SyncOnCloseType {
     /// Synce the file contents and the metadata.
     All,
 }
-
-pub fn sync_on_close(sync_on_close: SyncOnCloseType, file: &mut fs::File) -> io::Result<()> {
-    match sync_on_close {
-        SyncOnCloseType::None => Ok(()),
-        SyncOnCloseType::Data => file.sync_data(),
-        SyncOnCloseType::All => file.sync_all(),
-    }
-}
-
-pub fn sync_on_close_dyn(
-    sync_on_close: SyncOnCloseType,
-    file: &mut dyn WriteableTrait,
-) -> io::Result<()> {
-    match sync_on_close {
-        SyncOnCloseType::None => Ok(()),
-        SyncOnCloseType::Data => file.sync_data(),
-        SyncOnCloseType::All => file.sync_all(),
-    }
-}
-
-#[cfg(feature = "tokio")]
-pub async fn tokio_sync_on_close(
-    sync_on_close: SyncOnCloseType,
-    file: &mut tokio::fs::File,
-) -> io::Result<()> {
-    match sync_on_close {
-        SyncOnCloseType::None => Ok(()),
-        SyncOnCloseType::Data => file.sync_data().await,
-        SyncOnCloseType::All => file.sync_all().await,
-    }
-}

crates/polars-mem-engine/src/planner/lp.rs

@@ -372,8 +372,7 @@ fn create_physical_plan_impl(
                                 _ => panic!("enable filetype feature"),
                             }
 
-                            file.sync_on_close(unified_sink_args.sync_on_close)?;
-                            file.close()?;
+                            file.close(unified_sink_args.sync_on_close)?;
 
                             Ok(None)
                         }),

crates/polars-stream/src/nodes/io_sinks/csv.rs

@@ -110,8 +110,7 @@ impl SinkNode for CsvSinkNode {
                 }
             }
 
-            file.sync_on_close(sink_options.sync_on_close).await?;
-            file.close().await?;
+            file.close(sink_options.sync_on_close).await?;
 
             PolarsResult::Ok(())
         });

crates/polars-stream/src/nodes/io_sinks/ipc.rs

@@ -169,8 +169,7 @@ impl SinkNode for IpcSinkNode {
             writer.finish()?;
             drop(writer);
 
-            file.sync_on_close(sink_options.sync_on_close)?;
-            file.close()?;
+            file.close(sink_options.sync_on_close)?;
 
             PolarsResult::Ok(())
         });

crates/polars-stream/src/nodes/io_sinks/json.rs

@@ -80,8 +80,7 @@ impl SinkNode for NDJsonSinkNode {
                 }
             }
 
-            file.sync_on_close(sink_options.sync_on_close).await?;
-            file.close().await?;
+            file.close(sink_options.sync_on_close).await?;
 
             PolarsResult::Ok(())
         });

crates/polars-stream/src/nodes/io_sinks/parquet.rs

@@ -163,8 +163,7 @@ impl SinkNode for ParquetSinkNode {
             let file_size = writer.finish()?;
             drop(writer);
 
-            file.sync_on_close(sink_options.sync_on_close)?;
-            file.close()?;
+            file.close(sink_options.sync_on_close)?;
 
             output_file_size.store(file_size);
             PolarsResult::Ok(())