apache_avro/writer/
datum.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use bon::bon;
19use serde::Serialize;
20use std::io::Write;
21
22use crate::{
23    AvroResult, Schema,
24    encode::encode_internal,
25    error::Details,
26    schema::{NamesRef, ResolvedSchema},
27    serde::ser_schema::{Config, SchemaAwareSerializer},
28    types::Value,
29    util::is_human_readable,
30};
31
32/// Writer for writing raw Avro data.
33///
34/// This is most likely not what you need. Most users should use [`Writer`][crate::Writer],
35/// [`GenericSingleObjectWriter`][crate::GenericSingleObjectWriter], or
36/// [`SpecificSingleObjectWriter`][crate::SpecificSingleObjectWriter] instead.
37pub struct GenericDatumWriter<'s> {
38    schema: &'s Schema,
39    resolved: ResolvedSchema<'s>,
40    validate: bool,
41    human_readable: bool,
42    target_block_size: Option<usize>,
43}
44
45#[bon]
46impl<'s> GenericDatumWriter<'s> {
47    /// Configure a new writer.
48    #[builder]
49    pub fn new(
50        /// The schema for the data that will be written
51        #[builder(start_fn)]
52        schema: &'s Schema,
53        /// Already resolved schemata that will be used to resolve references in the writer's schema.
54        ///
55        /// You can also use [`Self::schemata`] instead.
56        resolved_schemata: Option<ResolvedSchema<'s>>,
57        /// Validate values against the writer schema before writing them.
58        ///
59        /// Defaults to `true`.
60        ///
61        /// Setting this to `false` and writing values that don't match the schema will make the
62        /// written data unreadable.
63        #[builder(default = true)]
64        validate: bool,
65        /// At what block size to start a new block (for arrays and maps).
66        ///
67        /// This is a minimum value, the block size will always be larger than this except for the last
68        /// block.
69        ///
70        /// When set to `None` all values will be written in a single block. This can be faster as no
71        /// intermediate buffer is used, but seeking through written data will be slower.
72        target_block_size: Option<usize>,
73        /// Should [`Serialize`] implementations pick a human readable representation.
74        ///
75        /// It is recommended to set this to `false`.
76        #[builder(default = is_human_readable())]
77        human_readable: bool,
78    ) -> AvroResult<Self> {
79        let resolved = if let Some(resolved) = resolved_schemata {
80            resolved
81        } else {
82            ResolvedSchema::try_from(schema)?
83        };
84        Ok(Self {
85            schema,
86            resolved,
87            validate,
88            human_readable,
89            target_block_size,
90        })
91    }
92}
93
94impl<'s, S: generic_datum_writer_builder::State> GenericDatumWriterBuilder<'s, S> {
95    /// Set the schemata that will be used to resolve any references in the schema.
96    ///
97    /// This is equivalent to `.resolved_schemata(ResolvedSchema::new_with_schemata(schemata)?)`.
98    /// If you already have a [`ResolvedSchema`], use that function instead.
99    pub fn schemata(
100        self,
101        schemata: Vec<&'s Schema>,
102    ) -> AvroResult<
103        GenericDatumWriterBuilder<'s, generic_datum_writer_builder::SetResolvedSchemata<S>>,
104    >
105    where
106        S::ResolvedSchemata: generic_datum_writer_builder::IsUnset,
107    {
108        let resolved = ResolvedSchema::new_with_schemata(schemata)?;
109        Ok(self.resolved_schemata(resolved))
110    }
111}
112
113impl GenericDatumWriter<'_> {
114    /// Write a value to the writer.
115    pub fn write_value<W: Write, V: Into<Value>>(
116        &self,
117        writer: &mut W,
118        value: V,
119    ) -> AvroResult<usize> {
120        let value = value.into();
121        self.write_value_ref(writer, &value)
122    }
123
124    /// Write a value to the writer.
125    pub fn write_value_ref<W: Write>(&self, writer: &mut W, value: &Value) -> AvroResult<usize> {
126        if self.validate
127            && value
128                .validate_internal(self.schema, self.resolved.get_names(), None)
129                .is_some()
130        {
131            return Err(Details::Validation.into());
132        }
133        encode_internal(value, self.schema, self.resolved.get_names(), None, writer)
134    }
135
136    /// Write a value to a [`Vec`].
137    pub fn write_value_to_vec<V: Into<Value>>(&self, value: V) -> AvroResult<Vec<u8>> {
138        let mut vec = Vec::new();
139        self.write_value(&mut vec, value)?;
140        Ok(vec)
141    }
142
143    /// Serialize `T` to the writer.
144    pub fn write_ser<W: Write, T: Serialize>(
145        &self,
146        writer: &mut W,
147        value: &T,
148    ) -> AvroResult<usize> {
149        let config = Config {
150            names: self.resolved.get_names(),
151            target_block_size: self.target_block_size,
152            human_readable: self.human_readable,
153        };
154        value.serialize(SchemaAwareSerializer::new(writer, self.schema, config)?)
155    }
156
157    /// Serialize `T` to a [`Vec`].
158    pub fn write_ser_to_vec<T: Serialize>(&self, value: &T) -> AvroResult<Vec<u8>> {
159        let mut vec = Vec::new();
160        self.write_ser(&mut vec, value)?;
161        Ok(vec)
162    }
163}
164
165/// Deprecated. Use [`GenericDatumWriter`] instead.
166///
167/// This is equivalent to:
168/// ```ignore
169/// GenericDatumWriter::builder(schema)
170///     .build()?
171///     .write_value_to_vec(value)
172/// ```
173///
174/// Encode a value into raw Avro data, also performs schema validation.
175///
176/// **NOTE**: This function has a quite small niche of usage and does NOT generate headers and sync
177/// markers; use [`Writer`] to be fully Avro-compatible if you don't know what
178/// you are doing, instead.
179///
180/// [`Writer`]: crate::Writer
181#[deprecated(since = "0.22.0", note = "Use GenericDatumWriter instead")]
182pub fn to_avro_datum<T: Into<Value>>(schema: &Schema, value: T) -> AvroResult<Vec<u8>> {
183    GenericDatumWriter::builder(schema)
184        .build()?
185        .write_value_to_vec(value)
186}
187
188/// Write the referenced [Serialize]able object to the provided [Write] object.
189///
190/// It is recommended to use [`GenericDatumWriter`] instead.
191///
192/// Returns a result with the number of bytes written.
193///
194/// **NOTE**: This function has a quite small niche of usage and does **NOT** generate headers and sync
195/// markers; use [`Writer::append_ser`] to be fully Avro-compatible
196/// if you don't know what you are doing, instead.
197///
198/// [`Writer::append_ser`]: crate::Writer::append_ser
199pub fn write_avro_datum_ref<T: Serialize, W: Write>(
200    schema: &Schema,
201    names: &NamesRef,
202    data: &T,
203    writer: &mut W,
204) -> AvroResult<usize> {
205    let config = Config {
206        names,
207        target_block_size: None,
208        human_readable: is_human_readable(),
209    };
210    data.serialize(SchemaAwareSerializer::new(writer, schema, config)?)
211}
212
213/// Deprecated. Use [`GenericDatumWriter`] instead.
214///
215/// This is equivalent to:
216/// ```ignore
217/// GenericDatumWriter::builder(schema)
218///     .schemata(schemata)?
219///     .build()?
220///     .write_value_to_vec(value)
221/// ```
222///
223/// Encode a value into raw Avro data, also performs schema validation.
224///
225/// If the provided `schema` is incomplete then its dependencies must be
226/// provided in `schemata`
227#[deprecated(since = "0.22.0", note = "Use GenericDatumWriter instead")]
228pub fn to_avro_datum_schemata<T: Into<Value>>(
229    schema: &Schema,
230    schemata: Vec<&Schema>,
231    value: T,
232) -> AvroResult<Vec<u8>> {
233    GenericDatumWriter::builder(schema)
234        .schemata(schemata)?
235        .build()?
236        .write_value_to_vec(value)
237}
238
239#[cfg(test)]
240mod tests {
241    use apache_avro_test_helper::TestResult;
242    use std::num::NonZero;
243
244    use super::*;
245    use crate::reader::datum::GenericDatumReader;
246    use crate::{
247        Days, Decimal, Duration, Millis, Months,
248        schema::{DecimalSchema, FixedSchema, InnerDecimalSchema, Name},
249        types::Record,
250        util::zig_i64,
251    };
252
253    const SCHEMA: &str = r#"
254    {
255      "type": "record",
256      "name": "test",
257      "fields": [
258        {
259          "name": "a",
260          "type": "long",
261          "default": 42
262        },
263        {
264          "name": "b",
265          "type": "string"
266        }
267      ]
268    }
269    "#;
270
271    const UNION_SCHEMA: &str = r#"["null", "long"]"#;
272
273    #[test]
274    fn test_to_avro_datum() -> TestResult {
275        let schema = Schema::parse_str(SCHEMA)?;
276        let mut record = Record::new(&schema).unwrap();
277        record.put("a", 27i64);
278        record.put("b", "foo");
279
280        let mut expected = Vec::new();
281        zig_i64(27, &mut expected)?;
282        zig_i64(3, &mut expected)?;
283        expected.extend(b"foo");
284
285        let written = GenericDatumWriter::builder(&schema)
286            .build()?
287            .write_value_to_vec(record)?;
288
289        assert_eq!(written, expected);
290
291        Ok(())
292    }
293
294    #[test]
295    fn avro_rs_193_write_avro_datum_ref() -> TestResult {
296        #[derive(Serialize)]
297        struct TestStruct {
298            a: i64,
299            b: String,
300        }
301
302        let schema = Schema::parse_str(SCHEMA)?;
303        let mut writer: Vec<u8> = Vec::new();
304        let data = TestStruct {
305            a: 27,
306            b: "foo".to_string(),
307        };
308
309        let mut expected = Vec::new();
310        zig_i64(27, &mut expected)?;
311        zig_i64(3, &mut expected)?;
312        expected.extend(b"foo");
313
314        let bytes = GenericDatumWriter::builder(&schema)
315            .build()?
316            .write_ser(&mut writer, &data)?;
317
318        assert_eq!(bytes, expected.len());
319        assert_eq!(writer, expected);
320
321        Ok(())
322    }
323
324    #[test]
325    fn test_union_not_null() -> TestResult {
326        let schema = Schema::parse_str(UNION_SCHEMA)?;
327        let union = Value::Union(1, Box::new(Value::Long(3)));
328
329        let mut expected = Vec::new();
330        zig_i64(1, &mut expected)?;
331        zig_i64(3, &mut expected)?;
332
333        let written = GenericDatumWriter::builder(&schema)
334            .build()?
335            .write_value_to_vec(union)?;
336        assert_eq!(written, expected);
337
338        Ok(())
339    }
340
341    #[test]
342    fn test_union_null() -> TestResult {
343        let schema = Schema::parse_str(UNION_SCHEMA)?;
344        let union = Value::Union(0, Box::new(Value::Null));
345
346        let mut expected = Vec::new();
347        zig_i64(0, &mut expected)?;
348
349        let written = GenericDatumWriter::builder(&schema)
350            .build()?
351            .write_value_to_vec(union)?;
352        assert_eq!(written, expected);
353
354        Ok(())
355    }
356
357    #[expect(
358        clippy::needless_pass_by_value,
359        reason = "Value does not implement PartialEq<&Value>"
360    )]
361    fn logical_type_test<T: Into<Value> + Clone>(
362        schema_str: &'static str,
363
364        expected_schema: &Schema,
365        value: Value,
366
367        raw_schema: &Schema,
368        raw_value: T,
369    ) -> TestResult {
370        let schema = Schema::parse_str(schema_str)?;
371        assert_eq!(&schema, expected_schema);
372        // The serialized format should be the same as the schema.
373        let ser = GenericDatumWriter::builder(&schema)
374            .build()?
375            .write_value_to_vec(value.clone())?;
376        let raw_ser = GenericDatumWriter::builder(raw_schema)
377            .build()?
378            .write_value_to_vec(raw_value)?;
379        assert_eq!(ser, raw_ser);
380
381        // Should deserialize from the schema into the logical type.
382        let mut r = ser.as_slice();
383        let de = GenericDatumReader::builder(&schema)
384            .build()?
385            .read_value(&mut r)?;
386        assert_eq!(de, value);
387        Ok(())
388    }
389
390    #[test]
391    fn date() -> TestResult {
392        logical_type_test(
393            r#"{"type": "int", "logicalType": "date"}"#,
394            &Schema::Date,
395            Value::Date(1_i32),
396            &Schema::Int,
397            1_i32,
398        )
399    }
400
401    #[test]
402    fn time_millis() -> TestResult {
403        logical_type_test(
404            r#"{"type": "int", "logicalType": "time-millis"}"#,
405            &Schema::TimeMillis,
406            Value::TimeMillis(1_i32),
407            &Schema::Int,
408            1_i32,
409        )
410    }
411
412    #[test]
413    fn time_micros() -> TestResult {
414        logical_type_test(
415            r#"{"type": "long", "logicalType": "time-micros"}"#,
416            &Schema::TimeMicros,
417            Value::TimeMicros(1_i64),
418            &Schema::Long,
419            1_i64,
420        )
421    }
422
423    #[test]
424    fn timestamp_millis() -> TestResult {
425        logical_type_test(
426            r#"{"type": "long", "logicalType": "timestamp-millis"}"#,
427            &Schema::TimestampMillis,
428            Value::TimestampMillis(1_i64),
429            &Schema::Long,
430            1_i64,
431        )
432    }
433
434    #[test]
435    fn timestamp_micros() -> TestResult {
436        logical_type_test(
437            r#"{"type": "long", "logicalType": "timestamp-micros"}"#,
438            &Schema::TimestampMicros,
439            Value::TimestampMicros(1_i64),
440            &Schema::Long,
441            1_i64,
442        )
443    }
444
445    #[test]
446    fn decimal_fixed() -> TestResult {
447        let size = 30;
448        let fixed = FixedSchema {
449            name: Name::new("decimal")?,
450            aliases: None,
451            doc: None,
452            size,
453            attributes: Default::default(),
454        };
455        let inner = InnerDecimalSchema::Fixed(fixed.clone());
456        let value = vec![0u8; size];
457        logical_type_test(
458            r#"{"type": {"type": "fixed", "size": 30, "name": "decimal"}, "logicalType": "decimal", "precision": 20, "scale": 5}"#,
459            &Schema::Decimal(DecimalSchema {
460                precision: NonZero::new(20).unwrap(),
461                scale: 5,
462                inner,
463            }),
464            Value::Decimal(Decimal::from(value.clone())),
465            &Schema::Fixed(fixed),
466            Value::Fixed(size, value),
467        )
468    }
469
470    #[test]
471    fn decimal_bytes() -> TestResult {
472        let value = vec![0u8; 10];
473        logical_type_test(
474            r#"{"type": "bytes", "logicalType": "decimal", "precision": 4, "scale": 3}"#,
475            &Schema::Decimal(DecimalSchema {
476                precision: NonZero::new(4).unwrap(),
477                scale: 3,
478                inner: InnerDecimalSchema::Bytes,
479            }),
480            Value::Decimal(Decimal::from(value.clone())),
481            &Schema::Bytes,
482            value,
483        )
484    }
485
486    #[test]
487    fn duration() -> TestResult {
488        let inner = Schema::Fixed(FixedSchema {
489            name: Name::new("duration")?,
490            aliases: None,
491            doc: None,
492            size: 12,
493            attributes: Default::default(),
494        });
495        let value = Value::Duration(Duration::new(
496            Months::new(256),
497            Days::new(512),
498            Millis::new(1024),
499        ));
500        logical_type_test(
501            r#"{"type": {"type": "fixed", "name": "duration", "size": 12}, "logicalType": "duration"}"#,
502            &Schema::Duration(FixedSchema {
503                name: Name::try_from("duration").expect("Name is valid"),
504                aliases: None,
505                doc: None,
506                size: 12,
507                attributes: Default::default(),
508            }),
509            value,
510            &inner,
511            Value::Fixed(12, vec![0, 1, 0, 0, 0, 2, 0, 0, 0, 4, 0, 0]),
512        )
513    }
514}
apache_avro/writer/datum.rs

apache_avro/writer/
datum.rs
