dongjoon-hyun commented on code in PR #59961:
URL: https://github.com/apache/spark/pull/59961#discussion_r4500198812
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -1014,32 +1089,44 @@ class NativePartitionReader(
"but Spark only reads batches without null rows: mark the values of
the columns as null " +
"instead.")
}
- val columns = root.getSchema.getFields.size
- children(batch, columns, "The batch", s"the schema has $columns columns")
- .zip(root.getSchema.getFields.asScala).foreach { case (column, field)
=>
+ val numColumns = root.getSchema.getFields.size
+ val arrays = children(batch, numColumns, "The batch", s"the schema has
$numColumns columns")
+ arrays.zip(root.getSchema.getFields.asScala).foreach { case (column,
field) =>
if (column.length < batch.length) {
throw error(name, "read", s"The column ${field.getName} has a length
of " +
s"${column.length}, but the batch has a length of ${batch.length}.")
}
checkColumn(column, field, field.getName)
}
+ arrays.map(_.null_count).toArray
}
override def get(): ColumnarBatch = batch
override def close(): Unit = synchronized {
if (!closed) {
closed = true
+ // A thread in next releases the reader once get_next returns, see
stopReading. Otherwise, a
+ // consumer may still read the batch that next returned, unless next was
called since.
if (!reading) {
- closeResources()
+ closeResources(freeBatch = !batchReturned)
Review Comment:
Under `DataSourceRDD`, this deferred path is the common case rather than the
exception. Whenever the task ends before the stream does (`LIMIT`,
`show()`/`take()`/`head()`, failed, killed or speculative tasks), the
completion listener calls `close()` while `batchReturned` is true, so
`closeResources(freeBatch = false)` only registers a Cleaner.
`PartitionIterator.close()` then sets `reader = null`
(`DataSourceRDD.scala:265`), so `next()` is never called again and the
free-at-once path in `startReading()` cannot run.
These off-heap buffers do not trigger a GC. On a large-heap executor where
long tasks let readers reach the old generation, each early-ended task can hold
up to one batch per partition (tens of MB for wide string batches) until an
old-gen collection, and native memory can grow until YARN/K8s kills the
container for exceeding `memoryOverhead`.
Since task-completion listeners run in reverse registration order, the
script-transformation listener already runs before `DataSourceRDD`'s close.
Could that listener wait for the writer thread instead, so that `close()` here
can always free the batch at once?
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeColumnVector.java:
##########
@@ -0,0 +1,345 @@
+/*
+ * 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. The strings, arrays, maps and child vectors that the vector
returns do not refer
+ * to the reader, so this holds for them only while the vector is reachable
too. 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 keeps the addresses of the buffers from one import to the
next, whereas
+ * {@link ArrowColumnVector} reads the buffers of its Arrow vector on each
access. So nothing may
+ * reallocate the buffers of the Arrow vectors of the reader between two
imports, e.g. by growing
+ * them: a consumer that writes them through {@link #getValueVector()}, as by
setting their value
+ * count, must keep it at the length of the batch.
+ *
+ * <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. Nothing
reads the field, but
+ // it must stay: each constructor takes the reader so that no column vector
misses it.
+ private final 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) {
+ if (vector instanceof BitVector) {
+ return new Booleans(vector, nullable, reader);
+ } else if (vector instanceof TinyIntVector) {
+ return new Bytes(vector, nullable, reader);
+ } else if (vector instanceof SmallIntVector) {
+ return new Shorts(vector, nullable, reader);
+ } else if (vector instanceof IntVector || vector instanceof DateDayVector)
{
+ return new Ints(vector, nullable, reader);
+ } else if (vector instanceof BigIntVector || vector instanceof
TimeStampMicroTZVector ||
+ vector instanceof TimeStampMicroVector) {
+ return new Longs(vector, nullable, reader);
+ } else if (vector instanceof Float4Vector) {
+ return new Floats(vector, nullable, reader);
+ } else if (vector instanceof Float8Vector) {
+ return new Doubles(vector, nullable, reader);
+ } else if (vector instanceof DecimalVector decimalVector) {
+ return new Decimals(decimalVector, nullable, reader);
+ } else if (vector instanceof VarCharVector) {
+ return new Strings(vector, nullable, reader);
+ } else if (vector instanceof ViewVarCharVector) {
+ return new StringViews(vector, nullable, reader);
+ } else {
+ return new NativeColumnVector(vector, reader);
+ }
+ }
+
+ NativeColumnVector(ValueVector vector, Object reader) {
+ super(vector);
+ this.reader = reader;
+ }
+
+ /**
+ * 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) {
+ }
+
+ /**
+ * Whether {@code vector} has a null, given the null_count {@code nullCount}
that the library
+ * exported for its array, or -1 if it is not known. 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.
+ */
+ static boolean hasNulls(ValueVector vector, long nullCount) {
+ return nullCount < 0 ? vector.getNullCount() > 0 : nullCount > 0;
+ }
+
+ /** 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, Object reader) {
+ super(vector, reader);
+ this.nullable = nullable;
+ }
+
+ @Override
+ void refresh(long nullCount) {
+ ValueVector vector = getValueVector();
+ boolean hasNulls = nullable && hasNulls(vector, nullCount);
+ validity = hasNulls ? vector.getValidityBuffer().memoryAddress() : 0;
+ data = vector.getDataBuffer().memoryAddress();
Review Comment:
The `Direct` getters read at the addresses cached in `refresh()`, without
the bounds and reference-count checks that `ArrowColumnVector`'s accessors get
from Arrow. So a misbehaving consumer now reads freed or out-of-range memory
where it used to get an exception:
- A consumer (e.g. a columnar plugin) that calls `ColumnarBatch.close()` or
`ColumnVector.close()` instead of `closeIfFreeable()`:
`ArrowColumnVector.close()` closes the shared Arrow vector, but `Longs.getLong`
and `Strings.getUTF8String` keep reading the old addresses, which gives garbage
values or a SIGSEGV instead of
`IndexOutOfBoundsException`/`IllegalStateException`.
- A retained `ColumnarBatchRow` read after `next()` with a `rowId` at or
beyond the new batch's length reads past the new buffers.
Also, the `UTF8String.fromAddress` results and the child
`ArrowColumnVector`s / `ColumnarArray`s do not keep the reader reachable, so
after a deferred close only the batch itself keeps the buffers alive.
A JVM crash is much harder to diagnose than an exception. Could we document
this contract explicitly on the class, and keep a guard where it is cheap?
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeColumnVector.java:
##########
@@ -0,0 +1,345 @@
+/*
+ * 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. The strings, arrays, maps and child vectors that the vector
returns do not refer
+ * to the reader, so this holds for them only while the vector is reachable
too. 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 keeps the addresses of the buffers from one import to the
next, whereas
+ * {@link ArrowColumnVector} reads the buffers of its Arrow vector on each
access. So nothing may
+ * reallocate the buffers of the Arrow vectors of the reader between two
imports, e.g. by growing
+ * them: a consumer that writes them through {@link #getValueVector()}, as by
setting their value
+ * count, must keep it at the length of the batch.
+ *
+ * <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. Nothing
reads the field, but
+ // it must stay: each constructor takes the reader so that no column vector
misses it.
+ private final 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) {
+ if (vector instanceof BitVector) {
+ return new Booleans(vector, nullable, reader);
+ } else if (vector instanceof TinyIntVector) {
+ return new Bytes(vector, nullable, reader);
+ } else if (vector instanceof SmallIntVector) {
+ return new Shorts(vector, nullable, reader);
+ } else if (vector instanceof IntVector || vector instanceof DateDayVector)
{
+ return new Ints(vector, nullable, reader);
+ } else if (vector instanceof BigIntVector || vector instanceof
TimeStampMicroTZVector ||
+ vector instanceof TimeStampMicroVector) {
+ return new Longs(vector, nullable, reader);
+ } else if (vector instanceof Float4Vector) {
+ return new Floats(vector, nullable, reader);
+ } else if (vector instanceof Float8Vector) {
+ return new Doubles(vector, nullable, reader);
+ } else if (vector instanceof DecimalVector decimalVector) {
+ return new Decimals(decimalVector, nullable, reader);
+ } else if (vector instanceof VarCharVector) {
+ return new Strings(vector, nullable, reader);
+ } else if (vector instanceof ViewVarCharVector) {
+ return new StringViews(vector, nullable, reader);
+ } else {
+ return new NativeColumnVector(vector, reader);
+ }
+ }
+
+ NativeColumnVector(ValueVector vector, Object reader) {
+ super(vector);
+ this.reader = reader;
+ }
+
+ /**
+ * 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) {
+ }
+
+ /**
+ * Whether {@code vector} has a null, given the null_count {@code nullCount}
that the library
+ * exported for its array, or -1 if it is not known. 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.
+ */
+ static boolean hasNulls(ValueVector vector, long nullCount) {
+ return nullCount < 0 ? vector.getNullCount() > 0 : nullCount > 0;
+ }
+
+ /** 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, Object reader) {
+ super(vector, reader);
+ this.nullable = nullable;
+ }
+
+ @Override
+ void refresh(long nullCount) {
+ ValueVector vector = getValueVector();
+ boolean hasNulls = nullable && hasNulls(vector, nullCount);
+ validity = hasNulls ? vector.getValidityBuffer().memoryAddress() : 0;
+ 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, boolean nullable, Object reader) {
+ super(vector, nullable, reader);
+ }
+
+ @Override
+ public boolean getBoolean(int rowId) {
+ return bit(data, rowId);
+ }
+ }
+
+ private static final class Bytes extends Direct {
+ Bytes(ValueVector vector, boolean nullable, Object reader) {
+ super(vector, nullable, reader);
+ }
+
+ @Override
+ public byte getByte(int rowId) {
+ return Platform.getByte(null, data + rowId);
+ }
+ }
+
+ private static final class Shorts extends Direct {
+ Shorts(ValueVector vector, boolean nullable, Object reader) {
+ super(vector, nullable, reader);
+ }
+
+ @Override
+ public short getShort(int rowId) {
+ return Platform.getShort(null, data + 2L * rowId);
+ }
+ }
+
+ private static final class Ints extends Direct {
+ Ints(ValueVector vector, boolean nullable, Object reader) {
+ super(vector, nullable, reader);
+ }
+
+ @Override
+ public int getInt(int rowId) {
+ return Platform.getInt(null, data + 4L * rowId);
+ }
+ }
+
+ private static final class Longs extends Direct {
+ Longs(ValueVector vector, boolean nullable, Object reader) {
+ super(vector, nullable, reader);
+ }
+
+ @Override
+ public long getLong(int rowId) {
+ return Platform.getLong(null, data + 8L * rowId);
+ }
+ }
+
+ private static final class Floats extends Direct {
+ Floats(ValueVector vector, boolean nullable, Object reader) {
+ super(vector, nullable, reader);
+ }
+
+ @Override
+ public float getFloat(int rowId) {
+ return Platform.getFloat(null, data + 4L * rowId);
+ }
+ }
+
+ private static final class Doubles extends Direct {
+ Doubles(ValueVector vector, boolean nullable, Object reader) {
+ super(vector, nullable, reader);
+ }
+
+ @Override
+ public double getDouble(int rowId) {
+ return Platform.getDouble(null, data + 8L * rowId);
+ }
+ }
+
+ private static final class Decimals extends Direct {
+ // Arrow stores the 128 bits of a decimal in native byte order, as
ArrowColumnVector reads it.
+ private static final int LOW_WORD = ByteOrder.nativeOrder() ==
ByteOrder.LITTLE_ENDIAN ? 0 : 8;
+ private static final int HIGH_WORD = 8 - LOW_WORD;
+
+ // The scale of the Arrow type.
+ private final int scale;
+
+ Decimals(DecimalVector vector, boolean nullable, Object reader) {
+ super(vector, nullable, reader);
+ this.scale = vector.getScale();
+ }
+
+ @Override
+ public Decimal getDecimal(int rowId, int precision, int scale) {
+ if (isNullAt(rowId)) return null;
+ // The integral conversions of a decimal that a long holds, such as
Decimal.toLong, only
+ // support the scales from 0 to 18, which those of Spark's own column
vectors have.
+ if (precision <= Decimal.MAX_LONG_DIGITS() && scale == this.scale &&
scale >= 0 &&
+ scale <= precision) {
+ long address = data + (long) DecimalVector.TYPE_WIDTH * rowId;
+ long unscaled = Platform.getLong(null, address + LOW_WORD);
+ if (Platform.getLong(null, address + HIGH_WORD) == unscaled >> 63) {
+ // Null if the value does not fit its precision.
+ Decimal decimal = new Decimal().setOrNull(unscaled, precision,
scale);
Review Comment:
Returning a long-backed `Decimal` only for top-level native columns makes
the same value convert differently depending on where it comes from. Struct
fields (through the inherited child `ArrowColumnVector`s), Python UDF output
and the other Arrow producers still get BigDecimal-backed decimals from
`ArrowColumnVector.SmallDecimalAccessor`, whose integral conversions go through
`double`:
- `DECIMAL(18,15)` value `127.999999999999999` under ANSI: `CAST(d AS
TINYINT)` on a top-level column returns 127, but `CAST(st.d AS TINYINT)` throws
`CAST_OVERFLOW`, because the double rounds to 128.0.
- `DECIMAL(18,9)` value `123456789.999999999`: `CAST AS INT` gives 123456789
top-level and 123456790 nested.
Since this is a correctness difference rather than only a performance one,
could we fix it at the root instead, in
`ArrowColumnVector.SmallDecimalAccessor` (or in the double path of
`Decimal.roundToNumeric`), so that all Arrow-backed columns behave the same?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeColumnVectorSuite.scala:
##########
@@ -0,0 +1,416 @@
+/*
+ * 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.math.{BigDecimal => JavaBigDecimal}
+import java.nio.charset.StandardCharsets
+
+import scala.util.{Failure, Success, Try, Using}
+
+import org.apache.arrow.memory.{BufferAllocator, RootAllocator}
+import org.apache.arrow.vector._
+import org.apache.arrow.vector.complex.StructVector
+import org.apache.arrow.vector.types.pojo.{ArrowType, FieldType}
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.sql.catalyst.expressions.UnsafeRow
+import org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter
+import org.apache.spark.sql.execution.vectorized.{OnHeapColumnVector,
ReusedArrowColumnVector}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{Decimal, DecimalType}
+import org.apache.spark.sql.vectorized.ArrowColumnVector
+
+/**
+ * Checks that a [[NativeColumnVector]] reads the nulls and values that an
[[ArrowColumnVector]]
+ * of the same Arrow vector reads.
+ */
+class NativeColumnVectorSuite extends SparkFunSuite {
+
+ // The rows of the batches, more than 8 so that the validity bitmaps have
several bytes, and the
+ // rows that are null.
+ private val rows = 0 until 20
+ private val nulls = Set(1, 8, 9, 15)
+
+ private def withAllocator(f: BufferAllocator => Unit): Unit = {
+ Using.resource(new RootAllocator())(f)
+ }
+
+ /** The column vector of `vector`, refreshed as the reader refreshes it
after an import. */
+ private def native(vector: ValueVector, nullable: Boolean = true):
NativeColumnVector = {
+ val column = NativeColumnVector.create(vector, nullable, null)
+ column.refresh(-1)
+ column
+ }
+
+ /** Checks that `column` reads the nulls and the values of `vector` as
ArrowColumnVector does. */
+ private def checkSame(vector: ValueVector, column: ArrowColumnVector)(
+ value: (ArrowColumnVector, Int) => Any): Unit = {
+ val expected = new ArrowColumnVector(vector)
+ (0 until vector.getValueCount).foreach { row =>
+ assert(column.isNullAt(row) === expected.isNullAt(row), s"row $row")
+ if (!expected.isNullAt(row)) {
+ assert(value(column, row) === value(expected, row), s"row $row")
+ }
+ }
+ }
+
+ /** Sets the rows of a vector: nulls, and `set` for the other rows. */
+ private def fill[V <: BaseFixedWidthVector](vector: V)(set: (V, Int) =>
Unit): V = {
+ vector.allocateNew(rows.size)
+ rows.foreach { row => if (nulls(row)) vector.setNull(row) else set(vector,
row) }
+ vector.setValueCount(rows.size)
+ vector
+ }
+
+ /** Checks a vector whose values the column vector reads at the addresses of
its buffers. */
+ private def checkDirect(vector: ValueVector)(value: (ArrowColumnVector, Int)
=> Any): Unit = {
+ Using.resource(vector) { vector =>
+ val column = native(vector)
+ assert(column.getClass !== classOf[NativeColumnVector],
vector.getClass.getName)
+ checkSame(vector, column)(value)
+ }
+ }
+
+ test("booleans, integers, floats, dates and timestamps") {
+ withAllocator { allocator =>
+ val bytes = Seq(Byte.MinValue, -1: Byte, 0: Byte, 1: Byte, Byte.MaxValue)
+ val shorts = Seq(Short.MinValue, -1: Short, 0: Short, 1: Short,
Short.MaxValue)
+ val ints = Seq(Int.MinValue, -1, 0, 1, Int.MaxValue)
+ val longs = Seq(Long.MinValue, -1L, 0L, 1L, Long.MaxValue)
+ val floats = Seq(0.0f, -0.0f, Float.NaN, Float.NegativeInfinity,
Float.MinValue, 1.5f)
+ val doubles = Seq(0.0, -0.0, Double.NaN, Double.PositiveInfinity,
Double.MaxValue, -2.5)
+
+ checkDirect(fill(new BitVector("b", allocator))((v, row) => v.set(row,
row % 3 % 2))) {
+ _.getBoolean(_)
+ }
+ checkDirect(fill(new TinyIntVector("i8", allocator)) { (v, row) =>
+ v.set(row, bytes(row % bytes.size))
+ })(_.getByte(_))
+ checkDirect(fill(new SmallIntVector("i16", allocator)) { (v, row) =>
+ v.set(row, shorts(row % shorts.size))
+ })(_.getShort(_))
+ checkDirect(fill(new IntVector("i32", allocator))((v, row) => v.set(row,
ints(row % 5)))) {
+ _.getInt(_)
+ }
+ checkDirect(fill(new DateDayVector("day", allocator))((v, row) =>
v.set(row, row - 9))) {
+ _.getInt(_)
+ }
+ checkDirect(fill(new BigIntVector("i64", allocator)) { (v, row) =>
+ v.set(row, longs(row % longs.size))
+ })(_.getLong(_))
+ checkDirect(fill(new TimeStampMicroTZVector("ts", allocator, "UTC")) {
(v, row) =>
+ v.set(row, longs(row % longs.size))
+ })(_.getLong(_))
+ checkDirect(fill(new TimeStampMicroVector("ntz", allocator)) { (v, row)
=>
+ v.set(row, longs(row % longs.size))
+ })(_.getLong(_))
+ // The bits, so that NaN equals NaN.
+ checkDirect(fill(new Float4Vector("f32", allocator)) { (v, row) =>
+ v.set(row, floats(row % floats.size))
+ })((column, row) =>
java.lang.Float.floatToRawIntBits(column.getFloat(row)))
+ checkDirect(fill(new Float8Vector("f64", allocator)) { (v, row) =>
+ v.set(row, doubles(row % doubles.size))
+ })((column, row) =>
java.lang.Double.doubleToRawLongBits(column.getDouble(row)))
+ }
+ }
+
+ test("strings and string views, in several batches") {
+ // Values that a view holds, of up to 12 bytes, and longer ones that a
data buffer holds, with
+ // characters of 2 bytes in UTF-8.
+ val (acute, grave) = (0xe9.toChar.toString, 0xe8.toChar.toString)
+ val short = Seq(Some(""), Some("a"), None, Some("twelve bytes"),
Some(acute + grave), None)
+ val long = Seq(Some("thirteen byte"), None, Some(acute * 40), Some("a long
value " * 20))
+ // Values that do not fit in one data buffer of a view vector, which holds
them in several.
+ val spread = Seq(Some("a" * 100000), None, Some(acute * 50000), Some("c" *
100000))
+ withAllocator { allocator =>
+ def check(vector: FieldVector)(set: (Int, Array[Byte]) => Unit, setNull:
Int => Unit)
+ : Unit = {
+ Using.resource(vector) { vector =>
+ val column = NativeColumnVector.create(vector, true, null)
+ assert(column.getClass !== classOf[NativeColumnVector],
vector.getClass.getName)
+ // Batches without and with data buffers, with values of both kinds,
and with values in
+ // several data buffers.
+ Seq(short, long, short ++ long, spread).foreach { values =>
+ vector.clear()
+ vector.allocateNew()
+ values.zipWithIndex.foreach {
+ case (Some(value), row) => set(row,
value.getBytes(StandardCharsets.UTF_8))
+ case (None, row) => setNull(row)
+ }
+ vector.setValueCount(values.size)
+ vector match {
+ case views: ViewVarCharVector if values eq spread =>
+ assert(views.getDataBuffers.size > 1)
+ case _ =>
+ }
+ column.refresh(-1)
+ checkSame(vector, column)(_.getUTF8String(_))
+ assert(column.getUTF8String(values.indexOf(None)) === null)
+ }
+ }
+ }
+ val strings = new VarCharVector("s", allocator)
+ check(strings)(strings.setSafe(_, _), strings.setNull)
+ val views = new ViewVarCharVector("v", allocator)
+ check(views)(views.setSafe(_, _), views.setNull)
+ }
+ }
+
+ private def withDecimalVector(precision: Int, scale: Int)(f: DecimalVector
=> Unit): Unit = {
+ withAllocator { allocator =>
+ Using.resource(new DecimalVector("decimal", allocator, precision,
scale))(f)
+ }
+ }
+
+ /**
+ * Checks that `column` returns the decimal of a row that `expected`, an
[[ArrowColumnVector]] of
+ * the same vector, returns, as Spark uses it, or fails as it does. Its
integral conversions are
+ * those of the decimal of Spark's own column vectors, which hold the
decimals that a long holds,
+ * of a scale from 0 to their precision, as longs, while ArrowColumnVector
holds BigDecimals,
+ * which Decimal rounds to bytes, shorts and ints through doubles.
+ */
+ private def checkDecimal(
+ expected: ArrowColumnVector,
+ column: ArrowColumnVector,
+ row: Int,
+ precision: Int,
+ scale: Int): Unit = {
+ def result(decimal: Decimal): (JavaBigDecimal, Int, Int, Long, UnsafeRow)
= {
+ val writer = new UnsafeRowWriter(1)
+ writer.write(0, decimal, precision, scale)
+ (decimal.toJavaBigDecimal, decimal.precision, decimal.scale,
decimal.toLong, writer.getRow)
+ }
+ def rounded(decimal: Decimal): Seq[Either[String, Long]] = {
+ Seq[Decimal => Long](_.roundToByte(), _.roundToShort(), _.roundToInt(),
_.roundToLong())
+ .map(round => Try(round(decimal)).toEither.left.map(e =>
s"${e.getClass}: ${e.getMessage}"))
+ }
+ def own(decimal: Decimal): Decimal = {
+ if (0 <= scale && scale <= precision && precision <=
Decimal.MAX_LONG_DIGITS) {
+ Decimal.createUnsafe(decimal.toUnscaledLong, precision, scale)
+ } else {
+ decimal
+ }
+ }
+ Try(result(expected.getDecimal(row, precision, scale))) match {
+ case Success(value) =>
+ val decimal = column.getDecimal(row, precision, scale)
+ assert(result(decimal) === value)
+ assert(rounded(decimal) === rounded(own(expected.getDecimal(row,
precision, scale))))
+ case Failure(error) =>
+ val actual = intercept[Exception](result(column.getDecimal(row,
precision, scale)))
+ assert(actual.getClass === error.getClass)
+ assert(actual.getMessage === error.getMessage)
+ }
+ }
+
+ test("decimals in and out of their precision, at other precisions and
scales") {
+ for (precision <- 1 to Decimal.MAX_LONG_DIGITS; scale <- Set(0, precision
/ 2, precision)) {
+ withDecimalVector(precision, scale) { vector =>
+ val limit = BigInt(10).pow(precision)
+ // The bounds of the precision, values out of it, which the byte
setter permits, and
+ // values that a long does not hold.
+ val values = Seq(BigInt(0), BigInt(1), BigInt(-1), limit - 1, -(limit
- 1), limit, -limit,
+ BigInt(Long.MaxValue), BigInt(Long.MinValue), (BigInt(1) << 64) + 7)
+ vector.allocateNew(values.size + 1)
+ values.zipWithIndex.foreach { case (value, row) =>
+ vector.setBigEndian(row, value.toByteArray)
+ }
+ vector.setNull(values.size)
+ vector.setValueCount(values.size + 1)
+ val expected = new ArrowColumnVector(vector)
+ val column = native(vector)
+ assert(column.getClass !== classOf[NativeColumnVector])
+ val requested = Seq((precision, scale), (Decimal.MAX_LONG_DIGITS,
scale), (38, scale),
+ (precision + 1, scale), (precision, scale + 1))
+ for (row <- values.indices; (p, s) <- requested) {
+ checkDecimal(expected, column, row, p, s)
+ }
+ assert(column.getDecimal(values.size, precision, scale) === null)
+ }
+ }
+ }
+
+ test("decimals of a negative scale") {
+ val conf = new SQLConf
+ val key = SQLConf.LEGACY_ALLOW_NEGATIVE_SCALE_OF_DECIMAL_ENABLED.key
+ SQLConf.withExistingConf(conf) {
+ conf.setConfString(key, "true")
+ withDecimalVector(5, -2) { vector =>
+ vector.allocateNew(1)
+ vector.set(0, new JavaBigDecimal("1.23E+4"))
+ vector.setValueCount(1)
+ val expected = new ArrowColumnVector(vector)
+ val column = native(vector)
+ Seq((5, -2), (8, 0), (5, -3)).foreach { case (precision, scale) =>
+ checkDecimal(expected, column, 0, precision, scale)
+ }
+ assert(column.getDecimal(0, 5, -2).toLong === 12300L)
+ conf.setConfString(key, "false")
+ checkDecimal(expected, column, 0, 5, -2)
+ }
+ }
+ }
+
+ test("decimals that a long holds round to integers as those of Spark's own
column vectors") {
+ // The decimals of Spark's own column vectors and of the native one, of
OnHeapColumnVector and
+ // NativeColumnVector, hold longs, which Decimal rounds exactly.
+ def decimals(value: String, scale: Int)(f: (Decimal, Decimal) => Unit):
Unit = {
+ val precision = Decimal.MAX_LONG_DIGITS
+ withDecimalVector(precision, scale) { vector =>
+ vector.allocateNew(1)
+ vector.set(0, new JavaBigDecimal(value))
+ vector.setValueCount(1)
+ Using.resource(new OnHeapColumnVector(1, DecimalType(precision,
scale))) { own =>
+ own.putDecimal(0, Decimal(value), precision)
+ f(own.getDecimal(0, precision, scale), native(vector).getDecimal(0,
precision, scale))
+ }
+ }
+ }
+ decimals("123456789.999999999", 9) { (own, column) =>
+ assert(column.roundToInt() === own.roundToInt())
+ }
+ decimals("127.999999999999999", 15) { (own, column) =>
+ assert(column.roundToByte() === own.roundToByte())
+ }
+ }
+
+ test("the decimals of the fields of a struct still hold BigDecimals") {
+ // ArrowColumnVector reads the fields of a struct, so that their decimals,
unlike those of a
+ // top-level column, hold BigDecimals, which Decimal rounds to ints
through doubles.
+ withAllocator { allocator =>
+ Using.resource(StructVector.empty("struct", allocator)) { struct =>
+ val field = FieldType.nullable(new ArrowType.Decimal(18, 9, 128))
+ val child = struct.addOrGet("d", field, classOf[DecimalVector])
+ struct.allocateNew()
+ child.setSafe(0, new JavaBigDecimal("123456789.999999999"))
+ struct.setIndexDefined(0)
+ child.setValueCount(1)
+ struct.setValueCount(1)
+ val top = native(child).getDecimal(0, 18, 9)
+ val nested = native(struct).getChild(0).getDecimal(0, 18, 9)
+ assert(top.toDebugString === "Decimal(compact, 123456789999999999, 18,
9)")
+ assert(top.roundToInt() === 123456789)
+ assert(nested.toDebugString === "Decimal(expanded,
123456789.999999999, 18, 9)")
+ assert(nested.roundToInt() === 123456790)
Review Comment:
This asserts the double-rounding result for the nested decimal (`123456790`,
while the exact truncation is `123456789`), which locks in the top-level vs.
nested inconsistency (see my comment on `NativeColumnVector.getDecimal`). A
later fix that makes `ArrowColumnVector` return long-backed decimals would then
fail this test. If we keep this behavior for now, could we at least note here
that the nested result is a known inconsistency rather than the intended
semantics?
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeColumnVector.java:
##########
@@ -0,0 +1,345 @@
+/*
+ * 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. The strings, arrays, maps and child vectors that the vector
returns do not refer
+ * to the reader, so this holds for them only while the vector is reachable
too. 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 keeps the addresses of the buffers from one import to the
next, whereas
+ * {@link ArrowColumnVector} reads the buffers of its Arrow vector on each
access. So nothing may
+ * reallocate the buffers of the Arrow vectors of the reader between two
imports, e.g. by growing
+ * them: a consumer that writes them through {@link #getValueVector()}, as by
setting their value
+ * count, must keep it at the length of the batch.
+ *
+ * <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. Nothing
reads the field, but
+ // it must stay: each constructor takes the reader so that no column vector
misses it.
+ private final 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) {
+ if (vector instanceof BitVector) {
+ return new Booleans(vector, nullable, reader);
+ } else if (vector instanceof TinyIntVector) {
+ return new Bytes(vector, nullable, reader);
+ } else if (vector instanceof SmallIntVector) {
+ return new Shorts(vector, nullable, reader);
+ } else if (vector instanceof IntVector || vector instanceof DateDayVector)
{
+ return new Ints(vector, nullable, reader);
+ } else if (vector instanceof BigIntVector || vector instanceof
TimeStampMicroTZVector ||
+ vector instanceof TimeStampMicroVector) {
+ return new Longs(vector, nullable, reader);
+ } else if (vector instanceof Float4Vector) {
+ return new Floats(vector, nullable, reader);
+ } else if (vector instanceof Float8Vector) {
+ return new Doubles(vector, nullable, reader);
+ } else if (vector instanceof DecimalVector decimalVector) {
+ return new Decimals(decimalVector, nullable, reader);
+ } else if (vector instanceof VarCharVector) {
+ return new Strings(vector, nullable, reader);
+ } else if (vector instanceof ViewVarCharVector) {
+ return new StringViews(vector, nullable, reader);
+ } else {
+ return new NativeColumnVector(vector, reader);
+ }
+ }
+
+ NativeColumnVector(ValueVector vector, Object reader) {
+ super(vector);
+ this.reader = reader;
+ }
+
+ /**
+ * 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) {
+ }
+
+ /**
+ * Whether {@code vector} has a null, given the null_count {@code nullCount}
that the library
+ * exported for its array, or -1 if it is not known. 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.
+ */
+ static boolean hasNulls(ValueVector vector, long nullCount) {
+ return nullCount < 0 ? vector.getNullCount() > 0 : nullCount > 0;
+ }
+
+ /** 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, Object reader) {
+ super(vector, reader);
+ this.nullable = nullable;
+ }
+
+ @Override
+ void refresh(long nullCount) {
+ ValueVector vector = getValueVector();
+ boolean hasNulls = nullable && hasNulls(vector, nullCount);
+ validity = hasNulls ? vector.getValidityBuffer().memoryAddress() : 0;
+ 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, boolean nullable, Object reader) {
+ super(vector, nullable, reader);
+ }
+
+ @Override
+ public boolean getBoolean(int rowId) {
+ return bit(data, rowId);
+ }
+ }
+
+ private static final class Bytes extends Direct {
+ Bytes(ValueVector vector, boolean nullable, Object reader) {
+ super(vector, nullable, reader);
+ }
+
+ @Override
+ public byte getByte(int rowId) {
+ return Platform.getByte(null, data + rowId);
+ }
+ }
+
+ private static final class Shorts extends Direct {
+ Shorts(ValueVector vector, boolean nullable, Object reader) {
+ super(vector, nullable, reader);
+ }
+
+ @Override
+ public short getShort(int rowId) {
+ return Platform.getShort(null, data + 2L * rowId);
+ }
+ }
+
+ private static final class Ints extends Direct {
+ Ints(ValueVector vector, boolean nullable, Object reader) {
+ super(vector, nullable, reader);
+ }
+
+ @Override
+ public int getInt(int rowId) {
+ return Platform.getInt(null, data + 4L * rowId);
+ }
+ }
+
+ private static final class Longs extends Direct {
+ Longs(ValueVector vector, boolean nullable, Object reader) {
+ super(vector, nullable, reader);
+ }
+
+ @Override
+ public long getLong(int rowId) {
+ return Platform.getLong(null, data + 8L * rowId);
+ }
+ }
+
+ private static final class Floats extends Direct {
+ Floats(ValueVector vector, boolean nullable, Object reader) {
+ super(vector, nullable, reader);
+ }
+
+ @Override
+ public float getFloat(int rowId) {
+ return Platform.getFloat(null, data + 4L * rowId);
+ }
+ }
+
+ private static final class Doubles extends Direct {
+ Doubles(ValueVector vector, boolean nullable, Object reader) {
+ super(vector, nullable, reader);
+ }
+
+ @Override
+ public double getDouble(int rowId) {
+ return Platform.getDouble(null, data + 8L * rowId);
+ }
+ }
+
+ private static final class Decimals extends Direct {
+ // Arrow stores the 128 bits of a decimal in native byte order, as
ArrowColumnVector reads it.
+ private static final int LOW_WORD = ByteOrder.nativeOrder() ==
ByteOrder.LITTLE_ENDIAN ? 0 : 8;
Review Comment:
`Decimals` re-implements `ArrowColumnVector.SmallDecimalAccessor`'s
low/high-word read and sign-extension check, and `StringViews` (line 307)
re-implements `StringViewAccessor`'s 16-byte view decoding. Keeping the same
Arrow layout logic in two places means that a fix to one copy (e.g. the
existing >2GB view-offset handling in `StringViewAccessor`, or endianness) can
miss the other. Have you considered caching the buffer addresses inside
`ArrowColumnVector`'s accessors instead? That would give every Arrow producer
(Python UDFs, the Arrow cache, etc.) the same speedup with a single
implementation.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/BaseScriptTransformationExec.scala:
##########
@@ -108,6 +108,24 @@ trait BaseScriptTransformationExec extends UnaryExecNode {
(outputStream, proc, inputStream, stderrBuffer)
}
+ /**
+ * Starts the thread that writes the input rows to the process, and destroys
the process once
+ * the task completes, as `PipedRDD` does. Otherwise, when the task
completes before it reads
+ * the whole output of the process, as under a LIMIT, the process blocks on
writing its output
+ * and the thread on writing the rest of the input, so that neither ends.
+ */
+ protected def startWriterThread(writerThread:
BaseScriptTransformationWriterThread): Unit = {
+ writerThread.taskContext.addTaskCompletionListener[Unit] { _ =>
+ if (writerThread.proc.isAlive) {
+ writerThread.proc.destroy()
+ }
+ if (writerThread.isAlive) {
+ writerThread.interrupt()
Review Comment:
This can kill the executor. `interrupt()` sets the interrupt flag of the
writer thread, and `proc.waitFor()` in the writer's `finally` block (`run()`,
line 368) then throws `InterruptedException`. The `case NonFatal(...)` there
does not match it, so it escapes `Utils.logUncaughtExceptions` and reaches
`SparkUncaughtExceptionHandler`, which calls `System.exit` on executors
(installed in `Executor.scala:300` for non-local mode). This is exactly what
the SPARK-25158 comment in `run()` warns about.
For example, `SELECT TRANSFORM(id) USING 'sort' AS (x) FROM t LIMIT 10`: the
writer writes all input, closes stdin and blocks in `proc.waitFor()` while
`sort` emits output. When the task completes after 10 rows, `proc.destroy()`
only sends SIGTERM asynchronously, and the immediate `interrupt()` makes
`waitFor()` throw before the process is reaped. The same happens on job
cancellation and task or speculative kills, and whenever the interrupt lands
while the writer is in a non-interruptible write or in a native `get_next`,
because the flag stays set until `waitFor()`.
Local-mode tests do not install this handler, so the existing tests cannot
catch it. Is the `interrupt()` needed at all, given that destroying the process
already makes the writer's next write fail? Otherwise, could we handle
`InterruptedException` in the writer's `finally` block?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -640,9 +639,26 @@ object NativeWriterCommitMessage {
* `PipedRDD`, while `DataSourceRDD` closes the reader once the task
completes. The library must
* not get the release of the stream during get_next, and the task context
must stay valid until
* then, so [[close]] then leaves the release to [[next]], once get_next
returns, rather than wait
- * for it and keep the task from completing. That [[next]] then returns false,
or throws the error
- * of its read. A [[next]] called once the reader is closed returns false
without calling the
- * library.
+ * for it and keep the task from completing. That [[next]] then throws the
error of its read, or
+ * fails as a [[next]] called once the reader is closed does, without reading
from the library: so
+ * such a thread does not take the close for the end of the data, whether or
not it was in
+ * get_next then. Script transformation, for example, then destroys its
process.
+ *
+ * Such a thread may also still read the batch that [[next]] returned last
when another thread
+ * closes the reader: script transformation gets its first batch on the task
thread, and reads its
+ * rows on a thread of its own. The column vectors read the buffers of the
batch at their
+ * addresses, so [[close]] frees these buffers at once only if [[next]] was
called again since it
+ * returned the batch, as at the end of the stream, after which consumers no
longer read the
+ * batch. Otherwise, it releases the stream, the dictionaries and the task at
once, since the
+ * arrays of a stream may outlive it, and their release does not use the task
(see
+ * `SparkDsTaskContext` in the header). A Cleaner then frees the buffers of
the batch and the
+ * allocator once the reader is unreachable, which each of its column vectors
keeps it from being,
+ * or once a thread calls [[next]] again, which then fails: that thread no
longer reads the batch.
+ * So a reader closed before the end of its stream, as with a LIMIT, and not
called again holds the
+ * buffers of one batch until a garbage collection, which these buffers do not
trigger: the JVM
+ * does not count them. Under `DataSourceRDD`, only a garbage collection frees
them: its
Review Comment:
The documented behavior here, that a thread asking `DataSourceRDD` for the
next batch "then fails without calling next", is in practice a
`NullPointerException`: `PartitionIterator.hasNext` calls `reader.next()`
(`DataSourceRDD.scala:275`) on the `reader` that `close()` set to null
(`DataSourceRDD.scala:265`). So this design depends on an incidental NPE in
another file.
With script transformation over a native source under `LIMIT`, that NPE
becomes `writerThread.exception` and is logged as the cause. If `DataSourceRDD`
ever returns `false` or keeps the reader instead, the freeing and failure
behavior documented here changes silently. Could we make this explicit in
`PartitionIterator`, or avoid depending on it?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -847,20 +890,47 @@ class NativePartitionReader(
}
override def next(): Boolean = {
- // A reader that was closed while it read has released the batch that it
read.
- startReading() && Utils.tryWithSafeFinally(readNext())(stopReading()) &&
synchronized(!closed)
+ startReading()
+ val hasBatch = Utils.tryWithSafeFinally(readNext())(stopReading())
+ synchronized {
+ // A reader that was closed while it read has released the batch that it
read, and fails as
+ // it does for a next called once it is closed (see the class doc).
+ if (closed) {
+ throw closedError()
+ }
+ batchReturned = hasBatch
+ batchReturned
+ }
}
- private def startReading(): Boolean = synchronized {
- reading = !closed
- reading
+ /**
+ * Fails once the reader is closed. In either case, the caller no longer
reads the batch that
+ * next returned last, if any, so a closed reader frees its buffers at once.
+ */
+ private def startReading(): Unit = synchronized {
+ if (closed) {
+ if (batchCleanable != null) {
+ // The Cleaner action does not refer to the batch, which then has no
rows, so that it is
+ // not read at the freed addresses.
+ batch.setNumRows(0)
Review Comment:
This frees the deferred batch for whichever thread calls `next()` on a
closed reader, which assumes that this thread is the batch's only consumer. If
the batch was fetched on one thread and handed to another (as script
transformation does with its first batch, per the class doc), and a thread
other than the one still reading the batch calls `next()` after `close()`,
`clean()` frees the buffers while the reading thread may be mid-iteration,
possibly with `numRows` already cached. That thread then reads freed native
memory. Today only `DataSourceRDD`'s single-caller pattern prevents this. Could
we document this single-consumer assumption here, or guard against it?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/BaseScriptTransformationExec.scala:
##########
@@ -108,6 +108,24 @@ trait BaseScriptTransformationExec extends UnaryExecNode {
(outputStream, proc, inputStream, stderrBuffer)
}
+ /**
+ * Starts the thread that writes the input rows to the process, and destroys
the process once
+ * the task completes, as `PipedRDD` does. Otherwise, when the task
completes before it reads
+ * the whole output of the process, as under a LIMIT, the process blocks on
writing its output
+ * and the thread on writing the rest of the input, so that neither ends.
+ */
+ protected def startWriterThread(writerThread:
BaseScriptTransformationWriterThread): Unit = {
+ writerThread.taskContext.addTaskCompletionListener[Unit] { _ =>
+ if (writerThread.proc.isAlive) {
+ writerThread.proc.destroy()
Review Comment:
Destroying the process at task completion also turns every normal
early-finishing `TRANSFORM` (under `LIMIT`/`show()`, with the script still
alive at task end) into ERROR logs. The writer fails with EPIPE, a
closed-reader error or an NPE and logs `Thread-...-Feed exit cause by:` with a
stack trace, and since `waitFor()` then returns 143 it also logs the stderr
buffer at ERROR. Before this PR the thread hung silently; after it, executor
logs fill with spurious ERRORs for successful queries. Could the writer thread
skip or downgrade these logs once the task has completed, e.g. by checking
`taskContext.isCompleted()`?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -1053,7 +1140,15 @@ 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)
+ batchCleanable = NativeHandle.cleaner.register(this, () =>
Review Comment:
The deferred free runs `root.close()`, and so the library's `ArrowArray`
release callbacks, on the single JVM-wide `NativeHandle.cleaner` thread. I see
that `spark_datasource.h` now documents that the release may happen from a
Cleaner thread, but a library whose array release is slow or blocks (frees a
large arena, takes a lock or the Python GIL) stalls this shared thread and
delays every other `NativeHandle` release in the JVM. Should the header also
ask libraries to keep the release of an array cheap and non-blocking, or should
this free run elsewhere?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/BaseScriptTransformationExec.scala:
##########
@@ -108,6 +108,24 @@ trait BaseScriptTransformationExec extends UnaryExecNode {
(outputStream, proc, inputStream, stderrBuffer)
}
+ /**
+ * Starts the thread that writes the input rows to the process, and destroys
the process once
+ * the task completes, as `PipedRDD` does. Otherwise, when the task
completes before it reads
+ * the whole output of the process, as under a LIMIT, the process blocks on
writing its output
+ * and the thread on writing the rest of the input, so that neither ends.
+ */
+ protected def startWriterThread(writerThread:
BaseScriptTransformationWriterThread): Unit = {
Review Comment:
This changes the behavior of script transformation for every input source
(both `SparkScriptTransformationExec` and `HiveScriptTransformationExec`): it
destroys the process and interrupts the writer thread at task completion. But
there is no test for it, nor for the cross-thread native-batch scenario that
motivates the Cleaner design. Could we add tests in the script-transformation
suites, e.g. `TRANSFORM` with `LIMIT` and a cancelled job, that check that the
process is destroyed and that the writer thread exits without an uncaught
exception? Since local mode does not install `SparkUncaughtExceptionHandler`,
such a test needs to assert on the writer thread's uncaught exception
explicitly to catch the `interrupt()` issue on line 123.
--
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]
