diff --git a/datafusion/src/physical_plan/repartition.rs b/datafusion/src/physical_plan/repartition.rs
index 56de364cc9951..1dd672aaea34a 100644
--- a/datafusion/src/physical_plan/repartition.rs
+++ b/datafusion/src/physical_plan/repartition.rs
@@ -536,7 +536,10 @@ mod tests {
         )
         .await?;
 
-        let total_rows: usize = output_partitions.iter().map(|x| x.len()).sum();
+        let total_rows: usize = output_partitions
+            .iter()
+            .map(|x| x.iter().map(|x| x.num_rows()).sum::<usize>())
+            .sum();
 
         assert_eq!(8, output_partitions.len());
         assert_eq!(total_rows, 8 * 50 * 3);