perf: [iceberg] per partition file scan task serialization

Author: parthchandra

docs/source/contributor-guide/iceberg_partition_optimization.md

@@ -29,6 +29,7 @@ Arrow FFI <ffi>
 JVM Shuffle <jvm_shuffle>
 Native Shuffle <native_shuffle>
 Parquet Scans <parquet_scans>
+Iceberg Partition Optimization <iceberg_partition_optimization>
 Development Guide <development>
 Debugging Guide <debugging>
 Benchmarking Guide <benchmarking>

docs/source/contributor-guide/index.md

@@ -946,3 +946,39 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_columnarToRowClose(
         Ok(())
     })
 }
+
+#[no_mangle]
+/// Get Iceberg partition tasks for the current partition via JNI callback.
+///
+/// Called by native planner to retrieve partition-specific FileScanTask bytes from
+/// the thread-local storage on the JVM side. This enables on-demand task retrieval
+/// instead of broadcasting all partition tasks to all executors.
+///
+/// # Safety
+/// This function is inherently unsafe since it deals with raw pointers passed from JNI.
+pub unsafe extern "system" fn Java_org_apache_comet_Native_getIcebergPartitionTasks<'local>(
+    mut e: JNIEnv<'local>,
+    _class: JClass,
+) -> jni::sys::jobject {
+    use jni::signature::ReturnType;
+
+    let jvm_classes = JVMClasses::get();
+
+    // Call Native.getIcebergPartitionTasksInternal() static method
+    let result = e.call_static_method_unchecked(
+        &jvm_classes.native.class,
+        jvm_classes
+            .native
+            .method_get_iceberg_partition_tasks_internal,
+        ReturnType::Array,
+        &[],
+    );
+
+    match result {
+        Ok(value) => match value.l() {
+            Ok(obj) => obj.as_raw(),
+            Err(_) => std::ptr::null_mut(),
+        },
+        Err(_) => std::ptr::null_mut(),
+    }
+}

native/core/src/execution/jni_api.rs

