Support co-partitioned range inner equi joins

Author: gene-bordegaray

datafusion/core/tests/physical_optimizer/sanity_checker.rs

@@ -19,9 +19,9 @@ use insta::assert_snapshot;
 use std::sync::Arc;
 
 use crate::physical_optimizer::test_utils::{
-    bounded_window_exec, global_limit_exec, local_limit_exec, memory_exec,
-    projection_exec, repartition_exec, sort_exec, sort_expr, sort_expr_options,
-    sort_merge_join_exec, sort_preserving_merge_exec, union_exec,
+    bounded_window_exec, global_limit_exec, hash_join_exec, local_limit_exec,
+    memory_exec, projection_exec, repartition_exec, sort_exec, sort_expr,
+    sort_expr_options, sort_merge_join_exec, sort_preserving_merge_exec, union_exec,
 };
 
 use arrow::compute::SortOptions;
@@ -30,8 +30,8 @@ use datafusion::datasource::stream::{FileStreamProvider, StreamConfig, StreamTab
 use datafusion::prelude::{CsvReadOptions, SessionContext};
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::{JoinType, Result, ScalarValue};
-use datafusion_physical_expr::Partitioning;
 use datafusion_physical_expr::expressions::{Literal, col};
+use datafusion_physical_expr::{Partitioning, RangePartitioning, SplitPoint};
 use datafusion_physical_expr_common::sort_expr::LexOrdering;
 use datafusion_physical_optimizer::PhysicalOptimizerRule;
 use datafusion_physical_optimizer::sanity_checker::SanityCheckPlan;
@@ -400,6 +400,50 @@ fn assert_sanity_check(plan: &Arc<dyn ExecutionPlan>, is_sane: bool) {
     );
 }
 
+fn range_partitioned_exec(
+    schema: &SchemaRef,
+    key: &str,
+    split_points: impl IntoIterator<Item = i32>,
+) -> Result<Arc<dyn ExecutionPlan>> {
+    let split_points = split_points
+        .into_iter()
+        .map(|value| SplitPoint::new(vec![ScalarValue::Int32(Some(value))]))
+        .collect();
+    let partitioning = Partitioning::Range(RangePartitioning::try_new(
+        [sort_expr(key, schema)].into(),
+        split_points,
+    )?);
+    RepartitionExec::try_new(memory_exec(schema), partitioning)
+        .map(|exec| Arc::new(exec) as Arc<dyn ExecutionPlan>)
+}
+
+#[test]
+fn test_partitioned_hash_join_requires_co_partitioned_children() -> Result<()> {
+    let schema = create_test_schema2();
+    let join_on = vec![(col("a", &schema)?, col("b", &schema)?)];
+    let right = range_partitioned_exec(&schema, "b", [10])?;
+
+    let valid_join = hash_join_exec(
+        range_partitioned_exec(&schema, "a", [10])?,
+        Arc::clone(&right),
+        join_on.clone(),
+        None,
+        &JoinType::Inner,
+    )?;
+    assert_sanity_check(&valid_join, true);
+
+    let invalid_join = hash_join_exec(
+        range_partitioned_exec(&schema, "a", [20])?,
+        right,
+        join_on,
+        None,
+        &JoinType::Inner,
+    )?;
+    assert_sanity_check(&invalid_join, false);
+
+    Ok(())
+}
+
 #[tokio::test]
 /// Tests that plan is valid when the sort requirements are satisfied.
 async fn test_bounded_window_agg_sort_requirement() -> Result<()> {

datafusion/physical-expr/src/partitioning.rs

@@ -244,50 +244,6 @@ impl RangePartitioning {
         self.split_points.len() + 1
     }
 
-    /// Returns true when `self` and `other` describe the same range partition
-    /// map.
-    ///
-    /// Single-partition range partitionings are always compatible. Otherwise,
-    /// the two partitionings must have identical split points and equivalent
-    /// ordering expressions with the same sort options.
-    pub fn compatible_with(
-        &self,
-        other: &Self,
-        eq_properties: &EquivalenceProperties,
-    ) -> bool {
-        if self.partition_count() == 1 && other.partition_count() == 1 {
-            return true;
-        }
-
-        if self.split_points != other.split_points
-            || self.ordering.len() != other.ordering.len()
-        {
-            return false;
-        }
-
-        if !self
-            .ordering
-            .iter()
-            .zip(other.ordering.iter())
-            .all(|(left, right)| left.options == right.options)
-        {
-            return false;
-        }
-
-        let left_exprs = self
-            .ordering
-            .iter()
-            .map(|sort_expr| Arc::clone(&sort_expr.expr))
-            .collect::<Vec<_>>();
-        let right_exprs = other
-            .ordering
-            .iter()
-            .map(|sort_expr| Arc::clone(&sort_expr.expr))
-            .collect::<Vec<_>>();
-
-        equivalent_exprs(&left_exprs, &right_exprs, eq_properties)
-    }
-
     /// Calculates the range partitioning after applying the given projection.
     ///
     /// Returns `None` if any range key cannot be projected or if projection
@@ -406,42 +362,6 @@ impl Partitioning {
         }
     }
 
