From 36b5c3dcd79b62db8304e893be2269bd4ec93172 Mon Sep 17 00:00:00 2001
From: Liang-Chi Hsieh <viirya@gmail.com>
Date: Tue, 16 Apr 2024 10:18:57 -0700
Subject: [PATCH] Add RecordBatchOptions for Expand

---
 core/src/execution/datafusion/operators/expand.rs | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/core/src/execution/datafusion/operators/expand.rs b/core/src/execution/datafusion/operators/expand.rs
index 5cf444b3b..ca3fdc1aa 100644
--- a/core/src/execution/datafusion/operators/expand.rs
+++ b/core/src/execution/datafusion/operators/expand.rs
@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use arrow_array::RecordBatch;
+use arrow_array::{RecordBatch, RecordBatchOptions};
 use arrow_schema::SchemaRef;
 use datafusion::{
     execution::TaskContext,
@@ -169,7 +169,9 @@ impl ExpandStream {
             Ok::<(), DataFusionError>(())
         })?;
 
-        RecordBatch::try_new(self.schema.clone(), columns).map_err(|e| e.into())
+        let options = RecordBatchOptions::new().with_row_count(Some(batch.num_rows()));
+        RecordBatch::try_new_with_options(self.schema.clone(), columns, &options)
+            .map_err(|e| e.into())
     }
 }