[AURON #1404] Support for Spark 4.0.1 Compatibility in Auron.

auron-build.sh

@@ -117,10 +117,10 @@ while [[ $# -gt 0 ]]; do
                 SPARK_VER="$2"
                 if [ "$SPARK_VER" = "3.0" ] || [ "$SPARK_VER" = "3.1" ] \
                   || [ "$SPARK_VER" = "3.2" ] || [ "$SPARK_VER" = "3.3" ] \
-                  || [ "$SPARK_VER" = "3.4" ] || [ "$SPARK_VER" = "3.5" ]; then
+                  || [ "$SPARK_VER" = "3.4" ] || [ "$SPARK_VER" = "3.5" ] || [ "$SPARK_VER" = "4.0" ]; then
                   echo "Building for Spark $SPARK_VER"
                 else
-                  echo "ERROR: Invalid Spark version: $SPARK_VER. The currently supported versions are: 3.0 / 3.1 / 3.2 / 3.3 / 3.4 / 3.5."
+                  echo "ERROR: Invalid Spark version: $SPARK_VER. The currently supported versions are: 3.0 / 3.1 / 3.2 / 3.3 / 3.4 / 3.5 / 4.0."
                   exit 1
                 fi
                 shift 2

pom.xml

@@ -621,6 +621,47 @@
       </properties>
     </profile>
 
+    <profile>
+      <id>spark-4.0</id>
+      <properties>
+        <shimName>spark-4.0</shimName>
+        <shimPkg>spark-extension-shims-spark4</shimPkg>
+        <scalaTestVersion>3.2.10</scalaTestVersion>
+        <scalafmtVersion>3.9.9</scalafmtVersion>
+        <sparkVersion>4.0.1</sparkVersion>
+      </properties>
+      <build>
+        <plugins>
+          <plugin>
+            <groupId>org.apache.maven.plugins</groupId>
+            <artifactId>maven-enforcer-plugin</artifactId>
+            <version>3.4.1</version>
+            <executions>
+              <execution>
+                <id>enforce-java-scala-version</id>
+                <goals>
+                  <goal>enforce</goal>
+                </goals>
+                <configuration>
+                  <rules>
+                    <requireJavaVersion>
+                      <version>[17,)</version>
+                      <message>Spark 4.0 requires JDK 17 or higher to compile!</message>
+                    </requireJavaVersion>
+                    <requireProperty>
+                      <property>scalaVersion</property>
+                      <regex>^2.13.*</regex>
+                      <regexMessage>Spark 4.0 requires Scala 2.13 and is not compatible with Scala 2.12!</regexMessage>
+                    </requireProperty>
+                  </rules>
+                </configuration>
+              </execution>
+            </executions>
+          </plugin>
+        </plugins>
+      </build>
+    </profile>
+
     <profile>
       <id>jdk-8</id>
       <activation>

spark-extension-shims-spark4/pom.xml

@@ -0,0 +1,112 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+  ~ 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.
+  -->
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+  <modelVersion>4.0.0</modelVersion>
+  <parent>
+    <groupId>org.apache.auron</groupId>
+    <artifactId>auron-parent_${scalaVersion}</artifactId>
+    <version>${project.version}</version>
+    <relativePath>../pom.xml</relativePath>
+  </parent>
+
+  <artifactId>spark-extension-shims-spark4_${scalaVersion}</artifactId>
+  <packaging>jar</packaging>
+
+  <properties>
+    <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
+  </properties>
+
+  <dependencies>
+    <dependency>
+      <groupId>org.apache.auron</groupId>
+      <artifactId>auron-common_${scalaVersion}</artifactId>
+      <version>${project.version}</version>
+    </dependency>
+    <dependency>
+      <groupId>org.apache.auron</groupId>
+      <artifactId>spark-extension_${scalaVersion}</artifactId>
+      <version>${project.version}</version>
+    </dependency>
+    <dependency>
+      <groupId>org.scala-lang.modules</groupId>
+      <artifactId>scala-java8-compat_${scalaVersion}</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>org.apache.spark</groupId>
+      <artifactId>spark-core_${scalaVersion}</artifactId>
+      <scope>provided</scope>
+    </dependency>
+    <dependency>
+      <groupId>org.apache.spark</groupId>
+      <artifactId>spark-hive_${scalaVersion}</artifactId>
+      <scope>provided</scope>
+    </dependency>
+    <dependency>
+      <groupId>org.apache.spark</groupId>
+      <artifactId>spark-sql_${scalaVersion}</artifactId>
+      <scope>provided</scope>
+    </dependency>
+    <dependency>
+      <groupId>org.apache.arrow</groupId>
+      <artifactId>arrow-c-data</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>org.apache.arrow</groupId>
+      <artifactId>arrow-compression</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>org.apache.arrow</groupId>
+      <artifactId>arrow-memory-unsafe</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>org.apache.arrow</groupId>
+      <artifactId>arrow-vector</artifactId>
+    </dependency>
+
+    <dependency>
+      <groupId>net.bytebuddy</groupId>
+      <artifactId>byte-buddy</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>net.bytebuddy</groupId>
+      <artifactId>byte-buddy-agent</artifactId>
+    </dependency>
+
+    <dependency>
+      <groupId>org.apache.auron</groupId>
+      <artifactId>auron-common_${scalaVersion}</artifactId>
+      <version>${project.version}</version>
+      <type>test-jar</type>
+    </dependency>
+    <dependency>
+      <groupId>org.apache.spark</groupId>
+      <artifactId>spark-core_${scalaVersion}</artifactId>
+      <type>test-jar</type>
+    </dependency>
+    <dependency>
+      <groupId>org.apache.spark</groupId>
+      <artifactId>spark-catalyst_${scalaVersion}</artifactId>
+      <type>test-jar</type>
+    </dependency>
+    <dependency>
+      <groupId>org.apache.spark</groupId>
+      <artifactId>spark-sql_${scalaVersion}</artifactId>
+      <type>test-jar</type>
+    </dependency>
+  </dependencies>
+</project>

spark-extension-shims-spark4/src/main/java/org/apache/spark/sql/auron/ForceApplyShuffledHashJoinInjector.java

@@ -0,0 +1,43 @@
+/*
+ * 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.auron;
+
+import static net.bytebuddy.matcher.ElementMatchers.named;
+
+import net.bytebuddy.ByteBuddy;
+import net.bytebuddy.agent.ByteBuddyAgent;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.dynamic.ClassFileLocator;
+import net.bytebuddy.dynamic.loading.ClassLoadingStrategy;
+import net.bytebuddy.implementation.MethodDelegation;
+import net.bytebuddy.pool.TypePool;
+
+public class ForceApplyShuffledHashJoinInjector {
+    public static void inject() {
+        ByteBuddyAgent.install();
+        ClassLoader contextClassLoader = Thread.currentThread().getContextClassLoader();
+        TypeDescription typeDescription = TypePool.Default.of(contextClassLoader)
+                .describe("org.apache.spark.sql.catalyst.optimizer.JoinSelectionHelper")
+                .resolve();
+        new ByteBuddy()
+                .redefine(typeDescription, ClassFileLocator.ForClassLoader.of(contextClassLoader))
+                .method(named("forceApplyShuffledHashJoin"))
+                .intercept(MethodDelegation.to(ForceApplyShuffledHashJoinInterceptor.class))
+                .make()
+                .load(contextClassLoader, ClassLoadingStrategy.Default.INJECTION);
+    }
+}

spark-extension-shims-spark4/src/main/java/org/apache/spark/sql/auron/ForceApplyShuffledHashJoinInterceptor.java

@@ -0,0 +1,32 @@
+/*
+ * 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.auron;
+
+import net.bytebuddy.implementation.bind.annotation.Argument;
+import net.bytebuddy.implementation.bind.annotation.RuntimeType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class ForceApplyShuffledHashJoinInterceptor {
+    private static final Logger logger = LoggerFactory.getLogger(ForceApplyShuffledHashJoinInterceptor.class);
+
+    @RuntimeType
+    public static Object intercept(@Argument(0) Object conf) {
+        logger.debug("calling JoinSelectionHelper.forceApplyShuffledHashJoin() intercepted by auron");
+        return AuronConf.FORCE_SHUFFLED_HASH_JOIN.booleanConf();
+    }
+}

spark-extension-shims-spark4/src/main/java/org/apache/spark/sql/auron/ValidateSparkPlanApplyInterceptor.java

@@ -0,0 +1,34 @@
+/*
+ * 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.auron;
+
+import net.bytebuddy.implementation.bind.annotation.Argument;
+import net.bytebuddy.implementation.bind.annotation.RuntimeType;
+import org.apache.spark.sql.execution.SparkPlan;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class ValidateSparkPlanApplyInterceptor {
+    private static final Logger logger = LoggerFactory.getLogger(ValidateSparkPlanApplyInterceptor.class);
+
+    @RuntimeType
+    public static Object intercept(@Argument(0) Object plan) {
+        logger.debug("calling ValidateSparkPlan.apply() intercepted by auron");
+        InterceptedValidateSparkPlan$.MODULE$.validate((SparkPlan) plan);
+        return plan;
+    }
+}

spark-extension-shims-spark4/src/main/java/org/apache/spark/sql/auron/ValidateSparkPlanInjector.java

@@ -0,0 +1,59 @@
+/*
+ * 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.auron;
+
+import static net.bytebuddy.matcher.ElementMatchers.named;
+
+import net.bytebuddy.ByteBuddy;
+import net.bytebuddy.agent.ByteBuddyAgent;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.dynamic.ClassFileLocator;
+import net.bytebuddy.dynamic.loading.ClassLoadingStrategy;
+import net.bytebuddy.implementation.MethodDelegation;
+import net.bytebuddy.pool.TypePool;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class ValidateSparkPlanInjector {
+
+    private static final Logger logger = LoggerFactory.getLogger(ValidateSparkPlanInjector.class);
+    private static boolean injected = false;
+
+    public static synchronized void inject() {
+        if (injected) {
+            logger.warn("ValidateSparkPlan already injected, skipping.");
+            return;
+        }
+        try {
+            ByteBuddyAgent.install();
+            ClassLoader contextClassLoader = Thread.currentThread().getContextClassLoader();
+            TypeDescription typeDescription = TypePool.Default.of(contextClassLoader)
+                    .describe("org.apache.spark.sql.execution.adaptive.ValidateSparkPlan$")
+                    .resolve();
+            new ByteBuddy()
+                    .redefine(typeDescription, ClassFileLocator.ForClassLoader.of(contextClassLoader))
+                    .method(named("apply"))
+                    .intercept(MethodDelegation.to(ValidateSparkPlanApplyInterceptor.class))
+                    .make()
+                    .load(contextClassLoader, ClassLoadingStrategy.Default.INJECTION);
+            logger.info("Successfully injected ValidateSparkPlan.");
+            injected = true;
+        } catch (TypePool.Resolution.NoSuchTypeException e) {
+            logger.debug("No such type of ValidateSparkPlan", e);
+        }
+    }
+}