fix: Support varying quantile values per group in group_by aggregation #25606
+521
−20
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Fixes #20951.
When using
group_by()with a quantile parameter that varies per group (e.g.,pl.col.quantile.first()), all groups incorrectly received the same quantile value instead of each group using its own.Reproduction
Cause
AggQuantileExpr::evaluate_on_groups()always calledget_quantile()which evaluates the quantile expression against the full dataframe, returning a single scalar. This worked for literal quantile values but failed when the quantile expression varied per group (e.g.,first()aggregation).Fix
Added
agg_varying_quantilewhich accepts a slice of quantile values (one per group) and computes quantile per group using the existing aggregation helpers.polars-core changes:
agg_helper_idx_on_all_with_idxand_agg_helper_slice_with_idxhelpers that pass the group index to closuresagg_varying_quantile_genericthat iterates over groups with their corresponding quantile valuesagg_varying_quantilemethods toFloat32Chunked,Float64Chunked, integerChunkedArray,Series, andColumnpolars-expr changes:
AggQuantileExpr::evaluate_on_groups()now detects whether the quantile is uniform (literal/scalar) or varies per group, and dispatches to the appropriate path