fix count aggregate function (evaluate only not null fields) (#453)

Co-authored-by: Sam Willis <sam.willis@gmail.com>

Author: fcannizzaro

.changeset/curly-doors-deny.md

@@ -0,0 +1,6 @@
+---
+"@tanstack/db-ivm": patch
+"@tanstack/db": patch
+---
+
+fix count aggregate function (evaluate only not null field values like SQL count)

packages/db-ivm/src/operators/groupBy.ts

@@ -165,13 +165,16 @@ export function sum<T>(
 /**
  * Creates a count aggregate function
  */
-export function count<T>(): AggregateFunction<T, number, number> {
+export function count<T>(
+  valueExtractor: (value: T) => any = (v) => v
+): AggregateFunction<T, number, number> {
   return {
-    preMap: () => 1,
+    // Count only not-null values (the `== null` comparison gives true for both null and undefined)
+    preMap: (data: T) => (valueExtractor(data) == null ? 0 : 1),
     reduce: (values: Array<[number, number]>) => {
       let totalCount = 0
-      for (const [_, multiplicity] of values) {
-        totalCount += multiplicity
+      for (const [nullMultiplier, multiplicity] of values) {
+        totalCount += nullMultiplier * multiplicity
       }
       return totalCount
     },

packages/db-ivm/tests/operators/groupBy.test.ts

@@ -304,6 +304,76 @@ describe(`Operators`, () => {
       expect(latestMessage.getInner()).toEqual(expectedDeleteResult)
     })
 
+    test(`with count (only not-null values)`, () => {
+      const graph = new D2()
+      const input = graph.newInput<{
+        category: string
+        amount: number | null
+      }>()
+      let latestMessage: any = null
+      const messages: Array<MultiSet<any>> = []
+
+      input.pipe(
+        groupBy(
+          (data) => ({
+            category: data.category,
+          }),
+          {
+            countNotNull: count((data) => data.amount),
+            count: count(),
+          }
+        ),
+        output((message) => {
+          latestMessage = message
+          messages.push(message)
+        })
+      )
+
+      graph.finalize()
+
+      // Initial data
+      input.sendData(
+        new MultiSet([
+          [{ category: `A`, amount: 10 }, 1],
+          [{ category: `B`, amount: 10 }, 1],
+          [{ category: `A`, amount: null }, 1],
+          [{ category: `B`, amount: null }, 1],
+        ])
+      )
+
+      graph.run()
+
+      // Verify we have the latest message
+      expect(latestMessage).not.toBeNull()
+
+      const expectedResult = [
+        [
+          [
+            `{"category":"A"}`,
+            {
+              category: `A`,
+              countNotNull: 1,
+              count: 2,
+            },
+          ],
+          1,
+        ],
+        [
+          [
+            `{"category":"B"}`,
+            {
+              category: `B`,
+              countNotNull: 1,
+              count: 2,
+            },
+          ],
+          1,
+        ],
+      ]
+
+      expect(latestMessage.getInner()).toEqual(expectedResult)
+    })
+
     test(`with avg and count aggregates`, () => {
       const graph = new D2()
       const input = graph.newInput<{

packages/db/src/query/compiler/group-by.ts

@@ -349,12 +349,17 @@ function getAggregateFunction(aggExpr: Aggregate) {
     return typeof value === `number` ? value : value != null ? Number(value) : 0
   }
 
+  // Create a raw value extractor function for the expression to aggregate
+  const rawValueExtractor = ([, namespacedRow]: [string, NamespacedRow]) => {
+    return compiledExpr(namespacedRow)
+  }
+
   // Return the appropriate aggregate function
   switch (aggExpr.name.toLowerCase()) {
     case `sum`:
       return sum(valueExtractor)
     case `count`:
-      return count() // count() doesn't need a value extractor
+      return count(rawValueExtractor)
     case `avg`:
       return avg(valueExtractor)
     case `min`:

packages/db/tests/query/group-by.test.ts

@@ -180,6 +180,40 @@ function createGroupByTests(autoIndex: `off` | `eager`): void {
         expect(customer3?.max_amount).toBe(250)
       })
 
+      test(`group by customer_id with count aggregation (not null only)`, () => {
+        const customerSummary = createLiveQueryCollection({
+          startSync: true,
+          query: (q) =>
+            q
+              .from({ orders: ordersCollection })
+              .groupBy(({ orders }) => orders.customer_id)
+              .select(({ orders }) => ({
+                customer_id: orders.customer_id,
+                total: count(orders.id),
+                onlyNotNull: count(orders.sales_rep_id),
+              })),
+        })
+
+        expect(customerSummary.size).toBe(3) // 3 customers
+
+        // Customer 1: orders 1, 2, 7 (total: 3, onlyNotNull: 3)
+        const customer1 = customerSummary.get(1)
+        expect(customer1?.total).toBe(3)
+        expect(customer1?.onlyNotNull).toBe(3)
+
+        // Customer 2: orders 3, 4 (total: 2, onlyNotNull: 2)
+        const customer2 = customerSummary.get(2)
+        expect(customer2).toBeDefined()
+        expect(customer2?.total).toBe(2)
+        expect(customer2?.onlyNotNull).toBe(2)
+
+        // Customer 3: orders 5, 6 (total: 2, onlyNotNull: 1)
+        const customer3 = customerSummary.get(3)
+        expect(customer3).toBeDefined()
+        expect(customer3?.total).toBe(2)
+        expect(customer3?.onlyNotNull).toBe(1)
+      })
+
       test(`group by status`, () => {
         const statusSummary = createLiveQueryCollection({
           startSync: true,