laminar_connectors/schema/json/decoder/
mod.rs1#![allow(clippy::disallowed_types)] use 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24pub enum UnknownFieldStrategy {
25 Ignore,
27 CollectExtra,
29 Reject,
31}
32
33#[derive(Debug, Clone, Copy, PartialEq, Eq)]
35pub enum TypeMismatchStrategy {
36 Null,
38 Coerce,
40 Reject,
42}
43
44impl TypeMismatchStrategy {
45 #[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#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
60pub enum EpochUnit {
61 Seconds,
63 #[default]
65 Millis,
66 Micros,
68 Nanos,
70}
71
72impl EpochUnit {
73 #[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 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#[derive(Debug, Clone)]
98enum ColumnExtraction {
99 DefaultPath,
101 CustomPath { segments: Vec<String> },
103}
104
105#[derive(Debug, Clone)]
107pub struct JsonDecoderConfig {
108 pub unknown_fields: UnknownFieldStrategy,
110
111 pub type_mismatch: TypeMismatchStrategy,
113
114 pub timestamp_formats: Vec<String>,
118
119 pub nested_as_jsonb: bool,
123
124 pub json_path: Option<Vec<String>>,
128
129 pub json_column_paths: HashMap<String, Vec<String>>,
133
134 pub numeric_timestamp_units: HashMap<String, EpochUnit>,
137
138 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 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 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
306pub struct JsonDecoder {
313 schema: SchemaRef,
315 config: JsonDecoderConfig,
317 field_indices: Vec<(String, usize)>,
319 mismatch_count: AtomicU64,
321 column_extractions: Vec<ColumnExtraction>,
323 column_epoch_units: Vec<EpochUnit>,
326 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 #[must_use]
350 pub fn new(schema: SchemaRef) -> Self {
351 Self::with_config(schema, JsonDecoderConfig::default())
352 }
353
354 #[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 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;