Test more files and add tolerance

Author: adamreeve

parquet/tests/metadata_memory.rs

@@ -19,6 +19,7 @@
 
 use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions};
 use parquet::encryption::decrypt::FileDecryptionProperties;
+use parquet::file::metadata::PageIndexPolicy;
 use std::alloc::{GlobalAlloc, Layout, System};
 use std::fs::File;
 use std::hint::black_box;
@@ -69,10 +70,24 @@ fn test_metadata_heap_memory() {
     // Run test cases sequentially so that heap allocations
     // are restricted to a single test case at a time.
     let test_data = arrow::util::test_util::parquet_test_data();
+    let reader_options =
+        ArrowReaderOptions::default().with_page_index_policy(PageIndexPolicy::Required);
 
     {
         let path = format!("{test_data}/alltypes_dictionary.parquet");
-        verify_metadata_heap_memory(&path, ArrowReaderOptions::default);
+        verify_metadata_heap_memory(&path, 0.0, || reader_options.clone());
+    }
+
+    {
+        // Calculated heap size doesn't match exactly, possibly due to extra overhead not accounted
+        // for in the HeapSize implementation for parquet::data_type::ByteArray.
+        let path = format!("{test_data}/alltypes_tiny_pages_plain.parquet");
+        verify_metadata_heap_memory(&path, 0.02, || reader_options.clone());
+    }
+
+    {
+        let path = format!("{test_data}/data_index_bloom_encoding_with_length.parquet");
+        verify_metadata_heap_memory(&path, 0.02, || reader_options.clone());
     }
 
     {
@@ -90,14 +105,16 @@ fn test_metadata_heap_memory() {
                 .with_column_key("float_field", column_2_key.into())
                 .build()
                 .unwrap();
-            ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties)
+            reader_options
+                .clone()
+                .with_file_decryption_properties(decryption_properties)
         };
 
-        verify_metadata_heap_memory(&path, get_options);
+        verify_metadata_heap_memory(&path, 0.0, get_options);
     }
 }
 
-fn verify_metadata_heap_memory<F>(path: &str, get_options: F)
+fn verify_metadata_heap_memory<F>(path: &str, rel_tol: f64, get_options: F)
 where
     F: FnOnce() -> ArrowReaderOptions,
 {
@@ -117,9 +134,23 @@ where
     black_box(metadata);
 
     assert!(metadata_heap_size > 0);
-    assert_eq!(
-        metadata_heap_size, allocated,
-        "Calculated heap size {} doesn't match allocated size {} for file {}",
-        metadata_heap_size, allocated, path
-    );
+    if rel_tol == 0.0 {
+        assert_eq!(
+            metadata_heap_size, allocated,
+            "Calculated heap size {} doesn't match the allocated size {} for file {}",
+            metadata_heap_size, allocated, path
+        );
+    } else {
+        assert!(rel_tol > 0.0 && rel_tol < 1.0);
+        let min_size = ((allocated as f64) * (1.0 - rel_tol)) as usize;
+        let max_size = ((allocated as f64) * (1.0 + rel_tol)) as usize;
+        assert!(
+            metadata_heap_size >= min_size && metadata_heap_size <= max_size,
+            "Calculated heap size {} doesn't match the allocated size {} within a relative tolerance of {} for file {}",
+            metadata_heap_size,
+            allocated,
+            rel_tol,
+            path
+        );
+    }
 }