[AURON #2178] [AURON #2179] Implement native support for first_value and last_value window functions

native-engine/auron-planner/proto/auron.proto

@@ -142,6 +142,8 @@ enum AggFunction {
   FIRST = 7;
   FIRST_IGNORES_NULL = 8;
   BLOOM_FILTER = 9;
+  LAST = 10;
+  LAST_IGNORES_NULL = 11;
   BRICKHOUSE_COLLECT = 1000;
   BRICKHOUSE_COMBINE_UNIQUE = 1001;
   UDAF = 1002;

native-engine/auron-planner/src/lib.rs

@@ -135,6 +135,8 @@ impl From<protobuf::AggFunction> for AggFunction {
             protobuf::AggFunction::CollectSet => AggFunction::CollectSet,
             protobuf::AggFunction::First => AggFunction::First,
             protobuf::AggFunction::FirstIgnoresNull => AggFunction::FirstIgnoresNull,
+            protobuf::AggFunction::Last => AggFunction::Last,
+            protobuf::AggFunction::LastIgnoresNull => AggFunction::LastIgnoresNull,
             protobuf::AggFunction::BloomFilter => AggFunction::BloomFilter,
             protobuf::AggFunction::BrickhouseCollect => AggFunction::BrickhouseCollect,
             protobuf::AggFunction::BrickhouseCombineUnique => AggFunction::BrickhouseCombineUnique,

native-engine/auron-planner/src/planner.rs

@@ -657,6 +657,12 @@ impl PhysicalPlanner {
                                 protobuf::AggFunction::FirstIgnoresNull => {
                                     WindowFunction::Agg(AggFunction::FirstIgnoresNull)
                                 }
+                                protobuf::AggFunction::Last => {
+                                    WindowFunction::Agg(AggFunction::Last)
+                                }
+                                protobuf::AggFunction::LastIgnoresNull => {
+                                    WindowFunction::Agg(AggFunction::LastIgnoresNull)
+                                }
                                 protobuf::AggFunction::BloomFilter => {
                                     WindowFunction::Agg(AggFunction::BloomFilter)
                                 }

native-engine/datafusion-ext-plans/src/agg/agg.rs

@@ -33,6 +33,8 @@ use crate::agg::{
     count::AggCount,
     first::AggFirst,
     first_ignores_null::AggFirstIgnoresNull,
+    last::AggLast,
+    last_ignores_null::AggLastIgnoresNull,
     maxmin::{AggMax, AggMin},
     spark_udaf_wrapper::SparkUDAFWrapper,
     sum::AggSum,
@@ -212,6 +214,14 @@ pub fn create_agg(
             let dt = children[0].data_type(input_schema)?;
             Arc::new(AggFirstIgnoresNull::try_new(children[0].clone(), dt)?)
         }
+        AggFunction::Last => {
+            let dt = children[0].data_type(input_schema)?;
+            Arc::new(AggLast::try_new(children[0].clone(), dt)?)
+        }
+        AggFunction::LastIgnoresNull => {
+            let dt = children[0].data_type(input_schema)?;
+            Arc::new(AggLastIgnoresNull::try_new(children[0].clone(), dt)?)
+        }
         AggFunction::BloomFilter => {
             let dt = children[0].data_type(input_schema)?;
             let empty_batch = RecordBatch::new_empty(Arc::new(Schema::empty()));

native-engine/datafusion-ext-plans/src/agg/last.rs

@@ -0,0 +1,235 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements.  See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::{
+    any::Any,
+    fmt::{Debug, Formatter},
+    sync::Arc,
+};
+
+use arrow::{array::*, datatypes::*};
+use datafusion::{
+    common::{Result, ScalarValue},
+    physical_expr::PhysicalExprRef,
+};
+use datafusion_ext_commons::{downcast_any, scalar_value::compacted_scalar_value_from_array};
+
+use crate::{
+    agg::{
+        Agg,
+        acc::{
+            AccBooleanColumn, AccBytes, AccBytesColumn, AccColumnRef, AccPrimColumn,
+            AccScalarValueColumn, create_acc_generic_column,
+        },
+        agg::IdxSelection,
+    },
+    idx_for_zipped,
+};
+
+pub struct AggLast {
+    child: PhysicalExprRef,
+    data_type: DataType,
+    acc_array_data_types: Vec<DataType>,
+}
+
+impl AggLast {
+    pub fn try_new(child: PhysicalExprRef, data_type: DataType) -> Result<Self> {
+        let acc_array_data_types = vec![data_type.clone()];
+        Ok(Self {
+            child,
+            data_type,
+            acc_array_data_types,
+        })
+    }
+}
+
+impl Debug for AggLast {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "Last({:?})", self.child)
+    }
+}
+
+impl Agg for AggLast {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn exprs(&self) -> Vec<PhysicalExprRef> {
+        vec![self.child.clone()]
+    }
+
+    fn with_new_exprs(&self, exprs: Vec<PhysicalExprRef>) -> Result<Arc<dyn Agg>> {
+        Ok(Arc::new(Self::try_new(
+            exprs[0].clone(),
+            self.data_type.clone(),
+        )?))
+    }
+
+    fn data_type(&self) -> &DataType {
+        &self.data_type
+    }
+
+    fn nullable(&self) -> bool {
+        true
+    }
+
+    fn create_acc_column(&self, num_rows: usize) -> AccColumnRef {
+        create_acc_generic_column(self.data_type.clone(), num_rows)
+    }
+
+    fn acc_array_data_types(&self) -> &[DataType] {
+        &self.acc_array_data_types
+    }
+
+    fn partial_update(
+        &self,
+        accs: &mut AccColumnRef,
+        acc_idx: IdxSelection<'_>,
+        partial_args: &[ArrayRef],
+        partial_arg_idx: IdxSelection<'_>,
+    ) -> Result<()> {
+        let partial_arg = &partial_args[0];
+        accs.ensure_size(acc_idx);
+
+        macro_rules! handle_bytes {
+            ($array:expr) => {{
+                let accs = downcast_any!(accs, mut AccBytesColumn)?;
+                let partial_arg = $array;
+                idx_for_zipped! {
+                    ((acc_idx, partial_arg_idx) in (acc_idx, partial_arg_idx)) => {
+                        if partial_arg.is_valid(partial_arg_idx) {
+                            accs.set_value(acc_idx, Some(AccBytes::from(partial_arg.value(partial_arg_idx).as_ref())));
+                        } else {
+                            accs.set_value(acc_idx, None);
+                        }
+                    }
+                }
+            }}
+        }
+
+        downcast_primitive_array! {
+            partial_arg => {
+                if let Ok(accs) = downcast_any!(accs, mut AccPrimColumn<_>) {
+                    idx_for_zipped! {
+                        ((acc_idx, partial_arg_idx) in (acc_idx, partial_arg_idx)) => {
+                            if partial_arg.is_valid(partial_arg_idx) {
+                                accs.set_value(acc_idx, Some(partial_arg.value(partial_arg_idx)));
+                            } else {
+                                accs.set_value(acc_idx, None);
+                            }
+                        }
+                    }
+                }
+            }
+            DataType::Boolean => {
+                let accs = downcast_any!(accs, mut AccBooleanColumn)?;
+                let partial_arg = downcast_any!(partial_arg, BooleanArray)?;
+                idx_for_zipped! {
+                    ((acc_idx, partial_arg_idx) in (acc_idx, partial_arg_idx)) => {
+                        if partial_arg.is_valid(partial_arg_idx) {
+                            accs.set_value(acc_idx, Some(partial_arg.value(partial_arg_idx)));
+                        } else {
+                            accs.set_value(acc_idx, None);
+                        }
+                    }
+                }
+            }
+            DataType::Utf8 => handle_bytes!(downcast_any!(partial_arg, StringArray)?),
+            DataType::Binary => handle_bytes!(downcast_any!(partial_arg, BinaryArray)?),
+            _other => {
+                let accs = downcast_any!(accs, mut AccScalarValueColumn)?;
+                idx_for_zipped! {
+                    ((acc_idx, partial_arg_idx) in (acc_idx, partial_arg_idx)) => {
+                        if partial_arg.is_valid(partial_arg_idx) {
+                            accs.set_value(acc_idx, compacted_scalar_value_from_array(partial_arg, partial_arg_idx)?);
+                        } else {
+                            accs.set_value(acc_idx, ScalarValue::try_from(&self.data_type)?);
+                        }
+                    }
+                }
+            }
+        }
+        Ok(())
+    }
+
+    fn partial_merge(
+        &self,
+        accs: &mut AccColumnRef,
+        acc_idx: IdxSelection<'_>,
+        merging_accs: &mut AccColumnRef,
+        merging_acc_idx: IdxSelection<'_>,
+    ) -> Result<()> {
+        accs.ensure_size(acc_idx);
+
+        // For last, always overwrite with the merging accumulator's value
+        macro_rules! handle_primitive {
+            ($ty:ty) => {{
+                type TNative = <$ty as ArrowPrimitiveType>::Native;
+                let accs = downcast_any!(accs, mut AccPrimColumn<TNative>)?;
+                let merging_accs = downcast_any!(merging_accs, mut AccPrimColumn<_>)?;
+                idx_for_zipped! {
+                    ((acc_idx, merging_acc_idx) in (acc_idx, merging_acc_idx)) => {
+                        accs.set_value(acc_idx, merging_accs.value(merging_acc_idx));
+                    }
+                }
+            }}
+        }
+
+        macro_rules! handle_boolean {
+            () => {{
+                let accs = downcast_any!(accs, mut AccBooleanColumn)?;
+                let merging_accs = downcast_any!(merging_accs, mut AccBooleanColumn)?;
+                idx_for_zipped! {
+                    ((acc_idx, merging_acc_idx) in (acc_idx, merging_acc_idx)) => {
+                        accs.set_value(acc_idx, merging_accs.value(merging_acc_idx));
+                    }
+                }
+            }};
+        }
+
+        macro_rules! handle_bytes {
+            () => {{
+                let accs = downcast_any!(accs, mut AccBytesColumn)?;
+                let merging_accs = downcast_any!(merging_accs, mut AccBytesColumn)?;
+                idx_for_zipped! {
+                    ((acc_idx, merging_acc_idx) in (acc_idx, merging_acc_idx)) => {
+                        accs.set_value(acc_idx, merging_accs.take_value(merging_acc_idx));
+                    }
+                }
+            }};
+        }
+
+        downcast_primitive! {
+            (&self.data_type) => (handle_primitive),
+            DataType::Boolean => handle_boolean!(),
+            DataType::Utf8 | DataType::Binary => handle_bytes!(),
+            DataType::Null => {}
+            _ => {
+                let accs = downcast_any!(accs, mut AccScalarValueColumn)?;
+                let merging_accs = downcast_any!(merging_accs, mut AccScalarValueColumn)?;
+                idx_for_zipped! {
+                    ((acc_idx, merging_acc_idx) in (acc_idx, merging_acc_idx)) => {
+                        accs.set_value(acc_idx, merging_accs.take_value(merging_acc_idx));
+                    }
+                }
+            }
+        }
+        Ok(())
+    }
+
+    fn final_merge(&self, accs: &mut AccColumnRef, acc_idx: IdxSelection<'_>) -> Result<ArrayRef> {
+        Ok(accs.freeze_to_arrays(acc_idx)?[0].clone())
+    }
+}