Skip to content

Commit 17e6c88

Browse files
minor: refactor with assert_or_internal_err!() in datafusion/optimizer (#18699)
## Which issue does this PR close? <!-- We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123. --> - Part of #18613 ## Rationale for this change <!-- Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed. Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes. --> ## What changes are included in this PR? <!-- There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR. --> ## Are these changes tested? <!-- We typically require tests for all PRs in order to: 1. Prevent the code from being accidentally broken by subsequent changes 2. Serve as another way to document the expected behavior of the code If tests are not included in your PR, please explain why (for example, are they covered by existing tests)? --> ## Are there any user-facing changes? <!-- If there are user-facing changes then we may require documentation to be updated before approving the PR. --> <!-- If there are any breaking changes to public APIs, please add the `api change` label. --> --------- Co-authored-by: Jeffrey Vo <jeffrey.vo.australia@gmail.com>
1 parent 0cfc1fe commit 17e6c88

File tree

6 files changed

+62
-46
lines changed

6 files changed

+62
-46
lines changed

datafusion/optimizer/src/decorrelate_predicate_subquery.rs

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,9 @@ use crate::{OptimizerConfig, OptimizerRule};
2727

2828
use datafusion_common::alias::AliasGenerator;
2929
use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
30-
use datafusion_common::{internal_err, plan_err, Column, Result};
30+
use datafusion_common::{
31+
assert_or_internal_err, plan_err, Column, DataFusionError, Result,
32+
};
3133
use datafusion_expr::expr::{Exists, InSubquery};
3234
use datafusion_expr::expr_rewriter::create_col_from_scalar_expr;
3335
use datafusion_expr::logical_plan::{JoinType, Subquery};
@@ -79,11 +81,10 @@ impl OptimizerRule for DecorrelatePredicateSubquery {
7981
.into_iter()
8082
.partition(has_subquery);
8183

82-
if with_subqueries.is_empty() {
83-
return internal_err!(
84-
"can not find expected subqueries in DecorrelatePredicateSubquery"
85-
);
86-
}
84+
assert_or_internal_err!(
85+
!with_subqueries.is_empty(),
86+
"can not find expected subqueries in DecorrelatePredicateSubquery"
87+
);
8788

8889
// iterate through all exists clauses in predicate, turning each into a join
8990
let mut cur_input = Arc::unwrap_or_clone(filter.input);

datafusion/optimizer/src/extract_equijoin_predicate.rs

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
use crate::optimizer::ApplyOrder;
2020
use crate::{OptimizerConfig, OptimizerRule};
2121
use datafusion_common::tree_node::Transformed;
22-
use datafusion_common::{internal_err, DFSchema};
22+
use datafusion_common::{assert_or_internal_err, DFSchema, DataFusionError};
2323
use datafusion_common::{NullEquality, Result};
2424
use datafusion_expr::utils::split_conjunction_owned;
2525
use datafusion_expr::utils::{can_hash, find_valid_equijoin_key_pair};
@@ -223,13 +223,12 @@ fn split_op_and_other_join_predicates(
223223
right_schema: &DFSchema,
224224
operator: Operator,
225225
) -> Result<(Vec<EquijoinPredicate>, Option<Expr>)> {
226-
if !matches!(operator, Operator::Eq | Operator::IsNotDistinctFrom) {
227-
return internal_err!(
228-
"split_op_and_other_join_predicates only supports 'Eq' or 'IsNotDistinctFrom' operators, \
229-
but received: {:?}",
230-
operator
231-
);
232-
}
226+
assert_or_internal_err!(
227+
matches!(operator, Operator::Eq | Operator::IsNotDistinctFrom),
228+
"split_op_and_other_join_predicates only supports 'Eq' or 'IsNotDistinctFrom' operators, \
229+
but received: {:?}",
230+
operator
231+
);
233232

234233
let exprs = split_conjunction_owned(filter);
235234

datafusion/optimizer/src/optimize_projections/mod.rs

Lines changed: 16 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -25,8 +25,9 @@ use std::collections::HashSet;
2525
use std::sync::Arc;
2626

2727
use datafusion_common::{
28-
get_required_group_by_exprs_indices, internal_datafusion_err, internal_err, Column,
29-
DFSchema, HashMap, JoinType, Result,
28+
assert_eq_or_internal_err, get_required_group_by_exprs_indices,
29+
internal_datafusion_err, internal_err, Column, DFSchema, DataFusionError, HashMap,
30+
JoinType, Result,
3031
};
3132
use datafusion_expr::expr::Alias;
3233
use datafusion_expr::{
@@ -341,11 +342,14 @@ fn optimize_projections(
341342
return Ok(Transformed::no(plan));
342343
};
343344
let children = extension.node.inputs();
344-
if children.len() != necessary_children_indices.len() {
345-
return internal_err!("Inconsistent length between children and necessary children indices. \
346-
Make sure `.necessary_children_exprs` implementation of the `UserDefinedLogicalNode` is \
347-
consistent with actual children length for the node.");
348-
}
345+
assert_eq_or_internal_err!(
346+
children.len(),
347+
necessary_children_indices.len(),
348+
"Inconsistent length between children and necessary children indices. \
349+
Make sure `.necessary_children_exprs` implementation of the \
350+
`UserDefinedLogicalNode` is consistent with actual children length \
351+
for the node."
352+
);
349353
children
350354
.into_iter()
351355
.zip(necessary_children_indices)
@@ -432,11 +436,11 @@ fn optimize_projections(
432436
// Required indices are currently ordered (child0, child1, ...)
433437
// but the loop pops off the last element, so we need to reverse the order
434438
child_required_indices.reverse();
435-
if child_required_indices.len() != plan.inputs().len() {
436-
return internal_err!(
437-
"OptimizeProjection: child_required_indices length mismatch with plan inputs"
438-
);
439-
}
439+
assert_eq_or_internal_err!(
440+
child_required_indices.len(),
441+
plan.inputs().len(),
442+
"OptimizeProjection: child_required_indices length mismatch with plan inputs"
443+
);
440444

441445
// Rewrite children of the plan
442446
let transformed_plan = plan.map_children(|child| {

datafusion/optimizer/src/push_down_filter.rs

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,8 @@ use datafusion_common::tree_node::{
2828
Transformed, TransformedResult, TreeNode, TreeNodeRecursion,
2929
};
3030
use datafusion_common::{
31-
internal_err, plan_err, qualified_name, Column, DFSchema, Result,
31+
assert_eq_or_internal_err, internal_err, plan_err, qualified_name, Column, DFSchema,
32+
DataFusionError, Result,
3233
};
3334
use datafusion_expr::expr::WindowFunction;
3435
use datafusion_expr::expr_rewriter::replace_col;
@@ -1135,12 +1136,13 @@ impl OptimizerRule for PushDownFilter {
11351136
let supported_filters = scan
11361137
.source
11371138
.supports_filters_pushdown(non_volatile_filters.as_slice())?;
1138-
if non_volatile_filters.len() != supported_filters.len() {
1139-
return internal_err!(
1140-
"Vec returned length: {} from supports_filters_pushdown is not the same size as the filters passed, which length is: {}",
1141-
supported_filters.len(),
1142-
non_volatile_filters.len());
1143-
}
1139+
assert_eq_or_internal_err!(
1140+
non_volatile_filters.len(),
1141+
supported_filters.len(),
1142+
"Vec returned length: {} from supports_filters_pushdown is not the same size as the filters passed, which length is: {}",
1143+
supported_filters.len(),
1144+
non_volatile_filters.len()
1145+
);
11441146

11451147
// Compose scan filters from non-volatile filters of `Exact` or `Inexact` pushdown type
11461148
let zip = non_volatile_filters.into_iter().zip(supported_filters);

datafusion/optimizer/src/scalar_subquery_to_join.rs

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,9 @@ use datafusion_common::alias::AliasGenerator;
3030
use datafusion_common::tree_node::{
3131
Transformed, TransformedResult, TreeNode, TreeNodeRecursion, TreeNodeRewriter,
3232
};
33-
use datafusion_common::{internal_err, plan_err, Column, Result, ScalarValue};
33+
use datafusion_common::{
34+
assert_or_internal_err, plan_err, Column, DataFusionError, Result, ScalarValue,
35+
};
3436
use datafusion_expr::expr_rewriter::create_col_from_scalar_expr;
3537
use datafusion_expr::logical_plan::{JoinType, Subquery};
3638
use datafusion_expr::utils::conjunction;
@@ -94,9 +96,10 @@ impl OptimizerRule for ScalarSubqueryToJoin {
9496
config.alias_generator(),
9597
)?;
9698

97-
if subqueries.is_empty() {
98-
return internal_err!("Expected subqueries not found in filter");
99-
}
99+
assert_or_internal_err!(
100+
!subqueries.is_empty(),
101+
"Expected subqueries not found in filter"
102+
);
100103

101104
// iterate through all subqueries in predicate, turning each into a left join
102105
let mut cur_input = filter.input.as_ref().clone();
@@ -154,9 +157,10 @@ impl OptimizerRule for ScalarSubqueryToJoin {
154157
all_subqueries.extend(subqueries);
155158
expr_to_rewrite_expr_map.insert(expr, rewrite_exprs);
156159
}
157-
if all_subqueries.is_empty() {
158-
return internal_err!("Expected subqueries not found in projection");
159-
}
160+
assert_or_internal_err!(
161+
!all_subqueries.is_empty(),
162+
"Expected subqueries not found in projection"
163+
);
160164
// iterate through all subqueries in predicate, turning each into a left join
161165
let mut cur_input = projection.input.as_ref().clone();
162166
for (subquery, alias) in all_subqueries {

datafusion/optimizer/src/single_distinct_to_groupby.rs

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ use crate::optimizer::ApplyOrder;
2323
use crate::{OptimizerConfig, OptimizerRule};
2424

2525
use datafusion_common::{
26-
internal_err, tree_node::Transformed, DataFusionError, HashSet, Result,
26+
assert_eq_or_internal_err, tree_node::Transformed, DataFusionError, HashSet, Result,
2727
};
2828
use datafusion_expr::builder::project;
2929
use datafusion_expr::expr::AggregateFunctionParams;
@@ -183,15 +183,21 @@ impl OptimizerRule for SingleDistinctToGroupBy {
183183
.map(|aggr_expr| match aggr_expr {
184184
Expr::AggregateFunction(AggregateFunction {
185185
func,
186-
params: AggregateFunctionParams { mut args, distinct, .. }
186+
params:
187+
AggregateFunctionParams {
188+
mut args, distinct, ..
189+
},
187190
}) => {
188191
if distinct {
189-
if args.len() != 1 {
190-
return internal_err!("DISTINCT aggregate should have exactly one argument");
191-
}
192+
assert_eq_or_internal_err!(
193+
args.len(),
194+
1,
195+
"DISTINCT aggregate should have exactly one argument"
196+
);
192197
let arg = args.swap_remove(0);
193198

194-
if group_fields_set.insert(arg.schema_name().to_string()) {
199+
if group_fields_set.insert(arg.schema_name().to_string())
200+
{
195201
inner_group_exprs
196202
.push(arg.alias(SINGLE_DISTINCT_ALIAS));
197203
}

0 commit comments

Comments
 (0)