o41-training-data 0.1.5

use crate::{config::Config, readers};
use anyhow::{Result, ensure};
use arrow::record_batch::RecordBatch;
use crossbeam_channel::{Receiver, bounded};
use serde::{Deserialize, Serialize};
use std::{
    collections::{HashMap, HashSet},
    sync::Arc,
};

#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub struct Position {
    pub file: usize,
    pub row: usize,
    pub epoch: u64,
    pub done: bool,
}
#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct Coordinate {
    pub stage: usize,
    pub source: usize,
    pub file: usize,
    pub row: usize,
    pub epoch: u64,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct State {
    pub version: u32,
    pub stage: usize,
    pub batch_index: u64,
    pub batches_yielded: u64,
    pub rows_yielded: u64,
    pub positions: Vec<Position>,
    pub mix_counter: u64,
    pub shuffle_counter: u64,
    pub buffer: Vec<Coordinate>,
    pub partial: Vec<Coordinate>,
}
impl State {
    pub fn new(config: &Config, stage: usize) -> Self {
        Self {
            version: 3,
            stage,
            batch_index: 0,
            batches_yielded: 0,
            rows_yielded: 0,
            positions: vec![Position::default(); config.stages[stage].sources.len()],
            mix_counter: 0,
            shuffle_counter: 0,
            buffer: Vec::new(),
            partial: Vec::new(),
        }
    }
    pub fn validate(&self, config: &Config) -> Result<()> {
        ensure!(
            self.version == 3 && self.stage < config.stages.len(),
            "unsupported cursor version or stage"
        );
        let stage = &config.stages[self.stage];
        ensure!(
            self.positions.len() == stage.sources.len(),
            "cursor source count mismatch"
        );
        for (p, g) in self.positions.iter().zip(&stage.sources) {
            ensure!(p.file <= g.files.len(), "cursor file out of bounds");
        }
        ensure!(
            self.partial.len() <= stage.batch_size,
            "cursor partial batch too large"
        );
        for c in self.buffer.iter().chain(&self.partial) {
            ensure!(
                c.stage < config.stages.len(),
                "cursor coordinate stage out of bounds"
            );
            let stage = &config.stages[c.stage];
            ensure!(
                c.source < stage.sources.len() && c.file < stage.sources[c.source].files.len(),
                "cursor coordinate out of bounds"
            );
        }
        Ok(())
    }
}
// UTF-8 names "mix" and "shuffle". All operations wrap at 64 bits.
// Stage keys are avalanched before combining, avoiding adjacent-counter shifts.
fn mix64(mut x: u64) -> u64 {
    x = (x ^ (x >> 30)).wrapping_mul(0xbf58476d1ce4e5b9);
    x = (x ^ (x >> 27)).wrapping_mul(0x94d049bb133111eb);
    x ^ (x >> 31)
}
pub fn random(seed: u64, stage: usize, name: &str, source: u64, counter: u64) -> u64 {
    let mut key = 14695981039346656037u64;
    for b in name.bytes() {
        key = (key ^ b as u64).wrapping_mul(1099511628211);
    }
    let golden = 0x9e3779b97f4a7c15u64;
    key = key.wrapping_add(seed)
        ^ mix64((stage as u64).wrapping_add(1).wrapping_mul(golden))
        ^ mix64(source.wrapping_mul(golden));
    mix64(key.wrapping_add(counter.wrapping_add(1).wrapping_mul(golden)))
}
type Key = (usize, usize, usize);
pub struct Engine {
    pub config: Arc<Config>,
    pub state: State,
    cache: HashMap<Key, readers::Shard>,
    pending: HashMap<Key, Receiver<Result<readers::Shard>>>,
    // None reads the still-cached shard; eviction copies only rows still needed.
    rows: HashMap<Coordinate, Option<RecordBatch>>,
    orders: Vec<Vec<Vec<usize>>>,
    window: Option<Key>,
    pool: rayon::ThreadPool,
    client: Arc<crate::storage::Client>,
}
pub struct Batch {
    pub columns: RecordBatch,
    pub stage: usize,
    pub index: u64,
    pub state: State,
}
impl Engine {
    pub fn new(config: Arc<Config>, state: State) -> Result<Self> {
        let pool = rayon::ThreadPoolBuilder::new()
            .num_threads(config.threads)
            .build()?;
        let runtime = Arc::new(
            tokio::runtime::Builder::new_multi_thread()
                .worker_threads(1)
                .enable_all()
                .build()?,
        );
        let client = Arc::new(crate::storage::Client::new(config.storage.clone(), runtime));
        for (stage_index, stage) in config.stages.iter().enumerate() {
            for (source_index, source) in stage.sources.iter().enumerate() {
                let fields = source.fields.as_deref().unwrap_or(&config.fields);
                let encoding = source
                    .encoding
                    .as_ref()
                    .or(stage.encoding.as_ref())
                    .unwrap();
                let mut found = false;
                for file in &source.files {
                    use anyhow::Context;
                    let sample = readers::read(file, encoding, &client, fields, 0, true, &config)
                        .with_context(|| {
                        format!("stage={stage_index} source={source_index}")
                    })?;
                    if sample.columns.num_rows() > 0 {
                        found = true;
                        break;
                    }
                }
                ensure!(
                    found,
                    "stage {stage_index} source {source_index} {}: no record for field {}",
                    source.files[0].uri,
                    fields[0].name
                );
            }
        }
        let orders = config
            .stages
            .iter()
            .enumerate()
            .map(|(stage_index, stage)| {
                stage
                    .sources
                    .iter()
                    .enumerate()
                    .map(|(source_index, source)| {
                        let mut order: Vec<_> = (0..source.files.len()).collect();
                        if stage.shuffle.is_some() {
                            for (draw, i) in (1..order.len()).rev().enumerate() {
                                let j = random(
                                    config.seed,
                                    stage_index,
                                    "shards",
                                    source_index as u64,
                                    draw as u64,
                                ) % (i as u64 + 1);
                                order.swap(i, j as usize);
                            }
                        }
                        order
                    })
                    .collect()
            })
            .collect();
        let mut engine = Self {
            config,
            state,
            cache: HashMap::new(),
            pending: HashMap::new(),
            rows: HashMap::new(),
            orders,
            window: None,
            pool,
            client,
        };
        engine.restore_rows()?;
        Ok(engine)
    }
    fn heads(&self) -> impl Iterator<Item = Key> + '_ {
        self.state
            .positions
            .iter()
            .enumerate()
            .filter(|(_, p)| !p.done)
            .map(|(source, p)| {
                (
                    self.state.stage,
                    source,
                    self.orders[self.state.stage][source][p.file],
                )
            })
    }
    fn schedule(
        &mut self,
        stage_index: usize,
        source: usize,
        file: usize,
        row: usize,
        required: bool,
    ) {
        let key = (stage_index, source, file);
        if self.cache.contains_key(&key) || self.pending.contains_key(&key) {
            return;
        }
        // Future shards cannot take the slots needed by other sources' current shards.
        let reserved = if required {
            0
        } else {
            self.heads()
                .filter(|k| !self.cache.contains_key(k) && !self.pending.contains_key(k))
                .count()
        };
        if self.cache.len() + self.pending.len() + reserved
            >= self.config.threads * self.config.lookahead
        {
            return;
        }
        let config = self.config.clone();
        let client = self.client.clone();
        let (send, recv) = bounded(1);
        self.pool.spawn(move || {
            let stage = &config.stages[stage_index];
            let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
                readers::read(
                    &stage.sources[source].files[file],
                    stage.sources[source]
                        .encoding
                        .as_ref()
                        .or(stage.encoding.as_ref())
                        .unwrap(),
                    &client,
                    stage.sources[source]
                        .fields
                        .as_deref()
                        .unwrap_or(&config.fields),
                    row,
                    false,
                    &config,
                )
            }))
            .unwrap_or_else(|_| Err(anyhow::anyhow!("shard reader panicked")));
            let result = result
                .map_err(|error| anyhow::anyhow!("stage={stage_index} source={source}: {error:#}"));
            let _ = send.send(result);
        });
        self.pending.insert(key, recv);
    }
    fn lookahead(&mut self, source: usize, file: usize) {
        if self.window == Some((self.state.stage, source, file)) {
            return;
        }
        self.window = Some((self.state.stage, source, file));
        let stage = &self.config.stages[self.state.stage];
        let (count, infinite) = (stage.sources[source].files.len(), stage.infinite);
        // At least the next shard, so one thread reads it during the current;
        // an infinite source continues into its next epoch.
        let window = self.config.threads.max(2);
        let ahead = if infinite {
            window.min(count)
        } else {
            window.min(count - file)
        };
        for k in 0..ahead {
            self.schedule(
                self.state.stage,
                source,
                self.orders[self.state.stage][source][(file + k) % count],
                if k == 0 {
                    self.state.positions[source].row
                } else {
                    0
                },
                k == 0,
            );
        }
    }
    fn load(&mut self, stage_index: usize, source: usize, file: usize, row: usize) -> Result<()> {
        let key = (stage_index, source, file);
        if self.cache.get(&key).is_some_and(|s| s.offset <= row) {
            return Ok(());
        }
        self.evict(key)?;
        if !self.pending.contains_key(&key) {
            ensure!(
                self.cache.len() + self.pending.len() < self.config.threads * self.config.lookahead,
                "reader slots exhausted"
            );
            self.schedule(stage_index, source, file, row, true);
        }
        let wait = crate::profile::Timer::new(7);
        let result = self
            .pending
            .remove(&key)
            .unwrap()
            .recv()
            .map_err(|_| anyhow::anyhow!("shard worker stopped"))??;
        drop(wait);
        if result.offset > row {
            return self.load(stage_index, source, file, row);
        }
        self.cache.insert(key, result);
        Ok(())
    }
    fn retain_row(&mut self, c: Coordinate) {
        self.rows.insert(c, None);
    }
    fn evict(&mut self, key: Key) -> Result<()> {
        if let Some(shard) = self.cache.get(&key) {
            for (c, row) in &mut self.rows {
                if row.is_none() && (c.stage, c.source, c.file) == key {
                    *row = Some(arrow::compute::interleave_record_batch(
                        &[&shard.columns],
                        &[(0, c.row - shard.offset)],
                    )?);
                }
            }
        }
        self.cache.remove(&key);
        Ok(())
    }
    fn restore_rows(&mut self) -> Result<()> {
        let mut groups = std::collections::BTreeMap::<Key, Vec<Coordinate>>::new();
        for c in self.state.buffer.iter().chain(&self.state.partial) {
            groups
                .entry((c.stage, c.source, c.file))
                .or_default()
                .push(c.clone());
        }
        let heads: HashSet<_> = self.heads().collect();
        let mut groups: Vec<_> = groups.into_iter().collect();
        // Restore old shards first; retain current shards for their unread suffixes.
        groups.sort_by_key(|(key, _)| heads.contains(key));
        for (key, coordinates) in groups {
            let mut start = coordinates.iter().map(|c| c.row).min().unwrap();
            if heads.contains(&key) {
                start = start.min(self.state.positions[key.1].row);
            }
            self.load(key.0, key.1, key.2, start)?;
            for c in coordinates {
                self.retain_row(c);
            }
            if !heads.contains(&key) {
                self.evict(key)?;
            }
        }
        Ok(())
    }
    fn source_next(&mut self, source: usize) -> Result<Option<Coordinate>> {
        let stage = self.state.stage;
        let count = self.orders[stage][source].len();
        let infinite = self.config.stages[stage].infinite;
        let mut empty = 0;
        loop {
            let p = &self.state.positions[source];
            if p.done {
                return Ok(None);
            }
            let (position, row, epoch) = (p.file, p.row, p.epoch);
            let file = self.orders[stage][source][position];
            self.lookahead(source, position);
            self.load(stage, source, file, row)?;
            let end = self.cache[&(stage, source, file)].num_rows();
            let result = if row < end {
                let c = Coordinate {
                    stage,
                    source,
                    file,
                    row,
                    epoch,
                };
                self.retain_row(c.clone());
                Some(c)
            } else {
                None
            };
            let p = &mut self.state.positions[source];
            p.row += usize::from(result.is_some());
            if p.row >= end {
                self.evict((stage, source, file))?;
                let p = &mut self.state.positions[source];
                p.row = 0;
                p.file += 1;
                if p.file == count {
                    p.epoch += 1;
                    if infinite {
                        p.file = 0;
                    } else {
                        p.done = true;
                    }
                }
            }
            if result.is_some() {
                return Ok(result);
            }
            empty += 1;
            ensure!(
                !infinite || empty <= count,
                "infinite source has no records"
            );
        }
    }
    fn mixed(&mut self) -> Result<Option<Coordinate>> {
        loop {
            let groups = &self.config.stages[self.state.stage].sources;
            let total: f64 = groups
                .iter()
                .zip(&self.state.positions)
                .filter(|(_, p)| !p.done)
                .map(|(g, _)| g.weight.unwrap_or(1.0))
                .sum();
            if total == 0.0 {
                return Ok(None);
            }
            let mut chosen = 0;
            if groups.len() > 1 {
                let draw = (random(
                    self.config.seed,
                    self.state.stage,
                    "mix",
                    0,
                    self.state.mix_counter,
                ) >> 11) as f64
                    * (1.0 / ((1u64 << 53) as f64));
                self.state.mix_counter = self.state.mix_counter.wrapping_add(1);
                let mut target = draw * total;
                for (i, (g, p)) in groups.iter().zip(&self.state.positions).enumerate() {
                    if p.done || g.weight.unwrap_or(1.0) == 0.0 {
                        continue;
                    }
                    chosen = i;
                    target -= g.weight.unwrap_or(1.0);
                    if target < 0.0 {
                        break;
                    }
                }
            }
            if let Some(c) = self.source_next(chosen)? {
                return Ok(Some(c));
            }
        }
    }
    fn coordinate(&mut self) -> Result<Option<Coordinate>> {
        let capacity = self.config.stages[self.state.stage]
            .shuffle
            .as_ref()
            .map_or(0, |s| s.buffer_rows);
        if capacity == 0 && self.state.buffer.is_empty() {
            return self.mixed();
        }
        while self.state.buffer.len() < capacity {
            if let Some(c) = self.mixed()? {
                self.state.buffer.push(c);
            } else {
                break;
            }
        }
        if self.state.buffer.is_empty() {
            return Ok(None);
        }
        let i = (random(
            self.config.seed,
            self.state.stage,
            "shuffle",
            0,
            self.state.shuffle_counter,
        ) % self.state.buffer.len() as u64) as usize;
        self.state.shuffle_counter = self.state.shuffle_counter.wrapping_add(1);
        let replacement = if self.state.buffer.len() <= capacity {
            self.mixed()?
        } else {
            None
        };
        let result = if let Some(c) = replacement {
            std::mem::replace(&mut self.state.buffer[i], c)
        } else {
            self.state.buffer.swap_remove(i)
        };
        Ok(Some(result))
    }
    fn retire(&mut self) {
        let keep: HashSet<_> = self
            .state
            .buffer
            .iter()
            .chain(&self.state.partial)
            .collect();
        self.rows.retain(|c, _| keep.contains(c));
    }
    fn materialize(&self, coordinates: &[Coordinate]) -> Result<RecordBatch> {
        let _assembly = crate::profile::Timer::new(6);
        let mut rows = Vec::new();
        let mut shards = HashMap::new();
        let mut order = Vec::with_capacity(coordinates.len());
        for c in coordinates {
            match &self.rows[c] {
                Some(row) => {
                    order.push((rows.len(), 0));
                    rows.push(row);
                }
                None => {
                    let key = (c.stage, c.source, c.file);
                    let shard = &self.cache[&key];
                    let at = *shards.entry(key).or_insert_with(|| {
                        let at = rows.len();
                        rows.push(&shard.columns);
                        at
                    });
                    order.push((at, c.row - shard.offset));
                }
            }
        }
        Ok(arrow::compute::interleave_record_batch(&rows, &order)?)
    }
    /// Assembles the next global batch and takes its index, regardless of rank.
    fn assemble(&mut self) -> Result<Option<(Vec<Coordinate>, u64)>> {
        let mut short_epochs = 0;
        loop {
            let stage = &self.config.stages[self.state.stage];
            let size = stage.batch_size;
            let drop_last = stage.drop_last;
            let epoch_boundary = stage.sources.len() == 1 && stage.shuffle.is_none();
            let mut coordinates = std::mem::take(&mut self.state.partial);
            while coordinates.len() < size {
                let Some(c) = self.coordinate()? else {
                    break;
                };
                if epoch_boundary
                    && coordinates
                        .first()
                        .is_some_and(|first| first.epoch != c.epoch)
                {
                    self.state.partial.push(c);
                    break;
                }
                coordinates.push(c);
            }
            if coordinates.len() < size && drop_last {
                if self.state.partial.is_empty() {
                    return Ok(None);
                }
                // Two short batches in a row mean a whole epoch fits in one.
                short_epochs += 1;
                ensure!(
                    short_epochs < 2,
                    "dropLast discards every batch: an epoch has fewer than batchSize rows"
                );
                continue;
            }
            if coordinates.is_empty() {
                return Ok(None);
            }
            let index = self.state.batch_index;
            self.state.batch_index += 1;
            return Ok(Some((coordinates, index)));
        }
    }
    pub fn next(&mut self) -> Result<Option<Batch>> {
        loop {
            let Some((coordinates, index)) = self.assemble()? else {
                return Ok(None);
            };
            if index % self.config.world_size != self.config.rank {
                self.retire();
                continue;
            }
            let columns = self.materialize(&coordinates)?;
            self.retire();
            // Finish the round of world_size batches before yielding, so every
            // rank yields the same count and a finite tail that cannot fill a
            // round is dropped on all ranks.
            while !self
                .state
                .batch_index
                .is_multiple_of(self.config.world_size)
            {
                if self.assemble()?.is_none() {
                    return Ok(None);
                }
                self.retire();
            }
            self.state.batches_yielded += 1;
            self.state.rows_yielded += coordinates.len() as u64;
            return Ok(Some(Batch {
                columns,
                stage: self.state.stage,
                index,
                state: self.state.clone(),
            }));
        }
    }
    /// Draws the global batches owned by other ranks up to the next multiple of
    /// world_size. Ranks that advance after the same number of local batches
    /// then switch stages from one identical cursor and stay partitioned.
    pub fn align(&mut self) -> Result<()> {
        while !self
            .state
            .batch_index
            .is_multiple_of(self.config.world_size)
        {
            if self.assemble()?.is_none() {
                break;
            }
            self.retire();
        }
        Ok(())
    }
}

impl Drop for Engine {
    fn drop(&mut self) {
        for (_, pending) in self.pending.drain() {
            let _ = pending.recv();
        }
    }
}