-
Notifications
You must be signed in to change notification settings - Fork 1.2k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Convert rank
/ dense_rank
and percent_rank
builtin functions to UDWF
#12718
Changes from 17 commits
05fd8bc
65d213c
d7f38fa
00e9044
ceba9d4
a3a4589
5a401ff
04085bb
07eb07f
937c030
f06bdc7
736fd6e
ceb594b
98dfbbf
f57c5d2
f1c7547
b9613e1
bd852e0
7ec1604
0c5d5be
38fc8ed
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change | ||||
---|---|---|---|---|---|---|
@@ -0,0 +1,200 @@ | ||||||
// 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. | ||||||
|
||||||
//! Defines physical expression for `dense_rank` that can evaluated at runtime during query execution | ||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Technically speaking this module contains the I have this PR checked out to resolve some logical conflicts anyways, so I will pushed a commit to change it in this PR too |
||||||
|
||||||
use std::any::Any; | ||||||
use std::fmt::Debug; | ||||||
use std::iter; | ||||||
use std::ops::Range; | ||||||
use std::sync::Arc; | ||||||
|
||||||
use crate::define_udwf_and_expr; | ||||||
use crate::rank::RankState; | ||||||
use datafusion_common::arrow::array::ArrayRef; | ||||||
use datafusion_common::arrow::array::UInt64Array; | ||||||
use datafusion_common::arrow::compute::SortOptions; | ||||||
use datafusion_common::arrow::datatypes::DataType; | ||||||
use datafusion_common::arrow::datatypes::Field; | ||||||
use datafusion_common::utils::get_row_at_idx; | ||||||
use datafusion_common::{Result, ScalarValue}; | ||||||
use datafusion_expr::{PartitionEvaluator, Signature, Volatility, WindowUDFImpl}; | ||||||
use datafusion_functions_window_common::field; | ||||||
use field::WindowUDFFieldArgs; | ||||||
|
||||||
define_udwf_and_expr!( | ||||||
DenseRank, | ||||||
dense_rank, | ||||||
"Returns the rank of each row within a window partition without gaps." | ||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. |
||||||
); | ||||||
|
||||||
/// dense_rank expression | ||||||
#[derive(Debug)] | ||||||
pub struct DenseRank { | ||||||
signature: Signature, | ||||||
} | ||||||
|
||||||
impl DenseRank { | ||||||
/// Create a new `dense_rank` function | ||||||
pub fn new() -> Self { | ||||||
Self { | ||||||
signature: Signature::any(0, Volatility::Immutable), | ||||||
} | ||||||
} | ||||||
} | ||||||
|
||||||
impl Default for DenseRank { | ||||||
fn default() -> Self { | ||||||
Self::new() | ||||||
} | ||||||
} | ||||||
|
||||||
impl WindowUDFImpl for DenseRank { | ||||||
fn as_any(&self) -> &dyn Any { | ||||||
self | ||||||
} | ||||||
|
||||||
fn name(&self) -> &str { | ||||||
"dense_rank" | ||||||
} | ||||||
|
||||||
fn signature(&self) -> &Signature { | ||||||
&self.signature | ||||||
} | ||||||
|
||||||
fn partition_evaluator(&self) -> Result<Box<dyn PartitionEvaluator>> { | ||||||
Ok(Box::<DenseRankEvaluator>::default()) | ||||||
} | ||||||
|
||||||
fn field(&self, field_args: WindowUDFFieldArgs) -> Result<Field> { | ||||||
Ok(Field::new(field_args.name(), DataType::UInt64, false)) | ||||||
} | ||||||
|
||||||
fn sort_options(&self) -> Option<SortOptions> { | ||||||
Some(SortOptions { | ||||||
descending: false, | ||||||
nulls_first: false, | ||||||
}) | ||||||
} | ||||||
} | ||||||
|
||||||
/// State for the `dense_rank` built-in window function. | ||||||
#[derive(Debug, Default)] | ||||||
struct DenseRankEvaluator { | ||||||
state: RankState, | ||||||
} | ||||||
|
||||||
impl PartitionEvaluator for DenseRankEvaluator { | ||||||
fn is_causal(&self) -> bool { | ||||||
// The dense_rank function doesn't need "future" values to emit results: | ||||||
true | ||||||
} | ||||||
|
||||||
fn evaluate( | ||||||
&mut self, | ||||||
values: &[ArrayRef], | ||||||
range: &Range<usize>, | ||||||
) -> Result<ScalarValue> { | ||||||
let row_idx = range.start; | ||||||
// There is no argument, values are order by column values (where rank is calculated) | ||||||
let range_columns = values; | ||||||
let last_rank_data = get_row_at_idx(range_columns, row_idx)?; | ||||||
let new_rank_encountered = | ||||||
if let Some(state_last_rank_data) = &self.state.last_rank_data { | ||||||
// if rank data changes, new rank is encountered | ||||||
state_last_rank_data != &last_rank_data | ||||||
} else { | ||||||
// First rank seen | ||||||
true | ||||||
}; | ||||||
|
||||||
if new_rank_encountered { | ||||||
self.state.last_rank_data = Some(last_rank_data); | ||||||
self.state.last_rank_boundary += self.state.current_group_count; | ||||||
self.state.current_group_count = 1; | ||||||
self.state.n_rank += 1; | ||||||
} else { | ||||||
// data is still in the same rank | ||||||
self.state.current_group_count += 1; | ||||||
} | ||||||
|
||||||
Ok(ScalarValue::UInt64(Some(self.state.n_rank as u64))) | ||||||
} | ||||||
|
||||||
fn evaluate_all_with_rank( | ||||||
&self, | ||||||
_num_rows: usize, | ||||||
ranks_in_partition: &[Range<usize>], | ||||||
) -> Result<ArrayRef> { | ||||||
let result = Arc::new(UInt64Array::from_iter_values( | ||||||
ranks_in_partition | ||||||
.iter() | ||||||
.zip(1u64..) | ||||||
.flat_map(|(range, rank)| { | ||||||
let len = range.end - range.start; | ||||||
iter::repeat(rank).take(len) | ||||||
}), | ||||||
)); | ||||||
|
||||||
Ok(result) | ||||||
} | ||||||
|
||||||
fn supports_bounded_execution(&self) -> bool { | ||||||
true | ||||||
} | ||||||
|
||||||
fn include_rank(&self) -> bool { | ||||||
true | ||||||
} | ||||||
} | ||||||
|
||||||
#[cfg(test)] | ||||||
mod tests { | ||||||
use super::*; | ||||||
use datafusion_common::cast::as_uint64_array; | ||||||
|
||||||
fn test_with_rank(expr: &DenseRank, expected: Vec<u64>) -> Result<()> { | ||||||
test_i32_result(expr, vec![0..2, 2..3, 3..6, 6..7, 7..8], expected) | ||||||
} | ||||||
|
||||||
#[allow(clippy::single_range_in_vec_init)] | ||||||
fn test_without_rank(expr: &DenseRank, expected: Vec<u64>) -> Result<()> { | ||||||
test_i32_result(expr, vec![0..8], expected) | ||||||
} | ||||||
|
||||||
fn test_i32_result( | ||||||
expr: &DenseRank, | ||||||
ranks: Vec<Range<usize>>, | ||||||
expected: Vec<u64>, | ||||||
) -> Result<()> { | ||||||
let result = expr | ||||||
.partition_evaluator()? | ||||||
.evaluate_all_with_rank(8, &ranks)?; | ||||||
let result = as_uint64_array(&result)?; | ||||||
let result = result.values(); | ||||||
assert_eq!(expected, *result); | ||||||
Ok(()) | ||||||
} | ||||||
|
||||||
#[test] | ||||||
fn test_dense_rank() -> Result<()> { | ||||||
let r = DenseRank::default(); | ||||||
test_without_rank(&r, vec![1; 8])?; | ||||||
test_with_rank(&r, vec![1, 1, 2, 3, 3, 3, 4, 5])?; | ||||||
Ok(()) | ||||||
} | ||||||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
the expr function have written docs for the percent_rank,
but when I try to import the percent_rank from datafusion_functions_window workspace,
its throwing error related to the circular depency
how to handle such cases ?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
There are a couple of ways to get around this. But here I noticed that the
first_value
function is a mock and not the realfirst_value
window function.So you can consider doing the same for
percent_rank
. Then you don't have to import anything. This is alright because the test is not testing the functionality ofpercent_rank
in the doc test.Another advantage of mocking in this specific case is that the change remains local to the doc test which I believe is good.