fix: use spill writer's schema instead of the first batch schema for spill files

Author: gruuya

datafusion/core/tests/union_nullable_spill.rs

@@ -0,0 +1,168 @@
+// 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 std::sync::Arc;
+
+use arrow::array::{Array, Decimal128Array, Int64Array, RecordBatch};
+use arrow::datatypes::{DataType, Field, Schema};
+use datafusion::datasource::memory::MemorySourceConfig;
+use datafusion_execution::config::SessionConfig;
+use datafusion_execution::memory_pool::FairSpillPool;
+use datafusion_execution::runtime_env::RuntimeEnvBuilder;
+use datafusion_physical_plan::repartition::RepartitionExec;
+use datafusion_physical_plan::union::UnionExec;
+use datafusion_physical_plan::{ExecutionPlan, Partitioning};
+use futures::StreamExt;
+
+/// Exercises spilling through UnionExec -> RepartitionExec where union children
+/// have mismatched nullability (one child's `val` is non-nullable, the other's
+/// is nullable with NULLs). A tiny FairSpillPool forces all batches to spill.
+///
+/// UnionExec returns child streams without schema coercion, so batches from
+/// different children carry different per-field nullability into the shared
+/// SpillPool. The IPC writer must use the SpillManager's canonical (nullable)
+/// schema — not the first batch's schema — so readback batches are valid.
+#[tokio::test]
+async fn test_union_repartition_spill_mixed_nullability() {
+    // Schema where `val` is NOT nullable (simulates literal-projection UNION branch)
+    let non_nullable_schema = Arc::new(Schema::new(vec![
+        Field::new("key", DataType::Int64, false),
+        Field::new("val", DataType::Decimal128(15, 7), false),
+    ]));
+
+    // Schema where `val` IS nullable (simulates table column-projection UNION branch with NULLs)
+    let nullable_schema = Arc::new(Schema::new(vec![
+        Field::new("key", DataType::Int64, false),
+        Field::new("val", DataType::Decimal128(15, 7), true),
+    ]));
+
+    // Generate many batches (4000 rows total) to ensure reliable spilling and coverage.
+    let num_batches: usize = 200;
+    let rows_per_batch: usize = 10;
+
+    let make_batches =
+        |schema: &Arc<Schema>, vals: Decimal128Array| -> Vec<RecordBatch> {
+            (0..num_batches)
+                .map(|i| {
+                    let start = (i * rows_per_batch) as i64;
+                    let keys: Vec<i64> = (start..start + rows_per_batch as i64).collect();
+                    RecordBatch::try_new(
+                        Arc::clone(schema),
+                        vec![Arc::new(Int64Array::from(keys)), Arc::new(vals.clone())],
+                    )
+                    .unwrap()
+                })
+                .collect()
+        };
+
+    let non_nullable_batches = make_batches(
+        &non_nullable_schema,
+        Decimal128Array::from(vec![0i128; rows_per_batch])
+            .with_precision_and_scale(15, 7)
+            .unwrap(),
+    );
+    let nullable_batches = make_batches(
+        &nullable_schema,
+        Decimal128Array::from(
+            (0..rows_per_batch)
+                .map(|j| {
+                    if j % 3 == 1 {
+                        None
+                    } else {
+                        Some(j as i128 * 10_000_000)
+                    }
+                })
+                .collect::<Vec<_>>(),
+        )
+        .with_precision_and_scale(15, 7)
+        .unwrap(),
+    );
+
+    // Create MemoryExec plans (1 partition each)
+    let non_nullable_exec = MemorySourceConfig::try_new_exec(
+        &[non_nullable_batches],
+        Arc::clone(&non_nullable_schema),
+        None,
+    )
+    .unwrap();
+
+    let nullable_exec = MemorySourceConfig::try_new_exec(
+        &[nullable_batches],
+        Arc::clone(&nullable_schema),
+        None,
+    )
+    .unwrap();
+
+    // UnionExec: declared schema coerces val to nullable
+    let union_exec = UnionExec::try_new(vec![non_nullable_exec, nullable_exec]).unwrap();
+
+    assert!(
+        union_exec.schema().field(1).is_nullable(),
+        "UnionExec schema should declare val as nullable"
+    );
+
+    // RepartitionExec: round-robin to 1 output partition.
+    // This ensures all batches from both children flow to the same output,
+    // sharing the same SpillPool and InProgressSpillFile.
+    let repartition =
+        RepartitionExec::try_new(union_exec, Partitioning::RoundRobinBatch(1)).unwrap();
+
+    // Use a tiny memory pool to force most batches to spill.
+    let session_config = SessionConfig::new().with_batch_size(2);
+    let runtime = RuntimeEnvBuilder::new()
+        .with_memory_pool(Arc::new(FairSpillPool::new(200)))
+        .build_arc()
+        .unwrap();
+    let task_ctx = Arc::new(
+        datafusion_execution::TaskContext::default()
+            .with_session_config(session_config)
+            .with_runtime(runtime),
+    );
+
+    // Execute partition 0 (the only output partition)
+    let mut stream = repartition.execute(0, task_ctx).unwrap();
+
+    let mut total_rows = 0usize;
+    let mut total_nulls = 0usize;
+    let mut batch_idx = 0usize;
+    while let Some(result) = stream.next().await {
+        let batch = result.unwrap();
+
+        RecordBatch::try_new(batch.schema(), batch.columns().to_vec()).unwrap_or_else(
+            |e| {
+                panic!(
+                    "Output batch {batch_idx} is invalid (schema says non-nullable \
+                     but data has nulls): {e}"
+                )
+            },
+        );
+
+        total_rows += batch.num_rows();
+        total_nulls += batch.column(1).null_count();
+        batch_idx += 1;
+    }
+
+    assert_eq!(
+        total_rows,
+        num_batches * rows_per_batch * 2,
+        "All rows from both UNION branches should be present"
+    );
+    assert!(
+        total_nulls > 0,
+        "Expected some null values in output (i.e. nullable batches were processed)"
+    );
+}

