feat: derive a distinct count from primary key and unique constraints

datafusion/datasource/src/file_scan_config/mod.rs

@@ -39,7 +39,8 @@ use arrow::datatypes::{DataType, Schema, SchemaRef};
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::tree_node::TreeNodeRecursion;
 use datafusion_common::{
-    Constraints, Result, ScalarValue, Statistics, internal_datafusion_err, internal_err,
+    Constraint, Constraints, Result, ScalarValue, Statistics, internal_datafusion_err,
+    internal_err,
 };
 use datafusion_execution::{
     SendableRecordBatchStream, TaskContext, object_store::ObjectStoreUrl,
@@ -542,9 +543,10 @@ impl FileScanConfigBuilder {
         } = self;
 
         let constraints = constraints.unwrap_or_default();
-        let statistics = statistics.unwrap_or_else(|| {
+        let mut statistics = statistics.unwrap_or_else(|| {
             Statistics::new_unknown(file_source.table_schema().table_schema())
         });
+        add_key_distinct_counts(&constraints, &mut statistics);
         let file_compression_type =
             file_compression_type.unwrap_or(FileCompressionType::UNCOMPRESSED);
 
@@ -568,6 +570,32 @@ impl FileScanConfigBuilder {
     }
 }
 
+/// Records that a key column holds one distinct value per row, which no file format
+/// stores. Single-column keys only: a composite key says nothing about its columns.
+fn add_key_distinct_counts(constraints: &Constraints, statistics: &mut Statistics) {
+    let num_rows = statistics.num_rows;
+    for constraint in constraints.iter() {
+        let (Constraint::PrimaryKey(indices) | Constraint::Unique(indices)) = constraint;
+        let [index] = indices[..] else {
+            continue;
+        };
+        let Some(column) = statistics.column_statistics.get_mut(index) else {
+            continue;
+        };
+        if column.distinct_count != Precision::Absent {
+            continue;
+        }
+        // A NULL is not a distinct value. A primary key has none; a unique column may
+        // repeat them, so an unknown count leaves the result inexact.
+        let nulls = match (constraint, column.null_count) {
+            (Constraint::PrimaryKey(_), Precision::Absent) => Precision::Exact(0),
+            (_, Precision::Absent) => Precision::Inexact(0),
+            (_, nulls) => nulls,
+        };
+        column.distinct_count = num_rows.sub(&nulls);
+    }
+}
+
 impl From<FileScanConfig> for FileScanConfigBuilder {
     fn from(config: FileScanConfig) -> Self {
         Self {
@@ -2133,6 +2161,117 @@ mod tests {
     }
 
     // sets default for configs that play no role in projections
+    fn config_with_constraints(
+        table_schema: TableSchema,
+        statistics: Statistics,
+        constraints: Vec<Constraint>,
+        projection: Option<Vec<usize>>,
+    ) -> FileScanConfig {
+        FileScanConfigBuilder::new(
+            ObjectStoreUrl::parse("test:///").unwrap(),
+            Arc::new(MockSource::new(table_schema)),
+        )
+        .with_statistics(statistics)
+        .with_constraints(Constraints::new_unverified(constraints))
+        .with_projection_indices(projection)
+        .unwrap()
+        .build()
+    }
+
+    /// A key column has one distinct value per row, which no file format records.
+    #[test]
+    fn key_columns_report_a_distinct_count() {
+        let file_schema = Arc::new(Schema::new(vec![
+            Field::new("id", DataType::Int32, false),
+            Field::new("part", DataType::Int32, false),
+            Field::new("code", DataType::Int32, true),
+        ]));
+        let table_schema = TableSchema::builder(Arc::clone(&file_schema)).build();
+        let mut statistics = Statistics::new_unknown(&file_schema);
+        statistics.num_rows = Precision::Exact(100);
+        // A unique column may repeat NULL, which is not a distinct value.
+        statistics.column_statistics[2].null_count = Precision::Exact(10);
+
+        let stats = |constraints| {
+            config_with_constraints(
+                table_schema.clone(),
+                statistics.clone(),
+                constraints,
+                None,
+            )
+            .statistics()
+        };
+
+        // A primary key cannot be null, so the count is as exact as the row count.
+        let primary_key = stats(vec![Constraint::PrimaryKey(vec![0])]);
+        assert_eq!(
+            primary_key.column_statistics[0].distinct_count,
+            Precision::Exact(100)
+        );
+        assert_eq!(
+            primary_key.column_statistics[1].distinct_count,
+            Precision::Absent
+        );
+
+        // The nulls a unique column may repeat are known here, so this is exact too.
+        let unique = stats(vec![Constraint::Unique(vec![2])]);
+        assert_eq!(
+            unique.column_statistics[2].distinct_count,
+            Precision::Exact(90)
+        );
+
+        // With an unknown null count it is not exact.
+        let mut unknown_nulls = statistics.clone();
+        unknown_nulls.column_statistics[2].null_count = Precision::Absent;
+        let unique = config_with_constraints(
+            table_schema.clone(),
+            unknown_nulls,
+            vec![Constraint::Unique(vec![2])],
+            None,
+        )
+        .statistics();
+        assert_eq!(
+            unique.column_statistics[2].distinct_count,
+            Precision::Inexact(100)
+        );
+
+        // A composite key leaves its columns alone: only the combination is unique.
+        let composite = stats(vec![Constraint::PrimaryKey(vec![0, 1])]);
+        assert_eq!(
+            composite.column_statistics[0].distinct_count,
+            Precision::Absent
+        );
+        assert_eq!(
+            composite.column_statistics[1].distinct_count,
+            Precision::Absent
+        );
+    }
+
+    /// The count has to reach the plan, which reads statistics through the projection.
+    #[test]
+    fn a_projected_scan_keeps_the_key_distinct_count() {
+        let file_schema = Arc::new(Schema::new(vec![
+            Field::new("value", DataType::Int32, true),
+            Field::new("id", DataType::Int32, false),
+        ]));
+        let table_schema = TableSchema::builder(Arc::clone(&file_schema)).build();
+        let mut statistics = Statistics::new_unknown(&file_schema);
+        statistics.num_rows = Precision::Exact(100);
+
+        let config = config_with_constraints(
+            table_schema,
+            statistics,
+            vec![Constraint::PrimaryKey(vec![1])],
+            Some(vec![1]),
+        );
+
+        let projected = config.partition_statistics(None).unwrap();
+        assert_eq!(
+            projected.column_statistics[0].distinct_count,
+            Precision::Exact(100)
+        );
+    }
+
     fn config_for_projection(
         file_schema: SchemaRef,
         projection: Option<Vec<usize>>,

datafusion/physical-plan/src/filter.rs

@@ -63,7 +63,7 @@ use datafusion_execution::TaskContext;
 use datafusion_expr::Operator;
 use datafusion_physical_expr::equivalence::ProjectionMapping;
 use datafusion_physical_expr::expressions::{
-    BinaryExpr, Column, IsNotNullExpr, Literal, lit,
+    BinaryExpr, Column, InListExpr, IsNotNullExpr, Literal, lit,
 };
 use datafusion_physical_expr::intervals::utils::check_support;
 use datafusion_physical_expr::utils::{collect_columns, reassign_expr_columns};
@@ -343,6 +343,11 @@ impl FilterExec {
         let input_num_rows = input_stats.num_rows;
         let input_total_byte_size = input_stats.total_byte_size;
 
+        // A column holding each of its values once, as a primary key or unique
+        // constraint says, matches one row per value asked for. No selectivity
+        // expresses that.
+        let match_limit = unique_match_limit(predicate, &input_stats);
+
         let (selectivity, num_rows, column_statistics) = if is_infeasible {
             // Contradictory predicate: no rows survive. Row-bounded counts are
             // zero; value statistics are undefined on an empty column.
@@ -406,6 +411,10 @@ impl FilterExec {
             }
         };
 
+        let num_rows = match (match_limit, num_rows.get_value()) {
+            (Some(limit), Some(rows)) if *rows > limit => Precision::Inexact(limit),
+            _ => num_rows,
+        };
         let total_byte_size =
             scale_byte_size_at_rows(input_total_byte_size, selectivity, num_rows);
 
@@ -948,6 +957,76 @@ impl EmbeddedProjection for FilterExec {
     }
 }
 
+/// The most rows a filter can match, when it restricts a column holding each value
+/// once to a fixed set of values: one row per value.
+fn unique_match_limit(
+    predicate: &Arc<dyn PhysicalExpr>,
+    statistics: &Statistics,
+) -> Option<usize> {
+    let mut limit: Option<usize> = None;
+    for expr in split_conjunction(predicate) {
+        let Some((index, values)) = restricted_column(expr) else {
+            continue;
+        };
+        let holds_once = statistics
+            .column_statistics
+            .get(index)
+            .is_some_and(|column| holds_each_value_once(column, &statistics.num_rows));
+        if !holds_once {
+            continue;
+        }
+        limit = Some(limit.map_or(values, |limit: usize| limit.min(values)));
+    }
+    limit
+}
+
+/// The column an expression restricts to a fixed set of values, and how many values
+/// that is. NULL is never one of them: it matches nothing.
+fn restricted_column(expr: &Arc<dyn PhysicalExpr>) -> Option<(usize, usize)> {
+    if let Some(in_list) = expr.downcast_ref::<InListExpr>() {
+        if in_list.negated() {
+            return None;
+        }
+        let column = in_list.expr().downcast_ref::<Column>()?;
+        let mut values: Vec<&ScalarValue> = vec![];
+        for expr in in_list.list() {
+            let value = expr.downcast_ref::<Literal>()?.value();
+            if !value.is_null() && !values.contains(&value) {
+                values.push(value);
+            }
+        }
+        return Some((column.index(), values.len()));
+    }
+
+    let binary = expr.downcast_ref::<BinaryExpr>()?;
+    if *binary.op() != Operator::Eq {
+        return None;
+    }
+    let (column, literal) = match (
+        binary.left().downcast_ref::<Column>(),
+        binary.right().downcast_ref::<Column>(),
+    ) {
+        (Some(column), None) => (column, binary.right()),
+        (None, Some(column)) => (column, binary.left()),
+        _ => return None,
+    };
+    let value = literal.downcast_ref::<Literal>()?.value();
+    (!value.is_null()).then_some((column.index(), 1))
+}
+
+/// Whether the column has as many distinct values as it has non-null rows, so each
+/// value appears once.
+fn holds_each_value_once(column: &ColumnStatistics, num_rows: &Precision<usize>) -> bool {
+    let (Some(rows), Some(distinct), Some(nulls)) = (
+        num_rows.get_value(),
+        column.distinct_count.get_value(),
+        column.null_count.get_value(),
+    ) else {
+        return false;
+    };
+    distinct.saturating_add(*nulls) >= *rows
+}
+
 /// Collects column equality information from `col = literal` predicates in a
 /// conjunction.
 ///
@@ -1462,6 +1541,100 @@ mod tests {
         Ok(())
     }
 
+    /// An equality on a column that holds each value once matches one row at most,
+    /// including on a type interval analysis cannot read, where the default
+    /// selectivity would otherwise apply.
+    #[tokio::test]
+    async fn test_filter_statistics_equality_on_a_unique_column() -> Result<()> {
+        let schema = Schema::new(vec![Field::new("id", DataType::Utf8, true)]);
+        let unique = ColumnStatistics {
+            null_count: Precision::Exact(0),
+            distinct_count: Precision::Exact(100),
+            ..Default::default()
+        };
+        let rows = |column: ColumnStatistics| -> Result<Precision<usize>> {
+            let input = Arc::new(StatisticsExec::new(
+                Statistics {
+                    num_rows: Precision::Exact(100),
+                    total_byte_size: Precision::Exact(800),
+                    column_statistics: vec![column],
+                },
+                schema.clone(),
+            ));
+            let predicate =
+                binary(col("id", &schema)?, Operator::Eq, lit("seven"), &schema)?;
+            let filter: Arc<dyn ExecutionPlan> =
+                Arc::new(FilterExec::try_new(predicate, input)?);
+            Ok(StatisticsContext::new()
+                .compute(filter.as_ref(), &StatisticsArgs::new())?
+                .num_rows)
+        };
+
+        assert_eq!(rows(unique.clone())?, Precision::Inexact(1));
+
+        // The nulls a unique column may repeat do not make it hold a value twice.
+        assert_eq!(
+            rows(ColumnStatistics {
+                null_count: Precision::Exact(10),
+                distinct_count: Precision::Exact(90),
+                ..unique.clone()
+            })?,
+            Precision::Inexact(1)
+        );
+
+        // Without a distinct count, the default selectivity applies as before.
+        assert_eq!(
+            rows(ColumnStatistics {
+                distinct_count: Precision::Absent,
+                ..unique
+            })?,
+            Precision::Inexact(20)
+        );
+
+        Ok(())
+    }
+
+    /// Asking a unique column for three values matches three rows at most.
+    #[tokio::test]
+    async fn test_filter_statistics_in_list_on_a_unique_column() -> Result<()> {
+        use datafusion_physical_expr::expressions::in_list;
+
+        let schema = Schema::new(vec![Field::new("id", DataType::Utf8, true)]);
+        let rows = |list: Vec<&str>, negated: bool| -> Result<Precision<usize>> {
+            let input = Arc::new(StatisticsExec::new(
+                Statistics {
+                    num_rows: Precision::Exact(100),
+                    total_byte_size: Precision::Exact(800),
+                    column_statistics: vec![ColumnStatistics {
+                        null_count: Precision::Exact(0),
+                        distinct_count: Precision::Exact(100),
+                        ..Default::default()
+                    }],
+                },
+                schema.clone(),
+            ));
+            let predicate = in_list(
+                col("id", &schema)?,
+                list.into_iter().map(|value| lit(value) as _).collect(),
+                &negated,
+                &schema,
+            )?;
+            let filter: Arc<dyn ExecutionPlan> =
+                Arc::new(FilterExec::try_new(predicate, input)?);
+            Ok(StatisticsContext::new()
+                .compute(filter.as_ref(), &StatisticsArgs::new())?
+                .num_rows)
+        };
+
+        assert_eq!(rows(vec!["a", "b", "c"], false)?, Precision::Inexact(3));
+        // Repeats ask for the same row twice.
+        assert_eq!(rows(vec!["a", "b", "a"], false)?, Precision::Inexact(2));
+        // `NOT IN` selects nearly everything, so the default applies.
+        assert_eq!(rows(vec!["a", "b", "c"], true)?, Precision::Inexact(20));
+
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_filter_statistics_basic_expr() -> Result<()> {
         // Table: