Skip to content

Commit

Permalink
Change n_columns API
Browse files Browse the repository at this point in the history
  • Loading branch information
alamb committed Jul 16, 2024
1 parent 9c34351 commit c9bd0af
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 10 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -144,8 +144,9 @@ impl PagePruningAccessPlanFilter {
return None;
}

if pp.required_columns().n_columns() > 1 {
if pp.required_columns().single_column().is_none() {
debug!("Ignoring multi-column page pruning predicate: {predicate}");
return None;
}

Some(pp)
Expand Down Expand Up @@ -196,7 +197,6 @@ impl PagePruningAccessPlanFilter {
// The selection for this particular row group
let mut overall_selection = None;
for predicate in page_index_predicates {

// find column index in the parquet schema
let col_idx = find_column_index(predicate, arrow_schema, parquet_schema);
let row_group_metadata = &groups[r];
Expand Down
22 changes: 14 additions & 8 deletions datafusion/core/src/physical_optimizer/pruning.rs
Original file line number Diff line number Diff line change
Expand Up @@ -738,16 +738,22 @@ impl RequiredColumns {
Self::default()
}

/// Returns number of unique columns
/// Returns Some(column) if this is a single column predicate.
///
/// Returns None if this is a multi-column predicate.
///
/// Examples:
/// * `a > 5 OR a < 10` returns `1`
/// * `a > 5 OR b < 10` returns `2`
pub(crate) fn n_columns(&self) -> usize {
self.iter()
.map(|(c, _s, _f)| c)
.collect::<HashSet<_>>()
.len()
/// * `a > 5 OR a < 10` returns `Some(a)`
/// * `a > 5 OR b < 10` returns `None`
/// * `true` returns None
pub(crate) fn single_column(&self) -> Option<&phys_expr::Column> {
let cols = self.iter().map(|(c, _s, _f)| c).collect::<HashSet<_>>();

if cols.len() == 1 {
cols.iter().next().copied()
} else {
None
}
}

/// Returns an iterator over items in columns (see doc on
Expand Down

0 comments on commit c9bd0af

Please sign in to comment.