From 000ceb76409e66230f9c5017a30fa3c9bb1e6575 Mon Sep 17 00:00:00 2001
From: jcsherin <jacob@protoship.io>
Date: Thu, 10 Oct 2024 15:55:33 +0530
Subject: [PATCH] Refactors shift offset computation

---
 datafusion/functions-window/src/lead_lag.rs | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/datafusion/functions-window/src/lead_lag.rs b/datafusion/functions-window/src/lead_lag.rs
index 434cf64013dd..9f9b7125bd92 100644
--- a/datafusion/functions-window/src/lead_lag.rs
+++ b/datafusion/functions-window/src/lead_lag.rs
@@ -88,10 +88,10 @@ impl WindowShiftKind {
         }
     }
 
-    fn shift_offset(&self, value: Option<i64>) -> i64 {
+    fn shift_offset(&self, value: i64) -> i64 {
         match self {
-            WindowShiftKind::Lag => value.unwrap_or(1),
-            WindowShiftKind::Lead => value.map(|v| v.neg()).unwrap_or(-1),
+            WindowShiftKind::Lag => value,
+            WindowShiftKind::Lead => value.neg(),
         }
     }
 }
@@ -146,11 +146,11 @@ impl WindowUDFImpl for WindowShift {
     ) -> Result<Box<dyn PartitionEvaluator>> {
         let shift_offset = scalar_at(&partition_evaluator_args, 1)?
             .map(get_signed_integer)
-            .map_or(Ok(None), |v| v.map(Some))
+            .unwrap_or(Ok(1))
             .map(|n| self.kind.shift_offset(n))
             .map(|offset| {
                 if partition_evaluator_args.is_reversed() {
-                    -offset
+                    offset.neg()
                 } else {
                     offset
                 }