[fix](be) Keep all samples when merging percentile states

be/src/util/percentile_util.h

@@ -47,10 +47,15 @@ class Counts {
 public:
     Counts() = default;
 
+    // Moves all samples of `other` into this state. `other` may itself be a merged state whose
+    // samples live in `_sorted_nums_vec`, and either side may still hold unsorted raw samples.
     void merge(Counts* other) {
-        if (other != nullptr && !other->_nums.empty()) {
-            _sorted_nums_vec.emplace_back(std::move(other->_nums));
+        _move_nums_to_sorted_vec();
+        other->_move_nums_to_sorted_vec();
+        for (auto& nums : other->_sorted_nums_vec) {
+            _sorted_nums_vec.emplace_back(std::move(nums));
         }
+        other->_sorted_nums_vec.clear();
     }
 
     void increment(Ty key, uint32_t i) {
@@ -66,16 +71,16 @@ class Counts {
     void increment_batch(const PaddedPODArray<Ty>& keys) { _nums.insert(keys.begin(), keys.end()); }
 
     void serialize(BufferWritable& buf) {
-        if (!_nums.empty()) {
+        if (_sorted_nums_vec.empty()) {
             pdqsort(_nums.begin(), _nums.end());
-            size_t size = _nums.size();
-            buf.write_binary(size);
-            buf.write(reinterpret_cast<const char*>(_nums.data()), sizeof(Ty) * size);
         } else {
-            // convert _sorted_nums_vec to _nums and do seiralize again
+            // merge all sorted runs (including the raw samples) into `_nums`
+            _move_nums_to_sorted_vec();
             _convert_sorted_num_vec_to_nums();
-            serialize(buf);
         }
+        size_t size = _nums.size();
+        buf.write_binary(size);
+        buf.write(reinterpret_cast<const char*>(_nums.data()), sizeof(Ty) * size);
     }
 
     void unserialize(BufferReadable& buf) {
@@ -87,9 +92,13 @@ class Counts {
     }
 
     double terminate(double quantile) {
+        if (!_sorted_nums_vec.empty()) {
+            _move_nums_to_sorted_vec();
+        }
         if (_sorted_nums_vec.size() <= 1) {
             if (_sorted_nums_vec.size() == 1) {
                 _nums = std::move(_sorted_nums_vec[0]);
+                _sorted_nums_vec.clear();
             }
 
             if (_nums.empty()) {
@@ -150,6 +159,17 @@ class Counts {
         auto operator<=>(const Node& other) const { return value <=> other.value; }
     };
 
+    void _move_nums_to_sorted_vec() {
+        if (_nums.empty()) {
+            return;
+        }
+        if (!std::is_sorted(_nums.begin(), _nums.end())) {
+            pdqsort(_nums.begin(), _nums.end());
+        }
+        _sorted_nums_vec.emplace_back(std::move(_nums));
+        DCHECK(_nums.empty());
+    }
+
     void _convert_sorted_num_vec_to_nums() {
         size_t rows = 0;
         for (const auto& i : _sorted_nums_vec) {

be/test/util/percentile_util_test.cpp

@@ -135,6 +135,90 @@ TEST_F(PercentileUtilTest, CountsSerializeMergedState) {
     EXPECT_DOUBLE_EQ(9.0, restored.terminate(1.0));
 }
 
+TEST_F(PercentileUtilTest, CountsMergeMergedStates) {
+    // Two merged states whose samples only live in sorted runs, e.g. partial states
+    // combined by different pipeline instances before the final merge.
+    Counts<int64_t> run0;
+    run0.increment(0);
+    Counts<int64_t> run1;
+    run1.increment(10);
+    run1.increment(0);
+    Counts<int64_t> left;
+    left.merge(&run0);
+    Counts<int64_t> right;
+    right.merge(&run1);
+
+    Counts<int64_t> merged;
+    merged.merge(&left);
+    merged.merge(&right);
+    // samples are [0, 0, 10], (3 - 1) * 0.625 = 1.25
+    EXPECT_DOUBLE_EQ(2.5, merged.terminate(0.625));
+
+    Counts<int64_t> merged_reverse;
+    Counts<int64_t> run2;
+    run2.increment(0);
+    Counts<int64_t> run3;
+    run3.increment(0);
+    run3.increment(10);
+    Counts<int64_t> left_reverse;
+    left_reverse.merge(&run2);
+    Counts<int64_t> right_reverse;
+    right_reverse.merge(&run3);
+    merged_reverse.merge(&right_reverse);
+    merged_reverse.merge(&left_reverse);
+    EXPECT_DOUBLE_EQ(2.5, merged_reverse.terminate(0.625));
+}
+
+TEST_F(PercentileUtilTest, CountsMergeRawStates) {
+    // Both sides hold unsorted raw samples, e.g. update states of different instances.
+    Counts<int64_t> dst;
+    dst.increment(9);
+    dst.increment(1);
+    Counts<int64_t> src;
+    src.increment(7);
+    src.increment(3);
+    src.increment(5);
+    dst.merge(&src);
+    // samples are [1, 3, 5, 7, 9]
+    EXPECT_DOUBLE_EQ(1.0, dst.terminate(0.0));
+    EXPECT_DOUBLE_EQ(4.0, dst.terminate(0.375));
+    EXPECT_DOUBLE_EQ(5.0, dst.terminate(0.5));
+    EXPECT_DOUBLE_EQ(8.0, dst.terminate(0.875));
+    EXPECT_DOUBLE_EQ(9.0, dst.terminate(1.0));
+
+    Counts<int64_t> single_run_dst;
+    single_run_dst.increment(4);
+    single_run_dst.increment(2);
+    Counts<int64_t> empty_src;
+    single_run_dst.merge(&empty_src);
+    EXPECT_DOUBLE_EQ(3.0, single_run_dst.terminate(0.5));
+}
+
+TEST_F(PercentileUtilTest, CountsSerializeMixedState) {
+    Counts<int64_t> src;
+    src.increment(8);
+    src.increment(2);
+    Counts<int64_t> merged;
+    merged.merge(&src);
+    // raw samples added after a merge must be kept together with the sorted runs
+    merged.increment(6);
+    merged.increment(4);
+
+    auto col = ColumnString::create();
+    BufferWritable writer(*col);
+    merged.serialize(writer);
+    writer.commit();
+
+    StringRef data(col->get_chars().data(), col->get_chars().size());
+    BufferReadable reader(data);
+    Counts<int64_t> restored;
+    restored.unserialize(reader);
+    // samples are [2, 4, 6, 8]
+    EXPECT_DOUBLE_EQ(2.0, restored.terminate(0.0));
+    EXPECT_DOUBLE_EQ(5.0, restored.terminate(0.5));
+    EXPECT_DOUBLE_EQ(8.0, restored.terminate(1.0));
+}
+
 TEST_F(PercentileUtilTest, CheckQuantileBoundary) {
     EXPECT_NO_THROW(check_quantile(0.0));
     EXPECT_NO_THROW(check_quantile(0.5));

regression-test/data/query_p0/aggregate/percentile_bucketed_agg_merge.out

@@ -0,0 +1,23 @@
+-- This file is automatically generated. You should know what you did if you want to edit this
+-- !union_merge --
+6.25
+
+-- !union_merge_group --
+0	6
+1	6.375
+
+-- !percentile_group --
+0	6	[0, 1, 5, 8]
+1	6.375	[0, 2.25, 4.5, 10]
+
+-- !union_merge_no_bucketed --
+6.25
+
+-- !union_merge_group_no_bucketed --
+0	6
+1	6.375
+
+-- !percentile_group_no_bucketed --
+0	6	[0, 1, 5, 8]
+1	6.375	[0, 2.25, 4.5, 10]
+

regression-test/suites/query_p0/aggregate/percentile_bucketed_agg_merge.groovy

@@ -0,0 +1,93 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+// Bucketed hash aggregation merges the per-instance states of the same group in the
+// source operator. PERCENTILE states merged there must keep every sample of both sides.
+suite("percentile_bucketed_agg_merge") {
+    sql "set experimental_enable_agg_state=true"
+    sql "set enable_bucketed_hash_agg=true"
+    sql "set be_number_for_test=1"
+    sql "set agg_phase=1"
+    sql "set parallel_pipeline_task_num=4"
+    sql "set bucketed_agg_min_input_rows=0"
+    sql "set bucketed_agg_max_group_keys=0"
+    sql "set bucketed_agg_high_card_threshold=1.0"
+
+    sql "DROP TABLE IF EXISTS percentile_bucketed_agg_merge_t"
+    sql """
+        CREATE TABLE percentile_bucketed_agg_merge_t (
+            id INT NOT NULL,
+            shard INT NOT NULL,
+            v INT NOT NULL
+        )
+        DUPLICATE KEY(id)
+        DISTRIBUTED BY HASH(id) BUCKETS 8
+        PROPERTIES ('replication_num' = '1')
+    """
+    sql """
+        INSERT INTO percentile_bucketed_agg_merge_t VALUES
+            (1, 0, 0), (2, 1, 0), (4, 1, 10),
+            (5, 0, 7), (6, 1, 3), (7, 0, 1), (8, 1, 9),
+            (9, 0, 5), (10, 1, 2), (11, 0, 8), (12, 1, 6)
+    """
+
+    explain {
+        sql """
+            SELECT shard, PERCENTILE_UNION(PERCENTILE_STATE(v, CAST(0.625 AS DOUBLE)))
+            FROM percentile_bucketed_agg_merge_t GROUP BY shard
+        """
+        contains("BUCKETED AGGREGATE")
+    }
+
+    // merged PERCENTILE_UNION states of the same shard come from different instances
+    order_qt_union_merge """
+        SELECT PERCENTILE_MERGE(s) FROM (
+            SELECT shard, PERCENTILE_UNION(PERCENTILE_STATE(v, CAST(0.625 AS DOUBLE))) s
+            FROM percentile_bucketed_agg_merge_t GROUP BY shard
+        ) q
+    """
+    order_qt_union_merge_group """
+        SELECT shard, PERCENTILE_MERGE(s) FROM (
+            SELECT shard, PERCENTILE_UNION(PERCENTILE_STATE(v, CAST(0.625 AS DOUBLE))) s
+            FROM percentile_bucketed_agg_merge_t GROUP BY shard
+        ) q GROUP BY shard
+    """
+    // raw (unmerged) update states of the same shard come from different instances
+    order_qt_percentile_group """
+        SELECT shard, PERCENTILE(v, 0.625), PERCENTILE_ARRAY(v, [0, 0.25, 0.5, 1])
+        FROM percentile_bucketed_agg_merge_t GROUP BY shard
+    """
+
+    // controls: the same queries without bucketed hash aggregation
+    sql "set enable_bucketed_hash_agg=false"
+    order_qt_union_merge_no_bucketed """
+        SELECT PERCENTILE_MERGE(s) FROM (
+            SELECT shard, PERCENTILE_UNION(PERCENTILE_STATE(v, CAST(0.625 AS DOUBLE))) s
+            FROM percentile_bucketed_agg_merge_t GROUP BY shard
+        ) q
+    """
+    order_qt_union_merge_group_no_bucketed """
+        SELECT shard, PERCENTILE_MERGE(s) FROM (
+            SELECT shard, PERCENTILE_UNION(PERCENTILE_STATE(v, CAST(0.625 AS DOUBLE))) s
+            FROM percentile_bucketed_agg_merge_t GROUP BY shard
+        ) q GROUP BY shard
+    """
+    order_qt_percentile_group_no_bucketed """
+        SELECT shard, PERCENTILE(v, 0.625), PERCENTILE_ARRAY(v, [0, 0.25, 0.5, 1])
+        FROM percentile_bucketed_agg_merge_t GROUP BY shard
+    """
+}