diff --git a/examples/avro_write.rs b/examples/avro_write.rs index 866ebe05f9c..6d970bd9e01 100644 --- a/examples/avro_write.rs +++ b/examples/avro_write.rs @@ -1,4 +1,4 @@ -use std::{fs::File, sync::Arc}; +use std::fs::File; use arrow2::{ array::{Array, Int32Array}, @@ -7,6 +7,39 @@ use arrow2::{ io::avro::write, }; +fn write_avro( + file: &mut W, + arrays: &[&dyn Array], + schema: &Schema, + compression: Option, +) -> Result<()> { + let avro_fields = write::to_avro_schema(schema)?; + + let mut serializers = arrays + .iter() + .zip(avro_fields.iter()) + .map(|(array, field)| write::new_serializer(*array, &field.schema)) + .collect::>(); + let mut block = write::Block::new(arrays[0].len(), vec![]); + + write::serialize(&mut serializers, &mut block)?; + + let mut compressed_block = write::CompressedBlock::default(); + + if let Some(compression) = compression { + write::compress(&block, &mut compressed_block, compression)?; + } else { + compressed_block.number_of_rows = block.number_of_rows; + std::mem::swap(&mut compressed_block.data, &mut block.data); + } + + write::write_metadata(file, avro_fields.clone(), compression)?; + + write::write_block(file, &compressed_block)?; + + Ok(()) +} + fn main() -> Result<()> { use std::env; let args: Vec = env::args().collect(); @@ -25,29 +58,8 @@ fn main() -> Result<()> { let field = Field::new("c1", array.data_type().clone(), true); let schema = Schema::new(vec![field]); - let avro_fields = write::to_avro_schema(&schema)?; - let mut file = File::create(path)?; - - let compression = None; - - write::write_metadata(&mut file, avro_fields.clone(), compression)?; - - let serializer = write::new_serializer(&array, &avro_fields[0].schema); - let mut block = write::Block::new(array.len(), vec![]); - - write::serialize(&mut [serializer], &mut block)?; - - let mut compressed_block = write::CompressedBlock::default(); - - if let Some(compression) = compression { - write::compress(&block, &mut compressed_block, compression)?; - } else { - compressed_block.number_of_rows = block.number_of_rows; - std::mem::swap(&mut compressed_block.data, &mut block.data); - } - - write::write_block(&mut file, &compressed_block)?; + write_avro(&mut file, &[(&array) as &dyn Array], &schema, None)?; Ok(()) } diff --git a/guide/src/SUMMARY.md b/guide/src/SUMMARY.md index a66c993096d..9719504ee56 100644 --- a/guide/src/SUMMARY.md +++ b/guide/src/SUMMARY.md @@ -17,3 +17,4 @@ - [Read Arrow stream](./io/ipc_stream_read.md) - [Write Arrow](./io/ipc_write.md) - [Read Avro](./io/avro_read.md) + - [Write Avro](./io/avro_write.md)