sheeter-engine 1.2.0

use crate::CanonicalValue;
use std::cmp::Ordering;
use std::collections::BinaryHeap;
use std::error::Error;
use std::fs::File;
use std::io::{BufRead, BufReader, BufWriter, Read, Write};

/// Sort specification parsed from `--sort "col [asc|desc][,col2...]"`.
///
/// # Examples
///
/// ```
/// use sheeter_engine::SortSpec;
///
/// let s = SortSpec::parse("age desc, name").unwrap();
/// assert_eq!(s.keys().len(), 2);
/// ```
#[derive(Debug, Clone)]
pub struct SortSpec {
    keys: Vec<(String, bool)>,
}

impl SortSpec {
    /// Parses a sort specification.
    ///
    /// # Errors
    /// Returns an error if no valid column is found.
    pub fn parse(input: &str) -> Result<Self, Box<dyn Error>> {
        let mut keys = Vec::new();
        for part in input.split(',') {
            let part = part.trim();
            if part.is_empty() {
                continue;
            }
            let mut words = part.split_whitespace();
            let col = match words.next() {
                Some(c) => c.to_string(),
                None => continue,
            };
            let asc = match words.next() {
                None | Some("asc") => true,
                Some("desc") => false,
                Some(other) => {
                    return Err(format!("unknown sort direction {other:?}").into());
                }
            };
            keys.push((col, asc));
        }
        if keys.is_empty() {
            return Err("sort spec has no columns".into());
        }
        Ok(Self { keys })
    }

    pub fn keys(&self) -> &[(String, bool)] {
        &self.keys
    }

    fn compare(&self, a: &[CanonicalValue], b: &[CanonicalValue], headers: &[String]) -> Ordering {
        for (col, asc) in &self.keys {
            let ia = headers.iter().position(|h| h == col).and_then(|i| a.get(i));
            let ib = headers.iter().position(|h| h == col).and_then(|i| b.get(i));
            let ord = compare_values(ia, ib);
            let ord = if *asc { ord } else { ord.reverse() };
            if ord != Ordering::Equal {
                return ord;
            }
        }
        Ordering::Equal
    }
}

fn compare_values(a: Option<&CanonicalValue>, b: Option<&CanonicalValue>) -> Ordering {
    match (a, b) {
        (None, None) => Ordering::Equal,
        (None, Some(_)) => Ordering::Less,
        (Some(_), None) => Ordering::Greater,
        (Some(CanonicalValue::Null), Some(CanonicalValue::Null)) => Ordering::Equal,
        (Some(CanonicalValue::Null), _) => Ordering::Less,
        (_, Some(CanonicalValue::Null)) => Ordering::Greater,
        (Some(x), Some(y)) => {
            if let (Some(xf), Some(yf)) = (as_f64(x), as_f64(y)) {
                return xf.partial_cmp(&yf).unwrap_or(Ordering::Equal);
            }
            as_string(x).cmp(&as_string(y))
        }
    }
}

fn as_f64(v: &CanonicalValue) -> Option<f64> {
    match v {
        CanonicalValue::Number(n) => n.as_f64(),
        CanonicalValue::String(s) => s.parse().ok(),
        CanonicalValue::Bytes(_) => None,
        _ => None,
    }
}

fn as_string(v: &CanonicalValue) -> String {
    match v {
        CanonicalValue::Null => String::new(),
        CanonicalValue::Bool(b) => b.to_string(),
        CanonicalValue::Number(n) => n.to_string(),
        CanonicalValue::String(s) => s.clone(),
        CanonicalValue::Bytes(b) => {
            use base64::Engine;
            base64::engine::general_purpose::STANDARD.encode(b)
        }
        _ => String::new(),
    }
}

