jayzhan211 commented on code in PR #27303:
URL: https://github.com/apache/datafusion/pull/27303#discussion_r4539092962
##########
datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs:
##########
@@ -624,8 +625,10 @@ mod tests {
vec![vec![(0..131072).collect()]],
first_value_udaf(),
|builder, schema| {
- Ok(builder
- .order_by(vec![PhysicalSortExpr::new_default(col("v",
schema)?)]))
+ Ok(builder.order_by(vec![
+ PhysicalSortExpr::new_default(col("v", schema)?),
+ PhysicalSortExpr::new_default(col("k", schema)?),
+ ]))
Review Comment:
@freakyzoidberg thanks for this. `FIRST_VALUE(v ORDER BY v)`, the query this
test was written for, now runs the new accumulator, so moving the test to
`ORDER BY v, k` leaves the new `shrink_to_fit` untested. If I empty it, all the
clustered tests still pass; with `ORDER BY v` restored it fails with
`ResourcesExhausted` (3.0 MB). Could you run the test for both orderings? With
the change below, each case fails when its own accumulator's `shrink_to_fit` is
emptied, and clippy is clean:
```diff
/// Like [`completed_groups_release_accumulator_scratch`], with the
native
- /// `FIRST_VALUE(v ORDER BY v, k)` accumulator. It keeps a bitmap of the
- /// groups each batch updated, which grows to one bit per group (16 KB
for
- /// 131072 groups), so it must release that bitmap too. The second
ORDER BY
- /// column keeps it off the single primitive key accumulator.
+ /// `FIRST_VALUE` accumulators: `ORDER BY v` runs the single primitive
key
+ /// one, `ORDER BY v, k` the generic one. Each keeps per-group scratch
for
+ /// the groups a batch updated (16 KB for 131072 groups), so it must
+ /// release it too.
@@
- let plan = aggregate_expr_over_sorted_input(
- AggregateMode::Partial,
- &["k"],
- vec![vec![(0..131072).collect()]],
- first_value_udaf(),
- |builder, schema| {
- Ok(builder.order_by(vec![
- PhysicalSortExpr::new_default(col("v", schema)?),
- PhysicalSortExpr::new_default(col("k", schema)?),
- ]))
- },
- )?;
- assert_eq!(plan.group_clustering_mode(),
&GroupClusteringMode::Full);
- assert_output_unchanged_under_limit(&plan, 16 * 1024).await
+ for order_by in [&["v"][..], &["v", "k"]] {
+ let plan = aggregate_expr_over_sorted_input(
+ AggregateMode::Partial,
+ &["k"],
+ vec![vec![(0..131072).collect()]],
+ first_value_udaf(),
+ |builder, schema| {
+ let order_by = order_by
+ .iter()
+ .map(|name|
Ok(PhysicalSortExpr::new_default(col(name, schema)?)))
+ .collect::<Result<Vec<_>>>()?;
+ Ok(builder.order_by(order_by))
+ },
+ )?;
+ assert_eq!(plan.group_clustering_mode(),
&GroupClusteringMode::Full);
+ assert_output_unchanged_under_limit(&plan, 16 * 1024)
+ .await
+ .map_err(|e| e.context(format!("ORDER BY {order_by:?}")))?;
+ }
+ Ok(())
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
