comphead commented on code in PR #5006:
URL: https://github.com/apache/datafusion-comet/pull/5006#discussion_r3640256925
##########
native/shuffle/src/writers/shuffle_block_writer.rs:
##########
@@ -32,42 +40,119 @@ pub enum CompressionCodec {
Snappy,
}
+/// Returns true if `data_type` is, or nests, a dictionary type.
+fn contains_dictionary(data_type: &DataType) -> bool {
+ match data_type {
+ DataType::Dictionary(_, _) => true,
+ DataType::List(f)
+ | DataType::LargeList(f)
+ | DataType::FixedSizeList(f, _)
+ | DataType::Map(f, _)
+ | DataType::RunEndEncoded(_, f) => contains_dictionary(f.data_type()),
+ DataType::Struct(fields) => fields.iter().any(|f|
contains_dictionary(f.data_type())),
+ DataType::Union(fields, _) => fields
+ .iter()
+ .any(|(_, f)| contains_dictionary(f.data_type())),
+ _ => false,
+ }
+}
+
/// Writes a record batch as a length-prefixed, compressed Arrow IPC block.
+///
+/// Each block is a self-contained Arrow IPC stream (schema message,
dictionary messages, record
+/// batch message, end-of-stream marker). For the common case of a schema with
no dictionary types,
+/// the schema flatbuffer is encoded once in [`Self::try_new`] and written
verbatim at the start of
+/// every block, rather than being re-serialized per block as
`StreamWriter::try_new` would do.
+/// Schemas that contain dictionary types fall back to `StreamWriter`, whose
dictionary-id
+/// bookkeeping ties schema and batch encoding together.
#[derive(Clone)]
pub struct ShuffleBlockWriter {
codec: CompressionCodec,
header_bytes: Vec<u8>,
+ schema: SchemaRef,
+ /// Pre-encoded Arrow IPC schema message, written at the start of every
block. Only used when
+ /// the schema has no dictionary types.
+ schema_message: Vec<u8>,
+ /// Whether the schema contains any dictionary types (see
[`Self::encode_ipc_stream`]).
+ has_dictionaries: bool,
}
impl ShuffleBlockWriter {
pub fn try_new(schema: &Schema, codec: CompressionCodec) -> Result<Self> {
- let header_bytes = Vec::with_capacity(20);
- let mut cursor = Cursor::new(header_bytes);
+ let mut header_bytes = Vec::with_capacity(20);
- // leave space for compressed message length
- cursor.seek_relative(8)?;
+ // leave space for compressed message length (filled in per block by
write_batch)
+ header_bytes.extend_from_slice(&[0u8; 8]);
// write number of columns because JVM side needs to know how many
addresses to allocate
let field_count = schema.fields().len();
- cursor.write_all(&field_count.to_le_bytes())?;
+ header_bytes.extend_from_slice(&field_count.to_le_bytes());
// write compression codec to header
- let codec_header = match &codec {
+ let codec_header: &[u8] = match &codec {
CompressionCodec::Snappy => b"SNAP",
CompressionCodec::Lz4Frame => b"LZ4_",
CompressionCodec::Zstd(_) => b"ZSTD",
CompressionCodec::None => b"NONE",
};
- cursor.write_all(codec_header)?;
-
- let header_bytes = cursor.into_inner();
+ header_bytes.extend_from_slice(codec_header);
+
+ // Pre-encode the IPC schema message once so it does not have to be
re-serialized per block.
+ let options = IpcWriteOptions::default();
+ let data_gen = IpcDataGenerator::default();
+ let mut dictionary_tracker = DictionaryTracker::new(true);
+ let encoded_schema = data_gen.schema_to_bytes_with_dictionary_tracker(
+ schema,
+ &mut dictionary_tracker,
+ &options,
+ );
+ let mut schema_message = Vec::new();
+ write_message(&mut schema_message, encoded_schema, &options)?;
+
+ let has_dictionaries = schema
+ .fields()
+ .iter()
+ .any(|f| contains_dictionary(f.data_type()));
Ok(Self {
codec,
header_bytes,
+ schema: Arc::new(schema.clone()),
+ schema_message,
+ has_dictionaries,
})
}
+ /// Serialize `batch` as a standalone Arrow IPC stream into `out`.
+ fn encode_ipc_stream<W: Write>(&self, batch: &RecordBatch, out: &mut W) ->
Result<()> {
+ if self.has_dictionaries {
+ // Dictionary encoding requires the schema and record batch to
share a dictionary
+ // tracker, so `StreamWriter` (which re-encodes the schema per
block) is used here.
+ let mut stream_writer = StreamWriter::try_new(out, &self.schema)?;
+ stream_writer.write(batch)?;
+ stream_writer.finish()?;
+ return Ok(());
+ }
+
+ // Fast path: reuse the pre-encoded schema message and write the
record batch manually.
+ let options = IpcWriteOptions::default();
+ let data_gen = IpcDataGenerator::default();
Review Comment:
would it make sense to cache and reuse these, instead of per batch
allocation?
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
