andygrove commented on code in PR #6371:
URL: https://github.com/apache/datafusion-comet/pull/6371#discussion_r4312099317
##########
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:
Done in a06a59a05b. The scaladoc of `getBatchFieldVectorsWithProviders` now
says that it may return vectors it allocated, a 32-bit copy or a materialized
constant, and that the caller releases them. `Utils.allocatedFor` tells them
apart from the batch's own vectors. `CachedBatchIpc.decodeDictionaries` no
longer relies on the guard: it adds them to the vectors it returns for the
caller to close, so a large-offset batch that reached it would not leak.
##########
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:
Done in a06a59a05b, using the shape you describe. `getBatchFieldVectors` and
`serializeBatches` close the vectors made for the batch if anything after the
column loop fails, which covers `combineDictionaryProviders` and `writeBatch`.
`decodeDictionaries` closes them if `lookupDictionary` or
`DictionaryEncoder.decode` fails, and on success it returns them with the
decoded ones. The batch's own vectors stay with their owner. Two tests: a
dictionary clash in `serializeBatches` next to a large-offset column, and a
failed compression in `CachedBatchIpc.serialize`. Each leaves the allocator
where it started, and each fails if I remove the new cleanup.
##########
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:
Done in a06a59a05b. The comment now says that
`CachedBatchIpc.matchesReaderLayout` also keeps such a batch off the direct
write path, because it refuses `LargeUtf8` against the reader's `Utf8`, so
either check alone is enough.
##########
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:
Done in a06a59a05b. It is now "coalesceBroadcastBatches coalesces batches of
mixed offset widths".
##########
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:
Agreed, the cleanup test asserts the same message and covers the refusal. I
removed this one in a06a59a05b.
--
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]
