Add expression partitioning enum variant

datafusion/ffi/src/physical_expr/partitioning.rs

@@ -45,6 +45,8 @@ impl From<&Partitioning> for FFI_Partitioning {
                     .collect();
                 Self::Hash(exprs, *size)
             }
+            // FFI does not yet expose expression partition metadata.
+            Partitioning::Expr(expr) => Self::UnknownPartitioning(expr.partition_count()),
             Partitioning::UnknownPartitioning(size) => Self::UnknownPartitioning(*size),
         }
     }

datafusion/physical-expr/src/lib.rs

@@ -58,7 +58,7 @@ pub use analysis::{AnalysisContext, ExprBoundaries, analyze};
 pub use equivalence::{
     AcrossPartitions, ConstExpr, EquivalenceProperties, calculate_union,
 };
-pub use partitioning::{Distribution, Partitioning};
+pub use partitioning::{Distribution, ExprPartitioning, Partitioning};
 pub use physical_expr::{
     add_offset_to_expr, add_offset_to_physical_sort_exprs, create_lex_ordering,
     create_ordering, create_physical_sort_expr, create_physical_sort_exprs,

datafusion/physical-expr/src/partitioning.rs

@@ -117,6 +117,8 @@ pub enum Partitioning {
     /// Allocate rows based on a hash of one of more expressions and the specified number of
     /// partitions
     Hash(Vec<Arc<dyn PhysicalExpr>>, usize),
+    /// Partition rows by source-declared expression domains
+    Expr(ExprPartitioning),
     /// Unknown partitioning scheme with a known number of partitions
     UnknownPartitioning(usize),
 }
@@ -133,13 +135,94 @@ impl Display for Partitioning {
                     .join(", ");
                 write!(f, "Hash([{phy_exprs_str}], {size})")
             }
+            Partitioning::Expr(expr) => write!(f, "{expr}"),
             Partitioning::UnknownPartitioning(size) => {
                 write!(f, "UnknownPartitioning({size})")
             }
         }
     }
 }
 
