Merge remote-tracking branch 'apache/main' into plan-stability-regression

Author: andygrove

common/src/main/scala/org/apache/comet/CometConf.scala

@@ -457,8 +457,10 @@ object CometConf extends ShimCometConf {
   val COMET_EXPLAIN_VERBOSE_ENABLED: ConfigEntry[Boolean] =
     conf("spark.comet.explain.verbose.enabled")
       .doc(
-        "When this setting is enabled, Comet will provide a verbose tree representation of " +
-          "the extended information.")
+        "When this setting is enabled, Comet's extended explain output will provide the full " +
+          "query plan annotated with fallback reasons as well as a summary of how much of " +
+          "the plan was accelerated by Comet. When this setting is disabled, a list of fallback " +
+          "reasons will be provided instead.")
       .booleanConf
       .createWithDefault(false)
 

dev/benchmarks/tpcbench.py

@@ -59,6 +59,7 @@ def main(benchmark: str, data_path: str, query_path: str, iterations: int, outpu
 
     for iteration in range(0, iterations):
         print(f"Starting iteration {iteration} of {iterations}")
+        iter_start_time = time.time()
 
         # Determine which queries to run
         if query_num is not None:

dev/diffs/iceberg/1.8.1.diff

@@ -17,7 +17,7 @@ index 7327b38905d..7967109f039 100644
        exclude group: 'org.apache.avro', module: 'avro'
        // already shaded by Parquet
 diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml
-index 04ffa8f4edc..cc0099ccc93 100644
+index 04ffa8f4edc..a909cd552c1 100644
 --- a/gradle/libs.versions.toml
 +++ b/gradle/libs.versions.toml
 @@ -34,6 +34,7 @@ azuresdk-bom = "1.2.31"
@@ -300,10 +300,10 @@ index 00000000000..ddf6c7de5ae
 +}
 diff --git a/parquet/src/main/java/org/apache/iceberg/parquet/CometVectorizedParquetReader.java b/parquet/src/main/java/org/apache/iceberg/parquet/CometVectorizedParquetReader.java
 new file mode 100644
-index 00000000000..88b195b76a2
+index 00000000000..a3cba401827
 --- /dev/null
 +++ b/parquet/src/main/java/org/apache/iceberg/parquet/CometVectorizedParquetReader.java
