dongjoon-hyun commented on code in PR #59961:
URL: https://github.com/apache/spark/pull/59961#discussion_r4463650728
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -722,6 +763,7 @@ class NativePartitionReader(
}
call(name, "read")(Data.importIntoVectorSchemaRoot(allocator, array,
root, dictionaries))
checkNulls()
+ columns.foreach(_.refresh())
Review Comment:
**[2] When `checkNulls()` throws, the column vectors keep the addresses of a
batch that the import has already released.** The import on L764 loads the new
batch into the root and releases the previous one, but `refresh()` runs only
after `checkNulls()` succeeds, and `batch.setNumRows` after that. So when a
library exports a null in a column that the read schema declares not nullable,
`next()` throws, but `get()` still returns the previous batch, with its row
count and with `NativeColumnVector`s that point at freed buffers. Spark's own
consumers stop at the exception, so this needs a caller that reads the batch
after a failed `next()`. But before this PR such a read went to the new buffers
through Arrow Java's checks, and now it reads freed memory, which can crash the
executor. Could `readNext` call `batch.setNumRows(0)` before the import, and
refresh the columns right after the import, before `checkNulls()`?
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeColumnVector.java:
##########
@@ -0,0 +1,327 @@
+/*
+ * 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 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 until the thread that read it calls
next again or closes
+ * the reader: a close on another thread releases them at once only if that
thread has ended, and
+ * otherwise leaves the release to that thread, or to a thread that waits for
it to end. 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.
+ *
+ * <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 column vector of a top-level Arrow vector of a reader, of a column
that the schema to
+ * read declares {@code nullable} or not.
+ */
+ static NativeColumnVector create(ValueVector vector, boolean nullable) {
+ if (vector instanceof BitVector) {
+ return new Booleans(vector, nullable);
+ } else if (vector instanceof TinyIntVector) {
+ return new Bytes(vector, nullable);
+ } else if (vector instanceof SmallIntVector) {
+ return new Shorts(vector, nullable);
+ } else if (vector instanceof IntVector || vector instanceof DateDayVector)
{
+ return new Ints(vector, nullable);
+ } else if (vector instanceof BigIntVector || vector instanceof
TimeStampMicroTZVector ||
+ vector instanceof TimeStampMicroVector) {
+ return new Longs(vector, nullable);
+ } else if (vector instanceof Float4Vector) {
+ return new Floats(vector, nullable);
+ } else if (vector instanceof Float8Vector) {
+ return new Doubles(vector, nullable);
+ } else if (vector instanceof DecimalVector decimalVector) {
+ return new Decimals(decimalVector, nullable);
+ } else if (vector instanceof VarCharVector) {
+ return new Strings(vector, nullable);
+ } else if (vector instanceof ViewVarCharVector viewVector) {
+ return new StringViews(viewVector, nullable);
+ } 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;
+ // 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.vector = vector;
+ this.nullable = nullable;
+ }
+
+ @Override
+ void refresh() {
+ // Arrow Java counts the nulls in the validity bitmap on each call.
+ boolean hasNulls = nullable && vector.getNullCount() > 0;
+ 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) {
+ super(vector, nullable);
+ }
+
+ @Override
+ public boolean getBoolean(int rowId) {
+ return bit(data, rowId);
+ }
+ }
+
+ private static final class Bytes extends Direct {
+ Bytes(ValueVector vector, boolean nullable) {
+ super(vector, nullable);
+ }
+
+ @Override
+ public byte getByte(int rowId) {
+ return Platform.getByte(null, data + rowId);
+ }
+ }
+
+ private static final class Shorts extends Direct {
+ Shorts(ValueVector vector, boolean nullable) {
+ super(vector, nullable);
+ }
+
+ @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) {
+ super(vector, nullable);
+ }
+
+ @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) {
+ super(vector, nullable);
+ }
+
+ @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) {
+ super(vector, nullable);
+ }
+
+ @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) {
+ super(vector, nullable);
+ }
+
+ @Override
+ public double getDouble(int rowId) {
+ return Platform.getDouble(null, data + 8L * rowId);
+ }
+ }
+
+ private static final class Decimals extends Direct {
+ // 10^i, for the precisions of the decimals that a long holds.
+ private static final long[] POW_10 = new long[Decimal.MAX_LONG_DIGITS() +
1];
Review Comment:
**[3] Nit: `Decimals` repeats the power-of-10 table and the bound check of
`Decimal`.** `Decimal.POW_10()` is the same table (`Decimal.scala:606`), and
`Decimal.setOrNull(unscaled, precision, scale)` returns null exactly when
`|unscaled| >= 10^precision`, for a precision of up to 18 (`Decimal.scala:94`).
Could the high-word check be followed by `Decimal d = new
Decimal().setOrNull(unscaled, precision, scale); if (d != null) return d;`, so
that the table, the static block and the bound on L254-256 can go?
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeColumnVector.java:
##########
@@ -0,0 +1,327 @@
+/*
+ * 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 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 until the thread that read it calls
next again or closes
+ * the reader: a close on another thread releases them at once only if that
thread has ended, and
+ * otherwise leaves the release to that thread, or to a thread that waits for
it to end. 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.
+ *
+ * <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 column vector of a top-level Arrow vector of a reader, of a column
that the schema to
+ * read declares {@code nullable} or not.
+ */
+ static NativeColumnVector create(ValueVector vector, boolean nullable) {
+ if (vector instanceof BitVector) {
+ return new Booleans(vector, nullable);
+ } else if (vector instanceof TinyIntVector) {
+ return new Bytes(vector, nullable);
+ } else if (vector instanceof SmallIntVector) {
+ return new Shorts(vector, nullable);
+ } else if (vector instanceof IntVector || vector instanceof DateDayVector)
{
+ return new Ints(vector, nullable);
+ } else if (vector instanceof BigIntVector || vector instanceof
TimeStampMicroTZVector ||
+ vector instanceof TimeStampMicroVector) {
+ return new Longs(vector, nullable);
+ } else if (vector instanceof Float4Vector) {
+ return new Floats(vector, nullable);
+ } else if (vector instanceof Float8Vector) {
+ return new Doubles(vector, nullable);
+ } else if (vector instanceof DecimalVector decimalVector) {
+ return new Decimals(decimalVector, nullable);
+ } else if (vector instanceof VarCharVector) {
+ return new Strings(vector, nullable);
+ } else if (vector instanceof ViewVarCharVector viewVector) {
+ return new StringViews(viewVector, nullable);
+ } 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;
+ // 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.vector = vector;
+ this.nullable = nullable;
+ }
+
+ @Override
+ void refresh() {
+ // Arrow Java counts the nulls in the validity bitmap on each call.
+ boolean hasNulls = nullable && vector.getNullCount() > 0;
+ 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) {
+ super(vector, nullable);
+ }
+
+ @Override
+ public boolean getBoolean(int rowId) {
+ return bit(data, rowId);
+ }
+ }
+
+ private static final class Bytes extends Direct {
+ Bytes(ValueVector vector, boolean nullable) {
+ super(vector, nullable);
+ }
+
+ @Override
+ public byte getByte(int rowId) {
+ return Platform.getByte(null, data + rowId);
+ }
+ }
+
+ private static final class Shorts extends Direct {
+ Shorts(ValueVector vector, boolean nullable) {
+ super(vector, nullable);
+ }
+
+ @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) {
+ super(vector, nullable);
+ }
+
+ @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) {
+ super(vector, nullable);
+ }
+
+ @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) {
+ super(vector, nullable);
+ }
+
+ @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) {
+ super(vector, nullable);
+ }
+
+ @Override
+ public double getDouble(int rowId) {
+ return Platform.getDouble(null, data + 8L * rowId);
+ }
+ }
+
+ private static final class Decimals extends Direct {
+ // 10^i, for the precisions of the decimals that a long holds.
+ private static final long[] POW_10 = new long[Decimal.MAX_LONG_DIGITS() +
1];
+ // 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;
+
+ static {
+ POW_10[0] = 1;
+ for (int i = 1; i < POW_10.length; i++) {
+ POW_10[i] = POW_10[i - 1] * 10;
+ }
+ }
+
+ // The scale of the Arrow type.
+ private final int scale;
+
+ Decimals(DecimalVector vector, boolean nullable) {
+ super(vector, nullable);
+ 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 + 16L * rowId;
+ long unscaled = Platform.getLong(null, address + LOW_WORD);
+ long bound = POW_10[precision];
+ if (Platform.getLong(null, address + HIGH_WORD) == unscaled >> 63 &&
+ -bound < unscaled && unscaled < bound) {
+ return Decimal.createUnsafe(unscaled, precision, scale);
+ }
+ }
+ // Wider decimals, other scales, and the values that do not fit their
precision, which
+ // ArrowColumnVector reads in full or rejects.
+ return super.getDecimal(rowId, precision, scale);
+ }
+ }
+
+ private static final class Strings extends Direct {
+ // The address of the offsets of the values of the batch.
+ private long offsets;
+
+ Strings(ValueVector vector, boolean nullable) {
+ super(vector, nullable);
+ }
+
+ @Override
+ void refresh() {
+ super.refresh();
+ offsets = vector.getOffsetBuffer().memoryAddress();
+ }
+
+ @Override
+ public UTF8String getUTF8String(int rowId) {
+ if (isNullAt(rowId)) return null;
+ long offset = offsets + 4L * rowId;
+ int start = Platform.getInt(null, offset);
+ int length = Platform.getInt(null, offset + 4) - start;
+ return UTF8String.fromAddress(null, data + start, length);
+ }
+ }
+
+ private static final class StringViews extends Direct {
+ private final ViewVarCharVector views;
+ // The addresses of the data buffers of the batch, which hold the values
that a view does not.
+ private long[] buffers = new long[0];
+
+ StringViews(ViewVarCharVector vector, boolean nullable) {
+ super(vector, nullable);
+ this.views = vector;
+ }
+
+ @Override
+ void refresh() {
+ super.refresh();
+ List<ArrowBuf> dataBuffers = views.getDataBuffers();
+ if (buffers.length != dataBuffers.size()) {
+ buffers = new long[dataBuffers.size()];
+ }
+ for (int i = 0; i < buffers.length; i++) {
+ buffers[i] = dataBuffers.get(i).memoryAddress();
+ }
+ }
+
+ @Override
+ public UTF8String getUTF8String(int rowId) {
+ if (isNullAt(rowId)) return null;
+ // A view of 16 bytes: the length of the value, then the value if it
fits in the other 12
+ // bytes, or its first 4 bytes, the index of the data buffer that holds
it, and its offset
+ // in that buffer.
+ long view = data + (long) BaseVariableWidthViewVector.ELEMENT_SIZE *
rowId;
+ int length = Platform.getInt(null, view);
+ if (length <= BaseVariableWidthViewVector.INLINE_SIZE) {
+ return UTF8String.fromAddress(null, view + 4, length);
Review Comment:
**[4] Nit: the view and decimal offsets are literals, although L318 and L320
already use Arrow's layout constants.** L318 and L320 use
`BaseVariableWidthViewVector.ELEMENT_SIZE` and `INLINE_SIZE`, but L321-324
hardcode `view + 4`, `+ 8` and `+ 12`, and L252 hardcodes `16L * rowId`.
`ArrowColumnVector.StringViewAccessor` reads the same fields with
`LENGTH_WIDTH`, `PREFIX_WIDTH` and `BUF_INDEX_WIDTH`
(`ArrowColumnVector.java:591-599`), and `SmallDecimalAccessor` with
`DecimalVector.TYPE_WIDTH` (`ArrowColumnVector.java:473`). Could these reads
use the same constants, so that a reader can check the unchecked reads against
the Arrow layout?
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeColumnVector.java:
##########
@@ -0,0 +1,327 @@
+/*
+ * 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 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 until the thread that read it calls
next again or closes
+ * the reader: a close on another thread releases them at once only if that
thread has ended, and
+ * otherwise leaves the release to that thread, or to a thread that waits for
it to end. 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.
+ *
+ * <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 column vector of a top-level Arrow vector of a reader, of a column
that the schema to
+ * read declares {@code nullable} or not.
+ */
+ static NativeColumnVector create(ValueVector vector, boolean nullable) {
+ if (vector instanceof BitVector) {
+ return new Booleans(vector, nullable);
+ } else if (vector instanceof TinyIntVector) {
+ return new Bytes(vector, nullable);
+ } else if (vector instanceof SmallIntVector) {
+ return new Shorts(vector, nullable);
+ } else if (vector instanceof IntVector || vector instanceof DateDayVector)
{
+ return new Ints(vector, nullable);
+ } else if (vector instanceof BigIntVector || vector instanceof
TimeStampMicroTZVector ||
+ vector instanceof TimeStampMicroVector) {
+ return new Longs(vector, nullable);
+ } else if (vector instanceof Float4Vector) {
+ return new Floats(vector, nullable);
+ } else if (vector instanceof Float8Vector) {
+ return new Doubles(vector, nullable);
+ } else if (vector instanceof DecimalVector decimalVector) {
+ return new Decimals(decimalVector, nullable);
+ } else if (vector instanceof VarCharVector) {
+ return new Strings(vector, nullable);
+ } else if (vector instanceof ViewVarCharVector viewVector) {
+ return new StringViews(viewVector, nullable);
+ } 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;
+ // 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.vector = vector;
+ this.nullable = nullable;
+ }
+
+ @Override
+ void refresh() {
+ // Arrow Java counts the nulls in the validity bitmap on each call.
+ boolean hasNulls = nullable && vector.getNullCount() > 0;
+ 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) {
+ super(vector, nullable);
+ }
+
+ @Override
+ public boolean getBoolean(int rowId) {
+ return bit(data, rowId);
+ }
+ }
+
+ private static final class Bytes extends Direct {
+ Bytes(ValueVector vector, boolean nullable) {
+ super(vector, nullable);
+ }
+
+ @Override
+ public byte getByte(int rowId) {
+ return Platform.getByte(null, data + rowId);
+ }
+ }
+
+ private static final class Shorts extends Direct {
+ Shorts(ValueVector vector, boolean nullable) {
+ super(vector, nullable);
+ }
+
+ @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) {
+ super(vector, nullable);
+ }
+
+ @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) {
+ super(vector, nullable);
+ }
+
+ @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) {
+ super(vector, nullable);
+ }
+
+ @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) {
+ super(vector, nullable);
+ }
+
+ @Override
+ public double getDouble(int rowId) {
+ return Platform.getDouble(null, data + 8L * rowId);
+ }
+ }
+
+ private static final class Decimals extends Direct {
+ // 10^i, for the precisions of the decimals that a long holds.
+ private static final long[] POW_10 = new long[Decimal.MAX_LONG_DIGITS() +
1];
+ // 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;
+
+ static {
+ POW_10[0] = 1;
+ for (int i = 1; i < POW_10.length; i++) {
+ POW_10[i] = POW_10[i - 1] * 10;
+ }
+ }
+
+ // The scale of the Arrow type.
+ private final int scale;
+
+ Decimals(DecimalVector vector, boolean nullable) {
+ super(vector, nullable);
+ this.scale = vector.getScale();
+ }
+
+ @Override
+ public Decimal getDecimal(int rowId, int precision, int scale) {
Review Comment:
**[1] Only top-level decimals become long-backed, so a nested decimal of the
same native source still casts differently.** Following up on @sarutak's
[thread](https://github.com/apache/spark/pull/59961#discussion_r4429239848):
the fields of structs and the elements of arrays and maps are still read by
plain `ArrowColumnVector` children (`ArrowColumnVector.java:258`, `:782`,
`:817-818`), whose `SmallDecimalAccessor` keeps BigDecimal-backed decimals on
purpose (SPARK-59443), and `Decimal.roundToNumeric` decides those on
`toDouble`. So with ANSI on, for a native source with `d DECIMAL(18,9)` and `st
STRUCT<d: DECIMAL(18,9)>` that both hold `123456789.999999999`, `SELECT CAST(d
AS INT), CAST(st.d AS INT)` returns 123456789 and 123456790. For DECIMAL(18,15)
`127.999999999999999`, `CAST(d AS TINYINT)` returns 127, while `CAST(st.d AS
TINYINT)` fails with CAST_OVERFLOW. Before this PR, both columns agreed. Could
nested decimals get the same long-backed reads, or, if that belongs in a fol
low-up, could a test pin the current nested result so that the difference
stays visible?
--
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]
