feat(cube): Customizer interfaces

Author: srh

datafusion-examples/examples/advanced_parquet_index.rs

@@ -612,8 +612,12 @@ 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,10 @@ 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 +82,10 @@ 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 +190,10 @@ 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,6 +17,8 @@
 
 //! Options related to how parquet files should be written
 
+use std::sync::Arc;
+
 use crate::{
     config::{ParquetOptions, TableParquetOptions},
     DataFusionError, Result,
@@ -51,18 +53,53 @@ 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 trait WriterPropertiesCustomizer: Sync + Send + std::fmt::Debug {
+    fn adjust_write_properties(
+        &self,
+        builder: WriterPropertiesBuilder,
+    ) -> WriterPropertiesBuilder;
+    fn allow_single_file_parallelism(&self) -> bool {
+        true
+    }
+}
+
+#[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,9 @@ 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 +546,18 @@ 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 +566,10 @@ 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)