src/znippy_pack.rs
//! **znippy converter** — pack collected Arrow batches (or raw byte blobs) into a
//! **znippy** archive for compact storage/transport, and read them back
//! BYTE-IDENTICAL. This is what korp uses to store/ship a collected dataset
//! compactly: a `Collector` lands Arrow `RecordBatch`es, and this module turns a
//! set of named tables into one self-describing, checksummed archive (and back).
//!
//! It **reuses the `znippy-common` pack format** — blake3-checksummed,
//! OpenZL-compressed (store-raw when compression doesn't shrink) blobs — so it
//! reinvents NO compression. Behind `--features znippy` (pulls the OpenZL C codec
//! via `znippy-common`, like `osm`/`skade-sink` are opt-in).
//!
//! Two layers:
//! * **blob** round-trip — [`pack_blobs`] / [`unpack_blob`] / [`unpack_all_blobs`]
//! store arbitrary named byte-blobs (e.g. a raw collected source payload) and
//! read them back byte-for-byte (blake3-verified on extract).
//! * **Arrow** round-trip — [`pack_tables`] / [`unpack_tables`] serialise each
//! named `RecordBatch` stream to Arrow IPC, pack those IPC blobs, and rebuild
//! the batches on read: the compact on-disk / transport form of a dataset.
//!
//! The znippy reader is path-based (positioned reads on demand), so a full
//! in-process round-trip stages the in-memory archive bytes on a temp file for
//! the read side; the pack side is fully in-memory (`create_archive_to_vec`).
use std::path::{Path, PathBuf};
use anyhow::{Context, Result, anyhow};
use arrow_array::RecordBatch;
use znippy_common::meta_sink_append::{create_archive, create_archive_to_vec};
use znippy_common::{ZnippyArchive, ZnippyReader};
/// The default OpenZL compression level for collected-data packs (a balanced
/// speed/ratio setting; store-raw still wins per-blob when it doesn't shrink).
pub const DEFAULT_LEVEL: i32 = 3;
/// The archive-entry suffix a packed Arrow table lands under (`{name}.arrow`).
pub const ARROW_SUFFIX: &str = ".arrow";
/// Serialize a slice of `RecordBatch`es (all sharing one schema) to Arrow IPC
/// **stream** bytes — one self-describing blob (schema + record batches). This is
/// the exact byte payload the archive stores for a table.
pub fn batches_to_ipc(batches: &[RecordBatch]) -> Result<Vec<u8>> {
let schema = batches
.first()
.map(|b| b.schema())
.ok_or_else(|| anyhow!("cannot serialize an empty batch set (no schema to write)"))?;
let mut buf = Vec::new();
{
let mut w = arrow_ipc::writer::StreamWriter::try_new(&mut buf, &schema)
.context("open Arrow IPC stream writer")?;
for b in batches {
w.write(b).context("write Arrow IPC batch")?;
}
w.finish().context("finish Arrow IPC stream")?;
}
Ok(buf)
}
/// Reconstruct `RecordBatch`es from Arrow IPC stream bytes (the inverse of
/// [`batches_to_ipc`]).
pub fn ipc_to_batches(ipc: &[u8]) -> Result<Vec<RecordBatch>> {
let reader = arrow_ipc::reader::StreamReader::try_new(std::io::Cursor::new(ipc), None)
.context("open Arrow IPC stream reader")?;
reader
.collect::<std::result::Result<Vec<_>, _>>()
.context("read Arrow IPC batches")
}
/// Pack named byte-blobs into a znippy archive **in memory** (no filesystem). The
/// returned `Vec<u8>` is byte-for-byte what [`pack_blobs_to_file`] would write.
pub fn pack_blobs(blobs: &[(String, Vec<u8>)]) -> Result<Vec<u8>> {
let (bytes, _report) =
create_archive_to_vec(blobs, DEFAULT_LEVEL).context("pack znippy archive in memory")?;
Ok(bytes)
}
/// Pack named byte-blobs into a znippy archive written to `path`.
pub fn pack_blobs_to_file(path: &Path, blobs: &[(String, Vec<u8>)]) -> Result<()> {
create_archive(path, blobs, DEFAULT_LEVEL)
.with_context(|| format!("write znippy archive to {}", path.display()))?;
Ok(())
}
/// Open a znippy archive at `path` and extract ONE blob by name, **blake3-verified**
/// (errors on any integrity mismatch).
pub fn unpack_blob_from_file(path: &Path, name: &str) -> Result<Vec<u8>> {
let ar = ZnippyArchive::open(path)
.with_context(|| format!("open znippy archive {}", path.display()))?;
ar.extract_file_verified(name)
.with_context(|| format!("extract `{name}`"))
}
/// Round-trip helper: stage in-memory archive bytes on a temp file, open it, and
/// extract ONE blob byte-identical (blake3-verified). The temp file is removed
/// when the returned bytes are produced.
pub fn unpack_blob(archive: &[u8], name: &str) -> Result<Vec<u8>> {
let staged = Staged::new(archive)?;
unpack_blob_from_file(staged.path(), name)
}
/// Extract EVERY blob from an in-memory archive as `(name, bytes)`, blake3-verified.
pub fn unpack_all_blobs(archive: &[u8]) -> Result<Vec<(String, Vec<u8>)>> {
let staged = Staged::new(archive)?;
let ar = ZnippyArchive::open(staged.path()).context("open znippy archive")?;
let mut out = Vec::new();
for name in ar.list_files().context("list archive files")? {
let bytes = ar
.extract_file_verified(&name)
.with_context(|| format!("extract `{name}`"))?;
out.push((name, bytes));
}
Ok(out)
}
/// Pack named Arrow tables (each a `RecordBatch` stream) into a znippy archive in
/// memory. Each table lands as one IPC blob at `{name}{ARROW_SUFFIX}`.
pub fn pack_tables(tables: &[(String, Vec<RecordBatch>)]) -> Result<Vec<u8>> {
let blobs: Vec<(String, Vec<u8>)> = tables
.iter()
.map(|(name, batches)| Ok((format!("{name}{ARROW_SUFFIX}"), batches_to_ipc(batches)?)))
.collect::<Result<Vec<_>>>()?;
pack_blobs(&blobs)
}
/// Unpack a znippy archive (in-memory bytes) back into named Arrow tables — the
/// inverse of [`pack_tables`]. Each `{name}{ARROW_SUFFIX}` blob is blake3-verified
/// then deserialised to its `RecordBatch`es.
pub fn unpack_tables(archive: &[u8]) -> Result<Vec<(String, Vec<RecordBatch>)>> {
let mut out = Vec::new();
for (fname, ipc) in unpack_all_blobs(archive)? {
let name = fname
.strip_suffix(ARROW_SUFFIX)
.unwrap_or(&fname)
.to_string();
out.push((name, ipc_to_batches(&ipc)?));
}
Ok(out)
}
/// A temp-file staging of in-memory archive bytes so the path-based znippy reader
/// can open it; the file (and its dir) are cleaned up on drop.
struct Staged {
_dir: tempfile::TempDir,
path: PathBuf,
}
impl Staged {
fn new(archive: &[u8]) -> Result<Self> {
let dir = tempfile::tempdir().context("temp dir for znippy archive")?;
let path = dir.path().join("korp-pack.znippy");
std::fs::write(&path, archive).context("stage znippy archive bytes")?;
Ok(Self { _dir: dir, path })
}
fn path(&self) -> &Path {
&self.path
}
}
#[cfg(test)]
mod tests {
use super::*;
use arrow_array::Float64Array;
use arrow_schema::{DataType, Field, Schema};
use std::sync::Arc;
fn sample_batch() -> RecordBatch {
let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new("lat", DataType::Float64, false),
]));
RecordBatch::try_new(
schema,
vec![
Arc::new(arrow_array::Int64Array::from(vec![1_i64, 2, 3])),
Arc::new(Float64Array::from(vec![59.3293_f64, 57.7089, 55.6050])),
],
)
.expect("batch")
}
#[test]
fn arrow_tables_round_trip_and_stored_blob_is_byte_identical() {
let rb = sample_batch();
let tables = vec![("collected".to_string(), vec![rb.clone()])];
let archive = pack_tables(&tables).expect("pack");
assert!(!archive.is_empty(), "archive carries bytes");
// Arrow round-trip: schema + values reconstruct.
let back = unpack_tables(&archive).expect("unpack");
assert_eq!(back.len(), 1);
assert_eq!(back[0].0, "collected", "table name preserved");
let got = &back[0].1[0];
assert_eq!(got.schema(), rb.schema(), "schema preserved");
assert_eq!(got.num_rows(), 3);
let lat = got
.column_by_name("lat")
.unwrap()
.as_any()
.downcast_ref::<Float64Array>()
.unwrap();
assert!(
(lat.value(0) - 59.3293).abs() < 1e-9,
"value survived round-trip"
);
// The STORED blob is byte-identical to the IPC we packed (blake3-verified).
let ipc = batches_to_ipc(&[rb]).expect("ipc");
let blob = unpack_blob(&archive, "collected.arrow").expect("blob");
assert_eq!(blob, ipc, "stored Arrow-IPC blob is byte-identical");
}
#[test]
fn raw_blob_round_trips_byte_identical() {
let raw = b"raw collected payload \x00\x01\x02\xff end".to_vec();
let archive = pack_blobs(&[("dataset.bin".to_string(), raw.clone())]).expect("pack");
let got = unpack_blob(&archive, "dataset.bin").expect("unpack");
assert_eq!(got, raw, "arbitrary bytes round-trip exactly");
}
}