rbt/json/
mod.rs

1//! `rbt::json`: `jshift` parse-avoiding JSONL path projection, field stamping, and filter kernels.
2
3use anyhow::{anyhow, Result};
4use arrow::array::{ArrayRef, BooleanBuilder, Float64Builder, Int64Builder, StringBuilder};
5use arrow::datatypes::{DataType, SchemaRef};
6use arrow::record_batch::RecordBatch;
7
8pub struct JsonExtractSpec {
9    pub paths: Vec<String>,
10}
11
12pub struct JShiftExtractor {
13    pub spec: JsonExtractSpec,
14}
15
16impl JShiftExtractor {
17    pub fn new(paths: Vec<String>) -> Self {
18        Self {
19            spec: JsonExtractSpec { paths },
20        }
21    }
22
23    /// Extracts target JSON paths from a JSONL byte stream (lines separated by `\n`)
24    /// and constructs a RecordBatch according to the provided target schema.
25    pub fn extract_jsonl(&self, jsonl_bytes: &[u8], schema: SchemaRef) -> Result<RecordBatch> {
26        let lines: Vec<&[u8]> = jsonl_bytes
27            .split(|&b| b == b'\n')
28            .map(trim_bytes) // Remove trailing \r or whitespaces
29            .filter(|line| !line.is_empty())
30            .collect();
31
32        let num_rows = lines.len();
33
34        // 1. Parse target paths using jshift
35        let parsed_paths: Vec<Vec<jshift::PathSegment<'_>>> = self
36            .spec
37            .paths
38            .iter()
39            .map(|p| {
40                jshift::try_parse_path(p)
41                    .map_err(|e| anyhow!("Failed to parse JSON path '{}': {:?}", p, e))
42            })
43            .collect::<Result<_>>()?;
44
45        // 2. Initialize builders for each field in schema
46        let mut builders: Vec<Box<dyn arrow::array::ArrayBuilder>> = Vec::new();
47        for field in schema.fields() {
48            match field.data_type() {
49                DataType::Int64 => builders.push(Box::new(Int64Builder::with_capacity(num_rows))),
50                DataType::Float64 => {
51                    builders.push(Box::new(Float64Builder::with_capacity(num_rows)))
52                }
53                DataType::Boolean => {
54                    builders.push(Box::new(BooleanBuilder::with_capacity(num_rows)))
55                }
56                DataType::Utf8 | DataType::LargeUtf8 => builders.push(Box::new(
57                    StringBuilder::with_capacity(num_rows, num_rows * 16),
58                )),
59                other => anyhow::bail!(
60                    "Unsupported target Arrow data type for JSON extraction: {:?}",
61                    other
62                ),
63            }
64        }
65
66        // Map field name in schema to the index in self.spec.paths
67        let path_indices: Vec<Option<usize>> = schema
68            .fields()
69            .iter()
70            .map(|field| self.spec.paths.iter().position(|p| p == field.name()))
71            .collect();
72
73        // 3. For each JSONL line, parse target fields
74        for line in lines {
75            for (col_idx, &path_idx) in path_indices.iter().enumerate() {
76                let builder = &mut builders[col_idx];
77                let field = schema.field(col_idx);
78
79                if let Some(p_idx) = path_idx {
80                    let path = &parsed_paths[p_idx];
81                    match jshift::find_value(line, path) {
82                        Ok(val_bytes) => {
83                            append_value(builder, field.data_type(), val_bytes)?;
84                        }
85                        Err(_) => {
86                            append_null(builder, field.data_type())?;
87                        }
88                    }
89                } else {
90                    append_null(builder, field.data_type())?;
91                }
92            }
93        }
94
95        // 4. Construct arrays and RecordBatch
96        let arrays: Vec<ArrayRef> = builders.into_iter().map(|mut b| b.finish()).collect();
97        let batch = RecordBatch::try_new(schema, arrays)?;
98        Ok(batch)
99    }
100}
101
102fn trim_bytes(mut s: &[u8]) -> &[u8] {
103    while !s.is_empty() && (s[0] == b' ' || s[0] == b'\t' || s[0] == b'\r' || s[0] == b'\n') {
104        s = &s[1..];
105    }
106    while !s.is_empty()
107        && (s[s.len() - 1] == b' '
108            || s[s.len() - 1] == b'\t'
109            || s[s.len() - 1] == b'\r'
110            || s[s.len() - 1] == b'\n')
111    {
112        s = &s[..s.len() - 1];
113    }
114    s
115}
116
117fn unescape_json_string(val_bytes: &[u8]) -> Result<String> {
118    if val_bytes.len() >= 2 && val_bytes[0] == b'"' && val_bytes[val_bytes.len() - 1] == b'"' {
119        let inner = &val_bytes[1..val_bytes.len() - 1];
120        let mut s = String::with_capacity(inner.len());
121        let mut chars = std::str::from_utf8(inner)?.chars();
122        while let Some(c) = chars.next() {
123            if c == '\\' {
124                if let Some(next_c) = chars.next() {
125                    match next_c {
126                        '"' => s.push('"'),
127                        '\\' => s.push('\\'),
128                        '/' => s.push('/'),
129                        'b' => s.push('\x08'),
130                        'f' => s.push('\x0c'),
131                        'n' => s.push('\n'),
132                        'r' => s.push('\r'),
133                        't' => s.push('\t'),
134                        other => {
135                            s.push('\\');
136                            s.push(other);
137                        }
138                    }
139                } else {
140                    s.push('\\');
141                }
142            } else {
143                s.push(c);
144            }
145        }
146        Ok(s)
147    } else {
148        Ok(std::str::from_utf8(val_bytes)?.to_string())
149    }
150}
151
152fn append_value(
153    builder: &mut Box<dyn arrow::array::ArrayBuilder>,
154    data_type: &DataType,
155    val_bytes: &[u8],
156) -> Result<()> {
157    match data_type {
158        DataType::Int64 => {
159            let s = std::str::from_utf8(val_bytes)?;
160            let val = s.trim().parse::<i64>()?;
161            builder
162                .as_any_mut()
163                .downcast_mut::<Int64Builder>()
164                .ok_or_else(|| anyhow!("Failed to downcast builder to Int64Builder"))?
165                .append_value(val);
166        }
167        DataType::Float64 => {
168            let s = std::str::from_utf8(val_bytes)?;
169            let val = s.trim().parse::<f64>()?;
170            builder
171                .as_any_mut()
172                .downcast_mut::<Float64Builder>()
173                .ok_or_else(|| anyhow!("Failed to downcast builder to Float64Builder"))?
174                .append_value(val);
175        }
176        DataType::Boolean => {
177            let val = match val_bytes {
178                b"true" => true,
179                b"false" => false,
180                other => {
181                    let s = std::str::from_utf8(other)?;
182                    s.trim().parse::<bool>()?
183                }
184            };
185            builder
186                .as_any_mut()
187                .downcast_mut::<BooleanBuilder>()
188                .ok_or_else(|| anyhow!("Failed to downcast builder to BooleanBuilder"))?
189                .append_value(val);
190        }
191        DataType::Utf8 | DataType::LargeUtf8 => {
192            let s = unescape_json_string(val_bytes)?;
193            builder
194                .as_any_mut()
195                .downcast_mut::<StringBuilder>()
196                .ok_or_else(|| anyhow!("Failed to downcast builder to StringBuilder"))?
197                .append_value(&s);
198        }
199        other => anyhow::bail!("Unsupported target Arrow data type: {:?}", other),
200    }
201    Ok(())
202}
203
204fn append_null(
205    builder: &mut Box<dyn arrow::array::ArrayBuilder>,
206    data_type: &DataType,
207) -> Result<()> {
208    match data_type {
209        DataType::Int64 => {
210            builder
211                .as_any_mut()
212                .downcast_mut::<Int64Builder>()
213                .ok_or_else(|| anyhow!("Failed to downcast builder to Int64Builder"))?
214                .append_null();
215        }
216        DataType::Float64 => {
217            builder
218                .as_any_mut()
219                .downcast_mut::<Float64Builder>()
220                .ok_or_else(|| anyhow!("Failed to downcast builder to Float64Builder"))?
221                .append_null();
222        }
223        DataType::Boolean => {
224            builder
225                .as_any_mut()
226                .downcast_mut::<BooleanBuilder>()
227                .ok_or_else(|| anyhow!("Failed to downcast builder to BooleanBuilder"))?
228                .append_null();
229        }
230        DataType::Utf8 | DataType::LargeUtf8 => {
231            builder
232                .as_any_mut()
233                .downcast_mut::<StringBuilder>()
234                .ok_or_else(|| anyhow!("Failed to downcast builder to StringBuilder"))?
235                .append_null();
236        }
237        other => anyhow::bail!("Unsupported target Arrow data type: {:?}", other),
238    }
239    Ok(())
240}
241
242#[cfg(test)]
243mod tests {
244    use super::*;
245    use arrow::array::Array;
246    use arrow::datatypes::{Field, Schema};
247    use std::sync::Arc;
248
249    #[test]
250    fn test_extract_jsonl() -> Result<()> {
251        let jsonl = b"
252            {\"id\": 1, \"name\": \"Alice\", \"active\": true, \"score\": 98.5}
253            {\"id\": 2, \"name\": \"Bob\", \"active\": false, \"score\": 85.0}
254            {\"id\": 3, \"name\": \"Charlie\", \"active\": true}
255        ";
256
257        let paths = vec![
258            "id".to_string(),
259            "name".to_string(),
260            "active".to_string(),
261            "score".to_string(),
262        ];
263        let extractor = JShiftExtractor::new(paths);
264
265        let schema = Arc::new(Schema::new(vec![
266            Field::new("id", DataType::Int64, true),
267            Field::new("name", DataType::Utf8, true),
268            Field::new("active", DataType::Boolean, true),
269            Field::new("score", DataType::Float64, true),
270        ]));
271
272        let batch = extractor.extract_jsonl(jsonl, schema)?;
273
274        assert_eq!(batch.num_rows(), 3);
275        assert_eq!(batch.num_columns(), 4);
276
277        // Verify values
278        let id_array = batch
279            .column(0)
280            .as_any()
281            .downcast_ref::<arrow::array::Int64Array>()
282            .unwrap();
283        assert_eq!(id_array.value(0), 1);
284        assert_eq!(id_array.value(1), 2);
285        assert_eq!(id_array.value(2), 3);
286
287        let name_array = batch
288            .column(1)
289            .as_any()
290            .downcast_ref::<arrow::array::StringArray>()
291            .unwrap();
292        assert_eq!(name_array.value(0), "Alice");
293        assert_eq!(name_array.value(1), "Bob");
294        assert_eq!(name_array.value(2), "Charlie");
295
296        let active_array = batch
297            .column(2)
298            .as_any()
299            .downcast_ref::<arrow::array::BooleanArray>()
300            .unwrap();
301        assert!(active_array.value(0));
302        assert!(!active_array.value(1));
303        assert!(active_array.value(2));
304
305        let score_array = batch
306            .column(3)
307            .as_any()
308            .downcast_ref::<arrow::array::Float64Array>()
309            .unwrap();
310        assert_eq!(score_array.value(0), 98.5);
311        assert_eq!(score_array.value(1), 85.0);
312        assert!(score_array.is_null(2)); // Charlie score is missing/null
313
314        Ok(())
315    }
316}
rbt/json/mod.rs

rbt/json/
mod.rs
