feat(expr): support shuffle approx_percentile

e2e_test/streaming/aggregate/approx_percentile.slt

@@ -0,0 +1,104 @@
+# Single phase approx percentile
+statement ok
+create table t(p_col double, grp_col int);
+
+statement ok
+insert into t select a, 1 from generate_series(-1000, 1000) t(a);
+
+statement ok
+flush;
+
+query I
+ select
+     approx_percentile(0.01, 0.01) within group (order by p_col) as p01,
+     approx_percentile(0.1, 0.01) within group (order by p_col) as p10,
+     approx_percentile(0.5, 0.01) within group (order by p_col) as p50,
+     approx_percentile(0.9, 0.01) within group (order by p_col) as p90,
+     approx_percentile(0.99, 0.01) within group (order by p_col) as p99
+ from t group by grp_col;
+----
+-982.5779489474152 -804.4614206837127 0 804.4614206837127 982.5779489474152
+
+query I
+select
+  percentile_disc(0.01) within group (order by p_col) as p01,
+  percentile_disc(0.1) within group (order by p_col) as p10,
+  percentile_disc(0.5) within group (order by p_col) as p50,
+  percentile_disc(0.9) within group (order by p_col) as p90,
+  percentile_disc(0.99) within group (order by p_col) as p99
+from t group by grp_col;
+----
+-980 -800 0 800 980
+
+statement ok
+create materialized view m1 as
+ select
+     approx_percentile(0.01, 0.01) within group (order by p_col) as p01,
+     approx_percentile(0.1, 0.01) within group (order by p_col) as p10,
+     approx_percentile(0.5, 0.01) within group (order by p_col) as p50,
+     approx_percentile(0.9, 0.01) within group (order by p_col) as p90,
+     approx_percentile(0.99, 0.01) within group (order by p_col) as p99
+ from t group by grp_col;
+
+query I
+select * from m1;
+----
+-982.5779489474152 -804.4614206837127 0 804.4614206837127 982.5779489474152
+
+# Test state encode / decode
+onlyif can-use-recover
+statement ok
+recover;
+
+onlyif can-use-recover
+sleep 10s
+
+query I
+select * from m1;
+----
+-982.5779489474152 -804.4614206837127 0 804.4614206837127 982.5779489474152
+
+# Test state encode / decode
+onlyif can-use-recover
+statement ok
+recover;
+
+onlyif can-use-recover
+sleep 10s
+
+query I
+select * from m1;
+----
+-982.5779489474152 -804.4614206837127 0 804.4614206837127 982.5779489474152
+
+# Test 0<x<1 values
+statement ok
+insert into t select 0.001, 1 from generate_series(1, 500);
+
+statement ok
+insert into t select 0.0001, 1 from generate_series(1, 501);
+
+statement ok
+flush;
+
+query I
+select * from m1;
+----
+-963.1209598593477 -699.3618972397041 0.00009999833511933609 699.3618972397041 963.1209598593477
+
+query I
+select
+  percentile_disc(0.01) within group (order by p_col) as p01,
+  percentile_disc(0.1) within group (order by p_col) as p10,
+  percentile_disc(0.5) within group (order by p_col) as p50,
+  percentile_disc(0.9) within group (order by p_col) as p90,
+  percentile_disc(0.99) within group (order by p_col) as p99
+from t group by grp_col;
+----
+-970 -700 0.0001 700 970
+
+statement ok
+drop materialized view m1;
+
+statement ok
+drop table t;

src/expr/core/src/aggregate/def.rs

@@ -425,7 +425,8 @@ pub mod agg_kinds {
                     | PbAggKind::BoolAnd
                     | PbAggKind::BoolOr
                     | PbAggKind::ApproxCountDistinct
-                    | PbAggKind::InternalLastSeenValue,
+                    | PbAggKind::InternalLastSeenValue
+                    | PbAggKind::ApproxPercentile,
             ) | AggKind::UserDefined(_)
         };
     }

src/expr/impl/Cargo.toml

@@ -36,6 +36,7 @@ arrow-udf-python = { workspace = true, optional = true }
 arrow-udf-wasm = { workspace = true, optional = true }
 async-trait = "0.1"
 auto_enums = { workspace = true }