-    /// Returns true when `self` and `other` describe compatible partition maps.
-    ///
-    /// Compatible partition maps can be used for partition-local behavior: if
-    /// this returns true, partition `i` from both partitionings can be treated
-    /// as covering the same partition domain. This is stricter than
-    /// [`Self::satisfaction`], which only answers whether this partitioning can
-    /// satisfy a required distribution.
-    pub fn compatible_with(
-        &self,
-        other: &Self,
-        eq_properties: &EquivalenceProperties,
-    ) -> bool {
-        if self.partition_count() == 1 && other.partition_count() == 1 {
-            return true;
-        }
-
-        match (self, other) {
-            (
-                Partitioning::Hash(left_exprs, left_count),
-                Partitioning::Hash(right_exprs, right_count),
-            ) => {
-                if left_count != right_count {
-                    return false;
-                }
-                if left_exprs.is_empty() || right_exprs.is_empty() {
-                    return false;
-                }
-                equivalent_exprs(left_exprs, right_exprs, eq_properties)
-            }
-            (Partitioning::Range(left), Partitioning::Range(right)) => {
-                left.compatible_with(right, eq_properties)
-            }
-            _ => false,
-        }
-    }
-
     /// Returns true if `subset_exprs` is a subset of `exprs`.
     /// For example: Hash(a, b) is subset of Hash(a) since a partition with all occurrences of
     /// a distinct (a) must also contain all occurrences of a distinct (a, b) with the same (a).
@@ -1304,157 +1224,6 @@ mod tests {
         Ok(())
     }
 
