cloud-fan commented on code in PR #59313: URL: https://github.com/apache/spark/pull/59313#discussion_r4260023379
########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourcePackage.scala: ########## @@ -0,0 +1,421 @@ +/* + * 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 com.google.common.cache.{Cache, CacheBuilder} + +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. + * @param distributedFileName the name of the copy of the package in the artifacts of the session, + * if the package is not itself an artifact of the session + */ +case class NativeDataSourcePackage( + path: String, + sha256: String, + manifest: NativeDataSourceManifest, + artifactUUID: Option[String] = None, + distributedFileName: 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] + + // Bounded, because each session stores its artifacts in a directory of its own, so a + // long-running server where many sessions add packages reads many different files. + private val MAX_CACHED_PACKAGES = 1000 + + private val packages: Cache[CacheKey, ReadResult] = + CacheBuilder.newBuilder().maximumSize(MAX_CACHED_PACKAGES).build[CacheKey, ReadResult]() + private val mapper = new ObjectMapper() + + /** Reads the package in the given file. The result is cached until the file changes. */ + def read(file: File): NativeDataSourcePackage = tryRead(file).fold(p => throw p.error, identity) + + /** + * Reads the package in the given file, or returns why it is invalid. The result is cached until + * the file changes, and an invalid package is logged when it is read. + */ + private[ffi] def tryRead(file: File): ReadResult = { + val canonical = try { + file.getCanonicalFile + } catch { + case e: IOException => + return Left(InvalidNativeDataSourcePackage( + invalidManifestError(file, s"The package cannot be read: $e", e), Nil)) + } + val key = CacheKey(canonical.getPath, canonical.length(), canonical.lastModified()) + Option(packages.getIfPresent(key)).getOrElse { + val result = parse(canonical) + result.swap.foreach { invalid => + logWarning(log"Invalid native data source package ${MDC(PATH, canonical)}.", + invalid.error) + } + packages.put(key, result) + result + } + } + + /** + * Forgets the packages read from the files under the given directory, such as the artifacts of + * a session that is cleaned up. + */ + def forgetPackagesIn(dir: File): Unit = { + val prefix = dir.getCanonicalPath + File.separator + packages.asMap().keySet().removeIf(_.path.startsWith(prefix)) + } + + /** Returns whether a package read from the file at the given canonical path is cached. */ + private[ffi] def isCached(canonicalPath: String): Boolean = { + packages.asMap().keySet().asScala.exists(_.path == canonicalPath) + } + + private def invalidManifestError(file: File, reason: String, cause: Throwable = null) = { + QueryExecutionErrors.invalidNativeDataSourcePackageError( + file.getPath, "INVALID_MANIFEST", Map("reason" -> reason), cause) + } + + private def parse(file: File): ReadResult = { + def invalid(reason: String, dataSources: Seq[String] = Nil, cause: Throwable = null) = { + Left(InvalidNativeDataSourcePackage(invalidManifestError(file, reason, cause), dataSources)) + } + try { + val manifest = Utils.tryWithResource(new ZipFile(file)) { zip => + Option(zip.getEntry(MANIFEST_NAME)) match { + case None => invalid(s"The package does not contain the manifest $MANIFEST_NAME.") + case Some(entry) => + val json = Utils.tryWithResource(zip.getInputStream(entry)) { in => + new String(in.readAllBytes(), StandardCharsets.UTF_8) + } + parseManifest(json) match { + case Left(reason) => invalid(reason) + case Right(manifest) => + manifest.libraries.values.find(zip.getEntry(_) == null) match { + case Some(library) => invalid( + s"The package does not contain the library $library.", manifest.dataSources) + case None => Right(manifest) + } + } + } + } + manifest.map(NativeDataSourcePackage(file.getPath, sha256(file), _)) + } catch { + case e: IOException => invalid(s"The package is not a valid zip file: $e", cause = e) + } + } + + private[ffi] def parseManifest(json: String): Either[String, NativeDataSourceManifest] = { + val root: JsonNode = try { + mapper.readTree(json) + } catch { + case e: IOException => return Left(s"The manifest is not valid JSON: ${e.getMessage}") + } + if (root == null || !root.isObject) { + return Left("The manifest must be a JSON object.") + } + val abiVersion = root.get("abiVersion") + if (abiVersion == null || !abiVersion.isInt) { + return Left("The manifest must have an integer field 'abiVersion'.") + } + val dataSources = root.get("dataSources") + if (dataSources == null || !dataSources.isArray || dataSources.isEmpty || + !dataSources.elements().asScala.forall(n => n.isTextual && n.asText().nonEmpty)) { + return Left("The manifest must have a field 'dataSources' with a non-empty array of " + + "data source names.") + } + val libraries = root.get("libraries") + if (libraries == null || !libraries.isObject || libraries.isEmpty || + !libraries.properties().asScala.forall(e => e.getValue.isTextual)) { + return Left("The manifest must have a field 'libraries' with an object that maps " + + "platforms to the paths of the libraries in the package.") + } + Right(NativeDataSourceManifest( + abiVersion.asInt(), + dataSources.elements().asScala.map(_.asText()).toSeq, + libraries.properties().asScala.map(e => e.getKey -> e.getValue.asText()).toMap)) + } + + private[ffi] def sha256(file: File): String = { + val digest = MessageDigest.getInstance("SHA-256") + Utils.tryWithResource(new FileInputStream(file)) { in => + val buffer = new Array[Byte](64 * 1024) + var read = in.read(buffer) + while (read != -1) { + digest.update(buffer, 0, read) + read = in.read(buffer) + } + } + HexFormat.of().formatHex(digest.digest()) + } +} + +/** + * A native data source library installed in the native library path, where + * [[NativeDataSourceRegistry]] finds it by its file name. It must be installed on every node. The + * SHA-256 checksum of the file identifies the build of the library, so that an executor can check + * that it loads the same build as the driver: the state that the driver serializes is only + * meant for the same build. + * + * @param path the path of the library on the driver + * @param sha256 the SHA-256 checksum of the library on the driver + */ +case class InstalledNativeLibrary(path: String, sha256: String) extends NativeLibraryLocation + +object InstalledNativeLibrary { + private case class CacheKey(path: String, length: Long, lastModified: Long) + + // Bounded, like the cache of the packages, although few libraries are installed. + private val checksums: Cache[CacheKey, String] = + CacheBuilder.newBuilder().maximumSize(1000).build[CacheKey, String]() + + /** The installed library in the given file. */ + def apply(file: File): InstalledNativeLibrary = InstalledNativeLibrary(file.getPath, sha256(file)) + + /** The SHA-256 checksum of the given file. The result is cached until the file changes. */ + private[ffi] def sha256(file: File): String = { + val canonical = file.getCanonicalFile + val key = CacheKey(canonical.getPath, canonical.length(), canonical.lastModified()) + Option(checksums.getIfPresent(key)).getOrElse { + val checksum = NativeDataSourcePackage.sha256(canonical) + checksums.put(key, checksum) + checksum + } + } +} + +/** The platform names used in [[NativeDataSourceManifest.libraries]]. */ +object NativePlatform { + /** The platform of this JVM, such as `linux-x86_64`, `linux-aarch64` or `osx-aarch64`. */ + lazy val current: String = name(System.getProperty("os.name"), System.getProperty("os.arch")) + + private[ffi] def name(osName: String, osArch: String): String = { + val os = osName.toLowerCase(Locale.ROOT) match { + case n if n.startsWith("linux") => "linux" + case n if n.startsWith("mac") || n.startsWith("darwin") => "osx" + case n if n.startsWith("windows") => "windows" + case n => n.replaceAll("[^a-z0-9]", "") + } + val arch = osArch.toLowerCase(Locale.ROOT) match { + case "amd64" | "x86_64" | "x64" => "x86_64" + case "aarch64" | "arm64" => "aarch64" + case a => a + } + s"$os-$arch" + } +} + +/** + * The native data source libraries loaded in this JVM. A library is loaded once, on first use, + * after it is extracted from its package if needed, and stays loaded until the JVM exits. + */ +object NativeLibraries extends Logging { + // The libraries of the packages, by the checksum of the package, and the installed libraries, + // by their canonical path: the JVM loads a library file in a single class loader. A library + // that is loaded but rejected stays loaded, so its error is kept: it cannot be loaded again. + private val packagedLibraries = new ConcurrentHashMap[String, Try[NativeLibrary]]() + private val installedLibraries = new ConcurrentHashMap[String, Try[NativeLibrary]]() + + private lazy val extractionDir: File = Utils.createTempDir(namePrefix = "native-datasources") + + /** Returns the library at the given location, and loads it if needed. */ + private[ffi] def get(location: NativeLibraryLocation): NativeLibrary = location match { + case pkg: NativeDataSourcePackage => + getOrLoad(packagedLibraries, pkg.sha256)(loadPackaged(pkg)) + case installed: InstalledNativeLibrary => + val file = locate(installed).getCanonicalFile + getOrLoad(installedLibraries, file.getPath)(loadInstalled(file)) Review Comment: **Non-blocking (P2):** `locate` picks the file whose current checksum matches, but `getOrLoad(installedLibraries, file.getPath)` then returns whatever was loaded from that path before, and the JVM cannot load the same path twice. So the same-build guarantee breaks after an in-place upgrade. For example, a long-running Spark Connect or Thrift server has loaded build A of `/usr/local/lib/libspark_datasource_x.so`, and an operator installs build B there. On the next query: - `lookupInstalledLibrary` sends `sha(B)`, because its checksum cache follows the length and modification time of the file. - The driver passes the check and serializes its state with build A. - An executor started after the upgrade loads build B and parses A's state. - An executor that loaded A earlier also passes the check and keeps running A. In native code that is undefined behavior, with no error. Packages avoid this because `packagedLibraries` is keyed by the package checksum. Could `NativeLibraries` remember the checksum of the build it loaded from each path, and fail clearly, asking for a restart, when another checksum is requested for that path? This extends an [existing discussion](https://github.com/apache/spark/pull/59313#pullrequestreview-5497999118). **Recommended change:** Record in NativeLibraries the checksum of the build loaded from each canonical path, and make get() fail with a clear error that names the path and both checksums and asks to restart the application when the requested checksum differs from the loaded one; document this under Installed Libraries. **Why this works:** get(installed) already resolves the file by checksum; keeping the loaded checksum next to the cached library and comparing it on every cache hit turns the silent mismatch into an error, on the driver before any state is serialized and on executors that loaded an older build. **Scope:** Tie each loaded installed library to the checksum of the build it was loaded from. **Compatibility:** Preserve behavior outside the stated repair boundary. **Risks:** A long-running driver must be restarted to use a library replaced in place; the docs need to say so. **Constraints:** The comparison must happen before get() returns a cached library, on both the driver and the executors. **Success:** After an installed library is replaced in place, a query on a JVM that loaded the old build fails with an error naming the path and the two checksums, instead of running mixed builds. A JVM that has not loaded the library loads the new build as before, and the same build at the same path keeps returning the cached library. **Verification:** - **Regression:** With an installed library loaded from a path, replacing the file with another build and resolving it again fails with the new error and does not return the old build. - **Behavior:** Resolving the same build from the same path again, and from another directory of the native library path, still returns the loaded library. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala: ########## @@ -0,0 +1,794 @@ +/* + * 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, CDataDictionaryProvider, Data} +import org.apache.arrow.vector.{FieldVector, VectorSchemaRoot} +import org.apache.arrow.vector.complex.MapVector +import org.apache.arrow.vector.types.pojo.{ArrowType, Field, Schema} + +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.unsafe.Platform +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)) Review Comment: **Non-blocking (P2):** `OffsetSeqLog` writes each offset as one line and reads the log back with `getLines()`, but these strings go there unchecked, and NativeBridge only says that offsets are JSON. Say a library returns valid but pretty-printed JSON, from `serde_json::to_string_pretty` or `nlohmann::json::dump(2)`. The query runs fine, but each micro-batch writes a multi-line entry. When the query restarts, the extra lines are read as more offsets, so it cannot restart from its checkpoint. A library that returns NULL without throwing produces `ColumnarStreamOffset(null)`. That fails internally when the offset log is written: with a NullPointerException in the version 2 format, or a MatchError in version 1. Neither error names the data source. Could we reject null, invalid JSON and line breaks here with `NATIVE_DATA_SOURCE_ERROR`, and state in NativeBridge that offsets must be compact single-line JSON? **Verification:** - **Regression:** A stream reader that returns pretty-printed JSON or null fails with NATIVE_DATA_SOURCE_ERROR instead of writing a multi-line offset or failing with an internal NullPointerException or MatchError when the offset log is written. - **Compatibility:** Compact JSON offsets reach streamPartitions and commitOffset unchanged, and a query restarts from its checkpoint. ########## sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java: ########## @@ -0,0 +1,442 @@ +/* + * 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.datasource; + +import org.apache.spark.annotation.Evolving; + +/** + * The binary interface between Spark and data sources implemented in native code, such as Rust + * or C++. + * <p> + * A native data source is a shared library that implements the {@code native} methods of this + * class as JNI functions, for example + * {@code Java_org_apache_spark_sql_datasource_NativeBridge_createDataSource}. Spark loads each + * library into its own copy of this class, so any number of libraries that export the same + * functions can be loaded side by side. Spark adapts a native data source to Data Source V2, with + * batch and micro-batch streaming reads, batch and streaming writes, and predicate, column and + * limit pushdown. + * <p> + * Spark finds the library of a data source automatically when it is installed like other native + * libraries and named after the data source: the library of {@code my_source} is + * {@code spark_datasource_my_source}, such as {@code libspark_datasource_my_source.so} in + * {@code /usr/local/lib} on Linux. Spark looks for it in the directories of + * {@code java.library.path}, which include {@code LD_LIBRARY_PATH} on Linux, and in the + * {@code lib} directory of each installation prefix in {@code PATH}. A library can also be + * distributed in a native data source package: a zip file with the extension {@code .sparkpkg} + * that contains a manifest named {@code spark-native-datasource.json} and the library built for + * one or more platforms, added to a session with {@code spark.addArtifact} or under the paths in + * {@code spark.sql.dataSource.native.paths}. See the "Native Data Sources" page of the Spark SQL + * guide for details and examples. + * + * <h2>Lifecycle</h2> + * Spark creates a data source with {@link #createDataSource} on the driver for each schema + * inference, scan and write, and closes it right after. + * <ul> + * <li><b>Batch reads.</b> On the driver, Spark creates a reader with {@link #createReader}, + * pushes operations down by calling, each at most once and in this order, + * {@link #pushPredicates}, {@link #pushLimit} and {@link #pruneColumns}, plans the partitions + * with {@link #partitions}, serializes the state of the reader with {@link #serializeReader}, + * and closes it. Then it calls {@link #read} on the executors, in one task for each + * partition.</li> + * <li><b>Streaming reads.</b> On the driver, Spark creates a stream reader with + * {@link #createStreamReader} and keeps it until the streaming query stops. For each + * micro-batch, it asks for the {@link #latestOffset}, plans the partitions with + * {@link #streamPartitions} and serializes the state with {@link #serializeStreamReader}, then + * calls {@link #read} on the executors. Offsets are JSON strings defined by the library. When + * a streaming query restarts, from its checkpoint, Spark creates a new stream reader and may + * call {@link #streamPartitions} again for the last range that it planned before, so the + * partitions of a range must read the same data each time.</li> + * <li><b>Writes.</b> On the driver, once it computed the data to write, right before it starts + * the tasks of the write (or of a micro-batch), Spark creates a writer with {@link #createWriter} + * or {@link #createStreamWriter}, and serializes its state with {@link #serializeWriter}. If + * computing the data fails, Spark creates no writer. On the executors, each task creates a data + * writer with {@link #createDataWriter}, calls {@link #write} for each batch of at most + * {@code spark.sql.execution.arrow.maxRecordsPerBatch} rows, which also ends once it holds + * {@code spark.sql.execution.arrow.maxBytesPerBatch} bytes, and calls {@link #commitDataWriter}, + * or {@link #abortDataWriter} if the task or the commit fails. Then the driver calls + * {@link #commit} with the messages of the tasks if they all succeeded, or {@link #abort} + * otherwise, including when {@link #commit} fails, and closes the writer. Several attempts of a + * task may write the same partition, after a failure or with speculative execution, each with a + * data writer of its own. Spark lets the attempts of a partition commit one at a time, and passes + * to {@link #commit} the message of one committed attempt for each partition, so the write + * consists of the data of these attempts only. It aborts the data writers of the other attempts, + * except on an executor that is lost, where a data writer is neither committed nor aborted, and + * may even have committed before its message reached the driver. A streaming write commits or + * aborts each micro-batch (epoch). Spark may commit the same epoch more than once, for example + * when a streaming query restarts after a failure, so committing an epoch that is already + * committed must have no effect, for exactly-once output. A micro-batch streaming query creates a + * new stream writer for each micro-batch, and Spark closes the previous stream writer of the + * query when the query creates a new one. Spark is not told when a streaming query stops, so it + * closes the last stream writer of the query once the writer is no longer referenced, or when the + * query is restarted.</li> + * </ul> + * + * <h2>Conventions</h2> + * <ul> + * <li><b>Handles.</b> The {@code create} methods return an opaque, non-zero handle, typically a + * pointer. Spark releases each handle exactly once with the matching {@code close} method, or + * with {@link #commitDataWriter} or {@link #abortDataWriter} for data writers, though not + * always right after the last use of a stream writer (see Writes above). A reader or + * writer handle must not depend on the data source handle it was created from, because Spark + * closes the data source handle right after creating it.</li> + * <li><b>State.</b> Partitions, the state of readers and writers, and commit messages are byte + * arrays in a format defined by the library. Spark sends them to the executors and back.</li> + * <li><b>Errors.</b> To report an error, throw a Java exception, for example with the JNI + * function {@code ThrowNew} on {@code java.lang.RuntimeException}, and return any value. Spark + * rethrows it as a {@code NATIVE_DATA_SOURCE_ERROR}.</li> + * <li><b>Arrow data.</b> Schemas and data are exchanged with the + * <a href="https://arrow.apache.org/docs/format/CDataInterface.html";>Arrow C data + * interface</a>. A {@code long} argument named {@code ...Address} is the address of a C struct + * allocated by Spark. For an output struct, the library initializes it and Spark takes + * ownership. For an input struct, the library may move it to take ownership; otherwise, Spark + * releases it after the call returns, so the library must not use it afterwards. The arrays + * that the library exports must have an offset of 0, at every level: Arrow Java ignores the + * offsets of the arrays it imports, so Spark rejects an array with another offset, such as a + * zero-copy slice. Copy a sliced array before exporting it. Spark also rejects an array with + * fewer values than its parent refers to, such as a column shorter than its batch. The buffers + * of an array must hold all of its values, as the C data interface requires: it does not carry + * the sizes of the buffers, so Spark cannot check them, and reads past the end of a buffer that + * is too short.</li> + * <li><b>Schemas.</b> Spark types map to Arrow types as in {@code spark.sql.execution.arrow}: + * for example, {@code bigint} is {@code int64}, {@code string} is {@code utf8}, + * {@code timestamp} is {@code timestamp[us]} with a time zone and {@code timestamp_ntz} is + * {@code timestamp[us]} without one. Data read from a library may also use + * {@code large_utf8}, {@code utf8_view} and the other variants of the same type, and any names + * for the fields of the entries of a map, but not unsigned integers, {@code decimal256}, list + * views or dictionary encoding.</li> + * <li><b>Options.</b> Option keys are passed in lower case, because Spark matches them + * case-insensitively.</li> + * <li><b>Optional methods.</b> A library only needs to export the functions for the operations + * it supports. If a function is missing, Spark treats the operation as unsupported, except for + * the pushdown methods and the {@code serialize} methods, which then push nothing down and + * return no state.</li> + * <li><b>Threads.</b> Spark calls the library from multiple threads, and calls on different + * handles may run concurrently. Calls on the same handle never run concurrently. A handle may + * be released from a thread other than the threads that used it, such as a JVM + * {@code Cleaner} thread when Spark releases a handle once it is no longer referenced, for + * example the last stream writer of a streaming query.</li> + * </ul> + * + * <h2>Compatibility</h2> + * The JVM binds the functions of a library by their names only, so the name of this class, the + * names of its methods and their signatures are the interface. Within an {@link #ABI_VERSION}: + * <ul> + * <li>the signature and the semantics of an existing method never change;</li> + * <li>new methods may be added, only as optional methods, so that a library that does not + * export them keeps working as before;</li> + * <li>new values may appear where this documentation says so, such as the names of pushed + * predicates.</li> + * </ul> + * An incompatible change uses a new {@link #ABI_VERSION}, with new names or a new class. A + * library reports the version it implements with {@link #abiVersion}, and Spark rejects a library + * of a version it does not support. Spark supports version 1 at least in all the releases of + * Spark 4 from 4.4.0. + * + * @since 4.4.0 + */ +@Evolving +public final class NativeBridge { + + /** + * The version of this interface. It changes when the interface changes incompatibly: see + * Compatibility above. + */ + public static final int ABI_VERSION = 1; + + private NativeBridge() {} + + /** + * Loads the library that implements the native methods of this copy of the class. Spark calls + * it on the copy of the class it creates for the library. + */ + static void load(String path) { + System.load(path); + } + + /** Returns the version of this interface that the library implements, {@link #ABI_VERSION}. */ + public static native int abiVersion(); + + // --------------------------------------------------------------------------------------------- + // Data source (driver) + // --------------------------------------------------------------------------------------------- + + /** + * Creates a data source with the options given by the user, such as + * {@code spark.read.option("key", "value")}. + * + * @param name the name of the data source, as listed in the package manifest. A library can + * implement several data sources. + * @param optionKeys the option keys, in lower case + * @param optionValues the option values, in the same order as the keys + * @return the data source handle + */ + public static native long createDataSource( + String name, String[] optionKeys, String[] optionValues); + + /** + * Exports the schema of the data. Spark calls it only when the user does not specify a schema. + * + * @param schemaAddress the address of the {@code ArrowSchema} to initialize with a struct type + * whose fields are the columns + */ + public static native void schema(long dataSource, long schemaAddress); + + /** + * Creates a reader for a batch scan. + * + * @param schemaAddress the address of the {@code ArrowSchema} of the schema to read: the one the + * user specified, or the one returned by {@link #schema} + * @return the reader handle + */ + public static native long createReader(long dataSource, long schemaAddress); + + /** + * Creates a reader for a micro-batch streaming scan. + * + * @param schemaAddress the address of the {@code ArrowSchema} of the schema to read: the one the + * user specified, or the one returned by {@link #schema} + * @return the stream reader handle + */ + public static native long createStreamReader(long dataSource, long schemaAddress); + + /** + * Creates a writer for a batch write. + * + * @param schemaAddress the address of the {@code ArrowSchema} of the data to write + * @param overwrite whether to replace the existing data (save mode "overwrite") instead of + * appending to it + * @return the writer handle + */ + public static native long createWriter(long dataSource, long schemaAddress, boolean overwrite); + + /** + * Creates a writer for a streaming write. + * + * @param schemaAddress the address of the {@code ArrowSchema} of the data to write + * @param overwrite whether to replace the existing data in each micro-batch (output mode + * "complete") instead of appending to it + * @return the writer handle + */ + public static native long createStreamWriter( + long dataSource, long schemaAddress, boolean overwrite); + + /** Releases a data source handle. */ + public static native void closeDataSource(long dataSource); + + // --------------------------------------------------------------------------------------------- + // Batch reader (driver) + // --------------------------------------------------------------------------------------------- + + /** + * Pushes predicates down to the reader. The predicates are combined with AND. + * <p> + * Each predicate is a JSON expression tree. An expression is one of: + * <ul> + * <li>{@code {"type": "column", "name": ["a", "b"]}}: a column. A nested field has a name + * part for each level.</li> + * <li>{@code {"type": "literal", "dataType": "bigint", "value": 1}}: a literal. The data type + * is a Spark SQL type name such as {@code int}, {@code string} or {@code decimal(10,2)}. The + * value is {@code null}, a JSON boolean or number for boolean and numeric types, a string + * for {@code string}, a base64 string for {@code binary}, a decimal string for + * {@code decimal}, the days since the epoch for {@code date}, and the microseconds since the + * epoch for {@code timestamp} and {@code timestamp_ntz}.</li> + * <li>{@code {"type": "function", "name": ">", "children": [...]}}: a predicate, with the + * name it has in {@code org.apache.spark.sql.connector.expressions.filter.Predicate}. The + * names are {@code =}, {@code <>}, {@code <=>} (equal, with null equal to null), {@code <}, + * {@code <=}, {@code >} and {@code >=} with two children, {@code IN} with the value to look + * for and then the values to look in, {@code IS_NULL} and {@code IS_NOT_NULL} with one + * child, {@code STARTS_WITH}, {@code ENDS_WITH} and {@code CONTAINS} with the string and the + * string to look for, {@code AND} and {@code OR} with two predicates, and {@code NOT} with + * one. The children of the other functions are columns, literals or predicates.</li> + * </ul> + * Spark does not pass predicates that it cannot express this way, such as predicates that use + * arithmetic or other functions, nor predicates on columns that contain strings with a + * non-binary collation, such as {@code UTF8_LCASE}: strings are compared by their UTF-8 bytes. Review Comment: **Non-blocking (P2):** This says that strings are compared by their UTF-8 bytes, but nothing says how a library gets those bytes from the `String`s of the interface. The JNI accessor that C and C++ code uses for that, `GetStringUTFChars`, returns modified UTF-8, which differs from UTF-8 for U+0000 and for supplementary characters. Both the `ToString` helper of the C++ example on the docs page and the test library use it. `NativePredicates` writes the JSON with a default `ObjectMapper`, which escapes control characters but writes a supplementary character as it is. So take a C++ library that accepts `name = '<literal>'` by comparing bytes. It gets the literal `x` followed by U+1F600 as `78 ed a0 bd ed b8 80`, while the row holds `78 f0 9f 98 80`. No row matches, and since the library reported the predicate as evaluated completely, Spark drops the filter and the query silently returns no rows. I checked both halves locally: what `GetStringUTFChars` returns for U+1F600 on JDK 17, and what Jackson 2.22 writes by default. Two small changes would close this: - Enable `JsonWriteFeature.ESCAPE_NON_ASCII` on the mapper of `NativePredicates`. The JSON becomes ASCII-only, and any JSON parser decodes it to the same characters. - Add a Strings convention: the arguments and results are Java strings, and `GetStringUTFChars` and `NewStringUTF` use modified UTF-8. This also covers the name, the options, the pruned column names and the offsets. **Verification:** - **Regression:** The predicate encoding test encodes a string literal with a supplementary character and a non-ASCII BMP character, checks that the JSON is ASCII-only, and checks that parsing it back gives the original literal. - **Inspection:** The Strings convention covers every string of the interface: the name and the options of createDataSource, the predicates, the pruned column names and the offsets. ########## sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java: ########## @@ -0,0 +1,442 @@ +/* + * 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.datasource; + +import org.apache.spark.annotation.Evolving; + +/** + * The binary interface between Spark and data sources implemented in native code, such as Rust + * or C++. + * <p> + * A native data source is a shared library that implements the {@code native} methods of this + * class as JNI functions, for example + * {@code Java_org_apache_spark_sql_datasource_NativeBridge_createDataSource}. Spark loads each + * library into its own copy of this class, so any number of libraries that export the same + * functions can be loaded side by side. Spark adapts a native data source to Data Source V2, with + * batch and micro-batch streaming reads, batch and streaming writes, and predicate, column and + * limit pushdown. + * <p> + * Spark finds the library of a data source automatically when it is installed like other native + * libraries and named after the data source: the library of {@code my_source} is + * {@code spark_datasource_my_source}, such as {@code libspark_datasource_my_source.so} in + * {@code /usr/local/lib} on Linux. Spark looks for it in the directories of + * {@code java.library.path}, which include {@code LD_LIBRARY_PATH} on Linux, and in the + * {@code lib} directory of each installation prefix in {@code PATH}. A library can also be + * distributed in a native data source package: a zip file with the extension {@code .sparkpkg} + * that contains a manifest named {@code spark-native-datasource.json} and the library built for + * one or more platforms, added to a session with {@code spark.addArtifact} or under the paths in + * {@code spark.sql.dataSource.native.paths}. See the "Native Data Sources" page of the Spark SQL + * guide for details and examples. + * + * <h2>Lifecycle</h2> + * Spark creates a data source with {@link #createDataSource} on the driver for each schema + * inference, scan and write, and closes it right after. + * <ul> + * <li><b>Batch reads.</b> On the driver, Spark creates a reader with {@link #createReader}, + * pushes operations down by calling, each at most once and in this order, + * {@link #pushPredicates}, {@link #pushLimit} and {@link #pruneColumns}, plans the partitions + * with {@link #partitions}, serializes the state of the reader with {@link #serializeReader}, + * and closes it. Then it calls {@link #read} on the executors, in one task for each + * partition.</li> + * <li><b>Streaming reads.</b> On the driver, Spark creates a stream reader with + * {@link #createStreamReader} and keeps it until the streaming query stops. For each + * micro-batch, it asks for the {@link #latestOffset}, plans the partitions with + * {@link #streamPartitions} and serializes the state with {@link #serializeStreamReader}, then + * calls {@link #read} on the executors. Offsets are JSON strings defined by the library. When + * a streaming query restarts, from its checkpoint, Spark creates a new stream reader and may + * call {@link #streamPartitions} again for the last range that it planned before, so the + * partitions of a range must read the same data each time.</li> + * <li><b>Writes.</b> On the driver, once it computed the data to write, right before it starts + * the tasks of the write (or of a micro-batch), Spark creates a writer with {@link #createWriter} + * or {@link #createStreamWriter}, and serializes its state with {@link #serializeWriter}. If + * computing the data fails, Spark creates no writer. On the executors, each task creates a data + * writer with {@link #createDataWriter}, calls {@link #write} for each batch of at most + * {@code spark.sql.execution.arrow.maxRecordsPerBatch} rows, which also ends once it holds + * {@code spark.sql.execution.arrow.maxBytesPerBatch} bytes, and calls {@link #commitDataWriter}, + * or {@link #abortDataWriter} if the task or the commit fails. Then the driver calls + * {@link #commit} with the messages of the tasks if they all succeeded, or {@link #abort} + * otherwise, including when {@link #commit} fails, and closes the writer. Several attempts of a + * task may write the same partition, after a failure or with speculative execution, each with a + * data writer of its own. Spark lets the attempts of a partition commit one at a time, and passes + * to {@link #commit} the message of one committed attempt for each partition, so the write + * consists of the data of these attempts only. It aborts the data writers of the other attempts, + * except on an executor that is lost, where a data writer is neither committed nor aborted, and + * may even have committed before its message reached the driver. A streaming write commits or + * aborts each micro-batch (epoch). Spark may commit the same epoch more than once, for example + * when a streaming query restarts after a failure, so committing an epoch that is already + * committed must have no effect, for exactly-once output. A micro-batch streaming query creates a + * new stream writer for each micro-batch, and Spark closes the previous stream writer of the + * query when the query creates a new one. Spark is not told when a streaming query stops, so it + * closes the last stream writer of the query once the writer is no longer referenced, or when the + * query is restarted.</li> + * </ul> + * + * <h2>Conventions</h2> + * <ul> + * <li><b>Handles.</b> The {@code create} methods return an opaque, non-zero handle, typically a + * pointer. Spark releases each handle exactly once with the matching {@code close} method, or + * with {@link #commitDataWriter} or {@link #abortDataWriter} for data writers, though not + * always right after the last use of a stream writer (see Writes above). A reader or + * writer handle must not depend on the data source handle it was created from, because Spark + * closes the data source handle right after creating it.</li> + * <li><b>State.</b> Partitions, the state of readers and writers, and commit messages are byte + * arrays in a format defined by the library. Spark sends them to the executors and back.</li> + * <li><b>Errors.</b> To report an error, throw a Java exception, for example with the JNI + * function {@code ThrowNew} on {@code java.lang.RuntimeException}, and return any value. Spark + * rethrows it as a {@code NATIVE_DATA_SOURCE_ERROR}.</li> + * <li><b>Arrow data.</b> Schemas and data are exchanged with the + * <a href="https://arrow.apache.org/docs/format/CDataInterface.html";>Arrow C data + * interface</a>. A {@code long} argument named {@code ...Address} is the address of a C struct + * allocated by Spark. For an output struct, the library initializes it and Spark takes + * ownership. For an input struct, the library may move it to take ownership; otherwise, Spark + * releases it after the call returns, so the library must not use it afterwards. The arrays + * that the library exports must have an offset of 0, at every level: Arrow Java ignores the + * offsets of the arrays it imports, so Spark rejects an array with another offset, such as a + * zero-copy slice. Copy a sliced array before exporting it. Spark also rejects an array with + * fewer values than its parent refers to, such as a column shorter than its batch. The buffers + * of an array must hold all of its values, as the C data interface requires: it does not carry + * the sizes of the buffers, so Spark cannot check them, and reads past the end of a buffer that + * is too short.</li> + * <li><b>Schemas.</b> Spark types map to Arrow types as in {@code spark.sql.execution.arrow}: + * for example, {@code bigint} is {@code int64}, {@code string} is {@code utf8}, + * {@code timestamp} is {@code timestamp[us]} with a time zone and {@code timestamp_ntz} is + * {@code timestamp[us]} without one. Data read from a library may also use + * {@code large_utf8}, {@code utf8_view} and the other variants of the same type, and any names + * for the fields of the entries of a map, but not unsigned integers, {@code decimal256}, list + * views or dictionary encoding.</li> + * <li><b>Options.</b> Option keys are passed in lower case, because Spark matches them + * case-insensitively.</li> + * <li><b>Optional methods.</b> A library only needs to export the functions for the operations + * it supports. If a function is missing, Spark treats the operation as unsupported, except for + * the pushdown methods and the {@code serialize} methods, which then push nothing down and + * return no state.</li> + * <li><b>Threads.</b> Spark calls the library from multiple threads, and calls on different + * handles may run concurrently. Calls on the same handle never run concurrently. A handle may + * be released from a thread other than the threads that used it, such as a JVM + * {@code Cleaner} thread when Spark releases a handle once it is no longer referenced, for + * example the last stream writer of a streaming query.</li> + * </ul> + * + * <h2>Compatibility</h2> + * The JVM binds the functions of a library by their names only, so the name of this class, the + * names of its methods and their signatures are the interface. Within an {@link #ABI_VERSION}: + * <ul> + * <li>the signature and the semantics of an existing method never change;</li> + * <li>new methods may be added, only as optional methods, so that a library that does not + * export them keeps working as before;</li> + * <li>new values may appear where this documentation says so, such as the names of pushed + * predicates.</li> + * </ul> + * An incompatible change uses a new {@link #ABI_VERSION}, with new names or a new class. A + * library reports the version it implements with {@link #abiVersion}, and Spark rejects a library + * of a version it does not support. Spark supports version 1 at least in all the releases of + * Spark 4 from 4.4.0. + * + * @since 4.4.0 + */ +@Evolving +public final class NativeBridge { + + /** + * The version of this interface. It changes when the interface changes incompatibly: see + * Compatibility above. + */ + public static final int ABI_VERSION = 1; + + private NativeBridge() {} + + /** + * Loads the library that implements the native methods of this copy of the class. Spark calls + * it on the copy of the class it creates for the library. + */ + static void load(String path) { + System.load(path); + } + + /** Returns the version of this interface that the library implements, {@link #ABI_VERSION}. */ + public static native int abiVersion(); + + // --------------------------------------------------------------------------------------------- + // Data source (driver) + // --------------------------------------------------------------------------------------------- + + /** + * Creates a data source with the options given by the user, such as + * {@code spark.read.option("key", "value")}. + * + * @param name the name of the data source, as listed in the package manifest. A library can + * implement several data sources. + * @param optionKeys the option keys, in lower case + * @param optionValues the option values, in the same order as the keys + * @return the data source handle + */ + public static native long createDataSource( + String name, String[] optionKeys, String[] optionValues); + + /** + * Exports the schema of the data. Spark calls it only when the user does not specify a schema. + * + * @param schemaAddress the address of the {@code ArrowSchema} to initialize with a struct type + * whose fields are the columns + */ + public static native void schema(long dataSource, long schemaAddress); + + /** + * Creates a reader for a batch scan. + * + * @param schemaAddress the address of the {@code ArrowSchema} of the schema to read: the one the + * user specified, or the one returned by {@link #schema} + * @return the reader handle + */ + public static native long createReader(long dataSource, long schemaAddress); + + /** + * Creates a reader for a micro-batch streaming scan. + * + * @param schemaAddress the address of the {@code ArrowSchema} of the schema to read: the one the + * user specified, or the one returned by {@link #schema} + * @return the stream reader handle + */ + public static native long createStreamReader(long dataSource, long schemaAddress); + + /** + * Creates a writer for a batch write. + * + * @param schemaAddress the address of the {@code ArrowSchema} of the data to write + * @param overwrite whether to replace the existing data (save mode "overwrite") instead of + * appending to it + * @return the writer handle + */ + public static native long createWriter(long dataSource, long schemaAddress, boolean overwrite); + + /** + * Creates a writer for a streaming write. + * + * @param schemaAddress the address of the {@code ArrowSchema} of the data to write + * @param overwrite whether to replace the existing data in each micro-batch (output mode + * "complete") instead of appending to it + * @return the writer handle + */ + public static native long createStreamWriter( + long dataSource, long schemaAddress, boolean overwrite); + + /** Releases a data source handle. */ + public static native void closeDataSource(long dataSource); + + // --------------------------------------------------------------------------------------------- + // Batch reader (driver) + // --------------------------------------------------------------------------------------------- + + /** + * Pushes predicates down to the reader. The predicates are combined with AND. + * <p> + * Each predicate is a JSON expression tree. An expression is one of: + * <ul> + * <li>{@code {"type": "column", "name": ["a", "b"]}}: a column. A nested field has a name + * part for each level.</li> + * <li>{@code {"type": "literal", "dataType": "bigint", "value": 1}}: a literal. The data type + * is a Spark SQL type name such as {@code int}, {@code string} or {@code decimal(10,2)}. The + * value is {@code null}, a JSON boolean or number for boolean and numeric types, a string + * for {@code string}, a base64 string for {@code binary}, a decimal string for + * {@code decimal}, the days since the epoch for {@code date}, and the microseconds since the + * epoch for {@code timestamp} and {@code timestamp_ntz}.</li> + * <li>{@code {"type": "function", "name": ">", "children": [...]}}: a predicate, with the + * name it has in {@code org.apache.spark.sql.connector.expressions.filter.Predicate}. The + * names are {@code =}, {@code <>}, {@code <=>} (equal, with null equal to null), {@code <}, + * {@code <=}, {@code >} and {@code >=} with two children, {@code IN} with the value to look + * for and then the values to look in, {@code IS_NULL} and {@code IS_NOT_NULL} with one + * child, {@code STARTS_WITH}, {@code ENDS_WITH} and {@code CONTAINS} with the string and the + * string to look for, {@code AND} and {@code OR} with two predicates, and {@code NOT} with + * one. The children of the other functions are columns, literals or predicates.</li> + * </ul> + * Spark does not pass predicates that it cannot express this way, such as predicates that use + * arithmetic or other functions, nor predicates on columns that contain strings with a + * non-binary collation, such as {@code UTF8_LCASE}: strings are compared by their UTF-8 bytes. + * A later version of Spark may pass predicates with other names without changing + * {@link #ABI_VERSION}, so a reader must not report a predicate as evaluated completely if it + * contains a name that the reader does not know, or a data type that it does not support. + * <p> + * A predicate that the reader evaluates completely must follow the semantics of Spark SQL. In + * particular, for floating-point columns, NaN is equal to NaN and greater than any other value, + * including positive infinity, and {@code -0.0} is equal to {@code 0.0}, unlike IEEE 754 + * comparisons. For example, {@code d > 1.0} is true for a NaN value of {@code d}. Comparisons + * with null are null, so a row only matches if the predicate is true. + * + * @return for each predicate, whether the reader evaluates it completely, so that Spark does + * not have to evaluate it again. The reader can still use the other predicates to skip + * data. + */ + public static native boolean[] pushPredicates(long reader, String[] predicates); + + /** + * Pushes a LIMIT down to the reader. + * + * @return whether the reader returns at most {@code limit} rows in total. Spark applies the + * limit again after reading either way. + */ + public static native boolean pushLimit(long reader, int limit); + + /** + * Prunes the columns to read. + * + * @param columnNames the names of the top-level columns that the query needs Review Comment: **Nit (P3):** Spark computes these names after `pushPredicates` has removed the fully evaluated predicates from the plan (`V2ScanRelationPushDown.pruneColumns`). So a column that only an accepted predicate uses is not in the list, and `count(*)` passes an empty list. For `SELECT name FROM t WHERE id > 5`, a reader accepts `id > 5` and then gets `pruneColumns(["name"])`. A library that applies its accepted predicates to the projected batches, which is natural with Arrow compute kernels, then either fails or returns rows that Spark no longer filters. The test library only works because it turns the predicates on `id` into a range up front. Could `pruneColumns`, `pushPredicates` and `read` say that the pruned columns may leave out the columns of accepted predicates, and may be empty, so a reader must apply its predicates before projecting? **Verification:** - **Inspection:** The Javadoc matches how column pruning computes the required columns after fully pushed filters are removed, including the empty case that a count exercises. ########## sql/core/src/main/java/org/apache/spark/sql/datasource/NativeBridge.java: ########## @@ -0,0 +1,442 @@ +/* + * 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.datasource; + +import org.apache.spark.annotation.Evolving; + +/** + * The binary interface between Spark and data sources implemented in native code, such as Rust + * or C++. + * <p> + * A native data source is a shared library that implements the {@code native} methods of this + * class as JNI functions, for example + * {@code Java_org_apache_spark_sql_datasource_NativeBridge_createDataSource}. Spark loads each + * library into its own copy of this class, so any number of libraries that export the same + * functions can be loaded side by side. Spark adapts a native data source to Data Source V2, with + * batch and micro-batch streaming reads, batch and streaming writes, and predicate, column and + * limit pushdown. + * <p> + * Spark finds the library of a data source automatically when it is installed like other native + * libraries and named after the data source: the library of {@code my_source} is + * {@code spark_datasource_my_source}, such as {@code libspark_datasource_my_source.so} in + * {@code /usr/local/lib} on Linux. Spark looks for it in the directories of + * {@code java.library.path}, which include {@code LD_LIBRARY_PATH} on Linux, and in the + * {@code lib} directory of each installation prefix in {@code PATH}. A library can also be + * distributed in a native data source package: a zip file with the extension {@code .sparkpkg} + * that contains a manifest named {@code spark-native-datasource.json} and the library built for + * one or more platforms, added to a session with {@code spark.addArtifact} or under the paths in + * {@code spark.sql.dataSource.native.paths}. See the "Native Data Sources" page of the Spark SQL + * guide for details and examples. + * + * <h2>Lifecycle</h2> + * Spark creates a data source with {@link #createDataSource} on the driver for each schema + * inference, scan and write, and closes it right after. + * <ul> + * <li><b>Batch reads.</b> On the driver, Spark creates a reader with {@link #createReader}, + * pushes operations down by calling, each at most once and in this order, + * {@link #pushPredicates}, {@link #pushLimit} and {@link #pruneColumns}, plans the partitions + * with {@link #partitions}, serializes the state of the reader with {@link #serializeReader}, + * and closes it. Then it calls {@link #read} on the executors, in one task for each + * partition.</li> + * <li><b>Streaming reads.</b> On the driver, Spark creates a stream reader with + * {@link #createStreamReader} and keeps it until the streaming query stops. For each + * micro-batch, it asks for the {@link #latestOffset}, plans the partitions with + * {@link #streamPartitions} and serializes the state with {@link #serializeStreamReader}, then + * calls {@link #read} on the executors. Offsets are JSON strings defined by the library. When + * a streaming query restarts, from its checkpoint, Spark creates a new stream reader and may + * call {@link #streamPartitions} again for the last range that it planned before, so the + * partitions of a range must read the same data each time.</li> + * <li><b>Writes.</b> On the driver, once it computed the data to write, right before it starts + * the tasks of the write (or of a micro-batch), Spark creates a writer with {@link #createWriter} + * or {@link #createStreamWriter}, and serializes its state with {@link #serializeWriter}. If + * computing the data fails, Spark creates no writer. On the executors, each task creates a data Review Comment: **Non-blocking (P2):** This only holds for the stages that adaptive execution runs before the write. `V2TableWriteExec.writeWithV2` calls `createBatchWriterFactory` after `query.execute()`, which returns the RDD of the final stage without computing it. So a projection that fails in the write tasks runs after `createWriter`, and so does every stage when AQE is off, or in a stateful micro-batch, where `MicroBatchExecution` disables AQE. That is why the test `create no writer when computing the data to write fails` needs `.repartition(2)`. Without it, the same `raise_error` write calls `createWriter` and then `abort`. Library authors implement against this text, so they might do irreversible work in `createWriter`, such as truncating for overwrite or starting a transaction, and expect it to be safe. Could this bullet, and the Scaladoc of `ColumnarBatchWrite` and `ColumnarStreamingWrite`, say that the writer is created right before the write tasks start? It would help to add that a failure after that, including in the tasks, leads to `abort`, and that irreversible work belongs in `commit`. This extends an [existing discussion](https://github.com/apache/spark/pull/59313#pullrequestreview-5480473897). **Verification:** - **Inspection:** The writer lifecycle statements match the call order of query.execute(), the writer factory and the write job, with and without adaptive execution and for stateful micro-batches. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala: ########## @@ -0,0 +1,276 @@ +/* + * 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, if it 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]]. Each + // lookup clears it first, so a provider never takes the result of an earlier lookup. + 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 + } + val (invalid, valid) = (configuredPackageFiles(conf) ++ artifactPackageFiles()) + .map(NativeDataSourcePackage.tryRead).partitionMap(identity) + // A package of a version of the interface that Spark does not support is unusable too. + val (unsupported, usable) = valid.partitionMap { pkg => + if (pkg.manifest.abiVersion == NativeBridge.ABI_VERSION) Right(pkg) else Left(pkg) + } + lookupPackage(name, usable) + .orElse(lookupInstalledLibrary(name)) + .orElse { + // An unusable package is only reported when no other package or installed library + // provides the data source: a stale copy does not hide a valid one. An invalid package + // whose manifest cannot be read is only logged, when it is read. + unsupported.find(_.manifest.dataSources.exists(_.equalsIgnoreCase(name))).foreach { pkg => + throw QueryExecutionErrors.invalidNativeDataSourcePackageError( + pkg.path, + "UNSUPPORTED_ABI_VERSION", + Map( + "version" -> pkg.manifest.abiVersion.toString, + "supported" -> NativeBridge.ABI_VERSION.toString)) + } + invalid.find(_.dataSources.exists(_.equalsIgnoreCase(name))).foreach(p => throw p.error) + None + } + } + + /** + * 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 the last [[lookupForProvider]] on this thread, if it was for the given + * name, and forgets it. + */ + def takeLookup(name: String): Option[(String, NativeLibraryLocation)] = { + val last = Option(lastLookup.get()) + lastLookup.remove() + last.collect { case (lookedUp, found) if lookedUp.equalsIgnoreCase(name) => found } + } + + /** Forgets the result of the last [[lookupForProvider]] on this thread. */ + def clearLookup(): Unit = lastLookup.remove() + + private def lookupPackage( + name: String, + packages: Seq[NativeDataSourcePackage]): Option[(String, NativeDataSourcePackage)] = { + val found = packages + .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) => 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) + findInstalledLibraries(libraryFileName(dataSource)).headOption + .map(file => dataSource -> InstalledNativeLibrary(file)) + } + + /** + * Finds the libraries with the given file name in the native library path of this node, in the + * order of the path. + */ + private[ffi] def findInstalledLibraries(fileName: String): Seq[File] = { + librarySearchPath(System.getProperty("java.library.path"), System.getenv("PATH")) + .map(new File(_, fileName)) + .filter(_.isFile) + } + + /** + * The file name of the installed library of the native data source with the given name, such + * as `libspark_datasource_my_source.so` on Linux, `libspark_datasource_my_source.dylib` on + * macOS and `spark_datasource_my_source.dll` on Windows. + */ + def libraryFileName(name: String): String = { + System.mapLibraryName(LIBRARY_NAME_PREFIX + name.toLowerCase(Locale.ROOT)) + } + + /** + * The native library path: the directories where native data source libraries are installed, + * in the order they are searched. They are the directories of `java.library.path`, where the + * JVM finds native libraries, which include `LD_LIBRARY_PATH` on Linux, `DYLD_LIBRARY_PATH` on + * macOS and `PATH` on Windows, and then the `lib` directory of each installation prefix in + * `PATH`, such as `/usr/local/lib` for `/usr/local/bin`. Relative directories are ignored. + */ + private[ffi] def librarySearchPath(javaLibraryPath: String, path: String): Seq[File] = { + def directories(value: String): Seq[File] = { + Option(value).toSeq.flatMap(_.split(File.pathSeparator)).map(new File(_)).filter(_.isAbsolute) + } + val prefixLibraries = directories(path) + .filter(_.getName == "bin") + .flatMap(bin => Option(bin.getParentFile)) + .map(new File(_, "lib")) + (directories(javaLibraryPath) ++ prefixLibraries).distinct + } + + /** + * Makes the library available to the executors of the active session, which runs the jobs of + * the scan or the write: for a streaming query, it is the session of the query, a clone of the + * session that started it. Returns where the executors find it. An installed library is already + * installed on every node. + */ + def distribute(location: NativeLibraryLocation): NativeLibraryLocation = location match { + case pkg: NativeDataSourcePackage => + SparkSession.getActiveSession match { + // In local mode, the executors run in the driver and read the package where it is. + case Some(session) if session.sparkContext.isLocal => pkg + case Some(session) => addToArtifacts(session, pkg) + case None => + logWarning(log"Cannot distribute the native data source package " + + log"${MDC(PATH, pkg.path)} to the executors: there is no active session. The " + + log"executors only find it if it is at the same path on every node.") + pkg + } + case installed: InstalledNativeLibrary => installed + } + + /** + * Makes the package available to the executors of the session, as one of its artifacts. A + * package that the session added with `spark.addArtifact` already is one. Any other package, + * such as a package under `spark.sql.dataSource.native.paths`, is copied to the artifacts under + * a name made of its file name and its checksum: the copy does not have the extension of a + * package, so it is not a package of the session that the lookups find, and two different + * packages with the same file name do not collide. + */ + private[ffi] def addToArtifacts( + session: SparkSession, + pkg: NativeDataSourcePackage): NativeDataSourcePackage = { + val artifactManager = session.artifactManager + val (added, artifactUUID) = artifactManager.getNativeDataSourcePackages + if (added.exists(_.getCanonicalPath == pkg.path)) { + pkg.copy(artifactUUID = artifactUUID) + } else { + val copyName = s"${pkg.fileName}.${pkg.sha256}" + if (!artifactManager.hasFileArtifact(copyName)) { Review Comment: **Non-blocking (P2):** This check-then-act is not atomic. `hasFileArtifact(copyName)` and `addLocalArtifacts` run without a lock, and `ArtifactManager.addArtifact` checks `Files.exists(target)`, copies with `Files.copy` (no `REPLACE_EXISTING`), and only then adds the path to `sparkContextRelativePaths`. So two queries of the same cluster session can first use a package from `spark.sql.dataSource.native.paths` at the same time: two Spark Connect operations, or two threads of a classic application. The second one either sees a partial copy and gets `ARTIFACT_ALREADY_EXISTS`, or also reaches `Files.copy` and gets a `FileAlreadyExistsException`. Either way its query fails during planning. A retry succeeds, so the failures are intermittent. Before this PR, this needed a client to add the same artifact twice at once; now Spark adds it implicitly on first use. Could we serialize this per session, for example by synchronizing on the session's artifact manager around the check and the add? **Verification:** - **Regression:** Concurrent distributions of the same configured package for one session all return the same distributed copy, with no ARTIFACT_ALREADY_EXISTS or FileAlreadyExistsException. ########## sql/api/src/main/scala/org/apache/spark/sql/Artifact.scala: ########## @@ -48,13 +52,15 @@ private[sql] object Artifact { newJarArtifact(targetFilePath, storage) case cf if cf.endsWith(".class") => newClassArtifact(targetFilePath, storage) + case pkg if pkg.endsWith(NATIVE_DATA_SOURCE_PACKAGE_EXTENSION) => Review Comment: **Nit (P3):** With this case, all five public `SparkSession` methods that add artifacts accept `.sparkpkg`: `addArtifact(path)`, `addArtifact(uri)` and `addArtifacts(uri*)` through `resolveArtifacts`, and `addArtifact(bytes, target)` and `addArtifact(source, target)` by the file name of the target. But their Scaladoc in `sql/api/.../SparkSession.scala` still says 'Currently it supports local files with extensions .jar and .class and Apache Ivy URIs' and 'Supported target file extensions are .jar and .class.' The new docs page tells Scala and Java users to add packages with `spark.addArtifact`. Could these five docs say what a `.sparkpkg` does, with a 4.4.0 note, like the PySpark `addArtifacts` docstring? **Verification:** - **Inspection:** The documented extensions of addArtifact match the extensions Artifact.newArtifactFromExtension accepts. ########## docs/sql-native-data-sources.md: ########## @@ -0,0 +1,602 @@ +--- +layout: global +title: Native Data Sources +displayTitle: Native Data Sources +license: | + 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. +--- + +* Table of contents +{:toc} + +## Overview + +A native data source is a data source implemented in native code, such as Rust or C++, without +any JVM code. It is a shared library that implements the +[`NativeBridge`](api/java/org/apache/spark/sql/datasource/NativeBridge.html) binary interface. +Spark finds the library automatically when it is installed like other native libraries, or when a +session adds it in a package, and loads it when a query uses the data source. Spark adapts it to +[Data Source V2](sql-v2-data-sources.html), so it supports: + +- batch and micro-batch streaming reads, +- batch and streaming writes, appending to or overwriting the existing data, +- predicate, column and limit pushdown. + +Data is exchanged only as columnar batches, through the +[Arrow C data interface](https://arrow.apache.org/docs/format/CDataInterface.html), so Spark reads +the data that a library returns without copying it. + +## How It Works + +A native data source library implements the `native` methods of `NativeBridge` as JNI functions, +named after the class, such as `Java_org_apache_spark_sql_datasource_NativeBridge_createDataSource`. +Spark loads each library into its own copy of the class, so any number of native data sources can +be loaded side by side. + +Spark calls the functions of the library as follows: + +- **Data sources**: for each schema inference, scan and write, Spark creates a data source on the + driver with `createDataSource`, which receives the options of the user, and closes it right + after. When the user does not specify a schema, Spark asks for it with `schema`. +- **Batch reads**: on the driver, Spark creates a reader with `createReader`, pushes operations down + with `pushPredicates`, `pushLimit` and `pruneColumns`, plans the partitions with `partitions`, + and serializes the state of the reader with `serializeReader`. Then it calls `read` on the + executors, in one task for each partition, and `read` exports an `ArrowArrayStream`. +- **Streaming reads**: on the driver, Spark creates a stream reader with `createStreamReader`. For + each micro-batch, it asks for the `latestOffset`, plans the partitions with `streamPartitions`, + and calls `read` on the executors. Offsets are JSON strings defined by the library. +- **Writes**: on the driver, Spark creates a writer with `createWriter` or `createStreamWriter`. On + the executors, each task creates a data writer with `createDataWriter`, writes the batches of its + data with `write`, and calls `commitDataWriter`, or `abortDataWriter` if the task or the commit + fails. Then the driver calls `commit` with the messages of the tasks, or `abort` if the write + failed. A micro-batch streaming query creates a new stream writer for each micro-batch, and Spark + closes the previous one when it creates the next one. Because Spark is not told when a streaming + query stops, the last stream writer of a query is closed once it is no longer referenced, or + when the query is restarted. + +The library defines the format of the partitions, of the state of its readers and writers, and of +the commit messages, which Spark sends to the executors and back as byte arrays. A library reports +an error by throwing a Java exception, for example with the JNI function `ThrowNew`, and Spark +reports it as a `NATIVE_DATA_SOURCE_ERROR`. + +A library only exports the functions of the operations it supports. If a function is missing, +Spark treats the operation as unsupported, except for the pushdown and `serialize` functions, which +are optional: + +| Operation | Required functions | Optional functions | +|-----------|--------------------|--------------------| +| All | `abiVersion`, `createDataSource`, `closeDataSource` | `schema` | +| Batch read | `createReader`, `partitions`, `read`, `closeReader` | `pushPredicates`, `pushLimit`, `pruneColumns`, `serializeReader` | +| Streaming read | `createStreamReader`, `initialOffset`, `latestOffset`, `streamPartitions`, `read`, `closeStreamReader` | `serializeStreamReader`, `commitOffset` | +| Batch write | `createWriter`, `commit`, `abort`, `closeWriter`, `createDataWriter`, `write`, `commitDataWriter`, `abortDataWriter` | `serializeWriter` | +| Streaming write | `createStreamWriter`, `commit`, `abort`, `closeWriter`, `createDataWriter`, `write`, `commitDataWriter`, `abortDataWriter` | `serializeWriter` | + +The [Javadoc of `NativeBridge`](api/java/org/apache/spark/sql/datasource/NativeBridge.html) +specifies each function, including the ownership of the Arrow structs and the JSON format of the +pushed predicates. A predicate that a library reports as evaluated completely must follow the +semantics of Spark SQL, for example for NaN values, which are equal to each other and greater +than any other value. The arrays that a library exports must have an offset of 0: copy a sliced +array, such as the result of `RecordBatch::slice` in arrow-rs, before exporting it. Spark also +checks that each array has at least as many values as its parent refers to, but it cannot check +that the buffers of the arrays hold those values, because the C data interface does not carry the +sizes of buffers: a library that exports a buffer that is too short makes Spark read past its end. + +The `get_next` function of the stream that `read` exports may block, for example on network I/O, +but Spark cannot interrupt native code: when a task is killed, Spark stops reading only once +`get_next` returns. Use timeouts for blocking operations, or enable the task reaper with +`spark.task.reaper.enabled` and `spark.task.reaper.killTimeout`, which kills an executor whose +killed task does not end. + +### Compatibility + +The JVM binds the functions of a library by their names only, so the names and the signatures of +the methods of `NativeBridge` are the interface. Its version, 1, is the `abiVersion` that a +library returns and that the manifest of a package declares. Within a version, the signature and +the semantics of an existing function never change, and new functions are only added as optional +functions, so a library keeps working with later releases of Spark. New values may appear where +the Javadoc says so: for example, later releases may push down predicates with other names, which +a library must not report as evaluated completely. An incompatible change comes with a new +version, with new names or a new class. Spark supports version 1 at least in all the releases of +Spark 4 from 4.4.0. + +### Limitations + +Version 1 of the interface does not cover these features of Data Source V2 yet. They can be added +later as optional functions: + +- statistics: Spark estimates the size of a native scan with `spark.sql.defaultSizeInBytes`, so a + native scan is never the broadcast side of a join, unless the query has a broadcast hint; Review Comment: **Nit (P3):** Two statements on this page do not match the code: 1. Here: with adaptive execution, which is on by default, a shuffle stage over a native scan reports its runtime size, and AQE can turn the join into a broadcast hash join below `spark.sql.adaptive.autoBroadcastJoinThreshold`. So 'never the broadcast side' only holds for the initial plan. 2. The Packages section says that Spark 'ignores, and logs' an unsupported package, but `tryRead` only logs the packages that fail to parse. A package of another `abiVersion` parses, and `lookup` drops it without a log unless its manifest lists the requested name. An operator who leaves an outdated package in a configured directory finds nothing in the logs. Could we reword both? The PR description repeats the logging sentence. **Verification:** - **Inspection:** The two statements match AQE's use of runtime statistics and the logging in NativeDataSourcePackage.tryRead and NativeDataSourceRegistry.lookup. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala: ########## @@ -0,0 +1,794 @@ +/* + * 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, CDataDictionaryProvider, Data} +import org.apache.arrow.vector.{FieldVector, VectorSchemaRoot} +import org.apache.arrow.vector.complex.MapVector +import org.apache.arrow.vector.types.pojo.{ArrowType, Field, Schema} + +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.unsafe.Platform +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() Review Comment: **Nit (P3):** This `finally` closes the handle when `library.partitions` or `serializeReader` throws, but leaves `plannedPartitions` null. The reader is shared through the cached `ColumnarScan.batch`, and the lazy vals above it (`ColumnarBatchScan.partitions`, `BatchScanExec.inputPartitions`) run again after an exception. So in classic mode, if `df.collect()` fails once on a transient library error (the test library's `fail=plan`), a second `df.collect()` enters `partitions()` again. It fails with `INTERNAL_ERROR: The native handle was already released.`, which blames Spark, and the DataFrame stays unusable after the cause is gone. Could the reader remember the first failure and rethrow it on later calls, while still releasing the handle? **Verification:** - **Regression:** With the test library's planning failure, two actions on the same DataFrame both fail with NATIVE_DATA_SOURCE_ERROR for the injected failure. ########## sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceClusterSuite.scala: ########## @@ -0,0 +1,147 @@ +/* + * 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.Files + +import scala.collection.mutable.ArrayBuffer +import scala.jdk.CollectionConverters._ + +import org.apache.spark.sql.{QueryTest, Row} +import org.apache.spark.sql.classic.{DataFrame, SparkSession} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.tags.ExtendedSQLTest +import org.apache.spark.util.Utils + +/** + * Tests native data sources in local-cluster mode, where the executors run in JVMs of their own: + * Spark distributes the packages to them, and they find, check and load the libraries + * themselves. [[NativeDataSourceSuite]] tests the rest in local mode, where the executors use the + * packages of the driver. + */ +@ExtendedSQLTest +class NativeDataSourceClusterSuite extends QueryTest { + import NativeDataSourceTestUtils._ + + protected var spark: SparkSession = _ + + // Jobs outside of SQL executions get the files of all the sessions, so the sessions are kept + // until the end of the suite, as in NativeDataSourceSuite. + private val sessions = ArrayBuffer.empty[SparkSession] + + private lazy val library = compile("test_native_datasource") + // Another build, so that its package has another checksum, which the executors did not load. + private lazy val otherLibrary = + compile("test_native_datasource_cluster", "-DNAME_PREFIX=\"cluster_\"") + + override def beforeAll(): Unit = { + super.beforeAll() + spark = SparkSession.builder() + .master("local-cluster[2,1,1024]") + .appName("NativeDataSourceClusterSuite") + .getOrCreate() + } + + override def afterAll(): Unit = { + try { + sessions.clear() + if (spark != null) { + spark.stop() + spark = null + } + } finally { + super.afterAll() + } + } + + private def nativeTest(name: String)(f: => Unit): Unit = test(name) { + assume(toolchain.isDefined, "A C++ compiler and the JNI headers are required.") + try f finally SparkSession.setActiveSession(spark) + } + + // Like Python data sources, native data sources are found in the active session. + private def newSession(): SparkSession = { + val session = spark.newSession() + SparkSession.setActiveSession(session) + sessions += session + session + } + + private def checkRows(df: => DataFrame, expected: Seq[Row]): Unit = { + QueryTest.checkAnswer(df, expected, checkToRDD = false) + } + + private def written(dir: File): Seq[String] = { + dir.listFiles().filter(_.getName.startsWith("part-")) + .flatMap(file => Files.readAllLines(file.toPath).asScala).toSeq.sorted + } + + nativeTest("read and write with a package added as an artifact") { + val session = newSession() + session.addArtifact(createPackage(library, dataSources("")).getPath) + // 4 partitions, so that both executors load the library. + val df = session.read.format("native_range") + .option("end", "20").option("partitions", "4").load().select("id") + checkRows(df, (0L until 20L).map(Row(_))) + + val dir = Utils.createTempDir() + session.range(10).repartition(4).selectExpr("id", "concat('v', id) AS name") + .write.format("native_sink").option("path", dir.getPath).mode("append").save() + assert(written(dir) == (0 until 10).map(i => s"$i,v$i").sorted) + + val streamDir = Utils.createTempDir() + val query = session.readStream.format("native_counter") + .option("max_offset", "12").option("step", "5").load() + .writeStream.format("native_sink") + .option("path", streamDir.getPath) + .option("checkpointLocation", Utils.createTempDir().getPath) + .start() + try { + query.processAllAvailable() + } finally { + query.stop() + } + assert(written(streamDir) == (0 until 12).map(_.toString).sorted) + } + + nativeTest("executors load the distributed copy of a configured package") { + val session = newSession() + val pkg = createPackage(otherLibrary, dataSources("cluster_")) + session.conf.set(SQLConf.NATIVE_DATA_SOURCE_PATHS.key, pkg.getPath) + // Prepares the package on the driver as a query would: Spark reads its manifest, loads its + // library and copies it to the artifacts of the session under a name with its checksum. + val found = NativeDataSourceRegistry.lookup("cluster_native_range", session.sessionState.conf) + .map(_._2).get.asInstanceOf[NativeDataSourcePackage] + NativeLibraries.get(found) + val distributed = + NativeDataSourceRegistry.distribute(found).asInstanceOf[NativeDataSourcePackage] + assert(distributed.artifactUUID.contains(session.sessionUUID)) + assert(distributed.distributedFileName.contains(s"${pkg.getName}.${found.sha256}")) + + // Then the package at the path of the driver becomes unusable for the executors, while the + // driver keeps the package it read, since the file keeps its length and modification time. + // So the query only succeeds if the executors find the distributed copy, and check that it + // is the same package by its checksum. + val lastModified = pkg.lastModified() + Files.write(pkg.toPath, new Array[Byte](pkg.length().toInt)) Review Comment: **Nit (P3):** The comment above says this checks the copy by its checksum, but zeroing the file makes `tryRead` reject it as an invalid zip on the executors before `locate` compares any checksum. If `.exists(_.sha256 == pkg.sha256)` in `NativeLibraries.locate` became `.isRight`, this test and all of NativeDataSourceSuite would still pass, although a node could then load an older valid package with the same file name. Installed libraries have such a test: `find installed libraries automatically` expects `NATIVE_DATA_SOURCE_LIBRARY_NOT_FOUND`. Could we add the same for packages, `NativeLibraries.get` on a valid package with another checksum expecting `NATIVE_DATA_SOURCE_PACKAGE_NOT_FOUND`, and reword this comment? This extends existing discussions [1](https://github.com/apache/spark/pull/59313#issuecomment-6139365807), [2](https://github.com/apache/spark/pull/59313#pullrequestreview-5497999118). **Verification:** - **Regression:** A package location whose checksum differs from the valid package at its path fails with NATIVE_DATA_SOURCE_PACKAGE_NOT_FOUND, and the test fails if locate accepts any valid package regardless of its checksum. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala: ########## @@ -0,0 +1,794 @@ +/* + * 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, CDataDictionaryProvider, Data} +import org.apache.arrow.vector.{FieldVector, VectorSchemaRoot} +import org.apache.arrow.vector.complex.MapVector +import org.apache.arrow.vector.types.pojo.{ArrowType, Field, Schema} + +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.unsafe.Platform +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 { Review Comment: **Nit (P3):** The Handles convention requires a non-zero handle, but nothing checks it, and 0 is also Spark's own 'released' value. A library can return 0 here because its handle table starts at 0, or because it returns 0 on failure without throwing. Then: - `new NativeHandle(0, ...)` makes the first `handle.get` throw `INTERNAL_ERROR: The native handle was already released.` - A 0 from `createDataWriter` makes the first write throw `The data writer was already committed or aborted.` - A 0 from `createDataSource` is passed on to the next call into the library. Could each create call fail with a `NATIVE_DATA_SOURCE_ERROR` that names the function when it gets 0? **Verification:** - **Regression:** A create function that returns 0 fails the scan or the write with NATIVE_DATA_SOURCE_ERROR, and the library is not called with the zero handle. -- 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]
