Trait datafusion_expr::PartitionEvaluator

pub trait PartitionEvaluator: Debug + Send {
    // Provided methods
    fn memoize(&mut self, _state: &mut WindowAggState) -> Result<()> { ... }
    fn get_range(&self, idx: usize, _n_rows: usize) -> Result<Range<usize>> { ... }
    fn evaluate_all(
        &mut self,
        values: &[ArrayRef],
        num_rows: usize
    ) -> Result<ArrayRef> { ... }
    fn evaluate(
        &mut self,
        _values: &[ArrayRef],
        _range: &Range<usize>
    ) -> Result<ScalarValue> { ... }
    fn evaluate_all_with_rank(
        &self,
        _num_rows: usize,
        _ranks_in_partition: &[Range<usize>]
    ) -> Result<ArrayRef> { ... }
    fn supports_bounded_execution(&self) -> bool { ... }
    fn uses_window_frame(&self) -> bool { ... }
    fn include_rank(&self) -> bool { ... }
}

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 the DataFusion runtime.

For example, evaluating window_func(val) OVER (PARTITION BY col) on the following data:

col | val
--- + ----
 A  | 10
 A  | 10
 C  | 20
 D  | 30
 D  | 30

Will instantiate three PartitionEvaluators, one each for the partitions defined by col=A, col=B, and col=C.

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 supports_bounded_execution, uses_window_frame, and include_rank,

When implementing a new PartitionEvaluator, implement corresponding evaluator according to table below.

Implementation Table

uses_window_frame	supports_bounded_execution	include_rank	function_to_implement
false (default)	false (default)	false (default)	evaluate_all
false	true	false	evaluate
false	true/false	true	evaluate_all_with_rank
true	true/false	true/false	evaluate

Provided Methods

fn memoize(&mut self, _state: &mut WindowAggState) -> Result<()>

When the window frame has a fixed beginning (e.g UNBOUNDED PRECEDING), some functions such as FIRST_VALUE, LAST_VALUE and NTH_VALUE do not need the (unbounded) input once they have seen a certain amount of input.

memoize is called after each input batch is processed, and such functions can save whatever they need and modify WindowAggState appropriately to allow rows to be pruned

fn get_range(&self, idx: usize, _n_rows: usize) -> Result<Range<usize>>

If uses_window_frame flag is false. This method is used to calculate required range for the window function during stateful execution.

Generally there is no required range, hence by default this returns smallest range(current row). e.g seeing current row is enough to calculate window result (such as row_number, rank, etc)

fn evaluate_all( &mut self, values: &[ArrayRef], num_rows: usize ) -> Result<ArrayRef>

Evaluate a window function on an entire input partition.

This function is called once per input partition for window functions that do not use values from the window frame, such as ROW_NUMBER, RANK, DENSE_RANK, PERCENT_RANK, CUME_DIST, LEAD, LAG).

It produces the result of all rows in a single pass. It expects to receive the entire partition as the value and must produce an output column with one output row for every input row.

num_rows is requied to correctly compute the output in case values.len() == 0

Implementing this function is an optimization: certain window functions are not affected by the window frame definition or the query doesn’t have a frame, and evaluate skips the (costly) window frame boundary calculation and the overhead of calling evaluate for each output row.

For example, the LAG built in window function does not use the values of its window frame (it can be computed in one shot on the entire partition with Self::evaluate_all regardless of the window defined in the OVER clause)

lag(x, 1) OVER (ORDER BY z ROWS BETWEEN 2 PRECEDING AND 3 FOLLOWING)

However, avg() computes the average in the window and thus does use its window frame

avg(x) OVER (PARTITION BY y ORDER BY z ROWS BETWEEN 2 PRECEDING AND 3 FOLLOWING)

fn evaluate( &mut self, _values: &[ArrayRef], _range: &Range<usize> ) -> Result<ScalarValue>

Evaluate window function on a range of rows in an input partition.x

This is the simplest and most general function to implement but also the least performant as it creates output one row at a time. It is typically much faster to implement stateful evaluation using one of the other specialized methods on this trait.

Returns a ScalarValue that is the value of the window function within range for the entire partition. Argument values contains the evaluation result of function arguments and evaluation results of ORDER BY expressions. If function has a single argument, values[1..] will contain ORDER BY expression results.

fn evaluate_all_with_rank( &self, _num_rows: usize, _ranks_in_partition: &[Range<usize>] ) -> Result<ArrayRef>

PartitionEvaluator::evaluate_all_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

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

[
  (0,1),
  (2,2),
  (3,4),
]

fn supports_bounded_execution(&self) -> bool

Can the window function be incrementally computed using bounded memory?

See the table on Self for what functions to implement

fn uses_window_frame(&self) -> bool

Does the window function use the values from the window frame, if one is specified?

See the table on Self for what functions to implement

fn include_rank(&self) -> bool

Can this function be evaluated with (only) rank

See the table on Self for what functions to implement

Implementors