use array ref

Author: Jiayu Liu

datafusion/src/physical_plan/expressions/row_number.rs

@@ -22,7 +22,7 @@ use crate::physical_plan::{
     window_functions::BuiltInWindowFunctionExpr, PhysicalExpr, WindowAccumulator,
 };
 use crate::scalar::ScalarValue;
-use arrow::array::ArrayRef;
+use arrow::array::{ArrayRef, UInt64Array};
 use arrow::datatypes::{DataType, Field};
 use std::any::Any;
 use std::sync::Arc;
@@ -89,13 +89,11 @@ impl WindowAccumulator for RowNumberAccumulator {
         &mut self,
         num_rows: usize,
         _values: &[ArrayRef],
-    ) -> Result<Option<Vec<ScalarValue>>> {
+    ) -> Result<Option<ArrayRef>> {
         let new_row_number = self.row_number + (num_rows as u64);
-        let result = (self.row_number..new_row_number)
-            .map(|i| ScalarValue::UInt64(Some(i)))
-            .collect();
+        let result = UInt64Array::from_iter_values(self.row_number..new_row_number);
         self.row_number = new_row_number;
-        Ok(Some(result))
+        Ok(Some(Arc::new(result)))
     }
 
     fn evaluate(&self) -> Result<Option<ScalarValue>> {

datafusion/src/physical_plan/mod.rs

@@ -539,12 +539,12 @@ pub trait WindowAccumulator: Send + Sync + Debug {
         &mut self,
         num_rows: usize,
         values: &[ArrayRef],
-    ) -> Result<Option<Vec<ScalarValue>>> {
+    ) -> Result<Option<ArrayRef>> {
         if values.is_empty() {
             return Ok(None);
         };
         // transpose columnar to row based so that we can apply window
-        let result: Vec<Option<ScalarValue>> = (0..num_rows)
+        let result: Option<Vec<ScalarValue>> = (0..num_rows)
             .map(|index| {
                 let v = values
                     .iter()
@@ -553,8 +553,15 @@ pub trait WindowAccumulator: Send + Sync + Debug {
                 self.scan(&v)
             })
             .into_iter()
-            .collect::<Result<Vec<Option<ScalarValue>>>>()?;
-        let result: Option<Vec<ScalarValue>> = result.into_iter().collect();
+            .into_iter()
+            .collect::<Result<Option<Vec<_>>>>()?;
+
+        let result: Option<ArrayRef> = if let Some(arr) = result {
+            Some(ScalarValue::iter_to_array(&arr)?)
+        } else {
+            None
+        };
+
         Ok(result)
     }
 

datafusion/src/physical_plan/windows.rs

@@ -29,7 +29,7 @@ use crate::physical_plan::{
 use crate::scalar::ScalarValue;
 use arrow::compute::concat;
 use arrow::{
-    array::ArrayRef,
+    array::{Array, ArrayRef},
     datatypes::{Field, Schema, SchemaRef},
     error::{ArrowError, Result as ArrowResult},
     record_batch::RecordBatch,
@@ -307,7 +307,7 @@ fn window_aggregate_batch(
     batch: &RecordBatch,
     window_accumulators: &mut [WindowAccumulatorItem],
     expressions: &[Vec<Arc<dyn PhysicalExpr>>],
-) -> Result<Vec<Option<Vec<ScalarValue>>>> {
+) -> Result<Vec<Option<ArrayRef>>> {
     // 1.1 iterate accumulators and respective expressions together
     // 1.2 evaluate expressions
     // 1.3 update / merge window accumulators with the expressions' values
@@ -364,7 +364,7 @@ async fn compute_window_aggregate(
     let expressions = Arc::new(expressions);
 
     // TODO each element shall have some size hint
-    let mut accumulator: Vec<Vec<ScalarValue>> =
+    let mut accumulator: Vec<Vec<ArrayRef>> =
         iter::repeat(vec![]).take(window_expr.len()).collect();
 
     let mut original_batches: Vec<RecordBatch> = vec![];
@@ -382,7 +382,7 @@ async fn compute_window_aggregate(
         accumulator.iter_mut().zip(batch_aggregated).for_each(
             |(acc_for_window, window_batch)| {
                 if let Some(data) = window_batch {
-                    acc_for_window.extend(data);
+                    acc_for_window.push(data);
                 }
             },
         );
@@ -400,13 +400,13 @@ async fn compute_window_aggregate(
                     scalar_value.to_array_of_size(total_num_rows)
                 }
                 (acc, None) if !acc.is_empty() => {
-                    return Err(DataFusionError::NotImplemented(
-                        "built in window function not yet implemented".to_owned(),
-                    ))
+                    let converted: Vec<&dyn Array> =
+                        acc.iter().map(|arc| arc.as_ref()).collect::<Vec<_>>();
+                    concat(&converted)?
                 }
                 _ => {
                     return Err(DataFusionError::Execution(
-                        "invalid window function behavior".to_owned(),
+                        "Invalid window function behavior".to_owned(),
                     ))
                 }
             };