use page row limit fix ut

Signed-off-by: yangjiang <yangjiang@ebay.com>
apache · Nov 5, 2022 · 2dd7cdf · 2dd7cdf
1 parent 4753240
commit 2dd7cdf
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 7 deletions.
diff --git a/datafusion/core/src/datasource/file_format/parquet.rs b/datafusion/core/src/datasource/file_format/parquet.rs
@@ -538,9 +538,7 @@ pub(crate) mod test_util {
             // All batches write in to one file, each batch must have same schema.
             let mut output = NamedTempFile::new().expect("creating temp file");
             let mut builder = WriterProperties::builder();
-            // todo https://github.com/apache/arrow-rs/issues/2941 release change to row limit.
-            builder = builder.set_data_pagesize_limit(1);
-            builder = builder.set_write_batch_size(1);
+            builder = builder.set_data_page_row_count_limit(2);
             let proper = builder.build();
             let mut writer =
                 ArrowWriter::try_new(&mut output, batches[0].schema(), Some(proper))

diff --git a/datafusion/core/src/physical_plan/file_format/parquet.rs b/datafusion/core/src/physical_plan/file_format/parquet.rs
@@ -1722,14 +1722,12 @@ mod tests {
 
         let metrics = rt.parquet_exec.metrics().unwrap();
 
-        // todo fix this  https://github.com/apache/arrow-rs/issues/2941 release change to row limit.
         // assert the batches and some metrics
         let expected = vec![
-            "+-----+", "| int |", "+-----+", "|     |", "| 1   |", "| 2   |", "| 3   |",
-            "| 4   |", "| 5   |", "+-----+",
+            "+-----+", "| int |", "+-----+", "| 3   |", "| 4   |", "| 5   |", "+-----+",
         ];
         assert_batches_sorted_eq!(expected, &rt.batches.unwrap());
-        assert_eq!(get_value(&metrics, "page_index_rows_filtered"), 0);
+        assert_eq!(get_value(&metrics, "page_index_rows_filtered"), 3);
         assert!(
             get_value(&metrics, "page_index_eval_time") > 0,
             "no eval time in metrics: {:#?}",