Skip to content
This repository has been archived by the owner on Feb 18, 2024. It is now read-only.

Commit

Permalink
Improved example and added guide entry
Browse files Browse the repository at this point in the history
  • Loading branch information
jorgecarleitao committed Dec 20, 2021
1 parent ceb6b91 commit 568c902
Show file tree
Hide file tree
Showing 2 changed files with 36 additions and 23 deletions.
58 changes: 35 additions & 23 deletions examples/avro_write.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use std::{fs::File, sync::Arc};
use std::fs::File;

use arrow2::{
array::{Array, Int32Array},
Expand All @@ -7,6 +7,39 @@ use arrow2::{
io::avro::write,
};

fn write_avro<W: std::io::Write>(
file: &mut W,
arrays: &[&dyn Array],
schema: &Schema,
compression: Option<write::Compression>,
) -> 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::<Vec<_>>();
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<String> = env::args().collect();
Expand All @@ -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(())
}
1 change: 1 addition & 0 deletions guide/src/SUMMARY.md
Original file line number Diff line number Diff line change
Expand Up @@ -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)

0 comments on commit 568c902

Please sign in to comment.