Skip to main content

faucet_core/file_format/
avro.rs

1//! Apache Avro Object Container Files (#719).
2//!
3//! # Reading
4//!
5//! An OCF embeds its writer schema, so a file decodes on its own. When a
6//! *reader* schema is supplied (`avro.schema` on a source) every file is
7//! resolved against it with Avro's schema-resolution rules — the way to
8//! project a subset of fields, rename through aliases, or read files written
9//! across a schema's evolution into one shape. [`ContainerDecoder`](super::container::ContainerDecoder)
10//! adds the multi-file rule: without an explicit reader schema the first
11//! file's writer schema becomes the reader schema for the rest, so a prefix
12//! whose later files cannot be resolved against the first fails with an error
13//! naming both files instead of producing rows of two shapes.
14//!
15//! # Logical types
16//!
17//! Mapped explicitly, never through a lossy numeric fallback:
18//!
19//! | Avro | JSON record | Arrow |
20//! |---|---|---|
21//! | `decimal(p, s)` | exact decimal string (`"12.30"`) | `Decimal128(p, s)` (`Decimal256` above 38 digits) |
22//! | `big-decimal` | decimal string | `Utf8` |
23//! | `date` | `"YYYY-MM-DD"` | `Date32` |
24//! | `time-millis` / `time-micros` | `"HH:MM:SS.fff"` / `"HH:MM:SS.ffffff"` | `Time32(ms)` / `Time64(µs)` |
25//! | `timestamp-{millis,micros,nanos}` | RFC 3339 UTC (`"…Z"`) at that precision | `Timestamp(unit, "UTC")` |
26//! | `local-timestamp-*` | naive ISO 8601 | `Timestamp(unit, None)` |
27//! | `uuid` | canonical string | `Utf8` |
28//! | `duration` | `{months, days, millis}` | `Struct` |
29//!
30//! `bytes` and `fixed` are standard base64 — the canonical binary rendering of
31//! [`crate::fidelity`], the same one Arrow binary columns get, so an Avro file
32//! and a Parquet file carrying the same bytes produce the same records. `NaN`
33//! and the infinities, which JSON cannot hold, become a counted `null`.
34//!
35//! A union of `null` and one type is that type, nullable. A union of several
36//! non-null types decodes to whichever branch the value holds; on the Arrow
37//! path — where a column needs one type — it is a `Utf8` column holding the
38//! value's JSON. A recursive record is likewise `Utf8` JSON on the Arrow path
39//! (Arrow has no recursive types) and a nested object on the record path.
40//!
41//! # Writing
42//!
43//! With `avro.schema` set the records are encoded against it; logical types
44//! are accepted in the shapes above (plus epoch integers). Without one the
45//! schema is inferred from the records being written: `integer` → `long`,
46//! `number` → `double`, objects → nested records, a field that is absent or
47//! null anywhere → `["null", T]` with a `null` default, and a field whose
48//! values mix types → `string`. Field names that are not valid Avro names are
49//! sanitized (`first-name` → `first_name`, a leading digit gains `_`), and the
50//! original name is kept as the field's `faucet.name` attribute.
51
52use 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
61/// Attribute holding a field's original name when it had to be sanitized.
62pub const ORIGINAL_NAME_ATTR: &str = "faucet.name";
63
64/// Name of the top-level record an inferred schema declares.
65pub const INFERRED_RECORD_NAME: &str = "faucet_record";
66
67/// The field a non-record top-level datum is wrapped in, so every record is an
68/// object whatever the file's root type.
69pub 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    /// The configured schema, parsed, if any.
82    pub fn parsed_schema(&self) -> Result<Option<Schema>, FaucetError> {
83        self.schema.as_ref().map(parse_schema).transpose()
84    }
85}
86
87/// Parse an Avro schema from its JSON form.
88pub 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
93/// Decode a whole OCF into records, resolved against `opts.schema` when set.
94pub 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
105/// The writer schema an OCF declares, read from its header only.
106pub 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
112/// Stream an OCF's records in chunks of `chunk` (at least 1), resolved against
113/// `reader_schema` when given. Returns the schema the records are shaped by.
114pub 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
123/// Checks a file's writer schema before any of its rows are decoded.
124pub(crate) type WriterCheck<'a> = dyn Fn(&Schema) -> Result<(), FaucetError> + 'a;
125
126/// [`read_records`] with a writer-schema check run before the first row.
127pub(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/// How a datum is rendered: plain JSON, or JSON the Arrow decoder can type
145/// against [`arrow_schema`] (multi-branch unions and recursion as JSON text).
146#[derive(Debug, Clone, Copy, PartialEq, Eq)]
147pub(crate) enum Mode {
148    Record,
149    Arrow,
150}
151
152/// Receives each decoded chunk with the schema it is shaped by.
153pub(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
191/// Every named type of a schema by full and short name, owned so a prepared
192/// [`AvroDatumCodec`] can keep it next to its schema.
193type Names = HashMap<String, Schema>;
194
195fn named(schema: &Schema) -> Names {
196    let mut out = HashMap::new();
197    collect_named(schema, &mut out);
198    out
199}
200
201fn collect_named(s: &Schema, out: &mut Names) {
202    match s {
203        Schema::Record(r) => {
204            out.insert(r.name.fullname(None), s.clone());
205            out.entry(r.name.name().to_string())
206                .or_insert_with(|| s.clone());
207            for f in &r.fields {
208                collect_named(&f.schema, out);
209            }
210        }
211        Schema::Enum(e) => {
212            out.insert(e.name.fullname(None), s.clone());
213            out.entry(e.name.name().to_string())
214                .or_insert_with(|| s.clone());
215        }
216        Schema::Fixed(x) | Schema::Duration(x) => {
217            out.insert(x.name.fullname(None), s.clone());
218            out.entry(x.name.name().to_string())
219                .or_insert_with(|| s.clone());
220        }
221        Schema::Array(a) => collect_named(&a.items, out),
222        Schema::Map(m) => collect_named(&m.types, out),
223        Schema::Union(u) => u.variants().iter().for_each(|v| collect_named(v, out)),
224        _ => {}
225    }
226}
227
228fn resolve<'a>(s: &'a Schema, names: &'a Names) -> Result<&'a Schema, FaucetError> {
229    match s {
230        Schema::Ref { name } => names
231            .get(&name.fullname(None))
232            .or_else(|| names.get(name.name()))
233            .ok_or_else(|| {
234                FaucetError::Source(format!(
235                    "avro: unresolved schema reference `{}`",
236                    name.fullname(None)
237                ))
238            }),
239        other => Ok(other),
240    }
241}
242
243fn ref_name(s: &Schema) -> Option<String> {
244    match s {
245        Schema::Ref { name } => Some(name.fullname(None)),
246        Schema::Record(r) => Some(r.name.fullname(None)),
247        _ => None,
248    }
249}
250
251/// A union needing one Arrow column for several value types.
252pub(crate) fn is_complex_union(u: &UnionSchema) -> bool {
253    u.variants()
254        .iter()
255        .filter(|v| !matches!(v, Schema::Null))
256        .count()
257        > 1
258}
259
260/// [`datum_to_json`] / [`datum_from_json_with`] against one schema whose
261/// named types are indexed once, for a codec that converts many datums (a
262/// Kafka topic's messages) — the free functions re-index the schema per call.
263#[derive(Debug, Clone)]
264pub struct AvroDatumCodec {
265    schema: Schema,
266    names: Names,
267}
268
269impl AvroDatumCodec {
270    /// Index `schema`'s named types.
271    pub fn new(schema: Schema) -> Self {
272        let names = named(&schema);
273        Self { schema, names }
274    }
275
276    /// The schema this codec converts against.
277    pub fn schema(&self) -> &Schema {
278        &self.schema
279    }
280
281    /// [`datum_to_json`] without re-indexing the schema.
282    pub fn to_json(&self, datum: &Av) -> Result<Value, FaucetError> {
283        to_json(
284            datum,
285            &self.schema,
286            &self.names,
287            Mode::Record,
288            &mut Vec::new(),
289        )
290    }
291
292    /// [`datum_from_json_with`] without re-indexing the schema.
293    pub fn from_json(&self, value: &Value, text: BytesText) -> Result<Av, String> {
294        from_json(value, &self.schema, &self.names, text, "")
295    }
296}
297
298/// Convert one decoded Avro datum to JSON with the logical-type mapping in the
299/// module docs (exact decimal strings, base64 bytes, ISO temporals). Unlike the
300/// file reader, a non-record datum is returned as-is rather than wrapped in
301/// [`VALUE_FIELD`] — the shape a single-message codec (a Kafka value) wants.
302pub fn datum_to_json(datum: &Av, schema: &Schema) -> Result<Value, FaucetError> {
303    let names = named(schema);
304    to_json(datum, schema, &names, Mode::Record, &mut Vec::new())
305}
306
307/// Build the Avro datum for `value` under `schema`, accepting the JSON shapes
308/// [`datum_to_json`] produces (the inverse mapping). The error names the path
309/// that did not fit.
310pub fn datum_from_json(value: &Value, schema: &Schema) -> Result<Av, FaucetError> {
311    datum_from_json_with(value, schema, BytesText::Base64)
312}
313
314/// How a JSON string is read for an Avro `bytes` / `fixed` value on encode.
315#[derive(
316    Debug,
317    Clone,
318    Copy,
319    Default,
320    PartialEq,
321    Eq,
322    serde::Serialize,
323    serde::Deserialize,
324    schemars::JsonSchema,
325)]
326#[serde(rename_all = "snake_case")]
327#[non_exhaustive]
328pub enum BytesText {
329    /// The string is standard base64, the shape [`datum_to_json`] produces.
330    #[default]
331    Base64,
332    /// The string is lowercase hex (the rendering before #885).
333    Hex,
334    /// The string's UTF-8 bytes are the value.
335    Utf8,
336}
337
338/// [`datum_from_json`] with an explicit reading of JSON strings for
339/// `bytes` / `fixed` values. A JSON array of byte numbers is accepted either way.
340pub fn datum_from_json_with(
341    value: &Value,
342    schema: &Schema,
343    text: BytesText,
344) -> Result<Av, FaucetError> {
345    let names = named(schema);
346    from_json(value, schema, &names, text, "")
347        .map_err(|e| FaucetError::Sink(format!("json->avro: {e}")))
348}
349
350fn root_to_json(v: &Av, schema: &Schema, names: &Names, mode: Mode) -> Result<Value, FaucetError> {
351    let mut stack = Vec::new();
352    let resolved = resolve(schema, names)?;
353    let j = to_json(v, schema, names, mode, &mut stack)?;
354    Ok(match (resolved, j) {
355        (Schema::Record(_), j @ Value::Object(_)) => j,
356        (_, j) => {
357            let mut m = Map::new();
358            m.insert(VALUE_FIELD.into(), j);
359            Value::Object(m)
360        }
361    })
362}
363
364fn to_json(
365    v: &Av,
366    schema: &Schema,
367    names: &Names,
368    mode: Mode,
369    stack: &mut Vec<String>,
370) -> Result<Value, FaucetError> {
371    let s = resolve(schema, names)?;
372    if mode == Mode::Arrow
373        && let Some(n) = ref_name(schema)
374        && stack.contains(&n)
375    {
376        let inner = to_json(v, s, names, Mode::Record, &mut Vec::new())?;
377        return Ok(Value::String(crate::util::canonical_json(&inner)));
378    }
379    Ok(match (v, s) {
380        (Av::Union(i, inner), Schema::Union(u)) => {
381            let branch = u.variants().get(*i as usize).ok_or_else(|| {
382                FaucetError::Source(format!("avro: union branch {i} out of range"))
383            })?;
384            if mode == Mode::Arrow && is_complex_union(u) {
385                if matches!(**inner, Av::Null) {
386                    return Ok(Value::Null);
387                }
388                let j = to_json(inner, branch, names, Mode::Record, &mut Vec::new())?;
389                return Ok(Value::String(crate::util::canonical_json(&j)));
390            }
391            to_json(inner, branch, names, mode, stack)?
392        }
393        (Av::Union(_, inner), other) => to_json(inner, other, names, mode, stack)?,
394        (Av::Null, _) => Value::Null,
395        (Av::Boolean(b), _) => Value::Bool(*b),
396        (Av::Int(i), _) => json!(i),
397        (Av::Long(i), _) => json!(i),
398        (Av::Float(f), _) if mode == Mode::Arrow && !f.is_finite() => {
399            non_finite_text(f64::from(*f))
400        }
401        (Av::Double(f), _) if mode == Mode::Arrow && !f.is_finite() => non_finite_text(*f),
402        (Av::Float(f), _) => crate::fidelity::float32(*f),
403        (Av::Double(f), _) => crate::fidelity::float64(*f),
404        (Av::Bytes(b), _) | (Av::Fixed(_, b), _) => crate::fidelity::binary(b),
405        (Av::String(t), _) => Value::String(t.clone()),
406        (Av::Enum(_, sym), _) => Value::String(sym.clone()),
407        (Av::Array(items), s) => {
408            let item_schema = match s {
409                Schema::Array(a) => &*a.items,
410                _ => s,
411            };
412            Value::Array(
413                items
414                    .iter()
415                    .map(|i| to_json(i, item_schema, names, mode, stack))
416                    .collect::<Result<_, _>>()?,
417            )
418        }
419        (Av::Map(m), s) => {
420            let value_schema = match s {
421                Schema::Map(m) => &*m.types,
422                _ => s,
423            };
424            let mut keys: Vec<&String> = m.keys().collect();
425            keys.sort();
426            let mut out = Map::new();
427            for k in keys {
428                out.insert(k.clone(), to_json(&m[k], value_schema, names, mode, stack)?);
429            }
430            Value::Object(out)
431        }
432        (Av::Record(fields), s) => {
433            let rec = match s {
434                Schema::Record(r) => Some(r),
435                _ => None,
436            };
437            let pushed = rec.map(|r| r.name.fullname(None));
438            if let Some(n) = &pushed {
439                stack.push(n.clone());
440            }
441            let mut out = Map::new();
442            for (name, fv) in fields {
443                let fs = rec
444                    .and_then(|r| r.lookup.get(name).map(|i| &r.fields[*i].schema))
445                    .unwrap_or(&Schema::Null);
446                out.insert(name.clone(), to_json(fv, fs, names, mode, stack)?);
447            }
448            if pushed.is_some() {
449                stack.pop();
450            }
451            Value::Object(out)
452        }
453        (Av::Date(d), _) => Value::String(date_string(*d)?),
454        (Av::Decimal(d), s) => {
455            let scale = match s {
456                Schema::Decimal(DecimalSchema { scale, .. }) => *scale,
457                _ => 0,
458            };
459            let bytes = <Vec<u8>>::try_from(d)
460                .map_err(|e| FaucetError::Source(format!("avro decimal: {e}")))?;
461            Value::String(decimal_to_string(&bytes, scale))
462        }
463        (Av::BigDecimal(d), _) => Value::String(d.to_string()),
464        (Av::TimeMillis(ms), _) => Value::String(time_string(i64::from(*ms) * 1_000_000, 3)?),
465        (Av::TimeMicros(us), _) => Value::String(time_string(us.saturating_mul(1_000), 6)?),
466        (Av::TimestampMillis(t), _) => {
467            Value::String(utc_string(*t, 1_000_000, SecondsFormat::Millis)?)
468        }
469        (Av::TimestampMicros(t), _) => Value::String(utc_string(*t, 1_000, SecondsFormat::Micros)?),
470        (Av::TimestampNanos(t), _) => Value::String(utc_string(*t, 1, SecondsFormat::Nanos)?),
471        (Av::LocalTimestampMillis(t), _) => Value::String(local_string(*t, 1_000_000, 3)?),
472        (Av::LocalTimestampMicros(t), _) => Value::String(local_string(*t, 1_000, 6)?),
473        (Av::LocalTimestampNanos(t), _) => Value::String(local_string(*t, 1, 9)?),
474        (Av::Duration(d), _) => json!({
475            "months": u32::from(d.months()),
476            "days": u32::from(d.days()),
477            "millis": u32::from(d.millis()),
478        }),
479        (Av::Uuid(u), _) => Value::String(u.to_string()),
480    })
481}
482
483/// On the Arrow path a non-finite float stays one: the JSON hop reads these
484/// words back into the `Float` column.
485fn non_finite_text(f: f64) -> Value {
486    Value::String(
487        if f.is_nan() {
488            "NaN"
489        } else if f > 0.0 {
490            "Infinity"
491        } else {
492            "-Infinity"
493        }
494        .into(),
495    )
496}
497
498fn unhex(s: &str) -> Result<Vec<u8>, String> {
499    if !s.len().is_multiple_of(2) {
500        return Err(format!("{s:?} is not hex (odd length)"));
501    }
502    (0..s.len())
503        .step_by(2)
504        .map(|i| u8::from_str_radix(&s[i..i + 2], 16).map_err(|_| format!("{s:?} is not hex")))
505        .collect()
506}
507
508const EPOCH: NaiveDate = NaiveDate::from_ymd_opt(1970, 1, 1).expect("valid epoch");
509
510fn date_string(days: i32) -> Result<String, FaucetError> {
511    EPOCH
512        .checked_add_signed(chrono::Duration::days(i64::from(days)))
513        .map(|d| d.format("%Y-%m-%d").to_string())
514        .ok_or_else(|| FaucetError::Source(format!("avro date {days} is out of range")))
515}
516
517fn time_string(nanos: i64, digits: usize) -> Result<String, FaucetError> {
518    let secs = u32::try_from(nanos.div_euclid(1_000_000_000)).ok();
519    let frac = nanos.rem_euclid(1_000_000_000) as u32;
520    let t = secs
521        .and_then(|s| NaiveTime::from_num_seconds_from_midnight_opt(s, frac))
522        .ok_or_else(|| FaucetError::Source(format!("avro time {nanos}ns is out of range")))?;
523    Ok(format!(
524        "{}.{:0w$}",
525        t.format("%H:%M:%S"),
526        t.nanosecond() / 10u32.pow(9 - digits as u32),
527        w = digits
528    ))
529}
530
531fn instant(t: i64, nanos_per_unit: i64) -> Option<DateTime<Utc>> {
532    let n = i128::from(t) * i128::from(nanos_per_unit);
533    let secs = i64::try_from(n.div_euclid(1_000_000_000)).ok()?;
534    DateTime::from_timestamp(secs, n.rem_euclid(1_000_000_000) as u32)
535}
536
537fn utc_string(t: i64, nanos_per_unit: i64, fmt: SecondsFormat) -> Result<String, FaucetError> {
538    instant(t, nanos_per_unit)
539        .map(|d| d.to_rfc3339_opts(fmt, true))
540        .ok_or_else(|| FaucetError::Source(format!("avro timestamp {t} is out of range")))
541}
542
543fn local_string(t: i64, nanos_per_unit: i64, digits: usize) -> Result<String, FaucetError> {
544    let d = instant(t, nanos_per_unit)
545        .ok_or_else(|| FaucetError::Source(format!("avro local timestamp {t} is out of range")))?
546        .naive_utc();
547    Ok(format!(
548        "{}.{:0w$}",
549        d.format("%Y-%m-%dT%H:%M:%S"),
550        d.nanosecond() / 10u32.pow(9 - digits as u32),
551        w = digits
552    ))
553}
554
555/// Render a big-endian two's-complement unscaled integer at `scale`.
556pub(crate) fn decimal_to_string(bytes: &[u8], scale: usize) -> String {
557    let negative = bytes.first().is_some_and(|b| b & 0x80 != 0);
558    let mut mag: Vec<u8> = if negative {
559        negate(bytes)
560    } else {
561        bytes.to_vec()
562    };
563    let mut digits = Vec::new();
564    while mag.iter().any(|b| *b != 0) {
565        let mut rem: u32 = 0;
566        for b in mag.iter_mut() {
567            let cur = (rem << 8) | u32::from(*b);
568            *b = (cur / 10) as u8;
569            rem = cur % 10;
570        }
571        digits.push(b'0' + rem as u8);
572    }
573    while digits.len() <= scale {
574        digits.push(b'0');
575    }
576    digits.reverse();
577    let (int, frac) = digits.split_at(digits.len() - scale);
578    let mut s = String::new();
579    if negative {
580        s.push('-');
581    }
582    s.push_str(std::str::from_utf8(int).expect("ascii digits"));
583    if scale > 0 {
584        s.push('.');
585        s.push_str(std::str::from_utf8(frac).expect("ascii digits"));
586    }
587    s
588}
589
590fn negate(bytes: &[u8]) -> Vec<u8> {
591    let mut out: Vec<u8> = bytes.iter().map(|b| !b).collect();
592    for b in out.iter_mut().rev() {
593        let (v, carry) = b.overflowing_add(1);
594        *b = v;
595        if !carry {
596            break;
597        }
598    }
599    out
600}
601
602/// Parse a decimal literal (optionally in exponent form) into the minimal
603/// big-endian two's-complement unscaled integer at `scale`.
604///
605/// Refuses digits beyond `scale` that are not zero — rounding a money column
606/// silently is not a formatting choice — and more than `precision` digits.
607pub(crate) fn decimal_from_str(
608    text: &str,
609    scale: usize,
610    precision: usize,
611) -> Result<Vec<u8>, String> {
612    let t = text.trim();
613    let (negative, body) = match t.strip_prefix('-') {
614        Some(rest) => (true, rest),
615        None => (false, t.strip_prefix('+').unwrap_or(t)),
616    };
617    let (mantissa, exp) = match body.find(['e', 'E']) {
618        Some(i) => (
619            &body[..i],
620            body[i + 1..]
621                .parse::<i64>()
622                .map_err(|_| format!("{text:?} is not a decimal"))?,
623        ),
624        None => (body, 0),
625    };
626    let (int, frac) = mantissa.split_once('.').unwrap_or((mantissa, ""));
627    if int.is_empty() && frac.is_empty()
628        || !int.bytes().chain(frac.bytes()).all(|b| b.is_ascii_digit())
629    {
630        return Err(format!("{text:?} is not a decimal"));
631    }
632    let mut digits: Vec<u8> = int.bytes().chain(frac.bytes()).map(|b| b - b'0').collect();
633    let out_of_range = || format!("{text:?} is out of range for precision {precision}");
634    let wanted = (int.len() as i64)
635        .checked_add(exp)
636        .and_then(|p| p.checked_add(scale as i64))
637        .ok_or_else(out_of_range)?;
638    // Check the significant digits the scaled value would need before padding,
639    // so an exponent such as `1e18` cannot drive an enormous allocation.
640    match digits.iter().position(|d| *d != 0) {
641        None => digits.clear(),
642        Some(first)
643            if wanted > first as i64 && (wanted - first as i64) as u128 > precision as u128 =>
644        {
645            return Err(format!(
646                "{text:?} needs {} digits, more than the precision {precision}",
647                wanted - first as i64
648            ));
649        }
650        Some(_) => {}
651    }
652    if digits.is_empty() {
653        // Zero at any exponent: nothing to scale.
654    } else if wanted < 0 {
655        return Err(format!("{text:?} has more than {scale} fractional digits"));
656    } else {
657        let wanted = wanted as usize;
658        if digits.len() > wanted {
659            if digits[wanted..].iter().any(|d| *d != 0) {
660                return Err(format!("{text:?} has more than {scale} fractional digits"));
661            }
662            digits.truncate(wanted);
663        } else {
664            digits.resize(wanted, 0);
665        }
666    }
667    let first = digits.iter().position(|d| *d != 0).unwrap_or(digits.len());
668    let digits = &digits[first..];
669    if digits.len() > precision {
670        return Err(format!(
671            "{text:?} needs {} digits, more than the precision {precision}",
672            digits.len()
673        ));
674    }
675    let mut mag: Vec<u8> = vec![0];
676    for d in digits {
677        let mut carry = u32::from(*d);
678        for b in mag.iter_mut().rev() {
679            let cur = u32::from(*b) * 10 + carry;
680            *b = (cur & 0xff) as u8;
681            carry = cur >> 8;
682        }
683        while carry > 0 {
684            mag.insert(0, (carry & 0xff) as u8);
685            carry >>= 8;
686        }
687    }
688    if mag[0] & 0x80 != 0 {
689        mag.insert(0, 0);
690    }
691    while mag.len() > 1 && mag[0] == 0 && mag[1] & 0x80 == 0 {
692        mag.remove(0);
693    }
694    Ok(if negative && mag.iter().any(|b| *b != 0) {
695        negate(&mag)
696    } else {
697        mag
698    })
699}
700
701fn sign_extend(bytes: Vec<u8>, size: usize) -> Result<Vec<u8>, String> {
702    if bytes.len() > size {
703        return Err(format!(
704            "decimal needs {} bytes, the fixed holds {size}",
705            bytes.len()
706        ));
707    }
708    let fill = if bytes.first().is_some_and(|b| b & 0x80 != 0) {
709        0xff
710    } else {
711        0
712    };
713    let mut out = vec![fill; size - bytes.len()];
714    out.extend(bytes);
715    Ok(out)
716}
717
718// ── Writing ──────────────────────────────────────────────────────────────────
719
720/// Encode records into one OCF, against `opts.schema` or an inferred schema.
721pub fn encode(records: &[Value], opts: &AvroOptions) -> Result<Vec<u8>, FaucetError> {
722    let schema = match opts.parsed_schema()? {
723        Some(s) => s,
724        None => infer_schema(records)?,
725    };
726    let names = named(&schema);
727    let mut w = apache_avro::Writer::with_codec(&schema, Vec::new(), codec(opts.codec))
728        .map_err(|e| FaucetError::Sink(format!("avro writer: {e}")))?;
729    for (i, r) in records.iter().enumerate() {
730        let datum = root_from_json(r, &schema, &names)
731            .map_err(|e| FaucetError::Sink(format!("avro: record {i}: {e}")))?;
732        w.append_value_ref(&datum)
733            .map_err(|e| FaucetError::Sink(format!("avro: record {i}: {e}")))?;
734    }
735    w.into_inner()
736        .map_err(|e| FaucetError::Sink(format!("avro writer: {e}")))
737}
738
739fn root_from_json(v: &Value, schema: &Schema, names: &Names) -> Result<Av, String> {
740    let resolved = resolve(schema, names).map_err(|e| e.to_string())?;
741    match resolved {
742        Schema::Record(_) => from_json(v, schema, names, BytesText::Base64, ""),
743        _ => {
744            let inner = v.get(VALUE_FIELD).unwrap_or(v);
745            from_json(inner, schema, names, BytesText::Base64, "")
746        }
747    }
748}
749
750fn at(path: &str) -> String {
751    if path.is_empty() {
752        "the record".into()
753    } else {
754        format!("`{path}`")
755    }
756}
757
758/// How well a union branch matches a JSON value's own type: `0` an exact
759/// match, `1` a compatible conversion, `2` the lenient text rendering a
760/// `string` branch gives any non-string value.
761fn union_branch_rank(v: &Value, b: &Schema) -> u8 {
762    match (v, b) {
763        (Value::String(_), Schema::String | Schema::Enum(_) | Schema::Uuid(_)) => 0,
764        (Value::Number(n), Schema::Int | Schema::Long) if n.as_i64().is_some() => 0,
765        (Value::Number(n), Schema::Float | Schema::Double) if n.as_i64().is_none() => 0,
766        (Value::Bool(_), Schema::Boolean) => 0,
767        (Value::Object(_), Schema::Record(_) | Schema::Map(_)) => 0,
768        (Value::Array(_), Schema::Array(_)) => 0,
769        (Value::String(_), _) => 1,
770        (_, Schema::String) => 2,
771        _ => 1,
772    }
773}
774
775fn from_json(
776    v: &Value,
777    schema: &Schema,
778    names: &Names,
779    text: BytesText,
780    path: &str,
781) -> Result<Av, String> {
782    let s = resolve(schema, names).map_err(|e| e.to_string())?;
783    let want = |what: &str| format!("{}: expected {what}, got {}", at(path), short(v));
784    Ok(match s {
785        Schema::Null => match v {
786            Value::Null => Av::Null,
787            _ => return Err(want("null")),
788        },
789        Schema::Boolean => Av::Boolean(v.as_bool().ok_or_else(|| want("a boolean"))?),
790        Schema::Int => {
791            let i = v.as_i64().ok_or_else(|| want("an integer"))?;
792            Av::Int(
793                i32::try_from(i)
794                    .map_err(|_| format!("{}: {i} does not fit an Avro int", at(path)))?,
795            )
796        }
797        Schema::Long => Av::Long(v.as_i64().ok_or_else(|| want("an integer"))?),
798        Schema::Float => Av::Float(float_from(v).ok_or_else(|| want("a number"))? as f32),
799        Schema::Double => Av::Double(float_from(v).ok_or_else(|| want("a number"))?),
800        Schema::Bytes => Av::Bytes(bytes_from(v, text).map_err(|e| format!("{}: {e}", at(path)))?),
801        Schema::String => Av::String(match v {
802            Value::String(t) => t.clone(),
803            Value::Null => return Err(want("a string")),
804            other => super::cell_text(other),
805        }),
806        Schema::Array(a) => {
807            let items = v.as_array().ok_or_else(|| want("an array"))?;
808            Av::Array(
809                items
810                    .iter()
811                    .enumerate()
812                    .map(|(i, x)| from_json(x, &a.items, names, text, &format!("{path}[{i}]")))
813                    .collect::<Result<_, _>>()?,
814            )
815        }
816        Schema::Map(m) => {
817            let obj = v.as_object().ok_or_else(|| want("an object"))?;
818            Av::Map(
819                obj.iter()
820                    .map(|(k, x)| {
821                        Ok((
822                            k.clone(),
823                            from_json(x, &m.types, names, text, &join(path, k))?,
824                        ))
825                    })
826                    .collect::<Result<_, String>>()?,
827            )
828        }
829        Schema::Record(r) => {
830            let obj = v.as_object().ok_or_else(|| want("an object"))?;
831            let mut fields = Vec::with_capacity(r.fields.len());
832            for f in &r.fields {
833                let key = source_name(f);
834                let fv = obj.get(key).unwrap_or(&Value::Null);
835                let datum = match (fv, accepts_null(&f.schema, names), &f.default) {
836                    (Value::Null, false, Some(default)) => {
837                        from_json(default, &f.schema, names, text, &join(path, key))?
838                    }
839                    (Value::Null, false, None) => {
840                        return Err(format!(
841                            "{}: required field is missing or null",
842                            at(&join(path, key))
843                        ));
844                    }
845                    _ => from_json(fv, &f.schema, names, text, &join(path, key))?,
846                };
847                fields.push((f.name.clone(), datum));
848            }
849            Av::Record(fields)
850        }
851        Schema::Union(u) => {
852            let variants = u.variants();
853            if v.is_null()
854                && let Some(i) = variants.iter().position(|b| matches!(b, Schema::Null))
855            {
856                return Ok(Av::Union(i as u32, Box::new(Av::Null)));
857            }
858            // Exact JSON-type branches first, the lenient string branch last,
859            // so `["string","long"]` writes 42 as a long, not as "42".
860            let mut order: Vec<(u8, usize)> = variants
861                .iter()
862                .enumerate()
863                .filter(|(_, b)| !matches!(b, Schema::Null))
864                .map(|(i, b)| (union_branch_rank(v, b), i))
865                .collect();
866            order.sort_by_key(|(rank, _)| *rank);
867            let mut last = None;
868            for (_, i) in order {
869                let b = &variants[i];
870                match from_json(v, b, names, text, path) {
871                    Ok(d) => return Ok(Av::Union(i as u32, Box::new(d))),
872                    Err(e) => last = Some(e),
873                }
874            }
875            return Err(last.unwrap_or_else(|| want("a value matching the union")));
876        }
877        Schema::Enum(e) => {
878            let sym = v.as_str().ok_or_else(|| want("an enum symbol"))?;
879            let i =
880                e.symbols.iter().position(|x| x == sym).ok_or_else(|| {
881                    format!("{}: {sym:?} is not one of {:?}", at(path), e.symbols)
882                })?;
883            Av::Enum(i as u32, sym.to_string())
884        }
885        Schema::Fixed(x) => {
886            let b = bytes_from(v, text).map_err(|e| format!("{}: {e}", at(path)))?;
887            if b.len() != x.size {
888                return Err(format!(
889                    "{}: fixed({}) got {} bytes",
890                    at(path),
891                    x.size,
892                    b.len()
893                ));
894            }
895            Av::Fixed(x.size, b)
896        }
897        Schema::Decimal(d) => {
898            let text = match v {
899                Value::String(t) => t.clone(),
900                Value::Number(n) => n.to_string(),
901                _ => return Err(want("a decimal string or number")),
902            };
903            let bytes = decimal_from_str(&text, d.scale, d.precision)
904                .map_err(|e| format!("{}: {e}", at(path)))?;
905            let bytes = match &d.inner {
906                apache_avro::schema::InnerDecimalSchema::Fixed(f) => {
907                    sign_extend(bytes, f.size).map_err(|e| format!("{}: {e}", at(path)))?
908                }
909                _ => bytes,
910            };
911            Av::Decimal(apache_avro::Decimal::from(bytes))
912        }
913        Schema::BigDecimal => {
914            let text = match v {
915                Value::String(t) => t.clone(),
916                Value::Number(n) => n.to_string(),
917                _ => return Err(want("a decimal string or number")),
918            };
919            Av::BigDecimal(
920                text.parse::<apache_avro::BigDecimal>()
921                    .map_err(|_| format!("{}: {text:?} is not a decimal", at(path)))?,
922            )
923        }
924        Schema::Uuid(_) => {
925            let t = v.as_str().ok_or_else(|| want("a UUID string"))?;
926            Av::Uuid(
927                apache_avro::Uuid::parse_str(t)
928                    .map_err(|_| format!("{}: {t:?} is not a UUID", at(path)))?,
929            )
930        }
931        Schema::Date => Av::Date(match v {
932            Value::String(t) => {
933                let d = NaiveDate::parse_from_str(t, "%Y-%m-%d")
934                    .map_err(|_| format!("{}: {t:?} is not YYYY-MM-DD", at(path)))?;
935                i32::try_from((d - EPOCH).num_days())
936                    .map_err(|_| format!("{}: {t:?} is out of range", at(path)))?
937            }
938            _ => i32::try_from(v.as_i64().ok_or_else(|| want("a date"))?)
939                .map_err(|_| format!("{}: date out of range", at(path)))?,
940        }),
941        Schema::TimeMillis => Av::TimeMillis(
942            i32::try_from(time_from(v, 1_000_000).map_err(|e| format!("{}: {e}", at(path)))?)
943                .map_err(|_| format!("{}: time out of range", at(path)))?,
944        ),
945        Schema::TimeMicros => {
946            Av::TimeMicros(time_from(v, 1_000).map_err(|e| format!("{}: {e}", at(path)))?)
947        }
948        Schema::TimestampMillis => Av::TimestampMillis(
949            ts_from(v, 1_000_000, true).map_err(|e| format!("{}: {e}", at(path)))?,
950        ),
951        Schema::TimestampMicros => {
952            Av::TimestampMicros(ts_from(v, 1_000, true).map_err(|e| format!("{}: {e}", at(path)))?)
953        }
954        Schema::TimestampNanos => {
955            Av::TimestampNanos(ts_from(v, 1, true).map_err(|e| format!("{}: {e}", at(path)))?)
956        }
957        Schema::LocalTimestampMillis => Av::LocalTimestampMillis(
958            ts_from(v, 1_000_000, false).map_err(|e| format!("{}: {e}", at(path)))?,
959        ),
960        Schema::LocalTimestampMicros => Av::LocalTimestampMicros(
961            ts_from(v, 1_000, false).map_err(|e| format!("{}: {e}", at(path)))?,
962        ),
963        Schema::LocalTimestampNanos => {
964            Av::LocalTimestampNanos(ts_from(v, 1, false).map_err(|e| format!("{}: {e}", at(path)))?)
965        }
966        Schema::Duration(_) => {
967            let part = |k: &str| -> Result<u32, String> {
968                v.get(k)
969                    .and_then(Value::as_u64)
970                    .and_then(|n| u32::try_from(n).ok())
971                    .ok_or_else(|| format!("{}: duration needs an unsigned `{k}`", at(path)))
972            };
973            Av::Duration(apache_avro::Duration::new(
974                apache_avro::Months::new(part("months")?),
975                apache_avro::Days::new(part("days")?),
976                apache_avro::Millis::new(part("millis")?),
977            ))
978        }
979        Schema::Ref { .. } => unreachable!("resolved above"),
980    })
981}
982
983fn source_name(f: &RecordField) -> &str {
984    f.custom_attributes
985        .get(ORIGINAL_NAME_ATTR)
986        .and_then(Value::as_str)
987        .unwrap_or(&f.name)
988}
989
990fn accepts_null(s: &Schema, names: &Names) -> bool {
991    match resolve(s, names) {
992        Ok(Schema::Null) => true,
993        Ok(Schema::Union(u)) => u.is_nullable(),
994        _ => false,
995    }
996}
997
998fn join(path: &str, k: &str) -> String {
999    if path.is_empty() {
1000        k.to_string()
1001    } else {
1002        format!("{path}.{k}")
1003    }
1004}
1005
1006fn short(v: &Value) -> String {
1007    let s = v.to_string();
1008    if s.len() > 60 {
1009        format!("{}…", &s[..s.floor_char_boundary(60)])
1010    } else {
1011        s
1012    }
1013}
1014
1015fn float_from(v: &Value) -> Option<f64> {
1016    match v {
1017        Value::Number(n) => n.as_f64(),
1018        Value::String(s) => match s.as_str() {
1019            "NaN" => Some(f64::NAN),
1020            "Infinity" => Some(f64::INFINITY),
1021            "-Infinity" => Some(f64::NEG_INFINITY),
1022            _ => None,
1023        },
1024        _ => None,
1025    }
1026}
1027
1028fn bytes_from(v: &Value, text: BytesText) -> Result<Vec<u8>, String> {
1029    match v {
1030        Value::String(s) if text == BytesText::Utf8 => Ok(s.as_bytes().to_vec()),
1031        Value::String(s) if text == BytesText::Hex => unhex(s),
1032        Value::String(s) => {
1033            use base64::Engine as _;
1034            base64::engine::general_purpose::STANDARD
1035                .decode(s)
1036                .map_err(|_| format!("{s:?} is not base64"))
1037        }
1038        Value::Array(items) => items
1039            .iter()
1040            .map(|i| {
1041                i.as_u64()
1042                    .and_then(|n| u8::try_from(n).ok())
1043                    .ok_or_else(|| "byte array holds a non-byte".to_string())
1044            })
1045            .collect(),
1046        other => Err(format!("expected base64 bytes, got {}", short(other))),
1047    }
1048}
1049
1050fn time_from(v: &Value, nanos_per_unit: i64) -> Result<i64, String> {
1051    match v {
1052        Value::String(t) => {
1053            let parsed = NaiveTime::parse_from_str(t, "%H:%M:%S%.f")
1054                .map_err(|_| format!("{t:?} is not HH:MM:SS[.fff]"))?;
1055            let nanos = i64::from(parsed.num_seconds_from_midnight()) * 1_000_000_000
1056                + i64::from(parsed.nanosecond());
1057            Ok(nanos / nanos_per_unit)
1058        }
1059        other => other
1060            .as_i64()
1061            .ok_or_else(|| format!("expected a time, got {}", short(other))),
1062    }
1063}
1064
1065fn ts_from(v: &Value, nanos_per_unit: i64, utc: bool) -> Result<i64, String> {
1066    let t = match v {
1067        Value::String(t) => t,
1068        other => {
1069            return other
1070                .as_i64()
1071                .ok_or_else(|| format!("expected a timestamp, got {}", short(other)));
1072        }
1073    };
1074    let dt: NaiveDateTime = match DateTime::parse_from_rfc3339(t) {
1075        Ok(d) if utc => d.with_timezone(&Utc).naive_utc(),
1076        Ok(d) => d.naive_local(),
1077        Err(_) => [
1078            "%Y-%m-%dT%H:%M:%S%.f",
1079            "%Y-%m-%d %H:%M:%S%.f",
1080            "%Y-%m-%dT%H:%M:%S",
1081            "%Y-%m-%d %H:%M:%S",
1082        ]
1083        .iter()
1084        .find_map(|f| NaiveDateTime::parse_from_str(t, f).ok())
1085        .ok_or_else(|| format!("{t:?} is not an RFC 3339 / ISO 8601 timestamp"))?,
1086    };
1087    let nanos = dt
1088        .and_utc()
1089        .timestamp_nanos_opt()
1090        .map(i128::from)
1091        .unwrap_or_else(|| {
1092            i128::from(dt.and_utc().timestamp()) * 1_000_000_000
1093                + i128::from(dt.and_utc().timestamp_subsec_nanos())
1094        });
1095    i64::try_from(nanos / i128::from(nanos_per_unit)).map_err(|_| format!("{t:?} is out of range"))
1096}
1097
1098// ── Schema inference ─────────────────────────────────────────────────────────
1099
1100/// Infer an Avro record schema from records (see the module docs for rules).
1101pub fn infer_schema(records: &[Value]) -> Result<Schema, FaucetError> {
1102    let json_schema = crate::schema::infer_schema(records);
1103    let order = super::header_union(records);
1104    let avro =
1105        record_json(INFERRED_RECORD_NAME, &json_schema, Some(&order)).map_err(FaucetError::Sink)?;
1106    Schema::parse(&avro)
1107        .map_err(|e| FaucetError::Sink(format!("avro: inferred schema is invalid: {e}")))
1108}
1109
1110fn record_json(name: &str, js: &Value, order: Option<&[String]>) -> Result<Value, String> {
1111    let props = js
1112        .get("properties")
1113        .and_then(Value::as_object)
1114        .cloned()
1115        .unwrap_or_default();
1116    let keys: Vec<String> = match order {
1117        Some(o) => o
1118            .iter()
1119            .filter(|k| props.contains_key(*k))
1120            .cloned()
1121            .collect(),
1122        None => {
1123            let mut k: Vec<String> = props.keys().cloned().collect();
1124            k.sort();
1125            k
1126        }
1127    };
1128    let mut seen: HashMap<String, String> = HashMap::new();
1129    let mut fields = Vec::with_capacity(keys.len());
1130    for key in keys {
1131        let clean = avro_name(&key);
1132        if let Some(prev) = seen.insert(clean.clone(), key.clone()) {
1133            return Err(format!(
1134                "avro: fields {prev:?} and {key:?} both sanitize to the Avro name {clean:?}"
1135            ));
1136        }
1137        let (ty, nullable) = type_json(&props[&key], &format!("{name}_{clean}"))?;
1138        let mut f = Map::new();
1139        f.insert("name".into(), Value::String(clean.clone()));
1140        if nullable {
1141            f.insert("type".into(), json!(["null", ty]));
1142            f.insert("default".into(), Value::Null);
1143        } else {
1144            f.insert("type".into(), ty);
1145        }
1146        if clean != key {
1147            f.insert(ORIGINAL_NAME_ATTR.into(), Value::String(key));
1148        }
1149        fields.push(Value::Object(f));
1150    }
1151    Ok(json!({"type": "record", "name": name, "fields": fields}))
1152}
1153
1154fn type_json(js: &Value, name: &str) -> Result<(Value, bool), String> {
1155    let types: HashSet<&str> = match js.get("type") {
1156        Some(Value::String(t)) => std::iter::once(t.as_str()).collect(),
1157        Some(Value::Array(a)) => a.iter().filter_map(Value::as_str).collect(),
1158        _ => HashSet::new(),
1159    };
1160    let nullable = types.contains("null");
1161    let non_null: Vec<&str> = types.iter().copied().filter(|t| *t != "null").collect();
1162    let ty = match non_null.as_slice() {
1163        [] => return Ok((json!("null"), false)),
1164        ["boolean"] => json!("boolean"),
1165        ["integer"] => json!("long"),
1166        ["number"] => json!("double"),
1167        ["string"] => json!("string"),
1168        ["array"] => {
1169            let items = match js.get("items") {
1170                Some(i) => {
1171                    let (t, n) = type_json(i, &format!("{name}_item"))?;
1172                    if n { json!(["null", t]) } else { t }
1173                }
1174                None => json!(["null", "string"]),
1175            };
1176            json!({"type": "array", "items": items})
1177        }
1178        ["object"] => record_json(name, js, None)?,
1179        _ => json!("string"),
1180    };
1181    Ok((ty, nullable))
1182}
1183
1184/// A valid Avro name for `s`: `[A-Za-z_][A-Za-z0-9_]*`.
1185pub fn avro_name(s: &str) -> String {
1186    let mut out: String = s
1187        .chars()
1188        .map(|c| {
1189            if c.is_ascii_alphanumeric() || c == '_' {
1190                c
1191            } else {
1192                '_'
1193            }
1194        })
1195        .collect();
1196    if out.is_empty() || out.starts_with(|c: char| c.is_ascii_digit()) {
1197        out.insert(0, '_');
1198    }
1199    out
1200}
1201
1202// ── Arrow ────────────────────────────────────────────────────────────────────
1203
1204/// The Arrow schema the records of `schema` decode into (see the module docs).
1205#[cfg(feature = "arrow")]
1206pub fn arrow_schema(schema: &Schema) -> Result<arrow::datatypes::SchemaRef, FaucetError> {
1207    use arrow::datatypes::{Field, Schema as ArrowSchema};
1208    let names = named(schema);
1209    let mut stack = Vec::new();
1210    let root = resolve(schema, &names)?;
1211    let fields: Vec<Field> = match root {
1212        Schema::Record(r) => {
1213            stack.push(r.name.fullname(None));
1214            r.fields
1215                .iter()
1216                .map(|f| {
1217                    let (dt, nullable) = arrow_type(&f.schema, &names, &mut stack)?;
1218                    Ok(Field::new(&f.name, dt, nullable))
1219                })
1220                .collect::<Result<_, FaucetError>>()?
1221        }
1222        other => {
1223            let (dt, nullable) = arrow_type(other, &names, &mut stack)?;
1224            vec![Field::new(VALUE_FIELD, dt, nullable)]
1225        }
1226    };
1227    Ok(std::sync::Arc::new(ArrowSchema::new(fields)))
1228}
1229
1230#[cfg(feature = "arrow")]
1231fn arrow_type(
1232    schema: &Schema,
1233    names: &Names,
1234    stack: &mut Vec<String>,
1235) -> Result<(arrow::datatypes::DataType, bool), FaucetError> {
1236    use arrow::datatypes::{DataType as D, Field, Fields, TimeUnit};
1237    use std::sync::Arc;
1238    if let Some(n) = ref_name(schema)
1239        && stack.contains(&n)
1240    {
1241        return Ok((D::Utf8, true));
1242    }
1243    let s = resolve(schema, names)?;
1244    let utc: Option<Arc<str>> = Some("UTC".into());
1245    Ok(match s {
1246        Schema::Null => (D::Null, true),
1247        Schema::Boolean => (D::Boolean, false),
1248        Schema::Int => (D::Int32, false),
1249        Schema::Long => (D::Int64, false),
1250        Schema::Float => (D::Float32, false),
1251        Schema::Double => (D::Float64, false),
1252        Schema::Bytes => (D::Binary, false),
1253        Schema::String | Schema::Enum(_) | Schema::Uuid(_) | Schema::BigDecimal => (D::Utf8, false),
1254        Schema::Fixed(f) => (
1255            D::FixedSizeBinary(i32::try_from(f.size).unwrap_or(i32::MAX)),
1256            false,
1257        ),
1258        Schema::Decimal(d) => match (u8::try_from(d.precision), i8::try_from(d.scale)) {
1259            (Ok(p), Ok(sc)) if p <= 38 => (D::Decimal128(p, sc), false),
1260            (Ok(p), Ok(sc)) if p <= 76 => (D::Decimal256(p, sc), false),
1261            _ => (D::Utf8, false),
1262        },
1263        Schema::Date => (D::Date32, false),
1264        Schema::TimeMillis => (D::Time32(TimeUnit::Millisecond), false),
1265        Schema::TimeMicros => (D::Time64(TimeUnit::Microsecond), false),
1266        Schema::TimestampMillis => (D::Timestamp(TimeUnit::Millisecond, utc), false),
1267        Schema::TimestampMicros => (D::Timestamp(TimeUnit::Microsecond, utc), false),
1268        Schema::TimestampNanos => (D::Timestamp(TimeUnit::Nanosecond, utc), false),
1269        Schema::LocalTimestampMillis => (D::Timestamp(TimeUnit::Millisecond, None), false),
1270        Schema::LocalTimestampMicros => (D::Timestamp(TimeUnit::Microsecond, None), false),
1271        Schema::LocalTimestampNanos => (D::Timestamp(TimeUnit::Nanosecond, None), false),
1272        Schema::Duration(_) => (
1273            D::Struct(Fields::from(vec![
1274                Field::new("months", D::Int64, false),
1275                Field::new("days", D::Int64, false),
1276                Field::new("millis", D::Int64, false),
1277            ])),
1278            false,
1279        ),
1280        Schema::Array(a) => {
1281            let (dt, n) = arrow_type(&a.items, names, stack)?;
1282            (D::List(Arc::new(Field::new("item", dt, n))), false)
1283        }
1284        Schema::Map(m) => {
1285            let (dt, n) = arrow_type(&m.types, names, stack)?;
1286            let entries = Field::new(
1287                "entries",
1288                D::Struct(Fields::from(vec![
1289                    Field::new("key", D::Utf8, false),
1290                    Field::new("value", dt, n),
1291                ])),
1292                false,
1293            );
1294            (D::Map(Arc::new(entries), false), false)
1295        }
1296        Schema::Record(r) => {
1297            stack.push(r.name.fullname(None));
1298            let fields = r
1299                .fields
1300                .iter()
1301                .map(|f| {
1302                    let (dt, n) = arrow_type(&f.schema, names, stack)?;
1303                    Ok(Field::new(&f.name, dt, n))
1304                })
1305                .collect::<Result<Vec<_>, FaucetError>>();
1306            stack.pop();
1307            (D::Struct(Fields::from(fields?)), false)
1308        }
1309        Schema::Union(u) => {
1310            let non_null: Vec<&Schema> = u
1311                .variants()
1312                .iter()
1313                .filter(|v| !matches!(v, Schema::Null))
1314                .collect();
1315            match non_null.as_slice() {
1316                [] => (D::Null, true),
1317                [one] => {
1318                    let (dt, _) = arrow_type(one, names, stack)?;
1319                    (dt, u.is_nullable())
1320                }
1321                _ => (D::Utf8, u.is_nullable()),
1322            }
1323        }
1324        Schema::Ref { .. } => unreachable!("resolved above"),
1325    })
1326}
1327
1328/// Stream an OCF as Arrow batches of at most `batch_size` rows (0 = one batch).
1329#[cfg(feature = "arrow")]
1330pub fn read_batches<R: std::io::Read>(
1331    reader: R,
1332    reader_schema: Option<&Schema>,
1333    batch_size: usize,
1334    f: &mut dyn FnMut(arrow::array::RecordBatch) -> Result<(), FaucetError>,
1335) -> Result<(Schema, arrow::datatypes::SchemaRef), FaucetError> {
1336    read_batches_checked(reader, reader_schema, None, batch_size, f)
1337}
1338
1339/// [`read_batches`] with a writer-schema check run before the first row.
1340#[cfg(feature = "arrow")]
1341pub(crate) fn read_batches_checked<R: std::io::Read>(
1342    reader: R,
1343    reader_schema: Option<&Schema>,
1344    check: Option<&WriterCheck<'_>>,
1345    batch_size: usize,
1346    f: &mut dyn FnMut(arrow::array::RecordBatch) -> Result<(), FaucetError>,
1347) -> Result<(Schema, arrow::datatypes::SchemaRef), FaucetError> {
1348    let chunk = if batch_size == 0 {
1349        usize::MAX
1350    } else {
1351        batch_size
1352    };
1353    let mut arrow: Option<arrow::datatypes::SchemaRef> = None;
1354    let schema = read_with(
1355        reader,
1356        reader_schema,
1357        check,
1358        chunk,
1359        Mode::Arrow,
1360        &mut |rows, schema| {
1361            let target = match &arrow {
1362                Some(a) => a.clone(),
1363                None => {
1364                    let a = arrow_schema(schema)?;
1365                    arrow = Some(a.clone());
1366                    a
1367                }
1368            };
1369            f(crate::columnar::values_to_record_batch(&rows, target)?)
1370        },
1371    )?;
1372    let arrow = match arrow {
1373        Some(a) => a,
1374        None => arrow_schema(&schema)?,
1375    };
1376    Ok((schema, arrow))
1377}
1378
1379#[cfg(test)]
1380mod tests {
1381    #[test]
1382    fn union_branch_rank_matches_an_array_to_an_array_branch() {
1383        let arr = Schema::parse_str(r#"{"type":"array","items":"long"}"#).unwrap();
1384        assert_eq!(union_branch_rank(&serde_json::json!([1]), &arr), 0);
1385        assert_eq!(union_branch_rank(&serde_json::json!(true), &arr), 1);
1386    }
1387    use super::*;
1388
1389    /// #862 L-034: a complex-union value stringified for an Arrow column is
1390    /// canonical JSON (sorted keys), whatever the record's field order.
1391    #[test]
1392    fn complex_union_text_is_canonical_json() {
1393        let schema = Schema::parse_str(
1394            r#"["null", "string", {"type":"record","name":"r","fields":[
1395                {"name":"z","type":"long"},{"name":"a","type":"long"}]}]"#,
1396        )
1397        .unwrap();
1398        let datum = Av::Union(
1399            2,
1400            Box::new(Av::Record(vec![
1401                ("z".into(), Av::Long(1)),
1402                ("a".into(), Av::Long(2)),
1403            ])),
1404        );
1405        let names = named(&schema);
1406        let v = to_json(&datum, &schema, &names, Mode::Arrow, &mut Vec::new()).unwrap();
1407        assert_eq!(v, json!(r#"{"a":2,"z":1}"#));
1408    }
1409
1410    #[test]
1411    fn a_union_prefers_the_branch_of_the_values_own_type() {
1412        let schema = Schema::parse_str(
1413            r#"{"type":"record","name":"R","fields":[
1414                {"name":"amount","type":["null","string","long"]},
1415                {"name":"ratio","type":["string","long","double"]},
1416                {"name":"flag","type":["string","boolean"]},
1417                {"name":"payload","type":["string",{"type":"record","name":"P","fields":[{"name":"x","type":"long"}]}]},
1418                {"name":"label","type":["long","string"]}]}"#,
1419        )
1420        .unwrap();
1421        let rec =
1422            json!({"amount": 42, "ratio": 1.5, "flag": true, "payload": {"x": 3}, "label": "a"});
1423        let datum = datum_from_json_with(&rec, &schema, BytesText::Base64).unwrap();
1424        let Av::Record(fields) = datum else { panic!() };
1425        assert_eq!(fields[0].1, Av::Union(2, Box::new(Av::Long(42))));
1426        assert_eq!(fields[1].1, Av::Union(2, Box::new(Av::Double(1.5))));
1427        assert_eq!(fields[2].1, Av::Union(1, Box::new(Av::Boolean(true))));
1428        assert!(matches!(&fields[3].1, Av::Union(1, b) if matches!(**b, Av::Record(_))));
1429        assert_eq!(fields[4].1, Av::Union(1, Box::new(Av::String("a".into()))));
1430        // A value no exact branch takes still falls back to the string branch.
1431        let lenient = Schema::parse_str(r#"["null","string","long"]"#).unwrap();
1432        assert_eq!(
1433            datum_from_json_with(&json!(true), &lenient, BytesText::Base64).unwrap(),
1434            Av::Union(1, Box::new(Av::String("true".into())))
1435        );
1436    }
1437
1438    #[test]
1439    fn a_prepared_codec_matches_the_free_functions() {
1440        let schema = Schema::parse_str(
1441            r#"{"type":"record","name":"R","fields":[
1442                {"name":"a","type":"long"},
1443                {"name":"n","type":["null",{"type":"record","name":"In","fields":[{"name":"x","type":"string"}]}]},
1444                {"name":"again","type":["null","In"]}]}"#,
1445        )
1446        .unwrap();
1447        let codec = AvroDatumCodec::new(schema.clone());
1448        assert_eq!(codec.schema(), &schema);
1449        let json = serde_json::json!({"a": 1, "n": {"x": "y"}, "again": {"x": "z"}});
1450        let datum = codec.from_json(&json, BytesText::Hex).unwrap();
1451        assert_eq!(
1452            datum,
1453            datum_from_json_with(&json, &schema, BytesText::Hex).unwrap()
1454        );
1455        assert_eq!(
1456            codec.to_json(&datum).unwrap(),
1457            datum_to_json(&datum, &schema).unwrap()
1458        );
1459        assert_eq!(codec.to_json(&datum).unwrap(), json);
1460    }
1461
1462    #[test]
1463    fn datum_conversion_keeps_logical_types_exact_and_round_trips() {
1464        let schema = Schema::parse_str(
1465            r#"{"type":"record","name":"r","fields":[
1466                {"name":"amount","type":{"type":"bytes","logicalType":"decimal","precision":10,"scale":2}},
1467                {"name":"raw","type":"bytes"},
1468                {"name":"day","type":{"type":"int","logicalType":"date"}}
1469            ]}"#,
1470        )
1471        .unwrap();
1472        let v = json!({"amount": "12.34", "raw": "0aff", "day": "2026-10-07"});
1473        let datum = datum_from_json(&v, &schema).unwrap();
1474        assert_eq!(datum_to_json(&datum, &schema).unwrap(), v);
1475        assert_eq!(
1476            datum_to_json(&Av::Long(7), &Schema::Long).unwrap(),
1477            json!(7),
1478            "a non-record datum is not wrapped"
1479        );
1480        let err = datum_from_json(&json!({"amount": 1}), &schema).unwrap_err();
1481        assert!(
1482            matches!(&err, FaucetError::Sink(m) if m.starts_with("json->avro: ")),
1483            "{err}"
1484        );
1485    }
1486
1487    fn logical_schema() -> Value {
1488        json!({
1489            "type": "record",
1490            "name": "payment",
1491            "namespace": "test",
1492            "fields": [
1493                {"name": "id", "type": "long"},
1494                {"name": "amount", "type": {"type": "bytes", "logicalType": "decimal", "precision": 10, "scale": 2}},
1495                {"name": "fee", "type": {"type": "fixed", "name": "fee_t", "size": 8, "logicalType": "decimal", "precision": 12, "scale": 3}},
1496                {"name": "day", "type": {"type": "int", "logicalType": "date"}},
1497                {"name": "at", "type": {"type": "long", "logicalType": "timestamp-micros"}},
1498                {"name": "at_ms", "type": {"type": "long", "logicalType": "timestamp-millis"}},
1499                {"name": "at_ns", "type": {"type": "long", "logicalType": "timestamp-nanos"}},
1500                {"name": "local", "type": {"type": "long", "logicalType": "local-timestamp-micros"}},
1501                {"name": "local_ms", "type": {"type": "long", "logicalType": "local-timestamp-millis"}},
1502                {"name": "local_ns", "type": {"type": "long", "logicalType": "local-timestamp-nanos"}},
1503                {"name": "t_ms", "type": {"type": "int", "logicalType": "time-millis"}},
1504                {"name": "t_us", "type": {"type": "long", "logicalType": "time-micros"}},
1505                {"name": "uid", "type": {"type": "string", "logicalType": "uuid"}},
1506                {"name": "span", "type": {"type": "fixed", "name": "dur", "size": 12, "logicalType": "duration"}},
1507                {"name": "note", "type": ["null", "string"], "default": null},
1508                {"name": "either", "type": ["null", "long", "string"], "default": null},
1509                {"name": "kind", "type": {"type": "enum", "name": "kind_t", "symbols": ["A", "B"]}},
1510                {"name": "raw", "type": "bytes"},
1511                {"name": "tags", "type": {"type": "array", "items": "string"}},
1512                {"name": "attrs", "type": {"type": "map", "values": "int"}},
1513                {"name": "ratio", "type": "float"},
1514                {"name": "score", "type": "double"},
1515                {"name": "ok", "type": "boolean"},
1516                {"name": "small", "type": "int"},
1517                {"name": "big", "type": {"type": "bytes", "logicalType": "big-decimal"}}
1518            ]
1519        })
1520    }
1521
1522    fn logical_record() -> Value {
1523        json!({
1524            "id": 7,
1525            "amount": "-1234.50",
1526            "fee": "0.125",
1527            "day": "2024-02-29",
1528            "at": "2024-02-29T12:34:56.123456Z",
1529            "at_ms": "2024-02-29T12:34:56.123Z",
1530            "at_ns": "2024-02-29T12:34:56.123456789Z",
1531            "local": "2024-02-29T12:34:56.123456",
1532            "local_ms": "2024-02-29T12:34:56.123",
1533            "local_ns": "2024-02-29T12:34:56.123456789",
1534            "t_ms": "01:02:03.004",
1535            "t_us": "01:02:03.000004",
1536            "uid": "6f1c2b1e-4b8a-4b8e-9a38-3c3f3f0f9e11",
1537            "span": {"months": 1, "days": 2, "millis": 3},
1538            "note": null,
1539            "either": "text",
1540            "kind": "B",
1541            "raw": "AP8Q",
1542            "tags": ["x", "y"],
1543            "attrs": {"a": 1, "b": 2},
1544            "ratio": 0.5,
1545            "score": "NaN",
1546            "ok": true,
1547            "small": 3,
1548            "big": "123456789012345678901234567890.5"
1549        })
1550    }
1551
1552    #[test]
1553    fn logical_types_round_trip_through_an_ocf() {
1554        for codec in [
1555            AvroCodec::Null,
1556            AvroCodec::Deflate,
1557            AvroCodec::Snappy,
1558            AvroCodec::Zstd,
1559        ] {
1560            let opts = AvroOptions {
1561                schema: Some(logical_schema()),
1562                codec,
1563            };
1564            let bytes = encode(&[logical_record()], &opts).expect("encode");
1565            let back = decode(&bytes, &AvroOptions::default()).expect("decode");
1566            let mut expected = logical_record();
1567            expected["score"] = Value::Null;
1568            assert_eq!(back, vec![expected], "{codec:?}");
1569        }
1570    }
1571
1572    #[test]
1573    fn epoch_integers_are_accepted_for_temporal_types() {
1574        let schema = json!({"type": "record", "name": "r", "fields": [
1575            {"name": "d", "type": {"type": "int", "logicalType": "date"}},
1576            {"name": "t", "type": {"type": "long", "logicalType": "timestamp-millis"}},
1577            {"name": "tm", "type": {"type": "int", "logicalType": "time-millis"}}
1578        ]});
1579        let opts = AvroOptions {
1580            schema: Some(schema),
1581            ..Default::default()
1582        };
1583        let bytes = encode(&[json!({"d": 1, "t": 1000, "tm": 1500})], &opts).expect("encode");
1584        let back = decode(&bytes, &AvroOptions::default()).expect("decode");
1585        assert_eq!(
1586            back[0],
1587            json!({"d": "1970-01-02", "t": "1970-01-01T00:00:01.000Z", "tm": "00:00:01.500"})
1588        );
1589    }
1590
1591    #[test]
1592    fn inferred_schema_widens_nullable_nested_and_mixed_fields() {
1593        let recs = vec![
1594            json!({"id": 1, "name": "a", "first-name": "x", "nested": {"k": 1.5}, "mixed": 1, "list": [1, 2], "empty": []}),
1595            json!({"id": 2, "name": null, "first-name": "y", "nested": {"k": 2}, "mixed": "two", "list": [], "empty": []}),
1596        ];
1597        let bytes = encode(&recs, &AvroOptions::default()).expect("encode");
1598        let schema = writer_schema(&bytes[..]).expect("schema");
1599        let text = serde_json::to_string(&schema).expect("json");
1600        assert!(text.contains("\"first_name\""), "{text}");
1601        assert!(text.contains(ORIGINAL_NAME_ATTR), "{text}");
1602        let back = decode(&bytes, &AvroOptions::default()).expect("decode");
1603        assert_eq!(back[0]["id"], json!(1));
1604        assert_eq!(back[1]["name"], Value::Null);
1605        assert_eq!(back[0]["first_name"], json!("x"));
1606        assert_eq!(back[1]["nested"], json!({"k": 2.0}));
1607        assert_eq!(back[0]["mixed"], json!("1"));
1608        assert_eq!(back[1]["mixed"], json!("two"));
1609        assert_eq!(back[0]["list"], json!([1, 2]));
1610        assert_eq!(back[0]["empty"], json!([]));
1611    }
1612
1613    #[test]
1614    fn inference_refuses_colliding_sanitized_names() {
1615        let err =
1616            encode(&[json!({"a-b": 1, "a_b": 2})], &AvroOptions::default()).expect_err("collision");
1617        assert!(err.to_string().contains("sanitize"), "{err}");
1618    }
1619
1620    #[test]
1621    fn avro_names_are_sanitized() {
1622        assert_eq!(avro_name("ok_1"), "ok_1");
1623        assert_eq!(avro_name("first name"), "first_name");
1624        assert_eq!(avro_name("1st"), "_1st");
1625        assert_eq!(avro_name(""), "_");
1626    }
1627
1628    #[test]
1629    fn reader_schema_projects_and_defaults_added_fields() {
1630        let bytes = encode(&[json!({"a": 1, "b": "x"})], &AvroOptions::default()).expect("encode");
1631        let reader = json!({"type": "record", "name": INFERRED_RECORD_NAME, "fields": [
1632            {"name": "b", "type": "string"},
1633            {"name": "c", "type": "long", "default": 9}
1634        ]});
1635        let back = decode(
1636            &bytes,
1637            &AvroOptions {
1638                schema: Some(reader),
1639                ..Default::default()
1640            },
1641        )
1642        .expect("decode");
1643        assert_eq!(back, vec![json!({"b": "x", "c": 9})]);
1644    }
1645
1646    #[test]
1647    fn a_non_record_root_is_wrapped_in_value() {
1648        let opts = AvroOptions {
1649            schema: Some(json!("long")),
1650            ..Default::default()
1651        };
1652        let bytes = encode(&[json!({"value": 5}), json!(6)], &opts).expect("encode");
1653        let back = decode(&bytes, &AvroOptions::default()).expect("decode");
1654        assert_eq!(back, vec![json!({"value": 5}), json!({"value": 6})]);
1655    }
1656
1657    #[test]
1658    fn recursive_records_decode_as_nested_objects() {
1659        let schema = json!({"type": "record", "name": "node", "fields": [
1660            {"name": "v", "type": "long"},
1661            {"name": "next", "type": ["null", "node"], "default": null}
1662        ]});
1663        let rec = json!({"v": 1, "next": {"v": 2, "next": null}});
1664        let opts = AvroOptions {
1665            schema: Some(schema),
1666            ..Default::default()
1667        };
1668        let bytes = encode(std::slice::from_ref(&rec), &opts).expect("encode");
1669        assert_eq!(
1670            decode(&bytes, &AvroOptions::default()).expect("decode"),
1671            vec![rec]
1672        );
1673    }
1674
1675    #[test]
1676    fn encoding_errors_name_the_field() {
1677        let schema = json!({"type": "record", "name": "r", "fields": [
1678            {"name": "n", "type": "int"},
1679            {"name": "req", "type": "string"}
1680        ]});
1681        let opts = AvroOptions {
1682            schema: Some(schema),
1683            ..Default::default()
1684        };
1685        let err = encode(&[json!({"n": "x", "req": "a"})], &opts).expect_err("type");
1686        assert!(err.to_string().contains("`n`"), "{err}");
1687        let err = encode(&[json!({"n": 1})], &opts).expect_err("missing");
1688        assert!(err.to_string().contains("required"), "{err}");
1689        let err = encode(&[json!({"n": 1i64 << 40, "req": "a"})], &opts).expect_err("range");
1690        assert!(err.to_string().contains("does not fit"), "{err}");
1691    }
1692
1693    #[test]
1694    fn invalid_schemas_and_bodies_are_typed_errors() {
1695        let opts = AvroOptions {
1696            schema: Some(json!({"type": "nope"})),
1697            ..Default::default()
1698        };
1699        assert!(matches!(opts.parsed_schema(), Err(FaucetError::Config(_))));
1700        let err = decode(b"not avro", &AvroOptions::default()).expect_err("garbage");
1701        assert!(err.to_string().contains("avro"), "{err}");
1702    }
1703
1704    #[test]
1705    fn decimals_render_and_parse_exactly() {
1706        let cases = [
1707            ("0", 0),
1708            ("-1", 0),
1709            ("12.30", 2),
1710            ("-0.05", 2),
1711            ("170141183460469231731687303715884105727", 0),
1712        ];
1713        for (text, scale) in cases {
1714            let bytes = decimal_from_str(text, scale, 60).expect(text);
1715            assert_eq!(decimal_to_string(&bytes, scale), text, "{text}");
1716        }
1717        assert_eq!(
1718            decimal_to_string(&decimal_from_str("1.5e1", 1, 5).expect("exp"), 1),
1719            "15.0"
1720        );
1721        assert_eq!(
1722            decimal_to_string(&decimal_from_str("1.230", 2, 5).expect("trailing zero"), 2),
1723            "1.23"
1724        );
1725        assert!(decimal_from_str("1.234", 2, 5).is_err());
1726        // A huge exponent is refused before any padding (CORE-35).
1727        for huge in ["1e18", "1e9223372036854775807", "-5e400000000000"] {
1728            let err = decimal_from_str(huge, 2, 10).unwrap_err();
1729            assert!(err.contains("precision"), "{huge}: {err}");
1730        }
1731        assert_eq!(
1732            decimal_to_string(&decimal_from_str("0e400000000000", 2, 5).expect("zero"), 2),
1733            "0.00"
1734        );
1735        assert!(decimal_from_str("1e-9223372036854775807", 2, 5).is_err());
1736        assert!(decimal_from_str("123456", 0, 5).is_err());
1737        assert!(decimal_from_str("abc", 0, 5).is_err());
1738        assert!(decimal_from_str("1e-9", 2, 5).is_err());
1739        assert_eq!(decimal_from_str("0e-9", 2, 5).expect("zero"), vec![0]);
1740        assert_eq!(
1741            sign_extend(vec![0xff], 3).expect("neg"),
1742            vec![0xff, 0xff, 0xff]
1743        );
1744        assert!(sign_extend(vec![1, 2, 3], 2).is_err());
1745    }
1746
1747    #[test]
1748    fn value_coercions_and_their_refusals() {
1749        let names = HashMap::new();
1750        assert!(
1751            matches!(from_json(&json!(1), &Schema::String, &names, BytesText::Hex, "s"), Ok(Av::String(s)) if s == "1")
1752        );
1753        assert!(from_json(&Value::Null, &Schema::String, &names, BytesText::Hex, "s").is_err());
1754        assert!(from_json(&json!(1), &Schema::Null, &names, BytesText::Hex, "s").is_err());
1755        assert!(
1756            matches!(from_json(&json!([1, 2]), &Schema::Bytes, &names, BytesText::Hex, "b"), Ok(Av::Bytes(b)) if b == vec![1, 2])
1757        );
1758        assert!(from_json(&json!("abc"), &Schema::Bytes, &names, BytesText::Hex, "b").is_err());
1759        assert!(from_json(&json!("zz"), &Schema::Bytes, &names, BytesText::Hex, "b").is_err());
1760        assert!(from_json(&json!([300]), &Schema::Bytes, &names, BytesText::Hex, "b").is_err());
1761        assert!(from_json(&json!(true), &Schema::Bytes, &names, BytesText::Hex, "b").is_err());
1762        assert!(
1763            from_json(
1764                &json!("Infinity"),
1765                &Schema::Double,
1766                &names,
1767                BytesText::Hex,
1768                "d"
1769            )
1770            .is_ok()
1771        );
1772        assert!(
1773            from_json(
1774                &json!("-Infinity"),
1775                &Schema::Double,
1776                &names,
1777                BytesText::Hex,
1778                "d"
1779            )
1780            .is_ok()
1781        );
1782        assert!(from_json(&json!("x"), &Schema::Double, &names, BytesText::Hex, "d").is_err());
1783        assert!(from_json(&json!("x"), &Schema::Date, &names, BytesText::Hex, "d").is_err());
1784        assert!(
1785            from_json(
1786                &json!("x"),
1787                &Schema::TimestampMicros,
1788                &names,
1789                BytesText::Hex,
1790                "d"
1791            )
1792            .is_err()
1793        );
1794        assert!(
1795            from_json(
1796                &json!("2024-01-01 00:00:00"),
1797                &Schema::TimestampMicros,
1798                &names,
1799                BytesText::Hex,
1800                "d"
1801            )
1802            .is_ok()
1803        );
1804        assert!(
1805            from_json(
1806                &json!("2024-01-01T00:00:00+02:00"),
1807                &Schema::LocalTimestampMicros,
1808                &names,
1809                BytesText::Hex,
1810                "d"
1811            )
1812            .is_ok()
1813        );
1814        assert!(
1815            from_json(
1816                &json!(true),
1817                &Schema::TimeMicros,
1818                &names,
1819                BytesText::Hex,
1820                "d"
1821            )
1822            .is_err()
1823        );
1824        assert!(
1825            from_json(
1826                &json!("bad"),
1827                &Schema::TimeMicros,
1828                &names,
1829                BytesText::Hex,
1830                "d"
1831            )
1832            .is_err()
1833        );
1834        assert!(from_json(&json!(5), &Schema::TimeMicros, &names, BytesText::Hex, "d").is_ok());
1835        assert!(
1836            from_json(
1837                &json!(true),
1838                &Schema::TimestampMicros,
1839                &names,
1840                BytesText::Hex,
1841                "d"
1842            )
1843            .is_err()
1844        );
1845        assert!(
1846            from_json(
1847                &json!({"months": 1}),
1848                &parse_schema(
1849                    &json!({"type": "fixed", "name": "d", "size": 12, "logicalType": "duration"})
1850                )
1851                .unwrap(),
1852                &names,
1853                BytesText::Hex,
1854                "d"
1855            )
1856            .is_err()
1857        );
1858        let e = parse_schema(&json!({"type": "enum", "name": "e", "symbols": ["A"]})).unwrap();
1859        assert!(from_json(&json!("Z"), &e, &named(&e), BytesText::Hex, "e").is_err());
1860        let f = parse_schema(&json!({"type": "fixed", "name": "f", "size": 2})).unwrap();
1861        assert!(from_json(&json!("00"), &f, &named(&f), BytesText::Hex, "f").is_err());
1862        let u = parse_schema(&json!({"type": "string", "logicalType": "uuid"})).unwrap();
1863        assert!(from_json(&json!("nope"), &u, &names, BytesText::Hex, "u").is_err());
1864        let bd = parse_schema(&json!({"type": "bytes", "logicalType": "big-decimal"})).unwrap();
1865        assert!(from_json(&json!(1.5), &bd, &names, BytesText::Hex, "b").is_ok());
1866        assert!(from_json(&json!("x"), &bd, &names, BytesText::Hex, "b").is_err());
1867        assert!(from_json(&json!(true), &bd, &names, BytesText::Hex, "b").is_err());
1868        let dec = parse_schema(
1869            &json!({"type": "bytes", "logicalType": "decimal", "precision": 4, "scale": 1}),
1870        )
1871        .unwrap();
1872        assert!(from_json(&json!(1.5), &dec, &names, BytesText::Hex, "d").is_ok());
1873        assert!(from_json(&json!(true), &dec, &names, BytesText::Hex, "d").is_err());
1874        let un = parse_schema(&json!(["long", "boolean"])).unwrap();
1875        assert!(from_json(&json!("x"), &un, &names, BytesText::Hex, "u").is_err());
1876        assert!(from_json(&Value::Null, &un, &names, BytesText::Hex, "u").is_err());
1877        assert!(short(&json!("x".repeat(100))).ends_with('…'));
1878    }
1879
1880    #[test]
1881    fn utf8_bytes_text_reads_strings_as_their_bytes() {
1882        let schema = parse_schema(&json!({"type": "record", "name": "r", "fields": [
1883            {"name": "raw", "type": "bytes"},
1884            {"name": "tag", "type": {"type": "fixed", "name": "t", "size": 4}}
1885        ]}))
1886        .unwrap();
1887        let v = json!({"raw": "cafe", "tag": "abcd"});
1888        assert_eq!(
1889            datum_from_json_with(&v, &schema, BytesText::Utf8).unwrap(),
1890            Av::Record(vec![
1891                ("raw".into(), Av::Bytes(b"cafe".to_vec())),
1892                ("tag".into(), Av::Fixed(4, b"abcd".to_vec())),
1893            ])
1894        );
1895        assert!(datum_from_json(&json!({"raw": "hello", "tag": "abcd"}), &schema).is_err());
1896        assert!(
1897            datum_from_json_with(
1898                &json!({"raw": [1, 2], "tag": "abcd"}),
1899                &schema,
1900                BytesText::Utf8
1901            )
1902            .is_ok()
1903        );
1904        assert_eq!(BytesText::default(), BytesText::Base64);
1905        assert_eq!(
1906            datum_from_json_with(
1907                &json!({"raw": "cafe", "tag": "61626364"}),
1908                &schema,
1909                BytesText::Hex
1910            )
1911            .unwrap(),
1912            Av::Record(vec![
1913                ("raw".into(), Av::Bytes(vec![0xca, 0xfe])),
1914                ("tag".into(), Av::Fixed(4, b"abcd".to_vec())),
1915            ])
1916        );
1917    }
1918
1919    #[test]
1920    fn out_of_range_temporal_values_are_errors() {
1921        assert!(date_string(i32::MAX).is_err());
1922        assert!(time_string(-1, 3).is_err());
1923        assert!(utc_string(i64::MAX, 1_000_000, SecondsFormat::Millis).is_err());
1924        assert!(local_string(i64::MAX, 1_000_000, 3).is_err());
1925        assert_eq!(crate::fidelity::float64(f64::NEG_INFINITY), Value::Null);
1926    }
1927
1928    #[cfg(feature = "arrow")]
1929    #[test]
1930    fn arrow_batches_carry_logical_types() {
1931        use arrow::datatypes::{DataType as D, TimeUnit};
1932        let opts = AvroOptions {
1933            schema: Some(logical_schema()),
1934            ..Default::default()
1935        };
1936        let rec = logical_record();
1937        let bytes = encode(&[rec.clone(), rec], &opts).expect("encode");
1938        let mut batches = Vec::new();
1939        let (_, schema) = read_batches(&bytes[..], None, 1, &mut |b| {
1940            batches.push(b);
1941            Ok(())
1942        })
1943        .expect("batches");
1944        assert_eq!(batches.len(), 2);
1945        let ty = |n: &str| schema.field_with_name(n).expect(n).data_type().clone();
1946        assert_eq!(ty("amount"), D::Decimal128(10, 2));
1947        assert_eq!(ty("day"), D::Date32);
1948        assert_eq!(
1949            ty("at"),
1950            D::Timestamp(TimeUnit::Microsecond, Some("UTC".into()))
1951        );
1952        assert_eq!(ty("local"), D::Timestamp(TimeUnit::Microsecond, None));
1953        assert_eq!(ty("t_ms"), D::Time32(TimeUnit::Millisecond));
1954        assert_eq!(ty("either"), D::Utf8);
1955        assert!(schema.field_with_name("note").unwrap().is_nullable());
1956        let back = crate::columnar::record_batch_to_values(&batches[0]).expect("values");
1957        let amount = batches[0]
1958            .column_by_name("amount")
1959            .and_then(|c| c.as_any().downcast_ref::<arrow::array::Decimal128Array>())
1960            .expect("decimal column");
1961        assert_eq!(amount.value(0), -123_450);
1962        assert_eq!(back[0]["day"], json!("2024-02-29"));
1963        assert_eq!(back[0]["either"], json!("\"text\""));
1964    }
1965
1966    #[cfg(feature = "arrow")]
1967    #[test]
1968    fn arrow_schema_handles_roots_recursion_and_wide_decimals() {
1969        use arrow::datatypes::DataType as D;
1970        let s = parse_schema(&json!("long")).unwrap();
1971        assert_eq!(arrow_schema(&s).unwrap().field(0).name(), VALUE_FIELD);
1972        let rec = parse_schema(&json!({"type": "record", "name": "node", "fields": [
1973            {"name": "next", "type": ["null", "node"], "default": null},
1974            {"name": "wide", "type": {"type": "bytes", "logicalType": "decimal", "precision": 50, "scale": 2}},
1975            {"name": "huge", "type": {"type": "bytes", "logicalType": "decimal", "precision": 90, "scale": 2}},
1976            {"name": "only_null", "type": ["null"]},
1977            {"name": "n", "type": "null"}
1978        ]}))
1979        .unwrap();
1980        let a = arrow_schema(&rec).unwrap();
1981        assert_eq!(a.field(0).data_type(), &D::Utf8);
1982        assert_eq!(a.field(1).data_type(), &D::Decimal256(50, 2));
1983        assert_eq!(a.field(2).data_type(), &D::Utf8);
1984        assert_eq!(a.field(3).data_type(), &D::Null);
1985        let opts = AvroOptions {
1986            schema: Some(json!({"type": "record", "name": "node", "fields": [
1987                {"name": "v", "type": "long"},
1988                {"name": "next", "type": ["null", "node"], "default": null}
1989            ]})),
1990            ..Default::default()
1991        };
1992        let bytes = encode(&[json!({"v": 1, "next": {"v": 2, "next": null}})], &opts).unwrap();
1993        let mut rows = Vec::new();
1994        read_batches(&bytes[..], None, 0, &mut |b| {
1995            rows.extend(crate::columnar::record_batch_to_values(&b)?);
1996            Ok(())
1997        })
1998        .unwrap();
1999        let next: Value = serde_json::from_str(rows[0]["next"].as_str().unwrap()).unwrap();
2000        assert_eq!(next, json!({"v": 2, "next": null}));
2001    }
2002}