Skip to content

refactor: Optimize required_columns from BTreeSet to Vec in struct PushdownChecker #19678

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
kosiew merged 3 commits into from
Jan 14, 2026
Merged

refactor: Optimize required_columns from BTreeSet to Vec in struct PushdownChecker #19678

Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
8c2462e
refactor: Optimize required_columns from BTreeSet to Vec in struct Pu…
kumarUjjawal Jan 7, 2026
2a51524
use dedup and update comments
kumarUjjawal Jan 12, 2026
63b2ab5
add comments for required_column
kumarUjjawal Jan 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 13 additions & 8 deletions datafusion/datasource-parquet/src/row_filter.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -276,7 +276,7 @@ struct PushdownChecker<'schema> {
/// Does the expression reference any columns not present in the file schema?
projected_columns: bool,
/// Indices into the file schema of columns required to evaluate the expression.
required_columns: BTreeSet<usize>,
required_columns: Vec<usize>,
Copy link
Contributor

@alamb alamb Jan 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

❤️

/// Tracks the nested column behavior found during traversal.
nested_behavior: NestedColumnSupport,
/// Whether nested list columns are supported by the predicate semantics.
Expand All @@ -290,7 +290,7 @@ impl<'schema> PushdownChecker<'schema> {
Self {
non_primitive_columns: false,
projected_columns: false,
required_columns: BTreeSet::default(),
required_columns: Vec::new(),
nested_behavior: NestedColumnSupport::PrimitiveOnly,
allow_list_columns,
file_schema,
Expand All @@ -307,7 +307,9 @@ impl<'schema> PushdownChecker<'schema> {
}
};

self.required_columns.insert(idx);
if !self.required_columns.contains(&idx) {
self.required_columns.push(idx);
}
Copy link
Contributor

@kosiew kosiew Jan 12, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it would help future contributors to

  • Add comment explaining linear search is acceptable for small n
    • OR switch to HashSet for O(1) deduplication if n might grow

Copy link
Contributor

@kosiew kosiew Jan 12, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

see suggested

fn into_sorted_columns
below

let data_type = self.file_schema.field(idx).data_type();

if DataType::is_nested(data_type) {
Expand Down Expand Up @@ -390,9 +392,8 @@ enum NestedColumnSupport {
Unsupported,
}

#[derive(Debug)]
struct PushdownColumns {
required_columns: BTreeSet<usize>,
required_columns: Vec<usize>,
Copy link
Contributor

@alamb alamb Jan 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

maybe it is worth a comment here explaining the assumption that required_columns are sorted and unique (non duplicate)

nested: NestedColumnSupport,
}

Expand All @@ -411,9 +412,13 @@ fn pushdown_columns(
let allow_list_columns = supports_list_predicates(expr);
let mut checker = PushdownChecker::new(file_schema, allow_list_columns);
expr.visit(&mut checker)?;
Ok((!checker.prevents_pushdown()).then_some(PushdownColumns {
required_columns: checker.required_columns,
nested: checker.nested_behavior,
let prevents_pushdown = checker.prevents_pushdown();
let nested = checker.nested_behavior;
let mut required_columns = checker.required_columns;
required_columns.sort_unstable();
Copy link
Contributor

@kosiew kosiew Jan 12, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

how about adding:

impl PushdownChecker {
    fn into_sorted_columns(mut self) -> PushdownColumns {
        self.required_columns.sort_unstable();
        self.required_columns.dedup(); // this removes the need for contains check
        PushdownColumns {
            required_columns: self.required_columns,
            nested: self.nested_behavior,
        }
    }

Ok((!prevents_pushdown).then_some(PushdownColumns {
required_columns,
nested,
}))
Copy link
Contributor

@kosiew kosiew Jan 12, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

and then rewriting this to:

Ok((!checker.prevents_pushdown())
    .then_some(checker.into_sorted_columns()))

}

Expand Down