diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 0a71607ead..7a304eeeb8 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -401,6 +401,7 @@ jobs: org.apache.comet.CometCodegenHOFSuite org.apache.comet.CometFuzzMathSuite org.apache.comet.CometCodegenFuzzSuite + org.apache.comet.CometScalaUDFClassLoaderSuite org.apache.comet.codegen.CometSpecializedGettersDispatchSuite org.apache.comet.CometStringDecodeSuite org.apache.comet.CometWidthBucketSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 2335588b70..46f672a56a 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -217,6 +217,7 @@ jobs: org.apache.comet.CometCodegenHOFSuite org.apache.comet.CometFuzzMathSuite org.apache.comet.CometCodegenFuzzSuite + org.apache.comet.CometScalaUDFClassLoaderSuite org.apache.comet.codegen.CometSpecializedGettersDispatchSuite org.apache.comet.CometStringDecodeSuite org.apache.comet.CometWidthBucketSuite diff --git a/docs/source/user-guide/latest/compatibility/index.md b/docs/source/user-guide/latest/compatibility/index.md index 5041856e5b..ae29a67bc1 100644 --- a/docs/source/user-guide/latest/compatibility/index.md +++ b/docs/source/user-guide/latest/compatibility/index.md @@ -121,10 +121,6 @@ divergence: raise raw Arrow errors that bypass `SparkErrorConverter` and surface as `CometNativeException` rather than `SparkArithmeticException` with the proper error class and query context ([#5072](https://github.com/apache/datafusion-comet/issues/5072)). -- `next_day` and `make_date` throw at the correct inputs but surface as `CometNativeException` - instead of `SparkIllegalArgumentException [ILLEGAL_DAY_OF_WEEK]` / - `SparkDateTimeException [DATETIME_FIELD_OUT_OF_BOUNDS.WITH_SUGGESTION]` - ([#5073](https://github.com/apache/datafusion-comet/issues/5073)). - Spark 4.2 introduced additional ANSI arithmetic overflow behavior differences that Comet does not yet track ([#4967](https://github.com/apache/datafusion-comet/issues/4967)). diff --git a/docs/source/user-guide/latest/scala_java_udfs.md b/docs/source/user-guide/latest/scala_java_udfs.md index 1ebf79250f..3a55982bbc 100644 --- a/docs/source/user-guide/latest/scala_java_udfs.md +++ b/docs/source/user-guide/latest/scala_java_udfs.md @@ -54,6 +54,7 @@ When a UDF is rejected, the reason surfaces through Comet's standard fallback di - Non-deterministic expressions referenced from the argument tree (`rand`, `uuid`, `monotonically_increasing_id`) produce per-partition sequences consistent with Spark. - `TaskContext.get()` inside the user function returns the driving Spark task's context. +- The Spark task thread's context ClassLoader is propagated to the thread that runs the user function, so functions defined in jars supplied with `--jars` / `spark.jars` resolve the same way they do under Spark's own execution. - The user function must be closure-serializable; the same function that works with Spark's executor execution works here. ## Known limitations diff --git a/native/common/src/error.rs b/native/common/src/error.rs index baeb3a119e..81d095658e 100644 --- a/native/common/src/error.rs +++ b/native/common/src/error.rs @@ -101,6 +101,12 @@ pub enum SparkError { #[error("[DATETIME_OVERFLOW] Datetime arithmetic overflow.")] DatetimeOverflow, + #[error("[ILLEGAL_DAY_OF_WEEK] Illegal input for day of week: {input}.")] + IllegalDayOfWeek { input: String }, + + #[error("[DATETIME_FIELD_OUT_OF_BOUNDS] {range_message}. If necessary set \"spark.sql.ansi.enabled\" to \"false\" to bypass this error.")] + DatetimeFieldOutOfBounds { range_message: String }, + #[error("[INVALID_ARRAY_INDEX] The index {index_value} is out of bounds. The array has {array_size} elements. Use the SQL function get() to tolerate accessing element at invalid index and return NULL instead. If necessary set \"spark.sql.ansi.enabled\" to \"false\" to bypass this error.")] InvalidArrayIndex { index_value: i32, array_size: i32 }, @@ -276,6 +282,8 @@ impl SparkError { "IntervalArithmeticOverflowWithoutSuggestion" } SparkError::DatetimeOverflow => "DatetimeOverflow", + SparkError::IllegalDayOfWeek { .. } => "IllegalDayOfWeek", + SparkError::DatetimeFieldOutOfBounds { .. } => "DatetimeFieldOutOfBounds", SparkError::InvalidArrayIndex { .. } => "InvalidArrayIndex", SparkError::InvalidElementAtIndex { .. } => "InvalidElementAtIndex", SparkError::InvalidBitmapPosition { .. } => "InvalidBitmapPosition", @@ -458,6 +466,16 @@ impl SparkError { "suggestedFunc": suggested_func, }) } + SparkError::IllegalDayOfWeek { input } => { + serde_json::json!({ + "string": input, + }) + } + SparkError::DatetimeFieldOutOfBounds { range_message } => { + serde_json::json!({ + "rangeMessage": range_message, + }) + } SparkError::InvalidFractionOfSecond { value } => { serde_json::json!({ "value": value, @@ -605,11 +623,17 @@ impl SparkError { // DateTimeException SparkError::InvalidInputInCastToDatetime { .. } | SparkError::CannotParseTimestamp { .. } - | SparkError::InvalidFractionOfSecond { .. } => "org/apache/spark/SparkDateTimeException", + | SparkError::InvalidFractionOfSecond { .. } + | SparkError::DatetimeFieldOutOfBounds { .. } => { + "org/apache/spark/SparkDateTimeException" + } // IllegalArgumentException SparkError::DatatypeCannotOrder { .. } - | SparkError::InvalidUtf8String { .. } => "org/apache/spark/SparkIllegalArgumentException", + | SparkError::InvalidUtf8String { .. } + | SparkError::IllegalDayOfWeek { .. } => { + "org/apache/spark/SparkIllegalArgumentException" + } // FileNotFound - will be converted to SparkFileNotFoundException by the shim SparkError::FileNotFound { .. } => "org/apache/spark/SparkException", @@ -693,6 +717,8 @@ impl SparkError { // DateTime errors SparkError::CannotParseTimestamp { .. } => Some("CANNOT_PARSE_TIMESTAMP"), SparkError::InvalidFractionOfSecond { .. } => Some("INVALID_FRACTION_OF_SECOND"), + SparkError::IllegalDayOfWeek { .. } => Some("ILLEGAL_DAY_OF_WEEK"), + SparkError::DatetimeFieldOutOfBounds { .. } => Some("DATETIME_FIELD_OUT_OF_BOUNDS"), // String/UTF8 errors SparkError::InvalidUtf8String { .. } => Some("INVALID_UTF8_STRING"), diff --git a/native/common/src/lib.rs b/native/common/src/lib.rs index b2d4a57431..c4153606a1 100644 --- a/native/common/src/lib.rs +++ b/native/common/src/lib.rs @@ -17,9 +17,11 @@ mod error; mod query_context; +mod schema; pub mod tracing; mod utils; pub use error::{decimal_overflow_error, SparkError, SparkErrorWithContext, SparkResult}; pub use query_context::{create_query_context_map, QueryContext, QueryContextMap}; +pub use schema::{cast_and_stamp_schema, widen_nested_nullability}; pub use utils::{bytes_to_i128, decode_utf8_spark_lossy}; diff --git a/native/common/src/schema.rs b/native/common/src/schema.rs new file mode 100644 index 0000000000..2b5995c8c2 --- /dev/null +++ b/native/common/src/schema.rs @@ -0,0 +1,378 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Helpers for reconciling the Arrow type an operator actually produced with the type Spark +//! catalyst declared. +//! +//! Arrow treats nested field nullability as part of `DataType` identity, so a column whose nested +//! `nullable` flags are narrower than the declared type is rejected by +//! `RecordBatch::try_new_with_options` even though the data is fine — a non-null child is a strict +//! subset of a nullable one. Every Comet boundary that stamps a declared schema onto a runtime +//! array therefore normalizes first instead of asserting. See +//! for the boundaries and +//! for the upstream drift itself. + +use arrow::array::{ArrayRef, RecordBatch, RecordBatchOptions}; +use arrow::compute::{cast_with_options, CastOptions}; +use arrow::datatypes::{DataType, FieldRef, SchemaRef}; +use datafusion::common::DataFusionError; +use std::sync::Arc; + +/// Builds a `RecordBatch` with `schema` from `columns`, casting any column whose Arrow type +/// differs from the declared field type. `operator` names the caller and is used only in error +/// messages. +/// +/// This is the normalizing counterpart to stamping `schema` on directly: it absorbs the +/// return-type drift — most commonly a nested `nullable` flag — that native kernels and non-Parquet +/// sources introduce, the same way `ScanExec` absorbs it at the FFI boundary. +/// +/// Reconciliation follows `schema` in both directions, so it will also narrow a nullable nested +/// child to non-null when that is what `schema` declares. That is not silently lossy: arrow's +/// `StructArray::try_new` rejects unmasked nulls under a non-nullable field, so data that cannot +/// survive the narrowing errors here rather than producing an array that misreports itself. +/// Callers that must not narrow should widen `schema` first — see [`widen_nested_nullability`], +/// which is what `SchemaAlignExec` and `ExpandExec` do. +pub fn cast_and_stamp_schema( + operator: &str, + schema: &SchemaRef, + mut columns: Vec, + num_rows: usize, +) -> Result { + if columns.len() != schema.fields().len() { + return Err(DataFusionError::Internal(format!( + "{operator} produced {} columns but its schema declares {}", + columns.len(), + schema.fields().len() + ))); + } + + for (idx, (column, field)) in columns.iter_mut().zip(schema.fields()).enumerate() { + if column.data_type() != field.data_type() { + *column = cast_with_options(column, field.data_type(), &CastOptions::default()) + .map_err(|e| cast_error(operator, schema, idx, column.data_type(), e))?; + } + } + + // Every column's type now matches its declared field — `cast_with_options` returns either an + // error or an array of exactly the requested type — so the stamp can only fail on row counts. + let options = RecordBatchOptions::new().with_row_count(Some(num_rows)); + RecordBatch::try_new_with_options(Arc::clone(schema), columns, &options).map_err(|e| { + DataFusionError::Context( + format!("{operator} cannot build a batch of {num_rows} rows with its declared schema"), + Box::new(DataFusionError::from(e)), + ) + }) +} + +/// Names the operator and the dotted path of the column that could not be reconciled, since +/// arrow's own message reports only `at column index N` and the two printed types may differ by a +/// single flag hundreds of characters in. +fn cast_error( + operator: &str, + schema: &SchemaRef, + idx: usize, + actual: &DataType, + source: arrow::error::ArrowError, +) -> DataFusionError { + let field = schema.field(idx); + let detail = describe_type_mismatch(field.name(), field.data_type(), actual) + .unwrap_or_else(|| format!("{}: expected {}", field.name(), field.data_type())); + DataFusionError::Context( + format!("{operator} cannot reconcile col[{idx}] with its declared schema at {detail}"), + Box::new(DataFusionError::from(source)), + ) +} + +/// Describes where `expected` and `actual` first diverge, as a dotted path rooted at `path`. +/// Returns `None` when the two types are equal. +/// +/// The nested arms cover the shapes Comet actually builds — see `make_all_fields_nullable` in the +/// planner and `to_arrow_datatype` in the serde layer, which walk the same set. +fn describe_type_mismatch(path: &str, expected: &DataType, actual: &DataType) -> Option { + if expected == actual { + return None; + } + match (expected, actual) { + (DataType::List(e), DataType::List(a)) + | (DataType::LargeList(e), DataType::LargeList(a)) => { + describe_field_mismatch(&format!("{path}.element"), e, a) + } + (DataType::FixedSizeList(e, e_len), DataType::FixedSizeList(a, a_len)) + if e_len == a_len => + { + describe_field_mismatch(&format!("{path}.element"), e, a) + } + (DataType::Map(e, e_sorted), DataType::Map(a, a_sorted)) if e_sorted == a_sorted => { + describe_field_mismatch(&format!("{path}.entries"), e, a) + } + (DataType::Struct(e), DataType::Struct(a)) if e.len() == a.len() => e + .iter() + .zip(a.iter()) + .find_map(|(e, a)| describe_field_mismatch(&format!("{path}.{}", e.name()), e, a)), + _ => Some(format!("{path}: expected {expected}, found {actual}")), + } +} + +fn describe_field_mismatch(path: &str, expected: &FieldRef, actual: &FieldRef) -> Option { + if expected.name() != actual.name() { + return Some(format!( + "{path}: expected field name '{}', found '{}'", + expected.name(), + actual.name() + )); + } + if expected.is_nullable() != actual.is_nullable() { + return Some(format!( + "{path}: expected {}, found {}", + nullability(expected), + nullability(actual) + )); + } + describe_type_mismatch(path, expected.data_type(), actual.data_type()) +} + +fn nullability(field: &FieldRef) -> String { + let qualifier = if field.is_nullable() { + "nullable" + } else { + "non-null" + }; + format!("{qualifier} {}", field.data_type()) +} + +/// Returns `base` with nested field nullability widened to also cover `other`, so that arrays of +/// either type can be stamped with the result. Shapes that do not line up (different struct field +/// counts, different list flavours, ...) are left as `base`; those are real type differences and +/// are handled by the cast in [`cast_and_stamp_schema`]. +/// +/// Related but not interchangeable: `make_all_fields_nullable` in the planner widens one type +/// unconditionally, and arrow's `Field::try_merge` unions two but errors on a leaf type difference +/// instead of tolerating it and does not recurse into maps or fixed-size lists. +pub fn widen_nested_nullability(base: &DataType, other: &DataType) -> DataType { + match (base, other) { + (DataType::List(b), DataType::List(o)) => DataType::List(widen_field(b, o)), + (DataType::LargeList(b), DataType::LargeList(o)) => DataType::LargeList(widen_field(b, o)), + (DataType::FixedSizeList(b, b_len), DataType::FixedSizeList(o, o_len)) + if b_len == o_len => + { + DataType::FixedSizeList(widen_field(b, o), *b_len) + } + (DataType::Map(b, b_sorted), DataType::Map(o, o_sorted)) if b_sorted == o_sorted => { + DataType::Map(widen_field(b, o), *b_sorted) + } + (DataType::Struct(b), DataType::Struct(o)) if b.len() == o.len() => DataType::Struct( + b.iter() + .zip(o.iter()) + .map(|(b, o)| widen_field(b, o)) + .collect(), + ), + _ => base.clone(), + } +} + +/// Widens a single field, keeping `base`'s name and metadata. A map's `entries` field must stay +/// non-nullable in Arrow, so only the nullability of fields that are already nullable on either +/// side is propagated — which is exactly `base.nullable || other.nullable`. +fn widen_field(base: &FieldRef, other: &FieldRef) -> FieldRef { + let data_type = widen_nested_nullability(base.data_type(), other.data_type()); + let nullable = base.is_nullable() || other.is_nullable(); + if nullable == base.is_nullable() && &data_type == base.data_type() { + return Arc::clone(base); + } + Arc::new( + base.as_ref() + .clone() + .with_data_type(data_type) + .with_nullable(nullable), + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{Array, Int32Array, Int64Array, ListArray, StringArray, StructArray}; + use arrow::buffer::OffsetBuffer; + use arrow::datatypes::{Field, Fields, Schema}; + + fn struct_field(nullable_child: bool) -> FieldRef { + Arc::new(Field::new_list_field( + DataType::Struct(Fields::from(vec![ + Field::new("id", DataType::Int64, true), + Field::new("flag", DataType::Boolean, nullable_child), + ])), + true, + )) + } + + /// The drift from the issue: `List(Struct(..non-null Boolean))` where catalyst declared the + /// child nullable. Only the child's `nullable` flag differs, so the path must pin it down. + #[test] + fn describes_nested_nullability_path() { + let expected = DataType::List(struct_field(true)); + let actual = DataType::List(struct_field(false)); + assert_eq!( + describe_type_mismatch("c0", &expected, &actual).unwrap(), + "c0.element.flag: expected nullable Boolean, found non-null Boolean" + ); + } + + #[test] + fn describes_leaf_type_difference() { + let expected = DataType::List(Arc::new(Field::new_list_field(DataType::Int64, true))); + let actual = DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true))); + assert_eq!( + describe_type_mismatch("c0", &expected, &actual).unwrap(), + "c0.element: expected Int64, found Int32" + ); + } + + #[test] + fn describes_nothing_for_equal_types() { + let dt = DataType::List(struct_field(true)); + assert_eq!(describe_type_mismatch("c0", &dt, &dt), None); + } + + #[test] + fn widens_nested_nullability_in_both_directions() { + let nullable_child = DataType::List(struct_field(true)); + let non_null_child = DataType::List(struct_field(false)); + assert_eq!( + widen_nested_nullability(&non_null_child, &nullable_child), + nullable_child + ); + assert_eq!( + widen_nested_nullability(&nullable_child, &non_null_child), + nullable_child + ); + } + + #[test] + fn widening_leaves_real_type_differences_alone() { + let int64 = DataType::List(Arc::new(Field::new_list_field(DataType::Int64, false))); + let int32 = DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true))); + // The element nullability still widens; the leaf type stays `base`'s for the cast to fix. + assert_eq!( + widen_nested_nullability(&int64, &int32), + DataType::List(Arc::new(Field::new_list_field(DataType::Int64, true))) + ); + } + + /// A `List(Struct(non-null Boolean))` array stamped with a schema declaring the child nullable + /// must be absorbed rather than rejected, and the values must survive unchanged. + #[test] + fn stamps_nested_nullability_drift() { + let actual = drifting_list_of_struct(); + let schema = Arc::new(Schema::new(vec![Field::new( + "c0", + DataType::List(struct_field(true)), + true, + )])); + + // Pin the reason this helper exists: stamping the declared schema on directly rejects the + // batch on the child's `nullable` flag alone. If arrow ever relaxes that, this assertion + // flips and every `cast_and_stamp_schema` call site can go back to a plain stamp. + let options = RecordBatchOptions::new().with_row_count(Some(2)); + assert!( + RecordBatch::try_new_with_options( + Arc::clone(&schema), + vec![Arc::clone(&actual)], + &options + ) + .is_err(), + "arrow no longer treats nested field nullability as part of DataType identity" + ); + + let batch = cast_and_stamp_schema("TestExec", &schema, vec![actual], 2).unwrap(); + assert_eq!(batch.num_rows(), 2); + assert_eq!(batch.schema(), schema); + let list = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(list.value(0).len(), 2); + assert_eq!(list.value(1).len(), 1); + let entries = list + .values() + .as_any() + .downcast_ref::() + .unwrap(); + let ids = entries + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(ids.values(), &[1, 2, 3]); + } + + /// `[[{1,true},{2,false}], [{3,true}]]` with a non-null `flag` child. + fn drifting_list_of_struct() -> ArrayRef { + let entries = StructArray::new( + Fields::from(vec![ + Field::new("id", DataType::Int64, true), + Field::new("flag", DataType::Boolean, false), + ]), + vec![ + Arc::new(Int64Array::from(vec![1, 2, 3])) as ArrayRef, + Arc::new(arrow::array::BooleanArray::from(vec![true, false, true])), + ], + None, + ); + Arc::new(ListArray::new( + struct_field(false), + OffsetBuffer::new(vec![0, 2, 3].into()), + Arc::new(entries), + None, + )) + } + + #[test] + fn stamps_equal_types_without_copying() { + let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, true)])); + let column: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + let batch = + cast_and_stamp_schema("TestExec", &schema, vec![Arc::clone(&column)], 3).unwrap(); + assert!(Arc::ptr_eq(batch.column(0), &column)); + } + + /// An unreconcilable difference must still name the operator and the column, not just an index. + #[test] + fn error_names_operator_and_column() { + let schema = Arc::new(Schema::new(vec![Field::new( + "payload", + DataType::Struct(Fields::from(vec![Field::new("id", DataType::Int64, true)])), + true, + )])); + let column: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"])); + let err = cast_and_stamp_schema("TestExec", &schema, vec![column], 2).unwrap_err(); + let msg = err.to_string(); + assert!(msg.contains("TestExec"), "{msg}"); + assert!(msg.contains("col[0]"), "{msg}"); + assert!(msg.contains("payload: expected Struct"), "{msg}"); + } + + #[test] + fn error_on_column_count_mismatch() { + let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, true)])); + let err = cast_and_stamp_schema("TestExec", &schema, vec![], 0).unwrap_err(); + assert!( + err.to_string() + .contains("produced 0 columns but its schema declares 1"), + "{err}" + ); + } +} diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index f798b80335..2e98bf9aa7 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -352,6 +352,11 @@ struct ExecutionContext { /// cheap to clone; the underlying `Global` releases its JNI global ref on drop /// via `jni`'s `Drop` impl. pub task_context: Option>>>, + /// Context `ClassLoader` of the driving Spark task thread, captured at `createPlan` time and + /// threaded into every JVM scalar UDF the planner builds; see `CometUdfBridge.evaluate` for why + /// it has to travel with the plan. `None` when no driving Spark task is present (unit tests, + /// direct native driver runs). Lifetime is as for `task_context` above. + pub class_loader: Option>>>, } /// Accept serialized query plan and return the address of the native query plan. @@ -379,6 +384,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( task_cpus: jlong, key_unwrapper_obj: JObject, task_context_obj: JObject, + class_loader_obj: JObject, ) -> jlong { try_unwrap_or_throw(&e, |env| { // Deserialize Spark configs @@ -500,15 +506,21 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( String::new() }; - // Capture the driving Spark task's TaskContext as a JNI global reference when - // non-null. The `Arc>` releases its global ref on drop, so cleanup - // is automatic when the ExecutionContext drops. + // Capture the driving Spark task's TaskContext and context ClassLoader as JNI global + // references when non-null. The `Arc>` releases its global ref on + // drop, so cleanup is automatic when the ExecutionContext drops. let task_context = if !task_context_obj.is_null() { Some(Arc::new(jni_new_global_ref!(env, task_context_obj)?)) } else { None }; + let class_loader = if !class_loader_obj.is_null() { + Some(Arc::new(jni_new_global_ref!(env, class_loader_obj)?)) + } else { + None + }; + let exec_context = Box::new(ExecutionContext { id, task_attempt_id, @@ -536,6 +548,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( ), tracing_event_name, task_context, + class_loader, }); Ok(Box::into_raw(exec_context) as i64) @@ -782,7 +795,8 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( PhysicalPlanner::new(Arc::clone(&exec_context.session_ctx), partition) .with_exec_id(exec_context_id) .with_sql_text_pool(&exec_context.spark_plan) - .with_task_context(exec_context.task_context.clone()); + .with_task_context(exec_context.task_context.clone()) + .with_class_loader(exec_context.class_loader.clone()); let (scans, shuffle_scans, root_op) = planner.create_plan( &exec_context.spark_plan, &mut exec_context.input_sources.clone(), diff --git a/native/core/src/execution/operators/expand.rs b/native/core/src/execution/operators/expand.rs index 8edc8c4d50..28d0c70225 100644 --- a/native/core/src/execution/operators/expand.rs +++ b/native/core/src/execution/operators/expand.rs @@ -15,8 +15,8 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{RecordBatch, RecordBatchOptions}; -use arrow::datatypes::SchemaRef; +use arrow::array::RecordBatch; +use arrow::datatypes::{Field, Schema, SchemaRef}; use datafusion::common::DataFusionError; use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; @@ -27,6 +27,7 @@ use datafusion::{ RecordBatchStream, SendableRecordBatchStream, }, }; +use datafusion_comet_common::{cast_and_stamp_schema, widen_nested_nullability}; use futures::{Stream, StreamExt}; use std::{ pin::Pin, @@ -65,6 +66,51 @@ impl ExpandExec { cache, } } + + /// Derives the operator's output schema from *every* projection, not just the first. + /// + /// Each projection produces the same output columns, so their types agree except where a + /// native kernel drifts on nested field nullability. Widening each column's nested `nullable` + /// flags across all projections gives a schema that is valid for all of them; deriving it from + /// `projections[0]` alone yields a schema that the remaining projections cannot be stamped + /// with. See . + pub fn build_schema( + projections: &[Vec>], + input_schema: &SchemaRef, + ) -> Result { + let mut data_types = projections + .first() + .ok_or_else(|| { + DataFusionError::Internal("Expand should have at least one projection".to_string()) + })? + .iter() + .map(|expr| expr.data_type(input_schema)) + .collect::, _>>()?; + + for projection in &projections[1..] { + if projection.len() != data_types.len() { + return Err(DataFusionError::Internal(format!( + "Expand projections produce differing column counts: {} and {}", + data_types.len(), + projection.len() + ))); + } + for (data_type, expr) in data_types.iter_mut().zip(projection) { + *data_type = widen_nested_nullability(data_type, &expr.data_type(input_schema)?); + } + } + + // `col_{idx}` is the name the planner used before this logic moved here, kept so that only + // the nullability derivation changes. A projection is an arbitrary expression with no + // natural output name, and parent operators bind to these columns by index (a bound + // reference takes its name from the child schema), so the name is positional either way. + let fields: Vec = data_types + .into_iter() + .enumerate() + .map(|(idx, dt)| Field::new(format!("col_{idx}"), dt, true)) + .collect(); + Ok(Arc::new(Schema::new(fields))) + } } impl DisplayAs for ExpandExec { @@ -174,9 +220,9 @@ impl ExpandStream { Ok::<(), DataFusionError>(()) })?; - let options = RecordBatchOptions::new().with_row_count(Some(batch.num_rows())); - RecordBatch::try_new_with_options(Arc::clone(&self.schema), columns, &options) - .map_err(|e| e.into()) + // A projection whose nested nullability is narrower than the operator's declared schema is + // reconciled here rather than rejected by the stamp. + cast_and_stamp_schema("CometExpandExec", &self.schema, columns, batch.num_rows()) } } @@ -214,3 +260,146 @@ impl RecordBatchStream for ExpandStream { Arc::clone(&self.schema) } } + +#[cfg(test)] +mod tests { + //! `ExpandExec` stamps its declared schema onto the output of every projection. Nested field + //! nullability is part of arrow's `DataType` identity, so a projection whose nested `nullable` + //! flags differ from the declared schema used to abort the task. These tests pin both halves of + //! the fix independently: `build_schema` widens the declared schema across all projections, and + //! the stamp reconciles anything still narrower. See + //! . + + use super::*; + use crate::execution::operators::nested_nullability_fixture::{ + list_of_struct, list_of_struct_type, + }; + use arrow::array::{Array, ListArray}; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::physical_expr::expressions::Column; + use datafusion::physical_plan::collect; + use datafusion::prelude::SessionContext; + + /// Child with two columns of the same logical `array>` type that disagree on + /// the `flag` field's nullability, so that projecting one per Expand projection reproduces the + /// drift. + fn drifting_child() -> Arc { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", list_of_struct_type(true), true), + Field::new("b", list_of_struct_type(false), true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![list_of_struct(true), list_of_struct(false)], + ) + .unwrap(); + MemorySourceConfig::try_new_exec(&[vec![batch]], schema, None).unwrap() + } + + fn drifting_projections() -> Vec>> { + vec![ + vec![Arc::new(Column::new("a", 0)) as Arc], + vec![Arc::new(Column::new("b", 1)) as Arc], + ] + } + + async fn expand_all(plan: Arc) -> Vec { + collect(plan, SessionContext::new().task_ctx()) + .await + .unwrap() + } + + /// Each output batch must carry the operator's declared schema and the values the projection + /// produced, unchanged. + fn assert_expanded(batches: &[RecordBatch], schema: &SchemaRef) { + assert_eq!(batches.len(), 2, "one batch per projection"); + for batch in batches { + assert_eq!(&batch.schema(), schema); + assert_eq!(batch.num_rows(), 2); + let list = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(list.value(0).len(), 2); + assert_eq!(list.value(1).len(), 1); + } + } + + /// `build_schema` must widen the nested `nullable` flags across all projections, not read them + /// off `projections[0]`. + #[test] + fn build_schema_widens_nested_nullability_across_projections() { + let child = drifting_child(); + // `projections[0]` alone would declare the widened type here, so drive the widening from + // the narrow side to prove it is not just projections[0] being echoed back. + let projections = vec![ + vec![Arc::new(Column::new("b", 1)) as Arc], + vec![Arc::new(Column::new("a", 0)) as Arc], + ]; + let schema = ExpandExec::build_schema(&projections, &child.schema()).unwrap(); + assert_eq!( + schema.field(0).data_type(), + &list_of_struct_type(true), + "the flag field must be nullable because projection[1] produces it nullable" + ); + } + + #[test] + fn build_schema_rejects_projections_of_differing_width() { + let child = drifting_child(); + let projections = vec![ + vec![Arc::new(Column::new("a", 0)) as Arc], + vec![ + Arc::new(Column::new("a", 0)) as Arc, + Arc::new(Column::new("b", 1)) as Arc, + ], + ]; + let err = ExpandExec::build_schema(&projections, &child.schema()).unwrap_err(); + assert!(err.to_string().contains("differing column counts"), "{err}"); + } + + /// End-to-end: projections that disagree on nested nullability both expand successfully under + /// the widened schema. + #[tokio::test] + async fn expands_projections_with_divergent_nested_nullability() { + let child = drifting_child(); + let projections = drifting_projections(); + let schema = ExpandExec::build_schema(&projections, &child.schema()).unwrap(); + let expand = Arc::new(ExpandExec::new(projections, child, Arc::clone(&schema))); + assert_expanded(&expand_all(expand).await, &schema); + } + + /// The stamp reconciles on its own: even given the narrow `projections[0]`-derived schema that + /// the planner used to build, the projection producing a non-null child no longer aborts. + #[tokio::test] + async fn stamp_reconciles_a_schema_narrower_than_a_projection() { + let child = drifting_child(); + let narrow = Arc::new(Schema::new(vec![Field::new( + "col_0", + list_of_struct_type(false), + true, + )])); + let expand = Arc::new(ExpandExec::new( + drifting_projections(), + child, + Arc::clone(&narrow), + )); + assert_expanded(&expand_all(expand).await, &narrow); + } + + /// When no projection drifts, the schema is exactly what `projections[0]` yields, and the stamp + /// is a no-op passthrough. + #[tokio::test] + async fn no_drift_is_unchanged() { + let child = drifting_child(); + let projections = vec![ + vec![Arc::new(Column::new("a", 0)) as Arc], + vec![Arc::new(Column::new("a", 0)) as Arc], + ]; + let schema = ExpandExec::build_schema(&projections, &child.schema()).unwrap(); + assert_eq!(schema.field(0).data_type(), &list_of_struct_type(true)); + let expand = Arc::new(ExpandExec::new(projections, child, Arc::clone(&schema))); + assert_expanded(&expand_all(expand).await, &schema); + } +} diff --git a/native/core/src/execution/operators/mod.rs b/native/core/src/execution/operators/mod.rs index 6fdc3b0486..b9b2b0fbd7 100644 --- a/native/core/src/execution/operators/mod.rs +++ b/native/core/src/execution/operators/mod.rs @@ -39,3 +39,52 @@ mod scan; mod shuffle_scan; pub use csv_scan::init_csv_datasource_exec; pub use shuffle_scan::ShuffleScanExec; + +/// Fixtures for the nested-nullability drift from +/// , shared by the `expand` and +/// `shuffle_scan` tests so the two boundaries are pinned against the same array. +#[cfg(test)] +pub(crate) mod nested_nullability_fixture { + use arrow::array::{ArrayRef, BooleanArray, Int64Array, ListArray, StructArray}; + use arrow::buffer::OffsetBuffer; + use arrow::datatypes::{DataType, Field, FieldRef, Fields}; + use std::sync::Arc; + + pub fn struct_fields(flag_nullable: bool) -> Fields { + Fields::from(vec![ + Field::new("id", DataType::Int64, true), + Field::new("flag", DataType::Boolean, flag_nullable), + ]) + } + + /// The element field of an `array>`, with `flag`'s nullability as given. + pub fn element_field(flag_nullable: bool) -> FieldRef { + Arc::new(Field::new_list_field( + DataType::Struct(struct_fields(flag_nullable)), + true, + )) + } + + pub fn list_of_struct_type(flag_nullable: bool) -> DataType { + DataType::List(element_field(flag_nullable)) + } + + /// `[[{1,true},{2,false}], [{3,true}]]`. Both variants hold identical values; only the `flag` + /// field's `nullable` flag differs, which is the whole of the drift being absorbed. + pub fn list_of_struct(flag_nullable: bool) -> ArrayRef { + let entries = StructArray::new( + struct_fields(flag_nullable), + vec![ + Arc::new(Int64Array::from(vec![1, 2, 3])) as ArrayRef, + Arc::new(BooleanArray::from(vec![true, false, true])), + ], + None, + ); + Arc::new(ListArray::new( + element_field(flag_nullable), + OffsetBuffer::new(vec![0, 2, 3].into()), + Arc::new(entries), + None, + )) + } +} diff --git a/native/core/src/execution/operators/shuffle_scan.rs b/native/core/src/execution/operators/shuffle_scan.rs index f3209b0c1b..71cfc29ae3 100644 --- a/native/core/src/execution/operators/shuffle_scan.rs +++ b/native/core/src/execution/operators/shuffle_scan.rs @@ -24,7 +24,7 @@ use crate::{ }; use arrow::array::ArrayRef; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; -use datafusion::common::{arrow_datafusion_err, Result as DataFusionResult}; +use datafusion::common::Result as DataFusionResult; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion::physical_plan::metrics::{ BaselineMetrics, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet, Time, @@ -34,6 +34,7 @@ use datafusion::{ physical_expr::*, physical_plan::{ExecutionPlan, *}, }; +use datafusion_comet_common::cast_and_stamp_schema; use futures::Stream; use jni::objects::{Global, JByteBuffer, JObject}; use std::{ @@ -318,14 +319,16 @@ impl Stream for ShuffleScanStream { InputBatch::EOF => Poll::Ready(None), InputBatch::Batch(columns, num_rows) => { self.baseline_metrics.record_output(*num_rows); - let options = - arrow::array::RecordBatchOptions::new().with_row_count(Some(*num_rows)); - let maybe_batch = arrow::array::RecordBatch::try_new_with_options( - self.shuffle_scan.schema(), + // Reconcile the decoded block with the catalyst-declared schema rather than + // stamping it on, so that nested field nullability drift is absorbed here the way + // `ScanExec` absorbs it at the FFI boundary. + // See https://github.com/apache/datafusion-comet/issues/5137. + let maybe_batch = cast_and_stamp_schema( + self.shuffle_scan.name(), + &self.shuffle_scan.schema, columns.clone(), - &options, - ) - .map_err(|e| arrow_datafusion_err!(e)); + *num_rows, + ); Poll::Ready(Some(maybe_batch)) } }; @@ -499,4 +502,79 @@ mod tests { assert_eq!(col1.value(2), "hello"); }); } + + /// A decoded shuffle block whose nested field nullability is narrower than the catalyst-declared + /// type must be reconciled, not rejected. `ShuffleScanExec` used to stamp the declared schema + /// straight onto the block, which aborted the task on a single nested `nullable` flag even + /// though a non-null child is a strict subset of a nullable one. + /// See . + #[test] + #[cfg_attr(miri, ignore)] + fn test_nested_nullability_drift_is_reconciled() { + use super::*; + use crate::execution::operators::nested_nullability_fixture::{ + list_of_struct, list_of_struct_type, + }; + use arrow::array::Array; + use datafusion::physical_plan::ExecutionPlan; + use futures::StreamExt; + + // The block carries `List(Struct("id": Int64, "flag": non-null Boolean))` while catalyst + // declared the `flag` child nullable. + let block_column = list_of_struct(false); + let declared = list_of_struct_type(true); + let mut scan = ShuffleScanExec::new( + super::super::super::planner::TEST_EXEC_CONTEXT_ID, + None, + vec![declared.clone()], + ) + .unwrap(); + scan.set_input_batch(InputBatch::new(vec![block_column], Some(2))); + + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + let ctx = Arc::new(TaskContext::default()); + let mut stream = scan.execute(0, ctx).unwrap(); + let batch = stream.next().await.unwrap().unwrap(); + + assert_eq!(batch.schema().field(0).data_type(), &declared); + assert_eq!(batch.num_rows(), 2); + // The values must survive the reconciliation untouched. `ArrayData` equality is + // logical, so this holds regardless of whether the cast materialized an all-valid + // null buffer for the widened child. + assert_eq!(batch.column(0).to_data(), list_of_struct(true).to_data()); + }); + } + + /// An unreconcilable column must name the operator and the column, since arrow's own message + /// reports only `at column index N`. + #[test] + #[cfg_attr(miri, ignore)] + fn test_unreconcilable_column_error_names_operator() { + use super::*; + use arrow::datatypes::Fields; + use datafusion::physical_plan::ExecutionPlan; + use futures::StreamExt; + + let declared = + DataType::Struct(Fields::from(vec![Field::new("id", DataType::Int64, true)])); + let mut scan = ShuffleScanExec::new( + super::super::super::planner::TEST_EXEC_CONTEXT_ID, + None, + vec![declared], + ) + .unwrap(); + let column: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"])); + scan.set_input_batch(InputBatch::new(vec![column], Some(2))); + + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + let ctx = Arc::new(TaskContext::default()); + let mut stream = scan.execute(0, ctx).unwrap(); + let err = stream.next().await.unwrap().unwrap_err().to_string(); + assert!(err.contains("ShuffleScanExec"), "{err}"); + assert!(err.contains("col[0]"), "{err}"); + assert!(err.contains("col_0: expected Struct"), "{err}"); + }); + } } diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index d4c97c509d..5cb305b4c6 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -250,6 +250,10 @@ pub struct PhysicalPlanner { /// Captured at `createPlan` time on `ExecutionContext`; see that struct for the /// propagation rationale. `None` when no driving Spark task is available. task_context: Option>>>, + /// Context `ClassLoader` of the driving Spark task thread, captured at `createPlan` time on + /// `ExecutionContext`; see that struct for the propagation rationale. `None` when no driving + /// Spark task is available. + class_loader: Option>>>, } impl Default for PhysicalPlanner { @@ -267,6 +271,7 @@ impl PhysicalPlanner { query_context_registry: datafusion_comet_spark_expr::create_query_context_map(), sql_text_pool: vec![], task_context: None, + class_loader: None, } } @@ -356,6 +361,17 @@ impl PhysicalPlanner { self } + /// Attach the driving Spark task thread's context `ClassLoader` as a global reference. Mirrors + /// `with_task_context`: called by the JNI `executePlan` entry with whatever was captured at + /// `createPlan` time, and cloned into every `JvmScalarUdfExpr` the planner builds. + pub fn with_class_loader( + mut self, + class_loader: Option>>>, + ) -> Self { + self.class_loader = class_loader; + self + } + /// Return session context of this planner. pub fn session_ctx(&self) -> &Arc { &self.session_ctx @@ -876,12 +892,14 @@ impl PhysicalPlanner { to_arrow_datatype(udf.return_type.as_ref().ok_or_else(|| { GeneralError("JvmScalarUdf missing return_type".to_string()) })?); - // Invariant: task_context is propagated for every JvmScalarUdfExpr built during - // normal execution. The TEST_EXEC_CONTEXT_ID path is the only context in which - // task_context may legitimately be None (unit tests, direct native driver runs). + // Invariant: task_context and class_loader are propagated for every + // JvmScalarUdfExpr built during normal execution. The TEST_EXEC_CONTEXT_ID path is + // the only context in which they may legitimately be None (unit tests, direct + // native driver runs). debug_assert!( - self.task_context.is_some() || self.exec_context_id == TEST_EXEC_CONTEXT_ID, - "task_context must be set for non-test execution" + (self.task_context.is_some() && self.class_loader.is_some()) + || self.exec_context_id == TEST_EXEC_CONTEXT_ID, + "task_context and class_loader must be set for non-test execution" ); Ok(Arc::new(JvmScalarUdfExpr::new( udf.class_name.clone(), @@ -889,6 +907,7 @@ impl PhysicalPlanner { return_type, udf.return_nullable, self.task_context.clone(), + self.class_loader.clone(), ))) } expr => Err(GeneralError(format!("Not implemented: {expr:?}"))), @@ -1862,21 +1881,7 @@ impl PhysicalPlanner { Ok::<(), ExecutionError>(()) })?; - assert!( - !projections.is_empty(), - "Expand should have at least one projection" - ); - - let datatypes = projections[0] - .iter() - .map(|expr| expr.data_type(&child.schema())) - .collect::, _>>()?; - let fields: Vec = datatypes - .iter() - .enumerate() - .map(|(idx, dt)| Field::new(format!("col_{idx}"), dt.clone(), true)) - .collect(); - let schema = Arc::new(Schema::new(fields)); + let schema = ExpandExec::build_schema(&projections, &child.schema())?; // `Expand` operator keeps the input batch and expands it to multiple output // batches. However, `ScanExec` will reuse input arrays for the next diff --git a/native/jni-bridge/src/comet_udf_bridge.rs b/native/jni-bridge/src/comet_udf_bridge.rs index e531d20cb1..bbd7465d00 100644 --- a/native/jni-bridge/src/comet_udf_bridge.rs +++ b/native/jni-bridge/src/comet_udf_bridge.rs @@ -41,7 +41,9 @@ impl<'a> CometUdfBridge<'a> { method_evaluate: env.get_static_method_id( JNIString::new(Self::JVM_CLASS), jni::jni_str!("evaluate"), - jni::jni_sig!("(Ljava/lang/String;[J[JJJILorg/apache/spark/TaskContext;)V"), + jni::jni_sig!( + "(Ljava/lang/String;[J[JJJILorg/apache/spark/TaskContext;Ljava/lang/ClassLoader;)V" + ), )?, method_evaluate_ret: ReturnType::Primitive(Primitive::Void), class, diff --git a/native/shuffle/src/schema_align.rs b/native/shuffle/src/schema_align.rs index 6c8d0bb97b..33ce490b69 100644 --- a/native/shuffle/src/schema_align.rs +++ b/native/shuffle/src/schema_align.rs @@ -22,21 +22,19 @@ //! return-type drift from DataFusion / `datafusion-spark` is self-healing. When a native plan's //! output crosses back to the JVM and feeds another native plan, the consuming `ScanExec` casts //! every imported column to the catalyst-declared type, so a wrong Arrow type never survives the -//! boundary. Shuffle is the lone exception, on two counts: +//! boundary. Shuffle is the lone exception because the writer hash-partitions on these columns, and +//! Spark's hash differs by type (e.g. `Int32` vs `Int64`), so a drifted type would route rows to +//! the wrong partition. A read-side cast cannot undo a wrong partition assignment, so the type must +//! be corrected before partitioning — which forces the alignment onto the writer input. //! -//! 1. The writer hash-partitions on these columns, and Spark's hash differs by type (e.g. `Int32` -//! vs `Int64`), so a drifted type would route rows to the wrong partition. A read-side cast -//! cannot undo a wrong partition assignment, so the type must be corrected before partitioning. -//! 2. The shuffle read path (`ShuffleScanExec`) does not cast; it stamps the catalyst schema onto -//! the decoded block and errors on any mismatch. The schema is serialized into the block on -//! write and trusted on read. +//! The read path (`ShuffleScanExec`) casts too, so a drift that only affects the batch's Arrow type +//! and not its partition assignment is absorbed there as well. See +//! . //! -//! Both force the alignment to happen on the writer input. See -//! for the running list of mismatched +//! See for the running list of mismatched //! functions. -use arrow::array::{ArrayRef, RecordBatch, RecordBatchOptions}; -use arrow::compute::{cast_with_options, CastOptions}; +use arrow::array::RecordBatch; use arrow::datatypes::{Field, Schema, SchemaRef}; use datafusion::common::DataFusionError; use datafusion::physical_expr::EquivalenceProperties; @@ -48,6 +46,7 @@ use datafusion::{ RecordBatchStream, SendableRecordBatchStream, }, }; +use datafusion_comet_common::cast_and_stamp_schema; use futures::{Stream, StreamExt}; use std::{ collections::HashSet, @@ -72,19 +71,9 @@ fn warn_dedup() -> &'static Mutex> { pub struct SchemaAlignExec { child: Arc, target_schema: SchemaRef, - column_actions: Arc>, cache: Arc, } -#[derive(Debug, Clone)] -enum ColumnAction { - /// Pass the input column through unchanged. Any nullability/metadata difference is - /// absorbed when the batch is re-stamped via `RecordBatch::try_new_with_options`. - Passthrough, - /// Cast the input column to the target data_type. - Cast, -} - impl SchemaAlignExec { /// Build a SchemaAlignExec that aligns `child`'s output to `expected`. Returns /// `Ok(child)` unchanged when no per-column reshape is needed; otherwise wraps `child` @@ -103,7 +92,6 @@ impl SchemaAlignExec { ))); } let mut needs_alignment = false; - let mut actions = Vec::with_capacity(actual.fields().len()); let mut target_fields = Vec::with_capacity(actual.fields().len()); for (idx, (actual_field, expected_field)) in actual .fields() @@ -111,8 +99,8 @@ impl SchemaAlignExec { .zip(expected.fields().iter()) .enumerate() { - let action = if actual_field.data_type() == expected_field.data_type() { - ColumnAction::Passthrough + let needs_cast = if actual_field.data_type() == expected_field.data_type() { + false } else { let signature = format!( "{}|{:?}|{:?}", @@ -130,10 +118,10 @@ impl SchemaAlignExec { expected_field.data_type() ); } - ColumnAction::Cast + true }; let target_nullable = actual_field.is_nullable() || expected_field.is_nullable(); - let field_changed = !matches!(action, ColumnAction::Passthrough) + let field_changed = needs_cast || target_nullable != actual_field.is_nullable() || expected_field.metadata() != actual_field.metadata() || expected_field.name() != actual_field.name(); @@ -148,7 +136,6 @@ impl SchemaAlignExec { ) .with_metadata(expected_field.metadata().clone()), ); - actions.push(action); } if !needs_alignment { return Ok(child); @@ -163,7 +150,6 @@ impl SchemaAlignExec { Ok(Arc::new(Self { child, target_schema, - column_actions: Arc::new(actions), cache, })) } @@ -204,7 +190,6 @@ impl ExecutionPlan for SchemaAlignExec { Ok(Arc::new(Self { child: new_child, target_schema: Arc::clone(&self.target_schema), - column_actions: Arc::clone(&self.column_actions), cache, })) } @@ -218,7 +203,6 @@ impl ExecutionPlan for SchemaAlignExec { Ok(Box::pin(SchemaAlignStream { child_stream, target_schema: Arc::clone(&self.target_schema), - column_actions: Arc::clone(&self.column_actions), })) } @@ -234,27 +218,17 @@ impl ExecutionPlan for SchemaAlignExec { struct SchemaAlignStream { child_stream: SendableRecordBatchStream, target_schema: SchemaRef, - column_actions: Arc>, } impl SchemaAlignStream { fn align(&self, batch: RecordBatch) -> Result { - let mut columns: Vec = Vec::with_capacity(batch.num_columns()); - for (idx, action) in self.column_actions.iter().enumerate() { - let column = batch.column(idx); - let aligned = match action { - ColumnAction::Passthrough => Arc::clone(column), - ColumnAction::Cast => cast_with_options( - column, - self.target_schema.field(idx).data_type(), - &CastOptions::default(), - )?, - }; - columns.push(aligned); - } - let options = RecordBatchOptions::new().with_row_count(Some(batch.num_rows())); - RecordBatch::try_new_with_options(Arc::clone(&self.target_schema), columns, &options) - .map_err(DataFusionError::from) + let num_rows = batch.num_rows(); + cast_and_stamp_schema( + "CometSchemaAlignExec", + &self.target_schema, + batch.columns().to_vec(), + num_rows, + ) } } diff --git a/native/spark-expr/src/datetime_funcs/make_date.rs b/native/spark-expr/src/datetime_funcs/make_date.rs index 02c9160587..b29094ba9f 100644 --- a/native/spark-expr/src/datetime_funcs/make_date.rs +++ b/native/spark-expr/src/datetime_funcs/make_date.rs @@ -25,6 +25,8 @@ use datafusion::logical_expr::{ }; use std::sync::Arc; +use crate::SparkError; + /// Spark-compatible make_date function. /// Creates a date from year, month, and day columns. /// For an invalid `(year, month, day)` triple Spark returns NULL when `spark.sql.ansi.enabled` is @@ -53,8 +55,8 @@ impl Default for SparkMakeDate { /// Build the error message Spark surfaces for an invalid date under ANSI mode. Spark wraps the /// `java.time.DateTimeException` raised by `LocalDate.of` (via `ansiDateTimeArgumentOutOfRange` / -/// `ansiDateTimeError`), so we reproduce `java.time`'s messages and validation order: month range, -/// then day range, then the day-vs-month check. +/// `ansiDateTimeError`), so we reproduce `java.time`'s messages and validation order: year range, +/// month range, day range, then the day-vs-month check. fn invalid_date_message(year: i32, month: i32, day: i32) -> String { const MONTH_NAMES: [&str; 12] = [ "JANUARY", @@ -70,6 +72,9 @@ fn invalid_date_message(year: i32, month: i32, day: i32) -> String { "NOVEMBER", "DECEMBER", ]; + if !(-999_999_999..=999_999_999).contains(&year) { + return format!("Invalid value for Year (valid values -999999999 - 999999999): {year}"); + } if !(1..=12).contains(&month) { return format!("Invalid value for MonthOfYear (valid values 1 - 12): {month}"); } @@ -178,7 +183,10 @@ impl ScalarUDFImpl for SparkMakeDate { Some(days) => builder.append_value(days), None => { if self.fail_on_error { - return Err(DataFusionError::Execution(invalid_date_message(y, m, d))); + return Err(SparkError::DatetimeFieldOutOfBounds { + range_message: invalid_date_message(y, m, d), + } + .into()); } builder.append_null(); } diff --git a/native/spark-expr/src/datetime_funcs/next_day.rs b/native/spark-expr/src/datetime_funcs/next_day.rs index df4c2f9096..fbd9defc5f 100644 --- a/native/spark-expr/src/datetime_funcs/next_day.rs +++ b/native/spark-expr/src/datetime_funcs/next_day.rs @@ -25,6 +25,8 @@ use datafusion::logical_expr::{ }; use std::sync::Arc; +use crate::SparkError; + /// Spark-compatible `next_day(start_date, day_of_week)` function. /// /// Returns the first date which is later than `start_date` and named as `day_of_week`. Unlike the @@ -146,9 +148,10 @@ impl ScalarUDFImpl for SparkNextDay { }, None => { if self.fail_on_error { - return Err(DataFusionError::Execution(format!( - "Illegal input for day of week: {day_of_week}" - ))); + return Err(SparkError::IllegalDayOfWeek { + input: day_of_week.to_string(), + } + .into()); } builder.append_null(); } diff --git a/native/spark-expr/src/jvm_udf/mod.rs b/native/spark-expr/src/jvm_udf/mod.rs index 0ca603ac9b..8e148ef40a 100644 --- a/native/spark-expr/src/jvm_udf/mod.rs +++ b/native/spark-expr/src/jvm_udf/mod.rs @@ -48,6 +48,11 @@ pub struct JvmScalarUdfExpr { /// Spark task is available; the bridge then leaves whatever `TaskContext.get()` already /// returns in place. task_context: Option>>>, + /// Context `ClassLoader` of the driving Spark task thread, captured at `createPlan` time and + /// threaded here by the planner. See `CometUdfBridge.evaluate`, which installs it for the + /// duration of the call. `None` when no driving Spark task is available (unit tests, direct + /// native driver runs); the bridge then installs nothing. + class_loader: Option>>>, } impl JvmScalarUdfExpr { @@ -57,6 +62,7 @@ impl JvmScalarUdfExpr { return_type: DataType, return_nullable: bool, task_context: Option>>>, + class_loader: Option>>>, ) -> Self { debug_assert!( !class_name.is_empty(), @@ -68,6 +74,7 @@ impl JvmScalarUdfExpr { return_type, return_nullable, task_context, + class_loader, } } } @@ -197,14 +204,17 @@ impl PhysicalExpr for JvmScalarUdfExpr { .set_region(env, 0, &in_sch_ptrs) .map_err(|e| CometError::JNI { source: e })?; - // Resolve the TaskContext reference once before building the arg array so the - // borrow lives until `call_static_method_unchecked` returns. When no TaskContext - // was propagated, pass a null object so the bridge's null-guard leaves the thread- - // local alone. - let null_task_context = JObject::null(); + // Resolve the TaskContext and ClassLoader references once before building the arg + // array so the borrows live until `call_static_method_unchecked` returns. Absent + // values are passed as a null object, which the bridge's null-guards skip. + let null_obj = JObject::null(); let task_context_ref: &JObject = match &self.task_context { Some(gref) => gref.as_obj(), - None => &null_task_context, + None => &null_obj, + }; + let class_loader_ref: &JObject = match &self.class_loader { + Some(gref) => gref.as_obj(), + None => &null_obj, }; let ret = unsafe { env.call_static_method_unchecked( @@ -219,6 +229,7 @@ impl PhysicalExpr for JvmScalarUdfExpr { JValue::Long(out_sch_ptr).as_jni(), JValue::Int(batch.num_rows() as i32).as_jni(), JValue::Object(task_context_ref).as_jni(), + JValue::Object(class_loader_ref).as_jni(), ], ) }; @@ -255,6 +266,7 @@ impl PhysicalExpr for JvmScalarUdfExpr { self.return_type.clone(), self.return_nullable, self.task_context.clone(), + self.class_loader.clone(), ))) } } diff --git a/native/spark-expr/src/math_funcs/negative.rs b/native/spark-expr/src/math_funcs/negative.rs index 650fa401ef..a0da6f52fb 100644 --- a/native/spark-expr/src/math_funcs/negative.rs +++ b/native/spark-expr/src/math_funcs/negative.rs @@ -18,9 +18,10 @@ use crate::arithmetic_overflow_error; use crate::SparkError; use arrow::array::RecordBatch; -use arrow::datatypes::IntervalDayTime; -use arrow::datatypes::{DataType, Schema}; -use arrow::{compute::kernels::numeric::neg_wrapping, datatypes::IntervalDayTimeType}; +use arrow::compute::kernels::numeric::{neg, neg_wrapping}; +use arrow::datatypes::IntervalDayTimeType; +use arrow::datatypes::{DataType, IntervalUnit, Schema}; +use arrow::error::ArrowError; use datafusion::common::{DataFusionError, Result, ScalarValue}; use datafusion::logical_expr::sort_properties::ExprProperties; use datafusion::{ @@ -59,24 +60,6 @@ impl PartialEq for NegativeExpr { } } -macro_rules! check_overflow { - ($array:expr, $array_type:ty, $min_val:expr, $type_name:expr) => {{ - let typed_array = $array - .as_any() - .downcast_ref::<$array_type>() - .expect(concat!(stringify!($array_type), " expected")); - for i in 0..typed_array.len() { - if typed_array.value(i) == $min_val { - if $type_name == "byte" || $type_name == "short" { - let value = format!("{:?} caused", typed_array.value(i)); - return Err(arithmetic_overflow_error(value.as_str()).into()); - } - return Err(arithmetic_overflow_error($type_name).into()); - } - } - }}; -} - impl NegativeExpr { /// Create new not expression pub fn new(arg: Arc, fail_on_error: bool) -> Self { @@ -95,6 +78,13 @@ impl std::fmt::Display for NegativeExpr { } } +fn map_neg_error(err: ArrowError, from_type: &'static str) -> DataFusionError { + match err { + ArrowError::ArithmeticOverflow(_) => arithmetic_overflow_error(from_type).into(), + other => DataFusionError::from(other), + } +} + impl PhysicalExpr for NegativeExpr { fn data_type(&self, input_schema: &Schema) -> Result { self.arg.data_type(input_schema) @@ -107,57 +97,50 @@ impl PhysicalExpr for NegativeExpr { fn evaluate(&self, batch: &RecordBatch) -> Result { let arg = self.arg.evaluate(batch)?; - // overflow checks only apply in ANSI mode - // datatypes supported are byte, short, integer, long, float, interval + // Overflow checks only apply in ANSI mode, and only the types listed in the + // match below have a Spark overflow message. Everything else (float, decimal, + // `Interval(MonthDayNano)`, ...) falls through to `neg_wrapping`. match arg { ColumnarValue::Array(array) => { - if self.fail_on_error { - match array.data_type() { - DataType::Int8 => { - check_overflow!(array, arrow::array::Int8Array, i8::MIN, "byte") - } - DataType::Int16 => { - check_overflow!(array, arrow::array::Int16Array, i16::MIN, "short") - } - DataType::Int32 => { - check_overflow!(array, arrow::array::Int32Array, i32::MIN, "integer") - } - DataType::Int64 => { - check_overflow!(array, arrow::array::Int64Array, i64::MIN, "long") - } - DataType::Interval(value) => match value { - arrow::datatypes::IntervalUnit::YearMonth => check_overflow!( - array, - arrow::array::IntervalYearMonthArray, - i32::MIN, - "interval" - ), - arrow::datatypes::IntervalUnit::DayTime => check_overflow!( - array, - arrow::array::IntervalDayTimeArray, - IntervalDayTime::MIN, - "interval" - ), - arrow::datatypes::IntervalUnit::MonthDayNano => { - // Overflow checks are not supported - } - }, - _ => { - // Overflow checks are not supported for other datatypes - } - } + if !self.fail_on_error { + return Ok(ColumnarValue::Array(neg_wrapping(array.as_ref())?)); } - let result = neg_wrapping(array.as_ref())?; - Ok(ColumnarValue::Array(result)) + // The shims render this as `{from_type} overflow` under + // `ARITHMETIC_OVERFLOW`. For byte/short that is byte-identical to Spark + // 4.x, which routes them through `MathUtils.negateExact` ("byte overflow" / + // "short overflow"). Spark 3.4/3.5 instead throw + // `_LEGACY_ERROR_TEMP_2043` ("- caused overflow."); that error + // class is out of reach here, so no string can match every version. + let from_type = match array.data_type() { + DataType::Int8 => "byte", + DataType::Int16 => "short", + DataType::Int32 => "integer", + DataType::Int64 => "long", + // `neg` checks each `DayTime` component, so either `days` or `ms` at + // `i32::MIN` overflows; routing intervals here maps the Arrow overflow + // error onto Spark's, matching the scalar path below. + DataType::Interval(IntervalUnit::YearMonth | IntervalUnit::DayTime) => { + "interval" + } + // Everything else falls through to `neg_wrapping`. Note it wraps only + // for integers: for any other type, `Interval(MonthDayNano)` included, + // it delegates to `neg`, so overflow is still detected -- it just + // surfaces as an Arrow error rather than a Spark one. + _ => return Ok(ColumnarValue::Array(neg_wrapping(array.as_ref())?)), + }; + Ok(ColumnarValue::Array( + neg(array.as_ref()).map_err(|e| map_neg_error(e, from_type))?, + )) } ColumnarValue::Scalar(scalar) => { if self.fail_on_error { match scalar { + // Keep scalar overflow type names aligned with the array path. ScalarValue::Int8(Some(i8::MIN)) => { - return Err(arithmetic_overflow_error(" caused").into()); + return Err(arithmetic_overflow_error("byte").into()); } ScalarValue::Int16(Some(i16::MIN)) => { - return Err(arithmetic_overflow_error(" caused").into()); + return Err(arithmetic_overflow_error("short").into()); } ScalarValue::Int32(Some(i32::MIN)) => { return Err(arithmetic_overflow_error("integer").into()); @@ -250,3 +233,217 @@ impl PhysicalExpr for NegativeExpr { Display::fmt(self, f) } } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::{array::*, buffer::NullBuffer, datatypes::*}; + use datafusion::{ + physical_expr::expressions::{Column, Literal}, + physical_plan::ColumnarValue, + }; + + fn eval_array(array: ArrayRef, fail_on_error: bool) -> Result { + let schema = Arc::new(Schema::new(vec![Field::new( + "a", + array.data_type().clone(), + true, + )])); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![array])?; + NegativeExpr::new(Arc::new(Column::new("a", 0)), fail_on_error).evaluate(&batch) + } + + fn eval_scalar(scalar: ScalarValue, fail_on_error: bool) -> Result { + let batch = RecordBatch::new_empty(Arc::new(Schema::empty())); + NegativeExpr::new(Arc::new(Literal::new(scalar)), fail_on_error).evaluate(&batch) + } + + fn assert_spark_overflow(err: DataFusionError, expected_from_type: &str) { + if let DataFusionError::External(ref e) = err { + if let Some(SparkError::ArithmeticOverflow { from_type }) = + e.downcast_ref::() + { + assert_eq!(from_type, expected_from_type); + return; + } + } + panic!( + "Expected SparkError::ArithmeticOverflow {{ from_type: {:?} }}, got: {:?}", + expected_from_type, err + ); + } + + /// Negate `[min, other]` in ANSI mode with slot 0 marked null, and assert the + /// null slot is skipped instead of raising a spurious overflow. + fn assert_null_min_slot_is_skipped( + min: T::Native, + other: T::Native, + negated_other: T::Native, + ) { + let nulls = NullBuffer::from(vec![false, true]); + let array: ArrayRef = Arc::new(PrimitiveArray::::new( + vec![min, other].into(), + Some(nulls), + )); + let ColumnarValue::Array(result) = eval_array(array, true).unwrap() else { + panic!("expected array result") + }; + let result = result.as_primitive::(); + assert!(result.is_null(0)); + assert_eq!(result.value(1), negated_other); + } + + /// A MIN sentinel left behind in a null slot (by filter, slice or FFI) must not raise a + /// spurious ANSI overflow: the overflow check has to consult the null buffer and skip + /// invalid slots. Each case below places `MIN` in a null slot. + #[test] + fn test_ansi_null_slot_with_min_values_does_not_overflow() { + assert_null_min_slot_is_skipped::(i8::MIN, 7, -7); + assert_null_min_slot_is_skipped::(i16::MIN, 7, -7); + assert_null_min_slot_is_skipped::(i32::MIN, 7, -7); + assert_null_min_slot_is_skipped::(i64::MIN, 7, -7); + assert_null_min_slot_is_skipped::(i32::MIN, 7, -7); + // `IntervalDayTime::MIN` is both components at `i32::MIN`. + assert_null_min_slot_is_skipped::( + IntervalDayTime::MIN, + IntervalDayTime::new(1, 2), + IntervalDayTime::new(-1, -2), + ); + } + + #[test] + fn test_ansi_valid_min_values_raise_exact_spark_overflow_errors() { + let arr_i8: ArrayRef = Arc::new(Int8Array::from(vec![i8::MIN])); + assert_spark_overflow(eval_array(arr_i8, true).unwrap_err(), "byte"); + + let arr_i16: ArrayRef = Arc::new(Int16Array::from(vec![i16::MIN])); + assert_spark_overflow(eval_array(arr_i16, true).unwrap_err(), "short"); + + let arr_i32: ArrayRef = Arc::new(Int32Array::from(vec![i32::MIN])); + assert_spark_overflow(eval_array(arr_i32, true).unwrap_err(), "integer"); + + let arr_i64: ArrayRef = Arc::new(Int64Array::from(vec![i64::MIN])); + assert_spark_overflow(eval_array(arr_i64, true).unwrap_err(), "long"); + + let arr_ym: ArrayRef = Arc::new(IntervalYearMonthArray::from(vec![i32::MIN])); + assert_spark_overflow(eval_array(arr_ym, true).unwrap_err(), "interval"); + + let arr_dt: ArrayRef = Arc::new(IntervalDayTimeArray::from(vec![IntervalDayTime::MIN])); + assert_spark_overflow(eval_array(arr_dt, true).unwrap_err(), "interval"); + } + + /// A single `DayTime` component at `i32::MIN` overflows: `neg` checks each component + /// (`neg_wrapping` delegates to `neg` for every non-integer type, + /// `downcast_integer! { ..., _ => neg(array) }`). These surface as Spark overflows like + /// every other ANSI overflow in this expression. + #[test] + fn test_ansi_interval_day_time_component_overflow_maps_to_spark_error() { + for value in [ + IntervalDayTime::new(i32::MIN, 0), + IntervalDayTime::new(0, i32::MIN), + ] { + let array: ArrayRef = Arc::new(IntervalDayTimeArray::from(vec![value])); + assert_spark_overflow(eval_array(array, true).unwrap_err(), "interval"); + } + } + + /// Negate `[min, null]` in legacy mode and assert `min` wraps to itself. + fn assert_legacy_wraps_min(min: T::Native) { + let array: ArrayRef = Arc::new(PrimitiveArray::::new( + vec![min, T::Native::default()].into(), + Some(NullBuffer::from(vec![true, false])), + )); + let ColumnarValue::Array(result) = eval_array(array, false).unwrap() else { + panic!("expected array result") + }; + let result = result.as_primitive::(); + assert_eq!(result.value(0), min); + assert!(result.is_null(1)); + } + + #[test] + fn test_legacy_mode_wraps_min_values() { + assert_legacy_wraps_min::(i8::MIN); + assert_legacy_wraps_min::(i16::MIN); + assert_legacy_wraps_min::(i32::MIN); + assert_legacy_wraps_min::(i64::MIN); + } + + #[test] + fn test_mixed_ordinary_values() { + let arr: ArrayRef = Arc::new(Int32Array::from(vec![Some(-7), Some(0), Some(12), None])); + + // ANSI mode + let ColumnarValue::Array(res_ansi) = eval_array(Arc::clone(&arr), true).unwrap() else { + panic!("expected array result") + }; + assert_eq!( + res_ansi.as_primitive::(), + &Int32Array::from(vec![Some(7), Some(0), Some(-12), None]) + ); + + // Legacy mode + let ColumnarValue::Array(res_legacy) = eval_array(Arc::clone(&arr), false).unwrap() else { + panic!("expected array result") + }; + assert_eq!( + res_legacy.as_primitive::(), + &Int32Array::from(vec![Some(7), Some(0), Some(-12), None]) + ); + } + + #[test] + fn test_interval_month_day_nano_preserves_existing_dispatch() { + let arr: ArrayRef = Arc::new(IntervalMonthDayNanoArray::from(vec![ + Some(IntervalMonthDayNano::new(1, 2, 3)), + None, + ])); + let ColumnarValue::Array(res) = eval_array(arr, true).unwrap() else { + panic!("expected array result") + }; + let p = res.as_primitive::(); + assert_eq!(p.value(0), IntervalMonthDayNano::new(-1, -2, -3)); + assert!(p.is_null(1)); + } + + #[test] + fn test_scalar_negation() { + // Valid scalar + let ColumnarValue::Scalar(res_valid) = + eval_scalar(ScalarValue::Int32(Some(42)), true).unwrap() + else { + panic!("expected scalar result") + }; + assert_eq!(res_valid, ScalarValue::Int32(Some(-42))); + + // Null scalar + let ColumnarValue::Scalar(res_null) = eval_scalar(ScalarValue::Int32(None), true).unwrap() + else { + panic!("expected scalar result") + }; + assert_eq!(res_null, ScalarValue::Int32(None)); + + // MIN scalar overflows in ANSI, with the same messages as the array path + for (scalar, from_type) in [ + (ScalarValue::Int8(Some(i8::MIN)), "byte"), + (ScalarValue::Int16(Some(i16::MIN)), "short"), + (ScalarValue::Int32(Some(i32::MIN)), "integer"), + (ScalarValue::Int64(Some(i64::MIN)), "long"), + ] { + assert_spark_overflow(eval_scalar(scalar, true).unwrap_err(), from_type); + } + } + + #[test] + fn test_map_neg_error_preserves_non_overflow_errors() { + let err = map_neg_error( + ArrowError::InvalidArgumentError("test custom error".to_string()), + "integer", + ); + assert!( + matches!(&err, DataFusionError::ArrowError(inner, _) + if matches!(inner.as_ref(), ArrowError::InvalidArgumentError(msg) if msg == "test custom error")), + "expected the ArrowError to pass through unchanged, got: {err:?}" + ); + } +} diff --git a/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java b/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java index 9e97ef2226..d8dea73135 100644 --- a/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java +++ b/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java @@ -31,6 +31,8 @@ import org.apache.spark.comet.CometTaskContextShim; import org.apache.spark.util.TaskCompletionListener; +import org.apache.comet.util.ClassLoaders; + /** * JNI entry point for native execution to invoke a {@link CometUDF}. Matches the static-method * pattern used by CometScalarSubquery so the native side can dispatch via @@ -91,6 +93,17 @@ public class CometUdfBridge { * left on a worker by a previous task. Its task attempt ID also keys the UDF-instance cache, * so a UDF holding per-task state in fields sees a consistent instance for every call within * the task regardless of which Tokio worker is polling. + * @param classLoader context ClassLoader captured on the driving Spark task thread, or {@code + * null} outside a Spark task. Installed as this thread's context ClassLoader for the duration + * of the call, with the prior value restored in {@code finally}. Tokio workers attach through + * JNI and an attached thread has no context ClassLoader, so without this every lookup falls + * back to the ClassLoader that loaded Comet, which never holds user jars ({@code --jars} / + * {@code spark.jars}). Both the {@code CometUDF} resolution below and the closure + * deserialization in {@code CometScalaUDFCodegen} depend on it. Installed per call rather + * than once when the worker attaches, because the Tokio runtime is process-global: one worker + * interleaves work from task attempts of different jobs, and under Spark Connect from + * sessions with different artifact ClassLoaders. The loader is a property of the plan, not of + * the thread. */ public static void evaluate( String udfClassName, @@ -99,7 +112,8 @@ public static void evaluate( long outArrayPtr, long outSchemaPtr, int numRows, - TaskContext taskContext) { + TaskContext taskContext, + ClassLoader classLoader) { assert udfClassName != null && !udfClassName.isEmpty() : "udfClassName must be non-empty"; assert inputArrayPtrs != null && inputSchemaPtrs != null : "input pointer arrays must be non-null"; @@ -111,13 +125,20 @@ public static void evaluate( // Save-and-restore rather than only-install-if-null: the propagated `taskContext` is the // ground truth for this call. Any value already on the thread is either (a) the same object - // on a Spark task thread, or (b) stale from a prior task on a reused Tokio worker. + // on a Spark task thread, or (b) stale from a prior task on a reused Tokio worker. The same + // reasoning applies to the propagated `classLoader`. TaskContext prior = TaskContext.get(); if (taskContext != null) { CometTaskContextShim.set(taskContext); assert TaskContext.get() == taskContext : "TaskContext install did not take effect on this thread"; } + Thread currentThread = Thread.currentThread(); + ClassLoader priorLoader = currentThread.getContextClassLoader(); + if (classLoader != null) { + currentThread.setContextClassLoader(classLoader); + } + try { evaluateInternal( udfClassName, @@ -128,6 +149,9 @@ public static void evaluate( numRows, taskContext); } finally { + // Unconditional: a no-op when nothing was installed, and it also undoes any change the + // user function made to the ClassLoader of a worker that outlives this call. + currentThread.setContextClassLoader(priorLoader); if (taskContext != null) { if (prior != null) { CometTaskContextShim.set(prior); @@ -175,14 +199,10 @@ private static void evaluateInternal( udfClassName, name -> { try { - // Resolve via the executor's context classloader so user-supplied UDF jars - // (added via spark.jars / --jars) are visible. - ClassLoader cl = Thread.currentThread().getContextClassLoader(); - if (cl == null) { - cl = CometUdfBridge.class.getClassLoader(); - } + // Resolves through the context ClassLoader installed by `evaluate`, so a + // user-supplied CometUDF shipped in a user jar is visible. return (CometUDF) - Class.forName(name, true, cl).getDeclaredConstructor().newInstance(); + ClassLoaders.loadClass(name).getDeclaredConstructor().newInstance(); } catch (ReflectiveOperationException e) { throw new RuntimeException("Failed to instantiate CometUDF: " + name, e); } diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index a3216404ac..57578235f7 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -125,8 +125,11 @@ class CometExecIterator( taskCPUs, keyUnwrapper, // Propagated to Tokio workers running JVM UDFs so they see this Spark task's - // TaskContext. See CometUdfBridge.evaluate. - TaskContext.get()) + // TaskContext and context ClassLoader. Read here because this class is only ever + // constructed on a Spark task thread (see `taskAttemptId` above); a JNI-attached Tokio + // worker has neither. See CometUdfBridge.evaluate. + TaskContext.get(), + Thread.currentThread().getContextClassLoader) } private var nextBatch: Option[ColumnarBatch] = None diff --git a/spark/src/main/scala/org/apache/comet/Native.scala b/spark/src/main/scala/org/apache/comet/Native.scala index 3cfa51b6e1..51b3e5e41b 100644 --- a/spark/src/main/scala/org/apache/comet/Native.scala +++ b/spark/src/main/scala/org/apache/comet/Native.scala @@ -48,6 +48,11 @@ class Native extends NativeBase { * @param taskMemoryManager * the task-level memory manager that is responsible for tracking memory usage across JVM and * native side. + * @param taskContext + * the driving Spark task's `TaskContext`, propagated to Tokio workers running JVM UDFs. + * @param classLoader + * the calling Spark task thread's context ClassLoader, propagated to Tokio workers running + * JVM UDFs. Must be read on the task thread; see `CometUdfBridge.evaluate`. * @return * the address to native query plan. */ @@ -70,7 +75,8 @@ class Native extends NativeBase { taskAttemptId: Long, taskCPUs: Long, keyUnwrapper: CometFileKeyUnwrapper, - taskContext: TaskContext): Long + taskContext: TaskContext, + classLoader: ClassLoader): Long // scalastyle:on /** diff --git a/spark/src/main/scala/org/apache/comet/serde/datetime.scala b/spark/src/main/scala/org/apache/comet/serde/datetime.scala index a2600bf688..0835a92dd7 100644 --- a/spark/src/main/scala/org/apache/comet/serde/datetime.scala +++ b/spark/src/main/scala/org/apache/comet/serde/datetime.scala @@ -450,12 +450,6 @@ object CometNextDay extends CometExpressionSerde[NextDay] { override def getIncompatibleReasons(): Seq[String] = DatetimeCollation.incompatibleReasons("next_day") - override def getCompatibleNotes(): Seq[String] = Seq( - "Under ANSI mode, an invalid `dayOfWeek` surfaces as `CometNativeException` rather than" + - " Spark's `SparkIllegalArgumentException` with error class `ILLEGAL_DAY_OF_WEEK`. The" + - " throw/NULL decision is correct; only the exception class and error class differ" + - " ([#5073](https://github.com/apache/datafusion-comet/issues/5073)).") - override def getSupportLevel(expr: NextDay): SupportLevel = { if (DatetimeCollation.hasNonDefaultCollation(expr)) { Incompatible(Some(collationReason)) @@ -483,11 +477,11 @@ object CometMakeDate extends CometExpressionSerde[MakeDate] { */ override def getCompatibleNotes(): Seq[String] = Seq( - "Under ANSI mode, an out-of-range `(year, month, day)` triple surfaces as" + - " `CometNativeException` rather than Spark's `SparkDateTimeException` with error class" + - " `DATETIME_FIELD_OUT_OF_BOUNDS.WITH_SUGGESTION`. The throw/NULL decision is correct;" + - " only the exception class and error class differ" + - " ([#5073](https://github.com/apache/datafusion-comet/issues/5073)).") + "Native `make_date` is limited to chrono's year range `[-262143, 262142]`; Spark accepts" + + " wider years (for example, `300000`), so Comet returns `NULL` or throws under ANSI mode" + + " for dates Spark accepts, and may incorrectly report valid dates as invalid (for example," + + " `300000-02-29` is falsely reported as not a leap year)" + + " ([#5208](https://github.com/apache/datafusion-comet/issues/5208)).") override def convert(expr: MakeDate, inputs: Seq[Attribute], binding: Boolean): Option[Expr] = { val childExpr = expr.children.map(exprToProtoInternal(_, inputs, binding)) diff --git a/spark/src/main/spark-3.4/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala b/spark/src/main/spark-3.4/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala index 09ac063cd2..6cb6448672 100644 --- a/spark/src/main/spark-3.4/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala +++ b/spark/src/main/spark-3.4/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala @@ -178,6 +178,16 @@ trait ShimSparkErrorConverter { Some( QueryExecutionErrors.ansiDateTimeParseError(new Exception(params("message").toString))) + case "IllegalDayOfWeek" => + Some( + QueryExecutionErrors + .ansiIllegalArgumentError(s"Illegal input for day of week: ${params("string")}")) + + case "DatetimeFieldOutOfBounds" => + Some( + QueryExecutionErrors.ansiDateTimeError( + new java.time.DateTimeException(params("rangeMessage").toString))) + case "InvalidFractionOfSecond" => Some(QueryExecutionErrors.invalidFractionOfSecondError()) diff --git a/spark/src/main/spark-3.5/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala b/spark/src/main/spark-3.5/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala index c502e4d55d..6b976e55de 100644 --- a/spark/src/main/spark-3.5/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala +++ b/spark/src/main/spark-3.5/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala @@ -175,6 +175,16 @@ trait ShimSparkErrorConverter { Some( QueryExecutionErrors.ansiDateTimeParseError(new Exception(params("message").toString))) + case "IllegalDayOfWeek" => + Some( + QueryExecutionErrors + .ansiIllegalArgumentError(s"Illegal input for day of week: ${params("string")}")) + + case "DatetimeFieldOutOfBounds" => + Some( + QueryExecutionErrors.ansiDateTimeError( + new java.time.DateTimeException(params("rangeMessage").toString))) + case "InvalidFractionOfSecond" => Some(QueryExecutionErrors.invalidFractionOfSecondError()) diff --git a/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala b/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala index 874a6af97c..7fb822b66b 100644 --- a/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala +++ b/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala @@ -23,8 +23,7 @@ import java.io.FileNotFoundException import scala.util.matching.Regex -import org.apache.spark.QueryContext -import org.apache.spark.SparkException +import org.apache.spark.{QueryContext, SparkException, SparkIllegalArgumentException} import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.sql.execution.datasources.SchemaColumnConvertNotSupportedException import org.apache.spark.sql.types._ @@ -201,6 +200,17 @@ trait ShimSparkErrorConverter { new Exception(params("message").toString), params("suggestedFunc").toString)) + case "IllegalDayOfWeek" => + Some( + new SparkIllegalArgumentException( + errorClass = "ILLEGAL_DAY_OF_WEEK", + messageParameters = Map("string" -> params("string").toString))) + + case "DatetimeFieldOutOfBounds" => + Some( + QueryExecutionErrors.ansiDateTimeArgumentOutOfRange( + new java.time.DateTimeException(params("rangeMessage").toString))) + case "InvalidFractionOfSecond" => Some(QueryExecutionErrors.invalidFractionOfSecondError(params("value").toString.toDouble)) diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/make_date_ansi.sql b/spark/src/test/resources/sql-tests/expressions/datetime/make_date_ansi.sql index 44880e96b1..75ca530429 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/make_date_ansi.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/make_date_ansi.sql @@ -16,10 +16,11 @@ -- under the License. -- ANSI mode: Spark's MakeDate wraps the java.time.DateTimeException raised by LocalDate.of in --- ansiDateTimeArgumentOutOfRange (4.0, DATETIME_FIELD_OUT_OF_BOUNDS) / ansiDateTimeError (3.4/3.5) --- when spark.sql.ansi.enabled=true. Comet's native SparkMakeDate now throws the same --- java.time-style message under ANSI instead of returning NULL. The expect_error patterns below --- are substrings of that message and match across Spark versions. +-- ansiDateTimeArgumentOutOfRange (4.x) / ansiDateTimeError (3.x) when +-- spark.sql.ansi.enabled=true. The DATETIME_FIELD_OUT_OF_BOUNDS subclass and message parameters +-- vary by Spark version. Comet's native SparkMakeDate throws the same java.time-style message +-- under ANSI instead of returning NULL. The expect_error patterns below are substrings of that +-- message and match across Spark versions. -- Config: spark.sql.ansi.enabled=true -- Sentinel: a valid date must still execute natively under ANSI. This guards against the diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/next_day_ansi.sql b/spark/src/test/resources/sql-tests/expressions/datetime/next_day_ansi.sql index f289f5b34d..f7f018b4eb 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/next_day_ansi.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/next_day_ansi.sql @@ -16,7 +16,8 @@ -- under the License. -- ANSI mode: Spark's NextDay throws on a malformed dayOfWeek (SparkIllegalArgumentException / --- ILLEGAL_DAY_OF_WEEK on 3.5+, IllegalArgumentException on 3.4) when spark.sql.ansi.enabled=true. +-- ILLEGAL_DAY_OF_WEEK on 4.0+, _LEGACY_ERROR_TEMP_2000 on 3.x) when +-- spark.sql.ansi.enabled=true. -- Comet's native next_day now throws the same "Illegal input for day of week" message under ANSI -- instead of returning NULL. -- Config: spark.sql.ansi.enabled=true diff --git a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala index 8091d9dd0a..549e7ab05e 100644 --- a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala @@ -19,8 +19,6 @@ package org.apache.comet -import java.time.{Duration, Period} - import scala.util.Random import org.apache.hadoop.fs.Path @@ -2326,11 +2324,19 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { CometConf.COMET_EXEC_ENABLED.key -> "true")(f) } - def checkOverflow(query: String, dtype: String): Unit = { + // Spark 3.4/3.5 throw `_LEGACY_ERROR_TEMP_2043` ("- caused overflow.") for byte and + // short. Spark 4.x routes them through `MathUtils.negateExact` and agrees with Comet, which + // always renders `SparkError::ArithmeticOverflow` with the Spark type name. + def sparkOverflowMsg(dtype: String): String = + if (isSpark40Plus) s"$dtype overflow" else "caused overflow" + + // Spark and Comet can render different overflow messages for the same operation, so assert each + // side's expected substring separately. + def checkOverflow(query: String, sparkExpected: String, cometExpected: String): Unit = { checkSparkAnswerMaybeThrows(sql(query)) match { case (Some(sparkException), Some(cometException)) => - assert(sparkException.getMessage.contains(dtype + " overflow")) - assert(cometException.getMessage.contains(dtype + " overflow")) + assert(sparkException.getMessage.contains(sparkExpected)) + assert(cometException.getMessage.contains(cometExpected)) case (None, None) => checkSparkAnswerAndOperator(sql(query)) case (None, Some(ex)) => fail("Comet threw an exception but Spark did not " + ex.getMessage) @@ -2339,44 +2345,50 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } - def runArrayTest(query: String, dtype: String, path: String): Unit = { + def runArrayTest( + query: String, + sparkExpected: String, + cometExpected: String, + path: String): Unit = { withParquetTable(path, "t") { withAnsiMode(enabled = false) { checkSparkAnswerAndOperator(sql(query)) } withAnsiMode(enabled = true) { - checkOverflow(query, dtype) + checkOverflow(query, sparkExpected, cometExpected) } } } withTempDir { dir => - // Array values test + // Array values test. Tuple is (file, data, Spark-expected substring, Comet-expected substring). val dataTypes = Seq( - ("array_test.parquet", Seq(Int.MaxValue, Int.MinValue).toDF("a"), "integer"), - ("long_array_test.parquet", Seq(Long.MaxValue, Long.MinValue).toDF("a"), "long"), - ("short_array_test.parquet", Seq(Short.MaxValue, Short.MinValue).toDF("a"), ""), - ("byte_array_test.parquet", Seq(Byte.MaxValue, Byte.MinValue).toDF("a"), "")) - - dataTypes.foreach { case (fileName, df, dtype) => + ( + "array_test.parquet", + Seq(Int.MaxValue, Int.MinValue).toDF("a"), + "integer overflow", + "integer overflow"), + ( + "long_array_test.parquet", + Seq(Long.MaxValue, Long.MinValue).toDF("a"), + "long overflow", + "long overflow"), + ( + "short_array_test.parquet", + Seq(Short.MaxValue, Short.MinValue).toDF("a"), + sparkOverflowMsg("short"), + "short overflow"), + ( + "byte_array_test.parquet", + Seq(Byte.MaxValue, Byte.MinValue).toDF("a"), + sparkOverflowMsg("byte"), + "byte overflow")) + + dataTypes.foreach { case (fileName, df, sparkExpected, cometExpected) => val path = new Path(dir.toURI.toString, fileName).toString df.write.mode("overwrite").parquet(path) val query = "select a, -a from t" - runArrayTest(query, dtype, path) - } - - withParquetTable((0 until 5).map(i => (i % 5, i % 3)), "tbl") { - withAnsiMode(enabled = true) { - // interval test without cast - val longDf = Seq(Long.MaxValue, Long.MaxValue, 2) - val yearMonthDf = Seq(Int.MaxValue, Int.MaxValue, 2) - .map(Period.ofMonths) - val dayTimeDf = Seq(106751991L, 106751991L, 2L) - .map(Duration.ofDays) - Seq(longDf, yearMonthDf, dayTimeDf).foreach { _ => - checkOverflow("select -(_1) FROM tbl", "") - } - } + runArrayTest(query, sparkExpected, cometExpected, path) } // scalar tests @@ -2387,19 +2399,32 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { CometConf.COMET_ENABLED.key -> "true", CometConf.COMET_EXEC_ENABLED.key -> "true") { for (n <- Seq("2147483647", "-2147483648")) { - checkOverflow(s"select -(cast(${n} as int)) FROM tbl", "integer") + checkOverflow( + s"select -(cast(${n} as int)) FROM tbl", + "integer overflow", + "integer overflow") } for (n <- Seq("32767", "-32768")) { - checkOverflow(s"select -(cast(${n} as short)) FROM tbl", "") + checkOverflow( + s"select -(cast(${n} as short)) FROM tbl", + sparkOverflowMsg("short"), + "short overflow") } for (n <- Seq("127", "-128")) { - checkOverflow(s"select -(cast(${n} as byte)) FROM tbl", "") + checkOverflow( + s"select -(cast(${n} as byte)) FROM tbl", + sparkOverflowMsg("byte"), + "byte overflow") } for (n <- Seq("9223372036854775807", "-9223372036854775808")) { - checkOverflow(s"select -(cast(${n} as long)) FROM tbl", "long") + checkOverflow( + s"select -(cast(${n} as long)) FROM tbl", + "long overflow", + "long overflow") } + // Float negation cannot overflow; confirm it stays native and returns the negated value. for (n <- Seq("3.4028235E38", "-3.4028235E38")) { - checkOverflow(s"select -(cast(${n} as float)) FROM tbl", "float") + checkSparkAnswerAndOperator(sql(s"select -(cast(${n} as float)) FROM tbl")) } } } diff --git a/spark/src/test/scala/org/apache/comet/CometScalaUDFClassLoaderSuite.scala b/spark/src/test/scala/org/apache/comet/CometScalaUDFClassLoaderSuite.scala new file mode 100644 index 0000000000..17fd0fea09 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/CometScalaUDFClassLoaderSuite.scala @@ -0,0 +1,195 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.comet + +import java.io.File +import java.net.URLClassLoader +import java.nio.charset.StandardCharsets.UTF_8 +import java.nio.file.{Files, Path, Paths} +import javax.tools.ToolProvider + +import org.apache.spark.SparkConf +import org.apache.spark.sql.{CometTestBase, Encoders} +import org.apache.spark.sql.catalyst.expressions.Expression + +/** + * Regression coverage for UDF closures whose capturing class lives in a user jar, which used to + * fail with: + * + * {{{ + * java.lang.ClassCastException: cannot assign instance of java.lang.invoke.SerializedLambda + * to field org.apache.spark.sql.catalyst.expressions.ScalaUDF.f of type scala.Function1 + * at org.apache.comet.udf.codegen.CometScalaUDFCodegen.lookupOrCompile + * }}} + * + * That exception is a masked `ClassNotFoundException`: when the deserializing ClassLoader cannot + * resolve a lambda's capturing class, `ObjectInputStream` records the CNFE against the object + * handle, therefore skips `SerializedLambda.readResolve`, and the raw `SerializedLambda` then + * fails the field-type check in `defaultCheckFieldValues`. The classloading rationale lives on + * `CometUdfBridge.evaluate`. + * + * The fixture puts the capturing class on `spark.executor.extraClassPath`, the local-mode + * equivalent of a `--jars` submission: `LocalSchedulerBackend` feeds it into the executor's + * `MutableURLClassLoader`, which becomes the task thread's context ClassLoader while staying + * invisible to the ClassLoader that loaded Comet. + */ +class CometScalaUDFClassLoaderSuite extends CometTestBase { + + import CometScalaUDFClassLoaderSuite._ + + override protected def sparkConf: SparkConf = + super.sparkConf + .set(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key, "true") + .set("spark.executor.extraClassPath", hiddenClassesDir.toString) + + private def withHiddenUdfTable(f: => Unit): Unit = { + spark.udf.register("hiddenUdf", hiddenFn) + withTable("t") { + sql("CREATE TABLE t (s STRING) USING parquet") + sql("INSERT INTO t VALUES ('a'), ('b'), (NULL)") + f + } + } + + test("fixture: hidden class reachable from task threads, not from Comet's ClassLoader") { + // Guards the tests below from passing vacuously. `classOf[Expression].getClassLoader` is the + // fallback `lookupOrCompile` uses when the calling thread has no context ClassLoader. + intercept[ClassNotFoundException] { + Class.forName(HiddenClassName, false, classOf[Expression].getClassLoader) + } + val reachable = spark + .range(4) + .repartition(2) + .mapPartitions(_ => Iterator(canLoadHiddenClass))(Encoders.scalaBoolean) + .collect() + assert(reachable.forall(identity), s"task threads could not load $HiddenClassName") + } + + // Both leaf shapes matter. With the native scan the dispatcher runs on a Tokio worker, which has + // no context ClassLoader of its own and depends on the propagated one; without it the dispatcher + // runs on the Spark task thread, which already has the executor's ClassLoader installed. + Seq(true, false).foreach { nativeScan => + test(s"ScalaUDF closure captured by a user-jar class, nativeScan=$nativeScan") { + withSQLConf(CometConf.COMET_NATIVE_SCAN_ENABLED.key -> nativeScan.toString) { + withHiddenUdfTable { + checkSparkAnswer(sql("SELECT hiddenUdf(s) FROM t")) + } + } + } + } + + test("the thread running the UDF sees the task thread's context ClassLoader") { + // Pins the propagation itself rather than its symptom: the UDF body reports whether the + // ClassLoader installed on whatever thread invoked it can reach the user jar. + spark.udf.register("loaderProbe", (_: String) => loaderReport()) + withHiddenUdfTable { + val reports = sql("SELECT loaderProbe(s) FROM t").collect().map(_.getString(0)).distinct + assert( + reports.forall(_.startsWith("loaded|")), + s"UDF thread could not load $HiddenClassName: ${reports.mkString(", ")}") + } + } +} + +object CometScalaUDFClassLoaderSuite { + + val HiddenClassName = "hidden.HiddenUdf" + + /** + * Directory holding the compiled capturing class. Compiled once per JVM, before the + * SparkSession starts, because `spark.executor.extraClassPath` is read at session creation. A + * directory works as a classpath entry, so there is no need to package a jar. + */ + lazy val hiddenClassesDir: Path = compileHiddenClass() + + /** + * The UDF, obtained through a ClassLoader over `hiddenClassesDir` alone. Deliberately not + * closed: the loaded class stays live in the registered UDF for the rest of the JVM's life. + */ + lazy val hiddenFn: String => String = + new URLClassLoader(Array(hiddenClassesDir.toUri.toURL), getClass.getClassLoader) + .loadClass(HiddenClassName) + .getMethod("make") + .invoke(null) + .asInstanceOf[String => String] + + /** Runs on Spark threads; kept on the companion so the closures capture nothing. */ + def canLoadHiddenClass: Boolean = + try { + Class.forName(HiddenClassName, false, Thread.currentThread().getContextClassLoader) + true + } catch { + case _: ClassNotFoundException => false + } + + /** + * Reports whether the current thread's context ClassLoader reaches the user jar, tagged with + * the thread name so a failure says which thread was missing it. + */ + def loaderReport(): String = { + val status = if (canLoadHiddenClass) "loaded" else "MISSING" + s"$status|${Thread.currentThread().getName}" + } + + private def compileHiddenClass(): Path = { + // createTempDirectory does not create its parent (java.io.tmpdir, pinned to target/tmp by the + // pom), which does not exist yet on a fresh checkout, so make it up front. + val tmpRoot = Files.createDirectories(Paths.get(System.getProperty("java.io.tmpdir"))) + val workDir = Files.createTempDirectory(tmpRoot, "comet-hidden-udf") + val src = workDir.resolve("HiddenUdf.java") + // The intersection cast is what makes the lambda serializable, and therefore what makes + // `hidden.HiddenUdf` the capturing class recorded in the SerializedLambda. + Files.write( + src, + """package hidden; + | + |import java.io.Serializable; + |import scala.Function1; + | + |public class HiddenUdf { + | public static Function1 make() { + | return (Function1 & Serializable) + | (Object o) -> (o == null ? null : "hidden:" + o); + | } + |} + |""".stripMargin.getBytes(UTF_8)) + + val classesDir = Files.createDirectories(workDir.resolve("classes")) + val compiler = ToolProvider.getSystemJavaCompiler + assert(compiler != null, "test must run on a JDK (needs the javax.tools compiler)") + // Only scala-library is needed. Handing javac the whole test classpath makes it open and index + // every jar on it, which costs more than the compile itself. + val classpath = Option(classOf[Function1[_, _]].getProtectionDomain.getCodeSource) + .map(_.getLocation.getPath) + .getOrElse(System.getProperty("java.class.path")) + val rc = + compiler.run(null, null, null, "-cp", classpath, "-d", classesDir.toString, src.toString) + assert(rc == 0, s"javac failed with exit code $rc") + + deleteOnExitRecursively(workDir.toFile) + classesDir + } + + /** Parents are registered before children, and deletion runs in reverse registration order. */ + private def deleteOnExitRecursively(file: File): Unit = { + file.deleteOnExit() + Option(file.listFiles()).foreach(_.foreach(deleteOnExitRecursively)) + } +} diff --git a/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala index 90d6992047..4df944527b 100644 --- a/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala @@ -21,6 +21,7 @@ package org.apache.comet import scala.util.Random +import org.apache.spark.SparkThrowable import org.apache.spark.sql.{CometTestBase, DataFrame, Row, SaveMode} import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute import org.apache.spark.sql.catalyst.expressions.{Days, Hours, Literal} @@ -39,6 +40,47 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH private val crossTimezones = Seq("UTC", "America/Los_Angeles", "Europe/London", "Asia/Tokyo") + private def causeChain(error: Throwable): Seq[Throwable] = + Iterator.iterate(error)(_.getCause).takeWhile(_ != null).toSeq + + private def deepestSparkThrowable(error: Throwable): SparkThrowable with Throwable = + causeChain(error) + .collect { case e: SparkThrowable with Throwable => e } + .lastOption + .getOrElse( + fail(s"No SparkThrowable in cause chain: ${causeChain(error).map(_.getClass.getName)}")) + + test("next_day and make_date ANSI errors match Spark exceptions") { + withSQLConf( + SQLConf.ANSI_ENABLED.key -> "true", + SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> + "org.apache.spark.sql.catalyst.optimizer.ConstantFolding") { + Seq( + "SELECT next_day(date('2024-01-01'), 'NOT_A_DAY')", + "SELECT make_date(2024, 13, 1)", + // 999999999 instead overflows the epoch-day conversion as a plain ArithmeticException. + "SELECT make_date(1000000000, 1, 1)", + "SELECT make_date(1000000000, 13, 0)") + .foreach { query => + val df = sql(query) + checkCometOperators(stripAQEPlan(df.queryExecution.executedPlan)) + + val (sparkError, cometError) = checkSparkAnswerMaybeThrows(df) + val sparkFailure = sparkError.getOrElse(fail(s"Spark did not fail for: $query")) + val cometFailure = cometError.getOrElse(fail(s"Comet did not fail for: $query")) + val expected = deepestSparkThrowable(sparkFailure) + val actual = deepestSparkThrowable(cometFailure) + + assert(actual.getClass == expected.getClass) + assert(actual.getErrorClass == expected.getErrorClass) + assert(actual.getSqlState == expected.getSqlState) + assert(actual.getMessageParameters == expected.getMessageParameters) + assert(actual.getMessage == expected.getMessage) + assert(!causeChain(cometFailure).exists(_.isInstanceOf[CometNativeException])) + } + } + } + test("trunc (TruncDate)") { val supportedFormats = CometTruncDate.supportedFormats val unsupportedFormats = Seq("invalid")