dongjoon-hyun commented on code in PR #59313: URL: https://github.com/apache/spark/pull/59313#discussion_r4236806540
########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala: ########## @@ -0,0 +1,373 @@ +/* + * 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, FileInputStream, IOException} +import java.nio.charset.StandardCharsets +import java.nio.file.{Files, StandardCopyOption} +import java.security.MessageDigest +import java.util.{HexFormat, Locale} +import java.util.concurrent.ConcurrentHashMap +import java.util.zip.ZipFile + +import scala.jdk.CollectionConverters._ +import scala.util.Try + +import com.fasterxml.jackson.databind.{JsonNode, ObjectMapper} + +import org.apache.spark.internal.Logging +import org.apache.spark.internal.LogKeys.PATH +import org.apache.spark.sql.Artifact +import org.apache.spark.sql.datasource.NativeBridge +import org.apache.spark.sql.errors.QueryExecutionErrors +import org.apache.spark.util.Utils + +/** + * The manifest of a native data source package, `spark-native-datasource.json`: + * {{{ + * { + * "abiVersion": 1, + * "dataSources": ["my_source"], + * "libraries": { + * "linux-x86_64": "linux-x86_64/libmy_source.so", + * "osx-aarch64": "osx-aarch64/libmy_source.dylib" + * } + * } + * }}} + * + * @param abiVersion the version of [[NativeBridge]] that the libraries implement + * @param dataSources the names of the data sources that the libraries implement + * @param libraries the path of the library in the package for each platform, as returned by + * [[NativePlatform.current]] + */ +case class NativeDataSourceManifest( + abiVersion: Int, + dataSources: Seq[String], + libraries: Map[String, String]) + +/** + * Where the library of a native data source is: in a [[NativeDataSourcePackage]], or installed in + * the native library path, as an [[InstalledNativeLibrary]]. It is sent to the executors, which + * load the library with [[NativeLibraries.get]]. + */ +sealed trait NativeLibraryLocation extends Serializable + +/** + * A native data source package: a zip file that contains a [[NativeDataSourceManifest]] and the + * libraries it lists. The SHA-256 checksum of the file identifies the package, so that an + * executor can check that it loads the same package as the driver. + * + * @param path the path of the package file on the driver + * @param artifactUUID the UUID of the session whose artifacts contain the package, if the session + * is isolated. Executors store the artifacts of such a session separately. + */ +case class NativeDataSourcePackage( + path: String, + sha256: String, + manifest: NativeDataSourceManifest, + artifactUUID: Option[String] = None) extends NativeLibraryLocation { + def fileName: String = new File(path).getName +} + +/** + * A package file that is not a valid native data source package. + * + * @param error the error that explains why the package is invalid + * @param dataSources the data sources listed in the manifest of the package, if it can be read + */ +case class InvalidNativeDataSourcePackage(error: Throwable, dataSources: Seq[String]) + +object NativeDataSourcePackage extends Logging { + val FILE_EXTENSION: String = Artifact.NATIVE_DATA_SOURCE_PACKAGE_EXTENSION + val MANIFEST_NAME = "spark-native-datasource.json" + + private case class CacheKey(path: String, length: Long, lastModified: Long) + + private[ffi] type ReadResult = Either[InvalidNativeDataSourcePackage, NativeDataSourcePackage] + + private val packages = new ConcurrentHashMap[CacheKey, ReadResult]() Review Comment: **[3] The package cache is never evicted** This cache is keyed by the canonical path, length and modification time, and never evicted. It now also keeps the invalid results, including their errors. Since each session stores its artifacts under its own directory, and each change of a file adds a new key, a long-running Spark Connect server where many sessions add packages keeps accumulating entries. Non-blocking: could we evict the entries of a session's artifacts when the session is cleaned up, or bound the cache? ########## 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() + } + + override def abort(messages: Array[WriterCommitMessage]): Unit = { + Utils.tryWithSafeFinally { + call(name, "abort a write to") { + library.abort(handle.get, -1L, NativeWriterCommitMessage.toBytes(messages)) + } + }(handle.close()) + } +} + +/** + * Writes to a native data source for a streaming write. Its handle is closed when the streaming + * query creates its next stream writer, see [[NativeStreamWriters]]. + */ +class NativeDataSourceStreamWriter( + location: NativeLibraryLocation, + name: String, + @transient private val handle: NativeHandle, + writerState: Array[Byte]) + extends ColumnarStreamWriter { + import NativeCalls._ + + @transient private lazy val library = NativeLibraries.get(location) + + override def createWriter( + partitionId: Int, + taskId: Long, + epochId: Long): DataWriter[ColumnarBatch] = { + new NativeDataWriter(location, name, writerState, partitionId, taskId, epochId) + } + + override def commit(epochId: Long, messages: Array[WriterCommitMessage]): Unit = { + call(name, "commit a write to") { + library.commit(handle.get, epochId, NativeWriterCommitMessage.toBytes(messages)) + } + } + + override def abort(epochId: Long, messages: Array[WriterCommitMessage]): Unit = { + call(name, "abort a write to") { + library.abort(handle.get, epochId, NativeWriterCommitMessage.toBytes(messages)) + } + } +} + +/** A partition planned by a native library. */ +case class NativeInputPartition(index: Int, bytes: Array[Byte]) extends InputPartition + +object NativeInputPartition { + def from(partitions: Array[Array[Byte]]): Array[InputPartition] = { + Option(partitions).getOrElse(Array.empty).zipWithIndex.map { case (bytes, i) => + NativeInputPartition(i, Option(bytes).getOrElse(Array.emptyByteArray)) + } + } +} + +/** The message returned by the data writer of a task of a native library. */ +case class NativeWriterCommitMessage(bytes: Array[Byte]) extends WriterCommitMessage + +object NativeWriterCommitMessage { + def toBytes(messages: Array[WriterCommitMessage]): Array[Array[Byte]] = messages.map { + case message: NativeWriterCommitMessage => message.bytes + case null => null + case other => + throw SparkException.internalError(s"Unexpected commit message: ${other.getClass.getName}") + } +} + +/** + * Reads a partition of a native data source: the Arrow stream exported by the library is + * imported into Arrow vectors that back the returned batch. + */ +class NativePartitionReader( + location: NativeLibraryLocation, + name: String, + readerState: Array[Byte], + partition: Array[Byte], + schema: StructType) + extends PartitionReader[ColumnarBatch] { + import NativeCalls._ + + private val allocator = + ArrowUtils.rootAllocator.newChildAllocator(s"native data source $name", 0, Long.MaxValue) + private var reader: ArrowReader = _ + private var root: VectorSchemaRoot = _ + private var batch: ColumnarBatch = _ + + try { + reader = openStream() + root = call(name, "read")(reader.getVectorSchemaRoot) + checkSchema() + batch = new ColumnarBatch( + root.getFieldVectors.asScala.map(new ArrowColumnVector(_): ColumnVector).toArray) + } catch { + case e: Throwable => + close() + throw e + } + + private def openStream(): ArrowReader = { + val library = NativeLibraries.get(location) + val stream = ArrowArrayStream.allocateNew(allocator) + try { + call(name, "read")(optional(library.read(readerState, partition, stream.memoryAddress()))) + .getOrElse { + throw new SparkUnsupportedOperationException( + "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", Map("description" -> name)) + } + Data.importArrayStream(allocator, stream) + } catch { + case e: Throwable => + stream.release() + throw e + } finally { + stream.close() + } + } + + private def checkSchema(): Unit = { + def isDictionaryEncoded(field: Field): Boolean = { + field.getDictionary != null || field.getChildren.asScala.exists(isDictionaryEncoded) + } + root.getSchema.getFields.asScala.find(isDictionaryEncoded).foreach { field => + throw error(name, "read", s"The column ${field.getName} is dictionary-encoded, which " + + "is not supported.") + } + val actual = ArrowUtils.fromArrowSchema(root.getSchema) + val matches = actual.length == schema.length && actual.zip(schema).forall { + case (actualField, field) => + actualField.name == field.name && + DataType.equalsIgnoreNullability(actualField.dataType, field.dataType) + } + if (!matches) { + throw error(name, "read", s"The data has the schema ${actual.toDDL}, but the expected " + + s"schema is ${schema.toDDL}.") + } + } + + override def next(): Boolean = { + val hasNext = call(name, "read")(reader.loadNextBatch()) + if (hasNext) { + batch.setNumRows(root.getRowCount) + } + hasNext + } + + override def get(): ColumnarBatch = batch + + override def close(): Unit = { + Utils.tryWithSafeFinally { + if (reader != null) { + reader.close() + } + }(Utils.tryLogNonFatalError(allocator.close())) + } +} + +/** Writes the batches of a task to the data writer of a native library. */ +class NativeDataWriter( + location: NativeLibraryLocation, + name: String, + writerState: Array[Byte], + partitionId: Int, + taskId: Long, + epochId: Long) + extends DataWriter[ColumnarBatch] { + import NativeCalls._ + + private val library = NativeLibraries.get(location) + // Zero once the data writer was committed or aborted, which releases it. It stays set if the + // commit fails, so that Spark aborts the data writer. + private var handle: Long = call(name, "write to") { + library.createDataWriter(writerState, partitionId, taskId, epochId) + } + + override def write(batch: ColumnarBatch): Unit = { + if (handle == 0) { + throw SparkException.internalError("The data writer was already committed or aborted.") + } + val vectors = (0 until batch.numCols()).map { i => + batch.column(i) match { + case vector: ArrowColumnVector => vector.getValueVector.asInstanceOf[FieldVector] + case vector => throw SparkException.internalError( + s"Expected an Arrow column vector, but got ${vector.getClass.getName}.") + } + } + val root = new VectorSchemaRoot( + vectors.map(_.getField).asJava, vectors.asJava, batch.numRows()) + val array = ArrowArray.allocateNew(ArrowUtils.rootAllocator) + val arrowSchema = ArrowSchema.allocateNew(ArrowUtils.rootAllocator) + try { + Data.exportVectorSchemaRoot(ArrowUtils.rootAllocator, root, null, array, arrowSchema) + call(name, "write to") { + library.write(handle, array.memoryAddress(), arrowSchema.memoryAddress()) + } + } finally { + // Releases the batch unless the library moved it to take ownership. + array.release() + array.close() + arrowSchema.release() + arrowSchema.close() + } + } + + override def commit(): WriterCommitMessage = { + if (handle == 0) { + throw SparkException.internalError("The data writer was already committed or aborted.") + } + val message = call(name, "write to")(library.commitDataWriter(handle)) + handle = 0 + NativeWriterCommitMessage(bytesOrEmpty(Some(message))) + } + + override def abort(): Unit = { + if (handle != 0) { + val dataWriter = handle + handle = 0 + call(name, "write to")(library.abortDataWriter(dataWriter)) + } + } + + override def close(): Unit = abort() +} + +/** + * A handle to a native object. It is released with `close()`, or once it becomes unreachable if + * it was never closed. + */ +final class NativeHandle(handle: Long, release: Long => Unit) { + private val state = new NativeHandle.State(handle, release) + private val cleanable = NativeHandle.cleaner.register(this, state) + + def get: Long = { + val value = state.handle.get() + if (value == 0) { + throw SparkException.internalError("The native handle was already released.") + } + value + } + + def close(): Unit = cleanable.clean() +} + +object NativeHandle extends Logging { + private val cleaner = Cleaner.create() + + // Must not refer to the handle, which would then never become unreachable. + private class State(value: Long, release: Long => Unit) extends Runnable { + val handle = new AtomicLong(value) + + override def run(): Unit = { + val value = handle.getAndSet(0) + if (value != 0) { + try { + release(value) + } catch { + case e @ (NonFatal(_) | _: LinkageError) => + logWarning("Failed to release a native data source handle.", e) + } + } + } + } +} + +/** + * The stream writers of the streaming queries, by query ID. The Data Source V2 streaming write + * has no hook for when the query stops, but a micro-batch streaming query creates a new stream + * writer for each micro-batch, after the previous micro-batch committed or aborted its write, and + * a restarted query creates a new one. So the previous stream writer of a query is closed when + * the query creates a new one, and the last one when it becomes unreachable. + */ +object NativeStreamWriters { Review Comment: **[4] No end-to-end test for closing the previous stream writer** `close the previous stream writer of a streaming query` calls `register` directly, but no streaming test checks that a micro-batch query actually closes the writer of the previous micro-batch, or that the `queryId` passed by `V2Writes` stays the same across micro-batches. If `V2Writes` or `MicroBatchExecution` changes, for example to pass the run ID or to reuse one `StreamingWrite` across micro-batches, stream writers would again only be released by GC, and no test would fail. Non-blocking: a `closeWriter` counter in the test library, checked after the multi-batch query in `streaming read and write`, would cover it. ########## 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) Review Comment: **[1] An invalid package that lists the name hides a valid one** This throws before the valid packages are considered, so a stale broken copy hides a valid package that provides the same data source. For example, a shared directory in `spark.sql.dataSource.native.paths` still has a broken `foo-1.sparkpkg` whose manifest lists `foo` but whose library entry is missing. The user adds a fixed `foo-2.sparkpkg` with `spark.addArtifact`, but `spark.read.format("foo")` keeps failing with `INVALID_NATIVE_DATA_SOURCE_PACKAGE` until someone removes the old file from the shared directory. Non-blocking: could we prefer a valid match when there is one, and only report the invalid package when nothing else provides the name? ########## 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)) Review Comment: **[2] The thread-local can keep a stale lookup on pooled threads** This is set whenever the lookup finds a native data source, but several callers of `lookupDataSource` never create a provider to take it: `ResolveSQLOnFile` (through `DataSource.providingClass`), the provider comparison in `PreprocessTableCreation`, and `DataStreamWriter` when the v2 writer is disabled. The entry then stays on the thread, e.g. a pooled Spark Connect gRPC thread, until a later lookup of the same name. It isn't wrong today, because every `NativeDataSourceV2` is created right after a successful lookup that overwrites the entry, but that invariant is implicit. Non-blocking: clearing the entry when a lookup does not find a native data source, or returning the result to the caller instead of using a thread-local, would make it more robust. -- 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]
