fix: preserve missing Parquet null counts

datafusion/core/tests/parquet/dynamic_row_group_pruning.rs

@@ -30,16 +30,190 @@
 //! because batch-arrival timing affects how soon the TopK heap fills,
 //! and we don't want this test to become flaky.
 
+use std::io::Write;
 use std::sync::Arc;
 
 use arrow::array::{ArrayRef, Int64Array, RecordBatch, StringArray};
 use arrow_schema::{DataType, Field, Schema};
 
-use datafusion::prelude::SessionConfig;
+use datafusion::physical_plan::collect;
+use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext};
+use datafusion_common::ScalarValue;
+use parquet::arrow::ArrowWriter;
+use parquet::file::metadata::{ParquetMetaData, ParquetMetaDataWriter};
+use parquet::file::properties::{EnabledStatistics, WriterProperties};
+use parquet::file::statistics::Statistics;
+use parquet::file::writer::TrackedWrite;
+use tempfile::NamedTempFile;
 
 use crate::parquet::Unit::RowGroup;
+use crate::parquet::utils::MetricsFinder;
 use crate::parquet::{ContextWithParquet, Scenario};
 
+/// Keep min/max but omit the middle row group's null count. Negate the
+/// values for DESC so the first group always establishes the winning bound.
+fn file_with_missing_null_count(descending: bool, has_null: bool) -> NamedTempFile {
+    let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, true)]));
+    let props = WriterProperties::builder()
+        .set_statistics_enabled(EnabledStatistics::Chunk)
+        .build();
+    let mut bytes = Vec::new();
+    let mut writer =
+        ArrowWriter::try_new(&mut bytes, schema.clone(), Some(props)).unwrap();
+    for values in [
+        vec![Some(0), Some(1), Some(2)],
+        vec![
+            if has_null { None } else { Some(100) },
+            Some(100),
+            Some(101),
+        ],
+        vec![Some(200), Some(201), Some(202)],
+    ] {
+        let values: Int64Array = values
+            .into_iter()
+            .map(|v| v.map(|v| if descending { -v } else { v }))
+            .collect();
+        let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(values)]).unwrap();
+        writer.write(&batch).unwrap();
+        writer.flush().unwrap();
+    }
+    let metadata = writer.close().unwrap();
+    let mut groups = metadata.row_groups().to_vec();
+    if has_null {
+        let column = groups[1].column(0).clone();
+        let Statistics::Int64(stats) = column.statistics().unwrap() else {
+            panic!("expected Int64 statistics");
+        };
+        let stats = Statistics::int64(
+            stats.min_opt().copied(),
+            stats.max_opt().copied(),
+            None,
+            None,
+            false,
+        );
+        let column = column.into_builder().set_statistics(stats).build().unwrap();
+        groups[1] = groups[1]
+            .clone()
+            .into_builder()
+            .set_column_metadata(vec![column])
+            .build()
+            .unwrap();
+    }
+    let metadata = ParquetMetaData::new(metadata.file_metadata().clone(), groups);
+
+    // Replace the footer, leaving the encoded rows and their offsets intact.
+    let footer_len =
+        u32::from_le_bytes(bytes[bytes.len() - 8..bytes.len() - 4].try_into().unwrap());
+    bytes.truncate(bytes.len() - 8 - footer_len as usize);
+    let mut file = tempfile::Builder::new()
+        .suffix(".parquet")
+        .tempfile()
+        .unwrap();
+    let mut output = TrackedWrite::new(file.as_file_mut());
+    output.write_all(&bytes).unwrap();
+    ParquetMetaDataWriter::new_with_tracked(output, &metadata)
+        .finish()
+        .unwrap();
+    file
+}
+
+#[tokio::test]
+async fn missing_null_count_preserves_null_filter_and_count() {
+    let file = file_with_missing_null_count(false, true);
+    let ctx = SessionContext::new();
+    ctx.register_parquet(
+        "t",
+        file.path().to_str().unwrap(),
+        ParquetReadOptions::default(),
+    )
+    .await
+    .unwrap();
+    for (sql, expected) in [
+        ("SELECT COUNT(v) FROM t", 8),
+        ("SELECT COUNT(*) FROM t WHERE v IS NULL", 1),
+    ] {
+        let batches = ctx.sql(sql).await.unwrap().collect().await.unwrap();
+        assert_eq!(
+            ScalarValue::try_from_array(batches[0].column(0), 0).unwrap(),
+            ScalarValue::Int64(Some(expected)),
+            "{sql}",
+        );
+    }
+}
+
+#[tokio::test]
+async fn dynamic_rg_pruning_preserves_missing_null_count() {
+    for descending in [false, true] {
+        for has_null in [false, true] {
+            let file = file_with_missing_null_count(descending, has_null);
+            for nulls_first in [false, true] {
+                for pushdown in [false, true] {
+                    let mut config = SessionConfig::new()
+                        .with_target_partitions(1)
+                        .with_batch_size(1)
+                        .with_parquet_page_index_pruning(false);
+                    // Keep the live filter on the row-group path: file-level
+                    // statistics can otherwise prune the entire remaining file.
+                    config.options_mut().execution.collect_statistics = false;
+                    config.options_mut().optimizer.enable_sort_pushdown = false;
+                    config
+                        .options_mut()
+                        .optimizer
+                        .enable_topk_dynamic_filter_pushdown = pushdown;
+                    config.options_mut().execution.parquet.pushdown_filters = false;
+                    let ctx = SessionContext::new_with_config(config);
+                    ctx.register_parquet(
+                        "t",
+                        file.path().to_str().unwrap(),
+                        ParquetReadOptions::default(),
+                    )
+                    .await
+                    .unwrap();
+                    let order = if descending { "DESC" } else { "ASC" };
+                    let null_order = if nulls_first { "FIRST" } else { "LAST" };
+                    let sql = format!(
+                        "SELECT v FROM t ORDER BY v {order} NULLS {null_order} LIMIT 1"
+                    );
+                    let plan = ctx
+                        .sql(&sql)
+                        .await
+                        .unwrap()
+                        .create_physical_plan()
+                        .await
+                        .unwrap();
+                    let batches = collect(plan.clone(), ctx.task_ctx()).await.unwrap();
+                    assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 1);
+                    let expected = if has_null && nulls_first {
+                        None
+                    } else {
+                        Some(0)
+                    };
+                    assert_eq!(
+                        ScalarValue::try_from_array(batches[0].column(0), 0).unwrap(),
+                        ScalarValue::Int64(expected),
+                        "{sql}, has_null={has_null}, pushdown={pushdown}",
+                    );
+                    let metrics = MetricsFinder::find_metrics(plan.as_ref()).unwrap();
+                    let pruned = metrics
+                        .sum_by_name("row_groups_pruned_dynamic_filter")
+                        .unwrap()
+                        .as_usize();
+                    if pushdown {
+                        assert!(
+                            pruned > 0,
+                            "runtime pruning must run: {sql}, has_null={has_null}\n{metrics}\n{}",
+                            datafusion::physical_plan::displayable(plan.as_ref())
+                                .indent(true),
+                        );
+                    } else {
+                        assert_eq!(pruned, 0);
+                    }
+                }
+            }
+        }
+    }
+}
+
 /// Build five `RecordBatch`es whose `v` column ranges are disjoint:
 /// batch `i` carries `v` values `[i*100, (i+1)*100)`. When written with
 /// `max_row_group_row_count = 100` each batch lands in its own row group.

