sail_scripts/kernel/src/compressor.rs

//! zstd Shard Compressor
//!
//! Compresses/decompresses Arrow shard files using zstd (Zstandard).
//! zstd is used by Meta, Linux kernel, and Firefox for its ideal balance of:
//!   - Extremely fast decompression (2–3 GB/s on modern CPUs)
//!   - Excellent compression ratio (better than gzip, faster than bzip2)
//!
//! Compressed shards use ~30–40% the disk space of raw Arrow.
//! At level 3 (default), compression is near-instantaneous.
//! Training I/O with on-the-fly decompression is still faster than streaming.
//!
//! Safety: zstd crate is a safe Rust wrapper. All buffer bounds are checked.

use pyo3::prelude::*;
use std::fs::File;
use std::io::{BufReader, BufWriter, Read, Write};
use crate::errors::KernelError;

// ── Constants ─────────────────────────────────────────────────────────────────

/// I/O chunk size for streaming compress/decompress (8 MB)
const CHUNK_SIZE: usize = 8 * 1024 * 1024;

// ── Public API ────────────────────────────────────────────────────────────────

/// Compress an Arrow shard file using zstd.
///
/// Args:
///     input_path (str): Path to the `.arrow` file to compress
///     output_path (str): Path for the output `.arrow.zst` file
///     level (int): zstd compression level (1–22, default 3)
///
/// Returns:
///     tuple[int, int]: (original_bytes, compressed_bytes) for reporting ratio
#[pyfunction]
#[pyo3(signature = (input_path, output_path, level = 3))]
pub fn compress_shard(
    py: Python<'_>,
    input_path: String,
    output_path: String,
    level: i32,
) -> PyResult<(u64, u64)> {
    py.allow_threads(move || {
        compress_internal(&input_path, &output_path, level)
            .map_err(|e| pyo3::exceptions::PyIOError::new_err(e.to_string()))
    })
}

/// Decompress a zstd-compressed Arrow shard file.
///
/// Args:
///     input_path (str): Path to the `.arrow.zst` file
///     output_path (str): Path for the decompressed `.arrow` file
///
/// Returns:
///     int: Number of bytes written
#[pyfunction]
pub fn decompress_shard(
    py: Python<'_>,
    input_path: String,
    output_path: String,
) -> PyResult<u64> {
    py.allow_threads(move || {
        decompress_internal(&input_path, &output_path)
            .map_err(|e| pyo3::exceptions::PyIOError::new_err(e.to_string()))
    })
}

// ── Internal implementation ───────────────────────────────────────────────────

fn compress_internal(
    input_path: &str,
    output_path: &str,
    level: i32,
) -> Result<(u64, u64), KernelError> {
    let clamped_level = level.clamp(1, 22);

    let input_file = File::open(input_path)
        .map_err(|e| KernelError::IoError(format!("Cannot open {}: {}", input_path, e)))?;
    let original_size = input_file.metadata()
        .map(|m| m.len())
        .unwrap_or(0);

    let mut reader = BufReader::with_capacity(CHUNK_SIZE, input_file);

    let output_file = File::create(output_path)
        .map_err(|e| KernelError::IoError(format!("Cannot create {}: {}", output_path, e)))?;
    let writer = BufWriter::with_capacity(CHUNK_SIZE, output_file);

    // Safe zstd encoder — bounds are handled inside the zstd crate
    let mut encoder = zstd::Encoder::new(writer, clamped_level)
        .map_err(|e| KernelError::CompressionError(e.to_string()))?;

    // Stream in chunks (prevents loading entire file into memory)
    let mut buf = vec![0u8; CHUNK_SIZE];
    loop {
        let n = reader.read(&mut buf)
            .map_err(|e| KernelError::IoError(e.to_string()))?;
        if n == 0 { break; }

        encoder.write_all(&buf[..n])
            .map_err(|e| KernelError::CompressionError(e.to_string()))?;
    }

    let mut compressed_writer = encoder.finish()
        .map_err(|e| KernelError::CompressionError(e.to_string()))?;

    compressed_writer.flush()
        .map_err(|e| KernelError::IoError(e.to_string()))?;

    let compressed_size = std::fs::metadata(output_path)
        .map(|m| m.len())
        .unwrap_or(0);

    Ok((original_size, compressed_size))
}

fn decompress_internal(
    input_path: &str,
    output_path: &str,
) -> Result<u64, KernelError> {
    let input_file = File::open(input_path)
        .map_err(|e| KernelError::IoError(format!("Cannot open {}: {}", input_path, e)))?;
    let reader = BufReader::with_capacity(CHUNK_SIZE, input_file);

    let output_file = File::create(output_path)
        .map_err(|e| KernelError::IoError(format!("Cannot create {}: {}", output_path, e)))?;
    let mut writer = BufWriter::with_capacity(CHUNK_SIZE, output_file);

    let mut decoder = zstd::Decoder::new(reader)
        .map_err(|e| KernelError::CompressionError(e.to_string()))?;

    let mut buf = vec![0u8; CHUNK_SIZE];
    let mut total_bytes = 0u64;

    loop {
        let n = decoder.read(&mut buf)
            .map_err(|e| KernelError::CompressionError(e.to_string()))?;
        if n == 0 { break; }

        writer.write_all(&buf[..n])
            .map_err(|e| KernelError::IoError(e.to_string()))?;
        total_bytes += n as u64;
    }

    writer.flush()
        .map_err(|e| KernelError::IoError(e.to_string()))?;

    Ok(total_bytes)
}

// ── Unit Tests ────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Write;

    #[test]
    fn test_compress_decompress_roundtrip() {
        let dir = tempfile::tempdir().expect("create temp dir");
        let input = dir.path().join("input.bin");
        let compressed = dir.path().join("input.bin.zst");
        let decompressed = dir.path().join("output.bin");

        // Write test data
        let test_data = b"Hello SAIL! This is a test of the kernel compressor.".repeat(1000);
        std::fs::write(&input, &test_data).unwrap();

        // Compress
        let (orig, comp) = compress_internal(
            input.to_str().unwrap(),
            compressed.to_str().unwrap(),
            DEFAULT_ZSTD_LEVEL,
        ).unwrap();

        assert!(comp < orig, "Compressed should be smaller than original");

        // Decompress
        let decompressed_bytes = decompress_internal(
            compressed.to_str().unwrap(),
            decompressed.to_str().unwrap(),
        ).unwrap();

        // Verify round-trip
        let result = std::fs::read(&decompressed).unwrap();
        assert_eq!(result, test_data, "Round-trip decompression must be lossless");
        assert_eq!(decompressed_bytes, test_data.len() as u64);
    }

    #[test]
    fn test_invalid_input_path() {
        let result = compress_internal("/nonexistent/path/file.bin", "/tmp/out.zst", 3);
        assert!(result.is_err(), "Should fail on missing input");
    }
}