sunchao commented on code in PR #6729:
URL: https://github.com/apache/datafusion-comet/pull/6729#discussion_r4198901342
##########
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:
[P2] Explicitly discard both new `setZero` return values. `ArrowBuf.setZero`
returns `ArrowBuf`, so this `Unit` helper and the fallback branch in
`clearValue` at line 399 produce `discarded non-Unit value` diagnostics. The
repository's strict Scala profile enables `-Ywarn-value-discard` and
`-Xfatal-warnings`, so `./mvnw -B test-compile -Pspark-3.5 -Pstrict-warnings
-DskipTests` fails compilation instead of passing the CI check. End both
expressions with an explicit `()` or equivalent intentional discard.
Evidence: Exact-head CI job
https://github.com/apache/datafusion-comet/actions/runs/37507425975/job/112423899780
reports errors at ArrowWriters.scala:135 and :399. Independently compiling the
exact head writer with Scala 2.12.18, Arrow 18.3.0 and the repository's strict
flags reproduced precisely these two warnings and a fatal compilation error.
The same isolated compilation using the specified base's ArrowWriters.scala
exited successfully.
--
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]
