dedupe/archive_io.rs
1//! Reading and writing whole archived sample files.
2//!
3//! Shared by the `recompress` and `dedupe` tools. Both read an archived file fully
4//! into memory with [`read_archive`] and write files in the current format with
5//! [`write_archive`], which uses the same schema and writer settings as the archiver.
6
7use std::{fs, path::Path, sync::Arc};
8
9use anyhow::{Result, anyhow, bail};
10use parquet::{
11 column::{reader::get_typed_column_reader, writer::ColumnWriter},
12 data_type::{DataType, DoubleType, Int64Type},
13 file::{
14 metadata::KeyValue,
15 reader::{FileReader, RowGroupReader, SerializedFileReader},
16 writer::SerializedFileWriter,
17 },
18 schema::parser::parse_message_type,
19};
20
21use crate::archive_format::{SAMPLE_SCHEMA, writer_properties};
22
23/// Rows per row group in rewritten files (1,048,576); matches the archiver.
24///
25/// Old files may have many small row groups. Merging them into large groups gives
26/// zstd and delta encoding longer runs to work on.
27pub const ROWS_PER_ROW_GROUP: usize = 1 << 20;
28/// Contents of one archived file, held fully in memory.
29pub struct Archive {
30 /// `timestamp_unix_ns` column, in file order (Unix nanoseconds).
31 pub timestamps: Vec<i64>,
32 /// `value` column, in file order, one per timestamp.
33 pub values: Vec<f64>,
34 /// File-level key-value metadata, such as the `calibration` JSON document.
35 pub metadata: Vec<KeyValue>,
36}
37
38/// Reads every value of one column in a row group.
39///
40/// Records are read in batches of up to the row group's row count until the reader
41/// returns none. The columns are `REQUIRED`, so no definition or repetition levels are
42/// requested.
43///
44/// # Arguments
45///
46/// * `row_group` - Row group to read from.
47/// * `index` - Column index (0 for `timestamp_unix_ns`, 1 for `value`).
48///
49/// # Returns
50///
51/// The column's values in file order.
52///
53/// # Errors
54///
55/// Returns an error if the column reader cannot be created or a page cannot be read
56/// or decoded.
57///
58/// # Panics
59///
60/// Panics if the physical type of column `index` does not match `T`.
61fn read_column<T: DataType>(row_group: &dyn RowGroupReader, index: usize) -> Result<Vec<T::T>> {
62 let rows = usize::try_from(row_group.metadata().num_rows()).unwrap_or(0);
63 let mut reader = get_typed_column_reader::<T>(row_group.get_column_reader(index)?);
64 let mut values = Vec::with_capacity(rows);
65 loop {
66 let (records, _, _) = reader.read_records(rows.max(1), None, None, &mut values)?;
67 if records == 0 {
68 break;
69 }
70 }
71 Ok(values)
72}
73
74/// Reads every row and the key-value metadata of an archived file.
75///
76/// Row groups are concatenated in order. A file without key-value metadata gives an
77/// empty `metadata` list.
78///
79/// # Arguments
80///
81/// * `reader` - Open reader for the file.
82///
83/// # Errors
84///
85/// Returns an error if the columns are not exactly `timestamp_unix_ns` and `value`, if
86/// a row group cannot be read, or if a row group's timestamp count, value count and
87/// row count disagree.
88pub fn read_archive(reader: &SerializedFileReader<fs::File>) -> Result<Archive> {
89 let schema = reader.metadata().file_metadata().schema_descr();
90 let names: Vec<&str> = schema.columns().iter().map(|c| c.name()).collect();
91 if names != ["timestamp_unix_ns", "value"] {
92 bail!("unexpected columns {names:?}");
93 }
94 let mut timestamps = Vec::new();
95 let mut values = Vec::new();
96 for i in 0..reader.num_row_groups() {
97 let rg = reader.get_row_group(i)?;
98 let ts = read_column::<Int64Type>(rg.as_ref(), 0)?;
99 let vs = read_column::<DoubleType>(rg.as_ref(), 1)?;
100 if ts.len() != vs.len() || ts.len() as i64 != rg.metadata().num_rows() {
101 bail!(
102 "row group {i}: {} timestamps, {} values",
103 ts.len(),
104 vs.len()
105 );
106 }
107 timestamps.extend(ts);
108 values.extend(vs);
109 }
110 let metadata = reader
111 .metadata()
112 .file_metadata()
113 .key_value_metadata()
114 .cloned()
115 .unwrap_or_default();
116 Ok(Archive {
117 timestamps,
118 values,
119 metadata,
120 })
121}
122
123/// Writes rows in the current archive format and syncs the file to disk.
124///
125/// Uses [`SAMPLE_SCHEMA`] and [`writer_properties`] (the same settings as the
126/// archiver) with the archive's original key-value metadata. Rows are split into row
127/// groups of [`ROWS_PER_ROW_GROUP`]. An existing file at `path` is truncated.
128///
129/// # Arguments
130///
131/// * `path` - Destination file, normally the temporary `*.recompress-tmp` path.
132/// * `archive` - Rows and metadata to write.
133///
134/// # Errors
135///
136/// Returns an error if the schema cannot be parsed, or if creating, writing, closing
137/// or syncing the file fails.
138pub fn write_archive(path: &Path, archive: &Archive) -> Result<()> {
139 let schema = Arc::new(parse_message_type(SAMPLE_SCHEMA)?);
140 let props = Arc::new(writer_properties(archive.metadata.clone()));
141 let mut writer = SerializedFileWriter::new(fs::File::create(path)?, schema, props)?;
142 for (ts, vs) in archive
143 .timestamps
144 .chunks(ROWS_PER_ROW_GROUP)
145 .zip(archive.values.chunks(ROWS_PER_ROW_GROUP))
146 {
147 let mut rg = writer.next_row_group()?;
148 let mut col = rg
149 .next_column()?
150 .ok_or_else(|| anyhow!("missing column 0"))?;
151 if let ColumnWriter::Int64ColumnWriter(w) = col.untyped() {
152 w.write_batch(ts, None, None)?;
153 }
154 col.close()?;
155 let mut col = rg
156 .next_column()?
157 .ok_or_else(|| anyhow!("missing column 1"))?;
158 if let ColumnWriter::DoubleColumnWriter(w) = col.untyped() {
159 w.write_batch(vs, None, None)?;
160 }
161 col.close()?;
162 rg.close()?;
163 }
164 writer.into_inner()?.sync_all()?;
165 Ok(())
166}