o41-training-data 0.1.6

use super::*;
use arrow::{array::Int64Array, ffi_stream::ArrowArrayStreamReader};
use std::{
    io::Write,
    time::{Duration, Instant},
};

fn write_jsonl(path: &std::path::Path, ids: std::ops::Range<i64>) {
    let mut gzip = flate2::write::GzEncoder::new(
        std::fs::File::create(path).unwrap(),
        flate2::Compression::default(),
    );
    for id in ids {
        writeln!(gzip, "{{\"id\":{id}}}").unwrap();
    }
    gzip.finish().unwrap();
}

unsafe fn pull(handle: *mut Loader) -> (u64, u64, Vec<i64>) {
    let mut stream = FFI_ArrowArrayStream::empty();
    let mut stage = 0;
    let mut index = 0;
    assert_eq!(
        unsafe { loader_next(handle, &mut stream, &mut stage, &mut index) },
        1
    );
    let mut reader = ArrowArrayStreamReader::try_new(stream).unwrap();
    let batch = reader.next().unwrap().unwrap();
    let ids = batch
        .column(0)
        .as_any()
        .downcast_ref::<Int64Array>()
        .unwrap()
        .values()
        .to_vec();
    (stage, index, ids)
}

#[test]
fn advance_with_exactly_eight_queued_batches() {
    let root =
        std::env::temp_dir().join(format!("o41-training-data-prefetch-{}", std::process::id()));
    std::fs::create_dir_all(&root).unwrap();
    let path = root.join("prefetch.jsonl.gz");
    write_jsonl(&path, 0..9);
    let stage = serde_json::json!({"encoding":"jsonl.gz", "sources":[{"files":[{"uri":path}]}], "batchSize":2,"shuffle":{"bufferRows":5},"infinite":true});
    let mut answers = Vec::new();
    for prefetch in [0, 8] {
        let config = serde_json::json!({"storage":{},"fields":[{"name":"id","from":"id","dtype":"i64"}],"stages":[stage,stage],"seed":7,"rank":0,"worldSize":1,"threads":2,"prefetchBatches":prefetch});
        let config = CString::new(config.to_string()).unwrap();
        // The handle and Arrow outputs live until their corresponding releases.
        unsafe {
            let handle = loader_open(config.as_ptr());
            assert!(!handle.is_null());
            for _ in 0..3 {
                pull(handle);
            }
            let deadline = Instant::now() + Duration::from_secs(5);
            while (*handle).worker.as_ref().unwrap().batches.len() < prefetch
                && Instant::now() < deadline
            {
                std::thread::sleep(Duration::from_millis(5));
            }
            assert_eq!((*handle).worker.as_ref().unwrap().batches.len(), prefetch);
            assert_eq!(loader_advance(handle, 1), 1);
            let result: Vec<_> = (0..12).map(|_| pull(handle)).collect();
            assert!(result.iter().all(|(stage, _, _)| *stage == 1));
            answers.push(result);
            loader_close(handle);
        }
    }
    assert_eq!(answers[0], answers[1]);
    std::fs::remove_dir_all(root).unwrap();
}