@@ -1144,14 +1144,69 @@ impl PhysicalPlanner {
                 let metadata_location = scan.metadata_location.clone();
 
                 debug_assert!(
-                    !scan.file_partitions.is_empty(),
-                    "IcebergScan must have at least one file partition. This indicates a bug in Scala serialization."
+                    scan.use_jni_task_retrieval || !scan.file_partitions.is_empty(),
+                    "IcebergScan must have use_jni_task_retrieval=true OR non-empty file_partitions. This indicates a bug in Scala serialization."
                 );
 
-                let tasks = parse_file_scan_tasks(
-                    scan,
-                    &scan.file_partitions[self.partition as usize].file_scan_tasks,
-                )?;
+                // Get file scan tasks either from JNI callback or from protobuf
+                let file_scan_tasks = if scan.use_jni_task_retrieval {
+                    // Call JNI to get partition-specific tasks from thread-local storage
+                    use crate::jvm_bridge::JVMClasses;
+                    use datafusion_comet_proto::spark_operator::IcebergFilePartition;
+                    use jni::objects::JByteArray;
+                    use jni::signature::ReturnType;
+                    use prost::Message;
+
+                    let mut env = JVMClasses::get_env()?;
+                    let jvm_classes = JVMClasses::get();
+
+                    // Call Native.getIcebergPartitionTasksInternal() directly
+                    let result = unsafe {
+                        env.call_static_method_unchecked(
+                            &jvm_classes.native.class,
+                            jvm_classes
+                                .native
+                                .method_get_iceberg_partition_tasks_internal,
+                            ReturnType::Array,
+                            &[],
+                        )
+                    };
+
+                    let result = result.map_err(|e| {
+                        ExecutionError::GeneralError(format!("JNI call failed: {}", e))
+                    })?;
+
+                    // Extract byte array from result
+                    let task_bytes_obj = result.l().map_err(|e| {
+                        ExecutionError::GeneralError(format!("Failed to extract JObject: {}", e))
+                    })?;
+                    let task_bytes_array: JByteArray = task_bytes_obj.into();
+
+                    if task_bytes_array.is_null() {
+                        return Err(ExecutionError::GeneralError(format!(
+                            "No partition tasks found for partition {} (JNI returned null). \
+                             This may indicate that partition tasks were not set in thread-local storage.",
+                            self.partition
+                        )));
+                    }
+
+                    // Convert JByteArray to Vec<u8>
+                    let task_bytes = env.convert_byte_array(&task_bytes_array).map_err(|e| {
+                        ExecutionError::GeneralError(format!("Failed to convert byte array: {}", e))
+                    })?;
+
+                    // Parse protobuf bytes into IcebergFilePartition
+                    let partition = IcebergFilePartition::decode(&task_bytes[..])?;
+
+                    partition.file_scan_tasks
+                } else {
+                    // Use tasks from protobuf (backward compatibility)
+                    scan.file_partitions[self.partition as usize]
+                        .file_scan_tasks
+                        .clone()
+                };
+
+                let tasks = parse_file_scan_tasks(scan, &file_scan_tasks)?;
                 let file_task_groups = vec![tasks];
 
                 let iceberg_scan = IcebergScanExec::new(

native/core/src/execution/planner.rs

@@ -174,11 +174,13 @@ pub use comet_exec::*;
 mod batch_iterator;
 mod comet_metric_node;
 mod comet_task_memory_manager;
+mod native;
 
 use crate::{errors::CometError, JAVA_VM};
 use batch_iterator::CometBatchIterator;
 pub use comet_metric_node::*;
 pub use comet_task_memory_manager::*;
+pub use native::*;
 
 /// The JVM classes that are used in the JNI calls.
 #[allow(dead_code)] // we need to keep references to Java items to prevent GC
@@ -207,6 +209,8 @@ pub struct JVMClasses<'a> {
     /// The CometTaskMemoryManager used for interacting with JVM side to
     /// acquire & release native memory.
     pub comet_task_memory_manager: CometTaskMemoryManager<'a>,
+    /// The Native object. Used for Iceberg partition task retrieval.
+    pub native: Native<'a>,
 }
 
 unsafe impl Send for JVMClasses<'_> {}
@@ -254,10 +258,71 @@ impl JVMClasses<'_> {
                 class_get_name_method,
                 throwable_get_message_method,
                 throwable_get_cause_method,
-                comet_metric_node: CometMetricNode::new(env).unwrap(),
-                comet_exec: CometExec::new(env).unwrap(),
-                comet_batch_iterator: CometBatchIterator::new(env).unwrap(),
-                comet_task_memory_manager: CometTaskMemoryManager::new(env).unwrap(),
+                comet_metric_node: {
+                    eprintln!(">> Initializing CometMetricNode...");
+                    match CometMetricNode::new(env) {
+                        Ok(node) => {
+                            eprintln!("   OK: CometMetricNode initialized");
+                            node
+                        }
+                        Err(e) => {
+                            eprintln!("   ERROR: CometMetricNode failed: {:?}", e);
+                            panic!("CometMetricNode initialization failed: {:?}", e);
+                        }
+                    }
+                },
+                comet_exec: {
+                    eprintln!(">> Initializing CometExec...");
+                    match CometExec::new(env) {
+                        Ok(exec) => {
+                            eprintln!("   OK: CometExec initialized");
+                            exec
+                        }
+                        Err(e) => {
+                            eprintln!("   ERROR: CometExec failed: {:?}", e);
+                            panic!("CometExec initialization failed: {:?}", e);
+                        }
+                    }
+                },
+                comet_batch_iterator: {
+                    eprintln!(">> Initializing CometBatchIterator...");
+                    match CometBatchIterator::new(env) {
+                        Ok(iter) => {
+                            eprintln!("   OK: CometBatchIterator initialized");
+                            iter
+                        }
+                        Err(e) => {
+                            eprintln!("   ERROR: CometBatchIterator failed: {:?}", e);
+                            panic!("CometBatchIterator initialization failed: {:?}", e);
+                        }
+                    }
+                },
+                comet_task_memory_manager: {
+                    eprintln!(">> Initializing CometTaskMemoryManager...");
+                    match CometTaskMemoryManager::new(env) {
+                        Ok(mgr) => {
+                            eprintln!("   OK: CometTaskMemoryManager initialized");
+                            mgr
+                        }
+                        Err(e) => {
+                            eprintln!("   ERROR: CometTaskMemoryManager failed: {:?}", e);
+                            panic!("CometTaskMemoryManager initialization failed: {:?}", e);
+                        }
+                    }
+                },
+                native: match Native::new(env) {
+                    Ok(native) => {
+                        eprintln!("✓ Successfully initialized Native JNI class");
+                        native
+                    }
+                    Err(e) => {
+                        eprintln!("✗ PANIC: Failed to initialize Native JNI class: {:?}", e);
+                        eprintln!("  Class name: org/apache/comet/NativeJNIBridge");
+                        eprintln!("  Method: getIcebergPartitionTasksInternal");
+                        eprintln!("  Signature: ()[B");
+                        panic!("Native JNI initialization failed: {:?}", e);
+                    }
+                },
             }
         });
     }

