Skip to main content

laminar_connectors/schema/json/decoder/
mod.rs

1//! JSON format decoder implementing [`FormatDecoder`].
2//!
3//! Converts raw JSON byte payloads into Arrow `RecordBatch`es.
4//! Constructed once at `CREATE SOURCE` time with a frozen Arrow schema;
5//! the decoder is stateless after construction so the Ring 1 hot path
6//! has zero schema lookups.
7#![allow(clippy::disallowed_types)] // cold path: schema management
8
9use std::collections::HashMap;
10use std::sync::atomic::{AtomicU64, Ordering};
11
12use arrow_array::RecordBatch;
13use arrow_schema::{DataType, SchemaRef};
14
15use crate::schema::error::{SchemaError, SchemaResult};
16use crate::schema::traits::FormatDecoder;
17use crate::schema::types::RawRecord;
18
19mod batch;
20mod value;
21
22/// Strategy for JSON fields not in the Arrow schema.
23#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24pub enum UnknownFieldStrategy {
25    /// Silently ignore unknown fields (default).
26    Ignore,
27    /// Collect unknown fields into an `_extra` `LargeBinary` (JSONB) column.
28    CollectExtra,
29    /// Return a decode error if any unknown field is encountered.
30    Reject,
31}
32
33/// Strategy for JSON values that don't match the expected Arrow type.
34#[derive(Debug, Clone, Copy, PartialEq, Eq)]
35pub enum TypeMismatchStrategy {
36    /// Insert null and increment the mismatch counter (default).
37    Null,
38    /// Attempt coercion (e.g., `"123"` → `123` for Int64 columns).
39    Coerce,
40    /// Return a decode error on the first mismatch.
41    Reject,
42}
43
44impl TypeMismatchStrategy {
45    /// Parse from a `WITH` option value (`schema.enforcement`).
46    #[must_use]
47    pub fn from_enforcement_str(s: &str) -> Option<Self> {
48        match s.to_lowercase().as_str() {
49            "coerce" => Some(Self::Coerce),
50            "strict" => Some(Self::Reject),
51            "permissive" => Some(Self::Null),
52            _ => None,
53        }
54    }
55}
56
57/// Epoch unit for *numeric* JSON timestamps feeding a `Timestamp` column
58/// (strings still go through `timestamp_formats`).
59#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
60pub enum EpochUnit {
61    /// Epoch seconds.
62    Seconds,
63    /// Epoch milliseconds (default).
64    #[default]
65    Millis,
66    /// Epoch microseconds.
67    Micros,
68    /// Epoch nanoseconds.
69    Nanos,
70}
71
72impl EpochUnit {
73    /// Parse a `json.column.<col>.epoch_unit` value.
74    #[must_use]
75    pub fn from_str_opt(s: &str) -> Option<Self> {
76        match s.trim().to_lowercase().as_str() {
77            "seconds" => Some(Self::Seconds),
78            "millis" => Some(Self::Millis),
79            "micros" => Some(Self::Micros),
80            "nanos" => Some(Self::Nanos),
81            _ => None,
82        }
83    }
84
85    /// Nanoseconds per one unit (used for lossless scaling).
86    const fn nanos_per(self) -> i64 {
87        match self {
88            Self::Seconds => 1_000_000_000,
89            Self::Millis => 1_000_000,
90            Self::Micros => 1_000,
91            Self::Nanos => 1,
92        }
93    }
94}
95
96/// Per-column extraction strategy, pre-computed at Ring 2.
97#[derive(Debug, Clone)]
98enum ColumnExtraction {
99    /// Extract field from the default path target (json.path or root).
100    DefaultPath,
101    /// Extract field from a custom absolute path from root.
102    CustomPath { segments: Vec<String> },
103}
104
105/// JSON decoder configuration.
106#[derive(Debug, Clone)]
107pub struct JsonDecoderConfig {
108    /// How to handle fields present in JSON but absent from the schema.
109    pub unknown_fields: UnknownFieldStrategy,
110
111    /// How to handle type mismatches.
112    pub type_mismatch: TypeMismatchStrategy,
113
114    /// Timestamp format patterns to try when parsing string values
115    /// into Timestamp columns. Tried in order; first match wins.
116    /// Use `"iso8601"` for RFC 3339 / ISO 8601 auto-detection.
117    pub timestamp_formats: Vec<String>,
118
119    /// Whether to encode nested objects as JSONB binary format
120    /// instead of JSON-serialized Utf8. When true, nested objects
121    /// become `LargeBinary` columns with JSONB encoding.
122    pub nested_as_jsonb: bool,
123
124    /// Dot-separated path to navigate before field extraction.
125    /// e.g. `"data"` → navigate into `{"data": {...}}` before extraction.
126    /// e.g. `"data.trade"` → navigate into `{"data":{"trade":{...}}}`.
127    pub json_path: Option<Vec<String>>,
128
129    /// Per-column absolute path overrides: `column_name` → path segments.
130    /// Parsed from `json.column.<name> = 'path.to.field'` options.
131    /// These paths are absolute from the root, ignoring `json_path`.
132    pub json_column_paths: HashMap<String, Vec<String>>,
133
134    /// Per-column numeric-timestamp epoch unit, from
135    /// `json.column.<name>.epoch_unit`. Absent ⇒ [`EpochUnit::Millis`].
136    pub numeric_timestamp_units: HashMap<String, EpochUnit>,
137
138    /// Column names for array-to-rows expansion.
139    /// When set, the target (after `json_path`) must be an array.
140    /// Each element becomes a row; positional names map array elements to columns.
141    pub json_explode: Option<Vec<String>>,
142}
143
144impl Default for JsonDecoderConfig {
145    fn default() -> Self {
146        Self {
147            unknown_fields: UnknownFieldStrategy::Ignore,
148            type_mismatch: TypeMismatchStrategy::Coerce,
149            timestamp_formats: vec![
150                "iso8601".into(),
151                "%Y-%m-%dT%H:%M:%S%.fZ".into(),
152                "%Y-%m-%dT%H:%M:%S%.f%:z".into(),
153                "%Y-%m-%d %H:%M:%S%.f".into(),
154                "%Y-%m-%d %H:%M:%S".into(),
155            ],
156            nested_as_jsonb: false,
157            json_path: None,
158            json_column_paths: HashMap::new(),
159            numeric_timestamp_units: HashMap::new(),
160            json_explode: None,
161        }
162    }
163}
164
165impl JsonDecoderConfig {
166    /// Build config from a [`ConnectorConfig`](crate::config::ConnectorConfig).
167    ///
168    /// Parses `json.path`, `json.column.*`, `json.explode`,
169    /// `schema.enforcement`, and `nested.as.jsonb` properties.
170    /// Epoch-unit options are validated against the frozen output schema.
171    /// Called once at Ring 2 (`CREATE SOURCE` time).
172    ///
173    /// # Errors
174    ///
175    /// Returns an invalid-configuration error when an epoch unit is unknown or
176    /// targets a missing or non-timestamp column.
177    pub fn from_connector_config(
178        config: &crate::config::ConnectorConfig,
179        schema: &SchemaRef,
180    ) -> SchemaResult<Self> {
181        let mut cfg = Self::default();
182
183        if let Some(path) = config.get("json.path") {
184            cfg.json_path = Some(parse_path_option("json.path", path)?);
185        }
186
187        let mut col_props: Vec<_> = config
188            .properties_with_prefix("json.column.")
189            .into_iter()
190            .collect();
191        col_props.sort_unstable_by(|left, right| left.0.cmp(&right.0));
192        for (col_name, val) in col_props {
193            // `json.column.<col>.epoch_unit = '<unit>'` configures numeric
194            // timestamp scaling for <col>; everything else is a path.
195            // (Column names are SQL identifiers — no dots — so the suffix
196            // split is unambiguous against dotted path *values*.)
197            if let Some(col) = col_name.strip_suffix(".epoch_unit") {
198                let unit =
199                    EpochUnit::from_str_opt(&val).ok_or_else(|| SchemaError::InvalidConfig {
200                        key: format!("json.column.{col}.epoch_unit"),
201                        message: format!(
202                            "invalid epoch unit '{val}'; expected seconds, millis, micros, or nanos"
203                        ),
204                    })?;
205                let field =
206                    schema
207                        .field_with_name(col)
208                        .map_err(|_| SchemaError::InvalidConfig {
209                            key: format!("json.column.{col}.epoch_unit"),
210                            message: format!("column '{col}' is not present in the source schema"),
211                        })?;
212                if !matches!(field.data_type(), DataType::Timestamp(_, _)) {
213                    return Err(SchemaError::InvalidConfig {
214                        key: format!("json.column.{col}.epoch_unit"),
215                        message: format!(
216                            "column '{col}' must be a Timestamp, found {:?}",
217                            field.data_type()
218                        ),
219                    });
220                }
221                cfg.numeric_timestamp_units.insert(col.to_string(), unit);
222                continue;
223            }
224            schema
225                .field_with_name(&col_name)
226                .map_err(|_| SchemaError::InvalidConfig {
227                    key: format!("json.column.{col_name}"),
228                    message: format!("column '{col_name}' is not present in the source schema"),
229                })?;
230            cfg.json_column_paths.insert(
231                col_name.clone(),
232                parse_path_option(&format!("json.column.{col_name}"), &val)?,
233            );
234        }
235
236        if let Some(explode) = config.get("json.explode") {
237            let columns: Vec<String> = explode
238                .split(',')
239                .map(str::trim)
240                .map(ToString::to_string)
241                .collect();
242            if columns.is_empty() || columns.iter().any(String::is_empty) {
243                return Err(SchemaError::InvalidConfig {
244                    key: "json.explode".into(),
245                    message: "expected one or more comma-separated schema columns".into(),
246                });
247            }
248            let mut seen = std::collections::HashSet::with_capacity(columns.len());
249            for column in &columns {
250                if !seen.insert(column) {
251                    return Err(SchemaError::InvalidConfig {
252                        key: "json.explode".into(),
253                        message: format!("column '{column}' is listed more than once"),
254                    });
255                }
256                schema
257                    .field_with_name(column)
258                    .map_err(|_| SchemaError::InvalidConfig {
259                        key: "json.explode".into(),
260                        message: format!("column '{column}' is not present in the source schema"),
261                    })?;
262            }
263            cfg.json_explode = Some(columns);
264        }
265
266        if let Some(enforcement) = config.get("schema.enforcement") {
267            cfg.type_mismatch = TypeMismatchStrategy::from_enforcement_str(enforcement)
268                .ok_or_else(|| SchemaError::InvalidConfig {
269                    key: "schema.enforcement".into(),
270                    message: format!(
271                        "invalid value '{enforcement}'; expected lenient, coerce, or strict"
272                    ),
273                })?;
274        }
275        if let Some(v) = config.get("nested.as.jsonb") {
276            cfg.nested_as_jsonb = match v.to_ascii_lowercase().as_str() {
277                "true" => true,
278                "false" => false,
279                _ => {
280                    return Err(SchemaError::InvalidConfig {
281                        key: "nested.as.jsonb".into(),
282                        message: format!("invalid boolean '{v}'; expected true or false"),
283                    });
284                }
285            };
286        }
287        Ok(cfg)
288    }
289}
290
291fn parse_path_option(key: &str, value: &str) -> Result<Vec<String>, SchemaError> {
292    let segments: Vec<String> = value
293        .split('.')
294        .map(str::trim)
295        .map(ToString::to_string)
296        .collect();
297    if segments.is_empty() || segments.iter().any(String::is_empty) {
298        return Err(SchemaError::InvalidConfig {
299            key: key.into(),
300            message: "path must contain non-empty dot-separated segments".into(),
301        });
302    }
303    Ok(segments)
304}
305
306/// Decodes JSON byte payloads into Arrow `RecordBatch`es.
307///
308/// # Ring Placement
309///
310/// - **Ring 1**: `decode_batch()` — parse JSON, build columnar Arrow output
311/// - **Ring 2**: Construction (`new` / `with_config`) — one-time setup
312pub struct JsonDecoder {
313    /// Frozen output schema.
314    schema: SchemaRef,
315    /// Decoder configuration.
316    config: JsonDecoderConfig,
317    /// Pre-computed field index map: field name → column index.
318    field_indices: Vec<(String, usize)>,
319    /// Cumulative type mismatch count (diagnostics).
320    mismatch_count: AtomicU64,
321    /// Ring 2 pre-computed: extraction strategy per schema column.
322    column_extractions: Vec<ColumnExtraction>,
323    /// Ring 2 pre-computed: numeric epoch unit per schema column
324    /// (aligned to `schema.fields()`; default [`EpochUnit::Millis`]).
325    column_epoch_units: Vec<EpochUnit>,
326    /// Ring 2 pre-computed: explode position → schema column index.
327    /// `Some` when `json.explode` is configured; each entry maps an
328    /// explode position to the schema column index (or `None` if
329    /// the explode name doesn't match any schema column).
330    explode_col_indices: Option<Vec<Option<usize>>>,
331}
332
333#[allow(clippy::missing_fields_in_debug)]
334impl std::fmt::Debug for JsonDecoder {
335    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
336        f.debug_struct("JsonDecoder")
337            .field("schema", &self.schema)
338            .field("config", &self.config)
339            .field(
340                "mismatch_count",
341                &self.mismatch_count.load(Ordering::Relaxed),
342            )
343            .finish()
344    }
345}
346
347impl JsonDecoder {
348    /// Creates a new JSON decoder for the given Arrow schema.
349    #[must_use]
350    pub fn new(schema: SchemaRef) -> Self {
351        Self::with_config(schema, JsonDecoderConfig::default())
352    }
353
354    /// Creates a new JSON decoder with custom configuration.
355    #[must_use]
356    pub fn with_config(schema: SchemaRef, config: JsonDecoderConfig) -> Self {
357        let field_indices: Vec<(String, usize)> = schema
358            .fields()
359            .iter()
360            .enumerate()
361            .map(|(i, f)| (f.name().clone(), i))
362            .collect();
363
364        let column_extractions: Vec<ColumnExtraction> = schema
365            .fields()
366            .iter()
367            .map(|f| {
368                let col_name = f.name();
369                if let Some(path_segments) = config.json_column_paths.get(col_name.as_str()) {
370                    ColumnExtraction::CustomPath {
371                        segments: path_segments.clone(),
372                    }
373                } else {
374                    ColumnExtraction::DefaultPath
375                }
376            })
377            .collect();
378
379        let column_epoch_units: Vec<EpochUnit> = schema
380            .fields()
381            .iter()
382            .map(|f| {
383                config
384                    .numeric_timestamp_units
385                    .get(f.name().as_str())
386                    .copied()
387                    .unwrap_or_default()
388            })
389            .collect();
390
391        let explode_col_indices = config.json_explode.as_ref().map(|names| {
392            names
393                .iter()
394                .map(|name| {
395                    field_indices
396                        .iter()
397                        .find(|(n, _)| n == name)
398                        .map(|(_, idx)| *idx)
399                })
400                .collect()
401        });
402
403        Self {
404            schema,
405            config,
406            field_indices,
407            mismatch_count: AtomicU64::new(0),
408            column_extractions,
409            column_epoch_units,
410            explode_col_indices,
411        }
412    }
413
414    /// Returns the cumulative type mismatch count.
415    pub fn mismatch_count(&self) -> u64 {
416        self.mismatch_count.load(Ordering::Relaxed)
417    }
418}
419
420impl FormatDecoder for JsonDecoder {
421    fn output_schema(&self) -> SchemaRef {
422        self.schema.clone()
423    }
424
425    fn decode_batch(&self, records: &[RawRecord]) -> SchemaResult<RecordBatch> {
426        let values: Vec<&[u8]> = records.iter().map(|r| r.value.as_slice()).collect();
427        self.decode_slices(&values)
428    }
429
430    fn format_name(&self) -> &'static str {
431        "json"
432    }
433}
434
435#[cfg(test)]
436mod tests;