[GitHub] [arrow-datafusion] Dandandan commented on a change in pull request #429: implement lead and lag built-in window function

[GitBox <gi...@apache.org>]

Dandandan commented on a change in pull request #429:
URL: https://github.com/apache/arrow-datafusion/pull/429#discussion_r639800391



##########
File path: datafusion/src/physical_plan/windows.rs
##########
@@ -421,18 +435,52 @@ async fn compute_window_aggregate(
     let aggregated_mapped = finalize_window_aggregation(&window_accumulators)
         .map_err(DataFusionError::into_arrow_external_error)?;
 
+    let window_shifts = window_accumulators
+        .iter()
+        .map(|acc| acc.window_shift())
+        .collect::<Vec<_>>();
+
     let mut columns: Vec<ArrayRef> = accumulator
         .iter()
         .zip(aggregated_mapped)
-        .map(|(acc, agg)| {
-            Ok(match (acc, agg) {
-                (acc, Some(scalar_value)) if acc.is_empty() => {
+        .zip(window_shifts)
+        .map(|((acc, agg), window_shift)| {
+            Ok(match (acc, agg, window_shift) {
+                (acc, Some(scalar_value), None) if acc.is_empty() => {
                     scalar_value.to_array_of_size(total_num_rows)
                 }
-                (acc, None) if !acc.is_empty() => {
+                (acc, None, window_shift) if !acc.is_empty() => {
                     let vec_array: Vec<&dyn Array> =
                         acc.iter().map(|arc| arc.as_ref()).collect();
-                    concat(&vec_array)?
+                    let arr: ArrayRef = concat(&vec_array)?;
+                    if arr.is_empty() {
+                        arr
+                    } else if arr.len() != total_num_rows {
+                        return Err(DataFusionError::Internal(format!(
+                            "Invalid concatenated array of length {}, expecting {}",
+                            arr.len(),
+                            total_num_rows
+                        )));
+                    } else {
+                        let data_type = arr.data_type();
+                        match window_shift {

Review comment:
       What about using: https://docs.rs/arrow/4.1.0/arrow/compute/kernels/window/fn.shift.html




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org