enhance: optimize count aggregation performance for nullable fields (#48055)

## Summary
- Optimize `addSingleGroupRawInput` (global aggregation): replace O(N)
row-by-row `ValidAt()` loop with O(1) `nullCount()`-based computation
- Optimize `addRawInput` (GROUP BY): add `nullCount()==0` fast path to
skip validity checks entirely; for columns with actual nulls, use direct
bitmap access instead of virtual `ValidAt()` calls

issue: #47566

## Test plan
- [ ] Existing aggregation E2E tests pass
- [ ] C++ unit tests for query aggregation pass
- [ ] Benchmark nullable vs non-nullable count aggregation performance

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Signed-off-by: MrPresent-Han <chun.han@gmail.com>
Co-authored-by: MrPresent-Han <chun.han@gmail.com>

internal/core/src/exec/operator/query-agg/CountAggregateBase.h

@@ -60,10 +60,21 @@ class CountAggregate : public SimpleNumericAggregate<bool, int64_t, int64_t> {
             input_column = std::dynamic_pointer_cast<ColumnVector>(input[0]);
             AssertInfo(input_column != nullptr,
                        "input[0] must be ColumnVector for count aggregation");
-            for (auto i = 0; i < input_column->size(); i++) {
-                if (input_column->ValidAt(i)) {
+            if (input_column->nullCount() == 0) {
+                for (auto i = 0; i < input_column->size(); i++) {
                     addToGroup(groups[i], 1);
                 }
+            } else {
+                // Safe: GetValidRawData() returns CustomBitset::data(), which
+                // exposes storage as uint64_t* per the project-wide bitset
+                // API contract (Bitset::data_impl, CustomBitset.h).
+                const uint64_t* validity = reinterpret_cast<const uint64_t*>(
+                    input_column->GetValidRawData());
+                for (auto i = 0; i < input_column->size(); i++) {
+                    if (validity[i >> 6] & (uint64_t(1) << (i & 63))) {
+                        addToGroup(groups[i], 1);
+                    }
+                }
             }
             return;
         }
@@ -89,15 +100,12 @@ class CountAggregate : public SimpleNumericAggregate<bool, int64_t, int64_t> {
         AssertInfo(column != nullptr,
                    "input[0] must be ColumnVector for count aggregation");
         if (column->IsBitmap()) {
-            // Use validity bitmap to count non-null values
             BitsetTypeView view(column->GetRawData(), column->size());
             addToGroup(group, view.size() - view.count());
         } else {
-            for (auto i = 0; i < column->size(); i++) {
-                if (column->ValidAt(i)) {
-                    addToGroup(group, 1);
-                }
-            }
+            addToGroup(
+                group,
+                column->size() - static_cast<int64_t>(column->nullCount()));
         }
     }
 

internal/core/unittest/test_query_group_by.cpp

@@ -17,6 +17,7 @@
 #include "plan/PlanNodeIdGenerator.h"
 #include "test_utils/storage_test_utils.h"
 #include "exec/expression/function/FunctionFactory.h"
+#include "exec/operator/query-agg/CountAggregateBase.h"
 #include "query/PlanImpl.h"
 #include "query/PlanNode.h"
 
@@ -791,4 +792,262 @@ TEST_P(QueryAggTest, RetrieveAggregationWithValidityBitmap) {
     // Verify the data values are present
     ASSERT_TRUE(field_data.has_scalars());
     EXPECT_GT(data_size, 0) << "Should have result data";
+}
+
+// ============================================================
+// Direct unit tests for CountAggregate optimized paths
+// ============================================================
+
+// Helper: allocate a group row buffer with int64_t accumulator at given offset.
+// Layout: [nullByte(1)] [padding(7)] [int64_t accumulator]
+// Total 16 bytes, accumulator at offset 8 (8-byte aligned).
+struct CountAggTestHelper {
+    static constexpr int32_t kAccOffset = 8;
+    static constexpr int32_t kNullByte = 0;
+    static constexpr uint8_t kNullMask = 1;
+    static constexpr int32_t kRowSizeOffset = 0;
+    static constexpr size_t kRowSize = 16;
+
+    static void
+    setupAggregate(milvus::exec::CountAggregate& agg) {
+        agg.setOffsets(kAccOffset, kNullByte, kNullMask, kRowSizeOffset);
+    }
+
+    static std::vector<char>
+    makeGroupRow() {
+        return std::vector<char>(kRowSize, 0);
+    }
+
+    static int64_t
+    getCount(const std::vector<char>& row) {
+        return *reinterpret_cast<const int64_t*>(row.data() + kAccOffset);
+    }
+};
+
+// addSingleGroupRawInput: all valid (nullCount == 0)
+TEST(CountAggregateTest, SingleGroupAllValid) {
+    milvus::exec::CountAggregate agg;
+    CountAggTestHelper::setupAggregate(agg);
+
+    auto row = CountAggTestHelper::makeGroupRow();
+    char* group = row.data();
+    const std::vector<milvus::vector_size_t> indices = {0};
+    agg.initializeNewGroups(&group, indices);
+
+    // 100 rows, all valid
+    auto col =
+        std::make_shared<milvus::ColumnVector>(milvus::DataType::INT64, 100);
+    std::vector<milvus::VectorPtr> input = {col};
+    agg.addSingleGroupRawInput(group, 100, input);
+
+    EXPECT_EQ(CountAggTestHelper::getCount(row), 100);
+}
+
+// addSingleGroupRawInput: all null
+TEST(CountAggregateTest, SingleGroupAllNull) {
+    milvus::exec::CountAggregate agg;
+    CountAggTestHelper::setupAggregate(agg);
+
+    auto row = CountAggTestHelper::makeGroupRow();
+    char* group = row.data();
+    const std::vector<milvus::vector_size_t> indices = {0};
+    agg.initializeNewGroups(&group, indices);
+
+    auto col =
+        std::make_shared<milvus::ColumnVector>(milvus::DataType::INT64, 100);
+    for (size_t i = 0; i < 100; i++) {
+        col->nullAt(i);
+    }
+    std::vector<milvus::VectorPtr> input = {col};
+    agg.addSingleGroupRawInput(group, 100, input);
+
+    EXPECT_EQ(CountAggTestHelper::getCount(row), 0);
+}
+
+// addSingleGroupRawInput: mixed nulls
+TEST(CountAggregateTest, SingleGroupMixedNulls) {
+    milvus::exec::CountAggregate agg;
+    CountAggTestHelper::setupAggregate(agg);
+
+    auto row = CountAggTestHelper::makeGroupRow();
+    char* group = row.data();
+    const std::vector<milvus::vector_size_t> indices = {0};
+    agg.initializeNewGroups(&group, indices);
+
+    // 100 rows, even indices are null
+    auto col =
+        std::make_shared<milvus::ColumnVector>(milvus::DataType::INT64, 100);
+    for (size_t i = 0; i < 100; i += 2) {
+        col->nullAt(i);
+    }
+    std::vector<milvus::VectorPtr> input = {col};
+    agg.addSingleGroupRawInput(group, 100, input);
+
+    EXPECT_EQ(CountAggTestHelper::getCount(row), 50);
+}
+
+// addSingleGroupRawInput: single row valid
+TEST(CountAggregateTest, SingleGroupSingleRow) {
+    milvus::exec::CountAggregate agg;
+    CountAggTestHelper::setupAggregate(agg);
+
+    auto row = CountAggTestHelper::makeGroupRow();
+    char* group = row.data();
+    const std::vector<milvus::vector_size_t> indices = {0};
+    agg.initializeNewGroups(&group, indices);
+
+    auto col =
+        std::make_shared<milvus::ColumnVector>(milvus::DataType::INT64, 1);
+    std::vector<milvus::VectorPtr> input = {col};
+    agg.addSingleGroupRawInput(group, 1, input);
+
+    EXPECT_EQ(CountAggTestHelper::getCount(row), 1);
+}
+
+// addSingleGroupRawInput: non-64-multiple size (e.g., 65 rows)
+TEST(CountAggregateTest, SingleGroupNon64Multiple) {
+    milvus::exec::CountAggregate agg;
+    CountAggTestHelper::setupAggregate(agg);
+
+    auto row = CountAggTestHelper::makeGroupRow();
+    char* group = row.data();
+    const std::vector<milvus::vector_size_t> indices = {0};
+    agg.initializeNewGroups(&group, indices);
+
+    // 65 rows, last row is null
+    auto col =
+        std::make_shared<milvus::ColumnVector>(milvus::DataType::INT64, 65);
+    col->nullAt(64);
+    std::vector<milvus::VectorPtr> input = {col};
+    agg.addSingleGroupRawInput(group, 65, input);
+
+    EXPECT_EQ(CountAggTestHelper::getCount(row), 64);
+}
+
+// addRawInput (GROUP BY): all valid, multiple groups
+TEST(CountAggregateTest, GroupByAllValid) {
+    milvus::exec::CountAggregate agg;
+    CountAggTestHelper::setupAggregate(agg);
+
+    // 3 groups
+    auto g0 = CountAggTestHelper::makeGroupRow();
+    auto g1 = CountAggTestHelper::makeGroupRow();
+    auto g2 = CountAggTestHelper::makeGroupRow();
+    std::vector<char*> groups_vec = {
+        g0.data(), g1.data(), g2.data(), g0.data(), g1.data(), g2.data()};
+    char** groups = groups_vec.data();
+
+    const std::vector<milvus::vector_size_t> indices = {0, 1, 2};
+    char* init_groups[] = {g0.data(), g1.data(), g2.data()};
+    agg.initializeNewGroups(init_groups, indices);
+
+    // 6 rows, all valid, round-robin across 3 groups
+    auto col =
+        std::make_shared<milvus::ColumnVector>(milvus::DataType::INT64, 6);
+    std::vector<milvus::VectorPtr> input = {col};
+    agg.addRawInput(groups, 6, input);
+
+    EXPECT_EQ(CountAggTestHelper::getCount(g0), 2);
+    EXPECT_EQ(CountAggTestHelper::getCount(g1), 2);
+    EXPECT_EQ(CountAggTestHelper::getCount(g2), 2);
+}
+
+// addRawInput (GROUP BY): mixed nulls with multiple groups
+TEST(CountAggregateTest, GroupByMixedNulls) {
+    milvus::exec::CountAggregate agg;
+    CountAggTestHelper::setupAggregate(agg);
+
+    auto g0 = CountAggTestHelper::makeGroupRow();
+    auto g1 = CountAggTestHelper::makeGroupRow();
+    // rows: [g0, g1, g0, g1, g0, g1]
+    // nulls: row 1 and 4 are null
+    std::vector<char*> groups_vec = {
+        g0.data(), g1.data(), g0.data(), g1.data(), g0.data(), g1.data()};
+    char** groups = groups_vec.data();
+
+    const std::vector<milvus::vector_size_t> indices = {0, 1};
+    char* init_groups[] = {g0.data(), g1.data()};
+    agg.initializeNewGroups(init_groups, indices);
+
+    auto col =
+        std::make_shared<milvus::ColumnVector>(milvus::DataType::INT64, 6);
+    col->nullAt(1);  // g1 row null
+    col->nullAt(4);  // g0 row null
+    std::vector<milvus::VectorPtr> input = {col};
+    agg.addRawInput(groups, 6, input);
+
+    // g0: rows 0,2,4 -> row 4 null -> count 2
+    // g1: rows 1,3,5 -> row 1 null -> count 2
+    EXPECT_EQ(CountAggTestHelper::getCount(g0), 2);
+    EXPECT_EQ(CountAggTestHelper::getCount(g1), 2);
+}
+
+// addRawInput (GROUP BY): all null
+TEST(CountAggregateTest, GroupByAllNull) {
+    milvus::exec::CountAggregate agg;
+    CountAggTestHelper::setupAggregate(agg);
+
+    auto g0 = CountAggTestHelper::makeGroupRow();
+    std::vector<char*> groups_vec = {g0.data(), g0.data(), g0.data()};
+    char** groups = groups_vec.data();
+
+    const std::vector<milvus::vector_size_t> indices = {0};
+    char* init_groups[] = {g0.data()};
+    agg.initializeNewGroups(init_groups, indices);
+
+    auto col =
+        std::make_shared<milvus::ColumnVector>(milvus::DataType::INT64, 3);
+    col->nullAt(0);
+    col->nullAt(1);
+    col->nullAt(2);
+    std::vector<milvus::VectorPtr> input = {col};
+    agg.addRawInput(groups, 3, input);
+
+    EXPECT_EQ(CountAggTestHelper::getCount(g0), 0);
+}
+
+// addRawInput (GROUP BY): non-64-multiple size with mixed nulls
+TEST(CountAggregateTest, GroupByNon64Multiple) {
+    milvus::exec::CountAggregate agg;
+    CountAggTestHelper::setupAggregate(agg);
+
+    auto g0 = CountAggTestHelper::makeGroupRow();
+    const std::vector<milvus::vector_size_t> indices = {0};
+    char* init_groups[] = {g0.data()};
+    agg.initializeNewGroups(init_groups, indices);
+
+    // 65 rows all in one group, row 63 and 64 are null
+    const int numRows = 65;
+    std::vector<char*> groups_vec(numRows, g0.data());
+    char** groups = groups_vec.data();
+
+    auto col = std::make_shared<milvus::ColumnVector>(milvus::DataType::INT64,
+                                                      numRows);
+    col->nullAt(63);
+    col->nullAt(64);
+    std::vector<milvus::VectorPtr> input = {col};
+    agg.addRawInput(groups, numRows, input);
+
+    EXPECT_EQ(CountAggTestHelper::getCount(g0), 63);
+}
+
+// Contract test: nullCount()-based path matches ValidAt() semantics
+TEST(CountAggregateTest, NullCountMatchesValidAt) {
+    // Verify that size() - nullCount() equals the number of ValidAt()==true
+    for (int size : {1, 7, 63, 64, 65, 127, 128, 129, 200}) {
+        auto col = std::make_shared<milvus::ColumnVector>(
+            milvus::DataType::INT64, size);
+        // null every 3rd element
+        for (int i = 0; i < size; i += 3) {
+            col->nullAt(i);
+        }
+        int64_t validAtCount = 0;
+        for (int i = 0; i < size; i++) {
+            if (col->ValidAt(i))
+                validAtCount++;
+        }
+        EXPECT_EQ(col->size() - static_cast<int64_t>(col->nullCount()),
+                  validAtCount)
+            << "Mismatch at size=" << size;
+    }
 }