Optional serialization/deserialization for WorkerChannel - #566
#566Optional serialization/deserialization for WorkerChannel#566gabotechs wants to merge 1 commit intomaindatafusion-contrib/datafusion-distributed:mainfrom gabrielmusat/optional-serdedatafusion-contrib/datafusion-distributed:gabrielmusat/optional-serdeCopy head branch name to clipboard
Conversation
a7fe44c to
40cf654
Compare
| .coordinator_channel(headers, coordinator_to_worker_stream) | ||
| .coordinator_channel( | ||
| headers, | ||
| set_plan_request, | ||
| coordinator_to_worker_stream, | ||
| &task_ctx, | ||
| metrics_set, | ||
| ) |
There was a problem hiding this comment.
Now we need to pass two extra things here:
- the
SetPlanRequest. Before, this was sent as the first message in the stream, this being an invariant of the coordinator->worker channel: the first message always needs to be theSetPlanRequest. This was an artificial invariant imposed by gRPC spec limitations: in gRPC you cannot send both a stream and a unary message at the same time. As this is specific to gRPC, handling this behavior was moved under thegrpcscope, and theWorkerChannelprotocol accepts now the two things separately theCoordinatorToWorkerStreamand theSetPlanRequest. - An
ExecutionPlanMetricsSet. Thecoordinator_channelimplementations now decide if the inner messages need to be deserialized or can be passed in-memory, so some metrics likeplan_bytes_sentare now implementation specific (an in-memory implementation will never even serialize to bytes). For this,ExecutionPlanMetricsSetis passed to give the chance to implementations to contribute their own metrics.
| } | ||
|
|
||
| impl ProducerHead { | ||
| pub(crate) fn to_spec(&self, cfg: &SessionConfig) -> Result<ProducerHeadSpec> { |
There was a problem hiding this comment.
Now, we can get read of the intermediary ProducerHeadSpec, as this existed just for bridging a gap that no longer exists: the original ProducerHead can just hold a reference to MaybeEncoded and handle both in-memory and serialized formats.
More information about what this was needed before in point 1 of the PR description in #512 (comment).
| } | ||
| ProducerHead::RepartitionExec { partitioning } => Arc::new(RepartitionExec::try_new( | ||
| input, | ||
| partitioning.try_decoded()?, |
There was a problem hiding this comment.
I don't love this. This is a runtime failure that was irrepresentable before, but now we need this try_decoded()? runtime check because we cannot proceed here if the ProducerHead contains a serialized Partitioning struct.
It allows to remove so much code that I think it's worth it.
| pb::coordinator_to_worker_msg::Inner::SetPlanRequest(_) => { | ||
| return Err(Status::invalid_argument( | ||
| "SetPlanRequest must be the first coordinator message", | ||
| )); | ||
| } |
There was a problem hiding this comment.
As the fact that SetPlanRequest is the first message in the stream is an invariant imposed by gRPC itself, it's now checked here, instead of the outside of the grpc scope
40cf654 to
fb1b599
Compare
Interesting, thanks. Is there some other motivation for this beyond the in-memory implementation? While the in-memory implementation is important for test/smoke coverage, I think that all practical implementations will need to serialize: ours certainly does. Knowing what the other use cases are would help justify the added complexity. |
|
It's exclusively for the in-memory implementation. For example, if the coordinating context decides that a stage does not need to run in more than one task, it can co-locate it in itself, using the in-memory WorkerChannel instead of paying serialization. My intention is that the in-memory implementation is not just something that runs for tests, it's part of the production machinery that triggers every time the coordinator or a worker needs to call itself. |
fb1b599 to
3e5d859
Compare
| /// Returns the decoded variant: | ||
| /// - If in `Decoded` state, it just passes through the content. | ||
| /// - If in `Encoded` state, it decodes using the provided callback. | ||
| pub(crate) fn decode_with(self, decode: impl FnOnce(Vec<u8>) -> Result<T>) -> Result<T> { | ||
| match self { | ||
| Self::Encoded(encoded) => decode(encoded), | ||
| Self::Decoded(decoded) => Ok(decoded), | ||
| } | ||
| } |
There was a problem hiding this comment.
So, this ends up shaped a lot like a OnceLock.
Could/should MaybeEncoded wrap a OnceLock to allow for in-place lazy decoding? A vague concern I have is that because decoding the MaybeEncoded doesn't cache the result, different consumers might end up decoding multiple times?
Adds some small changes to the
WorkerChannelcontract so that [de]serialization is optional. In-memory implementations might not want to pay the overhead of serializing plans.Details about the changes explained in individual comments guiding reviews.