nordisk/znippy

plugins/native/znippy-plugin-git/examples/explode_compress_ratio.rs

//! **What per-row compression of the eager exploded pool would buy, measured on a
//! real pool** — the number lever 3 needs before a byte of the writer changes.
//!
//! Arrow IPC's own body compression is per BUFFER, which kills the point read's
//! `pread` extent (one row cannot be lifted out of a compressed buffer), so the
//! candidate is compressing each row's payload on its own and inflating just
//! that row after the pread. This walks `objects.exploded` as the plain IPC
//! stream it is, compresses every payload with zlib (the codec git itself uses
//! per object, and already in this crate via `flate2`/zlib-rs — pure Rust),
//! and reports, per kind and per size class: rows, raw bytes, compressed bytes
//! at `fast` and `default`, and the inflate cost per row on a sample.
//!
//! Usage: `explode_compress_ratio <objects.exploded> [max_batches]`
use std::io::{Read, Write};
use std::time::Instant;

use anyhow::{Context, Result};
use znippy_common::arrow::array::{Array, LargeBinaryArray, UInt8Array};
use znippy_common::arrow::ipc::reader::StreamReader;

#[derive(Default, Clone, Copy)]
struct Acc {
    rows: u64,
    raw: u64,
    fast: u64,
    dflt: u64,
    /// Rows where `fast` did not reach 90 % of raw — compression not worth it.
    fast_useless: u64,
}

fn class(len: usize) -> usize {
    match len {
        0..=127 => 0,
        128..=511 => 1,
        512..=2047 => 2,
        2048..=8191 => 3,
        8192..=65535 => 4,
        _ => 5,
    }
}
const CLASS_NAMES: [&str; 6] = ["<128", "128-511", "512-2K", "2K-8K", "8K-64K", ">64K"];
const KIND_NAMES: [&str; 5] = ["?", "commit", "tree", "blob", "tag"];

fn zlib(data: &[u8], level: flate2::Compression) -> Vec<u8> {
    let mut e = flate2::write::ZlibEncoder::new(Vec::with_capacity(data.len() / 2 + 16), level);
    e.write_all(data).unwrap();
    e.finish().unwrap()
}

fn main() -> Result<()> {
    let mut args = std::env::args().skip(1);
    let path = args.next().context("usage: explode_compress_ratio <objects.exploded> [max_batches]")?;
    let max_batches: usize = args.next().map(|s| s.parse().unwrap_or(usize::MAX)).unwrap_or(usize::MAX);
    let f = std::fs::File::open(&path).with_context(|| format!("opening {path}"))?;
    let file_len = f.metadata()?.len();
    let reader = StreamReader::try_new(std::io::BufReader::with_capacity(8 << 20, f), None)
        .context("not an Arrow IPC stream")?;

    // [kind][class]
    let mut acc = [[Acc::default(); 6]; 5];
    let mut batches = 0usize;
    let mut sample: Vec<(Vec<u8>, usize)> = Vec::new(); // (compressed fast, raw len)
    let t0 = Instant::now();
    for batch in reader {
        let batch = batch?;
        let kinds = batch
            .column_by_name("object_type")
            .and_then(|c| c.as_any().downcast_ref::<UInt8Array>())
            .context("object_type column")?;
        let pay = batch
            .column_by_name("payload")
            .and_then(|c| c.as_any().downcast_ref::<LargeBinaryArray>())
            .context("payload column")?;
        for i in 0..batch.num_rows() {
            let k = (kinds.value(i) as usize).min(4);
            let p = pay.value(i);
            let c = class(p.len());
            let a = &mut acc[k][c];
            a.rows += 1;
            a.raw += p.len() as u64;
            let zf = zlib(p, flate2::Compression::fast());
            let zd = zlib(p, flate2::Compression::default());
            a.fast += zf.len() as u64;
            a.dflt += zd.len() as u64;
            if zf.len() * 10 > p.len() * 9 {
                a.fast_useless += 1;
            }
            if sample.len() < 20_000 && (i % 37 == 0) {
                sample.push((zf, p.len()));
            }
        }
        batches += 1;
        if batches >= max_batches {
            break;
        }
    }
    let walk = t0.elapsed();

    // Inflate cost per row, the read path's extra work.
    let t1 = Instant::now();
    let mut out = Vec::new();
    for (z, raw_len) in &sample {
        out.clear();
        let mut d = flate2::read::ZlibDecoder::new(z.as_slice());
        d.read_to_end(&mut out).unwrap();
        assert_eq!(out.len(), *raw_len);
    }
    let inflate = t1.elapsed();

    println!("file {file_len} B, {batches} batches walked in {:.1} s (includes two zlib passes per row)", walk.as_secs_f64());
    println!("inflate cost on {} sampled rows: {:.2} µs/row (mean)", sample.len(), inflate.as_secs_f64() * 1e6 / sample.len().max(1) as f64);
    println!();
    println!("{:<7} {:<8} {:>10} {:>14} {:>14} {:>6} {:>14} {:>6} {:>8}", "kind", "size", "rows", "raw", "zlib-fast", "×", "zlib-default", "×", "useless%");
    let mut tot = Acc::default();
    for k in 1..5 {
        for c in 0..6 {
            let a = acc[k][c];
            if a.rows == 0 { continue; }
            println!(
                "{:<7} {:<8} {:>10} {:>14} {:>14} {:>6.2} {:>14} {:>6.2} {:>7.1}%",
                KIND_NAMES[k], CLASS_NAMES[c], a.rows, a.raw, a.fast,
                a.raw as f64 / a.fast.max(1) as f64, a.dflt, a.raw as f64 / a.dflt.max(1) as f64,
                100.0 * a.fast_useless as f64 / a.rows as f64
            );
            tot.rows += a.rows; tot.raw += a.raw; tot.fast += a.fast; tot.dflt += a.dflt; tot.fast_useless += a.fast_useless;
        }
    }
    println!();
    println!(
        "TOTAL rows {} raw {} B; zlib-fast {} B ({:.2}×, saves {} B); zlib-default {} B ({:.2}×, saves {} B); rows where fast is useless {:.1}%",
        tot.rows, tot.raw, tot.fast, tot.raw as f64 / tot.fast.max(1) as f64, tot.raw.saturating_sub(tot.fast),
        tot.dflt, tot.raw as f64 / tot.dflt.max(1) as f64, tot.raw.saturating_sub(tot.dflt),
        100.0 * tot.fast_useless as f64 / tot.rows.max(1) as f64
    );
    Ok(())
}