Parquet encryption

parquet/Cargo.toml

@@ -30,8 +30,10 @@ edition = "2018"
 
 [dependencies]
 # update note: pin `parquet-format` to specific version until it does not break at minor
-# version, see ARROW-11187.
-parquet-format = "~2.6.1"
+# version, see ARROW-11187.  update: since this comment, it is now pinned at ~4.0.0 and
+# upstream arrow-rs parquet vendors it
+parquet-format = "~4.0.0"
+aes-gcm = "0.10.3"
 byteorder = "1"
 thrift = "0.13"
 snap = { version = "1.0", optional = true }
@@ -45,7 +47,9 @@ arrow = { path = "../arrow", version = "5.0.0", optional = true }
 base64 = { version = "0.13", optional = true }
 clap = { version = "2.33.3", optional = true }
 serde_json = { version = "1.0", features = ["preserve_order"], optional = true }
+serde = { version = "1.0.115", features = ["derive"] }
 rand = "0.8"
+sha3 = "0.10.8"
 
 [dev-dependencies]
 criterion = "0.3"

parquet/src/basic.rs

@@ -250,6 +250,15 @@ pub enum Encoding {
     ///
     /// The ids are encoded using the RLE encoding.
     RLE_DICTIONARY,
+
+    /// Encoding for floating-point data.
+    ///
+    /// K byte-streams are created where K is the size in bytes of the data type.
+    /// The individual bytes of an FP value are scattered to the corresponding stream and
+    /// the streams are concatenated.
+    /// This itself does not reduce the size of the data but can lead to better compression
+    /// afterwards.
+    BYTE_STREAM_SPLIT,
 }
 
 // ----------------------------------------------------------------------
@@ -701,6 +710,7 @@ impl convert::From<parquet::Encoding> for Encoding {
             parquet::Encoding::DeltaLengthByteArray => Encoding::DELTA_LENGTH_BYTE_ARRAY,
             parquet::Encoding::DeltaByteArray => Encoding::DELTA_BYTE_ARRAY,
             parquet::Encoding::RleDictionary => Encoding::RLE_DICTIONARY,
+            parquet::Encoding::ByteStreamSplit => Encoding::BYTE_STREAM_SPLIT,
         }
     }
 }
@@ -716,6 +726,7 @@ impl convert::From<Encoding> for parquet::Encoding {
             Encoding::DELTA_LENGTH_BYTE_ARRAY => parquet::Encoding::DeltaLengthByteArray,
             Encoding::DELTA_BYTE_ARRAY => parquet::Encoding::DeltaByteArray,
             Encoding::RLE_DICTIONARY => parquet::Encoding::RleDictionary,
+            Encoding::BYTE_STREAM_SPLIT => parquet::Encoding::ByteStreamSplit,
         }
     }
 }

parquet/src/column/page.rs

@@ -139,11 +139,14 @@ impl CompressedPage {
         self.uncompressed_size
     }
 
-    /// Returns compressed size in bytes.
+    /// Returns compressed size (but not encrypted size) in bytes.
     ///
