fix n_chunks_remaining when using --overwrite

paradigmxyz · Aug 29, 2023 · 451fb23 · 451fb23
1 parent 0d4501e
commit 451fb23
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/crates/freeze/src/types/files.rs b/crates/freeze/src/types/files.rs
@@ -1,7 +1,7 @@
 use polars::prelude::*;
 
 /// Options for file output
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct FileOutput {
     /// Path of directory where to save files
     pub output_dir: std::path::PathBuf,
@@ -22,7 +22,7 @@ pub struct FileOutput {
 }
 
 /// File format
-#[derive(Clone, Eq, PartialEq)]
+#[derive(Clone, Eq, PartialEq, Debug)]
 pub enum FileFormat {
     /// Parquet file format
     Parquet,

diff --git a/crates/freeze/src/types/queries.rs b/crates/freeze/src/types/queries.rs
@@ -37,6 +37,9 @@ pub struct MultiQuery {
 impl MultiQuery {
     /// get number of chunks that have not yet been collected
     pub fn get_n_chunks_remaining(&self, sink: &FileOutput) -> Result<u64, FreezeError> {
+        if sink.overwrite {
+            return Ok(self.chunks.len() as u64)
+        };
         let actual_files: HashSet<PathBuf> = list_files(&sink.output_dir)
             .map_err(|_e| {
                 FreezeError::CollectError(CollectError::CollectError(