gengliangwang commented on code in PR #59961: URL: https://github.com/apache/spark/pull/59961#discussion_r4488360807
########## sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeColumnVector.java: ########## @@ -0,0 +1,331 @@ +/* + * 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.nio.ByteOrder; +import java.util.List; + +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.vector.BaseVariableWidthViewVector; +import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.BitVector; +import org.apache.arrow.vector.DateDayVector; +import org.apache.arrow.vector.DecimalVector; +import org.apache.arrow.vector.Float4Vector; +import org.apache.arrow.vector.Float8Vector; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.SmallIntVector; +import org.apache.arrow.vector.TimeStampMicroTZVector; +import org.apache.arrow.vector.TimeStampMicroVector; +import org.apache.arrow.vector.TinyIntVector; +import org.apache.arrow.vector.ValueVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.ViewVarCharVector; + +import org.apache.spark.sql.execution.vectorized.ReusedArrowColumnVector; +import org.apache.spark.sql.types.Decimal; +import org.apache.spark.sql.vectorized.ArrowColumnVector; +import org.apache.spark.unsafe.Platform; +import org.apache.spark.unsafe.types.UTF8String; + +/** + * The column vector of a top-level column of a {@code NativePartitionReader}, which imports every + * batch into the same Arrow vectors. + * + * <p>After each import, the reader calls {@link #refresh(long)}, and the vector then reads + * booleans, integers, floats, dates, timestamps, decimals, and the strings of the Arrow types Utf8 + * and Utf8View at the addresses of the buffers of the batch. {@link ArrowColumnVector} reads each + * value through Arrow Java, which by default checks on each access the bounds and the reference + * count of the buffer, and whether the value is null. + * Spark does not need these checks here: it reads only the rows of the batch, whose buffers Arrow + * Java sizes by the lengths of the arrays, which the reader checks before it imports them, and + * the buffers of a batch stay allocated while the batch may be read: until next is called again, + * and if the reader is closed before then, until the reader is unreachable, which the vector keeps + * it from being. Spark reads a value of a nullable column only after {@link #isNullAt(int)}, and + * the reader checks that the columns that the schema declares not nullable have no null, so that + * the vector of such a column does not count the nulls of each batch. The vector of a nullable + * column takes the null_count that the library exported, and only counts the nulls when the + * library did not. + * + * <p>The vector returns the decimals that fit their precision in a long, of a scale from 0 to + * their precision, as the unscaled longs that Spark's own column vectors hold, which Spark adds as + * longs, rather than as BigDecimals. The other decimals, large strings, binary values, structs, + * lists and maps, and the fields of structs are read through {@link ArrowColumnVector}. + */ +class NativeColumnVector extends ReusedArrowColumnVector { Review Comment: The direct reads are only safe because of what this reader guarantees. It checks the length of every array before the import, and it rejects nulls in columns declared not nullable. It takes null_count from the C data interface, whose value is exact when it is not -1. And it keeps the buffers alive until no consumer can read them. ArrowOutputProcessorImpl loads IPC batches and gives none of these guarantees, so a refresh hook on ReusedArrowColumnVector would need each producer to provide them first. Otherwise its raw reads could go past the buffers. On the duplication: ArrowColumnVector's accessors decode through ArrowBuf, which is what gives them their bounds and refcount checks. Sharing one decoder would mean either moving those accessors to raw addresses, which drops the checks for every Arrow consumer, or adding a second path inside them. I think that is worth doing together with a second producer that provides the same guarantees, rather than here. ########## sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeColumnVector.java: ########## @@ -0,0 +1,331 @@ +/* + * 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.nio.ByteOrder; +import java.util.List; + +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.vector.BaseVariableWidthViewVector; +import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.BitVector; +import org.apache.arrow.vector.DateDayVector; +import org.apache.arrow.vector.DecimalVector; +import org.apache.arrow.vector.Float4Vector; +import org.apache.arrow.vector.Float8Vector; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.SmallIntVector; +import org.apache.arrow.vector.TimeStampMicroTZVector; +import org.apache.arrow.vector.TimeStampMicroVector; +import org.apache.arrow.vector.TinyIntVector; +import org.apache.arrow.vector.ValueVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.ViewVarCharVector; + +import org.apache.spark.sql.execution.vectorized.ReusedArrowColumnVector; +import org.apache.spark.sql.types.Decimal; +import org.apache.spark.sql.vectorized.ArrowColumnVector; +import org.apache.spark.unsafe.Platform; +import org.apache.spark.unsafe.types.UTF8String; + +/** + * The column vector of a top-level column of a {@code NativePartitionReader}, which imports every + * batch into the same Arrow vectors. + * + * <p>After each import, the reader calls {@link #refresh(long)}, and the vector then reads + * booleans, integers, floats, dates, timestamps, decimals, and the strings of the Arrow types Utf8 + * and Utf8View at the addresses of the buffers of the batch. {@link ArrowColumnVector} reads each + * value through Arrow Java, which by default checks on each access the bounds and the reference + * count of the buffer, and whether the value is null. + * Spark does not need these checks here: it reads only the rows of the batch, whose buffers Arrow + * Java sizes by the lengths of the arrays, which the reader checks before it imports them, and + * the buffers of a batch stay allocated while the batch may be read: until next is called again, + * and if the reader is closed before then, until the reader is unreachable, which the vector keeps + * it from being. Spark reads a value of a nullable column only after {@link #isNullAt(int)}, and + * the reader checks that the columns that the schema declares not nullable have no null, so that + * the vector of such a column does not count the nulls of each batch. The vector of a nullable + * column takes the null_count that the library exported, and only counts the nulls when the + * library did not. + * + * <p>The vector returns the decimals that fit their precision in a long, of a scale from 0 to + * their precision, as the unscaled longs that Spark's own column vectors hold, which Spark adds as + * longs, rather than as BigDecimals. The other decimals, large strings, binary values, structs, + * lists and maps, and the fields of structs are read through {@link ArrowColumnVector}. + */ +class NativeColumnVector extends ReusedArrowColumnVector { + + // The reader, which the vector keeps reachable: a reader closed while its batch may still be + // read frees the buffers of the batch only once it is unreachable. + private Object reader; + + /** + * The column vector of a top-level Arrow vector of {@code reader}, of a column that the schema to + * read declares {@code nullable} or not. + */ + static NativeColumnVector create(ValueVector vector, boolean nullable, Object reader) { + NativeColumnVector column; + if (vector instanceof BitVector) { + column = new Booleans(vector, nullable); + } else if (vector instanceof TinyIntVector) { + column = new Bytes(vector, nullable); + } else if (vector instanceof SmallIntVector) { + column = new Shorts(vector, nullable); + } else if (vector instanceof IntVector || vector instanceof DateDayVector) { + column = new Ints(vector, nullable); + } else if (vector instanceof BigIntVector || vector instanceof TimeStampMicroTZVector || + vector instanceof TimeStampMicroVector) { + column = new Longs(vector, nullable); + } else if (vector instanceof Float4Vector) { + column = new Floats(vector, nullable); + } else if (vector instanceof Float8Vector) { + column = new Doubles(vector, nullable); + } else if (vector instanceof DecimalVector decimalVector) { + column = new Decimals(decimalVector, nullable); + } else if (vector instanceof VarCharVector) { + column = new Strings(vector, nullable); + } else if (vector instanceof ViewVarCharVector) { + column = new StringViews(vector, nullable); + } else { + column = new NativeColumnVector(vector); + } + column.reader = reader; + return column; + } + + NativeColumnVector(ValueVector vector) { + super(vector); + } + + /** + * Reads the addresses of the buffers of the batch that the reader imported last, whose array of + * this column the library exported with a null_count of {@code nullCount}. + */ + void refresh(long nullCount) { + } + + /** A vector that reads its values and their validity at the addresses of its buffers. */ + private abstract static class Direct extends NativeColumnVector { + // Whether the schema to read declares the column nullable. The reader checks that the batches + // of a column that it does not declare nullable have no null. + private final boolean nullable; + // The address of the validity bitmap of the batch, or 0 if the batch has no null. + private long validity; + // The address of the data buffer of the batch. + long data; + + Direct(ValueVector vector, boolean nullable) { + super(vector); + this.nullable = nullable; + } + + @Override + void refresh(long nullCount) { + ValueVector vector = getValueVector(); + // The C data interface requires a null_count other than -1 to be exact, so the nulls are + // only counted, which Arrow Java does in the validity bitmap on each call, for a negative + // one. + boolean hasNulls = nullable && (nullCount < 0 ? vector.getNullCount() > 0 : nullCount > 0); + validity = hasNulls ? vector.getValidityBuffer().memoryAddress() : 0; + data = vector.getDataBuffer().memoryAddress(); Review Comment: Done in 4b5dc33eb0d1: The NativeColumnVector class doc now says that the vector caches the buffer addresses between imports, so nothing may reallocate those buffers in between, and that a consumer writing them through getValueVector() (e.g. setValueCount) must keep the length of the batch. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala: ########## @@ -1053,14 +1109,25 @@ class NativePartitionReader( if (task != null) { task.close() } - Utils.tryLogNonFatalError(allocator.close()) + if (freeBatch) { + Utils.tryLogNonFatalError(allocator.close()) + } else { + // The action must not refer to the reader, which would then never become unreachable. + val (batchRoot, batchAllocator) = (root, allocator) + NativePartitionReader.cleaner.register(this, () => + Utils.tryLogNonFatalError( + Utils.tryWithSafeFinally(batchRoot.close())(batchAllocator.close()))) + } } } } object NativePartitionReader { import NativeCalls._ + // Frees the buffers of the batches that may still be read when their readers are closed. + private val cleaner = Cleaner.create() Review Comment: Done in 4b5dc33eb0d1: Removed the second Cleaner: the reader now registers with NativeHandle.cleaner, which is private[ffi]. -- 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]
