andygrove commented on code in PR #6729:
URL: https://github.com/apache/datafusion-comet/pull/6729#discussion_r4199169121
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala:
##########
@@ -44,20 +45,124 @@ import org.apache.spark.unsafe.Platform
* directly.
*/
private[arrow] object ArrowWriter {
+
+ /**
+ * Prepares the vectors of `root` for a new batch, with room for
`fixedWidthCapacity` values in
+ * each top-level fixed-width vector, and returns a writer that appends to
them.
+ *
+ * A vector keeps its buffers, and the batch is written over the last one,
when nothing else
+ * references them. Otherwise it gets new buffers and leaves the old ones to
whoever holds them:
+ * a batch exported over the C Data Interface keeps its buffers until native
releases it, and a
+ * batch loaded into another vector until that vector lets go. Both retain
each buffer, which
+ * shows in its reference count. The count is per allocator, so loading the
buffers into a
+ * vector of another allocator does not show. A consumer that does must
close the batch after,
+ * as `ColumnarBatchArrowReader` does, which leaves these vectors no buffers
to keep.
+ *
+ * Data buffers are not zeroed, kept or new. The writers write every value
of a batch, and a
+ * zero for each null row they write one at a time, so a batch reads the
same either way.
+ */
def create(root: VectorSchemaRoot, fixedWidthCapacity: Int): ArrowWriter = {
require(fixedWidthCapacity >= 0, "Fixed-width capacity must be
non-negative")
val children = root.getFieldVectors().asScala.map { vector =>
vector match {
case fixedWidth: BaseFixedWidthVector =>
- fixedWidth.allocateNew(fixedWidthCapacity)
+ // The row path writes top-level fixed-width values without checking
capacity.
+ if (ownsBuffers(fixedWidth) && fixedWidth.getValueCapacity >=
fixedWidthCapacity) {
+ clearKeepingValues(fixedWidth)
+ } else {
+ allocateFixedWidth(fixedWidth, fixedWidthCapacity)
+ }
case _ =>
- vector.allocateNew()
+ if (ownsBuffers(vector)) clearKeepingValues(vector) else
vector.allocateNew()
}
createFieldWriter(vector)
}
new ArrowWriter(root, children.toArray)
}
+ /**
+ * Whether nothing but `vector` and its children references their buffers.
Arrow allocates a
+ * vector's validity and data or offset buffers as slices of one allocation,
each counting a
+ * reference to it, so references are counted per allocation. A vector whose
buffers are empty,
+ * as after `clear`, or of a type not handled here, has nothing to keep.
+ */
+ private def ownsBuffers(vector: FieldVector): Boolean = {
+ val references = new java.util.IdentityHashMap[ReferenceManager, Integer]()
+ def count(v: FieldVector): Boolean = {
+ val buffers = v match {
+ case f: BaseFixedWidthVector => Some(Seq(f.getValidityBuffer,
f.getDataBuffer))
+ case f: BaseVariableWidthVector =>
+ Some(Seq(f.getValidityBuffer, f.getOffsetBuffer, f.getDataBuffer))
+ case l: ListVector => Some(Seq(l.getValidityBuffer, l.getOffsetBuffer))
+ case s: StructVector => Some(Seq(s.getValidityBuffer))
+ case _: NullVector => Some(Seq.empty)
+ case _ => None
+ }
+ buffers.exists(_.forall { buffer =>
+ val manager = buffer.getReferenceManager
+ references.put(manager, references.getOrDefault(manager, 0).intValue +
1)
+ buffer.capacity > 0
+ }) && v.getChildrenFromFields.asScala.forall(count)
+ }
+ count(vector) && references.asScala.forall { case (manager, n) =>
+ manager.getRefCount == n.intValue
+ }
+ }
+
+ /**
+ * Empties `vector` and its children for a new batch, zeroing validity and
offsets as a new
+ * allocation does, but not the values, which the batch writes over.
Booleans are zeroed whole:
+ * their values are bits, and a byte holds some past the end of the batch.
+ */
+ private def clearKeepingValues(vector: FieldVector): Unit = {
+ vector match {
+ case bits: BitVector => bits.reset()
+ case f: BaseFixedWidthVector => zero(f.getValidityBuffer)
+ case f: BaseVariableWidthVector =>
+ zero(f.getValidityBuffer)
+ zero(f.getOffsetBuffer)
+ case l: ListVector =>
+ zero(l.getValidityBuffer)
+ zero(l.getOffsetBuffer)
+ l.setLastSet(-1)
+ case s: StructVector => zero(s.getValidityBuffer)
+ case _: NullVector =>
+ }
+ vector.getChildrenFromFields.asScala.foreach(clearKeepingValues)
+ vector.setValueCount(0)
+ }
+
+ private def zero(buffer: ArrowBuf): Unit = buffer.setZero(0, buffer.capacity)
Review Comment:
Fixed in 881e0ea82: both now discard the result with `val _ =`, as the rest
of the code does under this profile. `./mvnw -B test-compile -Pspark-3.5
-Pstrict-warnings -DskipTests` passes locally. The four `@nowarn` errors in the
shuffle managers went away with them. They only appeared because compilation
stopped at these two, before the warnings those annotations suppress were
reported.
--
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]
