viirya commented on code in PR #5560:
URL: https://github.com/apache/datafusion-comet/pull/5560#discussion_r3919493690
##########
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:
I verified this reachability analysis independently and it holds, with one
addition that I think argues for the cheaper of your two options.
`row.rs:1317,1329` explicitly disables dictionary encoding for array
elements and struct fields, so JVM shuffle can't produce it. I also checked the
stacked-`CometMapInBatchExec` route, since that feeds one runner's flattened
output into another: the output side calls `CometVector.getVector(vector,
null)` (`CometArrowPythonRunnerBase.scala:274`), and `getVector` does
`dictionaryProvider.lookup(...)` for any dictionary-encoded vector
(`CometVector.java:249`), so it would NPE there first — a pre-existing issue
outside this PR, but it means stacking can't quietly introduce a nested
dictionary either.
Since it's genuinely unreachable today, I'd favour the explicit check naming
the column over full recursion: recursion adds untestable complexity for an
unreachable path, while a named check converts a future FFI-path failure from
an NPE inside Arrow into something diagnosable.
--
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]
