andygrove commented on code in PR #5565:
URL: https://github.com/apache/datafusion-comet/pull/5565#discussion_r3895367901
##########
native/shuffle/src/ipc.rs:
##########
@@ -15,23 +15,65 @@
// 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! {
Review Comment:
The 8 MiB cap turns the unbounded retention into a bounded one, which is a
good fix, but the memory is still held for the life of the thread and still is
not visible to any reservation. Both production callers go through the
thread-local: the static JNI `decodeShuffleBlock` and `ShuffleScanStream`.
Those run on JVM task threads and tokio workers, all of which live as long as
the executor. A 16 core executor that decodes one zstd shuffle block ends up
sitting on roughly 128 MiB of native memory for the rest of its life, including
during stages that never shuffle.
`ShuffleScanStream` looks like a natural owner here. Could
`decode_shuffle_batch` take a `&mut ShuffleDecodeContext` held by the stream
and go through `read_ipc_compressed_with`? That would bound retention to the
operator rather than the thread, and it would leave the thread-local for
`Java_org_apache_comet_Native_decodeShuffleBlock`, which really has no handle
to hang a context off. Right now `read_ipc_compressed_with` and
`read_ipc_compressed_validated_with` are exported from `lib.rs` but only ever
called from tests, so this would also give them a real caller.
##########
native/shuffle/src/codec_context.rs:
##########
@@ -0,0 +1,140 @@
+// 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.
+
+use arrow::ipc::writer::CompressionContext;
+use std::io;
+use zstd::zstd_safe::{CCtx, CParameter, DCtx, ResetDirective};
+
+/// Largest zstd workspace worth caching between frames. Covers the commonly
configured
+/// levels; higher levels (tens to hundreds of MiB of window) fall back to a
fresh context
+/// per frame, which is what per-block encoding paid anyway.
+const MAX_RETAINED_ZSTD_CONTEXT_BYTES: usize = 8 * 1024 * 1024;
Review Comment:
I measured `CCtx::sizeof()` against the pinned zstd 1.5.7 build after one
streaming frame and the cap is closer to the edge than the comment suggests.
Levels 7 and 8 come in at 8,119,825 bytes against a cap of 8,388,608, so about
3% of headroom. Levels 1 through 6 are 1.31 to 5.24 MiB and level 9 jumps to
14.74 MiB, so anything at 9 or above never reuses at all. On the decode side a
level 19 frame leaves the `DCtx` at 8.47 MiB, which also misses the cap.
Two things that would help. Could the measured level to size table go in the
comment next to the constant, so the choice of 8 MiB is traceable and someone
bumping `zstd-sys` can see what they are moving? And could a test pin where the
boundary actually falls, say level 6 retains and level 9 does not? As it stands
a routine dependency bump could push levels 7 and 8 over the line and silently
disable the optimization for those users with every test still passing.
##########
native/shuffle/src/shuffle_writer.rs:
##########
@@ -577,6 +566,56 @@ mod test {
repartitioner.insert_batch(batch.clone()).await.unwrap();
}
+ /// The zstd context is reused within one encode burst but must not
survive past it: a
+ /// spill event and the final shuffle write each end with the context
released.
+ #[tokio::test]
+ #[cfg_attr(miri, ignore)] // miri can't call foreign function
`ZSTD_createCCtx`
+ async fn local_writer_releases_zstd_context_at_burst_boundaries() {
Review Comment:
This test only ever asserts `!holds_zstd_cctx()`, so it would pass just as
happily if the context were never created in the first place. That is true of
the suite generally. The tests establish that blocks round-trip correctly under
a shared context and that release happens at the right boundaries, but nothing
observes that N blocks produce fewer than N context creations, which is the
actual claim of the PR.
Would a test-only creation counter on `ShuffleCodecContext` work? Asserting
that one spill burst over two partitions creates exactly one context would pin
the reuse behaviour directly, and the same counter would let you pin the level
boundary from my comment on `codec_context.rs`.
##########
native/shuffle/src/writers/shuffle_block_writer.rs:
##########
@@ -269,9 +273,52 @@ impl ShuffleBlockWriter {
// write header
output.write_all(&self.header_bytes)?;
+ let encode_result =
Review Comment:
This is right given that `rss_codec_workspace` is charged per admitted
invocation, and I checked that it fires on the error paths as well as the
success path. The consequence though is that RSS allocates and frees a `CCtx`
per block exactly as it does on main, so the remote path gets none of the
benefit this PR is after. Small frames pushed to Celeborn are arguably the
shape where per-block context setup hurts most.
Was reserving the workspace once for the lifetime of the
`RssPartitionWriter` rather than per invocation considered? There is already
one writer per task, so the accounting would be a single up-front charge
instead of a repeated one. If that turns out to be awkward against the pusher's
admission model it would be worth saying so in the description, which currently
reads as though both paths benefit.
--
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]
