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..cc55b1314 --- /dev/null +++ b/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::>(); + 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, + 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"), + ); + 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 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::(), + 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); +} 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..6621b0ab7 --- /dev/null +++ b/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, + 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 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) + } +} 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..74097a5eb 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}, }; @@ -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,39 @@ fn execute_window( })) } +fn process_partial_group_limit_batch( + batch: RecordBatch, + window_ctx: &WindowContext, + processor: &mut WindowGroupLimitProcessor, +) -> Result { + 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( staging_batches: &mut Vec, window_ctx: &WindowContext, @@ -391,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( @@ -430,6 +492,34 @@ 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_range = 1..3; + let selected = select_batch_ranges(batch, std::slice::from_ref(&selected_range))?; + + 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), @@ -790,6 +880,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();