adamreeve commented on code in PR #7111:
URL: https://github.com/apache/arrow-rs/pull/7111#discussion_r1999942005
##########
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:
I see you've added a `PageModuleWriter` trait that tidies this up a bit
although this isn't really what I had in mind, I was hoping this logic could be
reused across the `ArrowPageWriter` and `SerializedPageWriter` rather than
having a similar trait in two different places. I'll see if I can make a change
to show what I mean.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
