From ccc9580f05f067823d367c6aa12e9995991502ae Mon Sep 17 00:00:00 2001 From: Simon Lin Date: Fri, 27 Sep 2024 16:28:32 +1000 Subject: [PATCH] c --- crates/polars-core/src/frame/group_by/proxy.rs | 2 +- crates/polars-expr/src/expressions/window.rs | 12 +++++------- 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/crates/polars-core/src/frame/group_by/proxy.rs b/crates/polars-core/src/frame/group_by/proxy.rs index adddf4e8dac7..d9aedc261faf 100644 --- a/crates/polars-core/src/frame/group_by/proxy.rs +++ b/crates/polars-core/src/frame/group_by/proxy.rs @@ -395,7 +395,7 @@ impl GroupsProxy { } } - pub fn is_sorted_flag(&self) -> bool { + pub(crate) fn is_sorted_flag(&self) -> bool { match self { GroupsProxy::Idx(groups) => groups.is_sorted_flag(), GroupsProxy::Slice { .. } => true, diff --git a/crates/polars-expr/src/expressions/window.rs b/crates/polars-expr/src/expressions/window.rs index ddff0f0975f6..b47d1744f662 100644 --- a/crates/polars-expr/src/expressions/window.rs +++ b/crates/polars-expr/src/expressions/window.rs @@ -332,16 +332,14 @@ impl WindowExpr { (WindowMapping::Join, AggState::AggregatedList(_)) => Ok(MapStrategy::Join), // no explicit aggregations, map over the groups //`(col("x").sum() * col("y")).over("groups")` - (WindowMapping::GroupsToRows, AggState::AggregatedList(_)) - if gb.get_groups().is_sorted_flag() => - { - if let GroupsProxy::Idx(_) = gb.get_groups() { - Ok(MapStrategy::Map) - } else { + (WindowMapping::GroupsToRows, AggState::AggregatedList(_)) => { + if let GroupsProxy::Slice { .. } = gb.get_groups() { + // Result can be directly exploded if the input was sorted. Ok(MapStrategy::Explode) + } else { + Ok(MapStrategy::Map) } }, - (WindowMapping::GroupsToRows, AggState::AggregatedList(_)) => Ok(MapStrategy::Map), // no aggregations, just return column // or an aggregation that has been flattened // we have to check which one