-    /// Note that it is assumed that buffer is compressed, but it may not be. In this
-    /// case compressed size will be equal to uncompressed size.
-    pub fn compressed_size(&self) -> usize {
+    /// Note that it is assumed that buffer is compressed, but it may not be. In this case
+    /// compressed size will be equal to uncompressed size.
+    ///
+    /// Other so-called "(total_)?compressed_size" fields include encryption overhead, when
+    /// applicable, which this does not.
+    pub fn compressed_unencrypted_size(&self) -> usize {
         self.compressed_page.buffer().len()
     }
 
@@ -206,7 +209,7 @@ pub trait PageWriter {
     ///
     /// This method is called for every compressed page we write into underlying buffer,
     /// either data page or dictionary page.
-    fn write_page(&mut self, page: CompressedPage) -> Result<PageWriteSpec>;
+    fn write_page(&mut self, page: CompressedPage, aad_page_ordinal: Option<u16>) -> Result<PageWriteSpec>;
 
     /// Writes column chunk metadata into the output stream/sink.
     ///
@@ -299,7 +302,7 @@ mod tests {
 
         assert_eq!(cpage.page_type(), PageType::DATA_PAGE);
         assert_eq!(cpage.uncompressed_size(), 5);
-        assert_eq!(cpage.compressed_size(), 3);
+        assert_eq!(cpage.compressed_unencrypted_size(), 3);
         assert_eq!(cpage.num_values(), 10);
         assert_eq!(cpage.encoding(), Encoding::PLAIN);
         assert_eq!(cpage.data(), &[0, 1, 2]);

parquet/src/column/writer.rs

@@ -168,6 +168,7 @@ pub struct ColumnWriterImpl<T: DataType> {
     descr: ColumnDescPtr,
     props: WriterPropertiesPtr,
     page_writer: Box<dyn PageWriter>,
+    page_ordinal: usize,
     has_dictionary: bool,
     dict_encoder: Option<DictEncoder<T>>,
     encoder: Box<dyn Encoder<T>>,
@@ -185,6 +186,8 @@ pub struct ColumnWriterImpl<T: DataType> {
     total_bytes_written: u64,
     total_rows_written: u64,
     total_uncompressed_size: u64,
+    // Includes encryption overhead -- the thrift definition field includes encryption overhead, and
+    // we keep its name here.
     total_compressed_size: u64,
     total_num_values: u64,
     dictionary_page_offset: Option<u64>,
@@ -231,10 +234,14 @@ impl<T: DataType> ColumnWriterImpl<T> {
         )
         .unwrap();
 
+        // We start counting pages from zero.
+        let page_ordinal: usize = 0;
+
         Self {
             descr,
             props,
             page_writer,
+            page_ordinal,
             has_dictionary,
             dict_encoder,
             encoder: fallback_encoder,
@@ -824,7 +831,10 @@ impl<T: DataType> ColumnWriterImpl<T> {
     /// Writes compressed data page into underlying sink and updates global metrics.
     #[inline]
     fn write_data_page(&mut self, page: CompressedPage) -> Result<()> {
-        let page_spec = self.page_writer.write_page(page)?;
+        let page_ordinal = self.page_ordinal;
+        let aad_page_ordinal: Option<u16> = Some(page_ordinal as u16);
+        self.page_ordinal += 1;
+        let page_spec = self.page_writer.write_page(page, aad_page_ordinal)?;
         self.update_metrics_for_page(page_spec);
         Ok(())
     }
@@ -858,7 +868,7 @@ impl<T: DataType> ColumnWriterImpl<T> {
             CompressedPage::new(dict_page, uncompressed_size)
         };
 
-        let page_spec = self.page_writer.write_page(compressed_page)?;
+        let page_spec = self.page_writer.write_page(compressed_page, None)?;
         self.update_metrics_for_page(page_spec);
         Ok(())
     }
@@ -1026,10 +1036,10 @@ fn has_dictionary_support(kind: Type, props: &WriterProperties) -> bool {
 mod tests {
     use rand::distributions::uniform::SampleUniform;
 
-    use crate::column::{
+    use crate::{column::{
         page::PageReader,
         reader::{get_column_reader, get_typed_column_reader, ColumnReaderImpl},
-    };
+    }, file::encryption::USUAL_ENCRYPTION_OVERHEAD};
     use crate::file::{
         properties::WriterProperties, reader::SerializedPageReader,
         writer::SerializedPageWriter,
@@ -1642,21 +1652,24 @@ mod tests {
         );
     }
 
+    const TEST_ROW_GROUP_ORDINAL: i16 = 1234;
+    const TEST_COLUMN_ORDINAL: u16 = 135;
+
     #[test]
     fn test_column_writer_add_data_pages_with_dict() {
         // ARROW-5129: Test verifies that we add data page in case of dictionary encoding
         // and no fallback occurred so far.
         let file = get_temp_file("test_column_writer_add_data_pages_with_dict", &[]);
         let sink = FileSink::new(&file);
-        let page_writer = Box::new(SerializedPageWriter::new(sink));
+        let page_writer = Box::new(SerializedPageWriter::new(sink, None, TEST_ROW_GROUP_ORDINAL, TEST_COLUMN_ORDINAL));
         let props = Arc::new(
             WriterProperties::builder()
                 .set_data_pagesize_limit(15) // actually each page will have size 15-18 bytes
                 .set_write_batch_size(3) // write 3 values at a time
                 .build(),
         );
         let data = &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
-        let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
+        let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props.clone());
         writer.write_batch(data, None, None).unwrap();
         let (bytes_written, _, _) = writer.close().unwrap();
 
@@ -1665,8 +1678,12 @@ mod tests {
         let mut page_reader = Box::new(
             SerializedPageReader::new(
                 source,
+                None,
+                TEST_ROW_GROUP_ORDINAL,
+                TEST_COLUMN_ORDINAL,
                 data.len() as i64,
                 Compression::UNCOMPRESSED,
+                props.dictionary_enabled(&ColumnPath::from("col")),
                 Int32Type::get_physical_type(),
             )
             .unwrap(),
@@ -1803,7 +1820,7 @@ mod tests {
     ) {
         let file = get_temp_file(file_name, &[]);
         let sink = FileSink::new(&file);
-        let page_writer = Box::new(SerializedPageWriter::new(sink));
+        let page_writer = Box::new(SerializedPageWriter::new(sink, None, TEST_ROW_GROUP_ORDINAL, TEST_COLUMN_ORDINAL));
 
         let max_def_level = match def_levels {
             Some(buf) => *buf.iter().max().unwrap_or(&0i16),
@@ -1823,11 +1840,12 @@ mod tests {
             max_batch_size = cmp::max(max_batch_size, levels.len());
         }
 
+        let props = Arc::new(props);
         let mut writer = get_test_column_writer::<T>(
             page_writer,
             max_def_level,
             max_rep_level,
-            Arc::new(props),
+            props.clone(),
         );
 
         let values_written = writer.write_batch(values, def_levels, rep_levels).unwrap();
@@ -1838,8 +1856,12 @@ mod tests {
         let page_reader = Box::new(
             SerializedPageReader::new(
                 source,
+                None,
+                TEST_ROW_GROUP_ORDINAL,
+                TEST_COLUMN_ORDINAL,
                 column_metadata.num_values(),
                 column_metadata.compression(),
+                props.dictionary_enabled(&ColumnPath::from("col")),
                 T::get_physical_type(),
             )
             .unwrap(),
@@ -1977,20 +1999,39 @@ mod tests {
 
     /// Returns page writer that collects pages without serializing them.
     fn get_test_page_writer() -> Box<dyn PageWriter> {
-        Box::new(TestPageWriter {})
+        Box::new(TestPageWriter {simulate_encrypted: false, last_page_ordinal: None})
     }
 
-    struct TestPageWriter {}
+    struct TestPageWriter {
+        /// Always false, currently -- enabling would just affect return values that get fed into
+        /// test assertions.
+        simulate_encrypted: bool,
+        last_page_ordinal: Option<u16>,
+    }
 
     impl PageWriter for TestPageWriter {
-        fn write_page(&mut self, page: CompressedPage) -> Result<PageWriteSpec> {
+        fn write_page(&mut self, page: CompressedPage, aad_page_ordinal: Option<u16>) -> Result<PageWriteSpec> {
+            // We're a bit loose in this assertion -- the caller could write or not write a dictionary page.
+            match aad_page_ordinal {
+                Some(n) if n != 0 => {
+                    assert_eq!(self.last_page_ordinal, Some(n - 1));
+                }
+                _ => {
+                    assert_eq!(None, self.last_page_ordinal);
+                }
+            }
+            self.last_page_ordinal = aad_page_ordinal;
+
+            // Note, the normal PageWriteSpec result would include PageMetaData overhead, and these
+            // values are thus not perfectly faked, but the only thing that looks at them are test
+            // assertions.
             let mut res = PageWriteSpec::new();
             res.page_type = page.page_type();
             res.uncompressed_size = page.uncompressed_size();
-            res.compressed_size = page.compressed_size();
+            res.compressed_size = self.simulate_encrypted as usize * USUAL_ENCRYPTION_OVERHEAD + page.compressed_unencrypted_size();
             res.num_values = page.num_values();
             res.offset = 0;
-            res.bytes_written = page.data().len() as u64;
+            res.bytes_written = (self.simulate_encrypted as usize * USUAL_ENCRYPTION_OVERHEAD + page.data().len()) as u64;
             Ok(res)
         }