Function join_agg_yields

datafusion/core/tests/execution/coop.rs:618–699

(
    #[values(false, true)] pretend_infinite: bool,
)

#[rstest]
#[tokio::test]
async fn join_agg_yields(
    #[values(false, true)] pretend_infinite: bool,
) -> Result<(), Box<dyn Error>> {
    // Session, schema, and a single 8 K‐row batch for each side
    let session_ctx = SessionContext::new();

    // on the right side, we’ll shift each value by +1 so that not everything joins,
    // but plenty of matching keys exist (e.g. 0 on left matches 1 on right, etc.)
    let infinite_left = make_lazy_exec_with_range("value", -10..10, false);
    let infinite_right =
        make_lazy_exec_with_range("value", 0..i64::MAX, pretend_infinite);

    // 2b) Create Join keys → join on “value” = “value”
    let left_keys: Vec<Arc<dyn PhysicalExpr>> = vec![Arc::new(Column::new("value", 0))];
    let right_keys: Vec<Arc<dyn PhysicalExpr>> = vec![Arc::new(Column::new("value", 0))];

    let part_left = Partitioning::Hash(left_keys, 1);
    let part_right = Partitioning::Hash(right_keys, 1);

    // Wrap each side in Repartition so they are both hashed into 1 partition
    let hashed_left = Arc::new(RepartitionExec::try_new(
        Arc::new(infinite_left),
        part_left,
    )?);
    let hashed_right = Arc::new(RepartitionExec::try_new(
        Arc::new(infinite_right),
        part_right,
    )?);

    // Build an Inner HashJoinExec → left.value = right.value
    let join = Arc::new(HashJoinExec::try_new(
        hashed_left,
        hashed_right,
        vec![(
            Arc::new(Column::new("value", 0)),
            Arc::new(Column::new("value", 0)),
        )],
        None,
        &JoinType::Inner,
        None,
        PartitionMode::CollectLeft,
        NullEquality::NullEqualsNull,
        false,
    )?);

    // Project only one column (“value” from the left side) because we just want to sum that
    let input_schema = join.schema();

    let proj_expr = vec![ProjectionExpr::new(
        Arc::new(Column::new_with_schema("value", &input_schema)?) as _,
        "value",
    )];

    let projection = Arc::new(ProjectionExec::try_new(proj_expr, join)?);
    let projection_schema = projection.schema();

    let output_fields = vec![Field::new("total", DataType::Int64, true)];
    let output_schema = Arc::new(Schema::new(output_fields));