From 9e77f9f9bb5b9ad45db24a18f976f7bd50cc54ec Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sat, 19 Sep 2026 04:25:42 +0000 Subject: [PATCH 1/7] 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/7] 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/7] 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/7] 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/7] 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" } From adb9d08750f2c60ca7eb44b0d7c4b636d379c082 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Tue, 29 Sep 2026 22:34:17 +0000 Subject: [PATCH 6/7] perf: carry completed join filters into lazy Union branches --- docs/source/contributor-guide/ffi.md | 20 + .../user-guide/latest/tuning/operators.md | 33 +- native/core/src/execution/jni_api.rs | 108 +++- .../core/src/execution/operators/broadcast.rs | 39 +- .../operators/dynamic_filter/join.rs | 80 ++- .../operators/dynamic_filter/join/tests.rs | 65 +++ .../execution/operators/dynamic_filter/mod.rs | 65 ++- .../dynamic_filter/parquet_reader.rs | 66 ++- native/core/src/execution/operators/mod.rs | 4 + .../operators/runtime_filter_projection.rs | 211 +++++++ native/core/src/execution/operators/scan.rs | 31 ++ .../src/execution/operators/union_filter.rs | 524 +++++++++++++++++ native/core/src/execution/planner.rs | 94 +++- native/proto/src/proto/operator.proto | 3 + .../scala/org/apache/comet/CometConf.scala | 13 + .../org/apache/comet/CometExecIterator.scala | 21 +- .../main/scala/org/apache/comet/Native.scala | 11 + .../spark/sql/comet/CometBroadcastInput.scala | 2 +- .../apache/spark/sql/comet/CometExecRDD.scala | 14 +- .../spark/sql/comet/CometMetricNode.scala | 35 +- .../spark/sql/comet/CometUnionInput.scala | 155 ++++++ .../arrow/CometNativeArrowSource.scala | 11 +- .../apache/spark/sql/comet/operators.scala | 29 + .../apache/comet/exec/CometJoinSuite.scala | 527 +++++++++++++++++- 24 files changed, 2100 insertions(+), 61 deletions(-) create mode 100644 native/core/src/execution/operators/runtime_filter_projection.rs create mode 100644 native/core/src/execution/operators/union_filter.rs create mode 100644 spark/src/main/scala/org/apache/spark/sql/comet/CometUnionInput.scala diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index 7027080bbf3..e55ad5ed0c9 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -260,6 +260,26 @@ resident until the JVM closes it. See [Crossing the FFI boundary](memory_managem | --------- | -------------------------------- | ---------------------------------------------------------- | | All cases | Native allocates, JVM references | JVM must call `close()` to trigger native release callback | +## Lazy Union Inputs + +An opt-in `CometUnionInput` is a task-scoped lazy iterator marker. Its native scan sets a +request flag, and the driving JNI thread opens the Spark iterator only after the join build +has published its completed filter. Native worker threads never open Spark Union iterators. + +`openStream` borrows an address containing an `Arc`. The JVM marker retains +its own owned handle before that callback returns. A bundle contains completed predicates, +leases on their accounted prepared builds, the task attempt, and the permitted native branch +root IDs. Each branch registers its handle before native planning; registration checks both +identities. The handle is released exactly once by the marker's close operation. + +Branch EOF does not close its native execution while the enclosing Union can retain its Arrow +buffers. EOF remains sticky, and the enclosing owner holds the registered cleanup. After the +outer outputs and native context drop, native teardown releases the C stream and invokes the +marker's cleanup on the driving thread. That cleanup closes children in reverse order and +releases the filter leases. Only then does native teardown wait for its memory reservations. +A failed cleanup still allows other owners to close; pending Java exceptions are captured and +cleared between callbacks and the first failure is rethrown after cleanup. + ## Further Reading - [Arrow C Data Interface Specification](https://arrow.apache.org/docs/format/CDataInterface.html) diff --git a/docs/source/user-guide/latest/tuning/operators.md b/docs/source/user-guide/latest/tuning/operators.md index 5995952e1e6..4993b7cd77c 100644 --- a/docs/source/user-guide/latest/tuning/operators.md +++ b/docs/source/user-guide/latest/tuning/operators.md @@ -54,11 +54,40 @@ adaptations cannot fail. Nested column pruning still reads only the requested st with supplied file statistics also skip reader attachment. These cases still use runtime filtering on decoded batches. -Filters stay within the task's native plan and do not propagate across Spark exchanges or JVM/Arrow -boundaries. A shuffled hash join can still filter probe batches after shuffle, but it cannot send +By default, filters stay within the task's native plan and do not propagate across Spark exchanges +or JVM/Arrow boundaries. The separate Union setting below enables a scoped exception. A shuffled hash join can still filter probe batches after shuffle, but it cannot send its filter back to an earlier scan stage. Compare the [runtime-filter and scan metrics](../metrics.md#hash-joins) with the setting disabled to distinguish reduced hash-probe work from reader I/O savings. +### Runtime Filters Across UNION ALL + +With both `spark.comet.exec.join.dynamicFilter.enabled=true` and +`spark.comet.exec.join.dynamicFilter.union.enabled=true`, an eligible broadcast inner join +can send its completed key filter into the native readers of a `UNION ALL` probe input. +Both settings default to false. + +For example, when a small set of selected account IDs joins the union of current and archived +transactions, each transaction reader can discard row groups outside those IDs. The original +join still checks every match and preserves duplicates. Spark retains its original Union +partitions and its single broadcast exchange; Comet opens each branch lazily after the join's +build is ready. + +This path accepts one direct signed integer join key and no join residual. Build columns must +be fixed-width scalars or plain UTF-8 strings; other build schemas keep ordinary execution. +Branches can rename or reorder columns and retain direct-column null checks. Computed projections, other residual +filters, limits, and intermediate joins stop reader propagation. Such branches can still be +filtered after their native output is produced. Existing per-file schema-conversion safeguards +remain in force. Missing or incomplete filters never delay opening a branch. + +A filter holds a lease on the same prepared build that the original join probes. This keeps its +storage charged until every consumer finishes, including early termination. Task-attempt and +native-root checks prevent a filter from reaching unrelated executions. Nested Union owners +release child plans and Arrow streams before releasing the leases. + +This draft depends on the prepared-build support in [PR #6037](https://github.com/apache/datafusion-comet/pull/6037) +and a compatible DataFusion release containing its prepared-build API. Runtime Union filtering +does not require enabling executor-wide broadcast reuse. + ## Adaptive Partial Aggregation Set `spark.comet.exec.aggregate.skipPartial.enabled=true` to let Comet bypass partial hash diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 4095a437c5b..709758be5d6 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -517,9 +517,14 @@ fn collect_op_names<'a>(op: &'a Operator, names: &mut std::collections::BTreeSet } /// Comet native execution context. Kept alive across JNI calls. +use super::operators::union_filter::{UnionFilterBundle, UnionFilterDomain}; + struct ExecutionContext { /// The id of the execution context. pub id: i64, + task_attempt_id: i64, + union_filter: Option>, + forwarded_union_domains: Vec>, /// The deserialized Spark plan pub spark_plan: Operator, /// The number of partitions @@ -747,6 +752,9 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( let exec_context = Box::new(ExecutionContext { id, + task_attempt_id, + union_filter: None, + forwarded_union_domains: vec![], spark_plan, partition_count: partition_count as usize, root_op: None, @@ -780,6 +788,70 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( }) } +/// Retain a Union domain bundle for a JVM stream owner. Handles are owned boxes, not registry keys. +/// # Safety +/// `handle` must address a live Arc supplied by openStream or retain. +#[no_mangle] +pub unsafe extern "system" fn Java_org_apache_comet_Native_retainUnionRuntimeFilter( + e: EnvUnowned, + _class: JClass, + handle: jlong, +) -> jlong { + try_unwrap_or_throw(&e, |_env| { + if handle == 0 { + return Err(CometError::NullPointer( + "Union runtime filter is null".into(), + )); + } + let bundle = &*(handle as *const Arc); + Ok(Box::into_raw(Box::new(Arc::clone(bundle))) as jlong) + }) +} + +/// # Safety +/// `handle` must be an owned handle returned by retain, released exactly once. +#[no_mangle] +pub unsafe extern "system" fn Java_org_apache_comet_Native_releaseUnionRuntimeFilter( + e: EnvUnowned, + _class: JClass, + handle: jlong, +) { + try_unwrap_or_throw(&e, |_env| { + if handle != 0 { + drop(Box::from_raw(handle as *mut Arc)); + } + Ok::<_, CometError>(()) + }) +} + +/// Attach completed domains before a child native pipeline starts. +/// # Safety +/// Both addresses must be live; the caller owns their lifetimes through this call. +#[no_mangle] +pub unsafe extern "system" fn Java_org_apache_comet_Native_setUnionRuntimeFilter( + e: EnvUnowned, + _class: JClass, + plan: jlong, + handle: jlong, +) { + try_unwrap_or_throw(&e, |_env| { + if plan == 0 || handle == 0 { + return Err(CometError::NullPointer( + "Union runtime filter context is null".into(), + )); + } + let context = get_execution_context(plan); + let bundle = &*(handle as *const Arc); + // Invalid or late registration is an optional-filter miss, never a query failure. + if context.root_op.is_none() + && bundle.accepts(context.spark_plan.plan_id, context.task_attempt_id) + { + context.union_filter = Some(Arc::clone(bundle)); + } + Ok::<_, CometError>(()) + }) +} + /// Binds one task-owned shuffle callback before native execution is initialized. /// /// Keeping callback registration separate preserves the existing `createPlan` JNI ABI for @@ -1280,6 +1352,27 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( &mut exec_context.input_sources.clone(), exec_context.partition_count, )?; + // A separate Union branch is planned without its parent join. Prepare the + // same metric-preserving projection/filter wrappers as an in-pipeline probe, + // preserving Spark metric ownership before the reader walker runs. + let root_op = if let Some(bundle) = exec_context + .union_filter + .as_ref() + .filter(|bundle| !bundle.domains.is_empty()) + { + let root_op = PhysicalPlanner::prepare_union_branch_for_runtime_reader(root_op); + let mut root = root_op.as_ref().clone(); + let attachment = super::operators::apply_union_filter_bundle( + Arc::clone(&root.native_plan), + bundle, + exec_context.session_ctx.copied_config().options(), + )?; + root.native_plan = attachment.plan; + exec_context.forwarded_union_domains = attachment.forwarded_domains; + Arc::new(root) + } else { + root_op + }; let physical_plan_time = start.elapsed(); exec_context.plan_creation_time += physical_plan_time; @@ -1431,7 +1524,18 @@ pub extern "system" fn Java_org_apache_comet_Native_releasePlan( // back what they hold the next time they yield, so wait for them. let id = execution_context.id; let plan_memory = Arc::clone(&execution_context.plan_memory); + let union_inputs = execution_context + .scans + .iter() + .filter_map(|scan| scan.union_input.clone()) + .collect::>(); drop(execution_context); + let mut union_inputs_closed = Ok(()); + for input in union_inputs { + // Try every owner even if an earlier callback failed. Branch plans and bundle + // leases must close before waiting for this plan's reservations to return. + union_inputs_closed = union_inputs_closed.and(input.close()); + } let held = plan_memory.wait_until_released(Instant::now() + PLAN_MEMORY_RELEASE_TIMEOUT); if held > 0 { warn!( @@ -1442,7 +1546,9 @@ pub extern "system" fn Java_org_apache_comet_Native_releasePlan( ); } - producer_stopped.and(metrics_flushed) + producer_stopped + .and(metrics_flushed) + .and(union_inputs_closed) }) } diff --git a/native/core/src/execution/operators/broadcast.rs b/native/core/src/execution/operators/broadcast.rs index 0457f181719..92747023977 100644 --- a/native/core/src/execution/operators/broadcast.rs +++ b/native/core/src/execution/operators/broadcast.rs @@ -465,7 +465,7 @@ fn utf8_copy_bytes(rows: usize, values: usize) -> Result { /// 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 { +pub(super) 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() { @@ -505,7 +505,7 @@ fn broadcast_copy_bytes(batch: &RecordBatch, schema: &Schema) -> Result { /// 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( +pub(super) fn copy_broadcast_columns( batch: &RecordBatch, schema: &Schema, reservation: &mut MemoryReservation, @@ -567,6 +567,7 @@ struct CachedBroadcastJoinExec { source: Arc, cache: Arc, metrics: ExecutionPlanMetricsSet, + union_targets: Vec, } impl DisplayAs for CachedBroadcastJoinExec { @@ -627,6 +628,7 @@ impl ExecutionPlan for CachedBroadcastJoinExec { let source = Arc::clone(&self.source); let cache = Arc::clone(&self.cache); let metrics = self.metrics.clone(); + let union_targets = self.union_targets.clone(); let hits = MetricBuilder::new(&self.metrics).counter("broadcast_build_cache_hits", partition); let misses = @@ -703,6 +705,25 @@ impl ExecutionPlan for CachedBroadcastJoinExec { } Err(error) => return Err(error), }; + // A failed cache admission keeps the task-local transport path sound. + if lease.is_none() && !union_targets.is_empty() { + return super::union_filter::prepare_union_join( + &join, + partition, + context, + union_targets, + metrics, + ); + } + let domains = lease + .as_ref() + .map(|build| { + union_targets + .iter() + .map(|target| target.publish(Arc::clone(build))) + .collect::>() + }) + .unwrap_or_default(); let result = execution.execute(partition, context); for metric in execution.metrics().unwrap_or_default().iter() { metrics.register(Arc::clone(metric)); @@ -712,6 +733,7 @@ impl ExecutionPlan for CachedBroadcastJoinExec { Some(prepared) => Ok(Box::pin(PreparedProbeStream { stream, _lease: prepared, + _domains: domains, }) as SendableRecordBatchStream), None => Ok(stream), } @@ -727,6 +749,7 @@ impl ExecutionPlan for CachedBroadcastJoinExec { struct PreparedProbeStream { stream: SendableRecordBatchStream, _lease: Arc, + _domains: Vec>, } impl Stream for PreparedProbeStream { @@ -750,18 +773,23 @@ impl RecordBatchStream for PreparedProbeStream { pub(crate) fn reuse_broadcast_build( plan: Arc, ) -> Result> { - reuse_broadcast_build_with_metrics(plan, ExecutionPlanMetricsSet::new()) + reuse_broadcast_build_with_transport(plan, ExecutionPlanMetricsSet::new(), vec![]) } /// 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( +pub(crate) fn reuse_broadcast_build_with_transport( plan: Arc, metrics: ExecutionPlanMetricsSet, + union_targets: Vec, ) -> Result> { if let Some(projection) = plan.downcast_ref::() { - let child = reuse_broadcast_build_with_metrics(Arc::clone(projection.input()), metrics)?; + let child = reuse_broadcast_build_with_transport( + Arc::clone(projection.input()), + metrics, + union_targets, + )?; return if Arc::ptr_eq(&child, projection.input()) { Ok(plan) } else { @@ -801,6 +829,7 @@ pub(crate) fn reuse_broadcast_build_with_metrics( cache: cache(source.generation), source, metrics, + union_targets, })) } diff --git a/native/core/src/execution/operators/dynamic_filter/join.rs b/native/core/src/execution/operators/dynamic_filter/join.rs index 9c147980108..60580ddfb65 100644 --- a/native/core/src/execution/operators/dynamic_filter/join.rs +++ b/native/core/src/execution/operators/dynamic_filter/join.rs @@ -45,7 +45,10 @@ use datafusion::physical_plan::{ }; use futures::StreamExt; -use super::super::broadcast::reuse_broadcast_build_with_metrics; +use super::super::broadcast::reuse_broadcast_build_with_transport; +use super::super::union_filter::{ + prepare_union_join, transport_build_supported, UnionFilterTarget, +}; use super::parquet_reader::try_attach_parquet_reader_filter; use super::DynamicFilterExec; @@ -58,6 +61,7 @@ pub(crate) struct DynamicFilterJoinExec { template: HashJoinExec, config: ConfigOptions, metrics: ExecutionPlanMetricsSet, + union_transport: bool, } /// Per-execution join state. The permanent plan keeps no live filter; this value @@ -65,6 +69,7 @@ pub(crate) struct DynamicFilterJoinExec { struct RuntimeDynamicFilterJoin { join: HashJoinExec, reader_filter_attached: bool, + union_targets: Vec, } impl DynamicFilterJoinExec { @@ -82,20 +87,37 @@ impl DynamicFilterJoinExec { template: join.builder().reset_state().build()?, config, metrics: ExecutionPlanMetricsSet::new(), + union_transport: false, }) } + pub(crate) fn with_union_transport(mut self, enabled: bool) -> Self { + self.union_transport = enabled; + self + } + fn build_runtime_join(&self) -> Result { let predicate = Arc::new(DynamicFilterPhysicalExpr::new( vec![Arc::clone(&self.template.on()[0].1)], lit(true), )); - let reader = try_attach_parquet_reader_filter( - self.template.right(), - Arc::clone(&predicate), - &self.config, - )?; - let reader_filter_attached = reader.is_some(); + let mut union_targets = vec![]; + let reader = + if self.union_transport && transport_build_supported(&self.template.left().schema()) { + super::parquet_reader::try_attach_parquet_reader_filter_with_transport( + self.template.right(), + Arc::clone(&predicate), + &self.config, + &mut union_targets, + )? + } else { + try_attach_parquet_reader_filter( + self.template.right(), + Arc::clone(&predicate), + &self.config, + )? + }; + let reader_filter_attached = reader.is_some() && union_targets.is_empty(); let consumer = Arc::new(DynamicFilterExec::new( reader.unwrap_or_else(|| Arc::clone(self.template.right())), Arc::clone(&predicate), @@ -114,14 +136,26 @@ impl DynamicFilterJoinExec { Ok(RuntimeDynamicFilterJoin { join, reader_filter_attached, + union_targets, }) } + #[cfg(test)] fn execute_runtime_join( &self, join: HashJoinExec, partition: usize, context: Arc, + ) -> Result { + self.execute_runtime_join_with_transport(join, partition, context, vec![]) + } + + fn execute_runtime_join_with_transport( + &self, + join: HashJoinExec, + partition: usize, + context: Arc, + targets: Vec, ) -> Result { // DataFusion can materialize one IN-list literal per build row, // despite admitting the list by packed-array bytes and distinct-key count. @@ -149,9 +183,22 @@ impl DynamicFilterJoinExec { 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); + let execution = reuse_broadcast_build_with_transport( + Arc::clone(&join), + self.metrics.clone(), + targets.clone(), + )?; + let result = if Arc::ptr_eq(&execution, &join) && !targets.is_empty() { + prepare_union_join( + join.downcast_ref::().unwrap(), + partition, + context, + targets, + self.metrics.clone(), + ) + } else { + 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() { @@ -233,13 +280,15 @@ impl ExecutionPlan for DynamicFilterJoinExec { .with_new_children(children)? .build()?; match Self::try_new(&join, &self.config)? { - Some(wrapper) => Ok(Arc::new(wrapper)), + Some(wrapper) => Ok(Arc::new(wrapper.with_union_transport(self.union_transport))), None => Ok(Arc::new(join)), } } fn reset_state(self: Arc) -> Result> { - Ok(Arc::new(Self::new(&self.template, self.config.clone())?)) + let mut reset = Self::new(&self.template, self.config.clone())?; + reset.union_transport = self.union_transport; + Ok(Arc::new(reset)) } fn execute( @@ -256,7 +305,12 @@ impl ExecutionPlan for DynamicFilterJoinExec { MetricBuilder::new(&self.metrics) .counter(attachment_metric, partition) .add(1); - self.execute_runtime_join(runtime.join, partition, context) + self.execute_runtime_join_with_transport( + runtime.join, + partition, + context, + runtime.union_targets, + ) } fn metrics(&self) -> Option { diff --git a/native/core/src/execution/operators/dynamic_filter/join/tests.rs b/native/core/src/execution/operators/dynamic_filter/join/tests.rs index 14242b4e4fd..c7fc86e5d24 100644 --- a/native/core/src/execution/operators/dynamic_filter/join/tests.rs +++ b/native/core/src/execution/operators/dynamic_filter/join/tests.rs @@ -1564,3 +1564,68 @@ fn wrapper_preserves_join_statistics_and_distribution() { ); } } + +#[tokio::test] +async fn union_branch_preparation_keeps_renamed_reader_metrics() { + use crate::execution::spark_plan::SparkPlan; + let session = Arc::new(SessionContext::new()); + let (_file, scan) = parquet_probe((0..4).collect(), &session, 1); + let scan_plan = Arc::new(SparkPlan::new(1, scan as _, vec![])); + let filter = Arc::new( + FilterExec::try_new( + Arc::new(IsNotNullExpr::new(Arc::new(Column::new("key", 0)))), + Arc::clone(&scan_plan.native_plan), + ) + .unwrap(), + ); + let filter_plan = Arc::new(SparkPlan::new(2, filter, vec![scan_plan])); + let projection = Arc::new( + ProjectionExec::try_new( + vec![( + Arc::new(Column::new("key", 0)) as Arc, + "renamed".to_owned(), + )], + Arc::clone(&filter_plan.native_plan), + ) + .unwrap(), + ); + let plan = Arc::new(SparkPlan::new(3, projection, vec![filter_plan])); + let prepared = PhysicalPlanner::prepare_union_branch_for_runtime_reader(plan); + let predicate = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::new(Column::new("renamed", 0))], + lit(true), + )); + predicate + .update(Arc::new(BinaryExpr::new( + Arc::new(Column::new("renamed", 0)), + Operator::Eq, + lit(3_i32), + ))) + .unwrap(); + predicate.mark_complete(); + let mut targets = vec![]; + let attached = super::super::parquet_reader::try_attach_parquet_reader_filter_with_transport( + &prepared.native_plan, + predicate, + session.copied_config().options(), + &mut targets, + ) + .unwrap() + .unwrap(); + assert!(targets.is_empty()); + let batches = collect(attached, session.task_ctx()).await.unwrap(); + assert_eq!(row_count(&batches), 1); + assert_eq!(batches[0].schema().field(0).name(), "renamed"); + assert_eq!( + prepared.native_plan.metrics().unwrap().output_rows(), + Some(1) + ); + assert_eq!( + prepared.children[0] + .native_plan + .metrics() + .unwrap() + .output_rows(), + Some(1) + ); +} diff --git a/native/core/src/execution/operators/dynamic_filter/mod.rs b/native/core/src/execution/operators/dynamic_filter/mod.rs index 3dc33ce15d4..83750d219e6 100644 --- a/native/core/src/execution/operators/dynamic_filter/mod.rs +++ b/native/core/src/execution/operators/dynamic_filter/mod.rs @@ -209,9 +209,72 @@ impl ExecutionPlan for DynamicFilterExec { } fn metrics(&self) -> Option { - Some(self.metrics.clone_inner()) + let metrics = ExecutionPlanMetricsSet::new(); + for metric in self.metrics.clone_inner().iter() { + metrics.register(Arc::clone(metric)); + } + if self.metric_prefix == "dynamic_filter_union" { + // Reader and join metrics can be registered after execute() starts. Forward live + // handles here instead of snapshotting an unexecuted branch's empty metric set. + for metric in self.input.metrics().unwrap_or_default().iter() { + metrics.register(Arc::clone(metric)); + } + } + Some(metrics.clone_inner()) } } #[cfg(test)] mod tests; + +pub(crate) struct AppliedUnionFilters { + pub plan: Arc, + pub forwarded_domains: Vec>, +} + +/// Attach transported domains only to this authorized native pipeline. Hold the bundle in its +/// ExecutionContext until both the stream and rewritten plan have been released. +pub(crate) fn apply_union_filter_bundle( + input: Arc, + bundle: &super::union_filter::UnionFilterBundle, + config: &datafusion::common::config::ConfigOptions, +) -> Result { + use super::union_filter::remap_domain; + let mut input = input; + let mut nested = vec![]; + for domain in &bundle.domains { + let metrics = ExecutionPlanMetricsSet::new(); + let predicate = match remap_domain(domain, &input.schema()) { + Ok(predicate) => predicate, + Err(_) => continue, + }; + let mut targets = vec![]; + let attachment = parquet_reader::try_attach_parquet_reader_filter_with_transport( + &input, + Arc::clone(&predicate), + config, + &mut targets, + )?; + let outcome = if attachment.is_none() { + "dynamic_filter_union_branches_skipped" + } else if targets.is_empty() { + "dynamic_filter_union_branches_attached" + } else { + "dynamic_filter_union_branches_forwarded" + }; + MetricBuilder::new(&metrics).counter(outcome, 0).add(1); + for target in targets { + nested.push(target.publish(Arc::clone(&domain.build))); + } + input = Arc::new(DynamicFilterExec::new( + attachment.unwrap_or(input), + predicate, + metrics.clone(), + "dynamic_filter_union", + )); + } + Ok(AppliedUnionFilters { + plan: input, + forwarded_domains: nested, + }) +} diff --git a/native/core/src/execution/operators/dynamic_filter/parquet_reader.rs b/native/core/src/execution/operators/dynamic_filter/parquet_reader.rs index f40fbfc6cb5..ca1e32c440a 100644 --- a/native/core/src/execution/operators/dynamic_filter/parquet_reader.rs +++ b/native/core/src/execution/operators/dynamic_filter/parquet_reader.rs @@ -31,6 +31,8 @@ use datafusion::physical_expr::utils::collect_columns; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::ExecutionPlan; +use super::super::union_filter::UnionFilterTarget; +use super::super::{CometProjectionExec, ScanExec}; use datafusion_comet_operators::CometFilterExec; mod schema_adapter; @@ -56,6 +58,29 @@ pub(super) fn try_attach_parquet_reader_filter( input: &Arc, predicate: Arc, config: &ConfigOptions, +) -> Result>> { + attach_reader_filter(input, predicate, config, &mut None) +} + +pub(crate) fn try_attach_parquet_reader_filter_with_transport( + input: &Arc, + predicate: Arc, + config: &ConfigOptions, + targets: &mut Vec, +) -> Result>> { + let start = targets.len(); + let result = attach_reader_filter(input, predicate, config, &mut Some(targets)); + if !matches!(&result, Ok(Some(_))) { + targets.truncate(start); + } + result +} + +fn attach_reader_filter( + input: &Arc, + predicate: Arc, + config: &ConfigOptions, + targets: &mut Option<&mut Vec>, ) -> Result>> { // Filtering before a fetch can change which rows are selected by its limit. if input.fetch().is_some() { @@ -81,7 +106,7 @@ pub(super) fn try_attach_parquet_reader_filter( return Ok(None); } let Some(reader) = - try_attach_parquet_reader_filter(filter.input(), Arc::clone(&predicate), config)? + attach_reader_filter(filter.input(), Arc::clone(&predicate), config, targets)? else { return Ok(None); }; @@ -95,6 +120,45 @@ pub(super) fn try_attach_parquet_reader_filter( } }; } + // A Union branch can rename/reorder columns. Every projected expression must + // be a direct column: otherwise early rejection could suppress its errors. + if let Some(projection) = input.downcast_ref::() { + if targets.is_none() { + return Ok(None); + } + let expressions = projection.projection().expr(); + if !expressions.iter().all(|expr| expr.expr.is::()) { + return Ok(None); + } + let mut children = Vec::new(); + for key in predicate.children() { + let Some(column) = key.downcast_ref::() else { + return Ok(None); + }; + let Some(expression) = expressions.get(column.index()) else { + return Ok(None); + }; + children.push(Arc::clone(&expression.expr)); + } + let rewritten = Arc::clone(&predicate).with_new_children(children)?; + let rewritten = Arc::downcast::(rewritten).map_err(|_| { + datafusion::common::internal_datafusion_err!("Dynamic filter changed type") + })?; + let Some(child) = attach_reader_filter(projection.input(), rewritten, config, targets)? + else { + return Ok(None); + }; + return Ok(Some(projection.with_execution_input(child)?)); + } + if let Some(scan) = input.downcast_ref::() { + if let (Some(source), Some(targets)) = (&scan.union_input, targets.as_mut()) { + targets.push(UnionFilterTarget { + source: Arc::clone(source), + predicate, + }); + return Ok(Some(Arc::clone(input))); + } + } let Some(scan) = input.downcast_ref::() else { log::debug!( "Join dynamic filter reader pushdown skipped: probe root is {}", diff --git a/native/core/src/execution/operators/mod.rs b/native/core/src/execution/operators/mod.rs index 201ea10dbc7..650eec8cb95 100644 --- a/native/core/src/execution/operators/mod.rs +++ b/native/core/src/execution/operators/mod.rs @@ -27,6 +27,10 @@ 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 union_filter; +pub(crate) use dynamic_filter::apply_union_filter_bundle; +mod runtime_filter_projection; +pub(crate) use runtime_filter_projection::CometProjectionExec; pub(crate) mod iceberg_common; pub use iceberg_common::clear_file_io_cache; mod iceberg_dictionary; diff --git a/native/core/src/execution/operators/runtime_filter_projection.rs b/native/core/src/execution/operators/runtime_filter_projection.rs new file mode 100644 index 00000000000..a0874903505 --- /dev/null +++ b/native/core/src/execution/operators/runtime_filter_projection.rs @@ -0,0 +1,211 @@ +// 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. + +//! A DataFusion projection whose metrics remain owned by its Spark plan node. +//! +//! Comet normally keeps a one-to-one Spark/native plan tree so native metric +//! handles map back to the corresponding Spark operator. Some execution-local +//! rewrites need to replace a projection's child. DataFusion gives that replacement +//! a new private metric set, so this adapter owns the stable metric set and +//! registers the handles from the projection that actually executes. + +use std::fmt::Formatter; +use std::sync::Arc; + +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::{internal_err, Result, Statistics}; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::execution_plan::{ + CardinalityEffect, ChildrenPropertiesMode, ReplaceChildrenOptions, +}; +use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; +use datafusion::physical_plan::projection::ProjectionExec; +use datafusion::physical_plan::statistics::{ChildStats, StatisticsArgs}; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream, +}; + +#[derive(Debug)] +pub(crate) struct CometProjectionExec { + projection: ProjectionExec, + metrics: ExecutionPlanMetricsSet, +} + +impl CometProjectionExec { + pub(crate) fn from_datafusion(projection: ProjectionExec) -> Self { + Self { + projection, + metrics: ExecutionPlanMetricsSet::new(), + } + } + + pub(crate) fn input(&self) -> &Arc { + self.projection.input() + } + + pub(crate) fn projection(&self) -> &ProjectionExec { + &self.projection + } + + fn replace_input( + &self, + input: Arc, + options: ReplaceChildrenOptions, + ) -> Result { + let replaced = Arc::new(self.projection.clone()).replace_children(vec![input], options)?; + let Some(projection) = replaced.downcast_ref::() else { + return internal_err!("ProjectionExec child replacement changed its plan type"); + }; + Ok(projection.clone()) + } + + /// Replace the child for one execution while keeping the metric identity + /// owned by the permanent Spark projection node. + pub(crate) fn with_execution_input( + &self, + input: Arc, + ) -> Result> { + Ok(Arc::new(Self { + projection: self.replace_input( + input, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?, + metrics: self.metrics.clone(), + })) + } +} + +impl DisplayAs for CometProjectionExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + self.projection.fmt_as(t, f) + } +} + +impl ExecutionPlan for CometProjectionExec { + fn name(&self) -> &str { + self.projection.name() + } + + fn properties(&self) -> &Arc { + self.projection.properties() + } + + fn children(&self) -> Vec<&Arc> { + vec![self.projection.input()] + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + self.projection.apply_expressions(f) + } + + fn benefits_from_input_partitioning(&self) -> Vec { + self.projection.benefits_from_input_partitioning() + } + + fn maintains_input_order(&self) -> Vec { + self.projection.maintains_input_order() + } + + fn replace_children( + self: Arc, + mut children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + if children.len() != 1 { + return internal_err!("CometProjectionExec requires one child"); + } + Ok(Arc::new(Self { + projection: self.replace_input(children.remove(0), options)?, + metrics: ExecutionPlanMetricsSet::new(), + })) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + // Every execution gets fresh DataFusion projection metrics. Register their + // handles on the stable set before returning the stream, including when + // opening the child fails. + let projection = self.replace_input( + Arc::clone(self.projection.input()), + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + let result = projection.execute(partition, context); + for metric in projection.metrics().unwrap_or_default().iter() { + self.metrics.register(Arc::clone(metric)); + } + result + } + + fn metrics(&self) -> Option { + Some(self.metrics.clone_inner()) + } + + fn child_stats_requests(&self, partition: Option) -> Vec { + self.projection.child_stats_requests(partition) + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + args: &StatisticsArgs, + ) -> Result> { + self.projection.statistics_from_inputs(input_stats, args) + } + + fn cardinality_effect(&self) -> CardinalityEffect { + self.projection.cardinality_effect() + } + + fn fetch(&self) -> Option { + self.projection.fetch() + } + + fn with_fetch(&self, fetch: Option) -> Option> { + let projection = self.projection.with_fetch(fetch)?; + let projection = projection.downcast_ref::()?.clone(); + Some(Arc::new(Self { + projection, + metrics: ExecutionPlanMetricsSet::new(), + })) + } + + fn with_preserve_order(&self, preserve_order: bool) -> Option> { + let projection = self.projection.with_preserve_order(preserve_order)?; + let projection = projection.downcast_ref::()?.clone(); + Some(Arc::new(Self { + projection, + metrics: ExecutionPlanMetricsSet::new(), + })) + } +} diff --git a/native/core/src/execution/operators/scan.rs b/native/core/src/execution/operators/scan.rs index 1e5be82badf..7cc57bc5752 100644 --- a/native/core/src/execution/operators/scan.rs +++ b/native/core/src/execution/operators/scan.rs @@ -51,6 +51,7 @@ pub struct ScanExec { /// The C Stream Interface reader. `None` only in unit tests that seed input via /// `set_input_batch`. pub input_source: Option>>, + pub(crate) union_input: Option>, pub input_source_description: String, pub data_types: Vec, pub schema: SchemaRef, @@ -90,6 +91,7 @@ impl ScanExec { Ok(Self { exec_context_id, input_source, + union_input: None, input_source_description: input_source_description.to_string(), data_types, batch: Arc::new(Mutex::new(None)), @@ -120,6 +122,32 @@ impl ScanExec { /// Pulls the next input batch from the upstream `ArrowArrayStreamReader` unless one is /// already buffered, then wakes the stream waiting for it. pub fn get_next_batch(&mut self) -> Result<(), CometError> { + if let Some(input) = &self.union_input { + if !input.requested() { + return Ok(()); + } + let mut current = self + .batch + .try_lock() + .map_err(|_| CometError::Internal("Union input batch contended".into()))?; + if current.is_some() { + return Ok(()); + } + *current = Some(match input.next()? { + None => InputBatch::EOF, + Some(batch) => InputBatch::new( + batch + .columns() + .iter() + .map(import_column) + .collect::, _>>()?, + Some(batch.num_rows()), + ), + }); + drop(current); + self.waker.wake(); + return Ok(()); + } if self.input_source.is_none() { // This is a unit test. Input batches are seeded via `set_input_batch`. return Ok(()); @@ -333,6 +361,9 @@ impl Stream for ScanStream<'_> { fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let mut timer = self.baseline_metrics.elapsed_compute().timer(); + if let Some(input) = &self.scan.union_input { + input.request(); + } let mut scan_batch = self.scan.batch.try_lock().unwrap(); let result = match &*scan_batch { diff --git a/native/core/src/execution/operators/union_filter.rs b/native/core/src/execution/operators/union_filter.rs new file mode 100644 index 00000000000..c5739af5065 --- /dev/null +++ b/native/core/src/execution/operators/union_filter.rs @@ -0,0 +1,524 @@ +// 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. + +//! Execution-owned transport of completed join domains through lazy Spark Union inputs. + +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Weak}; + +use super::broadcast::broadcast_copy_bytes; +use arrow::array::RecordBatch; +use arrow::datatypes::{DataType, Schema, SchemaRef}; +use arrow::ffi_stream::FFI_ArrowArrayStream; +use datafusion::common::{internal_err, Result}; +use datafusion::physical_expr::expressions::DynamicFilterPhysicalExpr; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::joins::PreparedHashJoinBuild; +use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; +use futures::FutureExt; +use jni::objects::{Global, JLongArray, JObject, JValue, ReleaseMode}; +use parking_lot::Mutex; + +use crate::errors::{CometError, CometResult}; +use crate::jvm_bridge::JVMClasses; +use arrow::ffi_stream::ArrowArrayStreamReader; + +/// The immutable build lease keeps every allocation referenced by membership expressions +/// charged until the last consumer has dropped its plan and predicate. +#[derive(Debug)] +pub(crate) struct UnionFilterDomain { + pub predicate: Arc, + pub build: Arc, + pub key_types: Vec, +} + +/// A handle is a Box>, never a process-wide lookup key. The JVM +/// marker retains its own box before leaving openStream and releases it with that stream. +#[derive(Debug)] +pub(crate) struct UnionFilterBundle { + pub domains: Vec>, + pub root_plan_ids: Vec, + pub task_attempt_id: i64, +} + +impl UnionFilterBundle { + pub fn accepts(&self, root_plan_id: u32, task_attempt_id: i64) -> bool { + self.task_attempt_id == task_attempt_id && self.root_plan_ids.contains(&root_plan_id) + } +} + +#[derive(Debug, Clone)] +pub(crate) struct UnionFilterTarget { + pub source: Arc, + pub predicate: Arc, +} + +impl UnionFilterTarget { + pub fn publish(&self, build: Arc) -> Arc { + let domain = Arc::new(UnionFilterDomain { + predicate: Arc::clone(&self.predicate), + build, + key_types: self + .predicate + .children() + .iter() + .map(|expr| { + let column = expr + .downcast_ref::() + .expect("transport target has direct columns"); + self.source.data_types[column.index()].clone() + }) + .collect(), + }); + let mut domains = self.source.domains.lock(); + domains.retain(|domain| domain.strong_count() > 0); + domains.push(Arc::downgrade(&domain)); + domain + } +} + +/// ScanExec still owns demand/batch handoff. Only its driving JNI feeder calls next; +/// Tokio polls merely set requested, and can never enter Spark task-local iterator code. +#[derive(Debug)] +pub(crate) struct UnionInput { + marker: Arc>>, + root_plan_ids: Vec, + task_attempt_id: i64, + data_types: Vec, + requested: AtomicBool, + domains: Mutex>>, + reader: Mutex>, + opened: AtomicBool, +} + +impl UnionInput { + pub fn try_new( + marker: Arc>>, + data_types: Vec, + ) -> CometResult>> { + let metadata = JVMClasses::with_env(|env| { + if !env.is_instance_of( + marker.as_obj(), + jni::strings::JNIString::new("org/apache/spark/sql/comet/CometUnionInput"), + )? { + return Ok::<_, CometError>(None); + } + let roots = env + .call_method( + marker.as_obj(), + jni::jni_str!("nativeRootPlanIds"), + jni::jni_sig!("()[J"), + &[], + )? + .l()?; + let roots = JLongArray::cast_local(env, roots)?; + let roots = unsafe { roots.get_elements(env, ReleaseMode::NoCopyBack)? }; + let root_plan_ids = roots + .iter() + .filter_map(|id| u32::try_from(*id).ok()) + .collect(); + drop(roots); + let task_attempt_id = env + .call_method( + marker.as_obj(), + jni::jni_str!("taskAttemptId"), + jni::jni_sig!("()J"), + &[], + )? + .j()?; + Ok(Some((root_plan_ids, task_attempt_id))) + })?; + Ok(metadata.map(|(root_plan_ids, task_attempt_id)| { + Arc::new(Self { + marker, + root_plan_ids, + task_attempt_id, + data_types, + requested: AtomicBool::new(false), + domains: Mutex::new(vec![]), + reader: Mutex::new(None), + opened: AtomicBool::new(false), + }) + })) + } + + pub fn request(&self) { + self.requested.store(true, Ordering::Release); + } + + pub fn requested(&self) -> bool { + self.requested.load(Ordering::Acquire) + } + + fn completed_bundle(&self) -> Arc { + // Zero bounded wait: a missing or incomplete producer never blocks the input. + // A hash join only requests its probe after publication, so its normal path is ready. + Arc::new(UnionFilterBundle { + domains: self + .domains + .lock() + .iter() + .filter_map(Weak::upgrade) + .filter(|domain| domain.predicate.wait_complete().now_or_never().is_some()) + .collect(), + root_plan_ids: self.root_plan_ids.clone(), + task_attempt_id: self.task_attempt_id, + }) + } + + pub fn close(&self) -> CometResult<()> { + // End C stream callbacks before entering JVM child cleanup. Do not hold the reader + // mutex while nested native plans recursively close their own Union inputs. + let reader = self.reader.lock().take(); + drop(reader); + JVMClasses::with_env(|env| { + let result = env.call_method( + self.marker.as_obj(), + jni::jni_str!("closeAfterExecution"), + jni::jni_sig!("()V"), + &[], + ); + // Clear a Java failure before returning it so later Union owners can still + // invoke their cleanup callbacks on this same JNI thread. + if let Some(exception) = crate::jvm_bridge::check_exception(env)? { + return Err(exception); + } + result?; + Ok(()) + }) + } + + pub fn next(&self) -> CometResult> { + let mut reader = self + .reader + .try_lock() + .ok_or_else(|| CometError::Internal("Union input reader contended".into()))?; + if !self.opened.swap(true, Ordering::AcqRel) { + let bundle = self.completed_bundle(); + let handle = Box::new(bundle); + let address = (&*handle as *const Arc) as i64; + // The marker clones this handle before returning; the temporary box can then drop. + let stream = JVMClasses::with_env(|env| { + let stream = env + .call_method( + self.marker.as_obj(), + jni::jni_str!("openStream"), + jni::jni_sig!("(J)Lorg/apache/arrow/c/ArrowArrayStream;"), + &[JValue::Long(address)], + )? + .l()?; + let address = env + .call_method( + &stream, + jni::jni_str!("memoryAddress"), + jni::jni_sig!("()J"), + &[], + )? + .j()?; + Ok::<_, CometError>(address) + })?; + *reader = Some(unsafe { + ArrowArrayStreamReader::from_raw(stream as *mut FFI_ArrowArrayStream) + }?); + } + let result = reader.as_mut().and_then(Iterator::next).transpose()?; + if result.is_none() { + reader.take(); + } + Ok(result) + } +} + +pub(crate) fn transport_build_supported(schema: &SchemaRef) -> bool { + schema.fields().iter().all(|field| { + let ty = field.data_type(); + ty.primitive_width().is_some() + || matches!( + ty, + DataType::Null | DataType::Boolean | DataType::FixedSizeBinary(_) | DataType::Utf8 + ) + }) +} + +pub(crate) fn remap_domain( + domain: &UnionFilterDomain, + schema: &SchemaRef, +) -> Result> { + remap_predicate(&domain.predicate, &domain.key_types, schema) +} + +fn remap_predicate( + predicate: &Arc, + key_types: &[DataType], + schema: &SchemaRef, +) -> Result> { + use datafusion::physical_expr::expressions::Column; + if predicate.children().len() != key_types.len() { + return internal_err!("Union runtime filter key metadata does not match the predicate"); + } + let mut children = vec![]; + for (expr, key_type) in predicate.children().into_iter().zip(key_types) { + let Some(column) = expr.downcast_ref::() else { + return internal_err!("Union runtime filter requires direct column bindings"); + }; + let Some(field) = schema.fields().get(column.index()) else { + return internal_err!("Union runtime filter column is outside consumer schema"); + }; + if field.data_type() != key_type + || !matches!( + field.data_type(), + DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 + ) + { + return internal_err!("Union runtime filter consumer key must be a signed integer"); + } + children.push(Arc::new(Column::new(field.name(), column.index())) as Arc); + } + let predicate = Arc::clone(predicate).with_new_children(children)?; + Arc::downcast::(predicate).map_err(|_| { + datafusion::common::internal_datafusion_err!("Dynamic filter remapping changed type") + }) +} + +/// Prepare one task-owned build before demanding any Union probe input. The same prepared +/// build is installed back into the original join; publication never builds a second domain. +pub(crate) fn prepare_union_join( + join: &datafusion::physical_plan::joins::HashJoinExec, + partition: usize, + context: Arc, + targets: Vec, + metrics: ExecutionPlanMetricsSet, +) -> Result { + use datafusion::execution::memory_pool::MemoryConsumer; + use datafusion::physical_plan::joins::PartitionMode; + use datafusion::physical_plan::metrics::MetricBuilder; + use datafusion::physical_plan::stream::RecordBatchStreamAdapter; + use datafusion::physical_plan::ExecutionPlan; + use futures::{StreamExt, TryStreamExt}; + + let join = join + .builder() + .with_partition_mode(PartitionMode::CollectLeft) + .build()?; + let schema = join.schema(); + // prepare_build does not expose its metrics. Preserve the task-local join's build + // counters here; the execution below only measures probing an already prepared build. + let build_time = MetricBuilder::new(&metrics).subset_time("build_time", partition); + let build_rows = MetricBuilder::new(&metrics).counter("build_input_rows", partition); + let build_batches = MetricBuilder::new(&metrics).counter("build_input_batches", partition); + let build_memory = MetricBuilder::new(&metrics).peak_memory_usage("build_mem_used", partition); + let stream = futures::stream::once(async move { + let build_timer = build_time.timer(); + let pool = Arc::clone(context.memory_pool()); + let reservation = MemoryConsumer::new("UnionFilterBuildCopy").register(&pool); + let input = join.left().execute(0, Arc::clone(&context))?; + let input_schema = input.schema(); + let input = input.inspect_ok(move |batch| { + // Count source batches, not the pieces produced by copy_batch_rows. + build_batches.add(1); + build_rows.add(batch.num_rows()); + }); + let copied = futures::stream::try_unfold( + (input, reservation, None::), + |(mut input, mut reservation, pending)| async move { + // PreparedHashJoinBuild has admitted the prior batch before polling again. + reservation.free(); + let already_split = pending.is_some(); + let batch = match pending { + Some(batch) => batch, + None => match input.next().await { + Some(batch) => batch?, + None => return Ok(None), + }, + }; + let rows = + copy_batch_rows(&batch, &batch.schema(), 32 * 1024 * 1024, already_split)?; + let pending = + (rows < batch.num_rows()).then(|| batch.slice(rows, batch.num_rows() - rows)); + let batch = batch.slice(0, rows); + let columns = super::broadcast::copy_broadcast_columns( + &batch, + &batch.schema(), + &mut reservation, + )?; + let batch = RecordBatch::try_new_with_options( + batch.schema(), + columns, + &arrow::array::RecordBatchOptions::new().with_row_count(Some(batch.num_rows())), + )?; + Ok(Some((batch, (input, reservation, pending)))) + }, + ); + let copied = Box::pin(RecordBatchStreamAdapter::new(input_schema, copied)); + let build = join + .prepare_build(copied, pool, Arc::clone(context.session_config().options())) + .await?; + build_memory.set(build.reserved_bytes()); + build_timer.done(); + let domains = targets + .iter() + .map(|target| target.publish(Arc::clone(&build))) + .collect::>(); + let execution = join + .builder() + .with_prepared_build(Arc::clone(&build)) + .build()?; + let stream = execution.execute(partition, context)?; + for metric in execution.metrics().unwrap_or_default().iter() { + if !matches!( + metric.value().name(), + "build_time" | "build_input_rows" | "build_input_batches" | "build_mem_used" + ) { + metrics.register(Arc::clone(metric)); + } + } + // Consumers can retain the prepared build after this stream has finished. + let output_schema = stream.schema(); + let stream = stream.map(move |batch| { + let _ = &domains; + batch + }); + Ok::<_, datafusion::common::DataFusionError>(Box::pin(RecordBatchStreamAdapter::new( + output_schema, + stream, + )) + as datafusion::physical_plan::SendableRecordBatchStream) + }) + .try_flatten(); + Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) +} + +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) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::Int32Array; + use arrow::datatypes::Field; + use datafusion::logical_expr::Operator; + use datafusion::physical_expr::expressions::{lit, BinaryExpr, Column}; + + #[test] + fn bundle_requires_both_authorized_root_and_task_attempt() { + let bundle = UnionFilterBundle { + domains: vec![], + root_plan_ids: vec![2, 7], + task_attempt_id: 11, + }; + assert!(bundle.accepts(2, 11)); + assert!(bundle.accepts(7, 11)); + assert!(!bundle.accepts(3, 11)); + assert!(!bundle.accepts(2, 12)); + assert!(!bundle.accepts(7, -1)); + } + + #[test] + fn branch_remapping_preserves_live_updates_and_rejects_incompatible_keys() { + let predicate = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::new(Column::new("union_key", 1))], + lit(true), + )); + let schema = Arc::new(Schema::new(vec![ + Field::new("payload", DataType::Int32, false), + Field::new("branch_key", DataType::Int32, false), + ])); + let remapped = remap_predicate(&predicate, &[DataType::Int32], &schema).unwrap(); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![5, 6])), + Arc::new(Int32Array::from(vec![1, 2])), + ], + ) + .unwrap(); + // Rebinding must share the producer's state, including a later completion. + predicate + .update(Arc::new(BinaryExpr::new( + Arc::new(Column::new("union_key", 1)), + Operator::Eq, + lit(2_i32), + ))) + .unwrap(); + predicate.mark_complete(); + assert!(remapped.wait_complete().now_or_never().is_some()); + let result = remapped.evaluate(&batch).unwrap().into_array(2).unwrap(); + assert_eq!( + result.as_ref(), + &arrow::array::BooleanArray::from(vec![false, true]) + ); + assert!(remap_predicate(&predicate, &[DataType::Int64], &schema).is_err()); + assert!(remap_predicate(&predicate, &[], &schema).is_err()); + let short = Arc::new(Schema::new(vec![Field::new( + "payload", + DataType::Int32, + false, + )])); + assert!(remap_predicate(&predicate, &[DataType::Int32], &short).is_err()); + } + + #[test] + fn copy_chunks_bound_large_handoffs_without_losing_rows() { + let schema = Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)])); + let mut batch = RecordBatch::try_new( + schema, + vec![Arc::new(Int32Array::from_iter_values(0..10000))], + ) + .unwrap(); + // A single row is still returned when fixed allocation overhead exceeds the target. + assert_eq!( + copy_batch_rows(&batch, &batch.schema(), 256, false).unwrap(), + 1 + ); + let target = 4096; + let mut rows = 0; + let mut already_split = false; + while batch.num_rows() > 0 { + let size = copy_batch_rows(&batch, &batch.schema(), target, already_split).unwrap(); + assert!(size > 0); + let chunk = batch.slice(0, size); + assert!(broadcast_copy_bytes(&chunk, &chunk.schema()).unwrap() <= target); + rows += size; + batch = batch.slice(size, batch.num_rows() - size); + already_split = true; + } + assert_eq!(rows, 10000); + } +} diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index ca9ae2f86f3..d4cd7db0e15 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -41,8 +41,8 @@ use crate::execution::operators::TopKReaderFilterExec; use crate::execution::operators::{reuse_broadcast_build, BroadcastInputExec}; use crate::execution::{ operators::{ - ExecutionError, MergeActionContext, MergeInstructionExec, MergeRowsExec, ScanExec, - ShuffleScanExec, + CometProjectionExec, ExecutionError, MergeActionContext, MergeInstructionExec, + MergeRowsExec, ScanExec, ShuffleScanExec, }, planner::expression_registry::ExpressionRegistry, planner::operator_registry::OperatorRegistry, @@ -1951,6 +1951,23 @@ impl PhysicalPlanner { )), )); } + if let Some(union_input) = super::operators::union_filter::UnionInput::try_new( + Arc::clone(&java_stream), + data_types.clone(), + )? { + let mut lazy_scan = + ScanExec::new(self.exec_context_id, None, &scan.source, data_types)?; + lazy_scan.union_input = Some(union_input); + return Ok(( + vec![lazy_scan.clone()], + vec![], + Arc::new(SparkPlan::new( + spark_plan.plan_id, + Arc::new(lazy_scan), + vec![], + )), + )); + } let address: i64 = JVMClasses::with_env(|env| unsafe { jni_call!(env, arrow_array_stream(java_stream.as_obj()).memory_address() -> i64) })?; @@ -2461,7 +2478,11 @@ impl PhysicalPlanner { } else { &mut join_params.left }; - *probe = Self::prepare_probe_filter_for_runtime_reader(Arc::clone(probe)); + *probe = if join.union_filter_transport_enabled { + Self::prepare_union_branch_for_runtime_reader(Arc::clone(probe)) + } else { + Self::prepare_probe_filter_for_runtime_reader(Arc::clone(probe)) + }; } let left = Arc::clone(&join_params.left.native_plan); @@ -2497,9 +2518,10 @@ impl PhysicalPlanner { // (which matches DataFusion's default), and swap_inputs would turn LeftAnti // into RightAnti, which DataFusion rejects with null_aware=true. if join.build_side == BuildSide::BuildLeft as i32 || join.null_aware_anti_join { - let hash_join = Self::apply_join_dynamic_filter( + let hash_join = Self::apply_join_dynamic_filter_with_transport( hash_join, join.dynamic_filter_enabled && !join.null_aware_anti_join, + join.union_filter_transport_enabled, self.session_ctx.copied_config().options(), )?; let hash_join = reuse_broadcast_build(hash_join)?; @@ -2515,9 +2537,10 @@ impl PhysicalPlanner { } else { let swapped_hash_join = hash_join.as_ref().swap_inputs(PartitionMode::Partitioned)?; - let swapped_hash_join = Self::apply_join_dynamic_filter( + let swapped_hash_join = Self::apply_join_dynamic_filter_with_transport( swapped_hash_join, join.dynamic_filter_enabled, + join.union_filter_transport_enabled, self.session_ctx.copied_config().options(), )?; let swapped_hash_join = reuse_broadcast_build(swapped_hash_join)?; @@ -2757,19 +2780,72 @@ impl PhysicalPlanner { Arc::new(prepared) } + /// Keep Spark metric identities while preparing an opt-in Union branch's safe reader path. + pub(crate) fn prepare_union_branch_for_runtime_reader(plan: Arc) -> Arc { + let filter = plan.native_plan.downcast_ref::(); + let projection = plan.native_plan.downcast_ref::(); + if filter.is_none() && projection.is_none() { + return plan; + } + let mut prepared = plan.as_ref().clone(); + let child = if plan.children.len() == 1 + && plan.native_plan.children().len() == 1 + && Arc::ptr_eq( + plan.native_plan.children()[0], + &plan.children[0].native_plan, + ) { + let child = + Self::prepare_union_branch_for_runtime_reader(Arc::clone(&plan.children[0])); + prepared.children = vec![Arc::clone(&child)]; + Arc::clone(&child.native_plan) + } else { + Arc::clone(plan.native_plan.children()[0]) + }; + let rewritten = if let Some(filter) = filter { + CometFilterExec::from_datafusion(filter.clone()).with_execution_input(child) + } else { + CometProjectionExec::from_datafusion(projection.unwrap().clone()) + .with_execution_input(child) + }; + match rewritten { + Ok(native_plan) => { + prepared.native_plan = native_plan; + Arc::new(prepared) + } + Err(error) => { + log::debug!("Runtime filter branch preparation skipped: {error}"); + plan + } + } + } + /// Attach after choosing the final build side, including the projection emitted /// by `swap_inputs`, without running DataFusion's physical optimizer. + #[cfg(test)] pub(crate) fn apply_join_dynamic_filter( plan: Arc, enabled: bool, config: &ConfigOptions, + ) -> Result, ExecutionError> { + Self::apply_join_dynamic_filter_with_transport(plan, enabled, false, config) + } + + fn apply_join_dynamic_filter_with_transport( + plan: Arc, + enabled: bool, + union_transport: bool, + config: &ConfigOptions, ) -> Result, ExecutionError> { if !enabled { return Ok(plan); } if let Some(projection) = plan.downcast_ref::() { - let child = - Self::apply_join_dynamic_filter(Arc::clone(projection.input()), enabled, config)?; + let child = Self::apply_join_dynamic_filter_with_transport( + Arc::clone(projection.input()), + enabled, + union_transport, + config, + )?; return if !Arc::ptr_eq(&child, projection.input()) { Ok(plan.replace_children( vec![child], @@ -2787,7 +2863,7 @@ impl PhysicalPlanner { return Ok(plan); }; match DynamicFilterJoinExec::try_new(join, config)? { - Some(wrapper) => Ok(Arc::new(wrapper)), + Some(wrapper) => Ok(Arc::new(wrapper.with_union_transport(union_transport))), None => Ok(plan), } } @@ -5726,6 +5802,7 @@ mod tests { sql_text_pool: vec![], children: vec![op_scan.clone(), op_scan.clone()], op_struct: Some(OpStruct::HashJoin(spark_operator::HashJoin { + union_filter_transport_enabled: false, left_join_keys: vec![create_bound_reference(0)], right_join_keys: vec![create_bound_reference(0)], join_type: 0, @@ -5761,6 +5838,7 @@ mod tests { let operator = Operator { children: vec![create_scan(), create_scan()], op_struct: Some(OpStruct::HashJoin(spark_operator::HashJoin { + union_filter_transport_enabled: false, left_join_keys: vec![create_bound_reference(0)], right_join_keys: vec![create_bound_reference(0)], join_type: spark_operator::JoinType::Inner as i32, diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 68b881cdcb1..fe7035b4a55 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -1030,6 +1030,9 @@ message HashJoin { // Opt-in completed-build filtering of probe batches and eligible native Parquet readers. // Older serialized plans leave it false. bool dynamic_filter_enabled = 10; + // Carry a completed broadcast join filter into separately planned Union readers. + // Older plans and shuffled joins leave this permission false. + bool union_filter_transport_enabled = 11; } message SortMergeJoin { diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 30cb709977e..d3c1d15d496 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -594,6 +594,19 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(false) + val COMET_EXEC_JOIN_DYNAMIC_FILTER_UNION_ENABLED: ConfigEntry[Boolean] = + conf(s"$COMET_EXEC_CONFIG_PREFIX.join.dynamicFilter.union.enabled") + .category(CATEGORY_EXEC) + .doc( + "Carry completed native broadcast-join filters across UNION ALL inputs into eligible " + + "branch readers. Requires join.dynamicFilter.enabled. Supports a single direct signed " + + "integer key with no join residual. Keeps the join, Union partitioning and broadcast " + + "exchange in place; filters retain their accounted build storage for the current " + + "task attempt. Existing native reader eligibility barriers still apply. Missing or " + + "ineligible filters proceed without reader pruning.") + .booleanConf + .createWithDefault(false) + val COMET_SCALA_UDF_CODEGEN_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.exec.scalaUDF.codegen.enabled") .category(CATEGORY_EXEC) diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index 22143b4f188..4d26941d3d0 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -32,7 +32,7 @@ import org.apache.spark._ import org.apache.spark.broadcast.Broadcast import org.apache.spark.internal.Logging import org.apache.spark.network.util.ByteUnit -import org.apache.spark.sql.comet.CometMetricNode +import org.apache.spark.sql.comet.{CometMetricNode, CometUnionInput} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.vectorized._ import org.apache.spark.util.SerializableConfiguration @@ -218,9 +218,8 @@ class CometExecIterator( // Register a task completion listener to ensure native resources are released // when the task is done. - TaskContext.get().addTaskCompletionListener[Unit] { _ => - this.close() - } + CometUnionInput.addCleanup(TaskContext.get())(close()) + CometUnionInput.registerExecution(plan) CometExecIterator.startMemoryUsageLog() @@ -262,8 +261,11 @@ class CometExecIterator( } } + private val deferCloseToUnion = CometUnionInput.isBranchExecution + private var exhausted = false + override def hasNext: Boolean = { - if (closed) return false + if (closed || exhausted) return false if (nextBatch.isDefined) { return true @@ -283,7 +285,10 @@ class CometExecIterator( if (nextBatch.isEmpty) { readPartitionOffsetsBeforeClose() - close() + exhausted = true + // An outer native operator may still retain this branch's Arrow buffers and build + // leases. Its Union owner already registered close() and releases us after its plan. + if (!deferCloseToUnion) close() false } else { true @@ -338,6 +343,10 @@ class CometExecIterator( // execution context frees this plan's task-shared memory pool reference and several JNI // global refs. attempt(nativeLib.releasePlan(plan)) + inputObjects.foreach { + case union: CometUnionInput => attempt(union.closeAfterExecution()) + case _ => + } // Run the diagnostics even when teardown failed: a failed teardown is exactly when the // non-zero memory usage warning below is most informative. diff --git a/spark/src/main/scala/org/apache/comet/Native.scala b/spark/src/main/scala/org/apache/comet/Native.scala index 89215e8a52f..81ccb97380a 100644 --- a/spark/src/main/scala/org/apache/comet/Native.scala +++ b/spark/src/main/scala/org/apache/comet/Native.scala @@ -79,6 +79,17 @@ class Native extends NativeBase { classLoader: ClassLoader): Long // scalastyle:on + /** Retain a Union filter lease while a lazy input's Arrow callbacks can open branch plans. */ + @native def retainUnionRuntimeFilter(handle: Long): Long + + /** Release one retained Union filter lease. The caller must consume its handle exactly once. */ + @native def releaseUnionRuntimeFilter(handle: Long): Unit + + /** + * Bind an attempt-scoped Union filter to an authorized branch before native execution starts. + */ + @native def setUnionRuntimeFilter(plan: Long, handle: Long): Unit + /** * Bind a task-owned remote shuffle callback to an existing native query plan. * 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 ce230a21b51..7af95680c30 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 @@ -48,7 +48,7 @@ class CometBroadcastInput private[comet] ( private val streams = ArrayBuffer.empty[CometBroadcastArrowStream] private var closed = false - context.addTaskCompletionListener[Unit](_ => close()) + CometUnionInput.addCleanup(context)(close()) def getBroadcastId(): Long = broadcast.id 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 6f1e6b469aa..27792d181f2 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 @@ -142,7 +142,7 @@ private[spark] class CometExecRDD( subqueries.foreach(sub => CometScalarSubquery.setSubquery(it.id, sub)) Option(context).foreach { ctx => - ctx.addTaskCompletionListener[Unit] { _ => + CometUnionInput.addCleanup(ctx) { subqueries.foreach(sub => CometScalarSubquery.removeSubquery(it.id, sub)) } } @@ -171,13 +171,9 @@ 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`), 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. + * a `CometShuffleBlockIterator`, a lazy `CometBroadcastInput` or `CometUnionInput`, or an + * `ArrowArrayStream`. Broadcasts without a supported memory owner open ordinary streams. Shared + * by [[CometExecRDD.compute]] and the native-shuffle path; lazy inputs stay unopened. */ def resolveInputObjects( inputRDDs: Seq[RDD[_]], @@ -200,6 +196,8 @@ 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[CometUnionInputRDD]) { + rdd.iterator(part, context).next().asInstanceOf[CometUnionInput] } else if (rdd.isInstanceOf[CometBroadcastInputRDD]) { val input = rdd.iterator(part, context).next().asInstanceOf[CometBroadcastInput] // Unsupported executor memory configuration must remain an ordinary Scan input. diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala index 17dc34dcb61..dc05fc1da62 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala @@ -30,6 +30,7 @@ import org.apache.spark.internal.Logging import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} +import org.apache.comet.CometConf import org.apache.comet.serde.Metric /** @@ -124,7 +125,7 @@ case class CometMetricNode(metrics: Map[String, SQLMetric], children: Seq[CometM */ def reportScanInputMetrics(ctx: TaskContext): Unit = { val seenMetrics = CometMetricNode.taskSeenMetrics(ctx).scanInput - val _ = ctx.addTaskCompletionListener[Unit] { _ => + CometUnionInput.addCleanup(ctx) { val scanLeaves = leafNodes.filter(_.metrics.contains("bytes_scanned")) def claimed(leaf: CometMetricNode, metricName: String): Long = leaf.metrics.get(metricName).fold(0L)(CometMetricNode.claimMetricValue(_, seenMetrics)) @@ -156,7 +157,7 @@ case class CometMetricNode(metrics: Map[String, SQLMetric], children: Seq[CometM * spill reports share. */ def reportNativeWriteOutputMetrics(ctx: TaskContext): Unit = { - val _ = ctx.addTaskCompletionListener[Unit] { _ => + CometUnionInput.addCleanup(ctx) { metrics.get("bytes_written").foreach { m => ctx.taskMetrics().outputMetrics.setBytesWritten(m.value) } @@ -181,7 +182,7 @@ case class CometMetricNode(metrics: Map[String, SQLMetric], children: Seq[CometM */ def reportSpillMetrics(ctx: TaskContext): Unit = { val seenMetrics = CometMetricNode.taskSeenMetrics(ctx) - val _ = ctx.addTaskCompletionListener[Unit] { _ => + CometUnionInput.addCleanup(ctx) { val diskBytesSpilled = sumMetricValues("spilled_bytes", seenMetrics.disk) if (diskBytesSpilled > 0L) { ctx.taskMetrics().incDiskBytesSpilled(diskBytesSpilled) @@ -312,7 +313,25 @@ object CometMetricNode { "output_rows" -> SQLMetrics.createMetric(sc, "number of output rows"), "elapsed_compute" -> SQLMetrics.createNanoTimingMetric( sc, - "total time (in ms) spent in this operator")) + "total time (in ms) spent in this operator")) ++ unionFilterMetrics(sc) + } + + def unionFilterMetrics(sc: SparkContext): Map[String, SQLMetric] = { + if (!CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_ENABLED.get() || + !CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_UNION_ENABLED.get()) { + return Map.empty + } + Seq( + "branches_attached" -> "Union branches with a reader filter", + "branches_forwarded" -> "Union branches forwarding a filter to nested Union inputs", + "branches_skipped" -> "Union branches without reader attachment", + "rows_evaluated" -> "Union branch rows evaluated by the filter", + "rows_pruned" -> "Union branch rows rejected by the filter", + "rows_bypassed" -> "Union branch rows bypassing the filter").map { + case (name, description) => + s"dynamic_filter_union_$name" -> SQLMetrics.createMetric(sc, description) + }.toMap + ("dynamic_filter_union_eval_time" -> + SQLMetrics.createNanoTimingMetric(sc, "Time evaluating Union filters")) } def aggregateMetrics(sc: SparkContext): Map[String, SQLMetric] = { @@ -332,7 +351,7 @@ object CometMetricNode { * regardless of which scan implementation is used. */ def baseScanMetrics(sc: SparkContext): Map[String, SQLMetric] = { - Map( + unionFilterMetrics(sc) ++ Map( "numOutputRows" -> SQLMetrics.createMetric(sc, "number of output rows"), "scanTime" -> SQLMetrics.createNanoTimingMetric(sc, "scan time")) } @@ -369,7 +388,7 @@ object CometMetricNode { * SQL Metrics from the native Datafusion reader. */ def nativeScanMetrics(sc: SparkContext): Map[String, SQLMetric] = { - Map( + unionFilterMetrics(sc) ++ Map( "output_rows" -> SQLMetrics.createMetric(sc, "number of output rows"), "time_elapsed_opening" -> SQLMetrics.createNanoTimingMetric(sc, "Wall clock time elapsed for file opening"), @@ -494,7 +513,7 @@ object CometMetricNode { * NestedLoopJoinExec). */ def joinMetrics(sc: SparkContext): Map[String, SQLMetric] = { - Map( + unionFilterMetrics(sc) ++ Map( "build_time" -> SQLMetrics.createNanoTimingMetric(sc, "Total time for collecting build-side of join"), "build_input_batches" -> @@ -533,7 +552,7 @@ object CometMetricNode { * SQL Metrics for DataFusion SortMergeJoin */ def sortMergeJoinMetrics(sc: SparkContext): Map[String, SQLMetric] = { - Map( + unionFilterMetrics(sc) ++ Map( "peak_mem_used" -> SQLMetrics.createSizeMetric(sc, "Memory used by build-side"), "input_batches" -> diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometUnionInput.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometUnionInput.scala new file mode 100644 index 00000000000..1903e93a741 --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometUnionInput.scala @@ -0,0 +1,155 @@ +/* + * 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 scala.collection.mutable.ArrayBuffer + +import org.apache.arrow.c.ArrowArrayStream +import org.apache.arrow.util.AutoCloseables +import org.apache.spark.{OneToOneDependency, Partition, TaskContext} +import org.apache.spark.comet.CometTaskContextShim +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.comet.execution.arrow.CometBroadcastArrowStream +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.vectorized.ColumnarBatch + +import org.apache.comet.Native + +/** Opens the original Union iterator only after the native join has completed its build. */ +class CometUnionInput private[comet] ( + iteratorFactory: () => Iterator[ColumnarBatch], + schema: StructType, + name: String, + context: TaskContext, + val nativeRootPlanIds: Array[Long]) + extends AutoCloseable { + + private val nativeLib = new Native() + private val children = ArrayBuffer.empty[AutoCloseable] + private var domainHandle = 0L + private var stream: CometBroadcastArrowStream = null + private var closed = false + + // Releasing the C stream stops pulls, but the outer iterator can still hold output buffers. + // Keep branch resources until it closes this input after releasing those outputs. + private val source = new Iterator[ColumnarBatch] with AutoCloseable { + private lazy val batches = iteratorFactory() + private var stopped = false + + override def hasNext: Boolean = !stopped && withScope(batches.hasNext) + + override def next(): ColumnarBatch = withScope(batches.next()) + + override def close(): Unit = { stopped = true } + } + + CometUnionInput.addCleanup(context)(close()) + + def taskAttemptId(): Long = context.taskAttemptId() + + /** Retain the borrowed JNI handle until Arrow has finished opening and consuming branches. */ + def openStream(handle: Long): ArrowArrayStream = withScope { + domainHandle = if (handle == 0L) 0L else nativeLib.retainUnionRuntimeFilter(handle) + stream = CometBroadcastArrowStream.open(source, schema, name) + stream.stream + } + + override def close(): Unit = synchronized { + if (!closed) { + closed = true + val lease: AutoCloseable = () => { + if (domainHandle != 0L) nativeLib.releaseUnionRuntimeFilter(domainHandle) + } + AutoCloseables.close((Seq(stream, source) ++ children.reverse :+ lease): _*) + children.clear() + } + } + + // Called after the enclosing native plan and exported outputs have been released. + def closeAfterExecution(): Unit = close() + + private def withScope[T](body: => T): T = { + val previousTask = TaskContext.get() + val previousInput = CometUnionInput.current.get() + CometTaskContextShim.set(context) + CometUnionInput.current.set(this) + try { + context.killTaskIfInterrupted() + body + } finally { + if (previousInput == null) CometUnionInput.current.remove() + else CometUnionInput.current.set(previousInput) + if (previousTask == null) CometTaskContextShim.unset() + else CometTaskContextShim.set(previousTask) + } + } +} + +object CometUnionInput { + // Used only during synchronous Spark iterator callbacks. Persistent ownership is on the input; + // native forwarding carries completed domains into nested Union inputs. + private val current = new ThreadLocal[CometUnionInput]() + + def isBranchExecution: Boolean = current.get() != null + + /** Lazy branch resources outlive the outer native plan and its exported output batches. */ + def addCleanup(context: TaskContext)(close: => Unit): Unit = { + val input = current.get() + if (input != null) input.children += (() => close) + else context.addTaskCompletionListener[Unit](_ => close) + } + + def registerExecution(plan: Long): Unit = { + Option(current.get()).foreach { input => + // Native validates the exact branch root and task attempt before retaining the lease. + if (input.domainHandle != 0L) { + input.nativeLib.setUnionRuntimeFilter(plan, input.domainHandle) + } + } + } +} + +/** Preserve the original Union partitions, dependencies and preferred locations. */ +private[comet] class CometUnionInputRDD( + batches: RDD[ColumnarBatch], + schema: StructType, + name: String, + nativeRootPlanIds: Array[Long]) + extends RDD[CometUnionInput](batches.context, Seq(new OneToOneDependency(batches))) { + + override val partitioner = batches.partitioner + + override protected def getPartitions: Array[Partition] = + firstParent[ColumnarBatch].partitions + + override def compute(split: Partition, context: TaskContext): Iterator[CometUnionInput] = { + val parent = firstParent[ColumnarBatch] + Iterator.single( + new CometUnionInput( + () => parent.iterator(split, context), + schema, + name, + context, + nativeRootPlanIds)) + } + + override protected def getPreferredLocations(split: Partition): Seq[String] = + firstParent[ColumnarBatch].preferredLocations(split) +} diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala index 19532cefa14..bc4f4f07563 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala @@ -29,6 +29,7 @@ import org.apache.arrow.vector.types.pojo.{Field, FieldType, Schema} import org.apache.spark.TaskContext import org.apache.spark.internal.Logging import org.apache.spark.rdd.RDD +import org.apache.spark.sql.comet.CometUnionInput import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.types.StructType @@ -232,7 +233,9 @@ object CometArrowStream extends Logging { * frees only the C struct, leaving the reader open and pinned for the life of the executor by * the JNI global ref that arrow-java holds on the stream's private data. Taking the stream * leaves a null release callback in the JVM's struct, and arrow-java's `release` skips a null - * callback, so a stream native took is still released only once. + * callback, so a stream native took is still released only once. Lazy Union branches instead + * register with their input owner, which closes them after the outer iterator has also released + * its exported output batches. */ def stream( name: String, @@ -256,8 +259,8 @@ object CometArrowStream extends Logging { } if (context != null) { val streamRef = arrowStream - context.addTaskCompletionListener[Unit] { _ => - // Release before close: `close` frees the struct that `release` reads. + CometUnionInput.addCleanup(context) { + // Release before close also handles streams native never took. streamRef.release() streamRef.close() allocator.close() @@ -293,7 +296,7 @@ object CometArrowStream extends Logging { throw t } if (context != null) { - context.addTaskCompletionListener[Unit] { _ => + CometUnionInput.addCleanup(context) { reader.close() allocator.close() } 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 6eac80f047e..2c6b080a988 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 @@ -1049,6 +1049,19 @@ abstract class CometNativeExec extends CometExec { */ def asArrowStreamRDD(plan: SparkPlan, partitionCount: Int, scanSlot: Int): RDD[_] = plan match { + case union: CometUnionExec + if CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_ENABLED.get(conf) && + CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_UNION_ENABLED.get(conf) => + def branchRoots(branch: SparkPlan): Seq[Long] = branch match { + case native: CometNativeExec => Seq(native.nativeOp.getPlanId) + case nested: CometUnionExec => nested.children.flatMap(branchRoots) + case _ => Seq.empty + } + new CometUnionInputRDD( + union.executeColumnar(), + union.schema, + union.nodeName, + union.children.flatMap(branchRoots).toArray) case s: CometNativeArrowSource => s.doExecuteAsArrowStream() case _ => @@ -2645,6 +2658,21 @@ case class CometHashAggregateExec( trait CometHashJoin { + private def unionFilterTransportEnabled(join: HashJoin): Boolean = { + val keysSupported = (join.leftKeys, join.rightKeys) match { + case (Seq(left: Attribute), Seq(right: Attribute)) if left.dataType == right.dataType => + left.dataType match { + case ByteType | ShortType | IntegerType | LongType => true + case _ => false + } + case _ => false + } + join.isInstanceOf[BroadcastHashJoinExec] && join.joinType == Inner && + join.condition.isEmpty && keysSupported && + CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_ENABLED.get(join.conf) && + CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_UNION_ENABLED.get(join.conf) + } + // Only BroadcastHashJoinExec can be null-aware (NOT IN subqueries). protected def isNullAware(join: HashJoin): Boolean = join match { case bhj: BroadcastHashJoinExec => bhj.isNullAwareAntiJoin @@ -2751,6 +2779,7 @@ trait CometHashJoin { else OperatorOuterClass.BuildSide.BuildRight) .setNullAwareAntiJoin(isNullAwareAntiJoin) .setDynamicFilterEnabled(CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_ENABLED.get(join.conf)) + .setUnionFilterTransportEnabled(unionFilterTransportEnabled(join)) condition.foreach(joinBuilder.setCondition) Some(builder.setHashJoin(joinBuilder).build()) } 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 e9eabf0533f..24f50a6d9f6 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -21,6 +21,8 @@ package org.apache.comet.exec import java.util.concurrent.{CountDownLatch, TimeUnit} +import scala.concurrent.duration.DurationInt + import org.scalactic.source.Position import org.scalatest.Tag @@ -33,7 +35,7 @@ import org.apache.spark.sql.catalyst.TableIdentifier import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation import org.apache.spark.sql.catalyst.expressions.{And, AttributeReference, DynamicPruningExpression, IsNotNull} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight} -import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometNativeScanExec, CometSortMergeJoinExec, CometUnionExec} +import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometNativeScanExec, CometSortExec, CometSortMergeJoinExec, CometUnionExec} import org.apache.spark.sql.execution.{LocalTableScanExec, SparkPlan} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec} import org.apache.spark.sql.execution.datasources.SchemaColumnConvertNotSupportedException @@ -41,8 +43,8 @@ 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, NativeBase} -import org.apache.comet.CometSparkSessionExtensions.isSpark35Plus +import org.apache.comet.{CometArrowAllocator, CometConf, NativeBase} +import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark41Plus} class CometJoinSuite extends CometTestBase { @@ -436,6 +438,523 @@ class CometJoinSuite extends CometTestBase { } } + private def unionFilterMetric(plan: SparkPlan, name: String): Long = + collect(plan) { case node: SparkPlan => + node.metrics.get(s"dynamic_filter_union_$name") + }.flatten.groupBy(_.id).values.map(values => math.max(0L, values.head.value)).sum + + for ((buildLeft, adaptive) <- Seq((false, false), (true, true))) { + test(s"union join filter reaches both readers: buildLeft=$buildLeft, AQE=$adaptive") { + withTempPath { firstPath => + withTempPath { secondPath => + withSQLConf( + CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_ENABLED.key -> "true", + CometConf.COMET_BROADCAST_REUSE_ENABLED.key -> "false", + CometConf.COMET_BATCH_SIZE.key -> "16", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.LEAF_NODE_DEFAULT_PARALLELISM.key -> "1") { + // Different physical column order and output names require branch-local key binding. + // Retain the same rows in both branches to verify UNION ALL multiplicity. + val rows = spark.range(0, 10000, 1, 1) + rows + .selectExpr( + "id AS unused", + "id AS payload", + "CASE WHEN id = 0 THEN CAST(NULL AS INT) ELSE CAST(id AS INT) END AS probe_key") + .write + .option("parquet.block.size", "1024") + .parquet(firstPath.getCanonicalPath) + rows + .selectExpr("CAST(id AS INT) AS other_key", "id AS unused", "id AS other_payload") + .write + .option("parquet.block.size", "1024") + .parquet(secondPath.getCanonicalPath) + for (path <- Seq(firstPath, secondPath); + file <- path.listFiles().filter(_.getName.endsWith(".parquet"))) { + val reader = ParquetFileReader.open( + HadoopInputFile + .fromPath(new Path(file.getAbsolutePath), spark.sessionState.newHadoopConf())) + try assert(reader.getRowGroups.size() > 1) + finally reader.close() + } + withParquetTable(firstPath.getCanonicalPath, "union_probe_a") { + withParquetTable(secondPath.getCanonicalPath, "union_probe_b") { + withParquetTable( + Seq((Some(2500), 11L), (Some(2500), 12L), (Some(2600), 13L), (None, 14L)), + "union_build") { + val probe = "(SELECT probe_key AS k, payload AS v FROM union_probe_a " + + "UNION ALL SELECT other_key, other_payload FROM union_probe_b) p" + val from = + if (buildLeft) s"union_build b JOIN $probe ON b._1 = p.k" + else s"$probe JOIN union_build b ON p.k = b._1" + val query = s"SELECT /*+ BROADCAST(b) */ p.k, p.v, b._2 FROM $from" + var baselineBytes = 0L + var baselinePartitions = 0 + var baselineSchema: StructType = null + for (enabled <- Seq(false, true)) { + withSQLConf(CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_UNION_ENABLED.key -> + enabled.toString) { + val df = sql(query) + val (_, plan) = checkSparkAnswerAndOperator( + df, + Seq(classOf[CometBroadcastHashJoinExec], classOf[CometUnionExec]), + classOf[ReusedExchangeExec]) + val joins = collect(plan) { case j: CometBroadcastHashJoinExec => j } + assert(joins.size == 1, plan.toString) + assert( + joins.forall(_.buildSide == (if (buildLeft) BuildLeft else BuildRight))) + assert(joins.map(_.metrics("output_rows").value).sum == 6L) + // Both Union partitions build the three non-null rows in one batch. + val joinMetrics = joins.head.metrics + assert(joinMetrics("build_input_rows").value == 6L) + assert(joinMetrics("build_input_batches").value == 2L) + assert(joinMetrics("build_mem_used").value > 0L) + assert(joinMetrics("build_time").value > 0L) + val probeScans = collect(plan) { + case scan: CometNativeScanExec + if scan.output.exists(a => Set("probe_key", "other_key")(a.name)) => + scan + } + assert(probeScans.size == 2, plan.toString) + val bytes = probeScans.map(_.metrics("bytes_scanned").value).sum + val partitions = df.rdd.getNumPartitions + if (enabled) { + assert(df.schema == baselineSchema) + assert(partitions == baselinePartitions) + assert(unionFilterMetric(plan, "branches_attached") == 2L, plan.toString) + assert(unionFilterMetric(plan, "branches_skipped") == 0L) + assert( + probeScans.forall(_.metrics("row_groups_pruned_statistics").value > 0L)) + assert( + bytes > 0L && bytes < baselineBytes, + s"Expected less reader I/O: baseline=$baselineBytes, enabled=$bytes") + info(s"Union reader bytes: baseline=$baselineBytes, enabled=$bytes, " + + s"buildLeft=$buildLeft, AQE=$adaptive") + // The single original join and its broadcast remain above Union, even after AQE. + val broadcasts = collect(plan) { case e: CometBroadcastExchangeExec => e } + assert(broadcasts.nonEmpty, plan.toString) + assert(broadcasts.map(_.id).distinct.size == 1, plan.toString) + } else { + assert( + joins.head.metrics("dynamic_filter_join_filters_attached").value == 0L) + baselineBytes = bytes + baselinePartitions = partitions + baselineSchema = df.schema + } + } + } + } + } + } + } + } + } + } + } + + test("union join filter releases lazy JVM branches after early termination") { + withSQLConf( + CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_ENABLED.key -> "true", + CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_UNION_ENABLED.key -> "true", + CometConf.COMET_BROADCAST_REUSE_ENABLED.key -> "false", + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "false", + CometConf.COMET_SPARK_TO_ARROW_ENABLED.key -> "true", + CometConf.COMET_BATCH_SIZE.key -> "8192", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> + ("org.apache.spark.sql.catalyst.optimizer.PushDownPredicates," + + "org.apache.spark.sql.catalyst.optimizer.PushPredicateThroughNonJoin," + + "org.apache.spark.sql.catalyst.optimizer.PushProjectionThroughUnion")) { + val probe = "(SELECT id AS k, id+1 AS v FROM range(0,10000,1,1) " + + "UNION ALL SELECT id AS k, id+2 AS v FROM range(0,10000,1,1)) p" + // Keep a native consumer above Union, both with and without a transported domain. + val queries = Seq( + ( + s"SELECT /*+ BROADCAST(b) */ p.k, p.v FROM $probe " + + "JOIN (SELECT id AS k FROM range(0,10000,1,1)) b ON p.k = b.k", + 1), + (s"SELECT p.k, p.v FROM $probe WHERE p.v > 0", 1), + // Sort retains batches after the inner native Filter reaches EOF. + ( + s"SELECT u.k, u.v FROM (SELECT p.k, p.v FROM $probe WHERE p.v > 0 " + + "UNION ALL SELECT id, id+3 FROM range(0,10000,1,1)) u SORT BY k, v", + 2)) + for ((query, unionCount) <- queries) { + val before = CometArrowAllocator.getAllocatedMemory + val df = sql(query) + val plan = df.queryExecution.executedPlan + assert(collect(plan) { case _: CometUnionExec => + 1 + }.size == unionCount) + if (unionCount == 2) { + assert(collect(plan) { case _: CometSortExec => 1 }.nonEmpty, plan.toString) + } + assert(df.take(1).toSeq == Seq(Row(0L, 1L))) + assert(CometArrowAllocator.getAllocatedMemory == before) + } + } + } + + test("union join filter cancellation releases lazy branches before a fresh query") { + withSQLConf( + CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_ENABLED.key -> "true", + CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_UNION_ENABLED.key -> "true", + CometConf.COMET_BROADCAST_REUSE_ENABLED.key -> "false", + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "false", + CometConf.COMET_SPARK_TO_ARROW_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + CometJoinSuite.unionEntered = new CountDownLatch(1) + CometJoinSuite.unionRelease = new CountDownLatch(1) + spark.udf.register( + "union_cancel_gate", + (id: Long) => { + Predef.require(org.apache.spark.TaskContext.get() != null) + CometJoinSuite.unionEntered.countDown() + Predef.require(CometJoinSuite.unionRelease.await(30, TimeUnit.SECONDS)) + id + }) + def query(key: Long): DataFrame = sql( + "SELECT /*+ BROADCAST(b) */ p.k, p.v FROM (" + + "SELECT id % 2 AS k, union_cancel_gate(id) AS v FROM range(0,100,1,1) " + + "UNION ALL SELECT id % 2 AS k, union_cancel_gate(id) AS v " + + s"FROM range(0,100,1,1)) p JOIN " + + s"(SELECT id AS k FROM range($key,${key + 1},1,1)) b ON p.k = b.k") + val cancelled = query(0L) + val plan = cancelled.queryExecution.executedPlan + assert(collect(plan) { case _: CometUnionExec => 1 }.size == 1) + assert(collect(plan) { case join: CometBroadcastHashJoinExec => + join.nativeOp.getHashJoin.getUnionFilterTransportEnabled + } == Seq(true)) + val failure = new java.util.concurrent.atomic.AtomicReference[Throwable]() + val group = java.util.UUID.randomUUID().toString + val before = CometArrowAllocator.getAllocatedMemory + val runner = new Thread(() => { + spark.sparkContext.setJobGroup(group, "cancel lazy Union", interruptOnCancel = true) + try cancelled.collect() + catch { case error: Throwable => failure.set(error) } + finally spark.sparkContext.clearJobGroup() + }) + runner.setDaemon(true) + runner.start() + try { + assert(CometJoinSuite.unionEntered.await(30, TimeUnit.SECONDS)) + spark.sparkContext.cancelJobGroup(group) + runner.join(30000) + assert(!runner.isAlive) + assert(failure.get() != null) + } finally { + CometJoinSuite.unionRelease.countDown() + spark.sparkContext.cancelJobGroup(group) + runner.join(30000) + } + try { + // Driver cancellation can complete before executor task cleanup. Wait for that + // cleanup before starting a fresh query; a leaked branch still fails this bound. + eventually(timeout(30.seconds), interval(100.millis)) { + assert(CometArrowAllocator.getAllocatedMemory == before) + } + // A fresh build selects the other keys after the failed attempt's lease is released. + val rows = query(1L).collect() + assert(rows.length == 100) + assert(rows.forall(row => row.getLong(0) == 1L && row.getLong(1) % 2L == 1L)) + assert(CometArrowAllocator.getAllocatedMemory == before) + } finally { + CometJoinSuite.unionEntered = null + CometJoinSuite.unionRelease = null + } + } + } + + test("union join filter preserves a fallible branch and filters the eligible branch") { + withSQLConf( + CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_ENABLED.key -> "true", + SQLConf.ANSI_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.EXCHANGE_REUSE_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.LEAF_NODE_DEFAULT_PARALLELISM.key -> "1") { + withParquetTable(Seq((1, 10), (2, 20)), "union_safe_probe") { + withParquetTable(Seq(Tuple1(1)), "union_fallible_build") { + for (raw <- Seq("bad", "40")) { + withParquetTable(Seq((1, "30"), (2, raw)), "union_fallible_probe") { + val query = "SELECT /*+ BROADCAST(b) */ p.k, p.payload FROM " + + "(SELECT _1 AS k, _2 AS payload FROM union_safe_probe UNION ALL " + + "SELECT _1, CAST(_2 AS INT) FROM union_fallible_probe) p " + + "JOIN union_fallible_build b ON p.k = b._1" + for (enabled <- Seq(false, true)) { + withSQLConf(CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_UNION_ENABLED.key -> + enabled.toString) { + val df = sql(query) + val plan = if (raw == "bad") { + // Preserve existing native evaluation before join rejection. Spark's + // codegen can defer a non-key cast, so do not require its error here. + val before = CometArrowAllocator.getAllocatedMemory + val failure = intercept[Exception](df.collect()) + assert(causeChain(failure).exists(e => + Option(e.getMessage).exists(_.contains("CAST_INVALID_INPUT")))) + assert(CometArrowAllocator.getAllocatedMemory == before) + df.queryExecution.executedPlan + } else { + checkSparkAnswerAndOperator(df, classOf[ReusedExchangeExec])._2 + } + val joins = collect(plan) { case j: CometBroadcastHashJoinExec => j } + assert(joins.size == 1, plan.toString) + if (enabled && raw != "bad") { + assert(unionFilterMetric(plan, "branches_attached") == 1L) + assert(unionFilterMetric(plan, "branches_skipped") == 1L) + } + } + } + } + } + } + } + } + } + + test("union join filter retains branch limits before membership rejection") { + withSQLConf( + CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.LEAF_NODE_DEFAULT_PARALLELISM.key -> "1") { + withTempPath { path => + Seq((0, 10), (1, 20)).toDF("k", "payload").coalesce(1).write.parquet(path.toString) + withParquetTable(path.toString, "union_limit_probe") { + withParquetTable(Seq(Tuple1(1)), "union_limit_build") { + val query = "SELECT /*+ BROADCAST(b) */ p.k, p.payload FROM " + + "((SELECT * FROM union_limit_probe LIMIT 1) UNION ALL " + + "(SELECT * FROM union_limit_probe LIMIT 1)) p " + + "JOIN union_limit_build b ON p.k = b._1" + for (enabled <- Seq(false, true)) { + withSQLConf(CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_UNION_ENABLED.key -> + enabled.toString) { + val (_, plan) = + checkSparkAnswerAndOperator(sql(query), classOf[ReusedExchangeExec]) + assert(collect(plan) { case _: CometBroadcastHashJoinExec => 1 }.size == 1) + assert(unionFilterMetric(plan, "branches_attached") == 0L) + assert(collect(plan) { case scan: CometNativeScanExec => scan } + .forall(_.metrics("row_groups_pruned_statistics").value == 0L)) + } + } + } + } + } + } + } + + test("union join filter reaches sequential bucketed readers within one task") { + assume(isSpark41Plus, "Partition-aware Union execution requires Spark 4.1") + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.LEAF_NODE_DEFAULT_PARALLELISM.key -> "1", + // Inferred null checks change Filter output nullability while its bucket partitioning + // retains the scan attributes, so Spark cannot prove equal Union partitioning. + SQLConf.CONSTRAINT_PROPAGATION_ENABLED.key -> "false", + "spark.sql.sources.bucketing.enabled" -> "true", + "spark.sql.sources.bucketing.autoBucketedScan.enabled" -> "false", + "spark.sql.unionOutputPartitioning" -> "true", + CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_ENABLED.key -> "true", + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "false", + CometConf.COMET_BROADCAST_REUSE_ENABLED.key -> "false", + CometConf.COMET_BATCH_SIZE.key -> "16") { + withTable("union_bucket_a", "union_bucket_b") { + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + for (table <- Seq("union_bucket_a", "union_bucket_b")) { + spark + .range(0, 10000, 1, numPartitions = 1) + .selectExpr("CAST(id AS INT) AS k", "CAST(id AS INT) AS payload") + .coalesce(1) + .write + .format("parquet") + .option("parquet.block.size", "1024") + .bucketBy(1, "k") + .sortBy("k") + .saveAsTable(table) + } + } + withParquetTable(Seq(Tuple1(2500), Tuple1(2500), Tuple1(2600)), "union_bucket_build") { + val query = "SELECT /*+ BROADCAST(b) */ p.k, p.payload FROM " + + "(SELECT k, payload FROM union_bucket_a UNION ALL " + + "SELECT k, payload FROM union_bucket_b) p " + + "JOIN union_bucket_build b ON p.k = b._1" + val expectedSchema = withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sql(query).schema + } + var baselineBytes = 0L + for (enabled <- Seq(false, true)) { + withSQLConf( + CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_UNION_ENABLED.key -> + enabled.toString) { + val df = sql(query) + val (_, plan) = checkSparkAnswerAndOperator(df) + assert(df.schema == expectedSchema) + val joins = collect(plan) { case join: CometBroadcastHashJoinExec => join } + val unions = collect(plan) { case union: CometUnionExec => union } + assert(joins.size == 1 && unions.size == 1, plan.toString) + assert(joins.head.metrics("output_rows").value == 6L) + val union = unions.head + assert(union.children.size == 2) + assert( + union.children.forall(child => child.outputPartitioning.numPartitions == 1), + union.toString) + assert( + union.outputPartitioning.numPartitions == 1, + s"${union.outputPartitioning}; ${union.children.map(_.outputPartitioning)}") + assert(union.executeColumnar().getNumPartitions == 1) + val scans = collect(union) { case scan: CometNativeScanExec => scan } + assert(scans.size == 2, plan.toString) + assert( + scans.forall(scan => + scan.originalPlan.bucketedScan && + scan.outputPartitioning.numPartitions == 1), + plan.toString) + val bytes = scans.map(_.metrics("bytes_scanned").value).sum + if (enabled) { + assert(unionFilterMetric(plan, "branches_attached") == 2L) + assert(unionFilterMetric(plan, "branches_skipped") == 0L) + assert(scans.forall(_.metrics("row_groups_pruned_statistics").value > 0L)) + assert(bytes > 0L && bytes < baselineBytes) + info(s"Sequential Union reader bytes: baseline=$baselineBytes, enabled=$bytes") + } else baselineBytes = bytes + } + } + } + } + } + } + + test("union join filter forwards a completed domain through nested Union owners") { + withSQLConf( + CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_ENABLED.key -> "true", + CometConf.COMET_BROADCAST_REUSE_ENABLED.key -> "false", + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "false", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> + ("org.apache.spark.sql.catalyst.optimizer.PushDownPredicates," + + "org.apache.spark.sql.catalyst.optimizer.PushPredicateThroughNonJoin," + + "org.apache.spark.sql.catalyst.optimizer.CombineUnions")) { + withTempPath { path => + spark + .range(0, 10000, 1, 1) + .selectExpr("CAST(id AS INT) AS k", "id AS v") + .write + .option("parquet.block.size", "1024") + .parquet(path.toString) + withParquetTable(path.toString, "union_nested_input") { + withParquetTable(Seq(Tuple1(2500), Tuple1(2600)), "union_nested_keys") { + val query = "SELECT /*+ BROADCAST(b) */ p.k, p.v FROM (" + + "SELECT k, v FROM (SELECT k, v FROM union_nested_input UNION ALL " + + "SELECT k, v FROM union_nested_input) inner_union WHERE v IS NOT NULL " + + "UNION ALL SELECT k, v FROM union_nested_input) p " + + "JOIN union_nested_keys b ON p.k = b._1" + var baselineBytes = 0L + for (enabled <- Seq(false, true)) { + withSQLConf(CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_UNION_ENABLED.key -> + enabled.toString) { + val (_, plan) = checkSparkAnswerAndOperator(sql(query)) + assert(collect(plan) { case _: CometUnionExec => 1 }.size == 2, plan.toString) + assert(collect(plan) { case _: CometBroadcastHashJoinExec => 1 }.size == 1) + val scans = collect(plan) { + case scan: CometNativeScanExec if scan.output.exists(_.name == "v") => scan + } + assert(scans.size == 3) + val bytes = scans.map(_.metrics("bytes_scanned").value).sum + if (enabled) { + assert(unionFilterMetric(plan, "branches_forwarded") == 2L) + assert(unionFilterMetric(plan, "branches_attached") == 3L) + assert(unionFilterMetric(plan, "branches_skipped") == 0L) + assert(bytes > 0 && bytes < baselineBytes) + } else baselineBytes = bytes + } + } + } + } + } + } + } + + test("union join filter preserves nested results across an intermediate join boundary") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.LEAF_NODE_DEFAULT_PARALLELISM.key -> "1", + "spark.sql.unionOutputPartitioning" -> "false", + CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_ENABLED.key -> "true", + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "false", + CometConf.COMET_BROADCAST_REUSE_ENABLED.key -> "false", + CometConf.COMET_BATCH_SIZE.key -> "16") { + withTempPath { path => + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark + .range(0, 10000, 1, numPartitions = 1) + .selectExpr("CAST(id AS INT) AS k", "CAST(id AS INT) AS payload") + .write + .option("parquet.block.size", "1024") + .parquet(path.getCanonicalPath) + } + withParquetTable(path.getCanonicalPath, "nested_union_probe") { + withParquetTable(Seq(Tuple1(2500), Tuple1(2600)), "nested_union_inner_build") { + withParquetTable( + Seq(Tuple1(2600), Tuple1(2600), Tuple1(2700)), + "nested_union_outer_build") { + val query = "SELECT /*+ BROADCAST(ob) */ u.k, u.payload FROM (" + + "SELECT /*+ BROADCAST(ib) */ p.k, p.payload FROM " + + "(SELECT k, payload FROM nested_union_probe UNION ALL " + + "SELECT k, payload FROM nested_union_probe) p " + + "JOIN nested_union_inner_build ib ON p.k = ib._1 UNION ALL " + + "SELECT k, payload FROM nested_union_probe) u " + + "JOIN nested_union_outer_build ob ON u.k = ob._1" + val expected = Seq.fill(6)(Row(2600, 2600)) :+ Row(2700, 2700) + val schema = withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sql(query).schema + } + var baselineBytes = 0L + for (enabled <- Seq(false, true)) { + withSQLConf(CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_UNION_ENABLED.key -> + enabled.toString) { + val df = sql(query) + checkAnswer(df, expected) + assert(df.schema == schema) + val plan = df.queryExecution.executedPlan + val joins = collect(plan) { case join: CometBroadcastHashJoinExec => join } + val unions = collect(plan) { case union: CometUnionExec => union } + assert(joins.size == 2 && unions.size == 2, plan.toString) + val outer = unions + .find(union => + collect(union) { case _: CometUnionExec => + 1 + }.size == 2) + .get + assert(outer.executeColumnar().getNumPartitions == 3) + val scans = collect(outer) { + case scan: CometNativeScanExec if scan.output.exists(_.name == "payload") => + scan + } + assert(scans.size == 3, plan.toString) + val bytes = scans.map(_.metrics("bytes_scanned").value).sum + if (enabled) { + // The outer domain stays above the intermediate join. Its own domain + // reaches the two inner readers; the third branch takes the outer domain. + assert(unionFilterMetric(plan, "branches_attached") == 3L) + assert(unionFilterMetric(plan, "branches_forwarded") == 0L) + assert(unionFilterMetric(plan, "branches_skipped") == 2L) + assert(scans.forall(_.metrics("row_groups_pruned_statistics").value > 0L)) + assert(bytes > 0L && bytes < baselineBytes) + info(s"Nested Union reader bytes: baseline=$baselineBytes, enabled=$bytes") + } else baselineBytes = bytes + } + } + } + } + } + } + } + } + test("join dynamic filter prunes a projected Parquet reader through null-check conjunctions") { withTempPath { probePath => withSQLConf( @@ -2007,4 +2526,6 @@ class CometJoinSuite extends CometTestBase { object CometJoinSuite { @volatile var probeGate: CountDownLatch = _ + @volatile var unionEntered: CountDownLatch = _ + @volatile var unionRelease: CountDownLatch = _ } From 234ebbbc0caa2820f291579650eed3f58dee1116 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sun, 4 Oct 2026 17:55:34 +0000 Subject: [PATCH 7/7] Clear strict Scala warnings in lazy Union input --- .../org/apache/spark/sql/comet/CometUnionInput.scala | 12 ++++++++---- .../scala/org/apache/spark/sql/comet/operators.scala | 2 +- .../scala/org/apache/comet/exec/CometJoinSuite.scala | 2 +- 3 files changed, 10 insertions(+), 6 deletions(-) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometUnionInput.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometUnionInput.scala index 1903e93a741..41536a1c01b 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometUnionInput.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometUnionInput.scala @@ -112,8 +112,12 @@ object CometUnionInput { /** Lazy branch resources outlive the outer native plan and its exported output batches. */ def addCleanup(context: TaskContext)(close: => Unit): Unit = { val input = current.get() - if (input != null) input.children += (() => close) - else context.addTaskCompletionListener[Unit](_ => close) + if (input != null) { + input.children += (() => close) + } else { + context.addTaskCompletionListener[Unit](_ => close) + } + () } def registerExecution(plan: Long): Unit = { @@ -130,7 +134,7 @@ object CometUnionInput { private[comet] class CometUnionInputRDD( batches: RDD[ColumnarBatch], schema: StructType, - name: String, + sourceName: String, nativeRootPlanIds: Array[Long]) extends RDD[CometUnionInput](batches.context, Seq(new OneToOneDependency(batches))) { @@ -145,7 +149,7 @@ private[comet] class CometUnionInputRDD( new CometUnionInput( () => parent.iterator(split, context), schema, - name, + sourceName, context, nativeRootPlanIds)) } 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 2c6b080a988..c4238274d66 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 @@ -1053,7 +1053,7 @@ abstract class CometNativeExec extends CometExec { if CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_ENABLED.get(conf) && CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_UNION_ENABLED.get(conf) => def branchRoots(branch: SparkPlan): Seq[Long] = branch match { - case native: CometNativeExec => Seq(native.nativeOp.getPlanId) + case native: CometNativeExec => Seq(native.nativeOp.getPlanId.toLong) case nested: CometUnionExec => nested.children.flatMap(branchRoots) case _ => Seq.empty } 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 24f50a6d9f6..b97e831ee3f 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -619,7 +619,7 @@ class CometJoinSuite extends CometTestBase { "SELECT /*+ BROADCAST(b) */ p.k, p.v FROM (" + "SELECT id % 2 AS k, union_cancel_gate(id) AS v FROM range(0,100,1,1) " + "UNION ALL SELECT id % 2 AS k, union_cancel_gate(id) AS v " + - s"FROM range(0,100,1,1)) p JOIN " + + "FROM range(0,100,1,1)) p JOIN " + s"(SELECT id AS k FROM range($key,${key + 1},1,1)) b ON p.k = b.k") val cancelled = query(0L) val plan = cancelled.queryExecution.executedPlan