address #57281 (comment)

Signed-off-by: Arthur <atte.book@gmail.com>

Author: ArthurBook

python/ray/data/aggregate.py

@@ -540,7 +540,7 @@ def __init__(
             zero_factory=lambda: list([0, 0]),  # noqa: C410
         )
 
-    def aggregate_block(self, block: Block) -> List[Union[int, float]]:
+    def aggregate_block(self, block: Block) -> Optional[List[Union[int, float]]]:
         block_acc = BlockAccessor.for_block(block)
         count = block_acc.count(self._target_col_name, self._ignore_nulls)
 
@@ -568,7 +568,7 @@ def finalize(self, accumulator: List[Union[int, float]]) -> Optional[float]:
         if accumulator[1] == 0:
             # If total_count is 0 (e.g., group was empty or all nulls ignored),
             # the mean is undefined. Return NaN
-            return np.nan
+            return np.nan[
 
         return accumulator[0] / accumulator[1]
 
@@ -724,7 +724,7 @@ def __init__(
             zero_factory=lambda: 0,
         )
 
-    def aggregate_block(self, block: Block) -> Union[int, float]:
+    def aggregate_block(self, block: Block) -> Optional[Union[int, float]]:
         block_accessor = BlockAccessor.for_block(block)
 
         max_ = block_accessor.max(self._target_col_name, self._ignore_nulls)