+bytes = "1"
 chrono = { version = "0.4", default-features = false, features = [
     "clock",
     "std",

src/expr/impl/src/aggregate/approx_percentile.rs

@@ -12,56 +12,210 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::BTreeMap;
+use std::mem::size_of;
 use std::ops::Range;
 
+use bytes::{Buf, Bytes};
 use risingwave_common::array::*;
+use risingwave_common::row::Row;
 use risingwave_common::types::*;
 use risingwave_common_estimate_size::EstimateSize;
 use risingwave_expr::aggregate::{AggCall, AggStateDyn, AggregateFunction, AggregateState};
 use risingwave_expr::{build_aggregate, Result};
 
-#[build_aggregate("approx_percentile(float8) -> float8")]
+/// TODO(kwannoel): for single phase agg, we can actually support `UDDSketch`.
+/// For two phase agg, we still use `DDSketch`.
+/// Then we also need to store the `relative_error` of the sketch, so we can report it
+/// in an internal table, if it changes.
+#[build_aggregate("approx_percentile(float8) -> float8", state = "bytea")]
 fn build(agg: &AggCall) -> Result<Box<dyn AggregateFunction>> {
-    let fraction = agg.direct_args[0]
+    let quantile = agg.direct_args[0]
         .literal()
-        .map(|x| (*x.as_float64()).into());
-    Ok(Box::new(ApproxPercentile { fraction }))
+        .map(|x| (*x.as_float64()).into())
+        .unwrap();
+    let relative_error: f64 = agg.direct_args[1]
+        .literal()
+        .map(|x| (*x.as_float64()).into())
+        .unwrap();
+    let base = (1.0 + relative_error) / (1.0 - relative_error);
+    Ok(Box::new(ApproxPercentile { quantile, base }))
 }
 
 #[allow(dead_code)]
 pub struct ApproxPercentile {
-    fraction: Option<f64>,
+    quantile: f64,
+    base: f64,
 }
 
-#[derive(Debug, Default, EstimateSize)]
-struct State(Vec<f64>);
+type BucketCount = u64;
+type BucketId = i32;
+type Count = u64;
+
+#[derive(Debug, Default)]
+struct State {
+    count: BucketCount,
+    pos_buckets: BTreeMap<BucketId, Count>,
+    zeros: Count,
+    neg_buckets: BTreeMap<BucketId, Count>,
+}
+
+impl EstimateSize for State {
+    fn estimated_heap_size(&self) -> usize {
+        let count_size = size_of::<BucketCount>();
+        let pos_buckets_size =
+            self.pos_buckets.len() * (size_of::<BucketId>() + size_of::<Count>());
+        let zero_bucket_size = size_of::<Count>();
+        let neg_buckets_size =
+            self.neg_buckets.len() * (size_of::<BucketId>() + size_of::<Count>());
+        count_size + pos_buckets_size + zero_bucket_size + neg_buckets_size
+    }
+}
 
 impl AggStateDyn for State {}
 
+impl ApproxPercentile {
+    fn add_datum(&self, state: &mut State, op: Op, datum: DatumRef<'_>) {
+        if let Some(value) = datum {
+            let prim_value = value.into_float64().into_inner();
+            let (non_neg, abs_value) = if prim_value < 0.0 {
+                (false, -prim_value)
+            } else {
+                (true, prim_value)
+            };
+            let bucket_id = abs_value.log(self.base).ceil() as BucketId;
+            match op {
+                Op::Delete | Op::UpdateDelete => {
+                    if abs_value == 0.0 {
+                        state.zeros -= 1;
+                    } else if non_neg {
+                        let count = state.pos_buckets.entry(bucket_id).or_insert(0);
+                        *count -= 1;
+                    } else {
+                        let count = state.neg_buckets.entry(bucket_id).or_insert(0);
+                        *count -= 1;
+                    }
+                    state.count -= 1;
+                }
+                Op::Insert | Op::UpdateInsert => {
+                    if abs_value == 0.0 {
+                        state.zeros += 1;
+                    } else if non_neg {
+                        let count = state.pos_buckets.entry(bucket_id).or_insert(0);
+                        *count += 1;
+                    } else {
+                        let count = state.neg_buckets.entry(bucket_id).or_insert(0);
+                        *count += 1;
+                    }
+                    state.count += 1;
+                }
+            }
+        };
+    }
+}
+
 #[async_trait::async_trait]
 impl AggregateFunction for ApproxPercentile {
     fn return_type(&self) -> DataType {
         DataType::Float64
     }
 
     fn create_state(&self) -> Result<AggregateState> {
-        todo!()
+        Ok(AggregateState::Any(Box::<State>::default()))
     }
 
-    async fn update(&self, _state: &mut AggregateState, _input: &StreamChunk) -> Result<()> {
-        todo!()
+    async fn update(&self, state: &mut AggregateState, input: &StreamChunk) -> Result<()> {
+        let state: &mut State = state.downcast_mut();
+        for (op, row) in input.rows() {
+            let datum = row.datum_at(0);
+            self.add_datum(state, op, datum);
+        }
+        Ok(())
     }
 
     async fn update_range(
         &self,
-        _state: &mut AggregateState,
-        _input: &StreamChunk,
-        _range: Range<usize>,
+        state: &mut AggregateState,
+        input: &StreamChunk,
+        range: Range<usize>,
     ) -> Result<()> {
-        todo!()
+        let state = state.downcast_mut();
+        for (op, row) in input.rows_in(range) {
+            self.add_datum(state, op, row.datum_at(0));
+        }
+        Ok(())
+    }
+
+    // TODO(kwannoel): Instead of iterating over all buckets, we can maintain the
+    // approximate quantile bucket on the fly.
+    async fn get_result(&self, state: &AggregateState) -> Result<Datum> {
+        let state = state.downcast_ref::<State>();
+        let quantile_count = (state.count as f64 * self.quantile) as u64;
+        let mut acc_count = 0;
+        for (bucket_id, count) in state.neg_buckets.iter().rev() {
+            acc_count += count;
+            if acc_count > quantile_count {
+                // approx value = -2 * y^i / (y + 1)
+                let approx_percentile = -2.0 * self.base.powi(*bucket_id) / (self.base + 1.0);
+                let approx_percentile = ScalarImpl::Float64(approx_percentile.into());
+                return Ok(Datum::from(approx_percentile));
+            }
+        }
+        acc_count += state.zeros;
+        if acc_count > quantile_count {
+            return Ok(Datum::from(ScalarImpl::Float64(0.0.into())));
+        }
+        for (bucket_id, count) in &state.pos_buckets {
+            acc_count += count;
+            if acc_count > quantile_count {
+                // approx value = 2 * y^i / (y + 1)
+                let approx_percentile = 2.0 * self.base.powi(*bucket_id) / (self.base + 1.0);
+                let approx_percentile = ScalarImpl::Float64(approx_percentile.into());
+                return Ok(Datum::from(approx_percentile));
+            }
+        }
+        return Ok(None);
+    }
+
+    fn encode_state(&self, state: &AggregateState) -> Result<Datum> {
+        let state = state.downcast_ref::<State>();
+        let mut encoded_state = Vec::with_capacity(state.estimated_heap_size());
+        encoded_state.extend_from_slice(&state.count.to_be_bytes());
+        encoded_state.extend_from_slice(&state.zeros.to_be_bytes());
+        let neg_buckets_size = state.neg_buckets.len() as u64;
+        encoded_state.extend_from_slice(&neg_buckets_size.to_be_bytes());
+        for (bucket_id, count) in &state.neg_buckets {
+            encoded_state.extend_from_slice(&bucket_id.to_be_bytes());
+            encoded_state.extend_from_slice(&count.to_be_bytes());
+        }
+        for (bucket_id, count) in &state.pos_buckets {
+            encoded_state.extend_from_slice(&bucket_id.to_be_bytes());
+            encoded_state.extend_from_slice(&count.to_be_bytes());
+        }
+        let encoded_scalar = ScalarImpl::Bytea(encoded_state.into());
+        Ok(Datum::from(encoded_scalar))
     }
 
-    async fn get_result(&self, _state: &AggregateState) -> Result<Datum> {
-        todo!()
+    fn decode_state(&self, datum: Datum) -> Result<AggregateState> {
+        let mut state = State::default();
+        let Some(scalar_state) = datum else {
+            return Ok(AggregateState::Any(Box::new(state)));
+        };
+        let encoded_state: Box<[u8]> = scalar_state.into_bytea();
+        let mut buf = Bytes::from(encoded_state);
+        state.count = buf.get_u64();
+        state.zeros = buf.get_u64();
+        let neg_buckets_size = buf.get_u64();
+        for _ in 0..neg_buckets_size {
+            let bucket_id = buf.get_i32();
+            let count = buf.get_u64();
+            state.neg_buckets.insert(bucket_id, count);
+        }
+        while !buf.is_empty() {
+            let bucket_id = buf.get_i32();
+            let count = buf.get_u64();
+            state.pos_buckets.insert(bucket_id, count);
+        }
+        Ok(AggregateState::Any(Box::new(state)))
     }
 }

src/frontend/src/optimizer/plan_node/generic/agg.rs

@@ -136,7 +136,6 @@ impl<PlanRef: GenericPlanRef> Agg<PlanRef> {
         self.agg_calls.iter().all(|c| {
             matches!(c.agg_kind, agg_kinds::single_value_state!())
                 || (matches!(c.agg_kind, agg_kinds::single_value_state_iff_in_append_only!() if stream_input_append_only))
-                || (matches!(c.agg_kind, AggKind::Builtin(PbAggKind::ApproxPercentile)))
         })
     }