o41-training-data 0.1.8

//! The memory bound a loader opens with holds what it allocates: Parquet shards
//! of wide list rows read whole, with the heap counted the whole time.
use arrow::{
    array::{ArrayRef, Int32Builder, ListBuilder},
    datatypes::{DataType, Field, Schema},
    record_batch::RecordBatch,
};
use o41_training_data::{Config, Loader};
use std::{
    alloc::{GlobalAlloc, Layout, System},
    sync::{
        Arc,
        atomic::{AtomicUsize, Ordering},
    },
};

struct Counting;
static LIVE: AtomicUsize = AtomicUsize::new(0);
static PEAK: AtomicUsize = AtomicUsize::new(0);
unsafe impl GlobalAlloc for Counting {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let live = LIVE.fetch_add(layout.size(), Ordering::Relaxed) + layout.size();
        PEAK.fetch_max(live, Ordering::Relaxed);
        unsafe { System.alloc(layout) }
    }
    unsafe fn dealloc(&self, pointer: *mut u8, layout: Layout) {
        LIVE.fetch_sub(layout.size(), Ordering::Relaxed);
        unsafe { System.dealloc(pointer, layout) }
    }
    unsafe fn realloc(&self, pointer: *mut u8, layout: Layout, size: usize) -> *mut u8 {
        // A growing block may be copied: the old and the new are both live.
        if size > layout.size() {
            let before = LIVE.fetch_add(size - layout.size(), Ordering::Relaxed);
            PEAK.fetch_max(before + size, Ordering::Relaxed);
        } else {
            LIVE.fetch_sub(layout.size() - size, Ordering::Relaxed);
        }
        unsafe { System.realloc(pointer, layout, size) }
    }
}
#[global_allocator]
static ALLOCATOR: Counting = Counting;

struct Case {
    widths: Vec<usize>,
    shards: usize,
    rows: usize,
    group_rows: usize,
    batch: usize,
}

/// The loader's bound and the most it held above where it started.
fn measure(case: &Case) -> (u64, u64) {
    let root = std::env::temp_dir().join(format!(
        "o41-memory-{}-{}-{}",
        std::process::id(),
        case.widths[0],
        case.batch
    ));
    std::fs::create_dir_all(&root).unwrap();
    let item = Arc::new(Field::new("item", DataType::Int32, true));
    let schema = Arc::new(Schema::new(
        (0..case.widths.len())
            .map(|i| Field::new(format!("f{i}"), DataType::List(item.clone()), false))
            .collect::<Vec<_>>(),
    ));
    let mut files = Vec::new();
    for shard in 0..case.shards {
        let path = root.join(format!("{shard}.parquet"));
        let properties = parquet::file::properties::WriterProperties::builder()
            .set_max_row_group_row_count(Some(case.group_rows))
            .set_compression(parquet::basic::Compression::SNAPPY)
            .build();
        let mut writer = parquet::arrow::ArrowWriter::try_new(
            std::fs::File::create(&path).unwrap(),
            schema.clone(),
            Some(properties),
        )
        .unwrap();
        for first in (0..case.rows).step_by(case.group_rows) {
            let columns = case
                .widths
                .iter()
                .enumerate()
                .map(|(column, &width)| {
                    let mut lists = ListBuilder::new(Int32Builder::new());
                    for row in first..(first + case.group_rows).min(case.rows) {
                        lists.values().append_slice(
                            &(0..width)
                                .map(|i| ((shard * 7919 + row * 31 + i + column) % 100_003) as i32)
                                .collect::<Vec<_>>(),
                        );
                        lists.append(true);
                    }
                    Arc::new(lists.finish()) as ArrayRef
                })
                .collect();
            writer
                .write(&RecordBatch::try_new(schema.clone(), columns).unwrap())
                .unwrap();
            writer.flush().unwrap();
        }
        writer.close().unwrap();
        files.push(serde_json::json!({"uri": path, "rows": case.rows}));
    }
    let fields: Vec<_> = (0..case.widths.len())
        .map(|i| serde_json::json!({"name": format!("f{i}"), "from": format!("f{i}"), "dtype": "i32", "shape": [case.widths[i]]}))
        .collect();
    let config: Config = serde_json::from_value(serde_json::json!({
        "fields": fields,
        "stages": [{"encoding": "parquet", "sources": [{"files": files}], "batchSize": case.batch, "infinite": false}],
        "seed": 0, "rank": 0, "worldSize": 1, "threads": 2, "prefetchBatches": 4, "lookahead": 2,
        "memoryLimitBytes": u64::MAX, "maxRecordBytes": 4 << 20, "maxRowGroupBytes": 64 << 20,
        "maxShardRows": 1 << 17,
    }))
    .unwrap();
    let bound = config.memory_bound().unwrap();
    let start = LIVE.load(Ordering::Relaxed);
    PEAK.store(start, Ordering::Relaxed);
    let mut loader = Loader::open(config, None).unwrap();
    let mut rows = 0;
    while let Some(batch) = loader.next().unwrap() {
        rows += batch.columns.num_rows();
    }
    assert_eq!(rows, case.shards * case.rows);
    drop(loader);
    std::fs::remove_dir_all(root).unwrap();
    (bound, (PEAK.load(Ordering::Relaxed) - start) as u64)
}

fn check(cases: &[Case]) {
    for case in cases {
        let (bound, peak) = measure(case);
        println!(
            "widths {:?} shards {} rows {} group {} batch {}: bound {bound} peak {peak}",
            case.widths, case.shards, case.rows, case.group_rows, case.batch
        );
        assert!(peak <= bound, "peak {peak} over bound {bound}");
    }
}

#[test]
fn the_bound_holds_wide_and_narrow_rows() {
    check(&[
        // More shards than reader slots, two rows a decode.
        Case {
            widths: vec![1 << 20; 2],
            shards: 6,
            rows: 8,
            group_rows: 2,
            batch: 2,
        },
        Case {
            widths: vec![512; 2],
            shards: 2,
            rows: 4096,
            group_rows: 1024,
            batch: 64,
        },
    ]);
}

/// The shards of the reports that found the bound 15 to 40 times the data,
/// a few GB of heap: `cargo test --release --test memory -- --ignored`.
#[test]
#[ignore]
fn the_bound_holds_the_reported_shards() {
    check(&[
        // One shard of 100 rows, three fields of 644,809.
        Case {
            widths: vec![644809; 3],
            shards: 1,
            rows: 100,
            group_rows: 4,
            batch: 1,
        },
        // Four shards of 25 rows, two fields of 2^20.
        Case {
            widths: vec![1 << 20; 2],
            shards: 4,
            rows: 25,
            group_rows: 2,
            batch: 1,
        },
        Case {
            widths: vec![1 << 20; 2],
            shards: 4,
            rows: 25,
            group_rows: 2,
            batch: 25,
        },
        // One shard of 80,000 rows of two fields of 5,544.
        Case {
            widths: vec![5544; 2],
            shards: 1,
            rows: 80000,
            group_rows: 8192,
            batch: 64,
        },
    ]);
}