Merge branch 'main' into windows_avg

Author: comphead

datafusion/functions-window/src/nth_value.rs

@@ -308,14 +308,22 @@ impl WindowUDFImpl for NthValue {
     }
 
     fn field(&self, field_args: WindowUDFFieldArgs) -> Result<FieldRef> {
-        let return_type = field_args
-            .input_fields()
-            .first()
-            .map(|f| f.data_type())
-            .cloned()
-            .unwrap_or(DataType::Null);
-
-        Ok(Field::new(field_args.name(), return_type, true).into())
+        let input_field =
+            field_args
+                .input_fields()
+                .first()
+                .cloned()
+                .unwrap_or_else(|| {
+                    Arc::new(Field::new(field_args.name(), DataType::Null, true))
+                });
+
+        // Clone the input field to preserve metadata, update name and nullability
+        Ok(input_field
+            .as_ref()
+            .clone()
+            .with_name(field_args.name())
+            .with_nullable(true)
+            .into())
     }
 
     fn reverse_expr(&self) -> ReversedUDWF {

datafusion/sqllogictest/test_files/metadata.slt

@@ -472,5 +472,51 @@ select arrow_metadata(with_metadata(id, 'unit', ''), 'unit') from table_with_met
 ----
 (empty)
 
+# Regression test: window functions should preserve field metadata
+# Test FIRST_VALUE window function preserves metadata
+query IT
+select
+  first_value(id) over (order by id asc nulls last) as fv,
+  arrow_metadata(first_value(id) over (order by id asc nulls last), 'metadata_key') as meta
+from table_with_metadata limit 1;
+----
+1 the id field
+
+# Test LAST_VALUE window function preserves metadata
+query IT
+select
+  last_value(id) over (order by id asc nulls last rows between unbounded preceding and unbounded following) as lv,
+  arrow_metadata(last_value(id) over (order by id asc nulls last rows between unbounded preceding and unbounded following), 'metadata_key') as meta
+from table_with_metadata limit 1;
+----
+NULL the id field
+
+# Test NTH_VALUE window function preserves metadata
+query IT
+select
+  nth_value(id, 2) over (order by id asc nulls last) as nv,
+  arrow_metadata(nth_value(id, 2) over (order by id asc nulls last), 'metadata_key') as meta
+from table_with_metadata limit 1;
+----
+NULL the id field
+
+# Test LEAD window function preserves metadata
+query IT
+select
+  lead(id) over (order by id asc nulls last) as ld,
+  arrow_metadata(lead(id) over (order by id asc nulls last), 'metadata_key') as meta
+from table_with_metadata limit 1;
+----
+3 the id field
+
+# Test LAG window function preserves metadata
+query IT
+select
+  lag(id) over (order by id asc nulls last) as lg,
+  arrow_metadata(lag(id) over (order by id asc nulls last), 'metadata_key') as meta
+from table_with_metadata limit 1;
+----
+NULL the id field
+
 statement ok
 drop table table_with_metadata;