Function partition_aware_union

datafusion/core/tests/dataframe/mod.rs:2470–2509

()

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

    let left_rows = left.clone().collect().await?;
    let right_rows = right.clone().collect().await?;
    let join1 =
        left.clone()
            .join(right.clone(), JoinType::Inner, &["c1"], &["c2_c1"], None)?;
    let join2 = left.join(right, JoinType::Inner, &["c1"], &["c2_c1"], None)?;

    let union = join1.union(join2)?;

    let union_rows = union.clone().collect().await?;

    assert_eq!(100, left_rows.iter().map(|x| x.num_rows()).sum::<usize>());
    assert_eq!(100, right_rows.iter().map(|x| x.num_rows()).sum::<usize>());
    assert_eq!(4016, union_rows.iter().map(|x| x.num_rows()).sum::<usize>());

    let physical_plan = union.create_physical_plan().await?;
    let default_partition_count = SessionConfig::new().target_partitions();

    // For partition aware union, the output partition count should not be changed.
    assert_eq!(
        physical_plan.output_partitioning().partition_count(),
        default_partition_count
    );
    // For partition aware union, the output partition is the same with the union's inputs
    for child in physical_plan.children() {
        assert_eq!(
            physical_plan.output_partitioning(),
            child.output_partitioning()
        );
    }

    Ok(())
}

#[tokio::test]
async fn non_partition_aware_union() -> Result<()> {