Simplify MemoryWriteExec

[tustvold]

Which issue does this PR close?

Closes #.

Rationale for this change

Simplifies the logic in #6049 to make it a bit easier to see what is actually going on.

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

[tustvold]

The previous logic had a special case for if the partitioning matched, this would effectively save an atomic increment per batch. Given we are polling a dyn Stream here, I am very skeptical there is any performance difference

[metesynnada]

There is a performance difference of batch_count * acquire_lock, if this is OK, we can move on to this. Do you think it would affect the benchmarks? I am not quite familiar with that part.

[tustvold]

unfold is incredibly useful. The only downside is you can't name the type, at least until we get existential types

[metesynnada]

Now I get what you mean, if we do not hold the async lock in the state as acquired, the folding becomes possible. This substantially shrinks the code size, cool pattern.

[tustvold]

We could easily hold the lock in the state if we wanted to, I just didn't think it warranted the added complexity. You could definitely do something like

let stream = futures::stream::unfold(state, |mut state| async move {
    let locked = state.1.write_owned().await;
    loop {
        let batch = match state.0.next().await {
            Some(Ok(batch)) => batch,
            Some(Err(e)) => return Some((Err(e), state)),
            None => return None,
        };
        locked.push(batch)
    }
});

Or even

let stream = futures::stream::unfold(state, |mut state| async move {
    let locked = state.1.write().await;
    loop {
        let batch = match state.0.next().await {
            Some(Ok(batch)) => batch,
            Some(Err(e)) => {
                drop(locked);
                return Some((Err(e), state))
            }
            None => {
                drop(locked);
                return None
            }
        };
        locked.push(batch)
    }
});

Ultimately an uncontended lock is not going to matter to performance, unless in a hot loop with no other branches

[metesynnada]

LGTM, both implementations are neat (partitions matching & not matching). Thanks for the PR.

[alamb]

Thank you @tustvold

This looks like a great simplification to me. cc @metesynnada and @ozankabak

[ozankabak]

Will take a look and measure if the special casing is really not useful or not

[alamb]

Will take a look and measure if the special casing is really not useful or not

@ozankabak have you had a chance to review the performance of this PR? If so, is it acceptable to merge?

[ozankabak]

Will do tomorrow and let you know

[ozankabak]

We did the benchmarks, @metesynnada will share them here today

[metesynnada]

TLDR; it is recommended to retain both versions of the lock mechanism to accommodate both scenarios where the partitions match and where they do not match.

So, we ran a benchmark to see how two different locking strategies perform when processing a stream of RecordBatch objects. Here are the two functions we compared:

1. lock_only_once:

fn lock_only_once(schema: SchemaRef, state: (SendableRecordBatchStream, Arc<RwLock<Vec<RecordBatch>>>)) -> SendableRecordBatchStream{
    let iter = futures::stream::unfold(state, |mut state| async move {
        let mut locked = state.1.write().await;
        loop {
            let batch = match state.0.next().await {
                Some(Ok(batch)) => batch,
                Some(Err(e)) => {
                    drop(locked);
                    return Some((Err(e), state))
                }
                None => {
                    drop(locked);
                    return None
                }
            };
            locked.push(batch)
        }
    });
    Box::pin(RecordBatchStreamAdapter::new(schema, iter))
}

2. lock_multiple:

fn lock_multiple(schema: SchemaRef, state: (SendableRecordBatchStream, Arc<RwLock<Vec<RecordBatch>>>)) -> SendableRecordBatchStream{
    let iter = futures::stream::unfold(state, |mut state| async move {
        loop {
            let batch = match state.0.next().await {
                Some(Ok(batch)) => batch,
                Some(Err(e)) => return Some((Err(e), state)),
                None => return None,
            };
            state.1.write().await.push(batch)
        }
    });
    Box::pin(RecordBatchStreamAdapter::new(schema, iter))
}

We tested these functions with an input size of 10,000, and here's what we found:

Multiple lock strategy (lock_multiple function):

• Mean execution time: 545.97 µs
• Range: [544.78 µs, 547.26 µs]
• A couple of outliers: 2 (2.00%) high mild

Single lock strategy (lock_only_once function):

• Mean execution time: 330.21 µs
• Range: [329.53 µs, 330.94 µs]
• A few outliers: 7 (7.00%) in total
  • 1 (1.00%) low mild
  • 3 (3.00%) high mild
  • 3 (3.00%) high severe

So, it looks like the "Single lock" strategy (using the lock_only_once function) works better than the "Multiple lock" strategy (using the lock_multiple function) for the given input size of 10,000, as we expected. The "Single lock" strategy takes about 39.5% less time to complete compared to the "Multiple lock" strategy.

[tustvold]

I can add the lock multiple approach if people feel strongly, although at 20ns per row that seems unlikely this would be visible in any practical workload - I wish our kernels were a similar order of magnitude 😅

[ozankabak]

Thank you! Since the code difference is small thanks to how you are leveraging unfold, let's add it and then this is good to go from our perspective.

[alamb]

I made the change suggested by @metesynnada by #6154 (comment) in 01fa417 and merged up from main to try and nudge this PR along.

[ozankabak]

Seems like we have a whitespace issue making cargo fmt fail

[alamb]

Thanks @metesynnada @tustvold and @ozankabak

Here is another PR that I think simplifies and speeds up this code even more: #6236