+/// Physical expression partitioning.
+///
+/// Partition `i` contains rows where `partition_exprs[i]` evaluates to true.
+/// The source declaring partitioning is responsible for ensuring that, for every
+/// row emitted, exactly one partition expression evaluates to true and that row
+/// is emitted by the corresponding partition. The expressions do not need to
+/// cover values that the plan cannot emit.
+///
+/// For example, a scan that can only emit rows for `2022` can declare two date
+/// partitions as:
+///
+/// ```text
+/// partition_exprs[0] = date >= 2022-01-01 AND date < 2022-07-01
+/// partition_exprs[1] = date >= 2022-07-01 AND date < 2023-01-01
+/// ```
+///
+/// This is valid even though values outside `2022` are not covered, as long as
+/// the source does not emit rows outside those ranges. It would not be valid
+/// for this plan to emit a row from `partition[i]` whose date is not within
+/// `partition_exprs[i]`, or to emit a row whose date matches multiple
+/// partition expressions.
+///
+/// More complex partitioning can be represented using normal expression
+/// composition. For example, one partition in a date and city range can be
+/// represented as `date >= 2021-01-01 AND date < 2022-07-01 AND city < 'Boston'`.
+///
+/// NOTE: Optimizer and execution behavior for this partitioning is intentionally
+/// not implemented and will be introduced incrementally.
+#[derive(Debug, Clone)]
+pub struct ExprPartitioning {
+    partition_exprs: Vec<Arc<dyn PhysicalExpr>>,
+}
+
+impl ExprPartitioning {
+    /// Creates expression partitioning metadata from one predicate expression
+    /// per partition.
+    pub fn new(partition_exprs: Vec<Arc<dyn PhysicalExpr>>) -> Self {
+        Self { partition_exprs }
+    }
+
+    /// Returns the partition predicate expressions.
+    pub fn partition_exprs(&self) -> &[Arc<dyn PhysicalExpr>] {
+        &self.partition_exprs
+    }
+
+    /// Returns the number of partitions.
+    pub fn partition_count(&self) -> usize {
+        self.partition_exprs.len()
+    }
+
+    fn project(
+        &self,
+        mapping: &ProjectionMapping,
+        input_eq_properties: &EquivalenceProperties,
+    ) -> Option<Self> {
+        let partition_exprs = input_eq_properties
+            .project_expressions(&self.partition_exprs, mapping)
+            .collect::<Option<Vec<_>>>()?;
+
+        Some(Self { partition_exprs })
+    }
+}
+
+impl Display for ExprPartitioning {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(
+            f,
+            "Expr({}, {})",
+            format_physical_expr_list(&self.partition_exprs),
+            self.partition_count()
+        )
+    }
+}
+
+impl PartialEq for ExprPartitioning {
+    fn eq(&self, other: &Self) -> bool {
+        physical_exprs_equal(&self.partition_exprs, &other.partition_exprs)
+    }
+}
+
 /// Represents how a [`Partitioning`] satisfies a [`Distribution`] requirement.
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum PartitioningSatisfaction {
@@ -167,6 +250,7 @@ impl Partitioning {
         use Partitioning::*;
         match self {
             RoundRobinBatch(n) | Hash(_, n) | UnknownPartitioning(n) => *n,
+            Expr(expr) => expr.partition_count(),
         }
     }
 
@@ -277,19 +361,27 @@ impl Partitioning {
         mapping: &ProjectionMapping,
         input_eq_properties: &EquivalenceProperties,
     ) -> Self {
-        if let Partitioning::Hash(exprs, part) = self {
-            let normalized_exprs = input_eq_properties
-                .project_expressions(exprs, mapping)
-                .zip(exprs)
-                .map(|(proj_expr, expr)| {
-                    proj_expr.unwrap_or_else(|| {
-                        Arc::new(UnKnownColumn::new(&expr.to_string()))
+        match self {
+            Partitioning::Hash(exprs, part) => {
+                let normalized_exprs = input_eq_properties
+                    .project_expressions(exprs, mapping)
+                    .zip(exprs)
+                    .map(|(proj_expr, expr)| {
+                        proj_expr.unwrap_or_else(|| {
+                            Arc::new(UnKnownColumn::new(&expr.to_string()))
+                        })
                     })
-                })
-                .collect();
-            Partitioning::Hash(normalized_exprs, *part)
-        } else {
-            self.clone()
+                    .collect();
+                Partitioning::Hash(normalized_exprs, *part)
+            }
+            Partitioning::Expr(expr) => {
+                if let Some(projected) = expr.project(mapping, input_eq_properties) {
+                    Partitioning::Expr(projected)
+                } else {
+                    Partitioning::UnknownPartitioning(expr.partition_count())
+                }
+            }
+            _ => self.clone(),
         }
     }
 }
@@ -306,6 +398,7 @@ impl PartialEq for Partitioning {
             {
                 true
             }
+            (Partitioning::Expr(left), Partitioning::Expr(right)) => left == right,
             _ => false,
         }
     }
@@ -845,4 +938,67 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_expr_partitioning_metadata() -> Result<()> {
+        let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)]));
+        let col_a: Arc<dyn PhysicalExpr> =
+            Arc::new(Column::new_with_schema("a", &schema)?);
+
+        let expr_partitioning =
+            ExprPartitioning::new(vec![Arc::clone(&col_a), Arc::clone(&col_a)]);
+        let partitioning = Partitioning::Expr(expr_partitioning);
+
+        assert_eq!(partitioning.partition_count(), 2);
+        assert_eq!(partitioning.to_string(), "Expr([a@0, a@0], 2)");
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_expr_partitioning_project_degrades_when_expr_dropped() -> Result<()> {
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("a", DataType::Int64, false),
+            Field::new("b", DataType::Int64, false),
+        ]));
+        let col_b: Arc<dyn PhysicalExpr> =
+            Arc::new(Column::new_with_schema("b", &schema)?);
+        let eq_properties = EquivalenceProperties::new(Arc::clone(&schema));
+        let projection_mapping = ProjectionMapping::from_indices(&[0], &schema)?;
+        let expr_partitioning = Partitioning::Expr(ExprPartitioning::new(vec![col_b]));
+
+        let projected = expr_partitioning.project(&projection_mapping, &eq_properties);
+        let Partitioning::UnknownPartitioning(partition_count) = projected else {
+            panic!("expected UnknownPartitioning, got {projected:?}");
+        };
+        assert_eq!(partition_count, 1);
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_multi_partition_expr_does_not_satisfy_hash_distribution() -> Result<()> {
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("a", DataType::Int64, false),
+            Field::new("b", DataType::Int64, false),
+        ]));
+        let col_a: Arc<dyn PhysicalExpr> =
+            Arc::new(Column::new_with_schema("a", &schema)?);
+        let col_b: Arc<dyn PhysicalExpr> =
+            Arc::new(Column::new_with_schema("b", &schema)?);
+
+        let eq_properties = EquivalenceProperties::new(Arc::clone(&schema));
+        let expr_partitioning = Partitioning::Expr(ExprPartitioning::new(vec![
+            Arc::clone(&col_a),
+            Arc::clone(&col_b),
+        ]));
+        let required = Distribution::HashPartitioned(vec![col_a, col_b]);
+
+        assert_eq!(
+            expr_partitioning.satisfaction(&required, &eq_properties, false),
+            PartitioningSatisfaction::NotSatisfied
+        );
+
+        Ok(())
+    }
 }

datafusion/physical-plan/src/joins/utils.rs

@@ -144,6 +144,11 @@ pub fn adjust_right_output_partitioning(
                 .collect::<Result<_>>()?;
             Partitioning::Hash(new_exprs, *size)
         }
+        Partitioning::Expr(_) => {
+            return not_impl_err!(
+                "Expression partitioning is not supported for join output partitioning"
+            );
+        }
         result => result.clone(),
     };
     Ok(result)

datafusion/physical-plan/src/lib.rs

