This is an automated email from the ASF dual-hosted git repository.

yiguolei pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/doris.git


The following commit(s) were added to refs/heads/master by this push:
     new 7b9693d17a0 [fix](be) Fix DataSketches HLL union accuracy and add 
configurable lg… (#67469)
7b9693d17a0 is described below

commit 7b9693d17a0de1e0c062dc006812a5963d169bc4
Author: nooneuse <[email protected]>
AuthorDate: Wed Sep 16 12:13:55 2026 +0800

    [fix](be) Fix DataSketches HLL union accuracy and add configurable lg… 
(#67469)
    
    ### What problem does this PR solve?
    
    Problem Summary:
    
    `DATASKETCHES_HLL_UNION_AGG` merges serialized Apache DataSketches HLL
    sketches. The previous implementation had two accuracy problems.
    
    First, Doris used `datasketches-cpp` 5.2.0, which contains an upstream
    regression in the lazy KxQ/`curMin` rebuild path. When an HLL-mode
    sketch was downsampled during union, the register array could be updated
    while the cached estimator state was still pending rebuild. Later
    updates, estimates, or serialization could then observe stale state.
    Depending on the merge order, this could produce an incorrect estimate,
    lose previously accumulated data, or serialize different results for
    equivalent unions.
    
    Second, Doris initialized the union limit from the first input sketch.
    This is unsafe for sparse LIST/SET sketches: they still contain exact
    coupons, so their configured `lgK` should not prematurely cap the union.
    Because the first input is nondeterministic in distributed aggregation,
    the effective precision and memory usage could depend on input order.
    
    This PR fixes both problems:
    
    - Updates the `datasketches-cpp` submodule from the 5.2.0 commit
    `de8553ba` to the exact upstream fix commit
    
[`46025e9`](https://github.com/apache/datasketches-cpp/commit/46025e9aeed8368b1184cbde9634dd99d0ee47c0).
    - Ignores empty sketches when initializing or merging an aggregate
    state.
    - Uses `MAX_LOG_K=21` as the internal upper bound for the one-argument
    form, preventing the first sparse sketch from unnecessarily reducing the
    available precision.
    - Adds an optional constant `lg_max_k` argument in the inclusive range
    `[7, 21]`, allowing users to choose a predictable precision and memory
    upper bound.
    
    ### Function behavior
    
    The function now supports both forms:
    
    ```sql
    DATASKETCHES_HLL_UNION_AGG(sketch)
    DATASKETCHES_HLL_UNION_AGG(sketch, lg_max_k)
    ```
    
    The aliases support the same optional argument:
    
    ```sql
    DS_HLL_ESTIMATE(sketch, lg_max_k)
    DATASKETCHES_HLL_ESTIMATE(sketch, lg_max_k)
    ```
    
    For the two-argument form, `lg_max_k` must be a constant integer in `[7,
    21]`. FE validates its type, constness, and range, and BE also validates
    the runtime range.
    
    `lg_max_k` is an upper bound rather than a request to force the final
    sketch to that precision. A dense input or intermediate state with a
    smaller effective `lgK` lowers the resulting union because DataSketches
    cannot restore precision by upsampling an already dense sketch.
    
    AggState types record the argument types but do not encode the literal
    value of `lg_max_k`. Therefore, states created with values such as `7`
    and `21` have the same SQL state type. When such states are combined,
    the result deterministically uses the minimum effective `lgK` carried by
    the serialized states.
    
    ### Compatibility and resource usage
    
    - The one-argument SQL signature is unchanged. Internally, its union cap
    is now `21`, but the resulting effective `lgK` can still be lower when a
    lower-precision dense sketch is consumed.
    - A cap of `21` does not immediately allocate a dense 2 MiB HLL array.
    LIST/SET sketches remain sparse. The higher memory usage appears only
    when an aggregate state reaches dense HLL mode at a high effective
    `lgK`, for example after accumulating enough exact coupons or consuming
    a dense high-`lgK` sketch.
    - Dense HLL_8 storage is approximately `2^lgK` bytes per aggregate
    state, excluding object and allocator overhead. `lgK=12` is
    approximately 4 KiB, while `lgK=21` is approximately 2 MiB.
    - Workloads that need a lower and predictable per-state memory ceiling
    should use the two-argument form explicitly.
    - The two-argument aggregate must only be used after every BE in the
    cluster has been upgraded. Older BEs do not recognize the new function
    arity.
    - The existing one-argument serialized state remains format-compatible,
    but an unupgraded BE still runs the old DataSketches implementation and
    does not contain the upstream accuracy fix.
    - No released `datasketches-cpp` version newer than 5.2.0 contains this
    fix, so the dependency is pinned to the exact reviewed upstream commit
    instead of following an upstream branch.
    
    ### Release note
    
    Fixed incorrect and merge-order-dependent results in
    `DATASKETCHES_HLL_UNION_AGG` for affected serialized HLL sketches. Added
    an optional constant `lg_max_k` argument in `[7, 21]` to control the
    union's precision and memory upper bound. The new two-argument form must
    only be used after all BEs have been upgraded.
    
    ### Check List (For Author)
    
    - Test <!-- At least one of them must be included. -->
        - [x] Regression test
        - [x] Unit Test
        - [ ] Manual test (add detailed scripts or steps below)
        - [ ] No need to test or manual test. Explain why:
    - [ ] This is a refactor/code format and no logic has been changed.
            - [ ] Previous test can cover this change.
            - [ ] No code files have been changed.
            - [ ] Other reason <!-- Add your reason?  -->
    
    - Behavior changed:
        - [ ] No.
        - [x] Yes. See "Important Behavior and Upgrade Notes"
    
    - Does this need documentation?
        - [ ] No.
    - [x] Yes. <!-- Add document PR link here. eg:
    https://github.com/apache/doris-website/pull/1214 -->
    
    ### Check List (For Reviewer who merge this PR)
    
    - [ ] Confirm the release note
    - [ ] Confirm test cases
    - [ ] Confirm document
    - [ ] Add branch pick label <!-- Add branch pick label that this PR
    should merge into -->
    
    ---------
    
    Co-authored-by: yuanyuhao <[email protected]>
    Co-authored-by: Tiewei Fang <[email protected]>
---
 ...gregate_function_datasketches_hll_union_agg.cpp |   2 +
 ...aggregate_function_datasketches_hll_union_agg.h |  81 ++++---
 .../agg_datasketches_hll_union_agg_test.cpp        | 238 +++++++++++++++++++++
 contrib/datasketches-cpp                           |   2 +-
 .../functions/agg/DataSketchesHllUnionAgg.java     |  51 ++++-
 .../rules/rewrite/EliminateAggCaseWhenTest.java    |  79 +++++++
 .../nereids/rules/rewrite/InferAggNotNullTest.java |  25 +++
 .../functions/agg/DataSketchesHllUnionAggTest.java | 190 ++++++++++++++++
 .../test_datasketches_hll_union_agg.out            |  24 +++
 ...st_datasketches_hll_union_agg_null_ignoring.out |  52 +++++
 .../test_datasketches_hll_union_agg.groovy         | 120 +++++++++++
 ...datasketches_hll_union_agg_null_ignoring.groovy | 116 ++++++++++
 12 files changed, 946 insertions(+), 34 deletions(-)

diff --git 
a/be/src/exprs/aggregate/aggregate_function_datasketches_hll_union_agg.cpp 
b/be/src/exprs/aggregate/aggregate_function_datasketches_hll_union_agg.cpp
index c9b7013e7a9..d8496cfb4fc 100644
--- a/be/src/exprs/aggregate/aggregate_function_datasketches_hll_union_agg.cpp
+++ b/be/src/exprs/aggregate/aggregate_function_datasketches_hll_union_agg.cpp
@@ -23,12 +23,14 @@
 #include "core/data_type/define_primitive_type.h"
 #include "exec/common/hash_table/hash.h" // IWYU pragma: keep
 #include "exprs/aggregate/aggregate_function_simple_factory.h"
+#include "exprs/aggregate/factory_helpers.h"
 #include "exprs/aggregate/helpers.h"
 namespace doris {
 template <template <PrimitiveType> class Data>
 AggregateFunctionPtr create_aggregate_function_datasketches_hll_union_agg(
         const std::string& name, const DataTypes& argument_types, const 
DataTypePtr& result_type,
         const bool result_is_nullable, const AggregateFunctionAttr& attr) {
+    assert_arity_range(name, argument_types, 1, 2);
     return creator_with_type_list<TYPE_STRING, TYPE_VARCHAR, 
TYPE_VARBINARY>::create<
             AggregateFunctionDataSketchesHllUnionAgg, Data>(argument_types, 
result_is_nullable,
                                                             attr);
diff --git 
a/be/src/exprs/aggregate/aggregate_function_datasketches_hll_union_agg.h 
b/be/src/exprs/aggregate/aggregate_function_datasketches_hll_union_agg.h
index 19d0a061814..c0575dadd4c 100644
--- a/be/src/exprs/aggregate/aggregate_function_datasketches_hll_union_agg.h
+++ b/be/src/exprs/aggregate/aggregate_function_datasketches_hll_union_agg.h
@@ -53,11 +53,9 @@ class ColumnDecimal;
 /// datasketches_hll_union_agg
 template <PrimitiveType T>
 struct AggregateFunctionHllSketchData {
-    /** We set the default LgK to 12,
-      * as this value is used as a performance baseline in the relevant 
documentation.
-      * (https://datasketches.apache.org/docs/HLL/HllPerformance.html)
-      */
-    static constexpr uint8_t DEFAULT_LOG_K = 12;
+    static constexpr uint8_t EMPTY_STATE_LOG_K = 12;
+    static constexpr uint8_t MIN_UNION_LOG_K = 7;
+    static constexpr uint8_t DEFAULT_UNION_LOG_K = 
datasketches::hll_constants::MAX_LOG_K;
     using Alloc = CustomStdAllocator<uint8_t>;
     using Sketch = datasketches::hll_sketch_alloc<Alloc>;
     using Union = datasketches::hll_union_alloc<Alloc>;
@@ -66,17 +64,12 @@ struct AggregateFunctionHllSketchData {
 
     static String get_name() { return "datasketches_hll_union_agg"; }
 
-    void merge(const Sketch& sketch_data) {
+    void merge(const Sketch& sketch_data, uint8_t lg_max_k) {
+        if (sketch_data.is_empty()) {
+            return;
+        }
         if (!hll_union_data.has_value()) {
-            /** We clamp max lg_k to [7, 21],
-              * considering that the code comment requires 7 to 21.
-              * See: datasketches-cpp/hll/include/hll.hpp:451
-              */
-            constexpr uint8_t MIN_UNION_LOG_K = 7;
-            const uint8_t union_lg_k =
-                    std::clamp<uint8_t>(sketch_data.get_lg_config_k(), 
MIN_UNION_LOG_K,
-                                        
datasketches::hll_constants::MAX_LOG_K);
-            hll_union_data.emplace(union_lg_k, Alloc());
+            hll_union_data.emplace(lg_max_k, Alloc());
         }
         try {
             hll_union_data->update(sketch_data);
@@ -91,12 +84,32 @@ struct AggregateFunctionHllSketchData {
                             "Internal error happened when update HLL sketch: 
unknown exception.");
         }
     }
-    void reset() {
-        if (hll_union_data.has_value()) {
-            hll_union_data->reset();
+    void merge(const Sketch& sketch_data) {
+        if (sketch_data.is_empty()) {
+            return;
         }
-        hll_union_data.reset();
+        const auto lg_max_k = std::max<uint8_t>(sketch_data.get_lg_config_k(), 
MIN_UNION_LOG_K);
+        if (hll_union_data.has_value() && lg_max_k < 
hll_union_data->get_lg_config_k()) {
+            try {
+                // Sparse sketch updates do not lower an existing union's lgK.
+                auto current = hll_union_data->get_result(datasketches::HLL_8);
+                hll_union_data.emplace(lg_max_k, Alloc());
+                hll_union_data->update(current);
+            } catch (const doris::Exception& e) {
+                throw Exception(e.code(), "Internal error happened when update 
HLL sketch: {}",
+                                e.to_string());
+            } catch (const std::exception& e) {
+                throw Exception(ErrorCode::INTERNAL_ERROR,
+                                "Internal error happened when update HLL 
sketch: {}", e.what());
+            } catch (...) {
+                throw Exception(
+                        ErrorCode::INTERNAL_ERROR,
+                        "Internal error happened when update HLL sketch: 
unknown exception.");
+            }
+        }
+        merge(sketch_data, lg_max_k);
     }
+    void reset() { hll_union_data.reset(); }
 
     void write_sketch(BufferWritable& buf, const Sketch& sk) const {
         auto serialized_bytes = sk.serialize_compact();
@@ -105,10 +118,7 @@ struct AggregateFunctionHllSketchData {
     }
     void write(BufferWritable& buf) const {
         if (!hll_union_data.has_value()) {
-            /** Using DEFAULT_LOG_K(12) here is surely sufficient,
-              * because in this case the union that actually needs to be 
serialized should contain no data.
-              */
-            Union u(DEFAULT_LOG_K, Alloc());
+            Union u(EMPTY_STATE_LOG_K, Alloc());
             write_sketch(buf, u.get_result());
             return;
         }
@@ -175,7 +185,7 @@ template <PrimitiveType T, typename Data>
 class AggregateFunctionDataSketchesHllUnionAgg final
         : public IAggregateFunctionDataHelper<Data,
                                               
AggregateFunctionDataSketchesHllUnionAgg<T, Data>>,
-          UnaryExpression,
+          VarargsExpression,
           NotNullableAggregateFunction {
 public:
     AggregateFunctionDataSketchesHllUnionAgg(const DataTypes& argument_types_)
@@ -186,7 +196,20 @@ public:
     void reset(AggregateDataPtr __restrict place) const override { 
this->data(place).reset(); }
     void add(AggregateDataPtr __restrict place, const IColumn** columns, 
ssize_t row_num,
              Arena&) const override {
-        add_one(this->data(place), *columns[0], row_num);
+        uint8_t lg_max_k = Data::DEFAULT_UNION_LOG_K;
+        if (this->argument_types.size() == 2) {
+            const auto value =
+                    assert_cast<const ColumnInt32&, 
TypeCheckOnRelease::DISABLE>(*columns[1])
+                            .get_element(row_num);
+            if (value < Data::MIN_UNION_LOG_K || value > 
datasketches::hll_constants::MAX_LOG_K) {
+                throw Exception(ErrorCode::INVALID_ARGUMENT,
+                                "lg_max_k must be between {} and {}, but was 
{}",
+                                Data::MIN_UNION_LOG_K, 
datasketches::hll_constants::MAX_LOG_K,
+                                value);
+            }
+            lg_max_k = static_cast<uint8_t>(value);
+        }
+        add_one(this->data(place), *columns[0], row_num, lg_max_k);
     }
     void merge(AggregateDataPtr __restrict place, ConstAggregateDataPtr rhs,
                Arena&) const override {
@@ -214,10 +237,14 @@ public:
             this->template check_argument_column_type<typename 
PrimitiveTypeTraits<T>::ColumnType>(
                     columns[0]);
         }
+        if (this->argument_types.size() == 2) {
+            this->template check_argument_column_type<ColumnInt32>(columns[1]);
+        }
     }
 
 private:
-    static void ALWAYS_INLINE add_one(Data& data, const IColumn& column, 
ssize_t row_num) {
+    static void ALWAYS_INLINE add_one(Data& data, const IColumn& column, 
ssize_t row_num,
+                                      uint8_t lg_max_k) {
         if constexpr (is_string_type(T) || is_varbinary(T)) {
             const auto& src_column = assert_cast<const typename 
PrimitiveTypeTraits<T>::ColumnType&,
                                                  
TypeCheckOnRelease::DISABLE>(column);
@@ -245,7 +272,7 @@ private:
                 }
             }();
 
-            data.merge(sketch_data);
+            data.merge(sketch_data, lg_max_k);
         }
     }
 };
diff --git a/be/test/exprs/aggregate/agg_datasketches_hll_union_agg_test.cpp 
b/be/test/exprs/aggregate/agg_datasketches_hll_union_agg_test.cpp
index eeaeb2a45dc..12b5df37841 100644
--- a/be/test/exprs/aggregate/agg_datasketches_hll_union_agg_test.cpp
+++ b/be/test/exprs/aggregate/agg_datasketches_hll_union_agg_test.cpp
@@ -17,6 +17,8 @@
 
 #include <gtest/gtest.h>
 
+#include <algorithm>
+#include <cmath>
 #include <hll.hpp>
 
 #include "agent/be_exec_version_manager.h"
@@ -41,10 +43,21 @@ void register_aggregate_function_datasketches_HLL_union_agg(
 
 class AggregateFunctionDataSketchesHllUnionAggTest : public ::testing::Test {
 protected:
+    using Data = AggregateFunctionHllSketchData<TYPE_STRING>;
+    using Sketch = Data::Sketch;
+
     void SetUp() override { arena = std::make_unique<Arena>(); }
 
     void TearDown() override { arena.reset(); }
 
+    static Sketch create_sketch(uint8_t lg_k, uint64_t start, uint64_t count) {
+        Sketch sketch(lg_k, datasketches::HLL_8, false, Data::Alloc());
+        for (uint64_t value = start; value < start + count; ++value) {
+            sketch.update(value);
+        }
+        return sketch;
+    }
+
     std::unique_ptr<Arena> arena;
 };
 
@@ -414,6 +427,17 @@ TEST_F(AggregateFunctionDataSketchesHllUnionAggTest, 
testFactoryCreateAndAliases
     ASSERT_NE(fn_alias_sr_estimate, nullptr);
     ASSERT_NE(fn_alias_datasketches_estimate, nullptr);
 
+    DataTypes capped_argument_types = {std::make_shared<DataTypeString>(),
+                                       std::make_shared<DataTypeInt32>()};
+    ASSERT_NE(factory.get("datasketches_hll_union_agg", capped_argument_types, 
nullptr, false,
+                          be_version),
+              nullptr);
+    ASSERT_NE(factory.get("ds_hll_estimate", capped_argument_types, nullptr, 
false, be_version),
+              nullptr);
+    ASSERT_NE(factory.get("datasketches_hll_estimate", capped_argument_types, 
nullptr, false,
+                          be_version),
+              nullptr);
+
     datasketches::hll_sketch sketch(8, datasketches::HLL_8);
     for (int i = 0; i < 7; ++i) sketch.update(i);
     const auto ser = sketch.serialize_compact();
@@ -426,6 +450,8 @@ TEST_F(AggregateFunctionDataSketchesHllUnionAggTest, 
testFactoryCreateAndAliases
         AggregateDataPtr place = arena->aligned_alloc(fn->size_of_data(), 
fn->align_of_data());
         fn->create(place);
         fn->add(place, columns, 0, *arena);
+        const auto& data = *reinterpret_cast<const Data*>(place);
+        EXPECT_EQ(data.hll_union_data.value().get_lg_config_k(), 
Data::DEFAULT_UNION_LOG_K);
         ColumnFloat64 result;
         fn->insert_result_into(place, result);
         fn->destroy(place);
@@ -536,6 +562,218 @@ TEST_F(AggregateFunctionDataSketchesHllUnionAggTest, 
testLowLgKSketchDoesNotRepo
     agg_func->destroy(place2);
 }
 
+TEST_F(AggregateFunctionDataSketchesHllUnionAggTest, 
testDefaultLgMaxKIsOrderIndependent) {
+    auto sparse = create_sketch(8, 0, 7);
+    auto dense = create_sketch(16, 1000, 10000);
+
+    EXPECT_EQ(Data::DEFAULT_UNION_LOG_K, 
datasketches::hll_constants::MAX_LOG_K);
+
+    Data sparse_first;
+    sparse_first.merge(sparse, Data::DEFAULT_UNION_LOG_K);
+    ASSERT_TRUE(sparse_first.hll_union_data.has_value());
+    EXPECT_EQ(sparse_first.hll_union_data->get_lg_config_k(), 
Data::DEFAULT_UNION_LOG_K);
+    sparse_first.merge(dense, Data::DEFAULT_UNION_LOG_K);
+    EXPECT_EQ(sparse_first.hll_union_data->get_lg_config_k(), 16);
+
+    Data dense_first;
+    dense_first.merge(dense, Data::DEFAULT_UNION_LOG_K);
+    dense_first.merge(sparse, Data::DEFAULT_UNION_LOG_K);
+    ASSERT_TRUE(dense_first.hll_union_data.has_value());
+    EXPECT_EQ(dense_first.hll_union_data->get_lg_config_k(), 16);
+}
+
+TEST_F(AggregateFunctionDataSketchesHllUnionAggTest, 
testEmptySketchDoesNotInitializeUnion) {
+    Sketch empty(4, datasketches::HLL_8, true, Data::Alloc());
+    Data data;
+
+    data.merge(empty, Data::DEFAULT_UNION_LOG_K);
+    EXPECT_FALSE(data.hll_union_data.has_value());
+
+    data.merge(create_sketch(12, 0, 10000), Data::DEFAULT_UNION_LOG_K);
+    ASSERT_TRUE(data.hll_union_data.has_value());
+    EXPECT_EQ(data.hll_union_data->get_lg_config_k(), 12);
+}
+
+TEST_F(AggregateFunctionDataSketchesHllUnionAggTest, 
testMixedLgKMergeKeepsAllInputs) {
+    constexpr uint64_t cardinality = 100000;
+    Data data;
+
+    data.merge(create_sketch(12, 0, cardinality), Data::DEFAULT_UNION_LOG_K);
+    data.merge(create_sketch(8, cardinality, cardinality), 
Data::DEFAULT_UNION_LOG_K);
+    data.merge(create_sketch(12, 2 * cardinality, cardinality), 
Data::DEFAULT_UNION_LOG_K);
+
+    ASSERT_TRUE(data.hll_union_data.has_value());
+    EXPECT_EQ(data.hll_union_data->get_lg_config_k(), 8);
+    EXPECT_NEAR(data.get_result(), 3 * cardinality, 0.2 * 3 * cardinality);
+}
+
+TEST_F(AggregateFunctionDataSketchesHllUnionAggTest,
+       testHigherLgMaxKTightensHighCardinalityErrorBounds) {
+    constexpr uint64_t cardinality_per_sketch = 100000;
+    constexpr double exact_cardinality = 2 * cardinality_per_sketch;
+    const auto first = create_sketch(16, 0, cardinality_per_sketch);
+    const auto second = create_sketch(16, cardinality_per_sketch, 
cardinality_per_sketch);
+
+    Data data_lg_k_7;
+    data_lg_k_7.merge(first, 7);
+    data_lg_k_7.merge(second, 7);
+    Data data_lg_k_12;
+    data_lg_k_12.merge(first, 12);
+    data_lg_k_12.merge(second, 12);
+    Data data_lg_k_16;
+    data_lg_k_16.merge(first, 16);
+    data_lg_k_16.merge(second, 16);
+
+    ASSERT_TRUE(data_lg_k_7.hll_union_data.has_value());
+    ASSERT_TRUE(data_lg_k_12.hll_union_data.has_value());
+    ASSERT_TRUE(data_lg_k_16.hll_union_data.has_value());
+
+    auto verify_accuracy = [&](const Data& data, uint8_t expected_lg_k) {
+        const auto& hll_union = data.hll_union_data.value();
+        EXPECT_EQ(hll_union.get_lg_config_k(), expected_lg_k);
+        EXPECT_LE(hll_union.get_lower_bound(3), exact_cardinality);
+        EXPECT_GE(hll_union.get_upper_bound(3), exact_cardinality);
+    };
+    verify_accuracy(data_lg_k_7, 7);
+    verify_accuracy(data_lg_k_12, 12);
+    verify_accuracy(data_lg_k_16, 16);
+
+    auto max_relative_error = [](uint8_t lg_k) {
+        return std::max(Data::Union::get_rel_err(false, true, lg_k, 3),
+                        std::abs(Data::Union::get_rel_err(true, true, lg_k, 
3)));
+    };
+    EXPECT_GT(max_relative_error(7), max_relative_error(12));
+    EXPECT_GT(max_relative_error(12), max_relative_error(16));
+}
+
+TEST_F(AggregateFunctionDataSketchesHllUnionAggTest,
+       testSparseStateSerializationPreservesPrecision) {
+    Data source;
+    source.merge(create_sketch(8, 0, 7), Data::DEFAULT_UNION_LOG_K);
+
+    auto buffer = ColumnString::create();
+    BufferWritable writer(*buffer);
+    source.write(writer);
+    writer.commit();
+
+    Data restored;
+    BufferReadable reader(buffer->get_data_at(0));
+    restored.read(reader);
+    ASSERT_TRUE(restored.hll_union_data.has_value());
+    EXPECT_EQ(restored.hll_union_data->get_lg_config_k(), 
Data::DEFAULT_UNION_LOG_K);
+
+    restored.merge(create_sketch(12, 1000, 10000), Data::DEFAULT_UNION_LOG_K);
+    EXPECT_EQ(restored.hll_union_data->get_lg_config_k(), 12);
+}
+
+TEST_F(AggregateFunctionDataSketchesHllUnionAggTest, 
testMixedCapStateMergeUsesMinimumLgK) {
+    constexpr uint64_t cardinality = 7;
+    Data low_cap_state;
+    low_cap_state.merge(create_sketch(21, 0, cardinality), 7);
+    ASSERT_TRUE(low_cap_state.hll_union_data.has_value());
+    EXPECT_EQ(low_cap_state.hll_union_data->get_lg_config_k(), 7);
+
+    Data high_cap_state;
+    high_cap_state.merge(create_sketch(21, 100, cardinality), 21);
+    ASSERT_TRUE(high_cap_state.hll_union_data.has_value());
+    EXPECT_EQ(high_cap_state.hll_union_data->get_lg_config_k(), 21);
+
+    const auto low_cap_sketch = 
low_cap_state.hll_union_data->get_result(datasketches::HLL_8);
+    const auto high_cap_sketch = 
high_cap_state.hll_union_data->get_result(datasketches::HLL_8);
+
+    Data high_then_low;
+    high_then_low.merge(high_cap_sketch);
+    high_then_low.merge(low_cap_sketch);
+    ASSERT_TRUE(high_then_low.hll_union_data.has_value());
+    EXPECT_EQ(high_then_low.hll_union_data->get_lg_config_k(), 7);
+
+    Data low_then_high;
+    low_then_high.merge(low_cap_sketch);
+    low_then_high.merge(high_cap_sketch);
+    ASSERT_TRUE(low_then_high.hll_union_data.has_value());
+    EXPECT_EQ(low_then_high.hll_union_data->get_lg_config_k(), 7);
+
+    EXPECT_NEAR(high_then_low.get_result(), 2 * cardinality, 1.0);
+    EXPECT_NEAR(low_then_high.get_result(), 2 * cardinality, 1.0);
+}
+
+TEST_F(AggregateFunctionDataSketchesHllUnionAggTest, 
testDeserializeLegacyDenseMaxLgKState) {
+    Sketch legacy_state(Data::DEFAULT_UNION_LOG_K, datasketches::HLL_8, true, 
Data::Alloc());
+    for (uint64_t value = 0; value < 10000; ++value) {
+        legacy_state.update(value);
+    }
+
+    auto buffer = ColumnString::create();
+    BufferWritable writer(*buffer);
+    const auto serialized = legacy_state.serialize_compact();
+    writer.write_binary(
+            StringRef(reinterpret_cast<const char*>(serialized.data()), 
serialized.size()));
+    writer.commit();
+
+    Data restored;
+    BufferReadable reader(buffer->get_data_at(0));
+    restored.read(reader);
+
+    ASSERT_TRUE(restored.hll_union_data.has_value());
+    EXPECT_EQ(restored.hll_union_data->get_lg_config_k(), 
Data::DEFAULT_UNION_LOG_K);
+    EXPECT_DOUBLE_EQ(restored.get_result(), legacy_state.get_estimate());
+}
+
+TEST_F(AggregateFunctionDataSketchesHllUnionAggTest, testExplicitLgMaxK) {
+    using AggFunc = AggregateFunctionDataSketchesHllUnionAgg<TYPE_STRING, 
Data>;
+    DataTypes argument_types = {std::make_shared<DataTypeString>(),
+                                std::make_shared<DataTypeInt32>()};
+    auto agg_func = std::make_shared<AggFunc>(argument_types);
+
+    const auto serialized = create_sketch(16, 0, 10000).serialize_compact();
+    auto sketches = ColumnString::create();
+    sketches->insert_data(reinterpret_cast<const char*>(serialized.data()), 
serialized.size());
+
+    for (int32_t value : {8, 16}) {
+        auto lg_max_k = ColumnInt32::create();
+        lg_max_k->insert_value(value);
+        const IColumn* columns[2] = {sketches.get(), lg_max_k.get()};
+
+        AggregateDataPtr place =
+                arena->aligned_alloc(agg_func->size_of_data(), 
agg_func->align_of_data());
+        agg_func->create(place);
+        agg_func->add(place, columns, 0, *arena);
+
+        const auto& data = *reinterpret_cast<const Data*>(place);
+        ASSERT_TRUE(data.hll_union_data.has_value());
+        EXPECT_EQ(data.hll_union_data->get_lg_config_k(), value);
+        agg_func->destroy(place);
+    }
+}
+
+TEST_F(AggregateFunctionDataSketchesHllUnionAggTest, testInvalidLgMaxK) {
+    using AggFunc = AggregateFunctionDataSketchesHllUnionAgg<TYPE_STRING, 
Data>;
+    DataTypes argument_types = {std::make_shared<DataTypeString>(),
+                                std::make_shared<DataTypeInt32>()};
+    auto agg_func = std::make_shared<AggFunc>(argument_types);
+
+    const auto serialized = create_sketch(12, 0, 10000).serialize_compact();
+    auto sketches = ColumnString::create();
+    sketches->insert_data(reinterpret_cast<const char*>(serialized.data()), 
serialized.size());
+
+    for (int32_t value : {6, 22}) {
+        auto lg_max_k = ColumnInt32::create();
+        lg_max_k->insert_value(value);
+        const IColumn* columns[2] = {sketches.get(), lg_max_k.get()};
+        AggregateDataPtr place =
+                arena->aligned_alloc(agg_func->size_of_data(), 
agg_func->align_of_data());
+        agg_func->create(place);
+
+        try {
+            agg_func->add(place, columns, 0, *arena);
+            ADD_FAILURE() << "Expected INVALID_ARGUMENT for lg_max_k=" << 
value;
+        } catch (const doris::Exception& e) {
+            EXPECT_EQ(e.code(), doris::ErrorCode::INVALID_ARGUMENT);
+        }
+        agg_func->destroy(place);
+    }
+}
+
 TEST_F(AggregateFunctionDataSketchesHllUnionAggTest, testAddEmptyStringThrows) 
{
     DataTypes argument_types = {std::make_shared<DataTypeString>()};
     auto agg_func = std::make_shared<AggregateFunctionDataSketchesHllUnionAgg<
diff --git a/contrib/datasketches-cpp b/contrib/datasketches-cpp
index de8553ba372..46025e9aeed 160000
--- a/contrib/datasketches-cpp
+++ b/contrib/datasketches-cpp
@@ -1 +1 @@
-Subproject commit de8553ba372e618382c2e7b44b0ffc9422b9458c
+Subproject commit 46025e9aeed8368b1184cbde9634dd99d0ee47c0
diff --git 
a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/DataSketchesHllUnionAgg.java
 
b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/DataSketchesHllUnionAgg.java
index 5091c45ca19..56d58322cd0 100644
--- 
a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/DataSketchesHllUnionAgg.java
+++ 
b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/DataSketchesHllUnionAgg.java
@@ -24,10 +24,11 @@ import 
org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSi
 import org.apache.doris.nereids.trees.expressions.functions.Function;
 import org.apache.doris.nereids.trees.expressions.functions.FunctionTrait;
 import org.apache.doris.nereids.trees.expressions.literal.DoubleLiteral;
-import org.apache.doris.nereids.trees.expressions.shape.UnaryExpression;
+import org.apache.doris.nereids.trees.expressions.literal.IntegerLikeLiteral;
 import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor;
 import org.apache.doris.nereids.types.DataType;
 import org.apache.doris.nereids.types.DoubleType;
+import org.apache.doris.nereids.types.IntegerType;
 import org.apache.doris.nereids.types.StringType;
 import org.apache.doris.nereids.types.VarBinaryType;
 import org.apache.doris.nereids.types.VarcharType;
@@ -39,14 +40,20 @@ import java.util.List;
 
 /** datasketches_hll_union_agg agg function. */
 public class DataSketchesHllUnionAgg extends NotNullableAggregateFunction
-        implements UnaryExpression, ExplicitlyCastableSignature, 
FunctionTrait, RollUpTrait,
+        implements ExplicitlyCastableSignature, FunctionTrait, RollUpTrait,
         NullIgnoringAggregateFunction {
     public static final List<FunctionSignature> SIGNATURES = ImmutableList.of(
             
FunctionSignature.ret(DoubleType.INSTANCE).args(StringType.INSTANCE),
             
FunctionSignature.ret(DoubleType.INSTANCE).args(VarcharType.SYSTEM_DEFAULT),
-            
FunctionSignature.ret(DoubleType.INSTANCE).args(VarBinaryType.INSTANCE)
+            
FunctionSignature.ret(DoubleType.INSTANCE).args(VarBinaryType.INSTANCE),
+            
FunctionSignature.ret(DoubleType.INSTANCE).args(StringType.INSTANCE, 
IntegerType.INSTANCE),
+            
FunctionSignature.ret(DoubleType.INSTANCE).args(VarcharType.SYSTEM_DEFAULT, 
IntegerType.INSTANCE),
+            
FunctionSignature.ret(DoubleType.INSTANCE).args(VarBinaryType.INSTANCE, 
IntegerType.INSTANCE)
     );
 
+    private static final int MIN_LG_MAX_K = 7;
+    private static final int MAX_LG_MAX_K = 21;
+
     /**
      * constructor with 1 argument.
      */
@@ -54,6 +61,11 @@ public class DataSketchesHllUnionAgg extends 
NotNullableAggregateFunction
         super("datasketches_hll_union_agg", arg);
     }
 
+    /** constructor with 2 arguments. */
+    public DataSketchesHllUnionAgg(Expression arg0, Expression arg1) {
+        super("datasketches_hll_union_agg", arg0, arg1);
+    }
+
     /**
      * constructor with 1 argument.
      */
@@ -61,6 +73,11 @@ public class DataSketchesHllUnionAgg extends 
NotNullableAggregateFunction
         this(arg);
     }
 
+    /** constructor with 2 arguments. */
+    public DataSketchesHllUnionAgg(boolean distinct, Expression arg0, 
Expression arg1) {
+        this(arg0, arg1);
+    }
+
     /** constructor for withChildren and reuse signature */
     protected DataSketchesHllUnionAgg(AggregateFunctionParams functionParams) {
         super(functionParams);
@@ -74,6 +91,26 @@ public class DataSketchesHllUnionAgg extends 
NotNullableAggregateFunction
             throw new AnalysisException(getName()
                 + " function's argument should be of STRING/VARCHAR/VARBINARY 
type, but was " + inputType);
         }
+        if (arity() == 2 && !getArgumentType(1).isIntegralType()) {
+            throw new AnalysisException(getName()
+                    + " requires lg_max_k to be a constant integer: " + 
this.toSql());
+        }
+    }
+
+    @Override
+    public void checkLegalityAfterRewrite() {
+        if (arity() == 1) {
+            return;
+        }
+        Expression lgMaxK = getArgument(1);
+        if (!(lgMaxK instanceof IntegerLikeLiteral)) {
+            throw new AnalysisException(getName() + " requires lg_max_k to be 
a constant integer: " + this.toSql());
+        }
+        long value = ((IntegerLikeLiteral) lgMaxK).getLongValue();
+        if (value < MIN_LG_MAX_K || value > MAX_LG_MAX_K) {
+            throw new AnalysisException(getName() + " requires lg_max_k to be 
between "
+                    + MIN_LG_MAX_K + " and " + MAX_LG_MAX_K + ", but was " + 
value);
+        }
     }
 
     @Override
@@ -88,8 +125,8 @@ public class DataSketchesHllUnionAgg extends 
NotNullableAggregateFunction
 
     @Override
     public DataSketchesHllUnionAgg withDistinctAndChildren(boolean distinct, 
List<Expression> children) {
-        Preconditions.checkArgument(children.size() == 1);
-        return new DataSketchesHllUnionAgg(getFunctionParams(distinct, 
children));
+        Preconditions.checkArgument(children.size() == 1 || children.size() == 
2);
+        return new DataSketchesHllUnionAgg(getFunctionParams(false, children));
     }
 
     @Override
@@ -99,7 +136,9 @@ public class DataSketchesHllUnionAgg extends 
NotNullableAggregateFunction
 
     @Override
     public Function constructRollUp(Expression param, Expression... varParams) 
{
-        return new 
DataSketchesHllUnionAgg(getFunctionParams(ImmutableList.of(param)));
+        return arity() == 1
+                ? new 
DataSketchesHllUnionAgg(getFunctionParams(ImmutableList.of(param)))
+                : new 
DataSketchesHllUnionAgg(getFunctionParams(ImmutableList.of(param, 
getArgument(1))));
     }
 
     @Override
diff --git 
a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/EliminateAggCaseWhenTest.java
 
b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/EliminateAggCaseWhenTest.java
new file mode 100644
index 00000000000..db3e06054ef
--- /dev/null
+++ 
b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/EliminateAggCaseWhenTest.java
@@ -0,0 +1,79 @@
+// 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.
+
+package org.apache.doris.nereids.rules.rewrite;
+
+import org.apache.doris.nereids.trees.expressions.Alias;
+import org.apache.doris.nereids.trees.expressions.EqualTo;
+import org.apache.doris.nereids.trees.expressions.Expression;
+import 
org.apache.doris.nereids.trees.expressions.functions.agg.DataSketchesHllUnionAgg;
+import org.apache.doris.nereids.trees.expressions.functions.scalar.If;
+import org.apache.doris.nereids.trees.expressions.literal.IntegerLiteral;
+import org.apache.doris.nereids.trees.expressions.literal.NullLiteral;
+import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate;
+import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan;
+import org.apache.doris.nereids.types.StringType;
+import org.apache.doris.nereids.util.MemoPatternMatchSupported;
+import org.apache.doris.nereids.util.MemoTestUtils;
+import org.apache.doris.nereids.util.PlanChecker;
+import org.apache.doris.nereids.util.PlanConstructor;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+import org.junit.jupiter.api.Test;
+
+class EliminateAggCaseWhenTest implements MemoPatternMatchSupported {
+    private final LogicalOlapScan scan1 = 
PlanConstructor.newLogicalOlapScan(0, "t1", 0);
+
+    @Test
+    void testEliminateSingleArgumentDataSketchesHllCaseWhen() {
+        Expression sketch = scan1.getOutput().get(1);
+        Expression predicate = new EqualTo(scan1.getOutput().get(0), new 
IntegerLiteral(1));
+        If conditionalSketch = new If(predicate, sketch, new 
NullLiteral(StringType.INSTANCE));
+        DataSketchesHllUnionAgg function = new 
DataSketchesHllUnionAgg(conditionalSketch);
+        LogicalAggregate<?> aggregate = new LogicalAggregate<>(
+                ImmutableList.of(), ImmutableList.of(new Alias(function, 
"estimate")), scan1);
+
+        PlanChecker.from(MemoTestUtils.createConnectContext(), aggregate)
+                .applyTopDown(new EliminateAggCaseWhen())
+                .matches(
+                        logicalAggregate(
+                                logicalFilter(logicalOlapScan()).when(filter ->
+                                        
filter.getConjuncts().equals(ImmutableSet.of(predicate)))
+                        ).when(agg -> agg.getAggregateFunctions()
+                                .equals(ImmutableSet.of(new 
DataSketchesHllUnionAgg(sketch))))
+                );
+    }
+
+    @Test
+    void testKeepDataSketchesHllCaseWhenWithLgMaxK() {
+        Expression sketch = scan1.getOutput().get(1);
+        Expression predicate = new EqualTo(scan1.getOutput().get(0), new 
IntegerLiteral(1));
+        If conditionalSketch = new If(predicate, sketch, new 
NullLiteral(StringType.INSTANCE));
+        DataSketchesHllUnionAgg function =
+                new DataSketchesHllUnionAgg(conditionalSketch, new 
IntegerLiteral(8));
+        LogicalAggregate<?> aggregate = new LogicalAggregate<>(
+                ImmutableList.of(), ImmutableList.of(new Alias(function, 
"estimate")), scan1);
+
+        PlanChecker.from(MemoTestUtils.createConnectContext(), aggregate)
+                .applyTopDown(new EliminateAggCaseWhen())
+                .matches(
+                        logicalAggregate(logicalOlapScan())
+                                .when(agg -> 
agg.getAggregateFunctions().equals(ImmutableSet.of(function)))
+                );
+    }
+}
diff --git 
a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/InferAggNotNullTest.java
 
b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/InferAggNotNullTest.java
index f0e709dc1b7..07ced8248ed 100644
--- 
a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/InferAggNotNullTest.java
+++ 
b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/InferAggNotNullTest.java
@@ -34,6 +34,7 @@ import 
org.apache.doris.nereids.trees.expressions.functions.agg.BitmapAgg;
 import org.apache.doris.nereids.trees.expressions.functions.agg.CollectList;
 import org.apache.doris.nereids.trees.expressions.functions.agg.Count;
 import org.apache.doris.nereids.trees.expressions.functions.agg.CountByEnum;
+import 
org.apache.doris.nereids.trees.expressions.functions.agg.DataSketchesHllUnionAgg;
 import org.apache.doris.nereids.trees.expressions.functions.agg.GroupConcat;
 import org.apache.doris.nereids.trees.expressions.functions.agg.MapAgg;
 import org.apache.doris.nereids.trees.expressions.functions.agg.MapAggV2;
@@ -42,6 +43,7 @@ import 
org.apache.doris.nereids.trees.expressions.functions.agg.Min;
 import 
org.apache.doris.nereids.trees.expressions.functions.agg.NullIgnoringAggregateFunction;
 import org.apache.doris.nereids.trees.expressions.functions.agg.Sum;
 import org.apache.doris.nereids.trees.expressions.functions.agg.Sum0;
+import org.apache.doris.nereids.trees.expressions.literal.IntegerLiteral;
 import org.apache.doris.nereids.trees.plans.RelationId;
 import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate;
 import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan;
@@ -54,6 +56,7 @@ import org.apache.doris.nereids.util.PlanConstructor;
 import org.apache.doris.thrift.TStorageType;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 
@@ -81,6 +84,28 @@ class InferAggNotNullTest implements 
MemoPatternMatchSupported {
                 );
     }
 
+    @Test
+    void testInferDataSketchesHllWithOptionalLgMaxK() {
+        Expression sketch = scan1.getOutput().get(1);
+        Not isNotNull = new Not(new IsNull(sketch), true);
+        for (DataSketchesHllUnionAgg function : ImmutableList.of(
+                new DataSketchesHllUnionAgg(sketch),
+                new DataSketchesHllUnionAgg(sketch, new IntegerLiteral(8)))) {
+            LogicalPlan plan = new LogicalPlanBuilder(scan1)
+                    .aggGroupUsingIndex(ImmutableList.of(), 
ImmutableList.of(new Alias(function, "estimate")))
+                    .build();
+
+            PlanChecker.from(MemoTestUtils.createConnectContext(), plan)
+                    .applyTopDown(new InferAggNotNull())
+                    .matches(
+                            logicalAggregate(
+                                    
logicalFilter(logicalOlapScan()).when(filter ->
+                                            
filter.getConjuncts().equals(ImmutableSet.of(isNotNull)))
+                            ).when(agg -> 
agg.getAggregateFunctions().equals(ImmutableSet.of(function)))
+                    );
+        }
+    }
+
     @Test
     void testNotInferWhenAggregateArgumentReturnsFalseForNullInput() {
         Expression isNotNull = new Not(new IsNull(scan1.getOutput().get(1)));
diff --git 
a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/expressions/functions/agg/DataSketchesHllUnionAggTest.java
 
b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/expressions/functions/agg/DataSketchesHllUnionAggTest.java
new file mode 100644
index 00000000000..0eed2e2521d
--- /dev/null
+++ 
b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/expressions/functions/agg/DataSketchesHllUnionAggTest.java
@@ -0,0 +1,190 @@
+// 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.
+
+package org.apache.doris.nereids.trees.expressions.functions.agg;
+
+import org.apache.doris.catalog.FunctionRegistry;
+import org.apache.doris.catalog.FunctionSignature;
+import org.apache.doris.nereids.exceptions.AnalysisException;
+import org.apache.doris.nereids.trees.expressions.Expression;
+import org.apache.doris.nereids.trees.expressions.SlotReference;
+import org.apache.doris.nereids.trees.expressions.functions.FunctionBuilder;
+import 
org.apache.doris.nereids.trees.expressions.functions.combinator.MergeCombinator;
+import 
org.apache.doris.nereids.trees.expressions.functions.combinator.StateCombinator;
+import 
org.apache.doris.nereids.trees.expressions.functions.combinator.UnionCombinator;
+import org.apache.doris.nereids.trees.expressions.literal.DecimalV3Literal;
+import org.apache.doris.nereids.trees.expressions.literal.DoubleLiteral;
+import org.apache.doris.nereids.trees.expressions.literal.IntegerLiteral;
+import org.apache.doris.nereids.trees.expressions.literal.NullLiteral;
+import 
org.apache.doris.nereids.trees.expressions.visitor.DefaultExpressionRewriter;
+import org.apache.doris.nereids.types.IntegerType;
+import org.apache.doris.nereids.types.StringType;
+
+import com.google.common.collect.ImmutableList;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import java.math.BigDecimal;
+import java.util.List;
+
+class DataSketchesHllUnionAggTest {
+    private static final SlotReference SKETCH = SlotReference.of("sketch", 
StringType.INSTANCE);
+
+    @Test
+    void testSignatures() {
+        List<FunctionSignature> signatures = new 
DataSketchesHllUnionAgg(SKETCH).getSignatures();
+
+        Assertions.assertEquals(6, signatures.size());
+        Assertions.assertEquals(3, signatures.stream().filter(signature -> 
signature.arity == 1).count());
+        Assertions.assertEquals(3, signatures.stream().filter(signature -> 
signature.arity == 2).count());
+        signatures.stream().filter(signature -> signature.arity == 2)
+                .forEach(signature -> 
Assertions.assertEquals(IntegerType.INSTANCE, signature.getArgType(1)));
+    }
+
+    @Test
+    void testDistinctIsIgnored() {
+        DataSketchesHllUnionAgg oneArgument = new 
DataSketchesHllUnionAgg(true, SKETCH);
+        DataSketchesHllUnionAgg twoArguments =
+                new DataSketchesHllUnionAgg(true, SKETCH, new 
IntegerLiteral(8));
+
+        for (DataSketchesHllUnionAgg function : ImmutableList.of(oneArgument, 
twoArguments)) {
+            DataSketchesHllUnionAgg rewritten = 
function.withDistinctAndChildren(true, function.children());
+            Assertions.assertFalse(function.isDistinct());
+            Assertions.assertFalse(rewritten.isDistinct());
+            Assertions.assertTrue(function.getDistinctArguments().isEmpty());
+            Assertions.assertTrue(rewritten.getDistinctArguments().isEmpty());
+        }
+    }
+
+    @Test
+    void testNullIgnoringContractSurvivesRewriting() {
+        for (DataSketchesHllUnionAgg function : ImmutableList.of(
+                new DataSketchesHllUnionAgg(SKETCH),
+                new DataSketchesHllUnionAgg(SKETCH, new IntegerLiteral(8)))) {
+            DataSketchesHllUnionAgg rewritten = 
function.withDistinctAndChildren(true, function.getArguments());
+
+            Assertions.assertInstanceOf(NullIgnoringAggregateFunction.class, 
function);
+            Assertions.assertInstanceOf(NullIgnoringAggregateFunction.class, 
rewritten);
+            Assertions.assertFalse(function.nullable());
+            Assertions.assertFalse(rewritten.nullable());
+            Assertions.assertEquals(new DoubleLiteral(0), 
function.resultForEmptyInput());
+            Assertions.assertEquals(function.resultForEmptyInput(), 
rewritten.resultForEmptyInput());
+            Assertions.assertEquals(function.getArguments(), 
rewritten.getArguments());
+        }
+    }
+
+    @Test
+    void testLgMaxKBoundaries() {
+        for (int value : new int[] {7, 21}) {
+            DataSketchesHllUnionAgg function =
+                    new DataSketchesHllUnionAgg(SKETCH, new 
IntegerLiteral(value));
+            
Assertions.assertDoesNotThrow(function::checkLegalityBeforeTypeCoercion);
+            Assertions.assertDoesNotThrow(function::checkLegalityAfterRewrite);
+        }
+    }
+
+    @Test
+    void testLgMaxKMustBeAConstantInteger() {
+        DataSketchesHllUnionAgg nonConstant = new DataSketchesHllUnionAgg(
+                SKETCH, SlotReference.of("lg_max_k", IntegerType.INSTANCE));
+        DataSketchesHllUnionAgg nonInteger = new DataSketchesHllUnionAgg(
+                SKETCH, new DecimalV3Literal(new BigDecimal("8.5")));
+        DataSketchesHllUnionAgg nullValue = new 
DataSketchesHllUnionAgg(SKETCH, new NullLiteral());
+
+        
Assertions.assertDoesNotThrow(nonConstant::checkLegalityBeforeTypeCoercion);
+        Assertions.assertThrows(AnalysisException.class, 
nonConstant::checkLegalityAfterRewrite);
+        Assertions.assertThrows(AnalysisException.class, 
nonInteger::checkLegalityBeforeTypeCoercion);
+        Assertions.assertThrows(AnalysisException.class, 
nullValue::checkLegalityBeforeTypeCoercion);
+    }
+
+    @Test
+    void testLgMaxKRange() {
+        for (int value : new int[] {6, 22}) {
+            DataSketchesHllUnionAgg function =
+                    new DataSketchesHllUnionAgg(SKETCH, new 
IntegerLiteral(value));
+            Assertions.assertThrows(AnalysisException.class, 
function::checkLegalityAfterRewrite);
+        }
+    }
+
+    @Test
+    void testTypedAggStateCanBuildMergeAndUnion() {
+        FunctionRegistry functionRegistry = new FunctionRegistry();
+        for (String functionName : ImmutableList.of(
+                "datasketches_hll_union_agg", "ds_hll_estimate", 
"datasketches_hll_estimate")) {
+            for (StateCombinator state : ImmutableList.of(
+                    buildState(functionRegistry, functionName + "_state", 
SKETCH),
+                    buildState(functionRegistry, functionName + "_state", 
SKETCH, new IntegerLiteral(8)))) {
+                Assertions.assertEquals(StateCombinator.class, 
state.getClass());
+                
Assertions.assertDoesNotThrow(state::checkLegalityAfterRewrite);
+                SlotReference stateSlot = new SlotReference("state", 
state.getDataType(), false);
+
+                String mergeName = functionName + "_merge";
+                FunctionBuilder mergeBuilder = 
functionRegistry.findFunctionBuilder(mergeName, stateSlot);
+                MergeCombinator merge = (MergeCombinator) 
mergeBuilder.build(mergeName, stateSlot).first;
+                
Assertions.assertDoesNotThrow(merge::checkLegalityBeforeTypeCoercion);
+
+                String unionName = functionName + "_union";
+                FunctionBuilder unionBuilder = 
functionRegistry.findFunctionBuilder(unionName, stateSlot);
+                UnionCombinator union = (UnionCombinator) 
unionBuilder.build(unionName, stateSlot).first;
+                
Assertions.assertDoesNotThrow(union::checkLegalityBeforeTypeCoercion);
+            }
+        }
+    }
+
+    @Test
+    void testStateCombinatorValidatesLgMaxKAfterRewrite() {
+        FunctionRegistry functionRegistry = new FunctionRegistry();
+        for (String functionName : ImmutableList.of(
+                "datasketches_hll_union_agg", "ds_hll_estimate", 
"datasketches_hll_estimate")) {
+            String stateName = functionName + "_state";
+            for (int value : new int[] {7, 21}) {
+                StateCombinator state = buildState(functionRegistry, 
stateName, SKETCH, new IntegerLiteral(value));
+                
Assertions.assertDoesNotThrow(state::checkLegalityAfterRewrite);
+            }
+            for (Expression invalidLgMaxK : ImmutableList.of(
+                    new IntegerLiteral(6),
+                    new IntegerLiteral(22),
+                    SlotReference.of("lg_max_k", IntegerType.INSTANCE))) {
+                StateCombinator state = buildState(functionRegistry, 
stateName, SKETCH, invalidLgMaxK);
+                Assertions.assertThrows(AnalysisException.class, 
state::checkLegalityAfterRewrite);
+            }
+        }
+    }
+
+    @Test
+    void testStateCombinatorRewritesAllChildren() {
+        FunctionRegistry functionRegistry = new FunctionRegistry();
+        StateCombinator state = buildState(
+                functionRegistry, "datasketches_hll_union_agg_state", SKETCH, 
new IntegerLiteral(8));
+        StateCombinator rewritten = (StateCombinator) state.accept(new 
DefaultExpressionRewriter<Void>() {
+            @Override
+            public Expression visitIntegerLiteral(IntegerLiteral 
integerLiteral, Void context) {
+                return new IntegerLiteral(22);
+            }
+        }, null);
+
+        Assertions.assertEquals(2, state.arity());
+        Assertions.assertThrows(AnalysisException.class, 
rewritten::checkLegalityAfterRewrite);
+    }
+
+    private static StateCombinator buildState(
+            FunctionRegistry functionRegistry, String stateName, Expression... 
stateArguments) {
+        List<Expression> arguments = ImmutableList.copyOf(stateArguments);
+        FunctionBuilder stateBuilder = 
functionRegistry.findFunctionBuilder(stateName, arguments);
+        return (StateCombinator) stateBuilder.build(stateName, 
arguments).first;
+    }
+}
diff --git 
a/regression-test/data/query_p0/sql_functions/aggregate_functions/test_datasketches_hll_union_agg.out
 
b/regression-test/data/query_p0/sql_functions/aggregate_functions/test_datasketches_hll_union_agg.out
index ae1424305c7..538ff8aaf8e 100644
--- 
a/regression-test/data/query_p0/sql_functions/aggregate_functions/test_datasketches_hll_union_agg.out
+++ 
b/regression-test/data/query_p0/sql_functions/aggregate_functions/test_datasketches_hll_union_agg.out
@@ -5,6 +5,21 @@
 -- !aliases --
 17     17      17
 
+-- !explicit_lg_max_k --
+18     17      17      17
+
+-- !default_typed_state_merge --
+17
+
+-- !typed_state_merge --
+17
+
+-- !typed_state_union --
+17
+
+-- !alias_typed_state_merge --
+17     17
+
 -- !group_by --
 1      7
 2      10
@@ -12,6 +27,12 @@
 -- !distinct --
 17     17
 
+-- !distinct_two_arguments_with_rollup --
+\N     17      17      2
+1      7       7       1
+2      10      10      1
+5      7       7       1
+
 -- !basic_union_varchar --
 17
 
@@ -26,3 +47,6 @@
 
 -- !empty_input --
 0
+
+-- !persisted_state --
+1      18
diff --git 
a/regression-test/data/query_p0/sql_functions/aggregate_functions/test_datasketches_hll_union_agg_null_ignoring.out
 
b/regression-test/data/query_p0/sql_functions/aggregate_functions/test_datasketches_hll_union_agg_null_ignoring.out
new file mode 100644
index 00000000000..8429cd8ddf2
--- /dev/null
+++ 
b/regression-test/data/query_p0/sql_functions/aggregate_functions/test_datasketches_hll_union_agg_null_ignoring.out
@@ -0,0 +1,52 @@
+-- This file is automatically generated. You should know what you did if you 
want to edit this
+-- !mixed_nullable --
+17     17      18
+
+-- !all_null --
+0      0       0
+
+-- !empty_input --
+0      0       0
+
+-- !single_case --
+17
+
+-- !single_case_without_rewrite --
+17
+
+-- !single_if --
+7
+
+-- !single_if_without_rewrite --
+7
+
+-- !single_case_all_null --
+0
+
+-- !single_case_all_null_without_rewrite --
+0
+
+-- !single_case_empty --
+0
+
+-- !single_case_empty_without_rewrite --
+0
+
+-- !two_argument_case --
+18
+
+-- !two_argument_case_without_rewrite --
+18
+
+-- !two_argument_if --
+7
+
+-- !two_argument_if_without_rewrite --
+7
+
+-- !two_argument_case_all_null --
+0
+
+-- !two_argument_case_all_null_without_rewrite --
+0
+
diff --git 
a/regression-test/suites/query_p0/sql_functions/aggregate_functions/test_datasketches_hll_union_agg.groovy
 
b/regression-test/suites/query_p0/sql_functions/aggregate_functions/test_datasketches_hll_union_agg.groovy
index 8d86b8c82d6..ec391d43090 100644
--- 
a/regression-test/suites/query_p0/sql_functions/aggregate_functions/test_datasketches_hll_union_agg.groovy
+++ 
b/regression-test/suites/query_p0/sql_functions/aggregate_functions/test_datasketches_hll_union_agg.groovy
@@ -16,6 +16,8 @@
 // under the License.
 
 suite("test_datasketches_hll_union_agg") {
+    sql "set enable_agg_state=true"
+
     def tableName = "test_datasketches_hll_union_agg_tbl"
     def varcharTableName = "test_datasketches_hll_union_agg_varchar_tbl"
     def emptyTableName = "test_datasketches_hll_union_agg_empty_tbl"
@@ -57,6 +59,54 @@ suite("test_datasketches_hll_union_agg") {
         FROM ${tableName}
     """
 
+    qt_explicit_lg_max_k """SELECT
+            CAST(ROUND(datasketches_hll_union_agg(sk, 7)) AS BIGINT),
+            CAST(ROUND(datasketches_hll_union_agg(sk, 21)) AS BIGINT),
+            CAST(ROUND(ds_hll_estimate(sk, 8)) AS BIGINT),
+            CAST(ROUND(datasketches_hll_estimate(sk, 8)) AS BIGINT)
+        FROM ${tableName}
+    """
+
+    order_qt_default_typed_state_merge """SELECT
+            CAST(ROUND(datasketches_hll_union_agg_merge(sk_state)) AS BIGINT)
+        FROM (
+            SELECT datasketches_hll_union_agg_state(sk) AS sk_state
+            FROM ${tableName}
+            WHERE id IN (1, 2)
+        ) states
+    """
+
+    order_qt_typed_state_merge """SELECT
+            CAST(ROUND(datasketches_hll_union_agg_merge(sk_state)) AS BIGINT)
+        FROM (
+            SELECT datasketches_hll_union_agg_state(sk, 8) AS sk_state
+            FROM ${tableName}
+            WHERE id IN (1, 2)
+        ) states
+    """
+
+    order_qt_typed_state_union """SELECT
+            CAST(ROUND(datasketches_hll_union_agg_merge(sk_state)) AS BIGINT)
+        FROM (
+            SELECT datasketches_hll_union_agg_union(sk_state) AS sk_state
+            FROM (
+                SELECT datasketches_hll_union_agg_state(sk, 8) AS sk_state
+                FROM ${tableName}
+                WHERE id IN (1, 2)
+            ) states
+        ) unioned
+    """
+
+    order_qt_alias_typed_state_merge """SELECT
+            CAST(ROUND(ds_hll_estimate_merge(sk_state)) AS BIGINT),
+            CAST(ROUND(datasketches_hll_estimate_merge(sk_state)) AS BIGINT)
+        FROM (
+            SELECT ds_hll_estimate_state(sk, 8) AS sk_state
+            FROM ${tableName}
+            WHERE id IN (1, 2)
+        ) states
+    """
+
     // 3) Group-by
     qt_group_by """SELECT id, CAST(ROUND(datasketches_hll_union_agg(sk)) AS 
BIGINT)
         FROM ${tableName}
@@ -73,6 +123,16 @@ suite("test_datasketches_hll_union_agg") {
         FROM ${tableName}
     """
 
+    order_qt_distinct_two_arguments_with_rollup """SELECT
+            id,
+            CAST(ROUND(datasketches_hll_union_agg(DISTINCT sk, 8)) AS BIGINT),
+            CAST(ROUND(datasketches_hll_union_agg(sk, 8)) AS BIGINT),
+            COUNT(DISTINCT sk)
+        FROM ${tableName}
+        WHERE id IN (1, 2, 5)
+        GROUP BY ROLLUP(id)
+    """
+
     // 4.1) Input type coverage: VARCHAR
     sql "DROP TABLE IF EXISTS ${varcharTableName}"
     sql """
@@ -149,6 +209,66 @@ suite("test_datasketches_hll_union_agg") {
         exception "CORRUPTION"
     }
 
+    test {
+        sql """SELECT datasketches_hll_union_agg(sk, 6) FROM ${tableName}"""
+        exception "requires lg_max_k to be between 7 and 21"
+    }
+    test {
+        sql """SELECT datasketches_hll_union_agg(sk, 22) FROM ${tableName}"""
+        exception "requires lg_max_k to be between 7 and 21"
+    }
+    test {
+        sql """SELECT datasketches_hll_union_agg(sk, NULL) FROM ${tableName}"""
+        exception "requires lg_max_k to be a constant integer"
+    }
+    test {
+        sql """SELECT datasketches_hll_union_agg(sk, 8.5) FROM ${tableName}"""
+        exception "requires lg_max_k to be a constant integer"
+    }
+    test {
+        sql """SELECT datasketches_hll_union_agg(sk, id) FROM ${tableName}"""
+        exception "requires lg_max_k to be a constant integer"
+    }
+
+    ["datasketches_hll_union_agg", "ds_hll_estimate", 
"datasketches_hll_estimate"].each { functionName ->
+        test {
+            sql """SELECT ${functionName}_state(sk, 6) FROM ${tableName}"""
+            exception "requires lg_max_k to be between 7 and 21"
+        }
+        test {
+            sql """SELECT ${functionName}_state(sk, 22) FROM ${tableName} 
WHERE id = 3"""
+            exception "requires lg_max_k to be between 7 and 21"
+        }
+    }
+    test {
+        sql """SELECT datasketches_hll_union_agg_state(sk, id) FROM 
${tableName}"""
+        exception "requires lg_max_k to be a constant integer"
+    }
+
+    sql "DROP TABLE IF EXISTS test_datasketches_hll_union_agg_state_tbl"
+    sql """
+        CREATE TABLE test_datasketches_hll_union_agg_state_tbl (
+            id INT,
+            sk_state AGG_STATE<datasketches_hll_union_agg(STRING, INT NOT 
NULL)> GENERIC
+        )
+        AGGREGATE KEY(id)
+        DISTRIBUTED BY HASH(id) BUCKETS 1
+        PROPERTIES (
+            "replication_num" = "1"
+        )
+    """
+    sql """INSERT INTO test_datasketches_hll_union_agg_state_tbl
+        SELECT 1, ds_hll_estimate_state(sk, 21) FROM ${tableName} WHERE id = 
2"""
+    sql """INSERT INTO test_datasketches_hll_union_agg_state_tbl
+        SELECT 1, datasketches_hll_estimate_state(sk, 7) FROM ${tableName} 
WHERE id = 1"""
+    order_qt_persisted_state """SELECT
+            id,
+            CAST(ROUND(datasketches_hll_union_agg_merge(sk_state)) AS BIGINT)
+        FROM test_datasketches_hll_union_agg_state_tbl
+        GROUP BY id
+        ORDER BY id
+    """
+
     // Empty string is a valid STRING value, but it is an invalid serialized 
DataSketches HLL sketch.
     // It should not fail at INSERT time; it should fail when the aggregate 
function reads it.
     sql "DROP TABLE IF EXISTS ${badTableName}"
diff --git 
a/regression-test/suites/query_p0/sql_functions/aggregate_functions/test_datasketches_hll_union_agg_null_ignoring.groovy
 
b/regression-test/suites/query_p0/sql_functions/aggregate_functions/test_datasketches_hll_union_agg_null_ignoring.groovy
new file mode 100644
index 00000000000..56d60f2412b
--- /dev/null
+++ 
b/regression-test/suites/query_p0/sql_functions/aggregate_functions/test_datasketches_hll_union_agg_null_ignoring.groovy
@@ -0,0 +1,116 @@
+// 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.
+
+suite("test_datasketches_hll_union_agg_null_ignoring") {
+    // The same lgK=8 sketches as test_datasketches_hll_union_agg: {0..6} and 
{20..29}.
+    def sk1Base64 = "AgEHCAMIBwjL18IEK/L7BoYv+Q11gWYHgbxdBntl5gj8LUIK"
+    def sk2Base64 = 
"AwEHCAUIAAkKAAAAIjvrBcS1nwfGGWoEyHokBO8t9wc1qTEENkcJB7hWqQxZf9QNnuSbGA=="
+
+    sql "DROP TABLE IF EXISTS test_datasketches_hll_union_agg_null_ignoring"
+    sql """
+        CREATE TABLE test_datasketches_hll_union_agg_null_ignoring (
+            id INT,
+            sk STRING
+        )
+        DISTRIBUTED BY HASH(id) BUCKETS 1
+        PROPERTIES ("replication_num" = "1")
+    """
+    sql """
+        INSERT INTO test_datasketches_hll_union_agg_null_ignoring VALUES
+            (1, from_base64('${sk1Base64}')),
+            (2, from_base64('${sk2Base64}')),
+            (3, NULL),
+            (NULL, from_base64('${sk2Base64}'))
+    """
+
+    order_qt_mixed_nullable """
+        SELECT CAST(ROUND(datasketches_hll_union_agg(sk)) AS BIGINT),
+               CAST(ROUND(datasketches_hll_union_agg(sk, 21)) AS BIGINT),
+               CAST(ROUND(datasketches_hll_union_agg(sk, 7)) AS BIGINT)
+        FROM test_datasketches_hll_union_agg_null_ignoring
+    """
+    order_qt_all_null """
+        SELECT CAST(ROUND(datasketches_hll_union_agg(sk)) AS BIGINT),
+               CAST(ROUND(datasketches_hll_union_agg(sk, 21)) AS BIGINT),
+               CAST(ROUND(datasketches_hll_union_agg(sk, 7)) AS BIGINT)
+        FROM test_datasketches_hll_union_agg_null_ignoring
+        WHERE id = 3
+    """
+    order_qt_empty_input """
+        SELECT CAST(ROUND(datasketches_hll_union_agg(sk)) AS BIGINT),
+               CAST(ROUND(datasketches_hll_union_agg(sk, 21)) AS BIGINT),
+               CAST(ROUND(datasketches_hll_union_agg(sk, 7)) AS BIGINT)
+        FROM test_datasketches_hll_union_agg_null_ignoring
+        WHERE id = 4
+    """
+
+    test {
+        sql """
+            SELECT datasketches_hll_union_agg(sk, 22)
+            FROM test_datasketches_hll_union_agg_null_ignoring
+            WHERE id = 3
+        """
+        exception "requires lg_max_k to be between 7 and 21"
+    }
+    test {
+        sql """
+            SELECT datasketches_hll_union_agg(sk, 6)
+            FROM test_datasketches_hll_union_agg_null_ignoring
+            WHERE id = 4
+        """
+        exception "requires lg_max_k to be between 7 and 21"
+    }
+
+    // Each query has one aggregate output so the single-argument rewrite can 
apply.
+    def queries = [
+        single_case: """
+            SELECT CAST(ROUND(datasketches_hll_union_agg(CASE WHEN id <= 2 
THEN sk END)) AS BIGINT)
+            FROM test_datasketches_hll_union_agg_null_ignoring
+        """,
+        single_if: """
+            SELECT CAST(ROUND(datasketches_hll_union_agg(IF(id = 1, sk, 
NULL))) AS BIGINT)
+            FROM test_datasketches_hll_union_agg_null_ignoring
+        """,
+        single_case_all_null: """
+            SELECT CAST(ROUND(datasketches_hll_union_agg(CASE WHEN id = 3 THEN 
sk END)) AS BIGINT)
+            FROM test_datasketches_hll_union_agg_null_ignoring
+        """,
+        single_case_empty: """
+            SELECT CAST(ROUND(datasketches_hll_union_agg(CASE WHEN id = 4 THEN 
sk END)) AS BIGINT)
+            FROM test_datasketches_hll_union_agg_null_ignoring
+        """,
+        // This union rounds to 18 with cap 7, but 17 if the precision 
argument is lost.
+        two_argument_case: """
+            SELECT CAST(ROUND(datasketches_hll_union_agg(CASE WHEN id <= 2 
THEN sk END, 7)) AS BIGINT)
+            FROM test_datasketches_hll_union_agg_null_ignoring
+        """,
+        two_argument_if: """
+            SELECT CAST(ROUND(datasketches_hll_union_agg(IF(id = 1, sk, NULL), 
21)) AS BIGINT)
+            FROM test_datasketches_hll_union_agg_null_ignoring
+        """,
+        two_argument_case_all_null: """
+            SELECT CAST(ROUND(datasketches_hll_union_agg(CASE WHEN id = 3 THEN 
sk END, 7)) AS BIGINT)
+            FROM test_datasketches_hll_union_agg_null_ignoring
+        """
+    ]
+    queries.each { tag, query ->
+        "order_qt_${tag}" query
+        // SET_VAR scopes the comparison to this query and leaves the session 
unchanged.
+        "order_qt_${tag}_without_rewrite" query.replaceFirst(
+                "SELECT", "SELECT /*+ 
SET_VAR(disable_nereids_rules=ELIMINATE_AGG_CASE_WHEN) */")
+    }
+}


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to