fix aggregation count function

tube/etl/indexers/aggregation/new_translator.py

@@ -20,7 +20,8 @@
     lit,
     sort_array,
     struct,
-    sum,
+    count,
+    when,
 )
 from copy import deepcopy
 
@@ -121,16 +122,18 @@ def aggregate_with_count_on_edge_tbl(self, node_name, df, edge_df, child):
                 break
 
         node_id = get_node_id_name(node_name)
+        child_id = get_node_id_name(child)
         if count_reducer is None:
             # if there is no reducer, group by parent key and get out empty value
             count_df = edge_df.select(node_id).drop_duplicates([node_id])
         else:
             # if there is reducer, group by parent key and get out the number of children
             # only non-leaf nodes goes through this step
+            child_prop_name = count_reducer.prop.name
             count_df = (
                 edge_df.groupBy(node_id)
-                .count()
-                .select(node_id, col("count").alias(count_reducer.prop.name))
+                .agg(count(when(col(child_id).isNotNull(), 1)).alias(child_prop_name))
+                .select(node_id, col(child_prop_name))
             )
             count_reducer.done = True
         # combine value lists new counted dataframe to existing one

tube/etl/indexers/base/translator.py

@@ -5,7 +5,7 @@
 
 from pyspark.sql.context import SQLContext
 from pyspark.sql.types import StructType, StructField, StringType
-from pyspark.sql.functions import col, min, sum, count, collect_set, collect_list, first
+from pyspark.sql.functions import col, min, sum, count, collect_set, collect_list, first, when
 
 from .lambdas import (
     extract_link,
@@ -278,7 +278,7 @@ def reducer_to_agg_func_expr(
         if func_name == "count":
             if is_merging:
                 return sum(col(value)).alias(col_alias)
-            return count(col(value)).alias(col_alias)
+            return count(when(col(value).isNotNull(), 1)).alias(col_alias)
         if func_name == "sum":
             return sum(col(value)).alias(col_alias)
         if func_name == "set":