andygrove commented on code in PR #5557:
URL: https://github.com/apache/datafusion-comet/pull/5557#discussion_r3889863626
##########
spark/src/test/resources/pyspark/benchmark_pyarrow_udf.py:
##########
@@ -84,6 +86,10 @@ def _build_spark() -> SparkSession:
.config("spark.plugins", "org.apache.spark.CometPlugin")
.config("spark.comet.enabled", "true")
.config("spark.comet.exec.enabled", "true")
+ .config(
+ "spark.sql.execution.arrow.useLargeVarTypes",
+ os.environ.get("BENCHMARK_LARGE_VAR_TYPES", "false"),
+ )
Review Comment:
`BENCHMARK_LARGE_VAR_TYPES` covers the widening path nicely. Could the
dictionary path get a number too? It looks like the more expensive of the two
new costs, and unlike widening it applies in the default
`useLargeVarTypes=false` mode as well.
On a single 8192-row string column with 64 distinct 32-byte values I
measured the decode at roughly 350 us per batch against roughly 13 us for the
whole plain serialization step, so about 25x to 30x. I also tried replacing
`DictionaryEncoder.decode` with a direct offset-driven copy to see whether
Arrow's per-row `getObject` boxing was the problem, and it came out slightly
slower (398 us versus 356 us). So the cost looks inherent to materializing the
values rather than something to micro-optimize, which seems worth knowing
before anyone tries.
The gap I would most like closed is that nothing in the repo would catch a
regression here. The three `WORKLOADS` all read straight from Parquet with no
shuffle in the plan, so none of them ever sees a `CometDictionaryVector`. Would
you add a fourth that repartitions on a low-cardinality string column with
`spark.comet.shuffle.mode=jvm` and `preferDictionary.ratio` set, and quote the
accelerated-versus-vanilla number for it? If the accelerated path still wins
end to end for dictionary input, that is the number that justifies removing the
fallback, and it belongs in the description next to the widening table. Note
`_mixed_with_strings` builds `concat('row_', id)`, which is all-distinct, so it
would need a low-cardinality column to be a useful dictionary case even with a
shuffle added.
##########
spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala:
##########
@@ -379,17 +449,27 @@ private[python] object CometArrowPythonRunnerBase {
buffers.add(structValidity)
buffers.addAll(sourceBatch.getBuffers)
- val wrappedBatch = new ArrowRecordBatch(
- numRows,
- nodes,
- buffers,
- sourceBatch.getBodyCompression,
- sourceBatch.getVariadicBufferCounts,
- true)
+ val widenedOffsets = if (useLargeVarTypes) new ArrayList[ArrowBuf]()
else null
Review Comment:
Small one: `widenedOffsets` is `null` when `useLargeVarTypes` is false, so
the same condition gets tested twice in two different ways, once as `if
(useLargeVarTypes)` to call `widenOffsets` and once as `if (widenedOffsets !=
null)` in the `finally`. Always allocating the list would drop the sentinel and
the null check for the cost of one empty `ArrayList` per batch, on a path that
already allocates an offset buffer per variable-width column when the flag is
on.
Related, the `useLargeVarTypes: Boolean = false` default on `serializeBatch`
is only reached from tests, since production always passes it explicitly.
##########
spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala:
##########
@@ -265,6 +285,317 @@ class CometArrowPythonRunnerSuite extends AnyFunSuite
with Matchers {
}
}
+ for {
+ failSerialization <- Seq(false, true)
+ useLargeVarTypes <- Seq(false, true)
+ } {
+ test(
+ "dictionary inputs materialize logical values " +
Review Comment:
Two coverage notes on this suite. Both cases pass for me, so this is about
where the guard lives rather than a bug.
These dictionary tests build their vectors from a plain `RootAllocator`, so
the reference manager under test is not the `ReferenceCountedArrowArray` that
real shuffle input carries. That is the same point that came up on #5368 for
the plain path, which you then covered with the `Data.exportVector` cases. I
drove real Comet JVM-shuffle batches through it (`shuffle.mode=jvm`,
`preferDictionary.ratio=1.01`, `repartition(2, "id")`), saw four genuine
`CometDictionaryVector` columns, and the decoded values matched `df.collect()`
in both offset modes. Worth one case, since that is the shape production
actually produces.
Separately, the four new large-type tests all use flat top-level vectors,
and `direct batches preserve nested list, struct, map, and null field layouts`
further down is not parameterized on `useLargeVarTypes`. The `bufferIndex` walk
in `widenOffsets` is the part of this change where a slip would silently
corrupt a later column instead of failing, and the pytest suite that does cover
nesting is not run by the per-Spark matrix jobs. Would you parameterize that
nested test the way you parameterized these? The arrangement that would catch
drift is a single batch holding `struct<utf8,int32,binary>`, `list<utf8>`,
`map<utf8,utf8>`, `list<struct<utf8>>`, a `NullVector` (zero field buffers), a
`FixedSizeBinaryVector` (two buffers, must not widen) and a trailing top-level
`utf8`. That round-trips exactly for me today.
##########
spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala:
##########
@@ -329,6 +348,55 @@ private[python] trait CometArrowPythonRunnerBase
private[python] object CometArrowPythonRunnerBase {
+ /**
+ * 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](
+ columns: Seq[CometDecodedVector],
+ allocator: BufferAllocator)(body: Seq[FieldVector] => T): T = {
+ val materialized = new ArrayList[FieldVector]()
+ try {
+ val vectors = columns.map {
+ case dictionaryVector: CometDictionaryVector =>
Review Comment:
This matches on the top-level column, so a dictionary-encoded field nested
inside a struct, list or map keeps its `DictionaryEncoding` in the advertised
schema while `widenOffsets` skips it as `Int32`. I built that shape by hand and
it fails with the same `NullPointerException: ... "provider" is null` you are
fixing for the top-level case.
I do not think Comet can produce it today. `make_builders` in
`native/shuffle/src/spark_unsafe/row.rs` passes `1.0` for list elements, struct
fields and map key/value builders, and `CometVector.getVector` dispatches on
the container types before it looks at `getDictionary()`. The one route I can
see is a chained UDF whose first worker returns a struct with a
dictionary-encoded child, which then becomes the second runner's input, which
is the shape your chained-UDF test already exercises with plain types.
Would it be worth either recursing into children here, or adding a `require`
that rejects a `getDictionary() != null` field anywhere in the tree, so that it
fails with a Comet message instead of an Arrow NPE?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
