sheeter_engine/
sort.rs

1use crate::CanonicalValue;
2use std::cmp::Ordering;
3use std::collections::BinaryHeap;
4use std::error::Error;
5use std::fs::File;
6use std::io::{BufRead, BufReader, BufWriter, Read, Write};
7
8/// Sort specification parsed from `--sort "col [asc|desc][,col2...]"`.
9///
10/// # Examples
11///
12/// ```
13/// use sheeter_engine::SortSpec;
14///
15/// let s = SortSpec::parse("age desc, name").unwrap();
16/// assert_eq!(s.keys().len(), 2);
17/// ```
18#[derive(Debug, Clone)]
19pub struct SortSpec {
20    keys: Vec<(String, bool)>,
21}
22
23impl SortSpec {
24    /// Parses a sort specification.
25    ///
26    /// # Errors
27    /// Returns an error if no valid column is found.
28    pub fn parse(input: &str) -> Result<Self, Box<dyn Error>> {
29        let mut keys = Vec::new();
30        for part in input.split(',') {
31            let part = part.trim();
32            if part.is_empty() {
33                continue;
34            }
35            let mut words = part.split_whitespace();
36            let col = match words.next() {
37                Some(c) => c.to_string(),
38                None => continue,
39            };
40            let asc = match words.next() {
41                None | Some("asc") => true,
42                Some("desc") => false,
43                Some(other) => {
44                    return Err(format!("unknown sort direction {other:?}").into());
45                }
46            };
47            keys.push((col, asc));
48        }
49        if keys.is_empty() {
50            return Err("sort spec has no columns".into());
51        }
52        Ok(Self { keys })
53    }
54
55    pub fn keys(&self) -> &[(String, bool)] {
56        &self.keys
57    }
58
59    fn compare(&self, a: &[CanonicalValue], b: &[CanonicalValue], headers: &[String]) -> Ordering {
60        for (col, asc) in &self.keys {
61            let ia = headers.iter().position(|h| h == col).and_then(|i| a.get(i));
62            let ib = headers.iter().position(|h| h == col).and_then(|i| b.get(i));
63            let ord = compare_values(ia, ib);
64            let ord = if *asc { ord } else { ord.reverse() };
65            if ord != Ordering::Equal {
66                return ord;
67            }
68        }
69        Ordering::Equal
70    }
71}
72
73fn compare_values(a: Option<&CanonicalValue>, b: Option<&CanonicalValue>) -> Ordering {
74    match (a, b) {
75        (None, None) => Ordering::Equal,
76        (None, Some(_)) => Ordering::Less,
77        (Some(_), None) => Ordering::Greater,
78        (Some(CanonicalValue::Null), Some(CanonicalValue::Null)) => Ordering::Equal,
79        (Some(CanonicalValue::Null), _) => Ordering::Less,
80        (_, Some(CanonicalValue::Null)) => Ordering::Greater,
81        (Some(x), Some(y)) => {
82            if let (Some(xf), Some(yf)) = (as_f64(x), as_f64(y)) {
83                return xf.partial_cmp(&yf).unwrap_or(Ordering::Equal);
84            }
85            as_string(x).cmp(&as_string(y))
86        }
87    }
88}
89
90fn as_f64(v: &CanonicalValue) -> Option<f64> {
91    match v {
92        CanonicalValue::Number(n) => n.as_f64(),
93        CanonicalValue::String(s) => s.parse().ok(),
94        _ => None,
95    }
96}
97
98fn as_string(v: &CanonicalValue) -> String {
99    match v {
100        CanonicalValue::Null => String::new(),
101        CanonicalValue::Bool(b) => b.to_string(),
102        CanonicalValue::Number(n) => n.to_string(),
103        CanonicalValue::String(s) => s.clone(),
104        _ => String::new(),
105    }
106}
107
108/// Sorts a delimited stream using bounded memory via external merge sort.
109///
110/// Reads `chunk_rows` rows at a time, sorts each chunk, spills it to a temp
111/// file, then k-way merges the temp files into the output. Peak memory is
112/// bounded by `chunk_rows * row_size`, independent of input size.
113///
114/// # Errors
115/// Returns an error if any I/O step fails or the sort spec references a
116/// missing column.
117///
118/// # Examples
119///
120/// ```
121/// use sheeter_engine::{external_sort_delimited, SortSpec, Transforms};
122///
123/// let input = b"n\n3\n1\n2\n";
124/// let mut out = Vec::new();
125/// external_sort_delimited(
126///     input.as_slice(), b',', &mut out, b',',
127///     &Transforms::default(),
128///     &SortSpec::parse("n").unwrap(),
129///     10,
130/// ).unwrap();
131/// assert_eq!(out, b"n\n1\n2\n3\n");
132/// ```
133pub fn external_sort_delimited(
134    reader: impl Read,
135    from_delim: u8,
136    writer: impl Write,
137    to_delim: u8,
138    transforms: &Transforms,
139    sort: &SortSpec,
140    chunk_rows: usize,
141) -> Result<(), Box<dyn Error>> {
142    let mut rdr = csv::ReaderBuilder::new()
143        .delimiter(from_delim)
144        .from_reader(reader);
145    let raw_headers: Vec<String> = rdr.headers()?.iter().map(|s| s.to_string()).collect();
146
147    let mut renamed = raw_headers.clone();
148    for (from, to) in &transforms.rename {
149        for h in renamed.iter_mut() {
150            if h == from {
151                *h = to.clone();
152                break;
153            }
154        }
155    }
156    let final_headers: Vec<String> = match &transforms.select {
157        Some(sel) => sel.clone(),
158        None => renamed,
159    };
160
161    let mut sorter = Sorter::new(sort.clone(), chunk_rows, final_headers.clone());
162
163    let limit = transforms.limit.unwrap_or(usize::MAX);
164    let mut accepted = 0usize;
165    for result in rdr.records() {
166        if accepted >= limit {
167            break;
168        }
169        let record = result?;
170        let mut row: Vec<(String, CanonicalValue)> = Vec::with_capacity(raw_headers.len());
171        for (i, h) in raw_headers.iter().enumerate() {
172            row.push((
173                h.clone(),
174                CanonicalValue::String(record.get(i).unwrap_or("").to_string()),
175            ));
176        }
177        if let Some(new_row) = transforms.apply_to_row(row) {
178            let values: Vec<CanonicalValue> = new_row.into_iter().map(|(_, v)| v).collect();
179            sorter.push(values)?;
180            accepted += 1;
181        }
182    }
183
184    sorter.finish(writer, to_delim, transforms.dedup.as_deref())
185}
186
187struct Sorter {
188    sort: SortSpec,
189    chunk_rows: usize,
190    headers: Vec<String>,
191    buffer: Vec<Vec<CanonicalValue>>,
192    temp_files: Vec<std::path::PathBuf>,
193}
194
195impl Sorter {
196    fn new(sort: SortSpec, chunk_rows: usize, headers: Vec<String>) -> Self {
197        Self {
198            sort,
199            chunk_rows: chunk_rows.max(1),
200            headers,
201            buffer: Vec::new(),
202            temp_files: Vec::new(),
203        }
204    }
205
206    fn push(&mut self, row: Vec<CanonicalValue>) -> Result<(), Box<dyn Error>> {
207        self.buffer.push(row);
208        if self.buffer.len() >= self.chunk_rows {
209            self.flush_chunk()?;
210        }
211        Ok(())
212    }
213
214    fn flush_chunk(&mut self) -> Result<(), Box<dyn Error>> {
215        if self.buffer.is_empty() {
216            return Ok(());
217        }
218        let mut chunk = std::mem::take(&mut self.buffer);
219        let sort = &self.sort;
220        let headers = &self.headers;
221        chunk.sort_by(|a, b| sort.compare(a, b, headers));
222
223        let tmp = tempfile::NamedTempFile::new()?;
224        let (file, path) = tmp.keep()?;
225        let mut w = BufWriter::new(file);
226        for row in &chunk {
227            let values: Vec<serde_json::Value> = row.iter().map(|v| v.to_json_value()).collect();
228            let line = serde_json::to_string(&values)?;
229            writeln!(w, "{line}")?;
230        }
231        w.flush()?;
232        self.temp_files.push(path);
233        Ok(())
234    }
235
236    fn finish(
237        mut self,
238        writer: impl Write,
239        to_delim: u8,
240        dedup: Option<&[String]>,
241    ) -> Result<(), Box<dyn Error>> {
242        self.flush_chunk()?;
243
244        let mut wtr = csv::WriterBuilder::new()
245            .delimiter(to_delim)
246            .from_writer(writer);
247        wtr.write_record(&self.headers)?;
248
249        let mut seen: std::collections::HashSet<String> = std::collections::HashSet::new();
250        let dedup_cols = dedup.map(|c| c.to_vec());
251        let headers_for_dedup = self.headers.clone();
252
253        if self.temp_files.is_empty() {
254            wtr.flush()?;
255            return Ok(());
256        }
257
258        let mut readers: Vec<std::io::Lines<BufReader<File>>> = Vec::new();
259        for path in &self.temp_files {
260            let f = File::open(path)?;
261            readers.push(BufReader::new(f).lines());
262        }
263
264        struct HeapEntry {
265            row: Vec<CanonicalValue>,
266            idx: usize,
267            sort: std::sync::Arc<SortSpec>,
268            headers: std::sync::Arc<Vec<String>>,
269        }
270        impl PartialEq for HeapEntry {
271            fn eq(&self, other: &Self) -> bool {
272                self.cmp(other) == Ordering::Equal
273            }
274        }
275        impl Eq for HeapEntry {}
276        impl PartialOrd for HeapEntry {
277            fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
278                Some(self.cmp(other))
279            }
280        }
281        impl Ord for HeapEntry {
282            fn cmp(&self, other: &Self) -> Ordering {
283                other.sort.compare(&self.row, &other.row, &self.headers)
284            }
285        }
286
287        let sort_arc = std::sync::Arc::new(self.sort.clone());
288        let headers_arc = std::sync::Arc::new(self.headers.clone());
289        let mut heap: BinaryHeap<HeapEntry> = BinaryHeap::new();
290
291        for (idx, reader) in readers.iter_mut().enumerate() {
292            if let Some(Ok(line)) = reader.next() {
293                let values: Vec<serde_json::Value> = serde_json::from_str(&line)?;
294                let row: Vec<CanonicalValue> = values
295                    .into_iter()
296                    .map(CanonicalValue::from_json_value)
297                    .collect();
298                heap.push(HeapEntry {
299                    row,
300                    idx,
301                    sort: sort_arc.clone(),
302                    headers: headers_arc.clone(),
303                });
304            }
305        }
306
307        while let Some(entry) = heap.pop() {
308            let dedup_key = if let Some(cols) = &dedup_cols {
309                let mut key = String::new();
310                for c in cols {
311                    if let Some(i) = headers_for_dedup.iter().position(|h| h == c) {
312                        key.push_str(&as_string(&entry.row[i]));
313                    }
314                    key.push('\x1f');
315                }
316                key
317            } else {
318                entry
319                    .row
320                    .iter()
321                    .map(as_string)
322                    .collect::<Vec<_>>()
323                    .join("\x1f")
324            };
325            if seen.insert(dedup_key) {
326                let record: Vec<String> = entry.row.iter().map(as_string).collect();
327                wtr.write_record(&record)?;
328            }
329            if let Some(Ok(line)) = readers[entry.idx].next() {
330                let values: Vec<serde_json::Value> = serde_json::from_str(&line)?;
331                let row: Vec<CanonicalValue> = values
332                    .into_iter()
333                    .map(CanonicalValue::from_json_value)
334                    .collect();
335                heap.push(HeapEntry {
336                    row,
337                    idx: entry.idx,
338                    sort: sort_arc.clone(),
339                    headers: headers_arc.clone(),
340                });
341            }
342        }
343
344        for path in &self.temp_files {
345            let _ = std::fs::remove_file(path);
346        }
347        wtr.flush()?;
348        Ok(())
349    }
350}
351
352use crate::Transforms;
sheeter_engine/sort.rs

sheeter_engine/
sort.rs
