diff --git a/docs/source/changelog/1.2.0.md b/docs/source/changelog/1.2.0.md new file mode 100644 index 00000000000..4303b074edc --- /dev/null +++ b/docs/source/changelog/1.2.0.md @@ -0,0 +1,29 @@ + + +# DataFusion Comet 1.2.0 Changelog (Unreleased) + +## Behaviour Changes + +Native integral `sequence` now has a per-invocation output byte ceiling, derived from the configured +memory pool and task slots by default in off-heap mode. Valid queries that exceed it can fail before +output allocation. Set `spark.comet.exec.sequence.maxBytesPerInvocation=-1` to restore the previous +absence of this byte ceiling. Existing per-row and batch element limits remain in effect. See the +[Upgrade Guide](../user-guide/latest/migration-guide.md#native-integral-sequence-byte-ceiling) for +scope and tuning options. diff --git a/docs/source/changelog/index.md b/docs/source/changelog/index.md index 792dfa2df3b..ba6f71552e3 100644 --- a/docs/source/changelog/index.md +++ b/docs/source/changelog/index.md @@ -24,6 +24,7 @@ Per-release change logs for Apache DataFusion Comet. ```{toctree} :maxdepth: 1 +1.2.0 (Unreleased) <1.2.0> 1.1.0 <1.1.0> 1.0.0 <1.0.0> 0.17.1 <0.17.1> diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index 184bf0de275..512f18f084f 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -80,7 +80,8 @@ page: | --------------------------------------- | ----------- | ------------------------------------------------------------- | ----------------- | | Spark execution + storage (on-heap) | JVM heap | `spark.executor.memory` and the unified memory manager | Yes | | Spark Tungsten (off-heap) | Off-heap | `spark.memory.offHeap.size` via `TaskMemoryManager` | Yes | -| Comet native (Rust global allocator) | Native heap | `memory_limit` (see below), enforced only via the memory pool | Reservations only | +| Comet native (Rust global allocator) | Native heap | Memory pools for declared reservations; no allocator-wide limit | Reservations only | +| Comet integral `sequence` output | Native heap | Logical per-invocation ceiling, without a pool reservation | No | | Comet JVM Arrow (`CometArrowAllocator`) | Off-heap | **Nothing**: a `RootAllocator(Long.MaxValue)` | No | | Comet JVM shuffle pages | Off-heap | `spark.memory.offHeap.size` via `TaskMemoryManager` | Yes | @@ -97,8 +98,10 @@ allocates the bytes and in who is able to see them. - **Native heap** is allocated by Rust through its global allocator, in the same process. No JVM code makes the call, so no JVM-side allocator or metric ever measures these bytes. The only layer that sees any of them is Comet's own memory pool, and then only the portion that operators - explicitly reserve. That portion is still charged to Spark, as the next paragraph describes; what - is never reserved is measured by nothing and budgeted by nobody. + explicitly reserve. That portion is still charged to Spark, as the next paragraph describes. + Unreserved allocations have no pool accounting. Individual guards, such as integral `sequence`'s + per-invocation ceiling, can refuse a logical output without making a reservation; they do not + bound total native usage. **The budget is shared even though the memory is not.** `memory_limit` is derived from `spark.memory.offHeap.size` (see [Where Comet's budget comes from](#where-comets-budget-comes-from)), @@ -257,6 +260,29 @@ back a slice of the off-heap pool for the memory Comet does not reserve, but it The only room Spark leaves for memory outside the pool is `spark.executor.memoryOverhead`. +### Integral sequence invocation ceiling + +`getMemoryConfig` also resolves `spark.comet.exec.sequence.maxBytesPerInvocation`. Unless explicitly +overridden, each native integral `sequence` invocation is limited to `memory_limit` divided by the +configured task slots: `max(1, floor(local or executor cores / CPUs allocated to the task))` when +the core count is known. Local-cluster uses its configured worker cores if `spark.executor.cores` +is unset. Standalone can assign cores without exposing their count in the executor configuration; +in that case the fallback is the whole `memory_limit`, not an inferred task share. Set +`spark.executor.cores` or an explicit sequence ceiling to choose a smaller limit. The actual task +CPU allocation and executor configuration respect resource profiles. This fixed ceiling does not +depend on live tasks, registered consumers or pool usage. On-heap mode has no default ceiling; +`-1` explicitly disables it and `0` is a real zero-byte ceiling. The resolved value crosses JNI in the existing +config map, without changing the memory pool or reserving anything from it. + +The kernel checks logical output values, list offsets and validity bytes after validating all +rows, before allocating its values and offsets. Row-length scratch storage and validity preparation +can precede that check. Allocator padding and shared input backing outside a validity slice do not +count. Earlier outputs, concurrent invocations and downstream sort, shuffle or explode do not +consume the ceiling. A scalar invocation still produces one list; generic scalar broadcasts and +copies of folded literals happen outside it. This bounds one expansion, and does not close the +accounting or lifetime gaps described below. See +[Limiting Integral Sequence Output](../user-guide/latest/tuning/memory.md#limiting-integral-sequence-output). + ### Resolving the pool type `parse_memory_pool_config` (`native/core/src/execution/memory_pools/config.rs`) turns the pool-type @@ -610,10 +636,11 @@ somewhat less abrupt. ## Open problems -Nothing described above prevents an executor from being OOM-killed. Every enforcement point Comet -has bounds _declared reservations_, and the sections above describe several structural reasons why -declared reservations are a lower bound on physical usage. The known gaps, roughly in order of how -much they matter: +Nothing described above guarantees that an executor avoids being OOM-killed. Memory pools bound +_declared reservations_, while the integral `sequence` ceiling bounds one logical output without +reserving it. Neither bounds total resident memory, and the sections above describe several reasons +why declared reservations are a lower bound on physical usage. The known gaps, roughly in order of +how much they matter: - **Real native usage is observed but not acted on.** Comet counts the bytes Rust's allocator has handed out, and each executor logs that count next to the pools' reservations. Nothing reads it diff --git a/docs/source/user-guide/latest/migration-guide.md b/docs/source/user-guide/latest/migration-guide.md index c7ceea0ee60..6d320a223be 100644 --- a/docs/source/user-guide/latest/migration-guide.md +++ b/docs/source/user-guide/latest/migration-guide.md @@ -59,8 +59,22 @@ key is removed. ## Upgrading to Comet 1.2.0 -Comet `1.2.0` makes no behavior changes that need a `spark.comet.legacy.*` key. The changes below -need none either, but check whether any of them applies to your deployment. +Review the following changes when upgrading to Comet `1.2.0`. + +### Native Integral Sequence Byte Ceiling + +Native integral `sequence` now checks each invocation's logical Arrow output size before allocating +its values and offsets. In off-heap mode the default ceiling comes from the configured memory pool +limit and task slots, so a valid query that previously completed can now fail when one invocation +exceeds that ceiling. On-heap test mode has no default ceiling. + +Set `spark.comet.exec.sequence.maxBytesPerInvocation=-1` to restore the pre-1.2.0 absence of this +byte ceiling. Existing Spark per-row and Arrow batch element limits still apply. + +To retain the new check, reduce `spark.comet.batchSize`, shorten the sequences or choose a larger +`spark.comet.exec.sequence.maxBytesPerInvocation`. Retained results, concurrent calls and later +scalar or folded-literal expansion do not consume an invocation's ceiling. See +[Limiting Integral Sequence Output](tuning/memory.md#limiting-integral-sequence-output). ### In-Memory Cache Enabled by Default diff --git a/docs/source/user-guide/latest/tuning/memory.md b/docs/source/user-guide/latest/tuning/memory.md index 2632369df24..eb129cc0498 100644 --- a/docs/source/user-guide/latest/tuning/memory.md +++ b/docs/source/user-guide/latest/tuning/memory.md @@ -272,6 +272,43 @@ narrow tables when memory is plentiful. `spark.comet.shuffle.jvm.batchSize` controls the batch size used when the JVM columnar shuffle writer flushes sorted spill files. A value larger than `spark.comet.batchSize` is capped at `spark.comet.batchSize`. +## Limiting Integral Sequence Output + +`spark.comet.exec.sequence.maxBytesPerInvocation` limits the logical Arrow output of each native +integral `sequence` invocation. The kernel validates all rows and checks this ceiling before +allocating the output values and offsets. A refusal fails the task with the required byte count and +the limit. Lower `spark.comet.batchSize` to fit fewer sequences in an invocation, shorten a sequence +that is too large on its own, or explicitly raise this setting. Byte units such as `64m` are accepted; +`-1` restores the pre-1.2.0 absence of this byte ceiling, and `0` allows no output bytes. Existing +Spark per-row and Arrow batch element limits still apply. See the +[Upgrade Guide](../migration-guide.md). + +When the setting is absent in off-heap mode and the core count is known, the default is: + +```text +task slots = max(1, floor(local or executor cores / CPUs allocated to the task)) +sequence ceiling = Comet's configured memory pool limit / task slots +``` + +The pool limit is resolved by `CometExecIterator.getMemoryConfig` from +`spark.memory.offHeap.size * spark.comet.exec.memoryPool.fraction`. Local mode uses the thread count +in `spark.master`; executors use `spark.executor.cores`. For `local-cluster[workers,cores,memory]`, +the configured worker cores are used when executor cores are unset. Standalone may assign cores +without exposing their count in the executor configuration. In that case the fallback is the whole +configured pool limit, not an inferred task share; set `spark.executor.cores` or an explicit +sequence ceiling to choose a smaller limit. + +The task's CPU allocation normally comes from `spark.task.cpus`, and may be overridden by a resource +profile; the executor configuration also reflects non-default profile allocations. The ceiling +does not change with live task counts or pool usage. On-heap test mode has no default ceiling; an +explicit setting still applies there. + +Logical bytes include the child values at their integral width, four bytes per list offset +(including the terminal offset), and one validity bit per row when null lists are present, rounded +up to whole bytes. Even an empty output has a four-byte terminal offset. Allocator padding, scratch +space and shared input backing outside a validity slice are excluded. Validity preparation and the +row-length scratch buffer may be allocated before the final check. + ## Limiting Spill Disk Usage Native operators that spill to disk (aggregate, sort, shuffle) are bounded by diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 262d8b8fd41..171a7719977 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -132,7 +132,7 @@ use crate::execution::tracing::{ use crate::execution::memory_pools::logging_pool::LoggingMemoryPool; use crate::execution::spark_config::{ - SparkConfig, COMET_DEBUG_ENABLED, COMET_DEBUG_MEMORY, + sequence_max_bytes, SparkConfig, COMET_DEBUG_ENABLED, COMET_DEBUG_MEMORY, COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED, COMET_EXEC_SPILL_COMPRESSION_CODEC, COMET_EXPLAIN_NATIVE_ENABLED, COMET_MAX_TEMP_DIRECTORY_SIZE, COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, COMET_STRICT_TIMESTAMP_LABELS, @@ -558,6 +558,8 @@ struct ExecutionContext { pub debug_native: bool, /// Whether boundaries fail on a relabelled timestamp instead of casting it (test-only) pub strict_timestamp_labels: bool, + /// Resolved logical output ceiling for each integral sequence invocation. + sequence_max_bytes: Option, /// Whether to write native plans with metrics to stdout pub explain_native: bool, /// Whether to log memory usage on each call to execute_plan @@ -710,6 +712,7 @@ struct PlanSettings { spark_config: HashMap, debug_native: bool, strict_timestamp_labels: bool, + sequence_max_bytes: Option, explain_native: bool, tracing_enabled: bool, max_temp_directory_size: u64, @@ -731,6 +734,7 @@ fn prepare_plan_settings(serialized_spark_configs: &[u8]) -> CometResult, } impl Default for PhysicalPlanner { @@ -379,6 +380,7 @@ impl PhysicalPlanner { shuffle_partition_pusher: None, float_operands: FloatOperands::Normalize, strict_timestamp_labels: false, + sequence_max_bytes: None, } } @@ -462,6 +464,11 @@ impl PhysicalPlanner { self } + pub fn with_sequence_max_bytes(mut self, max_bytes: Option) -> Self { + self.sequence_max_bytes = max_bytes; + self + } + pub(crate) fn strict_timestamp_labels(&self) -> bool { self.strict_timestamp_labels } @@ -4085,12 +4092,19 @@ impl PhysicalPlanner { } }; - let fun_expr = create_comet_physical_fun( - fun_name, - data_type.clone(), - &self.session_ctx.state(), - Some(expr.fail_on_error), - )?; + let fun_expr = if fun_name == "spark_sequence" { + Arc::new(ScalarUDF::new_from_impl(SparkSequence::new( + data_type.clone(), + self.sequence_max_bytes, + ))) + } else { + create_comet_physical_fun( + fun_name, + data_type.clone(), + &self.session_ctx.state(), + Some(expr.fail_on_error), + )? + }; let args = args .into_iter() diff --git a/native/core/src/execution/spark_config.rs b/native/core/src/execution/spark_config.rs index 7755f6d5dc4..1d316002a05 100644 --- a/native/core/src/execution/spark_config.rs +++ b/native/core/src/execution/spark_config.rs @@ -17,6 +17,10 @@ use std::collections::HashMap; +use crate::errors::{CometError, CometResult}; + +pub(crate) const COMET_EXEC_SEQUENCE_MAX_BYTES_PER_INVOCATION: &str = + "spark.comet.exec.sequence.maxBytesPerInvocation"; pub(crate) const COMET_TRACING_ENABLED: &str = "spark.comet.tracing.enabled"; pub(crate) const COMET_DEBUG_ENABLED: &str = "spark.comet.debug.enabled"; pub(crate) const COMET_EXPLAIN_NATIVE_ENABLED: &str = "spark.comet.explain.native.enabled"; @@ -32,6 +36,21 @@ pub(crate) const COMET_STRICT_TIMESTAMP_LABELS: &str = "spark.comet.testing.strictTimestampLabels.enabled"; pub(crate) const SPARK_EXECUTOR_CORES: &str = "spark.executor.cores"; +/// The JVM resolves the task's default and byte units before sending this value. +pub(crate) fn sequence_max_bytes(config: &HashMap) -> CometResult> { + match config + .get(COMET_EXEC_SEQUENCE_MAX_BYTES_PER_INVOCATION) + .map(String::as_str) + { + None | Some("-1") => Ok(None), + Some(value) => value.parse::().map(Some).map_err(|_| { + CometError::Config(format!( + "Invalid resolved {COMET_EXEC_SEQUENCE_MAX_BYTES_PER_INVOCATION}: {value}" + )) + }), + } +} + /// Comet configs read through this trait must be resolved by the JVM first: /// `CometExecIterator.serializeCometSQLConfs` sends booleans as `true` or `false` and sizes as a /// bare byte count. A config missing from its list arrives exactly as the user wrote it, and a @@ -61,3 +80,28 @@ impl SparkConfig for HashMap { .unwrap_or(default_value) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn sequence_ceiling_requires_resolved_bytes() { + let mut config = HashMap::new(); + assert_eq!(sequence_max_bytes(&config).unwrap(), None); + for (value, expected) in [("-1", None), ("0", Some(0)), ("1024", Some(1024))] { + config.insert( + COMET_EXEC_SEQUENCE_MAX_BYTES_PER_INVOCATION.into(), + value.into(), + ); + assert_eq!(sequence_max_bytes(&config).unwrap(), expected); + } + for value in ["-2", "1k", "invalid", "18446744073709551616"] { + config.insert( + COMET_EXEC_SEQUENCE_MAX_BYTES_PER_INVOCATION.into(), + value.into(), + ); + assert!(sequence_max_bytes(&config).is_err(), "{value}"); + } + } +} diff --git a/native/spark-expr/benches/sequence.rs b/native/spark-expr/benches/sequence.rs index 200500a0a7e..c06c53fa97f 100644 --- a/native/spark-expr/benches/sequence.rs +++ b/native/spark-expr/benches/sequence.rs @@ -20,7 +20,7 @@ use arrow::datatypes::{DataType, Field}; use criterion::{criterion_group, criterion_main, Criterion}; use datafusion::common::ScalarValue; use datafusion::physical_plan::ColumnarValue; -use datafusion_comet_spark_expr::spark_sequence; +use datafusion_comet_spark_expr::{spark_sequence, spark_sequence_with_limit}; use std::hint::black_box; use std::sync::Arc; @@ -57,6 +57,37 @@ fn criterion_benchmark(c: &mut Criterion) { let mut group = c.benchmark_group("sequence"); + // Keep a configured path beside the existing baseline shapes. The limit fits every output + // here, so this measures the successful guard without timing error formatting. + for (label, elems, null_every) in [ + ("short", 1, None), + ("sparse_nulls", 365, Some(10)), + ("all_nulls", 365, Some(1)), + ] { + let args = args_with_len(elems, null_every); + group.bench_function(format!("limited_{label}"), |b| { + b.iter(|| { + black_box( + spark_sequence_with_limit(&args, &return_type, Some(32 * 1024 * 1024)).unwrap(), + ) + }) + }); + } + + // Fixed guard overhead is most visible when scalar inputs produce one single-element list. + { + let args = vec![ + ColumnarValue::Scalar(ScalarValue::Int64(Some(0))), + ColumnarValue::Scalar(ScalarValue::Int64(Some(0))), + ]; + group.bench_function("limited_all_scalar_one_list_1_elems", |b| { + b.iter(|| { + black_box( + spark_sequence_with_limit(&args, &return_type, Some(32 * 1024 * 1024)).unwrap(), + ) + }) + }); + } // Short sequences: per-row overhead dominates. for elems in [1i64, 2, 5] { let args = args_with_len(elems, None); diff --git a/native/spark-expr/src/array_funcs/mod.rs b/native/spark-expr/src/array_funcs/mod.rs index ff0d55f0bdf..148d9936687 100644 --- a/native/spark-expr/src/array_funcs/mod.rs +++ b/native/spark-expr/src/array_funcs/mod.rs @@ -45,7 +45,7 @@ pub use get_array_struct_fields::GetArrayStructFields; pub use list_extract::ListExtract; pub use list_positions::ListPositionsExpr; pub use nested_comparison::{spark_comparison, spark_in_list, FloatOperands, SparkComparison}; -pub use sequence::spark_sequence; +pub use sequence::{spark_sequence, spark_sequence_with_limit, SparkSequence}; pub use size::{spark_size, SparkSizeFunc}; pub use sort_array::SparkSortArray; diff --git a/native/spark-expr/src/array_funcs/sequence.rs b/native/spark-expr/src/array_funcs/sequence.rs index 17d519c7ce6..b8f51c3fe5a 100644 --- a/native/spark-expr/src/array_funcs/sequence.rs +++ b/native/spark-expr/src/array_funcs/sequence.rs @@ -36,14 +36,64 @@ use arrow::datatypes::{ }; use datafusion::common::cast::as_primitive_array; use datafusion::common::{exec_err, internal_err, DataFusionError, Result, ScalarValue}; -use datafusion::logical_expr::ColumnarValue; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, +}; use crate::SparkError; /// Spark's ByteArrayMethods.MAX_ROUNDED_ARRAY_LENGTH (Integer.MAX_VALUE - 15). const MAX_ROUNDED_ARRAY_LENGTH: i64 = (i32::MAX - 15) as i64; +/// A limit on one invocation's logical output, independent of other calls and their results. +#[derive(Debug, Hash, Eq, PartialEq)] +pub struct SparkSequence { + signature: Signature, + data_type: DataType, + max_bytes: Option, +} + +impl SparkSequence { + pub fn new(data_type: DataType, max_bytes: Option) -> Self { + Self { + signature: Signature::variadic_any(Volatility::Immutable), + data_type, + max_bytes, + } + } +} + +impl ScalarUDFImpl for SparkSequence { + fn name(&self) -> &str { + "spark_sequence" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _: &[DataType]) -> Result { + Ok(self.data_type.clone()) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + // All-scalar arguments materialise one list. Generic scalar broadcasting and copies of + // folded literals happen outside this invocation and are not charged to its ceiling. + spark_sequence_with_limit(&args.args, &self.data_type, self.max_bytes) + } +} + pub fn spark_sequence(args: &[ColumnarValue], data_type: &DataType) -> Result { + spark_sequence_with_limit(args, data_type, None) +} + +/// Check logical Arrow output bytes before allocating values and offsets. This is not a pool +/// reservation: earlier results, concurrent invocations and downstream consumers share no tally. +pub fn spark_sequence_with_limit( + args: &[ColumnarValue], + data_type: &DataType, + max_bytes: Option, +) -> Result { let child_field = match data_type { DataType::List(field) => Arc::clone(field), other => return exec_err!("spark_sequence expects a List return type, got {other:?}"), @@ -77,10 +127,10 @@ pub fn spark_sequence(args: &[ColumnarValue], data_type: &DataType) -> Result sequence_integral::(args, num_rows, child_field), - DataType::Int16 => sequence_integral::(args, num_rows, child_field), - DataType::Int32 => sequence_integral::(args, num_rows, child_field), - DataType::Int64 => sequence_integral::(args, num_rows, child_field), + DataType::Int8 => sequence_integral::(args, num_rows, child_field, max_bytes), + DataType::Int16 => sequence_integral::(args, num_rows, child_field, max_bytes), + DataType::Int32 => sequence_integral::(args, num_rows, child_field, max_bytes), + DataType::Int64 => sequence_integral::(args, num_rows, child_field, max_bytes), other => exec_err!("spark_sequence does not support element type {other:?}"), }?; @@ -155,6 +205,7 @@ fn sequence_integral( args: &[ColumnarValue], num_rows: usize, child_field: FieldRef, + max_bytes: Option, ) -> Result where T::Native: Into, @@ -175,6 +226,7 @@ where .as_ref() .is_some_and(|nulls| nulls.null_count() == num_rows) { + check_output_size(num_rows, 0, size_of::(), true, max_bytes)?; return Ok(Arc::new(ListArray::new_null(child_field, num_rows))); } let row_step = |row: usize, start: i64, stop: i64| -> i64 { @@ -209,6 +261,13 @@ where ))); } + check_output_size( + num_rows, + total, + size_of::(), + nulls.is_some(), + max_bytes, + )?; let mut values: Vec = Vec::new(); values.try_reserve_exact(total).map_err(|_| { DataFusionError::External(Box::new(SparkError::SequenceBatchTooLarge { @@ -245,6 +304,46 @@ where )?)) } +fn check_output_size( + num_rows: usize, + total_elements: usize, + element_width: usize, + has_nulls: bool, + max_bytes: Option, +) -> Result<()> { + let Some(limit) = max_bytes else { + return Ok(()); + }; + // Count the logical buffers, including the terminal offset. Reused sliced validity may + // reference a larger allocation; bytes outside the slice and allocator padding are excluded. + let validity_bytes = if has_nulls { num_rows.div_ceil(8) } else { 0 }; + let output_bytes = total_elements + .checked_mul(element_width) + .and_then(|values| { + num_rows + .checked_add(1)? + .checked_mul(size_of::())? + .checked_add(values) + }) + .and_then(|bytes| bytes.checked_add(validity_bytes)) + .ok_or_else(|| { + DataFusionError::ResourcesExhausted( + "Comet's native `sequence` output byte count exceeds the addressable size. \ + Lower `spark.comet.batchSize` or shorten the sequences." + .to_string(), + ) + })?; + if output_bytes > limit { + return Err(DataFusionError::ResourcesExhausted(format!( + "Comet's native `sequence` invocation requires {output_bytes} output bytes, \ + exceeding its {limit}-byte limit. Lower `spark.comet.batchSize`, shorten the \ + sequences, or configure `spark.comet.exec.sequence.maxBytesPerInvocation` \ + (-1 disables the limit)." + ))); + } + Ok(()) +} + /// Match released Spark 3.4--4.1 Sequence.sequenceLength, including its wide-arithmetic /// fallback error when subtraction overflows even though the resulting length would fit. #[inline(always)] @@ -309,6 +408,7 @@ mod tests { use super::*; use arrow::array::{Int64Array, Int8Array}; use arrow::datatypes::Field; + use datafusion::config::ConfigOptions; fn list_of(elem: DataType) -> DataType { DataType::List(Arc::new(Field::new_list_field(elem, false))) @@ -336,6 +436,213 @@ mod tests { arr.as_any().downcast_ref::().unwrap().clone() } + fn invoke(udf: &SparkSequence, args: Vec, rows: usize) -> Result { + udf.invoke_with_args(ScalarFunctionArgs { + args, + arg_fields: vec![], + number_rows: rows, + return_field: Arc::new(Field::new("result", udf.data_type.clone(), true)), + config_options: Arc::new(ConfigOptions::default()), + }) + } + + #[test] + fn byte_ceiling_counts_each_integral_width_and_terminal_offset() { + macro_rules! check_width { + ($ty:ty) => {{ + let args = vec![ + ColumnarValue::Array(Arc::new(PrimitiveArray::<$ty>::from_iter_values([ + 0, 10, + ]))), + ColumnarValue::Array(Arc::new(PrimitiveArray::<$ty>::from_iter_values([ + 2, 11, + ]))), + ]; + let data_type = list_of(<$ty>::DATA_TYPE); + let bytes = 5 * size_of::<<$ty as ArrowPrimitiveType>::Native>() + 3 * 4; + let output = spark_sequence_with_limit(&args, &data_type, Some(bytes)) + .unwrap() + .into_array(2) + .unwrap(); + let expected = spark_sequence(&args, &data_type) + .unwrap() + .into_array(2) + .unwrap(); + assert_eq!(output.to_data(), expected.to_data()); + let error = + spark_sequence_with_limit(&args, &data_type, Some(bytes - 1)).unwrap_err(); + assert!(matches!(error, DataFusionError::ResourcesExhausted(_))); + assert!(error + .to_string() + .contains(&format!("requires {bytes} output bytes"))); + }}; + } + check_width!(Int8Type); + check_width!(Int16Type); + check_width!(Int32Type); + check_width!(Int64Type); + } + + #[test] + fn byte_ceiling_counts_logical_sliced_validity_and_all_null_outputs() { + let mut values = vec![Some(1i64); 32]; + values[8] = None; + values[15] = None; + let args = vec![ + ColumnarValue::Array(Arc::new(Int64Array::from(values).slice(7, 9))), + ColumnarValue::Scalar(ScalarValue::Int64(Some(1))), + ]; + // Seven values, ten offsets and nine validity bits. The slice starts within a byte + // and shares a larger input allocation; neither changes its logical output size. + let bytes = 7 * 8 + 10 * 4 + 2; + let output = spark_sequence_with_limit(&args, &list_of(DataType::Int64), Some(bytes)) + .unwrap() + .into_array(9) + .unwrap(); + assert_eq!(output.null_count(), 2); + assert!( + spark_sequence_with_limit(&args, &list_of(DataType::Int64), Some(bytes - 1)).is_err() + ); + + let args = vec![ + ColumnarValue::Array(Arc::new(Int64Array::from(vec![None; 9]))), + ColumnarValue::Scalar(ScalarValue::Int64(Some(i64::MAX))), + ColumnarValue::Scalar(ScalarValue::Int64(Some(-1))), + ]; + let output = spark_sequence_with_limit(&args, &list_of(DataType::Int64), Some(42)) + .unwrap() + .into_array(9) + .unwrap(); + assert_eq!(output.null_count(), 9); + assert!(spark_sequence_with_limit(&args, &list_of(DataType::Int64), Some(41)).is_err()); + + let args = vec![ + ColumnarValue::Scalar(ScalarValue::Int64(None)), + ColumnarValue::Scalar(ScalarValue::Int64(Some(1))), + ]; + assert!(matches!( + spark_sequence_with_limit(&args, &list_of(DataType::Int64), Some(9)).unwrap(), + ColumnarValue::Scalar(ScalarValue::List(list)) if list.is_null(0) + )); + assert!(spark_sequence_with_limit(&args, &list_of(DataType::Int64), Some(8)).is_err()); + } + + #[test] + fn empty_output_still_has_an_offset_and_zero_is_a_real_ceiling() { + let args = vec![ + ColumnarValue::Array(Arc::new(Int64Array::from(Vec::::new()))), + ColumnarValue::Scalar(ScalarValue::Int64(Some(1))), + ]; + let output = spark_sequence_with_limit(&args, &list_of(DataType::Int64), Some(4)) + .unwrap() + .into_array(0) + .unwrap(); + assert_eq!(as_primitive_list(&output).value_offsets(), &[0]); + for limit in [0, 3] { + assert!( + spark_sequence_with_limit(&args, &list_of(DataType::Int64), Some(limit)).is_err() + ); + } + } + + #[test] + fn ceiling_refuses_large_outputs_before_allocating_values_and_preserves_errors() { + let scalar = |value| ColumnarValue::Scalar(ScalarValue::Int64(Some(value))); + // This valid single sequence would allocate nearly 16 GiB without the ceiling. + let error = spark_sequence_with_limit( + &[scalar(0), scalar(MAX_ROUNDED_ARRAY_LENGTH - 1)], + &list_of(DataType::Int64), + Some(1024), + ) + .unwrap_err(); + assert!(matches!(error, DataFusionError::ResourcesExhausted(_))); + assert!(error.to_string().contains("1024-byte limit")); + + let starts = ColumnarValue::Array(Arc::new(Int64Array::from(vec![0, 5]))); + let stops = ColumnarValue::Array(Arc::new(Int64Array::from(vec![500, 1]))); + let error = spark_sequence_with_limit( + &[starts, stops, scalar(1)], + &list_of(DataType::Int64), + Some(0), + ) + .unwrap_err() + .to_string(); + assert!( + error.contains("Illegal sequence boundaries: 5 to 1 by 1"), + "{error}" + ); + + let starts = ColumnarValue::Array(Arc::new(Int64Array::from(vec![0, 0]))); + let error = spark_sequence_with_limit( + &[starts, scalar(1_073_741_823)], + &list_of(DataType::Int64), + Some(0), + ) + .unwrap_err() + .to_string(); + assert!(error.contains("2147483648 total elements"), "{error}"); + } + + #[test] + fn output_byte_arithmetic_is_checked_without_allocations() { + for (rows, elements, width, nulls) in [ + (usize::MAX, 0, 1, false), + (usize::MAX / 4, 0, 1, false), + (1, usize::MAX, 8, false), + (1, usize::MAX - 8, 1, true), + ] { + assert!(matches!( + check_output_size(rows, elements, width, nulls, Some(usize::MAX)), + Err(DataFusionError::ResourcesExhausted(_)) + )); + } + } + + #[test] + fn retained_result_does_not_reduce_the_next_invocations_ceiling() { + let udf = SparkSequence::new(list_of(DataType::Int64), Some(32)); + let args = vec![ + ColumnarValue::Array(Arc::new(Int64Array::from(vec![0]))), + ColumnarValue::Scalar(ScalarValue::Int64(Some(2))), + ]; + let first = invoke(&udf, args.clone(), 1).unwrap(); + let second = invoke(&udf, args, 1).unwrap(); + for output in [first, second] { + let output = output.into_array(1).unwrap(); + assert_eq!(row_values(&as_primitive_list(&output), 0), vec![0, 1, 2]); + } + } + + #[test] + fn scalar_invocations_remain_immutable_and_generic_broadcast_is_outside_the_ceiling() { + let data_type = list_of(DataType::Int64); + let udf = SparkSequence::new(data_type.clone(), Some(24)); + assert_eq!(udf.signature().volatility, Volatility::Immutable); + assert_ne!(udf, SparkSequence::new(data_type.clone(), Some(23))); + let distinct = std::collections::HashSet::from([ + SparkSequence::new(data_type.clone(), Some(24)), + SparkSequence::new(data_type.clone(), Some(23)), + SparkSequence::new(data_type.clone(), None), + ]); + assert_eq!(distinct.len(), 3); + let args = vec![ + ColumnarValue::Scalar(ScalarValue::Int64(Some(1))), + ColumnarValue::Scalar(ScalarValue::Int64(Some(2))), + ]; + for rows in [0, 1, 8192] { + let output = invoke(&udf, args.clone(), rows).unwrap(); + assert!(matches!(output, ColumnarValue::Scalar(_))); + } + // Materialising the scalar for many rows is a separate generic operation. + let broadcast = invoke(&udf, args.clone(), 32) + .unwrap() + .into_array(32) + .unwrap(); + assert_eq!(broadcast.len(), 32); + assert_eq!(as_primitive_list(&broadcast).values().len(), 64); + assert!(invoke(&SparkSequence::new(data_type, Some(23)), args, 1).is_err()); + } + fn row_values(list: &ListArray, row: usize) -> Vec { let v = list.value(row); as_primitive_array::(&v) diff --git a/native/spark-expr/src/comet_scalar_funcs.rs b/native/spark-expr/src/comet_scalar_funcs.rs index 36474561055..6031fecc9e6 100644 --- a/native/spark-expr/src/comet_scalar_funcs.rs +++ b/native/spark-expr/src/comet_scalar_funcs.rs @@ -31,13 +31,13 @@ use crate::math_funcs::SparkGreatestLeast; use crate::{ spark_ceil, spark_day_name, spark_decimal_div, spark_decimal_integral_div, spark_floor, spark_isnan, spark_lpad, spark_make_decimal, spark_month_name, spark_read_side_padding, - spark_round, spark_rpad, spark_sequence, spark_to_time, spark_unhex, spark_unscaled_value, - EvalMode, SparkArrayExtrema, SparkArrayPositionFunc, SparkArrayRemove, SparkArraySlice, + spark_round, spark_rpad, spark_to_time, spark_unhex, spark_unscaled_value, EvalMode, + SparkArrayExtrema, SparkArrayPositionFunc, SparkArrayRemove, SparkArraySlice, SparkArraysOverlap, SparkContains, SparkDateDiff, SparkDateFromUnixDate, SparkDateTrunc, SparkDayOfWeek, SparkFlatten, SparkFloatArrayContains, SparkIcebergBucket, SparkIcebergTemporalTransform, SparkIcebergTruncate, SparkMakeDate, SparkMakeInterval, - SparkMakeTime, SparkMapExtract, SparkNextDay, SparkSecondsToTimestamp, SparkSizeFunc, - SparkSortArray, SparkWeekDay, + SparkMakeTime, SparkMapExtract, SparkNextDay, SparkSecondsToTimestamp, SparkSequence, + SparkSizeFunc, SparkSortArray, SparkWeekDay, }; use arrow::datatypes::DataType; use datafusion::common::{DataFusionError, Result as DataFusionResult}; @@ -130,9 +130,9 @@ pub fn create_comet_physical_fun_with_eval_mode( "ceil" => { make_comet_scalar_udf!("ceil", spark_ceil, data_type) } - "spark_sequence" => { - make_comet_scalar_udf!("spark_sequence", spark_sequence, data_type) - } + "spark_sequence" => Ok(Arc::new(ScalarUDF::new_from_impl(SparkSequence::new( + data_type, None, + )))), "floor" => { make_comet_scalar_udf!("floor", spark_floor, data_type) } diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 9bd32816c64..23b65f2ddd2 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -579,6 +579,26 @@ object CometConf extends ShimCometConf { .bytesConf(ByteUnit.BYTE) .createOptional + val COMET_EXEC_SEQUENCE_MAX_BYTES_PER_INVOCATION: OptionalConfigEntry[Long] = + conf(s"$COMET_EXEC_CONFIG_PREFIX.sequence.maxBytesPerInvocation") + .category(CATEGORY_TUNING) + .doc( + "Maximum logical Arrow output bytes of one native integral sequence invocation, " + + "including values, list offsets and validity. When unset in off-heap mode, uses " + + "Comet's configured memory pool limit divided by the configured concurrent task " + + "slots (local or executor cores divided by the task's allocated CPUs, at least one). " + + "Local-cluster uses configured worker cores when executor cores are unset. For " + + "standalone with unknown executor cores, the fallback is the whole pool limit; " + + "set spark.executor.cores or this setting to choose a task ceiling. " + + "On-heap test mode is unbounded unless explicitly configured. Set -1 to disable " + + "the ceiling; 0 permits no output bytes. Earlier results and concurrent calls do " + + "not consume this limit. Scalar invocations produce one list; subsequent scalar " + + "broadcasts and copies of folded literals are outside this ceiling. This does not " + + s"reserve pool memory or bound executor memory usage. $TUNING_GUIDE.") + .bytesConf(ByteUnit.BYTE) + .checkValue(v => v >= -1, "Must be -1 (unbounded) or a non-negative byte count") + .createOptional + val COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED: ConfigEntry[Boolean] = conf(s"$COMET_EXEC_CONFIG_PREFIX.aggregate.skipPartial.enabled") .category(CATEGORY_EXEC) diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index 8997be2d7e9..d8791e5a5ec 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -114,8 +114,9 @@ class CometExecIterator( val conf = SparkEnv.get.conf val localDiskDirs = SparkEnv.get.blockManager.getLocalDiskDirs - // serialize Comet related Spark configs in protobuf format - val protobufSparkConfigs = CometExecIterator.serializeCometSQLConfs() + val memoryConfig = CometExecIterator.getMemoryConfig(conf, taskCPUs) + // Serialize Comet configs with the memory settings resolved for this task. + val protobufSparkConfigs = CometExecIterator.serializeCometSQLConfs(memoryConfig) // Create keyUnwrapper if encryption is enabled val keyUnwrapper = if (encryptedFilePaths.nonEmpty) { @@ -130,8 +131,6 @@ class CometExecIterator( null } - val memoryConfig = CometExecIterator.getMemoryConfig(conf) - val createdPlan = nativeLib.createPlan( id, inputObjects, @@ -861,7 +860,8 @@ object CometExecIterator extends Logging { private def cometSqlConfs: Map[String, String] = SQLConf.get.getAllConfs.filter(_._1.startsWith(CometConf.COMET_PREFIX)) - def serializeCometSQLConfs(): Array[Byte] = { + def serializeCometSQLConfs( + memoryConfig: MemoryConfig = getMemoryConfig(SparkEnv.get.conf)): Array[Byte] = { val builder = ConfigMap.newBuilder() cometSqlConfs.foreach { case (k, v) => if (k.startsWith(s"${CometConf.COMET_PREFIX}.datafusion.")) { @@ -876,6 +876,9 @@ object CometExecIterator extends Logging { // for tokio runtime thread count val executorCores = numDriverOrExecutorCores(SparkEnv.get.conf) builder.putEntries("spark.executor.cores", executorCores.toString) + builder.putEntries( + CometConf.COMET_EXEC_SEQUENCE_MAX_BYTES_PER_INVOCATION.key, + memoryConfig.sequenceMaxBytes.getOrElse(-1L).toString) // Any Comet config that the native side reads must be added here manually, resolved. // `cometSqlConfs` only carries values that were explicitly set, exactly as they were @@ -898,10 +901,14 @@ object CometExecIterator extends Logging { builder.build().toByteArray } - def getMemoryConfig(conf: SparkConf): MemoryConfig = { + def getMemoryConfig(conf: SparkConf): MemoryConfig = + getMemoryConfig(conf, conf.getInt("spark.task.cpus", 1).toLong) + + def getMemoryConfig(conf: SparkConf, taskCPUs: Long): MemoryConfig = { + require(taskCPUs > 0, "Task CPUs must be positive") // there are different paths for on-heap vs off-heap mode val offHeapMode = CometSparkSessionExtensions.isOffHeapEnabled(conf) - if (offHeapMode) { + val memoryConfig = if (offHeapMode) { // in off-heap mode, Comet uses unified memory management to share off-heap memory with Spark val offHeapSize = conf.getSizeAsBytes("spark.memory.offHeap.size") val memoryFraction = CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.get() @@ -920,6 +927,34 @@ object CometExecIterator extends Logging { logDebug("on-heap mode: native memory is unbounded and unaccounted") MemoryConfig(offHeapMode, memoryPoolType = "unbounded", memoryLimit = 0) } + val sequenceMaxBytes = CometConf.COMET_EXEC_SEQUENCE_MAX_BYTES_PER_INVOCATION.get() match { + case Some(-1L) => None + case Some(bytes) => Some(bytes) + case None if offHeapMode => + // This share is stable for the task: live task counts and pool usage do not affect it. + require( + memoryConfig.memoryLimit >= 0, + "The derived sequence memory limit must be non-negative") + Some(sequenceExecutorCores(conf).fold(memoryConfig.memoryLimit) { cores => + val taskSlots = math.max(1L, cores.toLong / taskCPUs) + memoryConfig.memoryLimit / taskSlots + }) + case None => None + } + memoryConfig.copy(sequenceMaxBytes = sequenceMaxBytes) + } + + private def sequenceExecutorCores(conf: SparkConf): Option[Int] = { + val localCluster = """local-cluster\[\s*[0-9]+\s*,\s*([0-9]+)\s*,\s*[0-9]+\s*]""".r + conf.get("spark.master") match { + case localCluster(workerCores) => + conf.getOption("spark.executor.cores").map(_.toInt).orElse(Some(workerCores.toInt)) + case master if master.startsWith("spark://") && !conf.contains("spark.executor.cores") => + // Standalone can assign all available worker cores without exposing that number here. + // Use the whole pool as a fallback, rather than claiming to infer a single task slot. + None + case _ => Some(numDriverOrExecutorCores(conf)) + } } private def numDriverOrExecutorCores(conf: SparkConf): Int = { @@ -950,4 +985,8 @@ object CometExecIterator extends Logging { } } -case class MemoryConfig(offHeapMode: Boolean, memoryPoolType: String, memoryLimit: Long) +case class MemoryConfig( + offHeapMode: Boolean, + memoryPoolType: String, + memoryLimit: Long, + sequenceMaxBytes: Option[Long] = None) diff --git a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala index c464746c8c7..926f3066904 100644 --- a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala @@ -27,6 +27,7 @@ import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.catalyst.expressions.{ArrayAppend, ArrayDistinct, ArrayExcept, ArrayInsert, ArrayIntersect, ArrayJoin, ArrayMax, ArrayMin, ArrayRepeat, ArrayUnion} import org.apache.spark.sql.catalyst.expressions.{ArrayContains, ArrayRemove} import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Cast, CreateArray, ElementAt, Literal, MonotonicallyIncreasingID} +import org.apache.spark.sql.comet.{CometProjectExec, CometSortExec} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf @@ -1699,6 +1700,53 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp } } + test("sequence invocation byte ceiling rejects a batch before materialising its output") { + val key = CometConf.COMET_EXEC_SEQUENCE_MAX_BYTES_PER_INVOCATION.key + withTable("t_seq_bytes") { + sql("CREATE TABLE t_seq_bytes(id BIGINT, a INT, b INT) USING parquet") + sql( + "INSERT INTO t_seq_bytes SELECT id, CAST(id AS INT), CAST(id + 4 AS INT) " + + "FROM range(0, 32, 1, 1)") + val query = "SELECT id, sequence(a, b) AS s FROM t_seq_bytes" + withSQLConf( + key -> "52", + CometConf.COMET_BATCH_SIZE.key -> "4", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val df = sql(query) + assert(df.queryExecution.executedPlan.collect { case p: CometProjectExec => p }.nonEmpty) + val error = intercept[SparkException](df.collect()) + assert(error.getMessage.contains("requires 100 output bytes"), error.getMessage) + assert(error.getMessage.contains("52-byte limit"), error.getMessage) + assert(error.getMessage.contains(key), error.getMessage) + } + // Two five-element integer lists need 40 value bytes and 12 offset bytes. Retaining all + // 16 batches downstream must not reduce any later invocation's 52-byte ceiling. + withSQLConf( + key -> "52", + CometConf.COMET_BATCH_SIZE.key -> "2", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + checkSparkAnswerAndOperator( + sql(s"$query ORDER BY id DESC"), + Seq(classOf[CometProjectExec], classOf[CometSortExec])) + } + withSQLConf(key -> "-1", CometConf.COMET_BATCH_SIZE.key -> "4") { + checkSparkAnswerAndOperator(sql(query), Seq(classOf[CometProjectExec])) + } + } + } + + test("folded sequence literals are outside the native invocation ceiling") { + withSQLConf(CometConf.COMET_EXEC_SEQUENCE_MAX_BYTES_PER_INVOCATION.key -> "0") { + withParquetDataFrame(Seq(Tuple1(1), Tuple1(2), Tuple1(3))) { frame => + // Spark folds this sequence to a literal before Comet sees it. The native projection + // broadcasts that literal; it never invokes the native sequence kernel. + val df = frame.selectExpr("sequence(1, 3) AS s") + assert(!df.queryExecution.optimizedPlan.toString.contains("sequence(")) + checkSparkAnswerAndOperator(df, Seq(classOf[CometProjectExec])) + } + } + } + // 8192 rows of sequence(0, 262143) put 8192 * 262144 = 2^31 elements in one batch, one past // the Int.MaxValue limit of Arrow's i32 list offsets. Spark builds one array per row and // completes; Comet fails the batch and points at spark.comet.batchSize. The lower batch size diff --git a/spark/src/test/scala/org/apache/comet/CometConfSuite.scala b/spark/src/test/scala/org/apache/comet/CometConfSuite.scala index 5c269589310..a985414db12 100644 --- a/spark/src/test/scala/org/apache/comet/CometConfSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometConfSuite.scala @@ -21,10 +21,123 @@ package org.apache.comet import org.scalatest.funsuite.AnyFunSuite +import org.apache.spark.SparkConf import org.apache.spark.sql.internal.SQLConf class CometConfSuite extends AnyFunSuite { + test("sequence invocation ceiling accepts byte units, zero and the explicit opt-out") { + val entry = CometConf.COMET_EXEC_SEQUENCE_MAX_BYTES_PER_INVOCATION + val conf = new SQLConf + assert(entry.get(conf).isEmpty) + Seq("-1" -> -1L, "0" -> 0L, "123" -> 123L, "2k" -> 2048L).foreach { case (input, expected) => + conf.setConfString(entry.key, input) + assert(entry.get(conf).contains(expected)) + } + Seq("-2", "invalid").foreach { input => + conf.setConfString(entry.key, input) + intercept[IllegalArgumentException](entry.get(conf)) + } + } + + test("sequence invocation ceiling uses configured task slots and the resolved pool limit") { + import CometExecIterator.getMemoryConfig + val poolBytes = 2L * 1024 * 1024 * 1024 + val offHeap = new SparkConf(false) + .set("spark.memory.offHeap.enabled", "true") + .set("spark.memory.offHeap.size", "2g") + .set("spark.executor.cores", "10") + .set("spark.task.cpus", "2") + val masters = Seq( + "local" -> 1, + "local[4]" -> 4, + "local[7,3]" -> 7, + "local[*]" -> Runtime.getRuntime.availableProcessors(), + "local[*,3]" -> Runtime.getRuntime.availableProcessors(), + "yarn" -> 10) + val sqlConf = new SQLConf + SQLConf.withExistingConf(sqlConf) { + masters.foreach { case (master, cores) => + val conf = offHeap.clone.set("spark.master", master) + assert( + getMemoryConfig(conf).sequenceMaxBytes.contains(poolBytes / math.max(1, cores / 2))) + } + // The actual allocation overrides spark.task.cpus; ten cores fit three three-CPU tasks. + assert( + getMemoryConfig(offHeap.clone.set("spark.master", "yarn"), taskCPUs = 3L).sequenceMaxBytes + .contains(poolBytes / 3)) + val conf = offHeap.clone.set("spark.master", "local[8]") + sqlConf.setConfString(CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key, "0.5") + Seq("fair_unified", "greedy_unified").foreach { poolType => + sqlConf.setConfString(CometConf.COMET_OFFHEAP_MEMORY_POOL_TYPE.key, poolType) + assert(getMemoryConfig(conf).sequenceMaxBytes.contains(poolBytes / 2 / 4)) + } + sqlConf.setConfString(CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key, "0") + assert(getMemoryConfig(conf).sequenceMaxBytes.contains(0L)) + sqlConf.setConfString(CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key, "-1") + intercept[IllegalArgumentException](getMemoryConfig(conf)) + } + } + + test( + "sequence invocation ceiling distinguishes configured cores from unknown standalone cores") { + import CometExecIterator.getMemoryConfig + val poolBytes = 2L * 1024 * 1024 * 1024 + val offHeap = new SparkConf(false) + .set("spark.memory.offHeap.enabled", "true") + .set("spark.memory.offHeap.size", "2g") + SQLConf.withExistingConf(new SQLConf) { + val localCluster = offHeap.clone.set("spark.master", "local-cluster[2,4,1024]") + assert(getMemoryConfig(localCluster).sequenceMaxBytes.contains(poolBytes / 4)) + assert( + getMemoryConfig( + offHeap.clone.set("spark.master", "local-cluster[ 2, 4, 1024 ]")).sequenceMaxBytes + .contains(poolBytes / 4)) + assert( + getMemoryConfig(localCluster.clone.set("spark.executor.cores", "2")).sequenceMaxBytes + .contains(poolBytes / 2)) + + val standalone = offHeap.clone.set("spark.master", "spark://worker:7077") + // No supported executor-local API exposes the assigned cores in this case. This is an + // explicit whole-pool fallback, independent of host processors or the task's CPU count. + assert(getMemoryConfig(standalone).sequenceMaxBytes.contains(poolBytes)) + assert(getMemoryConfig(standalone, taskCPUs = 3L).sequenceMaxBytes.contains(poolBytes)) + assert( + getMemoryConfig( + standalone.clone.set("spark.executor.cores", "8"), + taskCPUs = 2L).sequenceMaxBytes.contains(poolBytes / 4)) + + // YARN and Kubernetes retain their one-core default when executor cores are unset. + Seq("yarn", "k8s://https://kubernetes.default.svc:443").foreach { master => + assert( + getMemoryConfig(offHeap.clone.set("spark.master", master)).sequenceMaxBytes + .contains(poolBytes)) + } + } + } + + test("sequence invocation overrides also apply to otherwise unbounded on-heap test mode") { + import CometExecIterator.getMemoryConfig + val key = CometConf.COMET_EXEC_SEQUENCE_MAX_BYTES_PER_INVOCATION.key + val conf = new SparkConf(false) + .set("spark.master", "local[4]") + .set("spark.memory.offHeap.enabled", "false") + val sqlConf = new SQLConf + SQLConf.withExistingConf(sqlConf) { + assert(getMemoryConfig(conf).sequenceMaxBytes.isEmpty) + Seq("false", "true").foreach { offHeap => + val mode = conf.clone + .set("spark.memory.offHeap.enabled", offHeap) + .set("spark.memory.offHeap.size", "2g") + Seq("-1" -> None, "0" -> Some(0L), "2k" -> Some(2048L)).foreach { + case (value, expected) => + sqlConf.setConfString(key, value) + assert(getMemoryConfig(mode).sequenceMaxBytes == expected) + } + } + } + } + test("primary key wins over alternative when both are set") { val entry = CometConf .conf("spark.comet.testing.alias.primaryWins") diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index f40fd42d1bd..e2be2596585 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -118,6 +118,24 @@ class CometExecSuite extends CometTestBase { } } + test("SQLConf serde sends the resolved per-invocation sequence byte ceiling") { + val key = CometConf.COMET_EXEC_SEQUENCE_MAX_BYTES_PER_INVOCATION.key + val conf = new org.apache.spark.SparkConf(false) + .set("spark.master", "local[8]") + .set("spark.memory.offHeap.enabled", "true") + .set("spark.memory.offHeap.size", "2g") + def resolved: String = { + val memory = CometExecIterator.getMemoryConfig(conf, taskCPUs = 2L) + ConfigMap.parseFrom(CometExecIterator.serializeCometSQLConfs(memory)).getEntriesMap.get(key) + } + assert(resolved == "536870912") + Seq("2k" -> "2048", "0" -> "0", "-1" -> "-1").foreach { case (input, expected) => + withSQLConf(key -> input) { + assert(resolved == expected) + } + } + } + test("native sort spill files are compressed with the default spill codec") { val numRows = 20000 val compressibleValue = "native-sort-spill-compression-" * 8