[improvement](be) Optimize count on nullable column

be/src/exprs/aggregate/aggregate_function_count.h

@@ -37,6 +37,7 @@
 #include "core/data_type/data_type_number.h"
 #include "core/types.h"
 #include "exprs/aggregate/aggregate_function.h"
+#include "util/simd/bits.h"
 
 namespace doris {
 class Arena;
@@ -67,6 +68,11 @@ class AggregateFunctionCount final
         ++data(place).count;
     }
 
+    void add_batch_single_place(size_t batch_size, AggregateDataPtr place, const IColumn**,
+                                Arena&) const override {
+        data(place).count += batch_size;
+    }
+
     void reset(AggregateDataPtr place) const override {
         AggregateFunctionCount::data(place).count = 0;
     }
@@ -180,8 +186,7 @@ class AggregateFunctionCount final
     }
 };
 
-// TODO: Maybe AggregateFunctionCountNotNullUnary should be a subclass of AggregateFunctionCount
-// Simply count number of not-NULL values.
+// Used for unary count(nullable_expr). SQL count(expr) counts non-NULL values.
 class AggregateFunctionCountNotNullUnary final
         : public IAggregateFunctionDataHelper<AggregateFunctionCountData,
                                               AggregateFunctionCountNotNullUnary> {
@@ -202,6 +207,20 @@ class AggregateFunctionCountNotNullUnary final
                          .is_null_at(row_num);
     }
 
+    void add_batch_single_place(size_t batch_size, AggregateDataPtr place, const IColumn** columns,
+                                Arena&) const override {
+        const auto& nullable_column =
+                assert_cast<const ColumnNullable&, TypeCheckOnRelease::DISABLE>(*columns[0]);
+        const auto& null_map = nullable_column.get_null_map_data();
+        DCHECK_LE(batch_size, null_map.size());
+        if (!nullable_column.has_null(0, batch_size)) {
+            data(place).count += batch_size;
+            return;
+        }
+        data(place).count +=
+                simd::count_zero_num(reinterpret_cast<const int8_t*>(null_map.data()), batch_size);
+    }
+
     void reset(AggregateDataPtr place) const override { data(place).count = 0; }
 
     void merge(AggregateDataPtr __restrict place, ConstAggregateDataPtr rhs,

be/test/exprs/aggregate/agg_count_test.cpp

@@ -17,6 +17,7 @@
 
 #include <gtest/gtest.h>
 
+#include "core/data_type/data_type_nullable.h"
 #include "core/data_type/data_type_number.h"
 #include "exprs/aggregate/agg_function_test.h"
 
@@ -31,4 +32,24 @@ TEST_F(AggregateFunctionCountTest, test_int64) {
     execute(Block({ColumnHelper::create_column_with_name<DataTypeInt64>({1, 2, 3})}),
             ColumnHelper::create_column_with_name<DataTypeInt64>({3}));
 }
+
+TEST_F(AggregateFunctionCountTest, test_nullable_int64_without_null) {
+    create_agg("count", false,
+               {std::make_shared<DataTypeNullable>(std::make_shared<DataTypeInt64>())},
+               std::make_shared<DataTypeInt64>());
+
+    execute(Block({ColumnHelper::create_nullable_column_with_name<DataTypeInt64>({1, 2, 3, 4},
+                                                                                 {0, 0, 0, 0})}),
+            ColumnHelper::create_column_with_name<DataTypeInt64>({4}));
+}
+
+TEST_F(AggregateFunctionCountTest, test_nullable_int64_with_null) {
+    create_agg("count", false,
+               {std::make_shared<DataTypeNullable>(std::make_shared<DataTypeInt64>())},
+               std::make_shared<DataTypeInt64>());
+
+    execute(Block({ColumnHelper::create_nullable_column_with_name<DataTypeInt64>({1, 2, 3, 4, 5},
+                                                                                 {0, 1, 0, 1, 0})}),
+            ColumnHelper::create_column_with_name<DataTypeInt64>({3}));
+}
 } // namespace doris