WIP: ReaderOptionsCustomizer and WriterPropertiesCustomizer

Author: srh

datafusion-examples/examples/advanced_parquet_index.rs

@@ -612,8 +612,10 @@ impl AsyncFileReader for ParquetReaderWithCache {
 
     fn get_metadata(
         &mut self,
+        encryption_config: &Option<datafusion::parquet::file::encryption::ParquetEncryptionConfig>,
     ) -> BoxFuture<'_, datafusion::parquet::errors::Result<Arc<ParquetMetaData>>> {
         println!("get_metadata: {} returning cached metadata", self.filename);
+        assert!(encryption_config.is_none());
 
         // return the cached metadata so the parquet reader does not read it
         let metadata = self.metadata.clone();

datafusion/common/src/file_options/mod.rs

@@ -33,7 +33,7 @@ mod tests {
     use super::parquet_writer::ParquetWriterOptions;
     use crate::{
         config::{ConfigFileType, TableOptions},
-        file_options::{csv_writer::CsvWriterOptions, json_writer::JsonWriterOptions},
+        file_options::{csv_writer::CsvWriterOptions, json_writer::JsonWriterOptions, parquet_writer::WriterPropertiesConfig},
         parsers::CompressionTypeVariant,
         Result,
     };
@@ -79,7 +79,7 @@ mod tests {
         table_config.set_config_format(ConfigFileType::PARQUET);
         table_config.alter_with_string_hash_map(&option_map)?;
 
-        let parquet_options = ParquetWriterOptions::try_from(&table_config.parquet)?;
+        let parquet_options = ParquetWriterOptions::from_table_parquet_options(&table_config.parquet, WriterPropertiesConfig::noop().as_ref())?;
         let properties = parquet_options.writer_options();
 
         // Verify the expected options propagated down to parquet crate WriterProperties struct
@@ -184,7 +184,7 @@ mod tests {
         table_config.set_config_format(ConfigFileType::PARQUET);
         table_config.alter_with_string_hash_map(&option_map)?;
 
-        let parquet_options = ParquetWriterOptions::try_from(&table_config.parquet)?;
+        let parquet_options = ParquetWriterOptions::from_table_parquet_options(&table_config.parquet, WriterPropertiesConfig::noop().as_ref())?;
         let properties = parquet_options.writer_options();
 
         let col1 = ColumnPath::from(vec!["col1".to_owned()]);

datafusion/common/src/file_options/parquet_writer.rs

@@ -17,19 +17,18 @@
 
 //! Options related to how parquet files should be written
 
+use std::sync::Arc;
+
 use crate::{
     config::{ParquetOptions, TableParquetOptions},
     DataFusionError, Result,
 };
 
 use parquet::{
-    basic::{BrotliLevel, GzipLevel, ZstdLevel},
-    file::properties::{
+    basic::{BrotliLevel, GzipLevel, ZstdLevel}, file::properties::{
         EnabledStatistics, WriterProperties, WriterPropertiesBuilder, WriterVersion,
         DEFAULT_MAX_STATISTICS_SIZE, DEFAULT_STATISTICS_ENABLED,
-    },
-    format::KeyValue,
-    schema::types::ColumnPath,
+    }, format::KeyValue, schema::types::ColumnPath
 };
 
 /// Options for writing parquet files
@@ -51,18 +50,44 @@ impl ParquetWriterOptions {
     }
 }
 
-impl TryFrom<&TableParquetOptions> for ParquetWriterOptions {
-    type Error = DataFusionError;
-
-    fn try_from(parquet_table_options: &TableParquetOptions) -> Result<Self> {
+impl ParquetWriterOptions {
+    pub fn from_table_parquet_options(parquet_table_options: &TableParquetOptions, customizer: &dyn WriterPropertiesCustomizer) -> Result<Self> {
         // ParquetWriterOptions will have defaults for the remaining fields (e.g. sorting_columns)
+        let mut builder = WriterPropertiesBuilder::try_from(parquet_table_options)?;
+        builder = customizer.adjust_write_properties(builder);
+
         Ok(ParquetWriterOptions {
-            writer_options: WriterPropertiesBuilder::try_from(parquet_table_options)?
-                .build(),
+            writer_options: builder.build(),
         })
     }
 }
 
+pub struct WriterPropertiesConfig {
+    pub customizer: Arc<dyn WriterPropertiesCustomizer>,
+}
+
+impl WriterPropertiesConfig {
+    pub fn noop() -> Arc<dyn WriterPropertiesCustomizer> {
+        Arc::new(NoopWriterPropertiesCustomizer{})
+    }
+    pub fn todo() -> Arc<dyn WriterPropertiesCustomizer> {  // TODO upgrade DF: no
+        Self::noop()
+    }
+}
+
+pub trait WriterPropertiesCustomizer: Sync + Send + std::fmt::Debug {
+    fn adjust_write_properties(&self, builder: WriterPropertiesBuilder) -> WriterPropertiesBuilder;
+}
+
+#[derive(Debug)]
+pub struct NoopWriterPropertiesCustomizer;
+
+impl WriterPropertiesCustomizer for NoopWriterPropertiesCustomizer {
+    fn adjust_write_properties(&self, builder: WriterPropertiesBuilder) -> WriterPropertiesBuilder {
+        builder
+    }
+}
+
 impl TryFrom<&TableParquetOptions> for WriterPropertiesBuilder {
     type Error = DataFusionError;
 

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

@@ -34,6 +34,7 @@ use crate::execution::context::{SessionConfig, SessionState};
 
 use arrow::datatypes::{DataType, Schema, SchemaRef};
 use datafusion_common::config::TableOptions;
+use datafusion_common::file_options::parquet_writer::{WriterPropertiesConfig, WriterPropertiesCustomizer};
 use datafusion_common::{
     DEFAULT_ARROW_EXTENSION, DEFAULT_AVRO_EXTENSION, DEFAULT_CSV_EXTENSION,
     DEFAULT_JSON_EXTENSION, DEFAULT_PARQUET_EXTENSION,
@@ -543,6 +544,11 @@ impl ReadOptions<'_> for CsvReadOptions<'_> {
     }
 }
 
+/// Retrieves `WriterPropertiesConfig` from the `SessionConfig::extensions` field, constructing a Noop customizer if not present.
+pub fn get_writer_properties_customizer(config: &SessionConfig) -> Arc<dyn WriterPropertiesCustomizer> {
+    config.get_extension::<WriterPropertiesConfig>().map_or_else(|| WriterPropertiesConfig::noop(), |cfg| cfg.customizer.clone())
+}
+
 #[cfg(feature = "parquet")]
 #[async_trait]
 impl ReadOptions<'_> for ParquetReadOptions<'_> {
@@ -551,7 +557,8 @@ impl ReadOptions<'_> for ParquetReadOptions<'_> {
         config: &SessionConfig,
         table_options: TableOptions,
     ) -> ListingOptions {
-        let mut file_format = ParquetFormat::new().with_options(table_options.parquet);
+        let customizer = get_writer_properties_customizer(config);
+        let mut file_format = ParquetFormat::new().with_options(table_options.parquet).with_customizer(customizer);
 
         if let Some(parquet_pruning) = self.parquet_pruning {
             file_format = file_format.with_enable_pruning(parquet_pruning)