sheeter_engine/
ndjson.rs

1use crate::{CanonicalValue, Transforms};
2use std::error::Error;
3use std::io::{BufRead, BufReader, Read, Write};
4
5const SCHEMA_SAMPLE_ROWS: usize = 10_000;
6
7/// Streams NDJSON into CSV/TSV, sampling the first `SCHEMA_SAMPLE_ROWS`
8/// rows to discover the header union, then streaming the rest.
9///
10/// Keys that appear only after the sample window are silently dropped.
11/// Use `--select` to lock the schema explicitly if the input is known to
12/// be heterogeneous.
13///
14/// # Errors
15/// Returns an error if any line is malformed JSON or I/O fails.
16///
17/// # Examples
18///
19/// ```
20/// use sheeter_engine::{stream_ndjson_to_delimited, Transforms};
21///
22/// let input = b"{\"a\":1}\n{\"a\":2}\n";
23/// let mut out = Vec::new();
24/// stream_ndjson_to_delimited(input.as_slice(), &mut out, b',', &Transforms::default()).unwrap();
25/// assert_eq!(out, b"a\n1\n2\n");
26/// ```
27pub fn stream_ndjson_to_delimited(
28    reader: impl Read,
29    writer: impl Write,
30    to_delim: u8,
31    transforms: &Transforms,
32) -> Result<(), Box<dyn Error>> {
33    let mut buf = BufReader::new(reader);
34    let mut headers: Vec<String> = Vec::new();
35    let mut sample: Vec<CanonicalValue> = Vec::new();
36    let mut line = String::new();
37
38    loop {
39        line.clear();
40        if buf.read_line(&mut line)? == 0 {
41            break;
42        }
43        let trimmed = line.trim();
44        if trimmed.is_empty() {
45            continue;
46        }
47        let v: serde_json::Value = serde_json::from_str(trimmed)?;
48        let canon = CanonicalValue::from_json_value(v);
49        if let CanonicalValue::Object(o) = &canon {
50            for (k, _) in o {
51                if !headers.contains(k) {
52                    headers.push(k.clone());
53                }
54            }
55        }
56        sample.push(canon);
57        if sample.len() >= SCHEMA_SAMPLE_ROWS {
58            break;
59        }
60    }
61
62    let mut renamed = headers.clone();
63    for (from, to) in &transforms.rename {
64        for h in renamed.iter_mut() {
65            if h == from {
66                *h = to.clone();
67                break;
68            }
69        }
70    }
71    let final_headers: Vec<String> = match &transforms.select {
72        Some(sel) => sel.clone(),
73        None => renamed,
74    };
75
76    let mut wtr = csv::WriterBuilder::new()
77        .delimiter(to_delim)
78        .from_writer(writer);
79    wtr.write_record(&final_headers)?;
80
81    let limit = transforms.limit.unwrap_or(usize::MAX);
82    let mut emitted = 0usize;
83    let mut seen: std::collections::HashSet<String> = std::collections::HashSet::new();
84    let dedup_cols = transforms.dedup.clone();
85
86    for canon in sample {
87        if emitted >= limit {
88            break;
89        }
90        if let CanonicalValue::Object(row) = canon {
91            if let Some(new_row) = transforms.apply_to_row(row) {
92                write_row(
93                    &mut wtr,
94                    &new_row,
95                    &final_headers,
96                    &dedup_cols,
97                    &mut seen,
98                    &mut emitted,
99                )?;
100            }
101        }
102    }
103
104    while emitted < limit {
105        line.clear();
106        if buf.read_line(&mut line)? == 0 {
107            break;
108        }
109        let trimmed = line.trim();
110        if trimmed.is_empty() {
111            continue;
112        }
113        let v: serde_json::Value = serde_json::from_str(trimmed)?;
114        if let CanonicalValue::Object(row) = CanonicalValue::from_json_value(v) {
115            if let Some(new_row) = transforms.apply_to_row(row) {
116                write_row(
117                    &mut wtr,
118                    &new_row,
119                    &final_headers,
120                    &dedup_cols,
121                    &mut seen,
122                    &mut emitted,
123                )?;
124            }
125        }
126    }
127
128    wtr.flush()?;
129    Ok(())
130}
131
132/// Streams CSV/TSV into NDJSON, one line per row.
133///
134/// # Errors
135/// Returns an error if any row is malformed or I/O fails.
136///
137/// # Examples
138///
139/// ```
140/// use sheeter_engine::{stream_delimited_to_ndjson, Transforms};
141///
142/// let input = b"a\n1\n2\n";
143/// let mut out = Vec::new();
144/// stream_delimited_to_ndjson(input.as_slice(), b',', &mut out, &Transforms::default()).unwrap();
145/// assert_eq!(out, b"{\"a\":\"1\"}\n{\"a\":\"2\"}\n");
146/// ```
147pub fn stream_delimited_to_ndjson(
148    reader: impl Read,
149    from_delim: u8,
150    writer: impl Write,
151    transforms: &Transforms,
152) -> Result<(), Box<dyn Error>> {
153    let mut rdr = csv::ReaderBuilder::new()
154        .delimiter(from_delim)
155        .from_reader(reader);
156    let raw_headers: Vec<String> = rdr.headers()?.iter().map(|s| s.to_string()).collect();
157    let mut out = std::io::BufWriter::new(writer);
158    let limit = transforms.limit.unwrap_or(usize::MAX);
159    let mut emitted = 0usize;
160    let mut seen: std::collections::HashSet<String> = std::collections::HashSet::new();
161    let dedup_cols = transforms.dedup.clone();
162
163    for result in rdr.records() {
164        if emitted >= limit {
165            break;
166        }
167        let record = result?;
168        let mut row: Vec<(String, CanonicalValue)> = Vec::with_capacity(raw_headers.len());
169        for (i, h) in raw_headers.iter().enumerate() {
170            row.push((
171                h.clone(),
172                CanonicalValue::String(record.get(i).unwrap_or("").to_string()),
173            ));
174        }
175        let new_row = match transforms.apply_to_row(row) {
176            Some(r) => r,
177            None => continue,
178        };
179        if let Some(cols) = &dedup_cols {
180            let key = dedup_key(&new_row, cols);
181            if !seen.insert(key) {
182                continue;
183            }
184        }
185        let obj = CanonicalValue::Object(new_row);
186        let json = serde_json::to_string(&obj.to_json_value())?;
187        writeln!(out, "{json}")?;
188        emitted += 1;
189    }
190    out.flush()?;
191    Ok(())
192}
193
194/// Streams NDJSON to NDJSON, applying transforms per line.
195///
196/// # Errors
197/// Returns an error if any line is malformed or I/O fails.
198///
199/// # Examples
200///
201/// ```
202/// use sheeter_engine::{stream_ndjson_to_ndjson, Transforms, Filter};
203///
204/// let input = b"{\"a\":1}\n{\"a\":2}\n";
205/// let t = Transforms { filter: Some(Filter::parse("a > 1").unwrap()), ..Default::default() };
206/// let mut out = Vec::new();
207/// stream_ndjson_to_ndjson(input.as_slice(), &mut out, &t).unwrap();
208/// assert_eq!(out, b"{\"a\":2}\n");
209/// ```
210pub fn stream_ndjson_to_ndjson(
211    reader: impl Read,
212    writer: impl Write,
213    transforms: &Transforms,
214) -> Result<(), Box<dyn Error>> {
215    let buf = BufReader::new(reader);
216    let mut out = std::io::BufWriter::new(writer);
217    let limit = transforms.limit.unwrap_or(usize::MAX);
218    let mut emitted = 0usize;
219    let mut seen: std::collections::HashSet<String> = std::collections::HashSet::new();
220    let dedup_cols = transforms.dedup.clone();
221
222    for line in buf.lines() {
223        if emitted >= limit {
224            break;
225        }
226        let line = line?;
227        let trimmed = line.trim();
228        if trimmed.is_empty() {
229            continue;
230        }
231        let v: serde_json::Value = serde_json::from_str(trimmed)?;
232        let row = match CanonicalValue::from_json_value(v) {
233            CanonicalValue::Object(o) => o,
234            _ => continue,
235        };
236        let new_row = match transforms.apply_to_row(row) {
237            Some(r) => r,
238            None => continue,
239        };
240        if let Some(cols) = &dedup_cols {
241            let key = dedup_key(&new_row, cols);
242            if !seen.insert(key) {
243                continue;
244            }
245        }
246        let obj = CanonicalValue::Object(new_row);
247        let json = serde_json::to_string(&obj.to_json_value())?;
248        writeln!(out, "{json}")?;
249        emitted += 1;
250    }
251    out.flush()?;
252    Ok(())
253}
254
255fn write_row<W: Write>(
256    wtr: &mut csv::Writer<W>,
257    row: &[(String, CanonicalValue)],
258    final_headers: &[String],
259    dedup_cols: &Option<Vec<String>>,
260    seen: &mut std::collections::HashSet<String>,
261    emitted: &mut usize,
262) -> Result<(), Box<dyn Error>> {
263    if let Some(cols) = dedup_cols {
264        let key = dedup_key(row, cols);
265        if !seen.insert(key) {
266            return Ok(());
267        }
268    }
269    let values: Vec<String> = final_headers
270        .iter()
271        .map(|h| {
272            row.iter()
273                .find(|(k, _)| k == h)
274                .map(|(_, v)| canon_str(v))
275                .unwrap_or_default()
276        })
277        .collect();
278    wtr.write_record(&values)?;
279    *emitted += 1;
280    Ok(())
281}
282
283fn dedup_key(row: &[(String, CanonicalValue)], cols: &[String]) -> String {
284    if cols.is_empty() {
285        row.iter()
286            .map(|(_, v)| canon_str(v))
287            .collect::<Vec<_>>()
288            .join("\x1f")
289    } else {
290        cols.iter()
291            .map(|c| {
292                row.iter()
293                    .find(|(k, _)| k == c)
294                    .map(|(_, v)| canon_str(v))
295                    .unwrap_or_default()
296            })
297            .collect::<Vec<_>>()
298            .join("\x1f")
299    }
300}
301
302fn canon_str(v: &CanonicalValue) -> String {
303    match v {
304        CanonicalValue::Null => String::new(),
305        CanonicalValue::Bool(b) => b.to_string(),
306        CanonicalValue::Number(n) => n.to_string(),
307        CanonicalValue::String(s) => s.clone(),
308        CanonicalValue::Bytes(b) => {
309            use base64::Engine;
310            base64::engine::general_purpose::STANDARD.encode(b)
311        }
312        _ => String::new(),
313    }
314}
sheeter_engine/ndjson.rs

sheeter_engine/
ndjson.rs
