fix: parallel parquet can underflow when max_record_batch_rows < execution.batch_size

datafusion/core/src/dataframe/parquet.rs

@@ -166,4 +166,47 @@ mod tests {
 
         Ok(())
     }
+
+    #[tokio::test]
+    async fn write_parquet_with_small_rg_size() -> Result<()> {
+        let mut test_df = test_util::test_table().await?;
+        // make the test data larger so there are multiple batches
+        for _ in 0..7 {
+            test_df = test_df.clone().union(test_df)?;
+        }
+        let output_path = "file://local/test.parquet";
+
+        for rg_size in (1..7).step_by(5) {
+            let df = test_df.clone();
+            let tmp_dir = TempDir::new()?;
+            let local = Arc::new(LocalFileSystem::new_with_prefix(&tmp_dir)?);
+            let local_url = Url::parse("file://local").unwrap();
+            let ctx = &test_df.session_state;
+            ctx.runtime_env().register_object_store(&local_url, local);
+            let mut options = TableParquetOptions::default();
+            options.global.max_row_group_size = rg_size;
+            options.global.allow_single_file_parallelism = true;
+            df.write_parquet(
+                output_path,
+                DataFrameWriteOptions::new().with_single_file_output(true),
+                Some(options),
+            )
+            .await?;
+
+            // Check that file actually used the correct rg size
+            let file = std::fs::File::open(tmp_dir.into_path().join("test.parquet"))?;
+
+            let reader =
+                parquet::file::serialized_reader::SerializedFileReader::new(file)
+                    .unwrap();
+
+            let parquet_metadata = reader.metadata();
+
+            let written_rows = parquet_metadata.row_group(0).num_rows();
+
+            assert_eq!(written_rows as usize, rg_size);
+        }
+
+        Ok(())
+    }
 }

datafusion/core/src/datasource/file_format/parquet.rs

@@ -876,42 +876,47 @@ fn spawn_parquet_parallel_serialization_task(
             )?;
         let mut current_rg_rows = 0;
 
-        while let Some(rb) = data.recv().await {
-            if current_rg_rows + rb.num_rows() < max_row_group_rows {
-                send_arrays_to_col_writers(&col_array_channels, &rb, schema.clone())
-                    .await?;
-                current_rg_rows += rb.num_rows();
-            } else {
-                let rows_left = max_row_group_rows - current_rg_rows;
-                let a = rb.slice(0, rows_left);
-                send_arrays_to_col_writers(&col_array_channels, &a, schema.clone())
-                    .await?;
+        while let Some(mut rb) = data.recv().await {
+            // This loop allows the "else" block to repeatedly split the RecordBatch to handle the case
+            // when max_row_group_rows < execution.batch_size as an alternative to a recursive async
+            // function.
+            loop {
+                if current_rg_rows + rb.num_rows() < max_row_group_rows {
+                    send_arrays_to_col_writers(&col_array_channels, &rb, schema.clone())
+                        .await?;
+                    current_rg_rows += rb.num_rows();
+                    break;
+                } else {
+                    let rows_left = max_row_group_rows - current_rg_rows;
+                    let a = rb.slice(0, rows_left);
+                    send_arrays_to_col_writers(&col_array_channels, &a, schema.clone())
+                        .await?;
+
+                    // Signal the parallel column writers that the RowGroup is done, join and finalize RowGroup
+                    // on a separate task, so that we can immediately start on the next RG before waiting
+                    // for the current one to finish.
+                    drop(col_array_channels);
+                    let finalize_rg_task = spawn_rg_join_and_finalize_task(
+                        column_writer_handles,
+                        max_row_group_rows,
+                    );
+
+                    serialize_tx.send(finalize_rg_task).await.map_err(|_| {
+                        DataFusionError::Internal(
+                            "Unable to send closed RG to concat task!".into(),
+                        )
+                    })?;
 
-                // Signal the parallel column writers that the RowGroup is done, join and finalize RowGroup
-                // on a separate task, so that we can immediately start on the next RG before waiting
-                // for the current one to finish.
-                drop(col_array_channels);
-                let finalize_rg_task = spawn_rg_join_and_finalize_task(
-                    column_writer_handles,
-                    max_row_group_rows,
-                );
-
-                serialize_tx.send(finalize_rg_task).await.map_err(|_| {
-                    DataFusionError::Internal(
-                        "Unable to send closed RG to concat task!".into(),
-                    )
-                })?;
+                    current_rg_rows = 0;
+                    rb = rb.slice(rows_left, rb.num_rows() - rows_left);
 
-                let b = rb.slice(rows_left, rb.num_rows() - rows_left);
-                (column_writer_handles, col_array_channels) =
-                    spawn_column_parallel_row_group_writer(
-                        schema.clone(),
-                        writer_props.clone(),
-                        max_buffer_rb,
-                    )?;
-                send_arrays_to_col_writers(&col_array_channels, &b, schema.clone())
-                    .await?;
-                current_rg_rows = b.num_rows();
+                    (column_writer_handles, col_array_channels) =
+                        spawn_column_parallel_row_group_writer(
+                            schema.clone(),
+                            writer_props.clone(),
+                            max_buffer_rb,
+                        )?;
+                }
             }
         }