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        CanonicalValue::Bytes(_) => None,
95        _ => None,
96    }
97}
98
99fn as_string(v: &CanonicalValue) -> String {
100    match v {
101        CanonicalValue::Null => String::new(),
102        CanonicalValue::Bool(b) => b.to_string(),
103        CanonicalValue::Number(n) => n.to_string(),
104        CanonicalValue::String(s) => s.clone(),
105        CanonicalValue::Bytes(b) => {
106            use base64::Engine;
107            base64::engine::general_purpose::STANDARD.encode(b)
108        }
109        _ => String::new(),
110    }
111}
112
113/// Sorts a delimited stream using bounded memory via external merge sort.
114///
115/// Reads `chunk_rows` rows at a time, sorts each chunk, spills it to a temp
116/// file, then k-way merges the temp files into the output. Peak memory is
117/// bounded by `chunk_rows * row_size`, independent of input size.
118///
119/// # Errors
120/// Returns an error if any I/O step fails or the sort spec references a
121/// missing column.
122///
123/// # Examples
124///
125/// ```
126/// use sheeter_engine::{external_sort_delimited, SortSpec, Transforms};
127///
128/// let input = b"n\n3\n1\n2\n";
129/// let mut out = Vec::new();
130/// external_sort_delimited(
131///     input.as_slice(), b',', &mut out, b',',
132///     &Transforms::default(),
133///     &SortSpec::parse("n").unwrap(),
134///     10,
135/// ).unwrap();
136/// assert_eq!(out, b"n\n1\n2\n3\n");
137/// ```
138pub fn external_sort_delimited(
139    reader: impl Read,
140    from_delim: u8,
141    writer: impl Write,
142    to_delim: u8,
143    transforms: &Transforms,
144    sort: &SortSpec,
145    chunk_rows: usize,
146) -> Result<(), Box<dyn Error>> {
147    let mut rdr = csv::ReaderBuilder::new()
148        .delimiter(from_delim)
149        .from_reader(reader);
150    let raw_headers: Vec<String> = rdr.headers()?.iter().map(|s| s.to_string()).collect();
151
152    let mut renamed = raw_headers.clone();
153    for (from, to) in &transforms.rename {
154        for h in renamed.iter_mut() {
155            if h == from {
156                *h = to.clone();
157                break;
158            }
159        }
160    }
161    let final_headers: Vec<String> = match &transforms.select {
162        Some(sel) => sel.clone(),
163        None => renamed,
164    };
165
166    let mut sorter = Sorter::new(sort.clone(), chunk_rows, final_headers.clone());
167
168    let limit = transforms.limit.unwrap_or(usize::MAX);
169    let mut accepted = 0usize;
170    for result in rdr.records() {
171        if accepted >= limit {
172            break;
173        }
174        let record = result?;
175        let mut row: Vec<(String, CanonicalValue)> = Vec::with_capacity(raw_headers.len());
176        for (i, h) in raw_headers.iter().enumerate() {
177            row.push((
178                h.clone(),
179                CanonicalValue::String(record.get(i).unwrap_or("").to_string()),
180            ));
181        }
182        if let Some(new_row) = transforms.apply_to_row(row) {
183            let values: Vec<CanonicalValue> = new_row.into_iter().map(|(_, v)| v).collect();
184            sorter.push(values)?;
185            accepted += 1;
186        }
187    }
188
189    sorter.finish(writer, to_delim, transforms.dedup.as_deref())
190}
191
192struct Sorter {
193    sort: SortSpec,
194    chunk_rows: usize,
195    headers: Vec<String>,
196    buffer: Vec<Vec<CanonicalValue>>,
197    temp_files: Vec<std::path::PathBuf>,
198}
199
200impl Sorter {
201    fn new(sort: SortSpec, chunk_rows: usize, headers: Vec<String>) -> Self {
202        Self {
203            sort,
204            chunk_rows: chunk_rows.max(1),
205            headers,
206            buffer: Vec::new(),
207            temp_files: Vec::new(),
208        }
209    }
210
211    fn push(&mut self, row: Vec<CanonicalValue>) -> Result<(), Box<dyn Error>> {
212        self.buffer.push(row);
213        if self.buffer.len() >= self.chunk_rows {
214            self.flush_chunk()?;
215        }
216        Ok(())
217    }
218
219    fn flush_chunk(&mut self) -> Result<(), Box<dyn Error>> {
220        if self.buffer.is_empty() {
221            return Ok(());
222        }
223        let mut chunk = std::mem::take(&mut self.buffer);
224        let sort = &self.sort;
225        let headers = &self.headers;
226        chunk.sort_by(|a, b| sort.compare(a, b, headers));
227
228        let tmp = tempfile::NamedTempFile::new()?;
229        let (file, path) = tmp.keep()?;
230        let mut w = BufWriter::new(file);
231        for row in &chunk {
232            let values: Vec<serde_json::Value> = row.iter().map(|v| v.to_json_value()).collect();
233            let line = serde_json::to_string(&values)?;
234            writeln!(w, "{line}")?;
235        }
236        w.flush()?;
237        self.temp_files.push(path);
238        Ok(())
239    }
240
241    fn finish(
242        mut self,
243        writer: impl Write,
244        to_delim: u8,
245        dedup: Option<&[String]>,
246    ) -> Result<(), Box<dyn Error>> {
247        self.flush_chunk()?;
248
249        let mut wtr = csv::WriterBuilder::new()
250            .delimiter(to_delim)
251            .from_writer(writer);
252        wtr.write_record(&self.headers)?;
253
254        let mut seen: std::collections::HashSet<String> = std::collections::HashSet::new();
255        let dedup_cols = dedup.map(|c| c.to_vec());
256        let headers_for_dedup = self.headers.clone();
257
258        if self.temp_files.is_empty() {
259            wtr.flush()?;
260            return Ok(());
261        }
262
263        let mut readers: Vec<std::io::Lines<BufReader<File>>> = Vec::new();
264        for path in &self.temp_files {
265            let f = File::open(path)?;
266            readers.push(BufReader::new(f).lines());
267        }
268
269        struct HeapEntry {
270            row: Vec<CanonicalValue>,
271            idx: usize,
272            sort: std::sync::Arc<SortSpec>,
273            headers: std::sync::Arc<Vec<String>>,
274        }
275        impl PartialEq for HeapEntry {
276            fn eq(&self, other: &Self) -> bool {
277                self.cmp(other) == Ordering::Equal
278            }
279        }
280        impl Eq for HeapEntry {}
281        impl PartialOrd for HeapEntry {
282            fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
283                Some(self.cmp(other))
284            }
285        }
286        impl Ord for HeapEntry {
287            fn cmp(&self, other: &Self) -> Ordering {
288                self.sort.compare(&other.row, &self.row, &self.headers)
289            }
290        }
291
292        let sort_arc = std::sync::Arc::new(self.sort.clone());
293        let headers_arc = std::sync::Arc::new(self.headers.clone());
294        let mut heap: BinaryHeap<HeapEntry> = BinaryHeap::new();
295
296        for (idx, reader) in readers.iter_mut().enumerate() {
297            if let Some(Ok(line)) = reader.next() {
298                let values: Vec<serde_json::Value> = serde_json::from_str(&line)?;
299                let row: Vec<CanonicalValue> = values
300                    .into_iter()
301                    .map(CanonicalValue::from_json_value)
302                    .collect();
303                heap.push(HeapEntry {
304                    row,
305                    idx,
306                    sort: sort_arc.clone(),
307                    headers: headers_arc.clone(),
308                });
309            }
310        }
311
312        while let Some(entry) = heap.pop() {
313            let should_write = match &dedup_cols {
314                Some(cols) => {
315                    let key = if cols.is_empty() {
316                        entry
317                            .row
318                            .iter()
319                            .map(as_string)
320                            .collect::<Vec<_>>()
321                            .join("\x1f")
322                    } else {
323                        cols.iter()
324                            .map(|c| {
325                                headers_for_dedup
326                                    .iter()
327                                    .position(|h| h == c)
328                                    .map(|i| as_string(&entry.row[i]))
329                                    .unwrap_or_default()
330                            })
331                            .collect::<Vec<_>>()
332                            .join("\x1f")
333                    };
334                    seen.insert(key)
335                }
336                None => true,
337            };
338            if should_write {
339                let record: Vec<String> = entry.row.iter().map(as_string).collect();
340                wtr.write_record(&record)?;
341            }
342            if let Some(Ok(line)) = readers[entry.idx].next() {
343                let values: Vec<serde_json::Value> = serde_json::from_str(&line)?;
344                let row: Vec<CanonicalValue> = values
345                    .into_iter()
346                    .map(CanonicalValue::from_json_value)
347                    .collect();
348                heap.push(HeapEntry {
349                    row,
350                    idx: entry.idx,
351                    sort: sort_arc.clone(),
352                    headers: headers_arc.clone(),
353                });
354            }
355        }
356
357        for path in &self.temp_files {
358            let _ = std::fs::remove_file(path);
359        }
360        wtr.flush()?;
361        Ok(())
362    }
363}
364
365use crate::Transforms;
sheeter_engine/sort.rs

sheeter_engine/
sort.rs