-    #[test]
-    fn test_range_partitioning_compatible_with() -> Result<()> {
-        let fixture = PartitioningTestFixture::int64(&["a", "b"])?;
-        let mut eq_properties = fixture.eq_properties.clone();
-        eq_properties.add_equal_conditions(fixture.col(0), fixture.col(1))?;
-
-        let split_points = vec![int_split_point([10]), int_split_point([20])];
-        let range_a = fixture.range([0], split_points.clone());
-        let range_a_same = fixture.range([0], split_points.clone());
-        let range_b_equivalent = fixture.range([1], split_points.clone());
-        let range_b_different_split = fixture.range([1], vec![int_split_point([30])]);
-        let range_a_desc = RangePartitioning::try_new(
-            [fixture.range_sort_expr(0, SortOptions::new(true, false))].into(),
-            vec![int_split_point([10])],
-        )?;
-        let single_partition_range_a = fixture.range([0], vec![]);
-        let single_partition_range_b = fixture.range([1], vec![]);
-
-        assert!(range_a.compatible_with(&range_a_same, &fixture.eq_properties));
-        assert!(range_a.compatible_with(&range_b_equivalent, &eq_properties));
-        assert!(!range_a.compatible_with(&range_b_equivalent, &fixture.eq_properties));
-        assert!(!range_a.compatible_with(&range_b_different_split, &eq_properties));
-        assert!(!range_a.compatible_with(&range_a_desc, &eq_properties));
-        assert!(
-            single_partition_range_a
-                .compatible_with(&single_partition_range_b, &fixture.eq_properties)
-        );
-
-        assert!(
-            fixture
-                .range_partitioning([0], vec![int_split_point([10])])
-                .compatible_with(
-                    &fixture.range_partitioning([1], vec![int_split_point([10])]),
-                    &eq_properties
-                )
-        );
-        assert!(
-            !fixture
-                .range_partitioning([0], vec![int_split_point([10])])
-                .compatible_with(
-                    &fixture.range_partitioning([0], vec![int_split_point([20])]),
-                    &fixture.eq_properties
-                )
-        );
-        assert!(
-            !fixture
-                .range_partitioning([0], vec![int_split_point([10])])
-                .compatible_with(
-                    &fixture.hash_partitioning([0], 2),
-                    &fixture.eq_properties
-                )
-        );
-
-        Ok(())
-    }
-
-    #[test]
-    fn test_hash_partitioning_compatible_with() -> Result<()> {
-        let fixture = PartitioningTestFixture::int64(&["a", "b"])?;
-        let mut eq_properties = fixture.eq_properties.clone();
-        eq_properties.add_equal_conditions(fixture.col(0), fixture.col(1))?;
-
-        assert!(
-            fixture.hash_partitioning([0], 2).compatible_with(
-                &fixture.hash_partitioning([0], 2),
-                &fixture.eq_properties
-            )
-        );
-        assert!(
-            fixture
-                .hash_partitioning([0], 2)
-                .compatible_with(&fixture.hash_partitioning([1], 2), &eq_properties)
-        );
-        assert!(
-            !fixture.hash_partitioning([0], 2).compatible_with(
-                &fixture.hash_partitioning([1], 2),
-                &fixture.eq_properties
-            )
-        );
-        assert!(
-            !fixture.hash_partitioning([0], 2).compatible_with(
-                &fixture.hash_partitioning([0], 3),
-                &fixture.eq_properties
-            )
-        );
-        assert!(!fixture.hash_partitioning([0], 2).compatible_with(
-            &fixture.hash_partitioning([0, 1], 2),
-            &fixture.eq_properties
-        ));
-        assert!(
-            !Partitioning::Hash(vec![], 2)
-                .compatible_with(&Partitioning::Hash(vec![], 2), &fixture.eq_properties)
-        );
-        assert!(!fixture.hash_partitioning([0], 2).compatible_with(
-            &fixture.range_partitioning([0], vec![int_split_point([10])]),
-            &fixture.eq_properties
-        ));
-        assert!(
-            fixture.hash_partitioning([0], 1).compatible_with(
-                &Partitioning::RoundRobinBatch(1),
-                &fixture.eq_properties
-            )
-        );
-
-        Ok(())
-    }
-
-    #[test]
-    fn test_round_robin_partitioning_compatible_with() {
-        let eq_properties = EquivalenceProperties::new(Arc::new(Schema::empty()));
-
-        assert!(
-            Partitioning::RoundRobinBatch(1)
-                .compatible_with(&Partitioning::RoundRobinBatch(1), &eq_properties)
-        );
-        assert!(
-            !Partitioning::RoundRobinBatch(2)
-                .compatible_with(&Partitioning::RoundRobinBatch(2), &eq_properties)
-        );
-        assert!(
-            Partitioning::RoundRobinBatch(1)
-                .compatible_with(&Partitioning::UnknownPartitioning(1), &eq_properties)
-        );
-        assert!(
-            !Partitioning::RoundRobinBatch(2)
-                .compatible_with(&Partitioning::UnknownPartitioning(2), &eq_properties)
-        );
-    }
-
-    #[test]
-    fn test_unknown_partitioning_compatible_with() {
-        let eq_properties = EquivalenceProperties::new(Arc::new(Schema::empty()));
-
-        assert!(
-            Partitioning::UnknownPartitioning(1)
-                .compatible_with(&Partitioning::UnknownPartitioning(1), &eq_properties)
-        );
-        assert!(
-            !Partitioning::UnknownPartitioning(2)
-                .compatible_with(&Partitioning::UnknownPartitioning(2), &eq_properties)
-        );
-        assert!(
-            Partitioning::UnknownPartitioning(1)
-                .compatible_with(&Partitioning::RoundRobinBatch(1), &eq_properties)
-        );
-        assert!(
-            !Partitioning::UnknownPartitioning(2)
-                .compatible_with(&Partitioning::RoundRobinBatch(2), &eq_properties)
-        );
-    }
-
     #[test]
     fn test_multi_partition_range_does_not_satisfy_hash_distribution() -> Result<()> {
         let fixture = PartitioningTestFixture::int64(&["a", "b"])?;