test: compression fallback wins

Modify the compression fallback test to illustrate the benefit in
an admittedly differently contrived case. The heuristic borrowed from
parquet-java is still not ideal for all cases, so we'll need more
configurability.

Author: mzabaluev

parquet/src/arrow/arrow_writer/mod.rs

@@ -4857,6 +4857,15 @@ mod tests {
         assert_eq!(chunk_page_stats, file_page_stats);
     }
 
+    fn get_dict_page_size(meta: &ColumnChunkMetaData, data: Bytes) -> usize {
+        let mut reader = SerializedPageReader::new(Arc::new(data), meta, 0, None).unwrap();
+        let page = reader.get_next_page().unwrap().unwrap();
+        match page {
+            Page::DictionaryPage { buf, .. } => buf.len(),
+            _ => panic!("expected DictionaryPage"),
+        }
+    }
+
     #[test]
     fn test_different_dict_page_size_limit() {
         let array = Arc::new(Int64Array::from_iter(0..1024 * 1024));
@@ -4881,60 +4890,56 @@ mod tests {
         let col0_meta = metadata.row_group(0).column(0);
         let col1_meta = metadata.row_group(0).column(1);
 
-        let get_dict_page_size = move |meta: &ColumnChunkMetaData| {
-            let mut reader =
-                SerializedPageReader::new(Arc::new(data.clone()), meta, 0, None).unwrap();
-            let page = reader.get_next_page().unwrap().unwrap();
-            match page {
-                Page::DictionaryPage { buf, .. } => buf.len(),
-                _ => panic!("expected DictionaryPage"),
-            }
-        };
-
-        assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
-        assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
+        assert_eq!(get_dict_page_size(col0_meta, data.clone()), 1024 * 1024);
+        assert_eq!(get_dict_page_size(col1_meta, data.clone()), 1024 * 1024 * 4);
     }
 
     #[test]
     fn test_dict_page_size_decided_by_compression_fallback() {
-        let array = Arc::new(Int64Array::from_iter(0..1024 * 1024));
-        let schema = Arc::new(Schema::new(vec![
-            Field::new("col0", arrow_schema::DataType::Int64, false),
-            Field::new("col1", arrow_schema::DataType::Int64, false),
-        ]));
-        let batch =
-            arrow_array::RecordBatch::try_new(schema.clone(), vec![array.clone(), array]).unwrap();
+        // Generate values that are well dispersed across a range approximating (0..256 * 1024)
+        let array = Arc::new(Int32Array::from_iter(
+            (0i32..1024 * 1024).map(|x| x.wrapping_mul(163019) % 262139),
+        ));
+        let schema = Arc::new(Schema::new(vec![Field::new(
+            "col0",
+            arrow_schema::DataType::Int32,
+            false,
+        )]));
+        let batch = arrow_array::RecordBatch::try_new(schema.clone(), vec![array]).unwrap();
+
+        let props = WriterProperties::builder()
+            .set_dictionary_page_size_limit(1024 * 1024)
+            .build();
+        let mut writer = ArrowWriter::try_new(Vec::new(), schema.clone(), Some(props)).unwrap();
+        writer.write(&batch).unwrap();
+        let data = Bytes::from(writer.into_inner().unwrap());
+
+        // println!("file length, dictionary: {}", data.len());
+
+        let mut metadata = ParquetMetaDataReader::new();
+        metadata.try_parse(&data).unwrap();
+        let metadata = metadata.finish().unwrap();
+        let full_dict_meta = metadata.row_group(0).column(0);
+        assert_eq!(get_dict_page_size(full_dict_meta, data.clone()), 1_048_576);
 
         let props = WriterProperties::builder()
             .set_dictionary_page_size_limit(1024 * 1024)
-            .set_column_dictionary_page_size_limit(ColumnPath::from("col1"), 1024 * 1024 * 4)
             .set_column_dictionary_fallback(
-                ColumnPath::from("col1"),
+                ColumnPath::from("col0"),
                 DictionaryFallback::OnUnfavorableCompression,
             )
             .build();
-        let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
+        let mut writer = ArrowWriter::try_new(Vec::new(), schema.clone(), Some(props)).unwrap();
         writer.write(&batch).unwrap();
         let data = Bytes::from(writer.into_inner().unwrap());
 
+        // println!("file length, fallback: {}", data.len());
+
         let mut metadata = ParquetMetaDataReader::new();
         metadata.try_parse(&data).unwrap();
         let metadata = metadata.finish().unwrap();
-        let col0_meta = metadata.row_group(0).column(0);
-        let col1_meta = metadata.row_group(0).column(1);
-
-        let get_dict_page_size = move |meta: &ColumnChunkMetaData| {
-            let mut reader =
-                SerializedPageReader::new(Arc::new(data.clone()), meta, 0, None).unwrap();
-            let page = reader.get_next_page().unwrap().unwrap();
-            match page {
-                Page::DictionaryPage { buf, .. } => buf.len(),
-                _ => panic!("expected DictionaryPage"),
-            }
-        };
-
-        assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
-        assert_eq!(get_dict_page_size(col1_meta), 8192);
+        let fallback_meta = metadata.row_group(0).column(0);
+        assert_eq!(get_dict_page_size(fallback_meta, data.clone()), 4096);
     }
 
     struct WriteBatchesShape {