Skip to content

Commit 3b1738c

Browse files
committed
Made the nvl2 scalar UDF return an internal error when it reaches execution without being simplified to a CASE expression, removing the eager evaluation helpers that previously enforced eager semantics.
Updated the expr_api integration test to assert that unsimplified nvl2 evaluation now fails with the expected internal error message.
1 parent 76c5121 commit 3b1738c

File tree

2 files changed

+18
-58
lines changed

2 files changed

+18
-58
lines changed

datafusion/core/tests/expr_api/mod.rs

Lines changed: 15 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -322,33 +322,22 @@ async fn test_create_physical_expr() {
322322

323323
#[test]
324324
fn test_create_physical_expr_nvl2() {
325-
#[rustfmt::skip]
326-
evaluate_expr_test(
327-
nvl2(col("i"), lit(1i64), lit(0i64)),
328-
vec![
329-
"+------+",
330-
"| expr |",
331-
"+------+",
332-
"| 1 |",
333-
"| 0 |",
334-
"| 1 |",
335-
"+------+",
336-
],
337-
);
325+
let batch = &TEST_BATCH;
326+
let df_schema = DFSchema::try_from(batch.schema()).unwrap();
327+
let ctx = SessionContext::new();
338328

339-
#[rustfmt::skip]
340-
evaluate_expr_test(
341-
nvl2(lit(1i64), col("i"), lit(0i64)),
342-
vec![
343-
"+------+",
344-
"| expr |",
345-
"+------+",
346-
"| 10 |",
347-
"| |",
348-
"| 5 |",
349-
"+------+",
350-
],
351-
);
329+
let expect_err = |expr| {
330+
let physical_expr = ctx.create_physical_expr(expr, &df_schema).unwrap();
331+
let err = physical_expr.evaluate(batch).unwrap_err();
332+
assert!(
333+
err.to_string()
334+
.contains("nvl2 should have been simplified to case"),
335+
"unexpected error: {err:?}"
336+
);
337+
};
338+
339+
expect_err(nvl2(col("i"), lit(1i64), lit(0i64)));
340+
expect_err(nvl2(lit(1i64), col("i"), lit(0i64)));
352341
}
353342

354343
#[tokio::test]

datafusion/functions/src/core/nvl2.rs

Lines changed: 3 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -15,10 +15,8 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18-
use arrow::compute::is_not_null;
19-
use arrow::compute::kernels::zip::zip;
2018
use arrow::datatypes::{DataType, Field, FieldRef};
21-
use datafusion_common::{internal_err, plan_err, utils::take_function_args, Result};
19+
use datafusion_common::{internal_err, utils::take_function_args, Result};
2220
use datafusion_expr::{
2321
conditional_expressions::CaseBuilder,
2422
simplify::{ExprSimplifyResult, SimplifyInfo},
@@ -103,35 +101,8 @@ impl ScalarUDFImpl for NVL2Func {
103101
Ok(Field::new(self.name(), return_type, nullable).into())
104102
}
105103

106-
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
107-
let [test, if_non_null, if_null] = take_function_args(self.name(), args.args)?;
108-
109-
match test {
110-
ColumnarValue::Scalar(test_scalar) => {
111-
if test_scalar.is_null() {
112-
Ok(if_null)
113-
} else {
114-
Ok(if_non_null)
115-
}
116-
}
117-
ColumnarValue::Array(test_array) => {
118-
let len = test_array.len();
119-
120-
let if_non_null_array = match if_non_null {
121-
ColumnarValue::Array(array) => array,
122-
ColumnarValue::Scalar(scalar) => scalar.to_array_of_size(len)?,
123-
};
124-
125-
let if_null_array = match if_null {
126-
ColumnarValue::Array(array) => array,
127-
ColumnarValue::Scalar(scalar) => scalar.to_array_of_size(len)?,
128-
};
129-
130-
let mask = is_not_null(&test_array)?;
131-
let result = zip(&mask, &if_non_null_array, &if_null_array)?;
132-
Ok(ColumnarValue::Array(result))
133-
}
134-
}
104+
fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> {
105+
internal_err!("nvl2 should have been simplified to case")
135106
}
136107

137108
fn simplify(

0 commit comments

Comments
 (0)