o41-training-data 0.1.3

use crate::{
    columns,
    config::{Field, Source},
    storage::{Client, Input, Object},
};
use anyhow::{Result, ensure};
#[derive(Default)]
struct Location {
    record: usize,
    byte: u64,
    group: usize,
}
use arrow::record_batch::RecordBatch;
use bytes::{Buf, Bytes};
use rmpv::Value;
use std::{
    io::{BufRead, BufReader, Cursor, Read},
    sync::{Arc, Mutex},
};

/// Parquet tail read with the footer; it holds the metadata of most shards.
const TAIL: u64 = 64 * 1024;
/// Window for reads outside prefetched bytes, such as Parquet page headers.
const HEADER_WINDOW: u64 = 64 * 1024;
/// Bytes a sample reads ahead: a record, or a dictionary and first data page.
const SAMPLE_BYTES: u64 = 2 * 1024 * 1024;

fn validate(value: &Value) -> Result<()> {
    match value {
        Value::Map(m) => {
            for (k, v) in m {
                ensure!(k.as_str().is_some(), "map key must be UTF-8 string");
                validate(v)?;
            }
        }
        Value::Array(a) => {
            for v in a {
                validate(v)?;
            }
        }
        Value::String(s) => ensure!(s.as_str().is_some(), "invalid UTF-8"),
        _ => (),
    }
    Ok(())
}
fn msgpack(
    input: &mut impl Read,
    fields: &[Field],
    limit: usize,
    emit: &mut impl FnMut(Value) -> Result<bool>,
    location: &mut Location,
) -> Result<()> {
    loop {
        let mut header = [0u8; 8];
        if input.read(&mut header[..1])? == 0 {
            break;
        }
        input.read_exact(&mut header[1..])?;
        let size = u64::from_be_bytes(header);
        ensure!(
            size > 0 && size <= (256 * 1024 * 1024).min(limit as u64),
            "invalid frame length"
        );
        let mut compressed = vec![0; size as usize];
        input.read_exact(&mut compressed)?;
        let decompress = crate::profile::Timer::new(1);
        let frame_size = zstd::zstd_safe::find_frame_compressed_size(&compressed)
            .map_err(|e| anyhow::anyhow!("invalid zstd frame: {e}"))?;
        ensure!(
            frame_size == compressed.len(),
            "extra data after zstd frame"
        );
        let declared = zstd::zstd_safe::get_frame_content_size(&compressed)
            .map_err(|e| anyhow::anyhow!("invalid zstd content size: {e:?}"))?;
        ensure!(
            declared.is_none_or(|n| n <= limit as u64),
            "payload exceeds maxRecordBytes"
        );
        let mut payload = Vec::new();
        let mut decoder = zstd::stream::read::Decoder::new(compressed.as_slice())?;
        decoder.window_log_max(limit.max(1024).next_power_of_two().trailing_zeros())?;
        if let Err(error) = decoder.take(limit as u64 + 1).read_to_end(&mut payload) {
            location.byte += payload.len() as u64;
            return Err(error.into());
        }
        ensure!(payload.len() <= limit, "payload exceeds maxRecordBytes");
        drop(decompress);
        let parse = crate::profile::Timer::new(2);
        let mut cursor = Cursor::new(&payload);
        let row = (|| -> Result<Value> {
            let count = rmp::decode::read_map_len(&mut cursor)?;
            let mut values = Vec::new();
            for _ in 0..count {
                let key = rmpv::decode::read_value(&mut cursor)?;
                let name = key
                    .as_str()
                    .ok_or_else(|| anyhow::anyhow!("map key must be UTF-8 string"))?;
                if fields.iter().any(|f| f.from == name) {
                    values.push((key, rmpv::decode::read_value(&mut cursor)?));
                } else {
                    use serde::Deserialize;
                    serde::de::IgnoredAny::deserialize(&mut rmp_serde::Deserializer::new(
                        &mut cursor,
                    ))?;
                }
            }
            ensure!(
                cursor.position() == payload.len() as u64,
                "extra MessagePack data"
            );
            Ok(Value::Map(values))
        })();
        location.byte += cursor.position();
        let row = row?;
        ensure!(row.is_map(), "record must be a map");
        validate(&row)?;
        drop(parse);
        if !emit(row)? {
            break;
        }
        location.record += 1;
    }
    Ok(())
}
fn json_value(value: &jiter::JsonValue<'_>, strings: Option<&[Value]>) -> Result<Value> {
    use jiter::JsonValue as J;
    Ok(match value {
        J::Null => Value::Nil,
        J::Bool(b) => Value::Boolean(*b),
        J::Int(i) => Value::from(*i),
        J::BigInt(i) => {
            let digits = i.to_string();
            if let Ok(n) = digits.parse::<i64>() {
                Value::from(n)
            } else {
                Value::from(digits.parse::<u64>()?)
            }
        }
        J::Float(f) => Value::F64(*f),
        J::Str(s) => match strings {
            Some(values) => values[s.parse::<usize>()?].clone(),
            None => Value::from(s.as_ref()),
        },
        J::Array(a) => Value::Array(
            a.iter()
                .map(|v| json_value(v, strings))
                .collect::<Result<_>>()?,
        ),
        J::Object(m) => {
            let mut map = std::collections::BTreeMap::new();
            for (k, v) in m.iter() {
                let key = match strings {
                    Some(values) => values[k.parse::<usize>()?].as_str().ok_or_else(|| {
                        anyhow::anyhow!("Arrow field names require valid Unicode")
                    })?,
                    None => k.as_ref(),
                };
                map.insert(key, json_value(v, strings)?);
            }
            Value::Map(map.into_iter().map(|(k, v)| (Value::from(k), v)).collect())
        }
        #[allow(unreachable_patterns)]
        _ => Value::Nil,
    })
}
// serde_json's byte visitor preserves lone surrogates as WTF-8. Replace every
// string token with an index while jiter parses Python's non-finite numbers.
// Replacing all strings prevents placeholder collisions with user values.
struct StringBytes;
impl<'de> serde::de::Visitor<'de> for StringBytes {
    type Value = Vec<u8>;
    fn expecting(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str("a JSON string")
    }
    fn visit_bytes<E: serde::de::Error>(self, bytes: &[u8]) -> std::result::Result<Vec<u8>, E> {
        Ok(bytes.to_vec())
    }
}
fn json_tokens(line: &str) -> Result<(String, Vec<Value>)> {
    use serde::Deserializer;
    let mut out = String::with_capacity(line.len());
    let mut strings = Vec::new();
    let mut i = 0;
    let bytes = line.as_bytes();
    let mut start = 0;
    while i < bytes.len() {
        if bytes[i] != b'"' {
            i += 1;
            continue;
        }
        out.push_str(&line[start..i]);
        let token_start = i;
        i += 1;
        while i < bytes.len() {
            if bytes[i] == b'\\' {
                i = (i + 2).min(bytes.len());
            } else if bytes[i] == b'"' {
                i += 1;
                break;
            } else {
                i += 1;
            }
        }
        let mut decoder = serde_json::Deserializer::from_str(&line[token_start..i]);
        let decoded = decoder.deserialize_bytes(StringBytes)?;
        decoder.end()?;
        out.push_str(&format!("\"{}\"", strings.len()));
        strings.push(match String::from_utf8(decoded) {
            Ok(s) => Value::from(s),
            Err(e) => Value::Binary(e.into_bytes()),
        });
        start = i;
    }
    out.push_str(&line[start..]);
    Ok((out, strings))
}
fn selected_json(line: &str, fields: &[Field], strings: Option<&[Value]>) -> Result<Value> {
    let mut parser = jiter::Jiter::new(line.as_bytes()).with_allow_inf_nan();
    let mut key = parser.next_object()?;
    let mut values = Vec::new();
    while let Some(name) = key {
        // A key that is not valid Unicode cannot name a projected field.
        let name = match strings {
            Some(values) => values[name.parse::<usize>()?].as_str(),
            None => Some(name),
        };
        match name.filter(|name| fields.iter().any(|f| f.from == *name)) {
            Some(name) => {
                let name = name.to_owned();
                let value = {
                    let _parse = crate::profile::Timer::new(2);
                    parser.next_value()?
                };
                let _value = crate::profile::Timer::new(3);
                values.push((Value::from(name), json_value(&value, strings)?));
            }
            None => {
                let _parse = crate::profile::Timer::new(2);
                parser.next_skip()?;
            }
        }
        key = parser.next_key()?;
    }
    parser.finish()?;
    Ok(Value::Map(values))
}
fn jsonl(
    input: Input,
    fields: &[Field],
    limit: usize,
    emit: &mut impl FnMut(Value) -> Result<bool>,
    location: &mut Location,
) -> Result<()> {
    let mut decoder = flate2::bufread::GzDecoder::new(BufReader::new(input));
    {
        let mut lines = BufReader::new(&mut decoder);
        let mut bytes = Vec::new();
        loop {
            bytes.clear();
            let decompress = crate::profile::Timer::new(1);
            let read = (&mut lines)
                .take(limit as u64 + 1)
                .read_until(b'\n', &mut bytes);
            location.byte += bytes.len() as u64;
            let n = read?;
            let line = std::str::from_utf8(&bytes)?;
            ensure!(n <= limit, "JSON line exceeds maxRecordBytes");
            drop(decompress);
            if n == 0 {
                break;
            }
            {
                match selected_json(line, fields, None) {
                    Ok(row) => {
                        if !emit(row)? {
                            return Ok(());
                        }
                        location.record += 1;
                    }
                    Err(_) => {
                        // jiter rejects lone surrogates even in skipped strings.
                        // Normalize string tokens, then still skip unused arrays.
                        let Ok((normalized, strings)) = json_tokens(line) else {
                            continue;
                        };
                        match selected_json(&normalized, fields, Some(&strings)) {
                            Ok(row) => {
                                if !emit(row)? {
                                    return Ok(());
                                }
                                location.record += 1;
                            }
                            Err(error) => {
                                let mut syntax =
                                    jiter::Jiter::new(normalized.as_bytes()).with_allow_inf_nan();
                                if normalized.trim_start().starts_with('{')
                                    && syntax.next_skip().is_ok()
                                    && syntax.finish().is_ok()
                                {
                                    return Err(error);
                                }
                            }
                        }
                    }
                }
                continue;
            }
        }
    }
    ensure!(
        decoder.into_inner().fill_buf()?.is_empty(),
        "extra data after gzip member"
    );
    Ok(())
}

