perf: spill every shuffle partition of a task into one file

native/core/src/execution/jni_api.rs

@@ -100,7 +100,7 @@ use tokio::sync::mpsc;
 use crate::execution::memory_pools::{create_memory_pool, parse_memory_pool_config};
 use crate::execution::operators::{ScanExec, ShuffleScanExec};
 use crate::execution::shuffle::{
-    decode_remote_shuffle_batch, read_ipc_compressed, CompressionCodec,
+    decode_remote_shuffle_batch, read_ipc_compressed, CompressionCodec, ShuffleWriterExec,
 };
 use crate::execution::spark_plan::SparkPlan;
 
@@ -1155,6 +1155,60 @@ fn get_execution_context<'a>(id: i64) -> &'a mut ExecutionContext {
     }
 }
 
+/// Returns the partition offsets published by a finished native shuffle write.
+///
+/// The returned array holds `num_output_partitions + 1` offsets, the last being the total data
+/// file length.
+#[no_mangle]
+pub extern "system" fn Java_org_apache_comet_Native_getShufflePartitionOffsets(
+    e: EnvUnowned,
+    _class: JClass,
+    exec_context: jlong,
+) -> jlongArray {
+    try_unwrap_or_throw(&e, |env| {
+        let context = get_execution_context(exec_context);
+
+        let root_op = context.root_op.as_ref().ok_or_else(|| {
+            CometError::Internal(
+                "Cannot read shuffle partition offsets before the plan has been executed"
+                    .to_string(),
+            )
+        })?;
+
+        // `ExecutionPlan` has `Any` as a supertrait but no `as_any` method of its own, so upcast
+        // the trait object before downcasting to the writer.
+        let writer = (root_op.native_plan.as_ref() as &dyn std::any::Any)
+            .downcast_ref::<ShuffleWriterExec>()
+            .ok_or_else(|| {
+                CometError::Internal(
+                    "Shuffle partition offsets are only available on a native shuffle write plan"
+                        .to_string(),
+                )
+            })?;
+
+        let offsets = writer
+            .partition_offsets()
+            .ok_or_else(|| {
+                CometError::Internal(
+                    "Shuffle partition offsets are not published by a remote shuffle destination"
+                        .to_string(),
+                )
+            })?
+            .get()
+            .ok_or_else(|| {
+                CometError::Internal(
+                    "Shuffle writer has not published its partition offsets; the plan was not \
+                     drained to completion"
+                        .to_string(),
+                )
+            })?;
+
+        let long_array = env.new_long_array(offsets.len())?;
+        long_array.set_region(env, 0, offsets)?;
+        Ok(long_array.into_raw())
+    })
+}
+
 /// Used by Comet shuffle external sorter to write sorted records to disk.
 /// # Safety
 /// This function is inherently unsafe since it deals with raw pointers passed from JNI.

native/core/src/execution/planner.rs

@@ -47,7 +47,7 @@ use crate::execution::{
     planner::expression_registry::ExpressionRegistry,
     planner::operator_registry::OperatorRegistry,
     serde::{to_arrow_datatype, to_arrow_field},
-    shuffle::{SchemaAlignExec, ShuffleWriterDestination, ShuffleWriterExec},
+    shuffle::{PartitionOffsets, SchemaAlignExec, ShuffleWriterDestination, ShuffleWriterExec},
 };
 use crate::jvm_bridge::{jni_call, JVMClasses, ShufflePartitionPusher};
 use arrow::compute::CastOptions;
@@ -4143,7 +4143,7 @@ fn shuffle_writer_destination(
 
         return Ok(ShuffleWriterDestination::Local {
             output_data_file: writer.output_data_file.clone(),
-            output_index_file: writer.output_index_file.clone(),
+            partition_offsets: Arc::new(PartitionOffsets::default()),
         });
     };
 
@@ -4155,12 +4155,6 @@ fn shuffle_writer_destination(
                 ));
             }
 
-            if local.output_index_file.is_empty() {
-                return Err(GeneralError(
-                    "Local shuffle partition writer is missing its output index file".to_string(),
-                ));
-            }
-
             if !writer.output_data_file.is_empty()
                 && writer.output_data_file != local.output_data_file
             {
@@ -4171,16 +4165,6 @@ fn shuffle_writer_destination(
                 ));
             }
 
