Skip to content

Commit 5a64d83

Browse files
committed
refactor: including partitioning of arrow schema inserted into kv_metdata
1 parent 3893c8c commit 5a64d83

File tree

1 file changed

+21
-10
lines changed
  • datafusion/core/src/datasource/file_format

1 file changed

+21
-10
lines changed

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

Lines changed: 21 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -749,6 +749,26 @@ impl ParquetSink {
749749
}
750750
}
751751

752+
/// Create writer properties based upon configuration settings,
753+
/// including partitioning and the inclusion of arrow schema metadata.
754+
fn create_writer_props(&self) -> Result<WriterProperties> {
755+
let props = if !self.parquet_options.global.skip_arrow_metadata {
756+
let schema = if self.parquet_options.global.allow_single_file_parallelism {
757+
&self.get_writer_schema()
758+
} else {
759+
self.config.output_schema()
760+
};
761+
self.parquet_options
762+
.into_writer_properties_builder_with_arrow_schema(Some(schema))?
763+
.build()
764+
} else {
765+
self.parquet_options
766+
.into_writer_properties_builder()?
767+
.build()
768+
};
769+
Ok(props)
770+
}
771+
752772
/// Creates an AsyncArrowWriter which serializes a parquet file to an ObjectStore
753773
/// AsyncArrowWriters are used when individual parquet file serialization is not parallelized
754774
async fn create_async_arrow_writer(
@@ -791,16 +811,7 @@ impl DataSink for ParquetSink {
791811
data: SendableRecordBatchStream,
792812
context: &Arc<TaskContext>,
793813
) -> Result<u64> {
794-
let parquet_props = if !self.parquet_options.global.skip_arrow_metadata {
795-
let schema = self.config.output_schema();
796-
self.parquet_options
797-
.into_writer_properties_builder_with_arrow_schema(Some(schema))?
798-
.build()
799-
} else {
800-
self.parquet_options
801-
.into_writer_properties_builder()?
802-
.build()
803-
};
814+
let parquet_props = self.create_writer_props()?;
804815

805816
let object_store = context
806817
.runtime_env()

0 commit comments

Comments
 (0)