LuciferYang commented on code in PR #59313: URL: https://github.com/apache/spark/pull/59313#discussion_r4238891216
########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala: ########## @@ -0,0 +1,661 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.execution.datasources.v2.ffi + +import java.lang.ref.{Cleaner, WeakReference} +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.atomic.AtomicLong + +import scala.jdk.CollectionConverters._ +import scala.util.control.NonFatal + +import org.apache.arrow.c.{ArrowArray, ArrowArrayStream, ArrowSchema, Data} +import org.apache.arrow.vector.{FieldVector, VectorSchemaRoot} +import org.apache.arrow.vector.ipc.ArrowReader +import org.apache.arrow.vector.types.pojo.Field + +import org.apache.spark.{SparkException, SparkThrowable, SparkUnsupportedOperationException} +import org.apache.spark.internal.Logging +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader} +import org.apache.spark.sql.connector.write.{DataWriter, WriterCommitMessage} +import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} +import org.apache.spark.sql.execution.datasources.v2.columnar.{ColumnarDataSource, ColumnarReader, ColumnarStreamReader, ColumnarStreamWriter, ColumnarWriter} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{DataType, StructType} +import org.apache.spark.sql.util.{ArrowUtils, CaseInsensitiveStringMap} +import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector} +import org.apache.spark.util.Utils + +/** + * A [[ColumnarDataSource]] implemented by a native library, through + * `org.apache.spark.sql.datasource.NativeBridge`. + * + * @param location where the library is on the driver. The readers and the writers distribute it + * to the executors of the active session when they are created. + * @param name the name of the data source, as listed in the manifest of its package or in lower + * case for an installed library + */ +class NativeDataSource( + location: NativeLibraryLocation, + name: String, + options: CaseInsensitiveStringMap) + extends ColumnarDataSource { + import NativeCalls._ + + override def schema(): StructType = { + val library = NativeLibraries.get(location) + withDataSource(library, "infer the schema of") { dataSource => + call(name, "infer the schema of") { + optional(NativeArrow.importSchema(address => library.schema(dataSource, address))) + }.getOrElse(throw QueryCompilationErrors.dataSchemaNotSpecifiedError(name)) + } + } + + override def reader(schema: StructType): ColumnarReader = { + val executorLocation = NativeDataSourceRegistry.distribute(location) + val library = NativeLibraries.get(location) + val handle = create(library, "plan a scan of", "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED") { + dataSource => + NativeArrow.exportSchema(schema)(address => library.createReader(dataSource, address)) + } + new NativeDataSourceReader( + executorLocation, name, schema, new NativeHandle(handle, library.closeReader)) + } + + override def streamReader(schema: StructType): ColumnarStreamReader = { + val executorLocation = NativeDataSourceRegistry.distribute(location) + val library = NativeLibraries.get(location) + val handle = create(library, "plan a scan of", "DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED") { + dataSource => + NativeArrow.exportSchema(schema)(address => library.createStreamReader(dataSource, address)) + } + new NativeDataSourceStreamReader( + executorLocation, name, schema, new NativeHandle(handle, library.closeStreamReader)) + } + + override def writer(schema: StructType, overwrite: Boolean): ColumnarWriter = { + val executorLocation = NativeDataSourceRegistry.distribute(location) + val library = NativeLibraries.get(location) + val handle = new NativeHandle( + create(library, "plan a write to", "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED") { dataSource => + NativeArrow.exportSchema(schema) { address => + library.createWriter(dataSource, address, overwrite) + } + }, + library.closeWriter) + new NativeDataSourceWriter(executorLocation, name, handle, serializeWriter(library, handle)) + } + + override def streamWriter( + schema: StructType, + overwrite: Boolean, + queryId: String): ColumnarStreamWriter = { + val executorLocation = NativeDataSourceRegistry.distribute(location) + val library = NativeLibraries.get(location) + val handle = new NativeHandle( + create(library, "plan a write to", "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED") { + dataSource => + NativeArrow.exportSchema(schema) { address => + library.createStreamWriter(dataSource, address, overwrite) + } + }, + library.closeWriter) + val writer = new NativeDataSourceStreamWriter( + executorLocation, name, handle, serializeWriter(library, handle)) + NativeStreamWriters.register(queryId, handle) + writer + } + + /** Runs `f` with a new native data source, which is closed afterwards. */ + private def withDataSource[T](library: NativeLibrary, action: String)(f: Long => T): T = { + // The keys of a case-insensitive map are in lower case. + val entries = options.entrySet().asScala.toArray + val dataSource = call(name, "create") { + library.createDataSource(name, entries.map(_.getKey), entries.map(_.getValue)) + } + Utils.tryWithSafeFinally(f(dataSource)) { + call(name, action)(library.closeDataSource(dataSource)) + } + } + + /** Creates a reader or a writer with `f`, which fails if the library does not implement it. */ + private def create(library: NativeLibrary, action: String, unsupportedError: String)( + f: Long => Long): Long = { + withDataSource(library, action) { dataSource => + call(name, action)(optional(f(dataSource))).getOrElse { + throw new SparkUnsupportedOperationException(unsupportedError, Map("description" -> name)) + } + } + } + + private def serializeWriter(library: NativeLibrary, handle: NativeHandle): Array[Byte] = { + try { + bytesOrEmpty(call(name, "plan a write to")(optional(library.serializeWriter(handle.get)))) + } catch { + case e: Throwable => + handle.close() + throw e + } + } +} + +/** Reads a native data source for a batch scan. */ +class NativeDataSourceReader( + location: NativeLibraryLocation, + name: String, + schema: StructType, + @transient private val handle: NativeHandle) + extends ColumnarReader { + import NativeCalls._ + + // Set on the driver when the partitions are planned, and serialized to the executors. + private var readerState: Array[Byte] = Array.emptyByteArray + private var readSchema: StructType = schema + @transient private var plannedPartitions: Array[InputPartition] = _ + + @transient private lazy val library = NativeLibraries.get(location) + + override def pushPredicates(predicates: Array[Predicate]): Array[Predicate] = { + // Only the predicates that can be expressed in JSON are passed to the library. + val encoded = predicates.zipWithIndex.flatMap { case (predicate, i) => + NativePredicates.toJson(predicate, schema).map(_ -> i) + } + val pushed = if (encoded.isEmpty) { + Set.empty[Int] + } else { + call(name, "plan a scan of") { + optional(library.pushPredicates(handle.get, encoded.map(_._1))) + } match { + case Some(flags) if flags != null && flags.length == encoded.length => + encoded.map(_._2).zip(flags).collect { case (i, true) => i }.toSet + case Some(flags) => + throw error(name, "plan a scan of", s"pushPredicates returned " + + s"${Option(flags).map(_.length).getOrElse(0)} values for ${encoded.length} predicates.") + case None => Set.empty[Int] + } + } + predicates.zipWithIndex.collect { case (predicate, i) if !pushed.contains(i) => predicate } + } + + override def pushLimit(limit: Int): Boolean = { + call(name, "plan a scan of")(optional(library.pushLimit(handle.get, limit))).getOrElse(false) + } + + override def pruneColumns(requiredSchema: StructType): Boolean = { + val pruned = call(name, "plan a scan of") { + optional(library.pruneColumns(handle.get, requiredSchema.fieldNames)) + }.getOrElse(false) + if (pruned) { + readSchema = requiredSchema + } + pruned + } + + override def partitions(): Array[InputPartition] = synchronized { + if (plannedPartitions == null) { + try { + val partitions = call(name, "plan a scan of")(library.partitions(handle.get)) + readerState = bytesOrEmpty( + call(name, "plan a scan of")(optional(library.serializeReader(handle.get)))) + plannedPartitions = NativeInputPartition.from(partitions) + } finally { + // The executors only need the serialized state. + handle.close() + } + } + plannedPartitions + } + + override def read(partition: InputPartition): PartitionReader[ColumnarBatch] = { + new NativePartitionReader( + location, name, readerState, partition.asInstanceOf[NativeInputPartition].bytes, readSchema) + } +} + +/** Reads a native data source for a micro-batch streaming scan. */ +class NativeDataSourceStreamReader( + location: NativeLibraryLocation, + name: String, + schema: StructType, + @transient private val handle: NativeHandle) + extends ColumnarStreamReader { + import NativeCalls._ + + // Set on the driver whenever the partitions of a micro-batch are planned, and serialized to + // the executors with the partitions of that micro-batch. + private var readerState: Array[Byte] = Array.emptyByteArray + + @transient private lazy val library = NativeLibraries.get(location) + + override def initialOffset(): String = { + call(name, "plan a scan of")(library.initialOffset(handle.get)) + } + + override def latestOffset(): String = { + call(name, "plan a scan of")(library.latestOffset(handle.get)) + } + + override def partitions(start: String, end: String): Array[InputPartition] = { + val partitions = call(name, "plan a scan of") { + library.streamPartitions(handle.get, start, end) + } + readerState = bytesOrEmpty( + call(name, "plan a scan of")(optional(library.serializeStreamReader(handle.get)))) + NativeInputPartition.from(partitions) + } + + override def read(partition: InputPartition): PartitionReader[ColumnarBatch] = { + new NativePartitionReader( + location, name, readerState, partition.asInstanceOf[NativeInputPartition].bytes, schema) + } + + override def commit(end: String): Unit = { + call(name, "commit an offset of")(optional(library.commitOffset(handle.get, end))) + } + + override def stop(): Unit = handle.close() +} + +/** Writes to a native data source for a batch write. */ +class NativeDataSourceWriter( + location: NativeLibraryLocation, + name: String, + @transient private val handle: NativeHandle, + writerState: Array[Byte]) + extends ColumnarWriter { + import NativeCalls._ + + @transient private lazy val library = NativeLibraries.get(location) + + override def createWriter(partitionId: Int, taskId: Long): DataWriter[ColumnarBatch] = { + new NativeDataWriter(location, name, writerState, partitionId, taskId, epochId = -1L) + } + + override def commit(messages: Array[WriterCommitMessage]): Unit = { + call(name, "commit a write to") { + library.commit(handle.get, -1L, NativeWriterCommitMessage.toBytes(messages)) + } + // 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: the Arrow stream exported by the library is + * imported into Arrow vectors that back the returned batch. + */ +class NativePartitionReader( + location: NativeLibraryLocation, + name: String, + readerState: Array[Byte], + partition: Array[Byte], + schema: StructType) + extends PartitionReader[ColumnarBatch] { + import NativeCalls._ + + private val allocator = + ArrowUtils.rootAllocator.newChildAllocator(s"native data source $name", 0, Long.MaxValue) + private var reader: ArrowReader = _ + private var root: VectorSchemaRoot = _ + private var batch: ColumnarBatch = _ + + try { + reader = openStream() + root = call(name, "read")(reader.getVectorSchemaRoot) + checkSchema() + batch = new ColumnarBatch( + root.getFieldVectors.asScala.map(new ArrowColumnVector(_): ColumnVector).toArray) + } catch { + case e: Throwable => + close() + throw e + } + + private def openStream(): ArrowReader = { + val library = NativeLibraries.get(location) + val stream = ArrowArrayStream.allocateNew(allocator) + try { + call(name, "read")(optional(library.read(readerState, partition, stream.memoryAddress()))) + .getOrElse { + throw new SparkUnsupportedOperationException( + "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", Map("description" -> name)) + } + Data.importArrayStream(allocator, stream) + } catch { + case e: Throwable => + stream.release() + throw e + } finally { + stream.close() + } + } + + private def checkSchema(): Unit = { + def isDictionaryEncoded(field: Field): Boolean = { + field.getDictionary != null || field.getChildren.asScala.exists(isDictionaryEncoded) + } + root.getSchema.getFields.asScala.find(isDictionaryEncoded).foreach { field => + throw error(name, "read", s"The column ${field.getName} is dictionary-encoded, which " + + "is not supported.") + } + val actual = ArrowUtils.fromArrowSchema(root.getSchema) Review Comment: `checkSchema` compares the schemas after converting them to Spark types, so two Arrow layouts pass the check and then fail when `ArrowColumnVector` is built. One is a map whose entry fields aren't named `key`/`value`, which is what arrow-rs `MapBuilder` produces by default (`keys`/`values`). `MapAccessor` looks the children up by name and gets null, so every task fails with an NPE. The other is `decimal256`: it maps to a plain `DecimalType`, but `ArrowColumnVector` only handles `DecimalVector`, and the `_LEGACY_ERROR_TEMP_3160` it throws has an empty message. Neither error is a `NATIVE_DATA_SOURCE_ERROR`, and neither names the column. Could `checkSchema` also check these two (the entry field names, and a decimal bit width of 128), and report a `NATIVE_DATA_SOURCE_ERROR` that names the column? ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala: ########## @@ -0,0 +1,661 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.execution.datasources.v2.ffi + +import java.lang.ref.{Cleaner, WeakReference} +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.atomic.AtomicLong + +import scala.jdk.CollectionConverters._ +import scala.util.control.NonFatal + +import org.apache.arrow.c.{ArrowArray, ArrowArrayStream, ArrowSchema, Data} +import org.apache.arrow.vector.{FieldVector, VectorSchemaRoot} +import org.apache.arrow.vector.ipc.ArrowReader +import org.apache.arrow.vector.types.pojo.Field + +import org.apache.spark.{SparkException, SparkThrowable, SparkUnsupportedOperationException} +import org.apache.spark.internal.Logging +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader} +import org.apache.spark.sql.connector.write.{DataWriter, WriterCommitMessage} +import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} +import org.apache.spark.sql.execution.datasources.v2.columnar.{ColumnarDataSource, ColumnarReader, ColumnarStreamReader, ColumnarStreamWriter, ColumnarWriter} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{DataType, StructType} +import org.apache.spark.sql.util.{ArrowUtils, CaseInsensitiveStringMap} +import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector} +import org.apache.spark.util.Utils + +/** + * A [[ColumnarDataSource]] implemented by a native library, through + * `org.apache.spark.sql.datasource.NativeBridge`. + * + * @param location where the library is on the driver. The readers and the writers distribute it + * to the executors of the active session when they are created. + * @param name the name of the data source, as listed in the manifest of its package or in lower + * case for an installed library + */ +class NativeDataSource( + location: NativeLibraryLocation, + name: String, + options: CaseInsensitiveStringMap) + extends ColumnarDataSource { + import NativeCalls._ + + override def schema(): StructType = { + val library = NativeLibraries.get(location) + withDataSource(library, "infer the schema of") { dataSource => + call(name, "infer the schema of") { + optional(NativeArrow.importSchema(address => library.schema(dataSource, address))) + }.getOrElse(throw QueryCompilationErrors.dataSchemaNotSpecifiedError(name)) + } + } + + override def reader(schema: StructType): ColumnarReader = { + val executorLocation = NativeDataSourceRegistry.distribute(location) + val library = NativeLibraries.get(location) + val handle = create(library, "plan a scan of", "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED") { + dataSource => + NativeArrow.exportSchema(schema)(address => library.createReader(dataSource, address)) + } + new NativeDataSourceReader( + executorLocation, name, schema, new NativeHandle(handle, library.closeReader)) + } + + override def streamReader(schema: StructType): ColumnarStreamReader = { + val executorLocation = NativeDataSourceRegistry.distribute(location) + val library = NativeLibraries.get(location) + val handle = create(library, "plan a scan of", "DATA_SOURCE_MICRO_BATCH_SCAN_NOT_SUPPORTED") { + dataSource => + NativeArrow.exportSchema(schema)(address => library.createStreamReader(dataSource, address)) + } + new NativeDataSourceStreamReader( + executorLocation, name, schema, new NativeHandle(handle, library.closeStreamReader)) + } + + override def writer(schema: StructType, overwrite: Boolean): ColumnarWriter = { + val executorLocation = NativeDataSourceRegistry.distribute(location) + val library = NativeLibraries.get(location) + val handle = new NativeHandle( + create(library, "plan a write to", "DATA_SOURCE_BATCH_WRITE_NOT_SUPPORTED") { dataSource => + NativeArrow.exportSchema(schema) { address => + library.createWriter(dataSource, address, overwrite) + } + }, + library.closeWriter) + new NativeDataSourceWriter(executorLocation, name, handle, serializeWriter(library, handle)) + } + + override def streamWriter( + schema: StructType, + overwrite: Boolean, + queryId: String): ColumnarStreamWriter = { + val executorLocation = NativeDataSourceRegistry.distribute(location) + val library = NativeLibraries.get(location) + val handle = new NativeHandle( + create(library, "plan a write to", "DATA_SOURCE_STREAMING_WRITE_NOT_SUPPORTED") { + dataSource => + NativeArrow.exportSchema(schema) { address => + library.createStreamWriter(dataSource, address, overwrite) + } + }, + library.closeWriter) + val writer = new NativeDataSourceStreamWriter( + executorLocation, name, handle, serializeWriter(library, handle)) + NativeStreamWriters.register(queryId, handle) + writer + } + + /** Runs `f` with a new native data source, which is closed afterwards. */ + private def withDataSource[T](library: NativeLibrary, action: String)(f: Long => T): T = { + // The keys of a case-insensitive map are in lower case. + val entries = options.entrySet().asScala.toArray + val dataSource = call(name, "create") { + library.createDataSource(name, entries.map(_.getKey), entries.map(_.getValue)) + } + Utils.tryWithSafeFinally(f(dataSource)) { + call(name, action)(library.closeDataSource(dataSource)) + } + } + + /** Creates a reader or a writer with `f`, which fails if the library does not implement it. */ + private def create(library: NativeLibrary, action: String, unsupportedError: String)( + f: Long => Long): Long = { + withDataSource(library, action) { dataSource => + call(name, action)(optional(f(dataSource))).getOrElse { + throw new SparkUnsupportedOperationException(unsupportedError, Map("description" -> name)) + } + } + } + + private def serializeWriter(library: NativeLibrary, handle: NativeHandle): Array[Byte] = { + try { + bytesOrEmpty(call(name, "plan a write to")(optional(library.serializeWriter(handle.get)))) + } catch { + case e: Throwable => + handle.close() + throw e + } + } +} + +/** Reads a native data source for a batch scan. */ +class NativeDataSourceReader( + location: NativeLibraryLocation, + name: String, + schema: StructType, + @transient private val handle: NativeHandle) + extends ColumnarReader { + import NativeCalls._ + + // Set on the driver when the partitions are planned, and serialized to the executors. + private var readerState: Array[Byte] = Array.emptyByteArray + private var readSchema: StructType = schema + @transient private var plannedPartitions: Array[InputPartition] = _ + + @transient private lazy val library = NativeLibraries.get(location) + + override def pushPredicates(predicates: Array[Predicate]): Array[Predicate] = { + // Only the predicates that can be expressed in JSON are passed to the library. + val encoded = predicates.zipWithIndex.flatMap { case (predicate, i) => + NativePredicates.toJson(predicate, schema).map(_ -> i) + } + val pushed = if (encoded.isEmpty) { + Set.empty[Int] + } else { + call(name, "plan a scan of") { + optional(library.pushPredicates(handle.get, encoded.map(_._1))) + } match { + case Some(flags) if flags != null && flags.length == encoded.length => + encoded.map(_._2).zip(flags).collect { case (i, true) => i }.toSet + case Some(flags) => + throw error(name, "plan a scan of", s"pushPredicates returned " + + s"${Option(flags).map(_.length).getOrElse(0)} values for ${encoded.length} predicates.") + case None => Set.empty[Int] + } + } + predicates.zipWithIndex.collect { case (predicate, i) if !pushed.contains(i) => predicate } + } + + override def pushLimit(limit: Int): Boolean = { + call(name, "plan a scan of")(optional(library.pushLimit(handle.get, limit))).getOrElse(false) + } + + override def pruneColumns(requiredSchema: StructType): Boolean = { + val pruned = call(name, "plan a scan of") { + optional(library.pruneColumns(handle.get, requiredSchema.fieldNames)) + }.getOrElse(false) + if (pruned) { + readSchema = requiredSchema + } + pruned + } + + override def partitions(): Array[InputPartition] = synchronized { + if (plannedPartitions == null) { + try { + val partitions = call(name, "plan a scan of")(library.partitions(handle.get)) + readerState = bytesOrEmpty( + call(name, "plan a scan of")(optional(library.serializeReader(handle.get)))) + plannedPartitions = NativeInputPartition.from(partitions) + } finally { + // The executors only need the serialized state. + handle.close() + } + } + plannedPartitions + } + + override def read(partition: InputPartition): PartitionReader[ColumnarBatch] = { + new NativePartitionReader( + location, name, readerState, partition.asInstanceOf[NativeInputPartition].bytes, readSchema) + } +} + +/** Reads a native data source for a micro-batch streaming scan. */ +class NativeDataSourceStreamReader( + location: NativeLibraryLocation, + name: String, + schema: StructType, + @transient private val handle: NativeHandle) + extends ColumnarStreamReader { + import NativeCalls._ + + // Set on the driver whenever the partitions of a micro-batch are planned, and serialized to + // the executors with the partitions of that micro-batch. + private var readerState: Array[Byte] = Array.emptyByteArray + + @transient private lazy val library = NativeLibraries.get(location) + + override def initialOffset(): String = { + call(name, "plan a scan of")(library.initialOffset(handle.get)) + } + + override def latestOffset(): String = { + call(name, "plan a scan of")(library.latestOffset(handle.get)) + } + + override def partitions(start: String, end: String): Array[InputPartition] = { + val partitions = call(name, "plan a scan of") { + library.streamPartitions(handle.get, start, end) + } + readerState = bytesOrEmpty( + call(name, "plan a scan of")(optional(library.serializeStreamReader(handle.get)))) + NativeInputPartition.from(partitions) + } + + override def read(partition: InputPartition): PartitionReader[ColumnarBatch] = { + new NativePartitionReader( + location, name, readerState, partition.asInstanceOf[NativeInputPartition].bytes, schema) + } + + override def commit(end: String): Unit = { + call(name, "commit an offset of")(optional(library.commitOffset(handle.get, end))) + } + + override def stop(): Unit = handle.close() +} + +/** Writes to a native data source for a batch write. */ +class NativeDataSourceWriter( + location: NativeLibraryLocation, + name: String, + @transient private val handle: NativeHandle, + writerState: Array[Byte]) + extends ColumnarWriter { + import NativeCalls._ + + @transient private lazy val library = NativeLibraries.get(location) + + override def createWriter(partitionId: Int, taskId: Long): DataWriter[ColumnarBatch] = { + new NativeDataWriter(location, name, writerState, partitionId, taskId, epochId = -1L) + } + + override def commit(messages: Array[WriterCommitMessage]): Unit = { + call(name, "commit a write to") { + library.commit(handle.get, -1L, NativeWriterCommitMessage.toBytes(messages)) + } + // 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: the Arrow stream exported by the library is + * imported into Arrow vectors that back the returned batch. + */ +class NativePartitionReader( + location: NativeLibraryLocation, + name: String, + readerState: Array[Byte], + partition: Array[Byte], + schema: StructType) + extends PartitionReader[ColumnarBatch] { + import NativeCalls._ + + private val allocator = + ArrowUtils.rootAllocator.newChildAllocator(s"native data source $name", 0, Long.MaxValue) + private var reader: ArrowReader = _ + private var root: VectorSchemaRoot = _ + private var batch: ColumnarBatch = _ + + try { + reader = openStream() + root = call(name, "read")(reader.getVectorSchemaRoot) + checkSchema() + batch = new ColumnarBatch( + root.getFieldVectors.asScala.map(new ArrowColumnVector(_): ColumnVector).toArray) + } catch { + case e: Throwable => + close() + throw e + } + + private def openStream(): ArrowReader = { + val library = NativeLibraries.get(location) + val stream = ArrowArrayStream.allocateNew(allocator) + try { + call(name, "read")(optional(library.read(readerState, partition, stream.memoryAddress()))) + .getOrElse { + throw new SparkUnsupportedOperationException( + "DATA_SOURCE_BATCH_SCAN_NOT_SUPPORTED", Map("description" -> name)) + } + Data.importArrayStream(allocator, stream) Review Comment: Arrow Java ignores `ArrowArray.offset` when it imports C data: `ArrayImporter` uses `length` but never `offset`, so every buffer is read from position 0. If a library returns a zero-copy slice, such as `RecordBatch::slice` in arrow-rs, which exports a non-zero offset for each column, Spark reads the wrong rows without any error. On the PR head, I exported `[10, 11, 12, 13]`, set `offset=1, length=3`, and got `[10, 11, 12]` back. `NativeBridge` doesn't rule out non-zero offsets either. Could we require a zero offset in the Javadoc and check it before importing? The check could follow `ArrowArrayStreamReader.loadNextBatch`: call `getNext` into an `ArrowArray`, check `snapshot().offset` at every level, fail with `NATIVE_DATA_SOURCE_ERROR` if it isn't 0, and then call `Data.importIntoVectorSchemaRoot`. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala: ########## @@ -0,0 +1,211 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.execution.datasources.v2.columnar + +import scala.jdk.CollectionConverters._ + +import org.apache.arrow.vector.VectorSchemaRoot + +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.connector.metric.CustomTaskMetric +import org.apache.spark.sql.connector.write._ +import org.apache.spark.sql.connector.write.streaming.{StreamingDataWriterFactory, StreamingWrite} +import org.apache.spark.sql.execution.arrow.ArrowWriter +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.util.ArrowUtils +import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector} +import org.apache.spark.util.Utils + +class ColumnarWriteBuilder( + shortName: String, + info: LogicalWriteInfo, + createDataSource: () => ColumnarDataSource) + extends WriteBuilder with SupportsTruncate { + + private var overwrite = false + + override def truncate(): WriteBuilder = { + overwrite = true + this + } + + override def build(): Write = new ColumnarWrite(shortName, info, overwrite, createDataSource) +} + +class ColumnarWrite( + shortName: String, + info: LogicalWriteInfo, + overwrite: Boolean, + createDataSource: () => ColumnarDataSource) + extends Write { + + override def description(): String = shortName + + override def toBatch: BatchWrite = { + new ColumnarBatchWrite(createDataSource().writer(info.schema(), overwrite), info.schema()) Review Comment: The `NativeBridge` lifecycle says Spark calls `abort` when a write fails, but with AQE on (the default) an error in an upstream stage never reaches it. `toBatch` creates the library's writer right away, while the upstream stages run inside `query.execute()`, which in `writeWithV2` is outside the `try` whose `catch` calls `abort`. I tested with a JVM `ColumnarDataSource`: a write after `groupBy(...).agg(...)` that divides by zero in the shuffle stage gives `createWriter=1 commit=0 abort=0`. Without AQE, or without the shuffle, it gives `abort=1`. One option is to create and serialize the writer in `createBatchWriterFactory`, which runs after `query.execute()`. Alternatively, the Javadoc could say that a writer may be closed without a commit or an abort. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala: ########## @@ -0,0 +1,231 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.execution.datasources.v2.columnar + +import java.util +import java.util.Locale + +import org.apache.spark.SparkException +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.connector.catalog.{SupportsRead, SupportsWrite, Table, TableCapability} +import org.apache.spark.sql.connector.catalog.TableCapability.{BATCH_READ, BATCH_WRITE, MICRO_BATCH_READ, STREAMING_WRITE, TRUNCATE} +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.connector.read._ +import org.apache.spark.sql.connector.read.streaming.{MicroBatchStream, Offset} +import org.apache.spark.sql.connector.write.{LogicalWriteInfo, WriteBuilder} +import org.apache.spark.sql.internal.connector.SupportsMetadata +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.util.CaseInsensitiveStringMap +import org.apache.spark.sql.vectorized.ColumnarBatch + +/** + * The Data Source V2 [[Table]] of a [[ColumnarDataSource]]. It creates a new data source for each + * scan and each write, with their options on top of the options of the table. + * + * @param tableOptions the options of the table, given to `TableProvider.getTable`. They are the + * options of `CREATE TABLE ... USING ... OPTIONS (...)` for a table in a + * catalog, while the scans and the writes of the table only get the options + * of their query. + */ +class ColumnarTable( + shortName: String, + tableSchema: StructType, + tableOptions: util.Map[String, String], + createDataSource: CaseInsensitiveStringMap => ColumnarDataSource) + extends Table with SupportsRead with SupportsWrite { + + override def name(): String = shortName + + override def schema(): StructType = tableSchema + + override def capabilities(): util.Set[TableCapability] = + util.EnumSet.of(BATCH_READ, BATCH_WRITE, MICRO_BATCH_READ, STREAMING_WRITE, TRUNCATE) + + override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { + new ColumnarScanBuilder( + shortName, tableSchema, () => createDataSource(withTableOptions(options))) + } + + override def newWriteBuilder(info: LogicalWriteInfo): WriteBuilder = { + new ColumnarWriteBuilder( + shortName, info, () => createDataSource(withTableOptions(info.options()))) + } + + private def withTableOptions(options: CaseInsensitiveStringMap): CaseInsensitiveStringMap = { + val merged = new util.HashMap[String, String]() + tableOptions.forEach((key, value) => merged.put(key.toLowerCase(Locale.ROOT), value)) + // The keys of a case-insensitive map are in lower case. + options.forEach((key, value) => merged.put(key, value)) + new CaseInsensitiveStringMap(merged) + } +} + +class ColumnarScanBuilder( + shortName: String, + schema: StructType, + createDataSource: () => ColumnarDataSource) + extends ScanBuilder + with SupportsPushDownV2Filters + with SupportsPushDownLimit + with SupportsPushDownRequiredColumns { + + private lazy val dataSource = createDataSource() + + // Created on first use, because a streaming scan neither pushes anything down nor reads in + // batch, and its data source may not support batch scans at all. + private var reader: ColumnarReader = _ + + private var pushed: Array[Predicate] = Array.empty + private var pushedLimit: Option[Int] = None + private var readSchema: StructType = schema + + private def batchReader: ColumnarReader = { + if (reader == null) { + reader = dataSource.reader(schema) + } + reader + } + + override def pushPredicates(predicates: Array[Predicate]): Array[Predicate] = { + val remaining = batchReader.pushPredicates(predicates) + pushed = predicates.filterNot(remaining.contains) + remaining + } + + override def pushedPredicates(): Array[Predicate] = pushed + + override def pushLimit(limit: Int): Boolean = { + val isPushed = batchReader.pushLimit(limit) Review Comment: `NativeBridge` says `pushLimit` is called at most once, but `df.limit(10).offset(5)` calls it twice on the same reader. `pushDownLimitAndOffset` transforms top-down: the `LimitAndOffset` case pushes the limit once and keeps the `Offset`, because `isPartiallyPushed()` is always true, and the transform then reaches the `GlobalLimit` below and pushes it again. I counted the calls with a JVM reader on the PR head: `limit.offset` makes 2 calls per reader, `limit` and `offset.limit` make 1. A library that follows the contract and rejects the second call fails this query. Could `ColumnarScanBuilder.pushLimit` remember the first limit and its result, and return that result when it's called again with the same limit? A `limit(...).offset(...)` test would cover it. -- 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]
