sunchao commented on code in PR #5560:
URL: https://github.com/apache/datafusion-comet/pull/5560#discussion_r3991506148
##########
spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala:
##########
@@ -338,6 +348,170 @@ private[python] trait CometArrowPythonRunnerBase
private[python] object CometArrowPythonRunnerBase {
+ // A regular Arrow variable-width data buffer uses signed 32-bit offsets.
The Spark setting is
+ // already restricted to this range, but cap it here as a final guard for
direct test callers.
+ private val MaxDecodedBatchBytes = Int.MaxValue.toLong
+
+ private def dictionaryVector(column: CometDictionaryVector): FieldVector = {
+ val indices = column.getValueVector
+ val encoding = indices.getField.getDictionary
+ column.getDictionaryProvider.lookup(encoding.getId).getVector
+ }
+
+ private def initialDecodedBytes(values: FieldVector): Long =
+ values match {
+ case _: BaseVariableWidthVector => BaseVariableWidthVector.OFFSET_WIDTH
+ case _: BaseLargeVariableWidthVector =>
BaseLargeVariableWidthVector.OFFSET_WIDTH
+ case _ => 0L
+ }
+
+ /** Conservative logical bytes added by one decoded dictionary value. */
+ private def decodedValueBytes(
+ column: CometDictionaryVector,
+ values: FieldVector,
+ row: Int,
+ batchRow: Int): Long = {
+ val dictionaryIndex = if (column.isNullAt(row)) -1 else
column.indices.getInt(row)
+ val validityBytes = if ((batchRow & 7) == 0) 1L else 0L
+ values match {
+ case vector: BaseVariableWidthVector =>
+ val valueBytes = if (dictionaryIndex < 0) 0L else
vector.getValueLength(dictionaryIndex)
+ valueBytes + BaseVariableWidthVector.OFFSET_WIDTH + validityBytes
+ case vector: BaseLargeVariableWidthVector =>
+ val valueBytes = if (dictionaryIndex < 0) 0L else
vector.getValueLength(dictionaryIndex)
+ valueBytes + BaseLargeVariableWidthVector.OFFSET_WIDTH + validityBytes
+ case vector: BaseFixedWidthVector =>
+ vector.getBufferSizeFor(batchRow + 1).toLong -
+ vector.getBufferSizeFor(batchRow).toLong
+ case _: NullVector => 0L
+ case vector =>
+ // Comet's JVM shuffle currently dictionary-encodes only strings and
binary values.
+ // If another Arrow type reaches this path, the complete dictionary is
a safe upper
+ // bound for any one selected value and favors smaller batches over a
large allocation.
+ math.max(1L, vector.getBufferSize.toLong)
+ }
+ }
+
+ private def saturatedAdd(left: Long, right: Long): Long =
+ if (right >= Long.MaxValue - left) Long.MaxValue else left + right
+
+ /**
+ * Split a compact dictionary batch before decoding it.
+ *
+ * The byte estimate covers the temporary logical dictionary vectors. Plain
input vectors are
+ * already allocated and remain zero-copy when no dictionary column is
present. Every returned
+ * range is applied to all columns so rows stay aligned. A single oversized
row is allowed,
+ * matching Spark's Arrow batching contract.
+ */
+ private[python] def inputBatchRanges(
+ columns: Seq[CometDecodedVector],
+ numRows: Int,
+ maxRecordsPerBatch: Int,
+ maxBytesPerBatch: Long): Seq[(Int, Int)] = {
+ require(numRows >= 0, s"Input batch row count must be non-negative:
$numRows")
+
+ val dictionaries = columns.collect { case column: CometDictionaryVector =>
Review Comment:
Added the explicit check suggested here and in @viirya’s follow-up. Before
decoding, the runner recursively inspects child field metadata and rejects
dictionary-encoded descendants with their full path. The current shuffle
limitation is documented.
[`nested dictionaries fail with their full column path before materializing
input`](https://github.com/apache/datafusion-comet/blob/e75103f8b51bab3ec440bd9ee0fcb16b0ec4e069/spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala#L999)
passes locally on Spark 4.1. It checks the `outer.inner` error, that the
callback is never entered, and that source allocation/refcounts remain
unchanged.
##########
spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala:
##########
@@ -338,6 +348,170 @@ private[python] trait CometArrowPythonRunnerBase
private[python] object CometArrowPythonRunnerBase {
+ // A regular Arrow variable-width data buffer uses signed 32-bit offsets.
The Spark setting is
+ // already restricted to this range, but cap it here as a final guard for
direct test callers.
+ private val MaxDecodedBatchBytes = Int.MaxValue.toLong
+
+ private def dictionaryVector(column: CometDictionaryVector): FieldVector = {
+ val indices = column.getValueVector
+ val encoding = indices.getField.getDictionary
+ column.getDictionaryProvider.lookup(encoding.getId).getVector
+ }
+
+ private def initialDecodedBytes(values: FieldVector): Long =
+ values match {
+ case _: BaseVariableWidthVector => BaseVariableWidthVector.OFFSET_WIDTH
+ case _: BaseLargeVariableWidthVector =>
BaseLargeVariableWidthVector.OFFSET_WIDTH
+ case _ => 0L
+ }
+
+ /** Conservative logical bytes added by one decoded dictionary value. */
+ private def decodedValueBytes(
+ column: CometDictionaryVector,
+ values: FieldVector,
+ row: Int,
+ batchRow: Int): Long = {
+ val dictionaryIndex = if (column.isNullAt(row)) -1 else
column.indices.getInt(row)
+ val validityBytes = if ((batchRow & 7) == 0) 1L else 0L
+ values match {
+ case vector: BaseVariableWidthVector =>
+ val valueBytes = if (dictionaryIndex < 0) 0L else
vector.getValueLength(dictionaryIndex)
+ valueBytes + BaseVariableWidthVector.OFFSET_WIDTH + validityBytes
+ case vector: BaseLargeVariableWidthVector =>
+ val valueBytes = if (dictionaryIndex < 0) 0L else
vector.getValueLength(dictionaryIndex)
+ valueBytes + BaseLargeVariableWidthVector.OFFSET_WIDTH + validityBytes
+ case vector: BaseFixedWidthVector =>
+ vector.getBufferSizeFor(batchRow + 1).toLong -
+ vector.getBufferSizeFor(batchRow).toLong
+ case _: NullVector => 0L
+ case vector =>
+ // Comet's JVM shuffle currently dictionary-encodes only strings and
binary values.
+ // If another Arrow type reaches this path, the complete dictionary is
a safe upper
+ // bound for any one selected value and favors smaller batches over a
large allocation.
+ math.max(1L, vector.getBufferSize.toLong)
+ }
+ }
+
+ private def saturatedAdd(left: Long, right: Long): Long =
+ if (right >= Long.MaxValue - left) Long.MaxValue else left + right
+
+ /**
+ * Split a compact dictionary batch before decoding it.
+ *
+ * The byte estimate covers the temporary logical dictionary vectors. Plain
input vectors are
+ * already allocated and remain zero-copy when no dictionary column is
present. Every returned
+ * range is applied to all columns so rows stay aligned. A single oversized
row is allowed,
+ * matching Spark's Arrow batching contract.
+ */
+ private[python] def inputBatchRanges(
+ columns: Seq[CometDecodedVector],
+ numRows: Int,
+ maxRecordsPerBatch: Int,
+ maxBytesPerBatch: Long): Seq[(Int, Int)] = {
+ require(numRows >= 0, s"Input batch row count must be non-negative:
$numRows")
+
+ val dictionaries = columns.collect { case column: CometDictionaryVector =>
+ column -> dictionaryVector(column)
+ }
+ if (numRows == 0 || dictionaries.isEmpty) {
Review Comment:
Applied `maxRecordsPerBatch` uniformly, including plain-only and zero-column
inputs. The byte estimate still covers only decoded dictionary buffers, and
`pyarrow-udfs.md` now states that explicitly.
[`plain and zero-column inputs obey record limits without estimating decoded
bytes`](https://github.com/apache/datafusion-comet/blob/e75103f8b51bab3ec440bd9ee0fcb16b0ec4e069/spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala#L961)
passes locally on Spark 4.1, including `[3, 3, 3, 1]` ranges for ten rows. The
mixed-column IPC regression verifies that plain and dictionary columns stay
aligned through serialization.
##########
spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala:
##########
@@ -338,6 +348,170 @@ private[python] trait CometArrowPythonRunnerBase
private[python] object CometArrowPythonRunnerBase {
+ // A regular Arrow variable-width data buffer uses signed 32-bit offsets.
The Spark setting is
+ // already restricted to this range, but cap it here as a final guard for
direct test callers.
+ private val MaxDecodedBatchBytes = Int.MaxValue.toLong
+
+ private def dictionaryVector(column: CometDictionaryVector): FieldVector = {
+ val indices = column.getValueVector
+ val encoding = indices.getField.getDictionary
+ column.getDictionaryProvider.lookup(encoding.getId).getVector
+ }
+
+ private def initialDecodedBytes(values: FieldVector): Long =
+ values match {
+ case _: BaseVariableWidthVector => BaseVariableWidthVector.OFFSET_WIDTH
+ case _: BaseLargeVariableWidthVector =>
BaseLargeVariableWidthVector.OFFSET_WIDTH
+ case _ => 0L
+ }
+
+ /** Conservative logical bytes added by one decoded dictionary value. */
+ private def decodedValueBytes(
+ column: CometDictionaryVector,
+ values: FieldVector,
+ row: Int,
+ batchRow: Int): Long = {
+ val dictionaryIndex = if (column.isNullAt(row)) -1 else
column.indices.getInt(row)
+ val validityBytes = if ((batchRow & 7) == 0) 1L else 0L
+ values match {
+ case vector: BaseVariableWidthVector =>
+ val valueBytes = if (dictionaryIndex < 0) 0L else
vector.getValueLength(dictionaryIndex)
+ valueBytes + BaseVariableWidthVector.OFFSET_WIDTH + validityBytes
+ case vector: BaseLargeVariableWidthVector =>
+ val valueBytes = if (dictionaryIndex < 0) 0L else
vector.getValueLength(dictionaryIndex)
+ valueBytes + BaseLargeVariableWidthVector.OFFSET_WIDTH + validityBytes
+ case vector: BaseFixedWidthVector =>
+ vector.getBufferSizeFor(batchRow + 1).toLong -
+ vector.getBufferSizeFor(batchRow).toLong
+ case _: NullVector => 0L
+ case vector =>
+ // Comet's JVM shuffle currently dictionary-encodes only strings and
binary values.
+ // If another Arrow type reaches this path, the complete dictionary is
a safe upper
+ // bound for any one selected value and favors smaller batches over a
large allocation.
+ math.max(1L, vector.getBufferSize.toLong)
+ }
+ }
+
+ private def saturatedAdd(left: Long, right: Long): Long =
+ if (right >= Long.MaxValue - left) Long.MaxValue else left + right
+
+ /**
+ * Split a compact dictionary batch before decoding it.
+ *
+ * The byte estimate covers the temporary logical dictionary vectors. Plain
input vectors are
+ * already allocated and remain zero-copy when no dictionary column is
present. Every returned
+ * range is applied to all columns so rows stay aligned. A single oversized
row is allowed,
+ * matching Spark's Arrow batching contract.
+ */
+ private[python] def inputBatchRanges(
+ columns: Seq[CometDecodedVector],
+ numRows: Int,
+ maxRecordsPerBatch: Int,
+ maxBytesPerBatch: Long): Seq[(Int, Int)] = {
+ require(numRows >= 0, s"Input batch row count must be non-negative:
$numRows")
+
+ val dictionaries = columns.collect { case column: CometDictionaryVector =>
+ column -> dictionaryVector(column)
+ }
+ if (numRows == 0 || dictionaries.isEmpty) {
+ return Seq(0 -> numRows)
+ }
+
+ val recordLimit =
+ if (maxRecordsPerBatch > 0) maxRecordsPerBatch else Int.MaxValue
+ val byteLimit =
+ if (maxBytesPerBatch > 0) math.min(maxBytesPerBatch,
MaxDecodedBatchBytes)
+ else MaxDecodedBatchBytes
+ val initialBytes = dictionaries.foldLeft(0L) { case (bytes, (_, values)) =>
+ saturatedAdd(bytes, initialDecodedBytes(values))
+ }
+
+ val ranges = Seq.newBuilder[(Int, Int)]
+ var start = 0
+ var row = 0
+ var decodedBytes = initialBytes
+ while (row < numRows) {
+ var rowsInBatch = row - start
+ var rowBytes = dictionaries.foldLeft(0L) { case (bytes, (column,
values)) =>
+ saturatedAdd(bytes, decodedValueBytes(column, values, row,
rowsInBatch))
+ }
+ // Spark checks the configured byte limit before adding the next row, so
the row that
+ // crosses that soft limit stays in the current batch. The separate hard
check prevents a
+ // regular variable-width buffer from crossing Arrow's signed 32-bit
allocation ceiling.
+ val exceedsArrowLimit =
Review Comment:
Removed the redundant `decodedBytes >= MaxDecodedBatchBytes` disjunct and
`saturatedAdd`, and retained `rowBytes > MaxDecodedBatchBytes - decodedBytes`
as explained in @viirya’s reply. The configured limit is soft, so the crossing
row can otherwise combine with earlier rows into an oversized decode. The code
now includes that reachable example.
[`the hard Arrow ceiling splits gigabyte values with soft limits raised or
disabled`](https://github.com/apache/datafusion-comet/blob/e75103f8b51bab3ec440bd9ee0fcb16b0ec4e069/spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala#L981)
passes locally on Spark 4.1. It uses synthetic 1,100 MiB value lengths and
checks one-row ranges with raised or disabled soft limits, without allocating
gigabyte buffers.
##########
spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala:
##########
@@ -338,6 +348,170 @@ private[python] trait CometArrowPythonRunnerBase
private[python] object CometArrowPythonRunnerBase {
+ // A regular Arrow variable-width data buffer uses signed 32-bit offsets.
The Spark setting is
+ // already restricted to this range, but cap it here as a final guard for
direct test callers.
+ private val MaxDecodedBatchBytes = Int.MaxValue.toLong
+
+ private def dictionaryVector(column: CometDictionaryVector): FieldVector = {
+ val indices = column.getValueVector
+ val encoding = indices.getField.getDictionary
+ column.getDictionaryProvider.lookup(encoding.getId).getVector
+ }
+
+ private def initialDecodedBytes(values: FieldVector): Long =
+ values match {
+ case _: BaseVariableWidthVector => BaseVariableWidthVector.OFFSET_WIDTH
+ case _: BaseLargeVariableWidthVector =>
BaseLargeVariableWidthVector.OFFSET_WIDTH
+ case _ => 0L
+ }
+
+ /** Conservative logical bytes added by one decoded dictionary value. */
+ private def decodedValueBytes(
+ column: CometDictionaryVector,
+ values: FieldVector,
+ row: Int,
+ batchRow: Int): Long = {
+ val dictionaryIndex = if (column.isNullAt(row)) -1 else
column.indices.getInt(row)
+ val validityBytes = if ((batchRow & 7) == 0) 1L else 0L
+ values match {
+ case vector: BaseVariableWidthVector =>
+ val valueBytes = if (dictionaryIndex < 0) 0L else
vector.getValueLength(dictionaryIndex)
+ valueBytes + BaseVariableWidthVector.OFFSET_WIDTH + validityBytes
+ case vector: BaseLargeVariableWidthVector =>
+ val valueBytes = if (dictionaryIndex < 0) 0L else
vector.getValueLength(dictionaryIndex)
+ valueBytes + BaseLargeVariableWidthVector.OFFSET_WIDTH + validityBytes
+ case vector: BaseFixedWidthVector =>
+ vector.getBufferSizeFor(batchRow + 1).toLong -
+ vector.getBufferSizeFor(batchRow).toLong
+ case _: NullVector => 0L
+ case vector =>
+ // Comet's JVM shuffle currently dictionary-encodes only strings and
binary values.
+ // If another Arrow type reaches this path, the complete dictionary is
a safe upper
+ // bound for any one selected value and favors smaller batches over a
large allocation.
+ math.max(1L, vector.getBufferSize.toLong)
+ }
+ }
+
+ private def saturatedAdd(left: Long, right: Long): Long =
+ if (right >= Long.MaxValue - left) Long.MaxValue else left + right
+
+ /**
+ * Split a compact dictionary batch before decoding it.
+ *
+ * The byte estimate covers the temporary logical dictionary vectors. Plain
input vectors are
+ * already allocated and remain zero-copy when no dictionary column is
present. Every returned
+ * range is applied to all columns so rows stay aligned. A single oversized
row is allowed,
+ * matching Spark's Arrow batching contract.
+ */
+ private[python] def inputBatchRanges(
+ columns: Seq[CometDecodedVector],
+ numRows: Int,
+ maxRecordsPerBatch: Int,
+ maxBytesPerBatch: Long): Seq[(Int, Int)] = {
+ require(numRows >= 0, s"Input batch row count must be non-negative:
$numRows")
+
+ val dictionaries = columns.collect { case column: CometDictionaryVector =>
+ column -> dictionaryVector(column)
+ }
+ if (numRows == 0 || dictionaries.isEmpty) {
+ return Seq(0 -> numRows)
+ }
+
+ val recordLimit =
+ if (maxRecordsPerBatch > 0) maxRecordsPerBatch else Int.MaxValue
+ val byteLimit =
+ if (maxBytesPerBatch > 0) math.min(maxBytesPerBatch,
MaxDecodedBatchBytes)
+ else MaxDecodedBatchBytes
+ val initialBytes = dictionaries.foldLeft(0L) { case (bytes, (_, values)) =>
+ saturatedAdd(bytes, initialDecodedBytes(values))
+ }
+
+ val ranges = Seq.newBuilder[(Int, Int)]
+ var start = 0
+ var row = 0
+ var decodedBytes = initialBytes
+ while (row < numRows) {
+ var rowsInBatch = row - start
+ var rowBytes = dictionaries.foldLeft(0L) { case (bytes, (column,
values)) =>
+ saturatedAdd(bytes, decodedValueBytes(column, values, row,
rowsInBatch))
+ }
+ // Spark checks the configured byte limit before adding the next row, so
the row that
+ // crosses that soft limit stays in the current batch. The separate hard
check prevents a
+ // regular variable-width buffer from crossing Arrow's signed 32-bit
allocation ceiling.
+ val exceedsArrowLimit =
+ decodedBytes >= MaxDecodedBatchBytes ||
+ rowBytes > MaxDecodedBatchBytes - decodedBytes
+ if (rowsInBatch > 0 &&
+ (rowsInBatch >= recordLimit || decodedBytes >= byteLimit ||
exceedsArrowLimit)) {
+ ranges += start -> rowsInBatch
+ start = row
+ decodedBytes = initialBytes
+ rowsInBatch = 0
+ rowBytes = dictionaries.foldLeft(0L) { case (bytes, (column, values))
=>
+ saturatedAdd(bytes, decodedValueBytes(column, values, row,
rowsInBatch))
+ }
+ }
+ decodedBytes = saturatedAdd(decodedBytes, rowBytes)
+ row += 1
+ }
+ ranges += start -> (numRows - start)
+ ranges.result()
+ }
+
+ /** Materialize and visit each safely sized, row-aligned input range
synchronously. */
+ private[python] def foreachInputBatch(
+ columns: Seq[CometDecodedVector],
+ numRows: Int,
+ maxRecordsPerBatch: Int,
+ maxBytesPerBatch: Long,
+ allocator: BufferAllocator)(body: (Seq[FieldVector], Int) => Unit): Unit
= {
+ inputBatchRanges(columns, numRows, maxRecordsPerBatch,
maxBytesPerBatch).foreach {
+ case (0, length) if length == numRows =>
+ withMaterializedInputVectors(columns, allocator)(body(_, length))
+ case (offset, length) =>
+ val slices = new ArrayList[CometDecodedVector]()
+ try {
+ columns.foreach { column =>
+ slices.add(column.slice(offset,
length).asInstanceOf[CometDecodedVector])
+ }
+ withMaterializedInputVectors(slices.asScala.toSeq,
allocator)(body(_, length))
+ } finally {
+ slices.asScala.reverseIterator.foreach(_.close())
+ }
+ }
+ }
+
+ /**
+ * Supply logical Arrow vectors to the serializer for the duration of the
body.
+ *
+ * Plain Comet vectors already expose their logical values and remain
borrowed.
+ * Dictionary-backed shuffle columns expose only their integer indices
through getValueVector,
+ * so materialize those columns first. The temporary decoded vectors own
their buffers and are
+ * closed after the synchronous write, including schema and serialization
failures.
+ */
+ private[python] def withMaterializedInputVectors[T](
Review Comment:
Extracted `CometVectorUtils.withDecodedVectors` for both the Python runner
and `ColumnarBatchArrowReader`, plus `CometDictionaryVector.getDictionary` for
lookup in those paths and schema reconciliation. Cleanup closes every temporary
in reverse order and preserves the original failure, attaching cleanup failures
as suppressed exceptions.
The two
[`CometVectorUtilsSuite`](https://github.com/apache/datafusion-comet/blob/e75103f8b51bab3ec440bd9ee0fcb16b0ec4e069/spark/src/test/scala/org/apache/comet/vector/CometVectorUtilsSuite.scala#L38)
tests cover a later missing dictionary and a callback `IOException`, including
temporary cleanup, source refcounts, and original exception identity.
[`dictionary stream schema and reader preserve logical values after closing the
source`](https://github.com/apache/datafusion-comet/blob/e75103f8b51bab3ec440bd9ee0fcb16b0ec4e069/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala#L415)
covers the reader/schema path. Local Spark 4.1 validation passed all 44
targeted JVM tests and 139 Python worker tests (133 general + 6
dictionary-shuffle).
##########
spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala:
##########
@@ -336,6 +394,242 @@ class CometArrowPythonRunnerSuite extends AnyFunSuite
with Matchers {
}
}
+ for (failSerialization <- Seq(false, true)) {
+ test(s"dictionary inputs materialize logical values (failure:
$failSerialization)") {
+ val sourceAllocator = new RootAllocator(Long.MaxValue)
+ val writerAllocator = new RootAllocator(Long.MaxValue)
+ val intType = new ArrowType.Int(32, true)
+ val textEncoding = new DictionaryEncoding(11L, false, intType)
+ val binaryEncoding = new DictionaryEncoding(12L, false, intType)
+ val textValues = new VarCharVector("text", sourceAllocator)
+ val binaryValues = new VarBinaryVector("data", sourceAllocator)
+ val textIndices =
+ new IntVector("text", new FieldType(true, intType, textEncoding),
sourceAllocator)
+ val binaryIndices =
+ new IntVector("data", new FieldType(true, intType, binaryEncoding),
sourceAllocator)
+ val dictionaries = Map(
+ textEncoding.getId -> new Dictionary(textValues, textEncoding),
+ binaryEncoding.getId -> new Dictionary(binaryValues, binaryEncoding))
+ val provider = new DictionaryProvider {
+ override def lookup(id: Long): Dictionary = dictionaries(id)
+
+ override def getDictionaryIds: java.util.Set[java.lang.Long] =
+ dictionaries.keys.map(id => java.lang.Long.valueOf(id)).toSet.asJava
+ }
+ val columns = Seq[CometDecodedVector](
+ new CometDictionaryVector(
+ new CometPlainVector(textIndices),
+ new CometDictionary(new CometPlainVector(textValues)),
+ provider),
+ new CometDictionaryVector(
+ new CometPlainVector(binaryIndices),
+ new CometDictionary(new CometPlainVector(binaryValues)),
+ provider))
+ var failWrites = false
+ val output = new ByteArrayOutputStream() {
+ override def write(bytes: Array[Byte], offset: Int, length: Int): Unit
= {
+ if (failWrites) {
+ throw new IOException("injected dictionary IPC write failure")
+ }
+ super.write(bytes, offset, length)
+ }
+ }
+ try {
+ textValues.allocateNew()
+ Seq("same", "", "λ中文").zipWithIndex.foreach { case (value, index) =>
+ textValues.setSafe(index, value.getBytes(StandardCharsets.UTF_8))
+ }
+ textValues.setValueCount(3)
+ binaryValues.allocateNew()
+ Seq(Array[Byte](1, 2), Array.emptyByteArray, Array[Byte](0,
-1)).zipWithIndex.foreach {
+ case (value, index) => binaryValues.setSafe(index, value)
+ }
+ binaryValues.setValueCount(3)
+ textIndices.allocateNew()
+ Seq(0, 1, 0, 2).zipWithIndex.foreach { case (value, index) =>
+ textIndices.setSafe(index, value)
+ }
+ textIndices.setNull(2)
+ textIndices.setValueCount(4)
+ binaryIndices.allocateNew()
+ Seq(2, 0, 0, 1).zipWithIndex.foreach { case (value, index) =>
+ binaryIndices.setSafe(index, value)
+ }
+ binaryIndices.setNull(2)
+ binaryIndices.setValueCount(4)
+
+ val sourceVectors = Seq(textValues, binaryValues, textIndices,
binaryIndices)
+ val sourceBuffers = sourceVectors.flatMap(_.getFieldBuffers.asScala)
+ val sourceRefs = sourceBuffers.map(_.refCnt())
+ val sourceBytes = sourceAllocator.getAllocatedMemory
+
+ def writeDictionaryBatch(): Unit =
+ withMaterializedInputVectors(columns, writerAllocator) { vectors =>
+ vectors.map(_.getField.getDictionary) shouldBe Seq(null, null)
+ vectors.head.getObject(0).toString shouldBe "same"
+ vectors.head.getObject(1).toString shouldBe ""
+ vectors.head.isNull(2) shouldBe true
+ vectors.head.getObject(3).toString shouldBe "λ中文"
+ vectors(1).getObject(0).asInstanceOf[Array[Byte]] shouldBe
Array[Byte](0, -1)
+ vectors(1).isNull(2) shouldBe true
+
+ withWriter(vectors.map(_.getField), writerAllocator,
Channels.newChannel(output)) {
+ channel =>
+ failWrites = failSerialization
+ try {
+ serializeBatch(new WriteChannel(channel), vectors, 4,
writerAllocator)
+ } finally {
+ failWrites = false
+ }
+ }
+ }
+
+ if (failSerialization) {
+ val error = intercept[IOException](writeDictionaryBatch())
+ error.getMessage shouldBe "injected dictionary IPC write failure"
+ } else {
+ writeDictionaryBatch()
+ withReader(output.toByteArray) { reader =>
+ reader.loadNextBatch() shouldBe true
+ val struct =
reader.getVectorSchemaRoot.getVector(0).asInstanceOf[StructVector]
+ val resultText = struct.getChild("text")
+ val resultData = struct.getChild("data")
+ resultText.getField.getType shouldBe ArrowType.Utf8.INSTANCE
+ resultData.getField.getType shouldBe ArrowType.Binary.INSTANCE
+ resultText.getObject(0).toString shouldBe "same"
+ resultText.getObject(1).toString shouldBe ""
+ resultText.isNull(2) shouldBe true
+ resultText.getObject(3).toString shouldBe "λ中文"
+ resultData.getObject(0).asInstanceOf[Array[Byte]] shouldBe
Array[Byte](0, -1)
+ resultData.isNull(2) shouldBe true
+ reader.loadNextBatch() shouldBe false
+ }
+ }
+
+ writerAllocator.getAllocatedMemory shouldBe 0L
+ sourceAllocator.getAllocatedMemory shouldBe sourceBytes
+ sourceBuffers.map(_.refCnt()) shouldBe sourceRefs
+ textValues.getObject(0).toString shouldBe "same"
+ binaryValues.getObject(2).asInstanceOf[Array[Byte]] shouldBe
Array[Byte](0, -1)
+ } finally {
+ columns.foreach(_.close())
+ writerAllocator.close()
+ sourceAllocator.close()
+ }
+ }
+ }
+
+ test("dictionary inputs are sliced before decoding to the Arrow batch
limits") {
Review Comment:
Added [`dictionary slices keep plain nullable and nested columns aligned
through
IPC`](https://github.com/apache/datafusion-comet/blob/e75103f8b51bab3ec440bd9ee0fcb16b0ec4e069/spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala#L768)
with ten rows and dictionary text, nullable plain bigint, plain text, and
struct columns. It verifies `[3, 3, 3, 1]` batches, every value/null, and
source refcounts. Added zero-row/all-null IPC cases and 0/-1 limit cases as
well.
[`randomized dictionary ranges match the soft-limit oracle and cover every
input
row`](https://github.com/apache/datafusion-comet/blob/e75103f8b51bab3ec440bd9ee0fcb16b0ec4e069/spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala#L914)
directly checks `inputBatchRanges` against an independent oracle over 100
seeded configurations, including complete contiguous coverage and record
limits. All passed locally on Spark 4.1. All-null strings still need
offsets/validity, so that record-limit regression disables the byte soft limit.
##########
spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala:
##########
@@ -338,6 +348,170 @@ private[python] trait CometArrowPythonRunnerBase
private[python] object CometArrowPythonRunnerBase {
+ // A regular Arrow variable-width data buffer uses signed 32-bit offsets.
The Spark setting is
+ // already restricted to this range, but cap it here as a final guard for
direct test callers.
+ private val MaxDecodedBatchBytes = Int.MaxValue.toLong
+
+ private def dictionaryVector(column: CometDictionaryVector): FieldVector = {
+ val indices = column.getValueVector
+ val encoding = indices.getField.getDictionary
+ column.getDictionaryProvider.lookup(encoding.getId).getVector
+ }
+
+ private def initialDecodedBytes(values: FieldVector): Long =
+ values match {
+ case _: BaseVariableWidthVector => BaseVariableWidthVector.OFFSET_WIDTH
+ case _: BaseLargeVariableWidthVector =>
BaseLargeVariableWidthVector.OFFSET_WIDTH
+ case _ => 0L
+ }
+
+ /** Conservative logical bytes added by one decoded dictionary value. */
+ private def decodedValueBytes(
+ column: CometDictionaryVector,
+ values: FieldVector,
+ row: Int,
+ batchRow: Int): Long = {
+ val dictionaryIndex = if (column.isNullAt(row)) -1 else
column.indices.getInt(row)
+ val validityBytes = if ((batchRow & 7) == 0) 1L else 0L
+ values match {
+ case vector: BaseVariableWidthVector =>
+ val valueBytes = if (dictionaryIndex < 0) 0L else
vector.getValueLength(dictionaryIndex)
+ valueBytes + BaseVariableWidthVector.OFFSET_WIDTH + validityBytes
+ case vector: BaseLargeVariableWidthVector =>
+ val valueBytes = if (dictionaryIndex < 0) 0L else
vector.getValueLength(dictionaryIndex)
+ valueBytes + BaseLargeVariableWidthVector.OFFSET_WIDTH + validityBytes
+ case vector: BaseFixedWidthVector =>
+ vector.getBufferSizeFor(batchRow + 1).toLong -
+ vector.getBufferSizeFor(batchRow).toLong
+ case _: NullVector => 0L
+ case vector =>
+ // Comet's JVM shuffle currently dictionary-encodes only strings and
binary values.
+ // If another Arrow type reaches this path, the complete dictionary is
a safe upper
+ // bound for any one selected value and favors smaller batches over a
large allocation.
+ math.max(1L, vector.getBufferSize.toLong)
+ }
+ }
+
+ private def saturatedAdd(left: Long, right: Long): Long =
+ if (right >= Long.MaxValue - left) Long.MaxValue else left + right
+
+ /**
+ * Split a compact dictionary batch before decoding it.
+ *
+ * The byte estimate covers the temporary logical dictionary vectors. Plain
input vectors are
+ * already allocated and remain zero-copy when no dictionary column is
present. Every returned
+ * range is applied to all columns so rows stay aligned. A single oversized
row is allowed,
+ * matching Spark's Arrow batching contract.
+ */
+ private[python] def inputBatchRanges(
+ columns: Seq[CometDecodedVector],
+ numRows: Int,
+ maxRecordsPerBatch: Int,
+ maxBytesPerBatch: Long): Seq[(Int, Int)] = {
+ require(numRows >= 0, s"Input batch row count must be non-negative:
$numRows")
+
+ val dictionaries = columns.collect { case column: CometDictionaryVector =>
+ column -> dictionaryVector(column)
+ }
+ if (numRows == 0 || dictionaries.isEmpty) {
+ return Seq(0 -> numRows)
+ }
+
+ val recordLimit =
+ if (maxRecordsPerBatch > 0) maxRecordsPerBatch else Int.MaxValue
+ val byteLimit =
+ if (maxBytesPerBatch > 0) math.min(maxBytesPerBatch,
MaxDecodedBatchBytes)
+ else MaxDecodedBatchBytes
+ val initialBytes = dictionaries.foldLeft(0L) { case (bytes, (_, values)) =>
+ saturatedAdd(bytes, initialDecodedBytes(values))
+ }
+
+ val ranges = Seq.newBuilder[(Int, Int)]
+ var start = 0
+ var row = 0
+ var decodedBytes = initialBytes
+ while (row < numRows) {
+ var rowsInBatch = row - start
+ var rowBytes = dictionaries.foldLeft(0L) { case (bytes, (column,
values)) =>
+ saturatedAdd(bytes, decodedValueBytes(column, values, row,
rowsInBatch))
+ }
+ // Spark checks the configured byte limit before adding the next row, so
the row that
+ // crosses that soft limit stays in the current batch. The separate hard
check prevents a
+ // regular variable-width buffer from crossing Arrow's signed 32-bit
allocation ceiling.
+ val exceedsArrowLimit =
+ decodedBytes >= MaxDecodedBatchBytes ||
+ rowBytes > MaxDecodedBatchBytes - decodedBytes
+ if (rowsInBatch > 0 &&
Review Comment:
Clarified [the documented
limits](https://github.com/apache/datafusion-comet/blob/e75103f8b51bab3ec440bd9ee0fcb16b0ec4e069/docs/source/user-guide/latest/pyarrow-udfs.md#L219):
the estimate counts logical decoded dictionary buffers only, the crossing row
stays in its batch, and a single oversized row stays intact. The guard prevents
combining rows beyond the signed-32-bit estimate; it cannot split an individual
row or guarantee an allocation ceiling. The docs also explain capacity rounding
and that an allocation can approach twice its logical size, plus other memory
overhead.
[`the hard Arrow ceiling splits gigabyte values with soft limits raised or
disabled`](https://github.com/apache/datafusion-comet/blob/e75103f8b51bab3ec440bd9ee0fcb16b0ec4e069/spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala#L981)
passes locally on Spark 4.1; its synthetic lengths exercise the preventive
combining guard without allocating gigabyte buffers.
##########
spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala:
##########
@@ -338,6 +348,170 @@ private[python] trait CometArrowPythonRunnerBase
private[python] object CometArrowPythonRunnerBase {
+ // A regular Arrow variable-width data buffer uses signed 32-bit offsets.
The Spark setting is
+ // already restricted to this range, but cap it here as a final guard for
direct test callers.
+ private val MaxDecodedBatchBytes = Int.MaxValue.toLong
+
+ private def dictionaryVector(column: CometDictionaryVector): FieldVector = {
+ val indices = column.getValueVector
+ val encoding = indices.getField.getDictionary
+ column.getDictionaryProvider.lookup(encoding.getId).getVector
+ }
+
+ private def initialDecodedBytes(values: FieldVector): Long =
+ values match {
+ case _: BaseVariableWidthVector => BaseVariableWidthVector.OFFSET_WIDTH
+ case _: BaseLargeVariableWidthVector =>
BaseLargeVariableWidthVector.OFFSET_WIDTH
+ case _ => 0L
+ }
+
+ /** Conservative logical bytes added by one decoded dictionary value. */
+ private def decodedValueBytes(
+ column: CometDictionaryVector,
+ values: FieldVector,
+ row: Int,
+ batchRow: Int): Long = {
+ val dictionaryIndex = if (column.isNullAt(row)) -1 else
column.indices.getInt(row)
+ val validityBytes = if ((batchRow & 7) == 0) 1L else 0L
+ values match {
+ case vector: BaseVariableWidthVector =>
+ val valueBytes = if (dictionaryIndex < 0) 0L else
vector.getValueLength(dictionaryIndex)
+ valueBytes + BaseVariableWidthVector.OFFSET_WIDTH + validityBytes
+ case vector: BaseLargeVariableWidthVector =>
+ val valueBytes = if (dictionaryIndex < 0) 0L else
vector.getValueLength(dictionaryIndex)
+ valueBytes + BaseLargeVariableWidthVector.OFFSET_WIDTH + validityBytes
+ case vector: BaseFixedWidthVector =>
+ vector.getBufferSizeFor(batchRow + 1).toLong -
+ vector.getBufferSizeFor(batchRow).toLong
+ case _: NullVector => 0L
+ case vector =>
+ // Comet's JVM shuffle currently dictionary-encodes only strings and
binary values.
+ // If another Arrow type reaches this path, the complete dictionary is
a safe upper
+ // bound for any one selected value and favors smaller batches over a
large allocation.
+ math.max(1L, vector.getBufferSize.toLong)
+ }
+ }
+
+ private def saturatedAdd(left: Long, right: Long): Long =
+ if (right >= Long.MaxValue - left) Long.MaxValue else left + right
+
+ /**
+ * Split a compact dictionary batch before decoding it.
+ *
+ * The byte estimate covers the temporary logical dictionary vectors. Plain
input vectors are
+ * already allocated and remain zero-copy when no dictionary column is
present. Every returned
+ * range is applied to all columns so rows stay aligned. A single oversized
row is allowed,
+ * matching Spark's Arrow batching contract.
+ */
+ private[python] def inputBatchRanges(
Review Comment:
Implemented the dictionary-size upper bound and the while-loop fallback
described in the earlier performance thread. Selected value lengths are read
once per row; only bitmap cost is adjusted when that row starts a new range.
[`a conservative dictionary bound accepts small batches without reading row
indices`](https://github.com/apache/datafusion-comet/blob/e75103f8b51bab3ec440bd9ee0fcb16b0ec4e069/spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala#L947)
passes locally on Spark 4.1 and verifies zero row-index reads when the upper
bound accepts the batch, plus one read per row in the forced-split path.
--
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]
