fix: Add num_rows when building RecordBatch

core/src/execution/datafusion/shuffle_writer.rs

@@ -272,10 +272,11 @@ impl PartitionBuffer {
 
         // active -> staging
         let active = std::mem::take(&mut self.active);
+        let num_rows = self.num_active_rows;
         self.num_active_rows = 0;
         mem_diff -= self.active_slots_mem_size as isize;
 
-        let frozen_batch = make_batch(self.schema.clone(), active)?;
+        let frozen_batch = make_batch(self.schema.clone(), active, num_rows)?;
 
         let frozen_capacity_old = self.frozen.capacity();
         let mut cursor = Cursor::new(&mut self.frozen);
@@ -1148,9 +1149,11 @@ fn make_dict_builder(datatype: &DataType, capacity: usize) -> Box<dyn ArrayBuild
 fn make_batch(
     schema: SchemaRef,
     mut arrays: Vec<Box<dyn ArrayBuilder>>,
+    row_count: usize,
 ) -> ArrowResult<RecordBatch> {
     let columns = arrays.iter_mut().map(|array| array.finish()).collect();
-    RecordBatch::try_new(schema, columns)
+    let options = RecordBatchOptions::new().with_row_count(Option::from(row_count));
+    RecordBatch::try_new_with_options(schema, columns, &options)
 }
 
 /// Checksum algorithms for writing IPC bytes.

core/src/execution/shuffle/row.rs

@@ -37,7 +37,7 @@ use arrow_array::{
         StructBuilder, TimestampMicrosecondBuilder,
     },
     types::Int32Type,
-    Array, ArrayRef, RecordBatch,
+    Array, ArrayRef, RecordBatch, RecordBatchOptions,
 };
 use arrow_schema::{DataType, Field, Schema, TimeUnit};
 use jni::sys::{jint, jlong};
@@ -3347,7 +3347,7 @@ pub fn process_sorted_row_partition(
             .zip(schema.iter())
             .map(|(builder, datatype)| builder_to_array(builder, datatype, prefer_dictionary_ratio))
             .collect();
-        let batch = make_batch(array_refs?);
+        let batch = make_batch(array_refs?, n);
 
         let mut frozen: Vec<u8> = vec![];
         let mut cursor = Cursor::new(&mut frozen);
@@ -3420,7 +3420,7 @@ fn builder_to_array(
     }
 }
 
-fn make_batch(arrays: Vec<ArrayRef>) -> RecordBatch {
+fn make_batch(arrays: Vec<ArrayRef>, row_count: usize) -> RecordBatch {
     let mut dict_id = 0;
     let fields = arrays
         .iter()
@@ -3441,5 +3441,6 @@ fn make_batch(arrays: Vec<ArrayRef>) -> RecordBatch {
         })
         .collect::<Vec<_>>();
     let schema = Arc::new(Schema::new(fields));
-    RecordBatch::try_new(schema, arrays).unwrap()
+    let options = RecordBatchOptions::new().with_row_count(Option::from(row_count));
+    RecordBatch::try_new_with_options(schema, arrays, &options).unwrap()
 }

spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala

@@ -373,6 +373,7 @@ class CometExecSuite extends CometTestBase {
     withParquetDataFrame((0 until 5).map(i => (i, i + 1))) { df =>
       assert(df.where("_1 IS NOT NULL").count() == 5)
       checkSparkAnswerAndOperator(df)
+      assert(df.select().limit(2).count() === 2)
     }
   }