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
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 🚀
| /// # } | ||
| /// ``` | ||
| /// | ||
| /// [IPC Streaming Format]: https://arrow.apache.org/docs/format/Columnar.html#ipc-streaming-format |
There was a problem hiding this comment.
we can remove the comment/link to the IPC format
There was a problem hiding this comment.
Thanks for the review. The comment is removed in commit 6bea182
|
|
||
| // StreamEncoder and StreamWriter currently use separate encoding paths, so these tests | ||
| // verify the new sans-IO API preserves the existing IPC stream byte layout. | ||
| #[test] |
There was a problem hiding this comment.
Not strictly necessary, but since the originating issue for this PR wanted to use StreamWriter in an async context, could you add a test that mirrors the flow of an async writer (tokio is already a dev-dependency here) and then assert the final bytes against a sync StreamWriter?
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
| assert_eq!(reader.next().unwrap().unwrap(), batch); | ||
| assert!(reader.next().is_none()); |
| let capacity = batch | ||
| .columns() | ||
| .iter() | ||
| .map(|a| estimate_encoded_buffer_count(a.data_type())) | ||
| .sum(); | ||
| let mut encoded_buffers: Vec<EncodedBuffer> = Vec::with_capacity(capacity); | ||
| let (ipc_message, body_len, tail_pad) = self.record_batch_to_bytes( | ||
| batch, | ||
| write_options, | ||
| ipc_write_context, | ||
| &mut IpcBodySink::Collect(&mut encoded_buffers), | ||
| )?; | ||
|
|
||
| let alignment = write_options.alignment; | ||
| let alignment_mask = usize::from(alignment - 1); | ||
| let prefix_size = if write_options.write_legacy_ipc_format { | ||
| 4 | ||
| } else { | ||
| 8 | ||
| }; | ||
| let ipc_message_len = ipc_message.len(); | ||
| let aligned_size = (ipc_message_len + prefix_size + alignment_mask) & !alignment_mask; | ||
| push_continuation_buffer(out, write_options, (aligned_size - prefix_size) as i32)?; | ||
| out.push(Buffer::from(ipc_message)); | ||
| push_padding_buffer(out, aligned_size - ipc_message_len - prefix_size); | ||
| for enc in encoded_buffers { | ||
| let len = enc.len(); | ||
| let buffer = match enc { | ||
| EncodedBuffer::Raw(buffer) => buffer, | ||
| EncodedBuffer::Compressed(bytes) => Buffer::from(bytes), | ||
| }; | ||
| out.push(buffer); | ||
| push_padding_buffer(out, pad_to_alignment(alignment, len)); | ||
| } |
There was a problem hiding this comment.
I think the biggest issue I see with this PR is logic duplication. This is very similar to the existing approach.
With the IpcBodySink enum I think we can have a strong shared interface than this.
There was a problem hiding this comment.
Good catch. I refactored this to use a shared internal IpcMessageSink for the final framed IPC message emission. I kep IpcBodySink focused on materializing the record batch body buffers, and added IpcMessageSink one layer above it to handle where the complete framed IPC message goes: either a Write sink or ordered Buffers.
With that, IpcDataGenerator::write and encode_to_buffers now share the same write_to_sink path instead of duplicating dictionary/message framing logic.
| self.data_gen.encode_to_buffers( | ||
| batch, | ||
| &mut self.dictionary_tracker, | ||
| &self.write_options, | ||
| &mut self.ipc_write_context, | ||
| &mut out, | ||
| )?; | ||
| Ok(out) |
There was a problem hiding this comment.
Maybe instead of a new encode_to_buffers method we could reuse IpcDataGenerator::write? It already uses the IpcBodySink::collect() variant. it does mostly the same thing.
There was a problem hiding this comment.
Good point. I agree the two paths should share more of the implementation.
I looked at reusing IpcDataGenerator::write directly, but I don’t think that is quite the right boundary. write commits the output to a Write sink, while StreamEncoder needs to return ordered Buffers. Adapting the buffer path through Write would flatten the output and lose the buffer-preserving behavior this API is trying to expose.
I refactored this by adding an internal IpcMessageSink for the final framed-message emission. IpcBodySink still handles record-batch body materialization, and IpcMessageSink handles emitting the complete framed IPC message to either Write or Vec<Buffer>. This lets write and encode_to_buffers share the same pipeline without forcing the buffer API through Write.
| pub arrow_data: Vec<u8>, | ||
| } | ||
|
|
||
| fn encoded_data_to_buffers( |
There was a problem hiding this comment.
can you add a doc comment to these functions
| stream_reader.next().unwrap().unwrap() | ||
| } | ||
|
|
||
| fn encode_stream( |
There was a problem hiding this comment.
can you add a doc comment. maybe provide a small example of when to use this?
| assert_eq!(reader.next().unwrap().unwrap(), batch); | ||
| assert!(reader.next().is_none()); | ||
| } | ||
|
|
There was a problem hiding this comment.
I think the test here look fine
09d05ce to
42a7696
Compare
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>
42a7696 to
3f24353
Compare
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>
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
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.
| let root = message.finish(); | ||
| fbb.finish(root, None); | ||
|
|
||
| let data = fbb.finished_data(); | ||
| let metadata = fbb.finished_data(); | ||
| EncodedData { | ||
| ipc_message: data.to_vec(), | ||
| ipc_message: metadata.to_vec(), |
There was a problem hiding this comment.
nice, renamng data to metadata is a nice change
| ) -> Result<IpcWriteMetadata, ArrowError> { | ||
| let mut sink = IpcMessageSink::Buffers(out); | ||
| self.write_to_sink( | ||
| batch, | ||
| dictionary_tracker, | ||
| write_options, | ||
| ipc_write_context, | ||
| &mut sink, | ||
| ) | ||
| } |
There was a problem hiding this comment.
nice, this is a lot cleaner, and removes duplication
| fn encoded_data_to_buffers( | ||
| out: &mut Vec<Buffer>, | ||
| encoded: EncodedData, | ||
| write_options: &IpcWriteOptions, | ||
| ) -> Result<(usize, usize), ArrowError> { | ||
| let arrow_data_len = encoded.arrow_data.len(); | ||
| if arrow_data_len % usize::from(write_options.alignment) != 0 { | ||
| return Err(ArrowError::MemoryError( | ||
| "Arrow data not aligned".to_string(), | ||
| )); | ||
| } | ||
|
|
||
| let alignment_mask = usize::from(write_options.alignment - 1); | ||
| let flatbuf_size = encoded.ipc_message.len(); | ||
| let prefix_size = if write_options.write_legacy_ipc_format { | ||
| 4 | ||
| } else { | ||
| 8 | ||
| }; | ||
| let aligned_size = (flatbuf_size + prefix_size + alignment_mask) & !alignment_mask; | ||
| let padding_bytes = aligned_size - flatbuf_size - prefix_size; | ||
|
|
||
| push_continuation_buffer(out, write_options, (aligned_size - prefix_size) as i32)?; | ||
| if flatbuf_size > 0 { | ||
| out.push(Buffer::from(encoded.ipc_message)); | ||
| } | ||
| push_padding_buffer(out, padding_bytes); | ||
|
|
||
| let body_len = if arrow_data_len > 0 { | ||
| out.push(Buffer::from(encoded.arrow_data)); | ||
| arrow_data_len | ||
| } else { | ||
| 0 | ||
| }; | ||
|
|
||
| Ok((aligned_size, body_len)) | ||
| } | ||
|
|
||
| fn push_continuation_buffer( | ||
| out: &mut Vec<Buffer>, | ||
| write_options: &IpcWriteOptions, | ||
| total_len: i32, | ||
| ) -> Result<(), ArrowError> { | ||
| // Continuation bytes are generated stream framing, so they need a small owned buffer. | ||
| let mut buffer = Vec::with_capacity(8); | ||
| write_continuation(&mut buffer, write_options, total_len)?; | ||
| out.push(Buffer::from(buffer)); | ||
| Ok(()) | ||
| } | ||
|
|
||
| fn push_padding_buffer(out: &mut Vec<Buffer>, len: usize) { | ||
| if len > 0 { | ||
| out.push(Buffer::from(&PADDING[..len])); | ||
| } | ||
| } | ||
|
|
| let arrow_data_len = encoded.arrow_data.len(); | ||
| if arrow_data_len % usize::from(write_options.alignment) != 0 { | ||
| return Err(ArrowError::MemoryError( | ||
| "Arrow data not aligned".to_string(), | ||
| )); | ||
| } | ||
|
|
||
| let a = usize::from(write_options.alignment - 1); | ||
| let buffer = encoded.ipc_message; | ||
| let flatbuf_size = buffer.len(); | ||
| let prefix_size = if write_options.write_legacy_ipc_format { | ||
| 4 | ||
| } else { | ||
| 8 | ||
| }; | ||
| let aligned_size = (flatbuf_size + prefix_size + a) & !a; | ||
| let padding_bytes = aligned_size - flatbuf_size - prefix_size; | ||
|
|
||
| write_continuation( | ||
| &mut writer, | ||
| write_options, | ||
| (aligned_size - prefix_size) as i32, | ||
| )?; | ||
|
|
||
| // write the flatbuf | ||
| if flatbuf_size > 0 { | ||
| writer.write_all(&buffer)?; | ||
| } | ||
| // write padding | ||
| writer.write_all(&PADDING[..padding_bytes])?; | ||
|
|
||
| // write arrow data | ||
| let body_len = if arrow_data_len > 0 { | ||
| write_body_buffers(&mut writer, &encoded.arrow_data, write_options.alignment)? | ||
| } else { | ||
| 0 | ||
| }; | ||
|
|
||
| Ok((aligned_size, body_len)) | ||
| } | ||
|
|
||
| fn write_body_buffers<W: Write>( | ||
| mut writer: W, | ||
| data: &[u8], | ||
| alignment: u8, | ||
| ) -> Result<usize, ArrowError> { | ||
| let len = data.len(); | ||
| let pad_len = pad_to_alignment(alignment, len); | ||
| let total_len = len + pad_len; | ||
|
|
||
| // write body buffer | ||
| writer.write_all(data)?; | ||
| if pad_len > 0 { | ||
| writer.write_all(&PADDING[..pad_len])?; | ||
| } | ||
|
|
||
| Ok(total_len) | ||
| } | ||
|
|
||
| /// Write a record batch to the writer, writing the message size before the message | ||
| /// if the record batch is being written to a stream | ||
| fn write_continuation<W: Write>( | ||
| mut writer: W, | ||
| write_options: &IpcWriteOptions, | ||
| total_len: i32, | ||
| ) -> Result<usize, ArrowError> { | ||
| let mut written = 8; | ||
|
|
||
| // the version of the writer determines whether continuation markers should be added | ||
| match write_options.metadata_version { | ||
| crate::MetadataVersion::V1 | crate::MetadataVersion::V2 | crate::MetadataVersion::V3 => { | ||
| unreachable!("Options with the metadata version cannot be created") | ||
| } | ||
| crate::MetadataVersion::V4 => { | ||
| if !write_options.write_legacy_ipc_format { | ||
| // v0.15.0 format | ||
| writer.write_all(&CONTINUATION_MARKER)?; | ||
| written = 4; | ||
| } | ||
| writer.write_all(&total_len.to_le_bytes()[..])?; | ||
| } | ||
| crate::MetadataVersion::V5 => { | ||
| // write continuation marker and message length | ||
| writer.write_all(&CONTINUATION_MARKER)?; | ||
| writer.write_all(&total_len.to_le_bytes()[..])?; | ||
| } | ||
| z => panic!("Unsupported crate::MetadataVersion {z:?}"), | ||
| }; | ||
|
|
||
| Ok(written) |
| fn write_vec(&mut self, bytes: Vec<u8>) -> Result<(), ArrowError> { | ||
| if bytes.is_empty() { | ||
| return Ok(()); | ||
| } | ||
|
|
||
| match self { | ||
| Self::Writer(writer) => writer.write_all(&bytes)?, | ||
| Self::Buffers(out) => out.push(Buffer::from(bytes)), | ||
| } | ||
| Ok(()) | ||
| } |
There was a problem hiding this comment.
nice now callers dont need to handle this at each call site
| Ok(()) | ||
| } | ||
|
|
||
| fn write_encoded_buffer(&mut self, buffer: EncodedBuffer) -> Result<(), ArrowError> { |
There was a problem hiding this comment.
I think a doc comment above this would be nice, otherwise it makes sense to me.
There was a problem hiding this comment.
I added some doc comments to those methods that are mainly called from outside IpcMessageSink.
FYI：5cfcf8a
|
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 👍 |
Signed-off-by: Jiawei Zhao <Phoenix500526@163.com>
|
run benchmark ipc_writer |
|
🤖 Arrow criterion benchmark running (GKE) | trigger CPU Details (lscpu)Comparing issue/7812 (5cfcf8a) to 980667f (merge-base) diff File an issue against this benchmark runner |
|
🤖 Arrow criterion benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
looks to be mostly noise? its within +- 10. can we run it again to check if its reproducible |
Rich-T-kid
left a comment
Rich-T-kid
left a comment
There was a problem hiding this comment.
LGTM @Phoenix500526
@Jefffrey I think this is ready for a review
| criterion = { workspace = true } | ||
| tempfile = "3.3" | ||
| tokio = "1.43.0" | ||
| tokio = { version = "1.43.0", features = ["io-util", "macros", "rt"] } |
There was a problem hiding this comment.
Does this count as increasing our dependencies? Tokio was already imported, so I think this is fine
| } | ||
| } | ||
|
|
||
| /// Destination for a complete framed IPC message. |
There was a problem hiding this comment.
we should also add benchmarks for the stream encoder to catch regressions in the future.
There was a problem hiding this comment.
this should be a follow up so that its easier to review. I created this issue to track it #10373
4 participants
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.