dongjoon-hyun commented on code in PR #59313: URL: https://github.com/apache/spark/pull/59313#discussion_r4235638424
########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala: ########## @@ -0,0 +1,202 @@ +/* + * 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_" + + /** Returns whether the native data source with the given name exists. */ + def exists(name: String, conf: SQLConf): Boolean = lookup(name, conf).isDefined + + /** + * 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)) + } + + private def lookupPackage( + name: String, + files: Seq[File]): Option[(String, NativeDataSourcePackage)] = { + val found = files.map(_.getCanonicalFile).distinct.map(NativeDataSourcePackage.read) Review Comment: **[1] One invalid package breaks every lookup miss** `lookupPackage` parses every configured and artifact package before matching the name, so one invalid or unreadable `.sparkpkg` throws `INVALID_NATIVE_DATA_SOURCE_PACKAGE` for any name that misses the Java/Python lookup, even when that package does not provide the name. For example, with a corrupt `bad.sparkpkg` under `spark.sql.dataSource.native.paths` or added with `spark.addArtifact`: - `spark.read.format("rust_range")` fails, although `rust_range` comes from a different, valid package. - `spark.read.format("com.typo.Missing")` reports the package error instead of `DATA_SOURCE_NOT_FOUND`. - `spark.dataSource.register(...)` of any Python data source fails with `DATA_SOURCE_ALREADY_EXISTS`, because `DataSourceRegistration.checkDataSourceExists` treats every error other than `DATA_SOURCE_NOT_FOUND` as "exists". - `SELECT * FROM db.missing_table` goes through `ResolveSQLOnFile` -> `lookupDataSource("db")` and fails with `FAILED_TO_CREATE_PLAN_FOR_DIRECT_QUERY` instead of `TABLE_OR_VIEW_NOT_FOUND`. Could we skip (and log) invalid packages during the lookup, and only report the error when the invalid package is the one that matches the name? ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala: ########## @@ -0,0 +1,202 @@ +/* + * 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_" + + /** Returns whether the native data source with the given name exists. */ + def exists(name: String, conf: SQLConf): Boolean = lookup(name, conf).isDefined + + /** + * 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)) + } + + private def lookupPackage( + name: String, + files: Seq[File]): Option[(String, NativeDataSourcePackage)] = { + val found = files.map(_.getCanonicalFile).distinct.map(NativeDataSourcePackage.read) + .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. 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 + } + case installed: InstalledNativeLibrary => installed + } + + /** Adds the package as an artifact of the session, unless it already is one. */ + private[ffi] def addToArtifacts( + session: SparkSession, + pkg: NativeDataSourcePackage): NativeDataSourcePackage = { + val artifactManager = session.artifactManager + val (added, artifactUUID) = artifactManager.getNativeDataSourcePackages + if (!added.exists(_.getName == pkg.fileName)) { + artifactManager.addLocalArtifacts(Artifact.newFileArtifact( + Paths.get(pkg.fileName), new Artifact.LocalFile(Paths.get(pkg.path))) :: Nil) + } + pkg.copy(artifactUUID = artifactUUID) + } + + /** + * The local copies, on an executor, of a package distributed with [[distribute]]. The files of + * an isolated session are stored in a directory named after the UUID of the session. + */ + def distributedCopies(pkg: NativeDataSourcePackage): Seq[File] = { + if (SparkEnv.get == null) { + Nil + } else { + val root = new File(SparkFiles.getRootDirectory()) + pkg.artifactUUID.map(uuid => new File(new File(root, uuid), pkg.fileName)).toSeq :+ Review Comment: **[2] Streaming tasks look up the package under the wrong session UUID** `distribute` records the artifact UUID of the session that analyzed the query, but a streaming query runs its tasks in `sparkSessionForStream = sparkSession.cloneSession()`, which has its own `sessionUUID`. `ArtifactManager.clone` re-registers the files under the new UUID, so the executors store them in `SparkFiles.getRootDirectory()/<cloneUUID>`, which `distributedCopies` never checks. With an isolated session (e.g. Spark Connect) on a real cluster, `spark.addArtifact("x.sparkpkg")` followed by `spark.readStream.format("x").load().writeStream...start()` would fail with `NATIVE_DATA_SOURCE_PACKAGE_NOT_FOUND` on an executor that never ran a task of the original session. `NativeDataSourceConnectSuite` doesn't cover this, because it runs in local mode, where `distribute` returns the package as is. Could we resolve the UUID from the job artifact state of the running task on the executor, rather than capturing it at analysis time? ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceV2.scala: ########## @@ -0,0 +1,67 @@ +/* + * 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.util + +import org.apache.spark.sql.connector.catalog.Table +import org.apache.spark.sql.connector.expressions.Transform +import org.apache.spark.sql.errors.QueryCompilationErrors +import org.apache.spark.sql.execution.datasources.v2.NamedTableProvider +import org.apache.spark.sql.execution.datasources.v2.columnar.ColumnarTable +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.util.CaseInsensitiveStringMap + +/** + * The Data Source V2 provider of the native data sources: data sources implemented by a native + * library through `org.apache.spark.sql.datasource.NativeBridge`. + */ +class NativeDataSourceV2 extends NamedTableProvider { + private var shortName: String = _ + + override def setShortName(name: String): Unit = { + assert(shortName == null) + shortName = name + } + + // The name of the data source, and where its library is. Resolved on first use, which is while + // the query is analyzed, with the configuration of its session. + private lazy val resolved: (String, NativeLibraryLocation) = { + assert(shortName != null) + val (name, location) = NativeDataSourceRegistry.lookup(shortName, SQLConf.get).getOrElse { Review Comment: **[8] The registry lookup runs twice for each resolution** `DataSource.lookupDataSource` already calls `NativeDataSourceRegistry.exists` (a full `lookup`), and this calls `lookup` again. Each lookup lists the configured directories, canonicalizes and stats every package, and probes every library directory. Besides the duplicated file system work, a package added or removed between the two calls can resolve to a different location than the one `exists` saw. Could we pass the result of the first lookup to the provider, like `setShortName` passes the name, instead of looking it up again? ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala: ########## @@ -0,0 +1,617 @@ +/* + * 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 +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 + * @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 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( + location, name, schema, new NativeHandle(handle, library.closeReader)) + } + + override def streamReader(schema: StructType): ColumnarStreamReader = { + 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( + location, name, schema, new NativeHandle(handle, library.closeStreamReader)) + } + + override def writer(schema: StructType, overwrite: Boolean): ColumnarWriter = { + 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(location, name, handle, serializeWriter(library, handle)) + } + + override def streamWriter(schema: StructType, overwrite: Boolean): ColumnarStreamWriter = { + 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) + new NativeDataSourceStreamWriter(location, name, handle, serializeWriter(library, handle)) + } + + /** 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).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. */ +class NativeDataSourceStreamWriter( Review Comment: **[6] The stream writer handle is never closed explicitly** Unlike the batch writer, `NativeDataSourceStreamWriter` never calls `handle.close()`: `commit` and `abort` keep the handle, and `StreamingWrite` has no stop hook. `closeWriter` therefore only runs when the `Cleaner` collects the `NativeHandle`, which may be much later, or never on a quiet driver. The `NativeBridge` Javadoc says that "Spark releases each handle exactly once with the matching close method". Any connections, files or locks that the native writer holds stay open after the streaming query stops, and a restarted query on the same driver can run into them. Could we close the handle when the query stops, or at least document that a stream writer is released only by GC? ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala: ########## @@ -0,0 +1,92 @@ +/* + * 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.util.Base64 + +import com.fasterxml.jackson.databind.ObjectMapper +import com.fasterxml.jackson.databind.node.ObjectNode + +import org.apache.spark.sql.connector.expressions.{Expression, GeneralScalarExpression, Literal, NamedReference} +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.types._ + +/** + * Encodes Data Source V2 predicates as the JSON expression trees passed to + * `NativeBridge.pushPredicates`. + */ +object NativePredicates { + private val mapper = new ObjectMapper() + + /** Returns the JSON of the predicate, or None if it cannot be expressed in JSON. */ + def toJson(predicate: Predicate): Option[String] = { + encode(predicate).map(mapper.writeValueAsString) + } + + private def encode(expression: Expression): Option[ObjectNode] = expression match { + case reference: NamedReference => + val node = mapper.createObjectNode().put("type", "column") + val name = node.putArray("name") + reference.fieldNames().foreach(name.add) + Some(node) + + case literal: Literal[_] => + encodeLiteral(literal.value(), literal.dataType()) + + case function: GeneralScalarExpression => + val children = function.children().map(encode) + if (children.forall(_.isDefined)) { + val node = mapper.createObjectNode() + .put("type", "function") + .put("name", function.name()) + val array = node.putArray("children") + children.foreach(child => array.add(child.get)) + Some(node) + } else { + None + } + + case _ => None + } + + private def encodeLiteral(value: Any, dataType: DataType): Option[ObjectNode] = { + val node = mapper.createObjectNode() + .put("type", "literal") + .put("dataType", dataType.catalogString) + if (value == null) { + Some(node.putNull("value")) + } else { + val encoded = (value, dataType) match { + case (b: Boolean, BooleanType) => Some(node.put("value", b)) + case (n: Byte, ByteType) => Some(node.put("value", n.toInt)) + case (n: Short, ShortType) => Some(node.put("value", n)) + case (n: Int, IntegerType | DateType) => Some(node.put("value", n)) + case (n: Long, LongType | TimestampType | TimestampNTZType) => Some(node.put("value", n)) + case (n: Float, FloatType) if !n.isNaN && !n.isInfinite => Some(node.put("value", n)) Review Comment: **[4] Spark comparison semantics (NaN, collation) for pushed predicates** A library can report a predicate as fully evaluated, after which Spark drops it, but the JSON contract in `NativeBridge` doesn't state Spark's comparison semantics: - NaN literals are excluded here, but NaN column values are not. For `WHERE d > 1.0`, Spark treats NaN as greater than any value, so a NaN row qualifies; a library that uses IEEE comparisons returns false and silently drops the row. - The collation check only applies to literals, so `WHERE a = b` on two `UTF8_LCASE` (or other non-binary collated) columns is pushed. A byte-wise comparison then drops rows that differ only in case. Could we either document these semantics in the `NativeBridge` Javadoc (NaN ordering and equality, collations) or stop pushing predicates that compare floating-point or non-binary collated columns? ########## python/pyspark/sql/session.py: ########## @@ -2292,6 +2296,11 @@ def addArtifacts( "condition": "True together", }, ) + if not (file or pyfile or archive) and all(p.endswith(".sparkpkg") for p in path): Review Comment: **[3] Classic `addArtifact(..., file=True)` doesn't register the package** In classic mode, a `.sparkpkg` is routed to the JVM `ArtifactManager` only when no flag is set. With `file=True`, it goes through `SparkContext.addFile`, which `ArtifactManager.getNativeDataSourcePackages` never sees, so `spark.read.format(...)` then fails with `DATA_SOURCE_NOT_FOUND`. With Spark Connect, the same call does register the package (`elif file or name.endswith(".sparkpkg")` in `connect/client/artifact.py`). Also, a mixed call without flags such as `addArtifacts("a.sparkpkg", "b.py")` silently does nothing for either file. Shall we make both modes behave the same way, e.g. route `.sparkpkg` files to the JVM regardless of `file=True`, and handle mixed lists? ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSource.scala: ########## @@ -682,6 +683,7 @@ object DataSource extends Logging { val serviceLoader = ServiceLoader.load(classOf[DataSourceRegister], loader) lazy val isUserDefinedDataSource = SparkSession.getActiveSession.exists( _.sessionState.dataSourceManager.dataSourceExists(provider)) + lazy val isNativeDataSource = NativeDataSourceRegistry.exists(provider, conf) Review Comment: **[9] Installed native libraries now block Python data source registration** Since `lookupDataSource` now resolves the names of installed native libraries, `DataSourceRegistration.checkDataSourceExists` rejects registering a Python data source whose name matches any `libspark_datasource_<name>` under `java.library.path` or `<prefix>/lib` for a `<prefix>/bin` in `PATH`. For example, with `/usr/local/lib/libspark_datasource_foo.so` installed (possibly a stale copy left by another tool), `spark.dataSource.register(FooDataSource)` with the name `foo` fails with `DATA_SOURCE_ALREADY_EXISTS`, although a registered Python data source would take precedence over the native one in this lookup order. Whether it fails depends on files on the node that the user may not know about. Is this intended? If not, `checkDataSourceExists` could skip the native lookup, consistent with the precedence here. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala: ########## @@ -0,0 +1,226 @@ +/* + * 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.columnar + +import java.util +import java.util.Locale + +import org.apache.spark.SparkException +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.connector.catalog.{SupportsRead, SupportsWrite, Table, TableCapability} +import org.apache.spark.sql.connector.catalog.TableCapability.{BATCH_READ, BATCH_WRITE, MICRO_BATCH_READ, STREAMING_WRITE, TRUNCATE} +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.connector.read._ +import org.apache.spark.sql.connector.read.streaming.{MicroBatchStream, Offset} +import org.apache.spark.sql.connector.write.{LogicalWriteInfo, WriteBuilder} +import org.apache.spark.sql.internal.connector.SupportsMetadata +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.util.CaseInsensitiveStringMap +import org.apache.spark.sql.vectorized.ColumnarBatch + +/** + * The Data Source V2 [[Table]] of a [[ColumnarDataSource]]. It creates a new data source for each + * scan and each write, with their options on top of the options of the table. + * + * @param tableOptions the options of the table, given to `TableProvider.getTable`. They are the + * options of `CREATE TABLE ... USING ... OPTIONS (...)` for a table in a + * catalog, while the scans and the writes of the table only get the options + * of their query. + */ +class ColumnarTable( + shortName: String, + tableSchema: StructType, + tableOptions: util.Map[String, String], + createDataSource: CaseInsensitiveStringMap => ColumnarDataSource) + extends Table with SupportsRead with SupportsWrite { + + override def name(): String = shortName + + override def schema(): StructType = tableSchema + + override def capabilities(): util.Set[TableCapability] = + util.EnumSet.of(BATCH_READ, BATCH_WRITE, MICRO_BATCH_READ, STREAMING_WRITE, TRUNCATE) + + override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { + new ColumnarScanBuilder(tableSchema, () => createDataSource(withTableOptions(options))) + } + + override def newWriteBuilder(info: LogicalWriteInfo): WriteBuilder = { + new ColumnarWriteBuilder( + shortName, info, () => createDataSource(withTableOptions(info.options()))) + } + + private def withTableOptions(options: CaseInsensitiveStringMap): CaseInsensitiveStringMap = { + val merged = new util.HashMap[String, String]() + tableOptions.forEach((key, value) => merged.put(key.toLowerCase(Locale.ROOT), value)) + // The keys of a case-insensitive map are in lower case. + options.forEach((key, value) => merged.put(key, value)) + new CaseInsensitiveStringMap(merged) + } +} + +class ColumnarScanBuilder(schema: StructType, createDataSource: () => ColumnarDataSource) + extends ScanBuilder + with SupportsPushDownV2Filters + with SupportsPushDownLimit + with SupportsPushDownRequiredColumns { + + private lazy val dataSource = createDataSource() + + // Created on first use, because a streaming scan neither pushes anything down nor reads in + // batch, and its data source may not support batch scans at all. + private var reader: ColumnarReader = _ + + private var pushed: Array[Predicate] = Array.empty + private var pushedLimit: Option[Int] = None + private var readSchema: StructType = schema + + private def batchReader: ColumnarReader = { + if (reader == null) { + reader = dataSource.reader(schema) + } + reader + } + + override def pushPredicates(predicates: Array[Predicate]): Array[Predicate] = { + val remaining = batchReader.pushPredicates(predicates) + pushed = predicates.filterNot(remaining.contains) + remaining + } + + override def pushedPredicates(): Array[Predicate] = pushed + + override def pushLimit(limit: Int): Boolean = { + val isPushed = batchReader.pushLimit(limit) + if (isPushed) { + pushedLimit = Some(limit) + } + isPushed + } + + // Spark always applies the limit again: a data source is free to return more rows. + override def isPartiallyPushed(): Boolean = true + + override def pruneColumns(requiredSchema: StructType): Unit = { + // Only top-level columns are pruned by the data source. Spark prunes the nested fields of the + // columns it reads. + val required = + StructType(requiredSchema.fieldNames.flatMap(name => schema.find(_.name == name))) + if (required.length < schema.length && batchReader.pruneColumns(required)) { + readSchema = required + } + } + + override def build(): Scan = { + new ColumnarScan(schema, readSchema, () => batchReader, () => dataSource, pushed, pushedLimit) + } +} + +class ColumnarScan( + fullSchema: StructType, + prunedSchema: StructType, + batchReader: () => ColumnarReader, + dataSource: () => ColumnarDataSource, + pushedPredicates: Array[Predicate], + pushedLimit: Option[Int]) + extends Scan with SupportsMetadata { + + // Spark may call `toBatch` more than once for the same scan; the partitions are planned once. + private lazy val batch = new ColumnarBatchScan(batchReader()) + + override def readSchema(): StructType = prunedSchema + + // The class that implements the data source, such as NativeDataSource for a native one. + override def description(): String = dataSource().getClass.getSimpleName Review Comment: **[10] The scan description shows the class name instead of the data source name** This returns `NativeDataSource` for every native scan, so a query that joins `rust_range` and `cpp_range` shows two identical `BatchScan NativeDataSource` nodes in `EXPLAIN` and the UI. It also forces the lazy `dataSource` of the scan builder only to read its class name. Could we use the short name that `ColumnarTable` already has, as `ColumnarWrite.description()` does? ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala: ########## @@ -0,0 +1,202 @@ +/* + * 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_" + + /** Returns whether the native data source with the given name exists. */ + def exists(name: String, conf: SQLConf): Boolean = lookup(name, conf).isDefined + + /** + * 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)) + } + + private def lookupPackage( + name: String, + files: Seq[File]): Option[(String, NativeDataSourcePackage)] = { + val found = files.map(_.getCanonicalFile).distinct.map(NativeDataSourcePackage.read) + .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. 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 + } + case installed: InstalledNativeLibrary => installed + } + + /** Adds the package as an artifact of the session, unless it already is one. */ + private[ffi] def addToArtifacts( + session: SparkSession, + pkg: NativeDataSourcePackage): NativeDataSourcePackage = { + val artifactManager = session.artifactManager + val (added, artifactUUID) = artifactManager.getNativeDataSourcePackages + if (!added.exists(_.getName == pkg.fileName)) { Review Comment: **[5] Artifact dedup by file name can ship a different package** This compares only the file name. If the session already has a different artifact with the same file name, the configured package is not added, but the executors look for it by its SHA-256 checksum. For example, `spark.sql.dataSource.native.paths=/pkgs` contains `ds.sparkpkg` (data source `a`), and the session already added another `ds.sparkpkg` (data source `b`). A query on `a` resolves to `/pkgs/ds.sparkpkg`, this check skips adding it, and `NativeLibraries.locate` on the executors only finds the artifact copy, whose checksum differs, so it fails with `NATIVE_DATA_SOURCE_PACKAGE_NOT_FOUND`. Could we compare the checksums here and report a clear conflict when the names match but the contents differ? ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala: ########## @@ -0,0 +1,617 @@ +/* + * 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 +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 + * @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 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( + location, name, schema, new NativeHandle(handle, library.closeReader)) + } + + override def streamReader(schema: StructType): ColumnarStreamReader = { + 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( + location, name, schema, new NativeHandle(handle, library.closeStreamReader)) + } + + override def writer(schema: StructType, overwrite: Boolean): ColumnarWriter = { + 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(location, name, handle, serializeWriter(library, handle)) + } + + override def streamWriter(schema: StructType, overwrite: Boolean): ColumnarStreamWriter = { + 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) + new NativeDataSourceStreamWriter(location, name, handle, serializeWriter(library, handle)) + } + + /** 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).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. */ +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. + 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 = { + val dataWriter = handle + handle = 0 Review Comment: **[7] A failed `commitDataWriter` skips the task abort** The handle is reset to 0 before `commitDataWriter` is called. If the native commit throws, `DataWritingSparkTask` calls `dataWriter.abort()`, but it is a no-op because `handle == 0`, so `abortDataWriter` is never called for the partially written output of the task. The `NativeBridge` Javadoc doesn't say that a library must clean up by itself when `commitDataWriter` fails. Could we either document that `commitDataWriter` releases the handle and cleans up even when it throws, or call `abortDataWriter` when the commit fails? -- 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]
