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


##########
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;

Review Comment:
   I've pushed a commit to change this, let me know what you think: 
https://github.com/apache/arrow-rs/pull/7111/commits/4032b4f70c5407a1ce98358d1440cd00771543d5



-- 
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