fix bug in rank of single series

pola-rs · Oct 14, 2021 · b30de44 · b30de44
1 parent 692c6ac
commit b30de44
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 1 deletion.
diff --git a/polars/polars-core/src/chunked_array/ops/unique/rank.rs b/polars/polars-core/src/chunked_array/ops/unique/rank.rs
@@ -10,6 +10,13 @@ pub enum RankMethod {
 }
 
 pub(crate) fn rank(s: &Series, method: RankMethod) -> Series {
+    if s.len() == 1 {
+        return match method {
+            Average => Series::new(s.name(), &[1.0]),
+            _ => Series::new(s.name(), &[1u32]),
+        };
+    }
+
     // See: https://github.com/scipy/scipy/blob/v1.7.1/scipy/stats/stats.py#L8631-L8737
 
     let len = s.len();

diff --git a/polars/polars-core/src/utils.rs b/polars/polars-core/src/utils.rs
@@ -280,7 +280,7 @@ macro_rules! match_arrow_data_type_apply_macro {
             DataType::Int64 => $macro!(Int64Type $(, $opt_args)*),
             DataType::Float32 => $macro!(Float32Type $(, $opt_args)*),
             DataType::Float64 => $macro!(Float64Type $(, $opt_args)*),
-            _ => unimplemented!(),
+            dt => panic!("not implemented for dtype {:?}", dt),
         }
     }};
 }

diff --git a/polars/polars-lazy/src/test.rs b/polars/polars-lazy/src/test.rs
@@ -1925,3 +1925,20 @@ fn test_exclude_regex() -> Result<()> {
     assert_eq!(out.get_column_names(), &["A", "B"]);
     Ok(())
 }
+
+#[test]
+fn test_groupby_rank() -> Result<()> {
+    let df = fruits_cars();
+    let out = df
+        .lazy()
+        .stable_groupby([col("cars")])
+        .agg([col("B").rank(RankMethod::Dense)])
+        .collect()?;
+
+    let out = out.column("B")?;
+    let out = out.list()?.get(1).unwrap();
+    let out = out.u32()?;
+
+    assert_eq!(Vec::from(out), &[Some(1)]);
+    Ok(())
+}