1use super::{AvroCodec, AvroOptions};
53use crate::error::FaucetError;
54use apache_avro::schema::{DecimalSchema, RecordField, UnionSchema};
55use apache_avro::types::Value as Av;
56use apache_avro::{Codec, Schema};
57use chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime, SecondsFormat, Timelike, Utc};
58use serde_json::{Map, Value, json};
59use std::collections::{HashMap, HashSet};
60
61pub const ORIGINAL_NAME_ATTR: &str = "faucet.name";
63
64pub const INFERRED_RECORD_NAME: &str = "faucet_record";
66
67pub const VALUE_FIELD: &str = "value";
70
71fn codec(c: AvroCodec) -> Codec {
72 match c {
73 AvroCodec::Null => Codec::Null,
74 AvroCodec::Deflate => Codec::Deflate(Default::default()),
75 AvroCodec::Snappy => Codec::Snappy,
76 AvroCodec::Zstd => Codec::Zstandard(Default::default()),
77 }
78}
79
80impl AvroOptions {
81 pub fn parsed_schema(&self) -> Result<Option<Schema>, FaucetError> {
83 self.schema.as_ref().map(parse_schema).transpose()
84 }
85}
86
87pub fn parse_schema(v: &Value) -> Result<Schema, FaucetError> {
89 Schema::parse(v)
90 .map_err(|e| FaucetError::Config(format!("avro.schema is not a valid Avro schema: {e}")))
91}
92
93pub fn decode(bytes: &[u8], opts: &AvroOptions) -> Result<Vec<Value>, FaucetError> {
95 let reader_schema = opts.parsed_schema()?;
96 let mut out = Vec::new();
97 read_records(bytes, reader_schema.as_ref(), usize::MAX, &mut |chunk| {
98 out.extend(chunk);
99 Ok(())
100 })
101 .map_err(|e| FaucetError::Source(format!("avro: {e}")))?;
102 Ok(out)
103}
104
105pub fn writer_schema<R: std::io::Read>(reader: R) -> Result<Schema, FaucetError> {
107 let r = apache_avro::Reader::new(reader)
108 .map_err(|e| FaucetError::Source(format!("avro header: {e}")))?;
109 Ok(r.writer_schema().clone())
110}
111
112pub fn read_records<R: std::io::Read>(
115 reader: R,
116 reader_schema: Option<&Schema>,
117 chunk: usize,
118 f: &mut dyn FnMut(Vec<Value>) -> Result<(), FaucetError>,
119) -> Result<Schema, FaucetError> {
120 read_records_checked(reader, reader_schema, None, chunk, f)
121}
122
123pub(crate) type WriterCheck<'a> = dyn Fn(&Schema) -> Result<(), FaucetError> + 'a;
125
126pub(crate) fn read_records_checked<R: std::io::Read>(
128 reader: R,
129 reader_schema: Option<&Schema>,
130 check: Option<&WriterCheck<'_>>,
131 chunk: usize,
132 f: &mut dyn FnMut(Vec<Value>) -> Result<(), FaucetError>,
133) -> Result<Schema, FaucetError> {
134 read_with(
135 reader,
136 reader_schema,
137 check,
138 chunk,
139 Mode::Record,
140 &mut |chunk, _| f(chunk),
141 )
142}
143
144#[derive(Debug, Clone, Copy, PartialEq, Eq)]
147pub(crate) enum Mode {
148 Record,
149 Arrow,
150}
151
152pub(crate) type ChunkSink<'a> = dyn FnMut(Vec<Value>, &Schema) -> Result<(), FaucetError> + 'a;
154
155pub(crate) fn read_with<R: std::io::Read>(
156 reader: R,
157 reader_schema: Option<&Schema>,
158 check: Option<&WriterCheck<'_>>,
159 chunk: usize,
160 mode: Mode,
161 f: &mut ChunkSink<'_>,
162) -> Result<Schema, FaucetError> {
163 let builder = apache_avro::Reader::builder(reader);
164 let r = match reader_schema {
165 Some(s) => builder.reader_schema(s).build(),
166 None => builder.build(),
167 }
168 .map_err(|e| FaucetError::Source(format!("avro header: {e}")))?;
169 if let Some(check) = check {
170 check(r.writer_schema())?;
171 }
172 let schema = reader_schema
173 .cloned()
174 .unwrap_or_else(|| r.writer_schema().clone());
175 let names = named(&schema);
176 let chunk = chunk.max(1);
177 let mut buf = Vec::with_capacity(chunk.min(4096));
178 for datum in r {
179 let datum = datum.map_err(|e| FaucetError::Source(format!("avro datum: {e}")))?;
180 buf.push(root_to_json(&datum, &schema, &names, mode)?);
181 if buf.len() >= chunk {
182 f(std::mem::take(&mut buf), &schema)?;
183 }
184 }
185 if !buf.is_empty() {
186 f(buf, &schema)?;
187 }
188 Ok(schema)
189}
190
191type Names<'a> = HashMap<String, &'a Schema>;
192
193fn named(schema: &Schema) -> Names<'_> {
194 let mut out = HashMap::new();
195 collect_named(schema, &mut out);
196 out
197}
198
199fn collect_named<'a>(s: &'a Schema, out: &mut Names<'a>) {
200 match s {
201 Schema::Record(r) => {
202 out.insert(r.name.fullname(None), s);
203 out.entry(r.name.name().to_string()).or_insert(s);
204 for f in &r.fields {
205 collect_named(&f.schema, out);
206 }
207 }
208 Schema::Enum(e) => {
209 out.insert(e.name.fullname(None), s);
210 out.entry(e.name.name().to_string()).or_insert(s);
211 }
212 Schema::Fixed(x) | Schema::Duration(x) => {
213 out.insert(x.name.fullname(None), s);
214 out.entry(x.name.name().to_string()).or_insert(s);
215 }
216 Schema::Array(a) => collect_named(&a.items, out),
217 Schema::Map(m) => collect_named(&m.types, out),
218 Schema::Union(u) => u.variants().iter().for_each(|v| collect_named(v, out)),
219 _ => {}
220 }
221}
222
223fn resolve<'a>(s: &'a Schema, names: &Names<'a>) -> Result<&'a Schema, FaucetError> {
224 match s {
225 Schema::Ref { name } => names
226 .get(&name.fullname(None))
227 .or_else(|| names.get(name.name()))
228 .copied()
229 .ok_or_else(|| {
230 FaucetError::Source(format!(
231 "avro: unresolved schema reference `{}`",
232 name.fullname(None)
233 ))
234 }),
235 other => Ok(other),
236 }
237}
238
239fn ref_name(s: &Schema) -> Option<String> {
240 match s {
241 Schema::Ref { name } => Some(name.fullname(None)),
242 Schema::Record(r) => Some(r.name.fullname(None)),
243 _ => None,
244 }
245}
246
247pub(crate) fn is_complex_union(u: &UnionSchema) -> bool {
249 u.variants()
250 .iter()
251 .filter(|v| !matches!(v, Schema::Null))
252 .count()
253 > 1
254}
255
256pub fn datum_to_json(datum: &Av, schema: &Schema) -> Result<Value, FaucetError> {
261 let names = named(schema);
262 to_json(datum, schema, &names, Mode::Record, &mut Vec::new())
263}
264
265pub fn datum_from_json(value: &Value, schema: &Schema) -> Result<Av, String> {
269 let names = named(schema);
270 from_json(value, schema, &names, "")
271}
272
273fn root_to_json(
274 v: &Av,
275 schema: &Schema,
276 names: &Names<'_>,
277 mode: Mode,
278) -> Result<Value, FaucetError> {
279 let mut stack = Vec::new();
280 let resolved = resolve(schema, names)?;
281 let j = to_json(v, schema, names, mode, &mut stack)?;
282 Ok(match (resolved, j) {
283 (Schema::Record(_), j @ Value::Object(_)) => j,
284 (_, j) => {
285 let mut m = Map::new();
286 m.insert(VALUE_FIELD.into(), j);
287 Value::Object(m)
288 }
289 })
290}
291
292fn to_json(
293 v: &Av,
294 schema: &Schema,
295 names: &Names<'_>,
296 mode: Mode,
297 stack: &mut Vec<String>,
298) -> Result<Value, FaucetError> {
299 let s = resolve(schema, names)?;
300 if mode == Mode::Arrow
301 && let Some(n) = ref_name(schema)
302 && stack.contains(&n)
303 {
304 let inner = to_json(v, s, names, Mode::Record, &mut Vec::new())?;
305 return Ok(Value::String(inner.to_string()));
306 }
307 Ok(match (v, s) {
308 (Av::Union(i, inner), Schema::Union(u)) => {
309 let branch = u.variants().get(*i as usize).ok_or_else(|| {
310 FaucetError::Source(format!("avro: union branch {i} out of range"))
311 })?;
312 if mode == Mode::Arrow && is_complex_union(u) {
313 if matches!(**inner, Av::Null) {
314 return Ok(Value::Null);
315 }
316 let j = to_json(inner, branch, names, Mode::Record, &mut Vec::new())?;
317 return Ok(Value::String(j.to_string()));
318 }
319 to_json(inner, branch, names, mode, stack)?
320 }
321 (Av::Union(_, inner), other) => to_json(inner, other, names, mode, stack)?,
322 (Av::Null, _) => Value::Null,
323 (Av::Boolean(b), _) => Value::Bool(*b),
324 (Av::Int(i), _) => json!(i),
325 (Av::Long(i), _) => json!(i),
326 (Av::Float(f), _) => float_json(f64::from(*f)),
327 (Av::Double(f), _) => float_json(*f),
328 (Av::Bytes(b), _) | (Av::Fixed(_, b), _) => Value::String(hex(b)),
329 (Av::String(t), _) => Value::String(t.clone()),
330 (Av::Enum(_, sym), _) => Value::String(sym.clone()),
331 (Av::Array(items), s) => {
332 let item_schema = match s {
333 Schema::Array(a) => &*a.items,
334 _ => s,
335 };
336 Value::Array(
337 items
338 .iter()
339 .map(|i| to_json(i, item_schema, names, mode, stack))
340 .collect::<Result<_, _>>()?,
341 )
342 }
343 (Av::Map(m), s) => {
344 let value_schema = match s {
345 Schema::Map(m) => &*m.types,
346 _ => s,
347 };
348 let mut keys: Vec<&String> = m.keys().collect();
349 keys.sort();
350 let mut out = Map::new();
351 for k in keys {
352 out.insert(k.clone(), to_json(&m[k], value_schema, names, mode, stack)?);
353 }
354 Value::Object(out)
355 }
356 (Av::Record(fields), s) => {
357 let rec = match s {
358 Schema::Record(r) => Some(r),
359 _ => None,
360 };
361 let pushed = rec.map(|r| r.name.fullname(None));
362 if let Some(n) = &pushed {
363 stack.push(n.clone());
364 }
365 let mut out = Map::new();
366 for (name, fv) in fields {
367 let fs = rec
368 .and_then(|r| r.lookup.get(name).map(|i| &r.fields[*i].schema))
369 .unwrap_or(&Schema::Null);
370 out.insert(name.clone(), to_json(fv, fs, names, mode, stack)?);
371 }
372 if pushed.is_some() {
373 stack.pop();
374 }
375 Value::Object(out)
376 }
377 (Av::Date(d), _) => Value::String(date_string(*d)?),
378 (Av::Decimal(d), s) => {
379 let scale = match s {
380 Schema::Decimal(DecimalSchema { scale, .. }) => *scale,
381 _ => 0,
382 };
383 let bytes = <Vec<u8>>::try_from(d)
384 .map_err(|e| FaucetError::Source(format!("avro decimal: {e}")))?;
385 Value::String(decimal_to_string(&bytes, scale))
386 }
387 (Av::BigDecimal(d), _) => Value::String(d.to_string()),
388 (Av::TimeMillis(ms), _) => Value::String(time_string(i64::from(*ms) * 1_000_000, 3)?),
389 (Av::TimeMicros(us), _) => Value::String(time_string(us.saturating_mul(1_000), 6)?),
390 (Av::TimestampMillis(t), _) => {
391 Value::String(utc_string(*t, 1_000_000, SecondsFormat::Millis)?)
392 }
393 (Av::TimestampMicros(t), _) => Value::String(utc_string(*t, 1_000, SecondsFormat::Micros)?),
394 (Av::TimestampNanos(t), _) => Value::String(utc_string(*t, 1, SecondsFormat::Nanos)?),
395 (Av::LocalTimestampMillis(t), _) => Value::String(local_string(*t, 1_000_000, 3)?),
396 (Av::LocalTimestampMicros(t), _) => Value::String(local_string(*t, 1_000, 6)?),
397 (Av::LocalTimestampNanos(t), _) => Value::String(local_string(*t, 1, 9)?),
398 (Av::Duration(d), _) => json!({
399 "months": u32::from(d.months()),
400 "days": u32::from(d.days()),
401 "millis": u32::from(d.millis()),
402 }),
403 (Av::Uuid(u), _) => Value::String(u.to_string()),
404 })
405}
406
407fn float_json(f: f64) -> Value {
408 if f.is_nan() {
409 Value::String("NaN".into())
410 } else if f.is_infinite() {
411 Value::String(if f > 0.0 { "Infinity" } else { "-Infinity" }.into())
412 } else {
413 json!(f)
414 }
415}
416
417fn hex(b: &[u8]) -> String {
418 const DIGITS: &[u8; 16] = b"0123456789abcdef";
419 let mut s = String::with_capacity(b.len() * 2);
420 for byte in b {
421 s.push(DIGITS[(byte >> 4) as usize] as char);
422 s.push(DIGITS[(byte & 0xf) as usize] as char);
423 }
424 s
425}
426
427fn unhex(s: &str) -> Result<Vec<u8>, String> {
428 if !s.len().is_multiple_of(2) {
429 return Err(format!("{s:?} is not hex (odd length)"));
430 }
431 (0..s.len())
432 .step_by(2)
433 .map(|i| u8::from_str_radix(&s[i..i + 2], 16).map_err(|_| format!("{s:?} is not hex")))
434 .collect()
435}
436
437const EPOCH: NaiveDate = NaiveDate::from_ymd_opt(1970, 1, 1).expect("valid epoch");
438
439fn date_string(days: i32) -> Result<String, FaucetError> {
440 EPOCH
441 .checked_add_signed(chrono::Duration::days(i64::from(days)))
442 .map(|d| d.format("%Y-%m-%d").to_string())
443 .ok_or_else(|| FaucetError::Source(format!("avro date {days} is out of range")))
444}
445
446fn time_string(nanos: i64, digits: usize) -> Result<String, FaucetError> {
447 let secs = u32::try_from(nanos.div_euclid(1_000_000_000)).ok();
448 let frac = nanos.rem_euclid(1_000_000_000) as u32;
449 let t = secs
450 .and_then(|s| NaiveTime::from_num_seconds_from_midnight_opt(s, frac))
451 .ok_or_else(|| FaucetError::Source(format!("avro time {nanos}ns is out of range")))?;
452 Ok(format!(
453 "{}.{:0w$}",
454 t.format("%H:%M:%S"),
455 t.nanosecond() / 10u32.pow(9 - digits as u32),
456 w = digits
457 ))
458}
459
460fn instant(t: i64, nanos_per_unit: i64) -> Option<DateTime<Utc>> {
461 let n = i128::from(t) * i128::from(nanos_per_unit);
462 let secs = i64::try_from(n.div_euclid(1_000_000_000)).ok()?;
463 DateTime::from_timestamp(secs, n.rem_euclid(1_000_000_000) as u32)
464}
465
466fn utc_string(t: i64, nanos_per_unit: i64, fmt: SecondsFormat) -> Result<String, FaucetError> {
467 instant(t, nanos_per_unit)
468 .map(|d| d.to_rfc3339_opts(fmt, true))
469 .ok_or_else(|| FaucetError::Source(format!("avro timestamp {t} is out of range")))
470}
471
472fn local_string(t: i64, nanos_per_unit: i64, digits: usize) -> Result<String, FaucetError> {
473 let d = instant(t, nanos_per_unit)
474 .ok_or_else(|| FaucetError::Source(format!("avro local timestamp {t} is out of range")))?
475 .naive_utc();
476 Ok(format!(
477 "{}.{:0w$}",
478 d.format("%Y-%m-%dT%H:%M:%S"),
479 d.nanosecond() / 10u32.pow(9 - digits as u32),
480 w = digits
481 ))
482}
483
484pub(crate) fn decimal_to_string(bytes: &[u8], scale: usize) -> String {
486 let negative = bytes.first().is_some_and(|b| b & 0x80 != 0);
487 let mut mag: Vec<u8> = if negative {
488 negate(bytes)
489 } else {
490 bytes.to_vec()
491 };
492 let mut digits = Vec::new();
493 while mag.iter().any(|b| *b != 0) {
494 let mut rem: u32 = 0;
495 for b in mag.iter_mut() {
496 let cur = (rem << 8) | u32::from(*b);
497 *b = (cur / 10) as u8;
498 rem = cur % 10;
499 }
500 digits.push(b'0' + rem as u8);
501 }
502 while digits.len() <= scale {
503 digits.push(b'0');
504 }
505 digits.reverse();
506 let (int, frac) = digits.split_at(digits.len() - scale);
507 let mut s = String::new();
508 if negative {
509 s.push('-');
510 }
511 s.push_str(std::str::from_utf8(int).expect("ascii digits"));
512 if scale > 0 {
513 s.push('.');
514 s.push_str(std::str::from_utf8(frac).expect("ascii digits"));
515 }
516 s
517}
518
519fn negate(bytes: &[u8]) -> Vec<u8> {
520 let mut out: Vec<u8> = bytes.iter().map(|b| !b).collect();
521 for b in out.iter_mut().rev() {
522 let (v, carry) = b.overflowing_add(1);
523 *b = v;
524 if !carry {
525 break;
526 }
527 }
528 out
529}
530
531pub(crate) fn decimal_from_str(
537 text: &str,
538 scale: usize,
539 precision: usize,
540) -> Result<Vec<u8>, String> {
541 let t = text.trim();
542 let (negative, body) = match t.strip_prefix('-') {
543 Some(rest) => (true, rest),
544 None => (false, t.strip_prefix('+').unwrap_or(t)),
545 };
546 let (mantissa, exp) = match body.find(['e', 'E']) {
547 Some(i) => (
548 &body[..i],
549 body[i + 1..]
550 .parse::<i64>()
551 .map_err(|_| format!("{text:?} is not a decimal"))?,
552 ),
553 None => (body, 0),
554 };
555 let (int, frac) = mantissa.split_once('.').unwrap_or((mantissa, ""));
556 if int.is_empty() && frac.is_empty()
557 || !int.bytes().chain(frac.bytes()).all(|b| b.is_ascii_digit())
558 {
559 return Err(format!("{text:?} is not a decimal"));
560 }
561 let mut digits: Vec<u8> = int.bytes().chain(frac.bytes()).map(|b| b - b'0').collect();
562 let out_of_range = || format!("{text:?} is out of range for precision {precision}");
563 let wanted = (int.len() as i64)
564 .checked_add(exp)
565 .and_then(|p| p.checked_add(scale as i64))
566 .ok_or_else(out_of_range)?;
567 match digits.iter().position(|d| *d != 0) {
570 None => digits.clear(),
571 Some(first)
572 if wanted > first as i64 && (wanted - first as i64) as u128 > precision as u128 =>
573 {
574 return Err(format!(
575 "{text:?} needs {} digits, more than the precision {precision}",
576 wanted - first as i64
577 ));
578 }
579 Some(_) => {}
580 }
581 if digits.is_empty() {
582 } else if wanted < 0 {
584 return Err(format!("{text:?} has more than {scale} fractional digits"));
585 } else {
586 let wanted = wanted as usize;
587 if digits.len() > wanted {
588 if digits[wanted..].iter().any(|d| *d != 0) {
589 return Err(format!("{text:?} has more than {scale} fractional digits"));
590 }
591 digits.truncate(wanted);
592 } else {
593 digits.resize(wanted, 0);
594 }
595 }
596 let first = digits.iter().position(|d| *d != 0).unwrap_or(digits.len());
597 let digits = &digits[first..];
598 if digits.len() > precision {
599 return Err(format!(
600 "{text:?} needs {} digits, more than the precision {precision}",
601 digits.len()
602 ));
603 }
604 let mut mag: Vec<u8> = vec![0];
605 for d in digits {
606 let mut carry = u32::from(*d);
607 for b in mag.iter_mut().rev() {
608 let cur = u32::from(*b) * 10 + carry;
609 *b = (cur & 0xff) as u8;
610 carry = cur >> 8;
611 }
612 while carry > 0 {
613 mag.insert(0, (carry & 0xff) as u8);
614 carry >>= 8;
615 }
616 }
617 if mag[0] & 0x80 != 0 {
618 mag.insert(0, 0);
619 }
620 while mag.len() > 1 && mag[0] == 0 && mag[1] & 0x80 == 0 {
621 mag.remove(0);
622 }
623 Ok(if negative && mag.iter().any(|b| *b != 0) {
624 negate(&mag)
625 } else {
626 mag
627 })
628}
629
630fn sign_extend(bytes: Vec<u8>, size: usize) -> Result<Vec<u8>, String> {
631 if bytes.len() > size {
632 return Err(format!(
633 "decimal needs {} bytes, the fixed holds {size}",
634 bytes.len()
635 ));
636 }
637 let fill = if bytes.first().is_some_and(|b| b & 0x80 != 0) {
638 0xff
639 } else {
640 0
641 };
642 let mut out = vec![fill; size - bytes.len()];
643 out.extend(bytes);
644 Ok(out)
645}
646
647pub fn encode(records: &[Value], opts: &AvroOptions) -> Result<Vec<u8>, FaucetError> {
651 let schema = match opts.parsed_schema()? {
652 Some(s) => s,
653 None => infer_schema(records)?,
654 };
655 let names = named(&schema);
656 let mut w = apache_avro::Writer::with_codec(&schema, Vec::new(), codec(opts.codec))
657 .map_err(|e| FaucetError::Sink(format!("avro writer: {e}")))?;
658 for (i, r) in records.iter().enumerate() {
659 let datum = root_from_json(r, &schema, &names)
660 .map_err(|e| FaucetError::Sink(format!("avro: record {i}: {e}")))?;
661 w.append_value_ref(&datum)
662 .map_err(|e| FaucetError::Sink(format!("avro: record {i}: {e}")))?;
663 }
664 w.into_inner()
665 .map_err(|e| FaucetError::Sink(format!("avro writer: {e}")))
666}
667
668fn root_from_json(v: &Value, schema: &Schema, names: &Names<'_>) -> Result<Av, String> {
669 let resolved = resolve(schema, names).map_err(|e| e.to_string())?;
670 match resolved {
671 Schema::Record(_) => from_json(v, schema, names, ""),
672 _ => {
673 let inner = v.get(VALUE_FIELD).unwrap_or(v);
674 from_json(inner, schema, names, "")
675 }
676 }
677}
678
679fn at(path: &str) -> String {
680 if path.is_empty() {
681 "the record".into()
682 } else {
683 format!("`{path}`")
684 }
685}
686
687fn from_json(v: &Value, schema: &Schema, names: &Names<'_>, path: &str) -> Result<Av, String> {
688 let s = resolve(schema, names).map_err(|e| e.to_string())?;
689 let want = |what: &str| format!("{}: expected {what}, got {}", at(path), short(v));
690 Ok(match s {
691 Schema::Null => match v {
692 Value::Null => Av::Null,
693 _ => return Err(want("null")),
694 },
695 Schema::Boolean => Av::Boolean(v.as_bool().ok_or_else(|| want("a boolean"))?),
696 Schema::Int => {
697 let i = v.as_i64().ok_or_else(|| want("an integer"))?;
698 Av::Int(
699 i32::try_from(i)
700 .map_err(|_| format!("{}: {i} does not fit an Avro int", at(path)))?,
701 )
702 }
703 Schema::Long => Av::Long(v.as_i64().ok_or_else(|| want("an integer"))?),
704 Schema::Float => Av::Float(float_from(v).ok_or_else(|| want("a number"))? as f32),
705 Schema::Double => Av::Double(float_from(v).ok_or_else(|| want("a number"))?),
706 Schema::Bytes => Av::Bytes(bytes_from(v).map_err(|e| format!("{}: {e}", at(path)))?),
707 Schema::String => Av::String(match v {
708 Value::String(t) => t.clone(),
709 Value::Null => return Err(want("a string")),
710 other => super::cell_text(other),
711 }),
712 Schema::Array(a) => {
713 let items = v.as_array().ok_or_else(|| want("an array"))?;
714 Av::Array(
715 items
716 .iter()
717 .enumerate()
718 .map(|(i, x)| from_json(x, &a.items, names, &format!("{path}[{i}]")))
719 .collect::<Result<_, _>>()?,
720 )
721 }
722 Schema::Map(m) => {
723 let obj = v.as_object().ok_or_else(|| want("an object"))?;
724 Av::Map(
725 obj.iter()
726 .map(|(k, x)| Ok((k.clone(), from_json(x, &m.types, names, &join(path, k))?)))
727 .collect::<Result<_, String>>()?,
728 )
729 }
730 Schema::Record(r) => {
731 let obj = v.as_object().ok_or_else(|| want("an object"))?;
732 let mut fields = Vec::with_capacity(r.fields.len());
733 for f in &r.fields {
734 let key = source_name(f);
735 let fv = obj.get(key).unwrap_or(&Value::Null);
736 let datum = match (fv, accepts_null(&f.schema, names), &f.default) {
737 (Value::Null, false, Some(default)) => {
738 from_json(default, &f.schema, names, &join(path, key))?
739 }
740 (Value::Null, false, None) => {
741 return Err(format!(
742 "{}: required field is missing or null",
743 at(&join(path, key))
744 ));
745 }
746 _ => from_json(fv, &f.schema, names, &join(path, key))?,
747 };
748 fields.push((f.name.clone(), datum));
749 }
750 Av::Record(fields)
751 }
752 Schema::Union(u) => {
753 let variants = u.variants();
754 if v.is_null()
755 && let Some(i) = variants.iter().position(|b| matches!(b, Schema::Null))
756 {
757 return Ok(Av::Union(i as u32, Box::new(Av::Null)));
758 }
759 let mut last = None;
760 for (i, b) in variants.iter().enumerate() {
761 if matches!(b, Schema::Null) {
762 continue;
763 }
764 match from_json(v, b, names, path) {
765 Ok(d) => return Ok(Av::Union(i as u32, Box::new(d))),
766 Err(e) => last = Some(e),
767 }
768 }
769 return Err(last.unwrap_or_else(|| want("a value matching the union")));
770 }
771 Schema::Enum(e) => {
772 let sym = v.as_str().ok_or_else(|| want("an enum symbol"))?;
773 let i =
774 e.symbols.iter().position(|x| x == sym).ok_or_else(|| {
775 format!("{}: {sym:?} is not one of {:?}", at(path), e.symbols)
776 })?;
777 Av::Enum(i as u32, sym.to_string())
778 }
779 Schema::Fixed(x) => {
780 let b = bytes_from(v).map_err(|e| format!("{}: {e}", at(path)))?;
781 if b.len() != x.size {
782 return Err(format!(
783 "{}: fixed({}) got {} bytes",
784 at(path),
785 x.size,
786 b.len()
787 ));
788 }
789 Av::Fixed(x.size, b)
790 }
791 Schema::Decimal(d) => {
792 let text = match v {
793 Value::String(t) => t.clone(),
794 Value::Number(n) => n.to_string(),
795 _ => return Err(want("a decimal string or number")),
796 };
797 let bytes = decimal_from_str(&text, d.scale, d.precision)
798 .map_err(|e| format!("{}: {e}", at(path)))?;
799 let bytes = match &d.inner {
800 apache_avro::schema::InnerDecimalSchema::Fixed(f) => {
801 sign_extend(bytes, f.size).map_err(|e| format!("{}: {e}", at(path)))?
802 }
803 _ => bytes,
804 };
805 Av::Decimal(apache_avro::Decimal::from(bytes))
806 }
807 Schema::BigDecimal => {
808 let text = match v {
809 Value::String(t) => t.clone(),
810 Value::Number(n) => n.to_string(),
811 _ => return Err(want("a decimal string or number")),
812 };
813 Av::BigDecimal(
814 text.parse::<apache_avro::BigDecimal>()
815 .map_err(|_| format!("{}: {text:?} is not a decimal", at(path)))?,
816 )
817 }
818 Schema::Uuid(_) => {
819 let t = v.as_str().ok_or_else(|| want("a UUID string"))?;
820 Av::Uuid(
821 apache_avro::Uuid::parse_str(t)
822 .map_err(|_| format!("{}: {t:?} is not a UUID", at(path)))?,
823 )
824 }
825 Schema::Date => Av::Date(match v {
826 Value::String(t) => {
827 let d = NaiveDate::parse_from_str(t, "%Y-%m-%d")
828 .map_err(|_| format!("{}: {t:?} is not YYYY-MM-DD", at(path)))?;
829 i32::try_from((d - EPOCH).num_days())
830 .map_err(|_| format!("{}: {t:?} is out of range", at(path)))?
831 }
832 _ => i32::try_from(v.as_i64().ok_or_else(|| want("a date"))?)
833 .map_err(|_| format!("{}: date out of range", at(path)))?,
834 }),
835 Schema::TimeMillis => Av::TimeMillis(
836 i32::try_from(time_from(v, 1_000_000).map_err(|e| format!("{}: {e}", at(path)))?)
837 .map_err(|_| format!("{}: time out of range", at(path)))?,
838 ),
839 Schema::TimeMicros => {
840 Av::TimeMicros(time_from(v, 1_000).map_err(|e| format!("{}: {e}", at(path)))?)
841 }
842 Schema::TimestampMillis => Av::TimestampMillis(
843 ts_from(v, 1_000_000, true).map_err(|e| format!("{}: {e}", at(path)))?,
844 ),
845 Schema::TimestampMicros => {
846 Av::TimestampMicros(ts_from(v, 1_000, true).map_err(|e| format!("{}: {e}", at(path)))?)
847 }
848 Schema::TimestampNanos => {
849 Av::TimestampNanos(ts_from(v, 1, true).map_err(|e| format!("{}: {e}", at(path)))?)
850 }
851 Schema::LocalTimestampMillis => Av::LocalTimestampMillis(
852 ts_from(v, 1_000_000, false).map_err(|e| format!("{}: {e}", at(path)))?,
853 ),
854 Schema::LocalTimestampMicros => Av::LocalTimestampMicros(
855 ts_from(v, 1_000, false).map_err(|e| format!("{}: {e}", at(path)))?,
856 ),
857 Schema::LocalTimestampNanos => {
858 Av::LocalTimestampNanos(ts_from(v, 1, false).map_err(|e| format!("{}: {e}", at(path)))?)
859 }
860 Schema::Duration(_) => {
861 let part = |k: &str| -> Result<u32, String> {
862 v.get(k)
863 .and_then(Value::as_u64)
864 .and_then(|n| u32::try_from(n).ok())
865 .ok_or_else(|| format!("{}: duration needs an unsigned `{k}`", at(path)))
866 };
867 Av::Duration(apache_avro::Duration::new(
868 apache_avro::Months::new(part("months")?),
869 apache_avro::Days::new(part("days")?),
870 apache_avro::Millis::new(part("millis")?),
871 ))
872 }
873 Schema::Ref { .. } => unreachable!("resolved above"),
874 })
875}
876
877fn source_name(f: &RecordField) -> &str {
878 f.custom_attributes
879 .get(ORIGINAL_NAME_ATTR)
880 .and_then(Value::as_str)
881 .unwrap_or(&f.name)
882}
883
884fn accepts_null(s: &Schema, names: &Names<'_>) -> bool {
885 match resolve(s, names) {
886 Ok(Schema::Null) => true,
887 Ok(Schema::Union(u)) => u.is_nullable(),
888 _ => false,
889 }
890}
891
892fn join(path: &str, k: &str) -> String {
893 if path.is_empty() {
894 k.to_string()
895 } else {
896 format!("{path}.{k}")
897 }
898}
899
900fn short(v: &Value) -> String {
901 let s = v.to_string();
902 if s.len() > 60 {
903 format!("{}…", &s[..s.floor_char_boundary(60)])
904 } else {
905 s
906 }
907}
908
909fn float_from(v: &Value) -> Option<f64> {
910 match v {
911 Value::Number(n) => n.as_f64(),
912 Value::String(s) => match s.as_str() {
913 "NaN" => Some(f64::NAN),
914 "Infinity" => Some(f64::INFINITY),
915 "-Infinity" => Some(f64::NEG_INFINITY),
916 _ => None,
917 },
918 _ => None,
919 }
920}
921
922fn bytes_from(v: &Value) -> Result<Vec<u8>, String> {
923 match v {
924 Value::String(s) => unhex(s),
925 Value::Array(items) => items
926 .iter()
927 .map(|i| {
928 i.as_u64()
929 .and_then(|n| u8::try_from(n).ok())
930 .ok_or_else(|| "byte array holds a non-byte".to_string())
931 })
932 .collect(),
933 other => Err(format!("expected hex bytes, got {}", short(other))),
934 }
935}
936
937fn time_from(v: &Value, nanos_per_unit: i64) -> Result<i64, String> {
938 match v {
939 Value::String(t) => {
940 let parsed = NaiveTime::parse_from_str(t, "%H:%M:%S%.f")
941 .map_err(|_| format!("{t:?} is not HH:MM:SS[.fff]"))?;
942 let nanos = i64::from(parsed.num_seconds_from_midnight()) * 1_000_000_000
943 + i64::from(parsed.nanosecond());
944 Ok(nanos / nanos_per_unit)
945 }
946 other => other
947 .as_i64()
948 .ok_or_else(|| format!("expected a time, got {}", short(other))),
949 }
950}
951
952fn ts_from(v: &Value, nanos_per_unit: i64, utc: bool) -> Result<i64, String> {
953 let t = match v {
954 Value::String(t) => t,
955 other => {
956 return other
957 .as_i64()
958 .ok_or_else(|| format!("expected a timestamp, got {}", short(other)));
959 }
960 };
961 let dt: NaiveDateTime = match DateTime::parse_from_rfc3339(t) {
962 Ok(d) if utc => d.with_timezone(&Utc).naive_utc(),
963 Ok(d) => d.naive_local(),
964 Err(_) => [
965 "%Y-%m-%dT%H:%M:%S%.f",
966 "%Y-%m-%d %H:%M:%S%.f",
967 "%Y-%m-%dT%H:%M:%S",
968 "%Y-%m-%d %H:%M:%S",
969 ]
970 .iter()
971 .find_map(|f| NaiveDateTime::parse_from_str(t, f).ok())
972 .ok_or_else(|| format!("{t:?} is not an RFC 3339 / ISO 8601 timestamp"))?,
973 };
974 let nanos = dt
975 .and_utc()
976 .timestamp_nanos_opt()
977 .map(i128::from)
978 .unwrap_or_else(|| {
979 i128::from(dt.and_utc().timestamp()) * 1_000_000_000
980 + i128::from(dt.and_utc().timestamp_subsec_nanos())
981 });
982 i64::try_from(nanos / i128::from(nanos_per_unit)).map_err(|_| format!("{t:?} is out of range"))
983}
984
985pub fn infer_schema(records: &[Value]) -> Result<Schema, FaucetError> {
989 let json_schema = crate::schema::infer_schema(records);
990 let order = super::header_union(records);
991 let avro =
992 record_json(INFERRED_RECORD_NAME, &json_schema, Some(&order)).map_err(FaucetError::Sink)?;
993 Schema::parse(&avro)
994 .map_err(|e| FaucetError::Sink(format!("avro: inferred schema is invalid: {e}")))
995}
996
997fn record_json(name: &str, js: &Value, order: Option<&[String]>) -> Result<Value, String> {
998 let props = js
999 .get("properties")
1000 .and_then(Value::as_object)
1001 .cloned()
1002 .unwrap_or_default();
1003 let keys: Vec<String> = match order {
1004 Some(o) => o
1005 .iter()
1006 .filter(|k| props.contains_key(*k))
1007 .cloned()
1008 .collect(),
1009 None => {
1010 let mut k: Vec<String> = props.keys().cloned().collect();
1011 k.sort();
1012 k
1013 }
1014 };
1015 let mut seen: HashMap<String, String> = HashMap::new();
1016 let mut fields = Vec::with_capacity(keys.len());
1017 for key in keys {
1018 let clean = avro_name(&key);
1019 if let Some(prev) = seen.insert(clean.clone(), key.clone()) {
1020 return Err(format!(
1021 "avro: fields {prev:?} and {key:?} both sanitize to the Avro name {clean:?}"
1022 ));
1023 }
1024 let (ty, nullable) = type_json(&props[&key], &format!("{name}_{clean}"))?;
1025 let mut f = Map::new();
1026 f.insert("name".into(), Value::String(clean.clone()));
1027 if nullable {
1028 f.insert("type".into(), json!(["null", ty]));
1029 f.insert("default".into(), Value::Null);
1030 } else {
1031 f.insert("type".into(), ty);
1032 }
1033 if clean != key {
1034 f.insert(ORIGINAL_NAME_ATTR.into(), Value::String(key));
1035 }
1036 fields.push(Value::Object(f));
1037 }
1038 Ok(json!({"type": "record", "name": name, "fields": fields}))
1039}
1040
1041fn type_json(js: &Value, name: &str) -> Result<(Value, bool), String> {
1042 let types: HashSet<&str> = match js.get("type") {
1043 Some(Value::String(t)) => std::iter::once(t.as_str()).collect(),
1044 Some(Value::Array(a)) => a.iter().filter_map(Value::as_str).collect(),
1045 _ => HashSet::new(),
1046 };
1047 let nullable = types.contains("null");
1048 let non_null: Vec<&str> = types.iter().copied().filter(|t| *t != "null").collect();
1049 let ty = match non_null.as_slice() {
1050 [] => return Ok((json!("null"), false)),
1051 ["boolean"] => json!("boolean"),
1052 ["integer"] => json!("long"),
1053 ["number"] => json!("double"),
1054 ["string"] => json!("string"),
1055 ["array"] => {
1056 let items = match js.get("items") {
1057 Some(i) => {
1058 let (t, n) = type_json(i, &format!("{name}_item"))?;
1059 if n { json!(["null", t]) } else { t }
1060 }
1061 None => json!(["null", "string"]),
1062 };
1063 json!({"type": "array", "items": items})
1064 }
1065 ["object"] => record_json(name, js, None)?,
1066 _ => json!("string"),
1067 };
1068 Ok((ty, nullable))
1069}
1070
1071pub fn avro_name(s: &str) -> String {
1073 let mut out: String = s
1074 .chars()
1075 .map(|c| {
1076 if c.is_ascii_alphanumeric() || c == '_' {
1077 c
1078 } else {
1079 '_'
1080 }
1081 })
1082 .collect();
1083 if out.is_empty() || out.starts_with(|c: char| c.is_ascii_digit()) {
1084 out.insert(0, '_');
1085 }
1086 out
1087}
1088
1089#[cfg(feature = "arrow")]
1093pub fn arrow_schema(schema: &Schema) -> Result<arrow::datatypes::SchemaRef, FaucetError> {
1094 use arrow::datatypes::{Field, Schema as ArrowSchema};
1095 let names = named(schema);
1096 let mut stack = Vec::new();
1097 let root = resolve(schema, &names)?;
1098 let fields: Vec<Field> = match root {
1099 Schema::Record(r) => {
1100 stack.push(r.name.fullname(None));
1101 r.fields
1102 .iter()
1103 .map(|f| {
1104 let (dt, nullable) = arrow_type(&f.schema, &names, &mut stack)?;
1105 Ok(Field::new(&f.name, dt, nullable))
1106 })
1107 .collect::<Result<_, FaucetError>>()?
1108 }
1109 other => {
1110 let (dt, nullable) = arrow_type(other, &names, &mut stack)?;
1111 vec![Field::new(VALUE_FIELD, dt, nullable)]
1112 }
1113 };
1114 Ok(std::sync::Arc::new(ArrowSchema::new(fields)))
1115}
1116
1117#[cfg(feature = "arrow")]
1118fn arrow_type(
1119 schema: &Schema,
1120 names: &Names<'_>,
1121 stack: &mut Vec<String>,
1122) -> Result<(arrow::datatypes::DataType, bool), FaucetError> {
1123 use arrow::datatypes::{DataType as D, Field, Fields, TimeUnit};
1124 use std::sync::Arc;
1125 if let Some(n) = ref_name(schema)
1126 && stack.contains(&n)
1127 {
1128 return Ok((D::Utf8, true));
1129 }
1130 let s = resolve(schema, names)?;
1131 let utc: Option<Arc<str>> = Some("UTC".into());
1132 Ok(match s {
1133 Schema::Null => (D::Null, true),
1134 Schema::Boolean => (D::Boolean, false),
1135 Schema::Int => (D::Int32, false),
1136 Schema::Long => (D::Int64, false),
1137 Schema::Float => (D::Float32, false),
1138 Schema::Double => (D::Float64, false),
1139 Schema::Bytes => (D::Binary, false),
1140 Schema::String | Schema::Enum(_) | Schema::Uuid(_) | Schema::BigDecimal => (D::Utf8, false),
1141 Schema::Fixed(f) => (
1142 D::FixedSizeBinary(i32::try_from(f.size).unwrap_or(i32::MAX)),
1143 false,
1144 ),
1145 Schema::Decimal(d) => match (u8::try_from(d.precision), i8::try_from(d.scale)) {
1146 (Ok(p), Ok(sc)) if p <= 38 => (D::Decimal128(p, sc), false),
1147 (Ok(p), Ok(sc)) if p <= 76 => (D::Decimal256(p, sc), false),
1148 _ => (D::Utf8, false),
1149 },
1150 Schema::Date => (D::Date32, false),
1151 Schema::TimeMillis => (D::Time32(TimeUnit::Millisecond), false),
1152 Schema::TimeMicros => (D::Time64(TimeUnit::Microsecond), false),
1153 Schema::TimestampMillis => (D::Timestamp(TimeUnit::Millisecond, utc), false),
1154 Schema::TimestampMicros => (D::Timestamp(TimeUnit::Microsecond, utc), false),
1155 Schema::TimestampNanos => (D::Timestamp(TimeUnit::Nanosecond, utc), false),
1156 Schema::LocalTimestampMillis => (D::Timestamp(TimeUnit::Millisecond, None), false),
1157 Schema::LocalTimestampMicros => (D::Timestamp(TimeUnit::Microsecond, None), false),
1158 Schema::LocalTimestampNanos => (D::Timestamp(TimeUnit::Nanosecond, None), false),
1159 Schema::Duration(_) => (
1160 D::Struct(Fields::from(vec![
1161 Field::new("months", D::Int64, false),
1162 Field::new("days", D::Int64, false),
1163 Field::new("millis", D::Int64, false),
1164 ])),
1165 false,
1166 ),
1167 Schema::Array(a) => {
1168 let (dt, n) = arrow_type(&a.items, names, stack)?;
1169 (D::List(Arc::new(Field::new("item", dt, n))), false)
1170 }
1171 Schema::Map(m) => {
1172 let (dt, n) = arrow_type(&m.types, names, stack)?;
1173 let entries = Field::new(
1174 "entries",
1175 D::Struct(Fields::from(vec![
1176 Field::new("key", D::Utf8, false),
1177 Field::new("value", dt, n),
1178 ])),
1179 false,
1180 );
1181 (D::Map(Arc::new(entries), false), false)
1182 }
1183 Schema::Record(r) => {
1184 stack.push(r.name.fullname(None));
1185 let fields = r
1186 .fields
1187 .iter()
1188 .map(|f| {
1189 let (dt, n) = arrow_type(&f.schema, names, stack)?;
1190 Ok(Field::new(&f.name, dt, n))
1191 })
1192 .collect::<Result<Vec<_>, FaucetError>>();
1193 stack.pop();
1194 (D::Struct(Fields::from(fields?)), false)
1195 }
1196 Schema::Union(u) => {
1197 let non_null: Vec<&Schema> = u
1198 .variants()
1199 .iter()
1200 .filter(|v| !matches!(v, Schema::Null))
1201 .collect();
1202 match non_null.as_slice() {
1203 [] => (D::Null, true),
1204 [one] => {
1205 let (dt, _) = arrow_type(one, names, stack)?;
1206 (dt, u.is_nullable())
1207 }
1208 _ => (D::Utf8, u.is_nullable()),
1209 }
1210 }
1211 Schema::Ref { .. } => unreachable!("resolved above"),
1212 })
1213}
1214
1215#[cfg(feature = "arrow")]
1217pub fn read_batches<R: std::io::Read>(
1218 reader: R,
1219 reader_schema: Option<&Schema>,
1220 batch_size: usize,
1221 f: &mut dyn FnMut(arrow::array::RecordBatch) -> Result<(), FaucetError>,
1222) -> Result<(Schema, arrow::datatypes::SchemaRef), FaucetError> {
1223 read_batches_checked(reader, reader_schema, None, batch_size, f)
1224}
1225
1226#[cfg(feature = "arrow")]
1228pub(crate) fn read_batches_checked<R: std::io::Read>(
1229 reader: R,
1230 reader_schema: Option<&Schema>,
1231 check: Option<&WriterCheck<'_>>,
1232 batch_size: usize,
1233 f: &mut dyn FnMut(arrow::array::RecordBatch) -> Result<(), FaucetError>,
1234) -> Result<(Schema, arrow::datatypes::SchemaRef), FaucetError> {
1235 let chunk = if batch_size == 0 {
1236 usize::MAX
1237 } else {
1238 batch_size
1239 };
1240 let mut arrow: Option<arrow::datatypes::SchemaRef> = None;
1241 let schema = read_with(
1242 reader,
1243 reader_schema,
1244 check,
1245 chunk,
1246 Mode::Arrow,
1247 &mut |rows, schema| {
1248 let target = match &arrow {
1249 Some(a) => a.clone(),
1250 None => {
1251 let a = arrow_schema(schema)?;
1252 arrow = Some(a.clone());
1253 a
1254 }
1255 };
1256 f(crate::columnar::values_to_record_batch(&rows, target)?)
1257 },
1258 )?;
1259 let arrow = match arrow {
1260 Some(a) => a,
1261 None => arrow_schema(&schema)?,
1262 };
1263 Ok((schema, arrow))
1264}
1265
1266#[cfg(test)]
1267mod tests {
1268 use super::*;
1269
1270 #[test]
1271 fn datum_conversion_keeps_logical_types_exact_and_round_trips() {
1272 let schema = Schema::parse_str(
1273 r#"{"type":"record","name":"r","fields":[
1274 {"name":"amount","type":{"type":"bytes","logicalType":"decimal","precision":10,"scale":2}},
1275 {"name":"raw","type":"bytes"},
1276 {"name":"day","type":{"type":"int","logicalType":"date"}}
1277 ]}"#,
1278 )
1279 .unwrap();
1280 let v = json!({"amount": "12.34", "raw": "0aff", "day": "2026-10-07"});
1281 let datum = datum_from_json(&v, &schema).unwrap();
1282 assert_eq!(datum_to_json(&datum, &schema).unwrap(), v);
1283 assert_eq!(
1284 datum_to_json(&Av::Long(7), &Schema::Long).unwrap(),
1285 json!(7),
1286 "a non-record datum is not wrapped"
1287 );
1288 assert!(datum_from_json(&json!({"amount": 1}), &schema).is_err());
1289 }
1290
1291 fn logical_schema() -> Value {
1292 json!({
1293 "type": "record",
1294 "name": "payment",
1295 "namespace": "test",
1296 "fields": [
1297 {"name": "id", "type": "long"},
1298 {"name": "amount", "type": {"type": "bytes", "logicalType": "decimal", "precision": 10, "scale": 2}},
1299 {"name": "fee", "type": {"type": "fixed", "name": "fee_t", "size": 8, "logicalType": "decimal", "precision": 12, "scale": 3}},
1300 {"name": "day", "type": {"type": "int", "logicalType": "date"}},
1301 {"name": "at", "type": {"type": "long", "logicalType": "timestamp-micros"}},
1302 {"name": "at_ms", "type": {"type": "long", "logicalType": "timestamp-millis"}},
1303 {"name": "at_ns", "type": {"type": "long", "logicalType": "timestamp-nanos"}},
1304 {"name": "local", "type": {"type": "long", "logicalType": "local-timestamp-micros"}},
1305 {"name": "local_ms", "type": {"type": "long", "logicalType": "local-timestamp-millis"}},
1306 {"name": "local_ns", "type": {"type": "long", "logicalType": "local-timestamp-nanos"}},
1307 {"name": "t_ms", "type": {"type": "int", "logicalType": "time-millis"}},
1308 {"name": "t_us", "type": {"type": "long", "logicalType": "time-micros"}},
1309 {"name": "uid", "type": {"type": "string", "logicalType": "uuid"}},
1310 {"name": "span", "type": {"type": "fixed", "name": "dur", "size": 12, "logicalType": "duration"}},
1311 {"name": "note", "type": ["null", "string"], "default": null},
1312 {"name": "either", "type": ["null", "long", "string"], "default": null},
1313 {"name": "kind", "type": {"type": "enum", "name": "kind_t", "symbols": ["A", "B"]}},
1314 {"name": "raw", "type": "bytes"},
1315 {"name": "tags", "type": {"type": "array", "items": "string"}},
1316 {"name": "attrs", "type": {"type": "map", "values": "int"}},
1317 {"name": "ratio", "type": "float"},
1318 {"name": "score", "type": "double"},
1319 {"name": "ok", "type": "boolean"},
1320 {"name": "small", "type": "int"},
1321 {"name": "big", "type": {"type": "bytes", "logicalType": "big-decimal"}}
1322 ]
1323 })
1324 }
1325
1326 fn logical_record() -> Value {
1327 json!({
1328 "id": 7,
1329 "amount": "-1234.50",
1330 "fee": "0.125",
1331 "day": "2024-02-29",
1332 "at": "2024-02-29T12:34:56.123456Z",
1333 "at_ms": "2024-02-29T12:34:56.123Z",
1334 "at_ns": "2024-02-29T12:34:56.123456789Z",
1335 "local": "2024-02-29T12:34:56.123456",
1336 "local_ms": "2024-02-29T12:34:56.123",
1337 "local_ns": "2024-02-29T12:34:56.123456789",
1338 "t_ms": "01:02:03.004",
1339 "t_us": "01:02:03.000004",
1340 "uid": "6f1c2b1e-4b8a-4b8e-9a38-3c3f3f0f9e11",
1341 "span": {"months": 1, "days": 2, "millis": 3},
1342 "note": null,
1343 "either": "text",
1344 "kind": "B",
1345 "raw": "00ff10",
1346 "tags": ["x", "y"],
1347 "attrs": {"a": 1, "b": 2},
1348 "ratio": 0.5,
1349 "score": "NaN",
1350 "ok": true,
1351 "small": 3,
1352 "big": "123456789012345678901234567890.5"
1353 })
1354 }
1355
1356 #[test]
1357 fn logical_types_round_trip_through_an_ocf() {
1358 for codec in [
1359 AvroCodec::Null,
1360 AvroCodec::Deflate,
1361 AvroCodec::Snappy,
1362 AvroCodec::Zstd,
1363 ] {
1364 let opts = AvroOptions {
1365 schema: Some(logical_schema()),
1366 codec,
1367 };
1368 let bytes = encode(&[logical_record()], &opts).expect("encode");
1369 let back = decode(&bytes, &AvroOptions::default()).expect("decode");
1370 assert_eq!(back, vec![logical_record()], "{codec:?}");
1371 }
1372 }
1373
1374 #[test]
1375 fn epoch_integers_are_accepted_for_temporal_types() {
1376 let schema = json!({"type": "record", "name": "r", "fields": [
1377 {"name": "d", "type": {"type": "int", "logicalType": "date"}},
1378 {"name": "t", "type": {"type": "long", "logicalType": "timestamp-millis"}},
1379 {"name": "tm", "type": {"type": "int", "logicalType": "time-millis"}}
1380 ]});
1381 let opts = AvroOptions {
1382 schema: Some(schema),
1383 ..Default::default()
1384 };
1385 let bytes = encode(&[json!({"d": 1, "t": 1000, "tm": 1500})], &opts).expect("encode");
1386 let back = decode(&bytes, &AvroOptions::default()).expect("decode");
1387 assert_eq!(
1388 back[0],
1389 json!({"d": "1970-01-02", "t": "1970-01-01T00:00:01.000Z", "tm": "00:00:01.500"})
1390 );
1391 }
1392
1393 #[test]
1394 fn inferred_schema_widens_nullable_nested_and_mixed_fields() {
1395 let recs = vec![
1396 json!({"id": 1, "name": "a", "first-name": "x", "nested": {"k": 1.5}, "mixed": 1, "list": [1, 2], "empty": []}),
1397 json!({"id": 2, "name": null, "first-name": "y", "nested": {"k": 2}, "mixed": "two", "list": [], "empty": []}),
1398 ];
1399 let bytes = encode(&recs, &AvroOptions::default()).expect("encode");
1400 let schema = writer_schema(&bytes[..]).expect("schema");
1401 let text = serde_json::to_string(&schema).expect("json");
1402 assert!(text.contains("\"first_name\""), "{text}");
1403 assert!(text.contains(ORIGINAL_NAME_ATTR), "{text}");
1404 let back = decode(&bytes, &AvroOptions::default()).expect("decode");
1405 assert_eq!(back[0]["id"], json!(1));
1406 assert_eq!(back[1]["name"], Value::Null);
1407 assert_eq!(back[0]["first_name"], json!("x"));
1408 assert_eq!(back[1]["nested"], json!({"k": 2.0}));
1409 assert_eq!(back[0]["mixed"], json!("1"));
1410 assert_eq!(back[1]["mixed"], json!("two"));
1411 assert_eq!(back[0]["list"], json!([1, 2]));
1412 assert_eq!(back[0]["empty"], json!([]));
1413 }
1414
1415 #[test]
1416 fn inference_refuses_colliding_sanitized_names() {
1417 let err =
1418 encode(&[json!({"a-b": 1, "a_b": 2})], &AvroOptions::default()).expect_err("collision");
1419 assert!(err.to_string().contains("sanitize"), "{err}");
1420 }
1421
1422 #[test]
1423 fn avro_names_are_sanitized() {
1424 assert_eq!(avro_name("ok_1"), "ok_1");
1425 assert_eq!(avro_name("first name"), "first_name");
1426 assert_eq!(avro_name("1st"), "_1st");
1427 assert_eq!(avro_name(""), "_");
1428 }
1429
1430 #[test]
1431 fn reader_schema_projects_and_defaults_added_fields() {
1432 let bytes = encode(&[json!({"a": 1, "b": "x"})], &AvroOptions::default()).expect("encode");
1433 let reader = json!({"type": "record", "name": INFERRED_RECORD_NAME, "fields": [
1434 {"name": "b", "type": "string"},
1435 {"name": "c", "type": "long", "default": 9}
1436 ]});
1437 let back = decode(
1438 &bytes,
1439 &AvroOptions {
1440 schema: Some(reader),
1441 ..Default::default()
1442 },
1443 )
1444 .expect("decode");
1445 assert_eq!(back, vec![json!({"b": "x", "c": 9})]);
1446 }
1447
1448 #[test]
1449 fn a_non_record_root_is_wrapped_in_value() {
1450 let opts = AvroOptions {
1451 schema: Some(json!("long")),
1452 ..Default::default()
1453 };
1454 let bytes = encode(&[json!({"value": 5}), json!(6)], &opts).expect("encode");
1455 let back = decode(&bytes, &AvroOptions::default()).expect("decode");
1456 assert_eq!(back, vec![json!({"value": 5}), json!({"value": 6})]);
1457 }
1458
1459 #[test]
1460 fn recursive_records_decode_as_nested_objects() {
1461 let schema = json!({"type": "record", "name": "node", "fields": [
1462 {"name": "v", "type": "long"},
1463 {"name": "next", "type": ["null", "node"], "default": null}
1464 ]});
1465 let rec = json!({"v": 1, "next": {"v": 2, "next": null}});
1466 let opts = AvroOptions {
1467 schema: Some(schema),
1468 ..Default::default()
1469 };
1470 let bytes = encode(std::slice::from_ref(&rec), &opts).expect("encode");
1471 assert_eq!(
1472 decode(&bytes, &AvroOptions::default()).expect("decode"),
1473 vec![rec]
1474 );
1475 }
1476
1477 #[test]
1478 fn encoding_errors_name_the_field() {
1479 let schema = json!({"type": "record", "name": "r", "fields": [
1480 {"name": "n", "type": "int"},
1481 {"name": "req", "type": "string"}
1482 ]});
1483 let opts = AvroOptions {
1484 schema: Some(schema),
1485 ..Default::default()
1486 };
1487 let err = encode(&[json!({"n": "x", "req": "a"})], &opts).expect_err("type");
1488 assert!(err.to_string().contains("`n`"), "{err}");
1489 let err = encode(&[json!({"n": 1})], &opts).expect_err("missing");
1490 assert!(err.to_string().contains("required"), "{err}");
1491 let err = encode(&[json!({"n": 1i64 << 40, "req": "a"})], &opts).expect_err("range");
1492 assert!(err.to_string().contains("does not fit"), "{err}");
1493 }
1494
1495 #[test]
1496 fn invalid_schemas_and_bodies_are_typed_errors() {
1497 let opts = AvroOptions {
1498 schema: Some(json!({"type": "nope"})),
1499 ..Default::default()
1500 };
1501 assert!(matches!(opts.parsed_schema(), Err(FaucetError::Config(_))));
1502 let err = decode(b"not avro", &AvroOptions::default()).expect_err("garbage");
1503 assert!(err.to_string().contains("avro"), "{err}");
1504 }
1505
1506 #[test]
1507 fn decimals_render_and_parse_exactly() {
1508 let cases = [
1509 ("0", 0),
1510 ("-1", 0),
1511 ("12.30", 2),
1512 ("-0.05", 2),
1513 ("170141183460469231731687303715884105727", 0),
1514 ];
1515 for (text, scale) in cases {
1516 let bytes = decimal_from_str(text, scale, 60).expect(text);
1517 assert_eq!(decimal_to_string(&bytes, scale), text, "{text}");
1518 }
1519 assert_eq!(
1520 decimal_to_string(&decimal_from_str("1.5e1", 1, 5).expect("exp"), 1),
1521 "15.0"
1522 );
1523 assert_eq!(
1524 decimal_to_string(&decimal_from_str("1.230", 2, 5).expect("trailing zero"), 2),
1525 "1.23"
1526 );
1527 assert!(decimal_from_str("1.234", 2, 5).is_err());
1528 for huge in ["1e18", "1e9223372036854775807", "-5e400000000000"] {
1530 let err = decimal_from_str(huge, 2, 10).unwrap_err();
1531 assert!(err.contains("precision"), "{huge}: {err}");
1532 }
1533 assert_eq!(
1534 decimal_to_string(&decimal_from_str("0e400000000000", 2, 5).expect("zero"), 2),
1535 "0.00"
1536 );
1537 assert!(decimal_from_str("1e-9223372036854775807", 2, 5).is_err());
1538 assert!(decimal_from_str("123456", 0, 5).is_err());
1539 assert!(decimal_from_str("abc", 0, 5).is_err());
1540 assert!(decimal_from_str("1e-9", 2, 5).is_err());
1541 assert_eq!(decimal_from_str("0e-9", 2, 5).expect("zero"), vec![0]);
1542 assert_eq!(
1543 sign_extend(vec![0xff], 3).expect("neg"),
1544 vec![0xff, 0xff, 0xff]
1545 );
1546 assert!(sign_extend(vec![1, 2, 3], 2).is_err());
1547 }
1548
1549 #[test]
1550 fn value_coercions_and_their_refusals() {
1551 let names = HashMap::new();
1552 assert!(
1553 matches!(from_json(&json!(1), &Schema::String, &names, "s"), Ok(Av::String(s)) if s == "1")
1554 );
1555 assert!(from_json(&Value::Null, &Schema::String, &names, "s").is_err());
1556 assert!(from_json(&json!(1), &Schema::Null, &names, "s").is_err());
1557 assert!(
1558 matches!(from_json(&json!([1, 2]), &Schema::Bytes, &names, "b"), Ok(Av::Bytes(b)) if b == vec![1, 2])
1559 );
1560 assert!(from_json(&json!("abc"), &Schema::Bytes, &names, "b").is_err());
1561 assert!(from_json(&json!("zz"), &Schema::Bytes, &names, "b").is_err());
1562 assert!(from_json(&json!([300]), &Schema::Bytes, &names, "b").is_err());
1563 assert!(from_json(&json!(true), &Schema::Bytes, &names, "b").is_err());
1564 assert!(from_json(&json!("Infinity"), &Schema::Double, &names, "d").is_ok());
1565 assert!(from_json(&json!("-Infinity"), &Schema::Double, &names, "d").is_ok());
1566 assert!(from_json(&json!("x"), &Schema::Double, &names, "d").is_err());
1567 assert!(from_json(&json!("x"), &Schema::Date, &names, "d").is_err());
1568 assert!(from_json(&json!("x"), &Schema::TimestampMicros, &names, "d").is_err());
1569 assert!(
1570 from_json(
1571 &json!("2024-01-01 00:00:00"),
1572 &Schema::TimestampMicros,
1573 &names,
1574 "d"
1575 )
1576 .is_ok()
1577 );
1578 assert!(
1579 from_json(
1580 &json!("2024-01-01T00:00:00+02:00"),
1581 &Schema::LocalTimestampMicros,
1582 &names,
1583 "d"
1584 )
1585 .is_ok()
1586 );
1587 assert!(from_json(&json!(true), &Schema::TimeMicros, &names, "d").is_err());
1588 assert!(from_json(&json!("bad"), &Schema::TimeMicros, &names, "d").is_err());
1589 assert!(from_json(&json!(5), &Schema::TimeMicros, &names, "d").is_ok());
1590 assert!(from_json(&json!(true), &Schema::TimestampMicros, &names, "d").is_err());
1591 assert!(
1592 from_json(
1593 &json!({"months": 1}),
1594 &parse_schema(
1595 &json!({"type": "fixed", "name": "d", "size": 12, "logicalType": "duration"})
1596 )
1597 .unwrap(),
1598 &names,
1599 "d"
1600 )
1601 .is_err()
1602 );
1603 let e = parse_schema(&json!({"type": "enum", "name": "e", "symbols": ["A"]})).unwrap();
1604 assert!(from_json(&json!("Z"), &e, &named(&e), "e").is_err());
1605 let f = parse_schema(&json!({"type": "fixed", "name": "f", "size": 2})).unwrap();
1606 assert!(from_json(&json!("00"), &f, &named(&f), "f").is_err());
1607 let u = parse_schema(&json!({"type": "string", "logicalType": "uuid"})).unwrap();
1608 assert!(from_json(&json!("nope"), &u, &names, "u").is_err());
1609 let bd = parse_schema(&json!({"type": "bytes", "logicalType": "big-decimal"})).unwrap();
1610 assert!(from_json(&json!(1.5), &bd, &names, "b").is_ok());
1611 assert!(from_json(&json!("x"), &bd, &names, "b").is_err());
1612 assert!(from_json(&json!(true), &bd, &names, "b").is_err());
1613 let dec = parse_schema(
1614 &json!({"type": "bytes", "logicalType": "decimal", "precision": 4, "scale": 1}),
1615 )
1616 .unwrap();
1617 assert!(from_json(&json!(1.5), &dec, &names, "d").is_ok());
1618 assert!(from_json(&json!(true), &dec, &names, "d").is_err());
1619 let un = parse_schema(&json!(["long", "boolean"])).unwrap();
1620 assert!(from_json(&json!("x"), &un, &names, "u").is_err());
1621 assert!(from_json(&Value::Null, &un, &names, "u").is_err());
1622 assert!(short(&json!("x".repeat(100))).ends_with('…'));
1623 }
1624
1625 #[test]
1626 fn out_of_range_temporal_values_are_errors() {
1627 assert!(date_string(i32::MAX).is_err());
1628 assert!(time_string(-1, 3).is_err());
1629 assert!(utc_string(i64::MAX, 1_000_000, SecondsFormat::Millis).is_err());
1630 assert!(local_string(i64::MAX, 1_000_000, 3).is_err());
1631 assert_eq!(float_json(f64::NEG_INFINITY), json!("-Infinity"));
1632 }
1633
1634 #[cfg(feature = "arrow")]
1635 #[test]
1636 fn arrow_batches_carry_logical_types() {
1637 use arrow::datatypes::{DataType as D, TimeUnit};
1638 let opts = AvroOptions {
1639 schema: Some(logical_schema()),
1640 ..Default::default()
1641 };
1642 let rec = logical_record();
1643 let bytes = encode(&[rec.clone(), rec], &opts).expect("encode");
1644 let mut batches = Vec::new();
1645 let (_, schema) = read_batches(&bytes[..], None, 1, &mut |b| {
1646 batches.push(b);
1647 Ok(())
1648 })
1649 .expect("batches");
1650 assert_eq!(batches.len(), 2);
1651 let ty = |n: &str| schema.field_with_name(n).expect(n).data_type().clone();
1652 assert_eq!(ty("amount"), D::Decimal128(10, 2));
1653 assert_eq!(ty("day"), D::Date32);
1654 assert_eq!(
1655 ty("at"),
1656 D::Timestamp(TimeUnit::Microsecond, Some("UTC".into()))
1657 );
1658 assert_eq!(ty("local"), D::Timestamp(TimeUnit::Microsecond, None));
1659 assert_eq!(ty("t_ms"), D::Time32(TimeUnit::Millisecond));
1660 assert_eq!(ty("either"), D::Utf8);
1661 assert!(schema.field_with_name("note").unwrap().is_nullable());
1662 let back = crate::columnar::record_batch_to_values(&batches[0]).expect("values");
1663 let amount = batches[0]
1664 .column_by_name("amount")
1665 .and_then(|c| c.as_any().downcast_ref::<arrow::array::Decimal128Array>())
1666 .expect("decimal column");
1667 assert_eq!(amount.value(0), -123_450);
1668 assert_eq!(back[0]["day"], json!("2024-02-29"));
1669 assert_eq!(back[0]["either"], json!("\"text\""));
1670 }
1671
1672 #[cfg(feature = "arrow")]
1673 #[test]
1674 fn arrow_schema_handles_roots_recursion_and_wide_decimals() {
1675 use arrow::datatypes::DataType as D;
1676 let s = parse_schema(&json!("long")).unwrap();
1677 assert_eq!(arrow_schema(&s).unwrap().field(0).name(), VALUE_FIELD);
1678 let rec = parse_schema(&json!({"type": "record", "name": "node", "fields": [
1679 {"name": "next", "type": ["null", "node"], "default": null},
1680 {"name": "wide", "type": {"type": "bytes", "logicalType": "decimal", "precision": 50, "scale": 2}},
1681 {"name": "huge", "type": {"type": "bytes", "logicalType": "decimal", "precision": 90, "scale": 2}},
1682 {"name": "only_null", "type": ["null"]},
1683 {"name": "n", "type": "null"}
1684 ]}))
1685 .unwrap();
1686 let a = arrow_schema(&rec).unwrap();
1687 assert_eq!(a.field(0).data_type(), &D::Utf8);
1688 assert_eq!(a.field(1).data_type(), &D::Decimal256(50, 2));
1689 assert_eq!(a.field(2).data_type(), &D::Utf8);
1690 assert_eq!(a.field(3).data_type(), &D::Null);
1691 let opts = AvroOptions {
1692 schema: Some(json!({"type": "record", "name": "node", "fields": [
1693 {"name": "v", "type": "long"},
1694 {"name": "next", "type": ["null", "node"], "default": null}
1695 ]})),
1696 ..Default::default()
1697 };
1698 let bytes = encode(&[json!({"v": 1, "next": {"v": 2, "next": null}})], &opts).unwrap();
1699 let mut rows = Vec::new();
1700 read_batches(&bytes[..], None, 0, &mut |b| {
1701 rows.extend(crate::columnar::record_batch_to_values(&b)?);
1702 Ok(())
1703 })
1704 .unwrap();
1705 let next: Value = serde_json::from_str(rows[0]["next"].as_str().unwrap()).unwrap();
1706 assert_eq!(next, json!({"v": 2, "next": null}));
1707 }
1708}