mbutrovich commented on code in PR #6730:
URL: https://github.com/apache/datafusion-comet/pull/6730#discussion_r4210361552
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala:
##########
@@ -1031,31 +1190,57 @@ private[arrow] class DoubleWriter(val valueVector:
Float8Vector)
* that `getDecimal` would build is only built for a value that does not fit,
to fail as before.
*/
private[arrow] class DecimalWriter(val valueVector: DecimalVector, precision:
Int, scale: Int)
- extends FixedWidthArrowFieldWriter {
+ extends FixedWidthArrowFieldWriter(valueVector) {
import ArrowFieldWriter.LittleEndian
override def setValue(input: SpecializedGetters, ordinal: Int): Unit = {
ensureCapacity(count + 1)
setValueUnsafe(input, ordinal)
}
+ // Spark's unsafe rows and arrays hold up to 18 digits as the unscaled long,
and more as the
+ // unscaled bytes, which are read in place. An unsafe row's getDecimal does
not check the long
+ // against the precision, but an unsafe array's does, so a long that does
not fit is left to it.
override protected def setValueUnsafe(input: SpecializedGetters, ordinal:
Int): Unit = {
- if (precision > Decimal.MAX_LONG_DIGITS && LittleEndian &&
- (input.isInstanceOf[UnsafeRow] || input.isInstanceOf[UnsafeArrayData])) {
- val bytes = input.getBinary(ordinal)
- if (putBigEndianIfFits(count, bytes, Platform.BYTE_ARRAY_OFFSET.toLong,
bytes.length)) {
- BitVectorHelper.setBit(valueVector.getValidityBuffer, count.toLong)
- return
- }
- }
- val decimal = input.getDecimal(ordinal, precision, scale)
- if (decimal.changePrecision(precision, scale)) {
- setUnscaled(count, decimal)
+ val written = LittleEndian && (input match {
+ case row: UnsafeRow if precision <= Decimal.MAX_LONG_DIGITS =>
+ putLong(valueAddress, row.getLong(ordinal))
+ true
+ case array: UnsafeArrayData if precision <= Decimal.MAX_LONG_DIGITS =>
+ val unscaled = array.getLong(ordinal)
+ val fits = DecimalWriter.fitsPrecision(unscaled >> 63, unscaled,
precision)
+ if (fits) {
+ putLong(valueAddress, unscaled)
+ }
+ fits
Review Comment:
Could we add a test for this range check? The description says a value past
the precision in an unsafe array still fails as before, but the only overflow
test I found, [`a wide decimal past its precision still fails the row
path`](https://github.com/apache/datafusion-comet/blob/0182cc9c03d599836157adeef5a851ce6260cbaa/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowWriterSuite.scala#L490-L508),
uses an unsafe row. If I replace the `fitsPrecision` call here with `true`,
`CometArrowWriterSuite`, `CometStringWriterSuite` and `CometArrowStreamSuite`
still pass, so nothing would catch the array path writing an out-of-range long
without the error that `UnsafeArrayData.getDecimal` raises. Here is something
along the lines of the existing test, covering both the long-backed branch and
the wide `putUnsafeBytesIfFits` branch below it:
```scala
test("a decimal past its precision in an unsafe array still fails the row
path") {
// Written at a wider precision and read at a narrower one, so the array
holds more digits.
Seq(
(DecimalType(18, 2), DecimalType(5, 2), "1234567.89"),
(DecimalType(38, 0), DecimalType(20, 0),
"123456789012345678901")).foreach {
case (written, read, value) =>
val row = UnsafeProjection.create(new StructType().add("a",
ArrayType(written)))(
new GenericInternalRow(Array[Any](new GenericArrayData(Array[Any](
Decimal(new JavaBigDecimal(value), written.precision,
written.scale))))))
val allocator = new RootAllocator(Long.MaxValue)
val root = VectorSchemaRoot.create(
Utils.toArrowSchema(new StructType().add("a", ArrayType(read)),
"UTC"),
allocator)
try {
val writer = ArrowWriter.create(root, 1)
intercept[ArithmeticException](writer.write(row))
} finally {
root.close()
allocator.close()
}
}
}
```
--
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]
