From 77201b54ff3deb6be185fe0f2f2d5fc4cfe308d6 Mon Sep 17 00:00:00 2001 From: mazhengxuan Date: Wed, 5 Aug 2026 19:34:28 +0800 Subject: [PATCH 1/2] 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> { + let [data_type] = take_function_args(self.name(), arg_types)?; + + fn coerced_type(data_type: &DataType) -> Result { + 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 { 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); From 083616445cf6bf0c9d4ddf7727251613d71e4dab Mon Sep 17 00:00:00 2001 From: mazhengxuan Date: Thu, 6 Aug 2026 10:27:30 +0800 Subject: [PATCH 2/2] fix: handle null median in physical execution --- datafusion/functions-aggregate/src/median.rs | 60 ++++++++++++------- .../sqllogictest/test_files/aggregate.slt | 4 +- 2 files changed, 39 insertions(+), 25 deletions(-) diff --git a/datafusion/functions-aggregate/src/median.rs b/datafusion/functions-aggregate/src/median.rs index 383bf9e92e2ea..fb74da87c7fc8 100644 --- a/datafusion/functions-aggregate/src/median.rs +++ b/datafusion/functions-aggregate/src/median.rs @@ -40,19 +40,20 @@ use arrow::datatypes::{ }; use datafusion_common::hash_utils::RandomState; -use datafusion_common::utils::take_function_args; +use datafusion_common::types::{NativeType, logical_float64}; use datafusion_common::{ DataFusionError, Result, ScalarValue, assert_eq_or_internal_err, exec_datafusion_err, - exec_err, internal_datafusion_err, internal_err, + internal_datafusion_err, internal_err, }; use datafusion_expr::function::StateFieldsArgs; use datafusion_expr::{ - Accumulator, AggregateUDFImpl, Documentation, Signature, Volatility, - function::AccumulatorArgs, utils::format_state_name, + Accumulator, AggregateUDFImpl, Coercion, Documentation, Signature, TypeSignature, + TypeSignatureClass, Volatility, function::AccumulatorArgs, utils::format_state_name, }; use datafusion_expr::{EmitTo, GroupsAccumulator}; use datafusion_functions_aggregate_common::aggregate::groups_accumulator::accumulate::accumulate; use datafusion_functions_aggregate_common::aggregate::groups_accumulator::nulls::filtered_null_mask; +use datafusion_functions_aggregate_common::noop_accumulator::NoopAccumulator; use datafusion_functions_aggregate_common::utils::{GenericDistinctBuffer, Hashable}; use datafusion_macros::user_doc; use std::collections::HashMap; @@ -104,7 +105,22 @@ 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::user_defined(Volatility::Immutable), + 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, + ), } } } @@ -118,28 +134,22 @@ impl AggregateUDFImpl for Median { &self.signature } - fn coerce_types(&self, arg_types: &[DataType]) -> Result> { - let [data_type] = take_function_args(self.name(), arg_types)?; - - fn coerced_type(data_type: &DataType) -> Result { - 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 { Ok(arg_types[0].clone()) } fn state_fields(&self, args: StateFieldsArgs) -> Result> { + if args.input_fields[0].data_type().is_null() { + return Ok(vec![ + Field::new( + format_state_name(args.name, self.name()), + DataType::Null, + true, + ) + .into(), + ]); + } + //Intermediate state is a list of the elements we have collected so far let field = Field::new_list_field(args.input_fields[0].data_type().clone(), true); let state_name = if args.is_distinct { @@ -176,6 +186,10 @@ impl AggregateUDFImpl for Median { } let dt = acc_args.expr_fields[0].data_type().clone(); + if dt.is_null() { + return Ok(Box::new(NoopAccumulator::default())); + } + downcast_integer! { dt => (helper, dt), DataType::Float16 => helper!(Float16Type, dt), @@ -194,7 +208,7 @@ impl AggregateUDFImpl for Median { } fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool { - !args.is_distinct + !args.is_distinct && !args.expr_fields[0].data_type().is_null() } fn create_groups_accumulator( diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index 61a64c032c70a..1a3e3f5aa6653 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -1200,10 +1200,10 @@ select approx_median(NULL), arrow_typeof(approx_median(NULL)) from median_table; ---- NULL Null -query RT +query ?T select median(NULL), arrow_typeof(median(NULL)); ---- -NULL Float64 +NULL Null # median decimal statement ok