From 705fc113f4e8f59d11ecea875d1c70378932d191 Mon Sep 17 00:00:00 2001 From: msaf Date: Mon, 5 Oct 2026 08:46:57 +0100 Subject: [PATCH 01/14] fix: release each fetched shuffle block once and stop a killed task between blocks The coalesced and direct shuffle reads joined the fetched blocks with a SequenceInputStream. Its close() walks every remaining block, calling ShuffleBlockFetcherIterator.next() for each. When the task ends early, the fetcher's completion listener has already released the buffers in its results queue without removing them, so next() hands out a released buffer and closing its stream releases it again ("refCnt: 0, decrement: 1" from the reader's completion listener). With blocks still in flight, close() instead waits for them, and after an interrupt the channel runs that close() in the interrupting thread, so a killed task could not be stopped and the task reaper killed the executor. The blocks are now read by a stream that closes each block before fetching the next, never fetches on close, and checks for a kill at every block boundary. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../CometBlockStoreShuffleReader.scala | 87 +++++++- .../CometBlockStoreShuffleReaderSuite.scala | 198 ++++++++++++++++++ 2 files changed, 277 insertions(+), 8 deletions(-) create mode 100644 spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReaderSuite.scala diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala index 25b61634ac9..099cbee6bf7 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala @@ -19,7 +19,7 @@ package org.apache.spark.sql.comet.execution.shuffle -import java.io.InputStream +import java.io.{InputStream, IOException} import org.apache.spark.{InterruptibleIterator, MapOutputTracker, SparkEnv, TaskContext} import org.apache.spark.internal.{config, Logging} @@ -168,13 +168,8 @@ class CometBlockStoreShuffleReader[K, C]( * Returns the raw concatenated InputStream of all shuffle blocks, bypassing the decode step. * Used by ShuffleScan direct read path. */ - override def readAsRawStream(): InputStream = { - val streams = fetchIterator.map(_._2) - new java.io.SequenceInputStream(new java.util.Enumeration[InputStream] { - override def hasMoreElements: Boolean = streams.hasNext - override def nextElement(): InputStream = streams.next() - }) - } + override def readAsRawStream(): InputStream = + CometBlockStoreShuffleReader.concatenate(context, fetchIterator.map(_._2)) private def fetchContinuousBlocksInBatch: Boolean = { val conf = SparkEnv.get.conf @@ -204,3 +199,79 @@ class CometBlockStoreShuffleReader[K, C]( doBatchFetch } } + +private[shuffle] object CometBlockStoreShuffleReader { + def concatenate(context: TaskContext, blocks: Iterator[InputStream]): InputStream = + new FetchedBlocksInputStream(context, blocks) +} + +/** + * A block's stream releases the fetcher's current buffer when closed, so each block is closed + * before the next is fetched, and close() never fetches: the fetcher releases unread blocks. + */ +private final class FetchedBlocksInputStream(context: TaskContext, blocks: Iterator[InputStream]) + extends InputStream { + + private var current: InputStream = null + private var closed = false + + private def block(): InputStream = synchronized { + if (closed) { + throw new IOException("Shuffle block stream is closed") + } + if (current == null) { + context.killTaskIfInterrupted() + if (blocks.hasNext) { + current = blocks.next() + } + } + current + } + + private def finish(finished: InputStream): Unit = synchronized { + if (current eq finished) { + current = null + finished.close() + } + } + + override def read(): Int = { + var in = block() + while (in != null) { + val value = in.read() + if (value >= 0) { + return value + } + finish(in) + in = block() + } + -1 + } + + override def read(buffer: Array[Byte], offset: Int, length: Int): Int = { + if (length == 0) { + return 0 + } + var in = block() + while (in != null) { + val count = in.read(buffer, offset, length) + if (count >= 0) { + return count + } + finish(in) + in = block() + } + -1 + } + + override def close(): Unit = synchronized { + if (!closed) { + closed = true + val open = current + current = null + if (open != null) { + open.close() + } + } + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReaderSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReaderSuite.scala new file mode 100644 index 00000000000..5e8bc12eb1a --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReaderSuite.scala @@ -0,0 +1,198 @@ +/* + * 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.shuffle + +import java.io.InputStream +import java.nio.ByteBuffer +import java.nio.channels.Channels +import java.util.concurrent.{ConcurrentLinkedQueue, CountDownLatch, TimeUnit} + +import scala.collection.JavaConverters._ +import scala.concurrent.{Await, ExecutionContext, Future} +import scala.concurrent.duration._ + +import org.apache.spark.{SparkEnv, TaskContext, TaskKilledException} +import org.apache.spark.network.buffer.NettyManagedBuffer +import org.apache.spark.network.shuffle.{BlockFetchingListener, BlockStoreClient, DownloadFileManager} +import org.apache.spark.sql.CometTestBase +import org.apache.spark.storage.{BlockId, BlockManagerId, ShuffleBlockFetcherIterator, ShuffleBlockId} + +import io.netty.buffer.{ByteBuf, Unpooled} + +class CometBlockStoreShuffleReaderSuite extends CometTestBase { + + private val blockSize = 64 + private val remote = BlockManagerId("remote-exec", "remote-host", 7337) + private implicit val ec: ExecutionContext = ExecutionContext.global + + private class RemoteBlocks(delivered: Int => Boolean) extends BlockStoreClient { + val buffers = new ConcurrentLinkedQueue[ByteBuf]() + + override def fetchBlocks( + host: String, + port: Int, + execId: String, + blockIds: Array[String], + listener: BlockFetchingListener, + downloadFileManager: DownloadFileManager): Unit = { + blockIds.foreach { id => + val index = BlockId(id).asInstanceOf[ShuffleBlockId].mapId.toInt + if (delivered(index)) { + val buf = Unpooled.wrappedBuffer(Array.fill(blockSize)(index.toByte)) + buffers.add(buf) + listener.onBlockFetchSuccess(id, new NettyManagedBuffer(buf)) + buf.release() + } + } + } + + override def close(): Unit = {} + + def leaked: Seq[Int] = buffers.asScala.map(_.refCnt()).filter(_ != 0).toSeq + } + + private def blockStream( + context: TaskContext, + client: BlockStoreClient, + blocks: Int): InputStream = { + val fetcher = new ShuffleBlockFetcherIterator( + context, + client, + SparkEnv.get.blockManager, + SparkEnv.get.mapOutputTracker, + Iterator( + ( + remote, + (0 until blocks).map(i => (ShuffleBlockId(0, i, 0): BlockId, blockSize.toLong, i)))), + (_, in) => in, + 48L * 1024 * 1024, + Int.MaxValue, + Int.MaxValue, + Long.MaxValue, + 10, + false, + false, + false, + "ADLER32", + context.taskMetrics().createTempShuffleReadMetrics(), + false) + CometBlockStoreShuffleReader.concatenate(context, fetcher.toCompletionIterator.map(_._2)) + } + + private def readFully(in: InputStream, bytes: Int): Array[Byte] = { + val out = new Array[Byte](bytes) + var read = 0 + while (read < bytes) { + val n = in.read(out, read, bytes - read) + assert(n > 0) + read += n + } + out + } + + private def within[T](f: => T): T = Await.result(Future(f), 20.seconds) + + test("closing after the fetcher's cleanup releases every fetched buffer once") { + val context = TaskContext.empty() + val client = new RemoteBlocks(_ => true) + var stream: InputStream = null + context.addTaskCompletionListener[Unit](_ => stream.close()) + stream = blockStream(context, client, 4) + assert(stream.read() == 0) + context.markTaskCompleted(None) + assert(client.buffers.size == 4) + assert(client.leaked.isEmpty) + } + + test("closing after a failure in a later block releases every fetched buffer once") { + val context = TaskContext.empty() + val client = new RemoteBlocks(_ => true) + var stream: InputStream = null + context.addTaskCompletionListener[Unit](_ => stream.close()) + stream = blockStream(context, client, 5) + assert(readFully(stream, blockSize * 2 + 3).last == 2) + context.markTaskCompleted(Some(new RuntimeException("downstream failure"))) + assert(client.leaked.isEmpty) + } + + test("closing does not fetch the blocks that were not read") { + val context = TaskContext.empty() + val client = new RemoteBlocks(_ == 0) + val stream = blockStream(context, client, 3) + assert(stream.read() == 0) + within(stream.close()) + context.markTaskCompleted(None) + assert(client.leaked.isEmpty) + } + + test("a killed task stops reading at the next block") { + val context = TaskContext.empty() + val client = new RemoteBlocks(_ => true) + val stream = blockStream(context, client, 3) + assert(readFully(stream, blockSize).forall(_ == 0)) + context.markInterrupted("killed by test") + intercept[TaskKilledException](stream.read()) + context.markTaskCompleted(None) + within(stream.close()) + assert(client.leaked.isEmpty) + } + + test("a killed task does not wait for a block that never arrives") { + val context = TaskContext.empty() + val client = new RemoteBlocks(_ == 0) + val stream = blockStream(context, client, 2) + assert(readFully(stream, blockSize).forall(_ == 0)) + context.markInterrupted("killed by test") + within(intercept[TaskKilledException](stream.read())) + within(stream.close()) + context.markTaskCompleted(None) + assert(client.leaked.isEmpty) + } + + test("interrupting a reader blocked on a missing block does not block the interrupter") { + val context = TaskContext.empty() + val client = new RemoteBlocks(_ == 0) + val stream = blockStream(context, client, 3) + val channel = Channels.newChannel(stream) + val started = new CountDownLatch(1) + @volatile var failure: Throwable = null + val reader = new Thread(() => { + val buf = ByteBuffer.allocate(blockSize) + try { + while (buf.hasRemaining) channel.read(buf) + started.countDown() + buf.clear() + channel.read(buf) + } catch { + case t: Throwable => failure = t + } + }) + reader.setDaemon(true) + reader.start() + assert(started.await(20, TimeUnit.SECONDS)) + Thread.sleep(200) + within(reader.interrupt()) + reader.join(20000) + assert(!reader.isAlive) + assert(failure != null) + context.markTaskCompleted(None) + assert(client.leaked.isEmpty) + } +} From dfa6439e5e51ea95edcdb070d377512529ef3b44 Mon Sep 17 00:00:00 2001 From: msaf Date: Mon, 5 Oct 2026 09:16:58 +0100 Subject: [PATCH 02/14] fix: keep 32-bit offsets of repeated and coalesced nested columns from overflowing Explode repeats every column it does not unnest with take, bounding a build only by its row count, and never splits one input row. A row whose nested strings are large, repeated once per element of a long list, filled more than i32::MAX bytes of one offset buffer, e.g. a proposal's status_history[].merchant_variant_prices[].variant_id repeated for each of its target variants ("struct column 2 ("merchant_variant_prices") failed: ... offset overflow"). Explode now also bounds each build by what it adds to every 32-bit offset buffer of the repeated columns, at any depth, and unnests a row that alone would overflow one in parts, narrowing its lists to the elements each part outputs. The check costs one pass over the offset nodes per input batch; rows are only measured when the longest output length times a whole column could overflow. The shuffle read coalescer and the window's output coalescing joined batches by row count alone, so concatenating them could overflow the same way. They now stop joining before any offset buffer would. Co-Authored-By: Claude Opus 5.5 (1M context) --- native/common/src/lib.rs | 1 + native/common/src/offset_extents.rs | 230 ++++++++++ .../core/src/execution/operators/explode.rs | 400 +++++++++++++++++- .../operators/partition_aggregate_window.rs | 11 + native/shuffle/src/read_coalescer.rs | 96 +++++ 5 files changed, 725 insertions(+), 13 deletions(-) create mode 100644 native/common/src/offset_extents.rs diff --git a/native/common/src/lib.rs b/native/common/src/lib.rs index 12e9ac8a5ee..78521fad6d5 100644 --- a/native/common/src/lib.rs +++ b/native/common/src/lib.rs @@ -16,6 +16,7 @@ // under the License. mod error; +pub mod offset_extents; mod query_context; mod schema; pub mod struct_nulls; diff --git a/native/common/src/offset_extents.rs b/native/common/src/offset_extents.rs new file mode 100644 index 00000000000..3775a5213b6 --- /dev/null +++ b/native/common/src/offset_extents.rs @@ -0,0 +1,230 @@ +// 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. + +//! How much of each 32-bit offset buffer (`Utf8`, `Binary`, `List` and `Map` at any depth) an +//! array addresses, so that `concat` or a repeating `take` can be fed few enough rows that no +//! offset buffer of the result exceeds `i32::MAX`. Nodes are visited in a pre-order fixed by +//! the data type, so the extents of two arrays of one type line up position by position. + +use arrow::array::{Array, AsArray}; +use arrow::datatypes::DataType; + +/// The largest extent a 32-bit offset buffer can address. +pub const MAX_OFFSET_EXTENT: usize = i32::MAX as usize; + +/// The extent of every 32-bit offset node of `array`, in pre-order. Dictionary values are +/// included whole, since concatenating dictionaries can concatenate their values. +pub fn offset_extents(array: &dyn Array, extents: &mut Vec) { + visit(array, &[0, array.len()], true, &mut |bounds| { + extents.push(bounds[1] - bounds[0]) + }); +} + +/// Raises `row_extents[i]` to the largest extent row `i` of `array` has in any of its 32-bit +/// offset nodes. Repeating row `i` `n` times adds at most `n * row_extents[i]` to every such +/// node. Dictionary values are skipped: repeating keys does not copy them. +pub fn max_row_offset_extents(array: &dyn Array, row_extents: &mut [usize]) { + assert_eq!(array.len(), row_extents.len()); + let rows: Vec = (0..=array.len()).collect(); + visit(array, &rows, false, &mut |bounds| { + for (row, extent) in row_extents.iter_mut().enumerate() { + *extent = (*extent).max(bounds[row + 1] - bounds[row]); + } + }); +} + +/// Accumulates the extents of the batches that would be concatenated, to tell whether one more +/// keeps every offset node within [`MAX_OFFSET_EXTENT`]. +#[derive(Debug)] +pub struct OffsetBudget { + limit: usize, + used: Vec, + scratch: Vec, +} + +impl Default for OffsetBudget { + fn default() -> Self { + Self::with_limit(MAX_OFFSET_EXTENT) + } +} + +impl OffsetBudget { + pub fn with_limit(limit: usize) -> Self { + Self { + limit, + used: vec![], + scratch: vec![], + } + } + + /// Adds the columns' extents if they fit beside those already added, and reports whether + /// they did. An empty budget accepts anything. + pub fn try_add(&mut self, columns: &[impl AsRef]) -> bool { + self.scratch.clear(); + for column in columns { + offset_extents(column.as_ref(), &mut self.scratch); + } + if self.used.is_empty() { + std::mem::swap(&mut self.used, &mut self.scratch); + return true; + } + if self.used.len() != self.scratch.len() + || self + .used + .iter() + .zip(&self.scratch) + .any(|(used, extent)| used + extent > self.limit) + { + return false; + } + for (used, extent) in self.used.iter_mut().zip(&self.scratch) { + *used += extent; + } + true + } + + pub fn clear(&mut self) { + self.used.clear(); + } +} + +/// Calls `f` with the bounds, in the node's own index space, that `bounds` of `array` cover in +/// every 32-bit offset node below it. +fn visit(array: &dyn Array, bounds: &[usize], dictionaries: bool, f: &mut dyn FnMut(&[usize])) { + let through = + |offsets: &[i32]| -> Vec { bounds.iter().map(|&b| offsets[b] as usize).collect() }; + match array.data_type() { + DataType::Utf8 => f(&through(array.as_string::().value_offsets())), + DataType::Binary => f(&through(array.as_binary::().value_offsets())), + DataType::List(_) => { + let list = array.as_list::(); + let child = through(list.value_offsets()); + f(&child); + visit(list.values().as_ref(), &child, dictionaries, f); + } + DataType::Map(_, _) => { + let map = array.as_map(); + let child = through(map.value_offsets()); + f(&child); + visit(map.entries(), &child, dictionaries, f); + } + DataType::LargeList(_) => { + let list = array.as_list::(); + let offsets = list.value_offsets(); + let child: Vec = bounds.iter().map(|&b| offsets[b] as usize).collect(); + visit(list.values().as_ref(), &child, dictionaries, f); + } + DataType::FixedSizeList(_, size) => { + let size = *size as usize; + let child: Vec = bounds.iter().map(|&b| b * size).collect(); + visit( + array.as_fixed_size_list().values().as_ref(), + &child, + dictionaries, + f, + ); + } + DataType::Struct(_) => { + for column in array.as_struct().columns() { + visit(column.as_ref(), bounds, dictionaries, f); + } + } + DataType::Dictionary(_, _) if dictionaries => { + let values = array.as_any_dictionary().values(); + visit(values.as_ref(), &[0, values.len()], dictionaries, f); + } + _ => {} + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{ + ArrayRef, DictionaryArray, Int32Array, ListArray, MapArray, StringArray, StructArray, + }; + use arrow::buffer::OffsetBuffer; + use arrow::datatypes::{Field, Fields, Int32Type}; + use std::sync::Arc; + + fn nested() -> ArrayRef { + let names: ArrayRef = Arc::new(StringArray::from(vec!["a", "bb", "", "cccc", "dd"])); + let ids: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])); + let fields = Fields::from(vec![ + Field::new("name", DataType::Utf8, true), + Field::new("id", DataType::Int32, true), + ]); + let element = StructArray::new(fields.clone(), vec![names, ids], None); + Arc::new(ListArray::new( + Arc::new(Field::new_list_field(DataType::Struct(fields), true)), + OffsetBuffer::new(vec![0, 2, 2, 5].into()), + Arc::new(element), + None, + )) + } + + #[test] + fn extents_cover_every_offset_node_of_a_sliced_nested_array() { + let mut extents = vec![]; + offset_extents(nested().as_ref(), &mut extents); + assert_eq!(extents, vec![5, 9]); + + let mut extents = vec![]; + offset_extents(nested().slice(1, 2).as_ref(), &mut extents); + assert_eq!(extents, vec![3, 6]); + } + + #[test] + fn row_extents_take_the_largest_node_per_row() { + let mut rows = vec![0; 3]; + max_row_offset_extents(nested().as_ref(), &mut rows); + assert_eq!(rows, vec![3, 0, 6]); + + let mut rows = vec![0; 2]; + max_row_offset_extents(nested().slice(1, 2).as_ref(), &mut rows); + assert_eq!(rows, vec![0, 6]); + } + + #[test] + fn maps_and_dictionaries() { + let values = Int32Array::from(vec![1, 2]); + let map = + MapArray::new_from_strings(["k1", "k22"].into_iter(), &values, &[0, 1, 2]).unwrap(); + let mut extents = vec![]; + offset_extents(&map, &mut extents); + assert_eq!(extents, vec![2, 5]); + + let dict: DictionaryArray = vec!["xyz", "xyz", "q"].into_iter().collect(); + let mut extents = vec![]; + offset_extents(&dict, &mut extents); + assert_eq!(extents, vec![4]); + let mut rows = vec![0; 3]; + max_row_offset_extents(&dict, &mut rows); + assert_eq!(rows, vec![0, 0, 0]); + } + + #[test] + fn budget_refuses_what_would_overflow_and_restarts_when_cleared() { + let column = |s: &str| -> Vec { vec![Arc::new(StringArray::from(vec![s]))] }; + let mut budget = OffsetBudget::with_limit(10); + assert!(budget.try_add(&column("aaaaaa"))); + assert!(budget.try_add(&column("aaaa"))); + assert!(!budget.try_add(&column("a"))); + budget.clear(); + assert!(budget.try_add(&column("aaaaaaaaaaaaaaa"))); + } +} diff --git a/native/core/src/execution/operators/explode.rs b/native/core/src/execution/operators/explode.rs index 763a45eb36c..b6e630841b6 100644 --- a/native/core/src/execution/operators/explode.rs +++ b/native/core/src/execution/operators/explode.rs @@ -60,6 +60,7 @@ use arrow::array::{ LargeListArray, LargeListViewArray, ListArray, ListViewArray, PrimitiveArray, Scalar, StructArray, }; +use arrow::buffer::OffsetBuffer; use arrow::compute::kernels::length::length; use arrow::compute::kernels::zip::zip; use arrow::compute::{cast, is_not_null, kernels, sum}; @@ -89,6 +90,9 @@ use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, EmptyRecordBatchStream, ExecutionPlan, ExecutionPlanProperties, PlanProperties, RecordBatchStream, SendableRecordBatchStream, }; +use datafusion_comet_common::offset_extents::{ + max_row_offset_extents, offset_extents, MAX_OFFSET_EXTENT, +}; use futures::{Stream, StreamExt}; use std::cmp::{self, Ordering}; use std::pin::Pin; @@ -106,6 +110,7 @@ pub struct ExplodeExec { options: UnnestOptions, metrics: ExecutionPlanMetricsSet, cache: Arc, + offset_limit: usize, } impl ExplodeExec { @@ -131,9 +136,16 @@ impl ExplodeExec { options, metrics: ExecutionPlanMetricsSet::new(), cache: Arc::new(cache), + offset_limit: MAX_OFFSET_EXTENT, }) } + #[cfg(test)] + fn with_offset_limit(mut self, offset_limit: usize) -> Self { + self.offset_limit = offset_limit; + self + } + /// Compute the plan properties, keeping whatever the child guarantees about the columns /// that unnesting passes through untouched. /// @@ -273,6 +285,7 @@ impl ExecutionPlan for ExplodeExec { input_batches: MetricBuilder::new(&self.metrics).counter("input_batches", partition), input_rows: MetricBuilder::new(&self.metrics).counter("input_rows", partition), batch_size, + offset_limit: self.offset_limit, pending_input: None, }); @@ -308,6 +321,13 @@ struct PendingInput { /// case the whole remaining input is unnested in one call and only the output is split. /// See [`ExplodeStream::predict_output_lens`]. output_lens: Option>, + /// Per input row, the largest extent it has in any 32-bit offset buffer of the columns that + /// unnesting repeats, when repeating them as often as the input asks could overflow one. + /// See [`ExplodeStream::repeated_row_extents`]. + row_extents: Option>, + /// Output rows of the row at `row_offset` already unnested, when that row alone repeats + /// too much to be unnested in one build. + element_offset: usize, } impl PendingInput { @@ -316,29 +336,87 @@ impl PendingInput { } /// How many input rows to unnest next so the resulting batch holds at most `batch_size` - /// rows. + /// rows and repeats the other columns no further than `offset_limit`. /// /// Always returns at least 1 while rows remain, so the stream always makes progress: a - /// single input row is never split across output batches, so one row whose array is - /// longer than `batch_size` still produces one oversized build, which `BatchSplitStream` - /// slices down on the way out. - fn next_chunk_rows(&self, batch_size: usize) -> usize { + /// single input row is only split across builds when repeating it would overflow an offset + /// buffer (see [`Self::row_split`]), so one row whose array is longer than `batch_size` + /// still produces one oversized build, which `BatchSplitStream` slices down on the way out. + fn next_chunk_rows(&self, batch_size: usize, offset_limit: usize) -> usize { let Some(output_lens) = &self.output_lens else { return self.remaining_rows(); }; let lens = &output_lens.values()[self.row_offset..]; + let extents = self.row_extents.as_deref().map(|e| &e[self.row_offset..]); let batch_size = batch_size as i64; let mut output_rows = 0i64; + let mut repeated = 0usize; for (rows, len) in lens.iter().enumerate() { - if rows > 0 && output_rows + len > batch_size { + let added = extents.map_or(0, |e| e[rows].saturating_mul(*len as usize)); + if rows > 0 && (output_rows + len > batch_size || repeated + added > offset_limit) { return rows; } output_rows += len; + repeated += added; } lens.len() } + /// How many output rows of the row at `row_offset` to unnest per build, when repeating the + /// other columns for all of them at once would overflow a 32-bit offset buffer. + fn row_split(&self, offset_limit: usize) -> Option { + let extent = self.row_extents.as_ref()?[self.row_offset]; + let len = self.output_lens.as_ref()?.value(self.row_offset) as usize; + (self.element_offset > 0 || extent.saturating_mul(len) > offset_limit) + .then(|| (offset_limit / extent.max(1)).max(1)) + } + + /// The next `step` output rows of the row at `row_offset`, as that row with each unnested + /// list narrowed to the elements those output rows come from, and their count. + fn next_row_part( + &mut self, + step: usize, + list_columns: &[ListUnnest], + ) -> Result<(RecordBatch, PrimitiveArray)> { + let len = self.output_lens.as_ref().unwrap().value(self.row_offset) as usize; + let start = self.element_offset; + let end = (start + step).min(len); + let row = self.batch.slice(self.row_offset, 1); + let mut columns = row.columns().to_vec(); + for unnest in list_columns { + let list = columns[unnest.index_in_input_schema].as_list::(); + let DataType::List(field) = list.data_type() else { + return internal_err!("expected a List column to unnest, got {}", list.data_type()); + }; + let offsets = list.value_offsets(); + let (first, last) = (offsets[0] as usize, offsets[1] as usize); + let narrowed = ListArray::try_new( + Arc::clone(field), + OffsetBuffer::new( + vec![ + (first + start).min(last) as i32, + (first + end).min(last) as i32, + ] + .into(), + ), + Arc::clone(list.values()), + list.nulls().cloned(), + )?; + columns[unnest.index_in_input_schema] = Arc::new(narrowed); + } + if end == len { + self.row_offset += 1; + self.element_offset = 0; + } else { + self.element_offset = end; + } + Ok(( + RecordBatch::try_new(row.schema(), columns)?, + PrimitiveArray::from(vec![(end - start) as i64]), + )) + } + /// The per-row output lengths covering the next `rows` input rows, so the unnesting does /// not have to recompute what [`ExplodeStream::predict_output_lens`] already derived. fn chunk_lengths(&self, rows: usize) -> Option> { @@ -360,6 +438,8 @@ struct ExplodeStream { input_rows: Count, /// Target number of rows per output batch, from `datafusion.execution.batch_size`. batch_size: usize, + /// The extent a build may give any 32-bit offset buffer of the columns it repeats. + offset_limit: usize, /// Rows of the current input batch that have not been unnested yet. Unnesting one input /// batch can produce arbitrarily many output rows, so the input is consumed in chunks /// small enough that each chunk's output stays near `batch_size`. @@ -367,7 +447,8 @@ struct ExplodeStream { /// Note the scope of the memory bound this buys: chunking removes the input batch size /// from the peak, but not the length of an individual list. A single row whose list is /// longer than `batch_size`, and recursive unnesting (where the expansion cannot be - /// predicted up front), both still materialize their full expansion in one build. + /// predicted up front), both still materialize their full expansion in one build, unless + /// that build would overflow an offset buffer of a repeated column. pending_input: Option, } @@ -394,10 +475,20 @@ impl ExplodeStream { // always consumes at least one row, so it is dropped the moment it drains. debug_assert!(pending.remaining_rows() > 0); - let rows = pending.next_chunk_rows(self.batch_size); - let chunk = pending.batch.slice(pending.row_offset, rows); - let chunk_lengths = pending.chunk_lengths(rows); - pending.row_offset += rows; + let (chunk, chunk_lengths) = match pending.row_split(self.offset_limit) { + Some(step) => { + let (part, lengths) = + pending.next_row_part(step, &self.list_type_columns)?; + (part, Some(lengths)) + } + None => { + let rows = pending.next_chunk_rows(self.batch_size, self.offset_limit); + let chunk = pending.batch.slice(pending.row_offset, rows); + let chunk_lengths = pending.chunk_lengths(rows); + pending.row_offset += rows; + (chunk, chunk_lengths) + } + }; let drained = pending.remaining_rows() == 0; let timer = self.baseline_metrics.elapsed_compute().timer(); @@ -437,12 +528,15 @@ impl ExplodeStream { self.input_rows.add(batch.num_rows()); if batch.num_rows() > 0 { let timer = self.baseline_metrics.elapsed_compute().timer(); - let output_lens = self.predict_output_lens(&batch); + let output_lens = self.predict_output_lens(&batch)?; + let row_extents = self.repeated_row_extents(&batch, output_lens.as_ref()); timer.done(); self.pending_input = Some(PendingInput { batch, row_offset: 0, - output_lens: output_lens?, + output_lens, + row_extents, + element_offset: 0, }); } } @@ -504,6 +598,54 @@ impl ExplodeStream { let longest_length = find_longest_length(&list_arrays, &self.options)?; Ok(Some(longest_length.as_primitive::().clone())) } + + /// The per-row extents [`PendingInput::row_extents`] describes, or `None` when no build of + /// `batch` can overflow: every build then repeats each row at most the longest output length + /// times, which bounds what it adds to an offset buffer by that length times the buffer's + /// extent in the whole batch. Rows are only bounded or split for a single level of plain + /// `List` columns, each unnested once, which is every explode Comet plans. + fn repeated_row_extents( + &self, + batch: &RecordBatch, + output_lens: Option<&PrimitiveArray>, + ) -> Option> { + let output_lens = output_lens?; + let unnested: HashSet = self + .list_type_columns + .iter() + .map(|unnest| unnest.index_in_input_schema) + .collect(); + if unnested.len() != self.list_type_columns.len() + || unnested + .iter() + .any(|&index| !matches!(batch.column(index).data_type(), DataType::List(_))) + { + return None; + } + let repeated: Vec<&ArrayRef> = batch + .columns() + .iter() + .enumerate() + .filter(|(index, _)| !unnested.contains(index)) + .map(|(_, column)| column) + .collect(); + let longest = output_lens.values().iter().copied().max().unwrap_or(0) as usize; + let mut extents = vec![]; + for column in &repeated { + offset_extents(column.as_ref(), &mut extents); + } + if extents + .iter() + .all(|&e| e.saturating_mul(longest) <= self.offset_limit) + { + return None; + } + let mut row_extents = vec![0; batch.num_rows()]; + for column in repeated { + max_row_offset_extents(column.as_ref(), &mut row_extents); + } + Some(row_extents) + } } /// The per-row unnested length of a single `List` column: the row's list length, with NULL and @@ -1573,6 +1715,7 @@ mod tests { input_batches: MetricBuilder::new(&metrics).counter("input_batches", 0), input_rows: MetricBuilder::new(&metrics).counter("input_rows", 0), batch_size: 4, + offset_limit: MAX_OFFSET_EXTENT, pending_input: None, }; @@ -1877,4 +2020,235 @@ mod tests { ); assert_output_lens(with_nulls.slice(1, 3), [&[0, 0, 3], &[1, 0, 3], &[1, 1, 3]]); } + + /// One row whose `status_history` holds a `variant_id` of `bytes` bytes, alongside a list + /// of `variants` ints to explode, as in a proposal exploded by its target variants. + fn price_history_row( + values: &arrow::buffer::Buffer, + bytes: usize, + variants: i32, + ) -> RecordBatch { + let variant_id: ArrayRef = Arc::new(unsafe { + arrow::array::StringArray::new_unchecked( + OffsetBuffer::new(vec![0, bytes as i32].into()), + values.clone(), + None, + ) + }); + let price = arrow::datatypes::Fields::from(vec![ + Field::new("variant_id", DataType::Utf8, true), + Field::new("price", DataType::Int64, true), + ]); + let prices = StructArray::new( + price.clone(), + vec![variant_id, Arc::new(Int64Array::from(vec![1]))], + None, + ); + let price_list = Arc::new(Field::new_list_field(DataType::Struct(price), true)); + let merchant_variant_prices: ArrayRef = Arc::new(ListArray::new( + Arc::clone(&price_list), + OffsetBuffer::new(vec![0, 1].into()), + Arc::new(prices), + None, + )); + let status = arrow::datatypes::Fields::from(vec![ + Field::new("status", DataType::Utf8, true), + Field::new("updated_time", DataType::Int64, true), + Field::new("merchant_variant_prices", DataType::List(price_list), true), + ]); + let history = StructArray::new( + status.clone(), + vec![ + Arc::new(arrow::array::StringArray::from(vec!["pending"])), + Arc::new(Int64Array::from(vec![7])), + merchant_variant_prices, + ], + None, + ); + let history_list = Arc::new(Field::new_list_field(DataType::Struct(status), true)); + let status_history = ListArray::new( + Arc::clone(&history_list), + OffsetBuffer::new(vec![0, 1].into()), + Arc::new(history), + None, + ); + let targets = ListArray::from_iter_primitive::(vec![Some( + (0..variants).map(Some).collect::>(), + )]); + let schema = Arc::new(Schema::new(vec![ + Field::new("status_history", DataType::List(history_list), true), + Field::new("targets", targets.data_type().clone(), true), + ])); + RecordBatch::try_new(schema, vec![Arc::new(status_history), Arc::new(targets)]).unwrap() + } + + #[tokio::test] + async fn repeated_nested_strings_never_overflow_32_bit_offsets() { + let bytes = 600 << 20; + let values = arrow::buffer::Buffer::from(vec![b'v'; bytes]); + let input = price_history_row(&values, bytes, 4); + let output_schema = Arc::new(Schema::new(vec![ + input.schema().field(0).clone(), + Field::new("targets", DataType::Int32, true), + ])); + let source = + MemorySourceConfig::try_new_exec(&[vec![input.clone()]], input.schema(), None).unwrap(); + let explode = ExplodeExec::new( + source, + vec![ListUnnest { + index_in_input_schema: 1, + depth: 1, + }], + vec![], + output_schema, + UnnestOptions::new().with_null_handling(NullHandling::Drop), + ) + .unwrap(); + let mut stream = explode + .execute(0, Arc::new(TaskContext::default())) + .unwrap(); + let mut targets = vec![]; + while let Some(batch) = stream.next().await { + let batch = batch.unwrap(); + let history = batch.column(0).as_list::(); + assert_eq!(history.value_length(0), 1); + let row = history.value(0); + let mvp = row.as_struct().column(2).as_list::().value(0); + assert_eq!( + mvp.as_struct().column(0).as_string::().value_length(0) as usize, + bytes + ); + targets.extend(batch.column(1).as_primitive::().iter()); + } + assert_eq!(targets, seq(4)); + } + + async fn explode_with_offset_limit( + input: RecordBatch, + lists: &[usize], + output_schema: SchemaRef, + offset_limit: usize, + ) -> Vec { + let source = + MemorySourceConfig::try_new_exec(&[vec![input.clone()]], input.schema(), None).unwrap(); + let explode = ExplodeExec::new( + source, + lists + .iter() + .map(|&index_in_input_schema| ListUnnest { + index_in_input_schema, + depth: 1, + }) + .collect(), + vec![], + output_schema, + UnnestOptions::new().with_null_handling(NullHandling::Drop), + ) + .unwrap() + .with_offset_limit(offset_limit); + datafusion::physical_plan::common::collect( + explode + .execute(0, Arc::new(TaskContext::default())) + .unwrap(), + ) + .await + .unwrap() + } + + fn strings(batches: &[RecordBatch], column: usize) -> Vec { + batches + .iter() + .flat_map(|b| { + b.column(column) + .as_string::() + .iter() + .map(|v| v.unwrap().to_owned()) + .collect::>() + }) + .collect() + } + + #[tokio::test] + async fn bounds_repeated_offsets_across_rows_and_within_a_row() { + let names = arrow::array::StringArray::from(vec!["aaaa", "bbbbbbbbbb", "c"]); + let lists = list_batch(&[Some(3), Some(5), Some(2)]); + let schema = Arc::new(Schema::new(vec![ + Field::new("s", DataType::Utf8, true), + lists.schema().field(0).clone(), + ])); + let input = + RecordBatch::try_new(schema, vec![Arc::new(names), Arc::clone(lists.column(0))]) + .unwrap(); + let output_schema = Arc::new(Schema::new(vec![ + Field::new("s", DataType::Utf8, true), + Field::new("l", DataType::Int32, true), + ])); + let batches = explode_with_offset_limit(input, &[1], output_schema, 20).await; + assert_eq!(sizes(&batches), vec![3, 2, 2, 1, 2]); + for batch in &batches { + let s = batch.column(0).as_string::(); + assert!(s.value_offsets()[s.len()] - s.value_offsets()[0] <= 20); + } + let expected: Vec = [("aaaa", 3), ("bbbbbbbbbb", 5), ("c", 2)] + .iter() + .flat_map(|(s, n)| std::iter::repeat_n(s.to_string(), *n)) + .collect(); + assert_eq!(strings(&batches, 0), expected); + let values: Vec> = batches + .iter() + .flat_map(|b| { + b.column(1) + .as_primitive::() + .iter() + .collect::>() + }) + .collect(); + assert_eq!(values, seq(10)); + } + + #[tokio::test] + async fn splitting_a_row_keeps_parallel_lists_aligned() { + let names = arrow::array::StringArray::from(vec!["xxxxxxxx"]); + let positions = list_batch(&[Some(5)]); + let tags = ListArray::from_iter_primitive::(vec![Some(vec![ + Some(10), + Some(11), + Some(12), + ])]); + let schema = Arc::new(Schema::new(vec![ + Field::new("s", DataType::Utf8, true), + Field::new("p", positions.column(0).data_type().clone(), true), + Field::new("t", tags.data_type().clone(), true), + ])); + let input = RecordBatch::try_new( + schema, + vec![ + Arc::new(names), + Arc::clone(positions.column(0)), + Arc::new(tags), + ], + ) + .unwrap(); + let output_schema = Arc::new(Schema::new(vec![ + Field::new("s", DataType::Utf8, true), + Field::new("p", DataType::Int32, true), + Field::new("t", DataType::Int32, true), + ])); + let batches = explode_with_offset_limit(input, &[1, 2], output_schema, 16).await; + assert_eq!(sizes(&batches), vec![2, 2, 1]); + let column = |i: usize| -> Vec> { + batches + .iter() + .flat_map(|b| { + b.column(i) + .as_primitive::() + .iter() + .collect::>() + }) + .collect() + }; + assert_eq!(column(1), seq(5)); + assert_eq!(column(2), vec![Some(10), Some(11), Some(12), None, None]); + assert_eq!(strings(&batches, 0), vec!["xxxxxxxx".to_owned(); 5]); + } } diff --git a/native/core/src/execution/operators/partition_aggregate_window.rs b/native/core/src/execution/operators/partition_aggregate_window.rs index 13782df0383..dc6bcfa9ed2 100644 --- a/native/core/src/execution/operators/partition_aggregate_window.rs +++ b/native/core/src/execution/operators/partition_aggregate_window.rs @@ -48,6 +48,7 @@ use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, InputDistributionRequirements, InputOrderMode, PlanProperties, SendableRecordBatchStream, WindowExpr, }; +use datafusion_comet_common::offset_extents::OffsetBudget; use futures::{stream, StreamExt}; /// Window operator for expressions that need the whole partition, or every row after the @@ -464,6 +465,7 @@ impl ExecutionPlan for PartitionAggregateWindowExec { target_rows: context.session_config().batch_size().max(1), buffered: vec![], buffered_rows: 0, + buffered_offsets: OffsetBudget::default(), ready: VecDeque::new(), current_key: None, num_rows: 0, @@ -1139,6 +1141,7 @@ struct WindowState { target_rows: usize, buffered: Vec, buffered_rows: usize, + buffered_offsets: OffsetBudget, ready: VecDeque, } @@ -1605,6 +1608,13 @@ impl WindowState { }; } Some(batch) => { + if !self.buffered_offsets.try_add(batch.columns()) { + let buffered = self.flush()?; + self.buffered_offsets.try_add(batch.columns()); + self.buffered_rows = batch.num_rows(); + self.buffered.push(batch); + return Ok(buffered); + } self.buffered_rows += batch.num_rows(); self.buffered.push(batch); if self.buffered_rows >= self.target_rows { @@ -1622,6 +1632,7 @@ impl WindowState { } let batches = std::mem::take(&mut self.buffered); self.buffered_rows = 0; + self.buffered_offsets.clear(); Ok(Some(concat_batches(&self.schema, &batches)?)) } diff --git a/native/shuffle/src/read_coalescer.rs b/native/shuffle/src/read_coalescer.rs index 5006bd5f72f..d6c6d0d7532 100644 --- a/native/shuffle/src/read_coalescer.rs +++ b/native/shuffle/src/read_coalescer.rs @@ -18,6 +18,7 @@ use arrow::array::RecordBatch; use arrow::compute::concat_batches; use datafusion::error::Result; +use datafusion_comet_common::offset_extents::OffsetBudget; use std::sync::Arc; #[derive(Debug)] @@ -25,14 +26,20 @@ pub struct ShuffleReadCoalescer { target_rows: usize, pending: Vec, pending_rows: usize, + offsets: OffsetBudget, } impl ShuffleReadCoalescer { pub fn new(target_rows: usize) -> Self { + Self::with_offset_budget(target_rows, OffsetBudget::default()) + } + + fn with_offset_budget(target_rows: usize, offsets: OffsetBudget) -> Self { Self { target_rows: target_rows.max(1), pending: Vec::new(), pending_rows: 0, + offsets, } } @@ -46,8 +53,12 @@ impl ShuffleReadCoalescer { } let flushed = match self.pending.first() { Some(first) if !same_schema(first, &batch) => self.take()?, + Some(_) if !self.offsets.try_add(batch.columns()) => self.take()?, _ => None, }; + if self.pending.is_empty() { + self.offsets.try_add(batch.columns()); + } self.pending_rows += batch.num_rows(); self.pending.push(batch); if flushed.is_some() { @@ -65,6 +76,7 @@ impl ShuffleReadCoalescer { fn take(&mut self) -> Result> { self.pending_rows = 0; + self.offsets.clear(); match self.pending.len() { 0 => Ok(None), 1 => Ok(self.pending.pop()), @@ -89,6 +101,7 @@ mod tests { Array, ArrayRef, DictionaryArray, Int32Array, Int64Array, ListArray, StringArray, StructArray, }; + use arrow::buffer::OffsetBuffer; use arrow::datatypes::{DataType, Field, Fields, Int32Type, Schema}; use arrow::record_batch::RecordBatchOptions; @@ -235,6 +248,89 @@ mod tests { assert!(out.iter().all(|b| b.num_columns() == 0)); } + fn price_history(values: &arrow::buffer::Buffer, bytes: usize) -> RecordBatch { + let variant_id: ArrayRef = Arc::new(unsafe { + StringArray::new_unchecked( + OffsetBuffer::new(vec![0, bytes as i32].into()), + values.clone(), + None, + ) + }); + let price = Fields::from(vec![ + Field::new("variant_id", DataType::Utf8, true), + Field::new("price", DataType::Int64, true), + ]); + let prices = StructArray::new( + price.clone(), + vec![variant_id, Arc::new(Int64Array::from(vec![1]))], + None, + ); + let price_list = Arc::new(Field::new_list_field(DataType::Struct(price), true)); + let merchant_variant_prices: ArrayRef = Arc::new(ListArray::new( + Arc::clone(&price_list), + OffsetBuffer::new(vec![0, 1].into()), + Arc::new(prices), + None, + )); + let status = Fields::from(vec![ + Field::new("status", DataType::Utf8, true), + Field::new("updated_time", DataType::Int64, true), + Field::new("merchant_variant_prices", DataType::List(price_list), true), + ]); + let history = StructArray::new( + status.clone(), + vec![ + Arc::new(StringArray::from(vec!["pending"])), + Arc::new(Int64Array::from(vec![7])), + merchant_variant_prices, + ], + None, + ); + let history_list = Arc::new(Field::new_list_field(DataType::Struct(status), true)); + let status_history = ListArray::new( + Arc::clone(&history_list), + OffsetBuffer::new(vec![0, 1].into()), + Arc::new(history), + None, + ); + let schema = Arc::new(Schema::new(vec![Field::new( + "status_history", + DataType::List(history_list), + true, + )])); + RecordBatch::try_new(schema, vec![Arc::new(status_history)]).unwrap() + } + + #[test] + fn does_not_join_blocks_whose_nested_strings_overflow_32_bit_offsets() { + let bytes = (i32::MAX as usize) / 2 + (64 << 20); + let values = arrow::buffer::Buffer::from(vec![b'v'; bytes]); + let blocks: Vec<_> = (0..2).map(|_| price_history(&values, bytes)).collect(); + let out = drain(&mut ShuffleReadCoalescer::new(100), blocks.clone()); + assert_eq!(out, blocks); + } + + #[test] + fn joins_blocks_only_while_every_nested_offset_fits() { + let schema = schema(); + let blocks: Vec<_> = (0..10).map(|i| block(&schema, i * 4, 4)).collect(); + let expected = concat_batches(&schema, &blocks).unwrap(); + let mut coalescer = + ShuffleReadCoalescer::with_offset_budget(1000, OffsetBudget::with_limit(30)); + let out = drain(&mut coalescer, blocks); + assert!(out.len() > 1); + for batch in &out { + let p = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let tag = p.column(1).as_any().downcast_ref::().unwrap(); + assert!(tag.value_offsets()[tag.len()] - tag.value_offsets()[0] <= 30); + } + assert_eq!(concat_batches(&schema, &out).unwrap(), expected); + } + #[test] fn keeps_nulls_of_nested_columns() { let schema = schema(); From c9574d243f3ff80eeef86d6f66d777914f435900 Mon Sep 17 00:00:00 2001 From: msaf Date: Mon, 5 Oct 2026 09:17:07 +0100 Subject: [PATCH 03/14] fix: compare timestamps whose native producers label their time zone differently MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Native timestamp_seconds returned Timestamp(Microsecond, None), the Arrow type of TimestampNTZ, for Spark's TimestampType, so comparing it with a scan column or a timestamp literal failed with "Invalid comparison operation: Timestamp(µs) >= Timestamp(µs, "UTC")", and session-zone dependent expressions read it as wall-clock time. Port apache/datafusion-comet#6341, which labels its results as UTC. The Etc/UTC alias fix in date_trunc only covers that one label: a native date_trunc in any other session zone, or any other producer that labels its result with a zone other than UTC, still cannot be compared with UTC-labelled timestamps. Comparisons now relabel both sides as UTC when their zone labels differ. Spark only compares timestamps of one type and a TimestampType value is a UTC instant whatever its label, so this keeps every value. Co-Authored-By: Claude Opus 5.5 (1M context) --- native/core/src/execution/planner.rs | 79 ++++++- .../benches/seconds_to_timestamp.rs | 10 +- .../datetime_funcs/seconds_to_timestamp.rs | 199 ++++++++++++------ .../datetime/timestamp_seconds_timezone.sql | 73 +++++++ .../comet/CometTimestampComparisonSuite.scala | 137 ++++++++++++ 5 files changed, 430 insertions(+), 68 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/datetime/timestamp_seconds_timezone.sql create mode 100644 spark/src/test/scala/org/apache/comet/CometTimestampComparisonSuite.scala diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index d250b7e424f..3bc8e5a6915 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -1395,8 +1395,8 @@ impl PhysicalPlanner { /// ignore nullability. When a comparison's operands are nested types that differ only in /// nullability (e.g. a higher-order `transform` produces `List(non-null Struct)` while the /// other side is `List(nullable Struct)`), cast both to their nullability-union type so the - /// kernel accepts them. Non-comparison ops and non-nested or already-matching types are left - /// untouched. + /// kernel accepts them. Timestamps whose time zone labels differ are both relabelled as UTC. + /// Non-comparison ops and other or already-matching types are left untouched. pub fn reconcile_nested_comparison_types( left: Arc, right: Arc, @@ -1415,6 +1415,23 @@ impl PhysicalPlanner { (Ok(lt), Ok(rt)) => (lt, rt), _ => return (left, right), }; + if let (DataType::Timestamp(lu, ltz), DataType::Timestamp(ru, rtz)) = (<, &rt) { + // Spark only compares timestamps of one type, and a TimestampType value is a UTC + // instant whatever zone a native producer labelled it with (or none), so relabelling + // both sides as UTC keeps every value and lets Arrow compare them. + if lu != ru || ltz == rtz { + return (left, right); + } + let utc = DataType::Timestamp(*lu, Some("UTC".into())); + let relabel = |e: Arc, dt: &DataType| -> Arc { + if *dt == utc { + e + } else { + Arc::new(CastExpr::new(e, utc.clone(), None)) + } + }; + return (relabel(left, <), relabel(right, &rt)); + } // Only nested types route through `apply_cmp_for_nested`; primitives coerce fine. let nested = matches!( lt, @@ -7908,4 +7925,62 @@ mod tests { "unexpected error: {err}" ); } + + #[test] + fn comparisons_relabel_timestamps_whose_zones_differ() { + use arrow::array::{AsArray, RecordBatch, TimestampMicrosecondArray}; + use arrow::datatypes::TimeUnit::Microsecond; + use datafusion::logical_expr::Operator; + use datafusion::physical_expr::expressions::BinaryExpr; + + let micros = [Some(1_000_000), Some(2_000_000), None]; + let zones: [Option<&str>; 4] = [None, Some("UTC"), Some("Europe/Moscow"), Some("Etc/UTC")]; + let schema = Arc::new(Schema::new( + zones + .iter() + .enumerate() + .map(|(i, zone)| { + Field::new( + format!("c{i}"), + DataType::Timestamp(Microsecond, zone.map(Into::into)), + true, + ) + }) + .collect::>(), + )); + let columns = zones + .iter() + .enumerate() + .map(|(i, zone)| { + let shifted: Vec> = micros + .iter() + .map(|m| m.map(|m| m + i as i64 * 500_000 - 750_000)) + .collect(); + Arc::new(TimestampMicrosecondArray::from(shifted).with_timezone_opt(*zone)) + as ArrayRef + }) + .collect(); + let batch = RecordBatch::try_new(Arc::clone(&schema), columns).unwrap(); + for left in 0..zones.len() { + for right in 0..zones.len() { + let (l, r) = PhysicalPlanner::reconcile_nested_comparison_types( + Arc::new(Column::new(&format!("c{left}"), left)), + Arc::new(Column::new(&format!("c{right}"), right)), + &Operator::GtEq, + &schema, + ); + let result = BinaryExpr::new(l, Operator::GtEq, r) + .evaluate(&batch) + .unwrap() + .into_array(3) + .unwrap(); + let expected = [Some(left >= right), Some(left >= right), None]; + assert_eq!( + result.as_boolean().iter().collect::>(), + expected, + "c{left} >= c{right}" + ); + } + } + } } diff --git a/native/spark-expr/benches/seconds_to_timestamp.rs b/native/spark-expr/benches/seconds_to_timestamp.rs index 85a32ad65e6..28973a7402a 100644 --- a/native/spark-expr/benches/seconds_to_timestamp.rs +++ b/native/spark-expr/benches/seconds_to_timestamp.rs @@ -16,7 +16,7 @@ // under the License. use arrow::array::ArrayRef; -use arrow::datatypes::{DataType, Field, TimeUnit}; +use arrow::datatypes::Field; use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; use datafusion::common::config::ConfigOptions; use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl}; @@ -55,6 +55,8 @@ fn criterion_benchmark(c: &mut Criterion) { for rows in ROW_COUNTS { for (null_ratio, tag) in NULL_RATIOS { for (ty, arr) in inputs(rows, null_ratio) { + let return_type = udf.return_type(&[arr.data_type().clone()]).unwrap(); + let return_field = Arc::new(Field::new("result", return_type, true)); let args = vec![ColumnarValue::Array(arr)]; group.bench_with_input( BenchmarkId::from_parameter(format!("{ty}/{rows}/{tag}")), @@ -66,11 +68,7 @@ fn criterion_benchmark(c: &mut Criterion) { args: args.clone(), arg_fields: vec![], number_rows: rows, - return_field: Arc::new(Field::new( - "result", - DataType::Timestamp(TimeUnit::Microsecond, None), - true, - )), + return_field: Arc::clone(&return_field), config_options: Arc::new(ConfigOptions::default()), }) .unwrap(), diff --git a/native/spark-expr/src/datetime_funcs/seconds_to_timestamp.rs b/native/spark-expr/src/datetime_funcs/seconds_to_timestamp.rs index ab9388692b7..b4c2610591e 100644 --- a/native/spark-expr/src/datetime_funcs/seconds_to_timestamp.rs +++ b/native/spark-expr/src/datetime_funcs/seconds_to_timestamp.rs @@ -28,6 +28,11 @@ use std::sync::Arc; const MICROS_PER_SECOND: i64 = 1_000_000; +/// The result is Spark's `TimestampType`, which is `Timestamp(Microsecond, "UTC")` everywhere in a +/// native plan. Without the label, downstream kernels would read the result as `TimestampNTZType` +/// wall-clock time, and Arrow would reject comparing it with other timestamps. +const TIMEZONE: &str = "UTC"; + /// Spark-compatible seconds_to_timestamp (timestamp_seconds) function. /// Converts seconds since Unix epoch to a timestamp. #[derive(Debug, PartialEq, Eq, Hash)] @@ -69,71 +74,19 @@ impl ScalarUDFImpl for SparkSecondsToTimestamp { } fn return_type(&self, _: &[DataType]) -> Result { - Ok(DataType::Timestamp(TimeUnit::Microsecond, None)) + Ok(DataType::Timestamp( + TimeUnit::Microsecond, + Some(TIMEZONE.into()), + )) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { let [seconds] = take_function_args(self.name(), args.args)?; match seconds { - ColumnarValue::Array(arr) => { - // Handle Int32 input — no overflow possible since i32 * 1_000_000 fits in i64 - if let Some(int_array) = arr.as_any().downcast_ref::() { - let result: TimestampMicrosecondArray = - try_unary(int_array, |s| Ok((s as i64) * MICROS_PER_SECOND))?; - return Ok(ColumnarValue::Array(Arc::new(result))); - } - - // Handle Int64 input — error on overflow to match Spark's Math.multiplyExact - if let Some(int_array) = arr.as_any().downcast_ref::() { - let result: TimestampMicrosecondArray = try_unary(int_array, |s| { - s.checked_mul(MICROS_PER_SECOND).ok_or_else(|| { - arrow::error::ArrowError::ComputeError("long overflow".to_string()) - }) - })?; - return Ok(ColumnarValue::Array(Arc::new(result))); - } - - // Handle Float32 input — cast to f64 and use Float64 path - if let Some(float_array) = arr.as_any().downcast_ref::() { - let result: arrow::array::TimestampMicrosecondArray = float_array - .iter() - .map(|opt| { - opt.and_then(|s| { - let s = s as f64; - if s.is_nan() || s.is_infinite() { - None - } else { - Some((s * (MICROS_PER_SECOND as f64)) as i64) - } - }) - }) - .collect(); - return Ok(ColumnarValue::Array(Arc::new(result))); - } - - // Handle Float64 input — NaN and Infinity return null per Spark behavior - if let Some(float_array) = arr.as_any().downcast_ref::() { - let result: arrow::array::TimestampMicrosecondArray = float_array - .iter() - .map(|opt| { - opt.and_then(|s| { - if s.is_nan() || s.is_infinite() { - None - } else { - Some((s * (MICROS_PER_SECOND as f64)) as i64) - } - }) - }) - .collect(); - return Ok(ColumnarValue::Array(Arc::new(result))); - } - - Err(DataFusionError::Execution(format!( - "seconds_to_timestamp expects Int32, Int64, Float32 or Float64 input, got {:?}", - arr.data_type() - ))) - } + ColumnarValue::Array(arr) => Ok(ColumnarValue::Array(Arc::new( + seconds_to_micros(arr.as_ref())?.with_timezone(TIMEZONE), + ))), ColumnarValue::Scalar(scalar) => { let ts_micros = match &scalar { ScalarValue::Int32(Some(s)) => Some((*s as i64) * MICROS_PER_SECOND), @@ -175,7 +128,8 @@ impl ScalarUDFImpl for SparkSecondsToTimestamp { } }; Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond( - ts_micros, None, + ts_micros, + Some(TIMEZONE.into()), ))) } } @@ -185,3 +139,128 @@ impl ScalarUDFImpl for SparkSecondsToTimestamp { &self.aliases } } + +/// Converts an array of seconds to microseconds. The caller labels the result with [`TIMEZONE`]. +fn seconds_to_micros(arr: &dyn Array) -> Result { + // Handle Int32 input — no overflow possible since i32 * 1_000_000 fits in i64 + if let Some(int_array) = arr.as_any().downcast_ref::() { + return Ok(try_unary( + int_array, + |s| Ok((s as i64) * MICROS_PER_SECOND), + )?); + } + + // Handle Int64 input — error on overflow to match Spark's Math.multiplyExact + if let Some(int_array) = arr.as_any().downcast_ref::() { + return Ok(try_unary(int_array, |s| { + s.checked_mul(MICROS_PER_SECOND) + .ok_or_else(|| arrow::error::ArrowError::ComputeError("long overflow".to_string())) + })?); + } + + // Handle Float32 input — cast to f64 and use Float64 path + if let Some(float_array) = arr.as_any().downcast_ref::() { + return Ok(float_array + .iter() + .map(|opt| { + opt.and_then(|s| { + let s = s as f64; + if s.is_nan() || s.is_infinite() { + None + } else { + Some((s * (MICROS_PER_SECOND as f64)) as i64) + } + }) + }) + .collect()); + } + + // Handle Float64 input — NaN and Infinity return null per Spark behavior + if let Some(float_array) = arr.as_any().downcast_ref::() { + return Ok(float_array + .iter() + .map(|opt| { + opt.and_then(|s| { + if s.is_nan() || s.is_infinite() { + None + } else { + Some((s * (MICROS_PER_SECOND as f64)) as i64) + } + }) + }) + .collect()); + } + + Err(DataFusionError::Execution(format!( + "seconds_to_timestamp expects Int32, Int64, Float32 or Float64 input, got {:?}", + arr.data_type() + ))) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{ArrayRef, AsArray}; + use arrow::datatypes::{Field, TimestampMicrosecondType}; + use datafusion::config::ConfigOptions; + + fn utc_timestamp() -> DataType { + DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())) + } + + fn invoke(arg: ColumnarValue, number_rows: usize) -> ColumnarValue { + let udf = SparkSecondsToTimestamp::new(); + udf.invoke_with_args(ScalarFunctionArgs { + args: vec![arg], + arg_fields: vec![], + number_rows, + return_field: Arc::new(Field::new("result", utc_timestamp(), true)), + config_options: Arc::new(ConfigOptions::default()), + }) + .unwrap() + } + + /// Every input type yields the Arrow type `return_type` declares. + #[test] + fn arrays_are_labelled_utc() { + let udf = SparkSecondsToTimestamp::new(); + let inputs: Vec = vec![ + Arc::new(Int32Array::from(vec![Some(1800), None])), + Arc::new(Int64Array::from(vec![Some(1800), None])), + Arc::new(Float32Array::from(vec![Some(1800.0), None])), + Arc::new(Float64Array::from(vec![Some(1800.0), None])), + ]; + for input in inputs { + let input_type = input.data_type().clone(); + let declared = udf.return_type(std::slice::from_ref(&input_type)).unwrap(); + assert_eq!(declared, utc_timestamp(), "{input_type}"); + let ColumnarValue::Array(result) = invoke(ColumnarValue::Array(input), 2) else { + panic!("expected an array for {input_type}"); + }; + assert_eq!(result.data_type(), &declared, "{input_type}"); + let result = result.as_primitive::(); + assert_eq!(result.value(0), 1_800_000_000, "{input_type}"); + assert!(result.is_null(1), "{input_type}"); + } + } + + #[test] + fn scalars_are_labelled_utc() { + for (input, expected) in [ + (ScalarValue::Int32(Some(1800)), Some(1_800_000_000)), + (ScalarValue::Int64(Some(1800)), Some(1_800_000_000)), + (ScalarValue::Float64(Some(1800.5)), Some(1_800_500_000)), + (ScalarValue::Int64(None), None), + ] { + let ColumnarValue::Scalar(result) = invoke(ColumnarValue::Scalar(input.clone()), 1) + else { + panic!("expected a scalar for {input:?}"); + }; + assert_eq!( + result, + ScalarValue::TimestampMicrosecond(expected, Some("UTC".into())), + "{input:?}" + ); + } + } +} diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/timestamp_seconds_timezone.sql b/spark/src/test/resources/sql-tests/expressions/datetime/timestamp_seconds_timezone.sql new file mode 100644 index 00000000000..b16096ca790 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/datetime/timestamp_seconds_timezone.sql @@ -0,0 +1,73 @@ +-- 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. + +-- timestamp_seconds returns TimestampType. These queries use the result instead of only +-- projecting it. Expressions that depend on the session timezone must see an instant rather +-- than TIMESTAMP_NTZ wall-clock time, and comparisons, CASE and coalesce need the result to have +-- the same Arrow type as other timestamp columns. +-- https://github.com/apache/datafusion-comet/issues/6328 + +-- ConfigMatrix: spark.sql.session.timeZone=UTC,America/Los_Angeles,Asia/Kolkata + +statement +CREATE TABLE test_ts_seconds_tz(s bigint, i int, d double, ts timestamp) USING parquet + +-- ts equals timestamp_seconds(s) except in the row for 2024-03-10, where it is an hour later +statement +INSERT INTO test_ts_seconds_tz VALUES + (0, 0, 0.5, TIMESTAMP'1970-01-01 00:00:00Z'), + (1640995200, 1640995200, 1640995200.25, TIMESTAMP'2022-01-01 00:00:00Z'), + (-86400, -86400, -86400.5, TIMESTAMP'1969-12-31 00:00:00Z'), + (1710063000, 1710063000, 1710063000.0, TIMESTAMP'2024-03-10 10:30:00Z'), + (1730622600, 1730622600, 1730622600.0, TIMESTAMP'2024-11-03 08:30:00Z'), + (1719790200, 1719790200, 1719790200.0, TIMESTAMP'2024-06-30 23:30:00Z'), + (NULL, NULL, NULL, NULL) + +-- fields and casts that depend on the session timezone +query expect_native(timestamp_seconds) +SELECT s, hour(timestamp_seconds(s)), minute(timestamp_seconds(s)) FROM test_ts_seconds_tz + +query expect_native(timestamp_seconds) +SELECT s, CAST(timestamp_seconds(s) AS STRING), CAST(timestamp_seconds(s) AS DATE) FROM test_ts_seconds_tz + +query expect_native(timestamp_seconds) +SELECT s, CAST(timestamp_seconds(s) AS TIMESTAMP_NTZ) FROM test_ts_seconds_tz + +-- int and double inputs take their own native paths +query expect_native(timestamp_seconds) +SELECT i, hour(timestamp_seconds(i)), CAST(timestamp_seconds(i) AS STRING) FROM test_ts_seconds_tz + +query expect_native(timestamp_seconds) +SELECT d, CAST(timestamp_seconds(d) AS STRING) FROM test_ts_seconds_tz + +-- comparisons with a timestamp column and a timestamp literal +query expect_native(timestamp_seconds) +SELECT s, timestamp_seconds(s) = ts, timestamp_seconds(s) < ts FROM test_ts_seconds_tz + +query expect_native(timestamp_seconds) +SELECT s FROM test_ts_seconds_tz WHERE timestamp_seconds(s) >= TIMESTAMP'2022-01-01 00:00:00Z' + +-- conditionals that mix the result with a timestamp column +query expect_native(timestamp_seconds) +SELECT s, CASE WHEN s > 0 THEN timestamp_seconds(s) ELSE ts END FROM test_ts_seconds_tz + +query expect_native(timestamp_seconds) +SELECT s, coalesce(timestamp_seconds(s), ts) FROM test_ts_seconds_tz + +-- literal argument, which takes the scalar path because the suite disables constant folding +query expect_native(timestamp_seconds) +SELECT s, CAST(timestamp_seconds(1800) AS STRING), timestamp_seconds(0) = ts FROM test_ts_seconds_tz diff --git a/spark/src/test/scala/org/apache/comet/CometTimestampComparisonSuite.scala b/spark/src/test/scala/org/apache/comet/CometTimestampComparisonSuite.scala new file mode 100644 index 00000000000..067c3acf154 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/CometTimestampComparisonSuite.scala @@ -0,0 +1,137 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.comet + +import org.apache.spark.sql.CometTestBase +import org.apache.spark.sql.catalyst.expressions.{FromUTCTimestamp, TruncTimestamp} + +/** Comparisons of timestamps whose native producers label their time zone differently. */ +class CometTimestampComparisonSuite extends CometTestBase { + + private val zones = Seq("UTC", "Etc/UTC", "Europe/Moscow", "America/Los_Angeles") + + private def withEvents(f: => Unit): Unit = { + withTempPath { dir => + withSQLConf("spark.sql.session.timeZone" -> "UTC") { + spark + .sql("""SELECT * FROM VALUES + | (1710063000L, 1710063000123L, '2024-03-10 10:30:00.123456', + | TIMESTAMP'2024-03-10 10:30:00Z'), + | (1640995200L, 1640995200000L, '2022-01-01 00:00:00.000000', + | TIMESTAMP'2022-01-01 03:00:00Z'), + | (1730622600L, 1730622600999L, '2024-11-03 08:30:00.999000', + | TIMESTAMP'2024-11-03 01:00:00Z'), + | (100L, 100L, 'not a timestamp', TIMESTAMP'1970-01-01 00:00:00Z'), + | (NULL, NULL, NULL, NULL) + |AS t(s, ms, str, ts)""".stripMargin) + .write + .parquet(dir.getCanonicalPath) + } + spark.read.parquet(dir.getCanonicalPath).createOrReplaceTempView("events") + withTempView("events")(f) + } + } + + private def forZones(f: => Unit): Unit = inZones(zones)(f) + + private def inZones(zones: Seq[String])(f: => Unit): Unit = + zones.foreach { zone => + withSQLConf("spark.sql.session.timeZone" -> zone) { + withClue(s"session time zone $zone: ")(f) + } + } + + test("timestamp_seconds compares with timestamp literals and columns") { + withEvents { + forZones { + checkSparkAnswerAndOperator( + "SELECT s FROM events WHERE " + + "timestamp_seconds(s) BETWEEN TIMESTAMP '2020-01-01' AND TIMESTAMP '2030-01-01'") + checkSparkAnswerAndOperator( + "SELECT s, timestamp_seconds(s) >= ts, timestamp_seconds(s) = ts, " + + "timestamp_seconds(s) < ts FROM events") + checkSparkAnswerAndOperator( + "SELECT s, hour(timestamp_seconds(s)), CAST(timestamp_seconds(s) AS STRING) " + + "FROM events") + } + } + } + + test("timestamp_seconds, timestamp_millis and to_timestamp mix in conditionals") { + withEvents { + forZones { + checkSparkAnswer("""SELECT s, COALESCE( + | CASE WHEN timestamp_seconds(s) BETWEEN TIMESTAMP '2020-01-01' + | AND TIMESTAMP '2030-01-01' THEN timestamp_seconds(s) END, + | CASE WHEN timestamp_millis(ms) BETWEEN TIMESTAMP '2020-01-01' + | AND TIMESTAMP '2030-01-01' THEN timestamp_millis(ms) END, + | to_timestamp(str, 'yyyy-MM-dd HH:mm:ss.SSSSSS')) AS p_time + |FROM events""".stripMargin) + checkSparkAnswerAndOperator( + "SELECT s, CASE WHEN s > 1000 THEN timestamp_seconds(s) ELSE ts END FROM events") + } + } + } + + test("from_utc_timestamp results compare with to_timestamp results") { + withEvents { + withSQLConf(CometConf.getExprAllowIncompatConfigKey(classOf[FromUTCTimestamp]) -> "true") { + forZones { + checkSparkAnswer( + "SELECT a.s, b.s FROM events a JOIN events b ON a.s = b.s AND " + + "from_utc_timestamp(timestamp_millis(a.ms), 'Europe/Moscow') >= " + + "to_timestamp(b.str, 'yyyy-MM-dd HH:mm:ss.SSSSSS')") + checkSparkAnswer( + "SELECT s, from_utc_timestamp(timestamp_seconds(s), 'Europe/Moscow') >= ts, " + + "from_utc_timestamp(ts, 'Europe/Moscow') < timestamp_seconds(s) FROM events") + } + } + } + } + + test("date_trunc labelled with the session time zone compares with scan timestamps") { + withEvents { + withSQLConf( + CometConf.getExprEnabledConfigKey("TruncTimestamp") -> "true", + CometConf.getExprAllowIncompatConfigKey(classOf[TruncTimestamp]) -> "true") { + inZones(Seq("UTC", "Etc/UTC", "Europe/Moscow")) { + checkSparkAnswerAndOperator( + "SELECT ts, date_trunc('HOUR', ts) < ts, date_trunc('HOUR', ts) >= ts FROM events") + checkSparkAnswerAndOperator("SELECT ts FROM events WHERE date_trunc('HOUR', ts) < ts") + checkSparkAnswerAndOperator( + "SELECT s FROM events WHERE date_trunc('HOUR', timestamp_seconds(s)) <= ts") + } + } + } + } + + test("timestamp_ntz comparisons keep wall-clock semantics") { + withEvents { + forZones { + checkSparkAnswer( + "SELECT s, CAST(timestamp_seconds(s) AS TIMESTAMP_NTZ) >= " + + "TIMESTAMP_NTZ'2024-03-10 10:30:00', " + + "CAST(ts AS TIMESTAMP_NTZ) = CAST(timestamp_seconds(s) AS TIMESTAMP_NTZ), " + + "timestamp_seconds(s) > CAST(TIMESTAMP_NTZ'2024-03-10 10:30:00' AS TIMESTAMP) " + + "FROM events") + } + } + } +} From 29b0bb02b06d0aac31ae31dd83516010d178e8b6 Mon Sep 17 00:00:00 2001 From: msaf Date: Mon, 5 Oct 2026 11:05:19 +0100 Subject: [PATCH 04/14] fix: let a native sort's spill merge record what it needs past the sorter's memory A spill merges the batches the sorter buffered inside a SpillWorkspace that holds the sorter's reservation, twice the batches' size. The merge's builder keeps one sorted batch per run and its row cursors encode each batch's sort key, so when the key is most of the row the merge needs a little more than twice the data. Since d7d47f3bb counts the encoded rows a merge keeps, that remainder is asked of the pool, and with the task's Spark share used up the sort failed with "Additional allocation failed for ExternalSorter[0] ... Failed to acquire N bytes plus 0 bytes overcommitted", although the spill was about to release everything it holds. merchants_lifecycle (window over merchant_id, created_date, product_id) failed this way on all 10 attempts. Pristine 55.1.0 passes the same inputs because it does not count the reused rows: 0 of 284 failing vs 130 of 284 at the parent commit in a release grid. Everything the spill merge's children reserve already exists (sorted chunks, encoded rows, computed keys), so the workspace of a spill now grows its first parent with the infallible grow. In Comet the shortfall is carried as overcommit and repaid when the spill frees the workspace; the merge of spill files keeps the fallible workspace it sizes its passes with. Tests: a vendored sort of key-dominated rows in a pool it fills, and a Comet sort with a fixed Spark share and production's merge headroom ratio, which failed with the production message and now peaks within 2% of the share. Co-Authored-By: Claude Opus 5.5 (1M context) --- native/core/src/execution/jni_api.rs | 149 ++++++++++++++++++ .../src/sorts/sort.rs | 2 +- .../src/sorts/sort/comet_memory_tests.rs | 95 ++++++++++- .../src/sorts/spill_workspace.rs | 38 ++++- 4 files changed, 281 insertions(+), 3 deletions(-) diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index dbd6ffb69d3..b2b7ec52efd 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -3660,4 +3660,153 @@ mod native_sort_spill_tests { async fn sort_spills_before_output_above_the_threshold() { assert_spills_before_output(RowShape::KibKeyAndId, SPILL_BEFORE_OUTPUT_ROWS / 2).await; } + + /// Runs `plan` in one task whose Spark share stays at `share`, reads its output, checks + /// that all memory is handed back, and returns the rows and the most the pool reserved. + async fn run_with_fixed_share( + share: usize, + executor_cores: usize, + plan: Arc, + ) -> DataFusionResult<(usize, usize)> { + let off_heap_size = share * executor_cores; + let (pool, spark) = fair_unified_pool_with_fake_spark(off_heap_size, share); + let peak = Arc::new(PeakPool { + inner: pool, + peak: Default::default(), + }); + let pool: Arc = Arc::clone(&peak) as _; + let spill_dir = tempfile::tempdir().unwrap(); + // As in production, where the merge headroom is 10 MiB of a 4.25 GiB share. + let spark_config = HashMap::from([ + (SPARK_EXECUTOR_CORES.to_string(), executor_cores.to_string()), + ( + "spark.comet.datafusion.execution.sort_spill_reservation_bytes".to_string(), + (share / 400).to_string(), + ), + ]); + let session = prepare_datafusion_session_context( + 8192, + Arc::clone(&pool), + vec![spill_dir.path().to_string_lossy().into_owned()], + u64::MAX, + 1, + &spark_config, + &Operator::default(), + Some(off_heap_size), + ) + .unwrap(); + let mut stream = plan.execute(0, session.task_ctx())?; + let mut rows = 0; + while let Some(batch) = stream.next().await { + rows += batch?.num_rows(); + } + drop(stream); + assert_eq!(pool.reserved(), 0, "memory still reserved after the plan"); + assert_eq!(spark.held(), 0, "memory not handed back to Spark"); + Ok((rows, peak.peak())) + } + + fn streaming(source: Arc) -> Arc { + Arc::new( + StreamingTableExec::try_new( + Arc::clone(source.schema()), + vec![source], + None, + Vec::::new(), + false, + None, + ) + .unwrap(), + ) + } + + /// Rows whose sort key, two 24-character ids and a day, is nearly all of the row. + #[derive(Debug)] + struct KeyedRows { + schema: SchemaRef, + rows_per_batch: usize, + num_batches: usize, + } + + impl KeyedRows { + fn new(rows_per_batch: usize, num_batches: usize) -> Self { + Self { + schema: Arc::new(Schema::new(vec![ + Field::new("merchant_id", DataType::Utf8, false), + Field::new("created_date", DataType::Int32, false), + Field::new("product_id", DataType::Utf8, false), + Field::new("row", DataType::Int64, false), + ])), + rows_per_batch, + num_batches, + } + } + + fn ordering(&self) -> LexOrdering { + LexOrdering::new( + ["merchant_id", "created_date", "product_id"] + .map(|name| PhysicalSortExpr::new_default(col(name, &self.schema).unwrap())), + ) + .unwrap() + } + + fn batch(&self, index: usize) -> RecordBatch { + let start = (index * self.rows_per_batch) as u64; + let rows: Vec = (start..start + self.rows_per_batch as u64).collect(); + let id = |r: u64| format!("{:016x}{:08x}", mix(r), mix(r ^ 7) as u32); + RecordBatch::try_new( + Arc::clone(&self.schema), + vec![ + Arc::new(StringArray::from_iter_values( + rows.iter().map(|&r| id(r / 50)), + )), + Arc::new(Int32Array::from_iter_values( + rows.iter().map(|&r| (mix(r ^ 3) % 3650) as i32), + )), + Arc::new(StringArray::from_iter_values( + rows.iter().map(|&r| id(r ^ 11)), + )), + Arc::new(Int64Array::from_iter_values(rows.iter().map(|&r| r as i64))), + ], + ) + .unwrap() + } + } + + impl PartitionStream for KeyedRows { + fn schema(&self) -> &SchemaRef { + &self.schema + } + + fn execute(&self, _ctx: Arc) -> SendableRecordBatchStream { + let batches: Vec<_> = (0..self.num_batches).map(|i| Ok(self.batch(i))).collect(); + Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&self.schema), + futures::stream::iter(batches), + )) + } + } + + /// A spill merges the batches the sort buffered inside the memory it holds for them. + /// When the sort key is most of the row, the merge's batches and the rows it encodes + /// their keys into need a little more than that, and the merge asked Spark for it. With + /// the task's share used up the sort failed with "Additional allocation failed for + /// ExternalSorter[0] ... Failed to acquire N bytes plus 0 bytes overcommitted, only got N + /// bytes", as merchants_lifecycle did on every attempt. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn native_sort_by_a_key_that_is_most_of_the_row_spills_in_a_fixed_share() { + for share in [24, 28, 32, 36, 40].map(|mb| mb * MB) { + let source = Arc::new(KeyedRows::new(1024, 1200)); + let ordering = source.ordering(); + let sort: Arc = Arc::new(SortExec::new(ordering, streaming(source))); + match run_with_fixed_share(share, 8, Arc::clone(&sort)).await { + Ok((rows, peak)) => { + assert_eq!(rows, 1024 * 1200); + assert!(spill_count(&sort) > 0, "sort did not spill"); + assert!(peak <= share + share / 50, "{peak} reserved of {share}"); + } + Err(e) => panic!("native sort failed instead of spilling at {share}: {e}"), + } + } + } } diff --git a/native/vendor/datafusion-physical-plan/src/sorts/sort.rs b/native/vendor/datafusion-physical-plan/src/sorts/sort.rs index a3110f72c94..82d55594e5c 100644 --- a/native/vendor/datafusion-physical-plan/src/sorts/sort.rs +++ b/native/vendor/datafusion-physical-plan/src/sorts/sort.rs @@ -650,7 +650,7 @@ impl ExternalSorter { // merge instead of returning it to the pool and requesting it again, which fails // once the pool cannot grant what the sorter held. See `SpillWorkspace`. let buffered = self.reservation.size(); - let workspace = SpillWorkspace::new(vec![ + let workspace = SpillWorkspace::for_spill(vec![ self.reservation.take(), self.merge_reservation.take(), ]); diff --git a/native/vendor/datafusion-physical-plan/src/sorts/sort/comet_memory_tests.rs b/native/vendor/datafusion-physical-plan/src/sorts/sort/comet_memory_tests.rs index 5643a037a9b..8a4897a4ebe 100644 --- a/native/vendor/datafusion-physical-plan/src/sorts/sort/comet_memory_tests.rs +++ b/native/vendor/datafusion-physical-plan/src/sorts/sort/comet_memory_tests.rs @@ -26,9 +26,12 @@ use arrow::array::{ }; use arrow::datatypes::{DataType, Field, Int32Type, Int64Type, Schema}; use datafusion_execution::config::SessionConfig; -use datafusion_execution::memory_pool::{GreedyMemoryPool, MemoryLimit}; +use datafusion_execution::memory_pool::{ + GreedyMemoryPool, MemoryLimit, TrackConsumersPool, +}; use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_physical_expr::expressions::Column; +use std::num::NonZeroUsize; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; fn new_sorter( @@ -656,3 +659,93 @@ async fn small_view_batches_are_not_coalesced() -> Result<()> { assert!(sorter.small_batches.is_empty()); Ok(()) } + +fn key_dominated_batches( + batches: usize, + rows: usize, +) -> Result<(SchemaRef, LexOrdering, Vec)> { + let schema = Arc::new(Schema::new(vec![ + Field::new("m", DataType::Utf8, false), + Field::new("d", DataType::Int32, false), + Field::new("p", DataType::Utf8, false), + Field::new("v", DataType::Int64, false), + ])); + let mut state = 17u64; + let mut hex = || { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + format!("{:024x}", state >> 16) + }; + let input = (0..batches) + .map(|b| { + let m: Vec = (0..rows).map(|_| hex()).collect(); + let p: Vec = (0..rows).map(|_| hex()).collect(); + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(StringArray::from(m)), + Arc::new(Int32Array::from_iter_values( + (0..rows).map(|r| (r % 97) as i32), + )), + Arc::new(StringArray::from(p)), + Arc::new(Int64Array::from_iter_values( + (0..rows).map(|r| (b * rows + r) as i64), + )), + ], + ) + }) + .collect::>()?; + let ordering = [ + PhysicalSortExpr::new_default(crate::expressions::col("m", &schema)?), + PhysicalSortExpr::new_default(crate::expressions::col("d", &schema)?), + PhysicalSortExpr::new_default(crate::expressions::col("p", &schema)?), + ] + .into(); + Ok((schema, ordering, input)) +} + +/// A spill merges the batches the sorter buffered inside the memory it holds for them, +/// twice their size. When the sort key is most of each row, the merge's batches and the +/// rows it encodes their keys into take about that much on their own, and the merge +/// asked the pool for the rest. In a pool the sorter had filled that failed with +/// "Additional allocation failed for ExternalSorter[0]", although the spill was about to +/// release everything the sorter holds. +#[tokio::test] +async fn spill_merge_of_key_dominated_rows_does_not_fail() -> Result<()> { + let (schema, ordering, input) = key_dominated_batches(300, 1024)?; + let data: usize = input.iter().map(get_record_batch_memory_size).sum(); + let expected = sort_batch(&concat_batches(&schema, &input)?, &ordering, None)?; + for percent in [35, 38, 42, 46, 48] { + let pool: Arc = Arc::new(TrackConsumersPool::new( + GreedyMemoryPool::new(data * 2 * percent / 100), + NonZeroUsize::new(5).unwrap(), + )); + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool)) + .build_arc()?; + let mut sorter = ExternalSorter::new( + 0, + Arc::clone(&schema), + ordering.clone(), + 8192, + 64 * 1024, + 1024 * 1024, + SpillCompression::Uncompressed, + &ExecutionPlanMetricsSet::new(), + runtime, + )?; + for batch in &input { + sorter.insert_batch(batch.clone()).await?; + } + assert!(sorter.spill_count() > 0); + let output: Vec = sorter.sort().await?.try_collect().await?; + drop(sorter); + assert_eq!(pool.reserved(), 0); + let actual = concat_batches(&schema, &output)?; + for column in 0..3 { + assert_eq!(actual.column(column), expected.column(column)); + } + } + Ok(()) +} diff --git a/native/vendor/datafusion-physical-plan/src/sorts/spill_workspace.rs b/native/vendor/datafusion-physical-plan/src/sorts/spill_workspace.rs index 2d95d761e6a..271c24f2186 100644 --- a/native/vendor/datafusion-physical-plan/src/sorts/spill_workspace.rs +++ b/native/vendor/datafusion-physical-plan/src/sorts/spill_workspace.rs @@ -41,9 +41,14 @@ use parking_lot::Mutex; /// workspace, and only usage beyond it grows the first parent reservation, under the /// execution pool's limits. [`Self::close`] ends the retention, and [`Self::keep_at_most`] /// limits it. +/// +/// A workspace built with [`Self::for_spill`] grows its first parent with the infallible +/// `grow` instead, for a merge that writes a spill: everything its children reserve +/// already exists, and the spill it is writing releases all of it. #[derive(Debug)] pub(super) struct SpillWorkspace { state: Mutex, + fallible: bool, } #[derive(Debug)] @@ -101,6 +106,15 @@ impl Drop for WorkspaceLoan { impl SpillWorkspace { /// Takes over `parents`. The first one is grown if the children need more. pub(super) fn new(parents: Vec) -> Arc { + Self::with_growth(parents, true) + } + + /// [`Self::new`] for a merge that writes a spill, which never refuses its children. + pub(super) fn for_spill(parents: Vec) -> Arc { + Self::with_growth(parents, false) + } + + fn with_growth(parents: Vec, fallible: bool) -> Arc { assert!(!parents.is_empty()); Arc::new(Self { state: Mutex::new(State { @@ -108,6 +122,7 @@ impl SpillWorkspace { used: 0, keep: usize::MAX, }), + fallible, }) } @@ -195,7 +210,7 @@ impl MemoryPool for SpillWorkspace { let Some(used) = state.used.checked_add(additional) else { return resources_err!("Sort spill workspace overflow"); }; - state.cover(used, true) + state.cover(used, self.fallible) } fn reserved(&self) -> usize { @@ -270,6 +285,27 @@ mod tests { assert_eq!(parent.reserved(), 0); } + #[test] + fn a_spill_workspace_grows_its_first_parent_past_the_pool() { + let parent: Arc = Arc::new(GreedyMemoryPool::new(100)); + let sorter = MemoryConsumer::new("sorter").register(&parent); + sorter.try_grow(60).unwrap(); + let contender = MemoryConsumer::new("contender").register(&parent); + contender.try_grow(40).unwrap(); + let workspace = SpillWorkspace::for_spill(vec![sorter]); + let pool = Arc::clone(&workspace) as Arc; + let child = MemoryConsumer::new("child").register(&pool); + + child.try_grow(60).unwrap(); + child.try_grow(15).unwrap(); + assert_eq!(parent.reserved(), 115); + assert!(contender.try_grow(1).is_err()); + + workspace.close(); + drop(child); + assert_eq!(parent.reserved(), 40); + } + #[test] fn can_grow_counts_unused_workspace_and_leaves_reservations_unchanged() { let (parent, workspace, child) = setup(100, 40); From c6178ee56dbd4540ffa783b1758c63b55f127b00 Mon Sep 17 00:00:00 2001 From: msaf Date: Mon, 5 Oct 2026 11:24:17 +0100 Subject: [PATCH 05/14] fix: record what a fully ordered final aggregate holds instead of failing A final aggregate that spills replays its runs, merged, through a fully ordered OrderedFinalAggregateStream, which cannot spill and resized its reservation with try_resize after every batch. In a Comet task the native shuffle writer reading the aggregate buffers its output until Spark refuses it, so it soon holds the share the aggregate released, and Spark cannot make it spill; the spill merge feeding the replay also seats as many runs as the pool grants. The replay's next resize was refused and the task failed with "Additional allocation failed for FinalHashAggregateStream[0] ... ShuffleRepartitioner[0] consumed 4.2 GB" (538 times on 2026-10-05, among them ads_manual_product_rank_cohorts). DataFusion 55.1.0 does the same. A fully ordered table holds only the groups the last batch may continue, all of it already allocated, so record it with the infallible resize. Comet carries the shortfall as overcommit and refuses the next fallible request, here the shuffle writer's, until it is repaid, so the writer spills and the memory comes back. A partially ordered aggregate that cannot spill because temporary files are disabled keeps the error. DataFusion's test_sort_reservation_fails_during_spill expected the replay's error in a 500-byte pool; it now checks the query completes and releases everything. Tests: a vendored final aggregate read by a stand-in for the shuffle writer in pools of 4 to 12 MiB, and a Comet final aggregate under ShuffleWriterExec with a fixed Spark share, with and without a writer buffer limit; both failed with the production message and now stay within 5% of the share. Co-Authored-By: Claude Opus 5.5 (1M context) --- native/core/src/execution/jni_api.rs | 151 ++++++++++++++++++ .../src/aggregates/mod.rs | 15 +- .../src/aggregates/ordered_final_stream.rs | 22 ++- .../src/aggregates/starved_spill_tests.rs | 111 ++++++++++++- 4 files changed, 282 insertions(+), 17 deletions(-) diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index b2b7ec52efd..d7572f04078 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -3809,4 +3809,155 @@ mod native_sort_spill_tests { } } } + + /// Partial states of `count(v)` and `sum(v)` per `(a, b)` for 400k groups in random + /// order, as a final aggregate reads them from a shuffle. + #[derive(Debug)] + struct GroupStates { + schema: SchemaRef, + rows_per_batch: usize, + num_batches: usize, + } + + impl GroupStates { + fn new(rows_per_batch: usize, num_batches: usize) -> Self { + Self { + schema: Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Utf8, false), + Field::new("count_v[count]", DataType::Int64, false), + Field::new("sum_v[sum]", DataType::Int64, true), + ])), + rows_per_batch, + num_batches, + } + } + + fn batch(&self, index: usize) -> RecordBatch { + let start = (index * self.rows_per_batch) as u64; + let keys: Vec = (start..start + self.rows_per_batch as u64) + .map(|r| mix(r) % 400_000) + .collect(); + RecordBatch::try_new( + Arc::clone(&self.schema), + vec![ + Arc::new(Int32Array::from_iter_values( + keys.iter().map(|&k| (k % 1000) as i32), + )), + Arc::new(StringArray::from_iter_values( + keys.iter().map(|&k| format!("key-{k:020}")), + )), + Arc::new(Int64Array::from_iter_values(keys.iter().map(|_| 1))), + Arc::new(Int64Array::from_iter_values(keys.iter().map(|&k| k as i64))), + ], + ) + .unwrap() + } + } + + impl PartitionStream for GroupStates { + fn schema(&self) -> &SchemaRef { + &self.schema + } + + fn execute(&self, _ctx: Arc) -> SendableRecordBatchStream { + let batches: Vec<_> = (0..self.num_batches).map(|i| Ok(self.batch(i))).collect(); + Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&self.schema), + futures::stream::iter(batches), + )) + } + } + + /// The final aggregate of `GroupStates` written by the native shuffle writer, as the + /// stage of ads_manual_product_rank_cohorts that failed. + fn aggregate_into_shuffle( + max_buffer_bytes: Option, + output: &std::path::Path, + ) -> (Arc, Arc) { + use datafusion::functions_aggregate::count::count_udaf; + use datafusion::functions_aggregate::sum::sum_udaf; + use datafusion::physical_expr::aggregate::AggregateExprBuilder; + use datafusion::physical_plan::aggregates::{ + AggregateExec, AggregateMode, PhysicalGroupBy, + }; + use datafusion_comet_shuffle::{CometPartitioning, CompressionCodec, ShuffleWriterExec}; + + let source = Arc::new(GroupStates::new(4096, 200)); + let states = Arc::clone(&source.schema); + let raw = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Utf8, false), + Field::new("v", DataType::Int64, false), + ])); + let aggregate = |udaf, alias: &str| { + Arc::new( + AggregateExprBuilder::new(udaf, vec![col("v", &raw).unwrap()]) + .schema(Arc::clone(&raw)) + .alias(alias) + .build() + .unwrap(), + ) + }; + let group_by = PhysicalGroupBy::new_single( + ["a", "b"] + .map(|key| (col(key, &states).unwrap(), key.to_string())) + .to_vec(), + ); + let final_aggregate: Arc = Arc::new( + AggregateExec::try_new( + AggregateMode::Final, + group_by, + vec![ + aggregate(count_udaf(), "count_v"), + aggregate(sum_udaf(), "sum_v"), + ], + vec![None, None], + streaming(source), + raw, + ) + .unwrap(), + ); + let output_schema = final_aggregate.schema(); + let writer: Arc = Arc::new( + ShuffleWriterExec::try_new( + Arc::clone(&final_aggregate), + CometPartitioning::Hash(vec![col("b", &output_schema).unwrap()], 16), + CompressionCodec::None, + output.to_string_lossy().into_owned(), + false, + 1024 * 1024, + max_buffer_bytes, + ) + .unwrap(), + ); + (writer, final_aggregate) + } + + /// A final aggregate that spilled replays its runs through a fully ordered aggregate, + /// which cannot spill. The native shuffle writer above it buffers the replay's output + /// until Spark refuses it, so it soon holds the share the aggregate released, and + /// Spark cannot make it spill. Without a buffer limit the replay then failed with + /// "Additional allocation failed for FinalHashAggregateStream[0] ... + /// ShuffleRepartitioner[0] ... consumed", as ads_manual_product_rank_cohorts did. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn spilled_final_aggregate_feeding_the_native_shuffle_writer_finishes() { + let output = tempfile::tempdir().unwrap(); + for share in [8, 12, 16, 24].map(|mb| mb * MB) { + for max_buffer_bytes in [None, Some(share / 4)] { + let (writer, aggregate) = + aggregate_into_shuffle(max_buffer_bytes, &output.path().join("data")); + match run_with_fixed_share(share, 8, Arc::clone(&writer)).await { + Ok((_, peak)) => { + assert!(peak <= share + share / 20, "{peak} reserved of {share}") + } + Err(e) => panic!( + "aggregate failed instead of spilling at {share}, \ + max_buffer_bytes {max_buffer_bytes:?}: {e}" + ), + } + assert!(spill_count(&aggregate) > 0, "the aggregate did not spill"); + } + } + } } diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/mod.rs b/native/vendor/datafusion-physical-plan/src/aggregates/mod.rs index ebbf357fa4a..ddcf46a69c2 100644 --- a/native/vendor/datafusion-physical-plan/src/aggregates/mod.rs +++ b/native/vendor/datafusion-physical-plan/src/aggregates/mod.rs @@ -7012,16 +7012,11 @@ mod tests { let task_ctx = new_spill_ctx(1, 500); let result = collect(aggr.execute(0, Arc::clone(&task_ctx))?).await; - match &result { - Ok(_) => panic!("Expected ResourcesExhausted error but query succeeded"), - Err(e) => { - let root = e.find_root(); - assert!( - matches!(root, DataFusionError::ResourcesExhausted(_)), - "Expected ResourcesExhausted, got: {root}", - ); - } - } + // COMET PATCH: the spill succeeds, and the fully ordered replay that failed for + // the pool records the few groups it holds instead, see `OrderedFinalAggregateStream`. + let batches = result?; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 3); + assert_eq!(task_ctx.runtime_env().memory_pool.reserved(), 0); Ok(()) } diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/ordered_final_stream.rs b/native/vendor/datafusion-physical-plan/src/aggregates/ordered_final_stream.rs index 90892a4aea8..5696a6fdc10 100644 --- a/native/vendor/datafusion-physical-plan/src/aggregates/ordered_final_stream.rs +++ b/native/vendor/datafusion-physical-plan/src/aggregates/ordered_final_stream.rs @@ -67,6 +67,9 @@ pub(crate) struct OrderedFinalAggregateStream { reservation: MemoryReservation, baseline_metrics: BaselineMetrics, state: Option, + /// COMET PATCH: the group keys are fully ordered, so the table holds only the groups + /// the last batch may continue. + fully_ordered: bool, } /// Spill configuration and accumulated runs for partially ordered final @@ -370,6 +373,7 @@ impl OrderedFinalAggregateStream { table, spill_context, }), + fully_ordered: *input_order_mode == InputOrderMode::Sorted, }) } @@ -446,15 +450,21 @@ impl OrderedFinalAggregateStream { // Check memory reservation, and potentially spill. let timer = elapsed_compute.timer(); - let resize_result = - self.reservation - .try_resize(Self::reservation_size_for_table( - &table, - spill_context.as_deref(), - )); + let size = + Self::reservation_size_for_table(&table, spill_context.as_deref()); + let resize_result = self.reservation.try_resize(size); timer.done(); match resize_result { Ok(()) => {} + // COMET PATCH: a fully ordered table cannot spill and holds only memory + // already allocated for the groups the last batch may continue. Record + // it: Comet carries the shortfall as overcommit and refuses the next + // fallible request, such as the native shuffle writer's, until repaid. + Err(DataFusionError::ResourcesExhausted(_)) + if spill_context.is_none() && self.fully_ordered => + { + self.reservation.resize(size); + } Err(e @ DataFusionError::ResourcesExhausted(_)) => { let Some(spill_context) = spill_context else { // `None` means spilling is not supported, see comments diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/starved_spill_tests.rs b/native/vendor/datafusion-physical-plan/src/aggregates/starved_spill_tests.rs index 687d4c8586a..83aa06b6f9b 100644 --- a/native/vendor/datafusion-physical-plan/src/aggregates/starved_spill_tests.rs +++ b/native/vendor/datafusion-physical-plan/src/aggregates/starved_spill_tests.rs @@ -19,6 +19,7 @@ //! batch, and what the emptied table still holds must not fail the task. use std::fmt::{Debug, Formatter}; +use std::num::NonZeroUsize; use std::sync::{Arc, Mutex}; use arrow::array::{Int64Array, RecordBatch, StringArray, UInt32Array}; @@ -29,7 +30,7 @@ use datafusion_common::Result; use datafusion_execution::TaskContext; use datafusion_execution::config::SessionConfig; use datafusion_execution::memory_pool::{ - GreedyMemoryPool, MemoryConsumer, MemoryPool, MemoryReservation, + GreedyMemoryPool, MemoryConsumer, MemoryPool, MemoryReservation, TrackConsumersPool, }; use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_functions_aggregate::count::count_udaf; @@ -321,3 +322,111 @@ async fn ordered_final_aggregate_survives_a_starved_first_spill() -> Result<()> }) .await } + +/// Partial states of `count(v)` and `sum(v)` grouped by `(a, b)`, for 200k groups in +/// random order. +fn many_group_states(batches: u32, rows: u32) -> Result<(SchemaRef, Vec)> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::Utf8, false), + Field::new("count_v[count]", DataType::Int64, false), + Field::new("sum_v[sum]", DataType::Int64, true), + ])); + let mut state = 11u64; + let batches = (0..batches) + .map(|_| { + let keys: Vec = (0..rows) + .map(|_| { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + (state >> 33) % 200_000 + }) + .collect(); + Ok(RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(UInt32Array::from_iter_values( + keys.iter().map(|k| (k % 1000) as u32), + )), + Arc::new(StringArray::from_iter_values( + keys.iter().map(|k| format!("key-{k:020}")), + )), + Arc::new(Int64Array::from_iter_values(keys.iter().map(|_| 1))), + Arc::new(Int64Array::from_iter_values( + keys.iter().map(|k| *k as i64), + )), + ], + )?) + }) + .collect::>>()?; + Ok((schema, batches)) +} + +/// Reads `stream` as the native shuffle writer does: it keeps every batch reserved until +/// the pool refuses one, and then spills them all. +async fn read_as_shuffle_writer( + mut stream: SendableRecordBatchStream, + pool: &Arc, +) -> Result<(Vec, usize)> { + let writer = MemoryConsumer::new("ShuffleRepartitioner[0]") + .with_can_spill(true) + .register(pool); + let mut output = vec![]; + let mut spills = 0; + while let Some(batch) = stream.next().await { + let batch = batch?; + if writer.try_grow(batch.get_array_memory_size()).is_err() { + spills += 1; + writer.free(); + } + output.push(batch); + } + Ok((output, spills)) +} + +/// A final aggregate that spilled replays its runs through a fully ordered aggregate, +/// which cannot spill. The native shuffle writer reading it buffers its output until the +/// pool refuses it, so it soon holds nearly all the memory the aggregate released, and +/// the replay then failed with "Additional allocation failed for +/// FinalHashAggregateStream[0]" while the writer, which can spill, kept the rest. +#[tokio::test] +async fn spilled_final_aggregate_survives_a_shuffle_writer_that_takes_the_pool() +-> Result<()> { + let raw_schema = raw_schema(); + let (input_schema, batches) = many_group_states(100, 4096)?; + let unconstrained_input = TestMemoryExec::try_new_exec( + std::slice::from_ref(&batches), + Arc::clone(&input_schema), + None, + )?; + let unconstrained = aggregate( + AggregateMode::Final, + &["a", "b"], + unconstrained_input, + &raw_schema, + )?; + let expected = collect(unconstrained.execute(0, task_ctx(None)?)?).await?; + let expected = sorted_output(&unconstrained.schema(), &expected)?; + + for limit in [4 << 20, 6 << 20, 8 << 20, 12 << 20] { + let pool: Arc = Arc::new(TrackConsumersPool::new( + GreedyMemoryPool::new(limit), + NonZeroUsize::new(5).unwrap(), + )); + let input = TestMemoryExec::try_new_exec( + std::slice::from_ref(&batches), + Arc::clone(&input_schema), + None, + )?; + let starved = aggregate(AggregateMode::Final, &["a", "b"], input, &raw_schema)?; + let stream = starved.execute(0, task_ctx(Some(Arc::clone(&pool)))?)?; + let (actual, writer_spills) = read_as_shuffle_writer(stream, &pool).await?; + assert!(spill_count(&starved) > 0, "the aggregate must spill"); + assert!(writer_spills > 0, "the writer must spill"); + assert_eq!(sorted_output(&starved.schema(), &actual)?, expected); + drop(starved); + assert_eq!(pool.reserved(), 0); + } + Ok(()) +} From 3dbd273cd31783c390b605c6f0e7ba8bdfca201d Mon Sep 17 00:00:00 2001 From: msaf Date: Mon, 5 Oct 2026 12:31:07 +0100 Subject: [PATCH 06/14] fix: price array grouping keys and dispatched key expressions in the cost-based engine choice A native hash aggregate over grouping keys that hold arrays costs far more than the agg line predicts: star_order_2020's distinct over 18 key leaves, 7 of them in arrays, took 4.0 us per row in each native phase against 0.7 to 1.0 us in Spark, while the model priced Comet at 0.56 and kept it native. Two classes are added on top of agg: aggArrayKey per leaf of a key holding an array (half per phase), fitted on star_order_2020 and calib29's aggkeys, and codegenDispatch once per leaf of a key computed through the JVM codegen dispatcher, such as the transform normalizing doubles inside an array. Aggregates without such keys keep their prices. Co-Authored-By: Claude Opus 5.5 (1M context) --- docs/source/user-guide/latest/tuning.md | 2 + .../scala/org/apache/comet/CometConf.scala | 5 +- .../comet/rules/CostBasedEngineChoice.scala | 39 +++++--- .../apache/comet/rules/EngineCostTable.scala | 15 +++ .../org/apache/comet/rules/LeafColumns.scala | 8 ++ .../rules/CostBasedEngineChoiceSuite.scala | 99 +++++++++++++++++++ 6 files changed, 154 insertions(+), 14 deletions(-) diff --git a/docs/source/user-guide/latest/tuning.md b/docs/source/user-guide/latest/tuning.md index d0cecae67af..ce3a8006050 100644 --- a/docs/source/user-guide/latest/tuning.md +++ b/docs/source/user-guide/latest/tuning.md @@ -614,6 +614,8 @@ prices per row from a table of measurements: `c0 + k0*L + k1*L*min(L, 600)` ns n - `agg` for hash, object hash and sort aggregates, over the leaves of the grouping keys, half for each phase of a two-phase aggregate, plus the price of the class of each aggregate function (`aggDeclarative`, `aggCollectList`, `aggCollectSet`, `aggPercentile`, `aggPercentileApprox`, `aggOther`) and `aggObjectHash` for an object hash aggregate. + The leaves of grouping keys holding an array add `aggArrayKey`, also half for each phase, and the leaves of a grouping + key computed through the JVM codegen dispatcher add `codegenDispatch` once, in the phase that computes it. - `window` over the leaves of its input, plus `windowAggregate`, `windowOffset` or `windowRank` for each window function, at `L` the number of window functions; `wglPartial` and `wglFinal` for window group limits. - `expand`, per projection, and `generate`: free with Spark's whole-stage codegen, `expandNoCodegen` and diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 140008a8ab8..acc6e57215f 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -691,8 +691,9 @@ object CometConf extends ShimCometConf { "window, the number of functions). The classes are shuffleWrite, shuffleRead, sort, " + "sortSpill, smj, bhj, predicate, projectPassThrough, expression, agg, aggObjectHash, " + "aggDeclarative, aggCollectList, aggCollectSet, aggPercentile, aggPercentileApprox, " + - "aggOther, window, windowAggregate, windowOffset, windowRank, wglPartial, wglFinal, " + - "expand, generate, rowLocal, the comet-only c2r and r2c, and the spark-only " + + "aggOther, aggArrayKey, codegenDispatch, window, windowAggregate, windowOffset, " + + "windowRank, wglPartial, wglFinal, expand, generate, rowLocal, the comet-only c2r and " + + "r2c, and the spark-only " + "expressionOverScan, aggDeclarativeNoCodegen, expandNoCodegen and " + "generateNoCodegen. A row whose leaves are a fraction f inside structs, arrays or " + "maps costs (1 - f) times the flat price plus f times the nested one. The scalars " + diff --git a/spark/src/main/scala/org/apache/comet/rules/CostBasedEngineChoice.scala b/spark/src/main/scala/org/apache/comet/rules/CostBasedEngineChoice.scala index c6f39d24048..e9a89cff70c 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CostBasedEngineChoice.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CostBasedEngineChoice.scala @@ -25,7 +25,7 @@ import scala.collection.mutable import org.apache.spark.internal.Logging import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, NamedExpression} +import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, Expression, NamedExpression, ScalaUDF} import org.apache.spark.sql.catalyst.expressions.aggregate.{Complete, Partial} import org.apache.spark.sql.catalyst.plans.logical.statsEstimation.EstimationUtils import org.apache.spark.sql.catalyst.rules.Rule @@ -39,7 +39,7 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.comet.CometConf import org.apache.comet.CometSparkSessionExtensions.withFallbackReason import org.apache.comet.rules.BoundaryFormats._ -import org.apache.comet.serde.QueryPlanSerde +import org.apache.comet.serde.{CometCodegenDispatch, QueryPlanSerde} import org.apache.comet.shims.ShimCometWindowGroupLimit /** @@ -56,15 +56,17 @@ import org.apache.comet.shims.ShimCometWindowGroupLimit * those of its input; a filter's predicate those it references; an aggregate's those of its * grouping keys, at half the price for each phase of a two-phase aggregate. A project computes * once per leaf of the expressions it does not pass through; an expand copies once per - * projection. The functions of an aggregate cost the price of their class each, and those of a - * window the price of their class at the number of its window functions. In Spark, an operator - * whose rows, or its inputs', have more leaves than `codegenMaxFields` runs without whole-stage - * codegen and takes the classes of [[EngineCostTable.CostClass.withoutCodegen]], and a project - * whose input comes from a scan through filters, projects and conversions only passes its columns - * for free and computes at `expressionOverScan`. Shuffles cost their write and read over every - * leaf of the shuffled rows, scaled by their partitions, and their bytes beyond - * `perByteLeafAllowance` per leaf. A conversion to rows costs `c2r` over every leaf, and Comet's - * columnar shuffle also pays one `r2c` when it writes. + * projection. The functions of an aggregate cost the price of their class each, the leaves of its + * grouping keys that hold an array `aggArrayKey` on top of `agg`, and the leaves of a grouping + * key computed through the JVM codegen dispatcher `codegenDispatch` once; those of a window the + * price of their class at the number of its window functions. In Spark, an operator whose rows, + * or its inputs', have more leaves than `codegenMaxFields` runs without whole-stage codegen and + * takes the classes of [[EngineCostTable.CostClass.withoutCodegen]], and a project whose input + * comes from a scan through filters, projects and conversions only passes its columns for free + * and computes at `expressionOverScan`. Shuffles cost their write and read over every leaf of the + * shuffled rows, scaled by their partitions, and their bytes beyond `perByteLeafAllowance` per + * leaf. A conversion to rows costs `c2r` over every leaf, and Comet's columnar shuffle also pays + * one `r2c` when it writes. */ class EngineCostModel( val table: EngineCostTable, @@ -111,6 +113,13 @@ class EngineCostModel( } } + /** Whether Comet evaluates `expression` through the JVM codegen dispatcher. */ + def dispatchedThroughCodegen(expression: Expression): Boolean = expression.exists { + case _: ScalaUDF => true + case e => + QueryPlanSerde.exprSerdeMap.get(e.getClass).exists(_.isInstanceOf[CometCodegenDispatch[_]]) + } + private def aggregateShare(agg: BaseAggregateExec): Double = if (agg.aggregateExpressions.exists(_.mode == Complete)) 1.0 else 0.5 @@ -149,9 +158,15 @@ class EngineCostModel( } case (Agg, agg: BaseAggregateExec) => val share = aggregateShare(agg) + val keys = agg.groupingExpressions val functions = agg.aggregateExpressions.map(e => aggregateFunctionClass(e.aggregateFunction)) - Term(Agg, widthOfTypes(agg.groupingExpressions.map(_.dataType)), share) +: + val arrayKeys = widthOfTypes(keys.map(_.dataType).filter(LeafColumns.containsArray)) + val dispatched = widthOfTypes(keys.filter(dispatchedThroughCodegen).map(_.dataType)) + Seq( + Some(Term(Agg, widthOfTypes(keys.map(_.dataType)), share)), + Some(Term(AggArrayKey, arrayKeys, share)).filter(_.width.leaves > 0), + Some(Term(CodegenDispatch, dispatched)).filter(_.width.leaves > 0)).flatten ++ functions.distinct.map { c => Term(c, Width(functions.size, 0), share * functions.count(_ == c)) } diff --git a/spark/src/main/scala/org/apache/comet/rules/EngineCostTable.scala b/spark/src/main/scala/org/apache/comet/rules/EngineCostTable.scala index 3784ab11db1..1c4810f2cf5 100644 --- a/spark/src/main/scala/org/apache/comet/rules/EngineCostTable.scala +++ b/spark/src/main/scala/org/apache/comet/rules/EngineCostTable.scala @@ -186,6 +186,8 @@ object EngineCostTable { case object AggPercentile extends CostClass("aggPercentile") case object AggPercentileApprox extends CostClass("aggPercentileApprox") case object AggOther extends CostClass("aggOther") + case object AggArrayKey extends CostClass("aggArrayKey") + case object CodegenDispatch extends CostClass("codegenDispatch") case object Window extends CostClass("window") case object WindowAggregate extends CostClass("windowAggregate") case object WindowOffset extends CostClass("windowOffset") @@ -219,6 +221,8 @@ object EngineCostTable { AggPercentile, AggPercentileApprox, AggOther, + AggArrayKey, + CodegenDispatch, Window, WindowAggregate, WindowOffset, @@ -312,6 +316,15 @@ object EngineCostTable { * beyond `spark.sql.codegen.maxFields`), `aggCollectList`, `aggCollectSet`, * `aggPercentile`, `aggPercentileApprox` or `aggOther` (other imperative aggregates, by * analogy, unmeasured); an object hash aggregate adds `aggObjectHash`. + * - `aggArrayKey`: what each leaf of a grouping key holding an array adds to `agg`, partial + * and final together, fitted by least squares on two distincts without reduction: + * star_order_2020 (18 key leaves, 7 in arrays of 2.7 and 1.1 elements on average; Comet 4.0 + * us per row in each phase, Spark 0.71 to 0.98) and calib29's `aggkeys` over an array of + * 3-leaf structs of 2.7 elements (5 leaves, 3 in the array; Comet 2.4 us, Spark 1.7). + * `codegenDispatch`: once per leaf of a grouping key computed through the JVM codegen + * dispatcher, such as the `transform` normalizing the floating-point fields of an array, in + * the phase that computes it: Comet 0.65 us for 3 leaves on star_order_2020 (partial minus + * final), Spark 0.4 us in a local micro-benchmark (Spark's codegen of the same transform). * - `window`: the operator with one `row_number`, over the leaves of its input (Spark noisy, * maxrel 0.46). Each window function adds the price of its class at L = the number of * window functions of the operator: `windowAggregate` (measured on a running sum, other @@ -380,6 +393,8 @@ object EngineCostTable { both(AggPercentile, Line(130, 0, 0, 2400, 0)) ++ both(AggPercentileApprox, Line(270, 0, 0, 3900, 0)) ++ both(AggOther, Line(100, 0, 0, 1700, 0)) ++ + both(AggArrayKey, Line(0, 937, 0, 0, 76)) ++ + both(CodegenDispatch, Line(0, 217, 0, 0, 133)) ++ both(WindowAggregate, Line(380, 0, 0, 20, 0.8)) ++ both(WindowOffset, Line(60, 0, 0, 0, 0.25)) ++ both(WindowRank, Line(0, 0, 0, 0, 0)) ++ diff --git a/spark/src/main/scala/org/apache/comet/rules/LeafColumns.scala b/spark/src/main/scala/org/apache/comet/rules/LeafColumns.scala index efbea9567eb..30f8c3fbabf 100644 --- a/spark/src/main/scala/org/apache/comet/rules/LeafColumns.scala +++ b/spark/src/main/scala/org/apache/comet/rules/LeafColumns.scala @@ -40,6 +40,14 @@ object LeafColumns { case _ => false } + def containsArray(dataType: DataType): Boolean = dataType match { + case _: ArrayType => true + case struct: StructType => struct.fields.exists(f => containsArray(f.dataType)) + case map: MapType => containsArray(map.keyType) || containsArray(map.valueType) + case udt: UserDefinedType[_] => containsArray(udt.sqlType) + case _ => false + } + /** The attributes that none of `keys` references. */ def outside(attributes: Seq[Attribute], keys: Seq[Expression]): Seq[Attribute] = { val referenced = AttributeSet(keys.flatMap(_.references)) diff --git a/spark/src/test/scala/org/apache/comet/rules/CostBasedEngineChoiceSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CostBasedEngineChoiceSuite.scala index 3016af423cf..16a0ec8d8a0 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CostBasedEngineChoiceSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CostBasedEngineChoiceSuite.scala @@ -78,6 +78,42 @@ class CostBasedEngineChoiceSuite extends CometTestBase { } } + /** Runs `f` with view `star`, the schema of star_order_2020's purchases, 4 rows per order. */ + private def withStar(f: => Unit): Unit = { + withTempPath { dir => + spark + .range(2000) + .selectExpr( + "cast(id % 500 AS string) AS order_id", + "id AS event_ts", + "concat('c', cast(id % 37 AS string)) AS category_id", + "IF(id % 2 = 0, 'android', 'ios') AS os_type", + "array(named_struct('amount', cast(id % 13 AS double) / 4, 'type', 'coupon', " + + "'subType', cast(id % 3 AS string))) AS discounts", + "named_struct('name', cast(id % 7 AS string), 'activation_ts', id, " + + "'is_last_context', true, 'is_adtech_promoted', id % 5 = 0) AS last_context", + "array(named_struct('name', cast(id % 7 AS string), 'activation_ts', id, " + + "'is_last_context', true, 'is_adtech_promoted', false)) AS normalized_contexts", + "'joom' AS legal_entity", + "'joom' AS app_entity", + "'joom' AS app_entity_group", + "IF(id % 4 = 0, NULL, 'd') AS custom_domain") + .write + .parquet(dir.getCanonicalPath) + spark.read.parquet(dir.getCanonicalPath).createOrReplaceTempView("star") + withTempView("star")(f) + } + } + + private val starColumns = "order_id, category_id, os_type, discounts, last_context, " + + "normalized_contexts, legal_entity, app_entity, app_entity_group, custom_domain" + + /** The latest purchase of each order, deduplicated, as star_order_2020 reads them. */ + private def starQuery: DataFrame = + sql( + s"SELECT DISTINCT $starColumns FROM (SELECT *, rank() OVER (PARTITION BY order_id " + + "ORDER BY event_ts DESC) AS r FROM star) WHERE r = 1") + private def run(df: => DataFrame): SparkPlan = checkSparkAnswer(df)._2 private def run(query: String): SparkPlan = run(sql(query)) @@ -779,6 +815,69 @@ class CostBasedEngineChoiceSuite extends CometTestBase { } } + test("aggregate keys holding arrays and keys computed by the codegen dispatcher add classes") { + withStar { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", flag -> "false") { + val plan = runUnordered(sql(s"SELECT DISTINCT $starColumns FROM star")) + val aggs = nodes(plan).collect { case a: CometHashAggregateExec => a } + assert(aggs.size == 2, s"plan:\n$plan") + val (partial, finals) = aggs.partition(a => + model.dispatchedThroughCodegen( + a.originalPlan + .asInstanceOf[HashAggregateExec] + .groupingExpressions + .find(_.name == "discounts") + .get)) + assert(partial.size == 1 && finals.size == 1, s"plan:\n$plan") + val keys = Term(Agg, Width(18, 11), 0.5) + val arrays = Term(AggArrayKey, Width(7, 7), 0.5) + assert( + model.terms(partial.head, Engine.Comet) == + Seq(keys, arrays, Term(CodegenDispatch, Width(3, 3)))) + assert(model.terms(finals.head, Engine.Comet) == Seq(keys, arrays)) + val table = EngineCostTable.default + same( + model.operatorPrice(partial.head, Engine.Comet), + 0.5 * table.comet(Agg, Width(18, 11)) + 0.5 * 937 * 7 + 217 * 3) + same( + model.operatorPrice(partial.head, Engine.Spark), + 0.5 * table.spark(Agg, Width(18, 11)) + 0.5 * 76 * 7 + 133 * 3) + same( + model.operatorPrice(finals.head, Engine.Comet), + 0.5 * table.comet(Agg, Width(18, 11)) + 0.5 * 937 * 7) + + val flat = runUnordered(sql("SELECT DISTINCT order_id, category_id FROM star")) + nodes(flat).collect { case a: CometHashAggregateExec => a }.foreach { agg => + assert(model.terms(agg, Engine.Comet) == Seq(Term(Agg, Width(2, 0), 0.5))) + } + } + } + } + + for (aqe <- Seq("false", "true")) { + test( + s"a distinct over array keys after a window runs in Spark, as in star_order_2020 (AQE=$aqe)") { + withStar { + withAqe(aqe) { + def aggregates(plan: SparkPlan): (Int, Int) = + ( + count(plan) { case a: CometHashAggregateExec => a }, + count(plan) { case a: HashAggregateExec => a }) + val (off, on) = offAndOn(runUnordered(starQuery)) + assert(aggregates(off) == (2, 0), s"plan:\n$off") + assert(aggregates(on) == (0, 2), s"plan:\n$on") + withSQLConf( + flag -> "true", + costTable -> ("aggArrayKey.comet=0,0,0;aggArrayKey.spark=0,0;" + + "codegenDispatch.comet=0,0,0;codegenDispatch.spark=0,0")) { + val before = runUnordered(starQuery) + assert(aggregates(before) == (2, 0), s"plan:\n$before") + } + } + } + } + } + test("a window costs its line and the classes of its functions") { withTables { withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", flag -> "false") { From 55c5b6fe8b45b670a450d4a79479f0fdfd6e60a6 Mon Sep 17 00:00:00 2001 From: msaf Date: Mon, 5 Oct 2026 12:08:12 +0100 Subject: [PATCH 07/14] fix: stop a killed task's native plan instead of running it to completion Spark kills a task by marking its TaskContext and, only if asked to, interrupting its thread. Neither reaches a thread inside Native.executePlan: a plan without JVM input runs on a Tokio task while the task thread parks in blocking_recv until the next batch, which a plan rooted at a native shuffle writer only sends once its whole input is written, and a JVM-fed plan is polled on the task thread, where an operator that keeps polling a child that never returns Pending (an aggregate over a join multiplying its probe rows) runs without a break. JVM code called back from native (Comet's ScalaUDF kernels, the row to Arrow conversion) checked for a kill only between batches, or not at all. With spark.task.reaper.enabled the reaper then killed the whole executor after killTimeout: 43 thrift-server executors on 2026-10-05, 21 of them with the task thread inside executePlan under a native shuffle write and 20 more in JVM code it had called back, 19 of them ScalaUDF kernels. Each native plan now has a cancellation, set by cancelPlan from a watcher that polls the task's kill flag every 50 ms. It ends the task thread's wait on the producer task and on native I/O, aborts the producer task, makes every operator of the plan end its stream with an error at its next batch (a transparent wrapper added to each native operator), and stops a shuffle writer's spill or final write at its next batch. The UDF kernels and the row to Arrow reader check for a kill every 256 rows. A cancelled plan's failure is reported as a TaskKilledException. releasePlan aborts the producer task and waits up to a second for it to drop its stream, so that its memory is returned before the task ends. getShufflePartitionOffsets now downcasts through the wrapper. Co-Authored-By: Claude Opus 5.5 (1M context) --- native/common/src/cancellation.rs | 166 ++++++++ native/common/src/lib.rs | 1 + native/core/src/execution/jni_api.rs | 195 +++++++++- .../src/execution/operators/cancellable.rs | 366 ++++++++++++++++++ native/core/src/execution/operators/mod.rs | 2 + native/core/src/execution/spark_plan.rs | 5 +- .../src/partitioners/multi_partition.rs | 9 + .../partitioned_batch_iterator.rs | 40 +- native/shuffle/src/shuffle_writer.rs | 24 +- .../org/apache/comet/udf/CometUdfBridge.java | 3 + .../org/apache/comet/CometExecIterator.scala | 72 +++- .../main/scala/org/apache/comet/Native.scala | 10 + .../codegen/CometBatchKernelCodegen.scala | 4 + .../spark/comet/CometTaskContextShim.scala | 5 + .../execution/arrow/RowArrowReader.scala | 6 + .../sql/comet/CometNativeTaskKillSuite.scala | 217 +++++++++++ 16 files changed, 1084 insertions(+), 41 deletions(-) create mode 100644 native/common/src/cancellation.rs create mode 100644 native/core/src/execution/operators/cancellable.rs create mode 100644 spark/src/test/scala/org/apache/spark/sql/comet/CometNativeTaskKillSuite.scala diff --git a/native/common/src/cancellation.rs b/native/common/src/cancellation.rs new file mode 100644 index 00000000000..7aa60aca46f --- /dev/null +++ b/native/common/src/cancellation.rs @@ -0,0 +1,166 @@ +// 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. + +//! Cancellation of a native plan whose Spark task has been killed. The JVM sets it from outside +//! the task thread, and the plan's operators, its blocking waits and its Tokio tasks stop at +//! their next check instead of running the plan to completion. + +use datafusion::common::DataFusionError; +use datafusion::execution::TaskContext; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; +use std::task::{Context, Poll, Waker}; + +pub const CANCELLED_MESSAGE: &str = + "Comet native execution was cancelled because its Spark task was killed"; + +type Callback = Box; + +#[derive(Default)] +pub struct PlanCancellation { + cancelled: AtomicBool, + waiters: Mutex, +} + +#[derive(Default)] +struct Waiters { + wakers: Vec, + callbacks: Vec, +} + +impl std::fmt::Debug for PlanCancellation { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PlanCancellation") + .field("cancelled", &self.is_cancelled()) + .finish() + } +} + +impl PlanCancellation { + pub fn new() -> Arc { + Arc::new(Self::default()) + } + + /// The cancellation of the plan `context` runs for, if it has one. + pub fn of(context: &TaskContext) -> Option> { + context.session_config().get_extension::() + } + + pub fn cancel(&self) { + if self.cancelled.swap(true, Ordering::AcqRel) { + return; + } + let Waiters { wakers, callbacks } = std::mem::take(&mut *self.lock()); + wakers.into_iter().for_each(Waker::wake); + callbacks.into_iter().for_each(|callback| callback()); + } + + pub fn is_cancelled(&self) -> bool { + self.cancelled.load(Ordering::Acquire) + } + + pub fn check(&self) -> Result<(), DataFusionError> { + if self.is_cancelled() { + Err(cancelled_error()) + } else { + Ok(()) + } + } + + /// `Ready` once cancelled; otherwise wakes `cx` when it is. + pub fn poll_cancelled(&self, cx: &mut Context<'_>) -> Poll<()> { + if self.is_cancelled() { + return Poll::Ready(()); + } + let mut waiters = self.lock(); + if self.is_cancelled() { + return Poll::Ready(()); + } + if !waiters.wakers.iter().any(|w| w.will_wake(cx.waker())) { + waiters.wakers.push(cx.waker().clone()); + } + Poll::Pending + } + + /// Runs `callback` once the plan is cancelled, at once if it already is. + pub fn on_cancel(&self, callback: impl FnOnce() + Send + 'static) { + { + let mut waiters = self.lock(); + if !self.is_cancelled() { + waiters.callbacks.push(Box::new(callback)); + return; + } + } + callback() + } + + fn lock(&self) -> std::sync::MutexGuard<'_, Waiters> { + self.waiters.lock().unwrap_or_else(|e| e.into_inner()) + } +} + +pub fn cancelled_error() -> DataFusionError { + DataFusionError::Execution(CANCELLED_MESSAGE.to_string()) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::AtomicUsize; + use std::task::Wake; + + struct CountingWaker(AtomicUsize); + + impl Wake for CountingWaker { + fn wake(self: Arc) { + self.0.fetch_add(1, Ordering::SeqCst); + } + } + + #[test] + fn cancel_wakes_each_registered_waker_once_and_runs_callbacks() { + let cancellation = PlanCancellation::new(); + let counter = Arc::new(CountingWaker(AtomicUsize::new(0))); + let waker = Waker::from(Arc::clone(&counter)); + let mut cx = Context::from_waker(&waker); + assert!(cancellation.poll_cancelled(&mut cx).is_pending()); + assert!(cancellation.poll_cancelled(&mut cx).is_pending()); + let calls = Arc::new(AtomicUsize::new(0)); + let c = Arc::clone(&calls); + cancellation.on_cancel(move || { + c.fetch_add(1, Ordering::SeqCst); + }); + assert!(cancellation.check().is_ok()); + + cancellation.cancel(); + cancellation.cancel(); + assert_eq!(counter.0.load(Ordering::SeqCst), 1); + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert!(cancellation.poll_cancelled(&mut cx).is_ready()); + assert!(cancellation + .check() + .unwrap_err() + .to_string() + .contains(CANCELLED_MESSAGE)); + + let c = Arc::clone(&calls); + cancellation.on_cancel(move || { + c.fetch_add(1, Ordering::SeqCst); + }); + assert_eq!(calls.load(Ordering::SeqCst), 2); + } +} diff --git a/native/common/src/lib.rs b/native/common/src/lib.rs index 78521fad6d5..ffa9668bb92 100644 --- a/native/common/src/lib.rs +++ b/native/common/src/lib.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +pub mod cancellation; mod error; pub mod offset_extents; mod query_context; diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index d7572f04078..b0ee151a205 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -42,6 +42,7 @@ use datafusion::{ physical_plan::{display::DisplayableExecutionPlan, SendableRecordBatchStream}, prelude::{SessionConfig, SessionContext}, }; +use datafusion_comet_common::cancellation::{cancelled_error, PlanCancellation}; use datafusion_comet_common::decode_string_arrays; use datafusion_comet_proto::spark_expression::agg_expr::ExprStruct as AggExprStruct; use datafusion_comet_proto::spark_operator::{AggregateMode, Operator, ShuffleScan}; @@ -104,6 +105,7 @@ use std::{ }; use tokio::runtime::{Handle, Runtime}; use tokio::sync::mpsc; +use tokio::task::JoinHandle; use crate::execution::memory_pools::{create_memory_pool, parse_memory_pool_config}; use crate::execution::operators::{PartitionAggregateWindowEnabled, ScanExec, ShuffleScanExec}; @@ -134,6 +136,18 @@ use tikv_jemalloc_ctl::{epoch, stats}; static TOKIO_RUNTIME: Mutex> = Mutex::new(None); +/// The cancellation of every plan created and not yet released, by the address of its context, +/// so that `cancelPlan` reaches it without touching a context the task thread may be using. +static PLAN_CANCELLATIONS: OnceLock>>> = OnceLock::new(); + +fn plan_cancellations() -> &'static Mutex>> { + PLAN_CANCELLATIONS.get_or_init(|| Mutex::new(HashMap::new())) +} + +/// How long `releasePlan` waits for an aborted producer task to drop the plan's stream, and with +/// it the memory the stream holds, before the Spark task ends. +const PRODUCER_STOP_TIMEOUT: Duration = Duration::from_secs(1); + #[cfg(feature = "jemalloc")] fn log_jemalloc_usage() { let e = epoch::mib().unwrap(); @@ -525,6 +539,10 @@ struct ExecutionContext { pub stream: Option, /// Receives batches from a spawned tokio task (async I/O path) pub batch_receiver: Option>>, + /// The spawned tokio task that sends to `batch_receiver` + pub producer: Option>, + /// Set when the Spark task running this plan is killed + pub cancellation: Arc, /// Native metrics pub metrics: Arc>>, // The interval in milliseconds to update metrics @@ -679,8 +697,10 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( // We need to keep the session context alive. Some session state like temporary // dictionaries are stored in session context. If it is dropped, the temporary // dictionaries will be dropped as well. + let cancellation = PlanCancellation::new(); let session = prepare_datafusion_session_context( batch_size as usize, + Arc::clone(&cancellation), memory_pool, local_dirs_vec, max_temp_directory_size, @@ -742,6 +762,8 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( input_sources, stream: None, batch_receiver: None, + producer: None, + cancellation: Arc::clone(&cancellation), metrics, metrics_update_interval, metrics_last_update_time: Instant::now(), @@ -761,11 +783,33 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( memory_pool_registration, }); - Ok(Box::into_raw(exec_context) as i64) + let exec_context = Box::into_raw(exec_context) as i64; + plan_cancellations() + .lock() + .insert(exec_context, cancellation); + Ok(exec_context) }) }) } +/// Cancels a plan whose Spark task has been killed. Called from outside the task thread, at any +/// point before `releasePlan`: an `executePlan` in progress, or the next one, then fails at its +/// next check instead of running the plan to completion. +#[no_mangle] +pub extern "system" fn Java_org_apache_comet_Native_cancelPlan( + e: EnvUnowned, + _class: JClass, + exec_context: jlong, +) { + try_unwrap_or_throw(&e, |_| { + let cancellation = plan_cancellations().lock().get(&exec_context).cloned(); + if let Some(cancellation) = cancellation { + cancellation.cancel(); + } + Ok(()) + }) +} + /// Binds one task-owned shuffle callback before native execution is initialized. /// /// Keeping callback registration separate preserves the existing `createPlan` JNI ABI for @@ -860,6 +904,7 @@ fn configure_sort_spill_reservation( #[allow(clippy::too_many_arguments)] fn prepare_datafusion_session_context( batch_size: usize, + cancellation: Arc, memory_pool: Arc, local_dirs: Vec, max_temp_directory_size: u64, @@ -882,7 +927,8 @@ fn prepare_datafusion_session_context( // This DataFusion context is within the scope of an executing Spark Task. We want to set // its internal parallelism to the number of CPUs allocated to Spark Tasks. This can be // modified by changing spark.task.cpus in the Spark config. - .with_batch_size(batch_size); + .with_batch_size(batch_size) + .with_extension(cancellation); if let Some(limit) = off_heap_limit { let executor_cores = spark_config.get_usize(SPARK_EXECUTOR_CORES, 1); @@ -1110,12 +1156,16 @@ impl Wake for WakeFlag { /// this thread's budget unconstrained. async fn next_batch( stream: &mut S, + cancellation: &PlanCancellation, mut on_pending: impl FnMut() -> Result<(), CometError>, ) -> Result, CometError> where S: Stream> + Unpin, { poll_fn(|cx| { + if cancellation.poll_cancelled(cx).is_ready() { + return Poll::Ready(Err(cancelled_error().into())); + } let flag = Arc::new(WakeFlag { woken: AtomicBool::new(false), parent: cx.waker().clone(), @@ -1127,6 +1177,9 @@ where // `on_pending` calls into the JVM, which can run another Comet plan on this thread. // `block_in_place` exits the runtime context so that plan's `block_on` doesn't panic. tokio::task::block_in_place(&mut on_pending)?; + if cancellation.is_cancelled() { + return Poll::Ready(Err(cancelled_error().into())); + } if flag.woken.load(Ordering::Acquire) { // Poll again at once: a nested `block_on` may have taken the wake-up. cx.waker().wake_by_ref(); @@ -1136,6 +1189,30 @@ where .await } +/// Waits for the next batch from the plan's producer task, or for the plan to be cancelled. +fn recv_unless_cancelled( + rx: &mut mpsc::Receiver>, + cancellation: &PlanCancellation, +) -> CometResult>> { + get_runtime().block_on(poll_fn(|cx| { + if cancellation.poll_cancelled(cx).is_ready() { + return Poll::Ready(Err(cancelled_error().into())); + } + rx.poll_recv(cx).map(Ok) + })) +} + +/// Aborts the plan's producer task and waits a bounded time for it to end, so that the stream it +/// owns is dropped, and its memory returned, before the Spark task ends. An abort takes effect at +/// the task's next yield, which an operator busy inside one poll may not reach in time. +fn stop_producer(producer: JoinHandle<()>) { + producer.abort(); + let deadline = Instant::now() + PRODUCER_STOP_TIMEOUT; + while !producer.is_finished() && Instant::now() < deadline { + std::thread::sleep(Duration::from_millis(1)); + } +} + /// Accept serialized query plan and the addresses of Arrow Arrays from Spark, /// then execute the query. Return addresses of arrow vector. /// # Safety @@ -1215,7 +1292,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( // decreasing to 1 would serialize production and consumption. let (tx, rx) = mpsc::channel(2); let mut stream = stream; - get_runtime().spawn(async move { + let producer = get_runtime().spawn(async move { let result = std::panic::AssertUnwindSafe(async { while let Some(batch) = stream.next().await { if tx.send(batch).await.is_err() { @@ -1241,6 +1318,9 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( .await; } }); + let abort = producer.abort_handle(); + exec_context.cancellation.on_cancel(move || abort.abort()); + exec_context.producer = Some(producer); exec_context.batch_receiver = Some(rx); } else { exec_context.stream = Some(stream); @@ -1252,7 +1332,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( } if let Some(rx) = &mut exec_context.batch_receiver { - match rx.blocking_recv() { + match recv_unless_cancelled(rx, &exec_context.cancellation)? { Some(Ok(batch)) => { update_metrics(env, exec_context)?; return prepare_output( @@ -1278,7 +1358,8 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( // the stream by the end of the pull, waits on native I/O, and `next_batch` parks // until it completes. let mut stream = exec_context.stream.take().unwrap(); - let next = get_runtime().block_on(next_batch(&mut stream, || { + let cancellation = Arc::clone(&exec_context.cancellation); + let next = get_runtime().block_on(next_batch(&mut stream, &cancellation, || { pull_input_batches(exec_context)?; update_metrics_on_interval(env, exec_context) })); @@ -1337,11 +1418,16 @@ pub extern "system" fn Java_org_apache_comet_Native_releasePlan( "Comet execution context shouldn't be null!" ); + plan_cancellations().lock().remove(&exec_context); + // Reclaim ownership of the context up front so that it is always freed, even if updating // metrics below fails. Dropping it releases the memory pool and every JNI global ref the // context holds. let mut execution_context: Box = Box::from_raw(exec_context as *mut ExecutionContext); + if let Some(producer) = execution_context.producer.take() { + stop_producer(producer); + } // Unregister this context's pool and, when tracing, emit the remaining total for the // thread. Every context registers, but only a traced one writes counters, so the @@ -1454,9 +1540,8 @@ pub extern "system" fn Java_org_apache_comet_Native_getShufflePartitionOffsets( ) })?; - // `ExecutionPlan` has `Any` as a supertrait but no `as_any` method of its own, so upcast - // the trait object before downcasting to the writer. - let writer = (root_op.native_plan.as_ref() as &dyn std::any::Any) + let writer = root_op + .native_plan .downcast_ref::() .ok_or_else(|| { CometError::Internal( @@ -2601,11 +2686,15 @@ mod tests { .boxed(); let mut pulls = 0; let next = single_worker_runtime() - .block_on(without_a_lost_wake(next_batch(&mut stream, || { - // Every JVM-fed scan already holds a batch, so the pull wakes nothing. - pulls += 1; - Ok(()) - }))) + .block_on(without_a_lost_wake(next_batch( + &mut stream, + &PlanCancellation::default(), + || { + // Every JVM-fed scan already holds a batch, so the pull wakes nothing. + pulls += 1; + Ok(()) + }, + ))) .unwrap(); assert!(next.is_some()); assert!( @@ -2631,12 +2720,24 @@ mod tests { }; // Only the refill's wake gets the stream polled again. single_worker_runtime().block_on(without_a_lost_wake(async { - let first = next_batch(&mut stream, &mut pull).await.unwrap(); + let first = next_batch(&mut stream, &PlanCancellation::default(), &mut pull) + .await + .unwrap(); assert_eq!(first.unwrap().num_rows(), 3); assert_eq!(pulls.get(), 1); - assert!(next_batch(&mut stream, &mut pull).await.unwrap().is_none()); + assert!( + next_batch(&mut stream, &PlanCancellation::default(), &mut pull) + .await + .unwrap() + .is_none() + ); assert_eq!(pulls.get(), 2); - assert!(next_batch(&mut stream, &mut pull).await.unwrap().is_none()); + assert!( + next_batch(&mut stream, &PlanCancellation::default(), &mut pull) + .await + .unwrap() + .is_none() + ); assert_eq!(pulls.get(), 2); })); } @@ -2657,15 +2758,63 @@ mod tests { .boxed(); let mut pulls = 0; let next = runtime - .block_on(without_a_lost_wake(next_batch(&mut stream, || { - pulls += 1; - handle.block_on(tokio::time::sleep(Duration::from_millis(100))); - Ok(()) - }))) + .block_on(without_a_lost_wake(next_batch( + &mut stream, + &PlanCancellation::default(), + || { + pulls += 1; + handle.block_on(tokio::time::sleep(Duration::from_millis(100))); + Ok(()) + }, + ))) .unwrap(); assert!(next.is_some()); assert_eq!(pulls, 1); } + + fn cancel_after(cancellation: &Arc, delay: Duration) { + let cancellation = Arc::clone(cancellation); + std::thread::spawn(move || { + std::thread::sleep(delay); + cancellation.cancel(); + }); + } + + #[test] + fn next_batch_stops_waiting_on_native_io_once_cancelled() { + let mut stream = futures::stream::pending::>().boxed(); + let cancellation = PlanCancellation::new(); + cancel_after(&cancellation, Duration::from_millis(50)); + let start = Instant::now(); + let next = + single_worker_runtime().block_on(next_batch(&mut stream, &cancellation, || Ok(()))); + assert!(next.unwrap_err().to_string().contains("cancelled")); + assert!(start.elapsed() < Duration::from_secs(2)); + } + + #[test] + fn a_producer_busy_without_output_is_aborted_and_the_wait_ends_once_cancelled() { + let (tx, mut rx) = mpsc::channel::>(2); + let producer = get_runtime().spawn(async move { + let _tx = tx; + loop { + tokio::time::sleep(Duration::from_millis(1)).await; + } + }); + let cancellation = PlanCancellation::new(); + let abort = producer.abort_handle(); + cancellation.on_cancel(move || abort.abort()); + cancel_after(&cancellation, Duration::from_millis(50)); + let start = Instant::now(); + let next = recv_unless_cancelled(&mut rx, &cancellation); + assert!(next.unwrap_err().to_string().contains("cancelled")); + stop_producer(producer); + assert!(start.elapsed() < Duration::from_secs(2)); + assert!( + rx.blocking_recv().is_none(), + "the aborted producer still runs" + ); + } } #[cfg(test)] @@ -2894,6 +3043,7 @@ mod native_sort_spill_tests { )]); let session = prepare_datafusion_session_context( case.batch_size, + PlanCancellation::new(), Arc::clone(&pool), vec![spill_dir.path().to_string_lossy().into_owned()], u64::MAX, @@ -3164,6 +3314,7 @@ mod native_sort_spill_tests { HashMap::from([(SPARK_EXECUTOR_CORES.to_string(), executor_cores.to_string())]); let session = prepare_datafusion_session_context( batch_size, + PlanCancellation::new(), Arc::clone(&pool), vec![spill_dir.path().to_string_lossy().into_owned()], u64::MAX, @@ -3540,6 +3691,7 @@ mod native_sort_spill_tests { } let session = prepare_datafusion_session_context( 8192, + PlanCancellation::new(), Arc::clone(&pool), vec![spill_dir.path().to_string_lossy().into_owned()], u64::MAX, @@ -3686,6 +3838,7 @@ mod native_sort_spill_tests { ]); let session = prepare_datafusion_session_context( 8192, + PlanCancellation::new(), Arc::clone(&pool), vec![spill_dir.path().to_string_lossy().into_owned()], u64::MAX, diff --git a/native/core/src/execution/operators/cancellable.rs b/native/core/src/execution/operators/cancellable.rs new file mode 100644 index 00000000000..4d02211be50 --- /dev/null +++ b/native/core/src/execution/operators/cancellable.rs @@ -0,0 +1,366 @@ +// 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. + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion::common::config::ConfigOptions; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::{Result, Statistics}; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::{Distribution, OrderingRequirements, PhysicalExpr}; +use datafusion::physical_plan::execution_plan::CardinalityEffect; +use datafusion::physical_plan::metrics::MetricsSet; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, InputDistributionRequirements, PlanProperties, + RecordBatchStream, ReplaceChildrenOptions, SendableRecordBatchStream, +}; +use datafusion_comet_common::cancellation::{cancelled_error, PlanCancellation}; +use futures::{Stream, StreamExt}; +use std::any::Any; +use std::fmt; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +/// Makes the stream of the operator it wraps end with an error once the plan is cancelled. It is +/// otherwise invisible: it reports the wrapped operator's name, children, properties, metrics and +/// downcast identity, so a plan reads and rewrites as if it were not there. +/// +/// Every operator a Comet plan is built from is wrapped, so a cancelled plan stops at the next +/// batch any of its operators produces, even inside a parent that keeps polling its input without +/// ever yielding, such as an aggregate fed by a join that multiplies its probe rows. +#[derive(Debug)] +pub(crate) struct CancellableExec { + inner: Arc, +} + +impl CancellableExec { + pub(crate) fn wrap(inner: Arc) -> Arc { + if is_cancellable(inner.as_ref()) { + inner + } else { + Arc::new(Self { inner }) + } + } +} + +fn is_cancellable(plan: &dyn ExecutionPlan) -> bool { + (plan as &dyn Any).is::() +} + +impl DisplayAs for CancellableExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + self.inner.fmt_as(t, f) + } +} + +#[allow(deprecated)] +impl ExecutionPlan for CancellableExec { + fn name(&self) -> &str { + self.inner.name() + } + + fn downcast_delegate(&self) -> Option<&dyn ExecutionPlan> { + Some(self.inner.as_ref()) + } + + fn schema(&self) -> SchemaRef { + self.inner.schema() + } + + fn properties(&self) -> &Arc { + self.inner.properties() + } + + fn dynamic_expressions_produced(&self) -> Vec> { + self.inner.dynamic_expressions_produced() + } + + fn required_input_distribution(&self) -> Vec { + self.inner.required_input_distribution() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + self.inner.input_distribution_requirements() + } + + fn required_input_ordering(&self) -> Vec> { + self.inner.required_input_ordering() + } + + fn maintains_input_order(&self) -> Vec { + self.inner.maintains_input_order() + } + + fn benefits_from_input_partitioning(&self) -> Vec { + self.inner.benefits_from_input_partitioning() + } + + fn children(&self) -> Vec<&Arc> { + self.inner.children() + } + + fn replace_children( + self: Arc, + children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + Ok(Self::wrap( + Arc::clone(&self.inner).replace_children(children, options)?, + )) + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + self.inner.apply_expressions(f) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + Ok(Self::wrap( + Arc::clone(&self.inner).with_new_children(children)?, + )) + } + + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, + ) -> Result> { + Ok(Self::wrap( + Arc::clone(&self.inner).with_new_children_and_same_properties(children)?, + )) + } + + fn reset_state(self: Arc) -> Result> { + Ok(Self::wrap(Arc::clone(&self.inner).reset_state()?)) + } + + fn repartitioned( + &self, + target_partitions: usize, + config: &ConfigOptions, + ) -> Result>> { + Ok(self + .inner + .repartitioned(target_partitions, config)? + .map(Self::wrap)) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + let cancellation = PlanCancellation::of(&context); + let stream = self.inner.execute(partition, context)?; + Ok(match cancellation { + Some(cancellation) => Box::pin(CancellableStream { + inner: stream, + cancellation, + }), + None => stream, + }) + } + + fn metrics(&self) -> Option { + self.inner.metrics() + } + + fn partition_statistics(&self, partition: Option) -> Result> { + self.inner.partition_statistics(partition) + } + + fn supports_limit_pushdown(&self) -> bool { + self.inner.supports_limit_pushdown() + } + + fn with_fetch(&self, limit: Option) -> Option> { + self.inner.with_fetch(limit).map(Self::wrap) + } + + fn fetch(&self) -> Option { + self.inner.fetch() + } + + fn cardinality_effect(&self) -> CardinalityEffect { + self.inner.cardinality_effect() + } + + fn with_new_state(&self, state: Arc) -> Option> { + self.inner.with_new_state(state).map(Self::wrap) + } + + fn with_preserve_order(&self, preserve_order: bool) -> Option> { + self.inner + .with_preserve_order(preserve_order) + .map(Self::wrap) + } +} + +struct CancellableStream { + inner: SendableRecordBatchStream, + cancellation: Arc, +} + +impl Stream for CancellableStream { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + if self.cancellation.is_cancelled() { + return Poll::Ready(Some(Err(cancelled_error()))); + } + self.inner.poll_next_unpin(cx) + } +} + +impl RecordBatchStream for CancellableStream { + fn schema(&self) -> SchemaRef { + self.inner.schema() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::Int32Array; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion::common::{JoinType, NullEquality}; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::functions_aggregate::count::count_udaf; + use datafusion::physical_expr::aggregate::AggregateExprBuilder; + use datafusion::physical_expr::expressions::lit; + use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}; + use datafusion::physical_plan::filter::FilterExec; + use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode}; + use datafusion::physical_plan::{collect, expressions::col}; + use datafusion::physical_plan::{displayable, ChildrenPropertiesMode}; + use datafusion::prelude::{SessionConfig, SessionContext}; + use std::time::{Duration, Instant}; + + fn constant_keys(rows: usize, batches: usize) -> Arc { + let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![1; rows]))], + ) + .unwrap(); + MemorySourceConfig::try_new_exec(&[vec![batch; batches]], schema, None).unwrap() + } + + fn wrapped_join_count(rows: usize) -> Arc { + let build = CancellableExec::wrap(constant_keys(rows, 1)); + let probe = CancellableExec::wrap(constant_keys(rows, 1)); + let on = vec![( + col("k", &build.schema()).unwrap(), + col("k", &probe.schema()).unwrap(), + )]; + let join = CancellableExec::wrap(Arc::new( + HashJoinExec::try_new( + build, + probe, + on, + None, + &JoinType::Inner, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + false, + ) + .unwrap(), + )); + let schema = join.schema(); + let count = AggregateExprBuilder::new(count_udaf(), vec![lit(1)]) + .schema(Arc::clone(&schema)) + .alias("count") + .build() + .unwrap(); + CancellableExec::wrap(Arc::new( + AggregateExec::try_new( + AggregateMode::Single, + PhysicalGroupBy::new_single(vec![]), + vec![Arc::new(count)], + vec![None], + join, + schema, + ) + .unwrap(), + )) + } + + #[test] + fn a_wrapped_plan_displays_and_downcasts_as_the_plan_it_wraps() { + let source = constant_keys(3, 1); + let filter: Arc = + Arc::new(FilterExec::try_new(lit(true), Arc::clone(&source)).unwrap()); + let wrapped = CancellableExec::wrap(Arc::clone(&filter)); + assert!(wrapped.downcast_ref::().is_some()); + assert_eq!(wrapped.name(), filter.name()); + assert_eq!(wrapped.children().len(), 1); + assert_eq!( + displayable(wrapped.as_ref()).indent(true).to_string(), + displayable(filter.as_ref()).indent(true).to_string() + ); + let rewritten = Arc::clone(&wrapped) + .replace_children( + vec![source], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + .unwrap(); + assert!(is_cancellable(rewritten.as_ref())); + assert!(rewritten.downcast_ref::().is_some()); + let rewrapped = CancellableExec::wrap(Arc::clone(&wrapped)); + assert!(Arc::ptr_eq(&rewrapped, &wrapped)); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn a_cancelled_plan_stops_inside_an_operator_that_never_yields() { + let cancellation = PlanCancellation::new(); + let config = SessionConfig::new().with_extension(Arc::clone(&cancellation)); + let ctx = SessionContext::new_with_config(config).task_ctx(); + let plan = wrapped_join_count(40_000); + let canceller = Arc::clone(&cancellation); + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(300)).await; + canceller.cancel(); + }); + let start = Instant::now(); + let err = collect(plan, ctx).await.unwrap_err(); + assert!(err.to_string().contains("cancelled"), "{err}"); + assert!( + start.elapsed() < Duration::from_secs(2), + "took {:?}", + start.elapsed() + ); + } + + #[tokio::test] + async fn a_plan_without_cancellation_runs_unchanged() { + let ctx = SessionContext::new().task_ctx(); + let batches = collect(wrapped_join_count(100), ctx).await.unwrap(); + let count = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + assert_eq!(count, 10_000); + } +} diff --git a/native/core/src/execution/operators/mod.rs b/native/core/src/execution/operators/mod.rs index 78a611a9212..76826d8e2a7 100644 --- a/native/core/src/execution/operators/mod.rs +++ b/native/core/src/execution/operators/mod.rs @@ -25,6 +25,8 @@ pub use iceberg_scan::*; pub use scan::*; mod aligned_stream_reader; +mod cancellable; +pub(crate) use cancellable::CancellableExec; mod copy; mod dynamic_filter; pub(crate) use dynamic_filter::DynamicFilterJoinExec; diff --git a/native/core/src/execution/spark_plan.rs b/native/core/src/execution/spark_plan.rs index 2487c76a140..2e72257642e 100644 --- a/native/core/src/execution/spark_plan.rs +++ b/native/core/src/execution/spark_plan.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use crate::execution::operators::CancellableExec; use arrow::datatypes::SchemaRef; use datafusion::physical_plan::ExecutionPlan; use std::sync::Arc; @@ -44,7 +45,7 @@ impl SparkPlan { ) -> Self { Self { plan_id, - native_plan, + native_plan: CancellableExec::wrap(native_plan), children, additional_native_plans: vec![], } @@ -63,7 +64,7 @@ impl SparkPlan { } Self { plan_id, - native_plan, + native_plan: CancellableExec::wrap(native_plan), children, additional_native_plans: accum, } diff --git a/native/shuffle/src/partitioners/multi_partition.rs b/native/shuffle/src/partitioners/multi_partition.rs index 05b14b53243..8f3501ace98 100644 --- a/native/shuffle/src/partitioners/multi_partition.rs +++ b/native/shuffle/src/partitioners/multi_partition.rs @@ -28,6 +28,7 @@ use datafusion::common::utils::proxy::VecAllocExt; use datafusion::common::{DataFusionError, HashSet}; use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation}; use datafusion::execution::runtime_env::RuntimeEnv; +use datafusion_comet_common::cancellation::PlanCancellation; use datafusion_comet_common::tracing::{with_trace, with_trace_async}; use datafusion_comet_spark_expr::murmur3::create_murmur3_hashes; use itertools::Itertools; @@ -128,6 +129,7 @@ pub(crate) struct MultiPartitionShuffleRepartitioner { /// Rows this task has placed so far, which is the ordinal positional round robin keys on. /// A `u64` because it counts a whole task's input, not one batch. row_seq: u64, + cancellation: Option>, } /// Sum of the capacities of the backing buffers reachable from `batch` whose start address is @@ -253,9 +255,15 @@ impl MultiPartitionShuffleRepartitioner { tracing_enabled, pinned_buffers: HashSet::new(), row_seq: 0, + cancellation: None, }) } + pub(crate) fn with_cancellation(mut self, cancellation: Option>) -> Self { + self.cancellation = cancellation; + self + } + /// Shuffles rows in input batch into corresponding partition buffer. /// This function first calculates hashes for rows and then takes rows in same /// partition as a record batch which is appended into partition buffer. @@ -614,6 +622,7 @@ impl MultiPartitionShuffleRepartitioner { let empty = self.partition_indices.empty_like(); let indices = std::mem::replace(&mut self.partition_indices, empty); PartitionedBatchesProducer::new(buffered_batches, indices, self.batch_size) + .with_cancellation(self.cancellation.clone()) } pub(crate) fn spill(&mut self, unreserved_bytes: usize) -> datafusion::common::Result<()> { diff --git a/native/shuffle/src/partitioners/partitioned_batch_iterator.rs b/native/shuffle/src/partitioners/partitioned_batch_iterator.rs index 6b2a895ca52..5249f9c641b 100644 --- a/native/shuffle/src/partitioners/partitioned_batch_iterator.rs +++ b/native/shuffle/src/partitioners/partitioned_batch_iterator.rs @@ -21,6 +21,8 @@ use arrow::compute::{concat_batches, interleave_record_batch}; use datafusion::common::utils::proxy::VecAllocExt; use datafusion::common::DataFusionError; use datafusion::physical_plan::metrics::Time; +use datafusion_comet_common::cancellation::{cancelled_error, PlanCancellation}; +use std::sync::Arc; /// A contiguous run of rows within one buffered batch, bound for one output partition. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -91,6 +93,7 @@ pub(super) struct PartitionedBatchesProducer { buffered_batches: Vec, partition_indices: PartitionIndices, batch_size: usize, + cancellation: Option>, } impl PartitionedBatchesProducer { @@ -103,9 +106,17 @@ impl PartitionedBatchesProducer { partition_indices: indices, buffered_batches, batch_size, + cancellation: None, } } + /// Makes every partition's batches end with an error once the plan is cancelled, so that + /// a spill or the final write of a killed task stops at its next batch. + pub(super) fn with_cancellation(mut self, cancellation: Option>) -> Self { + self.cancellation = cancellation; + self + } + /// References to all buffered batches. Build this once per write cycle and share it /// across every partition's [`Self::produce`] call instead of rebuilding a fresh /// `Vec<&RecordBatch>` over all buffered batches for each partition. @@ -129,19 +140,23 @@ impl PartitionedBatchesProducer { self.buffered_batches.len(), "refs slice must cover every buffered batch" ); - match &self.partition_indices { - PartitionIndices::Rows(indices) => PartitionedBatchIterator::Rows(RowIterator::new( + let batches = match &self.partition_indices { + PartitionIndices::Rows(indices) => PartitionBatches::Rows(RowIterator::new( &indices[partition_id], refs, self.batch_size, interleave_time, )), - PartitionIndices::Runs(runs) => PartitionedBatchIterator::Runs(RunIterator::new( + PartitionIndices::Runs(runs) => PartitionBatches::Runs(RunIterator::new( &runs[partition_id], refs, self.batch_size, interleave_time, )), + }; + PartitionedBatchIterator { + batches, + cancellation: self.cancellation.as_deref(), } } } @@ -150,7 +165,12 @@ impl PartitionedBatchesProducer { /// /// One concrete type covering both index shapes, because [`crate::writers::PartitionWriter`] is /// generic over a single iterator type rather than taking a trait object. -pub(crate) enum PartitionedBatchIterator<'a> { +pub(crate) struct PartitionedBatchIterator<'a> { + batches: PartitionBatches<'a>, + cancellation: Option<&'a PlanCancellation>, +} + +enum PartitionBatches<'a> { Rows(RowIterator<'a>), Runs(RunIterator<'a>), } @@ -159,9 +179,15 @@ impl Iterator for PartitionedBatchIterator<'_> { type Item = datafusion::common::Result; fn next(&mut self) -> Option { - match self { - Self::Rows(iter) => iter.next(), - Self::Runs(iter) => iter.next(), + if self + .cancellation + .is_some_and(PlanCancellation::is_cancelled) + { + return Some(Err(cancelled_error())); + } + match &mut self.batches { + PartitionBatches::Rows(iter) => iter.next(), + PartitionBatches::Runs(iter) => iter.next(), } } } diff --git a/native/shuffle/src/shuffle_writer.rs b/native/shuffle/src/shuffle_writer.rs index 4d93e87a3f5..8fd49e205f7 100644 --- a/native/shuffle/src/shuffle_writer.rs +++ b/native/shuffle/src/shuffle_writer.rs @@ -41,6 +41,7 @@ use datafusion::{ DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream, }, }; +use datafusion_comet_common::cancellation::PlanCancellation; use datafusion_comet_jni_bridge::ShufflePartitionPusher; use futures::{StreamExt, TryStreamExt}; use std::{ @@ -475,16 +476,19 @@ fn create_repartitioner( writer, metrics, ))) } else { - Ok(Box::new(MultiPartitionShuffleRepartitioner::try_new( - partition, - writer, - partitioning, - metrics, - context.runtime_env(), - context.session_config().batch_size(), - tracing_enabled, - max_buffer_bytes, - )?)) + Ok(Box::new( + MultiPartitionShuffleRepartitioner::try_new( + partition, + writer, + partitioning, + metrics, + context.runtime_env(), + context.session_config().batch_size(), + tracing_enabled, + max_buffer_bytes, + )? + .with_cancellation(PlanCancellation::of(context)), + )) } } diff --git a/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java b/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java index d2881bfe449..c9f34ead075 100644 --- a/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java +++ b/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java @@ -128,6 +128,9 @@ public static void evaluate( // ground truth for this call. Any value already on the thread is either (a) the same object // on a Spark task thread, or (b) stale from a prior task on a reused Tokio worker. The same // reasoning applies to the propagated `classLoader`. + if (taskContext != null) { + taskContext.killTaskIfInterrupted(); + } TaskContext prior = TaskContext.get(); if (taskContext != null) { CometTaskContextShim.set(taskContext); diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index f29e5086abd..f36636b85b3 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -21,7 +21,7 @@ package org.apache.comet import java.lang.management.ManagementFactory import java.util.Locale -import java.util.concurrent.{Executors, ThreadFactory, TimeUnit} +import java.util.concurrent.{Executors, ScheduledFuture, ScheduledThreadPoolExecutor, ThreadFactory, TimeUnit} import java.util.concurrent.atomic.AtomicBoolean import scala.util.control.NonFatal @@ -30,6 +30,7 @@ import org.apache.arrow.c.ArrowArrayStream import org.apache.hadoop.conf.Configuration import org.apache.spark._ import org.apache.spark.broadcast.Broadcast +import org.apache.spark.comet.CometTaskContextShim import org.apache.spark.internal.Logging import org.apache.spark.network.util.ByteUnit import org.apache.spark.sql.comet.CometMetricNode @@ -223,12 +224,25 @@ class CometExecIterator( this.close() } + /** Guards [[planReleased]], so that the plan is never cancelled after it is released. */ + private val planLock = new Object + private var planReleased = false + + private val killWatch = CometExecIterator.watchForKill(TaskContext.get(), () => cancelPlan()) + + private def cancelPlan(): Unit = planLock.synchronized { + if (!planReleased) { + nativeLib.cancelPlan(plan) + } + } + CometExecIterator.startMemoryUsageLog() private def getNextBatch: Option[ColumnarBatch] = { assert(partitionIndex >= 0 && partitionIndex < numParts) val ctx = TaskContext.get() + CometTaskContextShim.killTaskIfInterrupted(ctx) try { val result = withTrace( @@ -247,6 +261,14 @@ class CometExecIterator( result } catch { + // A plan cancelled because its task was killed fails with whatever stopped it. Report the + // kill, so that Spark records the task as killed rather than failed. + case e: Throwable if ctx.isInterrupted() => + val killed = new TaskKilledException( + CometTaskContextShim.killReason(ctx).getOrElse("unknown reason")) + if (!e.isInstanceOf[TaskKilledException]) killed.addSuppressed(e) + throw killed + // Handle CometQueryExecutionException with JSON payload first case e: CometQueryExecutionException => logError(s"Native execution for task $taskAttemptId failed", e) @@ -334,6 +356,8 @@ class CometExecIterator( } attempt(nativeUtil.close()) shuffleBlockIterators.values.foreach(it => attempt(it.close())) + attempt(killWatch.cancel(false)) + attempt(planLock.synchronized { planReleased = true }) // Released last and exactly once, even if the teardown above failed: dropping the native // execution context frees this plan's task-shared memory pool reference and several JNI @@ -369,6 +393,52 @@ class CometExecIterator( object CometExecIterator extends Logging { + /** How often a running native plan checks whether its task has been killed. */ + private[comet] val KillCheckIntervalMs = 50L + + private lazy val killWatcher = { + val executor = new ScheduledThreadPoolExecutor( + 1, + new ThreadFactory { + override def newThread(runnable: Runnable): Thread = { + val thread = new Thread(runnable, "comet-native-kill-watcher") + thread.setDaemon(true) + thread + } + }) + executor.setRemoveOnCancelPolicy(true) + executor + } + + /** + * Cancels a task's native plan once the task is killed. Spark only marks a killed task and, if + * asked to, interrupts its thread; neither reaches a thread busy in native code, or parked + * there waiting on native work, so without this a killed task runs its plan to completion and + * the task reaper can kill the executor first. + */ + private def watchForKill(context: TaskContext, cancel: () => Unit): ScheduledFuture[_] = { + val cancelled = new AtomicBoolean(false) + killWatcher.scheduleWithFixedDelay( + new Runnable { + override def run(): Unit = { + if (context.isInterrupted() && cancelled.compareAndSet(false, true)) { + try { + cancel() + } catch { + case NonFatal(e) => + logWarning( + s"Failed to cancel the native plan of killed task " + + s"${context.taskAttemptId()}", + e) + } + } + } + }, + KillCheckIntervalMs, + KillCheckIntervalMs, + TimeUnit.MILLISECONDS) + } + private val memoryUsageLogStarted = new AtomicBoolean(false) /** Native plans running at the previous memory usage log. Only the log's own thread uses it. */ diff --git a/spark/src/main/scala/org/apache/comet/Native.scala b/spark/src/main/scala/org/apache/comet/Native.scala index d6e5bcf5bd7..c8076968521 100644 --- a/spark/src/main/scala/org/apache/comet/Native.scala +++ b/spark/src/main/scala/org/apache/comet/Native.scala @@ -138,6 +138,16 @@ class Native extends NativeBase { */ @native def releasePlan(plan: Long): Unit + /** + * Cancel a native query plan whose task has been killed: an `executePlan` in progress, or the + * next one, fails at its next check instead of running the plan to completion. Safe to call + * from any thread until the plan is released. + * + * @param plan + * the address to native query plan. + */ + @native def cancelPlan(plan: Long): Unit + /** * Used by Comet shuffle external sorter to write sorted records to disk. * diff --git a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala index 232ff9e135e..7fa7318c762 100644 --- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala +++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala @@ -327,7 +327,11 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come | // `splitExpressions` also take `InternalRow row` as a parameter; `this` flows | // implicitly via INPUT_ROW. | org.apache.spark.sql.catalyst.InternalRow row = this; + | org.apache.spark.TaskContext killCheckContext = org.apache.spark.TaskContext.get(); | for (int i = 0; i < numRows; i++) { + | if ((i & 255) == 0 && killCheckContext != null) { + | killCheckContext.killTaskIfInterrupted(); + | } | this.rowIdx = i; | $perRowBody | } diff --git a/spark/src/main/scala/org/apache/spark/comet/CometTaskContextShim.scala b/spark/src/main/scala/org/apache/spark/comet/CometTaskContextShim.scala index 9218fc5e78b..c4ce9dc9bf2 100644 --- a/spark/src/main/scala/org/apache/spark/comet/CometTaskContextShim.scala +++ b/spark/src/main/scala/org/apache/spark/comet/CometTaskContextShim.scala @@ -38,4 +38,9 @@ object CometTaskContextShim { def set(taskContext: TaskContext): Unit = TaskContext.setTaskContext(taskContext) def unset(): Unit = TaskContext.unset() + + /** Throws `TaskKilledException` if `taskContext` has been killed. */ + def killTaskIfInterrupted(taskContext: TaskContext): Unit = taskContext.killTaskIfInterrupted() + + def killReason(taskContext: TaskContext): Option[String] = taskContext.getKillReason() } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/RowArrowReader.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/RowArrowReader.scala index e4f484aa1e2..f0fd313252a 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/RowArrowReader.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/RowArrowReader.scala @@ -22,6 +22,7 @@ package org.apache.spark.sql.comet.execution.arrow import org.apache.arrow.memory.BufferAllocator import org.apache.arrow.vector.ipc.ArrowReader import org.apache.arrow.vector.types.pojo.Schema +import org.apache.spark.TaskContext import org.apache.spark.sql.catalyst.InternalRow /** @@ -43,6 +44,8 @@ private[comet] class RowArrowReader( require(maxRecordsPerBatch > 0, "Maximum records per batch must be positive") + private val taskContext = TaskContext.get() + override protected def readSchema(): Schema = arrowSchema override def bytesRead(): Long = 0L @@ -60,6 +63,9 @@ private[comet] class RowArrowReader( val writer = ArrowWriter.create(getVectorSchemaRoot, maxRecordsPerBatch) var rowCount = 0 while (rowIter.hasNext && rowCount < maxRecordsPerBatch) { + if ((rowCount & 255) == 0 && taskContext != null) { + taskContext.killTaskIfInterrupted() + } writer.write(rowIter.next()) rowCount += 1 } diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/CometNativeTaskKillSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/CometNativeTaskKillSuite.scala new file mode 100644 index 00000000000..3e3b5e7bdfd --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/CometNativeTaskKillSuite.scala @@ -0,0 +1,217 @@ +/* + * 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.UUID +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.atomic.AtomicReference + +import scala.collection.JavaConverters._ + +import org.apache.spark.{SparkConf, TaskKilled} +import org.apache.spark.rdd.RDD +import org.apache.spark.scheduler.{SparkListener, SparkListenerJobStart, SparkListenerTaskEnd, SparkListenerTaskStart} +import org.apache.spark.sql.{CometTestBase, DataFrame} +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.functions.{broadcast, col, sum, udf} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{LongType, StructField, StructType} + +import org.apache.comet.CometConf + +/** + * A killed task must stop well within `spark.task.reaper.killTimeout` while Comet native code is + * running its plan. Spark only marks the task killed, and interrupting the thread (not requested + * here, as when a failed job cancels its other stages) does not reach native code either. + */ +class CometNativeTaskKillSuite extends CometTestBase { + import CometNativeTaskKillSuite.busyWaitNanos + + private val killBoundMs = 2000L + + override protected def sparkConf: SparkConf = + super.sparkConf + .set("spark.task.reaper.enabled", "true") + .set("spark.task.reaper.killTimeout", "10s") + + private def nativeQueryConf(extra: (String, String)*)(f: => Unit): Unit = + withSQLConf( + extra ++ Seq( + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_MODE.key -> "native", + CometConf.COMET_EXEC_EXPLODE_ENABLED.key -> "true"): _*)(f) + + private def assertRunsNatively(df: DataFrame, operators: Class[_]*): Unit = { + val plan = df.queryExecution.executedPlan + operators.foreach { operator => + assert( + plan.collectFirst { case p if operator.isInstance(p) => p }.isDefined || + plan.collectWithSubqueries { case p if operator.isInstance(p) => p }.nonEmpty, + s"expected ${operator.getSimpleName} in\n$plan") + } + } + + /** + * Runs `df` without interrupting on cancel, cancels it once its tasks have run + * `runBeforeKillMs`, and checks that every task running then ends as killed within + * `killBoundMs`. + */ + private def assertKilledPromptly(df: DataFrame, runBeforeKillMs: Long = 1500L): Unit = { + val sc = spark.sparkContext + val group = s"comet-kill-${UUID.randomUUID()}" + val stages = ConcurrentHashMap.newKeySet[Int]() + val started = ConcurrentHashMap.newKeySet[Long]() + val ended = new ConcurrentHashMap[Long, (Long, AnyRef)]() + val listener = new SparkListener { + override def onJobStart(event: SparkListenerJobStart): Unit = { + if (Option(event.properties).exists(_.getProperty("spark.jobGroup.id") == group)) { + event.stageIds.foreach(stages.add) + } + } + override def onTaskStart(event: SparkListenerTaskStart): Unit = { + if (stages.contains(event.stageId)) started.add(event.taskInfo.taskId) + } + override def onTaskEnd(event: SparkListenerTaskEnd): Unit = { + if (stages.contains(event.stageId)) { + ended.put(event.taskInfo.taskId, (event.taskInfo.finishTime, event.reason)) + } + } + } + sc.addSparkListener(listener) + val failure = new AtomicReference[Throwable]() + val runner = new Thread(() => { + sc.setJobGroup(group, "kill test", interruptOnCancel = false) + try { + df.collect() + } catch { + case t: Throwable => failure.set(t) + } finally { + sc.clearJobGroup() + } + }) + runner.setDaemon(true) + try { + runner.start() + val startDeadline = System.currentTimeMillis() + 60000 + while (started.isEmpty && System.currentTimeMillis() < startDeadline) Thread.sleep(10) + assert(!started.isEmpty, "no task started") + Thread.sleep(runBeforeKillMs) + val running = started.asScala.toSet -- ended.keySet().asScala + assert(running.nonEmpty, "the query finished before it could be killed") + + val killedAt = System.currentTimeMillis() + sc.cancelJobGroup(group) + val endDeadline = killedAt + 30000 + while (running.exists(id => !ended.containsKey(id)) && + System.currentTimeMillis() < endDeadline) { + Thread.sleep(10) + } + val stillRunning = running.filterNot(ended.containsKey) + assert(stillRunning.isEmpty, s"tasks $stillRunning were still running 30 s after the kill") + val stopMs = running.map(id => ended.get(id)._1 - killedAt).max + logInfo(s"The killed tasks stopped within $stopMs ms") + assert(stopMs < killBoundMs, s"the killed tasks took $stopMs ms to stop") + running.foreach { id => + assert( + ended.get(id)._2.isInstanceOf[TaskKilled], + s"task $id ended with ${ended.get(id)._2}") + } + runner.join(30000) + assert(failure.get() != null, "the cancelled query did not fail") + } finally { + sc.removeSparkListener(listener) + } + assert(spark.range(10).count() == 10) + } + + test("a task killed while a native producer task runs its shuffle write stops promptly") { + nativeQueryConf(CometConf.COMET_BATCH_SIZE.key -> "16") { + withTempPath { path => + spark.range(0, 4000, 1, 2).write.parquet(path.getAbsolutePath) + withParquetTable(path.getAbsolutePath, "t") { + val df = sql( + "SELECT k % 7 AS g, count(*) FROM " + + "(SELECT explode(array_repeat(id, 1000000)) AS k FROM t) GROUP BY k % 7") + assertRunsNatively(df, classOf[CometNativeScanExec], classOf[CometExplodeExec]) + assertKilledPromptly(df) + } + } + } + } + + test("a task killed inside a native operator that never waits on its input stops promptly") { + nativeQueryConf() { + val build = spark.range(0, 400000, 1, 1).selectExpr("id * 0 AS k", "id AS b") + val df = spark + .range(0, 16384, 1, 2) + .selectExpr("id * 0 AS k", "id AS p") + .join(broadcast(build), "k") + .agg(sum(col("p") + col("b"))) + assertRunsNatively( + df, + classOf[CometSparkToColumnarExec], + classOf[CometBroadcastHashJoinExec]) + assertKilledPromptly(df) + } + } + + test("a task killed while native code waits on a slow JVM input stops promptly") { + nativeQueryConf() { + val schema = StructType(Seq(StructField("id", LongType))) + val rows: RDD[InternalRow] = spark.sparkContext + .parallelize(0 until 2, 2) + .mapPartitions { _ => + Iterator.range(0, 100000).map { i => + busyWaitNanos(500000L) + InternalRow(i.toLong) + } + } + val input = spark.internalCreateDataFrame(rows, schema) + val df = input.groupBy(col("id") % 7).count() + assertRunsNatively(df, classOf[CometSparkToColumnarExec]) + assertKilledPromptly(df) + } + } + + test("a task killed while native code evaluates a slow JVM UDF stops promptly") { + nativeQueryConf() { + val slow = udf { (id: java.lang.Long) => + busyWaitNanos(2000000L) + id + } + withTempPath { path => + spark.range(0, 20000, 1, 2).write.parquet(path.getAbsolutePath) + withParquetTable(path.getAbsolutePath, "t") { + val df = spark.table("t").select(slow(col("id")).as("v")).agg(sum("v")) + assertRunsNatively(df, classOf[CometNativeScanExec], classOf[CometProjectExec]) + assertKilledPromptly(df) + } + } + } + } +} + +object CometNativeTaskKillSuite { + def busyWaitNanos(nanos: Long): Unit = { + val end = System.nanoTime() + nanos + while (System.nanoTime() < end) {} + } +} From a7170c28787d539b2b1d1d3a18c33cfa6ac0e8bb Mon Sep 17 00:00:00 2001 From: msaf Date: Mon, 5 Oct 2026 15:18:09 +0100 Subject: [PATCH 08/14] fix: keep a native shuffle's plan out of the task binaries of the stages that read it The native shuffle writer's spec (child plan, metric tree, input RDDs and scan planning data) lived on CometShuffleDependency.nativeShuffleSpec. The dependency is serialized with its own map stage, but also with every stage that reads the shuffle, through CometShuffledBatchRDD.dependency. The spec's input RDDs pull in their upstream shuffle dependencies and specs in turn, so each reduce stage's task binary carried the plan and lineage of every native map stage upstream of it and grew with the depth of the query. Carry the spec on the map-side CometNativeShuffleInputRDD, which only the map stage serializes, and hand it to the writer through the input iterator. The dependency keeps it only on the driver. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../shuffle/CometCelebornShuffleManager.scala | 4 +- .../shuffle/CometNativeShuffleInputRDD.scala | 15 ++++-- .../shuffle/CometNativeShuffleWriter.scala | 11 +++- .../shuffle/CometShuffleDependency.scala | 2 +- .../shuffle/CometShuffleExchangeExec.scala | 7 +-- .../shuffle/CometShuffleManager.scala | 2 +- .../apache/spark/sql/comet/operators.scala | 2 +- ...ometCelebornNativeShuffleWriterSuite.scala | 9 ++-- .../CometCelebornShuffleReaderSuite.scala | 2 +- .../CometNativeShuffleInputRDDSuite.scala | 50 +++++++++++++++++++ 10 files changed, 84 insertions(+), 20 deletions(-) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala index c067bc6e49b..44cea79c896 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleManager.scala @@ -169,9 +169,7 @@ class CometCelebornShuffleManager private[shuffle] ( .put(resolved.celebornShuffleId, resolved.client) new CometNativeShuffleWriter[K, V]( - dependency.nativeShuffleSpec.getOrElse { - throw new IllegalStateException("Native Comet shuffle has no execution plan") - }, + Option(dependency.nativeShuffleSpec).flatten.orNull, dependency.outputPartitioning.getOrElse { throw new IllegalStateException("Native Comet shuffle has no output partitioning") }, diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala index 8770496bff2..5b0d17794fb 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala @@ -50,13 +50,15 @@ private[shuffle] class CometNativeShuffleInputRDD( sc, inputRDDs.map(rdd => new OneToOneDependency(rdd))) { + @volatile private[shuffle] var nativeShuffleSpec: NativeShuffleSpec = _ + /** * Give local fallback its own scheduling RDD over the original upstream inputs. Spark aborts * jobs whose RDD ancestry contains the failed stage's RDD, so reusing this instance or wrapping * it in a narrow dependency lets a late remote failure abort the local replacement as well. */ - private[shuffle] def copyForLocalShuffle(): CometNativeShuffleInputRDD = - new CometNativeShuffleInputRDD( + private[shuffle] def copyForLocalShuffle(): CometNativeShuffleInputRDD = { + val copy = new CometNativeShuffleInputRDD( context, inputRDDs, numPartitionsParam, @@ -64,6 +66,9 @@ private[shuffle] class CometNativeShuffleInputRDD( spillMetricNode, perPartitionByKey, positionalRoundRobin) + copy.nativeShuffleSpec = nativeShuffleSpec + copy + } /** * Spark's `isOrderSensitive` rule, applied to the RDD graph below the native plan. Spark only @@ -122,7 +127,8 @@ private[shuffle] class CometNativeShuffleInputRDD( partition.index, inputObjects, shuffleBlockIters, - partition.planDataByKey) + partition.planDataByKey, + nativeShuffleSpec) } override def getPreferredLocations(split: Partition): Seq[String] = { @@ -159,7 +165,8 @@ private[shuffle] class CometNativeShuffleInputIterator( val partitionIndex: Int, val inputObjects: Array[Object], val shuffleBlockIterators: Map[Int, CometShuffleBlockIterator], - val planDataByKey: Map[String, Array[Byte]]) + val planDataByKey: Map[String, Array[Byte]], + val nativeShuffleSpec: NativeShuffleSpec = null) extends Iterator[Product2[Int, ColumnarBatch]] { override def hasNext: Boolean = false diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala index c97dc62744b..c74c9c30a6e 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala @@ -68,6 +68,8 @@ class CometNativeShuffleWriter[K, V]( extends ShuffleWriter[K, V] with Logging { + private var activeSpec: NativeShuffleSpec = spec + var partitionLengths: Array[Long] = _ var mapStatus: MapStatus = _ private var stopped = false @@ -129,6 +131,11 @@ class CometNativeShuffleWriter[K, V]( "CometNativeShuffleInputIterator (produced by CometNativeShuffleInputRDD), got " + s"${other.getClass.getName}") } + Option(shuffleInputIter.nativeShuffleSpec).foreach(activeSpec = _) + if (activeSpec == null) { + throw new IllegalStateException("Native Comet shuffle has no execution plan") + } + val spec = activeSpec val partitionIdx = shuffleInputIter.partitionIndex val inputObjects = shuffleInputIter.inputObjects val shuffleBlockIters = shuffleInputIter.shuffleBlockIterators @@ -423,7 +430,7 @@ class CometNativeShuffleWriter[K, V]( CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_MAX_HASH_COLUMNS.get()) // Decided on the driver, from the shape of the plan fused into this writer; the executor // cannot re-derive it. See `CometShuffleExchangeExec.positionalRoundRobinSpec`. - spec.positionalRoundRobin.foreach { positional => + activeSpec.positionalRoundRobin.foreach { positional => partitioning.setPositional(true) partitioning.setPositionalGroupRows(positional.groupRows) // Per task, unlike the two above: which partition this mapper's first group goes to. @@ -458,7 +465,7 @@ class CometNativeShuffleWriter[K, V]( OperatorOuterClass.Operator .newBuilder() .setShuffleWriter(shuffleWriterBuilder) - .addChildren(spec.childNativeOp) + .addChildren(activeSpec.childNativeOp) .build() } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala index be91caa0ec9..ff36ad29329 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala @@ -86,7 +86,7 @@ class CometShuffleDependency[K: ClassTag, V: ClassTag, C: ClassTag]( val shuffleWriteMetrics: Map[String, SQLMetric] = Map.empty, val numParts: Int = 0, val rangePartitionBounds: Option[Seq[InternalRow]] = None, - val nativeShuffleSpec: Option[NativeShuffleSpec] = None, + @transient val nativeShuffleSpec: Option[NativeShuffleSpec] = None, val useLocalShuffle: Boolean = false, private[shuffle] val outputMetrics: Option[CometShuffleOutputMetrics] = None) extends ShuffleDependency[K, V, C]( diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala index 2f106e43222..b57b454adc2 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala @@ -986,9 +986,9 @@ object CometShuffleExchangeExec case e: Expression => e.collect { case s: ScalarSubquery => s } case _ => Nil } - // Drop the per-partition plan-data map off the spec that lands on the (non-transient) - // CometShuffleDependency.nativeShuffleSpec. Each partition's slice now rides on the thin RDD's - // Partition objects (see CometNativeShuffleInputRDD.getPartitions), so the full + // Drop the per-partition plan-data map off the spec that lands on the map-side + // CometNativeShuffleInputRDD.nativeShuffleSpec. Each partition's slice now rides on the thin + // RDD's Partition objects (see CometNativeShuffleInputRDD.getPartitions), so the full // O(numPartitions) map is dead weight here and would blow the 2GB ByteArrayOutputStream limit // at stage submission on very-high-partition-count jobs. NativeExecContext.perPartitionByKey is // also @transient (the structural guard against any build path), but we empty it explicitly @@ -997,6 +997,7 @@ object CometShuffleExchangeExec val augmentedSpec = spec.copy(execContext = spec.execContext.copy( subqueries = spec.execContext.subqueries ++ partitioningSubqueries, perPartitionByKey = Map.empty)) + thinRDD.nativeShuffleSpec = augmentedSpec // The code block below is mostly brought over from // ShuffleExchangeExec::prepareShuffleDependency diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleManager.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleManager.scala index 82a59707777..56ee2e6dfad 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleManager.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleManager.scala @@ -232,7 +232,7 @@ class CometShuffleManager(conf: SparkConf) extends ShuffleManager with Logging { case cometShuffleHandle: CometNativeShuffleHandle[K @unchecked, V @unchecked] => val dep = cometShuffleHandle.dependency.asInstanceOf[CometShuffleDependency[_, _, _]] new CometNativeShuffleWriter( - dep.nativeShuffleSpec.get, + Option(dep.nativeShuffleSpec).flatten.orNull, dep.outputPartitioning.get, dep.outputAttributes, dep.shuffleWriteMetrics, 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 ef2d5b2db19..a692e2fc454 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 @@ -767,7 +767,7 @@ private[comet] case class NativeExecContext( // counts it is huge. It is only read on the driver (to slice per partition onto each task's // Partition object - see CometNativeShuffleInputRDD / CometExecRDD); the executor reads its own // slice, never this map. Keeping it off the wire stops it from bloating the broadcast task - // binary when this context rides on the non-transient CometShuffleDependency.nativeShuffleSpec. + // binary when this context rides on the map-side CometNativeShuffleInputRDD.nativeShuffleSpec. @transient perPartitionByKey: Map[String, Array[Array[Byte]]], shuffleScanIndices: Set[Int], hasScanInput: Boolean) { diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornNativeShuffleWriterSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornNativeShuffleWriterSuite.scala index 048b4c09a6e..3f7c619ec0b 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornNativeShuffleWriterSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornNativeShuffleWriterSuite.scala @@ -208,7 +208,8 @@ class CometCelebornNativeShuffleWriterSuite extends CometTestBase { commitValidator = () => { validations += 1 true - }) + }, + spec = inputs.asInstanceOf[CometNativeShuffleInputIterator].nativeShuffleSpec) val plan = writer.buildUnifiedPlan("").getShuffleWriter assert(plan.getPartitionWriter.hasRss) assert(plan.getOutputDataFile.isEmpty) @@ -716,10 +717,10 @@ private[shuffle] object CometCelebornNativeShuffleWriterSuite { pusher: CelebornShufflePartitionPusher, commitAuthorized: Boolean = false, commitValidator: () => Boolean = () => true, - onSizeLimitExceeded: Throwable => Unit = _ => ()) - : CometNativeShuffleWriter[Int, ColumnarBatch] = + onSizeLimitExceeded: Throwable => Unit = _ => (), + spec: NativeShuffleSpec = null): CometNativeShuffleWriter[Int, ColumnarBatch] = new CometNativeShuffleWriter[Int, ColumnarBatch]( - dependency.nativeShuffleSpec.get, + Option(spec).orElse(Option(dependency.nativeShuffleSpec).flatten).orNull, dependency.outputPartitioning.get, dependency.outputAttributes, dependency.shuffleWriteMetrics, diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala index 54448ec2039..bd47204fc53 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala @@ -244,7 +244,7 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { dependency.rdd, (context: TaskContext, inputs: Iterator[Product2[Int, ColumnarBatch]]) => { val writer = new CometNativeShuffleWriter[Int, ColumnarBatch]( - dependency.nativeShuffleSpec.get, + Option(dependency.nativeShuffleSpec).flatten.orNull, dependency.outputPartitioning.get, dependency.outputAttributes, dependency.shuffleWriteMetrics, diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala index e2894d1891a..6e8c38cdfbe 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala @@ -255,4 +255,54 @@ class CometNativeShuffleInputRDDSuite extends CometTestBase { assert(part.planDataByKey.keySet == Set("scan-0")) assert(part.planDataByKey("scan-0").length == 1024) } + + test("the native shuffle plan rides on the map stage and stays off the reduce stage") { + val sc = spark.sparkContext + val ser = new JavaSerializer(sc.getConf).newInstance() + def build(planBytes: Int) = { + val metrics = Map( + "spilled_bytes" -> SQLMetrics.createSizeMetric(sc, "disk spilled bytes"), + "memory_spilled_bytes" -> SQLMetrics.createSizeMetric(sc, "memory spilled bytes")) + val rdd = new CometNativeShuffleInputRDD( + sc, + inputRDDs = Seq.empty, + numPartitionsParam = 4, + shuffleScanIndices = Set.empty, + spillMetricNode = CometMetricNode(metrics)) + val execContext = NativeExecContext( + inputs = Seq.empty, + numPartitions = 4, + subqueries = Seq.empty, + broadcastedHadoopConfForEncryption = None, + encryptedFilePaths = Seq.empty, + commonByKey = Map("scan-0" -> new Array[Byte](planBytes)), + perPartitionByKey = Map.empty, + shuffleScanIndices = Set.empty, + hasScanInput = false) + val spec = + NativeShuffleSpec(Operator.getDefaultInstance, CometMetricNode(Map.empty), execContext) + rdd.nativeShuffleSpec = spec + val dep = new CometShuffleDependency[Int, ColumnarBatch, ColumnarBatch]( + _rdd = rdd, + partitioner = new HashPartitioner(4), + decodeTime = SQLMetrics.createMetric(sc, "decode time"), + shuffleWriteMetrics = metrics, + nativeShuffleSpec = Some(spec)) + (rdd, dep) + } + + val (smallRdd, smallDep) = build(1024) + val (largeRdd, largeDep) = build(4 * 1024 * 1024) + val smallReduce = ser.serialize(new CometShuffledBatchRDD(smallDep, Map.empty)).limit() + val largeReduce = ser.serialize(new CometShuffledBatchRDD(largeDep, Map.empty)).limit() + assert(math.abs(largeReduce - smallReduce) < 1024, s"$smallReduce vs $largeReduce") + assert(ser.serialize((largeRdd, largeDep)).limit() > 4 * 1024 * 1024) + + val (mapRdd, mapDep) = ser.deserialize[( + CometNativeShuffleInputRDD, + CometShuffleDependency[Int, ColumnarBatch, ColumnarBatch])]( + ser.serialize((smallRdd, smallDep))) + assert(mapDep.nativeShuffleSpec == null) + assert(mapRdd.nativeShuffleSpec.execContext.commonByKey("scan-0").length == 1024) + } } From d0c3e79d2dafbcd253b94aa87f6544c9eeb1ee8d Mon Sep 17 00:00:00 2001 From: msaf Date: Mon, 5 Oct 2026 15:18:20 +0100 Subject: [PATCH 09/14] perf: intern the query text of scan planning data and native shuffle plans convertBlock pools QueryContext SQL text on the root of each native block, but three paths still shipped the whole query text once per expression: - a scan's common planning data (NativeScanCommon data filters, and the Delta contrib's DeltaSparkScan) is serialized from the scan's own operator and injected under the block root on executors; - the native shuffle writer's child plan is the un-interned nativeOp; - a Delta scan sits inside a ContribScan envelope the interner cannot open. With a long query these copies dominated the task binary: a union of 115 filtered scans with a 42 KB query text weighed 92 MB, a four-stage query with 13 KB of text 32 MB. Commons are now interned against the pool of the plan they are injected into (adding missing texts to it), the shuffle writer's child plan is the block's interned plan with its pool hoisted to the ShuffleWriter root, and PlanDataInjector.internScan lets the Delta injector intern its payload in the block plan. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../sql/comet/CometDeltaNativeScanExec.scala | 13 ++- .../sql/comet/DeltaPlanDataInjector.scala | 10 +++ .../delta/CometDeltaNativeScanSuite.scala | 50 +++++++++++ .../comet/serde/QueryContextInterner.scala | 61 +++++++++++++- .../spark/sql/comet/CometNativeScanExec.scala | 12 ++- .../shuffle/CometNativeShuffleWriter.scala | 1 + .../shuffle/CometShuffleDependency.scala | 3 +- .../shuffle/CometShuffleExchangeExec.scala | 16 +++- .../apache/spark/sql/comet/operators.scala | 82 +++++++++++++++++-- .../comet/QueryContextInternerSuite.scala | 26 ++++++ .../CometNativeShuffleInputRDDSuite.scala | 66 +++++++++++++++ 11 files changed, 325 insertions(+), 15 deletions(-) diff --git a/contrib/delta-spark/src/main/scala/org/apache/spark/sql/comet/CometDeltaNativeScanExec.scala b/contrib/delta-spark/src/main/scala/org/apache/spark/sql/comet/CometDeltaNativeScanExec.scala index 1c1f7e364c9..f8a922f8138 100644 --- a/contrib/delta-spark/src/main/scala/org/apache/spark/sql/comet/CometDeltaNativeScanExec.scala +++ b/contrib/delta-spark/src/main/scala/org/apache/spark/sql/comet/CometDeltaNativeScanExec.scala @@ -19,6 +19,8 @@ package org.apache.spark.sql.comet +import scala.jdk.CollectionConverters._ + import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.plans.QueryPlan @@ -30,7 +32,7 @@ import org.apache.spark.sql.types.StructType import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.comet.contrib.delta.DeltaSparkScanEnvelope -import org.apache.comet.serde.OperatorOuterClass +import org.apache.comet.serde.{OperatorOuterClass, QueryContextInterner} import org.apache.comet.serde.OperatorOuterClass.Operator /** @@ -148,12 +150,17 @@ case class CometDeltaNativeScanExec( override def doExecuteColumnar(): RDD[ColumnarBatch] = { val nativeMetrics = CometMetricNode.fromCometPlan(this) - val serializedPlan = CometExec.serializeNativePlan(nativeOp) + val plan = PlanDataInjector.internScans(QueryContextInterner.intern(nativeOp)) + val sqlTextPool = new QueryContextInterner.Pool(plan.getSqlTextPoolList.asScala.toSeq) + val commonByKey = + PlanDataInjector.internCommons(plan, Map(sourceKey -> commonData), sqlTextPool) + val serializedPlan = CometExec.serializeNativePlan( + plan.toBuilder.addAllSqlTextPool(sqlTextPool.added.asJava).build()) new CometExecRDD( sparkContext, Seq.empty, - Map(sourceKey -> commonData), + commonByKey, Map(sourceKey -> perPartitionData), serializedPlan, PlanDataInjector.planFingerprint(serializedPlan), diff --git a/contrib/delta-spark/src/main/scala/org/apache/spark/sql/comet/DeltaPlanDataInjector.scala b/contrib/delta-spark/src/main/scala/org/apache/spark/sql/comet/DeltaPlanDataInjector.scala index 6697b57bde8..36cddd1879d 100644 --- a/contrib/delta-spark/src/main/scala/org/apache/spark/sql/comet/DeltaPlanDataInjector.scala +++ b/contrib/delta-spark/src/main/scala/org/apache/spark/sql/comet/DeltaPlanDataInjector.scala @@ -63,6 +63,16 @@ class DeltaPlanDataInjector extends PlanDataInjector { op.toBuilder.setContribScan(DeltaSparkScanEnvelope.pack(scanBuilder.build())).build() } + + override def internScan(op: Operator, pool: QueryContextInterner.Pool): Operator = { + val scan = DeltaSparkScanEnvelope.unpack(op) + val common = QueryContextInterner.internScanCommon(scan.getCommon, pool) + if (common eq scan.getCommon) op + else { + val interned = scan.toBuilder.setCommon(common).build() + op.toBuilder.setContribScan(DeltaSparkScanEnvelope.pack(interned)).build() + } + } } object DeltaPlanDataInjector { diff --git a/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaNativeScanSuite.scala b/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaNativeScanSuite.scala index 2425f93b866..d59880dbf18 100644 --- a/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaNativeScanSuite.scala +++ b/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaNativeScanSuite.scala @@ -3881,4 +3881,54 @@ class CometDeltaNativeScanSuite extends CometDeltaTestBase { } } } + + test("stage task binaries over delta scans do not repeat the query text per expression") { + withTempDir { dir => + val paths = (0 until 20).map { i => + val path = new File(dir, s"t$i").getAbsolutePath + spark + .range(0, 200) + .selectExpr("id", s"id % ${i + 3} AS k", "concat('data/', cast(id AS string)) AS name") + .write + .format("delta") + .save(path) + path + } + val sql = paths + .map(p => + s"SELECT k, name FROM delta.`$p` WHERE name LIKE 'data/%' AND NOT endswith(name, '/') " + + "AND id % 7 <> 3 AND length(name) > 5") + .mkString("SELECT k, count(*) AS c FROM (\n", "\nUNION ALL\n", "\n) GROUP BY k") + def mapStageBytes(cometEnabled: Boolean): Long = { + var bytes = 0L + withSQLConf( + CometConf.COMET_ENABLED.key -> cometEnabled.toString, + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val df = spark.sql(sql) + val serializer = org.apache.spark.SparkEnv.get.closureSerializer.newInstance() + val deps = mutable.ArrayBuffer.empty[org.apache.spark.ShuffleDependency[_, _, _]] + def walk(rdd: org.apache.spark.rdd.RDD[_]): Unit = { + rdd.partitions + rdd.dependencies.foreach { + case d: org.apache.spark.ShuffleDependency[_, _, _] => deps += d + case d => walk(d.rdd) + } + } + walk(df.queryExecution.executedPlan.execute()) + deps.foreach { d => + walk(d.rdd) + bytes = math.max(bytes, serializer.serialize((d.rdd, d)).limit().toLong) + } + checkSparkAnswer(df) + } + bytes + } + val vanilla = mapStageBytes(cometEnabled = false) + val comet = mapStageBytes(cometEnabled = true) + // scalastyle:off println + println(s"delta task binary: vanilla=$vanilla comet=$comet sql=${sql.length}") + // scalastyle:on println + assert(comet < vanilla + 20L * sql.length * 3, s"vanilla=$vanilla comet=$comet") + } + } } diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryContextInterner.scala b/spark/src/main/scala/org/apache/comet/serde/QueryContextInterner.scala index 68a11f216b0..445fe0c29f6 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryContextInterner.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryContextInterner.scala @@ -21,8 +21,8 @@ package org.apache.comet.serde import scala.collection.mutable +import com.google.protobuf.{CodedInputStream, CodedOutputStream, Message, WireFormat} import com.google.protobuf.Descriptors.FieldDescriptor -import com.google.protobuf.Message import org.apache.comet.serde.OperatorOuterClass.Operator @@ -81,6 +81,65 @@ object QueryContextInterner { } } + def sqlTextPool(planBytes: Array[Byte]): IndexedSeq[String] = { + val in = CodedInputStream.newInstance(planBytes) + val pool = IndexedSeq.newBuilder[String] + var tag = in.readTag() + while (tag != 0) { + if (WireFormat.getTagFieldNumber(tag) == Operator.SQL_TEXT_POOL_FIELD_NUMBER) { + pool += in.readString() + } else { + in.skipField(tag) + } + tag = in.readTag() + } + pool.result() + } + + final class Pool(initial: Seq[String]) { + private val indexOf = mutable.HashMap(initial.zipWithIndex: _*) + private val extra = mutable.ArrayBuffer.empty[String] + + def index(text: String): Int = + indexOf.getOrElseUpdate( + text, { + extra += text + initial.size + extra.size - 1 + }) + + def added: Seq[String] = extra.toSeq + } + + def internAgainst[M <: Message](message: M, pool: Pool): M = { + var changed = false + val builder = message.toBuilder + walk( + builder, + ctx => + if (ctx.hasSqlTextIdx) { + Some(ctx) + } else { + changed = true + Some(ctx.toBuilder.clearSqlText().setSqlTextIdx(pool.index(ctx.getSqlText)).build()) + }) + if (changed) builder.build().asInstanceOf[M] else message + } + + def internScanCommon( + common: OperatorOuterClass.NativeScanCommon, + pool: Pool): OperatorOuterClass.NativeScanCommon = + internAgainst(common, pool) + + def appendToPool(planBytes: Array[Byte], texts: Seq[String]): Array[Byte] = { + if (texts.isEmpty) return planBytes + val out = new java.io.ByteArrayOutputStream(planBytes.length + texts.map(_.length + 8).sum) + out.write(planBytes) + val coded = CodedOutputStream.newInstance(out) + texts.foreach(coded.writeString(Operator.SQL_TEXT_POOL_FIELD_NUMBER, _)) + coded.flush() + out.toByteArray + } + /** * Returns `expr` with every `QueryContext` in it removed. * diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala index fad6ede8f21..b1ab7f4a014 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala @@ -19,6 +19,8 @@ package org.apache.spark.sql.comet +import scala.jdk.CollectionConverters._ + import org.apache.spark.{Partition, TaskContext} import org.apache.spark.rdd.RDD import org.apache.spark.sql.SparkSession @@ -39,6 +41,7 @@ import com.google.common.base.Objects import org.apache.comet.parquet.CometParquetUtils import org.apache.comet.serde.OperatorOuterClass.Operator +import org.apache.comet.serde.QueryContextInterner import org.apache.comet.serde.QueryPlanSerde.exprToProto import org.apache.comet.shims.ShimFileFormat @@ -239,7 +242,12 @@ case class CometNativeScanExec( override def doExecuteColumnar(): RDD[ColumnarBatch] = { val nativeMetrics = CometMetricNode.fromCometPlan(this) - val serializedPlan = CometExec.serializeNativePlan(nativeOp) + val plan = QueryContextInterner.intern(nativeOp) + val sqlTextPool = new QueryContextInterner.Pool(plan.getSqlTextPoolList.asScala.toSeq) + val commonByKey = + PlanDataInjector.internCommons(plan, Map(sourceKey -> commonData), sqlTextPool) + val serializedPlan = CometExec.serializeNativePlan( + plan.toBuilder.addAllSqlTextPool(sqlTextPool.added.asJava).build()) // Encryption config must be passed to each executor task val hadoopConf = relation.sparkSession.sessionState @@ -256,7 +264,7 @@ case class CometNativeScanExec( new CometExecRDD( sparkContext, Seq.empty, - Map(sourceKey -> commonData), + commonByKey, Map(sourceKey -> perPartitionData), serializedPlan, PlanDataInjector.planFingerprint(serializedPlan), diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala index c74c9c30a6e..b2ad736d1b4 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala @@ -465,6 +465,7 @@ class CometNativeShuffleWriter[K, V]( OperatorOuterClass.Operator .newBuilder() .setShuffleWriter(shuffleWriterBuilder) + .addAllSqlTextPool(activeSpec.sqlTextPool.asJava) .addChildren(activeSpec.childNativeOp) .build() } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala index ff36ad29329..f6ac9d45fc1 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleDependency.scala @@ -51,7 +51,8 @@ case class NativeShuffleSpec( * size so that a re-executed map task places rows exactly as the attempt it replaces. See * `CometShuffleExchangeExec.positionalRoundRobinSpec`. */ - positionalRoundRobin: Option[PositionalRoundRobin] = None) + positionalRoundRobin: Option[PositionalRoundRobin] = None, + sqlTextPool: Seq[String] = Seq.empty) /** * Parameters for positional round-robin placement, resolved on the driver. diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala index b57b454adc2..38febca1f98 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala @@ -222,8 +222,8 @@ case class CometShuffleExchangeExec( outputPartitioning, serializer, metrics, - NativeShuffleSpec( - nativeChild.nativeOp, + CometShuffleExchangeExec.nativeShuffleSpec( + nativeChild.internedNativeOp, nativeChildMetricNode, ctx, positionalRoundRobin)) @@ -945,6 +945,18 @@ object CometShuffleExchangeExec NativeShuffleSpec(scanOp, childMetricNode, ctx)) } + private[comet] def nativeShuffleSpec( + childNativeOp: OperatorOuterClass.Operator, + childMetricNode: CometMetricNode, + ctx: NativeExecContext, + positionalRoundRobin: Option[PositionalRoundRobin]): NativeShuffleSpec = + NativeShuffleSpec( + childNativeOp.toBuilder.clearSqlTextPool().build(), + childMetricNode, + ctx, + positionalRoundRobin, + childNativeOp.getSqlTextPoolList.asScala.toSeq ++ ctx.addedSqlTexts) + /** * Build a Comet native shuffle dependency for the [[CometShuffleExchangeExec]] case where the * shuffle is fed by a [[CometNativeExec]] child. The writer drives the unified 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 a692e2fc454..99a19be31eb 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 @@ -106,6 +106,8 @@ private[comet] trait PlanDataInjector { * operators need rebuilding. */ def inject(op: Operator, preparedCommon: Prepared, partitionBytes: Array[Byte]): Operator + + def internScan(op: Operator, pool: QueryContextInterner.Pool): Operator = op } /** @@ -412,6 +414,61 @@ private[comet] object PlanDataInjector extends Logging { bytes } + private[comet] def internScans(root: Operator): Operator = { + val pool = new QueryContextInterner.Pool(root.getSqlTextPoolList.asScala.toSeq) + def visit(op: Operator): Operator = { + val self = injectorsByKind + .get(op.getOpStructCase) + .flatMap(_.find(_.canInject(op))) + .map(_.internScan(op, pool)) + .getOrElse(op) + val children = self.getChildrenList + var builder: Operator.Builder = null + var i = 0 + while (i < children.size()) { + val child = children.get(i) + val visited = visit(child) + if (visited ne child) { + if (builder == null) builder = self.toBuilder + builder.setChildren(i, visited) + } + i += 1 + } + if (builder == null) self else builder.build() + } + val result = visit(root) + if (pool.added.isEmpty) result + else result.toBuilder.addAllSqlTextPool(pool.added.asJava).build() + } + + private[comet] def internCommons( + root: Operator, + commonByKey: Map[String, Array[Byte]], + pool: QueryContextInterner.Pool): Map[String, Array[Byte]] = { + if (commonByKey.isEmpty) return commonByKey + val interned = mutable.Map.empty[String, Array[Byte]] + def visit(op: Operator): Unit = { + injectorsByKind.get(op.getOpStructCase).flatMap(_.find(_.canInject(op))).foreach { + injector => + injector.getKey(op).foreach { key => + commonByKey.get(key).filterNot(_ => interned.contains(key)).foreach { bytes => + injector.prepareCommon(bytes) match { + case message: com.google.protobuf.Message => + val result = QueryContextInterner.internAgainst(message, pool) + if (result ne message) { + interned(key) = result.toByteArray + } + case _ => + } + } + } + } + op.getChildrenList.asScala.foreach(visit) + } + visit(root) + if (interned.isEmpty) commonByKey else commonByKey ++ interned + } + /** * Find all plan nodes with per-partition planning data in the plan tree. Returns two maps keyed * by a unique identifier: one for common data (shared across partitions) and one for @@ -770,7 +827,8 @@ private[comet] case class NativeExecContext( // binary when this context rides on the map-side CometNativeShuffleInputRDD.nativeShuffleSpec. @transient perPartitionByKey: Map[String, Array[Array[Byte]]], shuffleScanIndices: Set[Int], - hasScanInput: Boolean) { + hasScanInput: Boolean, + addedSqlTexts: Seq[String] = Seq.empty) { // Catch shape divergence (e.g. broadcast scans with different partition counts after DPP // filtering) at construction so consumers don't trip ArrayIndexOutOfBoundsException at // partition idx access time. @@ -794,6 +852,9 @@ abstract class CometNativeExec extends CometExec { /** The Comet native operator */ def nativeOp: Operator + private[comet] def internedNativeOp: Operator = + serializedPlanOpt.plan.map(Operator.parseFrom).getOrElse(nativeOp) + override protected def doPrepare(): Unit = prepareSubqueries(this) override lazy val metrics: Map[String, SQLMetric] = @@ -832,9 +893,11 @@ abstract class CometNativeExec extends CometExec { private[comet] def executeColumnarWithContext( ctx: NativeExecContext, nativeMetrics: CometMetricNode): RDD[ColumnarBatch] = { - val serializedPlan = serializedPlanOpt.plan.getOrElse( - throw new CometRuntimeException( - s"CometNativeExec should not be executed directly without a serialized plan: $this")) + val serializedPlan = QueryContextInterner.appendToPool( + serializedPlanOpt.plan.getOrElse( + throw new CometRuntimeException( + s"CometNativeExec should not be executed directly without a serialized plan: $this")), + ctx.addedSqlTexts) new CometExecRDD( sparkContext, @@ -895,7 +958,10 @@ abstract class CometNativeExec extends CometExec { } // Find planning data within this stage (stops at shuffle boundaries). - val (commonByKey, perPartitionByKey) = PlanDataInjector.findAllPlanData(this) + val (scanCommonByKey, perPartitionByKey) = PlanDataInjector.findAllPlanData(this) + val sqlTextPool = new QueryContextInterner.Pool( + serializedPlanOpt.plan.map(QueryContextInterner.sqlTextPool).getOrElse(Nil)) + val commonByKey = PlanDataInjector.internCommons(nativeOp, scanCommonByKey, sqlTextPool) // Collect the input batches from the child operators. Non-shuffle inputs become // RDD[ArrowArrayStream] (one stream per partition, exported via the C Stream Interface @@ -1054,6 +1120,7 @@ abstract class CometNativeExec extends CometExec { commonByKey = commonByKey, perPartitionByKey = perPartitionByKey, shuffleScanIndices = shuffleScanIndices, + addedSqlTexts = sqlTextPool.added, // A leaf Comet scan (`CometNativeScanExec`, `CometIcebergNativeScanExec`, or a contrib // leaf such as `CometDeltaNativeScanExec`) can contribute `bytes_scanned` / // `output_rows` to Spark's task-level input metrics, which drive the Input column on @@ -1139,7 +1206,10 @@ abstract class CometNativeExec extends CometExec { // Hoist duplicated QueryContext SQL text into a pool on the root operator. This is the // point where the whole native block is in hand, which is what the pool indices are // scoped to. - SerializedPlan(Some(CometExec.serializeNativePlan(QueryContextInterner.intern(nativeOp)))) + SerializedPlan( + Some( + CometExec.serializeNativePlan( + PlanDataInjector.internScans(QueryContextInterner.intern(nativeOp))))) case other: AnyRef => other case null => null } diff --git a/spark/src/test/scala/org/apache/comet/QueryContextInternerSuite.scala b/spark/src/test/scala/org/apache/comet/QueryContextInternerSuite.scala index fd33f249947..8a939e01e9e 100644 --- a/spark/src/test/scala/org/apache/comet/QueryContextInternerSuite.scala +++ b/spark/src/test/scala/org/apache/comet/QueryContextInternerSuite.scala @@ -134,6 +134,32 @@ class QueryContextInternerSuite extends CometTestBase { } } + test("ANSI error reports full SQL context from a native shuffle's map side and scan filter") { + withSQLConf("spark.sql.ansi.enabled" -> "true") { + withTempPath { dir => + spark + .range(1, 10) + .selectExpr("cast(id as int) as a", "cast(0 as int) as b", "id % 3 as k") + .write + .mode("overwrite") + .parquet(dir.getAbsolutePath) + withTempView("t_div") { + spark.read.parquet(dir.getAbsolutePath).createOrReplaceTempView("t_div") + Seq( + "select k, sum(a / b) as s from t_div group by k", + "select k, count(*) as c from t_div where a / b > 1 group by k").foreach { query => + val err = intercept[Exception](spark.sql(query).collect()) + val msg = Option(err.getMessage).getOrElse("") + assert(msg.contains("DIVIDE_BY_ZERO"), msg) + assert( + msg.contains("== SQL") && msg.contains("a / b"), + s"error message lost the SQL context:\n$msg") + } + } + } + } + } + test("ANSI error still reports full SQL context after interning") { // Deliberately its own integer-typed table rather than reusing `withTestTable`: integer // division routes through the native `CheckedBinaryExpr`, which is what consults the diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala index 6e8c38cdfbe..6d8aa2eb18f 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDDSuite.scala @@ -19,6 +19,9 @@ package org.apache.spark.sql.comet.execution.shuffle +import scala.collection.mutable +import scala.jdk.CollectionConverters._ + import org.apache.spark.{HashPartitioner, Partition, TaskContext} import org.apache.spark.rdd.{DeterministicLevel, RDD} import org.apache.spark.serializer.JavaSerializer @@ -27,6 +30,8 @@ import org.apache.spark.sql.comet.{CometMetricNode, NativeExecContext} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} import org.apache.spark.sql.vectorized.ColumnarBatch +import org.apache.comet.CometConf +import org.apache.comet.serde.{ExprOuterClass, OperatorOuterClass} import org.apache.comet.serde.OperatorOuterClass.Operator /** @@ -305,4 +310,65 @@ class CometNativeShuffleInputRDDSuite extends CometTestBase { assert(mapDep.nativeShuffleSpec == null) assert(mapRdd.nativeShuffleSpec.execContext.commonByKey("scan-0").length == 1024) } + + test("a native shuffle's plan and scan planning data refer to the writer's pool") { + def contexts(m: com.google.protobuf.Message): Seq[ExprOuterClass.QueryContext] = { + val own = m match { + case e: ExprOuterClass.Expr if e.hasQueryContext => Seq(e.getQueryContext) + case a: ExprOuterClass.AggExpr if a.hasQueryContext => Seq(a.getQueryContext) + case _ => Nil + } + own ++ m.getAllFields.values.asScala.toSeq.flatMap { + case child: com.google.protobuf.Message => contexts(child) + case list: java.util.List[_] => + list.asScala.toSeq.flatMap { + case child: com.google.protobuf.Message => contexts(child) + case _ => Nil + } + case _ => Nil + } + } + withTempPath { dir => + spark + .range(0, 100) + .selectExpr("cast(id % 17 as int) as k1", "cast(id % 1000 as int) as i1") + .write + .parquet(dir.getAbsolutePath) + withTempView("t1") { + spark.read.parquet(dir.getAbsolutePath).createOrReplaceTempView("t1") + withSQLConf(CometConf.COMET_SHUFFLE_MODE.key -> "native") { + val df = spark.sql( + "select k1, max(case when i1 > 100 then i1 + 1 else i1 * 2 end) as m " + + "from t1 where i1 between 1 and 1000 group by k1") + df.collect() + val rdds = mutable.ArrayBuffer.empty[RDD[_]] + def walk(rdd: RDD[_]): Unit = { + rdds += rdd + rdd.dependencies.foreach(d => walk(d.rdd)) + } + walk(df.queryExecution.executedPlan.execute()) + val specs = rdds.collect { + case r: CometNativeShuffleInputRDD if r.nativeShuffleSpec != null => + r.nativeShuffleSpec + } + assert(specs.nonEmpty, "expected a native shuffle fed by a native block") + specs.foreach { spec => + assert(spec.childNativeOp.getSqlTextPoolCount == 0) + assert(spec.sqlTextPool.size == 1) + val planContexts = contexts(spec.childNativeOp) + assert(planContexts.nonEmpty) + val commons = spec.execContext.commonByKey.values + .map(OperatorOuterClass.NativeScanCommon.parseFrom) + assert(commons.nonEmpty, "expected the scan to be fused into the shuffle's child") + val commonContexts = commons.flatMap(contexts) + assert(commonContexts.nonEmpty) + (planContexts ++ commonContexts).foreach { ctx => + assert(ctx.hasSqlTextIdx && ctx.getSqlText.isEmpty, ctx) + assert(ctx.getSqlTextIdx < spec.sqlTextPool.size) + } + } + } + } + } + } } From af29142a3876ea86d142cb3cb6328805b2bd3d6b Mon Sep 17 00:00:00 2001 From: msaf Date: Mon, 5 Oct 2026 15:18:20 +0100 Subject: [PATCH 10/14] test: bound stage task binary sizes against vanilla Spark Serializes every stage's task binary the way DAGScheduler does for a union of many filtered scans and a multi-stage query with a long query text, with and without AQE, and bounds each against the largest vanilla Spark stage plus one copy of the query text per native block. Also checks that a reduce stage does not carry the plan of the stage it reads. Co-Authored-By: Claude Opus 5.5 (1M context) --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + .../comet/exec/CometTaskBinarySizeSuite.scala | 201 ++++++++++++++++++ 3 files changed, 203 insertions(+) create mode 100644 spark/src/test/scala/org/apache/comet/exec/CometTaskBinarySizeSuite.scala diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 5ff917f1a8a..8b183a10bbc 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -527,6 +527,7 @@ jobs: org.apache.comet.exec.CometAggregateSuite org.apache.comet.exec.CometExec3_4PlusSuite org.apache.comet.exec.CometExecSuite + org.apache.comet.exec.CometTaskBinarySizeSuite org.apache.comet.exec.CometEmptyRelationExecSuite org.apache.comet.exec.CometInMemoryCacheSuite org.apache.comet.exec.CometInMemoryCacheKryoSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index c79bd3116ba..cb92fbcfd2d 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -175,6 +175,7 @@ jobs: org.apache.comet.exec.CometAggregateSuite org.apache.comet.exec.CometExec3_4PlusSuite org.apache.comet.exec.CometExecSuite + org.apache.comet.exec.CometTaskBinarySizeSuite org.apache.comet.exec.CometEmptyRelationExecSuite org.apache.comet.exec.CometInMemoryCacheSuite org.apache.comet.exec.CometInMemoryCacheKryoSuite diff --git a/spark/src/test/scala/org/apache/comet/exec/CometTaskBinarySizeSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometTaskBinarySizeSuite.scala new file mode 100644 index 00000000000..cbcb80f32c7 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/exec/CometTaskBinarySizeSuite.scala @@ -0,0 +1,201 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.comet.exec + +import scala.collection.mutable + +import org.apache.spark.{ShuffleDependency, SparkEnv, TaskContext} +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.{CometTestBase, DataFrame} +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.CometConf + +class CometTaskBinarySizeSuite extends CometTestBase { + + private case class StageBinary(name: String, bytes: Long) + + private def narrowLineage(rdd: RDD[_]): Seq[RDD[_]] = { + val seen = mutable.LinkedHashMap.empty[Int, RDD[_]] + def walk(r: RDD[_]): Unit = if (!seen.contains(r.id)) { + seen(r.id) = r + r.dependencies.foreach { + case _: ShuffleDependency[_, _, _] => + case d => walk(d.rdd) + } + } + walk(rdd) + seen.values.toSeq + } + + private def stageBinaries(df: DataFrame): Seq[StageBinary] = { + if (spark.conf.get(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key).toBoolean) df.collect() + val serializer = SparkEnv.get.closureSerializer.newInstance() + val func = (_: TaskContext, it: Iterator[_]) => it.size + val out = mutable.ArrayBuffer.empty[StageBinary] + val visited = mutable.Set.empty[Int] + def visit(rdd: RDD[_], dep: Option[ShuffleDependency[_, _, _]]): Unit = { + val lineage = narrowLineage(rdd) + lineage.foreach(_.partitions) + val payload: AnyRef = dep match { + case Some(d) => (rdd, d) + case None => (rdd, func) + } + out += StageBinary( + dep.map(d => s"shuffle ${d.shuffleId}").getOrElse("result"), + serializer.serialize(payload).limit().toLong) + lineage.flatMap(_.dependencies).foreach { + case d: ShuffleDependency[_, _, _] if visited.add(d.shuffleId) => visit(d.rdd, Some(d)) + case _ => + } + } + visit(df.queryExecution.executedPlan.execute(), None) + out.toSeq + } + + private def compare(sql: String, aqe: Boolean)( + check: (Seq[StageBinary], Seq[StageBinary]) => Unit): Unit = { + def binaries(cometEnabled: Boolean): Seq[StageBinary] = { + var result = Seq.empty[StageBinary] + withSQLConf( + CometConf.COMET_ENABLED.key -> cometEnabled.toString, + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) { + result = stageBinaries(spark.sql(sql)) + } + result + } + val vanilla = binaries(cometEnabled = false) + val comet = binaries(cometEnabled = true) + withClue(s"aqe=$aqe vanilla=$vanilla comet=$comet") { + check(vanilla, comet) + } + } + + private def writeInventory(dir: java.io.File, n: Int): Seq[String] = + (0 until n).map { i => + val path = new java.io.File(dir, s"inventory$i").getCanonicalPath + spark + .range(50) + .selectExpr( + s"'bucket$i' AS bucket", + "concat('data/versioned/', cast(id AS string)) AS name", + "cast(id AS decimal(20, 0)) AS size", + "IF(id % 5 = 0, timestamp'2026-01-01 00:00:00', NULL) AS timeDeleted") + .coalesce(1) + .write + .parquet(path) + path + } + + private def inventoryUnion(paths: Seq[String]): String = + paths + .map { p => + s"SELECT concat('gs://', bucket, '/', name) AS path, size FROM parquet.`$p` " + + "WHERE name IS NOT NULL AND timeDeleted IS NULL AND NOT endswith(name, '/') AND " + + "name LIKE 'data/versioned/%' AND " + + "startswith(concat('gs://', bucket, '/', name), 'gs://bucket1/data/')" + } + .mkString("SELECT path, count(*) AS c FROM (\n", "\nUNION ALL\n", "\n) GROUP BY path") + + test("a union of many scans carries the query text once per scan") { + withTempDir { dir => + val scans = 60 + val sql = inventoryUnion(writeInventory(dir, scans)) + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "200") { + checkSparkAnswer(sql) + Seq(false, true).foreach { aqe => + compare(sql, aqe) { (vanilla, comet) => + val bound = vanilla.map(_.bytes).max + scans * sql.length * 3 / 2 + comet.foreach(stage => assert(stage.bytes < bound, s"${stage.name} over $bound")) + } + } + } + } + } + + test("a reduce stage does not carry the plan of the stage it reads") { + withTempDir { dir => + val sql = inventoryUnion(writeInventory(dir, 30)) + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "200") { + compare(sql, aqe = false) { (vanilla, comet) => + val mapStage = comet.filter(_.name != "result").map(_.bytes).max + val result = comet.find(_.name == "result").get.bytes + val vanillaResult = vanilla.find(_.name == "result").get.bytes + assert(result < mapStage / 2) + assert(result < vanillaResult * 2) + } + } + } + } + + private def dbtLikeSql(events: String, dim: String, cols: Int): String = { + val derived = (0 until cols) + .map(i => + s"CASE WHEN c$i > ${i * 7} THEN c$i * ${i + 3} WHEN s${i % 4} LIKE 'x$i%' THEN " + + s"length(s${i % 4}) ELSE coalesce(c${(i + 1) % cols}, 0) END AS d$i") + .mkString(",\n ") + val aggs = (0 until cols).map(i => s"sum(d$i) AS sd$i, max(d$i) AS md$i").mkString(",\n ") + val outs = (0 until cols) + .map(i => s"coalesce(sd$i, 0) + coalesce(md$i, 0) * dim.w AS o$i") + .mkString(",\n ") + s"""WITH base AS ( + | SELECT k, s0, $derived + | FROM parquet.`$events` + | WHERE c0 IS NOT NULL AND s1 NOT LIKE '%zzz%' + |), agg AS ( + | SELECT k, s0, $aggs + | FROM base GROUP BY k, s0 + |), joined AS ( + | SELECT agg.k, agg.s0, $outs + | FROM agg JOIN parquet.`$dim` dim ON agg.k = dim.k + |) + |SELECT s0, ${(0 until cols).map(i => s"sum(o$i) AS t$i").mkString(", ")} + |FROM joined GROUP BY s0""".stripMargin + } + + test("stages of a multi-stage query do not grow with the query text") { + withTempDir { dir => + val cols = 40 + val events = new java.io.File(dir, "events").getCanonicalPath + spark + .range(5000) + .selectExpr( + Seq("id % 97 AS k") ++ (0 until cols).map(i => s"(id * ${i + 1}) % 1013 AS c$i") ++ + (0 until 4).map(i => s"concat('x', cast(id % ${i + 5} AS string)) AS s$i"): _*) + .repartition(4) + .write + .parquet(events) + val dim = new java.io.File(dir, "dim").getCanonicalPath + spark.range(97).selectExpr("id AS k", "id % 7 AS w").write.parquet(dim) + val sql = dbtLikeSql(events, dim, cols) + withSQLConf( + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "50") { + checkSparkAnswer(sql) + Seq(false, true).foreach { aqe => + compare(sql, aqe) { (vanilla, comet) => + val bound = vanilla.map(_.bytes).max + 2 * sql.length + comet.foreach(stage => assert(stage.bytes < bound, s"${stage.name} over $bound")) + } + } + } + } + } +} From 035ef5c912801d7928545a5c0451685ea3d7e8d8 Mon Sep 17 00:00:00 2001 From: msaf Date: Mon, 5 Oct 2026 15:19:40 +0100 Subject: [PATCH 11/14] ci: list the suites added by the joom-1.1-2 fixes Co-Authored-By: Claude Opus 5.5 (1M context) --- .github/workflows/pr_build_linux.yml | 3 +++ .github/workflows/pr_build_macos.yml | 3 +++ 2 files changed, 6 insertions(+) diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 8b183a10bbc..8e837389659 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -516,6 +516,7 @@ jobs: org.apache.comet.exec.CometShuffleEncryptionSuite org.apache.comet.exec.CometShuffleManagerSuite org.apache.comet.exec.CometShuffleReadCoalesceSuite + org.apache.spark.sql.comet.execution.shuffle.CometBlockStoreShuffleReaderSuite org.apache.comet.exec.CometAsyncShuffleSuite org.apache.comet.exec.DisableAQECometShuffleSuite org.apache.comet.exec.DisableAQECometAsyncShuffleSuite @@ -528,6 +529,7 @@ jobs: org.apache.comet.exec.CometExec3_4PlusSuite org.apache.comet.exec.CometExecSuite org.apache.comet.exec.CometTaskBinarySizeSuite + org.apache.spark.sql.comet.CometNativeTaskKillSuite org.apache.comet.exec.CometEmptyRelationExecSuite org.apache.comet.exec.CometInMemoryCacheSuite org.apache.comet.exec.CometInMemoryCacheKryoSuite @@ -603,6 +605,7 @@ jobs: org.apache.comet.CometVariantTypeSuite org.apache.comet.CometHashExpressionSuite org.apache.comet.CometTemporalExpressionSuite + org.apache.comet.CometTimestampComparisonSuite org.apache.comet.CometArrayExpressionSuite org.apache.comet.CometNativeCastSuite org.apache.comet.CometDateTimeUtilsSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index cb92fbcfd2d..2fcd137ded6 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -164,6 +164,7 @@ jobs: org.apache.comet.exec.CometShuffleEncryptionSuite org.apache.comet.exec.CometShuffleManagerSuite org.apache.comet.exec.CometShuffleReadCoalesceSuite + org.apache.spark.sql.comet.execution.shuffle.CometBlockStoreShuffleReaderSuite org.apache.comet.exec.CometAsyncShuffleSuite org.apache.comet.exec.DisableAQECometShuffleSuite org.apache.comet.exec.DisableAQECometAsyncShuffleSuite @@ -176,6 +177,7 @@ jobs: org.apache.comet.exec.CometExec3_4PlusSuite org.apache.comet.exec.CometExecSuite org.apache.comet.exec.CometTaskBinarySizeSuite + org.apache.spark.sql.comet.CometNativeTaskKillSuite org.apache.comet.exec.CometEmptyRelationExecSuite org.apache.comet.exec.CometInMemoryCacheSuite org.apache.comet.exec.CometInMemoryCacheKryoSuite @@ -251,6 +253,7 @@ jobs: org.apache.comet.CometVariantTypeSuite org.apache.comet.CometHashExpressionSuite org.apache.comet.CometTemporalExpressionSuite + org.apache.comet.CometTimestampComparisonSuite org.apache.comet.CometArrayExpressionSuite org.apache.comet.CometNativeCastSuite org.apache.comet.CometDateTimeUtilsSuite From 03c33a5ff2f9d65e2f60d7cb5d7a7cdb6dbdd225 Mon Sep 17 00:00:00 2001 From: msaf Date: Mon, 5 Oct 2026 16:11:25 +0100 Subject: [PATCH 12/14] style: drop a needless string interpolator flagged by scalafix Co-Authored-By: Claude Opus 5.5 (1M context) --- spark/src/main/scala/org/apache/comet/CometExecIterator.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index f36636b85b3..af825520317 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -427,7 +427,7 @@ object CometExecIterator extends Logging { } catch { case NonFatal(e) => logWarning( - s"Failed to cancel the native plan of killed task " + + "Failed to cancel the native plan of killed task " + s"${context.taskAttemptId()}", e) } From d783931227e4eefa62e956db8cb46cd9b6f885d2 Mon Sep 17 00:00:00 2001 From: msaf Date: Mon, 5 Oct 2026 17:14:47 +0100 Subject: [PATCH 13/14] test: make the timestamp and task kill suites independent of the Spark version Co-Authored-By: Claude Opus 5.5 (1M context) --- .../apache/comet/CometTimestampComparisonSuite.scala | 2 +- .../spark/sql/comet/CometNativeTaskKillSuite.scala | 11 +++++++---- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/CometTimestampComparisonSuite.scala b/spark/src/test/scala/org/apache/comet/CometTimestampComparisonSuite.scala index 067c3acf154..9b166717f55 100644 --- a/spark/src/test/scala/org/apache/comet/CometTimestampComparisonSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometTimestampComparisonSuite.scala @@ -53,7 +53,7 @@ class CometTimestampComparisonSuite extends CometTestBase { private def inZones(zones: Seq[String])(f: => Unit): Unit = zones.foreach { zone => - withSQLConf("spark.sql.session.timeZone" -> zone) { + withSQLConf("spark.sql.session.timeZone" -> zone, "spark.sql.ansi.enabled" -> "false") { withClue(s"session time zone $zone: ")(f) } } diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/CometNativeTaskKillSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/CometNativeTaskKillSuite.scala index 3e3b5e7bdfd..d0f8aabd380 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/CometNativeTaskKillSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/CometNativeTaskKillSuite.scala @@ -165,10 +165,13 @@ class CometNativeTaskKillSuite extends CometTestBase { .selectExpr("id * 0 AS k", "id AS p") .join(broadcast(build), "k") .agg(sum(col("p") + col("b"))) - assertRunsNatively( - df, - classOf[CometSparkToColumnarExec], - classOf[CometBroadcastHashJoinExec]) + assertRunsNatively(df, classOf[CometSparkToColumnarExec]) + val plan = df.queryExecution.executedPlan + assert( + plan.collectFirst { + case j @ (_: CometBroadcastHashJoinExec | _: CometBroadcastNestedLoopJoinExec) => j + }.isDefined, + s"expected a native broadcast join in\n$plan") assertKilledPromptly(df) } } From 6450a99f5e8cd52f32fd3760e7d84feed48041b5 Mon Sep 17 00:00:00 2001 From: msaf Date: Mon, 5 Oct 2026 18:15:39 +0100 Subject: [PATCH 14/14] fix: take a hash aggregate's groups out before their keys overflow 32-bit offsets A hash aggregate interns the bytes of every distinct Utf8 or Binary group key into one buffer addressed by i32 offsets, and only empties it when the memory pool refuses the table. With a large enough pool a task with more than 2 GiB of distinct key bytes fails with "offset overflow, buffer size > 2147483647" (multi-column keys) or panics in ArrowBytesMap (one key column). Stage 308 of ads_manual_product_rank_cohorts (SELECT DISTINCT date, product_id, country over 5.8 billion feed rows hash-partitioned by date into 16 tasks, ~167 million groups each) failed this way 24 times on 2026-10-05; a retry passed only when the pool spilled first. DataFusion 55.1.0 and its main branch do the same (apache/datafusion#23694, #24704). Each group values implementation now tells whether a batch's keys still fit the 32-bit offsets of what it holds: the byte builders by their buffer length plus the batch's value bytes, the single-column map by the bytes it interned, and the row-format ones, whose decoded offsets never exceed their encoded size, by that size plus the batch's largest offset extent at any depth. When a batch would not fit, the table takes out every group so far before interning it: a partial aggregate emits them as it does under memory pressure, a final or single aggregate sorts and spills them for the merge replay, and a partial-reduce aggregate emits them. Tests: SELECT DISTINCT over 3 GiB of distinct 1 MiB keys in a partial and a final aggregate on (Int32, Utf8) and a partial on Utf8 alone failed with the production error (and the ArrowBytesMap panic) and now return every group. Co-Authored-By: Claude Opus 5.5 (1M context) --- native/core/src/execution/jni_api.rs | 111 ++++++++++++++++++ .../aggregates/aggregate_hash_table/common.rs | 19 ++- .../aggregate_hash_table/final_table.rs | 2 +- .../partial_reduce_table.rs | 2 +- .../aggregate_hash_table/partial_table.rs | 2 +- .../aggregate_hash_table/single_table.rs | 2 +- .../src/aggregates/group_values/mod.rs | 102 +++++++++++++++- .../group_values/multi_group_by/bytes.rs | 18 +++ .../group_values/multi_group_by/mod.rs | 13 ++ .../group_values/multi_group_by/row_backed.rs | 9 ++ .../src/aggregates/group_values/row.rs | 14 ++- .../group_values/single_group_by/bytes.rs | 18 ++- .../src/aggregates/hash_stream.rs | 36 +++++- .../src/aggregates/partial_reduce_stream.rs | 25 ++-- .../src/aggregates/single_stream.rs | 22 +++- 15 files changed, 370 insertions(+), 25 deletions(-) diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index b0ee151a205..9002e83194b 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -4114,3 +4114,114 @@ mod native_sort_spill_tests { } } } + +#[cfg(test)] +mod aggregate_offset_overflow_tests { + use arrow::array::{Int32Array, StringArray}; + use arrow::buffer::{Buffer, OffsetBuffer}; + use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; + use arrow::record_batch::RecordBatch; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::execution::TaskContext; + use datafusion::physical_expr::expressions::col; + use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}; + use datafusion::physical_plan::ExecutionPlan; + use datafusion::prelude::SessionConfig; + use futures::StreamExt; + use std::sync::Arc; + + const KEY_BYTES: usize = 1 << 20; + const ROWS: usize = 1024; + const BATCHES: usize = 3; + + /// `(day, product_id)` batches of distinct 1 MiB keys, 1 GiB per batch and 3 GiB in all, + /// sharing one buffer of pseudo-random letters so that the input itself stays at 1 GiB. + fn wide_keys(schema: &SchemaRef) -> Vec { + let mut seed = 0x2545_f491_4f6c_dd1du64; + let letters: Vec = (0..ROWS * KEY_BYTES + BATCHES) + .map(|_| { + seed ^= seed << 13; + seed ^= seed >> 7; + seed ^= seed << 17; + b'a' + (seed % 26) as u8 + }) + .collect(); + let values = Buffer::from_vec(letters); + (0..BATCHES) + .map(|batch| { + let offsets = OffsetBuffer::new( + (0..=ROWS) + .map(|row| (batch + row * KEY_BYTES) as i32) + .collect(), + ); + RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(Int32Array::from(vec![20_000; ROWS])), + Arc::new(StringArray::try_new(offsets, values.clone(), None).unwrap()), + ], + ) + .unwrap() + }) + .collect() + } + + /// `SELECT DISTINCT` of `columns` in one aggregate of `mode`, with nothing but the + /// 32-bit offsets of its group keys to stop it from holding every group. + async fn distinct_wide_keys( + mode: AggregateMode, + columns: &[&str], + ) -> datafusion::common::Result { + let schema = Arc::new(Schema::new(vec![ + Field::new("day", DataType::Int32, false), + Field::new("product_id", DataType::Utf8, false), + ])); + let input = + MemorySourceConfig::try_new_exec(&[wide_keys(&schema)], Arc::clone(&schema), None)?; + let group_by = PhysicalGroupBy::new_single( + columns + .iter() + .map(|name| Ok((col(name, &schema)?, name.to_string()))) + .collect::>()?, + ); + let aggregate = AggregateExec::try_new(mode, group_by, vec![], vec![], input, schema)?; + let ctx = + TaskContext::default().with_session_config(SessionConfig::new().with_batch_size(64)); + let mut stream = aggregate.execute(0, Arc::new(ctx))?; + let mut rows = 0; + while let Some(batch) = stream.next().await { + rows += batch?.num_rows(); + } + Ok(rows) + } + + #[tokio::test] + async fn partial_aggregate_emits_before_group_keys_overflow_32_bit_offsets() { + assert_eq!( + distinct_wide_keys(AggregateMode::Partial, &["day", "product_id"]) + .await + .unwrap(), + ROWS * BATCHES + ); + } + + #[tokio::test] + async fn final_aggregate_spills_before_group_keys_overflow_32_bit_offsets() { + assert_eq!( + distinct_wide_keys(AggregateMode::Final, &["day", "product_id"]) + .await + .unwrap(), + ROWS * BATCHES + ); + } + + #[tokio::test] + async fn single_key_aggregate_emits_before_group_keys_overflow_32_bit_offsets() { + assert_eq!( + distinct_wide_keys(AggregateMode::Partial, &["product_id"]) + .await + .unwrap(), + ROWS * BATCHES + ); + } +} diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/common.rs b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/common.rs index 91e9d6555c3..5c6521cc354 100644 --- a/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -185,12 +185,27 @@ impl AggregateHashTable { /// Each aggregation mode chooses a different `aggregate_fn` according to its /// semantics. For example, partial aggregation takes raw inputs, and update them /// into stored partial states, so [`GroupsAccumulator::update_batch`] is used. + /// + /// If the batch's group keys could overflow the 32-bit offsets of the stored group + /// values, every group accumulated so far is first taken out as by + /// [`Self::take_state_batch`] and returned, for the caller to emit or spill. pub(super) fn aggregate_batch_inner( &mut self, batch: &RecordBatch, aggregate_fn: AggregateBatchFn, - ) -> Result<()> { + ) -> Result> { let evaluated_batch = self.evaluate_batch(batch)?; + let state = self.state.building(); + let taken = if !state.group_values.is_empty() + && !evaluated_batch + .grouping_set_args + .iter() + .all(|group_values| state.group_values.has_offset_room(group_values)) + { + self.take_state_batch()? + } else { + None + }; let state = self.state.building_mut(); let _timer = self.group_by_metrics.aggregation_time.timer(); @@ -210,7 +225,7 @@ impl AggregateHashTable { } } - Ok(()) + Ok(taken) } /// Materializes the full output once, then returns it downstream incrementally diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/final_table.rs b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/final_table.rs index b80e15d7f83..47758ecab16 100644 --- a/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/final_table.rs +++ b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/final_table.rs @@ -66,7 +66,7 @@ impl AggregateHashTable { pub(in crate::aggregates) fn aggregate_batch( &mut self, batch: &RecordBatch, - ) -> Result<()> { + ) -> Result> { self.aggregate_batch_inner(batch, HashAggregateAccumulator::merge_batch) } diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs index 4dfd6a74d18..32878b805f4 100644 --- a/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs +++ b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs @@ -60,7 +60,7 @@ impl AggregateHashTable { pub(in crate::aggregates) fn aggregate_batch( &mut self, batch: &RecordBatch, - ) -> Result<()> { + ) -> Result> { self.aggregate_batch_inner(batch, HashAggregateAccumulator::merge_batch) } diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs index a64fd32536e..f481dc6fca8 100644 --- a/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs +++ b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs @@ -105,7 +105,7 @@ impl AggregateHashTable { pub(in crate::aggregates) fn aggregate_batch( &mut self, batch: &RecordBatch, - ) -> Result<()> { + ) -> Result> { self.aggregate_batch_inner(batch, HashAggregateAccumulator::update_batch) } diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/single_table.rs b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/single_table.rs index 56d601c7932..1a81c6d51aa 100644 --- a/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/single_table.rs +++ b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/single_table.rs @@ -65,7 +65,7 @@ impl AggregateHashTable { pub(in crate::aggregates) fn aggregate_batch( &mut self, batch: &RecordBatch, - ) -> Result<()> { + ) -> Result> { self.aggregate_batch_inner(batch, HashAggregateAccumulator::update_batch) } diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/group_values/mod.rs b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/mod.rs index 1101d535311..aebb920fd65 100644 --- a/native/vendor/datafusion-physical-plan/src/aggregates/group_values/mod.rs +++ b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/mod.rs @@ -22,7 +22,7 @@ use arrow::array::types::{ Time64MicrosecondType, Time64NanosecondType, TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType, }; -use arrow::array::{ArrayRef, downcast_primitive}; +use arrow::array::{Array, ArrayRef, AsArray, downcast_primitive}; use arrow::datatypes::{DataType, SchemaRef, TimeUnit}; use datafusion_common::Result; @@ -115,6 +115,106 @@ pub trait GroupValues: Send { /// Clear the contents and shrink the capacity to the size of the batch (free up memory usage) fn clear_shrink(&mut self, num_rows: usize); + + /// Returns false if interning `cols` could take a 32-bit offset buffer of the stored + /// group values, or of the arrays they are emitted as, past `i32::MAX`. + fn has_offset_room(&self, _cols: &[ArrayRef]) -> bool { + true + } +} + +/// The largest extent `array` addresses in any of its 32-bit offset buffers (`Utf8`, +/// `Binary`, `List` and `Map` at any depth): bytes for the former, entries for the latter. +pub(crate) fn max_offset_extent(array: &dyn Array) -> usize { + max_extent(array, 0, array.len()) +} + +fn max_extent(array: &dyn Array, start: usize, end: usize) -> usize { + match array.data_type() { + DataType::Utf8 => { + let offsets = array.as_string::().value_offsets(); + (offsets[end] - offsets[start]) as usize + } + DataType::Binary => { + let offsets = array.as_binary::().value_offsets(); + (offsets[end] - offsets[start]) as usize + } + DataType::List(_) => { + let list = array.as_list::(); + let offsets = list.value_offsets(); + let (start, end) = (offsets[start] as usize, offsets[end] as usize); + (end - start).max(max_extent(list.values().as_ref(), start, end)) + } + DataType::Map(_, _) => { + let map = array.as_map(); + let offsets = map.value_offsets(); + let (start, end) = (offsets[start] as usize, offsets[end] as usize); + (end - start).max(max_extent(map.entries(), start, end)) + } + DataType::LargeList(_) => { + let list = array.as_list::(); + let offsets = list.value_offsets(); + max_extent( + list.values().as_ref(), + offsets[start] as usize, + offsets[end] as usize, + ) + } + DataType::FixedSizeList(_, size) => { + let size = *size as usize; + max_extent( + array.as_fixed_size_list().values().as_ref(), + start * size, + end * size, + ) + } + DataType::Struct(_) => array + .as_struct() + .columns() + .iter() + .map(|column| max_extent(column.as_ref(), start, end)) + .max() + .unwrap_or(0), + DataType::Dictionary(_, _) => { + let values = array.as_any_dictionary().values(); + max_extent(values.as_ref(), 0, values.len()).saturating_mul(end - start) + } + _ => 0, + } +} + +/// True if `data_type` has a 32-bit offset buffer at any depth. +pub(crate) fn has_offset_buffer(data_type: &DataType) -> bool { + match data_type { + DataType::Utf8 | DataType::Binary | DataType::List(_) | DataType::Map(_, _) => { + true + } + DataType::LargeList(field) + | DataType::FixedSizeList(field, _) + | DataType::ListView(field) + | DataType::LargeListView(field) => has_offset_buffer(field.data_type()), + DataType::Struct(fields) => fields + .iter() + .any(|field| has_offset_buffer(field.data_type())), + DataType::Dictionary(_, values) => has_offset_buffer(values), + DataType::RunEndEncoded(_, values) => has_offset_buffer(values.data_type()), + DataType::Union(fields, _) => fields + .iter() + .any(|(_, field)| has_offset_buffer(field.data_type())), + _ => false, + } +} + +/// Whether encoded group rows of `held_bytes`, after appending rows decoded from `cols`, +/// still decode into arrays whose 32-bit offsets fit. Every row-format encoding takes at least +/// one byte per value byte and per list entry, so the decoded extent of each offset buffer is +/// at most the encoded size. +pub(crate) fn rows_have_offset_room(held_bytes: usize, cols: &[ArrayRef]) -> bool { + let incoming = cols + .iter() + .map(|col| max_offset_extent(col.as_ref())) + .fold(0usize, usize::saturating_add); + held_bytes.saturating_add(incoming) <= i32::MAX as usize } /// Return a specialized implementation of [`GroupValues`] for the given schema. diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs index c83b1da4049..b41469c8c69 100644 --- a/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs +++ b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use crate::aggregates::group_values::max_offset_extent; use crate::aggregates::group_values::multi_group_by::{ GroupColumn, Nulls, nulls_equal_to, }; @@ -219,6 +220,10 @@ where } } +fn large_extent(offsets: &[i64]) -> usize { + (offsets[offsets.len() - 1] - offsets[0]) as usize +} + impl GroupColumn for ByteGroupValueBuilder where O: OffsetSizeTrait, @@ -244,6 +249,19 @@ where } } + fn has_offset_room(&self, column: &ArrayRef) -> bool { + let incoming = match column.data_type() { + DataType::LargeUtf8 => { + large_extent(column.as_string::().value_offsets()) + } + DataType::LargeBinary => { + large_extent(column.as_binary::().value_offsets()) + } + _ => max_offset_extent(column.as_ref()), + }; + self.buffer.len().saturating_add(incoming) <= self.max_buffer_size + } + fn append_val(&mut self, column: &ArrayRef, row: usize) -> Result<()> { // Sanity array type match self.output_type { diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/mod.rs b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/mod.rs index 5b474f3bae0..8c62adb3150 100644 --- a/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/mod.rs +++ b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/mod.rs @@ -112,6 +112,12 @@ pub trait GroupColumn: Send + Sync { /// Builds a new array from the first `n` stored rows, shifting the /// remaining rows to the start of the builder fn take_n(&mut self, n: usize) -> ArrayRef; + + /// Returns false if appending every row of `array` could take a 32-bit offset + /// buffer of this builder, or of the array it builds, past `i32::MAX` + fn has_offset_room(&self, _array: &ArrayRef) -> bool { + true + } } /// Determines if the nullability of the existing and new input array can be used @@ -1165,6 +1171,13 @@ impl GroupValues for GroupValuesColumn { group_values_size + self.map_size + self.hashes_buffer.allocated_size() } + fn has_offset_room(&self, cols: &[ArrayRef]) -> bool { + self.group_values + .iter() + .zip(cols) + .all(|(group_column, col)| group_column.has_offset_room(col)) + } + fn is_empty(&self) -> bool { self.len() == 0 } diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs index 1445a81f218..eec770ad886 100644 --- a/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs +++ b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs @@ -48,6 +48,7 @@ use crate::aggregates::group_values::multi_group_by::GroupColumn; use crate::aggregates::group_values::row::encode_array_if_necessary; +use crate::aggregates::group_values::{has_offset_buffer, rows_have_offset_room}; use arrow::array::{Array, ArrayRef, BooleanBufferBuilder}; use arrow::datatypes::DataType; @@ -290,6 +291,14 @@ impl GroupColumn for RowsGroupColumn { self.row_converter.size() + self.group_values.size() } + fn has_offset_room(&self, array: &ArrayRef) -> bool { + !has_offset_buffer(&self.output_type) + || rows_have_offset_room( + self.group_values.size(), + std::slice::from_ref(array), + ) + } + fn build(self: Box) -> ArrayRef { self.rows_to_array(&self.group_values) } diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/group_values/row.rs b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/row.rs index cbd7a609c5c..597b7ee7e67 100644 --- a/native/vendor/datafusion-physical-plan/src/aggregates/group_values/row.rs +++ b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/row.rs @@ -15,7 +15,9 @@ // specific language governing permissions and limitations // under the License. -use crate::aggregates::group_values::GroupValues; +use crate::aggregates::group_values::{ + GroupValues, has_offset_buffer, rows_have_offset_room, +}; use arrow::array::{ Array, ArrayRef, FixedSizeListArray, LargeListArray, LargeListViewArray, ListArray, ListViewArray, MapArray, PrimitiveArray, RunArray, StructArray, @@ -193,6 +195,16 @@ impl GroupValues for GroupValuesRows { + self.hashes_buffer.allocated_size() } + fn has_offset_room(&self, cols: &[ArrayRef]) -> bool { + let held = self.group_values.as_ref().map(|v| v.size()).unwrap_or(0); + !self + .schema + .fields() + .iter() + .any(|field| has_offset_buffer(field.data_type())) + || rows_have_offset_room(held, cols) + } + fn is_empty(&self) -> bool { self.len() == 0 } diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/group_values/single_group_by/bytes.rs b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/single_group_by/bytes.rs index b881a51b254..ab09b1c273f 100644 --- a/native/vendor/datafusion-physical-plan/src/aggregates/group_values/single_group_by/bytes.rs +++ b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/single_group_by/bytes.rs @@ -17,7 +17,7 @@ use std::mem::size_of; -use crate::aggregates::group_values::GroupValues; +use crate::aggregates::group_values::{GroupValues, max_offset_extent}; use arrow::array::{Array, ArrayRef, OffsetSizeTrait}; use datafusion_common::Result; @@ -33,6 +33,8 @@ pub struct GroupValuesBytes { map: ArrowBytesMap, /// The total number of groups so far (used to assign group_index) num_groups: usize, + /// Bytes of the distinct values interned so far + value_bytes: usize, } impl GroupValuesBytes { @@ -40,6 +42,7 @@ impl GroupValuesBytes { Self { map: ArrowBytesMap::new(output_type), num_groups: 0, + value_bytes: 0, } } } @@ -55,7 +58,8 @@ impl GroupValues for GroupValuesBytes { self.map.insert_if_new( arr, // called for each new group - |_value| { + |value| { + self.value_bytes += value.map_or(0, <[u8]>::len); // assign new group index on each insert let group_idx = self.num_groups; self.num_groups += 1; @@ -87,6 +91,7 @@ impl GroupValues for GroupValuesBytes { fn emit(&mut self, emit_to: EmitTo) -> Result> { // Reset the map to default, and convert it into a single array let map_contents = self.map.take().into_state(); + self.value_bytes = 0; let group_values = match emit_to { EmitTo::All => { @@ -124,5 +129,14 @@ impl GroupValues for GroupValuesBytes { // in theory we could potentially avoid this reallocation and clear the // contents of the maps, but for now we just reset the map from the beginning self.map.take(); + self.value_bytes = 0; + } + + fn has_offset_room(&self, cols: &[ArrayRef]) -> bool { + O::IS_LARGE + || self + .value_bytes + .saturating_add(max_offset_extent(cols[0].as_ref())) + <= i32::MAX as usize } } diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/hash_stream.rs b/native/vendor/datafusion-physical-plan/src/aggregates/hash_stream.rs index 3907eb34b82..d3b63c7083f 100644 --- a/native/vendor/datafusion-physical-plan/src/aggregates/hash_stream.rs +++ b/native/vendor/datafusion-physical-plan/src/aggregates/hash_stream.rs @@ -32,7 +32,9 @@ use std::task::{Context, Poll}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use datafusion_common::{DataFusionError, Result, internal_datafusion_err, internal_err}; +use datafusion_common::{ + DataFusionError, Result, exec_err, internal_datafusion_err, internal_err, +}; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; use datafusion_physical_expr::PhysicalSortExpr; @@ -355,7 +357,11 @@ impl FinalSpillContext { let Some(batch) = hash_table.take_state_batch()? else { return Ok(()); }; + self.spill_batch(batch) + } + /// Sorts and spills a batch of intermediate states taken from the table. + fn spill_batch(&mut self, batch: RecordBatch) -> Result<()> { let sorted_iter = IncrementalSortIterator::new(batch, self.spill_expr.clone(), self.batch_size); let spill_file = self @@ -578,8 +584,18 @@ impl PartialHashAggregateStream { let result = hash_table.aggregate_batch(&batch); timer.done(); - if let Err(e) = result { - return Self::break_with_err(e); + match result { + Ok(None) => {} + Ok(Some(taken)) => { + let _ = self.reservation.try_resize(hash_table.memory_size()); + return ControlFlow::Continue( + PartialHashAggregateState::EmittingOnMemoryPressure { + hash_table, + remaining_groups: taken, + }, + ); + } + Err(e) => return Self::break_with_err(e), } // -------------------------------- @@ -1129,7 +1145,7 @@ impl FinalHashAggregateStream { ) -> FinalHashAggregateStateTransition { let FinalHashAggregateState::ReadingInput { mut hash_table, - spill_context, + mut spill_context, } = original_state else { return Self::break_with_internal_err( @@ -1148,7 +1164,17 @@ impl FinalHashAggregateStream { Poll::Ready(Some(Ok(batch))) => { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); - let result = hash_table.aggregate_batch(&batch); + let result = hash_table.aggregate_batch(&batch).and_then(|taken| { + match (taken, spill_context.as_deref_mut()) { + (None, _) => Ok(()), + (Some(taken), Some(spill_context)) => { + spill_context.spill_batch(taken) + } + (Some(_), None) => exec_err!( + "Final hash aggregate cannot spill group keys that overflow 32-bit offsets because temporary files are not enabled in the DiskManager" + ), + } + }); timer.done(); if let Err(e) = result { diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/partial_reduce_stream.rs b/native/vendor/datafusion-physical-plan/src/aggregates/partial_reduce_stream.rs index 2f4535e66f4..c0ff4f255dc 100644 --- a/native/vendor/datafusion-physical-plan/src/aggregates/partial_reduce_stream.rs +++ b/native/vendor/datafusion-physical-plan/src/aggregates/partial_reduce_stream.rs @@ -218,12 +218,15 @@ impl PartialReduceHashAggregateStream { let result = original_state.hash_table_mut().aggregate_batch(&batch); timer.done(); - if let Err(e) = result { - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - original_state, - )); - } + let taken = match result { + Ok(taken) => taken, + Err(e) => { + return ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + original_state, + )); + } + }; if let Err(e) = self .reservation @@ -235,7 +238,15 @@ impl PartialReduceHashAggregateStream { )); } - ControlFlow::Continue(original_state) + match taken { + Some(taken) => ControlFlow::Break(( + Poll::Ready(Some( + Ok(taken.record_output(&self.baseline_metrics)), + )), + original_state, + )), + None => ControlFlow::Continue(original_state), + } } Poll::Ready(Some(Err(e))) => { ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)) diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/single_stream.rs b/native/vendor/datafusion-physical-plan/src/aggregates/single_stream.rs index 9541c3ca5ff..17525ccd0f1 100644 --- a/native/vendor/datafusion-physical-plan/src/aggregates/single_stream.rs +++ b/native/vendor/datafusion-physical-plan/src/aggregates/single_stream.rs @@ -28,7 +28,9 @@ use std::task::{Context, Poll}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use datafusion_common::{DataFusionError, Result, internal_datafusion_err, internal_err}; +use datafusion_common::{ + DataFusionError, Result, exec_err, internal_datafusion_err, internal_err, +}; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; use datafusion_physical_expr::PhysicalSortExpr; @@ -240,7 +242,11 @@ impl SingleSpillContext { let Some(batch) = hash_table.take_state_batch()? else { return Ok(()); }; + self.spill_batch(batch) + } + /// Sorts and spills a batch of intermediate states taken from the table. + fn spill_batch(&mut self, batch: RecordBatch) -> Result<()> { let sorted_iter = IncrementalSortIterator::new(batch, self.spill_expr.clone(), self.batch_size); let spill_file = self @@ -419,7 +425,7 @@ impl SingleHashAggregateStream { ) -> SingleHashAggregateStateTransition { let SingleHashAggregateState::ReadingInput { mut hash_table, - spill_context, + mut spill_context, } = original_state else { return Self::break_with_internal_err( @@ -438,7 +444,17 @@ impl SingleHashAggregateStream { Poll::Ready(Some(Ok(batch))) => { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); - let result = hash_table.aggregate_batch(&batch); + let result = hash_table.aggregate_batch(&batch).and_then(|taken| { + match (taken, spill_context.as_deref_mut()) { + (None, _) => Ok(()), + (Some(taken), Some(spill_context)) => { + spill_context.spill_batch(taken) + } + (Some(_), None) => exec_err!( + "Single hash aggregate cannot spill group keys that overflow 32-bit offsets because temporary files are not enabled in the DiskManager" + ), + } + }); timer.done(); if let Err(e) = result {