/// Sorts a delimited stream using bounded memory via external merge sort.
///
/// Reads `chunk_rows` rows at a time, sorts each chunk, spills it to a temp
/// file, then k-way merges the temp files into the output. Peak memory is
/// bounded by `chunk_rows * row_size`, independent of input size.
///
/// # Errors
/// Returns an error if any I/O step fails or the sort spec references a
/// missing column.
///
/// # Examples
///
/// ```
/// use sheeter_engine::{external_sort_delimited, SortSpec, Transforms};
///
/// let input = b"n\n3\n1\n2\n";
/// let mut out = Vec::new();
/// external_sort_delimited(
///     input.as_slice(), b',', &mut out, b',',
///     &Transforms::default(),
///     &SortSpec::parse("n").unwrap(),
///     10,
/// ).unwrap();
/// assert_eq!(out, b"n\n1\n2\n3\n");
/// ```
pub fn external_sort_delimited(
    reader: impl Read,
    from_delim: u8,
    writer: impl Write,
    to_delim: u8,
    transforms: &Transforms,
    sort: &SortSpec,
    chunk_rows: usize,
) -> Result<(), Box<dyn Error>> {
    let mut rdr = csv::ReaderBuilder::new()
        .delimiter(from_delim)
        .from_reader(reader);
    let raw_headers: Vec<String> = rdr.headers()?.iter().map(|s| s.to_string()).collect();

    let mut renamed = raw_headers.clone();
    for (from, to) in &transforms.rename {
        for h in renamed.iter_mut() {
            if h == from {
                *h = to.clone();
                break;
            }
        }
    }
    let final_headers: Vec<String> = match &transforms.select {
        Some(sel) => sel.clone(),
        None => renamed,
    };

    let mut sorter = Sorter::new(sort.clone(), chunk_rows, final_headers.clone());

    let limit = transforms.limit.unwrap_or(usize::MAX);
    let mut accepted = 0usize;
    for result in rdr.records() {
        if accepted >= limit {
            break;
        }
        let record = result?;
        let mut row: Vec<(String, CanonicalValue)> = Vec::with_capacity(raw_headers.len());
        for (i, h) in raw_headers.iter().enumerate() {
            row.push((
                h.clone(),
                CanonicalValue::String(record.get(i).unwrap_or("").to_string()),
            ));
        }
        if let Some(new_row) = transforms.apply_to_row(row) {
            let values: Vec<CanonicalValue> = new_row.into_iter().map(|(_, v)| v).collect();
            sorter.push(values)?;
            accepted += 1;
        }
    }

    sorter.finish(writer, to_delim, transforms.dedup.as_deref())
}

struct Sorter {
    sort: SortSpec,
    chunk_rows: usize,
    headers: Vec<String>,
    buffer: Vec<Vec<CanonicalValue>>,
    temp_files: Vec<std::path::PathBuf>,
}

impl Sorter {
    fn new(sort: SortSpec, chunk_rows: usize, headers: Vec<String>) -> Self {
        Self {
            sort,
            chunk_rows: chunk_rows.max(1),
            headers,
            buffer: Vec::new(),
            temp_files: Vec::new(),
        }
    }

    fn push(&mut self, row: Vec<CanonicalValue>) -> Result<(), Box<dyn Error>> {
        self.buffer.push(row);
        if self.buffer.len() >= self.chunk_rows {
            self.flush_chunk()?;
        }
        Ok(())
    }

    fn flush_chunk(&mut self) -> Result<(), Box<dyn Error>> {
        if self.buffer.is_empty() {
            return Ok(());
        }
        let mut chunk = std::mem::take(&mut self.buffer);
        let sort = &self.sort;
        let headers = &self.headers;
        chunk.sort_by(|a, b| sort.compare(a, b, headers));

        let tmp = tempfile::NamedTempFile::new()?;
        let (file, path) = tmp.keep()?;
        let mut w = BufWriter::new(file);
        for row in &chunk {
            let values: Vec<serde_json::Value> = row.iter().map(|v| v.to_json_value()).collect();
            let line = serde_json::to_string(&values)?;
            writeln!(w, "{line}")?;
        }
        w.flush()?;
        self.temp_files.push(path);
        Ok(())
    }

