Move schema and writer properties into ArrowRowGroupWriterFactory

Author: rok

parquet/src/arrow/arrow_writer/mod.rs

@@ -141,7 +141,7 @@ pub struct ArrowWriter<W: Write> {
     arrow_schema: SchemaRef,
 
     /// Creates new [`ArrowRowGroupWriter`] instances as required
-    pub row_group_writer_factory: ArrowRowGroupWriterFactory,
+    row_group_writer_factory: ArrowRowGroupWriterFactory,
 
     /// The length of arrays to write to each row group
     max_row_group_size: usize,
@@ -199,9 +199,14 @@ impl<W: Write + Send> ArrowWriter<W> {
         let max_row_group_size = props.max_row_group_size();
 
         let file_writer =
-            SerializedFileWriter::new(writer, schema.root_schema_ptr(), Arc::new(props))?;
+            SerializedFileWriter::new(writer, schema.root_schema_ptr(), Arc::new(props.clone()))?;
 
-        let row_group_writer_factory = ArrowRowGroupWriterFactory::new(&file_writer);
+        let row_group_writer_factory = ArrowRowGroupWriterFactory::new(
+            &file_writer,
+            schema,
+            arrow_schema.clone(),
+            props.into(),
+        );
 
         Ok(Self {
             writer: file_writer,
@@ -272,12 +277,10 @@ impl<W: Write + Send> ArrowWriter<W> {
 
         let in_progress = match &mut self.in_progress {
             Some(in_progress) => in_progress,
-            x => x.insert(self.row_group_writer_factory.create_row_group_writer(
-                self.writer.schema_descr(),
-                self.writer.properties(),
-                &self.arrow_schema,
-                self.writer.flushed_row_groups().len(),
-            )?),
+            x => x.insert(
+                self.row_group_writer_factory
+                    .create_row_group_writer(self.writer.flushed_row_groups().len())?,
+            ),
         };
 
         // If would exceed max_row_group_size, split batch
@@ -794,58 +797,61 @@ impl ArrowRowGroupWriter {
 /// This is used by [`ArrowWriter`] to create row group writers, but can be used
 /// directly for lower level API.
 pub struct ArrowRowGroupWriterFactory {
+    schema: SchemaDescriptor,
+    arrow_schema: SchemaRef,
+    props: WriterPropertiesPtr,
     #[cfg(feature = "encryption")]
     file_encryptor: Option<Arc<FileEncryptor>>,
 }
 
 impl ArrowRowGroupWriterFactory {
     /// Creates a new [`ArrowRowGroupWriterFactory`] using provided [`SerializedFileWriter`].
     #[cfg(feature = "encryption")]
-    pub fn new<W: Write + Send>(file_writer: &SerializedFileWriter<W>) -> Self {
+    pub fn new<W: Write + Send>(
+        file_writer: &SerializedFileWriter<W>,
+        schema: SchemaDescriptor,
+        arrow_schema: SchemaRef,
+        props: WriterPropertiesPtr,
+    ) -> Self {
         Self {
+            schema,
+            arrow_schema,
+            props,
             file_encryptor: file_writer.file_encryptor(),
         }
     }
 
     #[cfg(not(feature = "encryption"))]
-    pub fn new<W: Write + Send>(_file_writer: &SerializedFileWriter<W>) -> Self {
-        Self {}
+    pub fn new<W: Write + Send>(
+        _file_writer: &SerializedFileWriter<W>,
+        schema: SchemaDescriptor,
+        arrow_schema: SchemaRef,
+        props: WriterPropertiesPtr,
+    ) -> Self {
+        Self {
+            schema,
+            arrow_schema,
+            props,
+        }
     }
 
     /// Creates a new [`ArrowRowGroupWriter`] for the given parquet schema and writer properties.
     #[cfg(feature = "encryption")]
-    pub fn create_row_group_writer(
-        &self,
-        parquet: &SchemaDescriptor,
-        props: &WriterPropertiesPtr,
-        arrow: &SchemaRef,
-        row_group_index: usize,
-    ) -> Result<ArrowRowGroupWriter> {
-        let mut writers = Vec::with_capacity(arrow.fields.len());
-        let mut leaves = parquet.columns().iter();
-        let column_factory = ArrowColumnWriterFactory::new()
-            .with_file_encryptor(row_group_index, self.file_encryptor.clone());
-        for field in &arrow.fields {
-            column_factory.get_arrow_column_writer(
-                field.data_type(),
-                props,
-                &mut leaves,
-                &mut writers,
-            )?;
-        }
-        Ok(ArrowRowGroupWriter::new(writers, arrow))
+    pub fn create_row_group_writer(&self, row_group_index: usize) -> Result<ArrowRowGroupWriter> {
+        let writers = get_column_writers_with_encryptor(
+            &self.schema,
+            &self.props,
+            &self.arrow_schema,
+            self.file_encryptor.clone(),
+            row_group_index,
+        )?;
+        Ok(ArrowRowGroupWriter::new(writers, &self.arrow_schema))
     }
 
     #[cfg(not(feature = "encryption"))]
-    pub fn create_row_group_writer(
-        &self,
-        parquet: &SchemaDescriptor,
-        props: &WriterPropertiesPtr,
-        arrow: &SchemaRef,
-        _row_group_index: usize,
-    ) -> Result<ArrowRowGroupWriter> {
-        let writers = get_column_writers(parquet, props, arrow)?;
-        Ok(ArrowRowGroupWriter::new(writers, arrow))
+    pub fn create_row_group_writer(&self, _row_group_index: usize) -> Result<ArrowRowGroupWriter> {
+        let writers = get_column_writers(&self.schema, &self.props, &self.arrow_schema)?;
+        Ok(ArrowRowGroupWriter::new(writers, &self.arrow_schema))
     }
 }
 
@@ -871,7 +877,7 @@ pub fn get_column_writers(
 
 /// Returns the [`ArrowColumnWriter`] for a given schema and supports columnar encryption
 #[cfg(feature = "encryption")]
-pub fn get_column_writers_with_encryptor(
+fn get_column_writers_with_encryptor(
     parquet: &SchemaDescriptor,
     props: &WriterPropertiesPtr,
     arrow: &SchemaRef,
@@ -894,7 +900,7 @@ pub fn get_column_writers_with_encryptor(
 }
 
 /// Gets [`ArrowColumnWriter`] instances for different data types
-pub struct ArrowColumnWriterFactory {
+struct ArrowColumnWriterFactory {
     #[cfg(feature = "encryption")]
     row_group_index: usize,
     #[cfg(feature = "encryption")]
@@ -919,7 +925,7 @@ impl ArrowColumnWriterFactory {
     }
 
     #[cfg(feature = "encryption")]
-    fn with_file_encryptor(
+    pub fn with_file_encryptor(
         mut self,
         row_group_index: usize,
         file_encryptor: Option<Arc<FileEncryptor>>,

parquet/src/encryption/encrypt.rs

@@ -288,14 +288,14 @@ impl EncryptionPropertiesBuilder {
 
 #[derive(Debug)]
 /// The encryption configuration for a single Parquet file
-pub struct FileEncryptor {
+pub(crate) struct FileEncryptor {
     properties: FileEncryptionProperties,
     aad_file_unique: Vec<u8>,
     file_aad: Vec<u8>,
 }
 
 impl FileEncryptor {
-    pub fn new(properties: FileEncryptionProperties) -> Result<Self> {
+    pub(crate) fn new(properties: FileEncryptionProperties) -> Result<Self> {
         // Generate unique AAD for file
         let rng = SystemRandom::new();
         let mut aad_file_unique = vec![0u8; 8];

parquet/src/file/writer.rs

@@ -433,7 +433,7 @@ impl<W: Write + Send> SerializedFileWriter<W> {
 
     /// Get the file encryptor used by this instance to encrypt data
     #[cfg(feature = "encryption")]
-    pub fn file_encryptor(&self) -> Option<Arc<FileEncryptor>> {
+    pub(crate) fn file_encryptor(&self) -> Option<Arc<FileEncryptor>> {
         self.file_encryptor.clone()
     }
 }

parquet/tests/encryption/encryption.rs

@@ -1153,9 +1153,14 @@ async fn test_multi_threaded_encrypted_writing() {
     let mut file_writer =
         SerializedFileWriter::new(&temp_file, root_schema.clone(), props.clone()).unwrap();
 
-    let arrow_row_group_writer_factory = ArrowRowGroupWriterFactory::new(&file_writer);
+    let arrow_row_group_writer_factory = ArrowRowGroupWriterFactory::new(
+        &file_writer,
+        parquet_schema,
+        schema.clone(),
+        props.clone(),
+    );
     let arrow_row_group_writer = arrow_row_group_writer_factory
-        .create_row_group_writer(&parquet_schema, &props, &schema, 0)
+        .create_row_group_writer(0)
         .unwrap();
 
     // Get column writers with encryptor from ArrowRowGroupWriter