zhztheplayer commented on code in PR #59424: URL: https://github.com/apache/spark/pull/59424#discussion_r4293697066
########## 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: I felt this may cause bug since reusing a fixed `VectorSchemaRoot` can cause data corruption when the output of `NativeDataSource` is accumulated by the downstream operators ([example](https://github.com/apache/spark/blob/ca6524e34ce487bb61ab1fda4b440a86e2c2177b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala#L191-L192)), since the new data will replace the content of the old vectors in early `ColumnarBatch` instances. It might be worth trying the root-free `Data.importIntoVectorSchemaRoot` APIs that creates fresh roots internally? -- 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]
