comphead commented on code in PR #5006:
URL: https://github.com/apache/datafusion-comet/pull/5006#discussion_r3640303559
##########
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);
Review Comment:
we may want to comment why 20 is here
--
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]