-@@ -0,0 +1,255 @@
+@@ -0,0 +1,260 @@
 +/*
 + * Licensed to the Apache Software Foundation (ASF) under one
 + * or more contributor license agreements.  See the NOTICE file
@@ -446,6 +446,7 @@ index 00000000000..88b195b76a2
 +    private long valuesRead = 0;
 +    private T last = null;
 +    private final FileReader cometReader;
++    private ReadConf conf;
 +
 +    FileIterator(
 +        ReadConf conf,
@@ -470,6 +471,7 @@ index 00000000000..88b195b76a2
 +              length,
 +              fileEncryptionKey,
 +              fileAADPrefix);
++      this.conf = conf;
 +    }
 +
 +    private FileReader newCometReader(
@@ -556,6 +558,9 @@ index 00000000000..88b195b76a2
 +    public void close() throws IOException {
 +      model.close();
 +      cometReader.close();
++      if (conf != null && conf.reader() != null) {
++        conf.reader().close();
++      }
 +    }
 +  }
 +}

dev/diffs/iceberg/1.9.1.diff

@@ -291,10 +291,10 @@ index 00000000000..ddf6c7de5ae
 +}
 diff --git a/parquet/src/main/java/org/apache/iceberg/parquet/CometVectorizedParquetReader.java b/parquet/src/main/java/org/apache/iceberg/parquet/CometVectorizedParquetReader.java
 new file mode 100644
-index 00000000000..88b195b76a2
+index 00000000000..a3cba401827
 --- /dev/null
 +++ b/parquet/src/main/java/org/apache/iceberg/parquet/CometVectorizedParquetReader.java
-@@ -0,0 +1,255 @@
+@@ -0,0 +1,260 @@
 +/*
 + * Licensed to the Apache Software Foundation (ASF) under one
 + * or more contributor license agreements.  See the NOTICE file
@@ -437,6 +437,7 @@ index 00000000000..88b195b76a2
 +    private long valuesRead = 0;
 +    private T last = null;
 +    private final FileReader cometReader;
++    private ReadConf conf;
 +
 +    FileIterator(
 +        ReadConf conf,
@@ -461,6 +462,7 @@ index 00000000000..88b195b76a2
 +              length,
 +              fileEncryptionKey,
 +              fileAADPrefix);
++      this.conf = conf;
 +    }
 +
 +    private FileReader newCometReader(
@@ -547,6 +549,9 @@ index 00000000000..88b195b76a2
 +    public void close() throws IOException {
 +      model.close();
 +      cometReader.close();
++      if (conf != null && conf.reader() != null) {
++        conf.reader().close();
++      }
 +    }
 +  }
 +}

docs/source/user-guide/latest/configs.md

@@ -63,7 +63,7 @@ Comet provides the following configuration settings.
 | spark.comet.exec.union.enabled | Whether to enable union by default. | true |
 | spark.comet.exec.window.enabled | Whether to enable window by default. | true |
 | spark.comet.explain.native.enabled | When this setting is enabled, Comet will provide a tree representation of the native query plan before execution and again after execution, with metrics. | false |
-| spark.comet.explain.verbose.enabled | When this setting is enabled, Comet will provide a verbose tree representation of the extended information. | false |
+| spark.comet.explain.verbose.enabled | When this setting is enabled, Comet's extended explain output will provide the full query plan annotated with fallback reasons as well as a summary of how much of the plan was accelerated by Comet. When this setting is disabled, a list of fallback reasons will be provided instead. | false |
 | spark.comet.explainFallback.enabled | When this setting is enabled, Comet will provide logging explaining the reason(s) why a query stage cannot be executed natively. Set this to false to reduce the amount of logging. | false |
 | spark.comet.expression.allowIncompatible | Comet is not currently fully compatible with Spark for all expressions. Set this config to true to allow them anyway. For more information, refer to the Comet Compatibility Guide (https://datafusion.apache.org/comet/user-guide/compatibility.html). | false |
 | spark.comet.logFallbackReasons.enabled | When this setting is enabled, Comet will log warnings for all fallback reasons. | false |

native/core/src/execution/memory_pools/mod.rs

@@ -42,7 +42,8 @@ pub(crate) fn create_memory_pool(
     match memory_pool_config.pool_type {
         MemoryPoolType::Unified => {
             // Set Comet memory pool for native
-            let memory_pool = CometUnifiedMemoryPool::new(comet_task_memory_manager);
+            let memory_pool =
+                CometUnifiedMemoryPool::new(comet_task_memory_manager, task_attempt_id);
             Arc::new(TrackConsumersPool::new(
                 memory_pool,
                 NonZeroUsize::new(NUM_TRACKED_CONSUMERS).unwrap(),

native/core/src/execution/memory_pools/unified_pool.rs

@@ -40,6 +40,7 @@ use log::warn;
 pub struct CometUnifiedMemoryPool {
     task_memory_manager_handle: Arc<GlobalRef>,
     used: AtomicUsize,
+    task_attempt_id: i64,
 }
 
 impl Debug for CometUnifiedMemoryPool {
@@ -51,9 +52,13 @@ impl Debug for CometUnifiedMemoryPool {
 }
 
 impl CometUnifiedMemoryPool {
-    pub fn new(task_memory_manager_handle: Arc<GlobalRef>) -> CometUnifiedMemoryPool {
+    pub fn new(
+        task_memory_manager_handle: Arc<GlobalRef>,
+        task_attempt_id: i64,
+    ) -> CometUnifiedMemoryPool {
         Self {
             task_memory_manager_handle,
+            task_attempt_id,
             used: AtomicUsize::new(0),
         }
     }
@@ -82,7 +87,10 @@ impl Drop for CometUnifiedMemoryPool {
     fn drop(&mut self) {
         let used = self.used.load(Relaxed);
         if used != 0 {
-            warn!("CometUnifiedMemoryPool dropped with {used} bytes still reserved");
+            warn!(
+                "Task {} dropped CometUnifiedMemoryPool with {used} bytes still reserved",
+                self.task_attempt_id
+            );
         }
     }
 }
@@ -96,13 +104,20 @@ impl MemoryPool for CometUnifiedMemoryPool {
     }
 
     fn shrink(&self, _: &MemoryReservation, size: usize) {
-        self.release_to_spark(size)
-            .unwrap_or_else(|_| panic!("Failed to release {size} bytes"));
+        if let Err(e) = self.release_to_spark(size) {
+            panic!(
+                "Task {} failed to return {size} bytes to Spark: {e:?}",
+                self.task_attempt_id
+            );
+        }
         if let Err(prev) = self
             .used
             .fetch_update(Relaxed, Relaxed, |old| old.checked_sub(size))
         {
-            panic!("overflow when releasing {size} of {prev} bytes");
+            panic!(
+                "Task {} overflow when releasing {size} of {prev} bytes",
+                self.task_attempt_id
+            );
         }
     }
 
@@ -116,7 +131,8 @@ impl MemoryPool for CometUnifiedMemoryPool {
                 self.release_to_spark(acquired as usize)?;
 
                 return Err(resources_datafusion_err!(
-                    "Failed to acquire {} bytes, only got {}. Reserved: {}",
+                    "Task {} failed to acquire {} bytes, only got {}. Reserved: {}",
+                    self.task_attempt_id,
                     additional,
                     acquired,
                     self.reserved()
@@ -127,7 +143,8 @@ impl MemoryPool for CometUnifiedMemoryPool {
                 .fetch_update(Relaxed, Relaxed, |old| old.checked_add(acquired as usize))
             {
                 return Err(resources_datafusion_err!(
-                    "Failed to acquire {} bytes due to overflow. Reserved: {}",
+                    "Task {} failed to acquire {} bytes due to overflow. Reserved: {}",
+                    self.task_attempt_id,
                     additional,
                     prev
                 ));

spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java

@@ -40,22 +40,36 @@ public class CometTaskMemoryManager {
   /** The id uniquely identifies the native plan this memory manager is associated to */
   private final long id;
 
