[AURON #2507] Optimize partial WindowGroupLimit execution

native-engine/datafusion-ext-plans/benches/window_group_limit.rs

@@ -0,0 +1,194 @@
+// 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.
+
+#![feature(test)]
+
+extern crate test;
+
+use std::sync::Arc;
+
+use arrow::{
+    array::{ArrayRef, Int32Array},
+    datatypes::{DataType, Field, Schema},
+    record_batch::RecordBatch,
+};
+use datafusion::{
+    execution::TaskContext,
+    physical_expr::{PhysicalSortExpr, expressions::Column},
+    physical_plan::{ExecutionPlan, common, test::TestMemoryExec},
+    prelude::SessionContext,
+};
+use datafusion_ext_plans::{
+    window::{WindowExpr, WindowFunction, WindowRankType},
+    window_exec::WindowExec,
+};
+use test::{Bencher, black_box};
+
+fn create_batch(
+    num_partitions: usize,
+    rows_per_partition: usize,
+    peer_group_size: usize,
+) -> RecordBatch {
+    let num_rows = num_partitions * rows_per_partition;
+    let partitions = (0..num_rows)
+        .map(|row| (row / rows_per_partition) as i32)
+        .collect::<Vec<_>>();
+    let order_values = (0..num_rows)
+        .map(|row| ((row % rows_per_partition) / peer_group_size) as i32)
+        .collect::<Vec<_>>();
+    let payloads = (0..num_rows as i32).collect::<Vec<_>>();
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("partition", DataType::Int32, false),
+        Field::new("order", DataType::Int32, false),
+        Field::new("payload", DataType::Int32, false),
+    ]));
+    let columns: Vec<ArrayRef> = vec![
+        Arc::new(Int32Array::from(partitions)),
+        Arc::new(Int32Array::from(order_values)),
+        Arc::new(Int32Array::from(payloads)),
+    ];
+    RecordBatch::try_new(schema, columns).expect("benchmark batch should be valid")
+}
+
+fn create_window_group_limit(
+    rank_type: WindowRankType,
+    num_partitions: usize,
+    rows_per_partition: usize,
+    peer_group_size: usize,
+    group_limit: usize,
+) -> Arc<dyn ExecutionPlan> {
+    let batch = create_batch(num_partitions, rows_per_partition, peer_group_size);
+    let input = Arc::new(
+        TestMemoryExec::try_new(&[vec![batch.clone()]], batch.schema(), None)
+            .expect("benchmark input should be valid"),
+    );
+    let window_exprs = vec![WindowExpr::new(
+        WindowFunction::RankLike(rank_type),
+        vec![],
+        Arc::new(Field::new("rank", DataType::Int32, false)),
+        DataType::Int32,
+    )];
+    Arc::new(
+        WindowExec::try_new(
+            input,
+            window_exprs,
+            vec![Arc::new(Column::new("partition", 0))],
+            vec![PhysicalSortExpr {
+                expr: Arc::new(Column::new("order", 1)),
+                options: Default::default(),
+            }],
+            Some(group_limit),
+            false,
+        )
+        .expect("benchmark WindowGroupLimit should be valid"),
+    )
+}
+
+fn execute(
+    runtime: &tokio::runtime::Runtime,
+    task_ctx: &Arc<TaskContext>,
+    exec: &Arc<dyn ExecutionPlan>,
+) -> Vec<RecordBatch> {
+    runtime
+        .block_on(async {
+            let stream = exec
+                .execute(0, task_ctx.clone())
+                .expect("benchmark execution should start");
+            common::collect(stream).await
+        })
+        .expect("benchmark output should be collected")
+}
+
+fn expected_rows(
+    rank_type: WindowRankType,
+    num_partitions: usize,
+    rows_per_partition: usize,
+    peer_group_size: usize,
+    group_limit: usize,
+) -> usize {
+    let rows_per_partition = match rank_type {
+        WindowRankType::RowNumber => group_limit.min(rows_per_partition),
+        WindowRankType::Rank => group_limit
+            .div_ceil(peer_group_size)
+            .saturating_mul(peer_group_size)
+            .min(rows_per_partition),
+        WindowRankType::DenseRank => group_limit
+            .saturating_mul(peer_group_size)
+            .min(rows_per_partition),
+    };
+    num_partitions * rows_per_partition
+}
+
+fn bench_window_group_limit(
+    b: &mut Bencher,
+    rank_type: WindowRankType,
+    num_partitions: usize,
+    rows_per_partition: usize,
+    peer_group_size: usize,
+    group_limit: usize,
+) {
+    let runtime = tokio::runtime::Runtime::new().expect("benchmark runtime should be created");
+    let task_ctx = SessionContext::new().task_ctx();
+    let exec = create_window_group_limit(
+        rank_type,
+        num_partitions,
+        rows_per_partition,
+        peer_group_size,
+        group_limit,
+    );
+    let output = execute(&runtime, &task_ctx, &exec);
+    assert_eq!(
+        output.iter().map(RecordBatch::num_rows).sum::<usize>(),
+        expected_rows(
+            rank_type,
+            num_partitions,
+            rows_per_partition,
+            peer_group_size,
+            group_limit,
+        )
+    );
+
+    b.iter(|| black_box(execute(&runtime, &task_ctx, &exec)));
+}
+
+#[bench]
+fn window_group_limit_row_number_slice(b: &mut Bencher) {
+    bench_window_group_limit(b, WindowRankType::RowNumber, 1, 100_000, 10, 10);
+}
+
+#[bench]
+fn window_group_limit_rank_slice(b: &mut Bencher) {
+    bench_window_group_limit(b, WindowRankType::Rank, 1, 100_000, 10, 10);
+}
+
+#[bench]
+fn window_group_limit_dense_rank_slice(b: &mut Bencher) {
+    bench_window_group_limit(b, WindowRankType::DenseRank, 1, 100_000, 10, 10);
+}
+
+#[bench]
+fn window_group_limit_row_number(b: &mut Bencher) {
+    bench_window_group_limit(b, WindowRankType::RowNumber, 100, 1_000, 10, 10);
+}
+
+#[bench]
+fn window_group_limit_rank(b: &mut Bencher) {
+    bench_window_group_limit(b, WindowRankType::Rank, 100, 1_000, 10, 10);
+}
+
+#[bench]
+fn window_group_limit_dense_rank(b: &mut Bencher) {
+    bench_window_group_limit(b, WindowRankType::DenseRank, 100, 1_000, 10, 10);
+}

