Evaluate the window function against the batch. This function facilitates stateful, bounded-memory implementations.
(
&self,
partition_batches: &PartitionBatches,
window_agg_state: &mut PartitionWindowAggStates,
)
| 154 | /// Evaluate the window function against the batch. This function facilitates |
| 155 | /// stateful, bounded-memory implementations. |
| 156 | fn evaluate_stateful( |
| 157 | &self, |
| 158 | partition_batches: &PartitionBatches, |
| 159 | window_agg_state: &mut PartitionWindowAggStates, |
| 160 | ) -> Result<()> { |
| 161 | let field = self.expr.field()?; |
| 162 | let out_type = field.data_type(); |
| 163 | let sort_options = self.order_by.iter().map(|o| o.options).collect::<Vec<_>>(); |
| 164 | // create a WindowAggState to clone when `window_agg_state` does not contain the respective |
| 165 | // group, which is faster than potentially creating a new one at every iteration |
| 166 | let new_state = WindowAggState::new(out_type)?; |
| 167 | for (partition_row, partition_batch_state) in partition_batches.iter() { |
| 168 | let window_state = |
| 169 | if let Some(window_state) = window_agg_state.get_mut(partition_row) { |
| 170 | window_state |
| 171 | } else { |
| 172 | let evaluator = self.expr.create_evaluator()?; |
| 173 | window_agg_state |
| 174 | .entry(partition_row.clone()) |
| 175 | .or_insert(WindowState { |
| 176 | state: new_state.clone(), |
| 177 | window_fn: WindowFn::Builtin(evaluator), |
| 178 | }) |
| 179 | }; |
| 180 | let evaluator = match &mut window_state.window_fn { |
| 181 | WindowFn::Builtin(evaluator) => evaluator, |
| 182 | _ => unreachable!(), |
| 183 | }; |
| 184 | let state = &mut window_state.state; |
| 185 | |
| 186 | let batch_ref = &partition_batch_state.record_batch; |
| 187 | let mut values = self.evaluate_args(batch_ref)?; |
| 188 | let order_bys = if evaluator.uses_window_frame() || evaluator.include_rank() { |
| 189 | get_orderby_values(self.order_by_columns(batch_ref)?) |
| 190 | } else { |
| 191 | vec![] |
| 192 | }; |
| 193 | let n_args = values.len(); |
| 194 | values.extend(order_bys); |
| 195 | let order_bys_ref = &values[n_args..]; |
| 196 | |
| 197 | // We iterate on each row to perform a running calculation. |
| 198 | let record_batch = &partition_batch_state.record_batch; |
| 199 | let num_rows = record_batch.num_rows(); |
| 200 | let mut row_wise_results: Vec<ScalarValue> = vec![]; |
| 201 | let is_causal = if evaluator.uses_window_frame() { |
| 202 | self.window_frame.is_causal() |
| 203 | } else { |
| 204 | evaluator.is_causal() |
| 205 | }; |
| 206 | for idx in state.last_calculated_index..num_rows { |
| 207 | let frame_range = if evaluator.uses_window_frame() { |
| 208 | state |
| 209 | .window_frame_ctx |
| 210 | .get_or_insert_with(|| { |
| 211 | WindowFrameContext::new( |
| 212 | Arc::clone(&self.window_frame), |
| 213 | sort_options.clone(), |
no test coverage detected