From b99dae219c9c8bd522bae3aae8f43ff1826ee7c6 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Mon, 24 Aug 2026 11:53:32 -0700 Subject: [PATCH 1/2] fix: make wide date-to-timestamp casts safe --- native/common/src/error.rs | 24 +++ .../src/conversion_funcs/temporal.rs | 188 +++++++++++++----- .../apache/comet/SparkErrorConverter.scala | 9 +- .../codegen/CometBatchKernelCodegen.scala | 4 + .../apache/comet/expressions/CometCast.scala | 46 ++++- .../apache/comet/CometNativeCastSuite.scala | 174 +++++++++++++++- .../comet/SparkErrorConverterSuite.scala | 8 + 7 files changed, 390 insertions(+), 63 deletions(-) diff --git a/native/common/src/error.rs b/native/common/src/error.rs index 81d095658e5..3935774d5a6 100644 --- a/native/common/src/error.rs +++ b/native/common/src/error.rs @@ -69,6 +69,10 @@ pub enum SparkError { #[error("[ARITHMETIC_OVERFLOW] {from_type} overflow. If necessary set \"spark.sql.ansi.enabled\" to \"false\" to bypass this error.")] ArithmeticOverflow { from_type: String }, + // Spark's checked date/timestamp conversions throw this even with ANSI disabled. + #[error("long overflow")] + LongOverflow, + #[error("[ARITHMETIC_OVERFLOW] Overflow in integral divide. Use 'try_divide' to tolerate overflow and return NULL instead. If necessary set \"spark.sql.ansi.enabled\" to \"false\" to bypass this error.")] IntegralDivideOverflow, @@ -269,6 +273,7 @@ impl SparkError { SparkError::CastOverFlow { .. } => "CastOverFlow", SparkError::CannotParseDecimal => "CannotParseDecimal", SparkError::ArithmeticOverflow { .. } => "ArithmeticOverflow", + SparkError::LongOverflow => "LongOverflow", SparkError::IntegralDivideOverflow => "IntegralDivideOverflow", SparkError::DecimalSumOverflow { .. } => "DecimalSumOverflow", SparkError::DivideByZero => "DivideByZero", @@ -580,6 +585,8 @@ impl SparkError { /// Returns the appropriate Spark exception class for this error pub fn exception_class(&self) -> &'static str { match self { + SparkError::LongOverflow => "java/lang/ArithmeticException", + // ArithmeticException SparkError::DivideByZero | SparkError::RemainderByZero @@ -684,6 +691,7 @@ impl SparkError { SparkError::RemainderByZero => Some("REMAINDER_BY_ZERO"), SparkError::IntervalDividedByZero => Some("INTERVAL_DIVIDED_BY_ZERO"), SparkError::ArithmeticOverflow { .. } => Some("ARITHMETIC_OVERFLOW"), + SparkError::LongOverflow => None, SparkError::IntegralDivideOverflow => Some("ARITHMETIC_OVERFLOW"), SparkError::DecimalSumOverflow { .. } => Some("ARITHMETIC_OVERFLOW"), SparkError::BinaryArithmeticOverflow { .. } => Some("BINARY_ARITHMETIC_OVERFLOW"), @@ -898,6 +906,22 @@ mod tests { assert!(json.contains("\"errorClass\":\"REMAINDER_BY_ZERO\"")); } + #[test] + fn test_long_overflow_json() { + let error = SparkError::LongOverflow; + let parsed: serde_json::Value = serde_json::from_str(&error.to_json()).unwrap(); + assert_eq!( + parsed, + serde_json::json!({ + "errorType": "LongOverflow", + "errorClass": "", + "params": {}, + }) + ); + assert_eq!(error.exception_class(), "java/lang/ArithmeticException"); + assert_eq!(error.to_string(), "long overflow"); + } + #[test] fn test_binary_overflow_json() { let error = SparkError::BinaryArithmeticOverflow { diff --git a/native/spark-expr/src/conversion_funcs/temporal.rs b/native/spark-expr/src/conversion_funcs/temporal.rs index bcc2b025503..ad8d8b49991 100644 --- a/native/spark-expr/src/conversion_funcs/temporal.rs +++ b/native/spark-expr/src/conversion_funcs/temporal.rs @@ -15,13 +15,11 @@ // specific language governing permissions and limitations // under the License. -use crate::utils::resolve_local_datetime; -use crate::{SparkCastOptions, SparkResult}; -use arrow::array::timezone::Tz; +use crate::{EvalMode, SparkCastOptions, SparkError, SparkResult}; use arrow::array::{ArrayRef, AsArray, TimestampMicrosecondBuilder}; use arrow::datatypes::{DataType, Date32Type}; -use chrono::NaiveDate; -use std::str::FromStr; +use arrow::error::ArrowError; +use chrono::FixedOffset; use std::sync::Arc; pub(crate) fn is_df_cast_from_date_spark_compatible(to_type: &DataType) -> bool { @@ -43,47 +41,41 @@ pub(crate) fn cast_date_to_timestamp( let date_array = array_ref.as_primitive::(); let mut builder = TimestampMicrosecondBuilder::with_capacity(date_array.len()); - if target_tz.is_none() { - // TIMESTAMP_NTZ: pure day arithmetic, no session-TZ offset. - // Matches Spark: daysToMicros(d, ZoneOffset.UTC) - for date in date_array.iter() { - match date { - Some(d) => builder.append_value((d as i64) * 86_400 * 1_000_000), - None => builder.append_null(), - } - } + // Date32 has a wider range than both chrono's calendar and an i64 microsecond timestamp. + // Use day arithmetic for NTZ and fixed-offset zones. Region-zone casts run through Spark's + // JVM codegen dispatcher (see CometCast.canCastFromDate), which owns their full-range rules. + let offset_seconds = if target_tz.is_none() + || cast_options.timezone.is_empty() + || cast_options.timezone == "UTC" + { + 0 } else { - // TIMESTAMP: midnight in session TZ → UTC epoch μs - let tz_str = if cast_options.timezone.is_empty() { - "UTC" - } else { - cast_options.timezone.as_str() - }; - // safe to unwrap since we are falling back to UTC above - let tz = Tz::from_str(tz_str)?; - let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap(); - for date in date_array.iter() { - match date { - Some(d) => { - // safe to unwrap since chrono's range ( 262,143 yrs) is higher than - // number of years possible with days as i32 (~ 6 mil yrs) - // convert date in session timezone to timestamp in UTC - let naive_date = epoch + chrono::Duration::days(d as i64); - let local_midnight = naive_date.and_hms_opt(0, 0, 0).unwrap(); - // Use resolve_local_datetime to correctly handle DST transitions: - // - Single: normal case, uses the given offset - // - Ambiguous (fall back): uses the earlier/DST occurrence, matching Spark - // - None (spring forward gap at midnight, e.g. America/Sao_Paulo): uses the - // pre-transition offset to compute the correct UTC time, matching Spark's - // LocalDate.atStartOfDay(zoneId) behaviour. - let local_midnight_in_microsec = - resolve_local_datetime(&tz, local_midnight).timestamp_micros(); - builder.append_value(local_midnight_in_microsec); - } - None => { - builder.append_null(); + // Reject any unexpected region-zone call safely instead of constructing a chrono date. + cast_options + .timezone + .parse::() + .map_err(|_| { + ArrowError::ParseError(format!( + "Invalid timezone \"{}\": expected a fixed offset", + cast_options.timezone + )) + })? + .local_minus_utc() + }; + for date in date_array.iter() { + match date { + Some(d) => { + // Date32 days fit i64 seconds. Apply the offset before checking the microsecond + // range, since the offset can move midnight across a timestamp range boundary. + let seconds = i64::from(d) * 86_400 - i64::from(offset_seconds); + match seconds.checked_mul(1_000_000) { + Some(micros) => builder.append_value(micros), + None if cast_options.eval_mode == EvalMode::Try => builder.append_null(), + // Spark's daysToMicros uses Math.multiplyExact even in non-ANSI mode. + None => return Err(SparkError::LongOverflow), } } + None => builder.append_null(), } } Ok(Arc::new( @@ -102,7 +94,7 @@ mod tests { use arrow::array::{Array, ArrayRef}; use arrow::datatypes::TimestampMicrosecondType; - // verifying epoch , DST change dates (US) and a null value (comprehensive tests on spark side) + // Region-zone and DST behavior is covered by Spark's codegen fallback tests. let dates: ArrayRef = Arc::new(Date32Array::from(vec![ Some(0), Some(19723), @@ -129,24 +121,23 @@ mod tests { assert_eq!(ts.value(2), dst_date); assert!(ts.is_null(3)); - // validate LA timezone (follows Daylight savings) + // Negative fixed offset let result = cast_date_to_timestamp( &dates, - &SparkCastOptions::new(EvalMode::Legacy, "America/Los_Angeles", false), + &SparkCastOptions::new(EvalMode::Legacy, "-08:00", false), &target_tz, ) .unwrap(); let ts = result.as_primitive::(); assert_eq!(ts.value(0), eight_hours_ts); assert_eq!(ts.value(1), non_dst_date + eight_hours_ts); - // should adjust for DST - assert_eq!(ts.value(2), dst_date + seven_hours_ts); + assert_eq!(ts.value(2), dst_date + eight_hours_ts); assert!(ts.is_null(3)); - // Phoenix timezone (does not follow Daylight savings) + // A different fixed offset let result = cast_date_to_timestamp( &dates, - &SparkCastOptions::new(EvalMode::Legacy, "America/Phoenix", false), + &SparkCastOptions::new(EvalMode::Legacy, "-07:00", false), &target_tz, ) .unwrap(); @@ -157,6 +148,103 @@ mod tests { assert!(ts.is_null(3)); } + #[test] + fn test_cast_wide_dates_to_timestamp() { + use arrow::array::{Array, Date32Array}; + use arrow::datatypes::TimestampMicrosecondType; + + // +262143-01-01, both valid UTC-midnight timestamp boundaries, and null. + let dates: ArrayRef = Arc::new(Date32Array::from(vec![ + Some(95_026_237), + Some(106_751_991), + Some(-106_751_991), + None, + ])); + for target_tz in [None, Some("UTC".into())] { + let result = cast_date_to_timestamp( + &dates, + &SparkCastOptions::new(EvalMode::Legacy, "UTC", false), + &target_tz, + ) + .unwrap(); + let ts = result.as_primitive::(); + assert_eq!(ts.value(0), 8_210_266_876_800_000_000); + assert_eq!(ts.value(1), 9_223_372_022_400_000_000); + assert_eq!(ts.value(2), -9_223_372_022_400_000_000); + assert!(ts.is_null(3)); + assert_eq!(ts.timezone(), target_tz.as_deref()); + } + + let result = cast_date_to_timestamp( + &dates, + &SparkCastOptions::new(EvalMode::Legacy, "+01:00", false), + &Some("UTC".into()), + ) + .unwrap(); + let ts = result.as_primitive::(); + assert_eq!(ts.value(0), 8_210_266_873_200_000_000); + assert_eq!(ts.value(1), 9_223_372_018_800_000_000); + assert_eq!(ts.value(2), -9_223_372_026_000_000_000); + assert!(ts.is_null(3)); + } + + #[test] + fn test_cast_date_to_timestamp_overflow() { + use arrow::array::{Array, Date32Array}; + use arrow::datatypes::TimestampMicrosecondType; + + for (timezone, days) in [ + ("UTC", 108_853_388), // +300000-06-15 + ("UTC", 106_751_992), + ("UTC", -106_751_992), + ("UTC", i32::MAX), + ("UTC", i32::MIN), + ("-18:00", 106_751_991), + ("+18:00", -106_751_991), + ] { + for target_tz in [None, Some("UTC".into())] { + // NTZ ignores the session offset: these midnight boundaries remain valid. + if target_tz.is_none() && timezone != "UTC" { + continue; + } + let dates: ArrayRef = + Arc::new(Date32Array::from(vec![Some(0), Some(days), None, Some(1)])); + for mode in [EvalMode::Legacy, EvalMode::Ansi] { + let result = cast_date_to_timestamp( + &dates, + &SparkCastOptions::new(mode, timezone, false), + &target_tz, + ); + assert!(matches!(result, Err(SparkError::LongOverflow))); + } + let result = cast_date_to_timestamp( + &dates, + &SparkCastOptions::new(EvalMode::Try, timezone, false), + &target_tz, + ) + .unwrap(); + let ts = result.as_primitive::(); + assert!(!ts.is_null(0)); + assert!(ts.is_null(1)); + assert!(ts.is_null(2)); + assert!(!ts.is_null(3)); + } + } + } + + #[test] + fn test_cast_date_to_timestamp_rejects_region_timezone() { + use arrow::array::Date32Array; + + let dates: ArrayRef = Arc::new(Date32Array::from(vec![95_026_237])); + let result = cast_date_to_timestamp( + &dates, + &SparkCastOptions::new(EvalMode::Legacy, "America/Los_Angeles", false), + &Some("UTC".into()), + ); + assert!(result.is_err()); + } + #[test] fn test_cast_date_to_timestamp_ntz() { use crate::EvalMode; diff --git a/spark/src/main/scala/org/apache/comet/SparkErrorConverter.scala b/spark/src/main/scala/org/apache/comet/SparkErrorConverter.scala index a6bc21aca71..853b0fa5094 100644 --- a/spark/src/main/scala/org/apache/comet/SparkErrorConverter.scala +++ b/spark/src/main/scala/org/apache/comet/SparkErrorConverter.scala @@ -117,8 +117,13 @@ object SparkErrorConverter extends ShimSparkErrorConverter { val summary: String = errorJson.summary.getOrElse("") - // Delegate to version-specific shim - let conversion exceptions propagate - val optEx = convertErrorType(errorJson.errorType, errorClass, params, sparkContext, summary) + // Math.multiplyExact throws a plain JVM exception in every Spark version, without an + // ANSI error class or configuration advice. Delegate other errors to the version-specific shim. + val optEx = if (errorJson.errorType == "LongOverflow") { + Some(new ArithmeticException("long overflow")) + } else { + convertErrorType(errorJson.errorType, errorClass, params, sparkContext, summary) + } optEx match { case Some(exception) => // successfully converted - return the proper typed exception diff --git a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala index c29300f9ed6..c9f55fcdfac 100644 --- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala +++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala @@ -28,6 +28,7 @@ import org.apache.spark.sql.catalyst.expressions.codegen._ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ +import org.apache.comet.expressions.CometCast import org.apache.comet.shims.{CometExprTraitShim, CometTypeShim} /** @@ -117,6 +118,9 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come * nested-field count on `spark.sql.codegen.maxFields`. */ def canHandle(boundExpr: Expression): Option[String] = { + if (boundExpr.exists(CometCast.requiresSparkDateTimestampTryCast)) { + return Some("TRY_CAST date-to-timestamp nullability requires Spark row execution") + } if (!isSupportedDataType(boundExpr.dataType)) { return Some(s"codegen dispatch: unsupported output type ${boundExpr.dataType}") } diff --git a/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala b/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala index d29ef7cd3b7..171d5126116 100644 --- a/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala +++ b/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala @@ -87,7 +87,9 @@ object CometCast // would only duplicate the matrix and risk drifting from it. override def getSupportLevel(cast: Cast): SupportLevel = { - if (cast.child.isInstanceOf[Literal]) { + if (requiresSparkDateTimestampTryCast(cast)) { + Unsupported(Some("TRY_CAST date-to-timestamp nullability requires Spark row execution")) + } else if (cast.child.isInstanceOf[Literal]) { // A cast whose child is a literal is folded by Spark at planning time via `cast.eval()` // (see `convert`), so the cast never executes natively and the result matches Spark by // definition. `CometLiteral` then validates the resulting literal's data type, except @@ -102,6 +104,31 @@ object CometCast } } + /** + * Spark's cast nullability inference can miss date-to-timestamp overflow. A TRY cast can + * therefore produce a null with non-nullable metadata, including map keys and struct fields. + * Keep those casts out of both native Arrow arrays and the Arrow-based codegen dispatcher. + */ + private[comet] def requiresSparkDateTimestampTryCast(expr: Expression): Boolean = expr match { + case cast: Cast if evalMode(cast) == CometEvalMode.TRY => + (!cast.nullable || isComplexType(cast.dataType)) && + containsDateTimestampCast(cast.child.dataType, cast.dataType) + case _ => false + } + + private def containsDateTimestampCast(fromType: DataType, toType: DataType): Boolean = + (fromType, toType) match { + case (DataTypes.DateType, TimestampType | TimestampNTZType) => true + case (ArrayType(from, _), ArrayType(to, _)) => containsDateTimestampCast(from, to) + case (MapType(fromKey, fromValue, _), MapType(toKey, toValue, _)) => + containsDateTimestampCast(fromKey, toKey) || containsDateTimestampCast(fromValue, toValue) + case (from: StructType, to: StructType) => + from.fields.zip(to.fields).exists { case (a, b) => + containsDateTimestampCast(a.dataType, b.dataType) + } + case _ => false + } + override def convert( cast: Cast, inputs: Seq[Attribute], @@ -254,7 +281,7 @@ object CometCast isSupported(from_map.valueType, to_map.valueType, timeZoneId, evalMode) case other => other } - case (DataTypes.DateType, toType) => canCastFromDate(toType, evalMode) + case (DataTypes.DateType, toType) => canCastFromDate(toType, timeZoneId, evalMode) case _ => unsupported(fromType, toType) } } @@ -460,10 +487,21 @@ object CometCast case _ => Unsupported(Some(s"Cast from DecimalType to $toType is not supported")) } - private def canCastFromDate(toType: DataType, evalMode: CometEvalMode.Value): SupportLevel = + private def canCastFromDate( + toType: DataType, + timeZoneId: Option[String], + evalMode: CometEvalMode.Value): SupportLevel = toType match { case DataTypes.TimestampType => - Compatible() + val timezone = timeZoneId.getOrElse("UTC") + if (timezone == "UTC" || timezone.matches("[+-][0-9]{2}:[0-9]{2}")) { + Compatible() + } else { + // DateType extends beyond chrono's calendar. Let Spark handle region-zone rules + // across the entire date range, including future DST and historical offsets. + Unsupported( + Some("Date-to-timestamp casts are native only in UTC or +/-HH:MM timezones")) + } case _: TimestampNTZType => Compatible() case DataTypes.BooleanType | DataTypes.ByteType | DataTypes.ShortType | diff --git a/spark/src/test/scala/org/apache/comet/CometNativeCastSuite.scala b/spark/src/test/scala/org/apache/comet/CometNativeCastSuite.scala index 9bfa8f63774..148aba9a923 100644 --- a/spark/src/test/scala/org/apache/comet/CometNativeCastSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometNativeCastSuite.scala @@ -31,6 +31,7 @@ import org.apache.spark.SparkConf import org.apache.spark.sql.{CometTestBase, DataFrame, Row, SaveMode} import org.apache.spark.sql.catalyst.expressions.Cast import org.apache.spark.sql.catalyst.parser.ParseException +import org.apache.spark.sql.comet.CometProjectExec import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.functions.{col, monotonically_increasing_id} import org.apache.spark.sql.internal.SQLConf @@ -39,6 +40,7 @@ import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, import org.apache.comet.expressions.{CometCast, CometEvalMode} import org.apache.comet.rules.CometScanTypeChecker import org.apache.comet.serde.{Compatible, Incompatible, Unsupported} +import org.apache.comet.udf.codegen.CometScalaUDFCodegen /** * Spark-parity coverage for Comet's **native** `Cast` implementation, and for the @@ -50,18 +52,21 @@ import org.apache.comet.serde.{Compatible, Incompatible, Unsupported} * decisions themselves (`Compatible` / `Incompatible` / `Unsupported`, and the reason strings) * are asserted directly against `CometCast.isSupported`. * - * Out of scope: the JVM codegen dispatch path. `CometCast` mixes in `CodegenDispatchFallback`, so - * a cast that `isSupported` rejects still runs inside the Comet operator by way of the - * Arrow-direct codegen dispatcher, evaluating Spark's own `Cast` against Arrow vectors. That path - * is covered by `CometCodegenSuite` and friends. Where this suite touches it at all, it does so - * only to pin that an `Unsupported` cast keeps the enclosing operator native rather than falling - * back to Spark; it does not attempt to cover dispatch semantics. + * `CometCast` mixes in `CodegenDispatchFallback`, so a cast that `isSupported` rejects still runs + * inside the Comet operator by way of the Arrow-direct codegen dispatcher, evaluating Spark's own + * `Cast` against Arrow vectors. That path is primarily covered by `CometCodegenSuite` and + * friends. Targeted routing regressions here also pin that an `Unsupported` cast uses the + * dispatcher, retaining Spark parity and keeping the enclosing operator native rather than + * falling back to Spark. * * Because of that split, a cast marked `Unsupported` here is not untested overall — it is tested * through the dispatcher instead. Adding a native cast implementation therefore means moving a * pair out of the `Unsupported` assertions and into the parity matrix below. */ -class CometNativeCastSuite extends CometTestBase with AdaptiveSparkPlanHelper { +class CometNativeCastSuite + extends CometTestBase + with AdaptiveSparkPlanHelper + with CometCodegenAssertions { import testImplicits._ @@ -1506,6 +1511,152 @@ class CometNativeCastSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + test("cast DateType to timestamps preserves wide make_date results") { + // Store the components, not dates: this both avoids Parquet rebasing and prevents Spark + // from constant-folding make_date before it reaches the native cast. + val input: Seq[(Option[Int], Int, Int)] = Seq( + (Some(262143), 1, 1), + (Some(-262144), 1, 1), + (Some(1970), 1, 1), + (Some(2024), 11, 3), + (None, 1, 1)) + withParquetTable(input, "wide_dates") { + for { + (target, tz) <- Seq( + ("TIMESTAMP_NTZ", "America/Los_Angeles"), + ("TIMESTAMP", "UTC"), + ("TIMESTAMP", "+05:30"), + ("TIMESTAMP", "-08:00")) + ansi <- Seq("false", "true") + } { + withSQLConf( + SQLConf.SESSION_LOCAL_TIMEZONE.key -> tz, + SQLConf.ANSI_ENABLED.key -> ansi, + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true") { + assertNoCodegenRan { + checkSparkAnswerAndOperator( + s"SELECT CAST(make_date(_1, _2, _3) AS $target) FROM wide_dates") + } + } + } + + // Named regions and offsets with seconds must use Spark's timezone rules instead of + // constructing chrono dates. Neither the wide positive nor negative year fits chrono. + Seq("America/Los_Angeles", "GMT", "+05:30:15").foreach { tz => + withSQLConf( + SQLConf.SESSION_LOCAL_TIMEZONE.key -> tz, + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true") { + assertCodegenRan { + checkSparkAnswerAndOperator( + "SELECT CAST(make_date(_1, _2, _3) AS TIMESTAMP) FROM wide_dates") + } + } + } + withSQLConf( + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "America/Los_Angeles", + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false", + CometConf.COMET_STRICT_TESTING.key -> "true") { + checkSparkAnswer("SELECT CAST(make_date(_1, _2, _3) AS TIMESTAMP) FROM wide_dates") + } + } + } + + test("cast DateType to timestamps checks microsecond overflow") { + // +/-106751991 are the last whole days representable in microseconds at UTC. The next + // days overflow; a fixed offset can also push an otherwise valid boundary day out of range. + val input = withNulls(Seq(-106751992, -106751991, 0, 106751991, 106751992)) + .map(Tuple1(_)) + val wideDate = + "make_date(CASE WHEN _1 = 0 THEN 300000 WHEN _1 > 0 THEN 2024 END, 6, 15)" + withParquetTable(input, "date_boundaries") { + for { + (target, tz, overflowDays) <- Seq( + ("TIMESTAMP_NTZ", "America/Los_Angeles", Seq(-106751992, 106751992)), + ("TIMESTAMP", "UTC", Seq(-106751992, 106751992)), + ("TIMESTAMP", "+18:00", Seq(-106751991)), + ("TIMESTAMP", "-18:00", Seq(106751991))) + ansi <- Seq("false", "true") + } { + withSQLConf( + SQLConf.SESSION_LOCAL_TIMEZONE.key -> tz, + SQLConf.ANSI_ENABLED.key -> ansi, + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true") { + val overflowQueries = overflowDays.map { day => + s"SELECT CAST(date_from_unix_date(_1) AS $target) " + + s"FROM date_boundaries WHERE _1 = $day" + } ++ (if (target == "TIMESTAMP_NTZ") { + Seq(s"SELECT CAST($wideDate AS $target) FROM date_boundaries") + } else Seq.empty) + overflowQueries.foreach { query => + assertNoCodegenRan { + val df = sql(query) + checkCometOperators(stripAQEPlan(df.queryExecution.executedPlan)) + val (sparkError, cometError) = checkSparkAnswerMaybeThrows(df) + assert(sparkError.isDefined, s"Spark should overflow for $query in $tz") + assert(cometError.isDefined, s"Comet should overflow for $query in $tz") + Seq(sparkError.get, cometError.get).foreach { error => + assert( + causeChain(error).exists { cause => + cause.getClass == classOf[ArithmeticException] && + cause.getMessage == "long overflow" + }, + s"Expected ArithmeticException: long overflow, got $error") + } + } + } + // A mixed batch must retain valid values and return NULL only for overflowing rows + // and input NULLs. TRY_CAST does this regardless of the session ANSI setting. + // Compare LTZ as epoch seconds (these date casts are second-aligned): Spark's + // java.sql.Timestamp row conversion overflows when rebasing the extreme boundaries. + val casts = Seq( + s"TRY_CAST(date_from_unix_date(_1) AS $target)", + s"TRY_CAST($wideDate AS $target)").map { expr => + if (target == "TIMESTAMP") s"CAST($expr AS BIGINT)" else expr + } + assertNoCodegenRan { + checkSparkAnswerAndOperator( + s"SELECT _1, ${casts.mkString(", ")} " + + "FROM date_boundaries ORDER BY _1") + } + } + } + } + } + + test("cast DateType to timestamps falls back for unsafe TRY_CAST nullability") { + withParquetTable(Seq(Tuple1(300000), Tuple1(2024)), "try_wide_dates") { + val date = "make_date(_1, 6, 15)" + for { + target <- Seq("TIMESTAMP", "TIMESTAMP_NTZ") + codegen <- Seq("true", "false") + } { + withSQLConf( + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC", + SQLConf.ANSI_ENABLED.key -> "true", + "spark.sql.legacy.castComplexTypesToString.enabled" -> "true", + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> codegen, + CometConf.COMET_STRICT_TESTING.key -> "true") { + val expressions = Seq( + s"TRY_CAST(make_date(coalesce(_1, 300000), 6, 15) AS $target)", + s"TRY_CAST(map(1, $date) AS MAP)", + s"TRY_CAST(named_struct('d', $date) AS STRUCT)", + // The outer cast would otherwise dispatch. Its entire tree must be checked. + s"CAST(array(TRY_CAST(make_date(coalesce(_1, 300000), 6, 15) AS $target)) AS STRING)") ++ + (if (target == "TIMESTAMP") Seq(s"TRY_CAST(map($date, 1) AS MAP<$target, INT>)") + else Seq.empty) + expressions.foreach { expr => + assertNoCodegenRan { + val query = s"SELECT $expr FROM try_wide_dates" + val plan = stripAQEPlan(sql(query).queryExecution.executedPlan) + assert(!plan.exists(_.isInstanceOf[CometProjectExec])) + checkSparkAnswer(query) + } + } + } + } + } + } + // CAST from TimestampType ignore("cast TimestampType to BooleanType") { @@ -2358,6 +2509,15 @@ class CometNativeCastSuite extends CometTestBase with AdaptiveSparkPlanHelper { values.map(v => Some(v)) ++ Seq(None) } + private def assertNoCodegenRan(f: => Unit): Unit = { + CometScalaUDFCodegen.resetStats() + f + val stats = CometScalaUDFCodegen.stats() + assert( + stats.compileCount + stats.cacheHitCount == 0, + s"Expected a native cast, but the JVM codegen dispatcher ran: $stats") + } + private def castFallbackTest( input: DataFrame, toType: DataType, diff --git a/spark/src/test/scala/org/apache/comet/SparkErrorConverterSuite.scala b/spark/src/test/scala/org/apache/comet/SparkErrorConverterSuite.scala index 3b81a76ead5..7da4779b531 100644 --- a/spark/src/test/scala/org/apache/comet/SparkErrorConverterSuite.scala +++ b/spark/src/test/scala/org/apache/comet/SparkErrorConverterSuite.scala @@ -23,6 +23,14 @@ import org.scalatest.funsuite.AnyFunSuite class SparkErrorConverterSuite extends AnyFunSuite { + test("LongOverflow converts to a plain ArithmeticException") { + val json = """{"errorType":"LongOverflow","errorClass":"","params":{}}""" + val ex = SparkErrorConverter.convertToSparkException( + new org.apache.comet.exceptions.CometQueryExecutionException(json)) + assert(ex.getClass == classOf[ArithmeticException]) + assert(ex.getMessage == "long overflow") + } + test("CannotReadFile converts to a FAILED_READ_FILE SparkException naming the file") { val ex = SparkErrorConverter .convertErrorType( From 8b072abb394dfdd389687d60b71b68f8459a8123 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Mon, 24 Aug 2026 13:29:50 -0700 Subject: [PATCH 2/2] fix: skip hidden child values in nested casts --- .../spark-expr/src/conversion_funcs/cast.rs | 195 +++++++++++++++++- .../apache/comet/CometNativeCastSuite.scala | 50 +++++ 2 files changed, 240 insertions(+), 5 deletions(-) diff --git a/native/spark-expr/src/conversion_funcs/cast.rs b/native/spark-expr/src/conversion_funcs/cast.rs index b1b71e8267d..71444f45659 100644 --- a/native/spark-expr/src/conversion_funcs/cast.rs +++ b/native/spark-expr/src/conversion_funcs/cast.rs @@ -50,9 +50,9 @@ use arrow::error::ArrowError; use arrow::{ array::{ cast::AsArray, Array, ArrayRef, Int16Array, Int32Array, Int64Array, Int8Array, - OffsetSizeTrait, + OffsetSizeTrait, UInt64Array, }, - compute::{cast_with_options, CastOptions}, + compute::{cast_with_options, take, CastOptions}, record_batch::RecordBatch, util::display::FormatOptions, }; @@ -230,6 +230,7 @@ pub(crate) fn cast_array( return Ok(Arc::new(array)); } + let array = prepare_nested_cast_input(array)?; let array = array_with_timezone(array, cast_options.timezone.clone(), Some(to_type))?; let eval_mode = cast_options.eval_mode; @@ -418,6 +419,36 @@ pub(crate) fn cast_array( Ok(spark_cast_postprocess(cast_result?, &from_type, to_type)) } +/// Recursive casts must not evaluate child values hidden by a null parent or outside a slice. +fn prepare_nested_cast_input(array: ArrayRef) -> DataFusionResult { + let has_unused_values = match array.data_type() { + DataType::Struct(_) => false, + DataType::List(_) => { + let list = array.as_list::(); + list.value_offsets()[0] != 0 + || list.value_offsets()[list.len()] as usize != list.values().len() + } + DataType::Map(_, _) => { + let map = array.as_map(); + map.value_offsets()[0] != 0 + || map.value_offsets()[map.len()] as usize != map.entries().len() + } + _ => return Ok(array), + }; + if array.null_count() == 0 && !has_unused_values { + return Ok(array); + } + + // Nullable identity indices preserve the rows and their validity. Arrow's take propagates + // struct nulls to children and compacts list/map entries, without introducing null map keys + // or changing field nullability. Contiguous, non-null containers need no normalization. + let indices = UInt64Array::new( + (0..array.len() as u64).collect::>().into(), + array.nulls().cloned(), + ); + Ok(take(array.as_ref(), &indices, None)?) +} + /// Determines if DataFusion supports the given cast in a way that is /// compatible with Spark fn is_datafusion_spark_compatible(from_type: &DataType, to_type: &DataType) -> bool { @@ -847,9 +878,11 @@ fn cast_binary_to_string( #[cfg(test)] mod tests { use super::*; - use arrow::array::{BinaryArray, ListArray, NullArray, PrimitiveArray, StringArray}; - use arrow::buffer::OffsetBuffer; - use arrow::datatypes::{Field, Fields, Int32Type, TimestampMicrosecondType}; + use arrow::array::{ + BinaryArray, Date32Array, ListArray, NullArray, PrimitiveArray, StringArray, + }; + use arrow::buffer::{NullBuffer, OffsetBuffer}; + use arrow::datatypes::{Field, Fields, Int32Type, TimeUnit, TimestampMicrosecondType}; #[test] fn test_cast_to_dictionary_is_rejected() { @@ -1123,6 +1156,158 @@ mod tests { assert!(result.is_err()); } + fn nested_date_cast_inputs( + days: Vec, + nulls: Option, + ) -> Vec<(ArrayRef, DataType)> { + let offsets = OffsetBuffer::new((0..=days.len() as i32).collect::>().into()); + let dates: ArrayRef = Arc::new(Date32Array::from(days)); + let timestamp = DataType::Timestamp(TimeUnit::Microsecond, None); + let fields = Fields::from(vec![Field::new("d", DataType::Date32, false)]); + let entries_fields = Fields::from(vec![ + Field::new("key", DataType::Date32, false), + Field::new("value", DataType::Date32, false), + ]); + vec![ + ( + Arc::new(StructArray::new( + fields, + vec![Arc::clone(&dates)], + nulls.clone(), + )), + DataType::Struct(Fields::from(vec![Field::new( + "d", + timestamp.clone(), + false, + )])), + ), + ( + Arc::new(ListArray::new( + Arc::new(Field::new("element", DataType::Date32, false)), + offsets.clone(), + Arc::clone(&dates), + nulls.clone(), + )), + DataType::List(Arc::new(Field::new("element", timestamp.clone(), false))), + ), + ( + Arc::new(MapArray::new( + Arc::new(Field::new( + "entries", + DataType::Struct(entries_fields.clone()), + false, + )), + offsets, + StructArray::new(entries_fields, vec![Arc::clone(&dates), dates], None), + nulls, + false, + )), + DataType::Map( + Arc::new(Field::new( + "entries", + DataType::Struct(Fields::from(vec![ + Field::new( + "key", + DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), + false, + ), + Field::new("value", timestamp, false), + ])), + false, + )), + false, + ), + ), + ] + } + + fn assert_nested_timestamp_values(array: &ArrayRef, expected: &[i64]) { + array.to_data().validate_full().unwrap(); + let values = match array.data_type() { + DataType::Struct(_) => array.as_struct().column(0), + DataType::List(_) => array.as_list::().values(), + DataType::Map(_, _) => { + let keys = array + .as_map() + .keys() + .as_primitive::(); + assert_eq!(keys.null_count(), 0); + assert_eq!(keys.values().as_ref(), expected); + array.as_map().values() + } + _ => unreachable!(), + }; + assert_eq!( + values + .as_primitive::() + .iter() + .flatten() + .collect::>(), + expected, + ); + } + + #[test] + fn test_cast_nested_dates_ignores_null_parent_values() { + // Non-nullable children contain overflowing dates beneath null parent rows, as IF can + // produce without changing its child buffers. Map keys must remain non-nullable. + for (array, to_type) in nested_date_cast_inputs( + vec![i32::MAX, 0, i32::MIN, 1], + Some(NullBuffer::from(vec![false, true, false, true])), + ) { + for mode in [EvalMode::Legacy, EvalMode::Ansi, EvalMode::Try] { + let options = SparkCastOptions::new(mode, "UTC", false); + let casted = cast_array(Arc::clone(&array), &to_type, &options).unwrap(); + assert_eq!(casted.data_type(), &to_type); + assert_eq!(casted.len(), 4); + assert!(casted.is_null(0) && casted.is_null(2)); + assert!(!casted.is_null(1) && !casted.is_null(3)); + assert_nested_timestamp_values(&casted, &[0, 86_400_000_000]); + + let all_null = cast_array(array.slice(0, 1), &to_type, &options).unwrap(); + assert_eq!(all_null.len(), 1); + assert!(all_null.is_null(0)); + assert_nested_timestamp_values(&all_null, &[]); + } + } + } + + #[test] + fn test_cast_nested_dates_ignores_values_outside_slice() { + for (array, to_type) in nested_date_cast_inputs(vec![i32::MAX, 0, 1, i32::MIN], None) { + for mode in [EvalMode::Legacy, EvalMode::Ansi, EvalMode::Try] { + let options = SparkCastOptions::new(mode, "UTC", false); + // Lists/maps keep their original child buffer even when the parent is sliced. + let casted = cast_array(array.slice(1, 2), &to_type, &options).unwrap(); + assert_eq!(casted.len(), 2); + assert_eq!(casted.null_count(), 0); + assert_nested_timestamp_values(&casted, &[0, 86_400_000_000]); + + let empty = cast_array(array.slice(1, 0), &to_type, &options).unwrap(); + assert_eq!(empty.len(), 0); + assert_nested_timestamp_values(&empty, &[]); + } + } + } + + #[test] + fn test_cast_nested_dates_still_checks_visible_overflow() { + for (array, to_type) in nested_date_cast_inputs( + vec![i32::MAX, i32::MAX, 0], + Some(NullBuffer::from(vec![false, true, true])), + ) { + for mode in [EvalMode::Legacy, EvalMode::Ansi] { + let error = cast_array( + Arc::clone(&array), + &to_type, + &SparkCastOptions::new(mode, "UTC", false), + ) + .unwrap_err(); + assert!(error.to_string().contains("long overflow")); + } + } + } + #[test] fn test_cast_struct_to_struct_drop_column() { let a: ArrayRef = Arc::new(Int32Array::from(vec![ diff --git a/spark/src/test/scala/org/apache/comet/CometNativeCastSuite.scala b/spark/src/test/scala/org/apache/comet/CometNativeCastSuite.scala index 148aba9a923..19a31ccc3bc 100644 --- a/spark/src/test/scala/org/apache/comet/CometNativeCastSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometNativeCastSuite.scala @@ -1623,6 +1623,56 @@ class CometNativeCastSuite } } + test("cast DateType to timestamps ignores values in null containers") { + withTempPath { path => + // Materialize the containers to prevent constant folding. The native IF can null a parent + // while retaining its children; keep hidden overflow and visible valid dates in one batch. + withSQLConf( + CometConf.COMET_ENABLED.key -> "false", + "spark.sql.parquet.datetimeRebaseModeInWrite" -> "CORRECTED") { + val date = "make_date(year, 6, 15)" + Seq((false, 300000), (true, 2024), (false, -300000), (true, 1970)) + .toDF("flag", "year") + .selectExpr( + "flag", + s"named_struct('d', $date) AS s", + s"map(1, $date) AS m", + s"map($date, 1) AS k", + s"array(named_struct('d', $date)) AS a") + .coalesce(1) + .write + .parquet(path.toString) + } + withParquetTable(path.toString, "masked_dates") { + for { + target <- Seq("TIMESTAMP", "TIMESTAMP_NTZ") + ansi <- Seq("false", "true") + } { + withSQLConf( + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC", + SQLConf.ANSI_ENABLED.key -> ansi, + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false") { + val casts = Seq( + ("s", s"STRUCT"), + ("m", s"MAP"), + ("a", s"ARRAY>")) ++ + // Spark rejects DATE -> TIMESTAMP_NTZ map keys in legacy mode. + (if (target == "TIMESTAMP") Seq(("k", s"MAP<$target, INT>")) else Seq.empty) + casts.foreach { case (column, targetType) => + // Referencing the alias twice keeps Spark from pushing CAST inside IF. CASE uses + // a different native path that already masks the children before casting them. + val query = s"SELECT masked, CAST(masked AS $targetType) FROM " + + s"(SELECT IF(flag, $column, NULL) AS masked FROM masked_dates)" + assertNoCodegenRan { + checkSparkAnswerAndOperator(sql(query), Seq(classOf[CometProjectExec])) + } + } + } + } + } + } + } + test("cast DateType to timestamps falls back for unsafe TRY_CAST nullability") { withParquetTable(Seq(Tuple1(300000), Tuple1(2024)), "try_wide_dates") { val date = "make_date(_1, 6, 15)"