Pass constraints to MemoryExec

datafusion-cli/src/functions.rs

@@ -244,6 +244,7 @@ impl TableProvider for ParquetMetadataTable {
         Ok(Arc::new(MemoryExec::try_new(
             &[vec![self.batch.clone()]],
             TableProvider::schema(self),
+            None,
             projection.cloned(),
         )?))
     }

datafusion-examples/examples/remote_catalog.rs

@@ -339,6 +339,7 @@ impl TableProvider for RemoteTable {
         Ok(Arc::new(MemoryExec::try_new(
             &[batches],
             self.schema.clone(),
+            None,
             projection.cloned(),
         )?))
     }

datafusion-examples/examples/simple_udtf.rs

@@ -123,6 +123,7 @@ impl TableProvider for LocalCsvTable {
         Ok(Arc::new(MemoryExec::try_new(
             &[batches],
             TableProvider::schema(self),
+            None,
             projection.cloned(),
         )?))
     }

datafusion/core/benches/physical_plan.rs

@@ -59,6 +59,7 @@ fn sort_preserving_merge_operator(
         &batches.into_iter().map(|rb| vec![rb]).collect::<Vec<_>>(),
         schema,
         None,
+        None,
     )
     .unwrap();
     let merge = Arc::new(SortPreservingMergeExec::new(sort, Arc::new(exec)));

datafusion/core/benches/sort.rs

