Function verify_join_output_partitioning

datafusion/core/tests/dataframe/mod.rs:2558–2630

()

#[tokio::test]
async fn verify_join_output_partitioning() -> Result<()> {
    let left = test_table().await?.select_columns(&["c1", "c2"])?;
    let right = test_table_with_name("c2")
        .await?
        .select_columns(&["c1", "c2"])?
        .with_column_renamed("c2.c1", "c2_c1")?
        .with_column_renamed("c2.c2", "c2_c2")?;

    let all_join_types = vec![
        JoinType::Inner,
        JoinType::Left,
        JoinType::Right,
        JoinType::Full,
        JoinType::LeftSemi,
        JoinType::RightSemi,
        JoinType::LeftAnti,
        JoinType::RightAnti,
        JoinType::LeftMark,
        JoinType::RightMark,
    ];

    let default_partition_count = SessionConfig::new().target_partitions();

    for join_type in all_join_types {
        let join = left.clone().join(
            right.clone(),
            join_type,
            &["c1", "c2"],
            &["c2_c1", "c2_c2"],
            None,
        )?;
        let physical_plan = join.create_physical_plan().await?;
        let out_partitioning = physical_plan.output_partitioning();
        let join_schema = physical_plan.schema();

        match join_type {
            JoinType::Left
            | JoinType::LeftSemi
            | JoinType::LeftAnti
            | JoinType::LeftMark => {
                let left_exprs: Vec<Arc<dyn PhysicalExpr>> = vec![
                    Arc::new(Column::new_with_schema("c1", &join_schema)?),
                    Arc::new(Column::new_with_schema("c2", &join_schema)?),
                ];
                assert_eq!(
                    out_partitioning,
                    &Partitioning::Hash(left_exprs, default_partition_count)
                );
            }
            JoinType::Inner
            | JoinType::Right
            | JoinType::RightSemi
            | JoinType::RightAnti
            | JoinType::RightMark => {
                let right_exprs: Vec<Arc<dyn PhysicalExpr>> = vec![
                    Arc::new(Column::new_with_schema("c2_c1", &join_schema)?),
                    Arc::new(Column::new_with_schema("c2_c2", &join_schema)?),
                ];