sarutak commented on code in PR #59313: URL: https://github.com/apache/spark/pull/59313#discussion_r4236982595
########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala: ########## @@ -0,0 +1,655 @@ +/* + * 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.execution.datasources.v2.ffi + +import java.lang.ref.{Cleaner, WeakReference} +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.atomic.AtomicLong + +import scala.jdk.CollectionConverters._ +import scala.util.control.NonFatal + +import org.apache.arrow.c.{ArrowArray, ArrowArrayStream, ArrowSchema, Data} +import org.apache.arrow.vector.{FieldVector, VectorSchemaRoot} +import org.apache.arrow.vector.ipc.ArrowReader +import org.apache.arrow.vector.types.pojo.Field + +import org.apache.spark.{SparkException, SparkThrowable, SparkUnsupportedOperationException} +import org.apache.spark.internal.Logging +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader} +import org.apache.spark.sql.connector.write.{DataWriter, WriterCommitMessage} +import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} +import org.apache.spark.sql.execution.datasources.v2.columnar.{ColumnarDataSource, ColumnarReader, ColumnarStreamReader, ColumnarStreamWriter, ColumnarWriter} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{DataType, StructType} +import org.apache.spark.sql.util.{ArrowUtils, CaseInsensitiveStringMap} +import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector} +import org.apache.spark.util.Utils + +/** + * A [[ColumnarDataSource]] implemented by a native library, through + * `org.apache.spark.sql.datasource.NativeBridge`. + * + * @param location where the library is on the driver. The readers and the writers distribute it + * to the executors of the active session when they are created. + * @param name the name of the data source, as listed in the manifest of its package or in lower + * case for an installed library + */ +class NativeDataSource( + location: NativeLibraryLocation, + name: String, + options: CaseInsensitiveStringMap) + extends ColumnarDataSource { + import NativeCalls._ + + override def schema(): StructType = { + val library = NativeLibraries.get(location) + withDataSource(library, "infer the schema of") { dataSource => + call(name, "infer the schema of") { + optional(NativeArrow.importSchema(address => library.schema(dataSource, address))) + }.getOrElse(throw QueryCompilationErrors.dataSchemaNotSpecifiedError(name)) + } + } + + override def reader(schema: StructType): ColumnarReader = { + val executorLocation = NativeDataSourceRegistry.distribute(location) + val library = NativeLibraries.get(location) + val handle = create(library, "plan a scan of", "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED") { + dataSource => + NativeArrow.exportSchema(schema)(address => library.createReader(dataSource, address)) + } + new NativeDataSourceReader( + executorLocation, name, schema, new NativeHandle(handle, library.closeReader)) + } + + override def streamReader(schema: StructType): ColumnarStreamReader = { + val executorLocation = NativeDataSourceRegistry.distribute(location) + val library = NativeLibraries.get(location) + val handle = create(library, "plan a scan of", "DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED") { + dataSource => + NativeArrow.exportSchema(schema)(address => library.createStreamReader(dataSource, address)) + } + new NativeDataSourceStreamReader( + executorLocation, name, schema, new NativeHandle(handle, library.closeStreamReader)) + } + + override def writer(schema: StructType, overwrite: Boolean): ColumnarWriter = { + val executorLocation = NativeDataSourceRegistry.distribute(location) + val library = NativeLibraries.get(location) + val handle = new NativeHandle( + create(library, "plan a write to", "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED") { dataSource => + NativeArrow.exportSchema(schema) { address => + library.createWriter(dataSource, address, overwrite) + } + }, + library.closeWriter) + new NativeDataSourceWriter(executorLocation, name, handle, serializeWriter(library, handle)) + } + + override def streamWriter( + schema: StructType, + overwrite: Boolean, + queryId: String): ColumnarStreamWriter = { + val executorLocation = NativeDataSourceRegistry.distribute(location) + val library = NativeLibraries.get(location) + val handle = new NativeHandle( + create(library, "plan a write to", "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED") { + dataSource => + NativeArrow.exportSchema(schema) { address => + library.createStreamWriter(dataSource, address, overwrite) + } + }, + library.closeWriter) + val writer = new NativeDataSourceStreamWriter( + executorLocation, name, handle, serializeWriter(library, handle)) + NativeStreamWriters.register(queryId, handle) + writer + } + + /** Runs `f` with a new native data source, which is closed afterwards. */ + private def withDataSource[T](library: NativeLibrary, action: String)(f: Long => T): T = { + // The keys of a case-insensitive map are in lower case. + val entries = options.entrySet().asScala.toArray + val dataSource = call(name, "create") { + library.createDataSource(name, entries.map(_.getKey), entries.map(_.getValue)) + } + Utils.tryWithSafeFinally(f(dataSource)) { + call(name, action)(library.closeDataSource(dataSource)) + } + } + + /** Creates a reader or a writer with `f`, which fails if the library does not implement it. */ + private def create(library: NativeLibrary, action: String, unsupportedError: String)( + f: Long => Long): Long = { + withDataSource(library, action) { dataSource => + call(name, action)(optional(f(dataSource))).getOrElse { + throw new SparkUnsupportedOperationException(unsupportedError, Map("description" -> name)) + } + } + } + + private def serializeWriter(library: NativeLibrary, handle: NativeHandle): Array[Byte] = { + try { + bytesOrEmpty(call(name, "plan a write to")(optional(library.serializeWriter(handle.get)))) + } catch { + case e: Throwable => + handle.close() + throw e + } + } +} + +/** Reads a native data source for a batch scan. */ +class NativeDataSourceReader( + location: NativeLibraryLocation, + name: String, + schema: StructType, + @transient private val handle: NativeHandle) + extends ColumnarReader { + import NativeCalls._ + + // Set on the driver when the partitions are planned, and serialized to the executors. + private var readerState: Array[Byte] = Array.emptyByteArray + private var readSchema: StructType = schema + @transient private var plannedPartitions: Array[InputPartition] = _ + + @transient private lazy val library = NativeLibraries.get(location) + + override def pushPredicates(predicates: Array[Predicate]): Array[Predicate] = { + // Only the predicates that can be expressed in JSON are passed to the library. + val encoded = predicates.zipWithIndex.flatMap { case (predicate, i) => + NativePredicates.toJson(predicate, schema).map(_ -> i) + } + val pushed = if (encoded.isEmpty) { + Set.empty[Int] + } else { + call(name, "plan a scan of") { + optional(library.pushPredicates(handle.get, encoded.map(_._1))) + } match { + case Some(flags) if flags != null && flags.length == encoded.length => + encoded.map(_._2).zip(flags).collect { case (i, true) => i }.toSet + case Some(flags) => + throw error(name, "plan a scan of", s"pushPredicates returned " + + s"${Option(flags).map(_.length).getOrElse(0)} values for ${encoded.length} predicates.") + case None => Set.empty[Int] + } + } + predicates.zipWithIndex.collect { case (predicate, i) if !pushed.contains(i) => predicate } + } + + override def pushLimit(limit: Int): Boolean = { + call(name, "plan a scan of")(optional(library.pushLimit(handle.get, limit))).getOrElse(false) + } + + override def pruneColumns(requiredSchema: StructType): Boolean = { + val pruned = call(name, "plan a scan of") { + optional(library.pruneColumns(handle.get, requiredSchema.fieldNames)) + }.getOrElse(false) + if (pruned) { + readSchema = requiredSchema + } + pruned + } + + override def partitions(): Array[InputPartition] = synchronized { + if (plannedPartitions == null) { + try { + val partitions = call(name, "plan a scan of")(library.partitions(handle.get)) + readerState = bytesOrEmpty( + call(name, "plan a scan of")(optional(library.serializeReader(handle.get)))) + plannedPartitions = NativeInputPartition.from(partitions) + } finally { + // The executors only need the serialized state. + handle.close() + } + } + plannedPartitions + } + + override def read(partition: InputPartition): PartitionReader[ColumnarBatch] = { + new NativePartitionReader( + location, name, readerState, partition.asInstanceOf[NativeInputPartition].bytes, readSchema) + } +} + +/** Reads a native data source for a micro-batch streaming scan. */ +class NativeDataSourceStreamReader( + location: NativeLibraryLocation, + name: String, + schema: StructType, + @transient private val handle: NativeHandle) + extends ColumnarStreamReader { + import NativeCalls._ + + // Set on the driver whenever the partitions of a micro-batch are planned, and serialized to + // the executors with the partitions of that micro-batch. + private var readerState: Array[Byte] = Array.emptyByteArray + + @transient private lazy val library = NativeLibraries.get(location) + + override def initialOffset(): String = { + call(name, "plan a scan of")(library.initialOffset(handle.get)) + } + + override def latestOffset(): String = { + call(name, "plan a scan of")(library.latestOffset(handle.get)) + } + + override def partitions(start: String, end: String): Array[InputPartition] = { + val partitions = call(name, "plan a scan of") { + library.streamPartitions(handle.get, start, end) + } + readerState = bytesOrEmpty( + call(name, "plan a scan of")(optional(library.serializeStreamReader(handle.get)))) + NativeInputPartition.from(partitions) + } + + override def read(partition: InputPartition): PartitionReader[ColumnarBatch] = { + new NativePartitionReader( + location, name, readerState, partition.asInstanceOf[NativeInputPartition].bytes, schema) + } + + override def commit(end: String): Unit = { + call(name, "commit an offset of")(optional(library.commitOffset(handle.get, end))) + } + + override def stop(): Unit = handle.close() +} + +/** Writes to a native data source for a batch write. */ +class NativeDataSourceWriter( + location: NativeLibraryLocation, + name: String, + @transient private val handle: NativeHandle, + writerState: Array[Byte]) + extends ColumnarWriter { + import NativeCalls._ + + @transient private lazy val library = NativeLibraries.get(location) + + override def createWriter(partitionId: Int, taskId: Long): DataWriter[ColumnarBatch] = { + new NativeDataWriter(location, name, writerState, partitionId, taskId, epochId = -1L) + } + + override def commit(messages: Array[WriterCommitMessage]): Unit = { + call(name, "commit a write to") { + library.commit(handle.get, -1L, NativeWriterCommitMessage.toBytes(messages)) + } + handle.close() Review Comment: `abort` closes the handle in a `tryWithSafeFinally`, but `commit` closes it only after `library.commit` returns. If `library.commit` throws, `handle.close()` is skipped and the driver-side writer handle is left to the `Cleaner` (released on GC, or never on a quiet driver). This is a silent, non-deterministic delay in releasing the native writer handle rather than a leak that never resolves, but the asymmetry with `abort` looks unintended. Making `commit` symmetric is a one-line change: ```scala override def commit(messages: Array[WriterCommitMessage]): Unit = { Utils.tryWithSafeFinally { call(name, "commit a write to") { library.commit(handle.get, -1L, NativeWriterCommitMessage.toBytes(messages)) } }(handle.close()) } ``` (The streaming writer at L300 and after intentionally does not close here. Its handle is closed by `NativeStreamWriters.register` when the next micro-batch's writer is created, so this comment is only about the batch `NativeDataSourceWriter`.) ########## sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java: ########## @@ -0,0 +1,292 @@ +/* + * 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.execution.datasources.v2.ffi; + +import java.io.IOException; +import java.io.InputStream; +import java.lang.invoke.MethodHandle; +import java.lang.invoke.MethodHandles; +import java.lang.invoke.MethodType; +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Method; + +import org.apache.spark.sql.datasource.NativeBridge; + +/** + * A native data source library, loaded into its own copy of {@link NativeBridge}. + * <p> + * JNI links the native methods of a class to the libraries loaded by the class loader of that + * class, and a library belongs to the class loader of the class that loads it. Defining a copy of + * {@link NativeBridge} in a separate class loader for each library therefore lets all the + * libraries export the same JNI functions. Spark calls the methods of a copy through method + * handles, because the copy is a different class than the {@link NativeBridge} it is compiled + * against. Calling a method that the library does not export throws + * {@link UnsatisfiedLinkError}. + */ +final class NativeLibrary { + + private static final String BRIDGE_CLASS = NativeBridge.class.getName(); + + private final MethodHandle abiVersion; + private final MethodHandle createDataSource; + private final MethodHandle schema; + private final MethodHandle createReader; + private final MethodHandle createStreamReader; + private final MethodHandle createWriter; + private final MethodHandle createStreamWriter; + private final MethodHandle closeDataSource; + private final MethodHandle pushPredicates; + private final MethodHandle pushLimit; + private final MethodHandle pruneColumns; + private final MethodHandle partitions; + private final MethodHandle serializeReader; + private final MethodHandle closeReader; + private final MethodHandle initialOffset; + private final MethodHandle latestOffset; + private final MethodHandle streamPartitions; + private final MethodHandle serializeStreamReader; + private final MethodHandle commitOffset; + private final MethodHandle closeStreamReader; + private final MethodHandle read; + private final MethodHandle serializeWriter; + private final MethodHandle commit; + private final MethodHandle abort; + private final MethodHandle closeWriter; + private final MethodHandle createDataWriter; + private final MethodHandle write; + private final MethodHandle commitDataWriter; + private final MethodHandle abortDataWriter; + + private NativeLibrary(Class<?> bridge) throws ReflectiveOperationException { + Class<?> jlong = long.class; + Class<?> bytes = byte[].class; + Class<?> strings = String[].class; + abiVersion = find(bridge, "abiVersion", int.class); + createDataSource = find(bridge, "createDataSource", jlong, String.class, strings, strings); + schema = find(bridge, "schema", void.class, jlong, jlong); + createReader = find(bridge, "createReader", jlong, jlong, jlong); + createStreamReader = find(bridge, "createStreamReader", jlong, jlong, jlong); + createWriter = find(bridge, "createWriter", jlong, jlong, jlong, boolean.class); + createStreamWriter = find(bridge, "createStreamWriter", jlong, jlong, jlong, boolean.class); + closeDataSource = find(bridge, "closeDataSource", void.class, jlong); + pushPredicates = find(bridge, "pushPredicates", boolean[].class, jlong, strings); + pushLimit = find(bridge, "pushLimit", boolean.class, jlong, int.class); + pruneColumns = find(bridge, "pruneColumns", boolean.class, jlong, strings); + partitions = find(bridge, "partitions", byte[][].class, jlong); + serializeReader = find(bridge, "serializeReader", bytes, jlong); + closeReader = find(bridge, "closeReader", void.class, jlong); + initialOffset = find(bridge, "initialOffset", String.class, jlong); + latestOffset = find(bridge, "latestOffset", String.class, jlong); + streamPartitions = + find(bridge, "streamPartitions", byte[][].class, jlong, String.class, String.class); + serializeStreamReader = find(bridge, "serializeStreamReader", bytes, jlong); + commitOffset = find(bridge, "commitOffset", void.class, jlong, String.class); + closeStreamReader = find(bridge, "closeStreamReader", void.class, jlong); + read = find(bridge, "read", void.class, bytes, bytes, jlong); + serializeWriter = find(bridge, "serializeWriter", bytes, jlong); + commit = find(bridge, "commit", void.class, jlong, jlong, byte[][].class); + abort = find(bridge, "abort", void.class, jlong, jlong, byte[][].class); + closeWriter = find(bridge, "closeWriter", void.class, jlong); + createDataWriter = find(bridge, "createDataWriter", jlong, bytes, int.class, jlong, jlong); + write = find(bridge, "write", void.class, jlong, jlong, jlong); + commitDataWriter = find(bridge, "commitDataWriter", bytes, jlong); + abortDataWriter = find(bridge, "abortDataWriter", void.class, jlong); + } + + private static MethodHandle find( + Class<?> bridge, String name, Class<?> returnType, Class<?>... parameterTypes) + throws ReflectiveOperationException { + return MethodHandles.publicLookup() + .findStatic(bridge, name, MethodType.methodType(returnType, parameterTypes)); + } + + /** + * Loads the library at the given path into a new copy of {@link NativeBridge}. + */ + static NativeLibrary load(String path) throws Throwable { + Class<?> bridge = new BridgeClassLoader(NativeBridge.class.getClassLoader()) + .loadClass(BRIDGE_CLASS); + Method load = bridge.getDeclaredMethod("load", String.class); + load.setAccessible(true); Review Comment: Resolving the package-private `NativeBridge.load` reflectively and forcing access with `setAccessible(true)` works today, but it is a little fragile against future access-control tightening. Since `NativeBridge` is loaded in its own `BridgeClassLoader` copy specifically so libraries can coexist, a short comment on why the reflective access is needed (or making `load` accessible by contract) would help future readers. Non-blocking. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/internal/StaticSQLConf.scala: ########## @@ -464,6 +464,22 @@ object StaticSQLConf { .checkValue(_ >= 0, "The maximum total size must not be negative.") .createWithDefault(128 * 1024) // 128 KiB + val NATIVE_DATA_SOURCE_ENABLED = + buildStaticConf("spark.sql.dataSource.native.enabled") + .doc("When true, Spark finds and loads native data sources: data sources implemented in " + + "native code such as Rust or C++. The library of a data source, such as my_source, is " + + "found when it is installed in the native library path as spark_datasource_my_source, " + + "for example libspark_datasource_my_source.so in /usr/local/lib on Linux. The native " + + "library path consists of the directories of java.library.path, and the lib directory " + + "of each installation prefix in PATH. Libraries can also be distributed in native data " + + "source packages (.sparkpkg files), added with spark.addArtifact or under " + + "spark.sql.dataSource.native.paths. Loading a native data source runs its native code " + + "in the driver and executor processes.") + .version("4.4.0") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .booleanConf + .createWithDefault(true) Review Comment: With this on by default, Spark automatically discovers and loads any `spark_datasource_*` shared library found on `java.library.path` or under `<prefix>/lib` for a `<prefix>/bin` in `PATH`, and runs its native code in the driver and executor processes. That is a sensible end state, but for a new `@Evolving` feature that executes arbitrary native code, an opt-in default (`createWithDefault(false)` for the first release) would be more conservative and would make enabling it a deliberate choice. Non-blocking: just flagging the default for consideration. If `true` is intended, it may be worth a short note in `docs/sql-native-data-sources.md` on the security implication of loading native libraries by file-name convention from the library path. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala: ########## @@ -0,0 +1,243 @@ +/* + * 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.execution.datasources.v2.ffi + +import java.io.File +import java.nio.file.Paths +import java.util.Locale + +import org.apache.spark.{SparkEnv, SparkFiles} +import org.apache.spark.internal.Logging +import org.apache.spark.internal.LogKeys.PATH +import org.apache.spark.sql.Artifact +import org.apache.spark.sql.classic.SparkSession +import org.apache.spark.sql.datasource.NativeBridge +import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} +import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} +import org.apache.spark.util.Utils + +/** + * Finds the native data sources on the driver. Like the Python data sources installed in the + * Python path, the native data source libraries installed in the native library path are found + * automatically, by their file name: the library `spark_datasource_<name>`, such as + * `libspark_datasource_my_source.so` on Linux, implements the data source `<name>`. The active + * session can also add packages with `spark.addArtifact`, or with + * `spark.sql.dataSource.native.paths`, which take precedence over the installed libraries. + */ +object NativeDataSourceRegistry extends Logging { + + /** The prefix of the names of the installed native data source libraries. */ + val LIBRARY_NAME_PREFIX = "spark_datasource_" + + // The result of the last lookup by `DataSource.lookupDataSource` on this thread that found a + // native data source, by the name it looked up. The provider that the caller creates next for + // that name takes it, so that it does not look the name up again: see [[takeLookup]]. + private val lastLookup = new ThreadLocal[(String, (String, NativeLibraryLocation))] + + /** + * Finds the package or the installed library that provides the native data source with the + * given name, which is case-insensitive. Returns the name of the data source, as listed in the + * manifest of the package or in lower case for an installed library, and where its library is. + */ + def lookup(name: String, conf: SQLConf): Option[(String, NativeLibraryLocation)] = { + if (!conf.getConf(StaticSQLConf.NATIVE_DATA_SOURCE_ENABLED)) { + return None + } + lookupPackage(name, configuredPackageFiles(conf) ++ artifactPackageFiles()) + .orElse(lookupInstalledLibrary(name)) + } + + /** + * Like [[lookup]], and keeps the result for the provider that the caller creates next for the + * given name, which takes it with [[takeLookup]]. + */ + def lookupForProvider(name: String, conf: SQLConf): Option[(String, NativeLibraryLocation)] = { + val result = lookup(name, conf) + result.foreach(found => lastLookup.set(name -> found)) + result + } + + /** Takes the result of [[lookupForProvider]] for the given name on this thread, if any. */ + def takeLookup(name: String): Option[(String, NativeLibraryLocation)] = { + Option(lastLookup.get()).filter(_._1.equalsIgnoreCase(name)).map { case (_, found) => + lastLookup.remove() + found + } + } + + private def lookupPackage( + name: String, + files: Seq[File]): Option[(String, NativeDataSourcePackage)] = { + val (invalid, valid) = files.map(NativeDataSourcePackage.tryRead).partitionMap(identity) + // An invalid package does not prevent finding the data sources of the other packages, unless + // its manifest lists the data source. An invalid package is logged when it is read. + invalid.find(_.dataSources.exists(_.equalsIgnoreCase(name))).foreach(p => throw p.error) + val found = valid + .flatMap(pkg => pkg.manifest.dataSources.find(_.equalsIgnoreCase(name)).map(_ -> pkg)) + // The same package can be found more than once, for example when it is in a configured + // directory and was also added as an artifact. + .groupBy(_._2.sha256).values.map(_.head).toSeq + found match { + case Seq() => None + case Seq(result @ (_, pkg)) => + if (pkg.manifest.abiVersion != NativeBridge.ABI_VERSION) { + throw QueryExecutionErrors.invalidNativeDataSourcePackageError( + pkg.path, + "UNSUPPORTED_ABI_VERSION", + Map( + "version" -> pkg.manifest.abiVersion.toString, + "supported" -> NativeBridge.ABI_VERSION.toString)) + } + Some(result) + case _ => + throw QueryCompilationErrors.nativeDataSourcePackageConflictError( + name, found.map(_._2.path).sorted) + } + } + + private def lookupInstalledLibrary(name: String): Option[(String, InstalledNativeLibrary)] = { + // Other names, such as class names, are not the names of libraries. + if (!name.matches("[A-Za-z0-9_]+")) { + return None + } + val dataSource = name.toLowerCase(Locale.ROOT) + findInstalledLibrary(libraryFileName(dataSource)) + .map(file => dataSource -> InstalledNativeLibrary(file.getPath)) + } + + /** Finds the library with the given file name in the native library path of this node. */ + private[ffi] def findInstalledLibrary(fileName: String): Option[File] = { + librarySearchPath(System.getProperty("java.library.path"), System.getenv("PATH")) + .map(new File(_, fileName)) + .find(_.isFile) + } + + /** + * The file name of the installed library of the native data source with the given name, such + * as `libspark_datasource_my_source.so` on Linux, `libspark_datasource_my_source.dylib` on + * macOS and `spark_datasource_my_source.dll` on Windows. + */ + def libraryFileName(name: String): String = { + System.mapLibraryName(LIBRARY_NAME_PREFIX + name.toLowerCase(Locale.ROOT)) + } + + /** + * The native library path: the directories where native data source libraries are installed, + * in the order they are searched. They are the directories of `java.library.path`, where the + * JVM finds native libraries, which include `LD_LIBRARY_PATH` on Linux, `DYLD_LIBRARY_PATH` on + * macOS and `PATH` on Windows, and then the `lib` directory of each installation prefix in + * `PATH`, such as `/usr/local/lib` for `/usr/local/bin`. Relative directories are ignored. + */ + private[ffi] def librarySearchPath(javaLibraryPath: String, path: String): Seq[File] = { + def directories(value: String): Seq[File] = { + Option(value).toSeq.flatMap(_.split(File.pathSeparator)).map(new File(_)).filter(_.isAbsolute) + } + val prefixLibraries = directories(path) + .filter(_.getName == "bin") + .flatMap(bin => Option(bin.getParentFile)) + .map(new File(_, "lib")) + (directories(javaLibraryPath) ++ prefixLibraries).distinct + } + + /** + * Makes the library available to the executors of the active session, which runs the jobs of + * the scan or the write: for a streaming query, it is the session of the query, a clone of the + * session that started it. Returns where the executors find it. An installed library is already + * installed on every node. + */ + def distribute(location: NativeLibraryLocation): NativeLibraryLocation = location match { + case pkg: NativeDataSourcePackage => + SparkSession.getActiveSession match { + // In local mode, the executors run in the driver and read the package where it is. + case Some(session) if !session.sparkContext.isLocal => addToArtifacts(session, pkg) + case _ => pkg Review Comment: The `case _ => pkg` collapses two different situations: local mode (where returning the driver-side package is correct, as the comment says) and `getActiveSession == None` (where no package is distributed to the executors at all). In a non-local cluster, if `distribute` ever runs on a thread without an active session set, the package is not added to the artifacts, and on the executors `distributedCopies` will not find it. The job then fails with `NATIVE_DATA_SOURCE_PACKAGE_NOT_FOUND`, with no hint as to why. I could not construct a concrete path where the planning thread has no active session (`reader()`, `writer()`, `streamReader()`, and `streamWriter()` all run during planning on the driver, where the active session is normally set), so this is a defensive hardening rather than a demonstrated bug. But separating the two cases would make the intent explicit and avoid a silent skip: ```scala case Some(session) if session.sparkContext.isLocal => pkg case Some(session) => addToArtifacts(session, pkg) case None => logWarning(log"No active session while distributing the native data source " + log"package ${MDC(PATH, pkg.path)}; the executors may not find it.") pkg ``` Note the same `getActiveSession` dependency exists in `artifactPackageFiles()` (around L239), so if there is a real no-active-session path it would affect lookup as well as distribution. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