+  private final long taskAttemptId;
+
   public final TaskMemoryManager internal;
   private final NativeMemoryConsumer nativeMemoryConsumer;
   private final AtomicLong used = new AtomicLong();
 
-  public CometTaskMemoryManager(long id) {
+  public CometTaskMemoryManager(long id, long taskAttemptId) {
     this.id = id;
+    this.taskAttemptId = taskAttemptId;
     this.internal = TaskContext$.MODULE$.get().taskMemoryManager();
     this.nativeMemoryConsumer = new NativeMemoryConsumer();
   }
 
   // Called by Comet native through JNI.
   // Returns the actual amount of memory (in bytes) granted.
   public long acquireMemory(long size) {
+    if (logger.isTraceEnabled()) {
+      logger.trace("Task {} requested {} bytes", taskAttemptId, size);
+    }
     long acquired = internal.acquireExecutionMemory(size, nativeMemoryConsumer);
-    used.addAndGet(acquired);
+    long newUsed = used.addAndGet(acquired);
     if (acquired < size) {
+      logger.warn(
+          "Task {} requested {} bytes but only received {} bytes. Current allocation is {} and "
+              + "the total memory consumption is {} bytes.",
+          taskAttemptId,
+          size,
+          acquired,
+          newUsed,
+          internal.getMemoryConsumptionForThisTask());
       // If memory manager is not able to acquire the requested size, log memory usage
       internal.showMemoryUsage();
     }
@@ -64,10 +78,16 @@ public long acquireMemory(long size) {
 
   // Called by Comet native through JNI
   public void releaseMemory(long size) {
+    if (logger.isTraceEnabled()) {
+      logger.trace("Task {} released {} bytes", taskAttemptId, size);
+    }
     long newUsed = used.addAndGet(-size);
     if (newUsed < 0) {
       logger.error(
-          "Used memory is negative: " + newUsed + " after releasing memory chunk of: " + size);
+          "Task {} used memory is negative ({}) after releasing {} bytes",
+          taskAttemptId,
+          newUsed,
+          size);
     }
     internal.releaseExecutionMemory(size, nativeMemoryConsumer);
   }

spark/src/main/scala/org/apache/comet/CometExecIterator.scala

@@ -68,7 +68,8 @@ class CometExecIterator(
   private val memoryMXBean = ManagementFactory.getMemoryMXBean
   private val nativeLib = new Native()
   private val nativeUtil = new NativeUtil()
-  private val cometTaskMemoryManager = new CometTaskMemoryManager(id)
+  private val taskAttemptId = TaskContext.get().taskAttemptId
+  private val cometTaskMemoryManager = new CometTaskMemoryManager(id, taskAttemptId)
   private val cometBatchIterators = inputs.map { iterator =>
     new CometBatchIterator(iterator, nativeUtil)
   }.toArray
@@ -116,7 +117,7 @@ class CometExecIterator(
       memoryPoolType = COMET_EXEC_MEMORY_POOL_TYPE.get(),
       memoryLimit,
       memoryLimitPerTask,
-      taskAttemptId = TaskContext.get().taskAttemptId,
+      taskAttemptId,
       debug = COMET_DEBUG_ENABLED.get(),
       explain = COMET_EXPLAIN_NATIVE_ENABLED.get(),
       tracingEnabled)
@@ -175,6 +176,11 @@ class CometExecIterator(
         })
     } catch {
       case e: CometNativeException =>
+        // it is generally considered bad practice to log and then rethrow an
+        // exception, but it really helps debugging to be able to see which task
+        // threw the exception, so we log the exception with taskAttemptId here
+        logError(s"Native execution for task $taskAttemptId failed", e)
+
         val fileNotFoundPattern: Regex =
           ("""^External: Object at location (.+?) not found: No such file or directory """ +
             """\(os error \d+\)$""").r