Return proper number of expressions for nth_value_agg (#9044)

* Return proper number of expressions

* Fix bug

* Update datafusion/proto/tests/cases/roundtrip_physical_plan.rs

Co-authored-by: Andrew Lamb <[email protected]>

* Fix argument type

---------

Co-authored-by: Andrew Lamb <[email protected]>

Author: mustafasrepo

datafusion/physical-expr/src/aggregate/nth_value.rs

@@ -24,7 +24,7 @@ use std::sync::Arc;
 
 use crate::aggregate::array_agg_ordered::merge_ordered_arrays;
 use crate::aggregate::utils::{down_cast_any_ref, ordering_fields};
-use crate::expressions::format_state_name;
+use crate::expressions::{format_state_name, Literal};
 use crate::{
     reverse_order_bys, AggregateExpr, LexOrdering, PhysicalExpr, PhysicalSortExpr,
 };
@@ -117,7 +117,8 @@ impl AggregateExpr for NthValueAgg {
     }
 
     fn expressions(&self) -> Vec<Arc<dyn PhysicalExpr>> {
-        vec![self.expr.clone()]
+        let n = Arc::new(Literal::new(ScalarValue::Int64(Some(self.n)))) as _;
+        vec![self.expr.clone(), n]
     }
 
     fn order_bys(&self) -> Option<&[PhysicalSortExpr]> {
@@ -393,7 +394,9 @@ impl NthValueAccumulator {
         for index in 0..n_to_add {
             let row = get_row_at_idx(values, index)?;
             self.values.push_back(row[0].clone());
-            self.ordering_values.push_back(row[1..].to_vec());
+            // At index 1, we have n index argument.
+            // Ordering values cover starting from 2nd index to end
+            self.ordering_values.push_back(row[2..].to_vec());
         }
         Ok(())
     }

datafusion/proto/tests/cases/roundtrip_physical_plan.rs

@@ -36,6 +36,7 @@ use datafusion::logical_expr::{
 };
 use datafusion::parquet::file::properties::WriterProperties;
 use datafusion::physical_expr::expressions::Literal;
+use datafusion::physical_expr::expressions::NthValueAgg;
 use datafusion::physical_expr::window::SlidingAggregateWindowExpr;
 use datafusion::physical_expr::{PhysicalSortRequirement, ScalarFunctionExpr};
 use datafusion::physical_plan::aggregates::{
@@ -337,17 +338,16 @@ fn rountrip_aggregate() -> Result<()> {
             "AVG(b)".to_string(),
             DataType::Float64,
         ))],
-        // TODO: See <https://github.com/apache/arrow-datafusion/issues/9028>
-        // // NTH_VALUE
-        // vec![Arc::new(NthValueAgg::new(
-        //     col("b", &schema)?,
-        //     1,
-        //     "NTH_VALUE(b, 1)".to_string(),
-        //     DataType::Int64,
-        //     false,
-        //     Vec::new(),
-        //     Vec::new(),
-        // ))],
+        // NTH_VALUE
+        vec![Arc::new(NthValueAgg::new(
+            col("b", &schema)?,
+            1,
+            "NTH_VALUE(b, 1)".to_string(),
+            DataType::Int64,
+            false,
+            Vec::new(),
+            Vec::new(),
+        ))],
         // STRING_AGG
         vec![Arc::new(StringAgg::new(
             cast(col("b", &schema)?, &schema, DataType::Utf8)?,