From 9e77f9f9bb5b9ad45db24a18f976f7bd50cc54ec Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sat, 19 Sep 2026 04:25:42 +0000 Subject: [PATCH 1/5] perf: reuse prepared broadcast builds across executor tasks --- dev/ensure-jars-have-correct-contents.sh | 2 + .../user-guide/latest/tuning/operators.md | 36 + .../core/src/execution/operators/broadcast.rs | 1216 +++++++++++++++++ .../execution/operators/broadcast_cache.rs | 437 ++++++ .../operators/dynamic_filter/join.rs | 18 +- native/core/src/execution/operators/mod.rs | 3 + native/core/src/execution/operators/scan.rs | 2 +- native/core/src/execution/planner.rs | 20 +- native/core/src/lib.rs | 14 + .../java/org/apache/comet/NativeBase.java | 22 +- .../spark/CometBroadcastMemoryManager.java | 175 +++ .../scala/org/apache/comet/CometConf.scala | 24 + .../main/scala/org/apache/spark/Plugins.scala | 10 +- .../comet/CometBroadcastExchangeExec.scala | 11 +- .../spark/sql/comet/CometBroadcastInput.scala | 169 +++ .../sql/comet/CometBroadcastInputRDD.scala | 62 + .../apache/spark/sql/comet/CometExecRDD.scala | 18 +- .../arrow/CometBroadcastArrowStream.scala | 113 ++ .../apache/spark/sql/comet/operators.scala | 46 +- .../apache/comet/exec/CometJoinSuite.scala | 188 ++- .../spark/CometTaskMemoryManagerSuite.scala | 151 +- .../arrow/CometArrowStreamSuite.scala | 56 + 22 files changed, 2763 insertions(+), 30 deletions(-) create mode 100644 native/core/src/execution/operators/broadcast.rs create mode 100644 native/core/src/execution/operators/broadcast_cache.rs create mode 100644 spark/src/main/java/org/apache/spark/CometBroadcastMemoryManager.java create mode 100644 spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastInput.scala create mode 100644 spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastInputRDD.scala create mode 100644 spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometBroadcastArrowStream.scala diff --git a/dev/ensure-jars-have-correct-contents.sh b/dev/ensure-jars-have-correct-contents.sh index 9453539382a..11015821cb3 100755 --- a/dev/ensure-jars-have-correct-contents.sh +++ b/dev/ensure-jars-have-correct-contents.sh @@ -102,6 +102,8 @@ allowed_expr+="|^org/apache/spark/CometPlugin.class$" allowed_expr+="|^org/apache/spark/CometDriverPlugin.*$" allowed_expr+="|^org/apache/spark/CometExecutorPlugin.*$" allowed_expr+="|^org/apache/spark/CometSource.*$" +# Native broadcast storage accounting needs Spark-private memory APIs. +allowed_expr+='|^org/apache/spark/CometBroadcastMemoryManager\.class$' allowed_expr+="|^org/apache/spark/CometTaskMemoryManager.class$" allowed_expr+="|^org/apache/spark/CometTaskMemoryManager.*$" allowed_expr+="|^scala-collection-compat.properties$" diff --git a/docs/source/user-guide/latest/tuning/operators.md b/docs/source/user-guide/latest/tuning/operators.md index 644f35483ba..e00cbedccf2 100644 --- a/docs/source/user-guide/latest/tuning/operators.md +++ b/docs/source/user-guide/latest/tuning/operators.md @@ -163,3 +163,39 @@ native path with `spark.comet.expression.SortOrder.allowIncompatible=true`. `sort_array` sorts array elements rather than ordering rows. It follows Spark's floating-point ordering as well, so it also stays native with `spark.comet.exec.strictFloatingPoint=true`. + +## Reusing Broadcast Hash Builds (Experimental) + +`spark.comet.broadcast.reuse.enabled=true` allows tasks in one executor to share a prepared +native hash table for the same Spark broadcast. This avoids repeating Arrow decoding and hash-table +construction on a cache hit. It requires `org.apache.spark.CometPlugin` in `spark.plugins` and Spark +off-heap memory. The feature is disabled by default. + +The implementation covers non-spilling inner broadcast joins with direct column keys, +matching key types, and fixed-width or plain UTF8 build columns. Residual join conditions remain +local to each consuming task. Dictionary, view, binary, and nested build columns remain unsupported. +Unsupported joins keep ordinary task-local execution. Lookup is separated by executor generation +and uses the actual broadcast ID, build and probe schemas, and ordered keys. Cached builds omit +unused IN-list membership data; aliases of the same exchange can share a build. + +`spark.comet.broadcast.reuse.maxMemory` defaults to `1g`. It caps native prepared-build allocations +through Spark off-heap storage memory while tasks are preparing or using a build. When the last +task releases that build, its storage charge is returned. A later task wave may prepare the same +broadcast again; a later query without a broadcast hash join inherits no idle build charge. +The first flat-schema broadcast input marked for reuse fixes the executor's cap, even if its join +later proves ineligible and creates no build. A later input with a different cap uses ordinary +execution. Cache admission never waits for another join to release its memory. When admission +fails, Comet discards the partial preparation and opens a fresh uncached input for that task. +Preparation must also admit the temporary overlap between decoded native batches and DataFusion's +compact build batch; the memory needed to prepare a build can exceed its final retained size. + +This cap does not bound the existing driver broadcast representation or JVM decoder allocations. +The first task still decodes the existing broadcast format; bounded broadcast construction and +admission before JVM decoding are separate work. + +Spark join metrics expose cache hits, successful preparations, admission fallbacks, preparation +rows/bytes, and preparation time. Probe metrics remain per task. A successful preparation should +normally be followed by cache hits from other tasks using that broadcast. Validate elapsed query +time and executor memory with the feature both enabled and disabled before enabling it broadly; +reuse counts alone do not establish a query speedup. + diff --git a/native/core/src/execution/operators/broadcast.rs b/native/core/src/execution/operators/broadcast.rs new file mode 100644 index 00000000000..a491a6453b2 --- /dev/null +++ b/native/core/src/execution/operators/broadcast.rs @@ -0,0 +1,1216 @@ +// 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. + +//! Prepares a broadcast hash table once for tasks probing the same build at the +//! same time. `BroadcastInputExec` captures the Spark broadcast's identity, +//! schema and storage owner without decoding rows during planning. On a miss, +//! one task opens its JVM stream, copies admitted Arrow batches into native +//! buffers and asks DataFusion to prepare an immutable hash build. Other tasks +//! with a compatible join can use that build with their own probe and metrics. +//! +//! The lookup retains only weak references. Active probe streams keep the +//! prepared build alive. The last lease frees it and returns its Spark storage +//! charge while the owner is active; executor retirement can return the charge +//! earlier without revoking probes. Rejected admission or an ineligible join +//! uses an ordinary task-local hash join with a fresh broadcast stream. + +use std::fmt::{Debug, Display, Formatter}; +use std::mem::size_of; +use std::pin::Pin; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use arrow::array::{ + make_array, Array, ArrayRef, Capacities, MutableArrayData, RecordBatch, RecordBatchOptions, + StringArray, +}; +use arrow::compute::{cast_with_options, CastOptions}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::ffi_stream::{ArrowArrayStreamReader, FFI_ArrowArrayStream}; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::{ + internal_err, resources_datafusion_err, DataFusionError, JoinType, Result, +}; +use datafusion::execution::memory_pool::{ + MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation, +}; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; +use datafusion::physical_plan::execution_plan::{ + Boundedness, ChildrenPropertiesMode, EmissionType, ReplaceChildrenOptions, +}; +use datafusion::physical_plan::joins::PreparedHashJoinBuild; +use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode}; +use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricBuilder, MetricsSet}; +use datafusion::physical_plan::projection::ProjectionExec; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, RecordBatchStream, + SendableRecordBatchStream, +}; +use futures::{Stream, TryStreamExt}; +use jni::objects::{Global, JObject, JValue}; +use once_cell::sync::Lazy; +use parking_lot::Mutex; + +use super::broadcast_cache::{BuildCache, BuildKey}; +use super::scan::import_column; + +use crate::errors::{CometError, CometResult}; +use crate::jvm_bridge::JVMClasses; + +type PreparedCache = BuildCache; +type CacheGeneration = Option<(i64, Arc)>; + +// DataFusion charges each retained batch before polling for the next one. +// Bound the extra charge held by the producer during that handoff. +const MAX_COPY_BATCH_BYTES: usize = 32 * 1024 * 1024; + +// Lookup crosses Spark task/plan boundaries while builds are active. This +// executor-lifetime registry holds only weak references, capped at 64 keys. +// The registry fences SparkEnv generations; keys use actual Broadcast.id, not +// logical plan identity. Plugin shutdown retires the cache before SparkEnv. +static CACHE: Lazy> = Lazy::new(|| Mutex::new(None)); + +/// Returns this runtime's bounded cache, creating metadata only. Existing +/// probes retain their prepared builds after shutdown; a new runtime gets a new cache. +fn cache(generation: i64) -> Arc { + let (current, retired) = { + let mut cache = CACHE.lock(); + if let Some((existing, current)) = &*cache { + if *existing == generation { + return Arc::clone(current); + } + if *existing > generation { + // An old task racing environment replacement must never retire + // the newer environment or publish into its cache. + return Arc::new(BuildCache::new(0)); + } + } + let current = Arc::new(BuildCache::new(64)); + let retired = cache.replace((generation, Arc::clone(¤t))); + (current, retired) + }; + if let Some((_, retired)) = retired { + retired.clear(); + } + current +} + +/// Retires all native broadcast lookups before Spark releases storage accounting. +/// Prepared values still leased by outputs/tasks retain their original owner. +pub(crate) fn clear_broadcast_cache() { + let old = CACHE.lock().take(); + if let Some((_, old)) = old { + old.clear(); + } +} + +/// Charges temporary broadcast copies and the prepared hash table to Spark +/// executor storage instead of a task's memory manager. The captured JVM +/// owner enforces one cap across concurrent builds and live leases in this +/// SparkEnv generation, even when their lookup entries have been retired. +#[derive(Debug)] +struct BroadcastMemoryPool { + owner: Arc>>, + limit: usize, + used: AtomicUsize, +} + +impl BroadcastMemoryPool { + /// Requests an all-or-nothing storage grant without waiting for active probes. + fn acquire(&self, bytes: usize) -> CometResult { + JVMClasses::with_env(|env| { + Ok(env + .call_method( + self.owner.as_obj(), + jni::jni_str!("acquireMemory"), + jni::jni_sig!("(J)J"), + &[JValue::Long(bytes as i64)], + )? + .j()?) + }) + } + + /// Returns an existing grant to its captured environment. A retired JVM + /// owner accounts late releases locally without touching a replacement env. + fn release(&self, bytes: usize) -> CometResult<()> { + JVMClasses::with_env(|env| { + env.call_method( + self.owner.as_obj(), + jni::jni_str!("releaseMemory"), + jni::jni_sig!("(J)V"), + &[JValue::Long(bytes as i64)], + )?; + Ok(()) + }) + } +} + +impl Display for BroadcastMemoryPool { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + Debug::fmt(self, f) + } +} + +/// DataFusion reservations acquire storage in full or fail preparation. Shrinking +/// a reservation releases its grant through the captured SparkEnv owner. +impl MemoryPool for BroadcastMemoryPool { + fn name(&self) -> &str { + "CometBroadcastMemoryPool" + } + + fn grow(&self, reservation: &MemoryReservation, bytes: usize) { + self.try_grow(reservation, bytes) + .expect("broadcast allocations must be admitted"); + } + + fn shrink(&self, _: &MemoryReservation, bytes: usize) { + self.release(bytes) + .expect("broadcast storage grant release failed"); + self.used.fetch_sub(bytes, Ordering::Relaxed); + } + + fn try_grow(&self, _: &MemoryReservation, bytes: usize) -> Result<()> { + if bytes == 0 { + return Ok(()); + } + // Reject a build allocation above its byte cap before asking Spark. + // Spark's owner still checks the total across concurrent loader pools. + if bytes > self.limit.saturating_sub(self.reserved()) { + return Err(DataFusionError::ResourcesExhausted( + "Broadcast build allocation exceeds its byte cap".into(), + )); + } + if self.acquire(bytes)? == 0 { + return Err(resources_datafusion_err!( + "Broadcast cache storage admission denied" + )); + } + self.used.fetch_add(bytes, Ordering::Relaxed); + Ok(()) + } + + fn reserved(&self) -> usize { + self.used.load(Ordering::Relaxed) + } + fn memory_limit(&self) -> MemoryLimit { + MemoryLimit::Finite(self.limit) + } +} + +/// Task-owned broadcast input whose declared schema permits planning without +/// decoding Arrow. Each `open` creates a new stream, so a task can replay its +/// broadcast after a refused cache build. Prepared cache values never retain +/// the JVM input or its decoder. +#[derive(Clone, Debug)] +pub(crate) struct BroadcastInputExec { + input: Arc>>, + owner: Arc>>, + generation: i64, + broadcast_id: i64, + limit: usize, + properties: Arc, +} + +impl BroadcastInputExec { + /// Recognizes a JVM carrier and captures immutable metadata without opening + /// its stream. Ordinary Arrow inputs return `None`. Invalid carrier metadata + /// or JNI failure returns an error before ownership of any stream transfers. + pub fn try_new( + input: Arc>>, + types: &[DataType], + ) -> CometResult> { + let metadata = JVMClasses::with_env(|env| { + if !env.is_instance_of( + input.as_obj(), + jni::strings::JNIString::new("org/apache/spark/sql/comet/CometBroadcastInput"), + )? { + return Ok::<_, CometError>(None); + } + let broadcast_id = env + .call_method( + input.as_obj(), + jni::jni_str!("getBroadcastId"), + jni::jni_sig!("()J"), + &[], + )? + .j()?; + let owner = env + .call_method( + input.as_obj(), + jni::jni_str!("getMemoryManager"), + jni::jni_sig!("()Lorg/apache/spark/CometBroadcastMemoryManager;"), + &[], + )? + .l()?; + let generation = env + .call_method( + &owner, + jni::jni_str!("getGeneration"), + jni::jni_sig!("()J"), + &[], + )? + .j()?; + let limit = env + .call_method(&owner, jni::jni_str!("getLimit"), jni::jni_sig!("()J"), &[])? + .j()?; + let owner = Arc::new(jni_new_global_ref!(env, owner)?); + Ok(Some((owner, generation, broadcast_id, limit))) + })?; + let Some((owner, generation, broadcast_id, limit)) = metadata else { + return Ok(None); + }; + let limit = usize::try_from(limit) + .map_err(|_| CometError::Config("Invalid broadcast cache cap".into()))?; + let schema = Arc::new(Schema::new( + types + .iter() + .enumerate() + .map(|(i, dt)| Field::new(format!("col_{i}"), dt.clone(), true)) + .collect::>(), + )); + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Ok(Some(Self { + input, + owner, + generation, + broadcast_id, + limit, + properties, + })) + } + + /// Opens a fresh JVM decoder and transfers its C stream into this task's + /// native reader. Java retains only the emptied wrapper for completion cleanup. + fn open(&self) -> CometResult { + JVMClasses::with_env(|env| { + let stream = env + .call_method( + self.input.as_obj(), + jni::jni_str!("openStream"), + jni::jni_sig!("()Lorg/apache/arrow/c/ArrowArrayStream;"), + &[], + )? + .l()?; + let address = + unsafe { jni_call!(env, arrow_array_stream(&stream).memory_address() -> i64) }?; + // SAFETY: openStream returns a fresh owned C stream. from_raw moves + // its release callback out, leaving the Java wrapper empty. + unsafe { ArrowArrayStreamReader::from_raw(address as *mut FFI_ArrowArrayStream) } + .map_err(CometError::from) + }) + } + + /// Produces a lazy stream. With `pool`, every payload is deeply copied into + /// admitted native buffers; without it, this is ordinary task-local decoding. + fn stream(self: &Arc, pool: Option>) -> SendableRecordBatchStream { + let copy_reservation = pool + .as_ref() + .map(|pool| MemoryConsumer::new("Broadcast native copy").register(pool)); + Box::pin(BroadcastStream { + source: Arc::clone(self), + reader: None, + pending_batch: None, + copy_reservation, + }) + } +} + +impl DisplayAs for BroadcastInputExec { + fn fmt_as(&self, _: DisplayFormatType, f: &mut Formatter<'_>) -> std::fmt::Result { + write!( + f, + "CometBroadcastInputExec: broadcast={}", + self.broadcast_id + ) + } +} + +impl ExecutionPlan for BroadcastInputExec { + fn name(&self) -> &str { + "CometBroadcastInputExec" + } + fn properties(&self) -> &Arc { + &self.properties + } + fn children(&self) -> Vec<&Arc> { + vec![] + } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + if !children.is_empty() { + return internal_err!("Broadcast input has no children"); + } + Ok(self) + } + fn execute(&self, partition: usize, _: Arc) -> Result { + if partition != 0 { + return internal_err!("Broadcast input has one partition"); + } + Ok(Arc::new(self.clone()).stream(None)) + } +} + +/// Task-local lazy decoder stream. During preparation, its temporary copy +/// grant overlaps DataFusion's reservation of the previous batch until the +/// next poll. Ordinary joins read the same source without a copy grant. +/// Neither the stream nor its JVM reader is cached. +struct BroadcastStream { + source: Arc, + reader: Option, + pending_batch: Option, + copy_reservation: Option, +} + +impl BroadcastStream { + /// Releases the last temporary copy grant and reads or continues splitting + /// one batch. Cache preparation admits a conservative bound before copying + /// supported physical buffers. An ineligible encoding returns a resource + /// error so the caller can reopen the broadcast for ordinary execution. + fn next_batch(&mut self) -> Result> { + if let Some(reservation) = &mut self.copy_reservation { + reservation.free(); + } + if self.reader.is_none() { + self.reader = Some(self.source.open()?); + } + let already_split = self.pending_batch.is_some(); + let batch = match self.pending_batch.take() { + Some(batch) => batch, + None => { + let Some(batch) = self.reader.as_mut().unwrap().next() else { + return Ok(None); + }; + batch? + } + }; + let schema = self.source.schema(); + if batch.num_columns() != schema.fields().len() { + return internal_err!("Broadcast column count mismatch"); + } + let rows = if self.copy_reservation.is_some() { + copy_batch_rows(&batch, &schema, MAX_COPY_BATCH_BYTES, already_split)? + } else { + batch.num_rows() + }; + let batch = if rows < batch.num_rows() { + self.pending_batch = Some(batch.slice(rows, batch.num_rows() - rows)); + batch.slice(0, rows) + } else { + batch + }; + let columns = if let Some(reservation) = &mut self.copy_reservation { + copy_broadcast_columns(&batch, &schema, reservation)? + } else { + cast_uncached_broadcast_columns(&batch, &schema)? + }; + Ok(Some(RecordBatch::try_new_with_options( + schema, + columns, + &RecordBatchOptions::new().with_row_count(Some(batch.num_rows())), + )?)) + } +} + +/// Replay through ScanExec's UTF-8 decoding and dictionary unpacking before +/// ScanStream's safe casts. Cache-ineligible input still needs the ordinary +/// FFI validation, and its physical encoding may differ from the declared schema. +fn cast_uncached_broadcast_columns(batch: &RecordBatch, schema: &Schema) -> Result> { + let options = CastOptions::default(); + batch + .columns() + .iter() + .zip(schema.fields()) + .map(|(array, field)| { + let array = import_column(array)?; + if array.data_type() == field.data_type() { + Ok(array) + } else { + Ok(cast_with_options(&array, field.data_type(), &options)?) + } + }) + .collect() +} + +impl Stream for BroadcastStream { + type Item = Result; + fn poll_next(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + Poll::Ready(self.next_batch().transpose()) + } +} + +impl RecordBatchStream for BroadcastStream { + fn schema(&self) -> SchemaRef { + self.source.schema() + } +} + +/// Active Utf8 bytes include bytes behind null rows, but exclude a slice's +/// unrelated prefix/suffix. Read offsets only, before allocating the copy. +fn utf8_value_span(array: &StringArray) -> usize { + let offsets = array.value_offsets(); + (offsets[offsets.len() - 1] - offsets[0]) as usize +} + +/// Bound one explicitly sized Utf8 copy, including offsets, optional validity, +/// up to 64-byte padding per buffer, and the existing per-array metadata allowance. +fn utf8_copy_bytes(rows: usize, values: usize) -> Result { + rows.checked_add(1) + .and_then(|rows| rows.checked_mul(size_of::())) + .and_then(|bytes| bytes.checked_add(values)) + .and_then(|bytes| bytes.checked_add(rows.div_ceil(8))) + .and_then(|bytes| bytes.checked_add(3 * 64 + 1024)) + .ok_or_else(|| resources_datafusion_err!("Broadcast copy size overflow")) +} + +/// Bound copied bytes, not source capacity: slices share their original +/// buffers but own only their rows after copying. +fn broadcast_copy_bytes(batch: &RecordBatch, schema: &Schema) -> Result { + let mut bound = 0_usize; + for (column, field) in batch.columns().iter().zip(schema.fields()) { + if column.data_type() != field.data_type() { + return Err(resources_datafusion_err!( + "Broadcast physical encoding cannot be cached" + )); + } + let bytes = if column.data_type() == &DataType::Utf8 { + let array = column + .as_any() + .downcast_ref::() + .expect("Utf8 array"); + utf8_copy_bytes(array.len(), utf8_value_span(array))? + } else { + let rows = column.len(); + let values = match column.data_type() { + DataType::Boolean => Some(rows.div_ceil(8)), + DataType::Null => Some(0), + DataType::FixedSizeBinary(width) => usize::try_from(*width) + .ok() + .and_then(|width| rows.checked_mul(width)), + ty => ty + .primitive_width() + .and_then(|width| rows.checked_mul(width)), + }; + values + .and_then(|values| values.checked_add(rows.div_ceil(8))) + .and_then(|bytes| bytes.checked_add(3 * 64 + 1024)) + .ok_or_else(|| resources_datafusion_err!("Broadcast copy size overflow"))? + }; + bound = bound + .checked_add(bytes) + .ok_or_else(|| resources_datafusion_err!("Broadcast copy size overflow"))?; + } + Ok(bound) +} + +/// Largest prefix whose independent native copy fits one handoff. +/// Keep small batches whole to avoid extra concatenation work. Once a large +/// batch is split, continue until all its rows are copied. DataFusion separately +/// admits the compact output while the copied input batches are still live. +fn copy_batch_rows( + batch: &RecordBatch, + schema: &Schema, + target: usize, + already_split: bool, +) -> Result { + if batch.num_rows() == 0 { + return Ok(0); + } + let bytes = broadcast_copy_bytes(batch, schema)?; + if bytes <= target || (!already_split && bytes <= 2 * target) { + return Ok(batch.num_rows()); + } + let mut fits = 1; + let mut too_large = batch.num_rows(); + while fits + 1 < too_large { + let middle = fits + (too_large - fits) / 2; + if broadcast_copy_bytes(&batch.slice(0, middle), schema)? <= target { + fits = middle; + } else { + too_large = middle; + } + } + Ok(fits) +} + +/// Validate every column and admit all copies before allocating any of them. +/// Count aliases separately: each output column owns an independent copy. +fn copy_broadcast_columns( + batch: &RecordBatch, + schema: &Schema, + reservation: &mut MemoryReservation, +) -> Result> { + // FFI imports do not validate UTF-8. Check raw ArrayData before copying or + // constructing any string values, without allocating a decoded payload. + // Invalid Spark strings use ordinary replay, whose decoder preserves Spark + // semantics without charging an unbounded repair allocation to this cache. + for column in batch.columns() { + if column.data_type() == &DataType::Utf8 { + column.to_data().validate_data().map_err(|_| { + resources_datafusion_err!("Broadcast UTF-8 requires ordinary decoding") + })?; + } + } + let bound = broadcast_copy_bytes(batch, schema)?; + reservation.try_grow(bound)?; + batch + .columns() + .iter() + .map(|column| { + if let Some(array) = column.as_any().downcast_ref::() { + let data = array.to_data(); + // Generic `new(rows)` reserves `rows` value bytes and may double + // that capacity on growth. Explicit capacities instead match the + // admitted span, even for empty strings and sliced null values. + let mut copy = MutableArrayData::with_capacities( + vec![&data], + false, + Capacities::Binary(array.len(), Some(utf8_value_span(array))), + ); + copy.try_extend(0, 0, array.len())?; + Ok(make_array(copy.freeze())) + } else { + let data = column.to_data(); + let mut copy = MutableArrayData::new(vec![&data], false, column.len()); + copy.try_extend(0, 0, column.len())?; + Ok(make_array(copy.freeze())) + } + }) + .collect() +} + +/// Restricts retained payloads to arrays whose gather output owns its values. +fn cacheable_type(dt: &DataType) -> bool { + dt.is_primitive() + || matches!( + dt, + DataType::Boolean | DataType::Null | DataType::FixedSizeBinary(_) | DataType::Utf8 + ) +} + +/// Wraps an eligible inner hash join without caching its execution plan. Only +/// the immutable prepared build is shared; each task has its own probe child, +/// dynamic filter and execution metrics. +#[derive(Debug)] +struct CachedBroadcastJoinExec { + join: Arc, + source: Arc, + cache: Arc, + metrics: ExecutionPlanMetricsSet, +} + +impl DisplayAs for CachedBroadcastJoinExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter<'_>) -> std::fmt::Result { + self.join.fmt_as(t, f) + } +} + +impl ExecutionPlan for CachedBroadcastJoinExec { + fn name(&self) -> &str { + "CometCachedBroadcastHashJoinExec" + } + fn properties(&self) -> &Arc { + self.join.properties() + } + fn children(&self) -> Vec<&Arc> { + self.join.children() + } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + self.join.apply_expressions(f) + } + fn replace_children( + self: Arc, + children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + // A rewrite must re-establish exact build identity and compatibility. + Arc::clone(&self.join).replace_children(children, options) + } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn metrics(&self) -> Option { + Some(self.metrics.clone_inner()) + } + /// Defers lookup until this probe stream is polled. One caller prepares a + /// missing build while matching tasks wait or reuse an active build. Failed + /// admission uses a fresh stream over this task's broadcast for the ordinary + /// join; successful probes hold a lease until their output stream is dropped. + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + if partition != 0 { + return internal_err!("Cached broadcast join has one probe partition"); + } + let join = Arc::clone(&self.join); + let source = Arc::clone(&self.source); + let cache = Arc::clone(&self.cache); + let metrics = self.metrics.clone(); + let hits = + MetricBuilder::new(&self.metrics).counter("broadcast_build_cache_hits", partition); + let misses = + MetricBuilder::new(&self.metrics).counter("broadcast_build_cache_misses", partition); + let fallbacks = + MetricBuilder::new(&self.metrics).counter("broadcast_build_cache_fallbacks", partition); + let prepare_time = MetricBuilder::new(&self.metrics) + .subset_time("broadcast_build_prepare_time", partition); + let prepare_rows = + MetricBuilder::new(&self.metrics).counter("broadcast_build_prepare_rows", partition); + let prepare_bytes = + MetricBuilder::new(&self.metrics).gauge("broadcast_build_prepare_bytes", partition); + let key = BuildKey { + broadcast_id: source.broadcast_id, + schema: source.schema(), + probe_schema: join.right().schema(), + key_columns: join + .on() + .iter() + .map(|(left, _)| { + left.downcast_ref::() + .expect("eligible direct key") + .index() + }) + .collect(), + }; + let schema = self.schema(); + let stream = futures::stream::once(async move { + let pool: Arc = Arc::new(BroadcastMemoryPool { + owner: Arc::clone(&source.owner), + limit: source.limit, + used: AtomicUsize::new(0), + }); + let result = cache + .get_or_load(key, || async { + let _timer = prepare_time.timer(); + // Only runtime filter joins use membership, and they already + // disable IN-list construction. Direct joins need no list. + let mut config = context.session_config().options().as_ref().clone(); + config.optimizer.hash_join_inlist_pushdown_max_size = 0; + config + .optimizer + .hash_join_inlist_pushdown_max_distinct_values = 0; + let prepared = join + .prepare_build( + source.stream(Some(Arc::clone(&pool))), + pool, + Arc::new(config), + ) + .await?; + prepare_rows.add(prepared.num_rows()); + prepare_bytes.set(prepared.reserved_bytes()); + Ok(prepared) + }) + .await; + let (execution, lease) = match result { + Ok((prepared, hit)) => { + if hit { + hits.add(1); + } else { + misses.add(1); + } + let execution = Arc::new( + join.builder() + .with_prepared_build(Arc::clone(&prepared))? + .build()?, + ); + (execution, Some(prepared)) + } + Err(DataFusionError::ResourcesExhausted(_)) => { + // The source is replayable. Cache failure cannot consume this + // task's ordinary build input or wait on its other leases. + fallbacks.add(1); + (Arc::clone(&join), None) + } + Err(error) => return Err(error), + }; + let result = execution.execute(partition, context); + for metric in execution.metrics().unwrap_or_default().iter() { + metrics.register(Arc::clone(metric)); + } + let stream = result?; + match lease { + Some(prepared) => Ok(Box::pin(PreparedProbeStream { + stream, + _lease: prepared, + }) as SendableRecordBatchStream), + None => Ok(stream), + } + }) + .try_flatten(); + Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) + } +} + +/// Keeps the prepared build alive through the probe, even after its plan is +/// dropped. The final lease frees the table; Spark's owner returns its storage +/// charge then, or earlier on executor retirement. +struct PreparedProbeStream { + stream: SendableRecordBatchStream, + _lease: Arc, +} + +impl Stream for PreparedProbeStream { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { + self.stream.as_mut().poll_next(context) + } +} + +impl RecordBatchStream for PreparedProbeStream { + fn schema(&self) -> SchemaRef { + self.stream.schema() + } +} + +/// Wraps only a final, physical inner broadcast join with direct column keys and +/// flat payloads. Call after build-side swapping. A dynamic-filter wrapper must +/// instead call the shared-metrics variant on its fresh runtime join. Unsupported +/// joins retain the ordinary execution path. +pub(crate) fn reuse_broadcast_build( + plan: Arc, +) -> Result> { + reuse_broadcast_build_with_metrics(plan, ExecutionPlanMetricsSet::new()) +} + +/// Attaches reuse while registering asynchronous execution counters in the +/// caller's metric set. The set owns no execution state. An ineligible direct +/// join returns the same Arc so the caller can forward ordinary metrics itself. +pub(crate) fn reuse_broadcast_build_with_metrics( + plan: Arc, + metrics: ExecutionPlanMetricsSet, +) -> Result> { + if let Some(projection) = plan.downcast_ref::() { + let child = reuse_broadcast_build_with_metrics(Arc::clone(projection.input()), metrics)?; + return if Arc::ptr_eq(&child, projection.input()) { + Ok(plan) + } else { + plan.replace_children( + vec![child], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + }; + } + let Some(join) = plan.downcast_ref::() else { + return Ok(plan); + }; + let Some(source) = join.left().downcast_ref::() else { + return Ok(plan); + }; + if *join.join_type() != JoinType::Inner + || !join + .left() + .schema() + .fields() + .iter() + .all(|field| cacheable_type(field.data_type())) + || !join.on().iter().all(|(left, right)| { + left.downcast_ref::().is_some() && right.downcast_ref::().is_some() + }) + { + return Ok(plan); + } + let source = Arc::new(source.clone()); + let join = Arc::new( + join.builder() + .with_partition_mode(PartitionMode::CollectLeft) + .build()?, + ); + Ok(Arc::new(CachedBroadcastJoinExec { + join, + cache: cache(source.generation), + source, + metrics, + })) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{ + ArrayData, Decimal128Array, DictionaryArray, Int64Array, Int8Array, StringBuilder, + }; + use arrow::buffer::Buffer; + use arrow::datatypes::Int8Type; + use datafusion::common::utils::memory::get_record_batch_memory_size; + use datafusion::execution::memory_pool::GreedyMemoryPool; + use datafusion::physical_plan::empty::EmptyExec; + use datafusion::physical_plan::joins::HashJoinExecBuilder; + use std::ptr::NonNull; + use std::sync::Weak; + + fn batch(columns: Vec) -> RecordBatch { + let fields = columns + .iter() + .enumerate() + .map(|(i, array)| Field::new(format!("c{i}"), array.data_type().clone(), true)) + .collect::>(); + RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap() + } + + fn pool(limit: usize) -> Arc { + Arc::new(GreedyMemoryPool::new(limit)) + } + + #[tokio::test] + async fn large_broadcast_build_admits_compaction_and_releases_storage() -> Result<()> { + let rows = 280; + let payload = "x".repeat(256 * 1024); + let mut strings = StringBuilder::with_capacity(rows, rows * payload.len()); + for _ in 0..rows { + strings.append_value(&payload); + } + let source = batch(vec![ + Arc::new(Int64Array::from_iter_values(0..rows as i64)), + Arc::new(strings.finish()), + ]); + let schema = source.schema(); + let source_bytes = get_record_batch_memory_size(&source); + assert!(source_bytes > 64 * 1024 * 1024); + let bound = broadcast_copy_bytes(&source, &schema)?; + let limit = bound + source_bytes / 2; + let medium = source.slice(0, rows / 2); + assert!(broadcast_copy_bytes(&medium, &schema)? > MAX_COPY_BATCH_BYTES); + assert_eq!( + copy_batch_rows(&medium, &schema, MAX_COPY_BATCH_BYTES, false)?, + medium.num_rows() + ); + let join = HashJoinExecBuilder::new( + Arc::new(EmptyExec::new(Arc::clone(&schema))), + Arc::new(EmptyExec::new(Arc::new(Schema::new(vec![Field::new( + "c0", + DataType::Int64, + true, + )])))), + vec![( + Arc::new(Column::new("c0", 0)), + Arc::new(Column::new("c0", 0)), + )], + JoinType::Inner, + ) + .with_partition_mode(PartitionMode::CollectLeft) + .build()?; + let copied_stream = |target, pool: &Arc| { + let source = source.clone(); + let schema = Arc::clone(&schema); + let reservation = MemoryConsumer::new("test broadcast copy").register(pool); + let stream = futures::stream::try_unfold( + (source, 0, reservation), + move |(source, offset, mut reservation)| async move { + reservation.free(); + if offset == source.num_rows() { + return Ok::<_, DataFusionError>(None); + } + let remainder = source.slice(offset, source.num_rows() - offset); + let rows = + copy_batch_rows(&remainder, &remainder.schema(), target, offset != 0)?; + let columns = copy_broadcast_columns( + &remainder.slice(0, rows), + &remainder.schema(), + &mut reservation, + )?; + let copied = RecordBatch::try_new(remainder.schema(), columns)?; + Ok(Some((copied, (source, offset + rows, reservation)))) + }, + ); + Box::pin(RecordBatchStreamAdapter::new(schema, stream)) as SendableRecordBatchStream + }; + let constrained = pool(limit); + let whole = join + .prepare_build( + copied_stream(usize::MAX, &constrained), + Arc::clone(&constrained), + Arc::new(Default::default()), + ) + .await; + assert!(matches!(whole, Err(DataFusionError::ResourcesExhausted(_)))); + assert_eq!(constrained.reserved(), 0); + + // Public DataFusion compacts the copied inputs into one batch. Its + // reservation must also admit that output while inputs remain alive. + // Chunking the JVM-to-native handoff does not remove this overlap. + let chunked = join + .prepare_build( + copied_stream(MAX_COPY_BATCH_BYTES, &constrained), + Arc::clone(&constrained), + Arc::new(Default::default()), + ) + .await; + assert!(matches!( + chunked, + Err(DataFusionError::ResourcesExhausted(_)) + )); + assert_eq!(constrained.reserved(), 0); + + let compact_pool = pool(3 * bound); + let prepared = join + .prepare_build( + copied_stream(MAX_COPY_BATCH_BYTES, &compact_pool), + Arc::clone(&compact_pool), + Arc::new(Default::default()), + ) + .await?; + assert_eq!(prepared.num_rows(), rows); + assert_eq!(compact_pool.reserved(), prepared.reserved_bytes()); + assert!(prepared.reserved_bytes() < 3 * bound); + drop(prepared); + assert_eq!(compact_pool.reserved(), 0); + Ok(()) + } + + #[test] + fn uncached_broadcast_replay_uses_safe_scan_casts_for_schema_drift() { + let decimal: ArrayRef = Arc::new( + Decimal128Array::from(vec![Some(12300), Some(24400), None, Some(-100)]) + .with_precision_and_scale(38, 2) + .unwrap(), + ); + let unchanged: ArrayRef = Arc::new(Int64Array::from(vec![7, 8, 9, 10])); + let input = batch(vec![decimal, Arc::clone(&unchanged)]); + let schema = Schema::new(vec![ + Field::new("small_value", DataType::Int8, true), + Field::new("unchanged", DataType::Int64, true), + ]); + + // The physical decimal cannot be copied into the declared Int8 cache + // layout. An admission miss must replay this same input without error. + let pool = pool(usize::MAX); + let mut reservation = MemoryConsumer::new("test schema drift").register(&pool); + assert!(matches!( + copy_broadcast_columns(&input, &schema, &mut reservation), + Err(DataFusionError::ResourcesExhausted(_)) + )); + assert_eq!(reservation.size(), 0); + + let columns = cast_uncached_broadcast_columns(&input, &schema).unwrap(); + let small = columns[0].as_any().downcast_ref::().unwrap(); + assert_eq!( + small, + &Int8Array::from(vec![Some(123), None, None, Some(-1)]) + ); + assert!(Arc::ptr_eq(&columns[1], &unchanged)); + } + + /// Model unchecked FFI string buffers without ever reading them as &str. + fn unchecked_utf8(values: &[u8], offsets: &[i32]) -> ArrayRef { + let data = unsafe { + ArrayData::builder(DataType::Utf8) + .len(offsets.len() - 1) + .add_buffer(Buffer::from_slice_ref(offsets)) + .add_buffer(Buffer::from(values.to_vec())) + .build_unchecked() + }; + make_array(data) + } + + #[test] + fn invalid_broadcast_utf8_declines_cache_and_replays_spark_decoding() { + // The second case is valid UTF-8 as a whole but each row splits a codepoint. + for (values, expected) in [ + (&[0xff, b'a'][..], vec!["\u{FFFD}", "a"]), + (&[0xc3, 0xa9][..], vec!["\u{FFFD}", "\u{FFFD}"]), + ] { + let input = batch(vec![unchecked_utf8(values, &[0, 1, 2])]); + let pool = pool(usize::MAX); + let mut reservation = MemoryConsumer::new("test invalid UTF-8").register(&pool); + assert!(matches!( + copy_broadcast_columns(&input, &input.schema(), &mut reservation), + Err(DataFusionError::ResourcesExhausted(_)) + )); + assert_eq!(pool.reserved(), 0); + + let columns = cast_uncached_broadcast_columns(&input, &input.schema()).unwrap(); + let strings = columns[0].as_any().downcast_ref::().unwrap(); + assert_eq!( + strings.iter().collect::>(), + expected.into_iter().map(Some).collect::>() + ); + columns[0].to_data().validate_full().unwrap(); + } + } + + #[test] + fn uncached_broadcast_decodes_dictionary_strings_before_unpacking() { + let values = unchecked_utf8(&[0xff, b'a'], &[0, 1, 2]); + let dictionary = + DictionaryArray::::new(Int8Array::from(vec![Some(1), None, Some(0)]), values); + let input = batch(vec![Arc::new(dictionary)]); + let schema = Schema::new(vec![Field::new("decoded", DataType::Utf8, true)]); + let columns = cast_uncached_broadcast_columns(&input, &schema).unwrap(); + let strings = columns[0].as_any().downcast_ref::().unwrap(); + assert_eq!( + strings.iter().collect::>(), + vec![Some("a"), None, Some("\u{FFFD}")] + ); + columns[0].to_data().validate_full().unwrap(); + } + + #[test] + fn broadcast_utf8_copy_bounds_empty_and_short_values() { + for value in ["", "é"] { + let rows = 4096; + let source = StringArray::from(vec![value; rows]); + let bound = utf8_copy_bytes(rows, utf8_value_span(&source)).unwrap(); + let input = batch(vec![Arc::new(source)]); + let pool = pool(bound); + let mut reservation = MemoryConsumer::new("test broadcast copy").register(&pool); + let copies = copy_broadcast_columns(&input, &input.schema(), &mut reservation).unwrap(); + let copied = copies[0].as_any().downcast_ref::().unwrap(); + assert_eq!(copied.iter().collect::>(), vec![Some(value); rows]); + assert_eq!(pool.reserved(), bound); + assert!(copied.get_array_memory_size() <= bound); + if value.is_empty() { + assert_eq!(copied.values().capacity(), 0); + } + drop(copies); + drop(reservation); + assert_eq!(pool.reserved(), 0); + } + } + + struct ForeignUtf8 { + offsets: Vec, + values: Vec, + validity: Vec, + } + + fn foreign_utf8_slice() -> (StringArray, Weak) { + let owner = Arc::new(ForeignUtf8 { + offsets: vec![0, 6, 6, 17, 20, 24], + values: "prefixhidden-nullaétail".as_bytes().to_vec(), + validity: vec![0b0001_1011], + }); + let buffer = |ptr: *const u8, len: usize| { + // The immutable vectors remain allocated at these addresses until + // the last custom Buffer drops its Arc owner. + unsafe { + Buffer::from_custom_allocation( + NonNull::new(ptr.cast_mut()).unwrap(), + len, + Arc::::clone(&owner), + ) + } + }; + let data = ArrayData::builder(DataType::Utf8) + .len(5) + .add_buffer(buffer( + owner.offsets.as_ptr().cast(), + owner.offsets.len() * size_of::(), + )) + .add_buffer(buffer(owner.values.as_ptr(), owner.values.len())) + .null_bit_buffer(Some(buffer(owner.validity.as_ptr(), owner.validity.len()))) + .build() + .unwrap(); + (StringArray::from(data).slice(1, 3), Arc::downgrade(&owner)) + } + + #[test] + fn broadcast_utf8_copy_owns_sliced_foreign_buffers() { + let (source, owner) = foreign_utf8_slice(); + assert_eq!(utf8_value_span(&source), "hidden-nullaé".len()); + let bound = utf8_copy_bytes(source.len(), utf8_value_span(&source)).unwrap(); + let input = batch(vec![Arc::new(source)]); + let source_data = input.column(0).to_data(); + let source_ptrs = [ + source_data.buffers()[0].as_ptr(), + source_data.buffers()[1].as_ptr(), + source_data.nulls().unwrap().buffer().as_ptr(), + ]; + drop(source_data); + let pool = pool(bound); + let mut reservation = MemoryConsumer::new("test broadcast copy").register(&pool); + let copies = copy_broadcast_columns(&input, &input.schema(), &mut reservation).unwrap(); + assert_eq!(pool.reserved(), bound); + let strings = copies[0].as_any().downcast_ref::().unwrap(); + assert_eq!( + strings.iter().collect::>(), + vec![Some(""), None, Some("aé")] + ); + assert_eq!(strings.value_offsets(), &[0, 0, 11, 14]); + assert_eq!(strings.values().as_slice(), "hidden-nullaé".as_bytes()); + assert!(strings.get_array_memory_size() <= bound); + let data = strings.to_data(); + let copied_ptrs = [ + data.buffers()[0].as_ptr(), + data.buffers()[1].as_ptr(), + data.nulls().unwrap().buffer().as_ptr(), + ]; + for (source, copied) in source_ptrs.iter().zip(copied_ptrs.iter()) { + assert_ne!(source, copied); + } + drop(input); + assert!( + owner.upgrade().is_none(), + "Copied column retained the foreign task owner" + ); + drop(copies); + drop(reservation); + assert_eq!(pool.reserved(), 0); + } + + #[test] + fn broadcast_utf8_copy_denied_admission_keeps_input_replayable() { + let source: ArrayRef = Arc::new(StringArray::from(vec![Some("aé"), None, Some("")])); + let input = batch(vec![Arc::clone(&source), source]); + let bound = 2 * utf8_copy_bytes(3, "aé".len()).unwrap(); + let denied = pool(bound - 1); + let mut reservation = MemoryConsumer::new("test denied copy").register(&denied); + assert!(matches!( + copy_broadcast_columns(&input, &input.schema(), &mut reservation), + Err(DataFusionError::ResourcesExhausted(_)) + )); + assert_eq!(reservation.size(), 0); + assert_eq!(denied.reserved(), 0); + // Failure leaves the source intact for ordinary replay or a fresh grant. + let admitted = pool(bound); + let mut reservation = MemoryConsumer::new("test admitted copy").register(&admitted); + let copies = copy_broadcast_columns(&input, &input.schema(), &mut reservation).unwrap(); + for (source, copied) in input.columns().iter().zip(&copies) { + assert_eq!(source.to_data(), copied.to_data()); + } + drop(copies); + drop(reservation); + assert_eq!(admitted.reserved(), 0); + } +} diff --git a/native/core/src/execution/operators/broadcast_cache.rs b/native/core/src/execution/operators/broadcast_cache.rs new file mode 100644 index 00000000000..daac289c790 --- /dev/null +++ b/native/core/src/execution/operators/broadcast_cache.rs @@ -0,0 +1,437 @@ +// 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. + +//! Coordinates concurrent tasks preparing the same immutable broadcast build. +//! The first task for a compatible key loads it; other tasks wait and receive +//! their own strong references. This module owns lookup metadata, not Arrow +//! buffers or Spark storage grants. +//! +//! Successful entries keep only weak references, so the last active probe can +//! drop the prepared hash table. Its owner returns the Spark storage charge when +//! the last lease drops or, if earlier, on executor retirement. A later task +//! wave may build the same table again. Resource failures briefly reject +//! matching loads to avoid repeated partial work; other errors and canceled +//! loaders leave the key retryable. Retiring a SparkEnv generation clears its +//! lookups without revoking leases already held by active probes. + +use std::collections::HashMap; +use std::future::Future; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Weak}; +use std::time::{Duration, Instant}; + +use arrow::datatypes::SchemaRef; +use datafusion::common::{resources_datafusion_err, DataFusionError, Result}; +use futures::channel::oneshot; +use parking_lot::Mutex; + +/// Share failed admission across a task wave without permanently disabling reuse +/// when executor storage pressure changes. Rejections occupy bounded cache slots. +const RESOURCE_RETRY_DELAY: Duration = Duration::from_secs(5); + +/// Identifies a compatible materialization: the Spark broadcast, its declared +/// build and probe schemas, and ordered build key columns must all match. +/// SparkEnv generation is scoped by the enclosing cache, not by this key. +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +pub(super) struct BuildKey { + pub broadcast_id: i64, + pub schema: SchemaRef, + pub probe_schema: SchemaRef, + pub key_columns: Vec, +} + +/// Per-key state: one loader with waiters, a weak active build, or a short-lived +/// resource rejection. No entry owns a prepared build after loading finishes. +#[derive(Debug)] +enum Entry { + Loading(Vec>>>), + Ready(Weak), + Rejected(Instant), +} + +/// Bounds keys, unfinished loads, weak ready entries and temporary rejections +/// without retaining prepared builds. The payload byte cap belongs to the +/// allocation pool; callers own strong references returned by `get_or_load`. +#[derive(Debug)] +pub(super) struct BuildCache { + entries: Mutex>>, + max_entries: usize, + retired: AtomicBool, +} + +impl BuildCache { + /// Creates an empty cache. A zero limit disables admission. Does not allocate + /// payload memory or install any task/context ownership. + pub fn new(max_entries: usize) -> Self { + Self { + entries: Mutex::new(HashMap::new()), + max_entries, + retired: AtomicBool::new(false), + } + } + + /// Returns an active build or elects one loader for this task wave. Waiting + /// tasks receive strong references before the loader can drop its build. + /// A canceled load wakes waiters to retry; a resource error from the loader + /// is shared briefly so peers take ordinary fallback instead of repeating + /// partial preparation. Other errors remain uncached. Entry-capacity failure + /// occurs before `load` runs. The boolean is true on a hit or waited load. + pub async fn get_or_load(&self, key: BuildKey, load: F) -> Result<(Arc, bool)> + where + F: FnOnce() -> Fut, + Fut: Future>>, + { + enum Action { + Hit(Arc), + Wait(oneshot::Receiver>>), + Load, + } + let mut load = Some(load); + loop { + let action = { + let mut entries = self.entries.lock(); + if self.retired.load(Ordering::Relaxed) { + return Err(resources_datafusion_err!( + "Broadcast build cache is retired" + )); + } + match entries.get_mut(&key) { + Some(Entry::Ready(value)) => match value.upgrade() { + Some(value) => Action::Hit(value), + None => { + // The last probe released its build; this key may + // elect a new loader for a later task wave. + entries.remove(&key); + continue; + } + }, + Some(Entry::Rejected(retry_at)) => { + if Instant::now() < *retry_at { + return Err(DataFusionError::ResourcesExhausted( + "Broadcast build cache admission rejected".into(), + )); + } + // Expired rejections contain no payload or task ownership. + // Re-enter election so only one caller retries admission. + entries.remove(&key); + continue; + } + Some(Entry::Loading(waiters)) => { + let (sender, receiver) = oneshot::channel(); + waiters.push(sender); + Action::Wait(receiver) + } + None => { + if entries.len() >= self.max_entries { + // The cap also counts bookkeeping for old waves; + // reclaim dead weak refs and expired rejections. + let now = Instant::now(); + entries.retain(|_, entry| match entry { + Entry::Loading(_) => true, + Entry::Ready(value) => value.strong_count() != 0, + Entry::Rejected(retry_at) => *retry_at > now, + }); + if entries.len() >= self.max_entries { + return Err(resources_datafusion_err!( + "Broadcast build cache entry limit reached" + )); + } + } + entries.insert(key.clone(), Entry::Loading(vec![])); + Action::Load + } + } + }; + match action { + Action::Hit(value) => return Ok((value, true)), + Action::Wait(receiver) => { + if let Ok(Some(value)) = receiver.await { + if !self.retired.load(Ordering::Relaxed) { + return Ok((value, true)); + } + } + // Shared rejection or loader cancellation re-enters election. + } + Action::Load => { + let guard = LoadGuard { + cache: self, + key: key.clone(), + active: true, + }; + match load.take().expect("only the elected caller loads")().await { + Ok(value) => { + guard.publish(Entry::Ready(Arc::downgrade(&value)), Some(&value)); + return Ok((value, false)); + } + Err(DataFusionError::ResourcesExhausted(error)) => { + guard.publish( + Entry::Rejected(Instant::now() + RESOURCE_RETRY_DELAY), + None, + ); + return Err(DataFusionError::ResourcesExhausted(error)); + } + Err(error) => return Err(error), + } + } + } + } + } + + /// Retires all lookup entries on native-runtime shutdown. Active leases keep + /// their values and charges; loading callers cannot republish after retirement. + /// Dropping senders wakes waiting tasks without retaining any task future. + pub fn clear(&self) { + let entries = { + let mut entries = self.entries.lock(); + self.retired.store(true, Ordering::Relaxed); + std::mem::take(&mut *entries) + }; + drop(entries); + } +} + +/// A task-owned loading slot. Cancellation removes it and wakes waiters. The +/// guard becomes inactive before publishing a value or rejection. +struct LoadGuard<'a, T> { + cache: &'a BuildCache, + key: BuildKey, + active: bool, +} + +impl LoadGuard<'_, T> { + /// Publishes a weak lookup or rejection and hands the build to waiting tasks. + /// Runtime shutdown may already have retired the slot; canceled loaders + /// cannot publish into a newer environment. + fn publish(mut self, entry: Entry, value: Option<&Arc>) { + let removed = { + let mut entries = self.cache.entries.lock(); + let removed = if matches!(entries.get(&self.key), Some(Entry::Loading(_))) { + entries.insert(self.key.clone(), entry) + } else { + None + }; + self.active = false; + removed + }; + if let Some(Entry::Loading(waiters)) = removed { + for waiter in waiters { + let _ = waiter.send(value.cloned()); + } + } + } +} + +impl Drop for LoadGuard<'_, T> { + /// Retracts an unfinished load and wakes waiters. Published slots are untouched. + fn drop(&mut self) { + if !self.active { + return; + } + let removed = { + let mut entries = self.cache.entries.lock(); + if matches!(entries.get(&self.key), Some(Entry::Loading(_))) { + entries.remove(&self.key) + } else { + None + } + }; + drop(removed); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::datatypes::Schema; + use std::sync::atomic::{AtomicUsize, Ordering}; + + /// Builds a stable fake materialization key; tests own separate caches. + fn key(id: i64) -> BuildKey { + BuildKey { + broadcast_id: id, + schema: Arc::new(Schema::empty()), + probe_schema: Arc::new(Schema::empty()), + key_columns: vec![], + } + } + + #[tokio::test] + async fn waiting_probe_retains_build_then_next_wave_rebuilds() { + use futures::{poll, FutureExt}; + + let cache = BuildCache::new(1); + let loads = AtomicUsize::new(0); + let mut loader = Box::pin(cache.get_or_load(key(1), || async { + loads.fetch_add(1, Ordering::SeqCst); + tokio::task::yield_now().await; + Ok(Arc::new(7)) + })); + assert!(poll!(&mut loader).is_pending()); + let mut waiter = Box::pin(cache.get_or_load(key(1), || async { + loads.fetch_add(1, Ordering::SeqCst); + Ok(Arc::new(8)) + })); + assert!(poll!(&mut waiter).is_pending()); + + let (built, hit) = loader.await.unwrap(); + assert!(!hit); + let weak = Arc::downgrade(&built); + drop(built); + let (shared, hit) = waiter.now_or_never().unwrap().unwrap(); + assert!(hit); + assert_eq!(*shared, 7); + assert_eq!(loads.load(Ordering::SeqCst), 1); + drop(shared); + assert!(weak.upgrade().is_none()); + + let (next, hit) = cache + .get_or_load(key(1), || async { + loads.fetch_add(1, Ordering::SeqCst); + Ok(Arc::new(9)) + }) + .await + .unwrap(); + assert!(!hit); + assert_eq!(*next, 9); + assert_eq!(loads.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn released_build_frees_lookup_slot() { + let cache = BuildCache::new(2); + let (released, _) = cache + .get_or_load(key(1), || async { Ok(Arc::new(1)) }) + .await + .unwrap(); + drop(released); + let (active, _) = cache + .get_or_load(key(2), || async { Ok(Arc::new(2)) }) + .await + .unwrap(); + let (new, _) = cache + .get_or_load(key(3), || async { Ok(Arc::new(3)) }) + .await + .unwrap(); + let entries = cache.entries.lock(); + assert_eq!(entries.len(), 2); + assert!(!entries.contains_key(&key(1))); + assert!(entries.contains_key(&key(2))); + assert!(entries.contains_key(&key(3))); + assert_eq!(*active, 2); + assert_eq!(*new, 3); + } + + /// A failed loader rejects its waiters without repeating preparation. The + /// next task can load the same key after the short resource cooldown. + #[tokio::test] + async fn admission_failure_is_shared_then_retried() { + let cache = BuildCache::::new(1); + let loads = AtomicUsize::new(0); + let (failed, waiter) = futures::join!( + cache.get_or_load(key(1), || async { + loads.fetch_add(1, Ordering::SeqCst); + tokio::task::yield_now().await; + Err(DataFusionError::ResourcesExhausted("full".into())) + }), + cache.get_or_load(key(1), || async { + loads.fetch_add(1, Ordering::SeqCst); + Ok(Arc::new(7)) + }) + ); + assert!(matches!( + failed, + Err(DataFusionError::ResourcesExhausted(_)) + )); + assert!(matches!( + waiter, + Err(DataFusionError::ResourcesExhausted(_)) + )); + assert_eq!(loads.load(Ordering::SeqCst), 1); + assert!(cache + .get_or_load(key(1), || async { panic!("cooldown must bypass loading") }) + .await + .is_err()); + { + let mut entries = cache.entries.lock(); + let Entry::Rejected(retry_at) = entries.get_mut(&key(1)).unwrap() else { + panic!("resource denial must install a rejection"); + }; + *retry_at = Instant::now(); + } + let (value, hit) = cache + .get_or_load(key(1), || async { + loads.fetch_add(1, Ordering::SeqCst); + Ok(Arc::new(7)) + }) + .await + .unwrap(); + assert!(!hit); + assert_eq!(*value, 7); + assert_eq!(loads.load(Ordering::SeqCst), 2); + } + + /// Non-resource errors remain visible to their caller and are not cached; + /// an already-waiting consumer can immediately attempt its own load. + #[tokio::test] + async fn data_error_is_not_shared_as_resource_fallback() { + let cache = BuildCache::new(1); + let (failed, retried) = futures::join!( + cache.get_or_load(key(1), || async { + tokio::task::yield_now().await; + Err(DataFusionError::Execution("invalid input".into())) + }), + cache.get_or_load(key(1), || async { Ok(Arc::new(7)) }) + ); + assert!(matches!(failed, Err(DataFusionError::Execution(error)) + if error == "invalid input")); + let (value, hit) = retried.unwrap(); + assert!(!hit); + assert_eq!(*value, 7); + } + + #[tokio::test] + async fn canceled_loader_wakes_waiter_and_allows_retry() { + use futures::{poll, FutureExt}; + let cache = BuildCache::new(1); + let mut loader = Box::pin(cache.get_or_load(key(1), futures::future::pending)); + assert!(poll!(&mut loader).is_pending()); + let mut waiter = Box::pin(cache.get_or_load(key(1), || async { Ok(Arc::new(7)) })); + assert!(poll!(&mut waiter).is_pending()); + drop(loader); + assert_eq!(*waiter.now_or_never().unwrap().unwrap().0, 7); + } + + #[tokio::test] + async fn leases_survive_retirement_and_bound_admission() { + let cache = BuildCache::new(1); + let (lease, _) = cache + .get_or_load(key(1), || async { Ok(Arc::new(9)) }) + .await + .unwrap(); + assert!(cache + .get_or_load(key(2), || async { panic!("must not load") }) + .await + .is_err()); + cache.clear(); + assert_eq!(*lease, 9); + assert!(cache + .get_or_load(key(2), || async { panic!("retired cache cannot load") }) + .await + .is_err()); + assert!(cache.entries.lock().is_empty()); + } +} diff --git a/native/core/src/execution/operators/dynamic_filter/join.rs b/native/core/src/execution/operators/dynamic_filter/join.rs index ba2e2759608..9c147980108 100644 --- a/native/core/src/execution/operators/dynamic_filter/join.rs +++ b/native/core/src/execution/operators/dynamic_filter/join.rs @@ -45,6 +45,7 @@ use datafusion::physical_plan::{ }; use futures::StreamExt; +use super::super::broadcast::reuse_broadcast_build_with_metrics; use super::parquet_reader::try_attach_parquet_reader_filter; use super::DynamicFilterExec; @@ -145,12 +146,19 @@ impl DynamicFilterJoinExec { context.window_functions().clone(), context.runtime_env(), )); - let result = join.execute(partition, context); - // HashJoinExec registers its metrics synchronously in execute(). Keep the - // live counters, including on error, without retaining the producer plan. - for metric in join.metrics().unwrap_or_default().iter() { - self.metrics.register(Arc::clone(metric)); + let join: Arc = Arc::new(join); + // Attach reuse only after creating this execution's producer and filter. + // Cached execution publishes counters after its asynchronous cache lookup. + let execution = + reuse_broadcast_build_with_metrics(Arc::clone(&join), self.metrics.clone())?; + let result = execution.execute(partition, context); + if Arc::ptr_eq(&execution, &join) { + // An ordinary join registers synchronously, including on error. + for metric in join.metrics().unwrap_or_default().iter() { + self.metrics.register(Arc::clone(metric)); + } } + drop(execution); drop(join); let input = result?; // Drop execution state at EOF or error even if the caller retains the diff --git a/native/core/src/execution/operators/mod.rs b/native/core/src/execution/operators/mod.rs index a15a528f90b..201ea10dbc7 100644 --- a/native/core/src/execution/operators/mod.rs +++ b/native/core/src/execution/operators/mod.rs @@ -22,6 +22,9 @@ pub use crate::errors::ExecutionError; pub use iceberg_scan::*; pub use scan::*; +mod broadcast; +mod broadcast_cache; +pub(crate) use broadcast::{clear_broadcast_cache, reuse_broadcast_build, BroadcastInputExec}; mod dynamic_filter; pub(crate) use dynamic_filter::{DynamicFilterJoinExec, TopKReaderFilterExec}; pub(crate) mod iceberg_common; diff --git a/native/core/src/execution/operators/scan.rs b/native/core/src/execution/operators/scan.rs index e074f8ebf7f..1e5be82badf 100644 --- a/native/core/src/execution/operators/scan.rs +++ b/native/core/src/execution/operators/scan.rs @@ -179,7 +179,7 @@ impl ScanExec { /// Spark-rendered form, since arrow's `from_ffi` imports string buffers unchecked. Otherwise the /// column keeps the imported buffers without a copy: the stream transfers ownership by reference /// count, and the JVM decodes dictionaries before export. -fn import_column(col: &ArrayRef) -> Result { +pub(super) fn import_column(col: &ArrayRef) -> Result { Ok(decode_string_arrays(col)?) } diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 98f9c2a7cbc..ca9ae2f86f3 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -38,6 +38,7 @@ use crate::execution::operators::init_csv_datasource_exec; use crate::execution::operators::DynamicFilterJoinExec; use crate::execution::operators::IcebergScanExec; use crate::execution::operators::TopKReaderFilterExec; +use crate::execution::operators::{reuse_broadcast_build, BroadcastInputExec}; use crate::execution::{ operators::{ ExecutionError, MergeActionContext, MergeInstructionExec, MergeRowsExec, ScanExec, @@ -1934,6 +1935,22 @@ impl PhysicalPlanner { None } else { let java_stream = inputs.remove(0); + // A broadcast carrier exposes its schema without decoding. Keep it + // out of ordinary scans, which JNI eagerly pulls before execution; + // the cache winner alone opens this replayable input. + if let Some(broadcast) = + BroadcastInputExec::try_new(Arc::clone(&java_stream), &data_types)? + { + return Ok(( + vec![], + vec![], + Arc::new(SparkPlan::new( + spark_plan.plan_id, + Arc::new(broadcast), + vec![], + )), + )); + } let address: i64 = JVMClasses::with_env(|env| unsafe { jni_call!(env, arrow_array_stream(java_stream.as_obj()).memory_address() -> i64) })?; @@ -2485,6 +2502,7 @@ impl PhysicalPlanner { join.dynamic_filter_enabled && !join.null_aware_anti_join, self.session_ctx.copied_config().options(), )?; + let hash_join = reuse_broadcast_build(hash_join)?; Ok(( scans, shuffle_scans, @@ -2502,7 +2520,7 @@ impl PhysicalPlanner { join.dynamic_filter_enabled, self.session_ctx.copied_config().options(), )?; - + let swapped_hash_join = reuse_broadcast_build(swapped_hash_join)?; let mut additional_native_plans = vec![]; if swapped_hash_join.is::() { // a projection was added to the hash join diff --git a/native/core/src/lib.rs b/native/core/src/lib.rs index a05842cf031..e7dcf0d5cb4 100644 --- a/native/core/src/lib.rs +++ b/native/core/src/lib.rs @@ -183,6 +183,20 @@ pub extern "system" fn Java_org_apache_comet_NativeBase_release(_e: EnvUnowned, execution::jni_api::release_runtime(); } +#[no_mangle] +/// Retires broadcast lookups on every plugin shutdown without changing runtime ownership. +/// Call after native task shutdown and before retiring the JVM storage owner. Existing leases +/// retain their buffers; repeated calls are safe. JNI failures are returned as Java exceptions. +pub extern "system" fn Java_org_apache_comet_NativeBase_clearBroadcastCache( + e: EnvUnowned, + _class: JClass, +) { + try_unwrap_or_throw(&e, |_| { + execution::operators::clear_broadcast_cache(); + Ok(()) + }) +} + const LOG_PATTERN: &str = "{d(%y/%m/%d %H:%M:%S)} {l} {f}: {m}{n}"; /// JNI method to check if a specific feature is enabled in the native Rust code. diff --git a/spark/src/main/java/org/apache/comet/NativeBase.java b/spark/src/main/java/org/apache/comet/NativeBase.java index ab4eee856ba..f6a66d189c1 100644 --- a/spark/src/main/java/org/apache/comet/NativeBase.java +++ b/spark/src/main/java/org/apache/comet/NativeBase.java @@ -344,14 +344,30 @@ private static String resourceName() { /** Release native resources through JNI */ static native void release(); - /** Release native resources */ + /** + * Retire native broadcast cache lookups without releasing the runtime. Active native leases keep + * their buffers; call after task shutdown and before retiring their executor storage owner. + * Repeated calls are safe and native failures propagate to the caller. + */ + public static native void clearBroadcastCache(); + + /** + * Release the native runtime at most once, but retire broadcast cache ownership on every plugin + * shutdown. A later SparkContext may have populated another cache even after runtime release was + * recorded. Cache retirement still runs if release fails; unloaded libraries require no cleanup. + * Load or JNI failures propagate to the caller, which must finish executor-owner retirement. + */ public static void releaseNative() throws Throwable { if (!isLoaded()) { return; } - if (released.compareAndSet(false, true)) { - release(); + try { + if (released.compareAndSet(false, true)) { + release(); + } + } finally { + clearBroadcastCache(); } } diff --git a/spark/src/main/java/org/apache/spark/CometBroadcastMemoryManager.java b/spark/src/main/java/org/apache/spark/CometBroadcastMemoryManager.java new file mode 100644 index 00000000000..299ec61579c --- /dev/null +++ b/spark/src/main/java/org/apache/spark/CometBroadcastMemoryManager.java @@ -0,0 +1,175 @@ +/* + * 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.spark; + +/* + * Bridges native broadcast build allocations to Spark's executor storage memory. Creating an + * owner reserves no storage; native loaders request off-heap grants only when they build. Active + * probe streams can keep those grants across task boundaries. Each SparkEnv has its own owner + * generation, so retiring an executor never charges a replacement executor for old leases. + */ +import java.lang.ref.WeakReference; + +import org.apache.spark.memory.MemoryManager; +import org.apache.spark.memory.MemoryMode; +import org.apache.spark.storage.BlockId; +import org.apache.spark.storage.BroadcastBlockId; + +/** + * Executor-owned storage accounting for prepared native broadcasts. This class lives in + * org.apache.spark to use Spark's storage memory API. One cap covers loading and leased builds; + * CometPlugin clears the native cache before retiring this owner at shutdown. + */ +public final class CometBroadcastMemoryManager { + private static long nextGeneration; + private static CometBroadcastMemoryManager current; + + private final WeakReference environment; + private MemoryManager memoryManager; + private final BlockId blockId; + private final long generation; + private final long limit; + private long used; + private boolean retired; + + /** Capture an executor without retaining its SparkEnv after shutdown. */ + private CometBroadcastMemoryManager( + Object environment, MemoryManager memoryManager, long generation, long limit) { + this.environment = new WeakReference<>(environment); + this.memoryManager = memoryManager; + this.generation = generation; + this.limit = limit; + this.blockId = new BroadcastBlockId(generation, "comet-prepared-native"); + } + + /** + * Return an owner only when CometPlugin and Spark off-heap memory support reuse. Constructing the + * owner does not reserve storage; a null owner makes the caller open an ordinary stream. + */ + public static CometBroadcastMemoryManager getOrCreate(long maxBytes) { + SparkEnv env = SparkEnv.get(); + if (env == null) { + return null; + } + boolean hasCometPlugin = false; + for (String plugin : env.conf().get("spark.plugins", "").split(",")) { + if (plugin.trim().equals("org.apache.spark.CometPlugin")) { + hasCometPlugin = true; + break; + } + } + if (!hasCometPlugin) { + return null; + } + return getOrCreate( + env, + env.memoryManager(), + env.conf().getBoolean("spark.memory.offHeap.enabled", false), + maxBytes); + } + + /** + * Package-scope variant for lifecycle tests. One SparkEnv shares one cap: tasks with a different + * cap receive no owner rather than sharing an ambiguous budget. A new SparkEnv retires the old + * owner's Spark charge and starts a separate generation. + */ + static synchronized CometBroadcastMemoryManager getOrCreate( + Object environment, MemoryManager memoryManager, boolean offHeapEnabled, long maxBytes) { + if (current != null && current.environment.get() != environment) { + current.retire(); + current = null; + } + if (!offHeapEnabled) { + return null; + } + if (current != null) { + return !current.retired && current.limit == maxBytes ? current : null; + } + ++nextGeneration; + current = new CometBroadcastMemoryManager(environment, memoryManager, nextGeneration, maxBytes); + return current; + } + + /** + * Grant the whole native allocation or none, under one cap for all builds in this generation. + * Spark's usual off-heap storage admission may evict other cached Spark blocks. + */ + public synchronized long acquireMemory(long size) { + if (retired || size <= 0 || size > limit - used) { + return 0; + } + if (!memoryManager.acquireStorageMemory(blockId, size, MemoryMode.OFF_HEAP)) { + return 0; + } + used += size; + return size; + } + + /** + * Release a native lease. Retirement already returned the Spark storage charge, so late releases + * update only this owner's local count and cannot touch a replacement executor. + */ + public synchronized void releaseMemory(long size) { + if (size > used) { + throw new IllegalArgumentException("Broadcast memory release exceeds its outstanding grant"); + } + if (!retired && size != 0) { + memoryManager.releaseStorageMemory(size, MemoryMode.OFF_HEAP); + } + used -= size; + } + + /** Returns the outstanding native grants, including leases remaining after retirement. */ + public synchronized long getUsedMemory() { + return used; + } + + /** Returns this owner's immutable generation, which distinguishes replacement SparkEnvs. */ + public long getGeneration() { + return generation; + } + + /** Returns the immutable byte cap shared by all prepared builds using this owner. */ + public long getLimit() { + return limit; + } + + /** + * Return Spark's storage charge once, while keeping the outstanding count for native leases that + * release after executor shutdown or a SparkEnv replacement. + */ + private synchronized void retire() { + if (!retired) { + retired = true; + MemoryManager previousManager = memoryManager; + memoryManager = null; + if (used != 0) { + previousManager.releaseStorageMemory(used, MemoryMode.OFF_HEAP); + } + } + } + + /** Retire this executor's owner without reopening it for the same SparkEnv. */ + public static synchronized void shutdown() { + if (current != null) { + current.retire(); + } + } +} diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index f2f19ac5ce9..30cb709977e 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -511,6 +511,30 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(false) + val COMET_BROADCAST_REUSE_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.broadcast.reuse.enabled") + .category(CATEGORY_EXEC) + .doc( + "Experimental executor-wide reuse of prepared native broadcast hash joins. " + + "Requires org.apache.spark.CometPlugin in spark.plugins, Spark off-heap memory, " + + "and a supported fixed-width or plain UTF8 build schema. " + + "Unsupported joins and rejected cache admissions retain ordinary execution.") + .booleanConf + .createWithDefault(false) + + val COMET_BROADCAST_REUSE_MAX_MEMORY: ConfigEntry[Long] = + conf("spark.comet.broadcast.reuse.maxMemory") + .category(CATEGORY_EXEC) + .doc( + "Maximum bytes of prepared broadcast data held by active tasks per executor, charged " + + "to Spark off-heap storage memory. The charge is released after the last task drops " + + "the build. The first flat-schema broadcast input marked for reuse " + + "fixes this executor's cap, even if its join is later ineligible. Inputs requesting " + + "a different cap use ordinary execution.") + .bytesConf(ByteUnit.BYTE) + .checkValue(_ > 0, "Broadcast reuse memory limit must be positive") + .createWithDefault(1024L * 1024 * 1024) + val COMET_FORCE_SHJ: ConfigEntry[Boolean] = conf(s"$COMET_EXEC_CONFIG_PREFIX.forceShuffledHashJoin") .withAlternative(s"$COMET_EXEC_CONFIG_PREFIX.replaceSortMergeJoin") diff --git a/spark/src/main/scala/org/apache/spark/Plugins.scala b/spark/src/main/scala/org/apache/spark/Plugins.scala index e79074f45d2..b2baf66332f 100644 --- a/spark/src/main/scala/org/apache/spark/Plugins.scala +++ b/spark/src/main/scala/org/apache/spark/Plugins.scala @@ -413,6 +413,10 @@ class CometExecutorPlugin extends ExecutorPlugin with Logging { super.init(ctx, extraConf) } + /** + * Stops native work before retiring executor-owned broadcast grants. Retirement still runs if + * native shutdown throws; the original failure propagates unless retirement itself fails. + */ override def shutdown(): Unit = { logInfo("CometExecutorPlugin shutdown") @@ -420,7 +424,11 @@ class CometExecutorPlugin extends ExecutorPlugin with Logging { // already replaced it. CometExecutorPlugin.current.compareAndSet(context, null) - NativeBase.releaseNative() + try { + NativeBase.releaseNative() + } finally { + CometBroadcastMemoryManager.shutdown() + } super.shutdown() } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastExchangeExec.scala index 615f7bada26..fedd5a6c108 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastExchangeExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastExchangeExec.scala @@ -215,10 +215,13 @@ case class CometBroadcastExchangeExec( new CometBatchRDD(sparkContext, getNumPartitions(), broadcasted) } - // After https://issues.apache.org/jira/browse/SPARK-48195, Spark plan will cache created RDD. - // Since we may change the number of partitions in CometBatchRDD, - // we need a method that always creates a new CometBatchRDD. - def executeColumnar(numPartitions: Int): RDD[ColumnarBatch] = { + /** + * Create a fresh broadcast RDD aligned to the probe partitions. Spark can cache the RDD from + * doExecuteColumnar (SPARK-48195), whose partition count may differ. Keep the concrete type so + * a lazy native input can get each partition's actual Broadcast handle without decoding its + * batches. + */ + def executeColumnar(numPartitions: Int): CometBatchRDD = { if (isCanonicalizedPlan) { throw SparkException.internalError("A canonicalized plan is not supposed to be executed.") } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastInput.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastInput.scala new file mode 100644 index 00000000000..5b6677e00d0 --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastInput.scala @@ -0,0 +1,169 @@ +/* + * 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.spark.sql.comet + +/* + * Passes a Spark broadcast to native execution as a task-local handle, without decoding it on + * a cache hit. On a miss or ordinary-join fallback, native code can request fresh Arrow streams + * over the same broadcast. Each stream has its own decoder; task completion closes the exported + * wrappers and allocators after native execution drops their readers. + */ + +import scala.collection.mutable.ArrayBuffer + +import org.apache.arrow.c.ArrowArrayStream +import org.apache.arrow.util.AutoCloseables +import org.apache.spark.{CometBroadcastMemoryManager, TaskContext} +import org.apache.spark.broadcast.Broadcast +import org.apache.spark.comet.CometTaskContextShim +import org.apache.spark.sql.comet.execution.arrow.{ArrowReaderIterator, CometBroadcastArrowStream} +import org.apache.spark.sql.comet.util.Utils +import org.apache.spark.sql.types.{BooleanType, ByteType, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, ShortType, StringType, StructType, TimestampNTZType, TimestampType} +import org.apache.spark.sql.vectorized.ColumnarBatch +import org.apache.spark.util.io.ChunkedByteBuffer + +/** + * Task-local JNI input that identifies a real Spark broadcast without opening its payload. Native + * execution may open fresh streams on a cache miss or admission fallback. Each stream has + * independent decoder state; only Rust-owned copies may outlive this task. Register cleanup here, + * before CometExecIterator, so reverse task-listener order drops native readers first. + */ +class CometBroadcastInput private[comet] ( + broadcast: Broadcast[Array[ChunkedByteBuffer]], + schema: StructType, + name: String, + context: TaskContext, + memoryManager: CometBroadcastMemoryManager) + extends AutoCloseable { + + private val streams = ArrayBuffer.empty[CometBroadcastArrowStream] + private var closed = false + + context.addTaskCompletionListener[Unit](_ => close()) + + /** Return the actual materialization ID without evaluating Broadcast.value or decoding rows. */ + def getBroadcastId(): Long = broadcast.id + + /** Return the executor-lifetime owner, or null when the caller must use the uncached path. */ + def getMemoryManager(): CometBroadcastMemoryManager = memoryManager + + /** + * Export a new stream over this broadcast, preserving the captured task context on JNI threads. + * The caller moves the C stream into its native reader; this object retains its wrapper and + * allocator until task completion. Failure closes all resources from this attempt and leaves + * earlier streams intact. Calls after close or task cancellation fail before reading data. + */ + def openStream(): ArrowArrayStream = synchronized { + require(!closed, "Cannot open a completed broadcast input") + withTaskContext { + val source = new BroadcastBatchIterator + val stream = CometBroadcastArrowStream.open(source, schema, name) + streams += stream + stream.stream + } + } + + /** + * Release unclaimed exports and all wrapper/allocator storage after native readers are dropped. + * Idempotent; attempts every stream even if one cleanup fails. It never reads unopened data. + */ + override def close(): Unit = synchronized { + if (!closed) { + closed = true + val owned = streams.toArray + streams.clear() + AutoCloseables.close(owned: _*) + } + } + + /** + * Native callbacks may arrive on threads without Spark's task-local context. Install the + * captured context for decoder work and cancellation checks, then restore the caller's context. + */ + private def withTaskContext[T](body: => T): T = { + val previous = TaskContext.get() + CometTaskContextShim.set(context) + try { + context.killTaskIfInterrupted() + body + } finally { + if (previous == null) CometTaskContextShim.unset() + else CometTaskContextShim.set(previous) + } + } + + /** + * A fresh decoder cursor for one export. Marker construction and cache hits never read + * Broadcast.value; stream opening reconciles the schema and may decode the first chunk. An + * early drop closes the current IPC reader without draining remaining chunks. + */ + private class BroadcastBatchIterator extends Iterator[ColumnarBatch] with AutoCloseable { + private lazy val chunks = broadcast.value.iterator + private var batches: Iterator[ColumnarBatch] = Iterator.empty + private var stopped = false + + /** Advance across empty/exhausted IPC chunks; no new chunk is opened after close. */ + override def hasNext: Boolean = withTaskContext { + if (stopped) false + else { + while (!batches.hasNext && chunks.hasNext) { + batches = Utils.decodeBatches(chunks.next(), name) + } + batches.hasNext + } + } + + /** Return the next decoded batch, whose ownership passes to the Arrow stream reader. */ + override def next(): ColumnarBatch = { + if (!hasNext) throw new NoSuchElementException("No more broadcast batches") + batches.next() + } + + /** Close the current IPC reader if present; empty chunks have no reader to release. */ + private def closeCurrent(): Unit = { + batches match { + case reader: ArrowReaderIterator => reader.close() + case _ => + } + batches = Iterator.empty + } + + /** Idempotently close the current reader without opening or draining remaining chunks. */ + override def close(): Unit = { + if (!stopped) { + stopped = true + closeCurrent() + } + } + } +} + +private[comet] object CometBroadcastInput { + + /** Admit fixed-width fields and strings; native still checks the physical UTF8 encoding. */ + def supportsSchema(schema: StructType): Boolean = schema.fields.forall { field => + field.dataType match { + case BooleanType | ByteType | ShortType | IntegerType | LongType | FloatType | DoubleType | + DateType | TimestampType | TimestampNTZType | StringType | _: DecimalType => + true + case _ => false + } + } +} diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastInputRDD.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastInputRDD.scala new file mode 100644 index 00000000000..b1e3ae4e3a7 --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastInputRDD.scala @@ -0,0 +1,62 @@ +/* + * 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.spark.sql.comet + +/* + * Aligns a lazy broadcast input with each probe partition. The driver captures the schema and + * memory cap; executor compute passes on Spark's actual Broadcast handle without reading its + * payload. Without a reuse owner, CometExecRDD opens an ordinary Arrow stream. With an owner, + * native execution chooses whether to reuse a build or open a stream for this task. + */ + +import org.apache.spark.{CometBroadcastMemoryManager, OneToOneDependency, Partition, TaskContext} +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.vectorized.ColumnarBatch + +/** + * Broadcast input slots that preserve the actual Broadcast handle without constructing a decoder. + * The schema and memory cap are captured on the driver; executor SQLConf cannot alter admission. + */ +private[comet] class CometBroadcastInputRDD( + batches: CometBatchRDD, + schema: StructType, + maxBytes: Long, + name: String) + extends RDD[CometBroadcastInput](batches.context, Seq(new OneToOneDependency(batches))) { + + /** Reuse the parent's aligned broadcast partitions without evaluating their payloads. */ + override protected def getPartitions: Array[Partition] = + firstParent[ColumnarBatch].partitions + + /** + * Construct one lazy marker and its early cleanup listener; this never invokes Broadcast.value. + */ + override def compute(split: Partition, context: TaskContext): Iterator[CometBroadcastInput] = { + val partition = split.asInstanceOf[CometBatchPartition] + Iterator.single( + new CometBroadcastInput( + partition.value, + schema, + name, + context, + CometBroadcastMemoryManager.getOrCreate(maxBytes))) + } +} diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala index 1d876dfb83f..6f1e6b469aa 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala @@ -171,12 +171,13 @@ object CometExecRDD { /** * Resolve the per-partition native input slots for `createPlan`, in scan-input order. A slot is - * either a `CometShuffleBlockIterator` (for slots in `shuffleScanIndices`, fed by a - * `CometShuffledBatchRDD` consumed via the JNI block-iteration protocol) or the single - * `ArrowArrayStream` exported by a non-shuffle `RDD[ArrowArrayStream]`. Returned alongside the - * subset that are shuffle-block iterators, which `CometExecIterator` needs to drive block - * iteration. Shared by [[CometExecRDD.compute]] and the native-shuffle path so both classify - * and resolve slots identically. + * either a `CometShuffleBlockIterator` (for slots in `shuffleScanIndices`), a lazy + * `CometBroadcastInput` with an executor memory owner, or the single `ArrowArrayStream` + * exported by another input. A broadcast without a supported owner opens an ordinary stream. + * Returned alongside the shuffle-block iterators that `CometExecIterator` must drive. Shared by + * [[CometExecRDD.compute]] and the native-shuffle path so both classify and resolve slots + * identically. Resolution may allocate streams and register their task cleanup listeners; it + * never opens the payload of a successfully admitted broadcast marker. */ def resolveInputObjects( inputRDDs: Seq[RDD[_]], @@ -199,6 +200,11 @@ object CometExecRDD { s"Slot $idx is marked as a shuffle scan but the input RDD is " + s"${other.getClass.getName}, expected CometShuffledBatchRDD") } + } else if (rdd.isInstanceOf[CometBroadcastInputRDD]) { + val input = rdd.iterator(part, context).next().asInstanceOf[CometBroadcastInput] + // Unsupported executor memory configuration must remain an ordinary Scan input. + if (input.getMemoryManager() == null) input.openStream().asInstanceOf[Object] + else input.asInstanceOf[Object] } else { val streams = rdd.iterator(part, context).asInstanceOf[Iterator[ArrowArrayStream]] if (!streams.hasNext) { diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometBroadcastArrowStream.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometBroadcastArrowStream.scala new file mode 100644 index 00000000000..afaabc6284b --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometBroadcastArrowStream.scala @@ -0,0 +1,113 @@ +/* + * 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.spark.sql.comet.execution.arrow + +/* + * Owns one Arrow C stream export for a task-local broadcast input. Each export gets an independent + * reader and child allocator, with cleanup owned by the parent input's task listener. Native + * move-import transfers the C release callback out of the Java wrapper; closing the wrapper + * also releases any export native code never claimed. + */ + +import org.apache.arrow.c.{ArrowArrayStream, Data} +import org.apache.arrow.memory.BufferAllocator +import org.apache.arrow.util.AutoCloseables +import org.apache.arrow.vector.ipc.ArrowReader +import org.apache.spark.sql.comet.util.Utils +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.vectorized.ColumnarBatch + +import org.apache.comet.CometArrowAllocator + +/** + * Explicit owner for a lazily exported broadcast stream's wrapper and allocator. Its parent + * marker installs the task listener before native execution, so opening this owner must not + * install a later listener that could close the export before the native reader drops it. + */ +private[comet] class CometBroadcastArrowStream private ( + val stream: ArrowArrayStream, + allocator: BufferAllocator) + extends AutoCloseable { + private var closed = false + + /** + * Release an export if native never claimed it, then free the C wrapper and allocator. Native + * move-import clears the release pointer, making release a no-op for already claimed streams. + * The parent calls this after native drops its reader and batch references; repeated calls do + * nothing. + */ + override def close(): Unit = synchronized { + if (!closed) { + closed = true + var failure: Throwable = null + try stream.release() + catch { case error: Throwable => failure = error } + if (failure == null) AutoCloseables.close(stream, allocator) + else AutoCloseables.close(failure, stream, allocator) + if (failure != null) throw failure + } + } +} + +private[comet] object CometBroadcastArrowStream { + + /** + * Reconcile the source's actual Arrow schema and export an owned stream without task listeners. + * On any failure close the source, reader, wrapper and allocator before rethrowing the original + * error. Source ownership transfers to the exported reader's release callback on success. + */ + def open( + source: Iterator[ColumnarBatch] with AutoCloseable, + schema: StructType, + name: String): CometBroadcastArrowStream = { + val allocator = CometArrowAllocator.newChildAllocator(name, 0, Long.MaxValue) + var reader: ArrowReader = null + var stream: ArrowArrayStream = null + try { + val expected = Utils.toArrowSchema(schema, CometArrowStream.NATIVE_TIMEZONE) + val (actual, reconciled) = CometArrowStream.reconcileStreamSchema(name, expected, source) + reader = new ColumnarBatchArrowReader(allocator, actual, reconciled) { + private var readerClosed = false + + /** Make rollback and an export release safe when both encounter the same reader. */ + override def close(): Unit = synchronized { + if (!readerClosed) { + readerClosed = true + super.close() + } + } + + /** Release a partial decoder without draining it when native abandons the stream. */ + override protected def closeReadSource(): Unit = source.close() + } + stream = ArrowArrayStream.allocateNew(allocator) + Data.exportArrayStream(allocator, reader, stream) + new CometBroadcastArrowStream(stream, allocator) + } catch { + case failure: Throwable => + if (stream != null) { + try stream.release() + catch { case cleanup: Throwable => failure.addSuppressed(cleanup) } + } + AutoCloseables.close(failure, stream, reader, source, allocator) + throw failure + } + } +} diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index 70f6ce18749..6eac80f047e 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -1042,6 +1042,11 @@ abstract class CometNativeExec extends CometExec { case _ => None } + /** + * Build an input RDD in native scan-slot order, aligning broadcast partitions to the probe. + * Supported opt-in broadcasts retain their actual materialization as lazy input markers; + * other inputs keep the existing Arrow or direct-shuffle transport. + */ def asArrowStreamRDD(plan: SparkPlan, partitionCount: Int, scanSlot: Int): RDD[_] = plan match { case s: CometNativeArrowSource => @@ -1049,11 +1054,21 @@ abstract class CometNativeExec extends CometExec { case _ => asBroadcastExchange(plan) match { case Some(c) => - CometArrowStream.wrapColumnarBatchRDD( - c.executeColumnar(partitionCount), - c.schema, - CometArrowStream.NATIVE_TIMEZONE, - c.nodeName) + val batches = c.executeColumnar(partitionCount) + if (CometConf.COMET_BROADCAST_REUSE_ENABLED.get(conf) && + CometBroadcastInput.supportsSchema(c.schema)) { + new CometBroadcastInputRDD( + batches, + c.schema, + CometConf.COMET_BROADCAST_REUSE_MAX_MEMORY.get(conf), + c.nodeName) + } else { + CometArrowStream.wrapColumnarBatchRDD( + batches, + c.schema, + CometArrowStream.NATIVE_TIMEZONE, + c.nodeName) + } case None if isShuffleScanInput(plan) && shuffleScanIndices.contains(scanSlot) => // Direct-read shuffle: `CometShuffledBatchRDD` reaches native via // CometShuffleBlockIterator. Other shuffle slots fall through and get wrapped. @@ -3176,8 +3191,27 @@ case class CometBroadcastHashJoinExec( left, right) + /** + * Expose reuse counters before native planning decides whether this broadcast join is eligible. + * Ineligible joins leave them at zero; an eligible join reports its preparations, hits, or + * ordinary-join fallbacks alongside the existing per-task join metrics. + */ override lazy val metrics: Map[String, SQLMetric] = { - val joinMetrics = CometMetricNode.joinMetrics(sparkContext) + val joinMetrics = CometMetricNode.joinMetrics(sparkContext) ++ Map( + "broadcast_build_prepare_time" -> + SQLMetrics.createNanoTimingMetric( + sparkContext, + "Time preparing a reusable broadcast build"), + "broadcast_build_prepare_rows" -> + SQLMetrics.createMetric(sparkContext, "Rows prepared for broadcast reuse"), + "broadcast_build_prepare_bytes" -> + SQLMetrics.createSizeMetric(sparkContext, "Bytes retained by prepared broadcast builds"), + "broadcast_build_cache_hits" -> + SQLMetrics.createMetric(sparkContext, "Broadcast builds reused"), + "broadcast_build_cache_misses" -> + SQLMetrics.createMetric(sparkContext, "Broadcast builds prepared for reuse"), + "broadcast_build_cache_fallbacks" -> + SQLMetrics.createMetric(sparkContext, "Broadcast builds using uncached fallback")) if (nativeOp.getHashJoin.getDynamicFilterEnabled) { joinMetrics ++ CometMetricNode.joinDynamicFilterMetrics(sparkContext) } else { diff --git a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala index eec16205a11..0a651f539bc 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -19,13 +19,15 @@ package org.apache.comet.exec +import java.util.concurrent.{CountDownLatch, TimeUnit} + import org.scalactic.source.Position import org.scalatest.Tag import org.apache.hadoop.fs.Path import org.apache.parquet.hadoop.ParquetFileReader import org.apache.parquet.hadoop.util.HadoopInputFile -import org.apache.spark.SparkException +import org.apache.spark.{CometBroadcastMemoryManager, SparkConf, SparkException} import org.apache.spark.sql.{CometTestBase, DataFrame, Row} import org.apache.spark.sql.catalyst.TableIdentifier import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation @@ -39,13 +41,16 @@ import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExcha import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ArrayType, IntegerType, MetadataBuilder, StructField, StructType} -import org.apache.comet.CometConf +import org.apache.comet.{CometConf, NativeBase} import org.apache.comet.CometSparkSessionExtensions.isSpark35Plus class CometJoinSuite extends CometTestBase { import testImplicits._ + override protected def sparkConf: SparkConf = + super.sparkConf.set("spark.plugins", "org.apache.spark.CometPlugin") + override protected def test(testName: String, testTags: Tag*)(testFun: => Any)(implicit pos: Position): Unit = { super.test(testName, testTags: _*) { @@ -1152,6 +1157,181 @@ class CometJoinSuite extends CometTestBase { } } + test("prepared broadcast reuse preserves duplicate and null joins across task inputs") { + val build = + Seq[(java.lang.Long, Long)]((1L, 11L), (1L, 12L), (2L, 21L), (4L, 41L), (null, 99L)) + val probe = (0 until 240).map { i => + (if (i % 7 == 0) null else java.lang.Long.valueOf(i % 6), i.toLong) + } + withParquetTable(build, "reuse_build") { + withParquetTable(probe, "reuse_probe") { + withSQLConf( + CometConf.COMET_BROADCAST_REUSE_ENABLED.key -> "true", + CometConf.COMET_BROADCAST_REUSE_MAX_MEMORY.key -> "1g", + CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "native", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "12") { + val (_, cometPlan) = checkSparkAnswerAndOperator( + sql(""" + |SELECT /*+ BROADCAST(b) */ p._2, b._2 + |FROM (SELECT /*+ REPARTITION(12, _1) */ * FROM reuse_probe) p + |JOIN reuse_build b ON p._1 = b._1 + |DISTRIBUTE BY p._2 + |""".stripMargin), + Seq(classOf[CometBroadcastExchangeExec], classOf[CometBroadcastHashJoinExec])) + val joins = collect(cometPlan) { case join: CometBroadcastHashJoinExec => join } + assert(joins.size == 1) + val join = joins.head + val misses = join.metrics("broadcast_build_cache_misses").value + assert(misses > 0L) + assert(join.metrics("broadcast_build_cache_fallbacks").value == 0L) + assert(join.metrics("broadcast_build_prepare_rows").value == 4L * misses) + // Finishing the join leaves no idle native storage charge for later tasks. + val owner = CometBroadcastMemoryManager.getOrCreate(1024L * 1024 * 1024) + assert(owner != null) + assert(owner.getUsedMemory == 0L) + } + } + } + } + + test("prepared broadcast build is shared by concurrent Spark probe tasks") { + CometJoinSuite.probeGate = new CountDownLatch(2) + try { + withParquetTable(Seq((1L, 11L)), "reuse_hit_build") { + val probe = (101L to 120L).map(value => (1L, value)) + withParquetTable(probe, "reuse_hit_probe") { + withSQLConf( + CometConf.COMET_BROADCAST_REUSE_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val query = sql(""" + |SELECT /*+ BROADCAST(b) */ p._2, b._2 + |FROM (SELECT /*+ REPARTITION(2) */ * FROM reuse_hit_probe) p + |JOIN reuse_hit_build b ON p._1 = b._1 + |""".stripMargin) + val plan = query.queryExecution.executedPlan + val joins = collect(plan) { case join: CometBroadcastHashJoinExec => join } + val broadcasts = collect(plan) { case b: CometBroadcastExchangeExec => b } + assert(joins.size == 1) + assert(broadcasts.size == 1) + // Hold each native probe open after its first output row until two tasks overlap. + val output = plan + .execute() + .mapPartitions { rows => + rows.map { row => + val result = (row.getLong(0), row.getLong(1)) + val gate = CometJoinSuite.probeGate + gate.countDown() + Predef.require(gate.await(30, TimeUnit.SECONDS), "probe tasks did not overlap") + result + } + } + .collect() + .toSet + assert(output == probe.map { case (_, value) => (value, 11L) }.toSet) + assert(joins.head.metrics("broadcast_build_cache_misses").value == 1L) + assert(joins.head.metrics("broadcast_build_cache_hits").value == 1L) + } + } + } + } finally { + CometJoinSuite.probeGate = null + } + } + + test("prepared broadcast reuse handles composite string keys, payload and residual") { + val build = Seq[(java.lang.Long, String, String, Long, String)]( + (1L, "M", "College", 11L, "IN"), + (1L, "M", "College", 12L, "OH"), + (2L, "S", "Advanced Degree", 21L, "NJ"), + (3L, "W", "2 yr Degree", 31L, null), + (null, "M", "College", 91L, "IN"), + (4L, null, "College", 92L, "LA")) + val probe = (0 until 360).map { i => + val (key, marital, education) = i % 6 match { + case 0 => (java.lang.Long.valueOf(1L), "M", "College") + case 1 => (java.lang.Long.valueOf(2L), "S", "Advanced Degree") + case 2 => (java.lang.Long.valueOf(3L), "W", "2 yr Degree") + case 3 => (java.lang.Long.valueOf(1L), "M", "different") + case 4 => (null, "M", "College") + case _ => (java.lang.Long.valueOf(4L), null, "College") + } + (key, marital, education, i.toLong) + } + withParquetTable(build, "reuse_q85_returning_build") { + withParquetTable(probe, "reuse_q85_returning_probe") { + withSQLConf( + CometConf.COMET_BROADCAST_REUSE_ENABLED.key -> "true", + CometConf.COMET_BROADCAST_REUSE_MAX_MEMORY.key -> "1g", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "12") { + val (_, plan) = checkSparkAnswerAndOperator( + sql(""" + |SELECT /*+ BROADCAST(b) */ p._4, b._4, b._2, b._3, b._5 + |FROM (SELECT /*+ REPARTITION(12) */ * FROM reuse_q85_returning_probe) p + |JOIN reuse_q85_returning_build b + |ON p._1 = b._1 AND p._2 = b._2 AND p._3 = b._3 + |AND (b._5 = 'IN' OR p._4 > b._4) + |""".stripMargin), + Seq(classOf[CometBroadcastExchangeExec], classOf[CometBroadcastHashJoinExec])) + val joins = collect(plan) { case join: CometBroadcastHashJoinExec => join } + assert(joins.size == 1) + val join = joins.head + assert(join.condition.nonEmpty) + val misses = join.metrics("broadcast_build_cache_misses").value + assert(misses > 0L) + assert(join.metrics("broadcast_build_cache_fallbacks").value == 0L) + assert(join.metrics("broadcast_build_prepare_rows").value == 4L * misses) + } + } + } + } + + test("prepared broadcast reuse falls back for UTF8 residual joins when admission is denied") { + val build = + Seq[(String, java.lang.Long)](("a", 10L), ("a", 20L), ("b", 30L), ("c", null), (null, 99L)) + val probe = (0 until 240).map { i => + (if (i % 5 == 0) null else Seq("a", "b", "c", "missing")(i % 4), i % 35, i) + } + withParquetTable(build, "reuse_denied_utf8_build") { + withParquetTable(probe, "reuse_denied_utf8_probe") { + withSQLConf( + CometConf.COMET_BROADCAST_REUSE_ENABLED.key -> "true", + CometConf.COMET_BROADCAST_REUSE_MAX_MEMORY.key -> "1g", + CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "native", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "12") { + NativeBase.clearBroadcastCache() + val owner = CometBroadcastMemoryManager.getOrCreate(1024L * 1024 * 1024) + assert(owner != null) + val held = owner.acquireMemory(owner.getLimit()) + try { + assert(held == owner.getLimit()) + val (_, plan) = checkSparkAnswerAndOperator( + sql(""" + |SELECT /*+ BROADCAST(b) */ p._3, b._1, b._2 + |FROM (SELECT /*+ REPARTITION(12, _1) */ * FROM reuse_denied_utf8_probe) p + |JOIN reuse_denied_utf8_build b ON p._1 = b._1 AND p._2 < b._2 + |DISTRIBUTE BY p._3 + |""".stripMargin), + Seq(classOf[CometBroadcastExchangeExec], classOf[CometBroadcastHashJoinExec])) + val joins = collect(plan) { case join: CometBroadcastHashJoinExec => join } + assert(joins.size == 1) + val join = joins.head + assert(join.condition.nonEmpty) + assert(join.metrics("broadcast_build_cache_fallbacks").value > 0L) + assert(join.metrics("broadcast_build_cache_hits").value == 0L) + assert(join.metrics("broadcast_build_cache_misses").value == 0L) + } finally { + owner.releaseMemory(held) + } + } + } + } + } + test("Broadcast hash join build-side batch coalescing") { // Use many shuffle partitions to produce many small broadcast batches, // then verify that coalescing reduces the build-side batch count to 1 per task. @@ -1824,3 +2004,7 @@ class CometJoinSuite extends CometTestBase { } } } + +object CometJoinSuite { + @volatile var probeGate: CountDownLatch = _ +} diff --git a/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala b/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala index c9b6d3de7fc..5159d8c445e 100644 --- a/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala +++ b/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala @@ -20,11 +20,12 @@ package org.apache.spark import java.util.Properties +import java.util.concurrent.{Callable, CountDownLatch, Executors, TimeUnit} import org.apache.logging.log4j.{Level, LogManager} import org.apache.logging.log4j.core.{LogEvent, LoggerContext} import org.apache.spark.executor.TaskMetrics -import org.apache.spark.memory.{MemoryConsumer, TaskMemoryManager, TestMemoryManager} +import org.apache.spark.memory.{MemoryConsumer, MemoryManager, TaskMemoryManager, TestMemoryManager, UnifiedMemoryManager} class CometTaskMemoryManagerSuite extends SparkFunSuite { @@ -187,4 +188,152 @@ class CometTaskMemoryManagerSuite extends SparkFunSuite { field.setAccessible(true) field.get(manager).asInstanceOf[MemoryConsumer] } + + /** + * Runs f under the supplied executor memory manager and task identity. Completion releases all + * task execution memory and unsets TaskContext even if f throws; executor storage is + * independent. + */ + private def withTaskContext[T](memoryManager: MemoryManager, taskAttemptId: Long)( + f: => T): T = { + val taskMemoryManager = new TaskMemoryManager(memoryManager, taskAttemptId) + val taskContext = new TaskContextImpl( + stageId = 0, + stageAttemptNumber = 0, + partitionId = 0, + numPartitions = 1, + taskAttemptId = taskAttemptId, + attemptNumber = 0, + taskMemoryManager = taskMemoryManager, + localProperties = new Properties, + metricsSystem = null, + taskMetrics = TaskMetrics.empty, + cpus = 1, + resources = Map.empty) + TaskContext.setTaskContext(taskContext) + try { + f + } finally { + try { + taskMemoryManager.cleanUpAllAllocatedMemory() + } finally { + TaskContext.unset() + } + } + } + + /** Runs f with Spark off-heap storage and an isolated executor owner. */ + private def withBroadcastMemory[T](limit: Long, sparkLimit: Long = 4096L)( + f: (AnyRef, MemoryManager, CometBroadcastMemoryManager) => T): T = { + val environment = new Object + val conf = new SparkConf(false) + .set("spark.memory.offHeap.enabled", "true") + .set("spark.memory.offHeap.size", sparkLimit.toString) + val memoryManager = UnifiedMemoryManager(conf, 1) + val owner = + CometBroadcastMemoryManager.getOrCreate(environment, memoryManager, true, limit) + assert(owner != null) + try { + f(environment, memoryManager, owner) + } finally { + CometBroadcastMemoryManager.shutdown() + } + } + + test("prepared broadcast storage survives actual task-memory cleanup") { + withBroadcastMemory(1024L) { (environment, memoryManager, owner) => + withTaskContext(memoryManager, 4500000L) { + val taskOwner = new CometTaskMemoryManager(1L, 4500000L) + assert(taskOwner.acquireMemory(128L) == 128L) + assert(owner.acquireMemory(256L) == 256L) + assert(memoryManager.offHeapExecutionMemoryUsed == 128L) + assert(memoryManager.offHeapStorageMemoryUsed == 256L) + } + + assert(TaskContext.get() == null) + assert(memoryManager.offHeapExecutionMemoryUsed == 0L) + assert(memoryManager.offHeapStorageMemoryUsed == 256L) + withTaskContext(memoryManager, 4500001L) { + val reused = + CometBroadcastMemoryManager.getOrCreate(environment, memoryManager, true, 1024L) + assert(reused eq owner) + assert( + CometBroadcastMemoryManager.getOrCreate(environment, memoryManager, true, 512L) == null) + } + owner.releaseMemory(256L) + assert(memoryManager.offHeapStorageMemoryUsed == 0L) + } + } + + test("prepared broadcast admission enforces its cap before changing storage accounting") { + withBroadcastMemory(256L) { (_, memoryManager, owner) => + assert(owner.acquireMemory(192L) == 192L) + assert(owner.acquireMemory(65L) == 0L) + assert(owner.getUsedMemory == 192L) + assert(memoryManager.offHeapStorageMemoryUsed == 192L) + assert(owner.acquireMemory(64L) == 64L) + assert(owner.getUsedMemory == 256L) + owner.releaseMemory(256L) + assert(memoryManager.offHeapStorageMemoryUsed == 0L) + } + } + + test("prepared broadcast admission returns no grant when Spark storage cannot fit it") { + withBroadcastMemory(8192L) { (_, memoryManager, owner) => + // This request fits the native cap but exceeds the entire Spark off-heap pool. Spark's + // real storage admission rejects it before attempting MemoryStore eviction. + assert(owner.acquireMemory(4097L) == 0L) + assert(owner.getUsedMemory == 0L) + assert(memoryManager.offHeapStorageMemoryUsed == 0L) + assert(owner.acquireMemory(256L) == 256L) + owner.releaseMemory(256L) + assert(memoryManager.offHeapStorageMemoryUsed == 0L) + } + } + + test("prepared broadcast concurrent grants share one executor cap without TaskContext") { + withBroadcastMemory(256L) { (_, memoryManager, owner) => + val executor = Executors.newFixedThreadPool(2) + val start = new CountDownLatch(1) + try { + val requests = (0 until 2).map { _ => + executor.submit(new Callable[Long] { + override def call(): Long = { + assert(start.await(10L, TimeUnit.SECONDS)) + assert(TaskContext.get() == null) + owner.acquireMemory(200L) + } + }) + } + start.countDown() + val grants = requests.map(_.get(10L, TimeUnit.SECONDS)) + assert(grants.sorted == Seq(0L, 200L)) + assert(memoryManager.offHeapStorageMemoryUsed == 200L) + owner.releaseMemory(200L) + assert(owner.getUsedMemory == 0L) + } finally { + start.countDown() + executor.shutdownNow() + assert(executor.awaitTermination(10L, TimeUnit.SECONDS)) + } + } + } + + test("prepared broadcast retirement returns storage and isolates late lease releases") { + withBroadcastMemory(256L) { (_, memoryManager, owner) => + assert(owner.acquireMemory(192L) == 192L) + CometBroadcastMemoryManager.shutdown() + assert(memoryManager.offHeapStorageMemoryUsed == 0L) + + withBroadcastMemory(256L) { (_, nextMemoryManager, nextOwner) => + assert(nextOwner ne owner) + assert(nextOwner.acquireMemory(128L) == 128L) + owner.releaseMemory(192L) + assert(nextOwner.getUsedMemory == 128L) + assert(nextMemoryManager.offHeapStorageMemoryUsed == 128L) + nextOwner.releaseMemory(128L) + } + } + } + } diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala index ef73522577a..0779dd45047 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala @@ -28,6 +28,7 @@ import scala.util.Using import org.scalatest.funsuite.AnyFunSuite import org.scalatest.matchers.should.Matchers +import org.apache.arrow.c.ArrowArrayStream import org.apache.arrow.memory.{AllocationListener, RootAllocator} import org.apache.arrow.vector.{BaseFixedWidthVector, BaseValueVector, BigIntVector, BitVector, DecimalVector, IntervalMonthDayNanoVector, IntVector, VarCharVector, VectorLoader, VectorSchemaRoot, VectorUnloader} import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector} @@ -38,11 +39,16 @@ import org.apache.arrow.vector.types.pojo.{ArrowType, DictionaryEncoding, Field, import org.apache.spark.rdd.InputFileBlockHolder import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow, SpecializedGetters} import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, GenericArrayData} +import org.apache.spark.TaskContext +import org.apache.spark.broadcast.Broadcast +import org.apache.spark.comet.CometTaskContextShim +import org.apache.spark.sql.comet.CometBroadcastInput import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution.vectorized.{ConstantColumnVector, Dictionary, OffHeapColumnVector, OnHeapColumnVector} import org.apache.spark.sql.types.{ArrayType, BooleanType, ByteType, CalendarIntervalType, DataType, DateType, DayTimeIntervalType, Decimal, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, ShortType, StringType, StructField, StructType, TimestampNTZType, TimestampType, YearMonthIntervalType} import org.apache.spark.sql.vectorized.{ColumnarArray, ColumnarBatch, ColumnarMap, ColumnVector} import org.apache.spark.unsafe.types.{CalendarInterval, UTF8String} +import org.apache.spark.util.io.ChunkedByteBuffer import org.apache.comet.vector.{CometPlainVector, CometVector, NativeUtil} @@ -54,6 +60,56 @@ import org.apache.comet.vector.{CometPlainVector, CometVector, NativeUtil} */ class CometArrowStreamSuite extends AnyFunSuite with Matchers { + private class TestBroadcast(id: Long, read: () => Array[ChunkedByteBuffer]) + extends Broadcast[Array[ChunkedByteBuffer]](id) { + override protected def getValue(): Array[ChunkedByteBuffer] = read() + override protected def doUnpersist(blocking: Boolean): Unit = () + override protected def doDestroy(blocking: Boolean): Unit = () + } + + test("broadcast marker stays lazy, replays streams, and closes after native cleanup") { + val task = TaskContext.empty() + val previous = TaskContext.empty() + var reads = 0 + val broadcast = new TestBroadcast( + 72L, + () => { + TaskContext.get() should be theSameInstanceAs task + reads += 1 + Array.empty[ChunkedByteBuffer] + }) + val schema = StructType(Seq(StructField("key", LongType))) + val unopened = new CometBroadcastInput(broadcast, schema, "test-broadcast", task, null) + val input = new CometBroadcastInput(broadcast, schema, "test-broadcast", task, null) + unopened.getBroadcastId() shouldBe 72L + reads shouldBe 0 + + val prior = TaskContext.get() + CometTaskContextShim.set(previous) + try { + var first: ArrowArrayStream = null + var second: ArrowArrayStream = null + // This models the native-plan listener installed after constructing the input but before + // openStream. Both exports must still exist when native cleanup begins. + task.addTaskCompletionListener[Unit] { _ => + first.snapshot().release should not be 0L + second.snapshot().release should not be 0L + } + first = input.openStream() + second = input.openStream() + first.memoryAddress() should not be second.memoryAddress() + reads shouldBe 2 + TaskContext.get() should be theSameInstanceAs previous + task.markTaskCompleted(None) + reads shouldBe 2 + intercept[NullPointerException](first.memoryAddress()) + } finally { + input.close() + unopened.close() + if (prior == null) CometTaskContextShim.unset() else CometTaskContextShim.set(prior) + } + } + private def expectedSchema(types: (String, ArrowType)*): Schema = { val fields = types.map { case (name, t) => new Field(name, new FieldType(true, t, null), java.util.Collections.emptyList[Field]()) From 7a37da235067ed8f1323b8279dd83d5a6c138b78 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sat, 19 Sep 2026 04:52:04 +0000 Subject: [PATCH 2/5] Simplify broadcast build reuse and consolidate coverage --- .../user-guide/latest/tuning/operators.md | 5 +- .../core/src/execution/operators/broadcast.rs | 226 ++---------------- .../execution/operators/broadcast_cache.rs | 160 +++++-------- .../spark/CometBroadcastMemoryManager.java | 16 +- .../spark/sql/comet/CometBroadcastInput.scala | 55 ++--- .../sql/comet/CometBroadcastInputRDD.scala | 13 +- .../arrow/CometBroadcastArrowStream.scala | 34 +-- .../spark/CometTaskMemoryManagerSuite.scala | 75 ++---- 8 files changed, 132 insertions(+), 452 deletions(-) diff --git a/docs/source/user-guide/latest/tuning/operators.md b/docs/source/user-guide/latest/tuning/operators.md index e00cbedccf2..98d2ec2efce 100644 --- a/docs/source/user-guide/latest/tuning/operators.md +++ b/docs/source/user-guide/latest/tuning/operators.md @@ -174,9 +174,8 @@ off-heap memory. The feature is disabled by default. The implementation covers non-spilling inner broadcast joins with direct column keys, matching key types, and fixed-width or plain UTF8 build columns. Residual join conditions remain local to each consuming task. Dictionary, view, binary, and nested build columns remain unsupported. -Unsupported joins keep ordinary task-local execution. Lookup is separated by executor generation -and uses the actual broadcast ID, build and probe schemas, and ordered keys. Cached builds omit -unused IN-list membership data; aliases of the same exchange can share a build. +Unsupported joins keep ordinary task-local execution. Tasks can share a build when they use the +same broadcast, build schema, and ordered build keys, even if their probe schemas differ. `spark.comet.broadcast.reuse.maxMemory` defaults to `1g`. It caps native prepared-build allocations through Spark off-heap storage memory while tasks are preparing or using a build. When the last diff --git a/native/core/src/execution/operators/broadcast.rs b/native/core/src/execution/operators/broadcast.rs index a491a6453b2..0457f181719 100644 --- a/native/core/src/execution/operators/broadcast.rs +++ b/native/core/src/execution/operators/broadcast.rs @@ -78,10 +78,6 @@ use crate::jvm_bridge::JVMClasses; type PreparedCache = BuildCache; type CacheGeneration = Option<(i64, Arc)>; -// DataFusion charges each retained batch before polling for the next one. -// Bound the extra charge held by the producer during that handoff. -const MAX_COPY_BATCH_BYTES: usize = 32 * 1024 * 1024; - // Lookup crosses Spark task/plan boundaries while builds are active. This // executor-lifetime registry holds only weak references, capped at 64 keys. // The registry fences SparkEnv generations; keys use actual Broadcast.id, not @@ -191,13 +187,6 @@ impl MemoryPool for BroadcastMemoryPool { if bytes == 0 { return Ok(()); } - // Reject a build allocation above its byte cap before asking Spark. - // Spark's owner still checks the total across concurrent loader pools. - if bytes > self.limit.saturating_sub(self.reserved()) { - return Err(DataFusionError::ResourcesExhausted( - "Broadcast build allocation exceeds its byte cap".into(), - )); - } if self.acquire(bytes)? == 0 { return Err(resources_datafusion_err!( "Broadcast cache storage admission denied" @@ -332,7 +321,6 @@ impl BroadcastInputExec { Box::pin(BroadcastStream { source: Arc::clone(self), reader: None, - pending_batch: None, copy_reservation, }) } @@ -388,13 +376,12 @@ impl ExecutionPlan for BroadcastInputExec { struct BroadcastStream { source: Arc, reader: Option, - pending_batch: Option, copy_reservation: Option, } impl BroadcastStream { - /// Releases the last temporary copy grant and reads or continues splitting - /// one batch. Cache preparation admits a conservative bound before copying + /// Releases the last temporary copy grant and reads the next batch. + /// Cache preparation admits a conservative bound before copying /// supported physical buffers. An ineligible encoding returns a resource /// error so the caller can reopen the broadcast for ordinary execution. fn next_batch(&mut self) -> Result> { @@ -404,31 +391,14 @@ impl BroadcastStream { if self.reader.is_none() { self.reader = Some(self.source.open()?); } - let already_split = self.pending_batch.is_some(); - let batch = match self.pending_batch.take() { - Some(batch) => batch, - None => { - let Some(batch) = self.reader.as_mut().unwrap().next() else { - return Ok(None); - }; - batch? - } + let Some(batch) = self.reader.as_mut().unwrap().next() else { + return Ok(None); }; + let batch = batch?; let schema = self.source.schema(); if batch.num_columns() != schema.fields().len() { return internal_err!("Broadcast column count mismatch"); } - let rows = if self.copy_reservation.is_some() { - copy_batch_rows(&batch, &schema, MAX_COPY_BATCH_BYTES, already_split)? - } else { - batch.num_rows() - }; - let batch = if rows < batch.num_rows() { - self.pending_batch = Some(batch.slice(rows, batch.num_rows() - rows)); - batch.slice(0, rows) - } else { - batch - }; let columns = if let Some(reservation) = &mut self.copy_reservation { copy_broadcast_columns(&batch, &schema, reservation)? } else { @@ -533,36 +503,6 @@ fn broadcast_copy_bytes(batch: &RecordBatch, schema: &Schema) -> Result { Ok(bound) } -/// Largest prefix whose independent native copy fits one handoff. -/// Keep small batches whole to avoid extra concatenation work. Once a large -/// batch is split, continue until all its rows are copied. DataFusion separately -/// admits the compact output while the copied input batches are still live. -fn copy_batch_rows( - batch: &RecordBatch, - schema: &Schema, - target: usize, - already_split: bool, -) -> Result { - if batch.num_rows() == 0 { - return Ok(0); - } - let bytes = broadcast_copy_bytes(batch, schema)?; - if bytes <= target || (!already_split && bytes <= 2 * target) { - return Ok(batch.num_rows()); - } - let mut fits = 1; - let mut too_large = batch.num_rows(); - while fits + 1 < too_large { - let middle = fits + (too_large - fits) / 2; - if broadcast_copy_bytes(&batch.slice(0, middle), schema)? <= target { - fits = middle; - } else { - too_large = middle; - } - } - Ok(fits) -} - /// Validate every column and admit all copies before allocating any of them. /// Count aliases separately: each output column owns an independent copy. fn copy_broadcast_columns( @@ -702,7 +642,6 @@ impl ExecutionPlan for CachedBroadcastJoinExec { let key = BuildKey { broadcast_id: source.broadcast_id, schema: source.schema(), - probe_schema: join.right().schema(), key_columns: join .on() .iter() @@ -751,7 +690,7 @@ impl ExecutionPlan for CachedBroadcastJoinExec { } let execution = Arc::new( join.builder() - .with_prepared_build(Arc::clone(&prepared))? + .with_prepared_build(Arc::clone(&prepared)) .build()?, ); (execution, Some(prepared)) @@ -868,15 +807,10 @@ pub(crate) fn reuse_broadcast_build_with_metrics( #[cfg(test)] mod tests { use super::*; - use arrow::array::{ - ArrayData, Decimal128Array, DictionaryArray, Int64Array, Int8Array, StringBuilder, - }; + use arrow::array::{ArrayData, DictionaryArray, Int8Array}; use arrow::buffer::Buffer; use arrow::datatypes::Int8Type; - use datafusion::common::utils::memory::get_record_batch_memory_size; use datafusion::execution::memory_pool::GreedyMemoryPool; - use datafusion::physical_plan::empty::EmptyExec; - use datafusion::physical_plan::joins::HashJoinExecBuilder; use std::ptr::NonNull; use std::sync::Weak; @@ -893,145 +827,6 @@ mod tests { Arc::new(GreedyMemoryPool::new(limit)) } - #[tokio::test] - async fn large_broadcast_build_admits_compaction_and_releases_storage() -> Result<()> { - let rows = 280; - let payload = "x".repeat(256 * 1024); - let mut strings = StringBuilder::with_capacity(rows, rows * payload.len()); - for _ in 0..rows { - strings.append_value(&payload); - } - let source = batch(vec![ - Arc::new(Int64Array::from_iter_values(0..rows as i64)), - Arc::new(strings.finish()), - ]); - let schema = source.schema(); - let source_bytes = get_record_batch_memory_size(&source); - assert!(source_bytes > 64 * 1024 * 1024); - let bound = broadcast_copy_bytes(&source, &schema)?; - let limit = bound + source_bytes / 2; - let medium = source.slice(0, rows / 2); - assert!(broadcast_copy_bytes(&medium, &schema)? > MAX_COPY_BATCH_BYTES); - assert_eq!( - copy_batch_rows(&medium, &schema, MAX_COPY_BATCH_BYTES, false)?, - medium.num_rows() - ); - let join = HashJoinExecBuilder::new( - Arc::new(EmptyExec::new(Arc::clone(&schema))), - Arc::new(EmptyExec::new(Arc::new(Schema::new(vec![Field::new( - "c0", - DataType::Int64, - true, - )])))), - vec![( - Arc::new(Column::new("c0", 0)), - Arc::new(Column::new("c0", 0)), - )], - JoinType::Inner, - ) - .with_partition_mode(PartitionMode::CollectLeft) - .build()?; - let copied_stream = |target, pool: &Arc| { - let source = source.clone(); - let schema = Arc::clone(&schema); - let reservation = MemoryConsumer::new("test broadcast copy").register(pool); - let stream = futures::stream::try_unfold( - (source, 0, reservation), - move |(source, offset, mut reservation)| async move { - reservation.free(); - if offset == source.num_rows() { - return Ok::<_, DataFusionError>(None); - } - let remainder = source.slice(offset, source.num_rows() - offset); - let rows = - copy_batch_rows(&remainder, &remainder.schema(), target, offset != 0)?; - let columns = copy_broadcast_columns( - &remainder.slice(0, rows), - &remainder.schema(), - &mut reservation, - )?; - let copied = RecordBatch::try_new(remainder.schema(), columns)?; - Ok(Some((copied, (source, offset + rows, reservation)))) - }, - ); - Box::pin(RecordBatchStreamAdapter::new(schema, stream)) as SendableRecordBatchStream - }; - let constrained = pool(limit); - let whole = join - .prepare_build( - copied_stream(usize::MAX, &constrained), - Arc::clone(&constrained), - Arc::new(Default::default()), - ) - .await; - assert!(matches!(whole, Err(DataFusionError::ResourcesExhausted(_)))); - assert_eq!(constrained.reserved(), 0); - - // Public DataFusion compacts the copied inputs into one batch. Its - // reservation must also admit that output while inputs remain alive. - // Chunking the JVM-to-native handoff does not remove this overlap. - let chunked = join - .prepare_build( - copied_stream(MAX_COPY_BATCH_BYTES, &constrained), - Arc::clone(&constrained), - Arc::new(Default::default()), - ) - .await; - assert!(matches!( - chunked, - Err(DataFusionError::ResourcesExhausted(_)) - )); - assert_eq!(constrained.reserved(), 0); - - let compact_pool = pool(3 * bound); - let prepared = join - .prepare_build( - copied_stream(MAX_COPY_BATCH_BYTES, &compact_pool), - Arc::clone(&compact_pool), - Arc::new(Default::default()), - ) - .await?; - assert_eq!(prepared.num_rows(), rows); - assert_eq!(compact_pool.reserved(), prepared.reserved_bytes()); - assert!(prepared.reserved_bytes() < 3 * bound); - drop(prepared); - assert_eq!(compact_pool.reserved(), 0); - Ok(()) - } - - #[test] - fn uncached_broadcast_replay_uses_safe_scan_casts_for_schema_drift() { - let decimal: ArrayRef = Arc::new( - Decimal128Array::from(vec![Some(12300), Some(24400), None, Some(-100)]) - .with_precision_and_scale(38, 2) - .unwrap(), - ); - let unchanged: ArrayRef = Arc::new(Int64Array::from(vec![7, 8, 9, 10])); - let input = batch(vec![decimal, Arc::clone(&unchanged)]); - let schema = Schema::new(vec![ - Field::new("small_value", DataType::Int8, true), - Field::new("unchanged", DataType::Int64, true), - ]); - - // The physical decimal cannot be copied into the declared Int8 cache - // layout. An admission miss must replay this same input without error. - let pool = pool(usize::MAX); - let mut reservation = MemoryConsumer::new("test schema drift").register(&pool); - assert!(matches!( - copy_broadcast_columns(&input, &schema, &mut reservation), - Err(DataFusionError::ResourcesExhausted(_)) - )); - assert_eq!(reservation.size(), 0); - - let columns = cast_uncached_broadcast_columns(&input, &schema).unwrap(); - let small = columns[0].as_any().downcast_ref::().unwrap(); - assert_eq!( - small, - &Int8Array::from(vec![Some(123), None, None, Some(-1)]) - ); - assert!(Arc::ptr_eq(&columns[1], &unchanged)); - } - /// Model unchecked FFI string buffers without ever reading them as &str. fn unchecked_utf8(values: &[u8], offsets: &[i32]) -> ArrayRef { let data = unsafe { @@ -1077,6 +872,13 @@ mod tests { DictionaryArray::::new(Int8Array::from(vec![Some(1), None, Some(0)]), values); let input = batch(vec![Arc::new(dictionary)]); let schema = Schema::new(vec![Field::new("decoded", DataType::Utf8, true)]); + let pool = pool(usize::MAX); + let mut reservation = MemoryConsumer::new("test dictionary replay").register(&pool); + assert!(matches!( + copy_broadcast_columns(&input, &schema, &mut reservation), + Err(DataFusionError::ResourcesExhausted(_)) + )); + assert_eq!(pool.reserved(), 0); let columns = cast_uncached_broadcast_columns(&input, &schema).unwrap(); let strings = columns[0].as_any().downcast_ref::().unwrap(); assert_eq!( diff --git a/native/core/src/execution/operators/broadcast_cache.rs b/native/core/src/execution/operators/broadcast_cache.rs index daac289c790..c496c221f21 100644 --- a/native/core/src/execution/operators/broadcast_cache.rs +++ b/native/core/src/execution/operators/broadcast_cache.rs @@ -23,54 +23,47 @@ //! Successful entries keep only weak references, so the last active probe can //! drop the prepared hash table. Its owner returns the Spark storage charge when //! the last lease drops or, if earlier, on executor retirement. A later task -//! wave may build the same table again. Resource failures briefly reject -//! matching loads to avoid repeated partial work; other errors and canceled -//! loaders leave the key retryable. Retiring a SparkEnv generation clears its -//! lookups without revoking leases already held by active probes. +//! wave may build the same table again. Resource failures reject waiting tasks +//! together; subsequent tasks can retry. Errors and canceled loaders leave the +//! key retryable. Retiring a SparkEnv generation clears its lookups without +//! revoking leases already held by active probes. use std::collections::HashMap; use std::future::Future; -use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Weak}; -use std::time::{Duration, Instant}; use arrow::datatypes::SchemaRef; use datafusion::common::{resources_datafusion_err, DataFusionError, Result}; use futures::channel::oneshot; use parking_lot::Mutex; -/// Share failed admission across a task wave without permanently disabling reuse -/// when executor storage pressure changes. Rejections occupy bounded cache slots. -const RESOURCE_RETRY_DELAY: Duration = Duration::from_secs(5); - /// Identifies a compatible materialization: the Spark broadcast, its declared -/// build and probe schemas, and ordered build key columns must all match. +/// build schema and ordered build key columns must all match. Probe columns do +/// not affect the build. Comet always uses NullEqualsNothing for physical keys. /// SparkEnv generation is scoped by the enclosing cache, not by this key. #[derive(Clone, Debug, Eq, Hash, PartialEq)] pub(super) struct BuildKey { pub broadcast_id: i64, pub schema: SchemaRef, - pub probe_schema: SchemaRef, pub key_columns: Vec, } -/// Per-key state: one loader with waiters, a weak active build, or a short-lived -/// resource rejection. No entry owns a prepared build after loading finishes. +/// A loader hands each waiter a lease or an admission rejection. Successful +/// entries keep only weak references after loading finishes. #[derive(Debug)] enum Entry { Loading(Vec>>>), Ready(Weak), - Rejected(Instant), } -/// Bounds keys, unfinished loads, weak ready entries and temporary rejections +/// Bounds unfinished loads and weak ready entries /// without retaining prepared builds. The payload byte cap belongs to the /// allocation pool; callers own strong references returned by `get_or_load`. #[derive(Debug)] pub(super) struct BuildCache { - entries: Mutex>>, + // None means retired. Retirement and publication use this same lock. + entries: Mutex>>>, max_entries: usize, - retired: AtomicBool, } impl BuildCache { @@ -78,16 +71,15 @@ impl BuildCache { /// payload memory or install any task/context ownership. pub fn new(max_entries: usize) -> Self { Self { - entries: Mutex::new(HashMap::new()), + entries: Mutex::new(Some(HashMap::new())), max_entries, - retired: AtomicBool::new(false), } } /// Returns an active build or elects one loader for this task wave. Waiting /// tasks receive strong references before the loader can drop its build. /// A canceled load wakes waiters to retry; a resource error from the loader - /// is shared briefly so peers take ordinary fallback instead of repeating + /// is shared with waiters so peers take ordinary fallback instead of repeating /// partial preparation. Other errors remain uncached. Entry-capacity failure /// occurs before `load` runs. The boolean is true on a hit or waited load. pub async fn get_or_load(&self, key: BuildKey, load: F) -> Result<(Arc, bool)> @@ -100,15 +92,12 @@ impl BuildCache { Wait(oneshot::Receiver>>), Load, } - let mut load = Some(load); loop { let action = { - let mut entries = self.entries.lock(); - if self.retired.load(Ordering::Relaxed) { - return Err(resources_datafusion_err!( - "Broadcast build cache is retired" - )); - } + let mut state = self.entries.lock(); + let entries = state + .as_mut() + .ok_or_else(|| resources_datafusion_err!("Broadcast build cache is retired"))?; match entries.get_mut(&key) { Some(Entry::Ready(value)) => match value.upgrade() { Some(value) => Action::Hit(value), @@ -119,17 +108,6 @@ impl BuildCache { continue; } }, - Some(Entry::Rejected(retry_at)) => { - if Instant::now() < *retry_at { - return Err(DataFusionError::ResourcesExhausted( - "Broadcast build cache admission rejected".into(), - )); - } - // Expired rejections contain no payload or task ownership. - // Re-enter election so only one caller retries admission. - entries.remove(&key); - continue; - } Some(Entry::Loading(waiters)) => { let (sender, receiver) = oneshot::channel(); waiters.push(sender); @@ -137,13 +115,10 @@ impl BuildCache { } None => { if entries.len() >= self.max_entries { - // The cap also counts bookkeeping for old waves; - // reclaim dead weak refs and expired rejections. - let now = Instant::now(); + // Reclaim metadata for completed task waves. entries.retain(|_, entry| match entry { Entry::Loading(_) => true, Entry::Ready(value) => value.strong_count() != 0, - Entry::Rejected(retry_at) => *retry_at > now, }); if entries.len() >= self.max_entries { return Err(resources_datafusion_err!( @@ -158,30 +133,29 @@ impl BuildCache { }; match action { Action::Hit(value) => return Ok((value, true)), - Action::Wait(receiver) => { - if let Ok(Some(value)) = receiver.await { - if !self.retired.load(Ordering::Relaxed) { - return Ok((value, true)); - } + Action::Wait(receiver) => match receiver.await { + Ok(Some(value)) => return Ok((value, true)), + Ok(None) => { + return Err(resources_datafusion_err!( + "Broadcast build cache admission rejected" + )); } - // Shared rejection or loader cancellation re-enters election. - } + // Cancellation or a data error lets another task load. + Err(_) => continue, + }, Action::Load => { let guard = LoadGuard { cache: self, key: key.clone(), active: true, }; - match load.take().expect("only the elected caller loads")().await { + match load().await { Ok(value) => { - guard.publish(Entry::Ready(Arc::downgrade(&value)), Some(&value)); + guard.publish(Some(&value)); return Ok((value, false)); } Err(DataFusionError::ResourcesExhausted(error)) => { - guard.publish( - Entry::Rejected(Instant::now() + RESOURCE_RETRY_DELAY), - None, - ); + guard.publish(None); return Err(DataFusionError::ResourcesExhausted(error)); } Err(error) => return Err(error), @@ -191,15 +165,10 @@ impl BuildCache { } } - /// Retires all lookup entries on native-runtime shutdown. Active leases keep - /// their values and charges; loading callers cannot republish after retirement. - /// Dropping senders wakes waiting tasks without retaining any task future. + /// Retires lookups without revoking active builds. Dropping senders wakes + /// waiting tasks; loading callers cannot republish after retirement. pub fn clear(&self) { - let entries = { - let mut entries = self.entries.lock(); - self.retired.store(true, Ordering::Relaxed); - std::mem::take(&mut *entries) - }; + let entries = self.entries.lock().take(); drop(entries); } } @@ -213,17 +182,19 @@ struct LoadGuard<'a, T> { } impl LoadGuard<'_, T> { - /// Publishes a weak lookup or rejection and hands the build to waiting tasks. + /// Publishes a weak lookup and hands the build or rejection to waiting tasks. /// Runtime shutdown may already have retired the slot; canceled loaders /// cannot publish into a newer environment. - fn publish(mut self, entry: Entry, value: Option<&Arc>) { + fn publish(mut self, value: Option<&Arc>) { let removed = { - let mut entries = self.cache.entries.lock(); - let removed = if matches!(entries.get(&self.key), Some(Entry::Loading(_))) { - entries.insert(self.key.clone(), entry) - } else { - None - }; + let mut state = self.cache.entries.lock(); + let removed = state.as_mut().and_then(|entries| { + let removed = entries.remove(&self.key); + if let Some(value) = value { + entries.insert(self.key.clone(), Entry::Ready(Arc::downgrade(value))); + } + removed + }); self.active = false; removed }; @@ -242,12 +213,8 @@ impl Drop for LoadGuard<'_, T> { return; } let removed = { - let mut entries = self.cache.entries.lock(); - if matches!(entries.get(&self.key), Some(Entry::Loading(_))) { - entries.remove(&self.key) - } else { - None - } + let mut state = self.cache.entries.lock(); + state.as_mut().and_then(|entries| entries.remove(&self.key)) }; drop(removed); } @@ -264,7 +231,6 @@ mod tests { BuildKey { broadcast_id: id, schema: Arc::new(Schema::empty()), - probe_schema: Arc::new(Schema::empty()), key_columns: vec![], } } @@ -326,7 +292,8 @@ mod tests { .get_or_load(key(3), || async { Ok(Arc::new(3)) }) .await .unwrap(); - let entries = cache.entries.lock(); + let state = cache.entries.lock(); + let entries = state.as_ref().unwrap(); assert_eq!(entries.len(), 2); assert!(!entries.contains_key(&key(1))); assert!(entries.contains_key(&key(2))); @@ -336,7 +303,7 @@ mod tests { } /// A failed loader rejects its waiters without repeating preparation. The - /// next task can load the same key after the short resource cooldown. + /// next task can retry when executor memory pressure changes. #[tokio::test] async fn admission_failure_is_shared_then_retried() { let cache = BuildCache::::new(1); @@ -361,17 +328,6 @@ mod tests { Err(DataFusionError::ResourcesExhausted(_)) )); assert_eq!(loads.load(Ordering::SeqCst), 1); - assert!(cache - .get_or_load(key(1), || async { panic!("cooldown must bypass loading") }) - .await - .is_err()); - { - let mut entries = cache.entries.lock(); - let Entry::Rejected(retry_at) = entries.get_mut(&key(1)).unwrap() else { - panic!("resource denial must install a rejection"); - }; - *retry_at = Instant::now(); - } let (value, hit) = cache .get_or_load(key(1), || async { loads.fetch_add(1, Ordering::SeqCst); @@ -417,21 +373,35 @@ mod tests { #[tokio::test] async fn leases_survive_retirement_and_bound_admission() { - let cache = BuildCache::new(1); + use futures::{poll, FutureExt}; + + let cache = BuildCache::new(2); let (lease, _) = cache .get_or_load(key(1), || async { Ok(Arc::new(9)) }) .await .unwrap(); + let (complete, loading) = oneshot::channel(); + let mut loader = Box::pin(cache.get_or_load(key(2), || async { + loading.await.unwrap(); + Ok(Arc::new(10)) + })); + assert!(poll!(&mut loader).is_pending()); + let mut waiter = + Box::pin(cache.get_or_load(key(2), || async { panic!("retired waiter cannot load") })); + assert!(poll!(&mut waiter).is_pending()); assert!(cache - .get_or_load(key(2), || async { panic!("must not load") }) + .get_or_load(key(3), || async { panic!("must not load") }) .await .is_err()); cache.clear(); assert_eq!(*lease, 9); + assert!(waiter.now_or_never().unwrap().is_err()); + complete.send(()).unwrap(); + assert_eq!(*loader.await.unwrap().0, 10); assert!(cache .get_or_load(key(2), || async { panic!("retired cache cannot load") }) .await .is_err()); - assert!(cache.entries.lock().is_empty()); + assert!(cache.entries.lock().is_none()); } } diff --git a/spark/src/main/java/org/apache/spark/CometBroadcastMemoryManager.java b/spark/src/main/java/org/apache/spark/CometBroadcastMemoryManager.java index 299ec61579c..0ff366c17bb 100644 --- a/spark/src/main/java/org/apache/spark/CometBroadcastMemoryManager.java +++ b/spark/src/main/java/org/apache/spark/CometBroadcastMemoryManager.java @@ -19,12 +19,6 @@ package org.apache.spark; -/* - * Bridges native broadcast build allocations to Spark's executor storage memory. Creating an - * owner reserves no storage; native loaders request off-heap grants only when they build. Active - * probe streams can keep those grants across task boundaries. Each SparkEnv has its own owner - * generation, so retiring an executor never charges a replacement executor for old leases. - */ import java.lang.ref.WeakReference; import org.apache.spark.memory.MemoryManager; @@ -33,9 +27,9 @@ import org.apache.spark.storage.BroadcastBlockId; /** - * Executor-owned storage accounting for prepared native broadcasts. This class lives in - * org.apache.spark to use Spark's storage memory API. One cap covers loading and leased builds; - * CometPlugin clears the native cache before retiring this owner at shutdown. + * Charges prepared broadcasts to Spark off-heap storage across task boundaries. One cap covers + * loading and active builds in a SparkEnv generation; retired owners never charge a replacement + * environment. Lives in org.apache.spark to access Spark's storage memory API. */ public final class CometBroadcastMemoryManager { private static long nextGeneration; @@ -136,17 +130,15 @@ public synchronized void releaseMemory(long size) { used -= size; } - /** Returns the outstanding native grants, including leases remaining after retirement. */ + /** Includes outstanding native leases after retirement has returned their Spark charge. */ public synchronized long getUsedMemory() { return used; } - /** Returns this owner's immutable generation, which distinguishes replacement SparkEnvs. */ public long getGeneration() { return generation; } - /** Returns the immutable byte cap shared by all prepared builds using this owner. */ public long getLimit() { return limit; } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastInput.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastInput.scala index 5b6677e00d0..e8b36f34d71 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastInput.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastInput.scala @@ -19,13 +19,6 @@ package org.apache.spark.sql.comet -/* - * Passes a Spark broadcast to native execution as a task-local handle, without decoding it on - * a cache hit. On a miss or ordinary-join fallback, native code can request fresh Arrow streams - * over the same broadcast. Each stream has its own decoder; task completion closes the exported - * wrappers and allocators after native execution drops their readers. - */ - import scala.collection.mutable.ArrayBuffer import org.apache.arrow.c.ArrowArrayStream @@ -40,10 +33,10 @@ import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.spark.util.io.ChunkedByteBuffer /** - * Task-local JNI input that identifies a real Spark broadcast without opening its payload. Native - * execution may open fresh streams on a cache miss or admission fallback. Each stream has - * independent decoder state; only Rust-owned copies may outlive this task. Register cleanup here, - * before CometExecIterator, so reverse task-listener order drops native readers first. + * Task-local broadcast identity with replayable Arrow streams opened only on cache misses or + * fallback. Register cleanup before CometExecIterator: Spark's reverse listener order must drop + * native readers before their JVM stream wrappers and allocators. Only native copies outlive + * tasks. */ class CometBroadcastInput private[comet] ( broadcast: Broadcast[Array[ChunkedByteBuffer]], @@ -58,17 +51,14 @@ class CometBroadcastInput private[comet] ( context.addTaskCompletionListener[Unit](_ => close()) - /** Return the actual materialization ID without evaluating Broadcast.value or decoding rows. */ def getBroadcastId(): Long = broadcast.id - /** Return the executor-lifetime owner, or null when the caller must use the uncached path. */ def getMemoryManager(): CometBroadcastMemoryManager = memoryManager /** - * Export a new stream over this broadcast, preserving the captured task context on JNI threads. - * The caller moves the C stream into its native reader; this object retains its wrapper and - * allocator until task completion. Failure closes all resources from this attempt and leaves - * earlier streams intact. Calls after close or task cancellation fail before reading data. + * Native move-import claims the C stream; retain its wrapper and allocator until task + * completion. Each call opens an independent decoder, so failed cache admission can replay the + * broadcast. */ def openStream(): ArrowArrayStream = synchronized { require(!closed, "Cannot open a completed broadcast input") @@ -80,10 +70,6 @@ class CometBroadcastInput private[comet] ( } } - /** - * Release unclaimed exports and all wrapper/allocator storage after native readers are dropped. - * Idempotent; attempts every stream even if one cleanup fails. It never reads unopened data. - */ override def close(): Unit = synchronized { if (!closed) { closed = true @@ -94,8 +80,8 @@ class CometBroadcastInput private[comet] ( } /** - * Native callbacks may arrive on threads without Spark's task-local context. Install the - * captured context for decoder work and cancellation checks, then restore the caller's context. + * JNI callbacks can run outside Spark's executor thread. Decode and check cancellation with the + * captured task, then restore the calling thread's context. */ private def withTaskContext[T](body: => T): T = { val previous = TaskContext.get() @@ -110,16 +96,14 @@ class CometBroadcastInput private[comet] ( } /** - * A fresh decoder cursor for one export. Marker construction and cache hits never read - * Broadcast.value; stream opening reconciles the schema and may decode the first chunk. An - * early drop closes the current IPC reader without draining remaining chunks. + * Opening a stream may decode its first chunk for schema reconciliation. Early close releases + * that IPC reader without opening the remaining chunks. */ private class BroadcastBatchIterator extends Iterator[ColumnarBatch] with AutoCloseable { private lazy val chunks = broadcast.value.iterator private var batches: Iterator[ColumnarBatch] = Iterator.empty private var stopped = false - /** Advance across empty/exhausted IPC chunks; no new chunk is opened after close. */ override def hasNext: Boolean = withTaskContext { if (stopped) false else { @@ -130,26 +114,19 @@ class CometBroadcastInput private[comet] ( } } - /** Return the next decoded batch, whose ownership passes to the Arrow stream reader. */ override def next(): ColumnarBatch = { if (!hasNext) throw new NoSuchElementException("No more broadcast batches") batches.next() } - /** Close the current IPC reader if present; empty chunks have no reader to release. */ - private def closeCurrent(): Unit = { - batches match { - case reader: ArrowReaderIterator => reader.close() - case _ => - } - batches = Iterator.empty - } - - /** Idempotently close the current reader without opening or draining remaining chunks. */ override def close(): Unit = { if (!stopped) { stopped = true - closeCurrent() + batches match { + case reader: ArrowReaderIterator => reader.close() + case _ => + } + batches = Iterator.empty } } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastInputRDD.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastInputRDD.scala index b1e3ae4e3a7..ab5418df401 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastInputRDD.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastInputRDD.scala @@ -19,20 +19,13 @@ package org.apache.spark.sql.comet -/* - * Aligns a lazy broadcast input with each probe partition. The driver captures the schema and - * memory cap; executor compute passes on Spark's actual Broadcast handle without reading its - * payload. Without a reuse owner, CometExecRDD opens an ordinary Arrow stream. With an owner, - * native execution chooses whether to reuse a build or open a stream for this task. - */ - import org.apache.spark.{CometBroadcastMemoryManager, OneToOneDependency, Partition, TaskContext} import org.apache.spark.rdd.RDD import org.apache.spark.sql.types.StructType import org.apache.spark.sql.vectorized.ColumnarBatch /** - * Broadcast input slots that preserve the actual Broadcast handle without constructing a decoder. + * Align broadcast handles with probe partitions without computing the parent's decoded batches. * The schema and memory cap are captured on the driver; executor SQLConf cannot alter admission. */ private[comet] class CometBroadcastInputRDD( @@ -42,13 +35,9 @@ private[comet] class CometBroadcastInputRDD( name: String) extends RDD[CometBroadcastInput](batches.context, Seq(new OneToOneDependency(batches))) { - /** Reuse the parent's aligned broadcast partitions without evaluating their payloads. */ override protected def getPartitions: Array[Partition] = firstParent[ColumnarBatch].partitions - /** - * Construct one lazy marker and its early cleanup listener; this never invokes Broadcast.value. - */ override def compute(split: Partition, context: TaskContext): Iterator[CometBroadcastInput] = { val partition = split.asInstanceOf[CometBatchPartition] Iterator.single( diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometBroadcastArrowStream.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometBroadcastArrowStream.scala index afaabc6284b..c08a260add2 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometBroadcastArrowStream.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometBroadcastArrowStream.scala @@ -19,13 +19,6 @@ package org.apache.spark.sql.comet.execution.arrow -/* - * Owns one Arrow C stream export for a task-local broadcast input. Each export gets an independent - * reader and child allocator, with cleanup owned by the parent input's task listener. Native - * move-import transfers the C release callback out of the Java wrapper; closing the wrapper - * also releases any export native code never claimed. - */ - import org.apache.arrow.c.{ArrowArrayStream, Data} import org.apache.arrow.memory.BufferAllocator import org.apache.arrow.util.AutoCloseables @@ -37,9 +30,8 @@ import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.comet.CometArrowAllocator /** - * Explicit owner for a lazily exported broadcast stream's wrapper and allocator. Its parent - * marker installs the task listener before native execution, so opening this owner must not - * install a later listener that could close the export before the native reader drops it. + * Owns an exported broadcast stream's wrapper and allocator. Cleanup belongs to the input's early + * task listener; registering another listener here would run before native releases its reader. */ private[comet] class CometBroadcastArrowStream private ( val stream: ArrowArrayStream, @@ -48,20 +40,19 @@ private[comet] class CometBroadcastArrowStream private ( private var closed = false /** - * Release an export if native never claimed it, then free the C wrapper and allocator. Native - * move-import clears the release pointer, making release a no-op for already claimed streams. - * The parent calls this after native drops its reader and batch references; repeated calls do - * nothing. + * Native move-import clears the release pointer. Release therefore cleans up only unclaimed + * exports; the parent has already dropped native readers before closing the wrapper/allocator. */ override def close(): Unit = synchronized { if (!closed) { closed = true - var failure: Throwable = null try stream.release() - catch { case error: Throwable => failure = error } - if (failure == null) AutoCloseables.close(stream, allocator) - else AutoCloseables.close(failure, stream, allocator) - if (failure != null) throw failure + catch { + case failure: Throwable => + AutoCloseables.close(failure, stream, allocator) + throw failure + } + AutoCloseables.close(stream, allocator) } } } @@ -69,9 +60,8 @@ private[comet] class CometBroadcastArrowStream private ( private[comet] object CometBroadcastArrowStream { /** - * Reconcile the source's actual Arrow schema and export an owned stream without task listeners. - * On any failure close the source, reader, wrapper and allocator before rethrowing the original - * error. Source ownership transfers to the exported reader's release callback on success. + * Transfer source ownership to the export's release callback, unwinding partial acquisition on + * failure. The physical schema must match the decoded vectors, including fallback casts. */ def open( source: Iterator[ColumnarBatch] with AutoCloseable, diff --git a/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala b/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala index 5159d8c445e..328310ede04 100644 --- a/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala +++ b/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala @@ -156,13 +156,18 @@ class CometTaskMemoryManagerSuite extends SparkFunSuite { private def withTaskMemoryManager(f: TaskMemoryManager => Unit): Unit = { val memoryManager = new TestMemoryManager(new SparkConf()) memoryManager.limit(1024) - val taskMemoryManager = new TaskMemoryManager(memoryManager, 0L) + withTaskMemoryManager(memoryManager, 0L)(f) + } + + private def withTaskMemoryManager(memoryManager: MemoryManager, taskAttemptId: Long)( + f: TaskMemoryManager => Unit): Unit = { + val taskMemoryManager = new TaskMemoryManager(memoryManager, taskAttemptId) val taskContext = new TaskContextImpl( stageId = 0, stageAttemptNumber = 0, partitionId = 0, numPartitions = 1, - taskAttemptId = 0L, + taskAttemptId = taskAttemptId, attemptNumber = 0, taskMemoryManager = taskMemoryManager, localProperties = new Properties, @@ -189,46 +194,13 @@ class CometTaskMemoryManagerSuite extends SparkFunSuite { field.get(manager).asInstanceOf[MemoryConsumer] } - /** - * Runs f under the supplied executor memory manager and task identity. Completion releases all - * task execution memory and unsets TaskContext even if f throws; executor storage is - * independent. - */ - private def withTaskContext[T](memoryManager: MemoryManager, taskAttemptId: Long)( - f: => T): T = { - val taskMemoryManager = new TaskMemoryManager(memoryManager, taskAttemptId) - val taskContext = new TaskContextImpl( - stageId = 0, - stageAttemptNumber = 0, - partitionId = 0, - numPartitions = 1, - taskAttemptId = taskAttemptId, - attemptNumber = 0, - taskMemoryManager = taskMemoryManager, - localProperties = new Properties, - metricsSystem = null, - taskMetrics = TaskMetrics.empty, - cpus = 1, - resources = Map.empty) - TaskContext.setTaskContext(taskContext) - try { - f - } finally { - try { - taskMemoryManager.cleanUpAllAllocatedMemory() - } finally { - TaskContext.unset() - } - } - } - /** Runs f with Spark off-heap storage and an isolated executor owner. */ - private def withBroadcastMemory[T](limit: Long, sparkLimit: Long = 4096L)( - f: (AnyRef, MemoryManager, CometBroadcastMemoryManager) => T): T = { + private def withBroadcastMemory(limit: Long)( + f: (AnyRef, MemoryManager, CometBroadcastMemoryManager) => Unit): Unit = { val environment = new Object val conf = new SparkConf(false) .set("spark.memory.offHeap.enabled", "true") - .set("spark.memory.offHeap.size", sparkLimit.toString) + .set("spark.memory.offHeap.size", "4096") val memoryManager = UnifiedMemoryManager(conf, 1) val owner = CometBroadcastMemoryManager.getOrCreate(environment, memoryManager, true, limit) @@ -242,7 +214,7 @@ class CometTaskMemoryManagerSuite extends SparkFunSuite { test("prepared broadcast storage survives actual task-memory cleanup") { withBroadcastMemory(1024L) { (environment, memoryManager, owner) => - withTaskContext(memoryManager, 4500000L) { + withTaskMemoryManager(memoryManager, 4500000L) { _ => val taskOwner = new CometTaskMemoryManager(1L, 4500000L) assert(taskOwner.acquireMemory(128L) == 128L) assert(owner.acquireMemory(256L) == 256L) @@ -253,7 +225,7 @@ class CometTaskMemoryManagerSuite extends SparkFunSuite { assert(TaskContext.get() == null) assert(memoryManager.offHeapExecutionMemoryUsed == 0L) assert(memoryManager.offHeapStorageMemoryUsed == 256L) - withTaskContext(memoryManager, 4500001L) { + withTaskMemoryManager(memoryManager, 4500001L) { _ => val reused = CometBroadcastMemoryManager.getOrCreate(environment, memoryManager, true, 1024L) assert(reused eq owner) @@ -265,19 +237,6 @@ class CometTaskMemoryManagerSuite extends SparkFunSuite { } } - test("prepared broadcast admission enforces its cap before changing storage accounting") { - withBroadcastMemory(256L) { (_, memoryManager, owner) => - assert(owner.acquireMemory(192L) == 192L) - assert(owner.acquireMemory(65L) == 0L) - assert(owner.getUsedMemory == 192L) - assert(memoryManager.offHeapStorageMemoryUsed == 192L) - assert(owner.acquireMemory(64L) == 64L) - assert(owner.getUsedMemory == 256L) - owner.releaseMemory(256L) - assert(memoryManager.offHeapStorageMemoryUsed == 0L) - } - } - test("prepared broadcast admission returns no grant when Spark storage cannot fit it") { withBroadcastMemory(8192L) { (_, memoryManager, owner) => // This request fits the native cap but exceeds the entire Spark off-heap pool. Spark's @@ -285,9 +244,6 @@ class CometTaskMemoryManagerSuite extends SparkFunSuite { assert(owner.acquireMemory(4097L) == 0L) assert(owner.getUsedMemory == 0L) assert(memoryManager.offHeapStorageMemoryUsed == 0L) - assert(owner.acquireMemory(256L) == 256L) - owner.releaseMemory(256L) - assert(memoryManager.offHeapStorageMemoryUsed == 0L) } } @@ -309,8 +265,13 @@ class CometTaskMemoryManagerSuite extends SparkFunSuite { val grants = requests.map(_.get(10L, TimeUnit.SECONDS)) assert(grants.sorted == Seq(0L, 200L)) assert(memoryManager.offHeapStorageMemoryUsed == 200L) - owner.releaseMemory(200L) + assert(owner.acquireMemory(56L) == 56L) + assert(owner.acquireMemory(1L) == 0L) + assert(owner.getUsedMemory == 256L) + assert(memoryManager.offHeapStorageMemoryUsed == 256L) + owner.releaseMemory(256L) assert(owner.getUsedMemory == 0L) + assert(memoryManager.offHeapStorageMemoryUsed == 0L) } finally { start.countDown() executor.shutdownNow() From ddd4ccde108b3823d8f732d2c97742ef58449dee Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sat, 19 Sep 2026 05:27:53 +0000 Subject: [PATCH 3/5] Keep lazy broadcast execution on the Spark task thread --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + .../user-guide/latest/tuning/operators.md | 1 - native/core/src/execution/jni_api.rs | 4 +- .../spark/sql/comet/CometBroadcastInput.scala | 31 ++----- .../CometBroadcastInputLifecycleSuite.scala | 86 +++++++++++++++++++ .../arrow/CometArrowStreamSuite.scala | 11 ++- 7 files changed, 103 insertions(+), 32 deletions(-) create mode 100644 spark/src/test/scala/org/apache/spark/sql/comet/CometBroadcastInputLifecycleSuite.scala diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 44cd5c56a9f..9a4b96a2642 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -575,6 +575,7 @@ jobs: org.apache.spark.sql.CometCollationSuite org.apache.spark.sql.CometSortCollationSuite org.apache.comet.CometFuzzAggregateSuite + org.apache.spark.sql.comet.CometBroadcastInputLifecycleSuite org.apache.spark.sql.comet.execution.arrow.CometArrowStreamSuite org.apache.spark.sql.comet.execution.arrow.CachedBatchRowIteratorSuite org.apache.spark.sql.comet.execution.arrow.CometArrowWriterSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index dbc76c05d46..adeb92a625d 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -279,6 +279,7 @@ jobs: org.apache.spark.sql.CometCollationSuite org.apache.spark.sql.CometSortCollationSuite org.apache.comet.CometFuzzAggregateSuite + org.apache.spark.sql.comet.CometBroadcastInputLifecycleSuite org.apache.spark.sql.comet.execution.arrow.CometArrowStreamSuite org.apache.spark.sql.comet.execution.arrow.CachedBatchRowIteratorSuite org.apache.spark.sql.comet.execution.arrow.CometArrowWriterSuite diff --git a/docs/source/user-guide/latest/tuning/operators.md b/docs/source/user-guide/latest/tuning/operators.md index 98d2ec2efce..5995952e1e6 100644 --- a/docs/source/user-guide/latest/tuning/operators.md +++ b/docs/source/user-guide/latest/tuning/operators.md @@ -197,4 +197,3 @@ rows/bytes, and preparation time. Probe metrics remain per task. A successful pr normally be followed by cache hits from other tasks using that broadcast. Validate elapsed query time and executor memory with the feature both enabled and disabled before enabling it broadly; reuse counts alone do not establish a query speedup. - diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 361d3847da7..4095a437c5b 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -1297,7 +1297,9 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( // so we should always execute partition 0. let stream = root_op.native_plan.execute(0, task_ctx)?; - if exec_context.scans.is_empty() && exec_context.shuffle_scans.is_empty() { + // Lazy broadcasts are JVM sources too, despite not appearing in `scans`. + // Keep their readers on the Spark task thread and drop them before JVM cleanup. + if exec_context.input_sources.is_empty() { // No JVM data sources — spawn onto tokio so the executor // thread parks in blocking_recv instead of busy-polling. exec_context.batch_producer = diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastInput.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastInput.scala index e8b36f34d71..ce230a21b51 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastInput.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastInput.scala @@ -25,7 +25,6 @@ import org.apache.arrow.c.ArrowArrayStream import org.apache.arrow.util.AutoCloseables import org.apache.spark.{CometBroadcastMemoryManager, TaskContext} import org.apache.spark.broadcast.Broadcast -import org.apache.spark.comet.CometTaskContextShim import org.apache.spark.sql.comet.execution.arrow.{ArrowReaderIterator, CometBroadcastArrowStream} import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.types.{BooleanType, ByteType, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, ShortType, StringType, StructType, TimestampNTZType, TimestampType} @@ -62,12 +61,11 @@ class CometBroadcastInput private[comet] ( */ def openStream(): ArrowArrayStream = synchronized { require(!closed, "Cannot open a completed broadcast input") - withTaskContext { - val source = new BroadcastBatchIterator - val stream = CometBroadcastArrowStream.open(source, schema, name) - streams += stream - stream.stream - } + context.killTaskIfInterrupted() + val source = new BroadcastBatchIterator + val stream = CometBroadcastArrowStream.open(source, schema, name) + streams += stream + stream.stream } override def close(): Unit = synchronized { @@ -79,22 +77,6 @@ class CometBroadcastInput private[comet] ( } } - /** - * JNI callbacks can run outside Spark's executor thread. Decode and check cancellation with the - * captured task, then restore the calling thread's context. - */ - private def withTaskContext[T](body: => T): T = { - val previous = TaskContext.get() - CometTaskContextShim.set(context) - try { - context.killTaskIfInterrupted() - body - } finally { - if (previous == null) CometTaskContextShim.unset() - else CometTaskContextShim.set(previous) - } - } - /** * Opening a stream may decode its first chunk for schema reconciliation. Early close releases * that IPC reader without opening the remaining chunks. @@ -104,7 +86,8 @@ class CometBroadcastInput private[comet] ( private var batches: Iterator[ColumnarBatch] = Iterator.empty private var stopped = false - override def hasNext: Boolean = withTaskContext { + override def hasNext: Boolean = { + context.killTaskIfInterrupted() if (stopped) false else { while (!batches.hasNext && chunks.hasNext) { diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/CometBroadcastInputLifecycleSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/CometBroadcastInputLifecycleSuite.scala new file mode 100644 index 00000000000..b11a65c8c4a --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/CometBroadcastInputLifecycleSuite.scala @@ -0,0 +1,86 @@ +/* + * 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.spark.sql.comet + +import java.util.concurrent.atomic.AtomicReference + +import org.apache.spark.{CometBroadcastMemoryManager, SparkConf, TaskContext} +import org.apache.spark.broadcast.Broadcast +import org.apache.spark.sql.CometTestBase +import org.apache.spark.sql.catalyst.expressions.PrettyAttribute +import org.apache.spark.sql.types.{LongType, StructField, StructType} +import org.apache.spark.util.io.ChunkedByteBuffer + +import org.apache.comet.CometExecIterator + +class CometBroadcastInputLifecycleSuite extends CometTestBase { + + override protected def sparkConf: SparkConf = + super.sparkConf.set("spark.plugins", "org.apache.spark.CometPlugin") + + test("lazy broadcast inputs are polled on the driving Spark task thread") { + val results = spark.sparkContext + .parallelize(Seq(0), 1) + .mapPartitions { _ => + val task = TaskContext.get() + val drivingThread = Thread.currentThread() + val readThread = new AtomicReference[Thread]() + val broadcast = new Broadcast[Array[ChunkedByteBuffer]](73L) { + override protected def getValue(): Array[ChunkedByteBuffer] = { + readThread.set(Thread.currentThread()) + Array.empty[ChunkedByteBuffer] + } + override protected def doUnpersist(blocking: Boolean): Unit = () + override protected def doDestroy(blocking: Boolean): Unit = () + } + val owner = CometBroadcastMemoryManager.getOrCreate(1024L * 1024 * 1024) + Predef.require(owner != null, "test requires a broadcast storage owner") + val input = new CometBroadcastInput( + broadcast, + StructType(Seq(StructField("key", LongType))), + "broadcast-lifecycle-test", + task, + owner) + val plan = + CometExecUtils.getLimitNativePlan(Seq(PrettyAttribute("key", LongType)), 1).get + val iterator = new CometExecIterator( + id = task.taskAttemptId(), + inputObjects = Array(input.asInstanceOf[Object]), + numOutputCols = 1, + protobufQueryPlan = plan.toByteArray, + nativeMetrics = CometMetricNode(Map.empty), + numParts = 1, + partitionIndex = 0) + try { + Predef.require(!iterator.hasNext, "the empty broadcast must produce no rows") + // A lazy broadcast is absent from the ordinary scan list. It must still use the + // synchronous JNI path so native readers close before task-owned JVM allocators. + Predef.require( + readThread.get() eq drivingThread, + "lazy broadcast was polled outside the driving Spark task thread") + } finally { + iterator.close() + } + Iterator.single(1) + } + .collect() + assert(results.toSeq == Seq(1)) + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala index 0779dd45047..8ce209fbbf2 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala @@ -36,12 +36,12 @@ import org.apache.arrow.vector.dictionary.{Dictionary => ArrowDictionary} import org.apache.arrow.vector.dictionary.DictionaryProvider.MapDictionaryProvider import org.apache.arrow.vector.ipc.ArrowReader import org.apache.arrow.vector.types.pojo.{ArrowType, DictionaryEncoding, Field, FieldType, Schema} -import org.apache.spark.rdd.InputFileBlockHolder -import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow, SpecializedGetters} -import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, GenericArrayData} import org.apache.spark.TaskContext import org.apache.spark.broadcast.Broadcast import org.apache.spark.comet.CometTaskContextShim +import org.apache.spark.rdd.InputFileBlockHolder +import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow, SpecializedGetters} +import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, GenericArrayData} import org.apache.spark.sql.comet.CometBroadcastInput import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution.vectorized.{ConstantColumnVector, Dictionary, OffHeapColumnVector, OnHeapColumnVector} @@ -69,7 +69,6 @@ class CometArrowStreamSuite extends AnyFunSuite with Matchers { test("broadcast marker stays lazy, replays streams, and closes after native cleanup") { val task = TaskContext.empty() - val previous = TaskContext.empty() var reads = 0 val broadcast = new TestBroadcast( 72L, @@ -85,7 +84,7 @@ class CometArrowStreamSuite extends AnyFunSuite with Matchers { reads shouldBe 0 val prior = TaskContext.get() - CometTaskContextShim.set(previous) + CometTaskContextShim.set(task) try { var first: ArrowArrayStream = null var second: ArrowArrayStream = null @@ -99,7 +98,7 @@ class CometArrowStreamSuite extends AnyFunSuite with Matchers { second = input.openStream() first.memoryAddress() should not be second.memoryAddress() reads shouldBe 2 - TaskContext.get() should be theSameInstanceAs previous + TaskContext.get() should be theSameInstanceAs task task.markTaskCompleted(None) reads shouldBe 2 intercept[NullPointerException](first.memoryAddress()) From 19f59d90bac70b946d2345a300da25c6a3e437d2 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sun, 4 Oct 2026 17:55:30 +0000 Subject: [PATCH 4/5] Clear strict Scala warnings in broadcast input and tests --- .../org/apache/spark/sql/comet/CometBroadcastInputRDD.scala | 4 ++-- .../src/test/scala/org/apache/comet/exec/CometJoinSuite.scala | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastInputRDD.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastInputRDD.scala index ab5418df401..94e01050632 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastInputRDD.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastInputRDD.scala @@ -32,7 +32,7 @@ private[comet] class CometBroadcastInputRDD( batches: CometBatchRDD, schema: StructType, maxBytes: Long, - name: String) + sourceName: String) extends RDD[CometBroadcastInput](batches.context, Seq(new OneToOneDependency(batches))) { override protected def getPartitions: Array[Partition] = @@ -44,7 +44,7 @@ private[comet] class CometBroadcastInputRDD( new CometBroadcastInput( partition.value, schema, - name, + sourceName, context, CometBroadcastMemoryManager.getOrCreate(maxBytes))) } diff --git a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala index 0a651f539bc..e9eabf0533f 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -1161,7 +1161,7 @@ class CometJoinSuite extends CometTestBase { val build = Seq[(java.lang.Long, Long)]((1L, 11L), (1L, 12L), (2L, 21L), (4L, 41L), (null, 99L)) val probe = (0 until 240).map { i => - (if (i % 7 == 0) null else java.lang.Long.valueOf(i % 6), i.toLong) + (if (i % 7 == 0) null else java.lang.Long.valueOf((i % 6).toLong), i.toLong) } withParquetTable(build, "reuse_build") { withParquetTable(probe, "reuse_probe") { From 2b94d426d107187d4ade50c3fd88d557396b3bf9 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Mon, 5 Oct 2026 21:19:30 +0000 Subject: [PATCH 5/5] build: pin Arrow 59 prepared-build companion for draft validation --- native/Cargo.lock | 96 ++++++++++++++++------------------------------- native/Cargo.toml | 37 ++++++++++++++++++ 2 files changed, 69 insertions(+), 64 deletions(-) diff --git a/native/Cargo.lock b/native/Cargo.lock index fe72674cdb8..18c40f3df14 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock @@ -1867,8 +1867,7 @@ dependencies = [ [[package]] name = "datafusion" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "14313430439ef858ed5f52ca9c840e9eea9844da34d3716d05a476449634da5d" +source = "git+https://github.com/sunchao/arrow-datafusion?rev=c9a142d834b6b47c2ac15524171bbf84f3022f73#c9a142d834b6b47c2ac15524171bbf84f3022f73" dependencies = [ "arrow", "arrow-schema", @@ -1916,8 +1915,7 @@ dependencies = [ [[package]] name = "datafusion-catalog" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2c0756ebc770bcebae366dfd273dad50935f6103b4d212253cb17e2e53c5427f" +source = "git+https://github.com/sunchao/arrow-datafusion?rev=c9a142d834b6b47c2ac15524171bbf84f3022f73#c9a142d834b6b47c2ac15524171bbf84f3022f73" dependencies = [ "arrow", "async-trait", @@ -1941,8 +1939,7 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bbde2961698705d8402f0f382f08876fc6f1126ffde790ef17833cf0ab87cd59" +source = "git+https://github.com/sunchao/arrow-datafusion?rev=c9a142d834b6b47c2ac15524171bbf84f3022f73#c9a142d834b6b47c2ac15524171bbf84f3022f73" dependencies = [ "arrow", "async-trait", @@ -2146,8 +2143,7 @@ dependencies = [ [[package]] name = "datafusion-common" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a6e06fe711ff815c8cf3143158547aea599702fd43f76f5866c70d5a5b3693a0" +source = "git+https://github.com/sunchao/arrow-datafusion?rev=c9a142d834b6b47c2ac15524171bbf84f3022f73#c9a142d834b6b47c2ac15524171bbf84f3022f73" dependencies = [ "arrow", "arrow-ipc", @@ -2173,8 +2169,7 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f7a5107acff4e3ed50a7732acec06bba40e74ab49df021397ed250d60f04e42a" +source = "git+https://github.com/sunchao/arrow-datafusion?rev=c9a142d834b6b47c2ac15524171bbf84f3022f73#c9a142d834b6b47c2ac15524171bbf84f3022f73" dependencies = [ "futures", "log", @@ -2184,8 +2179,7 @@ dependencies = [ [[package]] name = "datafusion-datasource" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f69f45d2d700785be8df9d4b44410e1c8e4a856f10cc728b418a1066c0b594fb" +source = "git+https://github.com/sunchao/arrow-datafusion?rev=c9a142d834b6b47c2ac15524171bbf84f3022f73#c9a142d834b6b47c2ac15524171bbf84f3022f73" dependencies = [ "arrow", "async-compression", @@ -2220,8 +2214,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "81650af64894ef34b46d2ef818ef244b40a81d90818f3bc642bf8fbf6758f8e0" +source = "git+https://github.com/sunchao/arrow-datafusion?rev=c9a142d834b6b47c2ac15524171bbf84f3022f73#c9a142d834b6b47c2ac15524171bbf84f3022f73" dependencies = [ "arrow", "arrow-ipc", @@ -2244,8 +2237,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b5e2519bed59a6907c79298b9c21fe807bbc34e4ff8414577d699b0741083593" +source = "git+https://github.com/sunchao/arrow-datafusion?rev=c9a142d834b6b47c2ac15524171bbf84f3022f73#c9a142d834b6b47c2ac15524171bbf84f3022f73" dependencies = [ "arrow", "async-trait", @@ -2267,8 +2259,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "635e61290e171597f6e50427e0cb486eda468ab47e880fa24d15bbb0fa24ac7e" +source = "git+https://github.com/sunchao/arrow-datafusion?rev=c9a142d834b6b47c2ac15524171bbf84f3022f73#c9a142d834b6b47c2ac15524171bbf84f3022f73" dependencies = [ "arrow", "async-trait", @@ -2290,8 +2281,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "07cd518f85b3a027d913047b4ab0811111a2b734fdf84a51b78234455f10d8fa" +source = "git+https://github.com/sunchao/arrow-datafusion?rev=c9a142d834b6b47c2ac15524171bbf84f3022f73#c9a142d834b6b47c2ac15524171bbf84f3022f73" dependencies = [ "arrow", "arrow-schema", @@ -2322,14 +2312,12 @@ dependencies = [ [[package]] name = "datafusion-doc" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d44c5919b560862dd1cb20a5c75b3e48c3d9e49719c91d78da5c38eaad5495c4" +source = "git+https://github.com/sunchao/arrow-datafusion?rev=c9a142d834b6b47c2ac15524171bbf84f3022f73#c9a142d834b6b47c2ac15524171bbf84f3022f73" [[package]] name = "datafusion-execution" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7620ad535af8730186bdec05d720e40d9a4584b40c7eb75f72c43aa7978e5484" +source = "git+https://github.com/sunchao/arrow-datafusion?rev=c9a142d834b6b47c2ac15524171bbf84f3022f73#c9a142d834b6b47c2ac15524171bbf84f3022f73" dependencies = [ "arrow", "arrow-buffer", @@ -2355,8 +2343,7 @@ dependencies = [ [[package]] name = "datafusion-expr" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ea229af8b73e566fc6ebb1c8dd42ab250f8855c5304bd173f4b50ead78b64304" +source = "git+https://github.com/sunchao/arrow-datafusion?rev=c9a142d834b6b47c2ac15524171bbf84f3022f73#c9a142d834b6b47c2ac15524171bbf84f3022f73" dependencies = [ "arrow", "arrow-schema", @@ -2377,8 +2364,7 @@ dependencies = [ [[package]] name = "datafusion-expr-common" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6a18baa840f77760dcc9eaad281329bc57536c5c39e65e8649596f2a51216be7" +source = "git+https://github.com/sunchao/arrow-datafusion?rev=c9a142d834b6b47c2ac15524171bbf84f3022f73#c9a142d834b6b47c2ac15524171bbf84f3022f73" dependencies = [ "arrow", "datafusion-common", @@ -2389,8 +2375,7 @@ dependencies = [ [[package]] name = "datafusion-functions" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "763bbfeec40b55a8a961d7349c0fcfcec4f3f66cc234d6295250b95371b85e34" +source = "git+https://github.com/sunchao/arrow-datafusion?rev=c9a142d834b6b47c2ac15524171bbf84f3022f73#c9a142d834b6b47c2ac15524171bbf84f3022f73" dependencies = [ "arrow", "arrow-buffer", @@ -2421,8 +2406,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "46c0c2b2e856883ecf6a4fa8b23e7c6041183af8a68d0583dc95d1d21e4eaee0" +source = "git+https://github.com/sunchao/arrow-datafusion?rev=c9a142d834b6b47c2ac15524171bbf84f3022f73#c9a142d834b6b47c2ac15524171bbf84f3022f73" dependencies = [ "arrow", "datafusion-common", @@ -2442,8 +2426,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e35e42c87dfbff06934ea0e99cfed3440ac614c3f3bfb8c33be35ca0587d4140" +source = "git+https://github.com/sunchao/arrow-datafusion?rev=c9a142d834b6b47c2ac15524171bbf84f3022f73#c9a142d834b6b47c2ac15524171bbf84f3022f73" dependencies = [ "arrow", "datafusion-common", @@ -2454,8 +2437,7 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08079b0604cdcc43c7d6cf006cfd5d3ed9c235ade0c1ebe91085a62f875603df" +source = "git+https://github.com/sunchao/arrow-datafusion?rev=c9a142d834b6b47c2ac15524171bbf84f3022f73#c9a142d834b6b47c2ac15524171bbf84f3022f73" dependencies = [ "arrow", "arrow-ord", @@ -2479,8 +2461,7 @@ dependencies = [ [[package]] name = "datafusion-functions-table" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fe2070ecf5638a93c81c8babc7858339ef52165a9f54da1411afd69adf7d85ce" +source = "git+https://github.com/sunchao/arrow-datafusion?rev=c9a142d834b6b47c2ac15524171bbf84f3022f73#c9a142d834b6b47c2ac15524171bbf84f3022f73" dependencies = [ "arrow", "async-trait", @@ -2495,8 +2476,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0e9a164ebe1998c5c3f9e15c9a0c4e6b599a0c8739b5471bfb089752441963f1" +source = "git+https://github.com/sunchao/arrow-datafusion?rev=c9a142d834b6b47c2ac15524171bbf84f3022f73#c9a142d834b6b47c2ac15524171bbf84f3022f73" dependencies = [ "arrow", "datafusion-common", @@ -2512,8 +2492,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "98965b183f92201c16798994820ae032649c4821e089c4c80e7d0f4485d41df5" +source = "git+https://github.com/sunchao/arrow-datafusion?rev=c9a142d834b6b47c2ac15524171bbf84f3022f73#c9a142d834b6b47c2ac15524171bbf84f3022f73" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -2522,8 +2501,7 @@ dependencies = [ [[package]] name = "datafusion-macros" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a711bd5a06ec85d9683e142d28a98ad0679a7220669f5c1b8589e57e9b29e46" +source = "git+https://github.com/sunchao/arrow-datafusion?rev=c9a142d834b6b47c2ac15524171bbf84f3022f73#c9a142d834b6b47c2ac15524171bbf84f3022f73" dependencies = [ "datafusion-doc", "quote", @@ -2533,8 +2511,7 @@ dependencies = [ [[package]] name = "datafusion-optimizer" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ff3256b62ff44616ec83c7f04741d21124970833bb377ed56d9abc5ec6f707ec" +source = "git+https://github.com/sunchao/arrow-datafusion?rev=c9a142d834b6b47c2ac15524171bbf84f3022f73#c9a142d834b6b47c2ac15524171bbf84f3022f73" dependencies = [ "arrow", "chrono", @@ -2552,8 +2529,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c1e37a25b5aa1f3a1db114f0ee3605bf01639b83a831dcde35119ee61f0b11d6" +source = "git+https://github.com/sunchao/arrow-datafusion?rev=c9a142d834b6b47c2ac15524171bbf84f3022f73#c9a142d834b6b47c2ac15524171bbf84f3022f73" dependencies = [ "arrow", "arrow-schema", @@ -2574,8 +2550,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a4d12787d42e88b26b0d3ec006192fa23ceb832ef2a9e87439ce820a8e038cc" +source = "git+https://github.com/sunchao/arrow-datafusion?rev=c9a142d834b6b47c2ac15524171bbf84f3022f73#c9a142d834b6b47c2ac15524171bbf84f3022f73" dependencies = [ "arrow", "datafusion-common", @@ -2589,8 +2564,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d13bf7f50fd902f5afc46d50759f5185b5620c40b24054f80f7c648897d3e5a6" +source = "git+https://github.com/sunchao/arrow-datafusion?rev=c9a142d834b6b47c2ac15524171bbf84f3022f73#c9a142d834b6b47c2ac15524171bbf84f3022f73" dependencies = [ "arrow", "chrono", @@ -2606,8 +2580,7 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "11fbe7e6665072957d44ed9587676e17a1f9ab3acfadfc287dac23750585f18a" +source = "git+https://github.com/sunchao/arrow-datafusion?rev=c9a142d834b6b47c2ac15524171bbf84f3022f73#c9a142d834b6b47c2ac15524171bbf84f3022f73" dependencies = [ "arrow", "datafusion-common", @@ -2625,8 +2598,7 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1265d58e5bce07d154e642a51ff43033b576a6ae40d50a29ac2b9f311004eb52" +source = "git+https://github.com/sunchao/arrow-datafusion?rev=c9a142d834b6b47c2ac15524171bbf84f3022f73#c9a142d834b6b47c2ac15524171bbf84f3022f73" dependencies = [ "arrow", "arrow-data", @@ -2660,8 +2632,7 @@ dependencies = [ [[package]] name = "datafusion-pruning" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "74b06b333405c05015836ed36cea02cc29d1e2ee8fb414a77532d8346fb56492" +source = "git+https://github.com/sunchao/arrow-datafusion?rev=c9a142d834b6b47c2ac15524171bbf84f3022f73#c9a142d834b6b47c2ac15524171bbf84f3022f73" dependencies = [ "arrow", "datafusion-common", @@ -2676,8 +2647,7 @@ dependencies = [ [[package]] name = "datafusion-session" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e7616b17b1d2fee3e9b13c651eac91b768d153265e6ba88e06e0b4af7883f69c" +source = "git+https://github.com/sunchao/arrow-datafusion?rev=c9a142d834b6b47c2ac15524171bbf84f3022f73#c9a142d834b6b47c2ac15524171bbf84f3022f73" dependencies = [ "arrow-schema", "async-trait", @@ -2691,8 +2661,7 @@ dependencies = [ [[package]] name = "datafusion-spark" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f7809ecefb16fcb950c6b22ee3df99cb412dca6fae72c329d64c938ae53254f9" +source = "git+https://github.com/sunchao/arrow-datafusion?rev=c9a142d834b6b47c2ac15524171bbf84f3022f73#c9a142d834b6b47c2ac15524171bbf84f3022f73" dependencies = [ "arrow", "bigdecimal", @@ -2721,8 +2690,7 @@ dependencies = [ [[package]] name = "datafusion-sql" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5b4848be66dd008d55adbe1dbc8cbd6a2502588ad0c0b8483bab0f43fd4495d0" +source = "git+https://github.com/sunchao/arrow-datafusion?rev=c9a142d834b6b47c2ac15524171bbf84f3022f73#c9a142d834b6b47c2ac15524171bbf84f3022f73" dependencies = [ "arrow", "bigdecimal", diff --git a/native/Cargo.toml b/native/Cargo.toml index fa6ea26a217..1c4f0a0de7a 100644 --- a/native/Cargo.toml +++ b/native/Cargo.toml @@ -100,3 +100,40 @@ codegen-units = 16 # Parallel codegen (faster compile, slightly larger binary) debug-assertions = true panic = "unwind" # Allow panics to be caught and logged across FFI boundary # overflow-checks inherited as false from release + +# Temporary Arrow 59-compatible backport of apache/datafusion#25491 onto DataFusion 55.1. +# Patch the complete crate family so path dependencies share the same Rust types. +# Remove this pin when Comet upgrades to an official release containing prepared builds. +[patch.crates-io] +datafusion = { git = "https://github.com/sunchao/arrow-datafusion", rev = "c9a142d834b6b47c2ac15524171bbf84f3022f73" } +datafusion-catalog = { git = "https://github.com/sunchao/arrow-datafusion", rev = "c9a142d834b6b47c2ac15524171bbf84f3022f73" } +datafusion-catalog-listing = { git = "https://github.com/sunchao/arrow-datafusion", rev = "c9a142d834b6b47c2ac15524171bbf84f3022f73" } +datafusion-common = { git = "https://github.com/sunchao/arrow-datafusion", rev = "c9a142d834b6b47c2ac15524171bbf84f3022f73" } +datafusion-common-runtime = { git = "https://github.com/sunchao/arrow-datafusion", rev = "c9a142d834b6b47c2ac15524171bbf84f3022f73" } +datafusion-datasource = { git = "https://github.com/sunchao/arrow-datafusion", rev = "c9a142d834b6b47c2ac15524171bbf84f3022f73" } +datafusion-datasource-arrow = { git = "https://github.com/sunchao/arrow-datafusion", rev = "c9a142d834b6b47c2ac15524171bbf84f3022f73" } +datafusion-datasource-csv = { git = "https://github.com/sunchao/arrow-datafusion", rev = "c9a142d834b6b47c2ac15524171bbf84f3022f73" } +datafusion-datasource-json = { git = "https://github.com/sunchao/arrow-datafusion", rev = "c9a142d834b6b47c2ac15524171bbf84f3022f73" } +datafusion-datasource-parquet = { git = "https://github.com/sunchao/arrow-datafusion", rev = "c9a142d834b6b47c2ac15524171bbf84f3022f73" } +datafusion-doc = { git = "https://github.com/sunchao/arrow-datafusion", rev = "c9a142d834b6b47c2ac15524171bbf84f3022f73" } +datafusion-execution = { git = "https://github.com/sunchao/arrow-datafusion", rev = "c9a142d834b6b47c2ac15524171bbf84f3022f73" } +datafusion-expr = { git = "https://github.com/sunchao/arrow-datafusion", rev = "c9a142d834b6b47c2ac15524171bbf84f3022f73" } +datafusion-expr-common = { git = "https://github.com/sunchao/arrow-datafusion", rev = "c9a142d834b6b47c2ac15524171bbf84f3022f73" } +datafusion-functions = { git = "https://github.com/sunchao/arrow-datafusion", rev = "c9a142d834b6b47c2ac15524171bbf84f3022f73" } +datafusion-functions-aggregate = { git = "https://github.com/sunchao/arrow-datafusion", rev = "c9a142d834b6b47c2ac15524171bbf84f3022f73" } +datafusion-functions-aggregate-common = { git = "https://github.com/sunchao/arrow-datafusion", rev = "c9a142d834b6b47c2ac15524171bbf84f3022f73" } +datafusion-functions-nested = { git = "https://github.com/sunchao/arrow-datafusion", rev = "c9a142d834b6b47c2ac15524171bbf84f3022f73" } +datafusion-functions-table = { git = "https://github.com/sunchao/arrow-datafusion", rev = "c9a142d834b6b47c2ac15524171bbf84f3022f73" } +datafusion-functions-window = { git = "https://github.com/sunchao/arrow-datafusion", rev = "c9a142d834b6b47c2ac15524171bbf84f3022f73" } +datafusion-functions-window-common = { git = "https://github.com/sunchao/arrow-datafusion", rev = "c9a142d834b6b47c2ac15524171bbf84f3022f73" } +datafusion-macros = { git = "https://github.com/sunchao/arrow-datafusion", rev = "c9a142d834b6b47c2ac15524171bbf84f3022f73" } +datafusion-optimizer = { git = "https://github.com/sunchao/arrow-datafusion", rev = "c9a142d834b6b47c2ac15524171bbf84f3022f73" } +datafusion-physical-expr = { git = "https://github.com/sunchao/arrow-datafusion", rev = "c9a142d834b6b47c2ac15524171bbf84f3022f73" } +datafusion-physical-expr-adapter = { git = "https://github.com/sunchao/arrow-datafusion", rev = "c9a142d834b6b47c2ac15524171bbf84f3022f73" } +datafusion-physical-expr-common = { git = "https://github.com/sunchao/arrow-datafusion", rev = "c9a142d834b6b47c2ac15524171bbf84f3022f73" } +datafusion-physical-optimizer = { git = "https://github.com/sunchao/arrow-datafusion", rev = "c9a142d834b6b47c2ac15524171bbf84f3022f73" } +datafusion-physical-plan = { git = "https://github.com/sunchao/arrow-datafusion", rev = "c9a142d834b6b47c2ac15524171bbf84f3022f73" } +datafusion-pruning = { git = "https://github.com/sunchao/arrow-datafusion", rev = "c9a142d834b6b47c2ac15524171bbf84f3022f73" } +datafusion-session = { git = "https://github.com/sunchao/arrow-datafusion", rev = "c9a142d834b6b47c2ac15524171bbf84f3022f73" } +datafusion-spark = { git = "https://github.com/sunchao/arrow-datafusion", rev = "c9a142d834b6b47c2ac15524171bbf84f3022f73" } +datafusion-sql = { git = "https://github.com/sunchao/arrow-datafusion", rev = "c9a142d834b6b47c2ac15524171bbf84f3022f73" }