From 77201b54ff3deb6be185fe0f2f2d5fc4cfe308d6 Mon Sep 17 00:00:00 2001
From: mazhengxuan <mazhengxuan@didiglobal.com>
Date: Wed, 5 Aug 2026 19:34:28 +0800
Subject: [PATCH] fix: support untyped NULL input for median

---
 datafusion/functions-aggregate/src/median.rs  | 42 ++++++++++---------
 .../sqllogictest/test_files/aggregate.slt     |  5 +++
 2 files changed, 27 insertions(+), 20 deletions(-)

diff --git a/datafusion/functions-aggregate/src/median.rs b/datafusion/functions-aggregate/src/median.rs
index 81a3c076dffbe..383bf9e92e2ea 100644
--- a/datafusion/functions-aggregate/src/median.rs
+++ b/datafusion/functions-aggregate/src/median.rs
@@ -40,15 +40,15 @@ use arrow::datatypes::{
 };
 
 use datafusion_common::hash_utils::RandomState;
-use datafusion_common::types::{NativeType, logical_float64};
+use datafusion_common::utils::take_function_args;
 use datafusion_common::{
     DataFusionError, Result, ScalarValue, assert_eq_or_internal_err, exec_datafusion_err,
-    internal_datafusion_err, internal_err,
+    exec_err, internal_datafusion_err, internal_err,
 };
 use datafusion_expr::function::StateFieldsArgs;
 use datafusion_expr::{
-    Accumulator, AggregateUDFImpl, Coercion, Documentation, Signature, TypeSignature,
-    TypeSignatureClass, Volatility, function::AccumulatorArgs, utils::format_state_name,
+    Accumulator, AggregateUDFImpl, Documentation, Signature, Volatility,
+    function::AccumulatorArgs, utils::format_state_name,
 };
 use datafusion_expr::{EmitTo, GroupsAccumulator};
 use datafusion_functions_aggregate_common::aggregate::groups_accumulator::accumulate::accumulate;
@@ -104,22 +104,7 @@ impl Median {
             // Integer inputs are coerced to Float64 so the average of the two
             // middle values is not truncated. This matches DuckDB / PostgreSQL / Spark.
             // Float and Decimal inputs preserve their type.
-            signature: Signature::one_of(
-                vec![
-                    TypeSignature::Coercible(vec![Coercion::new_exact(
-                        TypeSignatureClass::Decimal,
-                    )]),
-                    TypeSignature::Coercible(vec![Coercion::new_exact(
-                        TypeSignatureClass::Float,
-                    )]),
-                    TypeSignature::Coercible(vec![Coercion::new_implicit(
-                        TypeSignatureClass::Native(logical_float64()),
-                        vec![TypeSignatureClass::Integer],
-                        NativeType::Float64,
-                    )]),
-                ],
-                Volatility::Immutable,
-            ),
+            signature: Signature::user_defined(Volatility::Immutable),
         }
     }
 }
@@ -133,6 +118,23 @@ impl AggregateUDFImpl for Median {
         &self.signature
     }
 
+    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
+        let [data_type] = take_function_args(self.name(), arg_types)?;
+
+        fn coerced_type(data_type: &DataType) -> Result<DataType> {
+            match data_type {
+                DataType::Dictionary(_, value_type) => coerced_type(value_type),
+                // Untyped NULL defaults to Float64, matching Signature::numeric.
+                DataType::Null => Ok(DataType::Float64),
+                data_type if data_type.is_integer() => Ok(DataType::Float64),
+                data_type if data_type.is_numeric() => Ok(data_type.clone()),
+                _ => exec_err!("Median not supported for {data_type}"),
+            }
+        }
+
+        Ok(vec![coerced_type(data_type)?])
+    }
+
     fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
         Ok(arg_types[0].clone())
     }
diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt
index 26b8a78f3921a..61a64c032c70a 100644
--- a/datafusion/sqllogictest/test_files/aggregate.slt
+++ b/datafusion/sqllogictest/test_files/aggregate.slt
@@ -1200,6 +1200,11 @@ select approx_median(NULL), arrow_typeof(approx_median(NULL)) from median_table;
 ----
 NULL Null
 
+query RT
+select median(NULL), arrow_typeof(median(NULL));
+----
+NULL Float64
+
 # median decimal
 statement ok
 create table t(c decimal(10, 4)) as values (0.0001), (0.0002), (0.0003), (0.0004), (0.0005), (0.0006);