viirya commented on code in PR #6371:
URL: https://github.com/apache/datafusion-comet/pull/6371#discussion_r4316162797
##########
spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala:
##########
@@ -233,29 +236,18 @@ class UtilsSuite extends CometTestBase {
test("isArrowBacked rejects large-offset Arrow vectors") {
// A CometPlainVector can wrap a LargeVarCharVector or
LargeVarBinaryVector -- an accelerated
- // mapInArrow returning pa.large_string() produces one -- but
getFieldVector rejects both. If
- // isArrowBacked accepted them, a caller would take the direct write path
and then fail, so it
- // must report false and let the caller convert the batch instead.
- val numRows = 2
- Seq[org.apache.arrow.vector.FieldVector](
- {
- val v = new org.apache.arrow.vector.LargeVarCharVector("s",
CometArrowAllocator)
- v.allocateNew(numRows)
- v.setSafe(0, "hello".getBytes("UTF-8"))
- v.setSafe(1, "world".getBytes("UTF-8"))
- v.setValueCount(numRows)
- v
- }, {
- val v = new org.apache.arrow.vector.LargeVarBinaryVector("b",
CometArrowAllocator)
- v.allocateNew(numRows)
- v.setSafe(0, "hello".getBytes("UTF-8"))
- v.setSafe(1, "world".getBytes("UTF-8"))
- v.setValueCount(numRows)
- v
- }).foreach { vector =>
+ // mapInArrow returning pa.large_string() produces one -- but
getFieldVector rejects both, and
+ // getBatchFieldVectors writes a 32-bit copy of them rather than the
vector itself. A caller
+ // taking the direct write path expects the batch's own vectors, so
isArrowBacked must report
Review Comment:
Nit: since `decodeDictionaries` now returns the copies for the caller to
close, "A caller taking the direct write path expects the batch's own vectors"
no longer holds. The direct path would release a narrowed copy correctly. What
keeps such a batch off it now is `isArrowBacked` itself and
`matchesReaderLayout`. Could this sentence go, leaving the part about either
check being enough?
##########
spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala:
##########
@@ -3039,6 +3039,57 @@ class CometInMemoryCacheSuite extends CometTestBase {
}
}
+ test("Comet in-memory cache releases the 32-bit copy it makes of a
large-offset column") {
Review Comment:
Nit: this calls `CachedBatchIpc.serialize` directly, so it skips the
`isArrowBacked` and `matchesReaderLayout` checks that keep a large-offset batch
off this path in practice. Could the comment say it tests the ownership
contract rather than a path the cache takes today? Otherwise a reader may
conclude that the cache writes large-offset columns by narrowing them.
##########
spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala:
##########
@@ -479,22 +488,59 @@ object Utils extends CometTypeShim with Logging {
def isArrowBacked(batch: ColumnarBatch): Boolean =
(0 until batch.numCols()).forall { i =>
batch.column(i) match {
- // Not every CometVector can be handed to getFieldVector: a
CometPlainVector can wrap a
+ // Not every CometVector is handed out as it stands: a
CometPlainVector can wrap a
// LargeVarCharVector or LargeVarBinaryVector (an accelerated
mapInArrow returning
- // pa.large_string(), for instance), which it rejects. Answering true
for those would
- // send a batch down the direct write path that then fails, so check
the vector itself
- // and let the caller convert instead.
+ // pa.large_string(), for instance), which getFieldVector rejects and
+ // getBatchFieldVectorsWithProviders replaces with a 32-bit copy. So
check the vector
+ // itself and let the caller convert such a batch.
case v: CometVector => isSupportedFieldVector(v.getValueVector)
case _ => false
}
}
+ /**
+ * Field vectors of `batch` and the dictionaries they refer to, as one
provider.
+ *
+ * Some of the vectors may have been allocated here rather than taken from
`batch`; see
+ * [[getBatchFieldVectorsWithProviders]]. They are released with the rest
when the caller clears
+ * the vectors it wrote, or by [[closeAfterFailure]] if the write fails.
+ */
def getBatchFieldVectors(
batch: ColumnarBatch): (Seq[FieldVector], Option[DictionaryProvider]) = {
val columns = getBatchFieldVectorsWithProviders(batch)
- (columns.map(_._1), combineDictionaryProviders(columns))
+ val vectors = columns.map(_._1)
+ try {
+ (vectors, combineDictionaryProviders(columns))
+ } catch {
+ case NonFatal(e) =>
+ closeAfterFailure(allocatedFor(batch, vectors), e)
+ throw e
+ }
+ }
+
+ /**
+ * The vectors of `vectors`, which were returned for `batch`, that were
allocated for it rather
+ * than taken from it: a 32-bit copy of a large-offset column, or a
materialized
+ * `ConstantColumnVector`. Whoever holds the returned vectors is the only
one who can release
+ * them, and `batch` still owns the rest.
+ */
+ def allocatedFor(batch: ColumnarBatch, vectors: Seq[FieldVector]):
Seq[FieldVector] = {
+ val borrowed = (0 until batch.numCols()).flatMap { i =>
+ batch.column(i) match {
+ case v: CometVector => Some(v.getValueVector)
+ case _ => None
+ }
+ }
+ vectors.filterNot(v => borrowed.exists(_ eq v))
Review Comment:
Nit: `decodeDictionaries` now calls this for every cached batch, and the
`exists` scan is quadratic in the column count. On the direct path it also
always comes back empty, since `isArrowBacked` rejects constants and
large-offset vectors. An identity set (`java.util.IdentityHashMap`-backed)
would keep it linear for wide tables. Not important at typical widths.
--
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]
