count() does not return nulls

2 files changed, 7 insertions, 32 deletions

diff --git a/ydb/core/kqp/ut/olap/kqp_olap_ut.cpp b/ydb/core/kqp/ut/olap/kqp_olap_ut.cpp
index 657cbda1c8..063f5d1864 100644
--- a/ydb/core/kqp/ut/olap/kqp_olap_ut.cpp
+++ b/ydb/core/kqp/ut/olap/kqp_olap_ut.cpp
@@ -1957,8 +1957,6 @@ Y_UNIT_TEST_SUITE(KqpOlap) {
     }
 
     Y_UNIT_TEST(Aggregation_Count_NullGroupBy) {
-        // Wait for KIKIMR-16831 fix
-        return;
         TAggregationTestCase testCase;
         testCase.SetQuery(R"(
                 SELECT
@@ -1979,8 +1977,6 @@ Y_UNIT_TEST_SUITE(KqpOlap) {
     }
 
     Y_UNIT_TEST(Aggregation_Count_NullMixGroupBy) {
-        // Wait for KIKIMR-16831 fix
-        return;
         TAggregationTestCase testCase;
         testCase.SetQuery(R"(
                 SELECT
@@ -2001,7 +1997,7 @@ Y_UNIT_TEST_SUITE(KqpOlap) {
     }
 
     Y_UNIT_TEST(Aggregation_Count_GroupByNull) {
-        // Wait for KIKIMR-16831 fix
+        // Wait for KIKIMR-16940 fix
         return;
         TAggregationTestCase testCase;
         testCase.SetQuery(R"(
@@ -2023,7 +2019,7 @@ Y_UNIT_TEST_SUITE(KqpOlap) {
     }
 
     Y_UNIT_TEST(Aggregation_Count_GroupByNullMix) {
-        // Wait for KIKIMR-16831 fix
+        // Wait for KIKIMR-16940 fix
         return;
         TAggregationTestCase testCase;
         testCase.SetQuery(R"(
diff --git a/ydb/library/arrow_clickhouse/AggregateFunctions/AggregateFunctionCount.h b/ydb/library/arrow_clickhouse/AggregateFunctions/AggregateFunctionCount.h
index 8432027620..d7157f31e4 100644
--- a/ydb/library/arrow_clickhouse/AggregateFunctions/AggregateFunctionCount.h
+++ b/ydb/library/arrow_clickhouse/AggregateFunctions/AggregateFunctionCount.h
@@ -18,24 +18,6 @@ namespace CH
 struct AggregateFunctionCountData
 {
     UInt64 count = 0;
-    bool has_value = false;
-
-    bool has() const
-    {
-        return has_value;
-    }
-
-    void inc()
-    {
-        has_value = true;
-        ++count;
-    }
-
-    void add(UInt64 value)
-    {
-        has_value = true;
-        count += value;
-    }
 };
 
 
@@ -56,7 +38,7 @@ public:
 
     void add(AggregateDataPtr __restrict place, const IColumn **, size_t, Arena *) const override
     {
-        data(place).inc();
+        ++data(place).count;
     }
 
     void addBatchSinglePlace(
@@ -71,25 +53,22 @@ public:
         {
             auto * condition_map = flags + column.offset();
             auto length = row_end - row_begin;
-            data(place).add(arrow::internal::CountSetBits(condition_map, row_begin, length));
+            data(place).count += arrow::internal::CountSetBits(condition_map, row_begin, length);
         }
         else
         {
-            data(place).add(row_end - row_begin);
+            data(place).count += row_end - row_begin;
         }
     }
 
     void merge(AggregateDataPtr __restrict place, ConstAggregateDataPtr rhs, Arena *) const override
     {
-        data(place).add(data(rhs).count);
+        data(place).count += data(rhs).count;
     }
 
     void insertResultInto(AggregateDataPtr __restrict place, MutableColumn & to, Arena *) const override
     {
-        if (data(place).has())
-            assert_cast<MutableColumnUInt64 &>(to).Append(data(place).count).ok();
-        else
-            assert_cast<MutableColumnUInt64 &>(to).AppendNull().ok();
+        assert_cast<MutableColumnUInt64 &>(to).Append(data(place).count).ok();
     }
 };