Skip to main content

recompress/
recompress.rs

1//! Offline tool that rewrites older archived Parquet files in the current archive format.
2//!
3//! The archiver writes one Parquet file per channel and rotation window under its
4//! Parquet root. Files written before the archiver switched to zstd with
5//! delta-encoded timestamps (see `archive_format::writer_properties`) are much larger
6//! than they need to be. This tool walks an archive tree and rewrites those files in
7//! place without changing a single stored value. It does not talk to NATS.
8//!
9//! # Usage
10//!
11//! ```text
12//! recompress <parquet_root> [--dry-run]
13//! ```
14//!
15//! * `<parquet_root>` - Directory to scan recursively. The first argument that does
16//!   not start with `--` is used. Required.
17//! * `--dry-run` - Write and verify each rewritten copy, then delete it and leave the
18//!   original in place. Reports the sizes a real run would produce.
19//!
20//! Other `--` arguments are ignored. No environment variables are read. The tool
21//! prints a progress line every 1000 files and a summary at the end, and exits with an
22//! error if any file failed.
23//!
24//! # Design
25//!
26//! Each `part-*.parquet` file that is not already zstd with delta-encoded
27//! timestamps is rewritten next to itself as `*.parquet.recompress-tmp`,
28//! synced, read back and compared value by value (bit for bit, so NaN counts
29//! as equal only to the identical NaN) together with its key-value metadata.
30//! Only then is the original replaced by an atomic rename, and the directory
31//! is synced, so every file is at all times either the complete old version or
32//! the verified new one. A failed file keeps its original.
33//!
34//! Unfinished (`.inprogress`) and quarantined files are never touched, and files
35//! with no row groups are counted as empty and left alone. The tool can be
36//! interrupted and rerun: finished files are skipped, and temporary copies left by
37//! an interrupted run are deleted on the next scan.
38
39use std::{
40    fs,
41    path::{Path, PathBuf},
42};
43
44use anyhow::{Context, Result, anyhow, bail};
45use parquet::{
46    basic::{Compression, Encoding},
47    file::reader::{FileReader, SerializedFileReader},
48};
49
50mod archive_format;
51mod archive_io;
52
53use archive_io::{Archive, read_archive, write_archive};
54
55/// Suffix appended to a file's full name for its temporary rewritten copy.
56///
57/// Any file with this suffix found during the scan is left over from an interrupted
58/// run and is deleted by [`collect_files`].
59const TMP_SUFFIX: &str = ".recompress-tmp";
60
61/// Result of processing one file with [`recompress_file`].
62#[derive(Debug, PartialEq, Eq)]
63enum Outcome {
64    /// The file was rewritten and verified (or, in a dry run, a verified copy was
65    /// written and then deleted).
66    Rewritten {
67        /// Size of the original file, in bytes.
68        bytes_before: u64,
69        /// Size of the verified rewritten copy, in bytes.
70        bytes_after: u64,
71        /// Number of rows in the file.
72        rows: usize,
73    },
74    /// The file already uses the current format and was left untouched.
75    AlreadyCurrent,
76    /// The file has no row groups, so there is nothing to compress. It is left
77    /// untouched. The previous archiver wrote such files when a rotation fired while
78    /// no data was arriving.
79    Empty,
80}
81
82/// Returns whether a file already uses the current archive format.
83///
84/// A file is current when it has at least one row group and every row group has
85/// exactly two columns, both zstd-compressed, with the timestamp column (column 0)
86/// using `DELTA_BINARY_PACKED` encoding. Only the footer metadata is inspected; no data
87/// pages are read.
88///
89/// A file with no row groups is never current; [`recompress_file`] reports it as
90/// [`Outcome::Empty`] before calling this.
91///
92/// # Arguments
93///
94/// * `reader` - Open reader for the file.
95fn is_current_format(reader: &SerializedFileReader<fs::File>) -> bool {
96    let meta = reader.metadata();
97    meta.num_row_groups() > 0
98        && (0..meta.num_row_groups()).all(|i| {
99            let rg = meta.row_group(i);
100            rg.num_columns() == 2
101                && (0..2).all(|c| matches!(rg.column(c).compression(), Compression::ZSTD(_)))
102                && rg
103                    .column(0)
104                    .encodings()
105                    .contains(&Encoding::DELTA_BINARY_PACKED)
106        })
107}
108
109/// Checks that two archives hold exactly the same rows and metadata.
110///
111/// Timestamps must be equal in order. Values are compared by their bit patterns
112/// (`f64::to_bits`), so a NaN matches only an identical NaN and `0.0` does not match
113/// `-0.0`. Key-value metadata must match key by key and value by value, in order.
114///
115/// # Arguments
116///
117/// * `original` - Archive read from the original file.
118/// * `rewritten` - Archive read back from the rewritten copy.
119///
120/// # Errors
121///
122/// Returns an error naming the first part that differs: timestamps, values or
123/// key-value metadata.
124fn verify_identical(original: &Archive, rewritten: &Archive) -> Result<()> {
125    if original.timestamps != rewritten.timestamps {
126        bail!("timestamps differ");
127    }
128    if original.values.len() != rewritten.values.len()
129        || original
130            .values
131            .iter()
132            .zip(&rewritten.values)
133            .any(|(a, b)| a.to_bits() != b.to_bits())
134    {
135        bail!("values differ");
136    }
137    let kv = |a: &Archive| {
138        a.metadata
139            .iter()
140            .map(|k| (k.key.clone(), k.value.clone()))
141            .collect::<Vec<_>>()
142    };
143    if kv(original) != kv(rewritten) {
144        bail!("key-value metadata differs");
145    }
146    Ok(())
147}
148
149/// Rewrites one file in the current format after verifying the copy.
150///
151/// Steps, in order:
152///
153/// 1. If the file has no row groups, return [`Outcome::Empty`]. If it is already
154///    current, return [`Outcome::AlreadyCurrent`]. Neither touches the file.
155/// 2. Read all rows and metadata into memory.
156/// 3. Write them to `<path>.recompress-tmp` with [`write_archive`] (which syncs it).
157/// 4. Reopen the copy, check it is in the current format, and compare it with the
158///    original using [`verify_identical`].
159/// 5. In a dry run, delete the copy. Otherwise rename it over the original and sync
160///    the parent directory so the rename is durable.
161///
162/// If any step from 3 to 4 fails, the temporary copy is deleted (best effort) and the
163/// original is left as it was.
164///
165/// # Arguments
166///
167/// * `path` - Path of a finished `part-*.parquet` file.
168/// * `dry_run` - When `true`, verify the rewrite but keep the original file.
169///
170/// # Returns
171///
172/// [`Outcome::Rewritten`] with the before and after sizes and row count (also in a dry
173/// run), [`Outcome::AlreadyCurrent`] or [`Outcome::Empty`].
174///
175/// # Errors
176///
177/// Returns an error if the file cannot be opened or read, if writing or verifying the
178/// copy fails, or if deleting the copy (dry run), renaming it or syncing the directory
179/// fails.
180fn recompress_file(path: &Path, dry_run: bool) -> Result<Outcome> {
181    let reader = SerializedFileReader::new(fs::File::open(path)?)?;
182    if reader.num_row_groups() == 0 {
183        return Ok(Outcome::Empty);
184    }
185    if is_current_format(&reader) {
186        return Ok(Outcome::AlreadyCurrent);
187    }
188    let original = read_archive(&reader)?;
189    let bytes_before = fs::metadata(path)?.len();
190    let rows = original.timestamps.len();
191
192    let tmp = PathBuf::from(format!("{}{TMP_SUFFIX}", path.display()));
193    let result = (|| -> Result<u64> {
194        write_archive(&tmp, &original)?;
195        let copy = SerializedFileReader::new(fs::File::open(&tmp)?)?;
196        if !is_current_format(&copy) {
197            bail!("rewritten file is not in the current format");
198        }
199        verify_identical(&original, &read_archive(&copy)?)?;
200        Ok(fs::metadata(&tmp)?.len())
201    })();
202    let bytes_after = match result {
203        Ok(size) => size,
204        Err(err) => {
205            let _ = fs::remove_file(&tmp);
206            return Err(err);
207        }
208    };
209
210    if dry_run {
211        fs::remove_file(&tmp)?;
212    } else {
213        fs::rename(&tmp, path)?;
214        if let Some(dir) = path.parent() {
215            fs::File::open(dir)?.sync_all()?;
216        }
217    }
218    Ok(Outcome::Rewritten {
219        bytes_before,
220        bytes_after,
221        rows,
222    })
223}
224
225/// Lists archived `part-*.parquet` files and removes stale temporary copies.
226///
227/// Walks `root` recursively. Files ending in [`TMP_SUFFIX`] are deleted and reported,
228/// since they can only be left over from an interrupted run. Files whose names start
229/// with `part-` and end with `.parquet` are appended to `files`. Everything else,
230/// including `.inprogress` and quarantined files, is ignored. Symbolic links to
231/// directories are not followed, and names that are not valid UTF-8 are skipped.
232///
233/// # Arguments
234///
235/// * `root` - Directory to scan.
236/// * `files` - Output list; matching paths are appended in directory order.
237///
238/// # Errors
239///
240/// Returns an error if a directory cannot be read, a file type cannot be determined,
241/// or a stale temporary file cannot be deleted.
242fn collect_files(root: &Path, files: &mut Vec<PathBuf>) -> Result<()> {
243    for entry in fs::read_dir(root).with_context(|| format!("reading {}", root.display()))? {
244        let entry = entry?;
245        let path = entry.path();
246        if entry.file_type()?.is_dir() {
247            collect_files(&path, files)?;
248            continue;
249        }
250        let Some(name) = path.file_name().and_then(|n| n.to_str()) else {
251            continue;
252        };
253        if name.ends_with(TMP_SUFFIX) {
254            fs::remove_file(&path)?;
255            println!("removed stale {}", path.display());
256        } else if name.starts_with("part-") && name.ends_with(".parquet") {
257            files.push(path);
258        }
259    }
260    Ok(())
261}
262
263/// Parses the command line, recompresses every archived file under the root and
264/// prints a summary.
265///
266/// Files are processed one at a time in sorted path order. A failure on one file is
267/// printed to stderr and does not stop the run.
268///
269/// # Errors
270///
271/// Returns an error if no `<parquet_root>` argument is given, if scanning the tree
272/// fails, or if any file failed (the originals of failed files are left untouched).
273fn main() -> Result<()> {
274    let args: Vec<String> = std::env::args().skip(1).collect();
275    let root = args
276        .iter()
277        .find(|a| !a.starts_with("--"))
278        .ok_or_else(|| anyhow!("usage: recompress <parquet_root> [--dry-run]"))?;
279    let dry_run = args.iter().any(|a| a == "--dry-run");
280
281    let mut files = Vec::new();
282    collect_files(Path::new(root), &mut files)?;
283    files.sort();
284    println!(
285        "{} archived files under {root}{}",
286        files.len(),
287        if dry_run { " (dry run)" } else { "" }
288    );
289
290    let (mut rewritten, mut current, mut empty, mut failed) = (0usize, 0usize, 0usize, 0usize);
291    let (mut before, mut after, mut rows) = (0u64, 0u64, 0usize);
292    for (i, path) in files.iter().enumerate() {
293        match recompress_file(path, dry_run) {
294            Ok(Outcome::Rewritten {
295                bytes_before,
296                bytes_after,
297                rows: n,
298            }) => {
299                rewritten += 1;
300                before += bytes_before;
301                after += bytes_after;
302                rows += n;
303            }
304            Ok(Outcome::AlreadyCurrent) => current += 1,
305            Ok(Outcome::Empty) => empty += 1,
306            Err(err) => {
307                failed += 1;
308                eprintln!("FAILED {}: {err:#}", path.display());
309            }
310        }
311        if (i + 1) % 1000 == 0 {
312            println!(
313                "{}/{} files; rewritten {rewritten}, {:.2} GB -> {:.2} GB",
314                i + 1,
315                files.len(),
316                before as f64 / 1e9,
317                after as f64 / 1e9
318            );
319        }
320    }
321    println!(
322        "done: rewritten {rewritten} ({rows} rows, {:.2} GB -> {:.2} GB), already current {current}, empty {empty}, failed {failed}",
323        before as f64 / 1e9,
324        after as f64 / 1e9
325    );
326    if failed > 0 {
327        bail!("{failed} file(s) failed; their originals were left untouched");
328    }
329    Ok(())
330}
331
332#[cfg(test)]
333mod tests {
334    use super::*;
335    use archive_format::SAMPLE_SCHEMA;
336    use parquet::{
337        column::writer::ColumnWriter,
338        file::{metadata::KeyValue, properties::WriterProperties, writer::SerializedFileWriter},
339        schema::parser::parse_message_type,
340    };
341    use std::sync::Arc;
342
343    /// Writes a file the way the archiver did before zstd: plain, many row groups.
344    fn write_old_format(path: &Path, groups: &[Vec<(i64, f64)>], calibration: &str) {
345        let schema = Arc::new(parse_message_type(SAMPLE_SCHEMA).unwrap());
346        let props = Arc::new(
347            WriterProperties::builder()
348                .set_key_value_metadata(Some(vec![KeyValue::new(
349                    "calibration".to_string(),
350                    calibration.to_string(),
351                )]))
352                .build(),
353        );
354        let mut writer =
355            SerializedFileWriter::new(fs::File::create(path).unwrap(), schema, props).unwrap();
356        for group in groups {
357            let mut rg = writer.next_row_group().unwrap();
358            let ts: Vec<i64> = group.iter().map(|r| r.0).collect();
359            let vs: Vec<f64> = group.iter().map(|r| r.1).collect();
360            let mut col = rg.next_column().unwrap().unwrap();
361            if let ColumnWriter::Int64ColumnWriter(w) = col.untyped() {
362                w.write_batch(&ts, None, None).unwrap();
363            }
364            col.close().unwrap();
365            let mut col = rg.next_column().unwrap().unwrap();
366            if let ColumnWriter::DoubleColumnWriter(w) = col.untyped() {
367                w.write_batch(&vs, None, None).unwrap();
368            }
369            col.close().unwrap();
370            rg.close().unwrap();
371        }
372        writer.close().unwrap();
373    }
374
375    /// An old-format file survives a dry run unchanged, is then rewritten into one
376    /// row group with identical contents, and is skipped on the next run.
377    #[test]
378    fn old_file_is_rewritten_identically_then_skipped() {
379        let dir = std::env::temp_dir().join(format!("recompress-test-{}", uuid::Uuid::new_v4()));
380        let day = dir.join("asset1001/2026-09-01/ch08");
381        fs::create_dir_all(&day).unwrap();
382        let path = day.join("part-0001.parquet");
383        let base = 1_788_220_800_000_000_000_i64;
384        let groups: Vec<Vec<(i64, f64)>> = (0..5)
385            .map(|g| {
386                (0..1000)
387                    .map(|i| {
388                        let k = g * 1000 + i;
389                        (
390                            base + k * 500_000,
391                            3.7 + ((k * 7919) % 97) as f64 * 0.000_315_6,
392                        )
393                    })
394                    .collect()
395            })
396            .collect();
397        let mut groups = groups;
398        groups[2][10].1 = f64::NAN;
399        groups[4][999].0 -= 3_000_000_000; // an out-of-order timestamp survives as-is
400        let calibration = r#"{"id":"tp3505","type":"linear","a":70.2,"b":-9.1}"#;
401        write_old_format(&path, &groups, calibration);
402        let original =
403            read_archive(&SerializedFileReader::new(fs::File::open(&path).unwrap()).unwrap())
404                .unwrap();
405
406        // A dry run verifies but leaves the original in place.
407        let dry = recompress_file(&path, true).unwrap();
408        assert!(matches!(dry, Outcome::Rewritten { rows: 5000, .. }));
409        let reader = SerializedFileReader::new(fs::File::open(&path).unwrap()).unwrap();
410        assert!(!is_current_format(&reader));
411
412        let outcome = recompress_file(&path, false).unwrap();
413        let Outcome::Rewritten {
414            bytes_before,
415            bytes_after,
416            rows,
417        } = outcome
418        else {
419            panic!("expected a rewrite");
420        };
421        assert_eq!(rows, 5000);
422        assert!(bytes_after < bytes_before);
423
424        let reader = SerializedFileReader::new(fs::File::open(&path).unwrap()).unwrap();
425        assert!(is_current_format(&reader));
426        assert_eq!(reader.num_row_groups(), 1);
427        verify_identical(&original, &read_archive(&reader).unwrap()).unwrap();
428        assert!(!PathBuf::from(format!("{}{TMP_SUFFIX}", path.display())).exists());
429
430        assert_eq!(
431            recompress_file(&path, false).unwrap(),
432            Outcome::AlreadyCurrent
433        );
434        fs::remove_dir_all(dir).unwrap();
435    }
436
437    /// Verification rejects any change to timestamps, values or metadata.
438    /// Checks that a file with no row groups is reported as empty and left untouched.
439    #[test]
440    fn empty_file_is_skipped_and_untouched() {
441        let dir = std::env::temp_dir().join(format!("recompress-empty-{}", uuid::Uuid::new_v4()));
442        fs::create_dir_all(&dir).unwrap();
443        let path = dir.join("part-0001.parquet");
444        write_old_format(&path, &[], r#"{"type":"identity"}"#);
445        let before = fs::read(&path).unwrap();
446        assert_eq!(recompress_file(&path, false).unwrap(), Outcome::Empty);
447        assert_eq!(fs::read(&path).unwrap(), before);
448        assert!(!PathBuf::from(format!("{}{TMP_SUFFIX}", path.display())).exists());
449        fs::remove_dir_all(dir).unwrap();
450    }
451
452    #[test]
453    fn verification_detects_any_difference() {
454        let a = Archive {
455            timestamps: vec![1, 2, 3],
456            values: vec![1.0, f64::NAN, 3.0],
457            metadata: vec![KeyValue::new("calibration".into(), "{}".to_string())],
458        };
459        let same = Archive {
460            timestamps: a.timestamps.clone(),
461            values: a.values.clone(),
462            metadata: a.metadata.clone(),
463        };
464        verify_identical(&a, &same).unwrap();
465        let mut b = Archive {
466            timestamps: vec![1, 2, 4],
467            ..same
468        };
469        assert!(verify_identical(&a, &b).is_err());
470        b.timestamps = a.timestamps.clone();
471        b.values[2] = 3.000_000_000_000_001;
472        assert!(verify_identical(&a, &b).is_err());
473        b.values = a.values.clone();
474        b.metadata = vec![KeyValue::new("calibration".into(), "{ }".to_string())];
475        assert!(verify_identical(&a, &b).is_err());
476    }
477
478    /// The scan deletes stale temporary copies and returns only finished part files.
479    #[test]
480    fn stale_temporary_files_are_removed_and_others_ignored() {
481        let dir = std::env::temp_dir().join(format!("recompress-walk-{}", uuid::Uuid::new_v4()));
482        let day = dir.join("asset1001/2026-09-01/ch08");
483        fs::create_dir_all(&day).unwrap();
484        for name in [
485            "part-0001.parquet",
486            "part-0002.parquet.inprogress",
487            "part-0003.parquet.inprogress.unfinished.quarantined-1-0",
488            "part-0004.parquet.recompress-tmp",
489            "notes.txt",
490        ] {
491            fs::write(day.join(name), b"x").unwrap();
492        }
493        let mut files = Vec::new();
494        collect_files(&dir, &mut files).unwrap();
495        assert_eq!(files, vec![day.join("part-0001.parquet")]);
496        assert!(!day.join("part-0004.parquet.recompress-tmp").exists());
497        assert!(day.join("part-0002.parquet.inprogress").exists());
498        fs::remove_dir_all(dir).unwrap();
499    }
500}