Skip to main content

meathook/
encode.rs

1//! Parquet encoding via `serde_arrow`: records stay plain
2//! `#[derive(Serialize)]` structs, no manual arrow builders.
3
4use arrow::datatypes::FieldRef;
5use parquet::arrow::ArrowWriter;
6use parquet::errors::ParquetError;
7use serde::Serialize;
8use serde::de::DeserializeOwned;
9use serde_arrow::schema::{SchemaLike, TracingOptions};
10
11/// Error encoding records to parquet.
12#[derive(Debug, thiserror::Error)]
13pub enum EncodeError {
14    #[error("failed to derive arrow schema from record type: {0}")]
15    Schema(#[source] serde_arrow::Error),
16    #[error("failed to build record batch: {0}")]
17    Batch(#[source] serde_arrow::Error),
18    #[error("failed to write parquet: {0}")]
19    Parquet(#[from] ParquetError),
20}
21
22/// Encode records into a parquet file held in memory.
23///
24/// The arrow schema is derived from `R` itself (not sampled from values, so
25/// an empty slice still produces a valid zero-row file), which is why
26/// `DeserializeOwned` is required alongside `Serialize`.
27///
28/// # Errors
29///
30/// Returns an error if schema derivation, record batch construction, or
31/// parquet writing fails.
32pub fn to_parquet<R: Serialize + DeserializeOwned>(records: &[R]) -> Result<Vec<u8>, EncodeError> {
33    let fields =
34        Vec::<FieldRef>::from_type::<R>(TracingOptions::default()).map_err(EncodeError::Schema)?;
35    let batch = serde_arrow::to_record_batch(&fields, &records).map_err(EncodeError::Batch)?;
36
37    let mut buf = vec![];
38    let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), None)?;
39    writer.write(&batch)?;
40    writer.close()?;
41    Ok(buf)
42}
43
44#[cfg(test)]
45mod tests {
46    use super::*;
47    use arrow::array::RecordBatch;
48    use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
49    use serde::Deserialize;
50
51    #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
52    struct Sample {
53        station_id: String,
54        timestamp: String,
55        value: f64,
56    }
57
58    #[test]
59    fn parquet_round_trip() {
60        let records = vec![
61            Sample {
62                station_id: "S100".into(),
63                timestamp: "2026-06-12T08:00:00+08:00".into(),
64                value: 29.4,
65            },
66            Sample {
67                station_id: "S117".into(),
68                timestamp: "2026-06-12T08:00:00+08:00".into(),
69                value: 30.1,
70            },
71        ];
72
73        let bytes = to_parquet(&records).unwrap();
74
75        let reader = ParquetRecordBatchReaderBuilder::try_new(bytes::Bytes::from(bytes))
76            .unwrap()
77            .build()
78            .unwrap();
79        let batches: Vec<_> = reader.collect::<Result<_, _>>().unwrap();
80        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
81
82        let round_tripped = serde_arrow::from_record_batch::<Vec<Sample>>(&batches[0]).unwrap();
83
84        assert_eq!(round_tripped, records);
85    }
86
87    #[test]
88    fn empty_slice_encodes_zero_row_file() {
89        let bytes = to_parquet::<Sample>(&[]).unwrap();
90        let reader = ParquetRecordBatchReaderBuilder::try_new(bytes::Bytes::from(bytes))
91            .unwrap()
92            .build()
93            .unwrap();
94        let rows: usize = reader.map(|b| b.unwrap().num_rows()).sum();
95        assert_eq!(rows, 0);
96    }
97}