-            if !writer.output_index_file.is_empty()
-                && writer.output_index_file != local.output_index_file
-            {
-                return Err(GeneralError(
-                    "Local shuffle partition writer output index file conflicts with the legacy \
-                     shuffle output index file"
-                        .to_string(),
-                ));
-            }
-
             if shuffle_partition_pusher.is_some() {
                 return Err(GeneralError(
                     "Local shuffle partition writer cannot use a remote shuffle callback"
@@ -4190,11 +4174,11 @@ fn shuffle_writer_destination(
 
             Ok(ShuffleWriterDestination::Local {
                 output_data_file: local.output_data_file.clone(),
-                output_index_file: local.output_index_file.clone(),
+                partition_offsets: Arc::new(PartitionOffsets::default()),
             })
         }
         Some(spark_operator::partition_writer::Writer::Rss(_)) => {
-            if !writer.output_data_file.is_empty() || !writer.output_index_file.is_empty() {
+            if !writer.output_data_file.is_empty() {
                 return Err(GeneralError(
                     "RSS shuffle partition writer cannot have local output files".to_string(),
                 ));
@@ -5155,15 +5139,11 @@ mod tests {
         }
     }
 
-    fn local_shuffle_partition_writer(
-        output_data_file: &str,
-        output_index_file: &str,
-    ) -> spark_operator::PartitionWriter {
+    fn local_shuffle_partition_writer(output_data_file: &str) -> spark_operator::PartitionWriter {
         spark_operator::PartitionWriter {
             writer: Some(spark_operator::partition_writer::Writer::Local(
                 spark_operator::LocalPartitionWriter {
                     output_data_file: output_data_file.to_string(),
-                    output_index_file: output_index_file.to_string(),
                 },
             )),
         }
@@ -5180,15 +5160,15 @@ mod tests {
     fn assert_local_shuffle_destination(
         writer: &spark_operator::ShuffleWriter,
         expected_data_file: &str,
-        expected_index_file: &str,
     ) {
         match super::shuffle_writer_destination(writer, None).unwrap() {
             ShuffleWriterDestination::Local {
                 output_data_file,
-                output_index_file,
+                partition_offsets,
             } => {
                 assert_eq!(output_data_file, expected_data_file);
-                assert_eq!(output_index_file, expected_index_file);
+                // A fresh destination has not run a writer yet, so nothing is published.
+                assert!(partition_offsets.get().is_none());
             }
             destination => panic!("expected a local shuffle destination, got {destination:?}"),
         }
@@ -5198,49 +5178,38 @@ mod tests {
     fn shuffle_partition_writer_legacy_paths_remain_supported() {
         let writer = spark_operator::ShuffleWriter {
             output_data_file: "legacy.data".to_string(),
-            output_index_file: "legacy.index".to_string(),
             ..Default::default()
         };
 
-        assert_local_shuffle_destination(&writer, "legacy.data", "legacy.index");
+        assert_local_shuffle_destination(&writer, "legacy.data");
     }
 
     #[test]
     fn shuffle_partition_writer_uses_nested_local_paths() {
         let writer = spark_operator::ShuffleWriter {
-            partition_writer: Some(local_shuffle_partition_writer(
-                "shuffle.data",
-                "shuffle.index",
-            )),
+            partition_writer: Some(local_shuffle_partition_writer("shuffle.data")),
             ..Default::default()
         };
 
-        assert_local_shuffle_destination(&writer, "shuffle.data", "shuffle.index");
+        assert_local_shuffle_destination(&writer, "shuffle.data");
     }
 
     #[test]
     fn shuffle_partition_writer_accepts_matching_legacy_paths() {
         let writer = spark_operator::ShuffleWriter {
             output_data_file: "shuffle.data".to_string(),
-            output_index_file: "shuffle.index".to_string(),
-            partition_writer: Some(local_shuffle_partition_writer(
-                "shuffle.data",
-                "shuffle.index",
-            )),
+            partition_writer: Some(local_shuffle_partition_writer("shuffle.data")),
             ..Default::default()
         };
 
-        assert_local_shuffle_destination(&writer, "shuffle.data", "shuffle.index");
+        assert_local_shuffle_destination(&writer, "shuffle.data");
     }
 
     #[test]
     fn shuffle_partition_writer_rejects_conflicting_legacy_data_path() {
         let writer = spark_operator::ShuffleWriter {
             output_data_file: "legacy.data".to_string(),
-            partition_writer: Some(local_shuffle_partition_writer(
-                "shuffle.data",
-                "shuffle.index",
-            )),
+            partition_writer: Some(local_shuffle_partition_writer("shuffle.data")),
             ..Default::default()
         };
 
@@ -5251,29 +5220,11 @@ mod tests {
         );
     }
 
-    #[test]
-    fn shuffle_partition_writer_rejects_conflicting_legacy_index_path() {
-        let writer = spark_operator::ShuffleWriter {
-            output_index_file: "legacy.index".to_string(),
-            partition_writer: Some(local_shuffle_partition_writer(
-                "shuffle.data",
-                "shuffle.index",
-            )),
-            ..Default::default()
-        };
-
-        let error = super::shuffle_writer_destination(&writer, None).unwrap_err();
-        assert!(
-            error.to_string().contains("output index file conflicts"),
-            "unexpected error: {error}"
-        );
-    }
-
     #[test]
     fn shuffle_partition_writer_rejects_empty_local_data_path() {
         let writer = spark_operator::ShuffleWriter {
             output_data_file: "legacy.data".to_string(),
-            partition_writer: Some(local_shuffle_partition_writer("", "shuffle.index")),
+            partition_writer: Some(local_shuffle_partition_writer("")),
             ..Default::default()
         };
 
@@ -5284,21 +5235,6 @@ mod tests {
         );
     }
 
-    #[test]
-    fn shuffle_partition_writer_rejects_empty_local_index_path() {
-        let writer = spark_operator::ShuffleWriter {
-            output_index_file: "legacy.index".to_string(),
-            partition_writer: Some(local_shuffle_partition_writer("shuffle.data", "")),
-            ..Default::default()
-        };
-
-        let error = super::shuffle_writer_destination(&writer, None).unwrap_err();
-        assert!(
-            error.to_string().contains("missing its output index file"),
-            "unexpected error: {error}"
-        );
-    }
-
     #[test]
     fn shuffle_partition_writer_rejects_missing_destination() {
         let writer = spark_operator::ShuffleWriter {
@@ -5686,28 +5622,10 @@ mod tests {
         );
     }
 
-    #[test]
-    fn shuffle_partition_writer_rejects_rss_with_legacy_index_path() {
-        let writer = spark_operator::ShuffleWriter {
-            output_index_file: "legacy.index".to_string(),
-            partition_writer: Some(rss_shuffle_partition_writer()),
-            ..Default::default()
-        };
-        let callback: Arc<dyn ShufflePartitionPusher> =
-            Arc::new(RecordingShufflePartitionPusher::default());
-
-        let error = super::shuffle_writer_destination(&writer, Some(&callback)).unwrap_err();
-        assert!(
-            error.to_string().contains("cannot have local output files"),
-            "unexpected error: {error}"
-        );
-    }
-
     #[test]
     fn shuffle_partition_writer_rejects_callback_for_legacy_local_destination() {
         let writer = spark_operator::ShuffleWriter {
             output_data_file: "legacy.data".to_string(),
-            output_index_file: "legacy.index".to_string(),
             ..Default::default()
         };
         let callback: Arc<dyn ShufflePartitionPusher> =
@@ -5725,10 +5643,7 @@ mod tests {
     #[test]
     fn shuffle_partition_writer_rejects_callback_for_explicit_local_destination() {
         let writer = spark_operator::ShuffleWriter {
-            partition_writer: Some(local_shuffle_partition_writer(
-                "shuffle.data",
-                "shuffle.index",
-            )),
+            partition_writer: Some(local_shuffle_partition_writer("shuffle.data")),
             ..Default::default()
         };
         let callback: Arc<dyn ShufflePartitionPusher> =

native/proto/src/lib.rs

@@ -62,21 +62,16 @@ mod tests {
     struct LegacyShuffleWriter {
         #[prost(string, tag = "3")]
         output_data_file: String,
-        #[prost(string, tag = "4")]
-        output_index_file: String,
     }
 
     fn local_shuffle_writer() -> ShuffleWriter {
         let output_data_file = "/tmp/shuffle.data".to_string();
-        let output_index_file = "/tmp/shuffle.index".to_string();
 
         ShuffleWriter {
             output_data_file: output_data_file.clone(),
-            output_index_file: output_index_file.clone(),
             partition_writer: Some(PartitionWriter {
                 writer: Some(partition_writer::Writer::Local(LocalPartitionWriter {
                     output_data_file,
-                    output_index_file,
                 })),
             }),
             ..Default::default()
@@ -89,14 +84,12 @@ mod tests {
         let decoded = ShuffleWriter::decode(encoded.as_slice()).unwrap();
 
         assert_eq!(decoded.output_data_file, "/tmp/shuffle.data");
-        assert_eq!(decoded.output_index_file, "/tmp/shuffle.index");
         let Some(partition_writer::Writer::Local(local)) =
             decoded.partition_writer.and_then(|writer| writer.writer)
         else {
             panic!("expected a local shuffle partition writer");
         };
         assert_eq!(local.output_data_file, "/tmp/shuffle.data");
-        assert_eq!(local.output_index_file, "/tmp/shuffle.index");
     }
 
     #[test]
@@ -121,19 +114,16 @@ mod tests {
         let decoded = LegacyShuffleWriter::decode(encoded.as_slice()).unwrap();
 
         assert_eq!(decoded.output_data_file, "/tmp/shuffle.data");
-        assert_eq!(decoded.output_index_file, "/tmp/shuffle.index");
     }
 
     #[test]
     fn new_shuffle_writer_decodes_legacy_plan_without_destination() {
         let legacy = LegacyShuffleWriter {
             output_data_file: "/tmp/legacy.data".to_string(),
-            output_index_file: "/tmp/legacy.index".to_string(),
         };
         let decoded = ShuffleWriter::decode(legacy.encode_to_vec().as_slice()).unwrap();
 
         assert_eq!(decoded.output_data_file, "/tmp/legacy.data");
-        assert_eq!(decoded.output_index_file, "/tmp/legacy.index");
         assert!(decoded.partition_writer.is_none());
     }
 }