feat: add partition writer destinations to native shuffle plans (2/n)

native/core/src/execution/planner.rs

@@ -1835,6 +1835,7 @@ impl PhysicalPlanner {
                     ))),
                 }?;
 
+                let (output_data_file, output_index_file) = shuffle_output_files(writer)?;
                 let write_buffer_size = writer.write_buffer_size as usize;
                 // Zero on the wire means the limit is disabled; normalize it here so the writer
                 // only ever sees a real limit or none at all.
@@ -1844,8 +1845,8 @@ impl PhysicalPlanner {
                     writer_input,
                     partitioning,
                     codec,
-                    writer.output_data_file.clone(),
-                    writer.output_index_file.clone(),
+                    output_data_file.to_string(),
+                    output_index_file.to_string(),
                     writer.tracing_enabled,
                     write_buffer_size,
                     max_buffer_bytes,
@@ -3891,6 +3892,65 @@ fn align_shuffle_writer_input(
         .map_err(|e| ExecutionError::DataFusionError(e.to_string()))
 }
 
+/// Resolves local shuffle output paths from the destination-aware writer descriptor.
+///
+/// Plans serialized before partition-writer descriptors were introduced only contain the
+/// top-level paths. New local plans populate both representations so older native binaries can
+/// still consume them; when both are present they must agree to avoid ambiguous destinations.
+fn shuffle_output_files(
+    writer: &spark_operator::ShuffleWriter,
+) -> Result<(&str, &str), ExecutionError> {
+    let Some(partition_writer) = writer.partition_writer.as_ref() else {
+        return Ok((&writer.output_data_file, &writer.output_index_file));
+    };
+
+    match partition_writer.writer.as_ref() {
+        Some(spark_operator::partition_writer::Writer::Local(local)) => {
+            if local.output_data_file.is_empty() {
+                return Err(GeneralError(
+                    "Local shuffle partition writer is missing its output data file".to_string(),
+                ));
+            }
+
+            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
+            {
+                return Err(GeneralError(
+                    "Local shuffle partition writer output data file conflicts with the legacy \
+                     shuffle output data file"
+                        .to_string(),
+                ));
+            }
+
+            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(),
+                ));
+            }
+
+            Ok((&local.output_data_file, &local.output_index_file))
+        }
+        Some(spark_operator::partition_writer::Writer::Rss(_)) => Err(GeneralError(
+            "RSS shuffle partition writers are not supported until remote shuffle execution is \
+             enabled"
+                .to_string(),
+        )),
+        None => Err(GeneralError(
+            "Shuffle partition writer has no destination".to_string(),
+        )),
+    }
+}
+
 /// Converts a protobuf PartitionValue to an iceberg Literal.
 ///
 fn partition_value_to_literal(
@@ -4723,6 +4783,166 @@ mod tests {
     };
     use datafusion_comet_spark_expr::EvalMode;
 
+    fn local_shuffle_partition_writer(
+        output_data_file: &str,
+        output_index_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(),
+                },
+            )),
+        }
+    }
+
+    #[test]
+    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_eq!(
+            super::shuffle_output_files(&writer).unwrap(),
+            ("legacy.data", "legacy.index")
+        );
+    }
+
+    #[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",
+            )),
+            ..Default::default()
+        };
+
+        assert_eq!(
+            super::shuffle_output_files(&writer).unwrap(),
+            ("shuffle.data", "shuffle.index")
+        );
+    }
+
+    #[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",
+            )),
+            ..Default::default()
+        };
+
+        assert_eq!(
+            super::shuffle_output_files(&writer).unwrap(),
+            ("shuffle.data", "shuffle.index")
+        );
+    }
+
+    #[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",
+            )),
+            ..Default::default()
+        };
+
+        let error = super::shuffle_output_files(&writer).unwrap_err();
+        assert!(
+            error.to_string().contains("output data file conflicts"),
+            "unexpected error: {error}"
+        );
+    }
+
+    #[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_output_files(&writer).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")),
+            ..Default::default()
+        };
+
+        let error = super::shuffle_output_files(&writer).unwrap_err();
+        assert!(
+            error.to_string().contains("missing its output data file"),
+            "unexpected error: {error}"
+        );
+    }
+
+    #[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_output_files(&writer).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 {
+            partition_writer: Some(spark_operator::PartitionWriter { writer: None }),
+            ..Default::default()
+        };
+
+        let error = super::shuffle_output_files(&writer).unwrap_err();
+        assert!(
+            error.to_string().contains("has no destination"),
+            "unexpected error: {error}"
+        );
+    }
+
+    #[test]
+    fn shuffle_partition_writer_rejects_rss_until_execution_is_supported() {
+        let writer = spark_operator::ShuffleWriter {
+            partition_writer: Some(spark_operator::PartitionWriter {
+                writer: Some(spark_operator::partition_writer::Writer::Rss(
+                    spark_operator::RssPartitionWriter {},
+                )),
+            }),
+            ..Default::default()
+        };
+
+        let error = super::shuffle_output_files(&writer).unwrap_err();
+        assert!(
+            error.to_string().contains("RSS shuffle partition writers"),
+            "unexpected error: {error}"
+        );
+    }
+
     #[test]
     fn test_unpack_dictionary_primitive() {
         let op_scan = Operator {

native/proto/src/lib.rs

@@ -50,3 +50,90 @@ pub mod spark_metric {
 pub mod spark_config {
     include!(concat!("generated", "/spark.spark_config.rs"));
 }
+
+#[cfg(test)]
+mod tests {
+    use super::spark_operator::{
+        partition_writer, LocalPartitionWriter, PartitionWriter, RssPartitionWriter, ShuffleWriter,
+    };
+    use prost::Message;
+
+    #[derive(Clone, PartialEq, prost::Message)]
+    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()
+        }
+    }
+
+    #[test]
+    fn shuffle_partition_writer_round_trips_local_destination() {
+        let encoded = local_shuffle_writer().encode_to_vec();
+        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]
+    fn shuffle_partition_writer_round_trips_rss_destination() {
+        let writer = ShuffleWriter {
+            partition_writer: Some(PartitionWriter {
+                writer: Some(partition_writer::Writer::Rss(RssPartitionWriter {})),
+            }),
+            ..Default::default()
+        };
+        let decoded = ShuffleWriter::decode(writer.encode_to_vec().as_slice()).unwrap();
+
+        assert!(matches!(
+            decoded.partition_writer.and_then(|writer| writer.writer),
+            Some(partition_writer::Writer::Rss(_))
+        ));
+    }
+
+    #[test]
+    fn legacy_shuffle_writer_decodes_new_plan_using_compatibility_paths() {
+        let encoded = local_shuffle_writer().encode_to_vec();
+        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());
+    }
+}

