1use 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
55const TMP_SUFFIX: &str = ".recompress-tmp";
60
61#[derive(Debug, PartialEq, Eq)]
63enum Outcome {
64 Rewritten {
67 bytes_before: u64,
69 bytes_after: u64,
71 rows: usize,
73 },
74 AlreadyCurrent,
76 Empty,
80}
81
82fn 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
109fn 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
149fn 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(©) {
197 bail!("rewritten file is not in the current format");
198 }
199 verify_identical(&original, &read_archive(©)?)?;
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
225fn 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
263fn 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 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 #[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; 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 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 #[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 #[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}