fix not expr

Author: xudong963

datafusion/core/tests/parquet/row_group_pruning.rs

@@ -207,16 +207,6 @@ impl RowGroupPruningTest {
             self.expected_files_pruned_by_statistics,
             "mismatched files_ranges_pruned_statistics",
         );
-        assert_eq!(
-            output.row_groups_matched_bloom_filter(),
-            self.expected_row_group_matched_by_bloom_filter,
-            "mismatched row_groups_matched_bloom_filter",
-        );
-        assert_eq!(
-            output.row_groups_pruned_bloom_filter(),
-            self.expected_row_group_pruned_by_bloom_filter,
-            "mismatched row_groups_pruned_bloom_filter",
-        );
         assert_eq!(
             output.limit_pruned_row_groups(),
             self.expected_limit_pruned_row_groups,
@@ -1746,12 +1736,13 @@ async fn test_limit_pruning() -> datafusion_common::error::Result<()> {
     // So 3 row groups are effectively pruned due to limit pruning.
 
     let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, false)]));
-    let query = "SELECT c1 FROM t WHERE c1 > 0 LIMIT 2";
+    let query = "SELECT c1 FROM t WHERE c1 >= 0 LIMIT 2";
 
     let batches = vec![
-        make_i32_batch("c1", vec![1, 2])?, // RG0: Fully matched, 2 rows
-        make_i32_batch("c1", vec![3, 4])?, // RG1: Fully matched, 2 rows
-        make_i32_batch("c1", vec![5, 6])?, // RG2: Fully matched, 2 rows
+        make_i32_batch("c1", vec![0, -2])?,
+        make_i32_batch("c1", vec![0, 0])?, // RG0: Fully matched, 2 rows
+        make_i32_batch("c1", vec![0, 0])?, // RG1: Fully matched, 2 rows
+        make_i32_batch("c1", vec![0, 0])?, // RG2: Fully matched, 2 rows
         make_i32_batch("c1", vec![-1, 0])?, // RG3: Pruned by statistics, 0 rows
     ];
 
@@ -1761,11 +1752,9 @@ async fn test_limit_pruning() -> datafusion_common::error::Result<()> {
         .with_expected_errors(Some(0))
         .with_expected_rows(2)
         .with_pruned_files(Some(0))
-        .with_matched_by_bloom_filter(Some(0))
-        .with_pruned_by_bloom_filter(Some(0))
-        .with_matched_by_stats(Some(3)) // RG0, RG1, RG2 are matched by stats (c1 > 0)
-        .with_pruned_by_stats(Some(1)) // RG3 is pruned by stats (c1 = [-1, 0] does not satisfy c1 > 0)
-        .with_limit_pruned_row_groups(Some(2)) // RG1, RG2 are pruned by limit. (RG3 is already pruned by stats)
+        .with_matched_by_stats(Some(5)) // RG0, RG1, RG2 are matched by stats (c1 > 0)
+        .with_pruned_by_stats(Some(0)) // RG3 is pruned by stats (c1 = [-1, 0] does not satisfy c1 > 0)
+        .with_limit_pruned_row_groups(Some(4)) // RG1, RG2 are pruned by limit. (RG3 is already pruned by stats)
         .test_row_group_prune_with_custom_data(schema, batches, 2)
         .await;
 

datafusion/datasource-parquet/src/row_group_filter.rs

@@ -25,6 +25,7 @@ use datafusion_common::pruning::PruningStatistics;
 use datafusion_common::{Column, Result, ScalarValue};
 use datafusion_datasource::FileRange;
 use datafusion_physical_expr::expressions::NotExpr;
+use datafusion_physical_expr::PhysicalExprSimplifier;
 use datafusion_pruning::PruningPredicate;
 use parquet::arrow::arrow_reader::statistics::StatisticsConverter;
 use parquet::arrow::parquet_column;
@@ -205,6 +206,10 @@ impl RowGroupAccessPlanFilter {
                     // Use NotExpr to create the inverted predicate
                     let inverted_expr =
                         Arc::new(NotExpr::new(predicate.orig_expr().clone()));
+                    // Simplify the NOT expression (e.g., NOT(c1 = 0) -> c1 != 0)
+                    // before building the pruning predicate
+                    let mut simplifier = PhysicalExprSimplifier::new(arrow_schema);
+                    let inverted_expr = simplifier.simplify(inverted_expr).unwrap();
                     if let Ok(inverted_predicate) = PruningPredicate::try_new(
                         inverted_expr,
                         predicate.schema().clone(),

datafusion/physical-expr/src/simplifier/mod.rs

@@ -24,8 +24,9 @@ use datafusion_common::{
 };
 use std::sync::Arc;
 
-use crate::PhysicalExpr;
+use crate::{simplifier::not::simplify_not_expr_recursive, PhysicalExpr};
 
+pub mod not;
 pub mod unwrap_cast;
 
 /// Simplifies physical expressions by applying various optimizations
@@ -56,6 +57,11 @@ impl<'a> TreeNodeRewriter for PhysicalExprSimplifier<'a> {
     type Node = Arc<dyn PhysicalExpr>;
 
     fn f_up(&mut self, node: Self::Node) -> Result<Transformed<Self::Node>> {
+        // Apply NOT expression simplification first
+        let not_simplified = simplify_not_expr_recursive(node, self.schema)?;
+        let node = not_simplified.data;
+        let transformed = not_simplified.transformed;
+
         // Apply unwrap cast optimization
         #[cfg(test)]
         let original_type = node.data_type(self.schema).unwrap();
@@ -66,7 +72,12 @@ impl<'a> TreeNodeRewriter for PhysicalExprSimplifier<'a> {
             original_type,
             "Simplified expression should have the same data type as the original"
         );
-        Ok(unwrapped)
+        // Combine transformation results
+        let final_transformed = transformed || unwrapped.transformed;
+        Ok(Transformed::new_transformed(
+            unwrapped.data,
+            final_transformed,
+        ))
     }
 }