@@ -167,7 +167,7 @@ impl BenchCase {
         let schema = partitions[0][0].schema();
         let sort = make_sort_exprs(schema.as_ref());
 
-        let exec = MemoryExec::try_new(partitions, schema, None).unwrap();
+        let exec = MemoryExec::try_new(partitions, schema, None, None).unwrap();
         let plan = Arc::new(SortPreservingMergeExec::new(sort, Arc::new(exec)));
 
         Self {
@@ -186,7 +186,7 @@ impl BenchCase {
         let schema = partitions[0][0].schema();
         let sort = make_sort_exprs(schema.as_ref());
 
-        let exec = MemoryExec::try_new(partitions, schema, None).unwrap();
+        let exec = MemoryExec::try_new(partitions, schema, None, None).unwrap();
         let exec =
             SortExec::new(sort.clone(), Arc::new(exec)).with_preserve_partitioning(true);
         let plan = Arc::new(SortPreservingMergeExec::new(sort, Arc::new(exec)));
@@ -208,7 +208,7 @@ impl BenchCase {
         let schema = partitions[0][0].schema();
         let sort = make_sort_exprs(schema.as_ref());
 
-        let exec = MemoryExec::try_new(partitions, schema, None).unwrap();
+        let exec = MemoryExec::try_new(partitions, schema, None, None).unwrap();
         let exec = Arc::new(CoalescePartitionsExec::new(Arc::new(exec)));
         let plan = Arc::new(SortExec::new(sort, exec));
 
@@ -229,7 +229,7 @@ impl BenchCase {
         let schema = partitions[0][0].schema();
         let sort = make_sort_exprs(schema.as_ref());
 
-        let exec = MemoryExec::try_new(partitions, schema, None).unwrap();
+        let exec = MemoryExec::try_new(partitions, schema, None, None).unwrap();
         let exec = SortExec::new(sort, Arc::new(exec)).with_preserve_partitioning(true);
         let plan = Arc::new(CoalescePartitionsExec::new(Arc::new(exec)));
 

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,8 @@ impl MemTable {
             }
         }
 
-        let exec = MemoryExec::try_new(&data, Arc::clone(&schema), None)?;
+        let exec =
+            MemoryExec::try_new(&data, Arc::clone(&schema), constraints.cloned(), None)?;
 
         if let Some(num_partitions) = output_partitions {
             let exec = RepartitionExec::try_new(
@@ -220,8 +222,12 @@ impl TableProvider for MemTable {
             partitions.push(inner_vec.clone())
         }
 
-        let mut exec =
-            MemoryExec::try_new(&partitions, self.schema(), projection.cloned())?;
+        let mut exec = MemoryExec::try_new(
+            &partitions,
+            self.schema(),
+            self.constraints().cloned(),
+            projection.cloned(),
+        )?;
 
         let show_sizes = state.config_options().explain.show_sizes;
         exec = exec.with_show_sizes(show_sizes);

datafusion/core/src/physical_optimizer/projection_pushdown.rs

@@ -216,6 +216,7 @@ fn try_swapping_with_memory(
             MemoryExec::try_new(
                 memory.partitions(),
                 memory.original_schema(),
+                Some(memory.constraints().to_owned()),
                 Some(new_projections),
             )
             .map(|e| Arc::new(e) as _)
@@ -1770,7 +1771,7 @@ mod tests {
             Field::new("e", DataType::Int32, true),
         ]));
 
-        Arc::new(MemoryExec::try_new(&[], schema, Some(vec![2, 0, 3, 4])).unwrap())
+        Arc::new(MemoryExec::try_new(&[], schema, None, Some(vec![2, 0, 3, 4])).unwrap())
     }
 
     #[test]

datafusion/core/src/physical_optimizer/test_utils.rs

@@ -216,7 +216,7 @@ pub fn coalesce_partitions_exec(input: Arc<dyn ExecutionPlan>) -> Arc<dyn Execut
 }
 
 pub(crate) fn memory_exec(schema: &SchemaRef) -> Arc<dyn ExecutionPlan> {
-    Arc::new(MemoryExec::try_new(&[vec![]], schema.clone(), None).unwrap())
+    Arc::new(MemoryExec::try_new(&[vec![]], schema.clone(), None, None).unwrap())
 }
 
 pub fn hash_join_exec(

datafusion/core/src/physical_planner.rs

@@ -1940,7 +1940,7 @@ impl DefaultPhysicalPlanner {
         let schema = record_batch.schema();
         let partitions = vec![vec![record_batch]];
         let projection = None;
-        let mem_exec = MemoryExec::try_new(&partitions, schema, projection)?;
+        let mem_exec = MemoryExec::try_new(&partitions, schema, None, projection)?;
         Ok(Arc::new(mem_exec))
     }
 

datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs

@@ -302,11 +302,12 @@ async fn run_aggregate_test(input1: Vec<RecordBatch>, group_by_columns: Vec<&str
 
     let concat_input_record = concat_batches(&schema, &input1).unwrap();
     let usual_source = Arc::new(
-        MemoryExec::try_new(&[vec![concat_input_record]], schema.clone(), None).unwrap(),
+        MemoryExec::try_new(&[vec![concat_input_record]], schema.clone(), None, None)
+            .unwrap(),
     );
 
     let running_source = Arc::new(
-        MemoryExec::try_new(&[input1.clone()], schema.clone(), None)
+        MemoryExec::try_new(&[input1.clone()], schema.clone(), None, None)
             .unwrap()
             .try_with_sort_information(vec![sort_keys])
             .unwrap(),

datafusion/core/tests/fuzz_cases/join_fuzz.rs

@@ -404,10 +404,12 @@ impl JoinFuzzTestCase {
     fn left_right(&self) -> (Arc<MemoryExec>, Arc<MemoryExec>) {
         let schema1 = self.input1[0].schema();
         let schema2 = self.input2[0].schema();
-        let left =
-            Arc::new(MemoryExec::try_new(&[self.input1.clone()], schema1, None).unwrap());
-        let right =
-            Arc::new(MemoryExec::try_new(&[self.input2.clone()], schema2, None).unwrap());
+        let left = Arc::new(
+            MemoryExec::try_new(&[self.input1.clone()], schema1, None, None).unwrap(),
+        );
+        let right = Arc::new(
+            MemoryExec::try_new(&[self.input2.clone()], schema2, None, None).unwrap(),
+        );
         (left, right)
     }
 

datafusion/core/tests/fuzz_cases/merge_fuzz.rs

@@ -116,7 +116,7 @@ async fn run_merge_test(input: Vec<Vec<RecordBatch>>) {
             },
         }]);
 
-        let exec = MemoryExec::try_new(&input, schema, None).unwrap();
+        let exec = MemoryExec::try_new(&input, schema, None, None).unwrap();
         let merge = Arc::new(SortPreservingMergeExec::new(sort, Arc::new(exec)));
 
         let session_config = SessionConfig::new().with_batch_size(batch_size);

datafusion/core/tests/fuzz_cases/sort_fuzz.rs

@@ -123,7 +123,7 @@ impl SortTest {
             },
         }]);
 
-        let exec = MemoryExec::try_new(&input, schema, None).unwrap();
+        let exec = MemoryExec::try_new(&input, schema, None, None).unwrap();
         let sort = Arc::new(SortExec::new(sort, Arc::new(exec)));
 
         let session_config = SessionConfig::new();

datafusion/core/tests/fuzz_cases/sort_preserving_repartition_fuzz.rs

@@ -357,7 +357,7 @@ mod sp_repartition_fuzz_tests {
         let concat_input_record = concat_batches(&schema, &input1).unwrap();
 
         let running_source = Arc::new(
-            MemoryExec::try_new(&[input1.clone()], schema.clone(), None)
+            MemoryExec::try_new(&[input1.clone()], schema.clone(), None, None)
                 .unwrap()
                 .try_with_sort_information(vec![sort_keys.clone()])
                 .unwrap(),

datafusion/core/tests/fuzz_cases/window_fuzz.rs

@@ -163,6 +163,7 @@ async fn bounded_window_causal_non_causal() -> Result<()> {
         &[batches.clone()],
         schema.clone(),
         None,
+        None,
     )?);
 
     // Different window functions to test causality
@@ -639,7 +640,7 @@ async fn run_window_test(
         },
     ]);
     let mut exec1 = Arc::new(
-        MemoryExec::try_new(&[vec![concat_input_record]], schema.clone(), None)?
+        MemoryExec::try_new(&[vec![concat_input_record]], schema.clone(), None, None)?
             .try_with_sort_information(vec![source_sort_keys.clone()])?,
     ) as _;
     // Table is ordered according to ORDER BY a, b, c In linear test we use PARTITION BY b, ORDER BY a
