Skip to content

Commit

Permalink
Add tests for nested decimal arrays (#2254)
Browse files Browse the repository at this point in the history
  • Loading branch information
tustvold committed Aug 1, 2022
1 parent b879977 commit 42b15a8
Showing 1 changed file with 82 additions and 0 deletions.
82 changes: 82 additions & 0 deletions parquet/src/arrow/arrow_reader.rs
Expand Up @@ -397,6 +397,7 @@ mod tests {
use tempfile::tempfile;

use arrow::array::*;
use arrow::buffer::Buffer;
use arrow::datatypes::{DataType as ArrowDataType, Field, Schema};
use arrow::error::Result as ArrowResult;
use arrow::record_batch::{RecordBatch, RecordBatchReader};
Expand Down Expand Up @@ -729,6 +730,87 @@ mod tests {
}
}

#[test]
fn test_decimal_nullable_struct() {
let decimals = Decimal128Array::from_iter_values([1, 2, 3, 4, 5, 6, 7, 8]);

let data = ArrayDataBuilder::new(ArrowDataType::Struct(vec![Field::new(
"decimals",
decimals.data_type().clone(),
false,
)]))
.len(8)
.null_bit_buffer(Some(Buffer::from(&[0b11101111])))
.child_data(vec![decimals.into_data()])
.build()
.unwrap();

let written = RecordBatch::try_from_iter([(
"struct",
Arc::new(StructArray::from(data)) as ArrayRef,
)])
.unwrap();

let mut buffer = Vec::with_capacity(1024);
let mut writer =
ArrowWriter::try_new(&mut buffer, written.schema(), None).unwrap();
writer.write(&written).unwrap();
writer.close().unwrap();

let read = ParquetFileArrowReader::try_new(Bytes::from(buffer))
.unwrap()
.get_record_reader(3)
.unwrap()
.collect::<ArrowResult<Vec<_>>>()
.unwrap();

assert_eq!(&written.slice(0, 3), &read[0]);
assert_eq!(&written.slice(3, 3), &read[1]);
assert_eq!(&written.slice(6, 2), &read[2]);
}

#[test]
#[ignore] // https://github.com/apache/arrow-rs/issues/2253
fn test_decimal_list() {
let decimals = Decimal128Array::from_iter_values([1, 2, 3, 4, 5, 6, 7, 8]);

// [[], [1], [2, 3], null, [4], null, [6, 7, 8]]
let data = ArrayDataBuilder::new(ArrowDataType::List(Box::new(Field::new(
"item",
decimals.data_type().clone(),
false,
))))
.len(7)
.add_buffer(Buffer::from_iter([0_i32, 0, 1, 3, 3, 4, 5, 8]))
.null_bit_buffer(Some(Buffer::from(&[0b01010111])))
.child_data(vec![decimals.into_data()])
.build()
.unwrap();

let written = RecordBatch::try_from_iter([(
"list",
Arc::new(ListArray::from(data)) as ArrayRef,
)])
.unwrap();

let mut buffer = Vec::with_capacity(1024);
let mut writer =
ArrowWriter::try_new(&mut buffer, written.schema(), None).unwrap();
writer.write(&written).unwrap();
writer.close().unwrap();

let read = ParquetFileArrowReader::try_new(Bytes::from(buffer))
.unwrap()
.get_record_reader(3)
.unwrap()
.collect::<ArrowResult<Vec<_>>>()
.unwrap();

assert_eq!(&written.slice(0, 3), &read[0]);
assert_eq!(&written.slice(3, 3), &read[1]);
assert_eq!(&written.slice(6, 1), &read[2]);
}

#[test]
fn test_read_decimal_file() {
use arrow::array::Decimal128Array;
Expand Down

0 comments on commit 42b15a8

Please sign in to comment.