diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 5ff917f1a8a..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 @@ -527,6 +528,8 @@ jobs: org.apache.comet.exec.CometAggregateSuite 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 @@ -602,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 c79bd3116ba..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 @@ -175,6 +176,8 @@ jobs: org.apache.comet.exec.CometAggregateSuite 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 @@ -250,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 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/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/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 12e9ac8a5ee..ffa9668bb92 100644 --- a/native/common/src/lib.rs +++ b/native/common/src/lib.rs @@ -15,7 +15,9 @@ // specific language governing permissions and limitations // under the License. +pub mod cancellation; 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/jni_api.rs b/native/core/src/execution/jni_api.rs index dbd6ffb69d3..9002e83194b 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, @@ -3660,4 +3812,416 @@ 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, + PlanCancellation::new(), + 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}"), + } + } + } + + /// 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"); + } + } + } +} + +#[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/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/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/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/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/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/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/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(); 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/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/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/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/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 { 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(()) +} 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); 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/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/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index f29e5086abd..af825520317 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( + "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/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/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/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/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/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/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/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..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 @@ -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,8 @@ class CometNativeShuffleWriter[K, V]( OperatorOuterClass.Operator .newBuilder() .setShuffleWriter(shuffleWriterBuilder) - .addChildren(spec.childNativeOp) + .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 be91caa0ec9..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. @@ -86,7 +87,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..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 @@ -986,9 +998,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 +1009,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..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 @@ -767,10 +824,11 @@ 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) { + 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/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..9b166717f55 --- /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, "spark.sql.ansi.enabled" -> "false") { + 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") + } + } + } +} 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/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")) + } + } + } + } + } +} 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") { 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..d0f8aabd380 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/CometNativeTaskKillSuite.scala @@ -0,0 +1,220 @@ +/* + * 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]) + 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) + } + } + + 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) {} + } +} 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) + } +} 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..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 /** @@ -255,4 +260,115 @@ 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) + } + + 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) + } + } + } + } + } + } }