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                self.sort.compare(&other.row, &self.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 should_write = match &dedup_cols {
309                Some(cols) => {
310                    let key = if cols.is_empty() {
311                        entry
312                            .row
313                            .iter()
314                            .map(as_string)
315                            .collect::<Vec<_>>()
316                            .join("\x1f")
317                    } else {
318                        cols.iter()
319                            .map(|c| {
320                                headers_for_dedup
321                                    .iter()
322                                    .position(|h| h == c)
323                                    .map(|i| as_string(&entry.row[i]))
324                                    .unwrap_or_default()
325                            })
326                            .collect::<Vec<_>>()
327                            .join("\x1f")
328                    };
329                    seen.insert(key)
330                }
331                None => true,
332            };
333            if should_write {
334                let record: Vec<String> = entry.row.iter().map(as_string).collect();
335                wtr.write_record(&record)?;
336            }
337            if let Some(Ok(line)) = readers[entry.idx].next() {
338                let values: Vec<serde_json::Value> = serde_json::from_str(&line)?;
339                let row: Vec<CanonicalValue> = values
340                    .into_iter()
341                    .map(CanonicalValue::from_json_value)
342                    .collect();
343                heap.push(HeapEntry {
344                    row,
345                    idx: entry.idx,
346                    sort: sort_arc.clone(),
347                    headers: headers_arc.clone(),
348                });
349            }
350        }
351
352        for path in &self.temp_files {
353            let _ = std::fs::remove_file(path);
354        }
355        wtr.flush()?;
356        Ok(())
357    }
358}
359
360use crate::Transforms;
sheeter_engine/sort.rs

sheeter_engine/
sort.rs
