From 7b1d19aeb320ac4264b50ea64c4c5f73dfe771ac Mon Sep 17 00:00:00 2001 From: Orson Peters Date: Thu, 26 Sep 2024 13:34:21 +0200 Subject: [PATCH] fix: Incorrect mode for sorted input --- crates/polars-ops/src/chunked_array/mode.rs | 31 ++++++++------------- 1 file changed, 12 insertions(+), 19 deletions(-) diff --git a/crates/polars-ops/src/chunked_array/mode.rs b/crates/polars-ops/src/chunked_array/mode.rs index a36b161775ca..1c981c9eb3e2 100644 --- a/crates/polars-ops/src/chunked_array/mode.rs +++ b/crates/polars-ops/src/chunked_array/mode.rs @@ -1,4 +1,3 @@ -use arrow::legacy::utils::CustomIterTools; use polars_core::prelude::*; use polars_core::{with_match_physical_integer_polars_type, POOL}; @@ -33,29 +32,23 @@ fn mode_64(ca: &Float64Chunked) -> PolarsResult { fn mode_indices(groups: GroupsProxy) -> Vec { match groups { GroupsProxy::Idx(groups) => { - let mut groups = groups.into_iter().collect_trusted::>(); - groups.sort_unstable_by_key(|k| k.1.len()); - let last = &groups.last().unwrap(); - let max_occur = last.1.len(); + let Some(max_len) = groups.iter().map(|g| g.1.len()).max() else { + return Vec::new(); + }; groups - .iter() - .rev() - .take_while(|v| v.1.len() == max_occur) - .map(|v| v.0) + .into_iter() + .filter(|g| g.1.len() == max_len) + .map(|g| g.0) .collect() }, GroupsProxy::Slice { groups, .. } => { - let last = groups.last().unwrap(); - let max_occur = last[1]; - + let Some(max_len) = groups.iter().map(|g| g[1]).max() else { + return Vec::new(); + }; groups - .iter() - .rev() - .take_while(|v| { - let len = v[1]; - len == max_occur - }) - .map(|v| v[0]) + .into_iter() + .filter(|g| g[1] == max_len) + .map(|g| g[0]) .collect() }, }