@@ -665,7 +666,7 @@ async fn run_window_test(
         vec![],
     )?) as _;
     let exec2 = Arc::new(
-        MemoryExec::try_new(&[input1.clone()], schema.clone(), None)?
+        MemoryExec::try_new(&[input1.clone()], schema.clone(), None, None)?
             .try_with_sort_information(vec![source_sort_keys.clone()])?,
     );
     let running_window_exec = Arc::new(BoundedWindowAggExec::try_new(

datafusion/core/tests/memory_limit/mod.rs

@@ -843,7 +843,7 @@ impl TableProvider for SortedTableProvider {
         _limit: Option<usize>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
         let mem_exec =
-            MemoryExec::try_new(&self.batches, self.schema(), projection.cloned())?
+            MemoryExec::try_new(&self.batches, self.schema(), None, projection.cloned())?
                 .try_with_sort_information(self.sort_information.clone())?;
 
         Ok(Arc::new(mem_exec))

datafusion/core/tests/physical_optimizer/limited_distinct_aggregation.rs

@@ -81,6 +81,7 @@ fn mock_data() -> Result<Arc<MemoryExec>> {
         &[vec![batch]],
         Arc::clone(&schema),
         None,
+        None,
     )?))
 }
 

datafusion/core/tests/user_defined/user_defined_table_functions.rs

@@ -156,6 +156,7 @@ impl TableProvider for SimpleCsvTable {
         Ok(Arc::new(MemoryExec::try_new(
             &[batches],
             TableProvider::schema(self),
+            None,
             projection.cloned(),
         )?))
     }

datafusion/physical-optimizer/src/aggregate_statistics.rs

@@ -252,6 +252,7 @@ mod tests {
             &[vec![batch]],
             Arc::clone(&schema),
             None,
+            None,
         )?))
     }
 

datafusion/physical-plan/benches/spm.rs

@@ -82,7 +82,7 @@ fn generate_spm_for_round_robin_tie_breaker(
         },
     ]);
 
-    let exec = MemoryExec::try_new(&partitiones, schema, None).unwrap();
+    let exec = MemoryExec::try_new(&partitiones, schema, None, None).unwrap();
     SortPreservingMergeExec::new(sort, Arc::new(exec))
         .with_round_robin_repartition(enable_round_robin_repartition)
 }

datafusion/physical-plan/src/aggregates/mod.rs

@@ -2185,6 +2185,7 @@ mod tests {
             ],
             Arc::clone(&schema),
             None,
+            None,
         )?);
         let aggregate_exec = Arc::new(AggregateExec::try_new(
             AggregateMode::Partial,
@@ -2415,6 +2416,7 @@ mod tests {
             &[input_batches],
             Arc::clone(&schema),
             None,
+            None,
         )?);
 
         let aggregate_exec = Arc::new(AggregateExec::try_new(
@@ -2538,6 +2540,7 @@ mod tests {
             &[vec![batch.clone()]],
             Arc::<Schema>::clone(&batch.schema()),
             None,
+            None,
         )?);
         let aggregate_exec = Arc::new(AggregateExec::try_new(
             AggregateMode::FinalPartitioned,
@@ -2607,6 +2610,7 @@ mod tests {
             &[input_data],
             Arc::clone(&schema),
             None,
+            None,
         )?);
         let aggregate_exec = Arc::new(AggregateExec::try_new(
             AggregateMode::Partial,
@@ -2697,6 +2701,7 @@ mod tests {
             &[input_data],
             Arc::clone(&schema),
             None,
+            None,
         )?);
         let aggregate_exec = Arc::new(AggregateExec::try_new(
             AggregateMode::Partial,