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