datafusion/datasource-parquet/src/metadata.rs

@@ -551,6 +551,10 @@ impl<'a> DFParquetMetadata<'a> {
                         file_metadata.schema_descr(),
                     ) {
                         Ok(stats_converter) => {
+                            // An omitted count must not become an exact zero in
+                            // file statistics used for pruning and aggregates.
+                            let stats_converter =
+                                stats_converter.with_missing_null_counts_as_zero(false);
                             let parquet_index = stats_converter.parquet_column_index();
                             if parquet_index.is_some_and(|index| {
                                 has_untrusted_min_max_order(
@@ -1366,6 +1370,43 @@ mod tests {
             ParquetMetaData::new(file_meta, row_groups)
         }
 
+        #[test]
+        fn test_statistics_preserve_missing_null_counts() {
+            let schema_descr = create_schema_descr(1);
+            let arrow_schema = create_arrow_schema(1);
+            for (null_counts, expected) in [
+                (vec![None], Precision::Absent),
+                (vec![Some(0), None], Precision::Inexact(0)),
+                (vec![Some(2), None], Precision::Inexact(2)),
+                (vec![Some(0), Some(0)], Precision::Exact(0)),
+            ] {
+                let row_groups = null_counts
+                    .into_iter()
+                    .map(|null_count| {
+                        create_row_group_with_stats(
+                            &schema_descr,
+                            vec![Some(ParquetStatistics::int32(
+                                Some(1),
+                                Some(10),
+                                None,
+                                null_count,
+                                false,
+                            ))],
+                            10,
+                        )
+                    })
+                    .collect();
+                let metadata =
+                    create_parquet_metadata(Arc::clone(&schema_descr), row_groups);
+                let statistics = DFParquetMetadata::statistics_from_parquet_metadata(
+                    &metadata,
+                    &arrow_schema,
+                )
+                .unwrap();
+                assert_eq!(statistics.column_statistics[0].null_count, expected);
+            }
+        }
+
         #[test]
         fn test_summarize_null_counts() {
             let schema_descr = create_schema_descr(1);

datafusion/datasource-parquet/src/push_decoder.rs

@@ -253,11 +253,6 @@ impl RowGroupPruner {
                 .map(Vec::as_slice),
             row_group_metadatas,
             arrow_schema: self.arrow_schema.as_ref(),
-            // Match the existing static row-group pruning behavior: when a
-            // statistic's null count is missing, treat it as zero. This is
-            // sound for runtime pruning because the predicate only needs to
-            // prove a row group *cannot* contain matching rows.
-            missing_null_counts_as_zero: true,
         };
 
         match pp.prune(&stats) {

datafusion/datasource-parquet/src/row_group_filter.rs

@@ -335,11 +335,6 @@ impl RowGroupAccessPlanFilter {
             column_orders,
             row_group_metadatas,
             arrow_schema,
-            // Preserve the existing row-group pruning behavior. This path only
-            // proves whether matching rows may exist, so it uses the
-            // StatisticsConverter default for older parquet-rs files where a
-            // missing null count can mean there are zero nulls.
-            missing_null_counts_as_zero: true,
         };
 
         // try to prune the row groups in a single call
@@ -407,11 +402,6 @@ impl RowGroupAccessPlanFilter {
                 .map(|&i| &groups[i])
                 .collect::<Vec<_>>(),
             arrow_schema,
-            // Fully matched row groups require a stronger proof: every row
-            // must pass the predicate. Missing null counts are unknown here;
-            // treating them as zero can incorrectly mark nullable row groups as
-            // fully matched and make limit pruning unsound.
-            missing_null_counts_as_zero: false,
         };
 
         let Ok(inverted_values) = inverted_predicate.prune(&inverted_pruning_stats)
@@ -495,7 +485,6 @@ pub(crate) struct RowGroupPruningStatistics<'a> {
     pub(crate) column_orders: Option<&'a [ColumnOrder]>,
     pub(crate) row_group_metadatas: Vec<&'a RowGroupMetaData>,
     pub(crate) arrow_schema: &'a Schema,
-    pub(crate) missing_null_counts_as_zero: bool,
 }
 
 impl<'a> RowGroupPruningStatistics<'a> {
@@ -510,7 +499,9 @@ impl<'a> RowGroupPruningStatistics<'a> {
             self.arrow_schema,
             self.parquet_schema,
         )?
-        .with_missing_null_counts_as_zero(self.missing_null_counts_as_zero))
+        // Missing counts cannot rule out nulls, either when pruning groups or
+        // when proving that every row matches the predicate.
+        .with_missing_null_counts_as_zero(false))
     }
 
     fn min_max_statistics_converter(
@@ -783,6 +774,48 @@ mod tests {
         assert_eq!(row_groups.is_fully_matched(), &vec![false, true, false]);
     }
 
+    #[test]
+    fn row_group_pruning_predicate_missing_null_count() {
+        let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, true)]));
+        let schema_descr = get_test_schema_descr(vec![PrimitiveTypeField::new(
+            "c1",
+            PhysicalType::INT32,
+        )]);
+        let groups = [None, Some(0), Some(1)].map(|null_count| {
+            get_row_group_meta_data(
+                &schema_descr,
+                vec![ParquetStatistics::int32(
+                    Some(100),
+                    Some(101),
+                    None,
+                    null_count,
+                    false,
+                )],
+            )
+        });
+
+        for (expr, expected) in [
+            (col("c1").is_null(), vec![0, 2]),
+            (col("c1").is_null().or(col("c1").lt(lit(0))), vec![0, 2]),
+            (col("c1").lt(lit(0)), vec![]),
+        ] {
+            let predicate = build_test_pruning_predicate(
+                logical2physical(&expr, &schema),
+                Arc::clone(&schema),
+            );
+            let mut filter =
+                RowGroupAccessPlanFilter::new(ParquetAccessPlan::new_all(groups.len()));
+            filter.prune_by_statistics(
+                &schema,
+                &schema_descr,
+                &groups,
+                &predicate,
+                &parquet_file_metrics(),
+            );
+            assert_eq!(filter.build().row_group_indexes(), expected, "{expr}");
+        }
+    }
+
     #[test]
     fn prune_by_limit_preserves_row_selection() {
         let field = PrimitiveTypeField::new("c1", PhysicalType::INT32);

datafusion/datasource-parquet/src/statistics_order_tests.rs

@@ -520,7 +520,6 @@ fn byte_array_order_guard_follows_parquet_type_not_arrow_representation() {
             column_orders: metadata.column_orders().map(Vec::as_slice),
             row_group_metadatas: file.metadata.row_groups().iter().collect(),
             arrow_schema: &schema,
-            missing_null_counts_as_zero: true,
         };
         for values in [stats.min_values(&column), stats.max_values(&column)] {
             let values = values.unwrap();