Skip to content

Simplify Parquet AST expressions for row group and page pruning #23992

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

Open
mhaseeb123 wants to merge 5 commits into
base: main
Choose a base branch
from
Open

Simplify Parquet AST expressions for row group and page pruning #23992

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
eb6b67b
Add Parquet pruning expression builder
mhaseeb123 Sep 9, 2026
8840ea3
Add a gtest
mhaseeb123 Sep 9, 2026
cb5b2cd
Merge branch 'main' into parquet-pruning-builder-v2
mhaseeb123 Sep 9, 2026
fd0fc26
Use spans instead of host_spans
mhaseeb123 Sep 9, 2026
6ee0308
Apply suggestions from coderabbit
mhaseeb123 Sep 9, 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
135 changes: 50 additions & 85 deletions cpp/src/io/parquet/bloom_filter_reader.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -178,20 +178,15 @@ struct bloom_filter_caster {
* @brief Converts AST expression to bloom filter membership (BloomfilterAST) expression.
* This is used in row group filtering based on equality predicate.
*/
class bloom_filter_expression_converter : public equality_literals_collector {
class bloom_filter_expression_converter : public parquet_expression_simplifier {
public:
bloom_filter_expression_converter(
ast::expression const& expr,
cudf::host_span<cudf::data_type const> output_dtypes,
cudf::host_span<std::vector<ast::literal*> const> equality_literals,
cuda::stream_ref stream)
: _equality_literals{equality_literals},
_always_true_scalar{std::make_unique<cudf::numeric_scalar<bool>>(true, true, stream)},
_always_true{std::make_unique<ast::literal>(*_always_true_scalar)}
cudf::host_span<std::vector<ast::literal*> const> equality_literals)
: parquet_expression_simplifier{std::span{output_dtypes.data(), output_dtypes.size()}},
_equality_literals{equality_literals}
{
// Set the output data types
_output_dtypes = output_dtypes;

// Compute and store columns literals offsets
_col_literals_offsets.reserve(static_cast<cudf::size_type>(_output_dtypes.size()) + 1);
_col_literals_offsets.emplace_back(0);
Expand All @@ -204,93 +199,61 @@ class bloom_filter_expression_converter : public equality_literals_collector {
static_cast<cudf::size_type>(col_literal_map.size());
});

// Add this visitor
expr.accept(*this);
_bloom_filter_expr = simplify_expr(expr);
}

/**
* @brief Delete equality literals getter as it's not needed in the derived class
* @brief Returns the AST to apply on bloom filter membership
*
* @return The membership expression, or std::nullopt if no row group can be pruned
*/
[[nodiscard]] std::vector<std::vector<ast::literal*>> get_equality_literals() && = delete;

// Bring all overloads of `visit` from equality_predicate_collector into scope
using equality_literals_collector::visit;
[[nodiscard]] simplified_expression_opt get_bloom_filter_expr() const
{
return _bloom_filter_expr;
}

protected:
/**
* @copydoc ast::detail::expression_transformer::visit(ast::operation const& )
* @copydoc parquet_expression_simplifier::simplify_comparison
*
* A bloom filter answers only "might this value be present", so equality is the one comparison
* it can evaluate. Every other node relaxes via the base class defaults, including `NOT`, whose
* membership answer cannot be complemented: `¬(some row is 5)` means "no row is 5", not "some
* row is not 5".
*/
std::reference_wrapper<ast::expression const> visit(ast::operation const& expr) override
[[nodiscard]] simplified_expression_opt simplify_comparison(ast::ast_operator op,
ast::column_reference const& col_ref,
ast::literal const& literal) override
{
using cudf::ast::ast_operator;

auto const input_op = expr.get_operator();
auto const operator_arity = cudf::ast::detail::ast_operator_arity(input_op);
if (op != ast_operator::EQUAL) { return std::nullopt; }

// Membership filters cannot evaluate unary operations. Visit operands and push always true
if (operator_arity == 1) {
std::ignore = this->visit_operands(expr.get_operands());
_bloom_filter_expr.push(ast::operation{ast_operator::IDENTITY, *_always_true});
return *_always_true;
}
auto const col_idx = col_ref.get_column_index();
auto const& equality_literals = _equality_literals[col_idx];

// Binary operation
auto const [op, lhs_kind, rhs_kind, col_ref, literal] = extract_binary_operands(expr);

// Push expressions for `col op lit` or `lit op col` forms
if (lhs_kind == operand_kind::COLUMN_REF and rhs_kind == operand_kind::LITERAL) {
col_ref->accept(*this);

if (op == ast_operator::EQUAL) {
auto const col_idx = col_ref->get_column_index();
auto const& equality_literals = _equality_literals[col_idx];
auto col_literal_offset = _col_literals_offsets[col_idx];
// Skip bloom filter probing for timestamp columns with empty vector of literals due to
// a timestamp scale mismatch — the literal can never match the native values.
if (cudf::is_timestamp(_output_dtypes[col_idx]) and equality_literals.empty()) {
return *_always_true;
}

auto const literal_iter =
std::find(equality_literals.cbegin(), equality_literals.cend(), literal);
CUDF_EXPECTS(literal_iter != equality_literals.end(),
"Bloom filter expression converter encountered an unexpected literal");

col_literal_offset += std::distance(equality_literals.cbegin(), literal_iter);
auto const& value = _bloom_filter_expr.push(ast::column_reference{col_literal_offset});
_bloom_filter_expr.push(ast::operation{ast_operator::IDENTITY, value});
} else {
_bloom_filter_expr.push(ast::operation{ast_operator::IDENTITY, *_always_true});
return *_always_true;
}
} // Visit operands and push expression for `expr op expr` form
else if (lhs_kind == operand_kind::EXPRESSION and rhs_kind == operand_kind::EXPRESSION) {
auto new_operands = visit_operands(expr.get_operands());
_bloom_filter_expr.push(ast::operation{op, new_operands.front(), new_operands.back()});
} // Push _always_true for `col op col`, `expr op col`, `expr op lit` forms
else {
_bloom_filter_expr.push(ast::operation{ast_operator::IDENTITY, *_always_true});
return *_always_true;
// Skip bloom filter probing for timestamp columns with empty vector of literals due to
// a timestamp scale mismatch — the literal can never match the native values.
if (cudf::is_timestamp(_output_dtypes[col_idx]) and equality_literals.empty()) {
return std::nullopt;
}

return _bloom_filter_expr.back();
}
auto const literal_iter =
std::find(equality_literals.cbegin(), equality_literals.cend(), &literal);
CUDF_EXPECTS(literal_iter != equality_literals.end(),
"Bloom filter expression converter encountered an unexpected literal");

/**
* @brief Returns the AST to apply on bloom filter membership.
*
* @return AST operation expression
*/
[[nodiscard]] std::reference_wrapper<ast::expression const> get_bloom_filter_expr() const
{
return _bloom_filter_expr.back();
auto const col_literal_offset =
_col_literals_offsets[col_idx] +
static_cast<cudf::size_type>(std::distance(equality_literals.cbegin(), literal_iter));
auto const& value = _tree.push(ast::column_reference{col_literal_offset});
return _tree.push(ast::operation{ast_operator::IDENTITY, value});
}

private:
std::vector<cudf::size_type> _col_literals_offsets;
cudf::host_span<std::vector<ast::literal*> const> _equality_literals;
ast::tree _bloom_filter_expr;
std::unique_ptr<cudf::numeric_scalar<bool>> _always_true_scalar;
std::unique_ptr<ast::literal> _always_true;
simplified_expression_opt _bloom_filter_expr;
};

} // namespace
Expand Down Expand Up @@ -413,6 +376,15 @@ std::optional<std::vector<std::vector<size_type>>> aggregate_reader_metadata::ap
std::reference_wrapper<ast::expression const> filter,
cuda::stream_ref stream) const
{
// Convert AST to BloomfilterAST expression with reference to bloom filter membership
// in above `bloom_filter_membership_table`
bloom_filter_expression_converter bloom_filter_expr_converter{
filter.get(), output_dtypes, {literals}};

// Return early if bloom filters cannot prune any row groups using the filter
auto const bloom_filter_expr = bloom_filter_expr_converter.get_bloom_filter_expr();
if (not bloom_filter_expr.has_value()) { return std::nullopt; }

// Number of input table columns
auto const num_input_columns = static_cast<cudf::size_type>(output_dtypes.size());

Expand Down Expand Up @@ -457,17 +429,10 @@ std::optional<std::vector<std::vector<size_type>>> aggregate_reader_metadata::ap
// Create a table from columns
auto bloom_filter_membership_table = cudf::table(std::move(bloom_filter_membership_columns));

// Convert AST to BloomfilterAST expression with reference to bloom filter membership
// in above `bloom_filter_membership_table`
bloom_filter_expression_converter bloom_filter_expr{
filter.get(), output_dtypes, {literals}, stream};

// Filter bloom filter membership table with the BloomfilterAST expression and collect
// filtered row group indices
return collect_filtered_row_group_indices(bloom_filter_membership_table,
bloom_filter_expr.get_bloom_filter_expr(),
input_row_group_indices,
stream);
return collect_filtered_row_group_indices(
bloom_filter_membership_table, bloom_filter_expr.value(), input_row_group_indices, stream);
}

