fix: tree-style explain for memory table shows wrong row count

datafusion/datasource/src/memory.rs

@@ -129,7 +129,12 @@ impl DataSource for MemorySourceConfig {
                 }
             }
             DisplayFormatType::TreeRender => {
-                let total_rows = self.partitions.iter().map(|b| b.len()).sum::<usize>();
+                let total_rows = self
+                    .partitions
+                    .iter()
+                    .flatten()
+                    .map(RecordBatch::num_rows)
+                    .sum::<usize>();
                 let total_bytes: usize = self
                     .partitions
                     .iter()

datafusion/sqllogictest/test_files/aggregate_memory_spill.slt

@@ -209,7 +209,7 @@ FROM (
 100000 5000050000
 
 # Assert spill happened in the final aggregation.
-# In multi-partitions configuration, 'spilled_rows' is not deterministic, so assert 
+# In multi-partitions configuration, 'spilled_rows' is not deterministic, so assert
 # the unit to be 'K'
 query TT
 EXPLAIN ANALYZE

datafusion/sqllogictest/test_files/array/init_data.slt.part

@@ -688,4 +688,3 @@ AS
     arrow_cast(column4, 'FixedSizeList(3, Float64)') AS column4
 FROM arrays_distance_table
 ;
-

datafusion/sqllogictest/test_files/array_add.slt

@@ -234,4 +234,4 @@ select array_add(array_add(a, b), c) from (values
 ----
 [111.0, 222.0]
 NULL
-[111.0, NULL]
+[111.0, NULL]

datafusion/sqllogictest/test_files/array_product.slt

@@ -142,4 +142,4 @@ select list_product(column1) from (values
 ) as t(column1);
 ----
 6
-NULL
+NULL

datafusion/sqllogictest/test_files/array_subtract.slt

@@ -234,4 +234,4 @@ select array_subtract(array_subtract(a, b), c) from (values
 ----
 [89.0, 178.0]
 NULL
-[89.0, NULL]
+[89.0, NULL]

datafusion/sqllogictest/test_files/explain_tree.slt

@@ -270,7 +270,7 @@ physical_plan
 08)│    --------------------   │
 09)│        bytes: 1024        │
 10)│       format: memory      │
-11)│          rows: 2          │
+11)│          rows: 6          │
 12)└───────────────────────────┘
 
 # 2 Joins
@@ -331,7 +331,7 @@ physical_plan
 09)│    --------------------   ││    --------------------   │
 10)│         bytes: 512        ││            on:            ├──────────────┐
 11)│       format: memory      ││    (int_col = int_col)    │              │
-12)│          rows: 1          ││                           │              │
+12)│          rows: 3          ││                           │              │
 13)└───────────────────────────┘└─────────────┬─────────────┘              │
 14)-----------------------------┌─────────────┴─────────────┐┌─────────────┴─────────────┐
 15)-----------------------------│       DataSourceExec      ││      RepartitionExec      │
@@ -570,7 +570,7 @@ physical_plan
 09)│    --------------------   │
 10)│         bytes: 512        │
 11)│       format: memory      │
-12)│          rows: 1          │
+12)│          rows: 3          │
 13)└───────────────────────────┘
 
 # Query with filter on json
@@ -927,7 +927,7 @@ physical_plan
 15)│    --------------------   │
 16)│         bytes: 512        │
 17)│       format: memory      │
-18)│          rows: 1          │
+18)│          rows: 3          │
 19)└───────────────────────────┘
 
 # Query with projection on json
@@ -1177,7 +1177,7 @@ physical_plan
 13)│    --------------------   ││    --------------------   │
 14)│        bytes: 5932        ││        bytes: 5932        │
 15)│       format: memory      ││       format: memory      │
-16)│          rows: 1          ││          rows: 1          │
+16)│          rows: 4          ││          rows: 4          │
 17)└───────────────────────────┘└───────────────────────────┘
 
 statement ok
@@ -1276,7 +1276,7 @@ physical_plan
 44)│    --------------------   ││    --------------------   │
 45)│         bytes: 288        ││         bytes: 280        │
 46)│       format: memory      ││       format: memory      │
-47)│          rows: 1          ││          rows: 1          │
+47)│          rows: 3          ││          rows: 3          │
 48)└───────────────────────────┘└───────────────────────────┘
 
 # Test explain tree for UnionExec
@@ -1295,14 +1295,14 @@ physical_plan
 06)│    --------------------   ││    --------------------   │
 07)│         bytes: 288        ││   id: CAST(id AS Int32)   │
 08)│       format: memory      ││         name: name        │
