From 5e4363141c7ef87b80ec2f3ffa98b16eb57cc758 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Wed, 9 Sep 2026 18:11:09 +0100 Subject: [PATCH] Cast stats to expected Arrow type using Arrow's cast semantics Signed-off-by: Adam Gutglick --- vortex-datafusion/src/persistent/format.rs | 127 +++++++++++++++++---- 1 file changed, 105 insertions(+), 22 deletions(-) diff --git a/vortex-datafusion/src/persistent/format.rs b/vortex-datafusion/src/persistent/format.rs index 6d05144e315..feb6c5c748b 100644 --- a/vortex-datafusion/src/persistent/format.rs +++ b/vortex-datafusion/src/persistent/format.rs @@ -5,6 +5,7 @@ use std::fmt::Debug; use std::fmt::Formatter; use std::sync::Arc; +use arrow_schema::DataType; use arrow_schema::Schema; use arrow_schema::SchemaRef; use async_trait::async_trait; @@ -671,36 +672,26 @@ impl FileFormat for VortexFormat { let column_size = stats_set.get_as::(Stat::UncompressedSizeInBytes, &PType::U64.into()); - let target_dtype = - session - .arrow() - .from_arrow_field(field.as_ref()) - .map_err(|e| { - DataFusionError::Execution(format!( - "Failed to derive Vortex DType for field {}: {e}", - field.name() - )) - })?; let min = scalar_stat_to_df( Stat::Min, stats_set.get(Stat::Min), stats_dtype, - &target_dtype, + field.data_type(), ); let max = scalar_stat_to_df( Stat::Max, stats_set.get(Stat::Max), stats_dtype, - &target_dtype, + field.data_type(), ); let null_count = stats_set.get_as::(Stat::NullCount, &PType::U64.into()); column_statistics.push(ColumnStatistics { null_count: null_count.to_df(), - min_value: min.to_df(), - max_value: max.to_df(), + min_value: min, + max_value: max, sum_value: DFPrecision::Absent, distinct_count: is_constant_to_distinct_count( stats_set.get_as::( @@ -782,20 +773,21 @@ fn scalar_stat_to_df( stat: Stat, value: Precision, stats_dtype: &DType, - target_dtype: &DType, -) -> Precision { + target_dtype: &DataType, +) -> DFPrecision { let Some(stat_dtype) = stat.dtype(stats_dtype) else { - return Precision::Absent; + return DFPrecision::Absent; }; value - .map(|stat_value| { - Scalar::try_new(stat_dtype, Some(stat_value))? - .cast(target_dtype)? + .and_then(|stat_value| { + let scalar = Scalar::try_new(stat_dtype, Some(stat_value)) + .ok()? .try_to_df() + .ok()?; + scalar.cast_to(target_dtype).ok() }) - .transpose() - .unwrap_or(Precision::Absent) + .to_df() } #[cfg(test)] @@ -806,6 +798,7 @@ mod tests { use arrow_array::Int32Array; use arrow_schema::DataType; use arrow_schema::Field; + use arrow_schema::TimeUnit; use datafusion_common::ScalarValue; use datafusion_common::config::ConfigOptions; use datafusion_expr::Operator; @@ -813,12 +806,102 @@ mod tests { use datafusion_physical_expr::expressions as df_expr; use datafusion_physical_expr::projection::ProjectionExprs; use datafusion_physical_plan::filter_pushdown::PushedDown; + use rstest::rstest; use vortex::expr::Expression; use super::*; use crate::common_tests::TestSessionContext; use crate::convert::DefaultExpressionConvertor; use crate::convert::ProcessedProjection; + use crate::convert::scalar_from_df; + + #[rstest] + #[case::timestamp_upscale( + ScalarValue::TimestampMillisecond(Some(1_234), None), + ScalarValue::TimestampMicrosecond(Some(1_234_000), None) + )] + #[case::timestamp_downscale( + ScalarValue::TimestampNanosecond(Some(-1_234_567), None), + ScalarValue::TimestampMicrosecond(Some(-1_234), None) + )] + #[case::timestamp_timezone( + ScalarValue::TimestampSecond(Some(123), Some("UTC".into())), + ScalarValue::TimestampMillisecond(Some(123_000), Some("UTC".into())) + )] + #[case::date_upscale(ScalarValue::Date32(Some(2)), ScalarValue::Date64(Some(172_800_000)))] + #[case::date_downscale(ScalarValue::Date64(Some(-172_800_000)), ScalarValue::Date32(Some(-2)))] + #[case::time_upscale( + ScalarValue::Time32Second(Some(123)), + ScalarValue::Time64Nanosecond(Some(123_000_000_000)) + )] + #[case::time_downscale( + ScalarValue::Time64Microsecond(Some(1_234_567)), + ScalarValue::Time32Millisecond(Some(1_234)) + )] + #[case::decimal_scale( + ScalarValue::Decimal32(Some(123), 5, 2), + ScalarValue::Decimal32(Some(1_230), 6, 3) + )] + #[case::integer_widening(ScalarValue::Int32(Some(123)), ScalarValue::Int64(Some(123)))] + #[case::string_view(ScalarValue::Utf8(Some("value".into())), ScalarValue::Utf8View(Some("value".into())))] + #[case::decimal_width( + ScalarValue::Decimal32(Some(123), 5, 2), + ScalarValue::Decimal128(Some(123), 5, 2) + )] + fn test_scalar_stat_to_df_cast( + #[case] value: ScalarValue, + #[case] expected: ScalarValue, + #[values(Stat::Min, Stat::Max)] stat: Stat, + #[values(true, false)] exact: bool, + ) -> VortexResult<()> { + let session = VortexSession::default(); + let scalar = scalar_from_df(&value, &session); + let value = scalar + .value() + .cloned() + .ok_or_else(|| vortex_err!("expected non-null scalar"))?; + let target_dtype = expected.data_type(); + let (value, expected) = if exact { + (Precision::Exact(value), DFPrecision::Exact(expected)) + } else { + (Precision::Inexact(value), DFPrecision::Inexact(expected)) + }; + + assert_eq!( + scalar_stat_to_df(stat, value, scalar.dtype(), &target_dtype), + expected + ); + Ok(()) + } + + #[rstest] + #[case::timestamp_overflow( + ScalarValue::TimestampSecond(Some(10_000_000_000), None), + DataType::Timestamp(TimeUnit::Nanosecond, None) + )] + #[case::timestamp_underflow( + ScalarValue::TimestampSecond(Some(-10_000_000_000), None), + DataType::Timestamp(TimeUnit::Nanosecond, None) + )] + #[case::integer_overflow(ScalarValue::Int64(Some(i64::MAX)), DataType::Int32)] + #[case::unsupported(ScalarValue::Boolean(Some(true)), DataType::Date32)] + fn test_scalar_stat_to_df_failed_cast( + #[case] value: ScalarValue, + #[case] target_dtype: DataType, + #[values(Stat::Min, Stat::Max)] stat: Stat, + ) -> VortexResult<()> { + let session = VortexSession::default(); + let scalar = scalar_from_df(&value, &session); + let value = scalar + .value() + .cloned() + .ok_or_else(|| vortex_err!("expected non-null scalar"))?; + assert_eq!( + scalar_stat_to_df(stat, Precision::Exact(value), scalar.dtype(), &target_dtype), + DFPrecision::Absent + ); + Ok(()) + } #[derive(Clone, Copy)] enum PushdownMode {