use arrow_array::cast::AsArray;
use arrow_array::{BooleanArray, RecordBatch, RecordBatchReader};
use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder};
use arrow_schema::{ArrowError, DataType as ArrowType, FieldRef, Schema, SchemaRef};
use arrow_select::filter::filter_record_batch;
pub use filter::{ArrowPredicate, ArrowPredicateFn, RowFilter};
use selection::MaskCursor;
pub use selection::{
MaskRunIter, RowSelection, RowSelectionCursor, RowSelectionPolicy, RowSelector,
};
use std::fmt::{Debug, Formatter};
use std::sync::Arc;
pub use crate::arrow::array_reader::RowGroups;
use crate::arrow::array_reader::{ArrayReader, ArrayReaderBuilder};
use crate::arrow::schema::{
ParquetField, parquet_to_arrow_schema_and_fields, virtual_type::is_virtual_column,
};
use crate::arrow::{FieldLevels, ProjectionMask, parquet_to_arrow_field_levels_with_virtual};
use crate::basic::{BloomFilterAlgorithm, BloomFilterCompression, BloomFilterHash};
use crate::bloom_filter::{
SBBF_HEADER_SIZE_ESTIMATE, Sbbf, chunk_read_bloom_filter_header_and_offset,
};
use crate::column::page::{PageIterator, PageReader};
#[cfg(feature = "encryption")]
use crate::encryption::decrypt::FileDecryptionProperties;
use crate::errors::{ParquetError, Result};
use crate::file::metadata::{
ColumnChunkMask, PageIndexPolicy, ParquetMetaData, ParquetMetaDataOptions,
ParquetMetaDataReader, ParquetStatisticsPolicy, RowGroupMetaData,
};
use crate::file::reader::{ChunkReader, SerializedPageReader};
use crate::schema::types::SchemaDescriptor;
use crate::arrow::arrow_reader::metrics::ArrowReaderMetrics;
pub use read_plan::{PredicateOptions, ReadPlan, ReadPlanBuilder};
mod filter;
pub mod metrics;
mod read_plan;
pub(crate) mod selection;
pub mod statistics;
pub const DEFAULT_BATCH_SIZE: usize = 1024;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RowGroupSelection {
pub(crate) row_group_index: usize,
pub(crate) selection: Option<RowSelection>,
}
impl RowGroupSelection {
pub fn new(row_group_index: usize, selection: Option<RowSelection>) -> Self {
Self {
row_group_index,
selection,
}
}
pub fn row_group_index(&self) -> usize {
self.row_group_index
}
pub fn selection(&self) -> Option<&RowSelection> {
self.selection.as_ref()
}
}
#[derive(Debug)]
pub(crate) enum RowGroupPlan {
Global {
row_groups: Option<Vec<usize>>,
selection: Option<RowSelection>,
},
PerRowGroup(Vec<RowGroupSelection>),
Conflicting,
}
impl RowGroupPlan {
fn set_row_groups(&mut self, new_row_groups: Vec<usize>) {
match self {
Self::Global { row_groups, .. } => *row_groups = Some(new_row_groups),
Self::PerRowGroup(_) => *self = Self::Conflicting,
Self::Conflicting => {}
}
}
fn set_row_selection(&mut self, new_selection: RowSelection) {
match self {
Self::Global { selection, .. } => *selection = Some(new_selection),
Self::PerRowGroup(_) => *self = Self::Conflicting,
Self::Conflicting => {}
}
}
pub(crate) fn set_row_group_selections(
&mut self,
row_group_selections: Vec<RowGroupSelection>,
) {
match self {
Self::Global {
row_groups: None,
selection: None,
}
| Self::PerRowGroup(_) => {
*self = Self::PerRowGroup(row_group_selections);
}
Self::Global { .. } => *self = Self::Conflicting,
Self::Conflicting => {}
}
}
pub(crate) fn conflict_error() -> ParquetError {
ParquetError::General(
"with_row_group_selections cannot be combined with with_row_groups or with_row_selection"
.to_string(),
)
}
fn into_global(self) -> Result<(Option<Vec<usize>>, Option<RowSelection>)> {
match self {
Self::Global {
row_groups,
selection,
} => Ok((row_groups, selection)),
Self::PerRowGroup(_) => Err(ParquetError::General(
"Row-group-local selections are not supported by the synchronous reader"
.to_string(),
)),
Self::Conflicting => Err(Self::conflict_error()),
}
}
}
pub struct ArrowReaderBuilder<T> {
pub(crate) input: T,
pub(crate) metadata: Arc<ParquetMetaData>,
pub(crate) schema: SchemaRef,
pub(crate) fields: Option<Arc<ParquetField>>,
pub(crate) batch_size: usize,
pub(crate) row_group_plan: RowGroupPlan,
pub(crate) projection: ProjectionMask,
pub(crate) filter: Option<RowFilter>,
pub(crate) row_selection_policy: RowSelectionPolicy,
pub(crate) limit: Option<usize>,
pub(crate) offset: Option<usize>,
pub(crate) metrics: ArrowReaderMetrics,
pub(crate) max_predicate_cache_size: usize,
}
impl<T: Debug> Debug for ArrowReaderBuilder<T> {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ArrowReaderBuilder<T>")
.field("input", &self.input)
.field("metadata", &self.metadata)
.field("schema", &self.schema)
.field("fields", &self.fields)
.field("batch_size", &self.batch_size)
.field("row_group_plan", &self.row_group_plan)
.field("projection", &self.projection)
.field("filter", &self.filter)
.field("row_selection_policy", &self.row_selection_policy)
.field("limit", &self.limit)
.field("offset", &self.offset)
.field("metrics", &self.metrics)
.finish()
}
}
impl<T> ArrowReaderBuilder<T> {
pub(crate) fn new_builder(input: T, metadata: ArrowReaderMetadata) -> Self {
Self {
input,
metadata: metadata.metadata,
schema: metadata.schema,
fields: metadata.fields,
batch_size: DEFAULT_BATCH_SIZE,
row_group_plan: RowGroupPlan::Global {
row_groups: None,
selection: None,
},
projection: ProjectionMask::all(),
filter: None,
row_selection_policy: RowSelectionPolicy::default(),
limit: None,
offset: None,
metrics: ArrowReaderMetrics::Disabled,
max_predicate_cache_size: 100 * 1024 * 1024, }
}
pub fn metadata(&self) -> &Arc<ParquetMetaData> {
&self.metadata
}
pub fn parquet_schema(&self) -> &SchemaDescriptor {
self.metadata.file_metadata().schema_descr()
}
pub fn schema(&self) -> &SchemaRef {
&self.schema
}
pub fn with_batch_size(self, batch_size: usize) -> Self {
let batch_size = batch_size.min(self.metadata.file_metadata().num_rows() as usize);
Self { batch_size, ..self }
}
pub fn with_row_groups(mut self, row_groups: Vec<usize>) -> Self {
self.row_group_plan.set_row_groups(row_groups);
self
}
pub fn with_projection(self, mask: ProjectionMask) -> Self {
Self {
projection: mask,
..self
}
}
pub fn with_row_selection_policy(self, policy: RowSelectionPolicy) -> Self {
Self {
row_selection_policy: policy,
..self
}
}
pub fn with_row_selection(mut self, selection: RowSelection) -> Self {
self.row_group_plan.set_row_selection(selection);
self
}
pub fn with_row_filter(self, filter: RowFilter) -> Self {
Self {
filter: Some(filter),
..self
}
}
pub fn with_limit(self, limit: usize) -> Self {
Self {
limit: Some(limit),
..self
}
}
pub fn with_offset(self, offset: usize) -> Self {
Self {
offset: Some(offset),
..self
}
}
pub fn with_metrics(self, metrics: ArrowReaderMetrics) -> Self {
Self { metrics, ..self }
}
pub fn with_max_predicate_cache_size(self, max_predicate_cache_size: usize) -> Self {
Self {
max_predicate_cache_size,
..self
}
}
}
#[derive(Debug, Clone, Default)]
pub struct ArrowReaderOptions {
skip_arrow_metadata: bool,
supplied_schema: Option<SchemaRef>,
column_index: PageIndexPolicy,
offset_index: PageIndexPolicy,
column_index_mask: ColumnChunkMask,
offset_index_mask: ColumnChunkMask,
metadata_options: ParquetMetaDataOptions,
#[cfg(feature = "encryption")]
pub(crate) file_decryption_properties: Option<Arc<FileDecryptionProperties>>,
virtual_columns: Vec<FieldRef>,
}
impl ArrowReaderOptions {
pub fn new() -> Self {
Self::default()
}
pub fn with_skip_arrow_metadata(self, skip_arrow_metadata: bool) -> Self {
Self {
skip_arrow_metadata,
..self
}
}
pub fn with_schema(self, schema: SchemaRef) -> Self {
Self {
supplied_schema: Some(schema),
skip_arrow_metadata: true,
..self
}
}
pub fn with_page_index_policy(self, policy: PageIndexPolicy) -> Self {
self.with_column_index_policy(policy)
.with_offset_index_policy(policy)
}
pub fn with_column_index_policy(mut self, policy: PageIndexPolicy) -> Self {
self.column_index = policy;
self
}
pub fn with_offset_index_policy(mut self, policy: PageIndexPolicy) -> Self {
self.offset_index = policy;
self
}
pub fn with_page_index_mask(self, mask: ColumnChunkMask) -> Self {
self.with_column_index_mask(mask.clone())
.with_offset_index_mask(mask)
}
pub fn with_column_index_mask(mut self, mask: ColumnChunkMask) -> Self {
self.column_index_mask = mask;
self
}
pub fn with_offset_index_mask(mut self, mask: ColumnChunkMask) -> Self {
self.offset_index_mask = mask;
self
}
pub fn with_parquet_schema(mut self, schema: Arc<SchemaDescriptor>) -> Self {
self.metadata_options.set_schema(schema);
self
}
pub fn with_encoding_stats_as_mask(mut self, val: bool) -> Self {
self.metadata_options.set_encoding_stats_as_mask(val);
self
}
pub fn with_encoding_stats_policy(mut self, policy: ParquetStatisticsPolicy) -> Self {
self.metadata_options.set_encoding_stats_policy(policy);
self
}
pub fn with_column_stats_policy(mut self, policy: ParquetStatisticsPolicy) -> Self {
self.metadata_options.set_column_stats_policy(policy);
self
}
pub fn with_size_stats_policy(mut self, policy: ParquetStatisticsPolicy) -> Self {
self.metadata_options.set_size_stats_policy(policy);
self
}
#[cfg(feature = "encryption")]
pub fn with_file_decryption_properties(
self,
file_decryption_properties: Arc<FileDecryptionProperties>,
) -> Self {
Self {
file_decryption_properties: Some(file_decryption_properties),
..self
}
}
pub fn with_virtual_columns(self, virtual_columns: Vec<FieldRef>) -> Result<Self> {
for field in &virtual_columns {
if !is_virtual_column(field) {
return Err(ParquetError::General(format!(
"Field '{}' is not a virtual column. Virtual columns must have extension type names starting with 'arrow.virtual.'",
field.name()
)));
}
}
Ok(Self {
virtual_columns,
..self
})
}
pub fn offset_index_policy(&self) -> PageIndexPolicy {
self.offset_index
}
pub fn column_index_policy(&self) -> PageIndexPolicy {
self.column_index
}
pub fn offset_index_mask(&self) -> &ColumnChunkMask {
&self.offset_index_mask
}
pub fn column_index_mask(&self) -> &ColumnChunkMask {
&self.column_index_mask
}
pub fn metadata_options(&self) -> &ParquetMetaDataOptions {
&self.metadata_options
}
#[cfg(feature = "encryption")]
pub fn file_decryption_properties(&self) -> Option<&Arc<FileDecryptionProperties>> {
self.file_decryption_properties.as_ref()
}
}
impl ParquetMetaDataReader {
pub fn with_arrow_reader_options(mut self, options: Option<&ArrowReaderOptions>) -> Self {
let Some(options) = options else { return self };
self = self.with_metadata_options(Some(options.metadata_options().clone()));
#[cfg(feature = "encryption")]
{
self = self.with_decryption_properties(
options.file_decryption_properties.as_ref().map(Arc::clone),
);
}
if options.column_index_policy() != PageIndexPolicy::Skip
|| options.offset_index_policy() != PageIndexPolicy::Skip
{
self = self
.with_column_index_policy(options.column_index_policy())
.with_offset_index_policy(options.offset_index_policy());
}
if !options.column_index_mask().is_all() {
self = self.with_column_index_mask(options.column_index_mask().clone());
}
if !options.offset_index_mask().is_all() {
self = self.with_offset_index_mask(options.offset_index_mask().clone());
}
self
}
}
#[derive(Debug, Clone)]
pub struct ArrowReaderMetadata {
pub(crate) metadata: Arc<ParquetMetaData>,
pub(crate) schema: SchemaRef,
pub(crate) fields: Option<Arc<ParquetField>>,
}
impl ArrowReaderMetadata {
pub fn load<T: ChunkReader>(reader: &T, options: ArrowReaderOptions) -> Result<Self> {
let metadata = ParquetMetaDataReader::new()
.with_column_index_policy(options.column_index_policy())
.with_offset_index_policy(options.offset_index_policy())
.with_column_index_mask(options.column_index_mask().clone())
.with_offset_index_mask(options.offset_index_mask().clone())
.with_metadata_options(Some(options.metadata_options.clone()));
#[cfg(feature = "encryption")]
let metadata = metadata.with_decryption_properties(
options.file_decryption_properties.as_ref().map(Arc::clone),
);
let metadata = metadata.parse_and_finish(reader)?;
Self::try_new(Arc::new(metadata), options)
}
pub fn try_new(metadata: Arc<ParquetMetaData>, options: ArrowReaderOptions) -> Result<Self> {
match options.supplied_schema {
Some(supplied_schema) => Self::with_supplied_schema(
metadata,
supplied_schema.clone(),
&options.virtual_columns,
),
None => {
let kv_metadata = match options.skip_arrow_metadata {
true => None,
false => metadata.file_metadata().key_value_metadata(),
};
let (schema, fields) = parquet_to_arrow_schema_and_fields(
metadata.file_metadata().schema_descr(),
ProjectionMask::all(),
kv_metadata,
&options.virtual_columns,
)?;
Ok(Self {
metadata,
schema: Arc::new(schema),
fields: fields.map(Arc::new),
})
}
}
}
fn with_supplied_schema(
metadata: Arc<ParquetMetaData>,
supplied_schema: SchemaRef,
virtual_columns: &[FieldRef],
) -> Result<Self> {
let parquet_schema = metadata.file_metadata().schema_descr();
let field_levels = parquet_to_arrow_field_levels_with_virtual(
parquet_schema,
ProjectionMask::all(),
Some(supplied_schema.fields()),
virtual_columns,
)?;
let fields = field_levels.fields;
let inferred_len = fields.len();
let supplied_len = supplied_schema.fields().len() + virtual_columns.len();
if inferred_len != supplied_len {
return Err(arrow_err!(format!(
"Incompatible supplied Arrow schema: expected {} columns received {}",
inferred_len, supplied_len
)));
}
let mut errors = Vec::new();
let field_iter = supplied_schema.fields().iter().zip(fields.iter());
for (field1, field2) in field_iter {
if field1.data_type() != field2.data_type() {
errors.push(format!(
"data type mismatch for field {}: requested {} but found {}",
field1.name(),
field1.data_type(),
field2.data_type()
));
}
if field1.is_nullable() != field2.is_nullable() {
errors.push(format!(
"nullability mismatch for field {}: expected {:?} but found {:?}",
field1.name(),
field1.is_nullable(),
field2.is_nullable()
));
}
if field1.metadata() != field2.metadata() {
errors.push(format!(
"metadata mismatch for field {}: expected {:?} but found {:?}",
field1.name(),
field1.metadata(),
field2.metadata()
));
}
}
if !errors.is_empty() {
let message = errors.join(", ");
return Err(ParquetError::ArrowError(format!(
"Incompatible supplied Arrow schema: {message}",
)));
}
let schema = if virtual_columns.is_empty() {
supplied_schema
} else {
Arc::new(Schema::new_with_metadata(
fields,
supplied_schema.metadata().clone(),
))
};
Ok(Self {
metadata,
schema,
fields: field_levels.levels.map(Arc::new),
})
}
pub fn metadata(&self) -> &Arc<ParquetMetaData> {
&self.metadata
}
pub fn parquet_schema(&self) -> &SchemaDescriptor {
self.metadata.file_metadata().schema_descr()
}
pub fn schema(&self) -> &SchemaRef {
&self.schema
}
}
#[doc(hidden)]
pub struct SyncReader<T: ChunkReader>(T);
impl<T: Debug + ChunkReader> Debug for SyncReader<T> {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_tuple("SyncReader").field(&self.0).finish()
}
}
pub type ParquetRecordBatchReaderBuilder<T> = ArrowReaderBuilder<SyncReader<T>>;
impl<T: ChunkReader + 'static> ParquetRecordBatchReaderBuilder<T> {
pub fn try_new(reader: T) -> Result<Self> {
Self::try_new_with_options(reader, Default::default())
}
pub fn try_new_with_options(reader: T, options: ArrowReaderOptions) -> Result<Self> {
let metadata = ArrowReaderMetadata::load(&reader, options)?;
Ok(Self::new_with_metadata(reader, metadata))
}
pub fn new_with_metadata(input: T, metadata: ArrowReaderMetadata) -> Self {
Self::new_builder(SyncReader(input), metadata)
}
pub fn get_row_group_column_bloom_filter(
&self,
row_group_idx: usize,
column_idx: usize,
) -> Result<Option<Sbbf>> {
let metadata = self.metadata.row_group(row_group_idx);
let column_metadata = metadata.column(column_idx);
let offset: u64 = if let Some(offset) = column_metadata.bloom_filter_offset() {
offset
.try_into()
.map_err(|_| ParquetError::General("Bloom filter offset is invalid".to_string()))?
} else {
return Ok(None);
};
let buffer = match column_metadata.bloom_filter_length() {
Some(length) => self.input.0.get_bytes(offset, length as usize),
None => self.input.0.get_bytes(offset, SBBF_HEADER_SIZE_ESTIMATE),
}?;
let (header, bitset_offset) =
chunk_read_bloom_filter_header_and_offset(offset, buffer.clone())?;
match header.algorithm {
BloomFilterAlgorithm::BLOCK => {
}
}
match header.compression {
BloomFilterCompression::UNCOMPRESSED => {
}
}
match header.hash {
BloomFilterHash::XXHASH => {
}
}
let bitset = match column_metadata.bloom_filter_length() {
Some(_) => {
let bitset_start = bitset_offset
.checked_sub(offset)
.and_then(|start| usize::try_from(start).ok())
.ok_or_else(|| {
ParquetError::General("Bloom filter offset is invalid".to_string())
})?;
buffer.slice(bitset_start..)
}
None => {
let bitset_length: usize = header.num_bytes.try_into().map_err(|_| {
ParquetError::General("Bloom filter length is invalid".to_string())
})?;
self.input.0.get_bytes(bitset_offset, bitset_length)?
}
};
Ok(Some(Sbbf::new(&bitset)))
}
pub fn build(self) -> Result<ParquetRecordBatchReader> {
let Self {
input,
metadata,
schema: _,
fields,
batch_size,
row_group_plan,
projection,
mut filter,
row_selection_policy,
limit,
offset,
metrics,
max_predicate_cache_size: _,
} = self;
let batch_size = batch_size.min(metadata.file_metadata().num_rows() as usize);
let (row_groups, selection) = row_group_plan.into_global()?;
let row_groups = row_groups.unwrap_or_else(|| (0..metadata.num_row_groups()).collect());
let reader = ReaderRowGroups {
reader: Arc::new(input.0),
metadata,
row_groups,
};
let mut plan_builder = ReadPlanBuilder::new(batch_size)
.with_selection(selection)
.with_row_selection_policy(row_selection_policy);
if let Some(filter) = filter.as_mut() {
for predicate in &mut filter.predicates {
if !plan_builder.selects_any() {
break;
}
let mut cache_projection = predicate.projection().clone();
cache_projection.intersect(&projection);
let array_reader = ArrayReaderBuilder::new(&reader, &metrics)
.with_batch_size(batch_size)
.with_parquet_metadata(&reader.metadata)
.build_array_reader(fields.as_deref(), predicate.projection())?;
plan_builder = plan_builder.with_predicate(array_reader, predicate.as_mut())?;
}
}
let array_reader = ArrayReaderBuilder::new(&reader, &metrics)
.with_batch_size(batch_size)
.with_parquet_metadata(&reader.metadata)
.build_array_reader(fields.as_deref(), &projection)?;
let read_plan = plan_builder
.limited(reader.num_rows())
.with_offset(offset)
.with_limit(limit)
.build_limited()
.build();
Ok(ParquetRecordBatchReader::new(array_reader, read_plan))
}
}
struct ReaderRowGroups<T: ChunkReader> {
reader: Arc<T>,
metadata: Arc<ParquetMetaData>,
row_groups: Vec<usize>,
}
impl<T: ChunkReader + 'static> RowGroups for ReaderRowGroups<T> {
fn num_rows(&self) -> usize {
let meta = self.metadata.row_groups();
self.row_groups
.iter()
.map(|x| meta[*x].num_rows() as usize)
.sum()
}
fn column_chunks(&self, i: usize) -> Result<Box<dyn PageIterator>> {
Ok(Box::new(ReaderPageIterator {
column_idx: i,
reader: self.reader.clone(),
metadata: self.metadata.clone(),
row_groups: self.row_groups.clone().into_iter(),
}))
}
fn row_groups(&self) -> Box<dyn Iterator<Item = &RowGroupMetaData> + '_> {
Box::new(
self.row_groups
.iter()
.map(move |i| self.metadata.row_group(*i)),
)
}
fn metadata(&self) -> &ParquetMetaData {
self.metadata.as_ref()
}
}
struct ReaderPageIterator<T: ChunkReader> {
reader: Arc<T>,
column_idx: usize,
row_groups: std::vec::IntoIter<usize>,
metadata: Arc<ParquetMetaData>,
}
impl<T: ChunkReader + 'static> ReaderPageIterator<T> {
fn next_page_reader(&self, rg_idx: usize) -> Result<SerializedPageReader<T>> {
let rg = self.metadata.row_group(rg_idx);
let column_chunk_metadata = rg.column(self.column_idx);
let page_locations = self
.metadata
.page_index()
.map(|i| i.page_locations(rg_idx, self.column_idx).cloned())
.unwrap_or(None);
let total_rows = rg.num_rows() as usize;
let reader = self.reader.clone();
SerializedPageReader::new(reader, column_chunk_metadata, total_rows, page_locations)?
.add_crypto_context(
rg_idx,
self.column_idx,
self.metadata.as_ref(),
column_chunk_metadata,
)
}
}
impl<T: ChunkReader + 'static> Iterator for ReaderPageIterator<T> {
type Item = Result<Box<dyn PageReader>>;
fn next(&mut self) -> Option<Self::Item> {
let rg_idx = self.row_groups.next()?;
let page_reader = self
.next_page_reader(rg_idx)
.map(|page_reader| Box::new(page_reader) as _);
Some(page_reader)
}
}
impl<T: ChunkReader + 'static> PageIterator for ReaderPageIterator<T> {}
pub struct ParquetRecordBatchReader {
array_reader: Box<dyn ArrayReader>,
schema: SchemaRef,
read_plan: ReadPlan,
}
#[derive(Default)]
enum FilterMaskAccumulator {
#[default]
Empty,
Single(BooleanBuffer),
Combined(BooleanBufferBuilder),
}
impl FilterMaskAccumulator {
fn append(&mut self, mask: BooleanBuffer) {
*self = match std::mem::take(self) {
Self::Empty => Self::Single(mask),
Self::Single(first) => {
let mut combined = BooleanBufferBuilder::new(first.len() + mask.len());
combined.append_buffer(&first);
combined.append_buffer(&mask);
Self::Combined(combined)
}
Self::Combined(mut combined) => {
combined.append_buffer(&mask);
Self::Combined(combined)
}
};
}
fn finish(self) -> Option<BooleanBuffer> {
match self {
Self::Empty => None,
Self::Single(mask) => Some(mask),
Self::Combined(combined) => Some(combined.build()),
}
}
}
fn consume_record_batch(array_reader: &mut dyn ArrayReader) -> Result<RecordBatch> {
let array = array_reader.consume_batch()?;
let struct_array = array.as_struct_opt().ok_or_else(|| {
ArrowError::ParquetError("Struct array reader should return struct array".to_string())
})?;
Ok(RecordBatch::from(struct_array))
}
fn read_mask_batch(
array_reader: &mut dyn ArrayReader,
mask_cursor: &mut MaskCursor,
batch_size: usize,
) -> Result<Option<RecordBatch>> {
let mut selected_rows = 0;
let mut filter_mask = FilterMaskAccumulator::default();
while selected_rows < batch_size && !mask_cursor.is_empty() {
let mask_chunk = mask_cursor.next_chunk(batch_size - selected_rows)?;
if mask_chunk.initial_skip > 0 {
let skipped = array_reader.skip_records(mask_chunk.initial_skip)?;
if skipped != mask_chunk.initial_skip {
return Err(general_err!(
"failed to skip rows, expected {}, got {}",
mask_chunk.initial_skip,
skipped
));
}
}
let mask = mask_cursor.mask_values_for(&mask_chunk)?;
let read = array_reader.read_records(mask_chunk.chunk_rows)?;
if read == 0 {
return Err(general_err!(
"reached end of column while expecting {} rows",
mask_chunk.chunk_rows
));
}
if read != mask_chunk.chunk_rows {
return Err(general_err!(
"insufficient rows read from array reader - expected {}, got {}",
mask_chunk.chunk_rows,
read
));
}
filter_mask.append(mask.values().clone());
selected_rows += mask_chunk.selected_rows;
}
if selected_rows == 0 {
return Ok(None);
}
let filter_mask = filter_mask
.finish()
.ok_or_else(|| general_err!("Internal Error: decoded Mask batch has no filter values"))?;
let batch = consume_record_batch(array_reader)?;
let filtered_batch = filter_record_batch(&batch, &BooleanArray::from(filter_mask))?;
if filtered_batch.num_rows() != selected_rows {
return Err(general_err!(
"filtered rows mismatch selection - expected {}, got {}",
selected_rows,
filtered_batch.num_rows()
));
}
Ok(Some(filtered_batch))
}
impl Debug for ParquetRecordBatchReader {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ParquetRecordBatchReader")
.field("array_reader", &"...")
.field("schema", &self.schema)
.field("read_plan", &self.read_plan)
.finish()
}
}
impl Iterator for ParquetRecordBatchReader {
type Item = Result<RecordBatch, ArrowError>;
fn next(&mut self) -> Option<Self::Item> {
self.next_inner()
.map_err(|arrow_err| arrow_err.into())
.transpose()
}
}
impl ParquetRecordBatchReader {
fn next_inner(&mut self) -> Result<Option<RecordBatch>> {
let mut read_records = 0;
let batch_size = self.batch_size();
if batch_size == 0 {
return Ok(None);
}
match self.read_plan.row_selection_cursor_mut() {
RowSelectionCursor::Mask(mask_cursor) => {
return read_mask_batch(self.array_reader.as_mut(), mask_cursor, batch_size);
}
RowSelectionCursor::Selectors(selectors_cursor) => {
while read_records < batch_size && !selectors_cursor.is_empty() {
let front = selectors_cursor.next_selector();
if front.skip {
let skipped = self.array_reader.skip_records(front.row_count)?;
if skipped != front.row_count {
return Err(general_err!(
"failed to skip rows, expected {}, got {}",
front.row_count,
skipped
));
}
continue;
}
if front.row_count == 0 {
continue;
}
let need_read = batch_size - read_records;
let to_read = match front.row_count.checked_sub(need_read) {
Some(remaining) if remaining != 0 => {
selectors_cursor.return_selector(RowSelector::select(remaining));
need_read
}
_ => front.row_count,
};
match self.array_reader.read_records(to_read)? {
0 => break,
rec => read_records += rec,
}
}
}
RowSelectionCursor::All => {
self.array_reader.read_records(batch_size)?;
}
}
let batch = consume_record_batch(self.array_reader.as_mut())?;
Ok(if batch.num_rows() > 0 {
Some(batch)
} else {
None
})
}
}
impl RecordBatchReader for ParquetRecordBatchReader {
fn schema(&self) -> SchemaRef {
self.schema.clone()
}
}
impl ParquetRecordBatchReader {
pub fn try_new<T: ChunkReader + 'static>(reader: T, batch_size: usize) -> Result<Self> {
ParquetRecordBatchReaderBuilder::try_new(reader)?
.with_batch_size(batch_size)
.build()
}
pub fn try_new_with_row_groups(
levels: &FieldLevels,
row_groups: &dyn RowGroups,
batch_size: usize,
selection: Option<RowSelection>,
) -> Result<Self> {
let metrics = ArrowReaderMetrics::disabled();
let array_reader = ArrayReaderBuilder::new(row_groups, &metrics)
.with_batch_size(batch_size)
.with_parquet_metadata(row_groups.metadata())
.build_array_reader(levels.levels.as_ref(), &ProjectionMask::all())?;
let read_plan = ReadPlanBuilder::new(batch_size)
.with_selection(selection)
.build();
Ok(Self {
array_reader,
schema: Arc::new(Schema::new(levels.fields.clone())),
read_plan,
})
}
pub(crate) fn new(array_reader: Box<dyn ArrayReader>, read_plan: ReadPlan) -> Self {
let schema = match array_reader.get_data_type() {
ArrowType::Struct(fields) => Schema::new(fields.clone()),
_ => unreachable!("Struct array reader's data type is not struct!"),
};
Self {
array_reader,
schema: Arc::new(schema),
read_plan,
}
}
#[inline(always)]
pub(crate) fn batch_size(&self) -> usize {
self.read_plan.batch_size()
}
}
#[cfg(test)]
pub(crate) mod tests;