adamreeve commented on code in PR #7111:
URL: https://github.com/apache/arrow-rs/pull/7111#discussion_r1999945331


##########
parquet/src/arrow/arrow_writer/mod.rs:
##########
@@ -457,20 +476,49 @@ type SharedColumnChunk = Arc<Mutex<ArrowColumnChunkData>>;
 #[derive(Default)]
 struct ArrowPageWriter {
     buffer: SharedColumnChunk,
+    #[cfg(feature = "encryption")]
+    page_encryptor: Option<PageEncryptor>,
+    #[cfg(not(feature = "encryption"))]
+    page_encryptor: Option<Never>,
+}
+
+#[cfg(feature = "encryption")]
+impl ArrowPageWriter {
+    pub fn with_encryptor(mut self, page_encryptor: Option<PageEncryptor>) -> 
Self {
+        self.page_encryptor = page_encryptor;
+        self
+    }
 }
 
 impl PageWriter for ArrowPageWriter {
     fn write_page(&mut self, page: CompressedPage) -> Result<PageWriteSpec> {
         let mut buf = self.buffer.try_lock().unwrap();
-        let page_header = page.to_thrift_header();
-        let header = {
-            let mut header = Vec::with_capacity(1024);
-            let mut protocol = TCompactOutputProtocol::new(&mut header);
-            page_header.write_to_out_protocol(&mut protocol)?;
-            Bytes::from(header)
+
+        let data = match self.page_encryptor.as_ref() {
+            #[cfg(feature = "encryption")]
+            Some(encryptor) => {
+                let encrypted_buffer = encryptor.encrypt_page(&page)?;
+                Bytes::from(encrypted_buffer)
+            }
+            _ => page.compressed_page().buffer().clone(),
+        };
+
+        let mut page_header = page.to_thrift_header();
+        page_header.compressed_page_size = data.len() as i32;
+
+        let mut header = Vec::with_capacity(1024);
+        match self.page_encryptor.as_ref() {

Review Comment:
   If the logic can't be shared easily then I don't think the new traits are 
useful, `serialize_page` could just be a normal method under `impl 
ArrowPageWriter` and `impl SerializedPageWriter`



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscr...@arrow.apache.org

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org

Reply via email to