Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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<Zstd<LEVEL>>`, 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
Expand Down
30 changes: 29 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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<LEVEL>`. Values outside that range do not implement the parquet
compression policy and fail to compile:

```rust
use meathook::{HfSink, ParquetEncoder, Zstd};

type DatasetParquet = ParquetEncoder<Zstd<3>>;

let sink = HfSink::<MyRecord>::new(client, repo, token)
.encoder(DatasetParquet::new());

let uncompressed = ParquetEncoder::default();
let zstd_level_1 = ParquetEncoder::<Zstd>::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.
Expand All @@ -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` |
Expand Down
240 changes: 232 additions & 8 deletions src/encode.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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")]
Expand Down Expand Up @@ -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<LEVEL>`](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<Compression, ParquetError>;
}

/// 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<Compression, ParquetError> {
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::<Zstd<23>>::new();
/// ```
#[cfg(feature = "parquet")]
#[derive(Debug, Clone, Copy, Default)]
pub struct Zstd<const LEVEL: u8 = 1>;

#[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<Compression, ParquetError> {
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::<Zstd>::new();
/// let zstd_3 = ParquetEncoder::<Zstd<3>>::new();
/// ```
#[cfg(feature = "parquet")]
#[derive(Debug, Clone, Copy, Default)]
pub struct ParquetEncoder;
#[derive(Debug, Clone, Copy)]
pub struct ParquetEncoder<C = Uncompressed> {
_compression: PhantomData<C>,
}

#[cfg(feature = "parquet")]
impl Encoder for ParquetEncoder {
impl<C: ParquetCompression> ParquetEncoder<C> {
/// Creates an encoder using the compression policy `C`.
#[must_use]
pub const fn new() -> Self {
Self {
_compression: PhantomData,
}
}

fn writer_properties() -> Result<WriterProperties, ParquetError> {
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<C: ParquetCompression> Encoder for ParquetEncoder<C> {
type Error = ParquetEncodeError;
const EXT: &'static str = "parquet";

Expand All @@ -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)
Expand Down Expand Up @@ -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<u8>) -> Vec<Compression> {
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()
Expand All @@ -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::<Sample>(&[]).unwrap();
let encoded = [
ParquetEncoder::default().encode::<Sample>(&[]).unwrap(),
ParquetEncoder::<Zstd>::new().encode::<Sample>(&[]).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::<Zstd>::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::<Result<_, _>>().unwrap();
let round_tripped = serde_arrow::from_record_batch::<Vec<Sample>>(&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::<Zstd<7>>::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::<Result<Vec<_>, _>>()
.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::<Zstd>::new().encode(&records).unwrap();

assert!(compressed.len() < uncompressed.len());
}
}

Expand Down
2 changes: 1 addition & 1 deletion src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Loading
Loading