You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "sunchao (via GitHub)" <gi...@apache.org> on 2026/08/30 18:24:29 UTC
[PR] fix: decode dictionary input for PyArrow UDFs [datafusion-comet]
sunchao opened a new pull request, #5560:
URL: https://github.com/apache/datafusion-comet/pull/5560
## Which issue does this PR close?
Extracted from #5557 while addressing #5555.
## Rationale for this change
A JVM Comet shuffle can encode repeated string and binary columns as dictionaries. The accelerated `mapInArrow` / `mapInPandas` runner currently takes the dictionary indices vector, advertises its dictionary metadata, and starts `ArrowStreamWriter` without a dictionary provider. Arrow then fails before the Python worker can receive the batch. This affects the existing accelerated path with `spark.sql.execution.arrow.useLargeVarTypes=false`, so it is an independent crash fix rather than part of large-offset support.
For example, a Parquet input with repeated string and binary values followed by `repartition(2, "id")`, using `spark.comet.shuffle.mode=jvm` and `spark.comet.shuffle.jvm.preferDictionary.ratio=1.01`, produces `CometDictionaryVector` input and fails on current `main`.
## What changes are included in this PR?
The runner now identifies top-level `CometDictionaryVector` columns and decodes their logical values into temporary vectors owned by the runner allocator. It derives the stream schema and performs the synchronous IPC write from those decoded vectors, then closes them on both success and failure. Plain Comet vectors keep the existing borrowed-buffer path and ownership model.
The regression coverage includes string and binary dictionaries, repeated, empty, Unicode, and null values, injected write failure, buffer reference counts, and allocator cleanup. A real-worker test exercises both `mapInArrow` and `mapInPandas` after an actual JVM Comet shuffle. The PyArrow workflow runs that module separately on each supported Spark worker so its shuffle-manager Spark session cannot inherit the main module's shuffle-disabled session.
The benchmark gains a low-cardinality JVM-shuffle workload and an optional `BENCHMARK_WORKLOAD` selector. The workload requires `CometColumnarExchange` in both modes and compares vanilla and accelerated Python execution on the same shuffled input.
## How are these changes tested?
Native code was built before the JVM checks; this PR does not change native code.
- Spark 4.0 / Scala 2.13 root-reactor package build: **11/11 focused JVM tests passed**, including dictionary success and injected-failure ownership cases.
- PySpark 4.0.4 / JDK 17 / PyArrow 25.0.1 / pandas 3.0.5: the existing **117 tests passed**, and the new dictionary-shuffle module passed **2/2 tests** in the separate process used by CI.
- Benchmark smoke: the new workload completed all four API/mode cases with 4,096 rows.
- A local `local[2]` debug-build timing with 2,000,000 rows, two warmups, and five measured iterations produced median end-to-end times of **2.453 s vanilla / 1.794 s accelerated (1.37x)** for `mapInArrow`, and **2.415 s / 1.797 s (1.34x)** for `mapInPandas`. These numbers show that the required dictionary materialization still preserves an end-to-end gain in this setup; they are not a general performance claim.
- Scala formatting/style, Python formatting/syntax, workflow YAML parsing, and `git diff --check` passed.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: decode dictionary input for PyArrow UDFs [datafusion-comet]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on code in PR #5560:
URL: https://github.com/apache/datafusion-comet/pull/5560#discussion_r3993592945
##########
spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala:
##########
@@ -178,44 +182,49 @@ private[python] trait CometArrowPythonRunnerBase
val cometBatch = currentGroup.next()
val startData = dataOut.size()
- val sourceVectors = (0 until cometBatch.numCols()).map { i =>
- cometBatch
- .column(i)
- .asInstanceOf[CometDecodedVector]
- .getValueVector
- .asInstanceOf[FieldVector]
+ val columns = (0 until cometBatch.numCols()).map { i =>
+ cometBatch.column(i).asInstanceOf[CometDecodedVector]
}
- val batchFields = sourceVectors.map(_.getField)
-
- if (arrowWriter == null) {
- // Build the schema-only struct root once from the first batch's child fields.
- // mapInArrow/mapInPandas exchange the columns under a single non-nullable struct.
- // Comet's FFI-imported vectors leave the Arrow Field name null, so restore the real
- // column names from the input schema (the worker reads columns by name, and shaded
- // Arrow rejects a null field name). Keep the field types and child structure as-is so
- // the advertised schema matches the source buffers. Keeping the type as-is also means
- // a TimestampType reaches the worker with Comet's UTC time zone
- // rather than the session zone vanilla Spark would label it with; this is a documented
- // limitation (see pyarrow-udfs.md), not a value difference, since the stored instant is
- // identical.
- val childNames = inputStructType.fieldNames
- streamFields = batchFields.zipWithIndex.map { case (field, i) =>
- renamed(field, childNames(i), forceNullable = true)
+ CometArrowPythonRunnerBase.foreachInputBatch(
+ columns,
+ cometBatch.numRows(),
+ arrowMaxRecordsPerBatch,
+ arrowMaxBytesPerBatch,
+ allocator) { (sourceVectors, numRows) =>
Review Comment:
Fixed in [279bbe303](https://github.com/apache/datafusion-comet/commit/279bbe303c06c3833a0a899f542847d37ead5321). The writer now retains the current source batch and remaining ranges, emits one slice per `writeNextInputToStream` call, and returns control to Spark's transport loop. It does not call either upstream iterator's `hasNext`/`next` until the final slice finishes, since `CometExecIterator.hasNext` can close the previous batch and reuse its buffers. Temporary slices/decoded vectors remain scoped to each write; upstream retains source ownership on completion, interruption, and failure.
Added regressions using the production writer and Spark's actual `DirectByteBufferOutputStream`, with Spark's threshold-based fill/drain loop. They cover both large slices and multiple small slices per drain, round-trip the complete IPC stream, check metrics and source lifetime, and exercise interruption and serialization failure with pending slices. Against the original runner, the two transport tests fail with 16,832,728 and 139,416 pending bytes; both pass with this fix. The tests assert pending-byte bounds directly rather than relying on `MaxDirectMemorySize` enforcement.
All 48 focused JVM tests pass on Spark 4.1.3. All 133 general PyArrow tests and 6 dictionary-shuffle tests also pass against the rebuilt packaged JAR with real Python workers. Hosted CI for the new head is pending.
##########
.github/workflows/pyarrow_udf_test.yml:
##########
@@ -28,23 +28,29 @@ on:
paths: &feature-paths
- "pom.xml"
- "common/pom.xml"
- - "common/src/main/scala/org/apache/comet/CometConf.scala"
+ - "native/shuffle/src/spark_unsafe/row.rs"
- "spark/pom.xml"
+ - "spark/src/main/java/org/apache/comet/vector/**"
+ - "spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/SpillWriter.java"
+ - "spark/src/main/scala/org/apache/comet/CometConf.scala"
- "spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala"
+ - "spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala"
+ - "spark/src/main/scala/org/apache/comet/vector/StreamReader.scala"
Review Comment:
Fixed in [279bbe303](https://github.com/apache/datafusion-comet/commit/279bbe303c06c3833a0a899f542847d37ead5321). Replaced the two individual Scala vector paths with `spark/src/main/scala/org/apache/comet/vector/**` in the shared path-filter anchor, so helper-only changes trigger both push and PR runs. Confirmed that the filter matches `CometVectorUtils.scala` as well as the previously covered files; Actionlint and YAML formatting pass.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: decode dictionary input for PyArrow UDFs [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: decode dictionary input for PyArrow UDFs [datafusion-comet]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #5560:
URL: https://github.com/apache/datafusion-comet/pull/5560#discussion_r3994139506
##########
spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala:
##########
@@ -178,44 +182,49 @@ private[python] trait CometArrowPythonRunnerBase
val cometBatch = currentGroup.next()
val startData = dataOut.size()
- val sourceVectors = (0 until cometBatch.numCols()).map { i =>
- cometBatch
- .column(i)
- .asInstanceOf[CometDecodedVector]
- .getValueVector
- .asInstanceOf[FieldVector]
+ val columns = (0 until cometBatch.numCols()).map { i =>
+ cometBatch.column(i).asInstanceOf[CometDecodedVector]
}
- val batchFields = sourceVectors.map(_.getField)
-
- if (arrowWriter == null) {
- // Build the schema-only struct root once from the first batch's child fields.
- // mapInArrow/mapInPandas exchange the columns under a single non-nullable struct.
- // Comet's FFI-imported vectors leave the Arrow Field name null, so restore the real
- // column names from the input schema (the worker reads columns by name, and shaded
- // Arrow rejects a null field name). Keep the field types and child structure as-is so
- // the advertised schema matches the source buffers. Keeping the type as-is also means
- // a TimestampType reaches the worker with Comet's UTC time zone
- // rather than the session zone vanilla Spark would label it with; this is a documented
- // limitation (see pyarrow-udfs.md), not a value difference, since the stored instant is
- // identical.
- val childNames = inputStructType.fieldNames
- streamFields = batchFields.zipWithIndex.map { case (field, i) =>
- renamed(field, childNames(i), forceNullable = true)
+ CometArrowPythonRunnerBase.foreachInputBatch(
+ columns,
+ cometBatch.numRows(),
+ arrowMaxRecordsPerBatch,
+ arrowMaxBytesPerBatch,
+ allocator) { (sourceVectors, numRows) =>
Review Comment:
Confirmed fixed in `279bbe303`. Returning after each slice gives Spark's transport loop an opportunity to drain, and retaining the source without calling upstream `hasNext` or `next` preserves its buffers until the final slice is serialized.
I ran both new transport regressions successfully. Against the previous runner, they fail with 16,832,728 and 139,416 pending bytes, respectively. The interruption and serialization-failure tests also pass. This addresses my blocker.
##########
.github/workflows/pyarrow_udf_test.yml:
##########
@@ -28,23 +28,29 @@ on:
paths: &feature-paths
- "pom.xml"
- "common/pom.xml"
- - "common/src/main/scala/org/apache/comet/CometConf.scala"
+ - "native/shuffle/src/spark_unsafe/row.rs"
- "spark/pom.xml"
+ - "spark/src/main/java/org/apache/comet/vector/**"
+ - "spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/SpillWriter.java"
+ - "spark/src/main/scala/org/apache/comet/CometConf.scala"
- "spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala"
+ - "spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala"
+ - "spark/src/main/scala/org/apache/comet/vector/StreamReader.scala"
Review Comment:
Confirmed. The Scala vector directory pattern now includes `CometVectorUtils.scala` and preserves coverage of the previously listed files. Since both push and pull-request triggers use the shared anchor, helper-only changes are covered.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: decode dictionary input for PyArrow UDFs [datafusion-comet]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on code in PR #5560:
URL: https://github.com/apache/datafusion-comet/pull/5560#discussion_r3991507909
##########
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 both requested regressions: [`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) verifies all four columns and nulls across `[3, 3, 3, 1]` batches, and [`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) calls `inputBatchRanges` directly over 100 seeded configurations. It checks exact agreement with the soft-limit oracle, contiguous complete coverage, and record limits. Both passed locally on Spark 4.1.
##########
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:
Reflowed [the description](https://github.com/apache/datafusion-comet/blob/e75103f8b51bab3ec440bd9ee0fcb16b0ec4e069/spark/src/main/scala/org/apache/comet/CometConf.scala#L597) so “whether to prefer dictionary encoding” stays together. The concatenated documentation string is unchanged, including “string or binary”; verified the string equality and formatting checks.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: decode dictionary input for PyArrow UDFs [datafusion-comet]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on PR #5560:
URL: https://github.com/apache/datafusion-comet/pull/5560#issuecomment-5655788733
Thanks @peterxcli, this looks like a useful follow-up. Keeping dictionaries in IPC is feasible with the existing direct `MessageSerializer` path and could reduce transport bytes and JVM decoding allocations.
For this PR, I'd keep the JVM-side decoding so the fix works with existing Spark Python workers. Moving decoding to Python would also require worker integration to split before decoding and preserve vanilla Spark's input types before user code or pandas conversion, plus dictionary ID/update handling across batches. It moves the bounded-decoding work rather than eliminating it.
I opened [#5906](https://github.com/apache/datafusion-comet/issues/5906) to track that improvement, including compatibility and lifetime tests and benchmarks of transport volume, JVM/Python peak memory, and throughput against this implementation.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: decode dictionary input for PyArrow UDFs [datafusion-comet]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on code in PR #5560:
URL: https://github.com/apache/datafusion-comet/pull/5560#discussion_r3991503568
##########
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:
Added the dictionary-entry maximum bound, which returns the whole batch when it provably fits both limits without reading row indices. The fallback now uses while loops over parallel arrays and computes each selected value length once; a split only adjusts bitmap accounting.
[`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: it asserts zero index reads for the accepted batch and exactly one per row when forced to split. The seeded range-oracle test also passes.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: decode dictionary input for PyArrow UDFs [datafusion-comet]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: decode dictionary input for PyArrow UDFs [datafusion-comet]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: decode dictionary input for PyArrow UDFs [datafusion-comet]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: decode dictionary input for PyArrow UDFs [datafusion-comet]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #5560:
URL: https://github.com/apache/datafusion-comet/pull/5560#discussion_r3991579134
##########
.github/workflows/pyarrow_udf_test.yml:
##########
@@ -28,23 +28,29 @@ on:
paths: &feature-paths
- "pom.xml"
- "common/pom.xml"
- - "common/src/main/scala/org/apache/comet/CometConf.scala"
+ - "native/shuffle/src/spark_unsafe/row.rs"
- "spark/pom.xml"
+ - "spark/src/main/java/org/apache/comet/vector/**"
+ - "spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/SpillWriter.java"
+ - "spark/src/main/scala/org/apache/comet/CometConf.scala"
- "spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala"
+ - "spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala"
+ - "spark/src/main/scala/org/apache/comet/vector/StreamReader.scala"
Review Comment:
**[P3] Include the new shared decoding helper in the workflow paths**
Now that the runner delegates decoding and cleanup to `spark/src/main/scala/org/apache/comet/vector/CometVectorUtils.scala`, could we add that file to this list, or cover the Scala vector directory?
The Java vector glob above does not match it, and the Scala entries currently list only `NativeUtil.scala` and `StreamReader.scala`. A future change confined to the shared decoding helper would therefore skip these Python worker regressions. This PR itself still triggers them through its other changed files.
##########
spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala:
##########
@@ -178,44 +182,49 @@ private[python] trait CometArrowPythonRunnerBase
val cometBatch = currentGroup.next()
val startData = dataOut.size()
- val sourceVectors = (0 until cometBatch.numCols()).map { i =>
- cometBatch
- .column(i)
- .asInstanceOf[CometDecodedVector]
- .getValueVector
- .asInstanceOf[FieldVector]
+ val columns = (0 until cometBatch.numCols()).map { i =>
+ cometBatch.column(i).asInstanceOf[CometDecodedVector]
}
- val batchFields = sourceVectors.map(_.getField)
-
- if (arrowWriter == null) {
- // Build the schema-only struct root once from the first batch's child fields.
- // mapInArrow/mapInPandas exchange the columns under a single non-nullable struct.
- // Comet's FFI-imported vectors leave the Arrow Field name null, so restore the real
- // column names from the input schema (the worker reads columns by name, and shaded
- // Arrow rejects a null field name). Keep the field types and child structure as-is so
- // the advertised schema matches the source buffers. Keeping the type as-is also means
- // a TimestampType reaches the worker with Comet's UTC time zone
- // rather than the session zone vanilla Spark would label it with; this is a documented
- // limitation (see pyarrow-udfs.md), not a value difference, since the stored instant is
- // identical.
- val childNames = inputStructType.fieldNames
- streamFields = batchFields.zipWithIndex.map { case (field, i) =>
- renamed(field, childNames(i), forceNullable = true)
+ CometArrowPythonRunnerBase.foreachInputBatch(
+ columns,
+ cometBatch.numRows(),
+ arrowMaxRecordsPerBatch,
+ arrowMaxBytesPerBatch,
+ allocator) { (sourceVectors, numRows) =>
Review Comment:
**[P1] Return to Spark between slices so the transport buffer can drain**
`foreachInputBatch` serializes every slice before this `writeNextInputToStream` call returns. In Spark 4.1, `dataOut` wraps a `DirectByteBufferOutputStream`; `ReaderInputStream.writeAdditionalInputToPythonWorker` only writes that buffer to the worker socket after this call returns ([Spark source](https://github.com/apache/spark/blob/v4.1.3/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala#L861)). Consequently, closing each decoded vector does not release the accumulated IPC bytes.
I reproduced this with 256 rows referencing one 64 KiB dictionary value and `maxRecordsPerBatch=1`. The source buffers occupied 83,968 bytes and the decoding allocator peaked at 114,688 bytes, but the transport buffer accumulated 16,832,512 bytes. With `-XX:MaxDirectMemorySize=8m`, this failed in `DirectByteBufferOutputStream.grow`; a control that drained/reset the buffer between slices passed, with at most 65,752 pending IPC bytes.
Could we retain the current source batch and remaining ranges in the writer, emit one slice per `writeNextInputToStream` call, and return so Spark can drain the transport buffer? The source must remain alive until all its slices have been consumed. Calling `flush()` on this stream would not drain it to the socket.
Please also add a regression covering transport buffering. The existing allocator-capped test proves that decoding is bounded, but does not detect accumulation outside the Arrow allocator.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: decode dictionary input for PyArrow UDFs [datafusion-comet]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: decode dictionary input for PyArrow UDFs [datafusion-comet]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao merged PR #5560:
URL: https://github.com/apache/datafusion-comet/pull/5560
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org