From 914e8f6fe8ffe9fd018f4f24c6df813fc4676e2e Mon Sep 17 00:00:00 2001 From: linfeng Date: Wed, 2 Sep 2026 22:55:01 +0800 Subject: [PATCH 1/4] perf: optimize partial window group limit execution --- .../benches/window_group_limit.rs | 137 ++++++++++++++ .../src/window/group_limit_processor.rs | 105 +++++++++++ .../datafusion-ext-plans/src/window/mod.rs | 10 +- .../datafusion-ext-plans/src/window_exec.rs | 171 +++++++++++++++++- 4 files changed, 414 insertions(+), 9 deletions(-) create mode 100644 native-engine/datafusion-ext-plans/benches/window_group_limit.rs create mode 100644 native-engine/datafusion-ext-plans/src/window/group_limit_processor.rs diff --git a/native-engine/datafusion-ext-plans/benches/window_group_limit.rs b/native-engine/datafusion-ext-plans/benches/window_group_limit.rs new file mode 100644 index 000000000..5e03236fa --- /dev/null +++ b/native-engine/datafusion-ext-plans/benches/window_group_limit.rs @@ -0,0 +1,137 @@ +// 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}; + +const NUM_PARTITIONS: usize = 100; +const ROWS_PER_PARTITION: usize = 1_000; +const PEER_GROUP_SIZE: usize = 10; +const GROUP_LIMIT: usize = 10; + +fn create_batch() -> RecordBatch { + let num_rows = NUM_PARTITIONS * ROWS_PER_PARTITION; + let partitions = (0..num_rows) + .map(|row| (row / ROWS_PER_PARTITION) as i32) + .collect::>(); + let order_values = (0..num_rows) + .map(|row| ((row % ROWS_PER_PARTITION) / PEER_GROUP_SIZE) as i32) + .collect::>(); + let payloads = (0..num_rows as i32).collect::>(); + 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 = 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) -> Arc { + let batch = create_batch(); + 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, + exec: &Arc, +) -> Vec { + 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 bench_window_group_limit(b: &mut Bencher, rank_type: WindowRankType) { + 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); + let expected_rows = match rank_type { + WindowRankType::RowNumber | WindowRankType::Rank => NUM_PARTITIONS * GROUP_LIMIT, + WindowRankType::DenseRank => NUM_PARTITIONS * GROUP_LIMIT * PEER_GROUP_SIZE, + }; + let output = execute(&runtime, &task_ctx, &exec); + assert_eq!( + output.iter().map(RecordBatch::num_rows).sum::(), + expected_rows + ); + + b.iter(|| black_box(execute(&runtime, &task_ctx, &exec))); +} + +macro_rules! benchmark { + ($name:ident, $rank_type:expr) => { + #[bench] + fn $name(b: &mut Bencher) { + bench_window_group_limit(b, $rank_type); + } + }; +} + +benchmark!(window_group_limit_row_number, WindowRankType::RowNumber); +benchmark!(window_group_limit_rank, WindowRankType::Rank); +benchmark!(window_group_limit_dense_rank, WindowRankType::DenseRank); diff --git a/native-engine/datafusion-ext-plans/src/window/group_limit_processor.rs b/native-engine/datafusion-ext-plans/src/window/group_limit_processor.rs new file mode 100644 index 000000000..0adbf098f --- /dev/null +++ b/native-engine/datafusion-ext-plans/src/window/group_limit_processor.rs @@ -0,0 +1,105 @@ +// 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 arrow::{ + array::{BooleanArray, BooleanBuilder}, + 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, + cur_order: Vec, + 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 { + 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 builder = BooleanBuilder::with_capacity(batch.num_rows()); + + 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(); + } + } + } + builder.append_value(self.cur_rank <= self.limit); + } + Ok(builder.finish()) + } +} diff --git a/native-engine/datafusion-ext-plans/src/window/mod.rs b/native-engine/datafusion-ext-plans/src/window/mod.rs index f22772b61..8c8904dc3 100644 --- a/native-engine/datafusion-ext-plans/src/window/mod.rs +++ b/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 { + match self.func { + WindowFunction::RankLike(rank_type) => Some(rank_type), + _ => None, + } + } + pub fn requires_full_partition(&self) -> bool { matches!( self.func, diff --git a/native-engine/datafusion-ext-plans/src/window_exec.rs b/native-engine/datafusion-ext-plans/src/window_exec.rs index 1edbe1b6f..208744145 100644 --- a/native-engine/datafusion-ext-plans/src/window_exec.rs +++ b/native-engine/datafusion-ext-plans/src/window_exec.rs @@ -38,7 +38,10 @@ use once_cell::sync::OnceCell; use crate::{ common::execution_context::{ExecutionContext, WrappedRecordBatchSender}, - window::{WindowExpr, WindowFunctionProcessor, window_context::WindowContext}, + window::{ + WindowExpr, WindowFunctionProcessor, group_limit_processor::WindowGroupLimitProcessor, + window_context::WindowContext, + }, }; #[derive(Debug)] @@ -215,11 +218,34 @@ fn execute_window( let elapsed_compute = exec_ctx.baseline_metrics().elapsed_compute().clone(); sender.exclude_time(&elapsed_compute); - let mut processors = window_ctx - .window_exprs - .iter() - .map(|expr: &WindowExpr| expr.create_processor(&window_ctx)) - .collect::>>()?; + let partial_group_limit = if window_ctx.output_window_cols { + None + } else { + window_ctx.group_limit + }; + let mut group_limit_processor = match partial_group_limit { + Some(limit) => match window_ctx.window_exprs.as_slice() { + [expr] => match expr.rank_type() { + Some(rank_type) => Some(WindowGroupLimitProcessor::new(rank_type, limit)), + None => return datafusion::common::internal_err!( + "WindowGroupLimit requires a row_number, rank, or dense_rank expression" + ), + }, + _ => return datafusion::common::internal_err!( + "WindowGroupLimit requires exactly one window expression" + ), + }, + None => None, + }; + let mut processors = if group_limit_processor.is_some() { + vec![] + } else { + window_ctx + .window_exprs + .iter() + .map(|expr: &WindowExpr| expr.create_processor(&window_ctx)) + .collect::>>()? + }; if window_ctx.requires_full_partition() { // Functions like percent_rank/lead need a complete window partition, @@ -294,8 +320,11 @@ fn execute_window( while let Some(batch) = input.next().await.transpose()? { let _timer = elapsed_compute.timer(); - let output_batch = - process_window_batch(batch, &window_ctx, processors.as_mut_slice())?; + let output_batch = if let Some(processor) = &mut group_limit_processor { + process_partial_group_limit_batch(batch, &window_ctx, processor)? + } else { + process_window_batch(batch, &window_ctx, processors.as_mut_slice())? + }; exec_ctx .baseline_metrics() .record_output(output_batch.num_rows()); @@ -305,6 +334,15 @@ fn execute_window( })) } +fn process_partial_group_limit_batch( + batch: RecordBatch, + window_ctx: &WindowContext, + processor: &mut WindowGroupLimitProcessor, +) -> Result { + let selection = processor.process_batch(window_ctx, &batch)?; + Ok(arrow::compute::filter_record_batch(&batch, &selection)?) +} + async fn flush_window_batches( staging_batches: &mut Vec, window_ctx: &WindowContext, @@ -790,6 +828,123 @@ mod test { Ok(()) } + #[test] + fn test_window_expr_rank_type() { + for rank_type in [ + WindowRankType::RowNumber, + WindowRankType::Rank, + WindowRankType::DenseRank, + ] { + let expr = WindowExpr::new( + WindowFunction::RankLike(rank_type), + vec![], + Arc::new(Field::new("rank", DataType::Int32, false)), + DataType::Int32, + ); + assert!(matches!(expr.rank_type(), Some(actual) if actual == rank_type)); + } + + let expr = WindowExpr::new( + WindowFunction::PercentRank, + vec![], + Arc::new(Field::new("percent_rank", DataType::Float64, false)), + DataType::Float64, + ); + assert!(expr.rank_type().is_none()); + } + + #[tokio::test] + async fn test_partial_window_group_limit_rank_like_functions() + -> Result<(), Box> { + let session_ctx = SessionContext::new(); + let task_ctx = session_ctx.task_ctx(); + + let cases = [ + ( + WindowRankType::RowNumber, + vec![ + "+----+----+----+", + "| a1 | b1 | c1 |", + "+----+----+----+", + "| 1 | 1 | 0 |", + "| 1 | 2 | 0 |", + "| 2 | 1 | 0 |", + "| 2 | 1 | 0 |", + "+----+----+----+", + ], + ), + ( + WindowRankType::Rank, + vec![ + "+----+----+----+", + "| a1 | b1 | c1 |", + "+----+----+----+", + "| 1 | 1 | 0 |", + "| 1 | 2 | 0 |", + "| 1 | 2 | 0 |", + "| 2 | 1 | 0 |", + "| 2 | 1 | 0 |", + "+----+----+----+", + ], + ), + ( + WindowRankType::DenseRank, + vec![ + "+----+----+----+", + "| a1 | b1 | c1 |", + "+----+----+----+", + "| 1 | 1 | 0 |", + "| 1 | 2 | 0 |", + "| 1 | 2 | 0 |", + "| 2 | 1 | 0 |", + "| 2 | 1 | 0 |", + "| 2 | 2 | 0 |", + "+----+----+----+", + ], + ), + ]; + + for (rank_type, expected) in cases { + let batch1 = build_table_i32( + ("a1", &vec![1, 1]), + ("b1", &vec![1, 2]), + ("c1", &vec![0; 2]), + )?; + let batch2 = build_table_i32( + ("a1", &vec![1, 1, 2, 2, 2, 2]), + ("b1", &vec![2, 3, 1, 1, 2, 3]), + ("c1", &vec![0; 6]), + )?; + let input = Arc::new(TestMemoryExec::try_new( + &[vec![batch1.clone(), batch2]], + batch1.schema(), + None, + )?); + let window_exprs = vec![WindowExpr::new( + WindowFunction::RankLike(rank_type), + vec![], + Arc::new(Field::new("rank", DataType::Int32, false)), + DataType::Int32, + )]; + let window = Arc::new(WindowExec::try_new( + input, + window_exprs, + vec![Arc::new(Column::new("a1", 0))], + vec![PhysicalSortExpr { + expr: Arc::new(Column::new("b1", 1)), + options: Default::default(), + }], + Some(2), + false, + )?); + + let stream = window.execute(0, task_ctx.clone())?; + let batches = datafusion::physical_plan::common::collect(stream).await?; + assert_batches_eq!(expected, &batches); + } + Ok(()) + } + #[tokio::test] async fn test_window_lead_across_batches() -> Result<(), Box> { let session_ctx = SessionContext::new(); From ec8d9fae91f10c79ce1dd15f9ae68352b349dc7f Mon Sep 17 00:00:00 2001 From: linfeng Date: Thu, 3 Sep 2026 22:44:46 +0800 Subject: [PATCH 2/4] perf: optimize partial window group limit filtering --- .../benches/window_group_limit.rs | 115 +++++++++++++----- .../src/window/group_limit_processor.rs | 23 ++-- .../datafusion-ext-plans/src/window_exec.rs | 63 +++++++++- 3 files changed, 158 insertions(+), 43 deletions(-) diff --git a/native-engine/datafusion-ext-plans/benches/window_group_limit.rs b/native-engine/datafusion-ext-plans/benches/window_group_limit.rs index 5e03236fa..c884a8ea1 100644 --- a/native-engine/datafusion-ext-plans/benches/window_group_limit.rs +++ b/native-engine/datafusion-ext-plans/benches/window_group_limit.rs @@ -36,18 +36,17 @@ use datafusion_ext_plans::{ }; use test::{Bencher, black_box}; -const NUM_PARTITIONS: usize = 100; -const ROWS_PER_PARTITION: usize = 1_000; -const PEER_GROUP_SIZE: usize = 10; -const GROUP_LIMIT: usize = 10; - -fn create_batch() -> RecordBatch { - let num_rows = NUM_PARTITIONS * ROWS_PER_PARTITION; +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) + .map(|row| (row / rows_per_partition) as i32) .collect::>(); let order_values = (0..num_rows) - .map(|row| ((row % ROWS_PER_PARTITION) / PEER_GROUP_SIZE) as i32) + .map(|row| ((row % rows_per_partition) / peer_group_size) as i32) .collect::>(); let payloads = (0..num_rows as i32).collect::>(); let schema = Arc::new(Schema::new(vec![ @@ -63,8 +62,14 @@ fn create_batch() -> RecordBatch { RecordBatch::try_new(schema, columns).expect("benchmark batch should be valid") } -fn create_window_group_limit(rank_type: WindowRankType) -> Arc { - let batch = create_batch(); +fn create_window_group_limit( + rank_type: WindowRankType, + num_partitions: usize, + rows_per_partition: usize, + peer_group_size: usize, + group_limit: usize, +) -> Arc { + 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"), @@ -84,7 +89,7 @@ fn create_window_group_limit(rank_type: WindowRankType) -> Arc 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); - let expected_rows = match rank_type { - WindowRankType::RowNumber | WindowRankType::Rank => NUM_PARTITIONS * GROUP_LIMIT, - WindowRankType::DenseRank => NUM_PARTITIONS * GROUP_LIMIT * PEER_GROUP_SIZE, - }; + 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::(), - expected_rows + expected_rows( + rank_type, + num_partitions, + rows_per_partition, + peer_group_size, + group_limit, + ) ); b.iter(|| black_box(execute(&runtime, &task_ctx, &exec))); } -macro_rules! benchmark { - ($name:ident, $rank_type:expr) => { - #[bench] - fn $name(b: &mut Bencher) { - bench_window_group_limit(b, $rank_type); - } - }; +#[bench] +fn window_group_limit_row_number_slice(b: &mut Bencher) { + bench_window_group_limit(b, WindowRankType::RowNumber, 1, 100_000, 10, 10); } -benchmark!(window_group_limit_row_number, WindowRankType::RowNumber); -benchmark!(window_group_limit_rank, WindowRankType::Rank); -benchmark!(window_group_limit_dense_rank, WindowRankType::DenseRank); +#[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); +} diff --git a/native-engine/datafusion-ext-plans/src/window/group_limit_processor.rs b/native-engine/datafusion-ext-plans/src/window/group_limit_processor.rs index 0adbf098f..6621b0ab7 100644 --- a/native-engine/datafusion-ext-plans/src/window/group_limit_processor.rs +++ b/native-engine/datafusion-ext-plans/src/window/group_limit_processor.rs @@ -13,10 +13,9 @@ // See the License for the specific language governing permissions and // limitations under the License. -use arrow::{ - array::{BooleanArray, BooleanBuilder}, - record_batch::RecordBatch, -}; +use std::ops::Range; + +use arrow::record_batch::RecordBatch; use datafusion::common::Result; use crate::window::{WindowRankType, window_context::WindowContext}; @@ -46,7 +45,7 @@ impl WindowGroupLimitProcessor { &mut self, context: &WindowContext, batch: &RecordBatch, - ) -> Result { + ) -> Result>> { let partition_rows = context.get_partition_rows(batch)?; let order_rows = match self.rank_type { WindowRankType::RowNumber => None, @@ -54,7 +53,8 @@ impl WindowGroupLimitProcessor { Some(context.get_order_rows(batch)?) } }; - let mut builder = BooleanBuilder::with_capacity(batch.num_rows()); + let mut selected_ranges = vec![]; + let mut selected_start = None; for row_idx in 0..batch.num_rows() { let same_partition = !context.has_partition() || { @@ -98,8 +98,15 @@ impl WindowGroupLimitProcessor { } } } - builder.append_value(self.cur_rank <= self.limit); + 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(builder.finish()) + Ok(selected_ranges) } } diff --git a/native-engine/datafusion-ext-plans/src/window_exec.rs b/native-engine/datafusion-ext-plans/src/window_exec.rs index 208744145..ab1d9fa14 100644 --- a/native-engine/datafusion-ext-plans/src/window_exec.rs +++ b/native-engine/datafusion-ext-plans/src/window_exec.rs @@ -13,11 +13,11 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::{any::Any, fmt::Formatter, mem, sync::Arc}; +use std::{any::Any, fmt::Formatter, mem, ops::Range, sync::Arc}; use arrow::{ - array::{Array, ArrayRef, Int32Array}, - compute::concat_batches, + array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, Int32Array}, + compute::{concat_batches, filter_record_batch}, datatypes::SchemaRef, record_batch::{RecordBatch, RecordBatchOptions}, }; @@ -339,8 +339,32 @@ fn process_partial_group_limit_batch( window_ctx: &WindowContext, processor: &mut WindowGroupLimitProcessor, ) -> Result { - let selection = processor.process_batch(window_ctx, &batch)?; - Ok(arrow::compute::filter_record_batch(&batch, &selection)?) + let selected_ranges = processor.process_batch(window_ctx, &batch)?; + select_batch_ranges(batch, &selected_ranges) +} + +fn select_batch_ranges( + batch: RecordBatch, + selected_ranges: &[Range], +) -> Result { + match selected_ranges { + [] => Ok(batch.slice(0, 0)), + [range] => Ok(batch.slice(range.start, range.len())), + ranges => { + let mut selection = BooleanBufferBuilder::new(batch.num_rows()); + let mut offset = 0; + for range in ranges { + selection.append_n(range.start - offset, false); + selection.append_n(range.len(), true); + offset = range.end; + } + selection.append_n(batch.num_rows() - offset, false); + Ok(filter_record_batch( + &batch, + &BooleanArray::new(selection.finish(), None), + )?) + } + } } async fn flush_window_batches( @@ -429,7 +453,7 @@ mod test { use crate::{ agg::AggFunction, window::{WindowExpr, WindowFunction, WindowRankType}, - window_exec::WindowExec, + window_exec::{WindowExec, select_batch_ranges}, }; fn build_table_i32( @@ -468,6 +492,33 @@ mod test { )?)) } + #[test] + fn test_select_batch_ranges_uses_zero_copy_slice_for_single_range() -> Result<()> { + let batch = build_table_i32( + ("a", &vec![1, 2, 3, 4]), + ("b", &vec![5, 6, 7, 8]), + ("c", &vec![9, 10, 11, 12]), + )?; + let input_buffer = batch.column(0).to_data().buffers()[0].data_ptr(); + + let selected = select_batch_ranges(batch, &[1..3])?; + + assert_eq!(selected.num_rows(), 2); + assert_eq!( + selected + .column(0) + .as_primitive::() + .values() + .as_ref(), + &[2, 3] + ); + assert_eq!( + selected.column(0).to_data().buffers()[0].data_ptr(), + input_buffer + ); + Ok(()) + } + fn build_nullable_utf8_table( a: (&str, &Vec), b: (&str, &Vec), From 7d01400540d508ef777f442502e704a10cd95ba2 Mon Sep 17 00:00:00 2001 From: linfeng <33561138+lyne7-sc@users.noreply.github.com> Date: Fri, 4 Sep 2026 14:47:13 +0800 Subject: [PATCH 3/4] correct typo --- .../datafusion-ext-plans/benches/window_group_limit.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/native-engine/datafusion-ext-plans/benches/window_group_limit.rs b/native-engine/datafusion-ext-plans/benches/window_group_limit.rs index c884a8ea1..cc55b1314 100644 --- a/native-engine/datafusion-ext-plans/benches/window_group_limit.rs +++ b/native-engine/datafusion-ext-plans/benches/window_group_limit.rs @@ -179,7 +179,7 @@ fn window_group_limit_dense_rank_slice(b: &mut Bencher) { } #[bench] -fn window_group_limit_row/_number(b: &mut Bencher) { +fn window_group_limit_row_number(b: &mut Bencher) { bench_window_group_limit(b, WindowRankType::RowNumber, 100, 1_000, 10, 10); } From e2541090c7d3293e44c6f5819505c1a374153329 Mon Sep 17 00:00:00 2001 From: linfeng <33561138+lyne7-sc@users.noreply.github.com> Date: Fri, 4 Sep 2026 16:15:51 +0800 Subject: [PATCH 4/4] fix clippy --- native-engine/datafusion-ext-plans/src/window_exec.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/native-engine/datafusion-ext-plans/src/window_exec.rs b/native-engine/datafusion-ext-plans/src/window_exec.rs index ab1d9fa14..74097a5eb 100644 --- a/native-engine/datafusion-ext-plans/src/window_exec.rs +++ b/native-engine/datafusion-ext-plans/src/window_exec.rs @@ -501,7 +501,8 @@ mod test { )?; let input_buffer = batch.column(0).to_data().buffers()[0].data_ptr(); - let selected = select_batch_ranges(batch, &[1..3])?; + let selected_range = 1..3; + let selected = select_batch_ranges(batch, std::slice::from_ref(&selected_range))?; assert_eq!(selected.num_rows(), 2); assert_eq!(