refactor: return to the WriterPropertiesBuilder::TryFrom<TableParquetOptions>, and separately add the arrow_schema to the kv_metadata on the TableParquetOptions

Author: wiedld

datafusion/common/src/config.rs

@@ -1502,6 +1502,20 @@ impl TableParquetOptions {
     pub fn new() -> Self {
         Self::default()
     }
+
+    /// Set whether the encoding of the arrow metadata should occur
+    /// during the writing of parquet.
+    ///
+    /// Default is to encode the arrow schema in the file kv_metadata.
+    pub fn with_skip_arrow_metadata(self, skip: bool) -> Self {
+        Self {
+            global: ParquetOptions {
+                skip_arrow_metadata: skip,
+                ..self.global
+            },
+            ..self
+        }
+    }
 }
 
 impl ConfigField for TableParquetOptions {

datafusion/common/src/file_options/mod.rs

@@ -39,7 +39,7 @@ mod tests {
 
     use parquet::{
         basic::{Compression, Encoding, ZstdLevel},
-        file::properties::{EnabledStatistics, WriterVersion},
+        file::properties::{EnabledStatistics, WriterPropertiesBuilder, WriterVersion},
         schema::types::ColumnPath,
     };
 
@@ -78,10 +78,8 @@ mod tests {
         table_config.set_config_format(ConfigFileType::PARQUET);
         table_config.alter_with_string_hash_map(&option_map)?;
 
-        let properties = table_config
-            .parquet
-            .into_writer_properties_builder()?
-            .build();
+        let properties =
+            WriterPropertiesBuilder::try_from(&table_config.parquet)?.build();
 
         // Verify the expected options propagated down to parquet crate WriterProperties struct
         assert_eq!(properties.max_row_group_size(), 123);
@@ -185,10 +183,8 @@ mod tests {
         table_config.set_config_format(ConfigFileType::PARQUET);
         table_config.alter_with_string_hash_map(&option_map)?;
 
-        let properties = table_config
-            .parquet
-            .into_writer_properties_builder()?
-            .build();
+        let properties =
+            WriterPropertiesBuilder::try_from(&table_config.parquet)?.build();
 
         let col1 = ColumnPath::from(vec!["col1".to_owned()]);
         let col2_nested = ColumnPath::from(vec!["col2".to_owned(), "nested".to_owned()]);

datafusion/common/src/file_options/parquet_writer.rs

@@ -59,52 +59,49 @@ impl ParquetWriterOptions {
 }
 
 impl TableParquetOptions {
-    #[deprecated(
-        since = "44.0.0",
-        note = "Please use `TableParquetOptions::into_writer_properties_builder` and `TableParquetOptions::into_writer_properties_builder_with_arrow_schema`"
-    )]
-    pub fn try_from(table_opts: &TableParquetOptions) -> Result<ParquetWriterOptions> {
+    /// Add the arrow schema to the parquet kv_metadata.
+    /// If already exists, then overwrites.
+    pub fn arrow_schema(&mut self, schema: &Arc<Schema>) {
+        self.key_value_metadata.insert(
+            ARROW_SCHEMA_META_KEY.into(),
+            Some(encode_arrow_schema(schema)),
+        );
+    }
+}
+
+impl TryFrom<&TableParquetOptions> for ParquetWriterOptions {
+    type Error = DataFusionError;
+
+    fn try_from(parquet_table_options: &TableParquetOptions) -> Result<Self> {
         // ParquetWriterOptions will have defaults for the remaining fields (e.g. sorting_columns)
         Ok(ParquetWriterOptions {
-            writer_options: table_opts.into_writer_properties_builder()?.build(),
+            writer_options: WriterPropertiesBuilder::try_from(parquet_table_options)?
+                .build(),
         })
     }
+}
 
-    /// Convert the session's [`TableParquetOptions`] into a single write action's [`WriterPropertiesBuilder`].
-    ///
-    /// The returned [`WriterPropertiesBuilder`] includes customizations applicable per column.
-    pub fn into_writer_properties_builder(&self) -> Result<WriterPropertiesBuilder> {
-        self.into_writer_properties_builder_with_arrow_schema(None)
-    }
+impl TryFrom<&TableParquetOptions> for WriterPropertiesBuilder {
+    type Error = DataFusionError;
 
     /// Convert the session's [`TableParquetOptions`] into a single write action's [`WriterPropertiesBuilder`].
     ///
-    /// The returned [`WriterPropertiesBuilder`] includes customizations applicable per column,
-    /// as well as the arrow schema encoded into the kv_meta at [`ARROW_SCHEMA_META_KEY`].
-    pub fn into_writer_properties_builder_with_arrow_schema(
-        &self,
-        to_encode: Option<&Arc<Schema>>,
-    ) -> Result<WriterPropertiesBuilder> {
+    /// The returned [`WriterPropertiesBuilder`] includes customizations applicable per column.
+    fn try_from(table_parquet_options: &TableParquetOptions) -> Result<Self> {
         // Table options include kv_metadata and col-specific options
         let TableParquetOptions {
             global,
             column_specific_options,
             key_value_metadata,
-        } = self;
+        } = table_parquet_options;
 
         let mut builder = global.into_writer_properties_builder()?;
 
         // add kv_meta, if any
-        let mut kv_meta = key_value_metadata.to_owned();
-        if let Some(schema) = to_encode {
-            kv_meta.insert(
-                ARROW_SCHEMA_META_KEY.into(),
-                Some(encode_arrow_schema(schema)),
-            );
-        }
-        if !kv_meta.is_empty() {
+        if !key_value_metadata.is_empty() {
             builder = builder.set_key_value_metadata(Some(
-                kv_meta
+                key_value_metadata
+                    .to_owned()
                     .drain()
                     .map(|(key, value)| KeyValue { key, value })
                     .collect(),
@@ -621,8 +618,7 @@ mod tests {
             key_value_metadata: [(key, value)].into(),
         };
 
-        let writer_props = table_parquet_opts
-            .into_writer_properties_builder()
+        let writer_props = WriterPropertiesBuilder::try_from(&table_parquet_opts)
             .unwrap()
             .build();
         assert_eq!(
@@ -649,10 +645,10 @@ mod tests {
         let default_writer_props = WriterProperties::new();
 
         // WriterProperties::try_from(TableParquetOptions::default), a.k.a. using datafusion's defaults
-        let from_datafusion_defaults = default_table_writer_opts
-            .into_writer_properties_builder()
-            .unwrap()
-            .build();
+        let from_datafusion_defaults =
+            WriterPropertiesBuilder::try_from(&default_table_writer_opts)
+                .unwrap()
+                .build();
 
         // Expected: how the defaults should not match
         assert_ne!(
@@ -705,10 +701,10 @@ mod tests {
         // the TableParquetOptions::default, with only the bloom filter turned on
         let mut default_table_writer_opts = TableParquetOptions::default();
         default_table_writer_opts.global.bloom_filter_on_write = true;
-        let from_datafusion_defaults = default_table_writer_opts
-            .into_writer_properties_builder()
-            .unwrap()
-            .build();
+        let from_datafusion_defaults =
+            WriterPropertiesBuilder::try_from(&default_table_writer_opts)
+                .unwrap()
+                .build();
 
         // the WriterProperties::default, with only the bloom filter turned on
         let default_writer_props = WriterProperties::builder()
@@ -733,10 +729,10 @@ mod tests {
         let mut default_table_writer_opts = TableParquetOptions::default();
         default_table_writer_opts.global.bloom_filter_on_write = true;
         default_table_writer_opts.global.bloom_filter_fpp = Some(0.42);
-        let from_datafusion_defaults = default_table_writer_opts
-            .into_writer_properties_builder()
-            .unwrap()
-            .build();
+        let from_datafusion_defaults =
+            WriterPropertiesBuilder::try_from(&default_table_writer_opts)
+                .unwrap()
+                .build();
 
         // the WriterProperties::default, with only fpp set
         let default_writer_props = WriterProperties::builder()
@@ -765,10 +761,10 @@ mod tests {
         let mut default_table_writer_opts = TableParquetOptions::default();
         default_table_writer_opts.global.bloom_filter_on_write = true;
         default_table_writer_opts.global.bloom_filter_ndv = Some(42);
-        let from_datafusion_defaults = default_table_writer_opts
-            .into_writer_properties_builder()
-            .unwrap()
-            .build();
+        let from_datafusion_defaults =
+            WriterPropertiesBuilder::try_from(&default_table_writer_opts)
+                .unwrap()
+                .build();
 
         // the WriterProperties::default, with only ndv set
         let default_writer_props = WriterProperties::builder()

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

@@ -73,7 +73,7 @@ use parquet::arrow::{
     arrow_to_parquet_schema, parquet_to_arrow_schema, AsyncArrowWriter,
 };
 use parquet::file::metadata::{ParquetMetaData, ParquetMetaDataReader, RowGroupMetaData};
-use parquet::file::properties::WriterProperties;
+use parquet::file::properties::{WriterProperties, WriterPropertiesBuilder};
 use parquet::file::writer::SerializedFileWriter;
 use parquet::format::FileMetaData;
 use tokio::io::{AsyncWrite, AsyncWriteExt};
@@ -752,24 +752,23 @@ impl ParquetSink {
     /// Create writer properties based upon configuration settings,
     /// including partitioning and the inclusion of arrow schema metadata.
     fn create_writer_props(&self) -> Result<WriterProperties> {
-        let props = if !self.parquet_options.global.skip_arrow_metadata {
-            let schema = if self.parquet_options.global.allow_single_file_parallelism {
-                // If parallelizing writes, we may be also be doing hive style partitioning
-                // into multiple files which impacts the schema per file.
-                // Refer to `self.get_writer_schema()`
-                &self.get_writer_schema()
-            } else {
-                self.config.output_schema()
-            };
-            self.parquet_options
-                .into_writer_properties_builder_with_arrow_schema(Some(schema))?
-                .build()
+        let schema = if self.parquet_options.global.allow_single_file_parallelism {
+            // If parallelizing writes, we may be also be doing hive style partitioning
+            // into multiple files which impacts the schema per file.
+            // Refer to `self.get_writer_schema()`
+            &self.get_writer_schema()
         } else {
-            self.parquet_options
-                .into_writer_properties_builder()?
-                .build()
+            self.config.output_schema()
         };
-        Ok(props)
+
+        // TODO: avoid this clone in follow up PR, where the writer properties & schema
+        // are calculated once on `ParquetSink::new`
+        let mut parquet_opts = self.parquet_options.clone();
+        if !self.parquet_options.global.skip_arrow_metadata {
+            parquet_opts.arrow_schema(schema);
+        }
+
+        Ok(WriterPropertiesBuilder::try_from(&parquet_opts)?.build())
     }
 
     /// Creates an AsyncArrowWriter which serializes a parquet file to an ObjectStore