From 213f1e466fe0f8affe37bcf2f890568838373d95 Mon Sep 17 00:00:00 2001 From: Thinh Bui Date: Sun, 29 Sep 2024 17:52:49 -0700 Subject: [PATCH 1/3] handle 0 and NULL value of NTH_VALUE function --- .../physical-expr/src/window/nth_value.rs | 29 +++++--------- datafusion/physical-plan/src/windows/mod.rs | 10 +++++ datafusion/sqllogictest/test_files/window.slt | 39 +++++++++++++++++++ 3 files changed, 59 insertions(+), 19 deletions(-) diff --git a/datafusion/physical-expr/src/window/nth_value.rs b/datafusion/physical-expr/src/window/nth_value.rs index 87c74579c639..d94983c5adf7 100644 --- a/datafusion/physical-expr/src/window/nth_value.rs +++ b/datafusion/physical-expr/src/window/nth_value.rs @@ -30,7 +30,7 @@ use crate::PhysicalExpr; use arrow::array::{Array, ArrayRef}; use arrow::datatypes::{DataType, Field}; use datafusion_common::Result; -use datafusion_common::{exec_err, ScalarValue}; +use datafusion_common::ScalarValue; use datafusion_expr::window_state::WindowAggState; use datafusion_expr::PartitionEvaluator; @@ -86,16 +86,13 @@ impl NthValue { n: i64, ignore_nulls: bool, ) -> Result { - match n { - 0 => exec_err!("NTH_VALUE expects n to be non-zero"), - _ => Ok(Self { - name: name.into(), - expr, - data_type, - kind: NthValueKind::Nth(n), - ignore_nulls, - }), - } + Ok(Self { + name: name.into(), + expr, + data_type, + kind: NthValueKind::Nth(n), + ignore_nulls, + }) } /// Get the NTH_VALUE kind @@ -188,10 +185,7 @@ impl PartitionEvaluator for NthValueEvaluator { // Negative index represents reverse direction. (n_range >= reverse_index, true) } - Ordering::Equal => { - // The case n = 0 is not valid for the NTH_VALUE function. - unreachable!(); - } + Ordering::Equal => (true, false), } } }; @@ -298,10 +292,7 @@ impl PartitionEvaluator for NthValueEvaluator { ) } } - Ordering::Equal => { - // The case n = 0 is not valid for the NTH_VALUE function. - unreachable!(); - } + Ordering::Equal => ScalarValue::try_from(arr.data_type()), } } } diff --git a/datafusion/physical-plan/src/windows/mod.rs b/datafusion/physical-plan/src/windows/mod.rs index b6f34ec69f68..eaccf80f0790 100644 --- a/datafusion/physical-plan/src/windows/mod.rs +++ b/datafusion/physical-plan/src/windows/mod.rs @@ -185,20 +185,30 @@ fn get_scalar_value_from_args( } fn get_signed_integer(value: ScalarValue) -> Result { + if value.is_null() { + return Ok(0); + } + if !value.data_type().is_integer() { return Err(DataFusionError::Execution( "Expected an integer value".to_string(), )); } + value.cast_to(&DataType::Int64)?.try_into() } fn get_unsigned_integer(value: ScalarValue) -> Result { + if value.is_null() { + return Ok(0); + } + if !value.data_type().is_integer() { return Err(DataFusionError::Execution( "Expected an integer value".to_string(), )); } + value.cast_to(&DataType::UInt64)?.try_into() } diff --git a/datafusion/sqllogictest/test_files/window.slt b/datafusion/sqllogictest/test_files/window.slt index 7fee84f9bcd9..cb6c6a5ace76 100644 --- a/datafusion/sqllogictest/test_files/window.slt +++ b/datafusion/sqllogictest/test_files/window.slt @@ -4894,3 +4894,42 @@ NULL a4 5 statement ok drop table t + +## test handle NULL and 0 value of nth_value +statement ok +CREATE TABLE t(v1 int, v2 int); + +statement ok +INSERT INTO t VALUES (1,1), (1,2),(1,3),(2,1),(2,2); + +query II +SELECT v1, NTH_VALUE(v2, null) OVER (PARTITION BY v1 ORDER BY v2) FROM t; +---- +1 NULL +1 NULL +1 NULL +2 NULL +2 NULL + +query II +SELECT v1, NTH_VALUE(v2, v2*null) OVER (PARTITION BY v1 ORDER BY v2) FROM t; +---- +1 NULL +1 NULL +1 NULL +2 NULL +2 NULL + +query II +SELECT v1, NTH_VALUE(v2, 0) OVER (PARTITION BY v1 ORDER BY v2) FROM t; +---- +1 NULL +1 NULL +1 NULL +2 NULL +2 NULL + +statement ok +DROP TABLE t; + +## end test handle NULL and 0 of NTH_VALUE \ No newline at end of file From 9bd8b5ca351533e6a5ab2dc0782c5df7af5c016b Mon Sep 17 00:00:00 2001 From: Thinh Bui Date: Mon, 30 Sep 2024 11:36:38 -0700 Subject: [PATCH 2/3] use exec_err --- datafusion/physical-plan/src/windows/mod.rs | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/datafusion/physical-plan/src/windows/mod.rs b/datafusion/physical-plan/src/windows/mod.rs index eaccf80f0790..f3ba923ab6aa 100644 --- a/datafusion/physical-plan/src/windows/mod.rs +++ b/datafusion/physical-plan/src/windows/mod.rs @@ -190,9 +190,7 @@ fn get_signed_integer(value: ScalarValue) -> Result { } if !value.data_type().is_integer() { - return Err(DataFusionError::Execution( - "Expected an integer value".to_string(), - )); + return exec_err!("Expected an integer value") } value.cast_to(&DataType::Int64)?.try_into() @@ -204,9 +202,7 @@ fn get_unsigned_integer(value: ScalarValue) -> Result { } if !value.data_type().is_integer() { - return Err(DataFusionError::Execution( - "Expected an integer value".to_string(), - )); + return exec_err!("Expected an integer value"); } value.cast_to(&DataType::UInt64)?.try_into() From a3c6b65bd05d5b35cda329d94d39d69ae57d6853 Mon Sep 17 00:00:00 2001 From: Thinh Bui Date: Mon, 30 Sep 2024 11:38:17 -0700 Subject: [PATCH 3/3] cargo fmt --- datafusion/physical-plan/src/windows/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/windows/mod.rs b/datafusion/physical-plan/src/windows/mod.rs index f3ba923ab6aa..6aafaad0ad77 100644 --- a/datafusion/physical-plan/src/windows/mod.rs +++ b/datafusion/physical-plan/src/windows/mod.rs @@ -190,7 +190,7 @@ fn get_signed_integer(value: ScalarValue) -> Result { } if !value.data_type().is_integer() { - return exec_err!("Expected an integer value") + return exec_err!("Expected an integer value"); } value.cast_to(&DataType::Int64)?.try_into()