andygrove commented on code in PR #5006:
URL: https://github.com/apache/datafusion-comet/pull/5006#discussion_r3645280019
##########
native/shuffle/src/writers/shuffle_block_writer.rs:
##########
@@ -32,42 +44,129 @@ pub enum CompressionCodec {
Snappy,
}
+/// How a writer encodes the Arrow IPC schema at the start of each block. The
two arms are mutually
+/// exclusive by schema shape, decided once in
[`ShuffleBlockWriter::try_new`], so the retained
+/// state never carries both a pre-encoded message and a schema at the same
time.
+#[derive(Clone)]
+enum SchemaEncoding {
+ /// Dictionary-free schema: the IPC schema message, pre-encoded once, is
written verbatim at the
+ /// start of every block instead of being re-serialized per block.
+ Precoded(Vec<u8>),
+ /// Schema containing dictionary types: the schema and record batch must
share a dictionary
+ /// tracker, so each block is encoded with `StreamWriter`, which
re-serializes the schema.
+ Fallback(SchemaRef),
+}
+
/// 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>,
+ /// IPC options shared by the schema and record-batch encoders so the two
can never diverge;
+ /// pinned to metadata version V5, which [`IPC_EOS`] depends on.
+ write_options: IpcWriteOptions,
+ schema_encoding: SchemaEncoding,
}
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);
+
+ // Shuffle blocks are always written with metadata version V5. Pin it
explicitly rather than
+ // relying on `IpcWriteOptions::default`, because IPC_EOS is only the
correct end-of-stream
+ // marker for V5. Alignment 64 and non-legacy framing match the arrow
defaults.
+ let write_options = IpcWriteOptions::try_new(64, false,
MetadataVersion::V5)?;
+
+ // `flattened_fields` walks the full nested field tree, so this
catches dictionary types
+ // nested under any container arrow itself recurses into when emitting
dictionaries.
+ let has_dictionaries = schema
+ .flattened_fields()
+ .iter()
+ .any(|f| matches!(f.data_type(), DataType::Dictionary(_, _)));
+
+ // 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.
+ let schema_encoding = if has_dictionaries {
+ SchemaEncoding::Fallback(Arc::new(schema.clone()))
+ } else {
+ 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,
+ &write_options,
+ );
+ let mut buf = Vec::new();
+ write_message(&mut buf, encoded_schema, &write_options)?;
+ SchemaEncoding::Precoded(buf)
+ };
Ok(Self {
codec,
header_bytes,
+ write_options,
+ schema_encoding,
})
}
+ /// Serialize `batch` as a standalone Arrow IPC stream into `out`.
+ fn encode_ipc_stream<W: Write>(&self, batch: &RecordBatch, out: &mut W) ->
Result<()> {
+ let schema_message = match &self.schema_encoding {
+ SchemaEncoding::Fallback(schema) => {
+ // 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_with_options(out, schema,
self.write_options.clone())?;
+ stream_writer.write(batch)?;
+ stream_writer.finish()?;
+ return Ok(());
+ }
+ SchemaEncoding::Precoded(schema_message) => schema_message,
+ };
+
+ // Fast path: reuse the pre-encoded schema message and write the
record batch manually.
+ let data_gen = IpcDataGenerator::default();
+ let mut dictionary_tracker = DictionaryTracker::new(true);
+ let mut compression_context = CompressionContext::default();
+ let (encoded_dictionaries, encoded_batch) = data_gen.encode(
+ batch,
+ &mut dictionary_tracker,
+ &self.write_options,
+ &mut compression_context,
+ )?;
+ debug_assert!(encoded_dictionaries.is_empty());
+
+ out.write_all(schema_message)?;
+ write_message(&mut *out, encoded_batch, &self.write_options)?;
+ out.write_all(&IPC_EOS)?;
Review Comment:
Good question, but flushing isn't needed here: for the compressed codecs,
`write_batch` calls `.finish()`/`.into_inner()` on the encoder after
`encode_ipc_stream` returns and before reading `stream_position()`, which fully
flushes into `output`.
The hang turned out to be unrelated. `ShuffleBlockWriter` is cloned once per
output partition, and this PR made that clone deep-copy the pre-encoded schema
`Vec`. `SPARK-48037` uses 16M+ partitions, so that allocated GBs and stalled
the executor (only the shard running `AdaptiveQueryExecSuite` hangs). Fixed in
the last commit by sharing the buffers via `Arc`.
--
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]
