sunchao commented on code in PR #5565:
URL: https://github.com/apache/datafusion-comet/pull/5565#discussion_r3891638153
##########
native/shuffle/src/writers/local/local_partition_writer.rs:
##########
@@ -64,6 +65,9 @@ enum DataOutput {
pub(crate) struct LocalPartitionWriter {
output_index_file: String,
data_output: DataOutput,
+ /// Compression state shared by every block this task writes; the
per-partition
+ /// `BufBatchWriter`s borrow it (see [`ShuffleCodecContext`]).
+ codec_context: ShuffleCodecContext,
Review Comment:
**[P2]** Could we account for this context's retained workspace for the full
local-writer lifetime, or release it at a spill or idle boundary?
`LocalPartitionWriter` has no memory reservation for the `CCtx`, while the
repartitioner frees its tracked reservation after spilling. With the locked
zstd build, the same `Encoder::with_context` path retained 72,082,969 bytes at
level 15 and 874,070,679 bytes at level 22 after `SessionOnly` reset. Those
levels are accepted by the current configuration. Concurrent tasks can
therefore keep large native allocations after the pool reports their buffered
memory as released. A size-based regression would also catch this because the
current test only checks that the context is present.
##########
native/shuffle/src/ipc.rs:
##########
@@ -15,23 +15,52 @@
// specific language governing permissions and limitations
// under the License.
+use crate::codec_context::ShuffleDecodeContext;
use arrow::array::RecordBatch;
use arrow::ipc::reader::StreamReader;
use datafusion::common::DataFusionError;
use datafusion::error::Result;
+use std::cell::RefCell;
use std::io::{Error, ErrorKind, Read};
+thread_local! {
+ /// Backs the entry points below. They're called from many JVM task
threads; a
+ /// thread-local gets each thread context reuse without changing any
caller.
+ static DECODE_CONTEXT: RefCell<ShuffleDecodeContext> =
Review Comment:
**[P2]** Could we move this decode context under reader or task ownership,
or release it when it exceeds a bounded size? `ResetDirective::SessionOnly`
preserves zstd's allocated window. With the locked zstd 1.5.7 build and the
same context sequence used here, a valid 17-byte level-22 frame made
`DCtx::sizeof()` grow from 95,992 to 134,707,000 bytes, and reset left it at
that size. Both production entry points use this thread-local context, so
executor worker threads retain the native allocation across tasks. At 32
threads that is about 4 GiB. The base path dropped the decoder per frame. It
might be worth adding a regression that verifies the workspace is released when
a reader closes and after a decode failure.
--
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]
