viirya commented on code in PR #6371:
URL: https://github.com/apache/datafusion-comet/pull/6371#discussion_r4308293598
##########
spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala:
##########
@@ -532,40 +533,130 @@ object Utils extends CometTypeShim with Logging {
def getBatchFieldVectorsWithProviders(
batch: ColumnarBatch): Seq[(FieldVector, Option[DictionaryProvider])] = {
val rows = batch.numRows()
- (0 until batch.numCols()).map { index =>
- batch.column(index) match {
- case a: CometVector =>
- val valueVector = a.getValueVector
- val provider =
- if (valueVector.getField.getDictionary != null)
Some(a.getDictionaryProvider)
- else None
-
- (getFieldVector(valueVector, "serialize"), provider)
-
- case cv: ConstantColumnVector =>
- // Spark wraps file-source partition columns and other per-batch
constants in
- // `ConstantColumnVector`. Materialise to an Arrow vector so the
serialisation path
- // doesn't reject the batch. "UTC" is intentional -- see
`ConstantColumnVectors`.
- val materialized = ConstantColumnVectors.materialize(
- cv,
- cv.dataType(),
- rows,
- s"_const_$index",
- org.apache.comet.CometArrowAllocator,
- "UTC")
- (materialized, None)
-
- case c =>
- throw new SparkException(
- s"Comet execution only takes Arrow Arrays, but got ${c.getClass}.
" +
- "This typically happens when a Comet scan falls back to Spark
due to unsupported " +
- "data types (e.g., complex types like structs, arrays, or maps).
" +
- "To resolve this, you can: " +
- "(1) enable spark.comet.scan.allowIncompatible=true to use a
compatible native " +
- "scan variant, or " +
- "(2) enable spark.comet.convert.parquet.enabled=true to convert
Spark Parquet " +
- "data to Arrow format automatically.")
+ // Vectors made here rather than taken from `batch`. The caller releases
them with the rest of
+ // what it writes, but if a later column fails it never gets them, so they
are closed here.
+ val made = scala.collection.mutable.ArrayBuffer.empty[FieldVector]
+ def keep(vector: FieldVector): FieldVector = {
+ made += vector
+ vector
+ }
+
+ try {
+ (0 until batch.numCols()).map { index =>
+ batch.column(index) match {
+ case a: CometVector =>
+ val valueVector = a.getValueVector
+ val provider =
+ if (valueVector.getField.getDictionary != null)
Some(a.getDictionaryProvider)
+ else None
+
+ val fieldVector = valueVector match {
+ case v: LargeVarCharVector => keep(narrowOffsets(v,
ArrowType.Utf8.INSTANCE))
Review Comment:
`getBatchFieldVectorsWithProviders` has a second caller,
`CachedBatchIpc.decodeDictionaries`. It closes only the vectors it decoded and
documents the rest as "returned as they are and stay owned by `batch`". That
stays true today only because `isArrowBacked` and `matchesReaderLayout` both
keep large-offset batches off the direct cache path. Could the scaladoc of this
method say it may now return vectors it allocated, which the caller releases by
clearing them? Then the next caller won't assume every returned vector is
borrowed.
##########
spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala:
##########
@@ -166,29 +173,16 @@ 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
Review Comment:
Nit: "A caller taking the direct write path expects the batch's own vectors"
reads as if the cache would misbehave on a narrowed copy. What actually keeps
such a batch off the direct path is also `matchesReaderLayout`, since
`LargeUtf8` doesn't match the reader's `Utf8`. Could the comment say which of
the two checks it relies on, or that both apply?
##########
spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala:
##########
@@ -197,4 +191,213 @@ 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 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)),
values.length)
+ assert(
+ roundTrip(batch) == Seq(
+ classOf[VarCharVector] -> values,
+ classOf[VarBinaryVector] -> values))
+ } finally {
+ strings.close()
+ binaries.close()
+ }
+ }
+
+ test("serializeBatches rebases large-offset vectors whose offsets start past
zero") {
+ // Offsets need not start at zero. A slice imported over the C Data
interface keeps its
+ // parent's offsets and data buffer, so they start past bytes that no
value covers. The copy
+ // holds only the bytes the values cover, with offsets counted from the
first of them.
+ val values = Seq("hello", null, "wörld")
+ val strings = startingPast("xx", new LargeVarCharVector("s",
CometArrowAllocator), values)
+ val binaries = startingPast("xx", new LargeVarBinaryVector("b",
CometArrowAllocator), values)
+ try {
+ assert(strings.getOffsetBuffer.getLong(0) == 2)
+ assert(binaries.getOffsetBuffer.getLong(0) == 2)
+ val batch =
+ new ColumnarBatch(Array(cometVector(strings), cometVector(binaries)),
values.length)
+ assert(
+ roundTrip(batch) == Seq(
+ classOf[VarCharVector] -> values,
+ classOf[VarBinaryVector] -> values))
+ } finally {
+ strings.close()
+ binaries.close()
+ }
+ }
+
+ test("serializeBatches writes empty large-offset vectors with 32-bit
offsets") {
+ // A vector that holds no values need not have an offset buffer, so none
is read.
+ val strings = new LargeVarCharVector("s", CometArrowAllocator)
+ val binaries = new LargeVarBinaryVector("b", CometArrowAllocator)
+ try {
+ strings.setValueCount(0)
+ binaries.setValueCount(0)
+ assert(strings.getOffsetBuffer.capacity == 0)
+ val batch = new ColumnarBatch(Array(cometVector(strings),
cometVector(binaries)), 0)
+ assert(
+ roundTrip(batch) == Seq(
+ classOf[VarCharVector] -> Seq.empty,
+ classOf[VarBinaryVector] -> Seq.empty))
+ } finally {
+ strings.close()
+ binaries.close()
+ }
+ }
+
+ test("coalesceBroadcastBatches appends a large-offset batch to a 32-bit
one") {
Review Comment:
Nit: the name says a large-offset batch is appended to a 32-bit one, but the
test serializes the large batch first. Something like "coalesceBroadcastBatches
coalesces batches of mixed offset widths" would match either order.
##########
spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala:
##########
@@ -197,4 +191,213 @@ 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 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)),
values.length)
+ assert(
+ roundTrip(batch) == Seq(
+ classOf[VarCharVector] -> values,
+ classOf[VarBinaryVector] -> values))
+ } finally {
+ strings.close()
+ binaries.close()
+ }
+ }
+
+ test("serializeBatches rebases large-offset vectors whose offsets start past
zero") {
+ // Offsets need not start at zero. A slice imported over the C Data
interface keeps its
+ // parent's offsets and data buffer, so they start past bytes that no
value covers. The copy
+ // holds only the bytes the values cover, with offsets counted from the
first of them.
+ val values = Seq("hello", null, "wörld")
+ val strings = startingPast("xx", new LargeVarCharVector("s",
CometArrowAllocator), values)
+ val binaries = startingPast("xx", new LargeVarBinaryVector("b",
CometArrowAllocator), values)
+ try {
+ assert(strings.getOffsetBuffer.getLong(0) == 2)
+ assert(binaries.getOffsetBuffer.getLong(0) == 2)
+ val batch =
+ new ColumnarBatch(Array(cometVector(strings), cometVector(binaries)),
values.length)
+ assert(
+ roundTrip(batch) == Seq(
+ classOf[VarCharVector] -> values,
+ classOf[VarBinaryVector] -> values))
+ } finally {
+ strings.close()
+ binaries.close()
+ }
+ }
+
+ test("serializeBatches writes empty large-offset vectors with 32-bit
offsets") {
+ // A vector that holds no values need not have an offset buffer, so none
is read.
+ val strings = new LargeVarCharVector("s", CometArrowAllocator)
+ val binaries = new LargeVarBinaryVector("b", CometArrowAllocator)
+ try {
+ strings.setValueCount(0)
+ binaries.setValueCount(0)
+ assert(strings.getOffsetBuffer.capacity == 0)
+ val batch = new ColumnarBatch(Array(cometVector(strings),
cometVector(binaries)), 0)
+ assert(
+ roundTrip(batch) == Seq(
+ classOf[VarCharVector] -> Seq.empty,
+ classOf[VarBinaryVector] -> Seq.empty))
+ } 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 = stringValues(out.column(0), out.numRows())
+ 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") {
Review Comment:
Nit: this test and `serializeBatches releases the vectors it made when a
later column fails` both build `tooBigForIntOffsets` and assert the same
message. The cleanup test already covers the refusal, so this one could go.
##########
spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala:
##########
@@ -532,40 +533,130 @@ object Utils extends CometTypeShim with Logging {
def getBatchFieldVectorsWithProviders(
batch: ColumnarBatch): Seq[(FieldVector, Option[DictionaryProvider])] = {
val rows = batch.numRows()
- (0 until batch.numCols()).map { index =>
- batch.column(index) match {
- case a: CometVector =>
- val valueVector = a.getValueVector
- val provider =
- if (valueVector.getField.getDictionary != null)
Some(a.getDictionaryProvider)
- else None
-
- (getFieldVector(valueVector, "serialize"), provider)
-
- case cv: ConstantColumnVector =>
- // Spark wraps file-source partition columns and other per-batch
constants in
- // `ConstantColumnVector`. Materialise to an Arrow vector so the
serialisation path
- // doesn't reject the batch. "UTC" is intentional -- see
`ConstantColumnVectors`.
- val materialized = ConstantColumnVectors.materialize(
- cv,
- cv.dataType(),
- rows,
- s"_const_$index",
- org.apache.comet.CometArrowAllocator,
- "UTC")
- (materialized, None)
-
- case c =>
- throw new SparkException(
- s"Comet execution only takes Arrow Arrays, but got ${c.getClass}.
" +
- "This typically happens when a Comet scan falls back to Spark
due to unsupported " +
- "data types (e.g., complex types like structs, arrays, or maps).
" +
- "To resolve this, you can: " +
- "(1) enable spark.comet.scan.allowIncompatible=true to use a
compatible native " +
- "scan variant, or " +
- "(2) enable spark.comet.convert.parquet.enabled=true to convert
Spark Parquet " +
- "data to Arrow format automatically.")
+ // Vectors made here rather than taken from `batch`. The caller releases
them with the rest of
+ // what it writes, but if a later column fails it never gets them, so they
are closed here.
+ val made = scala.collection.mutable.ArrayBuffer.empty[FieldVector]
+ def keep(vector: FieldVector): FieldVector = {
+ made += vector
+ vector
+ }
+
+ try {
+ (0 until batch.numCols()).map { index =>
+ batch.column(index) match {
+ case a: CometVector =>
+ val valueVector = a.getValueVector
+ val provider =
+ if (valueVector.getField.getDictionary != null)
Some(a.getDictionaryProvider)
+ else None
+
+ val fieldVector = valueVector match {
+ case v: LargeVarCharVector => keep(narrowOffsets(v,
ArrowType.Utf8.INSTANCE))
+ case v: LargeVarBinaryVector => keep(narrowOffsets(v,
ArrowType.Binary.INSTANCE))
+ case v => getFieldVector(v, "serialize")
+ }
+ (fieldVector, provider)
+
+ case cv: ConstantColumnVector =>
+ // Spark wraps file-source partition columns and other per-batch
constants in
+ // `ConstantColumnVector`. Materialise to an Arrow vector so the
serialisation path
+ // doesn't reject the batch. "UTC" is intentional -- see
`ConstantColumnVectors`.
+ val materialized = keep(
+ ConstantColumnVectors.materialize(
+ cv,
+ cv.dataType(),
+ rows,
+ s"_const_$index",
+ org.apache.comet.CometArrowAllocator,
+ "UTC"))
+ (materialized, None)
+
+ case c =>
+ throw new SparkException(
+ s"Comet execution only takes Arrow Arrays, but got
${c.getClass}. " +
+ "This typically happens when a Comet scan falls back to Spark
due to " +
+ "unsupported data types (e.g., complex types like structs,
arrays, or maps). " +
+ "To resolve this, you can: " +
+ "(1) enable spark.comet.scan.allowIncompatible=true to use a
compatible native " +
+ "scan variant, or " +
+ "(2) enable spark.comet.convert.parquet.enabled=true to
convert Spark Parquet " +
+ "data to Arrow format automatically.")
+ }
+ }
+ } catch {
+ case NonFatal(e) =>
Review Comment:
This covers a failure inside the column loop. A failure after this method
returns still drops the copies before any root is cleared, though. In
`serializeBatches` that means `combineDictionaryProviders` or
`writer.writeBatch`, and in `CachedBatchIpc` it means `lookupDictionary` or
`DictionaryEncoder.decode`. The gap already existed for materialized constants,
but a narrowed copy can be close to 2 GiB of `CometArrowAllocator` memory. Not
blocking. Would returning which vectors the caller owns, or putting the root's
lifecycle in `serializeBatches` in a `try`/`finally`, close it more completely
than tracking `made` here?
--
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]