@@ -37,7 +37,7 @@ pub use datafusion_expr::{Accumulator, ColumnarValue};
 use datafusion_physical_expr::PhysicalSortExpr;
 pub use datafusion_physical_expr::window::WindowExpr;
 pub use datafusion_physical_expr::{
-    Distribution, Partitioning, PhysicalExpr, expressions,
+    Distribution, ExprPartitioning, Partitioning, PhysicalExpr, expressions,
 };
 
 pub use crate::display::{DefaultDisplay, DisplayAs, DisplayFormatType, VerboseDisplay};

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

@@ -600,6 +600,11 @@ impl BatchPartitioner {
                     num_input_partitions,
                 ))
             }
+            Partitioning::Expr(_) => {
+                not_impl_err!(
+                    "Expression partitioning is not supported by RepartitionExec"
+                )
+            }
             other => {
                 not_impl_err!("Unsupported repartitioning scheme {other:?}")
             }
@@ -1260,6 +1265,11 @@ impl ExecutionPlan for RepartitionExec {
                 }
                 Partitioning::Hash(new_partitions, *size)
             }
+            Partitioning::Expr(_) => {
+                return not_impl_err!(
+                    "Expression partitioning is not supported for projection pushdown through RepartitionExec"
+                );
+            }
             others => others.clone(),
         };
 
@@ -1296,6 +1306,11 @@ impl ExecutionPlan for RepartitionExec {
         if !self.maintains_input_order()[0] {
             return Ok(SortOrderPushdownResult::Unsupported);
         }
+        if matches!(self.partitioning(), Partitioning::Expr(_)) {
+            return not_impl_err!(
+                "Expression partitioning is not supported for sort pushdown through RepartitionExec"
+            );
+        }
 
         // Delegate to the child and wrap with a new RepartitionExec
         self.input.try_pushdown_sort(order)?.try_map(|new_input| {
@@ -1319,6 +1334,11 @@ impl ExecutionPlan for RepartitionExec {
             RoundRobinBatch(_) => RoundRobinBatch(target_partitions),
             Hash(hash, _) => Hash(hash, target_partitions),
             UnknownPartitioning(_) => UnknownPartitioning(target_partitions),
+            Expr(_) => {
+                return not_impl_err!(
+                    "Expression partitioning is not supported for changing RepartitionExec partition counts"
+                );
+            }
         };
         Ok(Some(Arc::new(Self {
             input: Arc::clone(&self.input),
@@ -1447,6 +1467,11 @@ impl RepartitionExec {
                     num_input_partitions,
                 )
             }
+            Partitioning::Expr(_) => {
+                return not_impl_err!(
+                    "Expression partitioning is not supported by RepartitionExec"
+                );
+            }
             other => {
                 return not_impl_err!("Unsupported repartitioning scheme {other:?}");
             }
@@ -1863,6 +1888,7 @@ mod tests {
     use datafusion_common_runtime::JoinSet;
     use datafusion_execution::config::SessionConfig;
     use datafusion_execution::runtime_env::RuntimeEnvBuilder;
+    use datafusion_physical_expr::ExprPartitioning;
     use insta::assert_snapshot;
 
     #[test]
@@ -2155,6 +2181,34 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn unsupported_expr_partitioning() -> Result<()> {
+        let task_ctx = Arc::new(TaskContext::default());
+        let batch = RecordBatch::try_from_iter(vec![(
+            "my_awesome_field",
+            Arc::new(StringArray::from(vec!["foo", "bar"])) as ArrayRef,
+        )])?;
+
+        let schema = batch.schema();
+        let expr = col("my_awesome_field", &schema)?;
+        let input = MockExec::new(vec![Ok(batch)], Arc::clone(&schema));
+        let partitioning = Partitioning::Expr(ExprPartitioning::new(vec![expr]));
+        let exec = RepartitionExec::try_new(Arc::new(input), partitioning)?;
+        let output_stream = exec.execute(0, task_ctx)?;
+
+        let result_string = crate::common::collect(output_stream)
+            .await
+            .unwrap_err()
+            .to_string();
+        assert!(
+            result_string
+                .contains("Expression partitioning is not supported by RepartitionExec"),
+            "actual: {result_string}"
+        );
+
+        Ok(())
+    }
+
     #[tokio::test]
     async fn error_for_input_exec() {
         // This generates an error on a call to execute. The error

datafusion/physical-plan/src/sorts/sort_preserving_merge.rs

@@ -1505,11 +1505,7 @@ mod tests {
         let task_ctx = Arc::new(TaskContext::default());
         let schema = Schema::new(vec![Field::new("c1", DataType::UInt64, false)]);
         let properties = CongestedExec::compute_properties(Arc::new(schema.clone()));
-        let &partition_count = match properties.output_partitioning() {
-            Partitioning::RoundRobinBatch(partitions) => partitions,
-            Partitioning::Hash(_, partitions) => partitions,
-            Partitioning::UnknownPartitioning(partitions) => partitions,
-        };
+        let partition_count = properties.output_partitioning().partition_count();
         let source = CongestedExec {
             schema: schema.clone(),
             cache: Arc::new(properties),