xhumanoid commented on code in PR #5473: URL: https://github.com/apache/datafusion-comet/pull/5473#discussion_r3861137012
########## native/jni-bridge/src/shuffle_partition_pusher.rs: ########## @@ -0,0 +1,147 @@ +// 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 crate::{check_exception, errors::CometError, JVMClasses}; +use datafusion::common::{DataFusionError, Result}; +use jni::objects::{Global, JMethodID, JObject, JValue}; +use jni::signature::{Primitive, ReturnType}; +use jni::Env; + +/// Receives a complete encoded shuffle block for one output partition. +/// +/// Implementations must remain safe when invoked from native execution +/// threads that do not inherit Spark's task-local JVM state. +pub trait ShufflePartitionPusher: Send + Sync { + /// Sends one complete, length-prefixed Arrow IPC shuffle block. + fn push_partition_data(&self, partition_id: i32, data: &[u8]) -> Result<()>; +} + +/// Invokes a task-owned JVM shuffle callback from any native execution thread. +/// +/// The global callback reference keeps both the Java object and its class alive +/// for the lifetime of the cached method ID. No thread-local JNI environment +/// or Spark task context is retained between invocations. +pub struct JavaShufflePartitionPusher { + callback: Global<JObject<'static>>, + push_method: JMethodID, +} + +impl JavaShufflePartitionPusher { + /// Captures the callback while running on an attached JVM thread. + pub fn try_new(env: &mut Env<'_>, callback: &JObject<'_>) -> Result<Self> { + if callback.is_null() { + return Err(DataFusionError::Execution( + "Remote shuffle callback must not be null".to_string(), + )); + } + + let callback_class = env.get_object_class(callback).map_err(CometError::from)?; + let push_method = env + .get_method_id( + &callback_class, + jni::jni_str!("pushPartitionData"), + jni::jni_sig!("(I[BI)V"), + ) + .map_err(CometError::from)?; + let callback = env.new_global_ref(callback).map_err(CometError::from)?; + + Ok(Self { + callback, + push_method, + }) + } + + fn checked_payload_length(partition_id: i32, payload_length: usize) -> Result<i32> { + if partition_id < 0 { + return Err(DataFusionError::Execution(format!( + "Remote shuffle partition must be nonnegative, got {partition_id}" + ))); + } + + i32::try_from(payload_length).map_err(|_| { + DataFusionError::Execution(format!( + "Remote shuffle payload size {payload_length} exceeds the JVM array limit" + )) + }) Review Comment: in specification it Integer.MAX_VALUE, but HotSpot implementation details Integer.MAX_VALUE-8 so maybe make sense to test it as well https://github.com/openjdk/jdk/blob/master/src/java.base/share/classes/jdk/internal/util/ArraysSupport.java#L854-L866 ``` /** * A soft maximum array length imposed by array growth computations. * Some JVMs (such as HotSpot) have an implementation limit that will cause * * OutOfMemoryError("Requested array size exceeds VM limit") * * to be thrown if a request is made to allocate an array of some length near * Integer.MAX_VALUE, even if there is sufficient heap available. The actual * limit might depend on some JVM implementation-specific characteristics such * as the object header size. The soft maximum value is chosen conservatively so * as to be smaller than any implementation limit that is likely to be encountered. */ public static final int SOFT_MAX_ARRAY_LENGTH = Integer.MAX_VALUE - 8; ``` -- 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]
