viirya commented on code in PR #5560:
URL: https://github.com/apache/datafusion-comet/pull/5560#discussion_r3919496667
##########
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:
Agreed with @andygrove, and I'd rank this the most valuable follow-up in the
PR. The central guarantee of the change is that *every* column is sliced at the
same boundaries, but plain / nested / dictionary go through three different
`slice` implementations (`CometPlainVector:213`, `CometStructVector:62`,
`CometDictionaryVector:137`) and both slicing tests use all-dictionary column
sets — so if one of those stopped being sliced, nothing here goes red.
Separately: `inputBatchRanges` is deliberately `private[python]` but has
**no direct unit test** — it's only reached through `foreachInputBatch`, and
it's the subtlest arithmetic in the change. The randomized property he
describes (contiguous, starts at 0, sums to `numRows`, never exceeds the record
limit) is worth pointing straight at it.
##########
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:
Independent of the `exceedsArrowLimit` discussion: `rowsInBatch > 0` means a
single row is never split, which is the right Spark contract, but it also means
the 2GiB protection is one order weaker than the PR description claims ("adds a
hard guard for regular Arrow variable-width buffers"). A single logical value
above 2GiB still reaches `DictionaryEncoder.decode` and overflows. Extreme
under Spark's string limits, but worth one sentence in the comment or the doc
so the boundary is stated rather than left to be inferred.
Relatedly, `decodedValueBytes` is a *logical* estimate — offsets plus value
plus amortized validity bit. I checked the two branches are mutually consistent
(the `getBufferSizeFor(batchRow + 1) - getBufferSizeFor(batchRow)` delta
amortizes the validity byte the same way the var-width branch does), and it
errs conservative, which is the right direction. But Arrow rounds real
allocations up to powers of two, so actual memory can approach 2x the estimate
and `maxBytesPerBatch` isn't an actual memory ceiling for users. Worth noting
in `pyarrow-udfs.md`.
##########
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:
Seconding @andygrove's performance point, and I'd treat it as more than a
nice-to-have. With Comet's defaults this scan can *never* split —
`spark.comet.batchSize` and `spark.comet.shuffle.jvm.batchSize` are both 8192,
below `maxRecordsPerBatch` (10000), and an 8192-row batch is far under
`maxBytesPerBatch` (64MB on Spark 4.1, 256MB on 4.0). So the common case walks
every row and discards the result. The O(distinct) upper bound he describes is
a sound conservative estimate and short-circuits exactly that case.
Also worth noting the `rowBytes` fold appears twice (once before the split
decision, once recomputed after), so the split row's bytes are computed twice.
Rewriting as a `while` loop over parallel arrays addresses the cost and the
duplication together.
##########
spark/src/main/scala/org/apache/comet/CometConf.scala:
##########
@@ -596,9 +596,10 @@ object CometConf extends ShimCometConf {
.withAlternative("spark.comet.shuffle.preferDictionary.ratio")
.category(CATEGORY_SHUFFLE)
.doc(
- "The ratio of total values to distinct values in a string column to
decide whether to " +
+ "The ratio of total values to distinct values in a string or binary
column to decide " +
+ "whether to " +
Review Comment:
Nit: the reflow left `"whether to " +` as an orphan line, which reads oddly.
Rebreaking the string manually would be cleaner.
The wording change itself is correct —
`native/shuffle/src/spark_unsafe/row.rs:1451,1470` confirms both `Utf8` and
`Binary` are dictionary-encoded.
--
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]
