Phoenix500526 commented on code in PR #10277:
URL: https://github.com/apache/arrow-rs/pull/10277#discussion_r3601635412
##########
arrow-ipc/src/writer.rs:
##########
@@ -1564,6 +1624,137 @@ impl<W: Write> RecordBatchWriter for FileWriter<W> {
}
}
+/// Arrow IPC stream encoder.
+///
+/// Encodes Arrow [`RecordBatch`]es to byte buffers using the [IPC Streaming
Format],
+/// without performing any IO.
+///
+/// The returned [`Buffer`]s are ordered and should be written to the
destination
+/// stream in order. Uncompressed record batch body buffers can share the
original
+/// Arrow buffers instead of being copied into an intermediate contiguous
buffer.
+///
+/// # Example
+/// ```
+/// # use arrow_array::record_batch;
+/// # use arrow_ipc::writer::StreamEncoder;
+/// # use arrow_schema::ArrowError;
+/// # fn main() -> Result<(), ArrowError> {
+/// let batch = record_batch!(("a", Int32, [1, 2, 3]))?;
+///
+/// let mut encoder = StreamEncoder::try_new(&batch.schema())?;
+/// let mut stream = vec![];
+/// for buffer in encoder.encode(&batch)? {
+/// stream.extend_from_slice(buffer.as_slice());
+/// }
+/// for buffer in encoder.finish()? {
+/// stream.extend_from_slice(buffer.as_slice());
+/// }
+/// # Ok(())
+/// # }
+/// ```
+pub struct StreamEncoder {
+ schema: Schema,
+ /// IPC write options
+ write_options: IpcWriteOptions,
+ /// Whether the stream schema has been encoded
+ schema_encoded: bool,
+ /// Whether the end-of-stream marker has been encoded
+ finished: bool,
+ /// Keeps track of dictionaries that have been encoded
+ dictionary_tracker: DictionaryTracker,
+ data_gen: IpcDataGenerator,
+ ipc_write_context: IpcWriteContext,
+}
+
+impl StreamEncoder {
+ /// Try to create a new stream encoder.
+ pub fn try_new(schema: &Schema) -> Result<Self, ArrowError> {
+ let write_options = IpcWriteOptions::default();
+ Self::try_new_with_options(schema, write_options)
+ }
+
+ /// Try to create a new stream encoder with [`IpcWriteOptions`].
+ pub fn try_new_with_options(
+ schema: &Schema,
+ write_options: IpcWriteOptions,
+ ) -> Result<Self, ArrowError> {
+ ensure_supported_ipc_schema(schema)?;
+
+ Ok(Self {
+ schema: schema.clone(),
+ write_options,
+ schema_encoded: false,
+ finished: false,
+ dictionary_tracker: DictionaryTracker::new(false),
+ data_gen: IpcDataGenerator::default(),
+ ipc_write_context: IpcWriteContext::default(),
+ })
+ }
+
+ /// Encode a [`RecordBatch`] into buffers.
+ ///
+ /// The first call also includes the IPC stream schema message before the
+ /// record batch message. Later calls only include dictionary and record
+ /// batch messages.
+ ///
+ /// # Errors
+ ///
+ /// Returns an error if the encoder is already finished or encoding fails.
+ pub fn encode(&mut self, batch: &RecordBatch) -> Result<Vec<Buffer>,
ArrowError> {
+ if self.finished {
+ return Err(ArrowError::IpcError(
+ "Cannot encode record batch to stream encoder as it is
closed".to_string(),
+ ));
+ }
+
+ let mut out = vec![];
+ self.encode_schema(&mut out)?;
+ self.data_gen.encode_to_buffers(
+ batch,
+ &mut self.dictionary_tracker,
+ &self.write_options,
+ &mut self.ipc_write_context,
+ &mut out,
+ )?;
+ Ok(out)
Review 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 `Buffer`s. 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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
