build: Add spark-4.1 profile and shims

.github/workflows/pr_build_linux.yml

@@ -98,6 +98,11 @@ jobs:
             java_version: "17"
             maven_opts: "-Pspark-4.0"
             scan_impl: "native_comet"
+
+          - name: "Spark 4.1, JDK 17"
+            java_version: "17"
+            maven_opts: "-Pspark-4.1"
+            scan_impl: "native_comet"
         suite:
           - name: "fuzz"
             value: |

.github/workflows/spark_sql_test.yml

@@ -50,7 +50,7 @@ jobs:
     strategy:
       matrix:
         os: [ubuntu-24.04]
-        spark-version: [{short: '3.4', full: '3.4.3', java: 11}, {short: '3.5', full: '3.5.7', java: 11}, {short: '4.0', full: '4.0.1', java: 17}]
+        spark-version: [{short: '3.4', full: '3.4.3', java: 11}, {short: '3.5', full: '3.5.7', java: 11}, {short: '4.0', full: '4.0.1', java: 17}, {short: '4.1', full: '4.1.0', java: 17}]
         module:
           - {name: "catalyst", args1: "catalyst/test", args2: ""}
           - {name: "sql_core-1", args1: "", args2: sql/testOnly * -- -l org.apache.spark.tags.ExtendedSQLTest -l org.apache.spark.tags.SlowSQLTest}
@@ -59,10 +59,12 @@ jobs:
           - {name: "sql_hive-1", args1: "", args2: "hive/testOnly * -- -l org.apache.spark.tags.ExtendedHiveTest -l org.apache.spark.tags.SlowHiveTest"}
           - {name: "sql_hive-2", args1: "", args2: "hive/testOnly * -- -n org.apache.spark.tags.ExtendedHiveTest"}
           - {name: "sql_hive-3", args1: "", args2: "hive/testOnly * -- -n org.apache.spark.tags.SlowHiveTest"}
-        # Skip sql_hive-1 for Spark 4.0 due to https://github.com/apache/datafusion-comet/issues/2946
+        # Skip sql_hive-1 for Spark 4.0+ due to https://github.com/apache/datafusion-comet/issues/2946
         exclude:
           - spark-version: {short: '4.0', full: '4.0.1', java: 17}
             module: {name: "sql_hive-1", args1: "", args2: "hive/testOnly * -- -l org.apache.spark.tags.ExtendedHiveTest -l org.apache.spark.tags.SlowHiveTest"}
+          - spark-version: {short: '4.1', full: '4.1.0', java: 17}
+            module: { name: "sql_hive-1", args1: "", args2: "hive/testOnly * -- -l org.apache.spark.tags.ExtendedHiveTest -l org.apache.spark.tags.SlowHiveTest" }
       fail-fast: false
     name: spark-sql-${{ matrix.module.name }}/${{ matrix.os }}/spark-${{ matrix.spark-version.full }}/java-${{ matrix.spark-version.java }}
     runs-on: ${{ matrix.os }}

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

@@ -50,6 +50,7 @@ use only and should not be used in production yet.
 | Spark Version | Java Version | Scala Version | Comet Tests in CI | Spark SQL Tests in CI |
 | ------------- | ------------ | ------------- | ----------------- | --------------------- |
 | 4.0.1         | 17           | 2.13          | Yes               | Yes                   |
+| 4.1.0         | 17           | 2.13          | Yes               | Yes                   |
 
 Note that Comet may not fully work with proprietary forks of Apache Spark such as the Spark versions offered by
 Cloud Service Providers.

pom.xml

@@ -642,8 +642,29 @@ under the License.
         <parquet.version>1.15.2</parquet.version>
         <semanticdb.version>4.13.6</semanticdb.version>
         <slf4j.version>2.0.16</slf4j.version>