#[test]
fn resolve_checkpoint_and_resume_are_exact() {
    let root =
        std::env::temp_dir().join(format!("o41-training-data-resume-{}", std::process::id()));
    std::fs::create_dir_all(&root).unwrap();
    let a = root.join("a.jsonl.gz");
    let b = root.join("b.jsonl.gz");
    write_jsonl(&a, 0..7);
    write_jsonl(&b, 7..13);
    let shard = |path: &std::path::Path, records: usize| serde_json::json!({"uri": path, "records": records, "size_bytes": std::fs::metadata(path).unwrap().len(), "etag": ""});
    let manifest = root.join("data-manifest-v1.json");
    std::fs::write(
        &manifest,
        serde_json::json!({"version": 1, "data_format": "jsonl.gz", "splits": {"train": {
            "records": 13, "shards": [shard(&a, 7), shard(&b, 6)],
            "curriculum": {"easy": {"records": 7, "shards": [shard(&a, 7)]},
                           "all": {"records": 13, "shards": [shard(&a, 7), shard(&b, 6)]}}}}})
        .to_string(),
    )
    .unwrap();
    let config = serde_json::json!({
        "storage": {},
        "fields": [{"name":"id","from":"id","dtype":"i64"}],
        "stages": [{"sources": [{"manifestUri": manifest, "split": "train", "curriculumStages": ["easy", "all"]}],
                    "batchSize": 4, "infinite": true}],
        "seed": 1, "rank": 0, "worldSize": 1, "threads": 2, "prefetchBatches": 2
    });
    let config = CString::new(config.to_string()).unwrap();
    unsafe {
        let resolved = loader_resolve(config.as_ptr());
        assert!(
            !resolved.is_null(),
            "{:?}",
            CStr::from_ptr(loader_last_error(ptr::null_mut()))
        );
        let parsed: Config =
            serde_json::from_str(CStr::from_ptr(resolved).to_str().unwrap()).unwrap();
        let files: Vec<_> = parsed.stages[0].sources[0]
            .files
            .iter()
            .map(|f| (f.uri.clone(), f.rows))
            .collect();
        assert_eq!(
            files,
            vec![
                (a.to_string_lossy().into_owned(), Some(7)),
                (b.to_string_lossy().into_owned(), Some(6))
            ]
        );
        assert_eq!(
            parsed.stages[0].sources[0].encoding.as_deref(),
            Some("jsonl.gz")
        );

        let handle = loader_open(resolved);
        assert!(
            !handle.is_null(),
            "{:?}",
            CStr::from_ptr(loader_last_error(ptr::null_mut()))
        );
        loader_free_string(resolved);
        let first: Vec<_> = (0..3).map(|_| pull(handle)).collect();
        assert_eq!(first[0].2, vec![0, 1, 2, 3]);
        let mut out: *const c_char = ptr::null();
        assert_eq!(loader_checkpoint(handle, &mut out), 1);
        let checkpoint = CStr::from_ptr(out).to_str().unwrap().to_owned();
        let expected: Vec<_> = (0..5).map(|_| pull(handle)).collect();
        loader_close(handle);

        let document: serde_json::Value = serde_json::from_str(&checkpoint).unwrap();
        assert_eq!(document["batchesYielded"], 3);
        assert_eq!(document["rowsYielded"], 12);
        assert!(document["config"]["storage"].get("accessKeyId").is_none());
        let checkpoint = CString::new(checkpoint).unwrap();
        let storage = CString::new("{}").unwrap();
        let resumed = loader_resume(checkpoint.as_ptr(), storage.as_ptr());
        assert!(
            !resumed.is_null(),
            "{:?}",
            CStr::from_ptr(loader_last_error(ptr::null_mut()))
        );
        let actual: Vec<_> = (0..5).map(|_| pull(resumed)).collect();
        assert_eq!(actual, expected);
        loader_close(resumed);
    }
    std::fs::remove_dir_all(root).unwrap();
}

#[test]
fn vector_rows_keep_bits_across_eviction_shuffle_ranks_and_resume() {
    use arrow::array::{FixedSizeListArray, Float32Array};
    let root = std::env::temp_dir().join(format!("o41-row-assembly-{}", std::process::id()));
    std::fs::create_dir_all(&root).unwrap();
    let mut files = Vec::new();
    for file in 0..3 {
        let path = root.join(format!("{file}.jsonl.gz"));
        let mut gzip = flate2::write::GzEncoder::new(
            std::fs::File::create(&path).unwrap(),
            flate2::Compression::default(),
        );
        for id in file * 7..(file + 1) * 7 {
            writeln!(
                gzip,
                "{}",
                serde_json::json!({"id":id,"x":[-0.0f32,id as f32/4.0,-(id as f32)/8.0]})
            )
            .unwrap();
        }
        gzip.finish().unwrap();
        files.push(serde_json::json!({"uri":path}));
    }
    let check = |batch: &arrow::record_batch::RecordBatch| {
        let ids = batch
            .column(0)
            .as_any()
            .downcast_ref::<Int64Array>()
            .unwrap();
        let lists = batch
            .column(1)
            .as_any()
            .downcast_ref::<FixedSizeListArray>()
            .unwrap();
        let values = lists
            .values()
            .as_any()
            .downcast_ref::<Float32Array>()
            .unwrap();
        let bits: Vec<_> = values.values().iter().map(|v| v.to_bits()).collect();
        for (&id, actual) in ids.values().iter().zip(bits.as_chunks::<3>().0.iter()) {
            let expected = [-0.0f32, id as f32 / 4.0, -(id as f32) / 8.0].map(f32::to_bits);
            assert_eq!(*actual, expected);
        }
        (ids.values().to_vec(), bits)
    };
    let read = |loader: &mut Loader| {
        let batch = loader.next().unwrap().unwrap();
        (batch.index, check(&batch.columns))
    };
    for shuffle in [
        serde_json::Value::Null,
        serde_json::json!({"bufferRows":17}),
    ] {
        for rank in 0..2 {
            let config: Config = serde_json::from_value(serde_json::json!({
                "storage":{}, "fields":[{"name":"id","from":"id","dtype":"i64"},
                    {"name":"x","from":"x","dtype":"f32","shape":[3]}],
                "stages":[{"encoding":"jsonl.gz","sources":[{"files":files}],
                    "batchSize":5,"shuffle":shuffle,"infinite":true}],
                "seed":19,"rank":rank,"worldSize":2,"threads":2,"prefetchBatches":2
            }))
            .unwrap();
            let mut loader = Loader::open(config, None).unwrap();
            let retained = loader.next().unwrap().unwrap().columns;
            let saved = check(&retained);
            for _ in 0..7 {
                read(&mut loader);
            }
            let checkpoint = loader.checkpoint();
            let expected: Vec<_> = (0..12).map(|_| read(&mut loader)).collect();
            let serialized =
                serde_json::from_value(serde_json::to_value(&checkpoint).unwrap()).unwrap();
            for checkpoint in [checkpoint.clone(), serialized] {
                let mut resumed = Loader::resume(checkpoint, Storage::default()).unwrap();
                let actual: Vec<_> = (0..12).map(|_| read(&mut resumed)).collect();
                assert_eq!(actual, expected);
            }
            assert_eq!(check(&retained), saved);
        }
    }
    std::fs::remove_dir_all(root).unwrap();
}