equality_literals_collector::equality_literals_collector(
Expand Down
155 changes: 58 additions & 97 deletions cpp/src/io/parquet/experimental/dictionary_page_filter.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1313,23 +1313,20 @@ struct dictionary_caster {
}
};

using parquet::detail::parquet_expression_simplifier;
using parquet::detail::simplified_expression_opt;

/**
* @brief Converts AST expression to dictionary membership (DictionaryAST) expression.
* This is used in row group filtering based on equality predicate.
*/
class dictionary_expression_converter : public equality_literals_collector {
class dictionary_expression_converter : public parquet_expression_simplifier {
public:
dictionary_expression_converter(ast::expression const& expr,
cudf::host_span<cudf::data_type const> output_dtypes,
cudf::host_span<std::vector<ast::literal*> const> literals,
cuda::stream_ref stream)
: _literals{literals},
_always_true_scalar{std::make_unique<cudf::numeric_scalar<bool>>(true, true, stream)},
_always_true{std::make_unique<ast::literal>(*_always_true_scalar)}
std::span<cudf::data_type const> output_dtypes,
cudf::host_span<std::vector<ast::literal*> const> literals)
: parquet_expression_simplifier{output_dtypes}, _literals{literals}
{
// Set the output data types
_output_dtypes = output_dtypes;

// Compute and store columns literals offsets
_col_literals_offsets.reserve(static_cast<cudf::size_type>(_output_dtypes.size()) + 1);
_col_literals_offsets.emplace_back(0);
Expand All @@ -1343,100 +1340,59 @@ class dictionary_expression_converter : public equality_literals_collector {
static_cast<cudf::size_type>(col_literal_map.size());
});

// Add this visitor
expr.accept(*this);
_dictionary_expr = simplify_expr(expr);
}

