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

Commit

Permalink
adding test for LargeUtf8 and LargeBinary io
Browse files Browse the repository at this point in the history
  • Loading branch information
illumination-k committed Feb 10, 2022
1 parent e8aeb5e commit 6928229
Showing 1 changed file with 59 additions and 0 deletions.
59 changes: 59 additions & 0 deletions tests/it/io/avro/write.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
use std::fmt::Binary;

use arrow2::array::*;
use arrow2::chunk::Chunk;
use arrow2::datatypes::*;
Expand Down Expand Up @@ -140,3 +142,60 @@ fn snappy() -> Result<()> {
fn deflate() -> Result<()> {
roundtrip(Some(write::Compression::Deflate))
}

fn large_format_schema() -> Schema {
Schema::from(vec![
Field::new("large_utf8", DataType::LargeUtf8, false),
Field::new("large_utf8_nullable", DataType::LargeUtf8, true),
Field::new("large_binary", DataType::LargeBinary, false),
Field::new("large_binary_nullable", DataType::LargeBinary, true),
])
}

fn large_format_data() -> Chunk<Box<dyn Array>> {
let columns = vec![
Box::new(Utf8Array::<i64>::from_slice(&["a", "b"])) as Box<dyn Array>,
Box::new(Utf8Array::<i64>::from(&[Some("a"), None])),
Box::new(BinaryArray::<i64>::from_slice([b"foo", b"bar"])),
Box::new(BinaryArray::<i64>::from([Some(b"foo"), None])),
];
Chunk::new(columns)
}

fn large_format_expected_schema() -> Schema {
Schema::from(vec![
Field::new("large_utf8", DataType::Utf8, false),
Field::new("large_utf8_nullable", DataType::Utf8, true),
Field::new("large_binary", DataType::Binary, false),
Field::new("large_binary_nullable", DataType::Binary, true),
])
}

fn large_format_expected_data() -> Chunk<Box<dyn Array>> {
let columns = vec![
Box::new(Utf8Array::<i32>::from_slice(&["a", "b"])) as Box<dyn Array>,
Box::new(Utf8Array::<i32>::from(&[Some("a"), None])),
Box::new(BinaryArray::<i32>::from_slice([b"foo", b"bar"])),
Box::new(BinaryArray::<i32>::from([Some(b"foo"), None])),
];
Chunk::new(columns)
}

#[test]
fn check_large_format() -> Result<()> {
let write_schema = large_format_schema();
let write_data = large_format_data();

let data = write_avro(&write_data, &write_schema, None)?;
let (result, read_schame) = read_avro(&data, None)?;

let expected_schema = large_format_expected_schema();
assert_eq!(read_schame, expected_schema);

let expected_data = large_format_expected_data();
for (c1, c2) in result.columns().iter().zip(expected_data.columns().iter()) {
assert_eq!(c1.as_ref(), c2.as_ref());
}

Ok(())
}

0 comments on commit 6928229

Please sign in to comment.