ARROW-9790: [Rust][Parquet] Fix PrimitiveArrayReader boundary conditions

rust/parquet/src/arrow/array_reader.rs

@@ -136,10 +136,8 @@ impl<T: DataType> ArrayReader for PrimitiveArrayReader<T> {
  while records_read < batch_size {
  let records_to_read = batch_size - records_read;
 
+ // NB can be 0 if at end of page
  let records_read_once = self.record_reader.read_records(records_to_read)?;
- if records_read_once == 0 {
- break; // record reader has no record
- }
  records_read = records_read + records_read_once;
 
  // Record reader exhausted
@@ -934,7 +932,7 @@ mod tests {
  use super::*;
  use crate::arrow::converter::Utf8Converter;
  use crate::basic::{Encoding, Type as PhysicalType};
- use crate::column::page::Page;
+ use crate::column::page::{Page, PageReader};
  use crate::data_type::{ByteArray, DataType, Int32Type, Int64Type};
  use crate::errors::Result;
  use crate::file::reader::{FileReader, SerializedFileReader};
@@ -1000,6 +998,33 @@ mod tests {
  }
  }
 
+ #[test]
+ fn test_primitive_array_reader_empty_pages() {
+ // Construct column schema
+ let message_type = "
+ message test_schema {
+ REQUIRED INT32 leaf;
+ }
+ ";
+
+ let schema = parse_message_type(message_type)
+ .map(|t| Rc::new(SchemaDescriptor::new(Rc::new(t))))
+ .unwrap();
+
+ let column_desc = schema.column(0);
+ let page_iterator = EmptyPageIterator::new(schema.clone());
+
+ let mut array_reader = PrimitiveArrayReader::<Int32Type>::new(
+ Box::new(page_iterator),
+ column_desc.clone(),
+ )
+ .unwrap();
+
+ // expect no values to be read
+ let array = array_reader.next_batch(50).unwrap();
+ assert!(array.is_empty());
+ }
+
  #[test]
  fn test_primitive_array_reader_data() {
  // Construct column schema
@@ -1452,6 +1477,35 @@ mod tests {
  }
  }
 
+ /// Iterator for testing reading empty columns
+ struct EmptyPageIterator {
+ schema: SchemaDescPtr,
+ }
+
+ impl EmptyPageIterator {
+ fn new(schema: SchemaDescPtr) -> Self {
+ EmptyPageIterator { schema }
+ }
+ }
+
+ impl Iterator for EmptyPageIterator {
+ type Item = Result<Box<dyn PageReader>>;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ None
+ }
+ }
+
+ impl PageIterator for EmptyPageIterator {
+ fn schema(&mut self) -> Result<SchemaDescPtr> {
+ Ok(self.schema.clone())
+ }
+
+ fn column_schema(&mut self) -> Result<ColumnDescPtr> {
+ Ok(self.schema.column(0))
+ }
+ }
+
  #[test]
  fn test_struct_array_reader() {
  let array_1 = Arc::new(PrimitiveArray::<ArrowInt32>::from(vec![1, 2, 3, 4, 5]));

rust/parquet/src/arrow/arrow_reader.rs

@@ -304,6 +304,45 @@ mod tests {
  >(2, 100, 2, message_type, 15, 50, converter);
  }
 
+ #[test]
+ fn test_bool_single_column_reader_test_batch_size_divides_into_row_group_size() {
+ let message_type = "
+ message test_schema {
+ REQUIRED BOOLEAN leaf;
+ }
+ ";
+
+ // Use a batch size (5) so batches to fall on
+ // row group boundaries (25 rows in 3 row groups --> row
+ // groups of 10, 10, and 5) to test edge refilling edge cases.
+ let converter = FromConverter::new();
+ single_column_reader_test::<
+ BoolType,
+ BooleanArray,
+ FromConverter<Vec<Option<bool>>, BooleanArray>,
+ BoolType,
+ >(3, 25, 2, message_type, 5, 50, converter);
+ }
+
+ #[test]
+ fn test_bool_single_column_reader_test_batch_size_divides_into_row_group_size2() {
+ let message_type = "
+ message test_schema {
+ REQUIRED BOOLEAN leaf;
+ }
+ ";
+
+ // Ensure that every batch size (25) falls exactly a row group
+ // boundary (25 in this case) to test edge case.
+ let converter = FromConverter::new();
+ single_column_reader_test::<
+ BoolType,
+ BooleanArray,
+ FromConverter<Vec<Option<bool>>, BooleanArray>,
+ BoolType,
+ >(4, 100, 2, message_type, 25, 50, converter);
+ }
+
  struct RandFixedLenGen {}
 
  impl RandGen<FixedLenByteArrayType> for RandFixedLenGen {

rust/parquet/src/arrow/array_reader.rs

@@ -136,10 +136,8 @@ impl<T: DataType> ArrayReader for PrimitiveArrayReader<T> {
  while records_read < batch_size {
  let records_to_read = batch_size - records_read;
 
+ // NB can be 0 if at end of page
  let records_read_once = self.record_reader.read_records(records_to_read)?;
- if records_read_once == 0 {
- break; // record reader has no record
- }
  records_read = records_read + records_read_once;
 
  // Record reader exhausted
@@ -934,7 +932,7 @@ mod tests {
  use super::*;
  use crate::arrow::converter::Utf8Converter;
  use crate::basic::{Encoding, Type as PhysicalType};
- use crate::column::page::Page;
+ use crate::column::page::{Page, PageReader};
  use crate::data_type::{ByteArray, DataType, Int32Type, Int64Type};
  use crate::errors::Result;
  use crate::file::reader::{FileReader, SerializedFileReader};
@@ -1000,6 +998,33 @@ mod tests {
  }
  }
 
+ #[test]
+ fn test_primitive_array_reader_empty_pages() {
+ // Construct column schema
+ let message_type = "
+ message test_schema {
+ REQUIRED INT32 leaf;
+ }
+ ";
+
+ let schema = parse_message_type(message_type)
+ .map(|t| Rc::new(SchemaDescriptor::new(Rc::new(t))))
+ .unwrap();
+
+ let column_desc = schema.column(0);
+ let page_iterator = EmptyPageIterator::new(schema.clone());
+
+ let mut array_reader = PrimitiveArrayReader::<Int32Type>::new(
+ Box::new(page_iterator),
+ column_desc.clone(),
+ )
+ .unwrap();
+
+ // expect no values to be read
+ let array = array_reader.next_batch(50).unwrap();
+ assert!(array.is_empty());
+ }
+
  #[test]
  fn test_primitive_array_reader_data() {
  // Construct column schema
@@ -1452,6 +1477,35 @@ mod tests {
  }
  }
 
+ /// Iterator for testing reading empty columns
+ struct EmptyPageIterator {
+ schema: SchemaDescPtr,
+ }
+
+ impl EmptyPageIterator {
+ fn new(schema: SchemaDescPtr) -> Self {
+ EmptyPageIterator { schema }
+ }
+ }
+
+ impl Iterator for EmptyPageIterator {
+ type Item = Result<Box<dyn PageReader>>;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ None
+ }
+ }
+
+ impl PageIterator for EmptyPageIterator {
+ fn schema(&mut self) -> Result<SchemaDescPtr> {
+ Ok(self.schema.clone())
+ }
+
+ fn column_schema(&mut self) -> Result<ColumnDescPtr> {
+ Ok(self.schema.column(0))
+ }
+ }
+
  #[test]
  fn test_struct_array_reader() {
  let array_1 = Arc::new(PrimitiveArray::<ArrowInt32>::from(vec![1, 2, 3, 4, 5]));

rust/parquet/src/arrow/arrow_reader.rs

@@ -304,6 +304,45 @@ mod tests {
  >(2, 100, 2, message_type, 15, 50, converter);
  }
 
+ #[test]
+ fn test_bool_single_column_reader_test_batch_size_divides_into_row_group_size() {
+ let message_type = "
+ message test_schema {
+ REQUIRED BOOLEAN leaf;
+ }
+ ";
+
+ // Use a batch size (5) so batches to fall on
+ // row group boundaries (25 rows in 3 row groups --> row
+ // groups of 10, 10, and 5) to test edge refilling edge cases.
+ let converter = FromConverter::new();
+ single_column_reader_test::<
+ BoolType,
+ BooleanArray,
+ FromConverter<Vec<Option<bool>>, BooleanArray>,
+ BoolType,
+ >(3, 25, 2, message_type, 5, 50, converter);
+ }
+
+ #[test]
+ fn test_bool_single_column_reader_test_batch_size_divides_into_row_group_size2() {
+ let message_type = "
+ message test_schema {
+ REQUIRED BOOLEAN leaf;
+ }
+ ";
+
+ // Ensure that every batch size (25) falls exactly a row group
+ // boundary (25 in this case) to test edge case.
+ let converter = FromConverter::new();
+ single_column_reader_test::<
+ BoolType,
+ BooleanArray,
+ FromConverter<Vec<Option<bool>>, BooleanArray>,
+ BoolType,
+ >(4, 100, 2, message_type, 25, 50, converter);
+ }
+
  struct RandFixedLenGen {}
 
  impl RandGen<FixedLenByteArrayType> for RandFixedLenGen {

rust/parquet/src/arrow/array_reader.rs

@@ -136,10 +136,8 @@ impl<T: DataType> ArrayReader for PrimitiveArrayReader<T> {
  while records_read < batch_size {
  let records_to_read = batch_size - records_read;
 
+ // NB can be 0 if at end of page
  let records_read_once = self.record_reader.read_records(records_to_read)?;
- if records_read_once == 0 {
- break; // record reader has no record
- }
  records_read = records_read + records_read_once;
 
  // Record reader exhausted
@@ -934,7 +932,7 @@ mod tests {
  use super::*;
  use crate::arrow::converter::Utf8Converter;
  use crate::basic::{Encoding, Type as PhysicalType};
- use crate::column::page::Page;
+ use crate::column::page::{Page, PageReader};
  use crate::data_type::{ByteArray, DataType, Int32Type, Int64Type};
  use crate::errors::Result;
  use crate::file::reader::{FileReader, SerializedFileReader};
@@ -1000,6 +998,33 @@ mod tests {
  }
  }
 
+ #[test]
+ fn test_primitive_array_reader_empty_pages() {
+ // Construct column schema
+ let message_type = "
+ message test_schema {
+ REQUIRED INT32 leaf;
+ }
+ ";
+
+ let schema = parse_message_type(message_type)
+ .map(|t| Rc::new(SchemaDescriptor::new(Rc::new(t))))
+ .unwrap();
+
+ let column_desc = schema.column(0);
+ let page_iterator = EmptyPageIterator::new(schema.clone());
+
+ let mut array_reader = PrimitiveArrayReader::<Int32Type>::new(
+ Box::new(page_iterator),
+ column_desc.clone(),
+ )
+ .unwrap();
+
+ // expect no values to be read
+ let array = array_reader.next_batch(50).unwrap();
+ assert!(array.is_empty());
+ }
+
  #[test]
  fn test_primitive_array_reader_data() {
  // Construct column schema
@@ -1452,6 +1477,35 @@ mod tests {
  }
  }
 
+ /// Iterator for testing reading empty columns
+ struct EmptyPageIterator {
+ schema: SchemaDescPtr,
+ }
+
+ impl EmptyPageIterator {
+ fn new(schema: SchemaDescPtr) -> Self {
+ EmptyPageIterator { schema }
+ }
+ }
+
+ impl Iterator for EmptyPageIterator {
+ type Item = Result<Box<dyn PageReader>>;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ None
+ }
+ }
+
+ impl PageIterator for EmptyPageIterator {
+ fn schema(&mut self) -> Result<SchemaDescPtr> {
+ Ok(self.schema.clone())
+ }
+
+ fn column_schema(&mut self) -> Result<ColumnDescPtr> {
+ Ok(self.schema.column(0))
+ }
+ }
+
  #[test]
  fn test_struct_array_reader() {
  let array_1 = Arc::new(PrimitiveArray::<ArrowInt32>::from(vec![1, 2, 3, 4, 5]));

rust/parquet/src/arrow/arrow_reader.rs

@@ -304,6 +304,45 @@ mod tests {
  >(2, 100, 2, message_type, 15, 50, converter);
  }
 
+ #[test]
+ fn test_bool_single_column_reader_test_batch_size_divides_into_row_group_size() {
+ let message_type = "
+ message test_schema {
+ REQUIRED BOOLEAN leaf;
+ }
+ ";
+
+ // Use a batch size (5) so batches to fall on
+ // row group boundaries (25 rows in 3 row groups --> row
+ // groups of 10, 10, and 5) to test edge refilling edge cases.
+ let converter = FromConverter::new();
+ single_column_reader_test::<
+ BoolType,
+ BooleanArray,
+ FromConverter<Vec<Option<bool>>, BooleanArray>,
+ BoolType,
+ >(3, 25, 2, message_type, 5, 50, converter);
+ }
+
+ #[test]
+ fn test_bool_single_column_reader_test_batch_size_divides_into_row_group_size2() {
+ let message_type = "
+ message test_schema {
+ REQUIRED BOOLEAN leaf;
+ }
+ ";
+
+ // Ensure that every batch size (25) falls exactly a row group
+ // boundary (25 in this case) to test edge case.
+ let converter = FromConverter::new();
+ single_column_reader_test::<
+ BoolType,
+ BooleanArray,
+ FromConverter<Vec<Option<bool>>, BooleanArray>,
+ BoolType,
+ >(4, 100, 2, message_type, 25, 50, converter);
+ }
+
  struct RandFixedLenGen {}
 
  impl RandGen<FixedLenByteArrayType> for RandFixedLenGen {

rust/parquet/src/arrow/array_reader.rs

@@ -136,10 +136,8 @@ impl<T: DataType> ArrayReader for PrimitiveArrayReader<T> {
  while records_read < batch_size {
  let records_to_read = batch_size - records_read;
 
+ // NB can be 0 if at end of page
  let records_read_once = self.record_reader.read_records(records_to_read)?;
- if records_read_once == 0 {
- break; // record reader has no record
- }
  records_read = records_read + records_read_once;
 
  // Record reader exhausted
@@ -934,7 +932,7 @@ mod tests {
  use super::*;
  use crate::arrow::converter::Utf8Converter;
  use crate::basic::{Encoding, Type as PhysicalType};
- use crate::column::page::Page;
+ use crate::column::page::{Page, PageReader};
  use crate::data_type::{ByteArray, DataType, Int32Type, Int64Type};
  use crate::errors::Result;
  use crate::file::reader::{FileReader, SerializedFileReader};
@@ -1000,6 +998,33 @@ mod tests {
  }
  }
 
+ #[test]
+ fn test_primitive_array_reader_empty_pages() {
+ // Construct column schema
+ let message_type = "
+ message test_schema {
+ REQUIRED INT32 leaf;
+ }
+ ";
+
+ let schema = parse_message_type(message_type)
+ .map(|t| Rc::new(SchemaDescriptor::new(Rc::new(t))))
+ .unwrap();
+
+ let column_desc = schema.column(0);
+ let page_iterator = EmptyPageIterator::new(schema.clone());
+
+ let mut array_reader = PrimitiveArrayReader::<Int32Type>::new(
+ Box::new(page_iterator),
+ column_desc.clone(),
+ )
+ .unwrap();
+
+ // expect no values to be read
+ let array = array_reader.next_batch(50).unwrap();
+ assert!(array.is_empty());
+ }
+
  #[test]
  fn test_primitive_array_reader_data() {
  // Construct column schema
@@ -1452,6 +1477,35 @@ mod tests {
  }
  }
 
+ /// Iterator for testing reading empty columns
+ struct EmptyPageIterator {
+ schema: SchemaDescPtr,
+ }
+
+ impl EmptyPageIterator {
+ fn new(schema: SchemaDescPtr) -> Self {
+ EmptyPageIterator { schema }
+ }
+ }
+
+ impl Iterator for EmptyPageIterator {
+ type Item = Result<Box<dyn PageReader>>;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ None
+ }
+ }
+
+ impl PageIterator for EmptyPageIterator {
+ fn schema(&mut self) -> Result<SchemaDescPtr> {
+ Ok(self.schema.clone())
+ }
+
+ fn column_schema(&mut self) -> Result<ColumnDescPtr> {
+ Ok(self.schema.column(0))
+ }
+ }
+
  #[test]
  fn test_struct_array_reader() {
  let array_1 = Arc::new(PrimitiveArray::<ArrowInt32>::from(vec![1, 2, 3, 4, 5]));

rust/parquet/src/arrow/arrow_reader.rs

@@ -304,6 +304,45 @@ mod tests {
  >(2, 100, 2, message_type, 15, 50, converter);
  }
 
+ #[test]
+ fn test_bool_single_column_reader_test_batch_size_divides_into_row_group_size() {
+ let message_type = "
+ message test_schema {
+ REQUIRED BOOLEAN leaf;
+ }
+ ";
+
+ // Use a batch size (5) so batches to fall on
+ // row group boundaries (25 rows in 3 row groups --> row
+ // groups of 10, 10, and 5) to test edge refilling edge cases.
+ let converter = FromConverter::new();
+ single_column_reader_test::<
+ BoolType,
+ BooleanArray,
+ FromConverter<Vec<Option<bool>>, BooleanArray>,
+ BoolType,
+ >(3, 25, 2, message_type, 5, 50, converter);
+ }
+
+ #[test]
+ fn test_bool_single_column_reader_test_batch_size_divides_into_row_group_size2() {
+ let message_type = "
+ message test_schema {
+ REQUIRED BOOLEAN leaf;
+ }
+ ";
+
+ // Ensure that every batch size (25) falls exactly a row group
+ // boundary (25 in this case) to test edge case.
+ let converter = FromConverter::new();
+ single_column_reader_test::<
+ BoolType,
+ BooleanArray,
+ FromConverter<Vec<Option<bool>>, BooleanArray>,
+ BoolType,
+ >(4, 100, 2, message_type, 25, 50, converter);
+ }
+
  struct RandFixedLenGen {}
 
  impl RandGen<FixedLenByteArrayType> for RandFixedLenGen {

rust/parquet/src/arrow/array_reader.rs

@@ -136,10 +136,8 @@ impl<T: DataType> ArrayReader for PrimitiveArrayReader<T> {
  while records_read < batch_size {
  let records_to_read = batch_size - records_read;
 
+ // NB can be 0 if at end of page
  let records_read_once = self.record_reader.read_records(records_to_read)?;
- if records_read_once == 0 {
- break; // record reader has no record
- }
  records_read = records_read + records_read_once;
 
  // Record reader exhausted
@@ -934,7 +932,7 @@ mod tests {
  use super::*;
  use crate::arrow::converter::Utf8Converter;
  use crate::basic::{Encoding, Type as PhysicalType};
- use crate::column::page::Page;
+ use crate::column::page::{Page, PageReader};
  use crate::data_type::{ByteArray, DataType, Int32Type, Int64Type};
  use crate::errors::Result;
  use crate::file::reader::{FileReader, SerializedFileReader};
@@ -1000,6 +998,33 @@ mod tests {
  }
  }
 
+ #[test]
+ fn test_primitive_array_reader_empty_pages() {
+ // Construct column schema
+ let message_type = "
+ message test_schema {
+ REQUIRED INT32 leaf;
+ }
+ ";
+
+ let schema = parse_message_type(message_type)
+ .map(|t| Rc::new(SchemaDescriptor::new(Rc::new(t))))
+ .unwrap();
+
+ let column_desc = schema.column(0);
+ let page_iterator = EmptyPageIterator::new(schema.clone());
+
+ let mut array_reader = PrimitiveArrayReader::<Int32Type>::new(
+ Box::new(page_iterator),
+ column_desc.clone(),
+ )
+ .unwrap();
+
+ // expect no values to be read
+ let array = array_reader.next_batch(50).unwrap();
+ assert!(array.is_empty());
+ }
+
  #[test]
  fn test_primitive_array_reader_data() {
  // Construct column schema
@@ -1452,6 +1477,35 @@ mod tests {
  }
  }
 
+ /// Iterator for testing reading empty columns
+ struct EmptyPageIterator {
+ schema: SchemaDescPtr,
+ }
+
+ impl EmptyPageIterator {
+ fn new(schema: SchemaDescPtr) -> Self {
+ EmptyPageIterator { schema }
+ }
+ }
+
+ impl Iterator for EmptyPageIterator {
+ type Item = Result<Box<dyn PageReader>>;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ None
+ }
+ }
+
+ impl PageIterator for EmptyPageIterator {
+ fn schema(&mut self) -> Result<SchemaDescPtr> {
+ Ok(self.schema.clone())
+ }
+
+ fn column_schema(&mut self) -> Result<ColumnDescPtr> {
+ Ok(self.schema.column(0))
+ }
+ }
+
  #[test]
  fn test_struct_array_reader() {
  let array_1 = Arc::new(PrimitiveArray::<ArrowInt32>::from(vec![1, 2, 3, 4, 5]));

rust/parquet/src/arrow/arrow_reader.rs

@@ -304,6 +304,45 @@ mod tests {
  >(2, 100, 2, message_type, 15, 50, converter);
  }
 
+ #[test]
+ fn test_bool_single_column_reader_test_batch_size_divides_into_row_group_size() {
+ let message_type = "
+ message test_schema {
+ REQUIRED BOOLEAN leaf;
+ }
+ ";
+
+ // Use a batch size (5) so batches to fall on
+ // row group boundaries (25 rows in 3 row groups --> row
+ // groups of 10, 10, and 5) to test edge refilling edge cases.
+ let converter = FromConverter::new();
+ single_column_reader_test::<
+ BoolType,
+ BooleanArray,
+ FromConverter<Vec<Option<bool>>, BooleanArray>,
+ BoolType,
+ >(3, 25, 2, message_type, 5, 50, converter);
+ }
+
+ #[test]
+ fn test_bool_single_column_reader_test_batch_size_divides_into_row_group_size2() {
+ let message_type = "
+ message test_schema {
+ REQUIRED BOOLEAN leaf;
+ }
+ ";
+
+ // Ensure that every batch size (25) falls exactly a row group
+ // boundary (25 in this case) to test edge case.
+ let converter = FromConverter::new();
+ single_column_reader_test::<
+ BoolType,
+ BooleanArray,
+ FromConverter<Vec<Option<bool>>, BooleanArray>,
+ BoolType,
+ >(4, 100, 2, message_type, 25, 50, converter);
+ }
+
  struct RandFixedLenGen {}
 
  impl RandGen<FixedLenByteArrayType> for RandFixedLenGen {

rust/parquet/src/arrow/array_reader.rs

@@ -136,10 +136,8 @@ impl<T: DataType> ArrayReader for PrimitiveArrayReader<T> {
  while records_read < batch_size {
  let records_to_read = batch_size - records_read;
 
+ // NB can be 0 if at end of page
  let records_read_once = self.record_reader.read_records(records_to_read)?;
- if records_read_once == 0 {
- break; // record reader has no record
- }
  records_read = records_read + records_read_once;
 
  // Record reader exhausted
@@ -934,7 +932,7 @@ mod tests {
  use super::*;
  use crate::arrow::converter::Utf8Converter;
  use crate::basic::{Encoding, Type as PhysicalType};
- use crate::column::page::Page;
+ use crate::column::page::{Page, PageReader};
  use crate::data_type::{ByteArray, DataType, Int32Type, Int64Type};
  use crate::errors::Result;
  use crate::file::reader::{FileReader, SerializedFileReader};
@@ -1000,6 +998,33 @@ mod tests {
  }
  }
 
+ #[test]
+ fn test_primitive_array_reader_empty_pages() {
+ // Construct column schema
+ let message_type = "
+ message test_schema {
+ REQUIRED INT32 leaf;
+ }
+ ";
+
+ let schema = parse_message_type(message_type)
+ .map(|t| Rc::new(SchemaDescriptor::new(Rc::new(t))))
+ .unwrap();
+
+ let column_desc = schema.column(0);
+ let page_iterator = EmptyPageIterator::new(schema.clone());
+
+ let mut array_reader = PrimitiveArrayReader::<Int32Type>::new(
+ Box::new(page_iterator),
+ column_desc.clone(),
+ )
+ .unwrap();
+
+ // expect no values to be read
+ let array = array_reader.next_batch(50).unwrap();
+ assert!(array.is_empty());
+ }
+
  #[test]
  fn test_primitive_array_reader_data() {
  // Construct column schema
@@ -1452,6 +1477,35 @@ mod tests {
  }
  }
 
+ /// Iterator for testing reading empty columns
+ struct EmptyPageIterator {
+ schema: SchemaDescPtr,
+ }
+
+ impl EmptyPageIterator {
+ fn new(schema: SchemaDescPtr) -> Self {
+ EmptyPageIterator { schema }
+ }
+ }
+
+ impl Iterator for EmptyPageIterator {
+ type Item = Result<Box<dyn PageReader>>;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ None
+ }
+ }
+
+ impl PageIterator for EmptyPageIterator {
+ fn schema(&mut self) -> Result<SchemaDescPtr> {
+ Ok(self.schema.clone())
+ }
+
+ fn column_schema(&mut self) -> Result<ColumnDescPtr> {
+ Ok(self.schema.column(0))
+ }
+ }
+
  #[test]
  fn test_struct_array_reader() {
  let array_1 = Arc::new(PrimitiveArray::<ArrowInt32>::from(vec![1, 2, 3, 4, 5]));

rust/parquet/src/arrow/arrow_reader.rs

@@ -304,6 +304,45 @@ mod tests {
  >(2, 100, 2, message_type, 15, 50, converter);
  }
 
+ #[test]
+ fn test_bool_single_column_reader_test_batch_size_divides_into_row_group_size() {
+ let message_type = "
+ message test_schema {
+ REQUIRED BOOLEAN leaf;
+ }
+ ";
+
+ // Use a batch size (5) so batches to fall on
+ // row group boundaries (25 rows in 3 row groups --> row
+ // groups of 10, 10, and 5) to test edge refilling edge cases.
+ let converter = FromConverter::new();
+ single_column_reader_test::<
+ BoolType,
+ BooleanArray,
+ FromConverter<Vec<Option<bool>>, BooleanArray>,
+ BoolType,
+ >(3, 25, 2, message_type, 5, 50, converter);
+ }
+
+ #[test]
+ fn test_bool_single_column_reader_test_batch_size_divides_into_row_group_size2() {
+ let message_type = "
+ message test_schema {
+ REQUIRED BOOLEAN leaf;
+ }
+ ";
+
+ // Ensure that every batch size (25) falls exactly a row group
+ // boundary (25 in this case) to test edge case.
+ let converter = FromConverter::new();
+ single_column_reader_test::<
+ BoolType,
+ BooleanArray,
+ FromConverter<Vec<Option<bool>>, BooleanArray>,
+ BoolType,
+ >(4, 100, 2, message_type, 25, 50, converter);
+ }
+
  struct RandFixedLenGen {}
 
  impl RandGen<FixedLenByteArrayType> for RandFixedLenGen {

rust/parquet/src/arrow/array_reader.rs

@@ -136,10 +136,8 @@ impl<T: DataType> ArrayReader for PrimitiveArrayReader<T> {
  while records_read < batch_size {
  let records_to_read = batch_size - records_read;
 
+ // NB can be 0 if at end of page
  let records_read_once = self.record_reader.read_records(records_to_read)?;
- if records_read_once == 0 {
- break; // record reader has no record
- }
  records_read = records_read + records_read_once;
 
  // Record reader exhausted
@@ -934,7 +932,7 @@ mod tests {
  use super::*;
  use crate::arrow::converter::Utf8Converter;
  use crate::basic::{Encoding, Type as PhysicalType};
- use crate::column::page::Page;
+ use crate::column::page::{Page, PageReader};
  use crate::data_type::{ByteArray, DataType, Int32Type, Int64Type};
  use crate::errors::Result;
  use crate::file::reader::{FileReader, SerializedFileReader};
@@ -1000,6 +998,33 @@ mod tests {
  }
  }
 
+ #[test]
+ fn test_primitive_array_reader_empty_pages() {
+ // Construct column schema
+ let message_type = "
+ message test_schema {
+ REQUIRED INT32 leaf;
+ }
+ ";
+
+ let schema = parse_message_type(message_type)
+ .map(|t| Rc::new(SchemaDescriptor::new(Rc::new(t))))
+ .unwrap();
+
+ let column_desc = schema.column(0);
+ let page_iterator = EmptyPageIterator::new(schema.clone());
+
+ let mut array_reader = PrimitiveArrayReader::<Int32Type>::new(
+ Box::new(page_iterator),
+ column_desc.clone(),
+ )
+ .unwrap();
+
+ // expect no values to be read
+ let array = array_reader.next_batch(50).unwrap();
+ assert!(array.is_empty());
+ }
+
  #[test]
  fn test_primitive_array_reader_data() {
  // Construct column schema
@@ -1452,6 +1477,35 @@ mod tests {
  }
  }
 
+ /// Iterator for testing reading empty columns
+ struct EmptyPageIterator {
+ schema: SchemaDescPtr,
+ }
+
+ impl EmptyPageIterator {
+ fn new(schema: SchemaDescPtr) -> Self {
+ EmptyPageIterator { schema }
+ }
+ }
+
+ impl Iterator for EmptyPageIterator {
+ type Item = Result<Box<dyn PageReader>>;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ None
+ }
+ }
+
+ impl PageIterator for EmptyPageIterator {
+ fn schema(&mut self) -> Result<SchemaDescPtr> {
+ Ok(self.schema.clone())
+ }
+
+ fn column_schema(&mut self) -> Result<ColumnDescPtr> {
+ Ok(self.schema.column(0))
+ }
+ }
+
  #[test]
  fn test_struct_array_reader() {
  let array_1 = Arc::new(PrimitiveArray::<ArrowInt32>::from(vec![1, 2, 3, 4, 5]));

rust/parquet/src/arrow/arrow_reader.rs

@@ -304,6 +304,45 @@ mod tests {
  >(2, 100, 2, message_type, 15, 50, converter);
  }
 
+ #[test]
+ fn test_bool_single_column_reader_test_batch_size_divides_into_row_group_size() {
+ let message_type = "
+ message test_schema {
+ REQUIRED BOOLEAN leaf;
+ }
+ ";
+
+ // Use a batch size (5) so batches to fall on
+ // row group boundaries (25 rows in 3 row groups --> row
+ // groups of 10, 10, and 5) to test edge refilling edge cases.
+ let converter = FromConverter::new();
+ single_column_reader_test::<
+ BoolType,
+ BooleanArray,
+ FromConverter<Vec<Option<bool>>, BooleanArray>,
+ BoolType,
+ >(3, 25, 2, message_type, 5, 50, converter);
+ }
+
+ #[test]
+ fn test_bool_single_column_reader_test_batch_size_divides_into_row_group_size2() {
+ let message_type = "
+ message test_schema {
+ REQUIRED BOOLEAN leaf;
+ }
+ ";
+
+ // Ensure that every batch size (25) falls exactly a row group
+ // boundary (25 in this case) to test edge case.
+ let converter = FromConverter::new();
+ single_column_reader_test::<
+ BoolType,
+ BooleanArray,
+ FromConverter<Vec<Option<bool>>, BooleanArray>,
+ BoolType,
+ >(4, 100, 2, message_type, 25, 50, converter);
+ }
+
  struct RandFixedLenGen {}
 
  impl RandGen<FixedLenByteArrayType> for RandFixedLenGen {

rust/parquet/src/arrow/array_reader.rs

@@ -136,10 +136,8 @@ impl<T: DataType> ArrayReader for PrimitiveArrayReader<T> {
  while records_read < batch_size {
  let records_to_read = batch_size - records_read;
 
+ // NB can be 0 if at end of page
  let records_read_once = self.record_reader.read_records(records_to_read)?;
- if records_read_once == 0 {
- break; // record reader has no record
- }
  records_read = records_read + records_read_once;
 
  // Record reader exhausted
@@ -934,7 +932,7 @@ mod tests {
  use super::*;
  use crate::arrow::converter::Utf8Converter;
  use crate::basic::{Encoding, Type as PhysicalType};
- use crate::column::page::Page;
+ use crate::column::page::{Page, PageReader};
  use crate::data_type::{ByteArray, DataType, Int32Type, Int64Type};
  use crate::errors::Result;
  use crate::file::reader::{FileReader, SerializedFileReader};
@@ -1000,6 +998,33 @@ mod tests {
  }
  }
 
+ #[test]
+ fn test_primitive_array_reader_empty_pages() {
+ // Construct column schema
+ let message_type = "
+ message test_schema {
+ REQUIRED INT32 leaf;
+ }
+ ";
+
+ let schema = parse_message_type(message_type)
+ .map(|t| Rc::new(SchemaDescriptor::new(Rc::new(t))))
+ .unwrap();
+
+ let column_desc = schema.column(0);
+ let page_iterator = EmptyPageIterator::new(schema.clone());
+
+ let mut array_reader = PrimitiveArrayReader::<Int32Type>::new(
+ Box::new(page_iterator),
+ column_desc.clone(),
+ )
+ .unwrap();
+
+ // expect no values to be read
+ let array = array_reader.next_batch(50).unwrap();
+ assert!(array.is_empty());
+ }
+
  #[test]
  fn test_primitive_array_reader_data() {
  // Construct column schema
@@ -1452,6 +1477,35 @@ mod tests {
  }
  }
 
+ /// Iterator for testing reading empty columns
+ struct EmptyPageIterator {
+ schema: SchemaDescPtr,
+ }
+
+ impl EmptyPageIterator {
+ fn new(schema: SchemaDescPtr) -> Self {
+ EmptyPageIterator { schema }
+ }
+ }
+
+ impl Iterator for EmptyPageIterator {
+ type Item = Result<Box<dyn PageReader>>;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ None
+ }
+ }
+
+ impl PageIterator for EmptyPageIterator {
+ fn schema(&mut self) -> Result<SchemaDescPtr> {
+ Ok(self.schema.clone())
+ }
+
+ fn column_schema(&mut self) -> Result<ColumnDescPtr> {
+ Ok(self.schema.column(0))
+ }
+ }
+
  #[test]
  fn test_struct_array_reader() {
  let array_1 = Arc::new(PrimitiveArray::<ArrowInt32>::from(vec![1, 2, 3, 4, 5]));

rust/parquet/src/arrow/arrow_reader.rs

@@ -304,6 +304,45 @@ mod tests {
  >(2, 100, 2, message_type, 15, 50, converter);
  }
 
+ #[test]
+ fn test_bool_single_column_reader_test_batch_size_divides_into_row_group_size() {
+ let message_type = "
+ message test_schema {
+ REQUIRED BOOLEAN leaf;
+ }
+ ";
+
+ // Use a batch size (5) so batches to fall on
+ // row group boundaries (25 rows in 3 row groups --> row
+ // groups of 10, 10, and 5) to test edge refilling edge cases.
+ let converter = FromConverter::new();
+ single_column_reader_test::<
+ BoolType,
+ BooleanArray,
+ FromConverter<Vec<Option<bool>>, BooleanArray>,
+ BoolType,
+ >(3, 25, 2, message_type, 5, 50, converter);
+ }
+
+ #[test]
+ fn test_bool_single_column_reader_test_batch_size_divides_into_row_group_size2() {
+ let message_type = "
+ message test_schema {
+ REQUIRED BOOLEAN leaf;
+ }
+ ";
+
+ // Ensure that every batch size (25) falls exactly a row group
+ // boundary (25 in this case) to test edge case.
+ let converter = FromConverter::new();
+ single_column_reader_test::<
+ BoolType,
+ BooleanArray,
+ FromConverter<Vec<Option<bool>>, BooleanArray>,
+ BoolType,
+ >(4, 100, 2, message_type, 25, 50, converter);
+ }
+
  struct RandFixedLenGen {}
 
  impl RandGen<FixedLenByteArrayType> for RandFixedLenGen {