Fix: handle NULL input in lead/lag window function

apache · Oct 9, 2024 · dee4261 · dee4261
1 parent e00af2c
commit dee4261
Show file tree

Hide file tree

Showing 2 changed files with 81 additions and 9 deletions.
diff --git a/datafusion/physical-plan/src/windows/mod.rs b/datafusion/physical-plan/src/windows/mod.rs
@@ -257,31 +257,55 @@ fn create_built_in_window_expr(
             }
         }
         BuiltInWindowFunction::Lag => {
-            let arg = Arc::clone(&args[0]);
+            let mut arg = Arc::clone(&args[0]);
             let shift_offset = get_scalar_value_from_args(args, 1)?
                 .map(get_signed_integer)
                 .map_or(Ok(None), |v| v.map(Some))?;
-            let default_value =
-                get_casted_value(get_scalar_value_from_args(args, 2)?, out_data_type)?;
+            // If value is NULL, we use default data type as output data type, no need to cast data type
+            let default_value = match out_data_type {
+                DataType::Null => match get_scalar_value_from_args(args, 2)? {
+                    Some(value) => {
+                        let null_value = ScalarValue::try_from(value.data_type())?;
+                        arg = Arc::new(Literal::new(null_value));
+                        value
+                    }
+                    None => ScalarValue::try_from(DataType::Null)?,
+                },
+                _ => {
+                    get_casted_value(get_scalar_value_from_args(args, 2)?, out_data_type)?
+                }
+            };
             Arc::new(lag(
                 name,
-                out_data_type.clone(),
+                default_value.data_type().clone(),
                 arg,
                 shift_offset,
                 default_value,
                 ignore_nulls,
             ))
         }
         BuiltInWindowFunction::Lead => {
-            let arg = Arc::clone(&args[0]);
+            let mut arg = Arc::clone(&args[0]);
             let shift_offset = get_scalar_value_from_args(args, 1)?
                 .map(get_signed_integer)
                 .map_or(Ok(None), |v| v.map(Some))?;
-            let default_value =
-                get_casted_value(get_scalar_value_from_args(args, 2)?, out_data_type)?;
+            // If value is NULL, we use default data type as output data type, no need to cast data type
+            let default_value = match out_data_type {
+                DataType::Null => match get_scalar_value_from_args(args, 2)? {
+                    Some(value) => {
+                        let null_value = ScalarValue::try_from(value.data_type())?;
+                        arg = Arc::new(Literal::new(null_value));
+                        value
+                    }
+                    None => ScalarValue::try_from(DataType::Null)?,
+                },
+                _ => {
+                    get_casted_value(get_scalar_value_from_args(args, 2)?, out_data_type)?
+                }
+            };
             Arc::new(lead(
                 name,
-                out_data_type.clone(),
+                default_value.data_type().clone(),
                 arg,
                 shift_offset,
                 default_value,

diff --git a/datafusion/sqllogictest/test_files/window.slt b/datafusion/sqllogictest/test_files/window.slt
@@ -4932,4 +4932,52 @@ SELECT v1, NTH_VALUE(v2, 0) OVER (PARTITION BY v1 ORDER BY v2) FROM t;
 statement ok
 DROP TABLE t;
 
-## end test handle NULL and 0 of NTH_VALUE
+## end test handle NULL and 0 of NTH_VALUE
+
+## test handle NULL of lead
+
+statement ok
+create table t1(v1 int);
+
+statement ok
+insert into t1 values (1);
+
+query B
+SELECT LEAD(NULL, 0, false) OVER () FROM t1;
+----
+NULL
+
+query B
+SELECT LAG(NULL, 0, false) OVER () FROM t1;
+----
+NULL
+
+query B
+SELECT LEAD(NULL, 1, false) OVER () FROM t1;
+----
+false
+
+query B
+SELECT LAG(NULL, 1, false) OVER () FROM t1;
+----
+false
+
+statement ok
+insert into t1 values (2);
+
+query B
+SELECT LEAD(NULL, 1, false) OVER () FROM t1;
+----
+NULL
+false
+
+query B
+SELECT LAG(NULL, 1, false) OVER () FROM t1;
+----
+false
+NULL
+
+statement ok
+DROP TABLE t1;
+
+## end test handle NULL of lead