Rich-T-kid commented on code in PR #10277:
URL: https://github.com/apache/arrow-rs/pull/10277#discussion_r3596246137
##########
arrow-ipc/src/writer.rs:
##########
@@ -725,6 +724,67 @@ impl IpcDataGenerator {
})
}
+ /// Encode dictionary batches and the record batch to output buffers,
skipping the
+ /// intermediate body `Vec<u8>` allocations for uncompressed record batch
buffers.
+ fn encode_to_buffers(
+ &self,
+ batch: &RecordBatch,
+ dictionary_tracker: &mut DictionaryTracker,
+ write_options: &IpcWriteOptions,
+ ipc_write_context: &mut IpcWriteContext,
+ out: &mut Vec<Buffer>,
+ ) -> Result<IpcWriteMetadata, ArrowError> {
+ let encoded_dictionaries =
+ self.encode_all_dicts(batch, dictionary_tracker, write_options,
ipc_write_context)?;
+
+ let mut dictionary_block_sizes =
Vec::with_capacity(encoded_dictionaries.len());
+ for dict in encoded_dictionaries {
+ dictionary_block_sizes.push(encoded_data_to_buffers(out, dict,
write_options)?);
+ }
+
+ 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));
+ }
Review 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.
##########
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:
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.
##########
arrow-ipc/src/writer.rs:
##########
@@ -2523,6 +2771,111 @@ mod tests {
stream_reader.next().unwrap().unwrap()
}
+ fn encode_stream(
+ schema: &Schema,
+ batches: &[RecordBatch],
+ options: IpcWriteOptions,
+ ) -> Vec<u8> {
+ let mut encoder = StreamEncoder::try_new_with_options(schema,
options).unwrap();
+ let mut bytes = Vec::new();
+ for batch in batches {
+ for buffer in encoder.encode(batch).unwrap() {
+ bytes.extend_from_slice(buffer.as_slice());
+ }
+ }
+ for buffer in encoder.finish().unwrap() {
+ bytes.extend_from_slice(buffer.as_slice());
+ }
+ bytes
+ }
+
+ fn write_stream(schema: &Schema, batches: &[RecordBatch], options:
IpcWriteOptions) -> Vec<u8> {
+ let mut bytes = Vec::new();
+ let mut writer = StreamWriter::try_new_with_options(&mut bytes,
schema, options).unwrap();
+ for batch in batches {
+ writer.write(batch).unwrap();
+ }
+ writer.finish().unwrap();
+ bytes
+ }
+
+ // StreamEncoder and StreamWriter currently use separate encoding paths,
so these tests
+ // verify the new sans-IO API preserves the existing IPC stream byte
layout.
+ #[tokio::test]
+ async fn test_stream_encoder_async_writer_matches_stream_writer() {
+ use tokio::io::{AsyncReadExt, AsyncWriteExt};
+
+ let batch = record_batch!(("a", Int32, [1, 2, 3]), ("b", Utf8, ["x",
"y", "z"])).unwrap();
+ let options = IpcWriteOptions::default();
+ let expected = write_stream(
+ batch.schema_ref(),
+ std::slice::from_ref(&batch),
+ options.clone(),
+ );
+
+ let (mut sink, mut source) = tokio::io::duplex(64);
+ let read = tokio::spawn(async move {
+ let mut bytes = Vec::new();
+ source.read_to_end(&mut bytes).await.unwrap();
+ bytes
+ });
+
+ let mut encoder =
StreamEncoder::try_new_with_options(batch.schema_ref(), options).unwrap();
+ for buffer in encoder.encode(&batch).unwrap() {
+ sink.write_all(buffer.as_slice()).await.unwrap();
+ }
+ for buffer in encoder.finish().unwrap() {
+ sink.write_all(buffer.as_slice()).await.unwrap();
+ }
+ sink.shutdown().await.unwrap();
+
+ let encoded = read.await.unwrap();
+ assert_eq!(encoded, expected);
+
+ let mut reader = StreamReader::try_new(Cursor::new(encoded),
None).unwrap();
+ assert_eq!(reader.next().unwrap().unwrap(), batch);
+ assert!(reader.next().is_none());
+ }
+
+ #[test]
+ fn test_stream_encoder_empty_stream_matches_stream_writer() {
+ let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
+ let options = IpcWriteOptions::default();
+ let encoded = encode_stream(&schema, &[], options.clone());
+ let written = write_stream(&schema, &[], options);
+
+ assert_eq!(encoded, written);
+
+ let mut reader = StreamReader::try_new(Cursor::new(encoded),
None).unwrap();
+ assert!(reader.next().is_none());
+ }
+
+ #[test]
+ fn test_stream_encoder_dictionary_batches_match_stream_writer() {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "a",
+ DataType::Dictionary(Box::new(DataType::UInt8),
Box::new(DataType::Utf8)),
+ false,
+ )]));
+ let batch = RecordBatch::try_new(
+ schema.clone(),
+ vec![Arc::new(DictionaryArray::new(
+ UInt8Array::from_iter_values([0, 1, 0]),
+ Arc::new(StringArray::from_iter_values(["a", "b"])),
+ ))],
+ )
+ .unwrap();
+ let options = IpcWriteOptions::default();
+ let encoded = encode_stream(&schema, std::slice::from_ref(&batch),
options.clone());
+ let written = write_stream(&schema, std::slice::from_ref(&batch),
options);
+
+ assert_eq!(encoded, written);
+
+ let mut reader = StreamReader::try_new(Cursor::new(encoded),
None).unwrap();
+ assert_eq!(reader.next().unwrap().unwrap(), batch);
+ assert!(reader.next().is_none());
+ }
+
Review Comment:
I think the test here look fine
##########
arrow-ipc/src/writer.rs:
##########
@@ -2523,6 +2771,111 @@ mod tests {
stream_reader.next().unwrap().unwrap()
}
+ fn encode_stream(
+ schema: &Schema,
+ batches: &[RecordBatch],
+ options: IpcWriteOptions,
+ ) -> Vec<u8> {
+ let mut encoder = StreamEncoder::try_new_with_options(schema,
options).unwrap();
+ let mut bytes = Vec::new();
+ for batch in batches {
+ for buffer in encoder.encode(batch).unwrap() {
+ bytes.extend_from_slice(buffer.as_slice());
+ }
+ }
+ for buffer in encoder.finish().unwrap() {
+ bytes.extend_from_slice(buffer.as_slice());
+ }
+ bytes
+ }
+
+ fn write_stream(schema: &Schema, batches: &[RecordBatch], options:
IpcWriteOptions) -> Vec<u8> {
+ let mut bytes = Vec::new();
+ let mut writer = StreamWriter::try_new_with_options(&mut bytes,
schema, options).unwrap();
+ for batch in batches {
+ writer.write(batch).unwrap();
+ }
+ writer.finish().unwrap();
+ bytes
+ }
+
+ // StreamEncoder and StreamWriter currently use separate encoding paths,
so these tests
+ // verify the new sans-IO API preserves the existing IPC stream byte
layout.
+ #[tokio::test]
+ async fn test_stream_encoder_async_writer_matches_stream_writer() {
+ use tokio::io::{AsyncReadExt, AsyncWriteExt};
+
+ let batch = record_batch!(("a", Int32, [1, 2, 3]), ("b", Utf8, ["x",
"y", "z"])).unwrap();
+ let options = IpcWriteOptions::default();
+ let expected = write_stream(
+ batch.schema_ref(),
+ std::slice::from_ref(&batch),
+ options.clone(),
+ );
+
+ let (mut sink, mut source) = tokio::io::duplex(64);
+ let read = tokio::spawn(async move {
+ let mut bytes = Vec::new();
+ source.read_to_end(&mut bytes).await.unwrap();
+ bytes
+ });
+
+ let mut encoder =
StreamEncoder::try_new_with_options(batch.schema_ref(), options).unwrap();
+ for buffer in encoder.encode(&batch).unwrap() {
+ sink.write_all(buffer.as_slice()).await.unwrap();
+ }
+ for buffer in encoder.finish().unwrap() {
+ sink.write_all(buffer.as_slice()).await.unwrap();
+ }
+ sink.shutdown().await.unwrap();
+
+ let encoded = read.await.unwrap();
+ assert_eq!(encoded, expected);
+
+ let mut reader = StreamReader::try_new(Cursor::new(encoded),
None).unwrap();
+ assert_eq!(reader.next().unwrap().unwrap(), batch);
+ assert!(reader.next().is_none());
Review Comment:
nice
##########
arrow-ipc/src/writer.rs:
##########
@@ -1823,6 +2014,63 @@ pub struct EncodedData {
/// Arrow buffers to be written, should be an empty vec for schema messages
pub arrow_data: Vec<u8>,
}
+
+fn encoded_data_to_buffers(
Review Comment:
can you add a doc comment to these functions
##########
arrow-ipc/src/writer.rs:
##########
@@ -2523,6 +2771,111 @@ mod tests {
stream_reader.next().unwrap().unwrap()
}
+ fn encode_stream(
Review Comment:
can you add a doc comment. maybe provide a small example of when to use this?
--
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]
