diff --git a/datafusion/common/src/cast.rs b/datafusion/common/src/cast.rs index 28202c6684b50..7382bc9938739 100644 --- a/datafusion/common/src/cast.rs +++ b/datafusion/common/src/cast.rs @@ -286,3 +286,85 @@ pub fn as_generic_string_array( ) -> Result<&GenericStringArray> { Ok(downcast_value!(array, GenericStringArray, T)) } + +pub fn as_run_end_array(array: &dyn Array) -> Result<&RunArray> { + Ok(downcast_value!(array, RunArray, R)) +} + +pub fn run_end_array_values(array: &dyn Array) -> Result<&dyn Array> { + match array.data_type() { + DataType::RunEndEncoded(run_end_type, _) => match run_end_type.data_type() { + DataType::Int16 => { + let run_array = as_run_end_array::(array)?; + Ok(run_array.values()) + } + DataType::Int32 => { + let run_array = as_run_end_array::(array)?; + Ok(run_array.values()) + } + DataType::Int64 => { + let run_array = as_run_end_array::(array)?; + Ok(run_array.values()) + } + _ => Err(DataFusionError::Internal( + "Run end type not supported".to_string(), + )), + }, + _ => Err(DataFusionError::Internal(format!( + "Expected RunEndArray got {} instead ", + array.data_type() + ))), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_run_end_cast() { + // Create a values array + let values = Int16Array::from(vec![1, 2, 3]); + // Create run-ends array (indices where runs end) + let run_ends = Int16Array::from(vec![1, 2, 3]); + + // Create the RunEndEncoded array + let run_array = RunArray::::try_new(&run_ends, &values).unwrap(); + + // Test the cast + let result = as_run_end_array::(&run_array); + assert!(result.is_ok()); + } + + #[test] + fn test_valid_run_end_cast_with_nulls() { + // Create a values array with nulls + let values = Int32Array::from(vec![Some(10), None, Some(30)]); + // Create run-ends array (indices where runs end) + let run_ends = Int32Array::from(vec![1, 2, 3]); + + // Create the RunEndEncoded array + let run_array = RunArray::::try_new(&run_ends, &values).unwrap(); + + // Test the cast + let result = as_run_end_array::(&run_array); + assert!(result.is_ok()); + } + + #[test] + fn test_run_end_array_values() { + // Create a values array + let values = Int32Array::from(vec![10, 11, 20, 30]); + // Create run-ends array + let run_ends = Int32Array::from(vec![1, 2, 3, 4]); + + // Create the RunEndEncoded array + let run_array = RunArray::::try_new(&run_ends, &values).unwrap(); + + // Test getting the values + let result = run_end_array_values::(&run_array); + assert!(result.is_ok()); + let values_array = result.unwrap(); + assert_eq!(values_array.data_type(), &DataType::Int32); + } +}