use parquet::arrow::arrow_reader::RowSelection;
use roaring::RoaringBitmap;
use crate::{compile::CompileResult, selection::row_selection_to_roaring};
#[derive(Clone, Debug)]
pub struct PruneResult {
row_groups: Vec<usize>,
row_selection: Option<RowSelection>,
roaring: Option<RoaringBitmap>,
compile: CompileResult,
}
impl PruneResult {
pub fn row_groups(&self) -> &[usize] {
&self.row_groups
}
pub fn row_selection(&self) -> Option<&RowSelection> {
self.row_selection.as_ref()
}
pub fn roaring(&self) -> Option<&RoaringBitmap> {
self.roaring.as_ref()
}
pub fn compile_result(&self) -> &CompileResult {
&self.compile
}
pub fn into_roaring(self, total_rows: u64) -> Option<RoaringBitmap> {
if let Some(roaring) = self.roaring {
return Some(roaring);
}
self.row_selection
.and_then(|sel| row_selection_to_roaring(&sel, total_rows))
}
pub fn into_parts(
self,
) -> (
Vec<usize>,
Option<RowSelection>,
Option<RoaringBitmap>,
CompileResult,
) {
(
self.row_groups,
self.row_selection,
self.roaring,
self.compile,
)
}
}
impl PruneResult {
pub(super) fn new(
row_groups: Vec<usize>,
row_selection: Option<RowSelection>,
roaring: Option<RoaringBitmap>,
compile: CompileResult,
) -> Self {
Self {
row_groups,
row_selection,
roaring,
compile,
}
}
}