andygrove commented on code in PR #5560:
URL: https://github.com/apache/datafusion-comet/pull/5560#discussion_r3916497098
##########
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:
`inputBatchRanges` walks every row before any decoding happens, and with
Comet's defaults it can never split: `spark.comet.batchSize` and
`spark.comet.shuffle.jvm.batchSize` are both 8192, under
`spark.sql.execution.arrow.maxRecordsPerBatch` (10000), and a decoded 8192-row
batch is nowhere near `spark.sql.execution.arrow.maxBytesPerBatch` (64MB). I
measured it on an 8192-row batch with one dictionary column (64 distinct
~28-byte values): the scan takes 227-320 us against 630-731 us for the whole
ranges+decode+serialize path, so 34-43% of the work is thrown away.
Two things would help. First, a cheap upper bound that skips the scan
entirely when the batch provably fits: one pass over the *dictionary* for
`max(getValueLength)` is O(distinct values) rather than O(rows), and
`initialBytes + numRows * (maxLen + OFFSET_WIDTH + 1)` bounds the decoded size
from above. If that is under `byteLimit` and `numRows <= recordLimit`, return
`Seq(0 -> numRows)` without touching the indices. That took the same batch from
320 us to 1.0 us for me, and I checked it against the exact scan over 300
randomized configs (row counts, dictionary sizes, null rates, both limits) with
no disagreement on the 116 where it fired.
Second, the fallback scan itself: `dictionaries.foldLeft(0L)` boxes the
accumulator and destructures a tuple once per row per column, and the split
row's bytes are computed twice. A `while` loop over parallel arrays took it
from 320 us to 99 us on the same input.
##########
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:
`columns.collect { case column: CometDictionaryVector => ... }` only sees
top-level columns, so a struct, list or map whose child is dictionary-encoded
goes through the `case column =>` arm in `withMaterializedInputVectors`
unchanged and reaches `startWriter` with a `DictionaryEncoding` on the child
field and a null provider. I built a `Struct<child: Dictionary<Int32, Utf8>>`
and it fails with the same NPE this PR fixes:
```
java.lang.NullPointerException: Cannot invoke
"org.apache.arrow.vector.dictionary.DictionaryProvider.lookup(long)"
because "provider" is null
```
I do not think it is reachable today. `builder_to_array` in
`native/shuffle/src/spark_unsafe/row.rs` dictionary-encodes only top-level
`Utf8` and `Binary`, so the JVM shuffle cannot produce it. But
`CometStructVector` builds its children through `CometVector.getVector`, which
does create nested `CometDictionaryVector`s, so an FFI-imported batch could.
Would you either recurse into children or add an explicit check with a message
that names the column, so this surfaces as something diagnosable rather than an
NPE from inside Arrow?
##########
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:
With this early return the worker's batch boundaries depend on whether the
shuffle chose dictionary encoding for that column, which comes down to
`spark.comet.shuffle.jvm.preferDictionary.ratio` and the data's cardinality.
Vanilla Spark applies both limits to every `mapInArrow` / `mapInPandas` input
(`BatchedPythonArrowInput.writeSizedBatch`), and Comet's plain path already
exceeds `maxRecordsPerBatch` whenever `spark.comet.batchSize` is set above
10000.
I read the PR description and I see this is deliberate, and the plain path
serializing existing buffers is a real argument for not slicing it. Would it be
worth at least applying the record limit uniformly, since that one costs
nothing to check, and saying in `pyarrow-udfs.md` that the byte limit is
bounded only for dictionary inputs? Right now the doc says Comet "uses Spark's
Arrow record threshold and the decoded dictionary size against Spark's byte
threshold to split the compact batch", which reads as though both thresholds
are always honoured.
##########
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:
`byteLimit` is clamped to `MaxDecodedBatchBytes` a few lines above, so
`decodedBytes >= MaxDecodedBatchBytes` is strictly implied by the `decodedBytes
>= byteLimit` disjunct right next to it, and `rowBytes > MaxDecodedBatchBytes -
decodedBytes` can only fire if someone sets
`spark.sql.execution.arrow.maxBytesPerBatch` to within one row of 2GB, where
Arrow's own `OversizedAllocationException` gets there first. `saturatedAdd` is
similar: `decodedBytes` never exceeds `byteLimit` plus one row's worth, so a
`Long` cannot overflow.
Dropping `exceedsArrowLimit` and `saturatedAdd` and keeping just the clamp
would make the loop condition read as the one rule it actually implements, and
this is also the hot loop from my other comment.
##########
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:
`ColumnarBatchArrowReader.loadNextBatch` has the same block: match
`CometDictionaryVector`, look up the dictionary through the provider,
`DictionaryEncoder.decode` into the caller's allocator, close the temporaries
in a `finally`. The two have already drifted a little (`d.provider` there vs
`d.getDictionaryProvider` here, and the reader swallows exceptions from
`close()` while this one propagates them), and
`CometNativeArrowSource.actualFieldOf` has a third copy of the lookup half.
Would a shared helper next to `CometVector` be worth it, so the next person who
touches dictionary materialization only has one place to look?
##########
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:
Both new slicing tests use a column set made only of dictionary columns, so
the part of the change that I would most want pinned down is untested: that
every column is sliced at the same boundaries. Plain, nested and dictionary
columns go through three different `slice` implementations, and nothing here
would fail if one of them stopped being sliced.
I added a case locally with 10 rows, `maxRecordsPerBatch=3`, and columns
`[dictionary VarChar, plain BigInt with a null, plain VarChar, Struct<k:
bigint>]`. It passes on this head, batches come back `[3,3,3,1]` with every
column on the right row, so this is regression coverage rather than a bug.
Would you add it?
Three smaller ones in the same spirit, all passing today: an all-null
dictionary column (with every index null `decodedValueBytes` returns 0 for
every row, so the byte limit can never fire and only the record limit splits),
a zero-row dictionary batch, and `maxRecordsPerBatch` / `maxBytesPerBatch` at 0
and -1. A randomized property over `inputBatchRanges` asserting the ranges are
contiguous, start at 0, sum to `numRows`, and never exceed the record limit is
cheap and covers a lot of future off-by-one ground. I ran 200 configs and it
held.
--
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]
