andygrove commented on code in PR #6371:
URL: https://github.com/apache/datafusion-comet/pull/6371#discussion_r4260582738
##########
spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala:
##########
@@ -537,7 +538,12 @@ object Utils extends CometTypeShim with Logging {
if (valueVector.getField.getDictionary != null)
Some(a.getDictionaryProvider)
else None
- (getFieldVector(valueVector, "serialize"), provider)
+ val fieldVector = valueVector match {
+ case v: LargeVarCharVector => narrowOffsets(v,
ArrowType.Utf8.INSTANCE)
Review Comment:
Fixed in 1c8f68c048. `getBatchFieldVectorsWithProviders` now keeps track of
every vector it makes, the 32-bit copies and the materialized
`ConstantColumnVector`s, and closes them before rethrowing if a later column
fails. The batch's own vectors are left to their owner.
The new test `serializeBatches releases the vectors it made when a later
column fails` puts a constant, a small large-offset column and an oversized one
in one batch. Once the inputs are closed, it checks that `CometArrowAllocator`
is back where it started and that the small column's own vector is still
intact. Without the cleanup, 16402 bytes stay allocated.
##########
spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala:
##########
@@ -566,6 +572,68 @@ object Utils extends CometTypeShim with Logging {
}
}
+ /**
+ * A copy of `large` with 32-bit offsets, as a vector of `narrowType`.
+ *
+ * A PyArrow UDF can hand back a `large_string` or `large_binary` column,
which Comet reads but
+ * [[getFieldVector]] does not accept. The column is still `StringType` or
`BinaryType`, which
+ * Comet otherwise carries as `Utf8` or `Binary`, so it is written in that
form. Whatever reads
+ * the bytes back then gets the type it plans for, and the batches of one
broadcast share a
+ * schema whichever offset width each arrived with:
[[coalesceBroadcastBatches]] appends them to
+ * a root built from the first one's schema, and a native plan reading them
takes its input
+ * stream's schema from the first one too. Only a column holding more data
than 32-bit offsets
+ * can address is refused.
+ *
+ * Like a materialized `ConstantColumnVector`, the copy is released when the
caller clears the
+ * vectors it wrote. `large` itself stays with its owner.
+ */
+ private def narrowOffsets(
+ large: BaseLargeVariableWidthVector,
+ narrowType: ArrowType): FieldVector = {
+ val numValues = large.getValueCount
+ val largeOffsets = large.getOffsetBuffer
+ val largeWidth = BaseLargeVariableWidthVector.OFFSET_WIDTH.toLong
+ // An empty vector need not have an offset buffer to read.
+ val start = if (numValues == 0) 0L else largeOffsets.getLong(0)
Review Comment:
Good catch, none of them did. `serializeBatches rebases large-offset vectors
whose offsets start past zero` now loads your example with `loadFieldBuffers`:
`hello`, null and `wörld` at offsets `[2, 7, 7, 13]` over `xxhellowörld`, as
both string and binary. Dropping the `- start`, or copying the data from byte
0, now fails it. `serializeBatches writes empty large-offset vectors with
32-bit offsets` covers the zero-row case, with vectors that have no offset
buffer at all. Both are in 1c8f68c048.
##########
spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala:
##########
@@ -197,4 +202,94 @@ class UtilsSuite extends CometTestBase {
}
}
}
+
+ test("serializeBatches writes large-offset string and binary vectors with
32-bit offsets") {
+ // A PyArrow UDF returning pa.large_string() or pa.large_binary() hands
Comet a
+ // LargeVarCharVector or LargeVarBinaryVector. The column is still
StringType or BinaryType,
+ // so it is written as the Utf8 or Binary vector every reader of the bytes
expects for those.
+ val values = Seq("hello", null, "", "wörld")
+ val numRows = values.length
+ val strings = fill(new LargeVarCharVector("s", CometArrowAllocator),
values)
+ val binaries = fill(new LargeVarBinaryVector("b", CometArrowAllocator),
values)
+ try {
+ val batch =
+ new ColumnarBatch(Array(cometVector(strings), cometVector(binaries)),
numRows)
+
+ val (rowCount, buf) = Utils.serializeBatches(Iterator(batch)).next()
+ assert(rowCount == numRows)
+
+ val it = Utils.decodeBatches(buf, "test")
+ assert(it.hasNext)
+ val out = it.next()
+ val vectorClasses = (0 until out.numCols()).map(i =>
+ out.column(i).asInstanceOf[CometVector].getValueVector.getClass)
+ val gotStrings =
+ (0 until numRows).map(i =>
Option(out.column(0).getUTF8String(i)).map(_.toString).orNull)
+ val gotBinaries = (0 until numRows).map(i =>
+ Option(out.column(1).getBinary(i)).map(new String(_,
StandardCharsets.UTF_8)).orNull)
+ assert(!it.hasNext)
+
+ assert(vectorClasses == Seq(classOf[VarCharVector],
classOf[VarBinaryVector]))
+ assert(gotStrings == values)
+ assert(gotBinaries == values)
+ } finally {
+ strings.close()
+ binaries.close()
+ }
+ }
+
+ test("coalesceBroadcastBatches appends a large-offset batch to a 32-bit
one") {
+ // The coalescer appends every batch to a root built from the first one's
schema, and leaves
+ // the broadcast uncoalesced when a batch does not match it. A producer
can hand back
+ // large_string for one batch and string for the next, so both must
serialize alike.
+ val large = fill(new LargeVarCharVector("s", CometArrowAllocator),
Seq("a", null))
+ val regular = fill(new VarCharVector("s", CometArrowAllocator), Seq("b",
"c"))
+ try {
+ val batches = Seq(large, regular).map(v => new
ColumnarBatch(Array(cometVector(v)), 2))
+ val bufs =
Utils.serializeBatches(batches.iterator).map(_._2).toSeq.iterator
+ val (coalesced, batchCount, totalRows) =
Utils.coalesceBroadcastBatches(bufs)
+ // A batch count of zero is the fallback to the uncoalesced buffers.
+ assert(batchCount == 2)
+ assert(totalRows == 4)
+
+ val it = coalesced.iterator.flatMap(Utils.decodeBatches(_, "test"))
+ val out = it.next()
+ val got = (0 until out.numRows()).map(i =>
+ Option(out.column(0).getUTF8String(i)).map(_.toString).orNull)
+ assert(!it.hasNext)
+ assert(got == Seq("a", null, "b", "c"))
+ } finally {
+ large.close()
+ regular.close()
+ }
+ }
+
+ test("serializeBatches refuses a large-offset column too big for 32-bit
offsets") {
+ // The offsets alone claim 2 GiB of data. The size is checked before
anything is copied, so
+ // the data itself never has to exist.
+ val strings = fill(new LargeVarCharVector("s", CometArrowAllocator),
Seq("a"))
+ try {
+ strings.getOffsetBuffer.setLong(
+ BaseLargeVariableWidthVector.OFFSET_WIDTH,
+ Int.MaxValue + 1L)
+ val batch = new ColumnarBatch(Array(cometVector(strings)), 1)
+ val e =
intercept[SparkException](Utils.serializeBatches(Iterator(batch)).next())
+ assert(e.getMessage.contains("more than 32-bit offsets can address"),
e.getMessage)
+ } finally {
+ strings.close()
+ }
+ }
+
+ /** Writes `values` into `vector`, leaving a null entry unset so that it
reads back as null. */
+ private def fill[V <: VariableWidthFieldVector](vector: V, values:
Seq[String]): V = {
Review Comment:
Done in 1c8f68c048. That test builds both vectors with `fill` and wraps them
with `cometVector` now, and the round-trip tests share a `roundTrip` helper for
decoding and reading the values back.
--
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]
