viirya commented on code in PR #5560:
URL: https://github.com/apache/datafusion-comet/pull/5560#discussion_r3919489116
##########
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:
Partly disagreeing with this one. I think the three pieces should be split
rather than removed together.
**The first disjunct is dead, as you say.** `byteLimit = min(cfg,
MaxDecodedBatchBytes) <= MaxDecodedBatchBytes`, so `decodedBytes >=
MaxDecodedBatchBytes` implies `decodedBytes >= byteLimit` sitting right next to
it in the same `OR`. I removed only that disjunct and swept 4000 randomized
legal configs (row counts, per-value sizes up to 2e9, both limits): **0
disagreements**.
**The second disjunct is not dead.** The soft-limit semantics are what make
it reachable: `decodedBytes >= byteLimit` is checked *before* the row is added,
so the crossing row stays in and `decodedBytes` can reach `byteLimit - 1 +
maxSingleRowBytes`. Concretely, with `byteLimit == 2147483647`, `decodedBytes =
1073741828` and `rowBytes = 1073741824`: `decodedBytes >= byteLimit` is `false`
while `rowBytes > MaxDecodedBatchBytes - decodedBytes` is `1073741824 >
1073741819` → `true`. Four 1GiB rows give `[(0,1),(1,1),(2,1),(3,1)]` with the
guard and `[(0,2),(2,2)]` without it — a 2.0 GiB first batch. Over the same
4000 configs, dropping this disjunct changed **1552** of them, with a worst
observed batch of **3.68 GiB**.
Where the estimate of the trigger condition needs widening: I swept the
threshold and at `byteLimit` = 64MiB / 256MiB / 512MiB / 1GiB the guard and
no-guard results are **identical** (the soft limit splits each oversized row
off anyway). Divergence starts above roughly 1GiB. So it isn't "within one row
of 2GB" — it's the top half of the legal range.
`spark.sql.execution.arrow.maxBytesPerBatch` is `checkValue(x => x > 0 && x <=
Int.MaxValue)` (`SQLConf.scala:4040`), so anything up to 2147483647 is
accepted; the default is 256MB on Spark 4.0 and 64MB on 4.1.
I'd also not rely on Arrow throwing first: a 2GiB `Utf8` batch overflows
32-bit offsets, which is silent corruption rather than a reliable
`OversizedAllocationException`.
**`saturatedAdd` I agree is removable.** `getValueLength` and
`getBufferSize` both return `int`, so overflowing a `Long` would need ~4.3
billion dictionary columns in one batch.
Suggested resolution: drop the first disjunct and `saturatedAdd`, **keep**
`rowBytes > MaxDecodedBatchBytes - decodedBytes`, and add a comment naming the
reachable case (`maxBytesPerBatch` above ~1GiB with multi-hundred-MiB
dictionary values) plus a test pinning it. Nothing in the suite currently
exercises this branch, which is probably why it reads as dead.
--
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]