/// A Parquet object with some byte ranges already fetched. Reads inside them
/// are served from memory; any other read goes to storage.
#[derive(Clone)]
struct ParquetObject {
    object: Arc<Mutex<Object>>,
    size: u64,
    fetched: Arc<Vec<(u64, Bytes)>>,
}
impl ParquetObject {
    fn fetched(&self, start: u64) -> Option<(u64, &Bytes)> {
        self.fetched
            .iter()
            .find(|(offset, bytes)| *offset <= start && start < offset + bytes.len() as u64)
            .map(|(offset, bytes)| (*offset, bytes))
    }
}
impl parquet::file::reader::Length for ParquetObject {
    fn len(&self) -> u64 {
        self.size
    }
}
impl parquet::file::reader::ChunkReader for ParquetObject {
    type T = std::io::Chain<bytes::buf::Reader<Bytes>, Input>;
    fn get_read(&self, start: u64) -> parquet::errors::Result<Self::T> {
        let (head, end) = match self.fetched(start) {
            Some((offset, bytes)) => (
                bytes.slice((start - offset) as usize..),
                offset + bytes.len() as u64,
            ),
            None => (Bytes::new(), start),
        };
        Ok(head
            .reader()
            .chain(Input::at(self.object.clone(), end, HEADER_WINDOW)))
    }
    fn get_bytes(&self, start: u64, length: usize) -> parquet::errors::Result<Bytes> {
        if let Some((offset, bytes)) = self.fetched(start)
            && (start - offset) as usize + length <= bytes.len()
        {
            let from = (start - offset) as usize;
            return Ok(bytes.slice(from..from + length));
        }
        self.object
            .lock()
            .unwrap()
            .range(start, start + length as u64)
            .map_err(|e| parquet::errors::ParquetError::General(e.to_string()))
    }
}
fn trim(value: &mut Value, widths: &[usize]) {
    if let Some((width, inner)) = widths.split_first()
        && let Value::Array(values) = value
    {
        values.truncate(*width);
        values.shrink_to_fit();
        for value in values {
            trim(value, inner);
        }
    }
}
pub struct Shard {
    pub offset: usize,
    pub columns: RecordBatch,
}
impl Shard {
    pub fn num_rows(&self) -> usize {
        self.offset + self.columns.num_rows()
    }
}
pub fn read(
    source: &Source,
    encoding: &str,
    client: &Client,
    fields: &[Field],
    start: usize,
    sample: bool,
    config: &crate::config::Config,
) -> Result<Shard> {
    let mut location = Location::default();
    let result = (|| -> Result<Shard> {
        let object = Object::open(source, client)?;
        if encoding == "parquet" {
            return parquet_read(
                object,
                fields,
                start,
                sample,
                config,
                source.rows,
                &mut location,
            );
        }
        let mut chunks = Vec::new();
        let mut rows = Vec::new();
        let mut count = 0;
        let mut emit = |mut row: Value| -> Result<bool> {
            count += 1;
            ensure!(
                count <= source.rows.unwrap_or(config.max_shard_rows),
                "shard exceeds declared rows or maxShardRows"
            );
            if count <= start {
                return Ok(true);
            }
            // Retain only the declared prefix on every axis while accumulating
            // one builder chunk, so retained rows stay within the memory bound.
            let Value::Map(ref mut map) = row else {
                anyhow::bail!("record must be a map");
            };
            for (key, value) in map {
                let mut widths: Vec<usize> = Vec::new();
                for field in fields
                    .iter()
                    .filter(|f| Some(f.from.as_str()) == key.as_str())
                {
                    for (axis, width) in field.dimensions().into_iter().enumerate() {
                        match widths.get_mut(axis) {
                            Some(current) => *current = (*current).max(width),
                            None => widths.push(width),
                        }
                    }
                }
                trim(value, &widths);
            }
            rows.push(row);
            if rows.len() == 128 || sample {
                chunks.push(columns::project_values(
                    &rows.iter().collect::<Vec<_>>(),
                    fields,
                )?);
                rows.clear();
            }
            Ok(!sample)
        };
        // A sample reads one record; a full read streams the object ahead.
        let mut input = if sample {
            Input::at(Arc::new(Mutex::new(object)), 0, SAMPLE_BYTES)
        } else {
            Input::stream(object)?
        };
        if encoding == "msgpack.zstf" {
            msgpack(
                &mut input,
                fields,
                config.max_record_bytes,
                &mut emit,
                &mut location,
            )?;
        } else {
            jsonl(
                input,
                fields,
                config.max_record_bytes,
                &mut emit,
                &mut location,
            )?;
        }
        if !sample
            && encoding == "msgpack.zstf"
            && let Some(n) = source.rows
        {
            ensure!(count == n, "MessagePack row count mismatch");
        }
        if !rows.is_empty() {
            chunks.push(columns::project_values(
                &rows.iter().collect::<Vec<_>>(),
                fields,
            )?);
        }
        let schema = columns::project_values(&[], fields)?.schema();
        Ok(Shard {
            offset: start,
            columns: arrow::compute::concat_batches(&schema, chunks.iter())?,
        })
    })();
    crate::profile::shard(
        &source.uri,
        sample,
        location.byte,
        result
            .as_ref()
            .map_or(0, |s| s.columns.get_array_memory_size() as u64),
    );
    result.map_err(|error| {
        let position = if encoding == "parquet" {
            format!("row_group={} row={}", location.group, location.record)
        } else {
            format!("record={} byte_offset={}", location.record, location.byte)
        };
        anyhow::anyhow!(
            "shard={} {position}: {error:#}",
            serde_json::to_string(&source.uri).unwrap()
        )
    })
}