#[test]
fn f64_fields_keep_double_precision_and_pad() {
    use arrow::array::{FixedSizeListArray, Float64Array};
    let root = std::env::temp_dir().join(format!("o41-f64-{}", std::process::id()));
    std::fs::create_dir_all(&root).unwrap();
    let path = root.join("0.jsonl.gz");
    let mut gzip = flate2::write::GzEncoder::new(
        std::fs::File::create(&path).unwrap(),
        flate2::Compression::default(),
    );
    // Values no f32 holds exactly, and a row short of the declared width.
    let value = |id: i64| [0.0014102326241450737f64, id as f64 / 3.0];
    for id in 0..4 {
        let [a, b] = value(id);
        let x = if id % 2 == 0 { vec![a, b] } else { vec![a] };
        writeln!(gzip, "{}", serde_json::json!({"id":id,"x":x})).unwrap();
    }
    gzip.finish().unwrap();
    let config: Config = serde_json::from_value(serde_json::json!({
        "storage":{}, "fields":[{"name":"id","from":"id","dtype":"i64"},
            {"name":"x","from":"x","dtype":"f64","shape":[2],"pad":1.5}],
        "stages":[{"encoding":"jsonl.gz","sources":[{"files":[{"uri":path}]}],
            "batchSize":4,"infinite":false}],
        "seed":1,"rank":0,"worldSize":1,"threads":1,"prefetchBatches":1
    }))
    .unwrap();
    let mut loader = Loader::open(config, None).unwrap();
    let batch = loader.next().unwrap().unwrap().columns;
    let lists = batch
        .column(1)
        .as_any()
        .downcast_ref::<FixedSizeListArray>()
        .unwrap();
    let values = lists
        .values()
        .as_any()
        .downcast_ref::<Float64Array>()
        .unwrap();
    let actual: Vec<_> = values.values().iter().map(|v| v.to_bits()).collect();
    let expected: Vec<_> = (0..4)
        .flat_map(|id| {
            let [a, b] = value(id);
            [a, if id % 2 == 0 { b } else { 1.5 }]
        })
        .map(f64::to_bits)
        .collect();
    assert_eq!(actual, expected);
    assert_ne!(
        0.0014102326241450737f64 as f32 as f64,
        0.0014102326241450737f64
    );
    std::fs::remove_dir_all(root).unwrap();
}

