mbutrovich commented on code in PR #5006:
URL: https://github.com/apache/datafusion-comet/pull/5006#discussion_r3640745052
##########
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:
arrow-rs already provides this. `Schema::flattened_fields()` walks the full
nested field tree (via `Field::fields()` / `Field::_fields`), so the detection
at 103-106 becomes:
```rust
let has_dictionaries = schema
.flattened_fields()
.iter()
.any(|f| matches!(f.data_type(), DataType::Dictionary(_, _)));
```
and this function can be deleted. Beyond removing the hand-maintained match
arms, it closes a gap: `contains_dictionary` omits `ListView` and
`LargeListView`, but arrow's own dictionary traversal (`_encode_dictionaries`,
arrow-ipc/src/writer.rs:438,450) recurses into them. A
`ListView<Dictionary<..>>` column would be classified dictionary-free, take the
fast path, and trip the `encoded_dictionaries.is_empty()` assert (or drop
dictionary batches in release). Unreachable today since Comet emits no
ListView, but keying off arrow's own field walk keeps this from drifting on a
version bump, and it makes the `debug_assert` at line 157 sound rather than
dependent on this recursion staying in sync.
##########
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:
Following up on the single-field suggestion: `schema` and `schema_message`
are mutually exclusive by use (fast path reads only `schema_message`, fallback
reads only `schema`), so an enum states the invariant and drops the dead field:
```rust
enum SchemaEncoding {
Precoded(Vec<u8>), // dictionary-free: written verbatim per block
Fallback(SchemaRef), // dictionary schemas: StreamWriter re-encodes
}
```
Today `try_new` runs `Arc::new(schema.clone())` at line 130 unconditionally,
a deep clone of every field that is never read on the dictionary-free path.
With the enum the schema is only retained in the `Fallback` arm, and as an
`Arc::clone` rather than a deep clone.
##########
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:
Checked this against arrow-ipc 58.4: none of these need caching.
- `CompressionContext` is `#[derive(Default)]` over
`Option<zstd::bulk::Compressor>` (compression.rs:30-34), so `::default()` is
`None` and the zstd context is built lazily only when arrow's internal IPC
compression runs. We pass `IpcWriteOptions::default()` with
`batch_compression_type: None` (block compression is done by our own frame
encoders), so arrow never constructs it. No allocation here. (It was an eager
`ZSTD_createCCtx` in 58.1 and earlier, made lazy in 58.3.)
- `IpcDataGenerator` is zero-sized. `IpcWriteOptions::default()` is a few
scalar assignments, no allocation.
- `DictionaryTracker` must stay fresh per block, since each block is a
standalone stream and the tracker records which dict ids were already emitted.
So reuse would add interior mutability for no measurable gain. Worth
revisiting after the arrow 59.x bump: there `CompressionContext` becomes
`IpcWriteContext` with `reserve_scratch_with_capacity`, and reusing one
instance lets the `arrow_data` staging buffer grow once instead of per block.
That pairs with the zstd-context-reuse item still open in #5002.
##########
native/shuffle/src/writers/shuffle_block_writer.rs:
##########
@@ -16,12 +16,20 @@
// under the License.
use arrow::array::RecordBatch;
-use arrow::datatypes::Schema;
-use arrow::ipc::writer::StreamWriter;
+use arrow::datatypes::{DataType, Schema, SchemaRef};
+use arrow::ipc::writer::{
+ write_message, CompressionContext, DictionaryTracker, IpcDataGenerator,
IpcWriteOptions,
+ StreamWriter,
+};
use datafusion::common::DataFusionError;
use datafusion::error::Result;
use datafusion::physical_plan::metrics::Time;
-use std::io::{Cursor, Seek, SeekFrom, Write};
+use std::io::{Seek, SeekFrom, Write};
+use std::sync::Arc;
+
+/// Arrow IPC stream end-of-stream marker: the continuation marker
(`0xFFFFFFFF`) followed by a
+/// zero message length, matching what `StreamWriter::finish` emits for
metadata version V5.
+const IPC_EOS: [u8; 8] = [0xff, 0xff, 0xff, 0xff, 0x00, 0x00, 0x00, 0x00];
Review Comment:
Shuffle blocks are always V5 (`IpcWriteOptions::default()`), so hardcoding
the EOS is fine here. The comment already notes the V5 provenance; worth making
the constraint explicit so a future reader knows it is an assumption, not a
universal constant: these bytes are the continuation marker plus zero length
only for `MetadataVersion::V5` with `write_legacy_ipc_format = false`
(V4-legacy emits 4 zero bytes with no marker, per `write_continuation` in
arrow-ipc/src/writer.rs:1893-1922). If Comet ever uses arrow IPC elsewhere and
applies the same pre-encode trick under different options, this constant must
be revisited. A one-line `debug_assert!(matches!(options.metadata_version,
MetadataVersion::V5))` next to the write at line 161 pins the assumption at low
cost.
##########
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:
The schema is encoded with an `IpcWriteOptions` built in `try_new` (line
114) and the batch with a second one built here. They are identical today, but
storing one `IpcWriteOptions` on the writer and referencing it in both places
removes the chance of the schema and batch being encoded under divergent
options if a default is ever changed.
--
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]