    fn finish(
        mut self,
        writer: impl Write,
        to_delim: u8,
        dedup: Option<&[String]>,
    ) -> Result<(), Box<dyn Error>> {
        self.flush_chunk()?;

        let mut wtr = csv::WriterBuilder::new()
            .delimiter(to_delim)
            .from_writer(writer);
        wtr.write_record(&self.headers)?;

        let mut seen: std::collections::HashSet<String> = std::collections::HashSet::new();
        let dedup_cols = dedup.map(|c| c.to_vec());
        let headers_for_dedup = self.headers.clone();

        if self.temp_files.is_empty() {
            wtr.flush()?;
            return Ok(());
        }

        let mut readers: Vec<std::io::Lines<BufReader<File>>> = Vec::new();
        for path in &self.temp_files {
            let f = File::open(path)?;
            readers.push(BufReader::new(f).lines());
        }

        struct HeapEntry {
            row: Vec<CanonicalValue>,
            idx: usize,
            sort: std::sync::Arc<SortSpec>,
            headers: std::sync::Arc<Vec<String>>,
        }
        impl PartialEq for HeapEntry {
            fn eq(&self, other: &Self) -> bool {
                self.cmp(other) == Ordering::Equal
            }
        }
        impl Eq for HeapEntry {}
        impl PartialOrd for HeapEntry {
            fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
                Some(self.cmp(other))
            }
        }
        impl Ord for HeapEntry {
            fn cmp(&self, other: &Self) -> Ordering {
                self.sort.compare(&other.row, &self.row, &self.headers)
            }
        }

        let sort_arc = std::sync::Arc::new(self.sort.clone());
        let headers_arc = std::sync::Arc::new(self.headers.clone());
        let mut heap: BinaryHeap<HeapEntry> = BinaryHeap::new();

        for (idx, reader) in readers.iter_mut().enumerate() {
            if let Some(Ok(line)) = reader.next() {
                let values: Vec<serde_json::Value> = serde_json::from_str(&line)?;
                let row: Vec<CanonicalValue> = values
                    .into_iter()
                    .map(CanonicalValue::from_json_value)
                    .collect();
                heap.push(HeapEntry {
                    row,
                    idx,
                    sort: sort_arc.clone(),
                    headers: headers_arc.clone(),
                });
            }
        }

        while let Some(entry) = heap.pop() {
            let should_write = match &dedup_cols {
                Some(cols) => {
                    let key = if cols.is_empty() {
                        entry
                            .row
                            .iter()
                            .map(as_string)
                            .collect::<Vec<_>>()
                            .join("\x1f")
                    } else {
                        cols.iter()
                            .map(|c| {
                                headers_for_dedup
                                    .iter()
                                    .position(|h| h == c)
                                    .map(|i| as_string(&entry.row[i]))
                                    .unwrap_or_default()
                            })
                            .collect::<Vec<_>>()
                            .join("\x1f")
                    };
                    seen.insert(key)
                }
                None => true,
            };
            if should_write {
                let record: Vec<String> = entry.row.iter().map(as_string).collect();
                wtr.write_record(&record)?;
            }
            if let Some(Ok(line)) = readers[entry.idx].next() {
                let values: Vec<serde_json::Value> = serde_json::from_str(&line)?;
                let row: Vec<CanonicalValue> = values
                    .into_iter()
                    .map(CanonicalValue::from_json_value)
                    .collect();
                heap.push(HeapEntry {
                    row,
                    idx: entry.idx,
                    sort: sort_arc.clone(),
                    headers: headers_arc.clone(),
                });
            }
        }

        for path in &self.temp_files {
            let _ = std::fs::remove_file(path);
        }
        wtr.flush()?;
        Ok(())
    }
}

use crate::Transforms;