Fix COPY TO does not produce an output file for the empty set (comments added + multi-output file addressed)

Author: Bert Vermeiren

datafusion/datasource/src/write/demux.rs

@@ -40,9 +40,9 @@ use datafusion_common::cast::{
 };
 use datafusion_common::{exec_datafusion_err, internal_datafusion_err, not_impl_err};
 use datafusion_common_runtime::SpawnedTask;
-use datafusion_execution::TaskContext;
 
 use chrono::NaiveDate;
+use datafusion_execution::TaskContext;
 use futures::StreamExt;
 use object_store::path::Path;
 use rand::distr::SampleString;
@@ -68,6 +68,11 @@ pub type DemuxedStreamReceiver = UnboundedReceiver<(Path, RecordBatchReceiver)>;
 /// be written with the extension from the path. Otherwise the default extension
 /// will be used and the output will be split into multiple files.
 ///
+/// Output file guarantees:
+///  - Partitioned files: Files are created only for non-empty partitions.
+///  - Single-file output: At least one file is always written, even when the stream is empty.
+///  - Multi-file output: At least the minimum number of parallel files specified in [`datafusion_common::config::ExecutionOptions`] are always written, even when the stream is empty.
+///
 /// Examples of `base_output_path`
 ///  * `tmp/dataset/` -> is a folder since it ends in `/`
 ///  * `tmp/dataset` -> is still a folder since it does not end in `/` but has no valid file extension
@@ -171,10 +176,9 @@ async fn row_count_demuxer(
         max_rows_per_file
     };
 
-    // Single-file output requires creating at least one file stream in advance.
-    // If no record batches are present in the input stream,
-    // the file stream must still be created to produce a valid output file.
-    if single_file_output {
+    // ensure we have at least minimum_parallel_files open, even when
+    // the input stream is empty
+    for _ in 0..minimum_parallel_files {
         open_file_streams.push(create_new_file_stream(
             &base_output_path,
             &write_id,
@@ -189,20 +193,8 @@ async fn row_count_demuxer(
     }
 
     while let Some(rb) = input.next().await.transpose()? {
-        // ensure we have at least minimum_parallel_files open
-        if open_file_streams.len() < minimum_parallel_files {
-            open_file_streams.push(create_new_file_stream(
-                &base_output_path,
-                &write_id,
-                part_idx,
-                &file_extension,
-                single_file_output,
-                max_buffered_batches,
-                &mut tx,
-            )?);
-            row_counts.push(0);
-            part_idx += 1;
-        } else if row_counts[next_send_steam] >= max_rows_per_file {
+        // spill over to new file output if max rows reached
+        if row_counts[next_send_steam] >= max_rows_per_file {
             row_counts[next_send_steam] = 0;
             open_file_streams[next_send_steam] = create_new_file_stream(
                 &base_output_path,