-        <shims.majorVerSrc>spark-4.0</shims.majorVerSrc>
-        <shims.minorVerSrc>not-needed-yet</shims.minorVerSrc>
+        <shims.majorVerSrc>spark-4.x</shims.majorVerSrc>
+        <shims.minorVerSrc>spark-4.0</shims.minorVerSrc>
+        <!-- Use jdk17 by default -->
+        <java.version>17</java.version>
+        <maven.compiler.source>${java.version}</maven.compiler.source>
+        <maven.compiler.target>${java.version}</maven.compiler.target>
+      </properties>
+    </profile>
+
+    <profile>
+      <!-- FIXME: this is WIP. Tests may fail https://github.com/apache/datafusion-comet/issues/551 -->
+      <id>spark-4.1</id>
+      <properties>
+        <!-- Use Scala 2.13 by default -->
+        <scala.version>2.13.17</scala.version>
+        <scala.binary.version>2.13</scala.binary.version>
+        <spark.version>4.1.0</spark.version>
+        <spark.version.short>4.1</spark.version.short>
+        <parquet.version>1.16.0</parquet.version>
+        <semanticdb.version>4.13.9</semanticdb.version>
+        <slf4j.version>2.0.17</slf4j.version>
+        <shims.majorVerSrc>spark-4.x</shims.majorVerSrc>
+        <shims.minorVerSrc>spark-4.1</shims.minorVerSrc>
         <!-- Use jdk17 by default -->
         <java.version>17</java.version>
         <maven.compiler.source>${java.version}</maven.compiler.source>

spark/pom.xml

@@ -256,6 +256,32 @@ under the License.
         </dependency>
       </dependencies>
     </profile>
+
+    <profile>
+      <id>spark-4.1</id>
+      <dependencies>
+        <dependency>
+          <groupId>org.apache.iceberg</groupId>
+          <!-- TODO: Upgrade after iceberg-spark-runtime-4.1_2.13 release -->
+          <artifactId>iceberg-spark-runtime-4.0_${scala.binary.version}</artifactId>
+          <version>1.10.0</version>
+          <scope>test</scope>
+        </dependency>
+        <!-- Jetty 11.x for Spark 4.1 (jakarta.servlet) -->
+        <dependency>
+          <groupId>org.eclipse.jetty</groupId>
+          <artifactId>jetty-server</artifactId>
+          <version>11.0.24</version>
+          <scope>test</scope>
+        </dependency>
+        <dependency>
+          <groupId>org.eclipse.jetty</groupId>
+          <artifactId>jetty-servlet</artifactId>
+          <version>11.0.24</version>
+          <scope>test</scope>
+        </dependency>
+      </dependencies>
+    </profile>
   </profiles>
 
   <build>

spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometBypassMergeSortShuffleWriter.java

@@ -41,7 +41,6 @@
 import org.apache.spark.memory.TaskMemoryManager;
 import org.apache.spark.network.shuffle.checksum.ShuffleChecksumHelper;
 import org.apache.spark.scheduler.MapStatus;
-import org.apache.spark.scheduler.MapStatus$;
 import org.apache.spark.serializer.SerializerInstance;
 import org.apache.spark.shuffle.ShuffleWriteMetricsReporter;
 import org.apache.spark.shuffle.ShuffleWriter;
@@ -172,7 +171,7 @@ public void write(Iterator<Product2<K, V>> records) throws IOException {
                 .commitAllPartitions(ShuffleChecksumHelper.EMPTY_CHECKSUM_VALUE)
                 .getPartitionLengths();
         mapStatus =
-            MapStatus$.MODULE$.apply(blockManager.shuffleServerId(), partitionLengths, mapId);
+            new MapStatusBuilder(blockManager.shuffleServerId(), partitionLengths, mapId).build();
         return;
       }
       final long openStartTime = System.nanoTime();
@@ -261,7 +260,8 @@ public void write(Iterator<Product2<K, V>> records) throws IOException {
 
       // TODO: We probably can move checksum generation here when concatenating partition files
       partitionLengths = writePartitionedData(mapOutputWriter);
-      mapStatus = MapStatus$.MODULE$.apply(blockManager.shuffleServerId(), partitionLengths, mapId);
+      mapStatus =
+          new MapStatusBuilder(blockManager.shuffleServerId(), partitionLengths, mapId).build();
     } catch (Exception e) {
       try {
         mapOutputWriter.abort(e);

spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometUnsafeShuffleWriter.java

@@ -50,7 +50,6 @@
 import org.apache.spark.network.shuffle.checksum.ShuffleChecksumHelper;
 import org.apache.spark.network.util.LimitedInputStream;
 import org.apache.spark.scheduler.MapStatus;
-import org.apache.spark.scheduler.MapStatus$;
 import org.apache.spark.serializer.SerializationStream;
 import org.apache.spark.serializer.SerializerInstance;
 import org.apache.spark.shuffle.BaseShuffleHandle;
@@ -288,7 +287,8 @@ void closeAndWriteOutput() throws IOException {
         }
       }
     }
