Get column path from descriptor rather than Arrow schema

Author: adamreeve

parquet/src/arrow/arrow_writer/mod.rs

@@ -821,11 +821,8 @@ fn get_column_writers_with_encryptor(
 ) -> Result<Vec<ArrowColumnWriter>> {
     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,
-        file_encryptor,
-        arrow.clone(),
-    );
+    let column_factory =
+        ArrowColumnWriterFactory::new().with_file_encryptor(row_group_index, file_encryptor);
     for field in &arrow.fields {
         column_factory.get_arrow_column_writer(
             field.data_type(),
@@ -843,8 +840,6 @@ struct ArrowColumnWriterFactory {
     row_group_index: usize,
     #[cfg(feature = "encryption")]
     file_encryptor: Option<Arc<FileEncryptor>>,
-    #[cfg(feature = "encryption")]
-    schema_ref: Option<SchemaRef>,
 }
 
 impl ArrowColumnWriterFactory {
@@ -854,8 +849,6 @@ impl ArrowColumnWriterFactory {
             row_group_index: 0,
             #[cfg(feature = "encryption")]
             file_encryptor: None,
-            #[cfg(feature = "encryption")]
-            schema_ref: None,
         }
     }
 
@@ -864,28 +857,31 @@ impl ArrowColumnWriterFactory {
         mut self,
         row_group_index: usize,
         file_encryptor: Option<Arc<FileEncryptor>>,
-        schema_ref: SchemaRef,
     ) -> Self {
         self.row_group_index = row_group_index;
         self.file_encryptor = file_encryptor;
-        self.schema_ref = Some(schema_ref);
         self
     }
 
-    // todo: add column path
     #[cfg(feature = "encryption")]
-    fn create_page_writer(&self, column_index: usize) -> Box<ArrowPageWriter> {
+    fn create_page_writer(
+        &self,
+        column_descriptor: &ColumnDescPtr,
+        column_index: usize,
+    ) -> Box<ArrowPageWriter> {
         let page_encryptor = self.file_encryptor.as_ref().map(|fe| {
-            let binding = self.schema_ref.clone().unwrap();
-            let column_path = binding.field(column_index).name().as_bytes().to_vec();
-
+            let column_path = column_descriptor.path().string();
             PageEncryptor::new(fe.clone(), self.row_group_index, column_index, column_path)
         });
         Box::new(ArrowPageWriter::default().with_encryptor(page_encryptor))
     }
 
     #[cfg(not(feature = "encryption"))]
-    fn create_page_writer(&self, _column_index: usize) -> Box<ArrowPageWriter> {
+    fn create_page_writer(
+        &self,
+        _column_descriptor: &ColumnDescPtr,
+        _column_index: usize,
+    ) -> Box<ArrowPageWriter> {
         Box::<ArrowPageWriter>::default()
     }
 
@@ -897,7 +893,7 @@ impl ArrowColumnWriterFactory {
         out: &mut Vec<ArrowColumnWriter>,
     ) -> Result<()> {
         let col = |desc: &ColumnDescPtr| {
-            let page_writer = self.create_page_writer(out.len());
+            let page_writer = self.create_page_writer(desc, out.len());
             let chunk = page_writer.buffer.clone();
             let writer = get_column_writer(desc.clone(), props.clone(), page_writer);
             ArrowColumnWriter {
@@ -907,7 +903,7 @@ impl ArrowColumnWriterFactory {
         };
 
         let bytes = |desc: &ColumnDescPtr| {
-            let page_writer = self.create_page_writer(out.len());
+            let page_writer = self.create_page_writer(desc, out.len());
             let chunk = page_writer.buffer.clone();
             let writer = GenericColumnWriter::new(desc.clone(), props.clone(), page_writer);
             ArrowColumnWriter {

parquet/src/encryption/encrypt.rs

@@ -170,13 +170,14 @@ impl FileEncryptor {
         RingGcmBlockEncryptor::new(&self.properties.footer_key.key)
     }
 
-    pub(crate) fn get_column_encryptor(&self, column_path: &Vec<u8>) -> RingGcmBlockEncryptor {
+    pub(crate) fn get_column_encryptor(&self, column_path: &str) -> RingGcmBlockEncryptor {
         if self.properties.column_keys.is_empty() {
             return RingGcmBlockEncryptor::new(self.properties.footer_key.key());
         }
+        // TODO: Column paths should be stored as String
+        let column_path = column_path.as_bytes();
         match self.properties.column_keys.get(column_path) {
             None => todo!("Handle unencrypted columns"),
-            // None => RingGcmBlockEncryptor::new(self.properties.footer_key.key()),
             Some(column_key) => RingGcmBlockEncryptor::new(column_key.key()),
         }
     }

parquet/src/encryption/page_encryptor.rs

@@ -29,23 +29,23 @@ pub struct PageEncryptor {
     file_encryptor: Arc<FileEncryptor>,
     row_group_index: usize,
     column_index: usize,
+    column_path: String,
     page_index: usize,
-    column_path: Vec<u8>,
 }
 
 impl PageEncryptor {
     pub fn new(
         file_encryptor: Arc<FileEncryptor>,
         row_group_index: usize,
         column_index: usize,
-        column_path: Vec<u8>,
+        column_path: String,
     ) -> Self {
         Self {
             file_encryptor,
             row_group_index,
             column_index,
-            page_index: 0,
             column_path,
+            page_index: 0,
         }
     }
 
@@ -66,7 +66,7 @@ impl PageEncryptor {
             self.column_index,
             Some(self.page_index),
         )?;
-        let mut encryptor = self.file_encryptor.get_footer_encryptor();
+        let mut encryptor = self.file_encryptor.get_column_encryptor(&self.column_path);
         // todo: use column encryptor when needed
         // self.file_encryptor.get_column_encryptor(self.column_path.as_ref())
         let encrypted_buffer = encryptor.encrypt(page.data(), &aad);

parquet/src/file/writer.rs

@@ -596,7 +596,7 @@ impl<'a, W: Write + Send> SerializedRowGroupWriter<'a, W> {
                         file_encryptor,
                         row_group_index,
                         column_index,
-                        column.path().string().into_bytes(),
+                        column.path().string(),
                     )
                 });