gengliangwang commented on code in PR #59961: URL: https://github.com/apache/spark/pull/59961#discussion_r4429278858
########## sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeColumnVector.java: ########## @@ -0,0 +1,314 @@ +/* + * 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()}, and the vector then reads booleans, + * integers, floats, dates, timestamps, decimals and strings 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 until the reader imports the next batch or closes. 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. + * + * <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, binary values, structs, lists and maps, + * and the fields of structs are read through {@link ArrowColumnVector}. + */ +class NativeColumnVector extends ReusedArrowColumnVector { + + /** The column vector of a top-level Arrow vector of a reader. */ + static NativeColumnVector create(ValueVector vector) { + if (vector instanceof BitVector) { + return new Booleans(vector); + } else if (vector instanceof TinyIntVector) { + return new Bytes(vector); + } else if (vector instanceof SmallIntVector) { + return new Shorts(vector); + } else if (vector instanceof IntVector || vector instanceof DateDayVector) { + return new Ints(vector); + } else if (vector instanceof BigIntVector || vector instanceof TimeStampMicroTZVector || + vector instanceof TimeStampMicroVector) { + return new Longs(vector); + } else if (vector instanceof Float4Vector) { + return new Floats(vector); + } else if (vector instanceof Float8Vector) { + return new Doubles(vector); + } else if (vector instanceof DecimalVector decimalVector) { + return new Decimals(decimalVector); + } else if (vector instanceof VarCharVector) { + return new Strings(vector); + } else if (vector instanceof ViewVarCharVector viewVector) { + return new StringViews(viewVector); + } else { + return new NativeColumnVector(vector); + } + } + + NativeColumnVector(ValueVector vector) { + super(vector); + } + + /** Reads the addresses of the buffers of the batch that the reader imported last. */ + void refresh() { + } + + /** A vector that reads its values and their validity at the addresses of its buffers. */ + private abstract static class Direct extends NativeColumnVector { + final ValueVector vector; + // 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) { + super(vector); + this.vector = vector; + } + + @Override + void refresh() { + validity = vector.getNullCount() == 0 ? 0 : vector.getValidityBuffer().memoryAddress(); + data = vector.getDataBuffer().memoryAddress(); + } + + @Override + public final boolean isNullAt(int rowId) { + return validity != 0 && !bit(validity, rowId); + } + + /** The bit of a bitmap, whose bits Arrow orders from the least significant of each byte. */ + static boolean bit(long address, int index) { + return (Platform.getByte(null, address + (index >> 3)) & (1 << (index & 7))) != 0; + } + } + + private static final class Booleans extends Direct { + Booleans(ValueVector vector) { + super(vector); + } + + @Override + public boolean getBoolean(int rowId) { + return bit(data, rowId); + } + } + + private static final class Bytes extends Direct { + Bytes(ValueVector vector) { + super(vector); + } + + @Override + public byte getByte(int rowId) { + return Platform.getByte(null, data + rowId); + } + } + + private static final class Shorts extends Direct { + Shorts(ValueVector vector) { + super(vector); + } + + @Override + public short getShort(int rowId) { + return Platform.getShort(null, data + 2L * rowId); + } + } + + private static final class Ints extends Direct { + Ints(ValueVector vector) { + super(vector); + } + + @Override + public int getInt(int rowId) { + return Platform.getInt(null, data + 4L * rowId); + } + } + + private static final class Longs extends Direct { + Longs(ValueVector vector) { + super(vector); + } + + @Override + public long getLong(int rowId) { + return Platform.getLong(null, data + 8L * rowId); Review Comment: **[P0] Keep buffers alive while another thread consumes the batch** PipedRDD consumes its parent on a separate stdin writer thread and does not join that thread before task completion. DataSourceRDD's completion listener closes the reader, and NativePartitionReader immediately calls root.close() because stopReading() clears reading when next() returns, before batch consumption finishes. ColumnarToRowExec can subsequently request another row from that batch, so this getter dereferences the cached address after its allocation has been freed. Previously, a getter starting after close accessed Arrow's cleared buffers and failed through checked access. Retain the buffers until the consuming thread finishes. Already reported in https://github.com/apache/spark/pull/59961#discussion_r4429239839. **Failure scenario:** A native source returns a large BIGINT batch. Running df.rdd.pipe("head -1").collect() lets the task finish while the pipe writer still converts remaining batch rows. Reader cleanup frees the native allocation; subsequent getLong calls read freed memory and can crash the executor JVM. -- 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]