-    mapStatus = MapStatus$.MODULE$.apply(blockManager.shuffleServerId(), partitionLengths, mapId);
+    mapStatus =
+        new MapStatusBuilder(blockManager.shuffleServerId(), partitionLengths, mapId).build();
   }
 
   @VisibleForTesting

spark/src/main/scala/org/apache/comet/serde/aggregates.scala

@@ -30,7 +30,7 @@ import org.apache.comet.CometConf
 import org.apache.comet.CometConf.COMET_EXEC_STRICT_FLOATING_POINT
 import org.apache.comet.CometSparkSessionExtensions.withInfo
 import org.apache.comet.serde.QueryPlanSerde.{evalModeToProto, exprToProto, serializeDataType}
-import org.apache.comet.shims.CometEvalModeUtil
+import org.apache.comet.shims.{CometEvalModeUtil, CometSumShim}
 
 object CometMin extends CometAggregateExpressionSerde[Min] {
 
@@ -211,10 +211,10 @@ object CometAverage extends CometAggregateExpressionSerde[Average] {
   }
 }
 
-object CometSum extends CometAggregateExpressionSerde[Sum] {
+object CometSum extends CometAggregateExpressionSerde[Sum] with CometSumShim {
 
   override def getSupportLevel(sum: Sum): SupportLevel = {
-    sum.evalMode match {
+    sparkEvalMode(sum) match {
       case EvalMode.ANSI if !sum.dataType.isInstanceOf[DecimalType] =>
         Incompatible(Some("ANSI mode for non decimal inputs is not supported"))
       case EvalMode.TRY if !sum.dataType.isInstanceOf[DecimalType] =>
@@ -243,7 +243,8 @@ object CometSum extends CometAggregateExpressionSerde[Sum] {
       val builder = ExprOuterClass.Sum.newBuilder()
       builder.setChild(childExpr.get)
       builder.setDatatype(dataType.get)
-      builder.setEvalMode(evalModeToProto(CometEvalModeUtil.fromSparkEvalMode(sum.evalMode)))
+      builder.setEvalMode(
+        evalModeToProto(CometEvalModeUtil.fromSparkEvalMode(sparkEvalMode(sum))))
 
       Some(
         ExprOuterClass.AggExpr

spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleManager.scala

@@ -71,8 +71,14 @@ class CometShuffleManager(conf: SparkConf) extends ShuffleManager with Logging {
           case 2 =>
             c.newInstance(conf, null).asInstanceOf[IndexShuffleBlockResolver]
           case 3 =>
-            c.newInstance(conf, null, Collections.emptyMap())
-              .asInstanceOf[IndexShuffleBlockResolver]
+            val parameterTypes = c.getParameters.map(_.getType)
+            if (parameterTypes(2) == classOf[java.util.Map[Int, OpenHashSet[Long]]]) {
+              c.newInstance(conf, null, Collections.emptyMap())
+                .asInstanceOf[IndexShuffleBlockResolver]
+            } else {
+              c.newInstance(conf, null, new ConcurrentHashMap[Int, OpenHashSet[Long]]())
+                .asInstanceOf[IndexShuffleBlockResolver]
+            }
         }
       }
       .head

spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/MapStatusBuilder.scala

@@ -0,0 +1,31 @@
+/*
+ * 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 org.apache.spark.scheduler.MapStatus
+import org.apache.spark.storage.BlockManagerId
+
+class MapStatusBuilder(loc: BlockManagerId, uncompressedSizes: Array[Long], mapTaskId: Long) {
+
+  // This works for both Spark 4.1 and previous versions
+  def build(): MapStatus = {
+    MapStatus(loc, uncompressedSizes, mapTaskId)
+  }
+}

spark/src/main/spark-3.x/org/apache/comet/shims/CometSumShim.scala

@@ -0,0 +1,30 @@
+/*
+ * 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.shims
+
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate.Sum
+
+/**
+ * `CometSumShim` acts as a shim for parsing expression Sum from different Spark versions.
+ */
+trait CometSumShim extends CometExprShim {
+  protected def sparkEvalMode(s: Sum): EvalMode.Value = s.evalMode
+}

spark/src/main/spark-4.0/org/apache/comet/shims/CometSumShim.scala

@@ -0,0 +1,30 @@
+/*
+ * 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.shims
+
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate.Sum
+
+/**
+ * `CometSumShim` acts as a shim for parsing expression Sum from different Spark versions.
+ */
+trait CometSumShim extends CometExprShim {
+  protected def sparkEvalMode(s: Sum): EvalMode.Value = s.evalMode
+}

spark/src/main/spark-4.0/org/apache/comet/shims/ShimSQLConf.scala

@@ -20,8 +20,15 @@
 package org.apache.comet.shims
 
 import org.apache.spark.sql.internal.LegacyBehaviorPolicy
+import org.apache.spark.sql.internal.SQLConf
 
 trait ShimSQLConf {
   protected val LEGACY = LegacyBehaviorPolicy.LEGACY
   protected val CORRECTED = LegacyBehaviorPolicy.CORRECTED
+
+  def getBinaryOutputStyle: Option[SQLConf.BinaryOutputStyle.Value] = {
+    SQLConf.get
+      .getConf(SQLConf.BINARY_OUTPUT_STYLE)
+      .map(SQLConf.BinaryOutputStyle.withName)
+  }
 }

spark/src/main/spark-4.1/org/apache/comet/shims/CometSumShim.scala

@@ -0,0 +1,30 @@
+/*
+ * 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.shims
+
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate.Sum
+
+/**
+ * `CometSumShim` acts as a shim for parsing expression Sum from different Spark versions.
+ */
+trait CometSumShim extends CometExprShim {
+  protected def sparkEvalMode(s: Sum): EvalMode.Value = s.evalContext.evalMode
+}

spark/src/main/spark-4.1/org/apache/comet/shims/ShimSQLConf.scala

@@ -0,0 +1,33 @@
+/*
+ * 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.shims
+
+import org.apache.spark.sql.internal.LegacyBehaviorPolicy
+import org.apache.spark.sql.internal.SQLConf
+
+trait ShimSQLConf {
+  protected val LEGACY = LegacyBehaviorPolicy.LEGACY
+  protected val CORRECTED = LegacyBehaviorPolicy.CORRECTED
+
+  def getBinaryOutputStyle(): Option[SQLConf.BinaryOutputStyle.Value] = {
+    SQLConf.get
+      .getConf(SQLConf.BINARY_OUTPUT_STYLE)
+  }
+}

spark/src/main/spark-4.0/org/apache/comet/shims/CometExprShim.scala → spark/src/main/spark-4.x/org/apache/comet/shims/CometExprShim.scala

@@ -34,14 +34,12 @@ import org.apache.comet.serde.QueryPlanSerde.exprToProtoInternal
 /**
  * `CometExprShim` acts as a shim for parsing expressions from different Spark versions.
  */
-trait CometExprShim extends CommonStringExprs {
+trait CometExprShim extends CommonStringExprs with ShimSQLConf {
   protected def evalMode(c: Cast): CometEvalMode.Value =
     CometEvalModeUtil.fromSparkEvalMode(c.evalMode)
 
   protected def binaryOutputStyle: BinaryOutputStyle = {
-    SQLConf.get
-      .getConf(SQLConf.BINARY_OUTPUT_STYLE)
-      .map(SQLConf.BinaryOutputStyle.withName) match {
+    getBinaryOutputStyle match {
       case Some(SQLConf.BinaryOutputStyle.UTF8) => BinaryOutputStyle.UTF8
       case Some(SQLConf.BinaryOutputStyle.BASIC) => BinaryOutputStyle.BASIC
       case Some(SQLConf.BinaryOutputStyle.BASE64) => BinaryOutputStyle.BASE64