minimize load library code to pass github ci
diff --git a/Cargo.toml b/Cargo.toml index 8d3e883..5424ab2 100644 --- a/Cargo.toml +++ b/Cargo.toml
@@ -23,8 +23,8 @@ ] [profile.release] -lto = true -codegen-units = 1 +#lto = true +#codegen-units = 1 [profile.dev] overflow-checks = false
diff --git a/spark-extension/src/main/java/org/apache/spark/sql/blaze/JniBridge.java b/spark-extension/src/main/java/org/apache/spark/sql/blaze/JniBridge.java index 47eb784..6a65191 100644 --- a/spark-extension/src/main/java/org/apache/spark/sql/blaze/JniBridge.java +++ b/spark-extension/src/main/java/org/apache/spark/sql/blaze/JniBridge.java
@@ -29,10 +29,6 @@ public class JniBridge { public static final ConcurrentHashMap<String, Object> resourcesMap = new ConcurrentHashMap<>(); - static { - JniLoader.get().ensureLoaded(); - } - public static native long callNative( byte[] taskDefinition, long tokioPoolSize,
diff --git a/spark-extension/src/main/java/org/apache/spark/sql/blaze/JniLoader.java b/spark-extension/src/main/java/org/apache/spark/sql/blaze/JniLoader.java deleted file mode 100644 index a99dd42..0000000 --- a/spark-extension/src/main/java/org/apache/spark/sql/blaze/JniLoader.java +++ /dev/null
@@ -1,89 +0,0 @@ -/* - * 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.blaze; - -import java.io.File; -import java.io.FileNotFoundException; -import java.io.IOException; -import java.io.InputStream; -import java.nio.file.Files; -import java.nio.file.StandardCopyOption; -import java.util.*; -import org.apache.arrow.c.jni.JniWrapper; - -public class JniLoader { - private static final JniLoader INSTANCE = new JniLoader(Collections.singletonList("blaze")); - - public static JniLoader get() { - return INSTANCE; - } - - private final Set<String> librariesToLoad; - - private JniLoader(List<String> libraryNames) { - librariesToLoad = new HashSet<>(libraryNames); - } - - private boolean finished() { - return librariesToLoad.isEmpty(); - } - - /** If required JNI libraries are not loaded, then load them. */ - public void ensureLoaded() { - if (finished()) { - return; - } - loadRemaining(); - } - - private synchronized void loadRemaining() { - // The method is protected by a mutex via synchronized, if more than one thread - // race to call - // loadRemaining, at same time only one will do the actual loading and the - // others will wait for - // the mutex to be acquired then check on the remaining list: if there are - // libraries that were not - // successfully loaded then the mutex owner will try to load them again. - if (finished()) { - return; - } - List<String> libs = new ArrayList<>(librariesToLoad); - for (String lib : libs) { - load(lib); - librariesToLoad.remove(lib); - } - } - - private void load(String name) { - final String libraryToLoad = System.mapLibraryName(name); - try { - File temp = - File.createTempFile("jnilib-", ".tmp", new File(System.getProperty("java.io.tmpdir"))); - try (final InputStream is = - JniWrapper.class.getClassLoader().getResourceAsStream(libraryToLoad)) { - if (is == null) { - throw new FileNotFoundException(libraryToLoad); - } - Files.copy(is, temp.toPath(), StandardCopyOption.REPLACE_EXISTING); - System.load(temp.getAbsolutePath()); - } - } catch (IOException e) { - throw new IllegalStateException("error loading native libraries: " + e); - } - } -}
diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/blaze/BlazeSparkSessionExtension.scala b/spark-extension/src/main/scala/org/apache/spark/sql/blaze/BlazeSparkSessionExtension.scala index 6cf628b..7dada47 100644 --- a/spark-extension/src/main/scala/org/apache/spark/sql/blaze/BlazeSparkSessionExtension.scala +++ b/spark-extension/src/main/scala/org/apache/spark/sql/blaze/BlazeSparkSessionExtension.scala
@@ -17,9 +17,17 @@ package org.apache.spark.sql.blaze +import java.io.File +import java.io.FileNotFoundException +import java.io.InputStream +import java.io.IOException +import java.nio.file.Files +import java.nio.file.StandardCopyOption + import scala.annotation.tailrec import scala.collection.mutable.ArrayBuffer +import org.apache.arrow.c.jni.JniWrapper import org.apache.spark.sql.SparkSessionExtensions import org.apache.spark.SparkEnv import org.apache.spark.internal.Logging @@ -59,12 +67,12 @@ import org.apache.spark.sql.execution.ColumnarRule import org.apache.spark.sql.execution.RowToColumnarExec import org.apache.spark.sql.execution.UnaryExecNode -import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec import org.apache.spark.sql.execution.adaptive.BroadcastQueryStageExec import org.apache.spark.sql.execution.adaptive.CustomShuffleReaderExec import org.apache.spark.sql.execution.adaptive.ShuffleQueryStageExec import org.apache.spark.sql.execution.exchange.ReusedExchangeExec import org.apache.spark.sql.vectorized.ColumnarBatch +import org.apache.spark.util.Utils class BlazeSparkSessionExtension extends (SparkSessionExtensions => Unit) with Logging { override def apply(extensions: SparkSessionExtensions): Unit = { @@ -82,6 +90,25 @@ } } +object BlazeSparkSessionExtension { + var blazeNativeLibraryLoaded: Boolean = false + + def loadBlazeNativeLibrary(): Unit = + this.synchronized { + if (!blazeNativeLibraryLoaded) { + val libraryName = System.mapLibraryName("blaze") + Utils.tryWithResource( + classOf[BlazeSparkSessionExtension].getClassLoader.getResourceAsStream(libraryName)) { + is => + val temp = Files.createTempFile("jnilib-", ".tmp") + Files.copy(is, temp.toAbsolutePath, StandardCopyOption.REPLACE_EXISTING) + System.load(temp.toAbsolutePath.toString) + } + blazeNativeLibraryLoaded = true + } + } +} + case class BlazeQueryStagePrepOverrides(sparkSession: SparkSession) extends Rule[SparkPlan] with Logging {
diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/blaze/NativeSupports.scala b/spark-extension/src/main/scala/org/apache/spark/sql/blaze/NativeSupports.scala index 5d43714..7741f1e 100644 --- a/spark-extension/src/main/scala/org/apache/spark/sql/blaze/NativeSupports.scala +++ b/spark-extension/src/main/scala/org/apache/spark/sql/blaze/NativeSupports.scala
@@ -120,6 +120,9 @@ partition: Partition, context: TaskContext): Long = { + // make sure native library is loaded + BlazeSparkSessionExtension.loadBlazeNativeLibrary() + // do not use context.partitionId since it is not correct in Union plans. val partitionId = PartitionId .newBuilder()