[GitHub] [arrow-datafusion] alamb commented on a diff in pull request #3439: Ensure the row count is preserved when coalescing over empty records

[GitBox <gi...@apache.org>]

alamb commented on code in PR #3439:
URL: https://github.com/apache/arrow-datafusion/pull/3439#discussion_r968299337


##########
datafusion/core/tests/sql/mod.rs:
##########
@@ -766,13 +766,16 @@ async fn execute_to_batches(ctx: &SessionContext, sql: &str) -> Vec<RecordBatch>
     let logical_schema = plan.schema();
 
     let msg = format!("Optimizing logical plan for '{}': {:?}", sql, plan);
-    let plan = ctx
+    let optimized_logical_plan = ctx
         .optimize(&plan)
         .map_err(|e| format!("{:?} at {}", e, msg))
         .unwrap();
-    let optimized_logical_schema = plan.schema();
+    let optimized_logical_schema = optimized_logical_plan.schema();
 
-    let msg = format!("Creating physical plan for '{}': {:?}", sql, plan);
+    let msg = format!(
+        "Creating physical plan for '{}': {:?}",
+        sql, optimized_logical_plan
+    );
     let plan = ctx

Review Comment:
   It is a very subtle point here that the optimized logical plan is no longer actually run. 👍 It took me a few reads to see that. 
   
   A comment might help future readers. 
   



##########
datafusion/core/src/physical_plan/coalesce_batches.rs:
##########
@@ -291,7 +291,11 @@ pub fn concat_batches(
         batches.len(),
         row_count
     );
-    RecordBatch::try_new(schema.clone(), arrays)
+
+    let mut options = RecordBatchOptions::default();
+    options.row_count = Some(row_count);

Review Comment:
   As an FYI (not important) you can write this same pattern in what is likely more idiomatic rust (avoid `mut`) like:
   
   ```suggestion
       let options = RecordBatchOptions{
         row_count:Some(row_count),
         ..Default::default()
       };
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org