feat: Conditionally allow to keep partition_by columns when using PARTITIONED BY enhancement

datafusion/common/src/config.rs

@@ -303,6 +303,9 @@ config_namespace! {
         /// statistics into the same file groups.
         /// Currently experimental
         pub split_file_groups_by_statistics: bool, default = false
+
+        /// Should Datafusion keep the columns used for partition_by in the output RecordBatches
+        pub keep_partition_by_columns: bool, default = false
     }
 }
 
@@ -1282,6 +1285,11 @@ impl TableOptions {
             return ConfigField::set(self, key, value);
         }
 
+        // Only used for hive.keep_partition_by_columns
+        if prefix == "hive" {
+            return Ok(());
+        }
+
         let Some(e) = self.extensions.0.get_mut(prefix) else {
             return _config_err!("Could not find config namespace \"{prefix}\"");
         };

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

@@ -223,6 +223,7 @@ impl DataSink for ArrowFileSink {
             part_col,
             self.config.table_paths[0].clone(),
             "arrow".into(),
+            self.config.keep_partition_by_columns,
         );
 
         let mut file_write_tasks: JoinSet<std::result::Result<usize, DataFusionError>> =

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

@@ -548,7 +548,9 @@ impl ParquetSink {
     /// of hive style partitioning where some columns are removed from the
     /// underlying files.
     fn get_writer_schema(&self) -> Arc<Schema> {
-        if !self.config.table_partition_cols.is_empty() {
+        if !self.config.table_partition_cols.is_empty()
+            && !self.config.keep_partition_by_columns
+        {
             let schema = self.config.output_schema();
             let partition_names: Vec<_> = self
                 .config
@@ -638,6 +640,7 @@ impl DataSink for ParquetSink {
             part_col,
             self.config.table_paths[0].clone(),
             "parquet".into(),
+            self.config.keep_partition_by_columns,
         );
 
         let mut file_write_tasks: JoinSet<
@@ -1875,6 +1878,7 @@ mod tests {
             output_schema: schema.clone(),
             table_partition_cols: vec![],
             overwrite: true,
+            keep_partition_by_columns: false,
         };
         let parquet_sink = Arc::new(ParquetSink::new(
             file_sink_config,
@@ -1969,6 +1973,7 @@ mod tests {
             output_schema: schema.clone(),
             table_partition_cols: vec![("a".to_string(), DataType::Utf8)], // add partitioning
             overwrite: true,
+            keep_partition_by_columns: false,
         };
         let parquet_sink = Arc::new(ParquetSink::new(
             file_sink_config,

datafusion/core/src/datasource/file_format/write/demux.rs

@@ -75,6 +75,7 @@ pub(crate) fn start_demuxer_task(
     partition_by: Option<Vec<(String, DataType)>>,
     base_output_path: ListingTableUrl,
     file_extension: String,
+    keep_partition_by_columns: bool,
 ) -> (SpawnedTask<Result<()>>, DemuxedStreamReceiver) {
     let (tx, rx) = mpsc::unbounded_channel();
     let context = context.clone();
@@ -91,6 +92,7 @@ pub(crate) fn start_demuxer_task(
                     parts,
                     base_output_path,
                     file_extension,
+                    keep_partition_by_columns,
                 )
                 .await
             })
@@ -111,7 +113,7 @@ pub(crate) fn start_demuxer_task(
     (task, rx)
 }
 
-/// Dynamically partitions input stream to acheive desired maximum rows per file
+/// Dynamically partitions input stream to achieve desired maximum rows per file
 async fn row_count_demuxer(
     mut tx: UnboundedSender<(Path, Receiver<RecordBatch>)>,
     mut input: SendableRecordBatchStream,
@@ -240,6 +242,7 @@ async fn hive_style_partitions_demuxer(
     partition_by: Vec<(String, DataType)>,
     base_output_path: ListingTableUrl,
     file_extension: String,
+    keep_partition_by_columns: bool,
 ) -> Result<()> {
     let write_id =
         rand::distributions::Alphanumeric.sample_string(&mut rand::thread_rng(), 16);
@@ -298,9 +301,11 @@ async fn hive_style_partitions_demuxer(
                 }
             };
 
-            // remove partitions columns
-            let final_batch_to_send =
-                remove_partition_by_columns(&parted_batch, &partition_by)?;
+            let final_batch_to_send = if keep_partition_by_columns {
+                parted_batch
+            } else {
+                remove_partition_by_columns(&parted_batch, &partition_by)?
+            };
 
             // Finally send the partial batch partitioned by distinct value!
             part_tx.send(final_batch_to_send).await.map_err(|_| {

datafusion/core/src/datasource/file_format/write/orchestration.rs

@@ -224,6 +224,7 @@ pub(crate) async fn stateless_multipart_put(
         part_cols,
         base_output_path.clone(),
         file_extension,
+        config.keep_partition_by_columns,
     );
 
     let rb_buffer_size = &context

datafusion/core/src/datasource/listing/table.rs

@@ -917,6 +917,8 @@ impl TableProvider for ListingTable {
         .await?;
 
         let file_groups = file_list_stream.try_collect::<Vec<_>>().await?;
+        let keep_partition_by_columns =
+            state.config().options().execution.keep_partition_by_columns;
 
         // Sink related option, apart from format
         let config = FileSinkConfig {
@@ -926,6 +928,7 @@ impl TableProvider for ListingTable {
             output_schema: self.schema(),
             table_partition_cols: self.options.table_partition_cols.clone(),
             overwrite,
+            keep_partition_by_columns,
         };
 
         let unsorted: Vec<Vec<Expr>> = vec![];

datafusion/core/src/datasource/physical_plan/mod.rs

@@ -85,6 +85,8 @@ pub struct FileSinkConfig {
     pub table_partition_cols: Vec<(String, DataType)>,
     /// Controls whether existing data should be overwritten by this sink
     pub overwrite: bool,
+    /// Controls whether partition columns are kept for the file
+    pub keep_partition_by_columns: bool,
 }
 
 impl FileSinkConfig {

datafusion/core/src/physical_planner.rs

@@ -777,6 +777,16 @@ impl DefaultPhysicalPlanner {
                     .map(|s| (s.to_string(), arrow_schema::DataType::Null))
                     .collect::<Vec<_>>();
 
+                let keep_partition_by_columns = source_option_tuples
+                    .get("hive.keep_partition_by_columns")
+                    .map(|v| v.trim() == "true")
+                    .unwrap_or(false)
+                    || session_state
+                        .config()
+                        .options()
+                        .execution
+                        .keep_partition_by_columns;
+
                 // Set file sink related options
                 let config = FileSinkConfig {
                     object_store_url,
@@ -785,6 +795,7 @@ impl DefaultPhysicalPlanner {
                     output_schema: Arc::new(schema),
                     table_partition_cols,
                     overwrite: false,
+                    keep_partition_by_columns,
                 };
                 let mut table_options = session_state.default_table_options();
                 let sink_format: Arc<dyn FileFormat> = match format_options {

datafusion/expr/src/logical_plan/builder.rs

@@ -279,9 +279,9 @@ impl LogicalPlanBuilder {
         Ok(Self::from(LogicalPlan::Copy(CopyTo {
             input: Arc::new(input),
             output_url,
+            partition_by,
             format_options,
             options,
-            partition_by,
         })))
     }
 

datafusion/sql/src/parser.rs

@@ -1475,7 +1475,7 @@ mod tests {
     fn copy_to_multi_options() -> Result<(), ParserError> {
         // order of options is preserved
         let sql =
-            "COPY foo TO bar STORED AS parquet OPTIONS ('format.row_group_size' 55, 'format.compression' snappy)";
+            "COPY foo TO bar STORED AS parquet OPTIONS ('format.row_group_size' 55, 'format.compression' snappy, 'keep_partition_by_columns' true)";
 
         let expected_options = vec![
             (
@@ -1486,6 +1486,10 @@ mod tests {
                 "format.compression".to_string(),
                 Value::SingleQuotedString("snappy".to_string()),
             ),
+            (
+                "keep_partition_by_columns".to_string(),
+                Value::SingleQuotedString("true".to_string()),
+            ),
         ];
 
         let mut statements = DFParser::parse_sql(sql).unwrap();

datafusion/sql/src/statement.rs

@@ -888,7 +888,15 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
                 }
                 Some(v) => v,
             };
-            if !(&key.contains('.')) {
+
+            if key.to_lowercase().contains("keep_partition_by_columns") {
+                let renamed_key = if !&key.starts_with("hive.") {
+                    format!("hive.{}", key)
+                } else {
+                    key
+                };
+                options.insert(renamed_key.to_lowercase(), value_string.to_lowercase());
+            } else if !(&key.contains('.')) {
                 // If config does not belong to any namespace, assume it is
                 // a format option and apply the format prefix for backwards
                 // compatibility.

datafusion/sqllogictest/test_files/copy.slt

@@ -166,6 +166,23 @@ physical_plan
 01)DataSinkExec: sink=ParquetSink(file_groups=[])
 02)--MemoryExec: partitions=1, partition_sizes=[1]
 
+# Copy to directory as partitioned files with keep_partition_by_columns enabled
+query TT
+COPY (values ('1', 'a'), ('2', 'b'), ('3', 'c')) TO 'test_files/scratch/copy/partitioned_table4/' STORED AS parquet PARTITIONED BY (column1)
+OPTIONS (KEEP_PARTITION_BY_COLUMNS true);
+----
+3
+
+# validate generated file contains tables
+statement ok
+CREATE EXTERNAL TABLE validate_partitioned_parquet4 STORED AS PARQUET
+LOCATION 'test_files/scratch/copy/partitioned_table4/column1=1/*.parquet';
+
+query TT
+select column1, column2 from validate_partitioned_parquet4 order by column1,column2;
+----
+1 a
+
 # Copy more files to directory via query
 query IT
 COPY (select * from source_table UNION ALL select * from source_table) to 'test_files/scratch/copy/table/' STORED AS PARQUET;

docs/source/user-guide/sql/dml.md

@@ -39,7 +39,10 @@ TO '<i><b>file_name</i></b>'
 clause is not specified, it will be inferred from the file extension if possible.
 
 `PARTITIONED BY` specifies the columns to use for partitioning the output files into
-separate hive-style directories.
+separate hive-style directories. By default, columns used in `PARTITIONED BY` will be removed
+from the output format. If you want to keep the columns, you should provide the option
+`KEEP_PARTITION_BY_COLUMNS true`. `KEEP_PARTITION_BY_COLUMNS` flag can also be enabled
+through `ExecutionOptions` within `SessionConfig`.
 
 The output format is determined by the first match of the following rules:
 

docs/source/user-guide/sql/write_options.md

@@ -70,6 +70,16 @@ In this example, we write the entirety of `source_table` out to a folder of parq
 
 ## Available Options
 
+### Hive Specific Options
+
+The following options are available when writing hive-style partitioned data.
+
+| Option                    | Description                                                                        | Default Value |
+| ------------------------- | ---------------------------------------------------------------------------------- | ------------- |
+| KEEP_PARTITION_BY_COLUMNS | Flag to retain the columns in the output data when using `PARTITIONED BY` queries. | false         |
+
+Note: `KEEP_PARTITION_BY_COLUMNS` flag can also be enabled through `ExecutionOptions` within `SessionConfig`.
+
 ### JSON Format Specific Options
 
 The following options are available when writing JSON files. Note: If any unsupported option is specified, an error will be raised and the query will fail.