fix dispatch of quantile aggregations (#3234)

pola-rs · Apr 26, 2022 · 6c6d7ad · 6c6d7ad
1 parent 7ddff81
commit 6c6d7ad
Show file tree

Hide file tree

Showing 6 changed files with 25 additions and 8 deletions.
diff --git a/polars/polars-arrow/Cargo.toml b/polars/polars-arrow/Cargo.toml
@@ -9,7 +9,7 @@ description = "Arrow interfaces for Polars DataFrame library"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-arrow = { package = "arrow2", git = "https://github.com/jorgecarleitao/arrow2", rev = "e478619ef1f07c472b4e5daa35bfa08dbd87591d", default-features = false }
+arrow = { package = "arrow2", git = "https://github.com/jorgecarleitao/arrow2", rev = "f981955f0e6c6bb233020dc01b717112ff15647c", default-features = false }
 # arrow = { package = "arrow2", git = "https://github.com/ritchie46/arrow2", branch = "polars", default-features = false }
 # arrow = { package = "arrow2", version = "0.10", default-features = false, features = ["compute_concatenate"] }
 hashbrown = "0.12"

diff --git a/polars/polars-core/Cargo.toml b/polars/polars-core/Cargo.toml
@@ -171,7 +171,7 @@ thiserror = "^1.0"
 package = "arrow2"
 git = "https://github.com/jorgecarleitao/arrow2"
 # git = "https://github.com/ritchie46/arrow2"
-rev = "e478619ef1f07c472b4e5daa35bfa08dbd87591d"
+rev = "f981955f0e6c6bb233020dc01b717112ff15647c"
 # branch = "polars"
 version = "0.10"
 default-features = false

diff --git a/polars/polars-io/Cargo.toml b/polars/polars-io/Cargo.toml
@@ -34,7 +34,7 @@ private = ["polars-time/private"]
 [dependencies]
 ahash = "0.7"
 anyhow = "1.0"
-arrow = { package = "arrow2", git = "https://github.com/jorgecarleitao/arrow2", rev = "e478619ef1f07c472b4e5daa35bfa08dbd87591d", default-features = false }
+arrow = { package = "arrow2", git = "https://github.com/jorgecarleitao/arrow2", rev = "f981955f0e6c6bb233020dc01b717112ff15647c", default-features = false }
 # arrow = { package = "arrow2", git = "https://github.com/ritchie46/arrow2", branch = "polars", default-features = false }
 # arrow = { package = "arrow2", version = "0.10", default-features = false }
 csv-core = { version = "0.1.10", optional = true }

diff --git a/polars/polars-lazy/src/physical_plan/expressions/aggregation.rs b/polars/polars-lazy/src/physical_plan/expressions/aggregation.rs
@@ -362,11 +362,14 @@ impl PhysicalExpr for AggQuantileExpr {
     #[allow(clippy::ptr_arg)]
     fn evaluate_on_groups<'a>(
         &self,
-        _df: &DataFrame,
-        _groups: &'a GroupsProxy,
-        _state: &ExecutionState,
+        df: &DataFrame,
+        groups: &'a GroupsProxy,
+        state: &ExecutionState,
     ) -> Result<AggregationContext<'a>> {
-        unimplemented!()
+        let out = self.aggregate(df, groups, state)?.ok_or_else(|| {
+            PolarsError::ComputeError("Aggregation did not return a Series".into())
+        })?;
+        Ok(AggregationContext::new(out, Cow::Borrowed(groups), true))
     }
 
     fn to_field(&self, input_schema: &Schema) -> Result<Field> {

diff --git a/py-polars/Cargo.lock b/py-polars/Cargo.lock
diff --git a/py-polars/tests/test_window.py b/py-polars/tests/test_window.py
@@ -124,3 +124,17 @@ def test_no_panic_on_nan_3067() -> None:
         4.0,
         5.0,
     ]
+
+
+def test_quantile_as_window() -> None:
+    assert (
+        pl.DataFrame(
+            {
+                "group": [0, 0, 1, 1],
+                "value": [0, 1, 0, 2],
+            }
+        )
+        .select(pl.quantile("value", 0.9).over("group"))
+        .to_series()
+        .series_equal(pl.Series("value", [1.0, 1.0, 2.0, 2.0]))
+    )