diff --git a/CHANGELOG.md b/CHANGELOG.md index 4a1000b..a15e4d0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,24 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- Opt-in, type-level zstd parquet compression through + `ParquetEncoder>`, with `LEVEL` restricted at compile time to + `1..=22`. `Zstd` defaults to level 1 and `HfSink` remains uncompressed by + default. + +### Changed + +- **Breaking:** `ParquetEncoder` is now generic over a sealed + `ParquetCompression` policy and is no longer a unit value. Migrate + `ParquetEncoder.encode(&records)` to + `ParquetEncoder::default().encode(&records)` for the previous uncompressed + behavior. +- Changing compression changes the encoded-byte fingerprint in Hugging Face + object paths. Drain pending JSONL spool segments before changing compression + if replaying the same logical rows to a second path would be unacceptable. + ## [0.3.0](https://github.com/InfiniteUnion/meathook-rs/compare/v0.2.0...v0.3.0) - 2026-07-12 ### Added diff --git a/README.md b/README.md index d8dbd88..90bf9a3 100644 --- a/README.md +++ b/README.md @@ -149,6 +149,14 @@ Replaying a window produces the same bytes and overwrites the same file. Different payloads do not collide, including sub-hourly windows and windows split into chunks by `max_records`. +The compression setting is therefore part of replay identity. Before changing +an existing deployment from uncompressed parquet to zstd (or changing zstd +levels), drain its JSONL spool if duplicate rows are unacceptable. A segment +uploaded before a crash but replayed afterward with different compression has +the same logical records but different bytes, so it lands at a new content-hash +path. Existing uncompressed files remain valid and can coexist with new zstd +files. + The position of `JsonlStore` matters. Its protection starts when a batch reaches that tier and ends when the downstream sink accepts it. An outer `MemStore` reduces fsync traffic but leaves its current window in memory. A @@ -180,6 +188,26 @@ default, `JsonEncoder` is always available, and the `csv` feature adds `CsvEncoder`. Files use Hive-style partitions so the Hugging Face dataset viewer can read them. +Parquet stays uncompressed by default for compatibility. Compression is part +of the encoder type: use `Zstd` for level 1, or select a level from 1 through +22 with `Zstd`. Values outside that range do not implement the parquet +compression policy and fail to compile: + +```rust +use meathook::{HfSink, ParquetEncoder, Zstd}; + +type DatasetParquet = ParquetEncoder>; + +let sink = HfSink::::new(client, repo, token) + .encoder(DatasetParquet::new()); + +let uncompressed = ParquetEncoder::default(); +let zstd_level_1 = ParquetEncoder::::new(); +``` + +Compression is recorded per parquet column and is transparent to the Hugging +Face dataset viewer and parquet readers. + The sink retries transport errors, HTTP 429 responses, and 5xx responses with backoff. If those retries run out, the upstream tier keeps the records and tries again when it next flushes. @@ -195,7 +223,7 @@ forever. | Feature | Default | Implies | Adds | |---|---|---|---| -| `parquet` | Yes | Nothing | `Encoder` and `ParquetEncoder` using Arrow, Parquet, and serde_arrow | +| `parquet` | Yes | Nothing | `Encoder` and configurable uncompressed/zstd `ParquetEncoder` using Arrow, Parquet, and serde_arrow | | `csv` | No | Nothing | `CsvEncoder` | | `satay` | No | Nothing | `SatayCollector` for satay-generated API clients | | `huggingface` | Yes | `parquet`, `satay` | `HfSink` and the sans-IO `CommitAction` | diff --git a/src/encode.rs b/src/encode.rs index 29d0a6a..6dc174d 100644 --- a/src/encode.rs +++ b/src/encode.rs @@ -7,13 +7,19 @@ //! `csv` feature. use std::error; +#[cfg(feature = "parquet")] +use std::marker::PhantomData; #[cfg(feature = "parquet")] use arrow::datatypes::FieldRef; #[cfg(feature = "parquet")] use parquet::arrow::ArrowWriter; #[cfg(feature = "parquet")] +use parquet::basic::{Compression, ZstdLevel}; +#[cfg(feature = "parquet")] use parquet::errors::ParquetError; +#[cfg(feature = "parquet")] +use parquet::file::properties::WriterProperties; use serde::Serialize; use serde::de::DeserializeOwned; #[cfg(feature = "parquet")] @@ -70,17 +76,119 @@ pub enum ParquetEncodeError { Parquet(#[from] ParquetError), } +#[cfg(feature = "parquet")] +mod private { + pub trait Sealed {} +} + +/// A type-level parquet compression policy. +/// +/// This trait is sealed. The built-in policies are [`Uncompressed`] and +/// [`Zstd`](Zstd), where only levels `1..=22` implement this trait. +#[cfg(feature = "parquet")] +pub trait ParquetCompression: private::Sealed + Send + Sync + 'static { + #[doc(hidden)] + fn parquet_compression() -> Result; +} + +/// Type-level policy for uncompressed parquet output. +#[cfg(feature = "parquet")] +#[derive(Debug, Clone, Copy, Default)] +pub struct Uncompressed; + +#[cfg(feature = "parquet")] +impl private::Sealed for Uncompressed {} + +#[cfg(feature = "parquet")] +impl ParquetCompression for Uncompressed { + fn parquet_compression() -> Result { + Ok(Compression::UNCOMPRESSED) + } +} + +/// Type-level policy for zstd-compressed parquet output. +/// +/// `LEVEL` defaults to `1`. Only levels `1..=22` implement +/// [`ParquetCompression`], so an encoder with an invalid level cannot be +/// constructed. +/// +/// ```compile_fail +/// use meathook::{ParquetEncoder, Zstd}; +/// +/// let encoder = ParquetEncoder::>::new(); +/// ``` +#[cfg(feature = "parquet")] +#[derive(Debug, Clone, Copy, Default)] +pub struct Zstd; + +#[cfg(feature = "parquet")] +macro_rules! impl_zstd_levels { + ($($level:literal),* $(,)?) => { + $( + impl private::Sealed for Zstd<$level> {} + + impl ParquetCompression for Zstd<$level> { + fn parquet_compression() -> Result { + ZstdLevel::try_new($level).map(Compression::ZSTD) + } + } + )* + }; +} + +#[cfg(feature = "parquet")] +impl_zstd_levels!( + 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, +); + /// Encodes a window into a parquet file held in memory. /// /// The arrow schema is derived from `R` itself (not sampled from values, so /// an empty slice still produces a valid zero-row file), which is why /// `DeserializeOwned` is required alongside `Serialize`. +/// +/// The default compression policy is [`Uncompressed`]. Select zstd and its +/// level in the encoder type: +/// +/// ``` +/// use meathook::{ParquetEncoder, Zstd}; +/// +/// let uncompressed = ParquetEncoder::default(); +/// let zstd_1 = ParquetEncoder::::new(); +/// let zstd_3 = ParquetEncoder::>::new(); +/// ``` #[cfg(feature = "parquet")] -#[derive(Debug, Clone, Copy, Default)] -pub struct ParquetEncoder; +#[derive(Debug, Clone, Copy)] +pub struct ParquetEncoder { + _compression: PhantomData, +} #[cfg(feature = "parquet")] -impl Encoder for ParquetEncoder { +impl ParquetEncoder { + /// Creates an encoder using the compression policy `C`. + #[must_use] + pub const fn new() -> Self { + Self { + _compression: PhantomData, + } + } + + fn writer_properties() -> Result { + Ok(WriterProperties::builder() + .set_compression(C::parquet_compression()?) + .build()) + } +} + +#[cfg(feature = "parquet")] +impl Default for ParquetEncoder { + fn default() -> Self { + Self::new() + } +} + +#[cfg(feature = "parquet")] +impl Encoder for ParquetEncoder { type Error = ParquetEncodeError; const EXT: &'static str = "parquet"; @@ -98,7 +206,8 @@ impl Encoder for ParquetEncoder { serde_arrow::to_record_batch(&fields, &records).map_err(ParquetEncodeError::Batch)?; let mut buf = vec![]; - let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), None)?; + let mut writer = + ArrowWriter::try_new(&mut buf, batch.schema(), Some(Self::writer_properties()?))?; writer.write(&batch)?; writer.close()?; Ok(buf) @@ -176,14 +285,35 @@ mod tests { mod parquet_encoder { use arrow::array::RecordBatch; use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + use parquet::basic::{Compression, ZstdLevel}; + use parquet::file::metadata::{ColumnChunkMetaData, RowGroupMetaData}; use super::*; + fn compressions(bytes: Vec) -> Vec { + let builder = + ParquetRecordBatchReaderBuilder::try_new(bytes::Bytes::from(bytes)).unwrap(); + builder + .metadata() + .row_groups() + .iter() + .flat_map(RowGroupMetaData::columns) + .map(ColumnChunkMetaData::compression) + .collect() + } + #[test] fn parquet_round_trip() { let records = samples(); - let bytes = ParquetEncoder.encode(&records).unwrap(); + let encoder = ParquetEncoder::default(); + let bytes = encoder.encode(&records).unwrap(); + + assert!( + compressions(bytes.clone()) + .iter() + .all(|compression| *compression == Compression::UNCOMPRESSED) + ); let reader = ParquetRecordBatchReaderBuilder::try_new(bytes::Bytes::from(bytes)) .unwrap() @@ -197,15 +327,109 @@ mod tests { assert_eq!(round_tripped, records); } + #[test] + fn default_policy_uses_uncompressed_codec() { + let bytes = ParquetEncoder::default().encode(&samples()).unwrap(); + + assert!( + compressions(bytes) + .iter() + .all(|compression| *compression == Compression::UNCOMPRESSED) + ); + } + #[test] fn empty_slice_encodes_zero_row_file() { - let bytes = ParquetEncoder.encode::(&[]).unwrap(); + let encoded = [ + ParquetEncoder::default().encode::(&[]).unwrap(), + ParquetEncoder::::new().encode::(&[]).unwrap(), + ]; + + for bytes in encoded { + let reader = ParquetRecordBatchReaderBuilder::try_new(bytes::Bytes::from(bytes)) + .unwrap() + .build() + .unwrap(); + let rows: usize = reader.map(|b| b.unwrap().num_rows()).sum(); + assert_eq!(rows, 0); + } + } + + #[test] + fn zstd_default_level_round_trips() { + let records = samples(); + let bytes = ParquetEncoder::::new().encode(&records).unwrap(); + + assert!( + compressions(bytes.clone()) + .iter() + .all(|compression| { *compression == Compression::ZSTD(ZstdLevel::default()) }) + ); + let reader = ParquetRecordBatchReaderBuilder::try_new(bytes::Bytes::from(bytes)) .unwrap() .build() .unwrap(); - let rows: usize = reader.map(|b| b.unwrap().num_rows()).sum(); - assert_eq!(rows, 0); + let batches: Vec<_> = reader.collect::>().unwrap(); + let round_tripped = serde_arrow::from_record_batch::>(&batches[0]).unwrap(); + + assert_eq!(round_tripped, records); + } + + #[test] + fn zstd_explicit_level_is_applied() { + let level = ZstdLevel::try_new(7).unwrap(); + let encoder = ParquetEncoder::>::new(); + assert_eq!( + Zstd::<7>::parquet_compression().unwrap(), + Compression::ZSTD(level) + ); + + let bytes = encoder.encode(&samples()).unwrap(); + + assert!( + compressions(bytes.clone()) + .iter() + .all(|compression| matches!(compression, Compression::ZSTD(_))) + ); + ParquetRecordBatchReaderBuilder::try_new(bytes::Bytes::from(bytes)) + .unwrap() + .build() + .unwrap() + .collect::, _>>() + .unwrap(); + } + + #[test] + fn zstd_boundary_levels_are_applied() { + let policies = [ + Zstd::<1>::parquet_compression().unwrap(), + Zstd::<22>::parquet_compression().unwrap(), + ]; + + assert_eq!( + policies, + [ + Compression::ZSTD(ZstdLevel::try_new(1).unwrap()), + Compression::ZSTD(ZstdLevel::try_new(22).unwrap()), + ] + ); + } + + #[test] + fn zstd_shrinks_compressible_records() { + let records: Vec<_> = (0..4_096) + .map(|index| Sample { + station_id: format!("weather-station-{index:08}"), + timestamp: format!("2026-07-13T12:{:02}:{:02}+08:00", index / 60, index % 60), + value: 29.0 + f64::from(index % 10) / 10.0, + }) + .collect(); + + let uncompressed = ParquetEncoder::default().encode(&records).unwrap(); + let compressed = ParquetEncoder::::new().encode(&records).unwrap(); + + assert!(compressed.len() < uncompressed.len()); } } diff --git a/src/lib.rs b/src/lib.rs index 88386d5..e474c15 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -51,6 +51,6 @@ pub use store::{JsonlStore, JsonlStoreError, MemStore, Segment, Store}; pub use encode::{CsvEncoder, CsvError}; pub use encode::{Encoder, JsonEncoder}; #[cfg(feature = "parquet")] -pub use encode::{ParquetEncodeError, ParquetEncoder}; +pub use encode::{ParquetCompression, ParquetEncodeError, ParquetEncoder, Uncompressed, Zstd}; #[cfg(feature = "huggingface")] pub use sink::huggingface::{CommitGate, HfSink, HfSinkError}; diff --git a/src/sink/huggingface.rs b/src/sink/huggingface.rs index 96cac09..696643f 100644 --- a/src/sink/huggingface.rs +++ b/src/sink/huggingface.rs @@ -243,7 +243,7 @@ impl HfSink { branch: "main".to_owned(), token: token.into(), gate: None, - encoder: ParquetEncoder, + encoder: ParquetEncoder::default(), _record: PhantomData, } } @@ -491,7 +491,7 @@ mod tests { end: datetime!(2026-06-12 09:00 UTC), }; assert_eq!( - object_path(&meta, b"PARQUET", ParquetEncoder::EXT), + object_path(&meta, b"PARQUET", ::EXT), "data/air_temperature/2026-06-12/08-00-00-abe4fb8a17f5800b.parquet" ); } @@ -507,7 +507,7 @@ mod tests { object_path( &at(datetime!(2026-07-02 13:20 UTC)), b"PARQUET", - ParquetEncoder::EXT + ::EXT ), "data/pm25/2026-07-02/13-20-00-abe4fb8a17f5800b.parquet" ); @@ -517,7 +517,7 @@ mod tests { object_path( &at(datetime!(2026-07-02 13:20:30 UTC)), b"PARQUET", - ParquetEncoder::EXT + ::EXT ), "data/pm25/2026-07-02/13-20-30-abe4fb8a17f5800b.parquet" ); @@ -533,14 +533,17 @@ mod tests { start: datetime!(2026-07-02 13:20 UTC), end: datetime!(2026-07-02 13:30 UTC), }; - let first = object_path(&meta, b"chunk-1", ParquetEncoder::EXT); - let second = object_path(&meta, b"chunk-2", ParquetEncoder::EXT); + let first = object_path(&meta, b"chunk-1", ::EXT); + let second = object_path(&meta, b"chunk-2", ::EXT); assert_ne!(first, second); assert_eq!( first, "data/rainfall/2026-07-02/13-20-00-e58fad5f76c7ba24.parquet" ); - assert_eq!(object_path(&meta, b"chunk-1", ParquetEncoder::EXT), first); + assert_eq!( + object_path(&meta, b"chunk-1", ::EXT), + first + ); } #[test]