gengliangwang commented on code in PR #59961: URL: https://github.com/apache/spark/pull/59961#discussion_r4470203152
########## 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: Thanks, the difference is real. The fields of structs and the elements of arrays and maps are built inside `ArrowColumnVector` itself (`new ArrowColumnVector(child)` in the struct branch, and in `ArrowArray`/`ArrowMap` accessors), so `NativeColumnVector` has no hook to swap them out. Giving them long-backed reads means changing how `ArrowColumnVector` creates its children in catalyst. That also affects every other `ArrowColumnVector` user, which SPARK-59443 deliberately kept BigDecimal-backed, so I'd rather do it as a separate follow-up than widen this PR. For now I added a test to `NativeColumnVectorSuite` ("the decimals of the fields of a struct still hold BigDecimals"). It uses your example: DECIMAL(18,9) `123456789.999999999` read top-level is compact and `roundToInt` gives 123456789, while the same value as a struct field is expanded and gives 123456790. That keeps the current nested result pinned and visible until the follow-up changes it. Nested decimals behave as they did before this PR; only top-level ones now match Spark's own column vectors. ########## 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: Done in 5d9a78babea0: `readNext` now sets the batch's row count to 0 before the import and refreshes the column vectors right after it, before `checkNulls()`, so a failed `next()` leaves an empty batch whose vectors point at the new buffers. ########## 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: Done in 5d9a78babea0: After the high-word check, `Decimals` now returns `new Decimal().setOrNull(unscaled, precision, scale)` when it is non-null; the local `POW_10` table, its static block and the bound check are removed. ########## 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: Done in 5d9a78babea0: The decimal offset now uses `DecimalVector.TYPE_WIDTH`, and the string-view reads use `LENGTH_WIDTH`, `PREFIX_WIDTH` and `BUF_INDEX_WIDTH`, as `ArrowColumnVector.StringViewAccessor` does. -- 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]
