remove unnecessary computations

Author: berkaysynnada

datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs

@@ -67,7 +67,6 @@ use futures::stream::Stream;
 use futures::{ready, StreamExt};
 use hashbrown::hash_table::HashTable;
 use indexmap::IndexMap;
-use itertools::Itertools;
 use log::debug;
 
 /// Window execution plan
@@ -104,10 +103,8 @@ impl BoundedWindowAggExec {
         partition_keys: Vec<Arc<dyn PhysicalExpr>>,
         input_order_mode: InputOrderMode,
     ) -> Result<Self> {
-        let old_fields_len = input.schema().fields.len();
         let schema = create_schema(&input.schema(), &window_expr)?;
         let schema = Arc::new(schema);
-        let window_expr_indices = (old_fields_len..schema.fields.len()).collect_vec();
         let partition_by_exprs = window_expr[0].partition_by();
         let ordered_partition_by_indices = match &input_order_mode {
             InputOrderMode::Sorted => {
@@ -126,8 +123,7 @@ impl BoundedWindowAggExec {
                 vec![]
             }
         };
-        let cache =
-            Self::compute_properties(&input, &schema, &window_expr, window_expr_indices);
+        let cache = Self::compute_properties(&input, &schema, &window_expr);
         Ok(Self {
             input,
             window_expr,
@@ -195,15 +191,9 @@ impl BoundedWindowAggExec {
         input: &Arc<dyn ExecutionPlan>,
         schema: &SchemaRef,
         window_exprs: &[Arc<dyn WindowExpr>],
-        window_expr_indices: Vec<usize>,
     ) -> PlanProperties {
         // Calculate equivalence properties:
-        let eq_properties = window_equivalence_properties(
-            schema,
-            input,
-            window_exprs,
-            window_expr_indices,
-        );
+        let eq_properties = window_equivalence_properties(schema, input, window_exprs);
 
         // As we can have repartitioning using the partition keys, this can
         // be either one or more than one, depending on the presence of

datafusion/physical-plan/src/windows/mod.rs

@@ -337,21 +337,25 @@ pub(crate) fn window_equivalence_properties(
     schema: &SchemaRef,
     input: &Arc<dyn ExecutionPlan>,
     window_exprs: &[Arc<dyn WindowExpr>],
-    window_expr_indices: Vec<usize>,
 ) -> EquivalenceProperties {
     // We need to update the schema, so we can not directly use
     // `input.equivalence_properties()`.
     let mut window_eq_properties = EquivalenceProperties::new(Arc::clone(schema))
         .extend(input.equivalence_properties().clone());
 
+    let schema_len = schema.fields.len();
+    let window_expr_indices = (schema_len..(schema_len - window_exprs.len()))
+        .rev()
+        .collect::<Vec<_>>();
+
     for (i, expr) in window_exprs.iter().enumerate() {
-        let window_expr_index = window_expr_indices[i];
         if let Some(udf_window_expr) = expr.as_any().downcast_ref::<StandardWindowExpr>()
         {
             udf_window_expr.add_equal_orderings(&mut window_eq_properties);
         } else if let Some(aggregate_udf_window_expr) =
             expr.as_any().downcast_ref::<PlainAggregateWindowExpr>()
         {
+            let window_expr_index = window_expr_indices[i];
             aggregate_udf_window_expr
                 .add_equal_orderings(&mut window_eq_properties, window_expr_index);
         }

datafusion/physical-plan/src/windows/window_agg_exec.rs

@@ -47,7 +47,6 @@ use datafusion_execution::TaskContext;
 use datafusion_physical_expr_common::sort_expr::{LexOrdering, LexRequirement};
 
 use futures::{ready, Stream, StreamExt};
-use itertools::Itertools;
 
 /// Window execution plan
 #[derive(Debug, Clone)]
@@ -76,19 +75,12 @@ impl WindowAggExec {
         input: Arc<dyn ExecutionPlan>,
         partition_keys: Vec<Arc<dyn PhysicalExpr>>,
     ) -> Result<Self> {
-        let old_fields_len = input.schema().fields.len();
         let schema = create_schema(&input.schema(), &window_expr)?;
         let schema = Arc::new(schema);
-        let window_expr_indices = (old_fields_len..schema.fields.len()).collect_vec();
 
         let ordered_partition_by_indices =
             get_ordered_partition_by_indices(window_expr[0].partition_by(), &input);
-        let cache = Self::compute_properties(
-            Arc::clone(&schema),
-            &input,
-            &window_expr,
-            window_expr_indices,
-        );
+        let cache = Self::compute_properties(Arc::clone(&schema), &input, &window_expr);
         Ok(Self {
             input,
             window_expr,
@@ -129,15 +121,9 @@ impl WindowAggExec {
         schema: SchemaRef,
         input: &Arc<dyn ExecutionPlan>,
         window_exprs: &[Arc<dyn WindowExpr>],
-        window_expr_indices: Vec<usize>,
     ) -> PlanProperties {
         // Calculate equivalence properties:
-        let eq_properties = window_equivalence_properties(
-            &schema,
-            input,
-            window_exprs,
-            window_expr_indices,
-        );
+        let eq_properties = window_equivalence_properties(&schema, input, window_exprs);
 
         // Get output partitioning:
         // Because we can have repartitioning using the partition keys this