adamreeve commented on code in PR #7111:
URL: https://github.com/apache/arrow-rs/pull/7111#discussion_r1999969215
##########
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:
Yeah after looking at this in more detail, I don't think there's an easy way
to share this logic as `ArrowPageWriter` and `SerializedPageWriter` are too
different. I think we should get rid of the `PageModuleWriter` traits but
keeping separate `serialize_page` methods for with and without encryption
enabled would still be nice.
--
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
