1use crate::utils::{get_scalar_value_from_args, get_signed_integer};
21use arrow::datatypes::FieldRef;
22use datafusion_common::arrow::array::ArrayRef;
23use datafusion_common::arrow::datatypes::DataType;
24use datafusion_common::arrow::datatypes::Field;
25use datafusion_common::{arrow_datafusion_err, DataFusionError, Result, ScalarValue};
26use datafusion_expr::window_doc_sections::DOC_SECTION_ANALYTICAL;
27use datafusion_expr::{
28 Documentation, Literal, PartitionEvaluator, ReversedUDWF, Signature, TypeSignature,
29 Volatility, WindowUDFImpl,
30};
31use datafusion_functions_window_common::expr::ExpressionArgs;
32use datafusion_functions_window_common::field::WindowUDFFieldArgs;
33use datafusion_functions_window_common::partition::PartitionEvaluatorArgs;
34use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
35use std::any::Any;
36use std::cmp::min;
37use std::collections::VecDeque;
38use std::hash::Hash;
39use std::ops::{Neg, Range};
40use std::sync::{Arc, LazyLock};
41
42get_or_init_udwf!(
43 Lag,
44 lag,
45 "Returns the row value that precedes the current row by a specified \
46 offset within partition. If no such row exists, then returns the \
47 default value.",
48 WindowShift::lag
49);
50get_or_init_udwf!(
51 Lead,
52 lead,
53 "Returns the value from a row that follows the current row by a \
54 specified offset within the partition. If no such row exists, then \
55 returns the default value.",
56 WindowShift::lead
57);
58
59pub fn lag(
66 arg: datafusion_expr::Expr,
67 shift_offset: Option<i64>,
68 default_value: Option<ScalarValue>,
69) -> datafusion_expr::Expr {
70 let shift_offset_lit = shift_offset
71 .map(|v| v.lit())
72 .unwrap_or(ScalarValue::Null.lit());
73 let default_lit = default_value.unwrap_or(ScalarValue::Null).lit();
74
75 lag_udwf().call(vec![arg, shift_offset_lit, default_lit])
76}
77
78pub fn lead(
85 arg: datafusion_expr::Expr,
86 shift_offset: Option<i64>,
87 default_value: Option<ScalarValue>,
88) -> datafusion_expr::Expr {
89 let shift_offset_lit = shift_offset
90 .map(|v| v.lit())
91 .unwrap_or(ScalarValue::Null.lit());
92 let default_lit = default_value.unwrap_or(ScalarValue::Null).lit();
93
94 lead_udwf().call(vec![arg, shift_offset_lit, default_lit])
95}
96
97#[derive(Debug, PartialEq, Eq, Hash)]
98enum WindowShiftKind {
99 Lag,
100 Lead,
101}
102
103impl WindowShiftKind {
104 fn name(&self) -> &'static str {
105 match self {
106 WindowShiftKind::Lag => "lag",
107 WindowShiftKind::Lead => "lead",
108 }
109 }
110
111 fn shift_offset(&self, value: Option<i64>) -> i64 {
115 match self {
116 WindowShiftKind::Lag => value.unwrap_or(1),
117 WindowShiftKind::Lead => value.map(|v| v.neg()).unwrap_or(-1),
118 }
119 }
120}
121
122#[derive(Debug, PartialEq, Eq, Hash)]
124pub struct WindowShift {
125 signature: Signature,
126 kind: WindowShiftKind,
127}
128
129impl WindowShift {
130 fn new(kind: WindowShiftKind) -> Self {
131 Self {
132 signature: Signature::one_of(
133 vec![
134 TypeSignature::Any(1),
135 TypeSignature::Any(2),
136 TypeSignature::Any(3),
137 ],
138 Volatility::Immutable,
139 ),
140 kind,
141 }
142 }
143
144 pub fn lag() -> Self {
145 Self::new(WindowShiftKind::Lag)
146 }
147
148 pub fn lead() -> Self {
149 Self::new(WindowShiftKind::Lead)
150 }
151}
152
153static LAG_DOCUMENTATION: LazyLock<Documentation> = LazyLock::new(|| {
154 Documentation::builder(DOC_SECTION_ANALYTICAL, "Returns value evaluated at the row that is offset rows before the \
155 current row within the partition; if there is no such row, instead return default \
156 (which must be of the same type as value).", "lag(expression, offset, default)")
157 .with_argument("expression", "Expression to operate on")
158 .with_argument("offset", "Integer. Specifies how many rows back \
159 the value of expression should be retrieved. Defaults to 1.")
160 .with_argument("default", "The default value if the offset is \
161 not within the partition. Must be of the same type as expression.")
162 .with_sql_example(r#"
163```sql
164-- Example usage of the lag window function:
165SELECT employee_id,
166 salary,
167 lag(salary, 1, 0) OVER (ORDER BY employee_id) AS prev_salary
168FROM employees;
169
170+-------------+--------+-------------+
171| employee_id | salary | prev_salary |
172+-------------+--------+-------------+
173| 1 | 30000 | 0 |
174| 2 | 50000 | 30000 |
175| 3 | 70000 | 50000 |
176| 4 | 60000 | 70000 |
177+-------------+--------+-------------+
178```
179"#)
180 .build()
181});
182
183fn get_lag_doc() -> &'static Documentation {
184 &LAG_DOCUMENTATION
185}
186
187static LEAD_DOCUMENTATION: LazyLock<Documentation> = LazyLock::new(|| {
188 Documentation::builder(DOC_SECTION_ANALYTICAL,
189 "Returns value evaluated at the row that is offset rows after the \
190 current row within the partition; if there is no such row, instead return default \
191 (which must be of the same type as value).",
192 "lead(expression, offset, default)")
193 .with_argument("expression", "Expression to operate on")
194 .with_argument("offset", "Integer. Specifies how many rows \
195 forward the value of expression should be retrieved. Defaults to 1.")
196 .with_argument("default", "The default value if the offset is \
197 not within the partition. Must be of the same type as expression.")
198 .with_sql_example(r#"
199```sql
200-- Example usage of lead window function:
201SELECT
202 employee_id,
203 department,
204 salary,
205 lead(salary, 1, 0) OVER (PARTITION BY department ORDER BY salary) AS next_salary
206FROM employees;
207
208+-------------+-------------+--------+--------------+
209| employee_id | department | salary | next_salary |
210+-------------+-------------+--------+--------------+
211| 1 | Sales | 30000 | 50000 |
212| 2 | Sales | 50000 | 70000 |
213| 3 | Sales | 70000 | 0 |
214| 4 | Engineering | 40000 | 60000 |
215| 5 | Engineering | 60000 | 0 |
216+-------------+-------------+--------+--------------+
217```
218"#)
219 .build()
220});
221
222fn get_lead_doc() -> &'static Documentation {
223 &LEAD_DOCUMENTATION
224}
225
226impl WindowUDFImpl for WindowShift {
227 fn as_any(&self) -> &dyn Any {
228 self
229 }
230
231 fn name(&self) -> &str {
232 self.kind.name()
233 }
234
235 fn signature(&self) -> &Signature {
236 &self.signature
237 }
238
239 fn expressions(&self, expr_args: ExpressionArgs) -> Vec<Arc<dyn PhysicalExpr>> {
245 parse_expr(expr_args.input_exprs(), expr_args.input_fields())
246 .into_iter()
247 .collect::<Vec<_>>()
248 }
249
250 fn partition_evaluator(
251 &self,
252 partition_evaluator_args: PartitionEvaluatorArgs,
253 ) -> Result<Box<dyn PartitionEvaluator>> {
254 let shift_offset =
255 get_scalar_value_from_args(partition_evaluator_args.input_exprs(), 1)?
256 .map(get_signed_integer)
257 .map_or(Ok(None), |v| v.map(Some))
258 .map(|n| self.kind.shift_offset(n))
259 .map(|offset| {
260 if partition_evaluator_args.is_reversed() {
261 -offset
262 } else {
263 offset
264 }
265 })?;
266 let default_value = parse_default_value(
267 partition_evaluator_args.input_exprs(),
268 partition_evaluator_args.input_fields(),
269 )?;
270
271 Ok(Box::new(WindowShiftEvaluator {
272 shift_offset,
273 default_value,
274 ignore_nulls: partition_evaluator_args.ignore_nulls(),
275 non_null_offsets: VecDeque::new(),
276 }))
277 }
278
279 fn field(&self, field_args: WindowUDFFieldArgs) -> Result<FieldRef> {
280 let return_field = parse_expr_field(field_args.input_fields())?;
281
282 Ok(return_field
283 .as_ref()
284 .clone()
285 .with_name(field_args.name())
286 .into())
287 }
288
289 fn reverse_expr(&self) -> ReversedUDWF {
290 match self.kind {
291 WindowShiftKind::Lag => ReversedUDWF::Reversed(lag_udwf()),
292 WindowShiftKind::Lead => ReversedUDWF::Reversed(lead_udwf()),
293 }
294 }
295
296 fn documentation(&self) -> Option<&Documentation> {
297 match self.kind {
298 WindowShiftKind::Lag => Some(get_lag_doc()),
299 WindowShiftKind::Lead => Some(get_lead_doc()),
300 }
301 }
302}
303
304fn parse_expr(
317 input_exprs: &[Arc<dyn PhysicalExpr>],
318 input_fields: &[FieldRef],
319) -> Result<Arc<dyn PhysicalExpr>> {
320 assert!(!input_exprs.is_empty());
321 assert!(!input_fields.is_empty());
322
323 let expr = Arc::clone(input_exprs.first().unwrap());
324 let expr_field = input_fields.first().unwrap();
325
326 if !expr_field.data_type().is_null() {
328 return Ok(expr);
329 }
330
331 let default_value = get_scalar_value_from_args(input_exprs, 2)?;
332 default_value.map_or(Ok(expr), |value| {
333 ScalarValue::try_from(&value.data_type()).map(|v| {
334 Arc::new(datafusion_physical_expr::expressions::Literal::new(v))
335 as Arc<dyn PhysicalExpr>
336 })
337 })
338}
339
340static NULL_FIELD: LazyLock<FieldRef> =
341 LazyLock::new(|| Field::new("value", DataType::Null, true).into());
342
343fn parse_expr_field(input_fields: &[FieldRef]) -> Result<FieldRef> {
348 assert!(!input_fields.is_empty());
349 let expr_field = input_fields.first().unwrap_or(&NULL_FIELD);
350
351 if !expr_field.data_type().is_null() {
353 return Ok(expr_field.as_ref().clone().with_nullable(true).into());
354 }
355
356 let default_value_field = input_fields.get(2).unwrap_or(&NULL_FIELD);
357 Ok(default_value_field
358 .as_ref()
359 .clone()
360 .with_nullable(true)
361 .into())
362}
363
364fn parse_default_value(
367 input_exprs: &[Arc<dyn PhysicalExpr>],
368 input_types: &[FieldRef],
369) -> Result<ScalarValue> {
370 let expr_field = parse_expr_field(input_types)?;
371 let unparsed = get_scalar_value_from_args(input_exprs, 2)?;
372
373 unparsed
374 .filter(|v| !v.data_type().is_null())
375 .map(|v| v.cast_to(expr_field.data_type()))
376 .unwrap_or_else(|| ScalarValue::try_from(expr_field.data_type()))
377}
378
379#[derive(Debug)]
380struct WindowShiftEvaluator {
381 shift_offset: i64,
382 default_value: ScalarValue,
383 ignore_nulls: bool,
384 non_null_offsets: VecDeque<usize>,
386}
387
388impl WindowShiftEvaluator {
389 fn is_lag(&self) -> bool {
390 self.shift_offset > 0
392 }
393}
394
395fn evaluate_all_with_ignore_null(
397 array: &ArrayRef,
398 offset: i64,
399 default_value: &ScalarValue,
400 is_lag: bool,
401) -> Result<ArrayRef, DataFusionError> {
402 let valid_indices: Vec<usize> =
403 array.nulls().unwrap().valid_indices().collect::<Vec<_>>();
404 let direction = !is_lag;
405 let new_array_results: Result<Vec<_>, DataFusionError> = (0..array.len())
406 .map(|id| {
407 let result_index = match valid_indices.binary_search(&id) {
408 Ok(pos) => if direction {
409 pos.checked_add(offset as usize)
410 } else {
411 pos.checked_sub(offset.unsigned_abs() as usize)
412 }
413 .and_then(|new_pos| {
414 if new_pos < valid_indices.len() {
415 Some(valid_indices[new_pos])
416 } else {
417 None
418 }
419 }),
420 Err(pos) => if direction {
421 pos.checked_add(offset as usize)
422 } else if pos > 0 {
423 pos.checked_sub(offset.unsigned_abs() as usize)
424 } else {
425 None
426 }
427 .and_then(|new_pos| {
428 if new_pos < valid_indices.len() {
429 Some(valid_indices[new_pos])
430 } else {
431 None
432 }
433 }),
434 };
435
436 match result_index {
437 Some(index) => ScalarValue::try_from_array(array, index),
438 None => Ok(default_value.clone()),
439 }
440 })
441 .collect();
442
443 let new_array = new_array_results?;
444 ScalarValue::iter_to_array(new_array)
445}
446fn shift_with_default_value(
448 array: &ArrayRef,
449 offset: i64,
450 default_value: &ScalarValue,
451) -> Result<ArrayRef> {
452 use datafusion_common::arrow::compute::concat;
453
454 let value_len = array.len() as i64;
455 if offset == 0 {
456 Ok(Arc::clone(array))
457 } else if offset == i64::MIN || offset.abs() >= value_len {
458 default_value.to_array_of_size(value_len as usize)
459 } else {
460 let slice_offset = (-offset).clamp(0, value_len) as usize;
461 let length = array.len() - offset.unsigned_abs() as usize;
462 let slice = array.slice(slice_offset, length);
463
464 let nulls = offset.unsigned_abs() as usize;
466 let default_values = default_value.to_array_of_size(nulls)?;
467
468 if offset > 0 {
470 concat(&[default_values.as_ref(), slice.as_ref()])
471 .map_err(|e| arrow_datafusion_err!(e))
472 } else {
473 concat(&[slice.as_ref(), default_values.as_ref()])
474 .map_err(|e| arrow_datafusion_err!(e))
475 }
476 }
477}
478
479impl PartitionEvaluator for WindowShiftEvaluator {
480 fn get_range(&self, idx: usize, n_rows: usize) -> Result<Range<usize>> {
481 if self.is_lag() {
482 let start = if self.non_null_offsets.len() == self.shift_offset as usize {
483 let offset: usize = self.non_null_offsets.iter().sum();
485 idx.saturating_sub(offset)
486 } else if !self.ignore_nulls {
487 let offset = self.shift_offset as usize;
488 idx.saturating_sub(offset)
489 } else {
490 0
491 };
492 let end = idx + 1;
493 Ok(Range { start, end })
494 } else {
495 let end = if self.non_null_offsets.len() == (-self.shift_offset) as usize {
496 let offset: usize = self.non_null_offsets.iter().sum();
498 min(idx + offset + 1, n_rows)
499 } else if !self.ignore_nulls {
500 let offset = (-self.shift_offset) as usize;
501 min(idx + offset, n_rows)
502 } else {
503 n_rows
504 };
505 Ok(Range { start: idx, end })
506 }
507 }
508
509 fn is_causal(&self) -> bool {
510 self.is_lag()
512 }
513
514 fn evaluate(
515 &mut self,
516 values: &[ArrayRef],
517 range: &Range<usize>,
518 ) -> Result<ScalarValue> {
519 let array = &values[0];
520 let len = array.len();
521
522 let i = if self.is_lag() {
524 (range.end as i64 - self.shift_offset - 1) as usize
525 } else {
526 (range.start as i64 - self.shift_offset) as usize
528 };
529
530 let mut idx: Option<usize> = if i < len { Some(i) } else { None };
531
532 if self.ignore_nulls && self.is_lag() {
535 idx = if self.non_null_offsets.len() == self.shift_offset as usize {
538 let total_offset: usize = self.non_null_offsets.iter().sum();
539 Some(range.end - 1 - total_offset)
540 } else {
541 None
542 };
543
544 if array.is_valid(range.end - 1) {
546 self.non_null_offsets.push_back(1);
548 if self.non_null_offsets.len() > self.shift_offset as usize {
549 self.non_null_offsets.pop_front();
551 }
552 } else if !self.non_null_offsets.is_empty() {
553 let end_idx = self.non_null_offsets.len() - 1;
555 self.non_null_offsets[end_idx] += 1;
556 }
557 } else if self.ignore_nulls && !self.is_lag() {
558 let non_null_row_count = (-self.shift_offset) as usize;
561
562 if self.non_null_offsets.is_empty() {
563 let mut offset_val = 1;
565 for idx in range.start + 1..range.end {
566 if array.is_valid(idx) {
567 self.non_null_offsets.push_back(offset_val);
568 offset_val = 1;
569 } else {
570 offset_val += 1;
571 }
572 if self.non_null_offsets.len() == non_null_row_count + 1 {
575 break;
576 }
577 }
578 } else if range.end < len && array.is_valid(range.end) {
579 if array.is_valid(range.end) {
581 self.non_null_offsets.push_back(1);
583 } else {
584 let last_idx = self.non_null_offsets.len() - 1;
586 self.non_null_offsets[last_idx] += 1;
587 }
588 }
589
590 idx = if self.non_null_offsets.len() >= non_null_row_count {
592 let total_offset: usize =
593 self.non_null_offsets.iter().take(non_null_row_count).sum();
594 Some(range.start + total_offset)
595 } else {
596 None
597 };
598 if !self.non_null_offsets.is_empty() {
601 self.non_null_offsets[0] -= 1;
602 if self.non_null_offsets[0] == 0 {
603 self.non_null_offsets.pop_front();
605 }
606 }
607 }
608
609 #[allow(clippy::unnecessary_unwrap)]
615 if !(idx.is_none() || (self.ignore_nulls && array.is_null(idx.unwrap()))) {
616 ScalarValue::try_from_array(array, idx.unwrap())
617 } else {
618 Ok(self.default_value.clone())
619 }
620 }
621
622 fn evaluate_all(
623 &mut self,
624 values: &[ArrayRef],
625 _num_rows: usize,
626 ) -> Result<ArrayRef> {
627 let value = &values[0];
629 if !self.ignore_nulls {
630 shift_with_default_value(value, self.shift_offset, &self.default_value)
631 } else {
632 evaluate_all_with_ignore_null(
633 value,
634 self.shift_offset,
635 &self.default_value,
636 self.is_lag(),
637 )
638 }
639 }
640
641 fn supports_bounded_execution(&self) -> bool {
642 true
643 }
644}
645
646#[cfg(test)]
647mod tests {
648 use super::*;
649 use arrow::array::*;
650 use datafusion_common::cast::as_int32_array;
651 use datafusion_physical_expr::expressions::{Column, Literal};
652 use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
653
654 fn test_i32_result(
655 expr: WindowShift,
656 partition_evaluator_args: PartitionEvaluatorArgs,
657 expected: Int32Array,
658 ) -> Result<()> {
659 let arr: ArrayRef = Arc::new(Int32Array::from(vec![1, -2, 3, -4, 5, -6, 7, 8]));
660 let values = vec![arr];
661 let num_rows = values.len();
662 let result = expr
663 .partition_evaluator(partition_evaluator_args)?
664 .evaluate_all(&values, num_rows)?;
665 let result = as_int32_array(&result)?;
666 assert_eq!(expected, *result);
667 Ok(())
668 }
669
670 #[test]
671 fn lead_lag_get_range() -> Result<()> {
672 let lag_fn = WindowShiftEvaluator {
674 shift_offset: 2,
675 default_value: ScalarValue::Null,
676 ignore_nulls: false,
677 non_null_offsets: Default::default(),
678 };
679 assert_eq!(lag_fn.get_range(6, 10)?, Range { start: 4, end: 7 });
680 assert_eq!(lag_fn.get_range(0, 10)?, Range { start: 0, end: 1 });
681
682 let lag_fn = WindowShiftEvaluator {
684 shift_offset: 2,
685 default_value: ScalarValue::Null,
686 ignore_nulls: true,
687 non_null_offsets: vec![2, 2].into(), };
690 assert_eq!(lag_fn.get_range(6, 10)?, Range { start: 2, end: 7 });
691
692 let lead_fn = WindowShiftEvaluator {
694 shift_offset: -2,
695 default_value: ScalarValue::Null,
696 ignore_nulls: false,
697 non_null_offsets: Default::default(),
698 };
699 assert_eq!(lead_fn.get_range(6, 10)?, Range { start: 6, end: 8 });
700 assert_eq!(lead_fn.get_range(9, 10)?, Range { start: 9, end: 10 });
701
702 let lead_fn = WindowShiftEvaluator {
704 shift_offset: -2,
705 default_value: ScalarValue::Null,
706 ignore_nulls: true,
707 non_null_offsets: vec![2, 2].into(),
709 };
710 assert_eq!(lead_fn.get_range(4, 10)?, Range { start: 4, end: 9 });
711
712 Ok(())
713 }
714
715 #[test]
716 fn test_lead_window_shift() -> Result<()> {
717 let expr = Arc::new(Column::new("c3", 0)) as Arc<dyn PhysicalExpr>;
718
719 test_i32_result(
720 WindowShift::lead(),
721 PartitionEvaluatorArgs::new(
722 &[expr],
723 &[Field::new("f", DataType::Int32, true).into()],
724 false,
725 false,
726 ),
727 [
728 Some(-2),
729 Some(3),
730 Some(-4),
731 Some(5),
732 Some(-6),
733 Some(7),
734 Some(8),
735 None,
736 ]
737 .iter()
738 .collect::<Int32Array>(),
739 )
740 }
741
742 #[test]
743 fn test_lag_window_shift() -> Result<()> {
744 let expr = Arc::new(Column::new("c3", 0)) as Arc<dyn PhysicalExpr>;
745
746 test_i32_result(
747 WindowShift::lag(),
748 PartitionEvaluatorArgs::new(
749 &[expr],
750 &[Field::new("f", DataType::Int32, true).into()],
751 false,
752 false,
753 ),
754 [
755 None,
756 Some(1),
757 Some(-2),
758 Some(3),
759 Some(-4),
760 Some(5),
761 Some(-6),
762 Some(7),
763 ]
764 .iter()
765 .collect::<Int32Array>(),
766 )
767 }
768
769 #[test]
770 fn test_lag_with_default() -> Result<()> {
771 let expr = Arc::new(Column::new("c3", 0)) as Arc<dyn PhysicalExpr>;
772 let shift_offset =
773 Arc::new(Literal::new(ScalarValue::Int32(Some(1)))) as Arc<dyn PhysicalExpr>;
774 let default_value = Arc::new(Literal::new(ScalarValue::Int32(Some(100))))
775 as Arc<dyn PhysicalExpr>;
776
777 let input_exprs = &[expr, shift_offset, default_value];
778 let input_fields = [DataType::Int32, DataType::Int32, DataType::Int32]
779 .into_iter()
780 .map(|d| Field::new("f", d, true))
781 .map(Arc::new)
782 .collect::<Vec<_>>();
783
784 test_i32_result(
785 WindowShift::lag(),
786 PartitionEvaluatorArgs::new(input_exprs, &input_fields, false, false),
787 [
788 Some(100),
789 Some(1),
790 Some(-2),
791 Some(3),
792 Some(-4),
793 Some(5),
794 Some(-6),
795 Some(7),
796 ]
797 .iter()
798 .collect::<Int32Array>(),
799 )
800 }
801}