Skip to content

Commit 4198c5a

Browse files
authored
minor: refactor with assert_or_internal_err!() in datafusion/expr (#18731)
## 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. -->
1 parent 5791822 commit 4198c5a

File tree

4 files changed

+66
-60
lines changed

4 files changed

+66
-60
lines changed

datafusion/expr/src/expr_schema.rs

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -775,7 +775,9 @@ mod tests {
775775
use super::*;
776776
use crate::{col, lit, out_ref_col_with_metadata};
777777

778-
use datafusion_common::{internal_err, DFSchema, ScalarValue};
778+
use datafusion_common::{
779+
assert_or_internal_err, DFSchema, DataFusionError, ScalarValue,
780+
};
779781

780782
macro_rules! test_is_expr_nullable {
781783
($EXPR_TYPE:ident) => {{
@@ -1000,11 +1002,8 @@ mod tests {
10001002

10011003
impl ExprSchema for MockExprSchema {
10021004
fn nullable(&self, _col: &Column) -> Result<bool> {
1003-
if self.error_on_nullable {
1004-
internal_err!("nullable error")
1005-
} else {
1006-
Ok(self.field.is_nullable())
1007-
}
1005+
assert_or_internal_err!(!self.error_on_nullable, "nullable error");
1006+
Ok(self.field.is_nullable())
10081007
}
10091008

10101009
fn field_from_column(&self, _col: &Column) -> Result<&Field> {

datafusion/expr/src/logical_plan/invariants.rs

Lines changed: 9 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,9 @@
1616
// under the License.
1717

1818
use datafusion_common::{
19-
internal_err, plan_err,
19+
assert_or_internal_err, plan_err,
2020
tree_node::{TreeNode, TreeNodeRecursion},
21-
DFSchemaRef, Result,
21+
DFSchemaRef, DataFusionError, Result,
2222
};
2323

2424
use crate::{
@@ -114,15 +114,13 @@ fn assert_valid_semantic_plan(plan: &LogicalPlan) -> Result<()> {
114114
pub fn assert_expected_schema(schema: &DFSchemaRef, plan: &LogicalPlan) -> Result<()> {
115115
let compatible = plan.schema().logically_equivalent_names_and_types(schema);
116116

117-
if !compatible {
118-
internal_err!(
119-
"Failed due to a difference in schemas: original schema: {:?}, new schema: {:?}",
120-
schema,
121-
plan.schema()
122-
)
123-
} else {
124-
Ok(())
125-
}
117+
assert_or_internal_err!(
118+
compatible,
119+
"Failed due to a difference in schemas: original schema: {:?}, new schema: {:?}",
120+
schema,
121+
plan.schema()
122+
);
123+
Ok(())
126124
}
127125

128126
/// Asserts that the subqueries are structured properly with valid node placement.

datafusion/expr/src/logical_plan/plan.rs

Lines changed: 39 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -59,10 +59,10 @@ use datafusion_common::tree_node::{
5959
Transformed, TreeNode, TreeNodeContainer, TreeNodeRecursion,
6060
};
6161
use datafusion_common::{
62-
aggregate_functional_dependencies, internal_err, plan_err, Column, Constraints,
63-
DFSchema, DFSchemaRef, DataFusionError, Dependency, FunctionalDependence,
64-
FunctionalDependencies, NullEquality, ParamValues, Result, ScalarValue, Spans,
65-
TableReference, UnnestOptions,
62+
aggregate_functional_dependencies, assert_eq_or_internal_err, assert_or_internal_err,
63+
internal_err, plan_err, Column, Constraints, DFSchema, DFSchemaRef, DataFusionError,
64+
Dependency, FunctionalDependence, FunctionalDependencies, NullEquality, ParamValues,
65+
Result, ScalarValue, Spans, TableReference, UnnestOptions,
6666
};
6767
use indexmap::IndexSet;
6868

@@ -965,13 +965,13 @@ impl LogicalPlan {
965965
}
966966
LogicalPlan::Limit(Limit { skip, fetch, .. }) => {
967967
let old_expr_len = skip.iter().chain(fetch.iter()).count();
968-
if old_expr_len != expr.len() {
969-
return internal_err!(
970-
"Invalid number of new Limit expressions: expected {}, got {}",
971-
old_expr_len,
972-
expr.len()
973-
);
974-
}
968+
assert_eq_or_internal_err!(
969+
old_expr_len,
970+
expr.len(),
971+
"Invalid number of new Limit expressions: expected {}, got {}",
972+
old_expr_len,
973+
expr.len()
974+
);
975975
// `LogicalPlan::expressions()` returns in [skip, fetch] order, so we can pop from the end.
976976
let new_fetch = fetch.as_ref().and_then(|_| expr.pop());
977977
let new_skip = skip.as_ref().and_then(|_| expr.pop());
@@ -1158,43 +1158,47 @@ impl LogicalPlan {
11581158
#[inline]
11591159
#[expect(clippy::needless_pass_by_value)] // expr is moved intentionally to ensure it's not used again
11601160
fn assert_no_expressions(&self, expr: Vec<Expr>) -> Result<()> {
1161-
if !expr.is_empty() {
1162-
return internal_err!("{self:?} should have no exprs, got {:?}", expr);
1163-
}
1161+
assert_or_internal_err!(
1162+
expr.is_empty(),
1163+
"{self:?} should have no exprs, got {:?}",
1164+
expr
1165+
);
11641166
Ok(())
11651167
}
11661168

11671169
/// Helper for [Self::with_new_exprs] to use when no inputs are expected.
11681170
#[inline]
11691171
#[expect(clippy::needless_pass_by_value)] // inputs is moved intentionally to ensure it's not used again
11701172
fn assert_no_inputs(&self, inputs: Vec<LogicalPlan>) -> Result<()> {
1171-
if !inputs.is_empty() {
1172-
return internal_err!("{self:?} should have no inputs, got: {:?}", inputs);
1173-
}
1173+
assert_or_internal_err!(
1174+
inputs.is_empty(),
1175+
"{self:?} should have no inputs, got: {:?}",
1176+
inputs
1177+
);
11741178
Ok(())
11751179
}
11761180

11771181
/// Helper for [Self::with_new_exprs] to use when exactly one expression is expected.
11781182
#[inline]
11791183
fn only_expr(&self, mut expr: Vec<Expr>) -> Result<Expr> {
1180-
if expr.len() != 1 {
1181-
return internal_err!(
1182-
"{self:?} should have exactly one expr, got {:?}",
1183-
expr
1184-
);
1185-
}
1184+
assert_eq_or_internal_err!(
1185+
expr.len(),
1186+
1,
1187+
"{self:?} should have exactly one expr, got {:?}",
1188+
&expr
1189+
);
11861190
Ok(expr.remove(0))
11871191
}
11881192

11891193
/// Helper for [Self::with_new_exprs] to use when exactly one input is expected.
11901194
#[inline]
11911195
fn only_input(&self, mut inputs: Vec<LogicalPlan>) -> Result<LogicalPlan> {
1192-
if inputs.len() != 1 {
1193-
return internal_err!(
1194-
"{self:?} should have exactly one input, got {:?}",
1195-
inputs
1196-
);
1197-
}
1196+
assert_eq_or_internal_err!(
1197+
inputs.len(),
1198+
1,
1199+
"{self:?} should have exactly one input, got {:?}",
1200+
&inputs
1201+
);
11981202
Ok(inputs.remove(0))
11991203
}
12001204

@@ -1204,12 +1208,12 @@ impl LogicalPlan {
12041208
&self,
12051209
mut inputs: Vec<LogicalPlan>,
12061210
) -> Result<(LogicalPlan, LogicalPlan)> {
1207-
if inputs.len() != 2 {
1208-
return internal_err!(
1209-
"{self:?} should have exactly two inputs, got {:?}",
1210-
inputs
1211-
);
1212-
}
1211+
assert_eq_or_internal_err!(
1212+
inputs.len(),
1213+
2,
1214+
"{self:?} should have exactly two inputs, got {:?}",
1215+
&inputs
1216+
);
12131217
let right = inputs.remove(1);
12141218
let left = inputs.remove(0);
12151219
Ok((left, right))

datafusion/expr/src/udf.rs

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,10 @@ use crate::udf_eq::UdfEq;
2525
use crate::{ColumnarValue, Documentation, Expr, Signature};
2626
use arrow::datatypes::{DataType, Field, FieldRef};
2727
use datafusion_common::config::ConfigOptions;
28-
use datafusion_common::{not_impl_err, ExprSchema, Result, ScalarValue};
28+
use datafusion_common::{
29+
assert_or_internal_err, not_impl_err, DataFusionError, ExprSchema, Result,
30+
ScalarValue,
31+
};
2932
use datafusion_expr_common::dyn_eq::{DynEq, DynHash};
3033
use datafusion_expr_common::interval_arithmetic::Interval;
3134
use std::any::Any;
@@ -240,13 +243,15 @@ impl ScalarUDF {
240243
// This doesn't use debug_assert!, but it's meant to run anywhere except on production. It's same in spirit, thus conditioning on debug_assertions.
241244
#[cfg(debug_assertions)]
242245
{
243-
if &result.data_type() != return_field.data_type() {
244-
return datafusion_common::internal_err!("Function '{}' returned value of type '{:?}' while the following type was promised at planning time and expected: '{:?}'",
245-
self.name(),
246-
result.data_type(),
247-
return_field.data_type()
248-
);
249-
}
246+
let result_data_type = result.data_type();
247+
let expected_type = return_field.data_type();
248+
assert_or_internal_err!(
249+
result_data_type == *expected_type,
250+
"Function '{}' returned value of type '{:?}' while the following type was promised at planning time and expected: '{:?}'",
251+
self.name(),
252+
result_data_type,
253+
expected_type
254+
);
250255
// TODO verify return data is non-null when it was promised to be?
251256
}
252257
Ok(result)

0 commit comments

Comments
 (0)