HyukjinKwon commented on code in PR #59424: URL: https://github.com/apache/spark/pull/59424#discussion_r4299567529
########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala: ########## @@ -0,0 +1,980 @@ +/* + * 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 +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.atomic.AtomicLong + +import scala.jdk.CollectionConverters._ +import scala.util.control.NonFatal + +import com.fasterxml.jackson.core.JsonProcessingException +import com.fasterxml.jackson.databind.DeserializationFeature +import com.fasterxml.jackson.databind.json.JsonMapper +import org.apache.arrow.c.{ArrowArray, ArrowArrayStream, ArrowSchema, CDataDictionaryProvider, Data} +import org.apache.arrow.vector.{FieldVector, ValueVector, 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, TaskContext} +import org.apache.spark.internal.Logging +import org.apache.spark.sql.connector.catalog.TableCapability +import org.apache.spark.sql.connector.catalog.TableCapability.{BATCH_READ, BATCH_WRITE, MICRO_BATCH_READ, STREAMING_WRITE, TRUNCATE} +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.connector.read.{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.{ColumnarBatchPlanner, ColumnarDataSource, ColumnarOnlyPartitionReaderFactory, ColumnarPlannedBatch, ColumnarStreamReader, ColumnarStreamWriter, ColumnarWriter} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{DataType, StructField, 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 the C interface in + * `org/apache/spark/sql/datasource/spark_datasource.h`. + * + * @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.handle, address))) + }.getOrElse(throw QueryCompilationErrors.dataSchemaNotSpecifiedError(name)) + } + } + + /** + * Returns the batch writes that the data source declares, and always the reads and the + * streaming writes: Data Source V2 handles a table without those as a Data Source V1 source, + * which then fails with errors that do not say why. A scan or a streaming write that the data + * source does not declare fails when Spark plans it instead, on the driver. + */ + override def capabilities(): util.Set[TableCapability] = { + val library = NativeLibraries.get(location) + val declared = withDataSource(library, "check the capabilities of")(_.capabilities) + val capabilities = util.EnumSet.of(BATCH_READ, MICRO_BATCH_READ, STREAMING_WRITE) + if ((declared & NativeLibrary.BATCH_WRITE) != 0) { + capabilities.add(BATCH_WRITE) + } + if ((declared & NativeLibrary.TRUNCATE) != 0) { + capabilities.add(TRUNCATE) + } + capabilities + } + + override def batchPlanner(schema: StructType): ColumnarBatchPlanner = { + val executorLocation = NativeDataSourceRegistry.distribute(location) + val library = NativeLibraries.get(location) + val handle = create(library, "plan a scan of", NativeLibrary.BATCH_READ, + "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.BatchReadFunctions) { + dataSource => + NativeArrow.exportSchema(schema) { address => + library.createReader(dataSource.handle, address) + } + } + new NativeBatchPlanner( + library, 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", NativeLibrary.MICRO_BATCH_READ, + "DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.MicroBatchReadFunctions) { + dataSource => + NativeArrow.exportSchema(schema) { address => + library.createStreamReader(dataSource.handle, 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( + createWriter(library, schema, overwrite, queryId = null), 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(createWriter(library, schema, overwrite, queryId), 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: NativeLibrary.DataSourceHandle => T): T = { + // Checked first, so that a data source is never created without a way to release it. + requireFunctions(library, action, Seq(NativeLibrary.Function.CLOSE)) + // 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.handle)) + } + } + + /** + * Creates a reader or a writer with `f`, which fails if the data source does not declare the + * capability, or if the library does not implement the functions that it requires. + */ + private def create( + library: NativeLibrary, + action: String, + capability: Long, + unsupportedError: String, + functions: Seq[NativeLibrary.Function])( + f: NativeLibrary.DataSourceHandle => Long): Long = { + withDataSource(library, action) { dataSource => + if ((dataSource.capabilities & capability) == 0) { + throw new SparkUnsupportedOperationException(unsupportedError, Map("description" -> name)) + } + requireFunctions(library, action, functions) + call(name, action)(f(dataSource)) + } + } + + /** Creates a writer for a batch write, or for a streaming write with the ID of its query. */ + private def createWriter( + library: NativeLibrary, + schema: StructType, + overwrite: Boolean, + queryId: String): Long = { + val action = "plan a write to" + val (capability, unsupportedError, overwriteMode) = if (queryId == null) { + (NativeLibrary.BATCH_WRITE, "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED", + "the save mode overwrite") + } else { + (NativeLibrary.STREAMING_WRITE, "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED", + "the output mode complete") + } + create(library, action, capability, unsupportedError, NativeDataSource.WriteFunctions) { + dataSource => + if (overwrite && (dataSource.capabilities & NativeLibrary.TRUNCATE) == 0) { + throw error(name, action, s"The data source does not support $overwriteMode.") + } + val mode = if (overwrite) NativeLibrary.WRITE_TRUNCATE else NativeLibrary.WRITE_APPEND + NativeArrow.exportSchema(schema) { address => + library.createWriter(dataSource.handle, address, mode, queryId) + } + } + } + + private def requireFunctions( + library: NativeLibrary, + action: String, + functions: Seq[NativeLibrary.Function]): Unit = { + functions.find(!library.hasFunction(_)).foreach { function => + throw error(name, action, s"The library does not implement ${function.fieldName}.") + } + } + + 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 + } + } +} + +object NativeDataSource { + import NativeLibrary.Function._ + + // The functions that the operations of a data source require, besides create and close. The + // others are optional. They are checked on the driver, before the operation starts. + private val BatchReadFunctions = Seq(CREATE_READER, PARTITIONS, CLOSE_READER, READ) + private val MicroBatchReadFunctions = Seq(CREATE_STREAM_READER, INITIAL_OFFSET, LATEST_OFFSET, + STREAM_PARTITIONS, CLOSE_STREAM_READER, READ) + private val WriteFunctions = Seq(CREATE_WRITER, COMMIT, ABORT, CLOSE_WRITER, CREATE_DATA_WRITER, + WRITE, COMMIT_DATA_WRITER, ABORT_DATA_WRITER, CLOSE_DATA_WRITER) +} + +/** + * Plans a batch scan of a native data source, on the driver, with a reader of the library. The + * reader is released once the scan is planned, or when the planner becomes unreachable if the + * scan is never planned: the executors read the partitions with the state that it serialized. + * + * @param executorLocation where the executors find the library + */ +class NativeBatchPlanner( + library: NativeLibrary, + executorLocation: NativeLibraryLocation, + name: String, + schema: StructType, + handle: NativeHandle) + extends ColumnarBatchPlanner { + import NativeCalls._ + + private var readSchema: StructType = schema + + 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 used = if (encoded.isEmpty) { + Set.empty[Int] + } else { + // How the reader uses each predicate: whether it uses it completely or only to skip data + // does not matter, since Spark evaluates all the predicates again after reading. + call(name, "plan a scan of") { + optional(library.pushPredicates(handle.get, encoded.map(_._1))) + }.map { uses => + encoded.map(_._2).zip(uses).collect { + case (i, use) if use != NativeLibrary.PREDICATE_UNUSED => i + }.toSet + }.getOrElse(Set.empty[Int]) + } + predicates.zipWithIndex.collect { case (predicate, i) if used.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 plan(): ColumnarPlannedBatch = { + try { + val partitions = call(name, "plan a scan of")(library.partitions(handle.get)) + val readerState = bytesOrEmpty( + call(name, "plan a scan of")(optional(library.serializeReader(handle.get)))) + ColumnarPlannedBatch( + NativeInputPartition.from(partitions), + new NativePartitionReaderFactory(executorLocation, name, readerState, readSchema)) + } finally { + // The executors only need the serialized state. + handle.close() + } + } +} + +/** Reads the partitions of a planned batch scan of a native data source, on the executors. */ +class NativePartitionReaderFactory( + location: NativeLibraryLocation, + name: String, + readerState: Array[Byte], + readSchema: StructType) + extends ColumnarOnlyPartitionReaderFactory { + override def createColumnarReader(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 = { + checkOffset("initial_offset", call(name, "plan a scan of")(library.initialOffset(handle.get))) + } + + override def latestOffset(): String = { + checkOffset("latest_offset", call(name, "plan a scan of")(library.latestOffset(handle.get))) + } + + /** + * Checks an offset returned by the given function of the library: it must be JSON on a single + * line, because the offset log of the checkpoint stores each offset as a line. + */ + private def checkOffset(function: String, offset: String): String = { + if (offset.exists(c => c == '\n' || c == '\r')) { + throw error(name, "plan a scan of", s"$function returned an offset with a line break, " + + "but an offset must be JSON on a single line.") + } + val valid = try { + val json = NativeDataSourceStreamReader.mapper.readTree(offset) + json != null && !json.isMissingNode + } catch { + case _: JsonProcessingException => false + } + if (!valid) { + throw error(name, "plan a scan of", s"$function returned an offset that is not valid " + + s"JSON: $offset") + } + offset + } + + override def partitions(start: String, end: String): Array[InputPartition] = { + val partitions = call(name, "plan a scan of") { + library.streamPartitions(handle.get, start, end) + } + readerState = bytesOrEmpty( + call(name, "plan a scan of")(optional(library.serializeStreamReader(handle.get)))) + NativeInputPartition.from(partitions) + } + + override def read(partition: InputPartition): PartitionReader[ColumnarBatch] = { + new NativePartitionReader( + location, name, readerState, partition.asInstanceOf[NativeInputPartition].bytes, schema) + } + + override def commit(end: String): Unit = { + call(name, "commit an offset of")(optional(library.commitOffset(handle.get, end))) + } + + override def stop(): Unit = handle.close() +} + +object NativeDataSourceStreamReader { + // Parses the offsets of the libraries to check them, and rejects anything after the JSON value. + private val mapper = + JsonMapper.builder().enable(DeserializationFeature.FAIL_ON_TRAILING_TOKENS).build() +} + +/** Writes to a native data source for a batch write. */ +class NativeDataSourceWriter( + location: NativeLibraryLocation, + name: String, + @transient private val handle: NativeHandle, + writerState: Array[Byte]) + extends ColumnarWriter { + import NativeCalls._ + + @transient private lazy val library = NativeLibraries.get(location) + + override def createWriter(partitionId: Int, taskId: Long): DataWriter[ColumnarBatch] = { + new NativeDataWriter(location, name, writerState, partitionId, taskId, epochId = -1L) + } + + override def commit(messages: Array[WriterCommitMessage]): Unit = { + call(name, "commit a write to") { + library.commit(handle.get, -1L, NativeWriterCommitMessage.toBytes(messages)) + } + // Not closed if the commit fails: Spark then aborts the write, which closes it. + handle.close() + } + + override def abort(messages: Array[WriterCommitMessage]): Unit = { + Utils.tryWithSafeFinally { + call(name, "abort a write to") { + library.abort(handle.get, -1L, NativeWriterCommitMessage.toBytes(messages)) + } + }(handle.close()) + } +} + +/** + * Writes to a native data source for a streaming write. Its handle is closed when the streaming + * query creates its next stream writer, see [[NativeStreamWriters]]. + */ +class NativeDataSourceStreamWriter( + location: NativeLibraryLocation, + name: String, + @transient private val handle: NativeHandle, + writerState: Array[Byte]) + extends ColumnarStreamWriter { + import NativeCalls._ + + @transient private lazy val library = NativeLibraries.get(location) + + override def createWriter( + partitionId: Int, + taskId: Long, + epochId: Long): DataWriter[ColumnarBatch] = { + new NativeDataWriter(location, name, writerState, partitionId, taskId, epochId) + } + + override def commit(epochId: Long, messages: Array[WriterCommitMessage]): Unit = { + call(name, "commit a write to") { + library.commit(handle.get, epochId, NativeWriterCommitMessage.toBytes(messages)) + } + } + + override def abort(epochId: Long, messages: Array[WriterCommitMessage]): Unit = { + call(name, "abort a write to") { + library.abort(handle.get, epochId, NativeWriterCommitMessage.toBytes(messages)) + } + } +} + +/** A partition planned by a native library. */ +case class NativeInputPartition(index: Int, bytes: Array[Byte]) extends InputPartition + +object NativeInputPartition { + def from(partitions: Array[Array[Byte]]): Array[InputPartition] = { + Option(partitions).getOrElse(Array.empty).zipWithIndex.map { case (bytes, i) => + NativeInputPartition(i, Option(bytes).getOrElse(Array.emptyByteArray)) + } + } +} + +/** The message returned by the data writer of a task of a native library. */ +case class NativeWriterCommitMessage(bytes: Array[Byte]) extends WriterCommitMessage + +object NativeWriterCommitMessage { + def toBytes(messages: Array[WriterCommitMessage]): Array[Array[Byte]] = messages.map { + case message: NativeWriterCommitMessage => message.bytes + case null => null + case other => + throw SparkException.internalError(s"Unexpected commit message: ${other.getClass.getName}") + } +} + +/** + * Reads a partition of a native data source: each array of the Arrow stream exported by the + * library is checked, and imported into the Arrow vectors that back the returned batch. + */ +class NativePartitionReader( + location: NativeLibraryLocation, + name: String, + readerState: Array[Byte], + partition: Array[Byte], + schema: StructType) + extends PartitionReader[ColumnarBatch] { + import NativeCalls._ + + private val allocator = + ArrowUtils.rootAllocator.newChildAllocator(s"native data source $name", 0, Long.MaxValue) + // Read here rather than with `Data.importArrayStream`, to check each array before importing it. + private val stream = ArrowArrayStream.allocateNew(allocator) + private val dictionaries = new CDataDictionaryProvider() + // The task that reads the partition, which the library may use until the stream is released. + private var task: NativeLibrary.Task = _ + private var root: VectorSchemaRoot = _ + private var batch: ColumnarBatch = _ + + try { + openStream() + root = VectorSchemaRoot.create(importSchema(), allocator) + checkSchema() + batch = new ColumnarBatch(root.getFieldVectors.asScala.map(columnVector).toArray) + } catch { + case e: Throwable => + close() + throw e + } + + private def openStream(): Unit = { + val library = NativeLibraries.get(location) + task = NativeTask.create() + call(name, "read")(library.read(task, readerState, partition, stream.memoryAddress())) + } + + private def importSchema(): Schema = { + // Arrow calls the callbacks of the stream without checking them, and calling NULL would crash + // the JVM. + val snapshot = stream.snapshot() + if (snapshot.release == 0 || snapshot.get_schema == 0 || snapshot.get_next == 0 || + snapshot.get_last_error == 0) { + throw error(name, "read", "The library did not initialize the stream, or a callback of " + + "the stream is NULL.") + } + val arrowSchema = ArrowSchema.allocateNew(allocator) + try { + call(name, "read")(stream.getSchema(arrowSchema)) + val imported = call(name, "read")(Data.importSchema(allocator, arrowSchema, dictionaries)) + new Schema( + imported.getFields.asScala.map(withMapEntryNames).asJava, imported.getCustomMetadata) + } finally { + arrowSchema.close() + } + } + + /** + * Names the fields of the entries of the maps `key` and `value`, which ArrowColumnVector looks + * up by name. Arrow does not fix these names: arrow-rs names them `keys` and `values`, for + * example. The arrays are imported by position, so the names do not change how they are read. + */ + private def withMapEntryNames(field: Field): Field = { + val children = field.getChildren.asScala.map(withMapEntryNames).toSeq + val renamed = (field.getType, children) match { + case (_: ArrowType.Map, Seq(entries)) if entries.getChildren.size == 2 => + val Seq(key, value) = entries.getChildren.asScala.toSeq + Seq(new Field(entries.getName, entries.getFieldType, Seq( + new Field(MapVector.KEY_NAME, key.getFieldType, key.getChildren), + new Field(MapVector.VALUE_NAME, value.getFieldType, value.getChildren)).asJava)) + case _ => children + } + new Field(field.getName, field.getFieldType, renamed.asJava) + } + + private def checkSchema(): Unit = { + def isDictionaryEncoded(field: Field): Boolean = { + field.getDictionary != null || field.getChildren.asScala.exists(isDictionaryEncoded) + } + root.getSchema.getFields.asScala.find(isDictionaryEncoded).foreach { field => + throw error(name, "read", s"The column ${field.getName} is dictionary-encoded, which " + + "is not supported.") + } + // The Arrow types that convert to Spark types, but that Spark does not read: ArrowColumnVector + // cannot read decimals of other widths than 128 bits, and the lengths of the elements of a + // list view cannot be checked without a pass over all its lists (see checkArrays). + def unsupportedType(field: Field): Option[ArrowType] = field.getType match { + case decimal: ArrowType.Decimal if decimal.getBitWidth != 128 => Some(decimal) + case listView: ArrowType.ListView => Some(listView) + case _ => field.getChildren.asScala.view.flatMap(unsupportedType).headOption + } + root.getSchema.getFields.asScala.foreach { field => + unsupportedType(field).foreach { arrowType => + throw error(name, "read", s"The column ${field.getName} contains the Arrow type " + + s"$arrowType, which is not supported.") + } + } + val actual = StructType(root.getSchema.getFields.asScala.map { field => + val dataType = try { + ArrowUtils.fromArrowField(field) + } catch { + case e: SparkUnsupportedOperationException if e.getCondition == "UNSUPPORTED_ARROWTYPE" => + throw error(name, "read", s"The column ${field.getName} contains the Arrow type " + + s"${e.getMessageParameters.get("typeName")}, which is not supported.") + } + StructField(field.getName, dataType, field.isNullable) + }.toSeq) + // The schema to read as the library got it: Arrow does not carry everything, such as the + // collations of strings, the fields of intervals and user-defined types, which still apply to + // the data that Spark reads. + val expected = ArrowUtils.fromArrowSchema(ArrowUtils.toArrowSchema( + schema, "UTC", errorOnDuplicatedFieldNames = true, largeVarTypes = false)) + val matches = actual.length == expected.length && actual.zip(expected).forall { + case (actualField, field) => + actualField.name == field.name && + DataType.equalsIgnoreNullability(actualField.dataType, field.dataType) + } + if (!matches) { + throw error(name, "read", s"The data has the schema ${actual.toDDL}, but the expected " + + s"schema is ${expected.toDDL}.") + } + } + + private def columnVector(vector: FieldVector): ColumnVector = { + try { + new ReusedArrowColumnVector(vector) + } catch { + case NonFatal(e) => + throw QueryExecutionErrors.nativeDataSourceError("read", name, + s"Spark cannot read the column ${vector.getName} of the Arrow type " + + s"${vector.getField.getType}: ${message(e)}", e) + } + } + + override def next(): Boolean = { + val array = ArrowArray.allocateNew(allocator) + try { + call(name, "read")(stream.getNext(array)) + val snapshot = array.snapshot() + if (snapshot.release == 0) { + // The end of the stream. + false + } else { + try { + checkArrays(snapshot) + } catch { + case e: Throwable => + array.release() + throw e + } + call(name, "read")(Data.importIntoVectorSchemaRoot(allocator, array, root, dictionaries)) Review Comment: Thanks, this is the bug that @FelixYBW reported. A fresh root for each batch, as `Data.importVectorSchemaRoot` creates, would not fix it alone: the reader would still have to release each root at some point, and nothing tells it when a consumer is done with a batch. `ColumnarBatch` lets a reader reuse its vectors, and @sarutak showed that `ArrowCachedBatchSerializer`, which creates a root for each batch and closes the previous one in `next()`, fails in the same way. So the fix is in the consumer, which kept references to the pass-through columns: SPARK-60294, #59551, copies them. This PR keeps one root, and adds a test of Arrow UDFs over a native data source in 74aa0dffc534. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala: ########## @@ -0,0 +1,980 @@ +/* + * 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 +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.atomic.AtomicLong + +import scala.jdk.CollectionConverters._ +import scala.util.control.NonFatal + +import com.fasterxml.jackson.core.JsonProcessingException +import com.fasterxml.jackson.databind.DeserializationFeature +import com.fasterxml.jackson.databind.json.JsonMapper +import org.apache.arrow.c.{ArrowArray, ArrowArrayStream, ArrowSchema, CDataDictionaryProvider, Data} +import org.apache.arrow.vector.{FieldVector, ValueVector, 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, TaskContext} +import org.apache.spark.internal.Logging +import org.apache.spark.sql.connector.catalog.TableCapability +import org.apache.spark.sql.connector.catalog.TableCapability.{BATCH_READ, BATCH_WRITE, MICRO_BATCH_READ, STREAMING_WRITE, TRUNCATE} +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.connector.read.{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.{ColumnarBatchPlanner, ColumnarDataSource, ColumnarOnlyPartitionReaderFactory, ColumnarPlannedBatch, ColumnarStreamReader, ColumnarStreamWriter, ColumnarWriter} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{DataType, StructField, 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 the C interface in + * `org/apache/spark/sql/datasource/spark_datasource.h`. + * + * @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.handle, address))) + }.getOrElse(throw QueryCompilationErrors.dataSchemaNotSpecifiedError(name)) + } + } + + /** + * Returns the batch writes that the data source declares, and always the reads and the + * streaming writes: Data Source V2 handles a table without those as a Data Source V1 source, + * which then fails with errors that do not say why. A scan or a streaming write that the data + * source does not declare fails when Spark plans it instead, on the driver. + */ + override def capabilities(): util.Set[TableCapability] = { + val library = NativeLibraries.get(location) + val declared = withDataSource(library, "check the capabilities of")(_.capabilities) + val capabilities = util.EnumSet.of(BATCH_READ, MICRO_BATCH_READ, STREAMING_WRITE) + if ((declared & NativeLibrary.BATCH_WRITE) != 0) { + capabilities.add(BATCH_WRITE) + } + if ((declared & NativeLibrary.TRUNCATE) != 0) { + capabilities.add(TRUNCATE) + } + capabilities + } + + override def batchPlanner(schema: StructType): ColumnarBatchPlanner = { + val executorLocation = NativeDataSourceRegistry.distribute(location) + val library = NativeLibraries.get(location) + val handle = create(library, "plan a scan of", NativeLibrary.BATCH_READ, + "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.BatchReadFunctions) { + dataSource => + NativeArrow.exportSchema(schema) { address => + library.createReader(dataSource.handle, address) + } + } + new NativeBatchPlanner( + library, 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", NativeLibrary.MICRO_BATCH_READ, + "DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.MicroBatchReadFunctions) { + dataSource => + NativeArrow.exportSchema(schema) { address => + library.createStreamReader(dataSource.handle, 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( + createWriter(library, schema, overwrite, queryId = null), 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(createWriter(library, schema, overwrite, queryId), 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: NativeLibrary.DataSourceHandle => T): T = { + // Checked first, so that a data source is never created without a way to release it. + requireFunctions(library, action, Seq(NativeLibrary.Function.CLOSE)) + // 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.handle)) + } + } + + /** + * Creates a reader or a writer with `f`, which fails if the data source does not declare the + * capability, or if the library does not implement the functions that it requires. + */ + private def create( + library: NativeLibrary, + action: String, + capability: Long, + unsupportedError: String, + functions: Seq[NativeLibrary.Function])( + f: NativeLibrary.DataSourceHandle => Long): Long = { + withDataSource(library, action) { dataSource => + if ((dataSource.capabilities & capability) == 0) { + throw new SparkUnsupportedOperationException(unsupportedError, Map("description" -> name)) + } + requireFunctions(library, action, functions) + call(name, action)(f(dataSource)) + } + } + + /** Creates a writer for a batch write, or for a streaming write with the ID of its query. */ + private def createWriter( + library: NativeLibrary, + schema: StructType, + overwrite: Boolean, + queryId: String): Long = { + val action = "plan a write to" + val (capability, unsupportedError, overwriteMode) = if (queryId == null) { + (NativeLibrary.BATCH_WRITE, "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED", + "the save mode overwrite") + } else { + (NativeLibrary.STREAMING_WRITE, "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED", + "the output mode complete") + } + create(library, action, capability, unsupportedError, NativeDataSource.WriteFunctions) { + dataSource => + if (overwrite && (dataSource.capabilities & NativeLibrary.TRUNCATE) == 0) { + throw error(name, action, s"The data source does not support $overwriteMode.") + } + val mode = if (overwrite) NativeLibrary.WRITE_TRUNCATE else NativeLibrary.WRITE_APPEND + NativeArrow.exportSchema(schema) { address => + library.createWriter(dataSource.handle, address, mode, queryId) + } + } + } + + private def requireFunctions( + library: NativeLibrary, + action: String, + functions: Seq[NativeLibrary.Function]): Unit = { + functions.find(!library.hasFunction(_)).foreach { function => + throw error(name, action, s"The library does not implement ${function.fieldName}.") + } + } + + 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 + } + } +} + +object NativeDataSource { + import NativeLibrary.Function._ + + // The functions that the operations of a data source require, besides create and close. The + // others are optional. They are checked on the driver, before the operation starts. + private val BatchReadFunctions = Seq(CREATE_READER, PARTITIONS, CLOSE_READER, READ) + private val MicroBatchReadFunctions = Seq(CREATE_STREAM_READER, INITIAL_OFFSET, LATEST_OFFSET, + STREAM_PARTITIONS, CLOSE_STREAM_READER, READ) + private val WriteFunctions = Seq(CREATE_WRITER, COMMIT, ABORT, CLOSE_WRITER, CREATE_DATA_WRITER, + WRITE, COMMIT_DATA_WRITER, ABORT_DATA_WRITER, CLOSE_DATA_WRITER) +} + +/** + * Plans a batch scan of a native data source, on the driver, with a reader of the library. The + * reader is released once the scan is planned, or when the planner becomes unreachable if the + * scan is never planned: the executors read the partitions with the state that it serialized. + * + * @param executorLocation where the executors find the library + */ +class NativeBatchPlanner( + library: NativeLibrary, + executorLocation: NativeLibraryLocation, + name: String, + schema: StructType, + handle: NativeHandle) + extends ColumnarBatchPlanner { + import NativeCalls._ + + private var readSchema: StructType = schema + + 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 used = if (encoded.isEmpty) { + Set.empty[Int] + } else { + // How the reader uses each predicate: whether it uses it completely or only to skip data + // does not matter, since Spark evaluates all the predicates again after reading. + call(name, "plan a scan of") { + optional(library.pushPredicates(handle.get, encoded.map(_._1))) + }.map { uses => + encoded.map(_._2).zip(uses).collect { + case (i, use) if use != NativeLibrary.PREDICATE_UNUSED => i + }.toSet + }.getOrElse(Set.empty[Int]) + } + predicates.zipWithIndex.collect { case (predicate, i) if used.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 plan(): ColumnarPlannedBatch = { + try { + val partitions = call(name, "plan a scan of")(library.partitions(handle.get)) + val readerState = bytesOrEmpty( + call(name, "plan a scan of")(optional(library.serializeReader(handle.get)))) + ColumnarPlannedBatch( + NativeInputPartition.from(partitions), + new NativePartitionReaderFactory(executorLocation, name, readerState, readSchema)) + } finally { + // The executors only need the serialized state. + handle.close() + } + } +} + +/** Reads the partitions of a planned batch scan of a native data source, on the executors. */ +class NativePartitionReaderFactory( + location: NativeLibraryLocation, + name: String, + readerState: Array[Byte], + readSchema: StructType) + extends ColumnarOnlyPartitionReaderFactory { + override def createColumnarReader(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 = { + checkOffset("initial_offset", call(name, "plan a scan of")(library.initialOffset(handle.get))) + } + + override def latestOffset(): String = { + checkOffset("latest_offset", call(name, "plan a scan of")(library.latestOffset(handle.get))) + } + + /** + * Checks an offset returned by the given function of the library: it must be JSON on a single + * line, because the offset log of the checkpoint stores each offset as a line. + */ + private def checkOffset(function: String, offset: String): String = { + if (offset.exists(c => c == '\n' || c == '\r')) { + throw error(name, "plan a scan of", s"$function returned an offset with a line break, " + + "but an offset must be JSON on a single line.") + } + val valid = try { + val json = NativeDataSourceStreamReader.mapper.readTree(offset) + json != null && !json.isMissingNode + } catch { + case _: JsonProcessingException => false + } + if (!valid) { + throw error(name, "plan a scan of", s"$function returned an offset that is not valid " + + s"JSON: $offset") + } + offset + } + + override def partitions(start: String, end: String): Array[InputPartition] = { + val partitions = call(name, "plan a scan of") { + library.streamPartitions(handle.get, start, end) + } + readerState = bytesOrEmpty( + call(name, "plan a scan of")(optional(library.serializeStreamReader(handle.get)))) + NativeInputPartition.from(partitions) + } + + override def read(partition: InputPartition): PartitionReader[ColumnarBatch] = { + new NativePartitionReader( + location, name, readerState, partition.asInstanceOf[NativeInputPartition].bytes, schema) + } + + override def commit(end: String): Unit = { + call(name, "commit an offset of")(optional(library.commitOffset(handle.get, end))) + } + + override def stop(): Unit = handle.close() +} + +object NativeDataSourceStreamReader { + // Parses the offsets of the libraries to check them, and rejects anything after the JSON value. + private val mapper = + JsonMapper.builder().enable(DeserializationFeature.FAIL_ON_TRAILING_TOKENS).build() +} + +/** Writes to a native data source for a batch write. */ +class NativeDataSourceWriter( + location: NativeLibraryLocation, + name: String, + @transient private val handle: NativeHandle, + writerState: Array[Byte]) + extends ColumnarWriter { + import NativeCalls._ + + @transient private lazy val library = NativeLibraries.get(location) + + override def createWriter(partitionId: Int, taskId: Long): DataWriter[ColumnarBatch] = { + new NativeDataWriter(location, name, writerState, partitionId, taskId, epochId = -1L) + } + + override def commit(messages: Array[WriterCommitMessage]): Unit = { + call(name, "commit a write to") { + library.commit(handle.get, -1L, NativeWriterCommitMessage.toBytes(messages)) + } + // Not closed if the commit fails: Spark then aborts the write, which closes it. + handle.close() + } + + override def abort(messages: Array[WriterCommitMessage]): Unit = { + Utils.tryWithSafeFinally { + call(name, "abort a write to") { + library.abort(handle.get, -1L, NativeWriterCommitMessage.toBytes(messages)) + } + }(handle.close()) + } +} + +/** + * Writes to a native data source for a streaming write. Its handle is closed when the streaming + * query creates its next stream writer, see [[NativeStreamWriters]]. + */ +class NativeDataSourceStreamWriter( + location: NativeLibraryLocation, + name: String, + @transient private val handle: NativeHandle, + writerState: Array[Byte]) + extends ColumnarStreamWriter { + import NativeCalls._ + + @transient private lazy val library = NativeLibraries.get(location) + + override def createWriter( + partitionId: Int, + taskId: Long, + epochId: Long): DataWriter[ColumnarBatch] = { + new NativeDataWriter(location, name, writerState, partitionId, taskId, epochId) + } + + override def commit(epochId: Long, messages: Array[WriterCommitMessage]): Unit = { + call(name, "commit a write to") { + library.commit(handle.get, epochId, NativeWriterCommitMessage.toBytes(messages)) + } + } + + override def abort(epochId: Long, messages: Array[WriterCommitMessage]): Unit = { + call(name, "abort a write to") { + library.abort(handle.get, epochId, NativeWriterCommitMessage.toBytes(messages)) + } + } +} + +/** A partition planned by a native library. */ +case class NativeInputPartition(index: Int, bytes: Array[Byte]) extends InputPartition + +object NativeInputPartition { + def from(partitions: Array[Array[Byte]]): Array[InputPartition] = { + Option(partitions).getOrElse(Array.empty).zipWithIndex.map { case (bytes, i) => + NativeInputPartition(i, Option(bytes).getOrElse(Array.emptyByteArray)) + } + } +} + +/** The message returned by the data writer of a task of a native library. */ +case class NativeWriterCommitMessage(bytes: Array[Byte]) extends WriterCommitMessage + +object NativeWriterCommitMessage { + def toBytes(messages: Array[WriterCommitMessage]): Array[Array[Byte]] = messages.map { + case message: NativeWriterCommitMessage => message.bytes + case null => null + case other => + throw SparkException.internalError(s"Unexpected commit message: ${other.getClass.getName}") + } +} + +/** + * Reads a partition of a native data source: each array of the Arrow stream exported by the + * library is checked, and imported into the Arrow vectors that back the returned batch. + */ +class NativePartitionReader( + location: NativeLibraryLocation, + name: String, + readerState: Array[Byte], + partition: Array[Byte], + schema: StructType) + extends PartitionReader[ColumnarBatch] { + import NativeCalls._ + + private val allocator = + ArrowUtils.rootAllocator.newChildAllocator(s"native data source $name", 0, Long.MaxValue) + // Read here rather than with `Data.importArrayStream`, to check each array before importing it. + private val stream = ArrowArrayStream.allocateNew(allocator) + private val dictionaries = new CDataDictionaryProvider() + // The task that reads the partition, which the library may use until the stream is released. + private var task: NativeLibrary.Task = _ + private var root: VectorSchemaRoot = _ + private var batch: ColumnarBatch = _ + + try { + openStream() + root = VectorSchemaRoot.create(importSchema(), allocator) + checkSchema() + batch = new ColumnarBatch(root.getFieldVectors.asScala.map(columnVector).toArray) + } catch { + case e: Throwable => + close() + throw e + } + + private def openStream(): Unit = { + val library = NativeLibraries.get(location) + task = NativeTask.create() + call(name, "read")(library.read(task, readerState, partition, stream.memoryAddress())) + } + + private def importSchema(): Schema = { + // Arrow calls the callbacks of the stream without checking them, and calling NULL would crash + // the JVM. + val snapshot = stream.snapshot() + if (snapshot.release == 0 || snapshot.get_schema == 0 || snapshot.get_next == 0 || + snapshot.get_last_error == 0) { + throw error(name, "read", "The library did not initialize the stream, or a callback of " + + "the stream is NULL.") + } + val arrowSchema = ArrowSchema.allocateNew(allocator) + try { + call(name, "read")(stream.getSchema(arrowSchema)) + val imported = call(name, "read")(Data.importSchema(allocator, arrowSchema, dictionaries)) + new Schema( + imported.getFields.asScala.map(withMapEntryNames).asJava, imported.getCustomMetadata) + } finally { + arrowSchema.close() + } + } + + /** + * Names the fields of the entries of the maps `key` and `value`, which ArrowColumnVector looks + * up by name. Arrow does not fix these names: arrow-rs names them `keys` and `values`, for + * example. The arrays are imported by position, so the names do not change how they are read. + */ + private def withMapEntryNames(field: Field): Field = { + val children = field.getChildren.asScala.map(withMapEntryNames).toSeq + val renamed = (field.getType, children) match { + case (_: ArrowType.Map, Seq(entries)) if entries.getChildren.size == 2 => + val Seq(key, value) = entries.getChildren.asScala.toSeq + Seq(new Field(entries.getName, entries.getFieldType, Seq( + new Field(MapVector.KEY_NAME, key.getFieldType, key.getChildren), + new Field(MapVector.VALUE_NAME, value.getFieldType, value.getChildren)).asJava)) + case _ => children + } + new Field(field.getName, field.getFieldType, renamed.asJava) + } + + private def checkSchema(): Unit = { + def isDictionaryEncoded(field: Field): Boolean = { + field.getDictionary != null || field.getChildren.asScala.exists(isDictionaryEncoded) + } + root.getSchema.getFields.asScala.find(isDictionaryEncoded).foreach { field => + throw error(name, "read", s"The column ${field.getName} is dictionary-encoded, which " + + "is not supported.") + } + // The Arrow types that convert to Spark types, but that Spark does not read: ArrowColumnVector + // cannot read decimals of other widths than 128 bits, and the lengths of the elements of a + // list view cannot be checked without a pass over all its lists (see checkArrays). + def unsupportedType(field: Field): Option[ArrowType] = field.getType match { + case decimal: ArrowType.Decimal if decimal.getBitWidth != 128 => Some(decimal) + case listView: ArrowType.ListView => Some(listView) + case _ => field.getChildren.asScala.view.flatMap(unsupportedType).headOption + } + root.getSchema.getFields.asScala.foreach { field => + unsupportedType(field).foreach { arrowType => + throw error(name, "read", s"The column ${field.getName} contains the Arrow type " + + s"$arrowType, which is not supported.") + } + } + val actual = StructType(root.getSchema.getFields.asScala.map { field => + val dataType = try { + ArrowUtils.fromArrowField(field) + } catch { + case e: SparkUnsupportedOperationException if e.getCondition == "UNSUPPORTED_ARROWTYPE" => + throw error(name, "read", s"The column ${field.getName} contains the Arrow type " + + s"${e.getMessageParameters.get("typeName")}, which is not supported.") + } + StructField(field.getName, dataType, field.isNullable) + }.toSeq) + // The schema to read as the library got it: Arrow does not carry everything, such as the + // collations of strings, the fields of intervals and user-defined types, which still apply to + // the data that Spark reads. + val expected = ArrowUtils.fromArrowSchema(ArrowUtils.toArrowSchema( + schema, "UTC", errorOnDuplicatedFieldNames = true, largeVarTypes = false)) + val matches = actual.length == expected.length && actual.zip(expected).forall { + case (actualField, field) => + actualField.name == field.name && + DataType.equalsIgnoreNullability(actualField.dataType, field.dataType) + } + if (!matches) { + throw error(name, "read", s"The data has the schema ${actual.toDDL}, but the expected " + + s"schema is ${expected.toDDL}.") + } + } + + private def columnVector(vector: FieldVector): ColumnVector = { + try { + new ReusedArrowColumnVector(vector) + } catch { + case NonFatal(e) => + throw QueryExecutionErrors.nativeDataSourceError("read", name, + s"Spark cannot read the column ${vector.getName} of the Arrow type " + + s"${vector.getField.getType}: ${message(e)}", e) + } + } + + override def next(): Boolean = { + val array = ArrowArray.allocateNew(allocator) + try { + call(name, "read")(stream.getNext(array)) + val snapshot = array.snapshot() + if (snapshot.release == 0) { + // The end of the stream. + false + } else { + try { + checkArrays(snapshot) + } catch { + case e: Throwable => + array.release() + throw e + } + call(name, "read")(Data.importIntoVectorSchemaRoot(allocator, array, root, dictionaries)) + batch.setNumRows(root.getRowCount) + true + } + } finally { + array.close() + } + } + + /** + * Checks the array of a batch, and the arrays of its columns at every level, before importing + * them. Arrow Java trusts the arrays, and would read them wrongly otherwise: + * - An array must have an offset of 0. Arrow Java ignores the offsets of the arrays it imports, + * so it would read a sliced array from the start of its buffers. + * - An array must have at least as many values as its parent refers to: a column as many as + * the batch has rows, a field as many as its struct has values, and the elements of a list or + * a map as many as its last offset says. Arrow Java would extend a shorter array with nulls. + * Only the lengths that the library exported show this: the import sets the length of each + * array to what its parent refers to. + * The buffers are not checked, since the C data interface does not carry their sizes: Arrow + * Java takes them from the lengths, so the library must export buffers that hold every value. + */ + private def checkArrays(batch: ArrowArray.Snapshot): Unit = { + def children(array: ArrowArray.Snapshot): Seq[ArrowArray.Snapshot] = { + (0 until array.n_children.toInt).map { i => + ArrowArray.wrap(Platform.getLong(null, array.children + 8L * i)).snapshot() + } + } + // The last offset of a list or a map, in its second buffer. Arrow Java rejects the array if + // the buffer is missing. + def lastOffset(array: ArrowArray.Snapshot): Long = { + val offsets = if (array.length > 0 && array.n_buffers >= 2) { + Platform.getLong(null, array.buffers + 8L) + } else { + 0L + } + if (offsets == 0) 0 else Platform.getInt(null, offsets + 4L * array.length) + } + // The number of values that an array refers to in each of its children. + def referenced(array: ArrowArray.Snapshot, arrowType: ArrowType): Long = arrowType match { + case _: ArrowType.Struct => array.length + case _: ArrowType.List | _: ArrowType.Map => lastOffset(array) + // The other types that Spark reads have no children. List views are rejected before. + case _ => 0 + } + def checkColumn(array: ArrowArray.Snapshot, field: Field, column: String): Unit = { + if (array.offset != 0) { + throw error(name, "read", s"The column $column has an offset of ${array.offset}, but " + + "Spark only reads arrays with an offset of 0. Copy sliced arrays before exporting them.") + } + val length = referenced(array, field.getType) + children(array).zip(field.getChildren.asScala).foreach { case (child, childField) => + if (child.length < length) { + throw error(name, "read", s"The column $column contains an array with a length of " + + s"${child.length}, but its parent refers to $length values.") + } + checkColumn(child, childField, column) + } + } + if (batch.offset != 0) { + throw error(name, "read", s"The batch has an offset of ${batch.offset}, but Spark only " + + "reads arrays with an offset of 0. Copy sliced arrays before exporting them.") + } + children(batch).zip(root.getSchema.getFields.asScala).foreach { case (column, field) => Review Comment: Thanks, fixed in 74aa0dffc534. `checkArrays` now requires the batch, and each array at every level, to have as many children as its type has fields, with non-NULL pointers, before it reads them, so it does not read past the array of pointers either. The test library can now export batches without their columns (`short=children`), and `reject the Arrow data that Spark cannot read` checks that the read fails with `The batch has 0 child arrays, but the schema has 7 columns.` ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala: ########## @@ -0,0 +1,980 @@ +/* + * 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 +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.atomic.AtomicLong + +import scala.jdk.CollectionConverters._ +import scala.util.control.NonFatal + +import com.fasterxml.jackson.core.JsonProcessingException +import com.fasterxml.jackson.databind.DeserializationFeature +import com.fasterxml.jackson.databind.json.JsonMapper +import org.apache.arrow.c.{ArrowArray, ArrowArrayStream, ArrowSchema, CDataDictionaryProvider, Data} +import org.apache.arrow.vector.{FieldVector, ValueVector, 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, TaskContext} +import org.apache.spark.internal.Logging +import org.apache.spark.sql.connector.catalog.TableCapability +import org.apache.spark.sql.connector.catalog.TableCapability.{BATCH_READ, BATCH_WRITE, MICRO_BATCH_READ, STREAMING_WRITE, TRUNCATE} +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.connector.read.{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.{ColumnarBatchPlanner, ColumnarDataSource, ColumnarOnlyPartitionReaderFactory, ColumnarPlannedBatch, ColumnarStreamReader, ColumnarStreamWriter, ColumnarWriter} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{DataType, StructField, 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 the C interface in + * `org/apache/spark/sql/datasource/spark_datasource.h`. + * + * @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.handle, address))) + }.getOrElse(throw QueryCompilationErrors.dataSchemaNotSpecifiedError(name)) + } + } + + /** + * Returns the batch writes that the data source declares, and always the reads and the + * streaming writes: Data Source V2 handles a table without those as a Data Source V1 source, + * which then fails with errors that do not say why. A scan or a streaming write that the data + * source does not declare fails when Spark plans it instead, on the driver. + */ + override def capabilities(): util.Set[TableCapability] = { + val library = NativeLibraries.get(location) + val declared = withDataSource(library, "check the capabilities of")(_.capabilities) + val capabilities = util.EnumSet.of(BATCH_READ, MICRO_BATCH_READ, STREAMING_WRITE) + if ((declared & NativeLibrary.BATCH_WRITE) != 0) { + capabilities.add(BATCH_WRITE) + } + if ((declared & NativeLibrary.TRUNCATE) != 0) { + capabilities.add(TRUNCATE) + } + capabilities + } + + override def batchPlanner(schema: StructType): ColumnarBatchPlanner = { + val executorLocation = NativeDataSourceRegistry.distribute(location) + val library = NativeLibraries.get(location) + val handle = create(library, "plan a scan of", NativeLibrary.BATCH_READ, + "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.BatchReadFunctions) { + dataSource => + NativeArrow.exportSchema(schema) { address => + library.createReader(dataSource.handle, address) + } + } + new NativeBatchPlanner( + library, 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", NativeLibrary.MICRO_BATCH_READ, + "DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED", NativeDataSource.MicroBatchReadFunctions) { + dataSource => + NativeArrow.exportSchema(schema) { address => + library.createStreamReader(dataSource.handle, 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( + createWriter(library, schema, overwrite, queryId = null), 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(createWriter(library, schema, overwrite, queryId), 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: NativeLibrary.DataSourceHandle => T): T = { + // Checked first, so that a data source is never created without a way to release it. + requireFunctions(library, action, Seq(NativeLibrary.Function.CLOSE)) + // 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.handle)) + } + } + + /** + * Creates a reader or a writer with `f`, which fails if the data source does not declare the + * capability, or if the library does not implement the functions that it requires. + */ + private def create( + library: NativeLibrary, + action: String, + capability: Long, + unsupportedError: String, + functions: Seq[NativeLibrary.Function])( + f: NativeLibrary.DataSourceHandle => Long): Long = { + withDataSource(library, action) { dataSource => + if ((dataSource.capabilities & capability) == 0) { + throw new SparkUnsupportedOperationException(unsupportedError, Map("description" -> name)) + } + requireFunctions(library, action, functions) + call(name, action)(f(dataSource)) + } + } + + /** Creates a writer for a batch write, or for a streaming write with the ID of its query. */ + private def createWriter( + library: NativeLibrary, + schema: StructType, + overwrite: Boolean, + queryId: String): Long = { + val action = "plan a write to" + val (capability, unsupportedError, overwriteMode) = if (queryId == null) { + (NativeLibrary.BATCH_WRITE, "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED", + "the save mode overwrite") + } else { + (NativeLibrary.STREAMING_WRITE, "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED", + "the output mode complete") + } + create(library, action, capability, unsupportedError, NativeDataSource.WriteFunctions) { + dataSource => + if (overwrite && (dataSource.capabilities & NativeLibrary.TRUNCATE) == 0) { + throw error(name, action, s"The data source does not support $overwriteMode.") + } + val mode = if (overwrite) NativeLibrary.WRITE_TRUNCATE else NativeLibrary.WRITE_APPEND + NativeArrow.exportSchema(schema) { address => + library.createWriter(dataSource.handle, address, mode, queryId) + } + } + } + + private def requireFunctions( + library: NativeLibrary, + action: String, + functions: Seq[NativeLibrary.Function]): Unit = { + functions.find(!library.hasFunction(_)).foreach { function => + throw error(name, action, s"The library does not implement ${function.fieldName}.") + } + } + + 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 + } + } +} + +object NativeDataSource { + import NativeLibrary.Function._ + + // The functions that the operations of a data source require, besides create and close. The + // others are optional. They are checked on the driver, before the operation starts. + private val BatchReadFunctions = Seq(CREATE_READER, PARTITIONS, CLOSE_READER, READ) + private val MicroBatchReadFunctions = Seq(CREATE_STREAM_READER, INITIAL_OFFSET, LATEST_OFFSET, + STREAM_PARTITIONS, CLOSE_STREAM_READER, READ) + private val WriteFunctions = Seq(CREATE_WRITER, COMMIT, ABORT, CLOSE_WRITER, CREATE_DATA_WRITER, + WRITE, COMMIT_DATA_WRITER, ABORT_DATA_WRITER, CLOSE_DATA_WRITER) +} + +/** + * Plans a batch scan of a native data source, on the driver, with a reader of the library. The + * reader is released once the scan is planned, or when the planner becomes unreachable if the + * scan is never planned: the executors read the partitions with the state that it serialized. + * + * @param executorLocation where the executors find the library + */ +class NativeBatchPlanner( + library: NativeLibrary, + executorLocation: NativeLibraryLocation, + name: String, + schema: StructType, + handle: NativeHandle) + extends ColumnarBatchPlanner { + import NativeCalls._ + + private var readSchema: StructType = schema + + 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 used = if (encoded.isEmpty) { + Set.empty[Int] + } else { + // How the reader uses each predicate: whether it uses it completely or only to skip data + // does not matter, since Spark evaluates all the predicates again after reading. + call(name, "plan a scan of") { + optional(library.pushPredicates(handle.get, encoded.map(_._1))) + }.map { uses => + encoded.map(_._2).zip(uses).collect { + case (i, use) if use != NativeLibrary.PREDICATE_UNUSED => i + }.toSet + }.getOrElse(Set.empty[Int]) + } + predicates.zipWithIndex.collect { case (predicate, i) if used.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 plan(): ColumnarPlannedBatch = { + try { + val partitions = call(name, "plan a scan of")(library.partitions(handle.get)) + val readerState = bytesOrEmpty( + call(name, "plan a scan of")(optional(library.serializeReader(handle.get)))) + ColumnarPlannedBatch( + NativeInputPartition.from(partitions), + new NativePartitionReaderFactory(executorLocation, name, readerState, readSchema)) + } finally { + // The executors only need the serialized state. + handle.close() + } + } +} + +/** Reads the partitions of a planned batch scan of a native data source, on the executors. */ +class NativePartitionReaderFactory( + location: NativeLibraryLocation, + name: String, + readerState: Array[Byte], + readSchema: StructType) + extends ColumnarOnlyPartitionReaderFactory { + override def createColumnarReader(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 = { + checkOffset("initial_offset", call(name, "plan a scan of")(library.initialOffset(handle.get))) + } + + override def latestOffset(): String = { + checkOffset("latest_offset", call(name, "plan a scan of")(library.latestOffset(handle.get))) + } + + /** + * Checks an offset returned by the given function of the library: it must be JSON on a single + * line, because the offset log of the checkpoint stores each offset as a line. + */ + private def checkOffset(function: String, offset: String): String = { + if (offset.exists(c => c == '\n' || c == '\r')) { + throw error(name, "plan a scan of", s"$function returned an offset with a line break, " + + "but an offset must be JSON on a single line.") + } + val valid = try { + val json = NativeDataSourceStreamReader.mapper.readTree(offset) + json != null && !json.isMissingNode + } catch { + case _: JsonProcessingException => false + } + if (!valid) { + throw error(name, "plan a scan of", s"$function returned an offset that is not valid " + + s"JSON: $offset") + } + offset + } + + override def partitions(start: String, end: String): Array[InputPartition] = { + val partitions = call(name, "plan a scan of") { + library.streamPartitions(handle.get, start, end) + } + readerState = bytesOrEmpty( + call(name, "plan a scan of")(optional(library.serializeStreamReader(handle.get)))) + NativeInputPartition.from(partitions) + } + + override def read(partition: InputPartition): PartitionReader[ColumnarBatch] = { + new NativePartitionReader( + location, name, readerState, partition.asInstanceOf[NativeInputPartition].bytes, schema) + } + + override def commit(end: String): Unit = { + call(name, "commit an offset of")(optional(library.commitOffset(handle.get, end))) + } + + override def stop(): Unit = handle.close() +} + +object NativeDataSourceStreamReader { + // Parses the offsets of the libraries to check them, and rejects anything after the JSON value. + private val mapper = + JsonMapper.builder().enable(DeserializationFeature.FAIL_ON_TRAILING_TOKENS).build() +} + +/** Writes to a native data source for a batch write. */ +class NativeDataSourceWriter( + location: NativeLibraryLocation, + name: String, + @transient private val handle: NativeHandle, + writerState: Array[Byte]) + extends ColumnarWriter { + import NativeCalls._ + + @transient private lazy val library = NativeLibraries.get(location) + + override def createWriter(partitionId: Int, taskId: Long): DataWriter[ColumnarBatch] = { + new NativeDataWriter(location, name, writerState, partitionId, taskId, epochId = -1L) + } + + override def commit(messages: Array[WriterCommitMessage]): Unit = { + call(name, "commit a write to") { + library.commit(handle.get, -1L, NativeWriterCommitMessage.toBytes(messages)) + } + // Not closed if the commit fails: Spark then aborts the write, which closes it. + handle.close() + } + + override def abort(messages: Array[WriterCommitMessage]): Unit = { + Utils.tryWithSafeFinally { + call(name, "abort a write to") { + library.abort(handle.get, -1L, NativeWriterCommitMessage.toBytes(messages)) + } + }(handle.close()) + } +} + +/** + * Writes to a native data source for a streaming write. Its handle is closed when the streaming + * query creates its next stream writer, see [[NativeStreamWriters]]. + */ +class NativeDataSourceStreamWriter( + location: NativeLibraryLocation, + name: String, + @transient private val handle: NativeHandle, + writerState: Array[Byte]) + extends ColumnarStreamWriter { + import NativeCalls._ + + @transient private lazy val library = NativeLibraries.get(location) + + override def createWriter( + partitionId: Int, + taskId: Long, + epochId: Long): DataWriter[ColumnarBatch] = { + new NativeDataWriter(location, name, writerState, partitionId, taskId, epochId) + } + + override def commit(epochId: Long, messages: Array[WriterCommitMessage]): Unit = { + call(name, "commit a write to") { + library.commit(handle.get, epochId, NativeWriterCommitMessage.toBytes(messages)) + } + } + + override def abort(epochId: Long, messages: Array[WriterCommitMessage]): Unit = { + call(name, "abort a write to") { + library.abort(handle.get, epochId, NativeWriterCommitMessage.toBytes(messages)) + } + } +} + +/** A partition planned by a native library. */ +case class NativeInputPartition(index: Int, bytes: Array[Byte]) extends InputPartition + +object NativeInputPartition { + def from(partitions: Array[Array[Byte]]): Array[InputPartition] = { + Option(partitions).getOrElse(Array.empty).zipWithIndex.map { case (bytes, i) => + NativeInputPartition(i, Option(bytes).getOrElse(Array.emptyByteArray)) + } + } +} + +/** The message returned by the data writer of a task of a native library. */ +case class NativeWriterCommitMessage(bytes: Array[Byte]) extends WriterCommitMessage + +object NativeWriterCommitMessage { + def toBytes(messages: Array[WriterCommitMessage]): Array[Array[Byte]] = messages.map { + case message: NativeWriterCommitMessage => message.bytes + case null => null + case other => + throw SparkException.internalError(s"Unexpected commit message: ${other.getClass.getName}") + } +} + +/** + * Reads a partition of a native data source: each array of the Arrow stream exported by the + * library is checked, and imported into the Arrow vectors that back the returned batch. + */ +class NativePartitionReader( + location: NativeLibraryLocation, + name: String, + readerState: Array[Byte], + partition: Array[Byte], + schema: StructType) + extends PartitionReader[ColumnarBatch] { + import NativeCalls._ + + private val allocator = + ArrowUtils.rootAllocator.newChildAllocator(s"native data source $name", 0, Long.MaxValue) + // Read here rather than with `Data.importArrayStream`, to check each array before importing it. + private val stream = ArrowArrayStream.allocateNew(allocator) + private val dictionaries = new CDataDictionaryProvider() + // The task that reads the partition, which the library may use until the stream is released. + private var task: NativeLibrary.Task = _ + private var root: VectorSchemaRoot = _ + private var batch: ColumnarBatch = _ + + try { + openStream() + root = VectorSchemaRoot.create(importSchema(), allocator) + checkSchema() + batch = new ColumnarBatch(root.getFieldVectors.asScala.map(columnVector).toArray) + } catch { + case e: Throwable => + close() + throw e + } + + private def openStream(): Unit = { + val library = NativeLibraries.get(location) + task = NativeTask.create() + call(name, "read")(library.read(task, readerState, partition, stream.memoryAddress())) + } + + private def importSchema(): Schema = { + // Arrow calls the callbacks of the stream without checking them, and calling NULL would crash + // the JVM. + val snapshot = stream.snapshot() + if (snapshot.release == 0 || snapshot.get_schema == 0 || snapshot.get_next == 0 || + snapshot.get_last_error == 0) { + throw error(name, "read", "The library did not initialize the stream, or a callback of " + + "the stream is NULL.") + } + val arrowSchema = ArrowSchema.allocateNew(allocator) + try { + call(name, "read")(stream.getSchema(arrowSchema)) + val imported = call(name, "read")(Data.importSchema(allocator, arrowSchema, dictionaries)) + new Schema( + imported.getFields.asScala.map(withMapEntryNames).asJava, imported.getCustomMetadata) + } finally { + arrowSchema.close() + } + } + + /** + * Names the fields of the entries of the maps `key` and `value`, which ArrowColumnVector looks + * up by name. Arrow does not fix these names: arrow-rs names them `keys` and `values`, for + * example. The arrays are imported by position, so the names do not change how they are read. + */ + private def withMapEntryNames(field: Field): Field = { + val children = field.getChildren.asScala.map(withMapEntryNames).toSeq + val renamed = (field.getType, children) match { + case (_: ArrowType.Map, Seq(entries)) if entries.getChildren.size == 2 => + val Seq(key, value) = entries.getChildren.asScala.toSeq + Seq(new Field(entries.getName, entries.getFieldType, Seq( + new Field(MapVector.KEY_NAME, key.getFieldType, key.getChildren), + new Field(MapVector.VALUE_NAME, value.getFieldType, value.getChildren)).asJava)) + case _ => children + } + new Field(field.getName, field.getFieldType, renamed.asJava) + } + + private def checkSchema(): Unit = { + def isDictionaryEncoded(field: Field): Boolean = { + field.getDictionary != null || field.getChildren.asScala.exists(isDictionaryEncoded) + } + root.getSchema.getFields.asScala.find(isDictionaryEncoded).foreach { field => + throw error(name, "read", s"The column ${field.getName} is dictionary-encoded, which " + + "is not supported.") + } + // The Arrow types that convert to Spark types, but that Spark does not read: ArrowColumnVector + // cannot read decimals of other widths than 128 bits, and the lengths of the elements of a + // list view cannot be checked without a pass over all its lists (see checkArrays). + def unsupportedType(field: Field): Option[ArrowType] = field.getType match { + case decimal: ArrowType.Decimal if decimal.getBitWidth != 128 => Some(decimal) + case listView: ArrowType.ListView => Some(listView) + case _ => field.getChildren.asScala.view.flatMap(unsupportedType).headOption + } + root.getSchema.getFields.asScala.foreach { field => + unsupportedType(field).foreach { arrowType => + throw error(name, "read", s"The column ${field.getName} contains the Arrow type " + + s"$arrowType, which is not supported.") + } + } + val actual = StructType(root.getSchema.getFields.asScala.map { field => + val dataType = try { + ArrowUtils.fromArrowField(field) + } catch { + case e: SparkUnsupportedOperationException if e.getCondition == "UNSUPPORTED_ARROWTYPE" => + throw error(name, "read", s"The column ${field.getName} contains the Arrow type " + + s"${e.getMessageParameters.get("typeName")}, which is not supported.") + } + StructField(field.getName, dataType, field.isNullable) + }.toSeq) + // The schema to read as the library got it: Arrow does not carry everything, such as the + // collations of strings, the fields of intervals and user-defined types, which still apply to + // the data that Spark reads. + val expected = ArrowUtils.fromArrowSchema(ArrowUtils.toArrowSchema( + schema, "UTC", errorOnDuplicatedFieldNames = true, largeVarTypes = false)) + val matches = actual.length == expected.length && actual.zip(expected).forall { + case (actualField, field) => + actualField.name == field.name && + DataType.equalsIgnoreNullability(actualField.dataType, field.dataType) Review Comment: Thanks, fixed in 74aa0dffc534. After importing a batch, Spark now checks that the data agrees with the nullability of the schema to read, at every level: a column, a struct field, an element or a map value that the schema declares not nullable must have no null where its parent is not null, and map keys none at all. It starts from the null counts of the vectors, so only a column that has nulls is scanned further. A null fails the read with `NATIVE_DATA_SOURCE_ERROR`, such as `The column name has a null, but the schema to read declares it not nullable.`, and the header says that the data must agree with the nullability of the schema. I checked the data rather than making the inferred schema nullable, since a user can also declare a column `NOT NULL` in the schema to read, and the schema of a table in a catalog also applies to its writes. Test: `reject the data that has nulls where the schema to read has none`. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativePredicates.scala: ########## @@ -0,0 +1,136 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.execution.datasources.v2.ffi + +import java.util.Base64 + +import com.fasterxml.jackson.core.json.JsonWriteFeature +import com.fasterxml.jackson.databind.json.JsonMapper +import com.fasterxml.jackson.databind.node.ObjectNode + +import org.apache.spark.sql.connector.expressions.{Expression, GeneralScalarExpression, Literal, NamedReference} +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.types._ + +/** + * Encodes Data Source V2 predicates as the JSON expression trees passed to + * `push_predicates` of `spark_datasource.h`. + */ +object NativePredicates { + // Writes ASCII only, escaping the other characters, which any JSON parser decodes. + private val mapper = JsonMapper.builder().enable(JsonWriteFeature.ESCAPE_NON_ASCII).build() + + /** + * The names of the predicates that are passed to the libraries, as listed in + * `push_predicates`. They are part of the interface, so a predicate with any other + * name, such as an arithmetic function, is not passed: its semantics, for example on overflow, + * would be part of the interface too, and the names of Data Source V2 grow with new releases. + */ + val FUNCTION_NAMES: Set[String] = Set( + "=", "<>", "<=>", "<", "<=", ">", ">=", "IN", "IS_NULL", "IS_NOT_NULL", + "STARTS_WITH", "ENDS_WITH", "CONTAINS", "AND", "OR", "NOT") + + /** + * Returns the JSON of the predicate on the columns of the given schema, or None if it cannot + * be expressed in JSON with the functions in [[FUNCTION_NAMES]]. + */ + def toJson(predicate: Predicate, schema: StructType): Option[String] = { + encode(predicate, schema).map(mapper.writeValueAsString) + } + + private def encode(expression: Expression, schema: StructType): Option[ObjectNode] = { + expression match { + case reference: NamedReference => + encodeColumn(reference, schema) + + case literal: Literal[_] => + encodeLiteral(literal.value(), literal.dataType()) + + case function: GeneralScalarExpression if FUNCTION_NAMES.contains(function.name()) => + val children = function.children().map(encode(_, schema)) + if (children.forall(_.isDefined)) { + val node = mapper.createObjectNode() + .put("type", "function") + .put("name", function.name()) + val array = node.putArray("children") + children.foreach(child => array.add(child.get)) + Some(node) + } else { + None + } + + case _ => None + } + } + + private def encodeColumn(reference: NamedReference, schema: StructType): Option[ObjectNode] = { + // A source compares strings by their bytes, so a column that contains strings with a + // non-binary collation, such as UTF8_LCASE, is not pushed down. Neither is a column that is + // not in the schema. + val hasNonBinaryCollation = columnType(reference.fieldNames().toSeq, schema).forall { + _.existsRecursively { + case st: StringType => !st.isUTF8BinaryCollation + case _ => false + } + } + if (hasNonBinaryCollation) { + None + } else { + val node = mapper.createObjectNode().put("type", "column") + val name = node.putArray("name") + reference.fieldNames().foreach(name.add) + Some(node) + } + } + + /** The type of the column or the nested field with the given name parts. */ + private def columnType(names: Seq[String], dataType: DataType): Option[DataType] = { + (names, dataType) match { + case (Seq(), _) => Some(dataType) + case (name +: rest, struct: StructType) => + struct.fields.find(_.name == name).flatMap(field => columnType(rest, field.dataType)) + case _ => None + } + } + + private def encodeLiteral(value: Any, dataType: DataType): Option[ObjectNode] = { + val node = mapper.createObjectNode() + .put("type", "literal") + .put("dataType", dataType.catalogString) + if (value == null) { + Some(node.putNull("value")) + } else { + val encoded = (value, dataType) match { + case (b: Boolean, BooleanType) => Some(node.put("value", b)) + case (n: Byte, ByteType) => Some(node.put("value", n.toInt)) + case (n: Short, ShortType) => Some(node.put("value", n)) + case (n: Int, IntegerType | DateType) => Some(node.put("value", n)) + case (n: Long, LongType | TimestampType | TimestampNTZType) => Some(node.put("value", n)) + case (n: Float, FloatType) if !n.isNaN && !n.isInfinite => Some(node.put("value", n)) Review Comment: Thanks, both fixed in 74aa0dffc534: - A float literal is now written as the double of the same value, `0.10000000149011612` for `0.1f`, which a reader that parses JSON numbers as doubles gets exactly, and which a reader that parses it as a float reads back as the same float. The header says to read a literal at the precision of its data type, and that the values of the integer types, dates and timestamps are exact 64-bit integers. `predicate-cases.json` has a float column now, with `f = CAST(0.1 AS FLOAT)` and `f > CAST(0.1 AS FLOAT)`, and writes the floats of its rows the same way. - A string literal that is not valid UTF-8 is no longer pushed, and the header says so. ########## sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h: ########## @@ -0,0 +1,688 @@ +/* + * 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. + */ + +/* + * spark_datasource.h: the binary interface between Spark and data sources implemented in native + * code, such as Rust or C++. Since Spark 4.4.0. + * + * A native data source is a shared library that exports a single C function, spark_ds_init, + * which returns the table of the functions that implement the data source, a SparkDsLibrary. + * Spark calls them with the Foreign Function and Memory API of Java, so native data sources + * require Java 22 or later. Each library has its own table, so any number of libraries 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. + * + * This header is self-contained: it only depends on the C standard library, and defines the + * structs of the Arrow C data interface unless they are already defined, as by Arrow C++. A library + * can copy it into its source tree, or use the copy in the include directory of a Spark + * distribution. A library built against it works with all the releases of Spark that support its + * SPARK_DS_ABI_VERSION. + * + * 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 my_source is + * spark_datasource_my_source, such as libspark_datasource_my_source.so in /usr/local/lib on Linux. + * Spark looks for it in the directories of java.library.path, which include LD_LIBRARY_PATH on + * Linux, and in the lib directory of each installation prefix in PATH. A library can also be + * distributed in a native data source package: a zip file with the extension .sparkpkg that + * contains a manifest named spark-native-datasource.json and the library built for one or more + * platforms, added to a session with spark.addArtifact or under the paths in + * spark.sql.dataSource.native.paths. See the "Native Data Sources" page of the Spark SQL guide for + * details and examples. + * + * Lifecycle + * --------- + * Spark loads a library once in each JVM, and calls spark_ds_init before any other function. It + * creates a data source with create on the driver for each schema inference, scan and write, and + * closes it right after. create also returns the capabilities of the data source: Spark reports + * the operations that a data source does not declare as unsupported, without calling the library. + * Spark also creates a data source while it analyzes a query, to reject a batch write that the + * data source does not support before it computes the data to write. That data source has the + * options of the table, those of CREATE TABLE for a table in a catalog, which the options of a + * write can override, so Spark checks the capabilities of each operation again. + * + * - Batch reads. On the driver, Spark creates a reader with create_reader, pushes operations down + * by calling, each at most once and in this order, push_predicates, push_limit and + * prune_columns, plans the partitions with partitions, serializes the state of the reader with + * serialize_reader, and closes it. Then it calls read on the executors, in one task for each + * partition. If Spark never plans the partitions, for example to only explain the plan of the + * query, it closes the reader once the reader is no longer referenced. + * - Streaming reads. On the driver, Spark creates a stream reader with create_stream_reader and + * keeps it until the streaming query stops. For each micro-batch, it asks for the latest_offset, + * plans the partitions with stream_partitions and serializes the state with + * serialize_stream_reader, then calls read on the executors. Offsets are defined by the library, + * as compact JSON on a single line: Spark stores each offset as a line of the offset log of the + * checkpoint, so it rejects an offset that is not valid JSON or contains a line break. When a + * streaming query restarts, from its checkpoint, Spark creates a new stream reader and may call + * stream_partitions again for the last range that it planned before, so the partitions of a + * range must read the same data each time. + * - Writes. On the driver, right before it starts the tasks of the write (or of a micro-batch), + * Spark creates a writer with create_writer, and serializes its state with serialize_writer. + * With adaptive query execution, Spark first runs the stages of the data to write that end with + * a shuffle, and creates no writer if they fail. The rest of the data, such as its last stage, + * is computed in the tasks of the write, so a writer may still be aborted because computing the + * data failed: do irreversible work, such as removing the existing data for an overwrite, in + * commit. On the executors, each task creates a data writer with create_data_writer, calls write + * for each batch of at most spark.sql.execution.arrow.maxRecordsPerBatch rows, which also ends + * once it holds spark.sql.execution.arrow.maxBytesPerBatch bytes, and calls commit_data_writer, + * or abort_data_writer if the task or the commit fails, before it closes the data writer. Then + * the driver calls commit with the messages of the tasks if they all succeeded, or abort + * otherwise, including when 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 + * 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. The ID of the query, which create_writer gets for + * a streaming write, stays the same when the query restarts from its checkpoint, so the ID of the + * query and the epoch ID identify the output of a micro-batch, also when other queries write to + * the same destination. 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. + * + * Conventions + * ----------- + * - Errors. The functions that return an int return 0 on success, or any other value on failure. + * They then set the message that last_error returns, and Spark reports it as a + * NATIVE_DATA_SOURCE_ERROR. The close functions cannot fail, and they may run after the query, + * from another thread: do the work that can fail a query, such as flushing data, in commit. A + * C++ exception or a Rust panic must not cross the interface: catch it and return an error + * instead. + * - Handles. create, create_reader, create_stream_reader, create_writer and create_data_writer + * return an opaque, non-NULL handle in their last argument, typically a pointer: Spark reports a + * NULL handle as an error. Spark releases each handle exactly once, with the close function of + * its kind: close, close_reader, close_stream_reader, close_writer or close_data_writer, and no + * other function releases a handle. A handle is not always released right after its last use: a + * batch reader whose partitions are never planned and the last stream writer of a query are + * released once they are no longer referenced (see Lifecycle 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. + * - Memory. What Spark passes to a function, such as strings, byte arrays, SparkDsCreateInfo and + * SparkDsWriteInfo, is only valid during the call, except for SparkDsHost and + * SparkDsTaskContext, which document how long they live. Strings are UTF-8 and NUL-terminated. + * The JSON of the predicates is ASCII: it escapes the other characters, which any JSON parser + * decodes. The library returns strings and byte arrays by passing them to a SparkDsSink, which + * copies them. + * - State. 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. + * - Arrow data. Schemas and data are exchanged with the Arrow C data interface + * (https://arrow.apache.org/docs/format/CDataInterface.html), as structs 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. + * - Schemas. Spark types map to Arrow types as in spark.sql.execution.arrow: for example, bigint + * is int64, string is utf8, timestamp is timestamp[us] with a time zone and timestamp_ntz is + * timestamp[us] without one. Data read from a library may also use large_utf8 and utf8_view for + * strings, large_binary and binary_view for binary data, and any names for the fields of the + * entries of a map. Spark rejects the other Arrow types, such as unsigned integers, decimal256, + * large lists, fixed-size lists and list views, and dictionary encoding. + * - Options. Option keys are passed in lower case, because Spark matches them + * case-insensitively. + * - Optional functions. A library only sets the functions of the operations it supports in its + * table, and leaves the others NULL. A data source must not declare a capability whose + * functions are NULL: Spark then fails the operation. The pushdown functions and the serialize + * functions are optional: without them, Spark pushes nothing down, and the reader or writer has + * no state. So are schema, without which the user must specify the schema, commit_offset and + * last_error. + * - Threads. 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 Cleaner thread when Spark + * releases a handle once it is no longer referenced, for example the last stream writer of a + * streaming query. A SparkDsSink must be called on the thread of the call that it was passed + * to, during that call. The functions of SparkDsHost and SparkDsTaskContext, and the callbacks + * of the stream that read initializes, can be called from any thread. The callbacks of the + * stream never run concurrently. + * + * Compatibility + * ------------- + * The function table and the structs that Spark passes to the library start with their size, so + * that they can grow. Within a SPARK_DS_ABI_VERSION: + * - fields are only appended to the structs, including SparkDsLibrary, and the existing fields + * never change, nor do the semantics of the existing functions; + * - Spark only reads the functions that fit in the size of the table of a library, so a library + * built against an earlier header keeps working, with the functions that it does not know as + * NULL. A library must only read the fields that fit in the size of a struct that Spark passes; + * - new values may appear where this documentation says so, such as the names of pushed + * predicates. Spark only uses a new capability or write mode with a data source that declares + * it. + * An incompatible change uses a new SPARK_DS_ABI_VERSION: spark_ds_init gets the newest version + * that Spark supports, the library returns a table of a version that it implements, and Spark + * rejects a table of a version that it does not support. Spark supports version 1 at least in all + * the releases of Spark 4 from 4.4.0. + */ +#ifndef SPARK_DATASOURCE_H +#define SPARK_DATASOURCE_H + +#include <stdbool.h> +#include <stdint.h> + +#ifdef __cplusplus +extern "C" { +#endif + +/* The version of this interface. */ +#define SPARK_DS_ABI_VERSION 1 + +/* Exports spark_ds_init from the library, also when the library hides its other symbols. */ +#if defined(_WIN32) +#define SPARK_DS_EXPORT __declspec(dllexport) +#else +#define SPARK_DS_EXPORT __attribute__((visibility("default"))) +#endif + +/* The Arrow C data interface: https://arrow.apache.org/docs/format/CDataInterface.html */ +#ifndef ARROW_C_DATA_INTERFACE +#define ARROW_C_DATA_INTERFACE + +#define ARROW_FLAG_DICTIONARY_ORDERED 1 +#define ARROW_FLAG_NULLABLE 2 +#define ARROW_FLAG_MAP_KEYS_SORTED 4 + +struct ArrowSchema { + const char* format; + const char* name; + const char* metadata; + int64_t flags; + int64_t n_children; + struct ArrowSchema** children; + struct ArrowSchema* dictionary; + void (*release)(struct ArrowSchema*); + void* private_data; +}; + +struct ArrowArray { + int64_t length; + int64_t null_count; + int64_t offset; + int64_t n_buffers; + int64_t n_children; + const void** buffers; + struct ArrowArray** children; + struct ArrowArray* dictionary; + void (*release)(struct ArrowArray*); + void* private_data; +}; + +#endif /* ARROW_C_DATA_INTERFACE */ + +#ifndef ARROW_C_STREAM_INTERFACE +#define ARROW_C_STREAM_INTERFACE + +struct ArrowArrayStream { + int (*get_schema)(struct ArrowArrayStream*, struct ArrowSchema* out); + int (*get_next)(struct ArrowArrayStream*, struct ArrowArray* out); + const char* (*get_last_error)(struct ArrowArrayStream*); + void (*release)(struct ArrowArrayStream*); + void* private_data; +}; + +#endif /* ARROW_C_STREAM_INTERFACE */ + +/* + * Receives a string or a byte array that the library returns, and copies it. `sink` is the + * argument passed with the sink to the function that returns it. + */ +typedef void (*SparkDsSink)(void* sink, const uint8_t* data, int64_t length); + +/* --------------------------------------------------------------------------------------------- + * What Spark passes to the library + * --------------------------------------------------------------------------------------------- */ + +/* The levels of the messages that a library logs with SparkDsHost.log. */ +#define SPARK_DS_LOG_ERROR 1 +#define SPARK_DS_LOG_WARN 2 +#define SPARK_DS_LOG_INFO 3 +#define SPARK_DS_LOG_DEBUG 4 +#define SPARK_DS_LOG_TRACE 5 + +/* Spark, as spark_ds_init gets it. It stays valid until the JVM exits. */ +typedef struct SparkDsHost { + /* + * The size of this struct in the header that Spark was built with. A later version of Spark may + * append fields, so a library must only read the fields that fit in it. + */ + int64_t size; + /* The newest version of this interface that Spark supports. */ + int32_t abi_version; + /* The version of Spark, such as "4.4.0". */ + const char* spark_version; + /* + * Logs a message, a NUL-terminated UTF-8 string, at one of the SPARK_DS_LOG_* levels: in the log + * of the driver or of the executor, with the path of the library. + */ + void (*log)(int32_t level, const char* message); +} SparkDsHost; + +/* A data source, as create gets it. */ +typedef struct SparkDsCreateInfo { + /* The size of this struct in the header that Spark was built with, as in SparkDsHost. */ + int64_t size; + /* + * The name of the data source, as listed in the package manifest, or in lower case for an + * installed library. A library can implement several data sources. + */ + const char* name; + /* + * The num_options options given by the user, such as spark.read.option("key", "value"), with + * the keys in lower case. + */ + const char* const* option_keys; + const char* const* option_values; + int64_t num_options; +} SparkDsCreateInfo; + +/* + * The task that calls read or create_data_writer, on an executor. It stays valid until the stream + * that read returns is released, or until the data writer is closed. + */ +typedef struct SparkDsTaskContext { + /* The size of this struct in the header that Spark was built with, as in SparkDsHost. */ + int64_t size; + /* The ID of the stage of the task, and the attempt number of the stage. */ + int32_t stage_id; + int32_t stage_attempt_number; + /* The partition that the task reads or writes, and the attempt number of the task. */ + int32_t partition_id; + int32_t attempt_number; + /* The ID of the task attempt, unique in the Spark application. */ + int64_t task_attempt_id; + /* + * Returns whether the task was killed, for example because its query was canceled. Spark cannot + * interrupt native code, so a read or a write that runs for long, or blocks, should check it + * regularly, and fail once it is true. + */ + bool (*is_interrupted)(const struct SparkDsTaskContext* task); +} SparkDsTaskContext; + +/* The modes of a write. */ +#define SPARK_DS_WRITE_APPEND 0 +#define SPARK_DS_WRITE_TRUNCATE 1 + +/* A write, as create_writer gets it. */ +typedef struct SparkDsWriteInfo { + /* The size of this struct in the header that Spark was built with, as in SparkDsHost. */ + int64_t size; + /* + * SPARK_DS_WRITE_APPEND to add the data to the existing data, or SPARK_DS_WRITE_TRUNCATE to + * replace the existing data: with the save mode "overwrite" of a batch write, and in each + * micro-batch with the output mode "complete" of a streaming write. Spark only passes + * SPARK_DS_WRITE_TRUNCATE to a data source with the capability SPARK_DS_TRUNCATE. + */ + int32_t mode; + /* + * The ID of the streaming query of a streaming write, or NULL for a batch write. It stays the + * same when the query restarts from its checkpoint, so with the epoch IDs, it identifies the + * commits of the query: see Writes above. + */ + const char* query_id; +} SparkDsWriteInfo; + +/* --------------------------------------------------------------------------------------------- + * What the library returns to Spark + * --------------------------------------------------------------------------------------------- */ + +/* The capabilities of a data source, which create returns: the operations that it supports. */ +#define SPARK_DS_BATCH_READ 1 /* batch reads, with create_reader */ +#define SPARK_DS_MICRO_BATCH_READ 2 /* micro-batch streaming reads, with create_stream_reader */ +#define SPARK_DS_BATCH_WRITE 4 /* batch writes, with create_writer */ +#define SPARK_DS_STREAMING_WRITE 8 /* streaming writes, with create_writer */ +#define SPARK_DS_TRUNCATE 16 /* writes with the mode SPARK_DS_WRITE_TRUNCATE */ + +/* How a reader uses a predicate pushed with push_predicates. */ +#define SPARK_DS_PREDICATE_UNUSED 0 /* not at all */ +/* only to skip data, so it may return rows for which the predicate is not true */ +#define SPARK_DS_PREDICATE_PARTIAL 1 +/* to only return the rows for which the predicate is true */ +#define SPARK_DS_PREDICATE_COMPLETE 2 + +/* + * The functions of a library, which spark_ds_init returns. A library sets the functions of the + * operations it supports, and leaves the others NULL: see "Optional functions" above. + */ +typedef struct SparkDsLibrary { + /* + * The size of this struct in the header that the library was built with. Spark only reads the + * functions that fit in it. + */ + int64_t size; + /* The version of this interface that the table implements, SPARK_DS_ABI_VERSION. */ + int32_t abi_version; + + /* + * Returns the message of the last error on the calling thread, as a NUL-terminated UTF-8 string + * that stays valid until the next call on that thread. Optional: without it, Spark only reports + * the value that the function returned. + */ + const char* (*last_error)(void); + + /* ------------------------------------------------------------------------------------------- + * Data source (driver) + * ------------------------------------------------------------------------------------------- */ + + /* + * Creates a data source. + * + * info: the name of the data source, and the options given by the user + * capabilities: set to the capabilities of the data source, SPARK_DS_* flags combined with |, + * which may depend on its name and its options. It is 0 when create is called. + * source: set to the data source handle + */ + int (*create)(const SparkDsCreateInfo* info, uint64_t* capabilities, void** source); + + /* + * Exports the schema of the data. Spark calls it only when the user does not specify a schema. + * + * out: initialized with a struct type whose fields are the columns + */ + int (*schema)(void* source, struct ArrowSchema* out); + + /* + * Creates a reader for a batch scan. + * + * schema: the schema to read: the one the user specified, or the one returned by schema + * reader: set to the reader handle + */ + int (*create_reader)(void* source, struct ArrowSchema* schema, void** reader); + + /* + * Creates a reader for a micro-batch streaming scan. + * + * schema: the schema to read: the one the user specified, or the one returned by schema + * reader: set to the stream reader handle + */ + int (*create_stream_reader)(void* source, struct ArrowSchema* schema, void** reader); + + /* + * Creates a writer for a batch write, or for a micro-batch of a streaming write. + * + * schema: the schema of the data to write + * info: the mode of the write, and the ID of the query of a streaming write + * writer: set to the writer handle + */ + int (*create_writer)(void* source, struct ArrowSchema* schema, const SparkDsWriteInfo* info, + void** writer); + + /* Releases a data source handle. */ + void (*close)(void* source); + + /* ------------------------------------------------------------------------------------------- + * Batch reader (driver) + * ------------------------------------------------------------------------------------------- */ + + /* + * Pushes predicates down to the reader. The predicates are combined with AND. + * + * Each predicate is a JSON expression tree. An expression is one of: + * - {"type": "column", "name": ["a", "b"]}: a column. A nested field has a name part for each + * level. + * - {"type": "literal", "dataType": "bigint", "value": 1}: a literal. The data type is a Spark + * SQL type name such as int, string or decimal(10,2). The value is null, a JSON boolean or + * number for boolean and numeric types, a string for string, a base64 string for binary, a + * decimal string for decimal, the days since the epoch for date, and the microseconds since + * the epoch for timestamp and timestamp_ntz. + * - {"type": "function", "name": ">", "children": [...]}: a predicate, with the name it has in + * org.apache.spark.sql.connector.expressions.filter.Predicate. The names are =, <>, <=> + * (equal, with null equal to null), <, <=, > and >= with two children, IN with the value to + * look for and then the values to look in, IS_NULL and IS_NOT_NULL with one child, + * STARTS_WITH, ENDS_WITH and CONTAINS with the string and the string to look for, AND and OR + * with two predicates, and NOT with one. The children of the other functions are columns, + * literals or predicates. + * 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 UTF8_LCASE: strings are compared by their UTF-8 bytes. A later + * version of Spark may pass predicates with other names without changing SPARK_DS_ABI_VERSION, + * so a reader must ignore a predicate that contains a name that it does not know, or a data + * type that it does not support. + * + * Spark evaluates all the predicates again after reading, so a reader may use them only to skip + * data, for example the files whose statistics show that no row matches. But it must never drop + * or skip a row for which the predicates are true under the semantics of Spark SQL, which Spark + * cannot detect. In particular: + * - For floating-point columns, NaN is equal to NaN and greater than any other value, including + * positive infinity, and -0.0 is equal to 0.0. IEEE 754 comparisons treat NaN differently, and + * total orders, such as f64::total_cmp in Rust or the comparison kernels of arrow-rs, also + * order -0.0 before 0.0. For example, d > 1.0 is true for a NaN value of d, and d = 0.0 for + * -0.0. + * - Comparisons with null are null, and a row only matches if the predicates are true. For + * example, NOT (i IN (1, NULL)) matches no row. + * - Strings are compared by their UTF-8 bytes, which order the characters outside the Basic + * Multilingual Plane after all the others, unlike comparisons of UTF-16 strings. + * - To use only a part of a predicate, a reader may take the parts that it does not support as + * true, but not under a NOT, where that would drop rows. Review Comment: Thanks, fixed in 74aa0dffc534. The rule now says that a reader may only take a part as true when every function between that part and the root of the predicate is AND or OR, with your example, and `predicate-cases.json` has `(id > 3) = (id < 3)`, which is true for the row whose id is 3. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceRegistry.scala: ########## @@ -0,0 +1,279 @@ +/* + * 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.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 == NativeLibrary.ABI_VERSION) Right(pkg) else Left(pkg) Review Comment: Thanks, settled in 74aa0dffc534 with your second option: `abiVersion` is the oldest version that the library implements. Spark uses a package whose `abiVersion` is at most the newest version that it supports, and `spark_ds_init` returns a table of a version that both support. So a library that implements 1 and 2 declares 1, and a later Spark that supports only 2 and 3 still loads it, and gets a table of version 2. The docs promise that Spark ignores the fields of the manifest that it does not know, so a later release can add one, for example for the newest version. The manifest must also have a positive `abiVersion`, and names without U+0000. -- 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]
