Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Cleanup parquet tests #3116

Merged
merged 2 commits into from
Nov 22, 2022
Merged

Cleanup parquet tests #3116

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
51dc726
Cleanup parquet tests
tustvold Nov 15, 2022
13f5d07
Merge remote-tracking branch 'upstream/master' into cleanup-parquet-t…
tustvold Nov 21, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
117 changes: 49 additions & 68 deletions parquet/src/file/writer.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -648,14 +648,15 @@ mod tests {
use crate::basic::{Compression, Encoding, LogicalType, Repetition, Type};
use crate::column::page::PageReader;
use crate::compression::{create_codec, Codec, CodecOptionsBuilder};
use crate::data_type::Int32Type;
use crate::data_type::{BoolType, Int32Type};
use crate::file::reader::ChunkReader;
use crate::file::{
properties::{ReaderProperties, WriterProperties, WriterVersion},
reader::{FileReader, SerializedFileReader, SerializedPageReader},
statistics::{from_thrift, to_thrift, Statistics},
};
use crate::format::SortingColumn;
use crate::record::RowAccessor;
use crate::record::{Row, RowAccessor};
use crate::schema::types::{ColumnDescriptor, ColumnPath};
use crate::util::memory::ByteBufferPtr;

Expand Down Expand Up @@ -1163,16 +1164,35 @@ mod tests {
assert_eq!(to_thrift(left.statistics()), to_thrift(right.statistics()));
}

/// File write-read roundtrip.
/// `data` consists of arrays of values for each row group.
fn test_file_roundtrip(
file: File,
/// Tests roundtrip of i32 data written using `W` and read using `R`
fn test_roundtrip_i32<W, R>(
file: W,
data: Vec<Vec<i32>>,
) -> crate::format::FileMetaData {
) -> crate::format::FileMetaData
where
W: Write,
R: ChunkReader + From<W> + 'static,
{
test_roundtrip::<W, R, Int32Type, _>(file, data, |r| r.get_int(0).unwrap())
}

/// Tests roundtrip of data of type `D` written using `W` and read using `R`
/// and the provided `values` function
fn test_roundtrip<W, R, D, F>(
mut file: W,
data: Vec<Vec<D::T>>,
value: F,
) -> crate::format::FileMetaData
where
W: Write,
R: ChunkReader + From<W> + 'static,
D: DataType,
F: Fn(Row) -> D::T,
{
let schema = Arc::new(
types::Type::group_type_builder("schema")
.with_fields(&mut vec![Arc::new(
types::Type::primitive_type_builder("col1", Type::INT32)
types::Type::primitive_type_builder("col1", D::get_physical_type())
.with_repetition(Repetition::REQUIRED)
.build()
.unwrap(),
Expand All @@ -1181,16 +1201,15 @@ mod tests {
.unwrap(),
);
let props = Arc::new(WriterProperties::builder().build());
let mut file_writer = assert_send(
SerializedFileWriter::new(file.try_clone().unwrap(), schema, props).unwrap(),
);
let mut file_writer =
SerializedFileWriter::new(&mut file, schema, props).unwrap();
let mut rows: i64 = 0;

for (idx, subset) in data.iter().enumerate() {
let mut row_group_writer = file_writer.next_row_group().unwrap();
if let Some(mut writer) = row_group_writer.next_column().unwrap() {
rows += writer
.typed::<Int32Type>()
.typed::<D>()
.write_batch(&subset[..], None, None)
.unwrap() as i64;
writer.close().unwrap();
Expand All @@ -1202,7 +1221,7 @@ mod tests {
}
let file_metadata = file_writer.close().unwrap();

let reader = assert_send(SerializedFileReader::new(file).unwrap());
let reader = SerializedFileReader::new(R::from(file)).unwrap();
assert_eq!(reader.num_row_groups(), data.len());
assert_eq!(
reader.metadata().file_metadata().num_rows(),
Expand All @@ -1212,16 +1231,19 @@ mod tests {
for (i, item) in data.iter().enumerate().take(reader.num_row_groups()) {
let row_group_reader = reader.get_row_group(i).unwrap();
let iter = row_group_reader.get_row_iter(None).unwrap();
let res = iter
.map(|elem| elem.get_int(0).unwrap())
.collect::<Vec<i32>>();
let res: Vec<_> = iter.map(&value).collect();
assert_eq!(res, *item);
}
file_metadata
}

fn assert_send<T: Send>(t: T) -> T {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is this check covered elsewhere ?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, all the async plumbing requires it

t
/// File write-read roundtrip.
/// `data` consists of arrays of values for each row group.
fn test_file_roundtrip(
file: File,
data: Vec<Vec<i32>>,
) -> crate::format::FileMetaData {
test_roundtrip_i32::<File, File>(file, data)
}

#[test]
Expand All @@ -1245,58 +1267,17 @@ mod tests {
}

fn test_bytes_roundtrip(data: Vec<Vec<i32>>) {
let mut buffer = vec![];

let schema = Arc::new(
types::Type::group_type_builder("schema")
.with_fields(&mut vec![Arc::new(
types::Type::primitive_type_builder("col1", Type::INT32)
.with_repetition(Repetition::REQUIRED)
.build()
.unwrap(),
)])
.build()
.unwrap(),
);

let mut rows: i64 = 0;
{
let props = Arc::new(WriterProperties::builder().build());
let mut writer =
SerializedFileWriter::new(&mut buffer, schema, props).unwrap();

for subset in &data {
let mut row_group_writer = writer.next_row_group().unwrap();
if let Some(mut writer) = row_group_writer.next_column().unwrap() {
rows += writer
.typed::<Int32Type>()
.write_batch(&subset[..], None, None)
.unwrap() as i64;

writer.close().unwrap();
}
row_group_writer.close().unwrap();
}
writer.close().unwrap();
}

let reading_cursor = Bytes::from(buffer);
let reader = SerializedFileReader::new(reading_cursor).unwrap();
test_roundtrip_i32::<Vec<u8>, Bytes>(Vec::with_capacity(1024), data);
}

assert_eq!(reader.num_row_groups(), data.len());
assert_eq!(
reader.metadata().file_metadata().num_rows(),
rows,
"row count in metadata not equal to number of rows written"
#[test]
fn test_boolean_roundtrip() {
let my_bool_values: Vec<_> = (0..2049).map(|idx| idx % 2 == 0).collect();
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I double checked that 2049 is the same size as parquet/tests/boolean_writer.rs

test_roundtrip::<Vec<u8>, Bytes, BoolType, _>(
Vec::with_capacity(1024),
vec![my_bool_values],
|r| r.get_bool(0).unwrap(),
);
for (i, item) in data.iter().enumerate().take(reader.num_row_groups()) {
let row_group_reader = reader.get_row_group(i).unwrap();
let iter = row_group_reader.get_row_iter(None).unwrap();
let res = iter
.map(|elem| elem.get_int(0).unwrap())
.collect::<Vec<i32>>();
assert_eq!(res, *item);
}
}

#[test]
Expand Down
89 changes: 0 additions & 89 deletions parquet/tests/boolean_writer.rs
View file
Open in desktop

This file was deleted.