/**
* @brief Delete equality literals getter as it's not needed in the derived class
* @brief Returns the AST to apply on dictionary membership
*
* @return The membership expression, or std::nullopt if no row group can be pruned
*/
[[nodiscard]] std::vector<std::vector<ast::literal*>> get_equality_literals() && = delete;

// Bring all overloads of `visit` from equality_predicate_collector into scope
using equality_literals_collector::visit;
[[nodiscard]] simplified_expression_opt get_dictionary_expr() const { return _dictionary_expr; }

protected:
/**
* @copydoc ast::detail::expression_transformer::visit(ast::operation const& )
* @copydoc parquet::detail::parquet_expression_simplifier::simplify_comparison
*
* Dictionary membership supports equality and inequality comparisons. Unsupported operations
* return std::nullopt.
*/
std::reference_wrapper<ast::expression const> visit(ast::operation const& expr) override
[[nodiscard]] simplified_expression_opt simplify_comparison(ast::ast_operator op,
ast::column_reference const& col_ref,
ast::literal const& literal) override
{
using cudf::ast::ast_operator;
using parquet::detail::extract_binary_operands;
using parquet::detail::operand_kind;

auto const input_op = expr.get_operator();
auto const operator_arity = cudf::ast::detail::ast_operator_arity(input_op);

// Membership filters cannot evaluate unary operations. Visit operands and push always true
if (operator_arity == 1) {
std::ignore = this->visit_operands(expr.get_operands());
_dictionary_expr.push(ast::operation{ast_operator::IDENTITY, *_always_true});
return *_always_true;
}

// Binary operation
auto const [op, lhs_kind, rhs_kind, col_ref, literal] = extract_binary_operands(expr);

// Push expressions for `col op lit` or `lit op col` forms
if (lhs_kind == operand_kind::COLUMN_REF and rhs_kind == operand_kind::LITERAL) {
col_ref->accept(*this);

if (op == ast_operator::EQUAL or op == ast_operator::NOT_EQUAL) {
auto const col_idx = col_ref->get_column_index();
auto const& equality_literals = _literals[col_idx];
auto col_literal_offset = _col_literals_offsets[col_idx];
auto const literal_iter =
std::find(equality_literals.cbegin(), equality_literals.cend(), literal);
CUDF_EXPECTS(literal_iter != equality_literals.end(),
"Dictionary expression converter encountered an unexpected literal");
col_literal_offset += std::distance(equality_literals.cbegin(), literal_iter);

auto const& value = _dictionary_expr.push(ast::column_reference{col_literal_offset});

if (op == ast_operator::NOT_EQUAL) {
// For NOT_EQUAL operator, simply evaluate boolean is_false(value) expression as
// NOT(value). The value indicates if the row group should be pruned (if the literal is
// present in the hash set and it's the only value in the hash set)
_dictionary_expr.push(ast::operation{ast_operator::NOT, value});
} else {
// For EQUAL operator, evaluate boolean is_true(value) expression as IDENTITY(value)
// The value indicates if the row group should be kept (if the literal is present in the
// hash set)
_dictionary_expr.push(ast::operation{ast_operator::IDENTITY, value});
}
} // For all other expressions, push the `_always_true` expression
else {
_dictionary_expr.push(ast::operation{ast_operator::IDENTITY, *_always_true});
return *_always_true;
}
} // Visit operands and push expression for `expr op expr` form
else if (lhs_kind == operand_kind::EXPRESSION and rhs_kind == operand_kind::EXPRESSION) {
auto new_operands = visit_operands(expr.get_operands());
_dictionary_expr.push(ast::operation{op, new_operands.front(), new_operands.back()});
} // Push _always_true for `col op col`, `expr op col`, `expr op lit` forms
else {
_dictionary_expr.push(ast::operation{ast_operator::IDENTITY, *_always_true});
return *_always_true;
if (op != ast_operator::EQUAL and op != ast_operator::NOT_EQUAL) { return std::nullopt; }

auto const col_idx = col_ref.get_column_index();
auto const& equality_literals = _literals[col_idx];
auto const literal_iter =
std::find(equality_literals.cbegin(), equality_literals.cend(), &literal);
CUDF_EXPECTS(literal_iter != equality_literals.end(),
"Dictionary expression converter encountered an unexpected literal");
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Outdated

auto const col_literal_offset =
_col_literals_offsets[col_idx] +
static_cast<cudf::size_type>(std::distance(equality_literals.cbegin(), literal_iter));
auto const& value = _tree.push(ast::column_reference{col_literal_offset});

if (op == ast_operator::NOT_EQUAL) {
// For NOT_EQUAL operator, simply evaluate boolean is_false(value) expression as
// NOT(value). The value indicates if the row group should be pruned (if the literal is
// present in the hash set and it's the only value in the hash set)
return _tree.push(ast::operation{ast_operator::NOT, value});
}
return _dictionary_expr.back();
}

/**
* @brief Returns the AST to apply on dictionary membership.
*
* @return AST operation expression
*/
[[nodiscard]] std::reference_wrapper<ast::expression const> get_dictionary_expr() const
{
return _dictionary_expr.back();
// For EQUAL operator, evaluate boolean is_true(value) expression as IDENTITY(value)
// The value indicates if the row group should be kept (if the literal is present in the
// hash set)
return _tree.push(ast::operation{ast_operator::IDENTITY, value});
}

private:
std::vector<cudf::size_type> _col_literals_offsets;
cudf::host_span<std::vector<ast::literal*> const> _literals;
ast::tree _dictionary_expr;
std::unique_ptr<cudf::numeric_scalar<bool>> _always_true_scalar;
std::unique_ptr<ast::literal> _always_true;
simplified_expression_opt _dictionary_expr;
};

} // namespace
Expand All @@ -1454,6 +1410,19 @@ aggregate_reader_metadata::apply_dictionary_filter(
std::reference_wrapper<ast::expression const> filter,
cuda::stream_ref stream) const
{
CUDF_FUNC_RANGE();

// Convert AST to DictionaryAST expression with reference to dictionary membership
// in above `dictionary_membership_table`
dictionary_expression_converter dictionary_expr_converter{
filter.get(),
output_dtypes,
cudf::host_span<std::vector<ast::literal*> const>{literals.data(), literals.size()}};

// Dictionary membership cannot filter anything in the filter, all row groups survive
auto const dictionary_expr = dictionary_expr_converter.get_dictionary_expr();
if (not dictionary_expr.has_value()) { return std::nullopt; }

// Number of input table columns
auto const num_input_columns = static_cast<cudf::size_type>(output_dtypes.size());
// Number of columns with dictionaries
Expand Down Expand Up @@ -1519,19 +1488,11 @@ aggregate_reader_metadata::apply_dictionary_filter(
// Create a table from columns
auto const dictionary_membership_table = cudf::table(std::move(dictionary_membership_columns));

// Convert AST to DictionaryAST expression with reference to dictionary membership
// in above `dictionary_membership_table`
dictionary_expression_converter dictionary_expr{
filter.get(),
cudf::host_span<cudf::data_type const>{output_dtypes.data(), output_dtypes.size()},
cudf::host_span<std::vector<ast::literal*> const>{literals.data(), literals.size()},
stream};

// Filter dictionary membership table with the DictionaryAST expression and collect
// filtered row group indices
return parquet::detail::collect_filtered_row_group_indices(
dictionary_membership_table,
dictionary_expr.get_dictionary_expr(),
dictionary_expr.value(),
cudf::host_span<std::vector<size_type> const>{input_row_group_indices.data(),
input_row_group_indices.size()},
stream);
Expand Down
Loading
Loading