native/core/src/jvm_bridge/mod.rs

@@ -0,0 +1,75 @@
+// 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 jni::{
+    errors::Result as JniResult,
+    objects::{JClass, JStaticMethodID},
+    signature::ReturnType,
+    JNIEnv,
+};
+
+/// A struct that holds all the JNI methods and fields for JVM Native object.
+pub struct Native<'a> {
+    pub class: JClass<'a>,
+    pub method_get_iceberg_partition_tasks_internal: JStaticMethodID,
+    pub method_get_iceberg_partition_tasks_internal_ret: ReturnType,
+}
+
+impl<'a> Native<'a> {
+    pub const JVM_CLASS: &'static str = "org/apache/comet/NativeJNIBridge";
+
+    pub fn new(env: &mut JNIEnv<'a>) -> JniResult<Native<'a>> {
+        eprintln!("→ Initializing Native JNI class...");
+        eprintln!("  Looking up class: {}", Self::JVM_CLASS);
+
+        let class = match env.find_class(Self::JVM_CLASS) {
+            Ok(c) => {
+                eprintln!("  ✓ Found class: {}", Self::JVM_CLASS);
+                c
+            }
+            Err(e) => {
+                eprintln!("  ✗ Failed to find class: {}", Self::JVM_CLASS);
+                eprintln!("  Error: {:?}", e);
+                return Err(e);
+            }
+        };
+
+        eprintln!("  Looking up method: getIcebergPartitionTasksInternal with signature ()[B");
+        let method = match env.get_static_method_id(
+            Self::JVM_CLASS,
+            "getIcebergPartitionTasksInternal",
+            "()[B",
+        ) {
+            Ok(m) => {
+                eprintln!("  ✓ Found method: getIcebergPartitionTasksInternal");
+                m
+            }
+            Err(e) => {
+                eprintln!("  ✗ Failed to find method: getIcebergPartitionTasksInternal");
+                eprintln!("  Error: {:?}", e);
+                return Err(e);
+            }
+        };
+
+        eprintln!("✓ Native JNI class initialized successfully");
+        Ok(Native {
+            method_get_iceberg_partition_tasks_internal: method,
+            method_get_iceberg_partition_tasks_internal_ret: ReturnType::Array,
+            class,
+        })
+    }
+}

native/core/src/jvm_bridge/native.rs

@@ -164,6 +164,7 @@ message IcebergScan {
   map<string, string> catalog_properties = 2;
 
   // Pre-planned file scan tasks grouped by Spark partition
+  // DEPRECATED when use_jni_task_retrieval=true: Tasks retrieved via JNI callback instead
   repeated IcebergFilePartition file_partitions = 3;
 
   // Table metadata file path for FileIO initialization
@@ -178,6 +179,10 @@ message IcebergScan {
   repeated PartitionData partition_data_pool = 10;
   repeated DeleteFileList delete_files_pool = 11;
   repeated spark.spark_expression.Expr residual_pool = 12;
+
+  // When true, file_partitions is empty and tasks are retrieved via JNI callback
+  // This optimizes network/memory usage by avoiding broadcast of all partition tasks
+  bool use_jni_task_retrieval = 13;
 }
 
 // Helper message for deduplicating field ID lists

native/proto/src/proto/operator.proto

@@ -0,0 +1,41 @@
+/*
+ * 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.
+ */
+
+package org.apache.comet;
+
+/**
+ * JNI bridge for accessing Scala object methods as static methods.
+ *
+ * <p>This class provides static methods that can be called from native code via JNI, delegating to
+ * the Scala Native object singleton.
+ */
+public class NativeJNIBridge {
+
+  /**
+   * Gets Iceberg partition tasks for the current thread (JNI-accessible static method).
+   *
+   * <p>This method is called by native Rust code via JNI to retrieve partition-specific tasks
+   * during Iceberg scan execution.
+   *
+   * @return Serialized protobuf bytes containing IcebergFileScanTask messages, or null if not set
+   */
+  public static byte[] getIcebergPartitionTasksInternal() {
+    return Native$.MODULE$.getIcebergPartitionTasksInternal();
+  }
+}