Fix fully matched row groups with null counts

Author: xudong963

datafusion/datasource-parquet/src/row_group_filter.rs

@@ -24,8 +24,10 @@ use arrow::datatypes::Schema;
 use datafusion_common::pruning::PruningStatistics;
 use datafusion_common::{Column, Result, ScalarValue};
 use datafusion_datasource::FileRange;
-use datafusion_physical_expr::PhysicalExprSimplifier;
-use datafusion_physical_expr::expressions::NotExpr;
+use datafusion_expr::Operator;
+use datafusion_physical_expr::expressions::{BinaryExpr, IsNullExpr, NotExpr};
+use datafusion_physical_expr::utils::collect_columns;
+use datafusion_physical_expr::{PhysicalExpr, PhysicalExprSimplifier};
 use datafusion_pruning::PruningPredicate;
 use parquet::arrow::arrow_reader::statistics::StatisticsConverter;
 use parquet::basic::Type;
@@ -272,6 +274,7 @@ impl RowGroupAccessPlanFilter {
             parquet_schema,
             row_group_metadatas,
             arrow_schema,
+            missing_null_counts_as_zero: true,
         };
 
         // try to prune the row groups in a single call
@@ -327,10 +330,33 @@ impl RowGroupAccessPlanFilter {
             return;
         }
 
-        // Use NotExpr to create the inverted predicate
-        let inverted_expr = Arc::new(NotExpr::new(Arc::clone(predicate.orig_expr())));
+        let mut inverted_expr: Arc<dyn PhysicalExpr> =
+            Arc::new(NotExpr::new(Arc::clone(predicate.orig_expr())));
 
-        // Simplify the NOT expression (e.g., NOT(c1 = 0) -> c1 != 0)
+        // Rows where the predicate evaluates to NULL do not pass the filter.
+        // Include NULL checks in the inverted expression so a row group is only
+        // considered fully matched when every referenced column is known non-null.
+        // This is conservative for null-accepting predicates, but fully matched
+        // row groups must not have false positives.
+        let mut columns = collect_columns(predicate.orig_expr())
+            .into_iter()
+            .filter(|column| arrow_schema.field(column.index()).is_nullable())
+            .collect::<Vec<_>>();
+        columns.sort_by(|a, b| {
+            a.index()
+                .cmp(&b.index())
+                .then_with(|| a.name().cmp(b.name()))
+        });
+
+        for column in columns {
+            inverted_expr = Arc::new(BinaryExpr::new(
+                inverted_expr,
+                Operator::Or,
+                Arc::new(IsNullExpr::new(Arc::new(column))),
+            ));
+        }
+
+        // Simplify the inverted expression (e.g., NOT(c1 = 0) -> c1 != 0)
         // before building the pruning predicate
         let simplifier = PhysicalExprSimplifier::new(arrow_schema);
         let Ok(inverted_expr) = simplifier.simplify(inverted_expr) else {
@@ -350,6 +376,7 @@ impl RowGroupAccessPlanFilter {
                 .map(|&i| &groups[i])
                 .collect::<Vec<_>>(),
             arrow_schema,
+            missing_null_counts_as_zero: false,
         };
 
         let Ok(inverted_values) = inverted_predicate.prune(&inverted_pruning_stats)
@@ -582,6 +609,7 @@ struct RowGroupPruningStatistics<'a> {
     parquet_schema: &'a SchemaDescriptor,
     row_group_metadatas: Vec<&'a RowGroupMetaData>,
     arrow_schema: &'a Schema,
+    missing_null_counts_as_zero: bool,
 }
 
 impl<'a> RowGroupPruningStatistics<'a> {
@@ -598,7 +626,8 @@ impl<'a> RowGroupPruningStatistics<'a> {
             &column.name,
             self.arrow_schema,
             self.parquet_schema,
-        )?)
+        )?
+        .with_missing_null_counts_as_zero(self.missing_null_counts_as_zero))
     }
 }
 
@@ -767,6 +796,65 @@ mod tests {
         assert_pruned(row_groups, ExpectedPruning::Some(vec![1]))
     }
 
+    #[test]
+    fn row_group_fully_matched_requires_known_non_null_predicate_columns() {
+        use datafusion_expr::{col, lit};
+
+        let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, true)]));
+        let expr = logical2physical(&col("c1").gt(lit(15)), &schema);
+        let pruning_predicate = PruningPredicate::try_new(expr, schema.clone()).unwrap();
+
+        let field = PrimitiveTypeField::new("c1", PhysicalType::INT32);
+        let schema_descr = get_test_schema_descr(vec![field]);
+
+        // All three row groups have non-null values in the predicate range,
+        // so none are pruned. Only the second row group can be proven fully
+        // matched because it is the only one with a known zero null count.
+        let rg_with_null = get_row_group_meta_data(
+            &schema_descr,
+            vec![ParquetStatistics::int32(
+                Some(16),
+                Some(20),
+                None,
+                Some(1),
+                false,
+            )],
+        );
+        let rg_without_null = get_row_group_meta_data(
+            &schema_descr,
+            vec![ParquetStatistics::int32(
+                Some(16),
+                Some(20),
+                None,
+                Some(0),
+                false,
+            )],
+        );
+        let rg_unknown_null_count = get_row_group_meta_data(
+            &schema_descr,
+            vec![ParquetStatistics::int32(
+                Some(16),
+                Some(20),
+                None,
+                None,
+                false,
+            )],
+        );
+
+        let metrics = parquet_file_metrics();
+        let mut row_groups = RowGroupAccessPlanFilter::new(ParquetAccessPlan::new_all(3));
+        row_groups.prune_by_statistics(
+            &schema,
+            &schema_descr,
+            &[rg_with_null, rg_without_null, rg_unknown_null_count],
+            &pruning_predicate,
+            &metrics,
+        );
+
+        assert_eq!(row_groups.access_plan.row_group_indexes(), vec![0, 1, 2]);
+        assert_eq!(row_groups.is_fully_matched(), &vec![false, true, false]);
+    }
+
     #[test]
     fn row_group_pruning_predicate_missing_stats() {
         use datafusion_expr::{col, lit};