#[test]
fn parquet_row_group_over_the_byte_budget_reads_in_slices() {
    use arrow::{
        array::{FixedSizeListArray, Int32Array, Int32Builder, ListBuilder},
        datatypes::{DataType, Field, Schema},
    };
    // Two 900-cell int32 list columns in 64 KiB pages after a dictionary page
    // of 400 KB: a row group of 750 rows holds 3.6 MB of them, one of 250 rows
    // 1.6 MB.
    let root = std::env::temp_dir().join(format!("o41-row-group-{}", std::process::id()));
    std::fs::create_dir_all(&root).unwrap();
    let path = root.join("mazes.parquet");
    let cell =
        |row: usize, column: usize, i: usize| ((row * 7919 + i * 31 + column) % 100_003) as i32;
    let item = Arc::new(Field::new("item", DataType::Int32, true));
    let schema = Arc::new(Schema::new(vec![
        Field::new("a", DataType::List(item.clone()), false),
        Field::new("b", DataType::List(item), false),
    ]));
    let properties = parquet::file::properties::WriterProperties::builder()
        .set_max_row_group_row_count(Some(750))
        .set_data_page_size_limit(64 * 1024)
        .build();
    let mut writer = parquet::arrow::ArrowWriter::try_new(
        std::fs::File::create(&path).unwrap(),
        schema.clone(),
        Some(properties),
    )
    .unwrap();
    let columns = (0..2)
        .map(|column| {
            let mut lists = ListBuilder::new(Int32Builder::new());
            for row in 0..1000 {
                lists
                    .values()
                    .append_slice(&(0..900).map(|i| cell(row, column, i)).collect::<Vec<_>>());
                lists.append(true);
            }
            Arc::new(lists.finish()) as arrow::array::ArrayRef
        })
        .collect();
    writer
        .write(&arrow::record_batch::RecordBatch::try_new(schema, columns).unwrap())
        .unwrap();
    writer.close().unwrap();

    let runtime = Arc::new(
        tokio::runtime::Builder::new_multi_thread()
            .worker_threads(1)
            .enable_all()
            .build()
            .unwrap(),
    );
    let client = storage::Client::new(Default::default(), runtime);
    let source: config::Source = serde_json::from_value(serde_json::json!({"uri":path})).unwrap();
    let fields = vec![
        serde_json::json!({"name":"a","from":"a","dtype":"i32","shape":[900]}),
        serde_json::json!({"name":"b","from":"b","dtype":"i32","shape":[900]}),
    ];
    let read = |budget: usize, start: usize, sample: bool| {
        let config: config::Config = serde_json::from_value(serde_json::json!({
            "fields": fields, "stages": [], "seed": 0, "rank": 0, "worldSize": 1,
            "threads": 1, "prefetchBatches": 1, "maxRowGroupBytes": budget,
        }))
        .unwrap();
        readers::read(
            &source,
            "parquet",
            &client,
            &config.fields,
            start,
            sample,
            &config,
        )
        .map(|shard| shard.columns)
    };
    let check = |batch: &arrow::record_batch::RecordBatch, start: usize| {
        assert_eq!(batch.num_rows(), 1000 - start);
        for column in 0..2 {
            let lists = batch
                .column(column)
                .as_any()
                .downcast_ref::<FixedSizeListArray>()
                .unwrap();
            let values = lists
                .values()
                .as_any()
                .downcast_ref::<Int32Array>()
                .unwrap();
            for row in [0, 1, 749 - start, 750 - start, 999 - start] {
                for i in [0, 450, 899] {
                    assert_eq!(values.value(row * 900 + i), cell(start + row, column, i));
                }
            }
        }
    };
    // A 3 MiB budget reads the first row group in slices of 512 KiB a column
    // and fetches the second whole; a 1 GiB budget fetches both whole. A
    // sample reads the first row through the slices too.
    for start in [0, 123] {
        let sliced = read(3 * 1024 * 1024, start, false).unwrap();
        check(&sliced, start);
        assert_eq!(sliced, read(1 << 30, start, false).unwrap());
    }
    let sample = read(3 * 1024 * 1024, 0, true).unwrap();
    assert_eq!(sample.num_rows(), 1);
    assert_eq!(sample, read(1 << 30, 0, true).unwrap());
    // A page larger than a slice is refused, by the sample at start as by a
    // read, naming the sizes and the remedy.
    for sample in [true, false] {
        let error = read(64 * 1024, 0, sample).unwrap_err().to_string();
        assert!(
            error.contains(
                "a Parquet page of 400012 bytes in a row group over 65536 bytes, where a page \
                 may hold 10922; write smaller row groups, such as pyarrow's row_group_size=4096"
            ),
            "{error}"
        );
    }
    std::fs::remove_dir_all(root).unwrap();
}