Propagate table constraints through physical plans to optimize sort operations

datafusion/common/src/functional_dependencies.rs

@@ -60,6 +60,41 @@ impl Constraints {
     pub fn is_empty(&self) -> bool {
         self.inner.is_empty()
     }
+
+    /// Projects constraints using the given projection indices.
+    /// Returns None if any of the constraint columns are not included in the projection.
+    pub fn project(&self, proj_indices: &[usize]) -> Option<Self> {
+        let projected = self
+            .inner
+            .iter()
+            .filter_map(|constraint| {
+                match constraint {
+                    Constraint::PrimaryKey(indices) => {
+                        let new_indices =
+                            update_elements_with_matching_indices(indices, proj_indices);
+                        // Only keep constraint if all columns are preserved
+                        (new_indices.len() == indices.len())
+                            .then_some(Constraint::PrimaryKey(new_indices))
+                    }
+                    Constraint::Unique(indices) => {
+                        let new_indices =
+                            update_elements_with_matching_indices(indices, proj_indices);
+                        // Only keep constraint if all columns are preserved
+                        (new_indices.len() == indices.len())
+                            .then_some(Constraint::Unique(new_indices))
+                    }
+                }
+            })
+            .collect::<Vec<_>>();
+
+        (!projected.is_empty()).then_some(Constraints::new_unverified(projected))
+    }
+}
+
+impl Default for Constraints {
+    fn default() -> Self {
+        Constraints::empty()
+    }
 }
 
 impl IntoIterator for Constraints {
@@ -73,13 +108,13 @@ impl IntoIterator for Constraints {
 
 impl Display for Constraints {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        let pk: Vec<String> = self.inner.iter().map(|c| format!("{:?}", c)).collect();
+        let pk = self
+            .inner
+            .iter()
+            .map(|c| format!("{:?}", c))
+            .collect::<Vec<_>>();
         let pk = pk.join(", ");
-        if !pk.is_empty() {
-            write!(f, " constraints=[{pk}]")
-        } else {
-            write!(f, "")
-        }
+        write!(f, "constraints=[{pk}]")
     }
 }
 
@@ -599,6 +634,24 @@ mod tests {
         assert_eq!(iter.next(), None);
     }
 
+    #[test]
+    fn test_project_constraints() {
+        let constraints = Constraints::new_unverified(vec![
+            Constraint::PrimaryKey(vec![1, 2]),
+            Constraint::Unique(vec![0, 3]),
+        ]);
+
+        // Project keeping columns 1,2,3
+        let projected = constraints.project(&[1, 2, 3]).unwrap();
+        assert_eq!(
+            projected,
+            Constraints::new_unverified(vec![Constraint::PrimaryKey(vec![0, 1])])
+        );
+
+        // Project keeping only column 0 - should return None as no constraints are preserved
+        assert!(constraints.project(&[0]).is_none());
+    }
+
     #[test]
     fn test_get_updated_id_keys() {
         let fund_dependencies =

datafusion/core/src/datasource/listing/table.rs

@@ -935,6 +935,7 @@ impl TableProvider for ListingTable {
                 session_state,
                 FileScanConfig::new(object_store_url, Arc::clone(&self.file_schema))
                     .with_file_groups(partitioned_file_lists)
+                    .with_constraints(self.constraints.clone())
                     .with_statistics(statistics)
                     .with_projection(projection.cloned())
                     .with_limit(limit)

datafusion/core/src/datasource/memory.rs

@@ -132,6 +132,7 @@ impl MemTable {
         state: &SessionState,
     ) -> Result<Self> {
         let schema = t.schema();
+        let constraints = t.constraints();
         let exec = t.scan(state, None, &[], None).await?;
         let partition_count = exec.output_partitioning().partition_count();
 
@@ -162,7 +163,10 @@ impl MemTable {
             }
         }
 
-        let exec = MemoryExec::try_new(&data, Arc::clone(&schema), None)?;
+        let mut exec = MemoryExec::try_new(&data, Arc::clone(&schema), None)?;
+        if let Some(cons) = constraints {
+            exec = exec.with_constraints(cons.clone());
+        }
 
         if let Some(num_partitions) = output_partitions {
             let exec = RepartitionExec::try_new(

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

@@ -35,7 +35,7 @@ use arrow::buffer::Buffer;
 use arrow_ipc::reader::FileDecoder;
 use arrow_schema::SchemaRef;
 use datafusion_common::config::ConfigOptions;
-use datafusion_common::Statistics;
+use datafusion_common::{Constraints, Statistics};
 use datafusion_execution::TaskContext;
 use datafusion_physical_expr::{EquivalenceProperties, LexOrdering};
 use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
@@ -60,11 +60,16 @@ pub struct ArrowExec {
 impl ArrowExec {
     /// Create a new Arrow reader execution plan provided base configurations
     pub fn new(base_config: FileScanConfig) -> Self {
-        let (projected_schema, projected_statistics, projected_output_ordering) =
-            base_config.project();
+        let (
+            projected_schema,
+            projected_constraints,
+            projected_statistics,
+            projected_output_ordering,
+        ) = base_config.project();
         let cache = Self::compute_properties(
             Arc::clone(&projected_schema),
             &projected_output_ordering,
+            projected_constraints,
             &base_config,
         );
         Self {
@@ -88,12 +93,14 @@ impl ArrowExec {
     /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc.
     fn compute_properties(
         schema: SchemaRef,
-        projected_output_ordering: &[LexOrdering],
+        output_ordering: &[LexOrdering],
+        constraints: Constraints,
         file_scan_config: &FileScanConfig,
     ) -> PlanProperties {
         // Equivalence Properties
         let eq_properties =
-            EquivalenceProperties::new_with_orderings(schema, projected_output_ordering);
+            EquivalenceProperties::new_with_orderings(schema, output_ordering)
+                .with_constraints(constraints);
 
         PlanProperties::new(
             eq_properties,

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

@@ -29,6 +29,7 @@ use crate::physical_plan::{
 };
 
 use arrow::datatypes::SchemaRef;
+use datafusion_common::Constraints;
 use datafusion_execution::TaskContext;
 use datafusion_physical_expr::{EquivalenceProperties, LexOrdering};
 use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
@@ -48,11 +49,16 @@ pub struct AvroExec {
 impl AvroExec {
     /// Create a new Avro reader execution plan provided base configurations
     pub fn new(base_config: FileScanConfig) -> Self {
-        let (projected_schema, projected_statistics, projected_output_ordering) =
-            base_config.project();
+        let (
+            projected_schema,
+            projected_constraints,
+            projected_statistics,
+            projected_output_ordering,
+        ) = base_config.project();
         let cache = Self::compute_properties(
             Arc::clone(&projected_schema),
             &projected_output_ordering,
+            projected_constraints,
             &base_config,
         );
         Self {
@@ -73,10 +79,12 @@ impl AvroExec {
     fn compute_properties(
         schema: SchemaRef,
         orderings: &[LexOrdering],
+        constraints: Constraints,
         file_scan_config: &FileScanConfig,
     ) -> PlanProperties {
         // Equivalence Properties
-        let eq_properties = EquivalenceProperties::new_with_orderings(schema, orderings);
+        let eq_properties = EquivalenceProperties::new_with_orderings(schema, orderings)
+            .with_constraints(constraints);
         let n_partitions = file_scan_config.file_groups.len();
 
         PlanProperties::new(

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

@@ -40,6 +40,7 @@ use crate::physical_plan::{
 use arrow::csv;
 use arrow::datatypes::SchemaRef;
 use datafusion_common::config::ConfigOptions;
+use datafusion_common::Constraints;
 use datafusion_execution::TaskContext;
 use datafusion_physical_expr::{EquivalenceProperties, LexOrdering};
 
@@ -209,11 +210,16 @@ impl CsvExecBuilder {
             newlines_in_values,
         } = self;
 
-        let (projected_schema, projected_statistics, projected_output_ordering) =
-            base_config.project();
+        let (
+            projected_schema,
+            projected_constraints,
+            projected_statistics,
+            projected_output_ordering,
+        ) = base_config.project();
         let cache = CsvExec::compute_properties(
             projected_schema,
             &projected_output_ordering,
+            projected_constraints,
             &base_config,
         );
 
@@ -320,10 +326,12 @@ impl CsvExec {
     fn compute_properties(
         schema: SchemaRef,
         orderings: &[LexOrdering],
+        constraints: Constraints,
         file_scan_config: &FileScanConfig,
     ) -> PlanProperties {
         // Equivalence Properties
-        let eq_properties = EquivalenceProperties::new_with_orderings(schema, orderings);
+        let eq_properties = EquivalenceProperties::new_with_orderings(schema, orderings)
+            .with_constraints(constraints);
 
         PlanProperties::new(
             eq_properties,

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

@@ -33,7 +33,9 @@ use arrow::datatypes::{ArrowNativeType, UInt16Type};
 use arrow_array::{ArrayRef, DictionaryArray, RecordBatch, RecordBatchOptions};
 use arrow_schema::{DataType, Field, Schema, SchemaRef};
 use datafusion_common::stats::Precision;
-use datafusion_common::{exec_err, ColumnStatistics, DataFusionError, Statistics};
+use datafusion_common::{
+    exec_err, ColumnStatistics, Constraints, DataFusionError, Statistics,
+};
 use datafusion_physical_expr::LexOrdering;
 
 use log::warn;
@@ -114,6 +116,8 @@ pub struct FileScanConfig {
     /// concurrently, however files *within* a partition will be read
     /// sequentially, one after the next.
     pub file_groups: Vec<Vec<PartitionedFile>>,
+    /// Table constraints
+    pub constraints: Constraints,
     /// Estimated overall statistics of the files, taking `filters` into account.
     /// Defaults to [`Statistics::new_unknown`].
     pub statistics: Statistics,
@@ -146,6 +150,7 @@ impl FileScanConfig {
             object_store_url,
             file_schema,
             file_groups: vec![],
+            constraints: Constraints::empty(),
             statistics,
             projection: None,
             limit: None,
@@ -154,6 +159,12 @@ impl FileScanConfig {
         }
     }
 
+    /// Set the table constraints of the files
+    pub fn with_constraints(mut self, constraints: Constraints) -> Self {
+        self.constraints = constraints;
+        self
+    }
+
     /// Set the statistics of the files
     pub fn with_statistics(mut self, statistics: Statistics) -> Self {
         self.statistics = statistics;
@@ -210,30 +221,31 @@ impl FileScanConfig {
         self
     }
 
-    /// Project the schema and the statistics on the given column indices
-    pub fn project(&self) -> (SchemaRef, Statistics, Vec<LexOrdering>) {
+    /// Project the schema, constraints, and the statistics on the given column indices
+    pub fn project(&self) -> (SchemaRef, Constraints, Statistics, Vec<LexOrdering>) {
         if self.projection.is_none() && self.table_partition_cols.is_empty() {
             return (
                 Arc::clone(&self.file_schema),
+                self.constraints.clone(),
                 self.statistics.clone(),
                 self.output_ordering.clone(),
             );
         }
 
-        let proj_iter: Box<dyn Iterator<Item = usize>> = match &self.projection {
-            Some(proj) => Box::new(proj.iter().copied()),
-            None => Box::new(
-                0..(self.file_schema.fields().len() + self.table_partition_cols.len()),
-            ),
+        let proj_indices = if let Some(proj) = &self.projection {
+            proj
+        } else {
+            let len = self.file_schema.fields().len() + self.table_partition_cols.len();
+            &(0..len).collect::<Vec<_>>()
         };
 
         let mut table_fields = vec![];
         let mut table_cols_stats = vec![];
-        for idx in proj_iter {
-            if idx < self.file_schema.fields().len() {
-                let field = self.file_schema.field(idx);
+        for idx in proj_indices {
+            if *idx < self.file_schema.fields().len() {
+                let field = self.file_schema.field(*idx);
                 table_fields.push(field.clone());
-                table_cols_stats.push(self.statistics.column_statistics[idx].clone())
+                table_cols_stats.push(self.statistics.column_statistics[*idx].clone())
             } else {
                 let partition_idx = idx - self.file_schema.fields().len();
                 table_fields.push(self.table_partition_cols[partition_idx].to_owned());
@@ -254,10 +266,20 @@ impl FileScanConfig {
             self.file_schema.metadata().clone(),
         ));
 
+        let projected_constraints = self
+            .constraints
+            .project(proj_indices)
+            .unwrap_or_else(Constraints::empty);
+
         let projected_output_ordering =
             get_projected_output_ordering(self, &projected_schema);
 
-        (projected_schema, table_stats, projected_output_ordering)
+        (
+            projected_schema,
+            projected_constraints,
+            table_stats,
+            projected_output_ordering,
+        )
     }
 
     #[cfg_attr(not(feature = "avro"), allow(unused))] // Only used by avro
@@ -635,7 +657,7 @@ mod tests {
             )]),
         );
 
-        let (proj_schema, proj_statistics, _) = conf.project();
+        let (proj_schema, _, proj_statistics, _) = conf.project();
         assert_eq!(proj_schema.fields().len(), file_schema.fields().len() + 1);
         assert_eq!(
             proj_schema.field(file_schema.fields().len()).name(),
@@ -675,7 +697,7 @@ mod tests {
         );
 
         // verify the proj_schema includes the last column and exactly the same the field it is defined
-        let (proj_schema, _proj_statistics, _) = conf.project();
+        let (proj_schema, _, _, _) = conf.project();
         assert_eq!(proj_schema.fields().len(), file_schema.fields().len() + 1);
         assert_eq!(
             *proj_schema.field(file_schema.fields().len()),
@@ -708,7 +730,7 @@ mod tests {
             )]),
         );
 
-        let (proj_schema, proj_statistics, _) = conf.project();
+        let (proj_schema, _, proj_statistics, _) = conf.project();
         assert_eq!(
             columns(&proj_schema),
             vec!["date".to_owned(), "c1".to_owned()]