native-engine/datafusion-ext-plans/src/window/group_limit_processor.rs

@@ -0,0 +1,112 @@
+// 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::ops::Range;
+
+use arrow::record_batch::RecordBatch;
+use datafusion::common::Result;
+
+use crate::window::{WindowRankType, window_context::WindowContext};
+
+pub(crate) struct WindowGroupLimitProcessor {
+    rank_type: WindowRankType,
+    limit: i32,
+    cur_partition: Vec<u8>,
+    cur_order: Vec<u8>,
+    cur_rank: i32,
+    cur_equals: i32,
+}
+
+impl WindowGroupLimitProcessor {
+    pub(crate) fn new(rank_type: WindowRankType, limit: usize) -> Self {
+        Self {
+            rank_type,
+            limit: i32::try_from(limit).unwrap_or(i32::MAX),
+            cur_partition: vec![],
+            cur_order: vec![],
+            cur_rank: 0,
+            cur_equals: 1,
+        }
+    }
+
+    pub(crate) fn process_batch(
+        &mut self,
+        context: &WindowContext,
+        batch: &RecordBatch,
+    ) -> Result<Vec<Range<usize>>> {
+        let partition_rows = context.get_partition_rows(batch)?;
+        let order_rows = match self.rank_type {
+            WindowRankType::RowNumber => None,
+            WindowRankType::Rank | WindowRankType::DenseRank => {
+                Some(context.get_order_rows(batch)?)
+            }
+        };
+        let mut selected_ranges = vec![];
+        let mut selected_start = None;
+
+        for row_idx in 0..batch.num_rows() {
+            let same_partition = !context.has_partition() || {
+                let partition_row = partition_rows.row(row_idx);
+                if partition_row.as_ref() != self.cur_partition {
+                    self.cur_partition = partition_row.as_ref().into();
+                    false
+                } else {
+                    true
+                }
+            };
+
+            match self.rank_type {
+                WindowRankType::RowNumber => {
+                    if !same_partition {
+                        self.cur_rank = 0;
+                    }
+                    self.cur_rank += 1;
+                }
+                WindowRankType::Rank | WindowRankType::DenseRank => {
+                    let order_row = order_rows
+                        .as_ref()
+                        .expect("rank and dense_rank must have order rows")
+                        .row(row_idx);
+                    if same_partition {
+                        if order_row.as_ref() == self.cur_order {
+                            self.cur_equals += 1;
+                        } else {
+                            self.cur_rank += match self.rank_type {
+                                WindowRankType::Rank => self.cur_equals,
+                                WindowRankType::DenseRank => 1,
+                                WindowRankType::RowNumber => unreachable!(),
+                            };
+                            self.cur_equals = 1;
+                            self.cur_order = order_row.as_ref().into();
+                        }
+                    } else {
+                        self.cur_rank = 1;
+                        self.cur_equals = 1;
+                        self.cur_order = order_row.as_ref().into();
+                    }
+                }
+            }
+            if self.cur_rank <= self.limit {
+                selected_start.get_or_insert(row_idx);
+            } else if let Some(start) = selected_start.take() {
+                selected_ranges.push(start..row_idx);
+            }
+        }
+        if let Some(start) = selected_start {
+            selected_ranges.push(start..batch.num_rows());
+        }
+        Ok(selected_ranges)
+    }
+}

native-engine/datafusion-ext-plans/src/window/mod.rs

@@ -32,6 +32,7 @@ use crate::{
     },
 };
 
+pub(crate) mod group_limit_processor;
 pub mod processors;
 pub mod window_context;
 
@@ -45,7 +46,7 @@ pub enum WindowFunction {
     Agg(AggFunction),
 }
 
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum WindowRankType {
     RowNumber,
     Rank,
@@ -112,6 +113,13 @@ impl WindowExpr {
         }
     }
 
+    pub fn rank_type(&self) -> Option<WindowRankType> {
+        match self.func {
+            WindowFunction::RankLike(rank_type) => Some(rank_type),
+            _ => None,
+        }
+    }
+
     pub fn requires_full_partition(&self) -> bool {
         matches!(
             self.func,