feat(arrow-ipc): add sans-IO stream encoder - #10277
Conversation
|
I'll try and give this a review within a few days |
Rich-T-kid
left a comment
There was a problem hiding this comment.
I took a brief glance, looks pretty good to me so far. I can take a more in depth look tomorrow. I left a suggestion 🚀
234c1f6 to
94eaf13
Compare
There was a problem hiding this comment.
thank you @Phoenix500526, I think there's value in having a sans-I/O stream encoder, but we need to be careful about the duplication we're adding to arrow-ipc. The file is already hard to understand without much context, and adding duplicate logic in two places makes it hard to maintain.
The tests showcase a good way this can be used, but I'd like the implementation to be leaner. I left a few suggestion let me know what you think
3f24353 to
3a6e3c9
Compare
|
@Phoenix500526 is this good for review? |
@Rich-T-kid Yes, this is ready for review. Thanks for checking. |
Rich-T-kid
left a comment
There was a problem hiding this comment.
I think this looks pretty good. I left a couple of suggestions. Thank you for splitting up the commits into separate independent chunks!
Would be nice to validate that this PR didn't accidentally slow down the Stream writer somehow.
|
It would be nice to double check that we didnt accidntly slow down the StreamWriter as well. @alamb could you please run the arrow-ipc benchmarks? Thank you 👍 |
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as duplicate.
This comment was marked as duplicate.
This comment was marked as outdated.
This comment was marked as outdated.
|
looks to be mostly noise? its within +- 10. can we run it again to check if its reproducible |
Rich-T-kid
left a comment
There was a problem hiding this comment.
LGTM @Phoenix500526
@Jefffrey I think this is ready for a review
This comment was marked as outdated.
This comment was marked as outdated.
1 similar comment
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as duplicate.
This comment was marked as duplicate.
This comment was marked as duplicate.
This comment was marked as duplicate.
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as outdated.
|
benchmarks seem to suggest a consistent regression 🤔 |
5cfcf8a to
84eb07f
Compare
|
run benchmark ipc_writer |
This comment was marked as duplicate.
This comment was marked as duplicate.
|
🤖 Arrow criterion benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
run benchmark ipc_writer |
This comment was marked as duplicate.
This comment was marked as duplicate.
|
🤖 Arrow criterion benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
numbers looking better now 👀 |
|
|
run benchmark ipc_writer |
This comment was marked as duplicate.
This comment was marked as duplicate.
|
🤖 Arrow criterion benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
| Ok(()) | ||
| } | ||
|
|
||
| fn write_record_batch( |
There was a problem hiding this comment.
do we still need this custom impl, or can rely on the default?
There was a problem hiding this comment.
I've tried, but that makes the performance regress again. I extracted a shared metadata_layout method to reduce the duplication in commit 1b502ea
|
MIRI was fixed in #10433, merging up to get fix |
alamb
left a comment
There was a problem hiding this comment.
Thank you @Jefffrey and @Phoenix500526 and @Rich-T-kid
I think this is a very nicely written and commented PR. Since most of the PR is internal implementation / not exposed publically most of my comments could be done as a follow on PR
However, I do think it is worth considering a slightly different StreamEncoder API (comments below) that don't force buffering the output buffers into a Vec<Buffer>
| } | ||
|
|
||
| #[inline] | ||
| fn metadata_layout(metadata_len: usize, write_options: &IpcWriteOptions) -> MetadataLayout { |
There was a problem hiding this comment.
As a follow on, this might naturally be encapsualted as MetadataLayout::new type constructor
| trait IpcMessageSink { | ||
| fn write_slice(&mut self, bytes: &[u8]) -> Result<(), ArrowError>; | ||
|
|
||
| fn write_vec(&mut self, bytes: Vec<u8>) -> Result<(), ArrowError> { |
| /// Destination for a complete framed IPC message. | ||
| /// | ||
| /// This emits the stream/file framing around the serialized FlatBuffer | ||
| /// [`crate::Message`] metadata plus its optional body buffers. |
There was a problem hiding this comment.
It might also be worth mentioning to point out that the default implementation will copy data
| options: IpcWriteOptions, | ||
| ) -> Vec<u8> { | ||
| let mut encoder = StreamEncoder::try_new_with_options(schema, options).unwrap(); | ||
| let mut bytes = Vec::new(); |
There was a problem hiding this comment.
since Vec implements Write why not just use the Write impl directly?
| /// # Errors | ||
| /// | ||
| /// Returns an error if encoding fails. | ||
| pub fn encode(&mut self, batch: &RecordBatch) -> Result<Vec<Buffer>, ArrowError> { |
There was a problem hiding this comment.
I think this API is quite good and supports a usecase like async writer quite well.
My only suggestion / feedback is that this API basically requires buffering an entire record batch in RAM and now requies copying into several intermediate Vecs (e.g. for the metadata header, etc)
Another API we might consider would be to make the SteamEncoder templated on IpcMessageSink so it can bass the buffers directly
Something like
struct StreamEncoder<IpcMessageSink>` {
...
/// encode the batches, calling methods like `append_vec()`, `append_buffer`, append_slice, etc
pub fn encode(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
}However, this does still feel like there is "IO" in the encoder 🤔 -- on the other hand you could implement an IpcMessageSync that just buffers the data to get the same effect.
There was a problem hiding this comment.
Thanks @alamb , I agree this is worth exploring. I’m hesitant to change this PR to make StreamEncoder generic over a sink because that would be a larger API design change and makes the sans-IO boundary less obvious. I’ll keep the current encode -> Vec<Buffer> API in this PR and open a follow-up issue to track a possible sink-based / encode_to API. I open an issue #10445 to track it.
| } | ||
|
|
||
| /// Writes the IPC continuation marker and metadata length prefix. | ||
| fn write_continuation( |
There was a problem hiding this comment.
I found the mixing of methods you are supposed to override (like write_slice and write_vec) and methods that you aren't (like write_continuation) to be somewhat confusing.
It would be easier to read this code I think if we separated out the two types of methods
Perhaps something like this
/// The only things a sink actually needs to customize.
trait IpcMessageSink {
fn write_slice(&mut self, bytes: &[u8]) -> Result<(), ArrowError>;
fn write_vec(&mut self, bytes: Vec<u8>) -> Result<(), ArrowError> {
self.write_slice(&bytes)
}
fn write_encoded_buffer(&mut self, buffer: EncodedBuffer) -> Result<(), ArrowError> {
self.write_slice(buffer.as_slice())
}
}
/// The shared framing code
trait IpcMessageSinkExt: IpcMessageSink {
fn write_padding(&mut self, len: usize) -> Result<(), ArrowError> {
self.write_slice(&PADDING[..len])
}
fn write_continuation(&mut self, write_options: &IpcWriteOptions, metadata_len: i32) -> Result<(), ArrowError> {
// ... unchanged body ...
}
fn write_encoded_data(&mut self, encoded: EncodedData, write_options: &IpcWriteOptions) -> Result<(usize, usize), ArrowError> {
// ... unchanged body, calling self.write_vec / self.write_padding ...
}
fn write_record_batch(&mut self, metadata: Vec<u8>, encoded_buffers: Vec<EncodedBuffer>, body_len: usize, tail_pad: usize, write_options: &IpcWriteOptions) ->
Result<(usize, usize), ArrowError> {
// ... unchanged body ...
}
fn write_eos(&mut self, write_options: &IpcWriteOptions) -> Result<(), ArrowError> {
self.write_continuation(write_options, 0)
}
}
/// Blanket-implemented, so no impl of `IpcMessageSink` can ever override these.
impl<T: IpcMessageSink + ?Sized> IpcMessageSinkExt for T {}Introduce StreamEncoder for IPC streaming without requiring a std::io::Write sink. The encoder owns stream state, emits ordered Buffer chunks, and preserves the low-copy path for uncompressed record batch body buffers. Add byte-for-byte compatibility tests against StreamWriter for normal batches, empty streams, and dictionary batches. Closes apache#7812 Signed-off-by: Jiawei Zhao <Phoenix500526@163.com>
Exercise StreamEncoder with a Tokio AsyncWrite flow and compare the bytes with StreamWriter. This keeps the PR tied to the async writer use case from the original issue. Refs apache#7812 Signed-off-by: Jiawei Zhao <Phoenix500526@163.com>
Route writer and buffer output through a shared internal sink. This avoids duplicated IPC framing while preserving Buffer output. Signed-off-by: Jiawei Zhao <Phoenix500526@163.com>
Signed-off-by: Jiawei Zhao <Phoenix500526@163.com>
Avoid erasing W: Write on the writer path while preserving shared IPC framing for buffer output. Refs apache#7812 Signed-off-by: Jiawei Zhao <Phoenix500526@163.com>
Implement the sink trait for Write directly so the writer path no longer needs a wrapper sink. Consume StreamEncoder in finish to remove the closed-state check. Refs apache#7812 Signed-off-by: Jiawei Zhao <Phoenix500526@163.com>
Extract the IPC metadata padding calculation so writer and buffer sinks use the same framing layout. Signed-off-by: Jiawei Zhao <Phoenix500526@163.com>
Separate sink-specific write primitives from shared IPC framing helpers while preserving the record batch writer specialization. Signed-off-by: Jiawei Zhao <Phoenix500526@163.com>
66a7cb1 to
ee44c9c
Compare
alamb
left a comment
There was a problem hiding this comment.
Looks good -- thank you @Phoenix500526
|
Thanks also @Jefffrey and @Rich-T-kid for the review |
Which issue does this PR close?
Rationale for this change
StreamWriter currently requires a std::io::Write sink, which is awkward for async or chunk-oriented destinations such as object stores. This PR adds a sans-IO IPC stream encoder so callers can encode Arrow IPC stream data into ordered Buffer chunks and send those chunks through their own IO layer.
What changes are included in this PR?
This PR adds StreamEncoder, a stateful IPC stream encoder that:
Are these changes tested?
Yes.
Added tests compare StreamEncoder output byte-for-byte with StreamWriter output for:
a normal record batch stream
an empty stream
a stream containing dictionary batches
Are there any user-facing changes?
Yes. This adds a new public arrow_ipc::writer::StreamEncoder API.
There are no breaking changes. Existing StreamWriter behavior is unchanged.