andygrove commented on code in PR #5006:
URL: https://github.com/apache/datafusion-comet/pull/5006#discussion_r3641321996
##########
native/shuffle/src/writers/shuffle_block_writer.rs:
##########
@@ -32,42 +40,128 @@ pub enum CompressionCodec {
Snappy,
}
+/// Returns true if `data_type` is, or nests, a dictionary type.
+fn contains_dictionary(data_type: &DataType) -> bool {
Review Comment:
Done in ca61876. `contains_dictionary` is deleted; detection is now:
```rust
let has_dictionaries = schema
.flattened_fields()
.iter()
.any(|f| matches!(f.data_type(), DataType::Dictionary(_, _)));
```
Confirmed `flattened_fields` is present in the arrow-schema we pin (58.3.0,
`schema.rs:364`). This also closes the `ListView`/`LargeListView` gap you
pointed at, so the `encoded_dictionaries.is_empty()` debug_assert no longer
depends on our own recursion staying in sync with arrow's dictionary traversal
across version bumps.
##########
native/shuffle/src/writers/shuffle_block_writer.rs:
##########
@@ -32,42 +40,128 @@ 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,
Review Comment:
Done in ca61876. Both fields are replaced by:
```rust
enum SchemaEncoding {
Precoded(Vec<u8>), // dictionary-free: written verbatim per block
Fallback(SchemaRef), // dictionary schemas: StreamWriter re-encodes
}
```
The `Arc::new(schema.clone())` now happens only in the `Fallback` arm, so
the dictionary-free path no longer deep-clones a schema it never reads.
One correction to the detail: `try_new` takes `&Schema`, not a `SchemaRef`,
so the fallback arm is still a deep clone rather than an `Arc::clone`. It is
now paid only on the rare dictionary path, and once per writer rather than per
block, so I left the signature alone.
##########
native/shuffle/src/writers/shuffle_block_writer.rs:
##########
@@ -32,42 +40,128 @@ 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 verbatim at the start of
every block.
+ ///
+ /// `None` indicates the schema contains dictionary types, whose
dictionary-id bookkeeping ties
+ /// schema and batch encoding together, so the schema cannot be reused
across blocks and
+ /// [`Self::encode_ipc_stream`] falls back to `StreamWriter`.
+ schema_message: Option<Vec<u8>>,
}
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);
+ // Header layout: 8-byte block length placeholder + 8-byte field count
(usize) + 4-byte
+ // codec tag = 20 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);
+
+ let has_dictionaries = schema
+ .fields()
+ .iter()
+ .any(|f| contains_dictionary(f.data_type()));
+
+ // For dictionary-free schemas, pre-encode the IPC schema message once
so it does not have
+ // to be re-serialized per block. Dictionary schemas use the
`StreamWriter` fallback and
+ // leave this `None`.
+ let schema_message = if has_dictionaries {
+ None
+ } else {
+ 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 buf = Vec::new();
+ write_message(&mut buf, encoded_schema, &options)?;
+ Some(buf)
+ };
Ok(Self {
codec,
header_bytes,
+ schema: Arc::new(schema.clone()),
+ schema_message,
})
}
+ /// Serialize `batch` as a standalone Arrow IPC stream into `out`.
+ fn encode_ipc_stream<W: Write>(&self, batch: &RecordBatch, out: &mut W) ->
Result<()> {
+ let Some(schema_message) = &self.schema_message else {
+ // 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();
Review Comment:
Done in ca61876. A single `write_options: IpcWriteOptions` is stored on the
writer and used by all three encode sites: the schema encode in `try_new`, the
batch encode in `encode_ipc_stream`, and the dictionary fallback (switched from
`StreamWriter::try_new` to `try_new_with_options`). There is no longer a second
options value that could drift from the first.
--
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]
