Merge branch 'apache:main' into main

Author: coderfender

.github/workflows/benchmark.yml renamed to .github/workflows/benchmark-tpcds.yml

@@ -63,7 +63,8 @@ jobs:
           key: ${{ runner.os }}-java-maven-${{ hashFiles('**/pom.xml') }}
           restore-keys: |
             ${{ runner.os }}-java-maven-
-
+      - name: Build Comet
+        run: make release
       - name: Cache TPC-DS generated data
         id: cache-tpcds-sf-1
         uses: actions/cache@v4
@@ -76,17 +77,6 @@ jobs:
         with:
           repository: databricks/tpcds-kit
           path: ./tpcds-kit
-      - name: Build Comet
-        run: make release
-      - name: Upload Comet native lib
-        uses: actions/upload-artifact@v4
-        with:
-          name: libcomet-${{ github.run_id }}
-          path: |
-            native/target/release/libcomet.so
-            native/target/release/libcomet.dylib
-          retention-days: 1 # remove the artifact after 1 day, only valid for this workflow
-          overwrite: true
       - name: Build tpcds-kit
         if: steps.cache-tpcds-sf-1.outputs.cache-hit != 'true'
         run: |
@@ -132,11 +122,8 @@ jobs:
           path: ./tpcds-sf-1
           key: tpcds-${{ hashFiles('.github/workflows/benchmark.yml') }}
           fail-on-cache-miss: true # it's always be cached as it should be generated by pre-step if not existed
-      - name: Download Comet native lib
-        uses: actions/download-artifact@v5
-        with:
-          name: libcomet-${{ github.run_id }}
-          path: native/target/release
+      - name: Build Comet
+        run: make release
       - name: Run TPC-DS queries (Sort merge join)
         if: matrix.join == 'sort_merge'
         run: |

.github/workflows/benchmark-tpch.yml

@@ -71,15 +71,6 @@ jobs:
           key: tpch-${{ hashFiles('.github/workflows/benchmark-tpch.yml') }}
       - name: Build Comet
         run: make release
-      - name: Upload Comet native lib
-        uses: actions/upload-artifact@v4
-        with:
-          name: libcomet-${{ github.run_id }}
-          path: |
-            native/target/release/libcomet.so
-            native/target/release/libcomet.dylib
-          retention-days: 1 # remove the artifact after 1 day, only valid for this workflow
-          overwrite: true
       - name: Generate TPC-H (SF=1) table data
         if: steps.cache-tpch-sf-1.outputs.cache-hit != 'true'
         run: |
@@ -115,11 +106,8 @@ jobs:
           path: ./tpch
           key: tpch-${{ hashFiles('.github/workflows/benchmark-tpch.yml') }}
           fail-on-cache-miss: true # it's always be cached as it should be generated by pre-step if not existed
-      - name: Download Comet native lib
-        uses: actions/download-artifact@v5
-        with:
-          name: libcomet-${{ github.run_id }}
-          path: native/target/release
+      - name: Build Comet
+        run: make release
       - name: Run TPC-H queries
         run: |
           SPARK_HOME=`pwd` SPARK_TPCH_DATA=`pwd`/tpch/sf1_parquet ./mvnw -B -Prelease -Dsuites=org.apache.spark.sql.CometTPCHQuerySuite test

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/jni_api.rs

@@ -41,6 +41,8 @@ use datafusion::{
 };
 use datafusion_comet_proto::spark_operator::Operator;
 use datafusion_spark::function::bitwise::bit_get::SparkBitGet;
+use datafusion_spark::function::datetime::date_add::SparkDateAdd;
+use datafusion_spark::function::datetime::date_sub::SparkDateSub;
 use datafusion_spark::function::hash::sha2::SparkSha2;
 use datafusion_spark::function::math::expm1::SparkExpm1;
 use datafusion_spark::function::string::char::CharFunc;
@@ -303,6 +305,8 @@ fn prepare_datafusion_session_context(
     session_ctx.register_udf(ScalarUDF::new_from_impl(SparkSha2::default()));
     session_ctx.register_udf(ScalarUDF::new_from_impl(CharFunc::default()));
     session_ctx.register_udf(ScalarUDF::new_from_impl(SparkBitGet::default()));
+    session_ctx.register_udf(ScalarUDF::new_from_impl(SparkDateAdd::default()));
+    session_ctx.register_udf(ScalarUDF::new_from_impl(SparkDateSub::default()));
 
     // Must be the last one to override existing functions with the same name
     datafusion_comet_spark_expr::register_all_comet_functions(&mut session_ctx)?;

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
                 ));