Introduce way to customize prefix of multi file outputs

Author: DoumanAsh

datafusion/common/src/config.rs

@@ -649,6 +649,12 @@ config_namespace! {
         /// # Default
         /// `false` — ANSI SQL mode is disabled by default.
         pub enable_ansi_mode: bool, default = false
+
+        /// Prefix to use when generating file name in multi file output
+        ///
+        /// When prefix is non-empty string, this prefix will be used to generate file name as
+        /// `{partitioned_file_prefix_name}{datafusion generated suffix}`
+        pub partitioned_file_prefix_name: String, default = String::new()
     }
 }
 

datafusion/datasource-csv/src/source.rs

@@ -453,16 +453,17 @@ pub async fn plan_to_csv(
     let parsed = ListingTableUrl::parse(path)?;
     let object_store_url = parsed.object_store();
     let store = task_ctx.runtime_env().object_store(&object_store_url)?;
-    let writer_buffer_size = task_ctx
-        .session_config()
-        .options()
-        .execution
-        .objectstore_writer_buffer_size;
+    let exec_options = &task_ctx.session_config().options().execution;
+    let writer_buffer_size = exec_options.objectstore_writer_buffer_size;
     let mut join_set = JoinSet::new();
     for i in 0..plan.output_partitioning().partition_count() {
         let storeref = Arc::clone(&store);
         let plan: Arc<dyn ExecutionPlan> = Arc::clone(&plan);
-        let filename = format!("{}/part-{i}.csv", parsed.prefix());
+        let filename = format!(
+            "{}/{}part-{i}.csv",
+            exec_options.partitioned_file_prefix_name,
+            parsed.prefix()
+        );
         let file = object_store::path::Path::parse(filename)?;
 
         let mut stream = plan.execute(i, Arc::clone(&task_ctx))?;

datafusion/datasource-json/src/source.rs

@@ -275,16 +275,17 @@ pub async fn plan_to_json(
     let parsed = ListingTableUrl::parse(path)?;
     let object_store_url = parsed.object_store();
     let store = task_ctx.runtime_env().object_store(&object_store_url)?;
-    let writer_buffer_size = task_ctx
-        .session_config()
-        .options()
-        .execution
-        .objectstore_writer_buffer_size;
+    let exec_options = &task_ctx.session_config().options().execution;
+    let writer_buffer_size = exec_options.objectstore_writer_buffer_size;
     let mut join_set = JoinSet::new();
     for i in 0..plan.output_partitioning().partition_count() {
         let storeref = Arc::clone(&store);
         let plan: Arc<dyn ExecutionPlan> = Arc::clone(&plan);
-        let filename = format!("{}/part-{i}.json", parsed.prefix());
+        let filename = format!(
+            "{}/{}part-{i}.json",
+            exec_options.partitioned_file_prefix_name,
+            parsed.prefix()
+        );
         let file = object_store::path::Path::parse(filename)?;
 
         let mut stream = plan.execute(i, Arc::clone(&task_ctx))?;

datafusion/datasource-parquet/src/writer.rs

@@ -39,21 +39,22 @@ pub async fn plan_to_parquet(
     let object_store_url = parsed.object_store();
     let store = task_ctx.runtime_env().object_store(&object_store_url)?;
     let mut join_set = JoinSet::new();
+    let exec_options = &task_ctx.session_config().options().execution;
     for i in 0..plan.output_partitioning().partition_count() {
         let plan: Arc<dyn ExecutionPlan> = Arc::clone(&plan);
-        let filename = format!("{}/part-{i}.parquet", parsed.prefix());
+        let filename = format!(
+            "{}/{}part-{i}.parquet",
+            exec_options.partitioned_file_prefix_name,
+            parsed.prefix()
+        );
         let file = Path::parse(filename)?;
         let propclone = writer_properties.clone();
 
         let storeref = Arc::clone(&store);
         let buf_writer = BufWriter::with_capacity(
             storeref,
             file.clone(),
-            task_ctx
-                .session_config()
-                .options()
-                .execution
-                .objectstore_writer_buffer_size,
+            exec_options.objectstore_writer_buffer_size,
         );
         let mut stream = plan.execute(i, Arc::clone(&task_ctx))?;
         join_set.spawn(async move {

datafusion/datasource/src/write/demux.rs

@@ -156,7 +156,8 @@ async fn row_count_demuxer(
     let max_buffered_batches = exec_options.max_buffered_batches_per_output_file;
     let minimum_parallel_files = exec_options.minimum_parallel_output_files;
     let mut part_idx = 0;
-    let write_id = rand::distr::Alphanumeric.sample_string(&mut rand::rng(), 16);
+    let mut write_id = exec_options.partitioned_file_prefix_name.clone();
+    rand::distr::Alphanumeric.append_string(&mut rand::rng(), &mut write_id, 16);
 
     let mut open_file_streams = Vec::with_capacity(minimum_parallel_files);
 
@@ -283,9 +284,10 @@ async fn hive_style_partitions_demuxer(
     file_extension: String,
     keep_partition_by_columns: bool,
 ) -> Result<()> {
-    let write_id = rand::distr::Alphanumeric.sample_string(&mut rand::rng(), 16);
-
     let exec_options = &context.session_config().options().execution;
+    let mut write_id = exec_options.partitioned_file_prefix_name.clone();
+    rand::distr::Alphanumeric.append_string(&mut rand::rng(), &mut write_id, 16);
+
     let max_buffered_recordbatches = exec_options.max_buffered_batches_per_output_file;
 
     // To support non string partition col types, cast the type to &str first