better test

Author: adriangb

datafusion/core/src/datasource/physical_plan/parquet.rs

@@ -100,6 +100,7 @@ mod tests {
         predicate: Option<Expr>,
         pushdown_predicate: bool,
         page_index_predicate: bool,
+        bloom_filters: bool,
     }
 
     impl RoundTrip {
@@ -132,6 +133,11 @@ mod tests {
             self
         }
 
+        fn with_bloom_filters(mut self) -> Self {
+            self.bloom_filters = true;
+            self
+        }
+
         /// run the test, returning only the resulting RecordBatches
         async fn round_trip_to_batches(
             self,
@@ -156,10 +162,20 @@ mod tests {
                 source = source
                     .with_pushdown_filters(true)
                     .with_reorder_filters(true);
+            } else {
+                source = source.with_pushdown_filters(false);
             }
 
             if self.page_index_predicate {
                 source = source.with_enable_page_index(true);
+            } else {
+                source = source.with_enable_page_index(false);
+            }
+
+            if self.bloom_filters {
+                source = source.with_bloom_filter_on_read(true);
+            } else {
+                source = source.with_bloom_filter_on_read(false);
             }
 
             source.with_schema(Arc::clone(&file_schema))
@@ -859,23 +875,30 @@ mod tests {
 
         // Predicate should prune all row groups
         let filter = col("c1").eq(lit(ScalarValue::Utf8(Some("aaa".to_string()))));
-        let read = RoundTrip::new()
+        let rt = RoundTrip::new()
             .with_predicate(filter)
             .with_schema(schema.clone())
-            .round_trip_to_batches(vec![batch.clone()])
-            .await
-            .unwrap();
-        assert_eq!(read.len(), 0);
+            .round_trip(vec![batch.clone()])
+            .await;
+        // There should be no predicate evaluation errors
+        let metrics = rt.parquet_exec.metrics().unwrap();
+        assert_eq!(get_value(&metrics, "predicate_evaluation_errors"), 0);
+        assert_eq!(get_value(&metrics, "pushdown_rows_matched"), 0);
+        assert_eq!(rt.batches.unwrap().len(), 0);
 
         // Predicate should prune no row groups
         let filter = col("c1").eq(lit(ScalarValue::Utf8(Some("foo".to_string()))));
-        let read = RoundTrip::new()
+        let rt = RoundTrip::new()
             .with_predicate(filter)
             .with_schema(schema)
-            .round_trip_to_batches(vec![batch])
-            .await
-            .unwrap();
-        assert_eq!(read.len(), 1);
+            .round_trip(vec![batch])
+            .await;
+        // There should be no predicate evaluation errors
+        let metrics = rt.parquet_exec.metrics().unwrap();
+        assert_eq!(get_value(&metrics, "predicate_evaluation_errors"), 0);
+        assert_eq!(get_value(&metrics, "pushdown_rows_matched"), 0);
+        let read = rt.batches.unwrap().iter().map(|b| b.num_rows()).sum::<usize>();
+        assert_eq!(read, 2, "Expected 2 rows to match the predicate");
     }
 
     #[tokio::test]
@@ -889,25 +912,28 @@ mod tests {
 
         // Predicate should prune all row groups
         let filter = col("c1").eq(lit(ScalarValue::UInt64(Some(5))));
-        let read = RoundTrip::new()
+        let rt = RoundTrip::new()
             .with_predicate(filter)
             .with_schema(schema.clone())
-            .round_trip_to_batches(vec![batch.clone()])
-            .await
-            .unwrap();
-        assert_eq!(read.len(), 0);
+            .round_trip(vec![batch.clone()])
+            .await;
+        // There should be no predicate evaluation errors
+        let metrics = rt.parquet_exec.metrics().unwrap();
+        assert_eq!(get_value(&metrics, "predicate_evaluation_errors"), 0);
+        assert_eq!(rt.batches.unwrap().len(), 0);
 
-        // TODO: this is failing on main, and has been for a long time!
-        // See https://github.com/apache/datafusion/pull/16086#discussion_r2094817127
-        // // Predicate should prune no row groups
-        // let filter = col("c1").eq(lit(ScalarValue::UInt64(Some(1))));
-        // let read = RoundTrip::new()
-        //     .with_predicate(filter)
-        //     .with_schema(schema)
-        //     .round_trip_to_batches(vec![batch])
-        //     .await
-        //     .unwrap();
-        // assert_eq!(read.len(), 1);
+        // Predicate should prune no row groups
+        let filter = col("c1").eq(lit(ScalarValue::UInt64(Some(1))));
+        let rt = RoundTrip::new()
+            .with_predicate(filter)
+            .with_schema(schema)
+            .round_trip(vec![batch])
+            .await;
+        // There should be no predicate evaluation errors
+        let metrics = rt.parquet_exec.metrics().unwrap();
+        assert_eq!(get_value(&metrics, "predicate_evaluation_errors"), 0);
+        let read = rt.batches.unwrap().iter().map(|b| b.num_rows()).sum::<usize>();
+        assert_eq!(read, 2, "Expected 2 rows to match the predicate");
     }
 
     #[tokio::test]
@@ -1692,6 +1718,7 @@ mod tests {
         let rt = RoundTrip::new()
             .with_predicate(filter.clone())
             .with_pushdown_predicate()
+            .with_bloom_filters()
             .round_trip(vec![batch1])
             .await;