comphead commented on code in PR #7354: URL: https://github.com/apache/datafusion-comet/pull/7354#discussion_r4470288642
########## native/shuffle/src/decompress.rs: ########## @@ -0,0 +1,401 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Decompressors kept from one shuffle block to the next, see [`Codecs`]. + +use crate::codec_context::MAX_RETAINED_ZSTD_CONTEXT_BYTES; +use crate::ipc::decode_error; +use datafusion::error::Result; +use lz4_flex::frame::FrameDecoder as Lz4FrameDecoder; +use snap::read::FrameDecoder as SnappyFrameDecoder; +use std::io::{Error, ErrorKind, Read}; +use std::marker::PhantomData; +use std::ptr; +use zstd::zstd_safe::{DCtx, ResetDirective}; + +/// Decompressors kept from one block to the next. +/// +/// Building a decoder per block was most of the fixed cost of decoding a small block. zstd +/// allocates a decompression context with its window buffers, lz4_flex allocates block-sized +/// input and output buffers as it reads the frame, and snap allocates and zero-fills about +/// 140 KiB of them up front. Kept on the thread, the lz4_flex and snap decoders reuse those +/// buffers for every following frame, and the zstd context keeps its workspace across the +/// session reset that readies it for the next frame. +/// +/// Retention per thread is bounded. An lz4 decoder's buffers are sized to its frame's block +/// size, and lz4_flex expects them to stay that size, so one decoder is kept per block size the +/// thread has seen, 64 KiB to 4 MiB, at most about 11 MiB in all; Comet's writer picks the size +/// from the batch, so a shuffle normally uses one or two. A frame of any other layout (linked +/// blocks, or the legacy format) gets a decoder for that block only. snap's buffers are fixed at Review Comment: lz4_flex resolves `BlockSize::Auto` from the length of the first `write` call (`begin_frame` calls `BlockSize::from_buf_length`), and in `ShuffleBlockWriter::encode_ipc_stream` that first write is the precoded schema message, or the 8-byte continuation marker from `StreamWriter` on the dictionary path. So Comet's frames use 64 KiB blocks whatever the batch size, unless the encoded schema itself is larger than 64 KiB, and `Auto` never picks 1 MiB. A scratch program against lz4_flex 0.14.0 agrees: a 3,000-byte first write followed by 8 MiB gives block size code 4. Could this say that the size follows the schema message rather than the batch? The test comment on `decompressors_are_built_once_per_thread` ("the writer gives the 3-row batch 64 KiB blocks and the 8192-row one 256 KiB blocks") describes the test's `encode` helper, which writes the whole payload in one `write_all`, rather than Comet's writer. ########## native/shuffle/src/ipc.rs: ########## @@ -1421,6 +1438,144 @@ mod tests { } } + const COMPRESSED_CODECS: [&[u8; 4]; 3] = [b"LZ4_", b"ZSTD", b"SNAP"]; + + /// One decompressor per codec serves every block the thread decodes, across changes of + /// codec, of schema, and of validation. lz4 keeps one per block size: the writer gives the + /// 3-row batch 64 KiB blocks and the 8192-row one 256 KiB blocks. + #[test] + #[cfg_attr(miri, ignore)] // Miri cannot call Zstd's C FFI. + fn decompressors_are_built_once_per_thread() { + reset_codecs(); + for _ in 0..3 { + for batch in [mixed_batch(), wide_batch(8192)] { + for codec in CODECS { + let block = block_for(&batch, codec); + assert_eq!(read_ipc_compressed(&block).unwrap(), batch); + assert_eq!(read_ipc_compressed_validated(&block).unwrap(), batch); + } + } + } + assert_eq!( + codec_creations(), + CodecCreations { + lz4: 2, + snappy: 1, + zstd: 1, + } + ); + } + + /// A decode that fails inside the compressed stream leaves no state behind: the next block + /// of that codec decodes correctly. The lz4 and snappy decoders are replaced, since they + /// would otherwise resume the broken frame; the zstd context is reset in place. + #[test] + #[cfg_attr(miri, ignore)] // Miri cannot call Zstd's C FFI. + fn failed_decode_does_not_poison_the_next_block() { + let batch = wide_batch(2048); + reset_codecs(); + for codec in COMPRESSED_CODECS { + let block = block_for(&batch, codec); + assert_eq!(read_ipc_compressed(&block).unwrap(), batch); + let before = codec_creations(); + // cut inside the compressed data, early and late in the frame + for cut in [24, block.len() * 2 / 3] { + assert!( + read_ipc_compressed(&block[..cut]).is_err(), + "{codec:?}: truncated at {cut}" + ); + assert_eq!(read_ipc_compressed(&block).unwrap(), batch, "{codec:?}"); + } + let after = codec_creations(); + let replaced = match codec { + b"LZ4_" => after.lz4 - before.lz4, + b"SNAP" => after.snappy - before.snappy, + b"ZSTD" => after.zstd - before.zstd, + _ => unreachable!(), + }; + assert_eq!(replaced, if codec == b"ZSTD" { 0 } else { 2 }, "{codec:?}"); + } + } + + /// A kept snappy decoder must not relax what a fresh one requires: every block opens with + /// the stream identifier. + #[test] + fn snappy_block_without_stream_identifier_is_rejected_warm() { + let batch = mixed_batch(); + let block = block_for(&batch, b"SNAP"); + assert_eq!(read_ipc_compressed(&block).unwrap(), batch); + + // the 10-byte stream identifier chunk follows the codec tag + assert_eq!(&block[4..14], b"\xff\x06\x00\x00sNaPpY"); + let mut headless = block[..4].to_vec(); + headless.extend_from_slice(&block[14..]); + let error = read_ipc_compressed(&headless).unwrap_err().to_string(); + assert!(error.contains("snappy stream header"), "{error}"); + assert_eq!(read_ipc_compressed(&block).unwrap(), batch); + } + + /// lz4 decoders are kept per block size, so frames of two sizes each keep their own, and a + /// frame of another layout (linked blocks) decodes with a decoder that is not kept. + #[test] + fn lz4_decoders_are_kept_per_block_size_and_not_for_other_layouts() { + use lz4_flex::frame::{BlockMode, BlockSize, FrameEncoder, FrameInfo}; + let batch = mixed_batch(); + let payload = ipc_bytes(&batch); + let encode_with = |frame_info: FrameInfo| { + let mut block = b"LZ4_".to_vec(); + let mut writer = FrameEncoder::with_frame_info(frame_info, &mut block); + writer.write_all(&payload).unwrap(); + writer.finish().unwrap(); + block + }; + let small = encode(b"LZ4_", &payload); + let large = encode_with(FrameInfo::new().block_size(BlockSize::Max4MB)); + let linked = encode_with(FrameInfo::new().block_mode(BlockMode::Linked)); + reset_codecs(); + for (block, expected_creations, case) in [ + (&small, 1, "first 64 KiB frame"), + (&large, 2, "first 4 MiB frame"), + (&small, 2, "64 KiB decoder kept"), + (&large, 2, "4 MiB decoder kept"), + (&linked, 3, "linked frame gets its own"), + (&linked, 4, "which is not kept"), + (&small, 4, "64 KiB decoder still kept"), + ] { + assert_eq!(read_ipc_compressed(block).unwrap(), batch, "{case}"); + assert_eq!(codec_creations().lz4, expected_creations, "{case}"); + } + } + + /// Warm, a compressed block allocates no more than the same block uncompressed: the batch + /// itself. The decompressor's buffers and context come from the thread. + #[test] + #[cfg_attr(miri, ignore)] // Miri cannot call Zstd's C FFI. + fn warm_compressed_decode_allocates_no_more_than_an_uncompressed_one() { Review Comment: libzstd allocates its context and window buffers with the C `malloc`, which the counting global allocator does not see. A scratch program against zstd 0.13.3 counts 0 Rust-side bytes for `DCtx::try_create` and for a cold 2 MB decode, while `dctx.sizeof()` reports about 1 MB afterwards. So the `ZSTD` case here passes whether or not the context is reused. Would it be clearer to run this test over `LZ4_` and `SNAP` only, and leave zstd reuse to `decompressors_are_built_once_per_thread`, which checks it through the creation counter? ########## native/shuffle/src/decompress.rs: ########## @@ -0,0 +1,401 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Decompressors kept from one shuffle block to the next, see [`Codecs`]. + +use crate::codec_context::MAX_RETAINED_ZSTD_CONTEXT_BYTES; +use crate::ipc::decode_error; +use datafusion::error::Result; +use lz4_flex::frame::FrameDecoder as Lz4FrameDecoder; +use snap::read::FrameDecoder as SnappyFrameDecoder; +use std::io::{Error, ErrorKind, Read}; +use std::marker::PhantomData; +use std::ptr; +use zstd::zstd_safe::{DCtx, ResetDirective}; + +/// Decompressors kept from one block to the next. +/// +/// Building a decoder per block was most of the fixed cost of decoding a small block. zstd +/// allocates a decompression context with its window buffers, lz4_flex allocates block-sized +/// input and output buffers as it reads the frame, and snap allocates and zero-fills about +/// 140 KiB of them up front. Kept on the thread, the lz4_flex and snap decoders reuse those +/// buffers for every following frame, and the zstd context keeps its workspace across the +/// session reset that readies it for the next frame. +/// +/// Retention per thread is bounded. An lz4 decoder's buffers are sized to its frame's block +/// size, and lz4_flex expects them to stay that size, so one decoder is kept per block size the +/// thread has seen, 64 KiB to 4 MiB, at most about 11 MiB in all; Comet's writer picks the size +/// from the batch, so a shuffle normally uses one or two. A frame of any other layout (linked +/// blocks, or the legacy format) gets a decoder for that block only. snap's buffers are fixed at +/// about 140 KiB. A zstd context is dropped once it outgrows +/// [`MAX_RETAINED_ZSTD_CONTEXT_BYTES`], as the writer drops its compression context. +#[derive(Default)] +pub(crate) struct Codecs { + /// One per block size code, see [`lz4_retention_slot`]. + lz4: [Option<Lz4Decoder>; LZ4_BLOCK_SIZES], + /// For lz4 frames of any other layout; never kept past the block. + lz4_other: Option<Lz4Decoder>, Review Comment: `lz4_other` is always `None` between blocks, and it is the only reason `Loan` carries `keep`. For a frame that `lz4_retention_slot` rejects, could `ipc.rs` decode with `Lz4FrameDecoder::new(RequireLz4EndMark(&mut encoded))` over the borrowed slice, as before this PR? That decoder needs no loan, so the raw pointer would only serve the decoders that are kept, and `keep`, `lz4_other` and the `keep` false case in the `Loan` test could go. ########## native/shuffle/src/decompress.rs: ########## @@ -0,0 +1,401 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Decompressors kept from one shuffle block to the next, see [`Codecs`]. + +use crate::codec_context::MAX_RETAINED_ZSTD_CONTEXT_BYTES; +use crate::ipc::decode_error; +use datafusion::error::Result; +use lz4_flex::frame::FrameDecoder as Lz4FrameDecoder; +use snap::read::FrameDecoder as SnappyFrameDecoder; +use std::io::{Error, ErrorKind, Read}; +use std::marker::PhantomData; +use std::ptr; +use zstd::zstd_safe::{DCtx, ResetDirective}; + +/// Decompressors kept from one block to the next. +/// +/// Building a decoder per block was most of the fixed cost of decoding a small block. zstd +/// allocates a decompression context with its window buffers, lz4_flex allocates block-sized +/// input and output buffers as it reads the frame, and snap allocates and zero-fills about +/// 140 KiB of them up front. Kept on the thread, the lz4_flex and snap decoders reuse those +/// buffers for every following frame, and the zstd context keeps its workspace across the +/// session reset that readies it for the next frame. +/// +/// Retention per thread is bounded. An lz4 decoder's buffers are sized to its frame's block +/// size, and lz4_flex expects them to stay that size, so one decoder is kept per block size the +/// thread has seen, 64 KiB to 4 MiB, at most about 11 MiB in all; Comet's writer picks the size +/// from the batch, so a shuffle normally uses one or two. A frame of any other layout (linked +/// blocks, or the legacy format) gets a decoder for that block only. snap's buffers are fixed at +/// about 140 KiB. A zstd context is dropped once it outgrows +/// [`MAX_RETAINED_ZSTD_CONTEXT_BYTES`], as the writer drops its compression context. +#[derive(Default)] +pub(crate) struct Codecs { + /// One per block size code, see [`lz4_retention_slot`]. + lz4: [Option<Lz4Decoder>; LZ4_BLOCK_SIZES], + /// For lz4 frames of any other layout; never kept past the block. + lz4_other: Option<Lz4Decoder>, + snappy: Option<SnappyDecoder>, + zstd: Option<DCtx<'static>>, + #[cfg(test)] + pub(crate) creations: CodecCreations, +} + +/// The LZ4 frame block maximum sizes lz4_flex decodes: codes 4 to 7, 64 KiB to 4 MiB. +const LZ4_BLOCK_SIZES: usize = 4; + +pub(crate) type Lz4Decoder = Lz4FrameDecoder<RequireLz4EndMark<BlockInput>>; +pub(crate) type SnappyDecoder = SnappyFrameDecoder<BlockInput>; + +impl Codecs { + pub(crate) fn lend_lz4<'a, 'b>(&'a mut self, block: &'b [u8]) -> Loan<'a, 'b, Lz4Decoder> { + let (slot, keep) = match lz4_retention_slot(block) { + Some(index) => (&mut self.lz4[index], true), + None => (&mut self.lz4_other, false), + }; + #[cfg(test)] + if slot.is_none() { + self.creations.lz4 += 1; + } + Loan::new(slot, block, keep) + } + + pub(crate) fn lend_snappy<'a, 'b>( + &'a mut self, + block: &'b [u8], + ) -> Loan<'a, 'b, SnappyDecoder> { + #[cfg(test)] + if self.snappy.is_none() { + self.creations.snappy += 1; + } + Loan::new(&mut self.snappy, block, true) + } + + /// The shared zstd context, reset for one frame. The reset also clears whatever a failed + /// decode left behind, so the context survives errors. + pub(crate) fn zstd_dctx(&mut self) -> Result<&mut DCtx<'static>> { + if self.zstd.is_none() { + #[cfg(test)] + { + self.creations.zstd += 1; + } + let dctx = DCtx::try_create() + .ok_or_else(|| decode_error("failed to allocate zstd decompression context"))?; + self.zstd = Some(dctx); + } + let dctx = self.zstd.as_mut().expect("created above"); + dctx.reset(ResetDirective::SessionOnly) + .map_err(|code| decode_error(zstd::zstd_safe::get_error_name(code)))?; + Ok(dctx) + } + + /// Drops the zstd context when its workspace, sized by the frames it has decoded, outgrew + /// [`MAX_RETAINED_ZSTD_CONTEXT_BYTES`]; the next zstd block creates a new one. + pub(crate) fn release_zstd_if_oversized(&mut self) { + if self + .zstd + .as_ref() + .is_some_and(|dctx| dctx.sizeof() > MAX_RETAINED_ZSTD_CONTEXT_BYTES) + { + self.zstd = None; + } + } +} + +/// Whether a snappy block opens with the stream identifier chunk (type `0xff`), or is empty. +/// +/// snap checks for the identifier only on a decoder's first chunk, so a retained decoder would +/// accept a later block that lacks it. A fresh decoder does not, and neither does the reader. +/// An empty block passes, to fail as an empty IPC stream the way it always has. +pub(crate) fn snappy_block_opens_with_stream_identifier(block: &[u8]) -> bool { + const STREAM_IDENTIFIER_CHUNK: u8 = 0xff; + block + .first() + .is_none_or(|&chunk_type| chunk_type == STREAM_IDENTIFIER_CHUNK) +} + +/// Which retained decoder an LZ4 frame goes to: one per block maximum size, for a standard frame +/// with independent blocks, the layout Comet's writer produces. `None` for any other frame. +/// +/// lz4_flex sizes a decoder's buffers to its frame's block size and asserts, in debug builds, +/// that they keep that size, so frames of different sizes cannot share a decoder. The header +/// is the 4-byte magic number, then `FLG`, whose bit 5 marks independent blocks, then `BD`, +/// whose bits 6 to 4 code the block size: 4 for 64 KiB through 7 for 4 MiB. +pub(crate) fn lz4_retention_slot(frame: &[u8]) -> Option<usize> { + const MAGIC: [u8; 4] = 0x184D2204_u32.to_le_bytes(); + const INDEPENDENT_BLOCKS: u8 = 1 << 5; + const BLOCK_SIZE_MASK: u8 = 0b0111_0000; + const SMALLEST_CODE: u8 = 4; + let (magic, header) = frame.split_at_checked(MAGIC.len())?; + let &[flg, bd, ..] = header else { + return None; + }; + if magic != MAGIC || flg & INDEPENDENT_BLOCKS == 0 { + return None; + } + let code = (bd & BLOCK_SIZE_MASK) >> 4; + (SMALLEST_CODE..SMALLEST_CODE + LZ4_BLOCK_SIZES as u8) + .contains(&code) + .then(|| usize::from(code - SMALLEST_CODE)) +} + +/// A decoder kept in [`Codecs`], reading the block it is lent through a [`BlockInput`]. +pub(crate) trait RetainedDecoder { + fn new() -> Self; + fn input(&mut self) -> &mut BlockInput; +} + +impl RetainedDecoder for Lz4Decoder { + fn new() -> Self { + Lz4FrameDecoder::new(RequireLz4EndMark(BlockInput::default())) + } + + fn input(&mut self) -> &mut BlockInput { + &mut self.get_mut().0 + } +} + +impl RetainedDecoder for SnappyDecoder { + fn new() -> Self { + SnappyFrameDecoder::new(BlockInput::default()) + } + + fn input(&mut self) -> &mut BlockInput { + self.get_mut() + } +} + +/// The compressed bytes of the block being decoded, read by a decoder that outlives the block. +/// +/// The lz4_flex and snap frame decoders own their reader. A reader that borrowed the block would +/// put the block's lifetime into the decoder's type, and a decoder kept in a thread-local has to +/// be `'static`. So the block is lent as a raw pointer instead, for the duration of one decode, +/// by a [`Loan`]: lending points the input at the block, and the loan ends by clearing it, +/// before the borrow it was made from can end. Outside a loan the input is empty, and a read +/// returns nothing without touching the pointer. +pub(crate) struct BlockInput { + /// The unread rest of the lent block; dangling and empty outside a loan. + unread: *const u8, + len: usize, +} + +impl Default for BlockInput { + fn default() -> Self { + Self { + unread: ptr::null(), + len: 0, + } + } +} + +impl BlockInput { + fn lend(&mut self, block: &[u8]) { + self.unread = block.as_ptr(); + self.len = block.len(); + } + + /// Ends the loan, returning how many bytes of the block were left unread. + fn revoke(&mut self) -> usize { + self.unread = ptr::null(); + std::mem::take(&mut self.len) + } +} + +impl Read for BlockInput { + fn read(&mut self, buffer: &mut [u8]) -> std::io::Result<usize> { + let count = buffer.len().min(self.len); + if count == 0 { + return Ok(0); + } + // SAFETY: `len` is nonzero only during a loan, when `unread..unread + len` is the unread + // tail of the slice `lend` was given, and the `Loan` holding that slice's lifetime ends + // the loan before the slice can go away. + let bytes = unsafe { std::slice::from_raw_parts(self.unread, count) }; + buffer[..count].copy_from_slice(bytes); + // SAFETY: `count <= len`, so the result stays within the same slice. + self.unread = unsafe { self.unread.add(count) }; + self.len -= count; + Ok(count) + } +} + +/// A block lent to a retained decoder for one decode, see [`BlockInput`]. +/// +/// The decoder is built on the first loan. [`Loan::settle`] ends the loan and says whether the +/// decode finished cleanly. A decoder that did not, because of an error, or of a panic that +/// unwinds through this guard's drop, is discarded, since it would carry the broken frame's +/// state into the next block. A decoder lent with `keep` false is discarded either way. +pub(crate) struct Loan<'a, 'b, D: RetainedDecoder> { + slot: &'a mut Option<D>, + keep: bool, + settled: bool, + block: PhantomData<&'b [u8]>, +} + +impl<'a, 'b, D: RetainedDecoder> Loan<'a, 'b, D> { + fn new(slot: &'a mut Option<D>, block: &'b [u8], keep: bool) -> Self { + slot.get_or_insert_with(D::new).input().lend(block); + Self { + slot, + keep, + settled: false, + block: PhantomData, + } + } + + pub(crate) fn decoder(&mut self) -> &mut D { + self.slot.as_mut().expect("lent on construction") + } Review Comment: `decoder()` hands out `&mut D` while the block is lent, so code elsewhere in the crate could `std::mem::replace(loan.decoder(), D::new())` and keep the lent decoder past the loan. Its `BlockInput` would still point into the block, and a later `read` would dereference freed memory with no `unsafe` at the call site. The `SAFETY` comment in `BlockInput::read` relies on the decoder never leaving the slot, which nothing enforces today. Would it work to implement `Read` for `Loan` by forwarding to the decoder, and pass `Streamed(&mut loan)` in `ipc.rs`? Then the decoder cannot be moved out, and the tests in this module can still reach it through the slot. `mem::forget` on a `Loan` is already harmless, since the next `Loan::new` lends again before any read. -- 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]