fn parquet_read(
    object: Object,
    fields: &[Field],
    start: usize,
    sample: bool,
    config: &crate::config::Config,
    declared_rows: Option<usize>,
    location: &mut Location,
) -> Result<Shard> {
    use parquet::arrow::{
        ProjectionMask,
        arrow_reader::{
            ArrowReaderMetadata, ArrowReaderOptions, ParquetRecordBatchReaderBuilder, RowSelection,
            RowSelector,
        },
    };
    let size = object.size;
    let object = Arc::new(Mutex::new(object));
    let input = |fetched| ParquetObject {
        object: object.clone(),
        size,
        fetched: Arc::new(fetched),
    };
    let tail = size.saturating_sub(TAIL);
    let footer = object.lock().unwrap().range(tail, size)?;
    let metadata =
        ArrowReaderMetadata::load(&input(vec![(tail, footer)]), ArrowReaderOptions::default())?;
    let total = metadata.metadata().file_metadata().num_rows() as usize;
    ensure!(
        declared_rows.is_none_or(|n| n == total),
        "Parquet row count mismatch"
    );
    ensure!(
        total <= config.max_shard_rows && start <= total,
        "Parquet exceeds maxShardRows or cursor bounds"
    );
    let indices: Vec<_> = metadata
        .schema()
        .fields()
        .iter()
        .enumerate()
        .filter(|(_, f)| fields.iter().any(|p| p.from == *f.name()))
        .map(|(i, _)| i)
        .collect();
    let mask = ProjectionMask::roots(metadata.metadata().file_metadata().schema_descr(), indices);
    // Plan the row groups to read with the bytes of their projected column
    // chunks, then fetch those ahead of decoding.
    let mut skipped = 0;
    let mut plan = Vec::new();
    for (i, group) in metadata.metadata().row_groups().iter().enumerate() {
        let rows = group.num_rows() as usize;
        if skipped + rows <= start {
            skipped += rows;
            continue;
        }
        location.group = i;
        location.record = start.saturating_sub(skipped);
        let columns: Vec<_> = group
            .columns()
            .iter()
            .enumerate()
            .filter(|(i, _)| mask.leaf_included(*i))
            .map(|(_, c)| c)
            .collect();
        let bytes: i64 = columns.iter().map(|c| c.uncompressed_size()).sum();
        ensure!(
            bytes >= 0 && bytes as usize <= config.max_row_group_bytes,
            "Parquet row group exceeds maxRowGroupBytes"
        );
        let mut ranges = Vec::new();
        for c in columns {
            let offset = c.dictionary_page_offset().unwrap_or(c.data_page_offset());
            let length = c.compressed_size();
            ensure!(offset >= 0 && length >= 0, "invalid Parquet column chunk");
            let length = if sample {
                length.min(SAMPLE_BYTES as i64)
            } else {
                length
            };
            ranges.push(offset as u64..(offset + length) as u64);
        }
        plan.push((i, location.record, rows, ranges));
        if sample {
            break;
        }
        skipped += rows;
    }
    let mut fetched = object
        .lock()
        .unwrap()
        .prefetch(plan.iter().map(|(.., ranges)| ranges.clone()).collect())?;
    let mut chunks = Vec::new();
    for (i, first, rows, ranges) in plan {
        location.group = i;
        location.record = first;
        let bodies = fetched
            .next()
            .ok_or_else(|| anyhow::anyhow!("Parquet prefetch stopped"))??;
        let selection = RowSelection::from(vec![
            RowSelector::skip(first),
            RowSelector::select(if sample { 1 } else { rows - first }),
        ]);
        let reader = ParquetRecordBatchReaderBuilder::new_with_metadata(
            input(ranges.iter().map(|r| r.start).zip(bodies).collect()),
            metadata.clone(),
        )
        .with_projection(mask.clone())
        .with_row_groups(vec![i])
        .with_row_selection(selection)
        .with_batch_size(if sample { 1 } else { 128 })
        .build()?;
        for batch in reader {
            let batch = batch?;
            chunks.push(columns::project(&batch, fields)?);
            location.record += batch.num_rows();
            if sample {
                break;
            }
        }
    }
    let schema = columns::project_values(&[], fields)?.schema();
    Ok(Shard {
        offset: start,
        columns: arrow::compute::concat_batches(&schema, chunks.iter())?,
    })
}