rbt_json/
lib.rs

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

rbt_json/
lib.rs
