rich_data/
record.rs

1//! JSON records as rows, shared by the [`jsonl`](crate::jsonl) and
2//! [`serialize`](crate::serialize) adapters.
3
4use std::collections::{BTreeSet, HashMap, HashSet, VecDeque};
5
6use serde_json::{Map, Value as Json};
7
8use crate::{DataError, DataType, Field, Row, RowSource, Schema, Value};
9
10/// How many levels of nested arrays and objects the schema describes; below
11/// that a value's type is `any`.
12pub const MAX_DEPTH: usize = 32;
13
14/// How many records a streaming source reads ahead to choose its columns and
15/// schema, unless told otherwise.
16pub const DEFAULT_SAMPLE: usize = 1000;
17
18/// How many keys seen after the sample a source names in
19/// [`RecordSource::unknown_keys`].
20pub const MAX_UNKNOWN_KEYS: usize = 100;
21
22/// Rows from a stream of JSON records.
23///
24/// The source reads the first `sample` records ahead: their keys, in the
25/// order first seen, are the columns, and their values decide the schema
26/// (`integer` and `float` widen to `float`; disagreeing types become `any`;
27/// a key some records lack, or that holds `null`, is nullable). A record
28/// that is not an object is one `value` column. Scalars become cells as they
29/// are (`true`/`false` as text, since [`Value`] has no boolean; whole
30/// numbers in a `float` column as floats), and arrays and objects become
31/// their compact JSON text.
32///
33/// Keys first seen after the sample have no column; their values are
34/// dropped and their names kept in [`unknown_keys`](Self::unknown_keys), so
35/// a caller can say so: the first [`MAX_UNKNOWN_KEYS`] of them, so a stream
36/// whose every record has a new key stays bounded
37/// ([`more_unknown_keys`](Self::more_unknown_keys) says there were more).
38/// Read everything at once (the adapters' `read`) to sample every record.
39#[derive(Debug)]
40pub struct RecordSource<I> {
41    records: I,
42    buffer: VecDeque<Map<String, Json>>,
43    columns: Vec<String>,
44    /// `columns`, to look keys up in.
45    names: HashSet<String>,
46    schema: Schema,
47    unknown: BTreeSet<String>,
48    more_unknown: bool,
49}
50
51impl<I> RecordSource<I>
52where
53    I: Iterator<Item = Result<Json, DataError>>,
54{
55    /// Read ahead `sample` records (at least one) from `records`.
56    pub(crate) fn new(mut records: I, sample: usize) -> Result<Self, DataError> {
57        let mut buffer = VecDeque::new();
58        let mut fields: Vec<Field> = Vec::new();
59        // Each field's position in `fields`, the record it was first seen in
60        // and how many records held it: a hash and counts rather than a scan
61        // of every field per key, which made a record of many keys quadratic.
62        let mut index: HashMap<String, usize> = HashMap::new();
63        let mut first: Vec<usize> = Vec::new();
64        let mut held: Vec<usize> = Vec::new();
65        for number in 0..sample.max(1) {
66            let Some(record) = records.next() else {
67                break;
68            };
69            let record = object(record?);
70            for (key, value) in &record {
71                let seen = field_of(key, value, 0);
72                match index.get(key) {
73                    Some(&at) => {
74                        fields[at] = merge(fields[at].clone(), seen);
75                        held[at] += 1;
76                    }
77                    // Absent from every earlier record.
78                    None => {
79                        index.insert(key.clone(), fields.len());
80                        fields.push(seen.nullable(number > 0 || value.is_null()));
81                        first.push(number);
82                        held.push(1);
83                    }
84                }
85            }
86            buffer.push_back(record);
87        }
88        // A column some record since its first lacks may be null.
89        let read = buffer.len();
90        for ((field, first), held) in fields.iter_mut().zip(first).zip(held) {
91            if held < read - first {
92                *field = field.clone().nullable(true);
93            }
94        }
95        Ok(RecordSource {
96            records,
97            buffer,
98            names: index.into_keys().collect(),
99            columns: fields.iter().map(|f| f.name().to_string()).collect(),
100            schema: Schema::new(fields),
101            unknown: BTreeSet::new(),
102            more_unknown: false,
103        })
104    }
105
106    /// Keys seen after the sample, which have no column (so far): the first
107    /// [`MAX_UNKNOWN_KEYS`] seen, sorted.
108    pub fn unknown_keys(&self) -> impl Iterator<Item = &str> {
109        self.unknown.iter().map(String::as_str)
110    }
111
112    /// Whether more keys than [`unknown_keys`](Self::unknown_keys) holds were
113    /// seen after the sample.
114    pub fn more_unknown_keys(&self) -> bool {
115        self.more_unknown
116    }
117
118    fn row(&mut self, record: Map<String, Json>) -> Row {
119        for key in record.keys() {
120            if !self.names.contains(key) && !self.unknown.contains(key) {
121                if self.unknown.len() < MAX_UNKNOWN_KEYS {
122                    self.unknown.insert(key.clone());
123                } else {
124                    self.more_unknown = true;
125                }
126            }
127        }
128        self.columns
129            .iter()
130            .zip(self.schema.fields())
131            .map(
132                |(column, field)| match record.get(column).map_or(Value::Null, cell) {
133                    // A float column's whole numbers are floats too.
134                    Value::Int(n) if field.data_type() == &DataType::Float => {
135                        Value::Float(n as f64)
136                    }
137                    other => other,
138                },
139            )
140            .collect()
141    }
142}
143
144impl<I> RowSource for RecordSource<I>
145where
146    I: Iterator<Item = Result<Json, DataError>>,
147{
148    fn columns(&self) -> &[String] {
149        &self.columns
150    }
151
152    fn schema(&self) -> Option<&Schema> {
153        Some(&self.schema)
154    }
155
156    fn next_row(&mut self) -> Option<Result<Row, DataError>> {
157        let record = match self.buffer.pop_front() {
158            Some(record) => record,
159            None => match self.records.next()? {
160                Ok(record) => object(record),
161                Err(error) => return Some(Err(error)),
162            },
163        };
164        Some(Ok(self.row(record)))
165    }
166}
167
168/// A record as an object; anything else is `{"value": record}`.
169fn object(record: Json) -> Map<String, Json> {
170    match record {
171        Json::Object(map) => map,
172        other => Map::from_iter([("value".to_string(), other)]),
173    }
174}
175
176/// A JSON value as a cell.
177pub(crate) fn cell(value: &Json) -> Value {
178    match value {
179        Json::Null => Value::Null,
180        Json::Bool(b) => Value::Str(b.to_string()),
181        Json::Number(n) => match n.as_i64() {
182            Some(i) => Value::Int(i),
183            None => n.as_f64().map_or(Value::Null, Value::Float),
184        },
185        Json::String(s) => Value::Str(s.clone()),
186        other => Value::Str(other.to_string()),
187    }
188}
189
190/// The field one value describes: non-null unless it is `null`.
191fn field_of(name: &str, value: &Json, depth: usize) -> Field {
192    let data_type = match value {
193        Json::Null => return Field::new(name, DataType::Unknown),
194        _ if depth >= MAX_DEPTH => DataType::Any,
195        Json::Bool(_) => DataType::Boolean,
196        Json::Number(n) if n.is_f64() => DataType::Float,
197        Json::Number(_) => DataType::Integer,
198        Json::String(_) => DataType::String,
199        Json::Array(items) => {
200            let item = items
201                .iter()
202                .map(|item| field_of("item", item, depth + 1))
203                .reduce(merge)
204                .unwrap_or_else(|| Field::new("item", DataType::Unknown));
205            DataType::List(Box::new(item))
206        }
207        Json::Object(map) => {
208            let fields = map.iter().map(|(k, v)| field_of(k, v, depth + 1)).collect();
209            DataType::Struct(fields)
210        }
211    };
212    Field::new(name, data_type).nullable(false)
213}
214
215/// Two observations of one field, merged.
216fn merge(a: Field, b: Field) -> Field {
217    let nullable = a.is_nullable() || b.is_nullable();
218    let data_type = match (a.data_type().clone(), b.data_type().clone()) {
219        (DataType::Unknown, other) | (other, DataType::Unknown) => other,
220        (x, y) if x == y => x,
221        (DataType::Integer, DataType::Float) | (DataType::Float, DataType::Integer) => {
222            DataType::Float
223        }
224        (DataType::List(x), DataType::List(y)) => DataType::List(Box::new(merge(*x, *y))),
225        (DataType::Struct(xs), DataType::Struct(ys)) => {
226            // By name through a hash, not a scan of `ys` per field of `xs`.
227            let at: HashMap<String, usize> = ys
228                .iter()
229                .enumerate()
230                .map(|(i, y)| (y.name().to_string(), i))
231                .collect();
232            let mut ys: Vec<Option<Field>> = ys.into_iter().map(Some).collect();
233            let mut fields: Vec<Field> = xs
234                .into_iter()
235                .map(|x| match at.get(x.name()).and_then(|&i| ys[i].take()) {
236                    Some(y) => merge(x, y),
237                    None => x.nullable(true),
238                })
239                .collect();
240            fields.extend(ys.into_iter().flatten().map(|y| y.nullable(true)));
241            DataType::Struct(fields)
242        }
243        _ => DataType::Any,
244    };
245    a.with_type(data_type).nullable(nullable)
246}
247
248#[cfg(test)]
249mod tests {
250    use super::*;
251    use serde_json::json;
252
253    fn source(
254        records: Vec<Json>,
255        sample: usize,
256    ) -> RecordSource<impl Iterator<Item = Result<Json, DataError>>> {
257        RecordSource::new(records.into_iter().map(Ok), sample).unwrap()
258    }
259
260    #[test]
261    fn the_sample_decides_columns_and_types() {
262        let source = source(
263            vec![
264                json!({"id": 1, "score": 2, "tags": ["a"], "meta": {"x": 1}}),
265                json!({"id": 2, "score": 2.5, "tags": [], "meta": {"y": "z"}, "extra": null}),
266            ],
267            10,
268        );
269        assert_eq!(source.columns(), ["id", "score", "tags", "meta", "extra"]);
270        let schema = source.schema().unwrap();
271        let types: Vec<String> = schema
272            .fields()
273            .iter()
274            .map(|f| f.data_type().to_string())
275            .collect();
276        assert_eq!(
277            types,
278            ["integer", "float", "list<string>", "struct", "unknown"]
279        );
280        assert!(!schema.fields()[0].is_nullable());
281        assert!(schema.fields()[4].is_nullable());
282        let meta: Vec<(&str, bool)> = schema.fields()[3]
283            .children()
284            .map(|f| (f.name(), f.is_nullable()))
285            .collect();
286        assert_eq!(meta, [("x", true), ("y", true)]);
287    }
288
289    #[test]
290    fn keys_after_the_sample_are_reported_not_invented() {
291        let mut source = source(
292            vec![json!({"a": 1}), json!({"a": "x", "b": true}), json!(7)],
293            1,
294        );
295        assert_eq!(source.columns(), ["a"]);
296        assert_eq!(source.next_row().unwrap().unwrap(), [Value::Int(1)]);
297        assert_eq!(source.next_row().unwrap().unwrap(), [Value::from("x")]);
298        // A scalar record is a `value` column, which this source lacks.
299        assert_eq!(source.next_row().unwrap().unwrap(), [Value::Null]);
300        assert!(source.next_row().is_none());
301        assert_eq!(source.unknown_keys().collect::<Vec<_>>(), ["b", "value"]);
302    }
303
304    #[test]
305    fn cells_keep_scalars_and_write_nested_values_as_json() {
306        assert_eq!(cell(&json!(true)), Value::from("true"));
307        assert_eq!(cell(&json!(1.5)), Value::Float(1.5));
308        assert_eq!(cell(&json!(u64::MAX)), Value::Float(u64::MAX as f64));
309        assert_eq!(cell(&json!({"a": [1, 2]})), Value::from(r#"{"a":[1,2]}"#));
310    }
311}
rich_data/record.rs

rich_data/
record.rs