native/proto/src/proto/operator.proto

@@ -686,8 +686,28 @@ enum CompressionCodec {
   Gzip = 4;
 }
 
+// Selects where encoded shuffle partitions are written.
+message PartitionWriter {
+  oneof writer {
+    LocalPartitionWriter local = 1;
+    RssPartitionWriter rss = 2;
+  }
+}
+
+// Local shuffle output consists of a data file and its partition-offset index.
+message LocalPartitionWriter {
+  string output_data_file = 1;
+  string output_index_file = 2;
+}
+
+// Marker for remote shuffle output. The task-owned callback is bound outside
+// the serialized plan and is never transmitted through protobuf.
+message RssPartitionWriter {}
+
 message ShuffleWriter {
   spark.spark_partitioning.Partitioning partitioning = 1;
+  // Retained for compatibility with native binaries that predate partition_writer.
+  // Local plans also carry these paths in partition_writer.local.
   string output_data_file = 3;
   string output_index_file = 4;
   CompressionCodec codec = 5;
@@ -705,6 +725,8 @@ message ShuffleWriter {
   // Maximum number of bytes that the writer buffers in memory before spilling to disk.
   // Zero means no limit, in which case spilling is driven only by memory pool pressure.
   uint64 max_buffer_bytes = 10;
+  // Explicit output destination. When absent, use the legacy output file fields.
+  PartitionWriter partition_writer = 11;
 }
 
 message ParquetWriter {