-
Notifications
You must be signed in to change notification settings - Fork 1.3k
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
Minor: Add additional docstrings to Window function implementations #6592
Changes from 5 commits
5899adf
e1d6663
0d6e587
b6ef1aa
5e0217b
d56a5b2
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 |
---|---|---|
|
@@ -15,7 +15,7 @@ | |
// specific language governing permissions and limitations | ||
// under the License. | ||
|
||
//! partition evaluation module | ||
//! Partition evaluation module | ||
|
||
use crate::window::window_expr::BuiltinWindowState; | ||
use crate::window::WindowAggState; | ||
|
@@ -25,24 +25,97 @@ use datafusion_common::{DataFusionError, ScalarValue}; | |
use std::fmt::Debug; | ||
use std::ops::Range; | ||
|
||
/// Partition evaluator | ||
/// Partition evaluator for Window Functions | ||
/// | ||
/// # Background | ||
/// | ||
/// An implementation of this trait is created and used for each | ||
/// partition defined by an `OVER` clause and is instantiated by | ||
/// [`BuiltInWindowFunctionExpr::create_evaluator`] | ||
/// | ||
/// For example, evaluating `window_func(val) OVER (PARTITION BY col)` | ||
/// on the following data: | ||
/// | ||
/// ```text | ||
/// col | val | ||
/// --- + ---- | ||
/// A | 10 | ||
/// A | 10 | ||
/// C | 20 | ||
/// D | 30 | ||
/// D | 30 | ||
/// ``` | ||
/// | ||
/// Will instantiate three `PartitionEvaluator`s, one each for the | ||
/// partitions defined by `col=A`, `col=B`, and `col=C`. | ||
/// | ||
/// ```text | ||
/// col | val | ||
/// --- + ---- | ||
/// A | 10 <--- partition 1 | ||
/// A | 10 | ||
/// | ||
/// col | val | ||
/// --- + ---- | ||
/// C | 20 <--- partition 2 | ||
/// | ||
/// col | val | ||
/// --- + ---- | ||
/// D | 30 <--- partition 3 | ||
/// D | 30 | ||
/// ``` | ||
/// | ||
/// Different methods on this trait will be called depending on the | ||
/// capabilities described by [`BuiltInWindowFunctionExpr`]: | ||
/// | ||
/// # Stateless `PartitionEvaluator` | ||
/// | ||
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. @mustafasrepo / @ozankabak if you have time to help me describe more clearly what Stateful and Stateless 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. Some builtin window functions use window frame information inside the window expression (those are
Currently, we have support for bounded(stateful) execution for 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. Thank you @mustafasrepo -- this is super helpful. I am incorporating this information into this PR |
||
/// In this case, [`Self::evaluate`], [`Self::evaluate_with_rank`] or | ||
/// [`Self::evaluate_inside_range`] is called with values for the | ||
/// entire partition. | ||
/// | ||
/// # Stateful `PartitionEvaluator` | ||
/// | ||
/// In this case, [`Self::evaluate_stateful`] is called to calculate | ||
/// the results of the window function incrementally for each new | ||
/// batch, saving and restoring any state needed to do so as | ||
/// [`BuiltinWindowState`]. | ||
/// | ||
/// For example, when computing `ROW_NUMBER` incrementally, | ||
/// [`Self::evaluate_stateful`] will be called multiple times with | ||
/// different batches. For all batches after the first, the output | ||
/// `row_number` must start from last `row_number` produced for the | ||
/// previous batch. The previous row number is saved and restored as | ||
/// the state. | ||
/// | ||
/// [`BuiltInWindowFunctionExpr`]: crate::window::BuiltInWindowFunctionExpr | ||
/// [`BuiltInWindowFunctionExpr::create_evaluator`]: crate::window::BuiltInWindowFunctionExpr::create_evaluator | ||
pub trait PartitionEvaluator: Debug + Send { | ||
/// Whether the evaluator should be evaluated with rank | ||
/// Can this evaluator be evaluated with (only) rank | ||
/// | ||
/// If `include_rank` is true, then [`Self::evaluate_with_rank`] | ||
/// will be called for each partition, which includes the | ||
/// `rank`. | ||
fn include_rank(&self) -> bool { | ||
false | ||
} | ||
|
||
/// Returns state of the Built-in Window Function | ||
/// Returns the internal state of the window function | ||
/// | ||
/// Only used for stateful evaluation | ||
fn state(&self) -> Result<BuiltinWindowState> { | ||
// If we do not use state we just return Default | ||
Ok(BuiltinWindowState::Default) | ||
} | ||
|
||
/// Updates the internal state for Built-in window function | ||
// state is useful to update internal state for Built-in window function. | ||
// idx is the index of last row for which result is calculated. | ||
// range_columns is the result of order by column values. It is used to calculate rank boundaries | ||
// sort_partition_points is the boundaries of each rank in the range_column. It is used to update rank. | ||
/// Updates the internal state for window function | ||
/// | ||
/// Only used for stateful evaluation | ||
/// | ||
/// `state`: is useful to update internal state for window function. | ||
/// `idx`: is the index of last row for which result is calculated. | ||
/// `range_columns`: is the result of order by column values. It is used to calculate rank boundaries | ||
/// `sort_partition_points`: is the boundaries of each rank in the range_column. It is used to update rank. | ||
fn update_state( | ||
&mut self, | ||
_state: &WindowAggState, | ||
|
@@ -54,36 +127,72 @@ pub trait PartitionEvaluator: Debug + Send { | |
Ok(()) | ||
} | ||
|
||
/// Sets the internal state for window function | ||
/// | ||
/// Only used for stateful evaluation | ||
fn set_state(&mut self, _state: &BuiltinWindowState) -> Result<()> { | ||
Err(DataFusionError::NotImplemented( | ||
"set_state is not implemented for this window function".to_string(), | ||
)) | ||
} | ||
|
||
/// Gets the range where Built-in window function result is calculated. | ||
// idx is the index of last row for which result is calculated. | ||
// n_rows is the number of rows of the input record batch (Used during bound check) | ||
/// Gets the range where the window function result is calculated. | ||
/// | ||
/// `idx`: is the index of last row for which result is calculated. | ||
/// `n_rows`: is the number of rows of the input record batch (Used during bounds check) | ||
fn get_range(&self, _idx: usize, _n_rows: usize) -> Result<Range<usize>> { | ||
Err(DataFusionError::NotImplemented( | ||
"get_range is not implemented for this window function".to_string(), | ||
)) | ||
} | ||
|
||
/// Evaluate the partition evaluator against the partition | ||
/// Called for window functions that *do not use* values from the | ||
/// the window frame, such as `ROW_NUMBER`, `RANK`, `DENSE_RANK`, | ||
/// `PERCENT_RANK`, `CUME_DIST`, `LEAD`, `LAG`). | ||
fn evaluate(&self, _values: &[ArrayRef], _num_rows: usize) -> Result<ArrayRef> { | ||
Err(DataFusionError::NotImplemented( | ||
"evaluate is not implemented by default".into(), | ||
)) | ||
} | ||
|
||
/// Evaluate window function result inside given range | ||
/// Evaluate window function result inside given range. | ||
/// | ||
/// Only used for stateful evaluation | ||
fn evaluate_stateful(&mut self, _values: &[ArrayRef]) -> Result<ScalarValue> { | ||
Err(DataFusionError::NotImplemented( | ||
"evaluate_stateful is not implemented by default".into(), | ||
)) | ||
} | ||
|
||
/// evaluate the partition evaluator against the partition but with rank | ||
/// [`PartitionEvaluator::evaluate_with_rank`] is called for window | ||
/// functions that only need the rank of a row within its window | ||
/// frame. | ||
/// | ||
/// Evaluate the partition evaluator against the partition using | ||
/// the row ranks. For example, `RANK(col)` produces | ||
/// | ||
/// ```text | ||
/// col | rank | ||
/// --- + ---- | ||
/// A | 1 | ||
/// A | 1 | ||
/// C | 3 | ||
/// D | 4 | ||
/// D | 5 | ||
/// ``` | ||
/// | ||
/// For this case, `num_rows` would be `5` and the | ||
/// `ranks_in_partition` would be called with | ||
/// | ||
/// ```text | ||
/// [ | ||
/// (0,1), | ||
/// (2,2), | ||
/// (3,4), | ||
/// ] | ||
/// ``` | ||
/// | ||
/// See [`Self::include_rank`] for more details | ||
fn evaluate_with_rank( | ||
&self, | ||
_num_rows: usize, | ||
|
@@ -94,7 +203,11 @@ pub trait PartitionEvaluator: Debug + Send { | |
)) | ||
} | ||
|
||
/// evaluate window function result inside given range | ||
/// Called for window functions that use values from window frame, | ||
/// such as `FIRST_VALUE`, `LAST_VALUE`, `NTH_VALUE` and produce a | ||
/// single value for every row in the partition. | ||
/// | ||
/// Returns a [`ScalarValue`] that is the value of the window function for the entire partition | ||
fn evaluate_inside_range( | ||
&self, | ||
_values: &[ArrayRef], | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -36,12 +36,10 @@ use crate::{ | |
expressions::PhysicalSortExpr, reverse_order_bys, AggregateExpr, PhysicalExpr, | ||
}; | ||
|
||
/// A window expr that takes the form of an aggregate function | ||
/// Aggregate Window Expressions that have the form | ||
/// `OVER({ROWS | RANGE| GROUPS} BETWEEN UNBOUNDED PRECEDING AND ...)` | ||
/// e.g cumulative window frames uses `PlainAggregateWindowExpr`. Where as Aggregate Window Expressions | ||
/// that have the form `OVER({ROWS | RANGE| GROUPS} BETWEEN M {PRECEDING| FOLLOWING} AND ...)` | ||
/// e.g sliding window frames uses `SlidingAggregateWindowExpr`. | ||
/// A window expr that takes the form of an aggregate function that | ||
/// can be incrementally computed over sliding windows. | ||
/// | ||
/// See comments on [`WindowExpr`] for more details. | ||
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. Consolidated this description into |
||
#[derive(Debug)] | ||
pub struct SlidingAggregateWindowExpr { | ||
aggregate: Arc<dyn AggregateExpr>, | ||
|
@@ -72,10 +70,11 @@ impl SlidingAggregateWindowExpr { | |
} | ||
} | ||
|
||
/// peer based evaluation based on the fact that batch is pre-sorted given the sort columns | ||
/// and then per partition point we'll evaluate the peer group (e.g. SUM or MAX gives the same | ||
/// results for peers) and concatenate the results. | ||
|
||
/// Incrementally update window function using the fact that batch is | ||
/// pre-sorted given the sort columns and then per partition point. | ||
/// | ||
/// Evaluates the peer group (e.g. `SUM` or `MAX` gives the same results | ||
/// for peers) and concatenate the results. | ||
impl WindowExpr for SlidingAggregateWindowExpr { | ||
/// Return a reference to Any that can be used for downcasting | ||
fn as_any(&self) -> &dyn Any { | ||
|
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.
Maybe we can use
re-sorting
instead ofresorting
.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.
Done in d56a5b2