datafusion/physical-plan/src/spill/in_progress_spill_file.rs

@@ -62,7 +62,11 @@ impl InProgressSpillFile {
             ));
         }
         if self.writer.is_none() {
-            let schema = batch.schema();
+            // Use the SpillManager's declared schema rather than the batch's schema.
+            // Individual batches may have different schemas (e.g., different nullability)
+            // when they come from different branches of a UnionExec. The SpillManager's
+            // schema represents the canonical schema that all batches should conform to.
+            let schema = Arc::clone(self.spill_writer.schema());
             if let Some(in_progress_file) = &mut self.in_progress_file {
                 self.writer = Some(IPCStreamWriter::new(
                     in_progress_file.path(),
@@ -138,3 +142,121 @@ impl InProgressSpillFile {
         Ok(self.in_progress_file.take())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use arrow::array::{Array, Decimal128Array, Int64Array};
+    use arrow_schema::{DataType, Field, Schema};
+    use datafusion_execution::runtime_env::RuntimeEnvBuilder;
+    use datafusion_physical_expr_common::metrics::{
+        ExecutionPlanMetricsSet, SpillMetrics,
+    };
+    use futures::StreamExt;
+
+    /// Unit-level test: proves that InProgressSpillFile uses the SpillManager's
+    /// declared schema for the IPC writer, so readback batches always have the
+    /// correct schema even when input batches have mismatched nullability.
+    ///
+    /// Scenario:
+    ///   - SpillManager declares schema with nullable `val`
+    ///   - First appended batch has non-nullable `val` (simulates literal-projection UNION branch)
+    ///   - Second appended batch has nullable `val` with NULLs (simulates table UNION branch)
+    ///   - On readback, both batches must have the nullable schema
+    #[tokio::test]
+    async fn test_spill_file_uses_spill_manager_schema() {
+        let nullable_schema = Arc::new(Schema::new(vec![
+            Field::new("key", DataType::Int64, false),
+            Field::new("val", DataType::Decimal128(15, 7), true),
+        ]));
+        let non_nullable_schema = Arc::new(Schema::new(vec![
+            Field::new("key", DataType::Int64, false),
+            Field::new("val", DataType::Decimal128(15, 7), false),
+        ]));
+
+        let runtime = Arc::new(RuntimeEnvBuilder::new().build().unwrap());
+        let metrics_set = ExecutionPlanMetricsSet::new();
+        let spill_metrics = SpillMetrics::new(&metrics_set, 0);
+        let spill_manager = Arc::new(SpillManager::new(
+            runtime,
+            spill_metrics,
+            Arc::clone(&nullable_schema),
+        ));
+
+        let mut in_progress = spill_manager.create_in_progress_file("test").unwrap();
+
+        // First batch: non-nullable val (simulates literal-0 UNION branch)
+        let non_nullable_batch = RecordBatch::try_new(
+            Arc::clone(&non_nullable_schema),
+            vec![
+                Arc::new(Int64Array::from(vec![1, 2, 3])),
+                Arc::new(
+                    Decimal128Array::from(vec![0i128; 3])
+                        .with_precision_and_scale(15, 7)
+                        .unwrap(),
+                ),
+            ],
+        )
+        .unwrap();
+        in_progress.append_batch(&non_nullable_batch).unwrap();
+
+        // Second batch: nullable val with NULLs (simulates table UNION branch)
+        let nullable_batch = RecordBatch::try_new(
+            Arc::clone(&nullable_schema),
+            vec![
+                Arc::new(Int64Array::from(vec![4, 5, 6])),
+                Arc::new(
+                    Decimal128Array::from(vec![
+                        Some(10_000_000i128),
+                        None,
+                        Some(30_000_000i128),
+                    ])
+                    .with_precision_and_scale(15, 7)
+                    .unwrap(),
+                ),
+            ],
+        )
+        .unwrap();
+        in_progress.append_batch(&nullable_batch).unwrap();
+
+        let spill_file = in_progress.finish().unwrap().unwrap();
+
+        // Read back
+        let mut stream = spill_manager
+            .read_spill_as_stream(spill_file, None)
+            .unwrap();
+
+        assert!(
+            stream.schema().field(1).is_nullable(),
+            "Stream schema should be nullable"
+        );
+
+        let mut batches = vec![];
+        while let Some(result) = stream.next().await {
+            batches.push(result.unwrap());
+        }
+        assert_eq!(batches.len(), 2);
+
+        // Both readback batches must have the SpillManager's nullable schema
+        for (i, batch) in batches.iter().enumerate() {
+            assert!(
+                batch.schema().field(1).is_nullable(),
+                "Readback batch {i} should have nullable schema from SpillManager"
+            );
+        }
+
+        // The second batch must preserve its NULL data
+        let val_col = batches[1]
+            .column(1)
+            .as_any()
+            .downcast_ref::<Decimal128Array>()
+            .unwrap();
+        assert_eq!(val_col.null_count(), 1, "Second batch should have 1 null");
+
+        // Rebuilding the batch with its own schema must succeed (would fail if
+        // schema said non-nullable but data contained nulls)
+        RecordBatch::try_new(batches[1].schema(), batches[1].columns().to_vec()).expect(
+            "Readback batch should be valid: schema should match data nullability",
+        );
+    }
+}