-09)│          rows: 1          ││                           │
+09)│          rows: 3          ││                           │
 10)└───────────────────────────┘└─────────────┬─────────────┘
 11)-----------------------------┌─────────────┴─────────────┐
 12)-----------------------------│       DataSourceExec      │
 13)-----------------------------│    --------------------   │
 14)-----------------------------│         bytes: 280        │
 15)-----------------------------│       format: memory      │
-16)-----------------------------│          rows: 1          │
+16)-----------------------------│          rows: 3          │
 17)-----------------------------└───────────────────────────┘
 
 # cleanup
@@ -1646,7 +1646,7 @@ physical_plan
 13)│    --------------------   │
 14)│        bytes: 2576        │
 15)│       format: memory      │
-16)│          rows: 1          │
+16)│          rows: 3          │
 17)└───────────────────────────┘
 
 query TT
@@ -1669,7 +1669,7 @@ physical_plan
 13)│    --------------------   │
 14)│        bytes: 2576        │
 15)│       format: memory      │
-16)│          rows: 1          │
+16)│          rows: 3          │
 17)└───────────────────────────┘
 
 query TT
@@ -1692,7 +1692,7 @@ physical_plan
 13)│    --------------------   │
 14)│        bytes: 2576        │
 15)│       format: memory      │
-16)│          rows: 1          │
+16)│          rows: 3          │
 17)└───────────────────────────┘
 
 
@@ -1754,7 +1754,7 @@ physical_plan
 48)│    --------------------   │
 49)│         bytes: 160        │
 50)│       format: memory      │
-51)│          rows: 1          │
+51)│          rows: 10         │
 52)└───────────────────────────┘
 
 # clean up

datafusion/sqllogictest/test_files/insert.slt

@@ -463,7 +463,7 @@ drop table unsigned_bigint_test
 
 # Config reset
 
-# The SLT runner sets `target_partitions` to 4 instead of using the default, so 
+# The SLT runner sets `target_partitions` to 4 instead of using the default, so
 # reset it explicitly.
 statement ok
 set datafusion.execution.target_partitions = 4;

datafusion/sqllogictest/test_files/join_is_not_distinct_from.slt

@@ -113,7 +113,7 @@ physical_plan
 17)│    --------------------   ││    --------------------   │
 18)│         bytes: 288        ││         bytes: 288        │
 19)│       format: memory      ││       format: memory      │
-20)│          rows: 1          ││          rows: 1          │
+20)│          rows: 4          ││          rows: 4          │
 21)└───────────────────────────┘└───────────────────────────┘
 
 statement ok

datafusion/sqllogictest/test_files/limit.slt

@@ -401,7 +401,7 @@ physical_plan
 09)│    --------------------   │
 10)│         bytes: 160        │
 11)│       format: memory      │
-12)│          rows: 1          │
+12)│          rows: 10         │
 13)└───────────────────────────┘
 
 # generate BIGINT data from 1 to 1000 in multiple partitions

datafusion/sqllogictest/test_files/monotonic_projection_test.slt

@@ -314,4 +314,3 @@ DROP TABLE t2_int_float;
 
 statement ok
 DROP TABLE t3_int_bigint;
-

datafusion/sqllogictest/test_files/ordered_aggregate_spill.slt

@@ -73,7 +73,7 @@ GROUP BY v1
 
 # ==================================================================================
 # Input is partially ordered by group keys (input order by (a), query is 'group by a,b')
-# 
+#
 # Try different memory limits, ensure result is the same, but spill count differ
 
 # HACK: check `spilled_bytes=x KB` to ensure it has spilled. If it has not spilled,

datafusion/sqllogictest/test_files/repartition.slt

@@ -140,7 +140,7 @@ FROM t1 WHERE ((false > (v1 = v1)) IS DISTINCT FROM true);
 
 # Config reset
 
-# The SLT runner sets `target_partitions` to 4 instead of using the default, so 
+# The SLT runner sets `target_partitions` to 4 instead of using the default, so
 # reset it explicitly.
 statement ok
 set datafusion.execution.target_partitions = 4;

datafusion/sqllogictest/test_files/spark/math/hypot.slt

@@ -132,4 +132,4 @@ NaN
 query R
 SELECT hypot('Infinity'::double, '-Infinity'::double);
 ----
-Infinity
+Infinity

datafusion/sqllogictest/test_files/struct.slt

@@ -1732,4 +1732,3 @@ drop view leaf_view;
 
 statement ok
 drop table leaf_base;
-

datafusion/sqllogictest/test_files/window_topn.slt

@@ -283,7 +283,7 @@ physical_plan
 54)│    --------------------   │
 55)│         bytes: 480        │
 56)│       format: memory      │
-57)│          rows: 1          │
+57)│          rows: 10         │
 58)└───────────────────────────┘
 
 statement ok