From c424d4477ceae2babd6e900dfa5ecb3ee37717f9 Mon Sep 17 00:00:00 2001
From: Tim Saucer <timsaucer@gmail.com>
Date: Fri, 20 Jun 2025 10:08:55 -0400
Subject: [PATCH] Column indices were not computed correctly, causing a panic

---
 .../ffi/src/udwf/partition_evaluator_args.rs  | 27 ++++++++++++-------
 1 file changed, 17 insertions(+), 10 deletions(-)

diff --git a/datafusion/ffi/src/udwf/partition_evaluator_args.rs b/datafusion/ffi/src/udwf/partition_evaluator_args.rs
index e74d47aa1a16..dffeb23741b6 100644
--- a/datafusion/ffi/src/udwf/partition_evaluator_args.rs
+++ b/datafusion/ffi/src/udwf/partition_evaluator_args.rs
@@ -75,17 +75,24 @@ impl TryFrom<PartitionEvaluatorArgs<'_>> for FFI_PartitionEvaluatorArgs {
             })
             .collect();
 
-        let max_column = required_columns.keys().max().unwrap_or(&0).to_owned();
-        let fields: Vec<_> = (0..max_column)
-            .map(|idx| match required_columns.get(&idx) {
-                Some((name, data_type)) => Field::new(*name, (*data_type).clone(), true),
-                None => Field::new(
-                    format!("ffi_partition_evaluator_col_{idx}"),
-                    DataType::Null,
-                    true,
-                ),
+        let max_column = required_columns.keys().max();
+        let fields: Vec<_> = max_column
+            .map(|max_column| {
+                (0..(max_column + 1))
+                    .map(|idx| match required_columns.get(&idx) {
+                        Some((name, data_type)) => {
+                            Field::new(*name, (*data_type).clone(), true)
+                        }
+                        None => Field::new(
+                            format!("ffi_partition_evaluator_col_{idx}"),
+                            DataType::Null,
+                            true,
+                        ),
+                    })
+                    .collect()
             })
-            .collect();
+            .unwrap_or_default();
+
         let schema = Arc::new(Schema::new(fields));
 
         let codec = DefaultPhysicalExtensionCodec {};