diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index 36328468f34..395b5c45909 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -241,6 +241,16 @@ level and shares every other buffer. `move_to_spark` applies it to executed batc or sliced booleans reach the JVM misaligned ([#6288](https://github.com/apache/datafusion-comet/issues/6288)). +`zero_offsets` leaves the offsets of a sliced string, binary, list or map array as they are. arrow-rs slices those +arrays by moving the start of their offsets, so the JVM imports them with offsets that start past zero and with the +rows before the slice still at the front of the data buffer, or of the child vector for a list or map. Reading values +through the offsets is correct, but Arrow Java 18.3's `VectorSchemaRootAppender` assumes offsets that start at zero and +folds those rows into the first value it appends, and `VectorUnloader` writes them into an IPC stream. So +`Utils.serializeBatches` rebases a broadcast batch's offsets with `VectorSchemaRoot.slice` before +`coalesceBroadcastBatches` appends it ([#6779](https://github.com/apache/datafusion-comet/issues/6779)), and code that +appends imported vectors has to rebase them too. The in-memory cache serializer still writes the earlier rows +([#6995](https://github.com/apache/datafusion-comet/issues/6995)). + ### Ownership and Lifecycle Native allocates the data, and the JVM references it without copying: diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala index c884344ef15..98b70a1da78 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala @@ -27,7 +27,7 @@ import scala.jdk.CollectionConverters._ import org.apache.arrow.c.CDataDictionaryProvider import org.apache.arrow.vector._ -import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector} +import org.apache.arrow.vector.complex.{BaseRepeatedValueVector, ListVector, MapVector, StructVector} import org.apache.arrow.vector.dictionary.{Dictionary, DictionaryProvider} import org.apache.arrow.vector.dictionary.DictionaryProvider.MapDictionaryProvider import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter} @@ -267,11 +267,22 @@ object Utils extends CometTypeShim with Logging { } val provider = batchProviderOpt.getOrElse(dictionaryProvider) - val writer = new ArrowStreamWriter(root, provider, Channels.newChannel(out)) - writer.start() - writer.writeBatch() - root.clear() - writer.close() + // VectorSchemaRootAppender, which coalesces broadcast batches, assumes offsets that start + // at zero, and Arrow Java writes the bytes before the first offset too. A slice of the root + // rebases the offsets to zero and leaves those bytes out. + val written = + if (fieldVectors.forall(offsetsStartAtZero)) root else root.slice(0, root.getRowCount) + try { + val writer = new ArrowStreamWriter(written, provider, Channels.newChannel(out)) + writer.start() + writer.writeBatch() + root.clear() + writer.close() + } finally { + // The slice holds its own references on the batch's buffers, so it is closed even when the + // write fails, as it does when a cancelled broadcast interrupts the task. + if (written ne root) written.close() + } if (out.size() > 0) { (batch.numRows().toLong, cbbos.toChunkedByteBuffer) @@ -281,6 +292,23 @@ object Utils extends CometTypeShim with Logging { } } + /** + * Whether every offset buffer in `vector`, its children included, starts at zero. + * + * arrow-rs slices a string, binary, list or map array by moving the start of its offsets, and + * the C Data import keeps them, so a vector imported from a sliced native array has offsets + * that start past zero. The rows before the slice stay at the front of its data buffer, or of + * its child vector for a list or map. + */ + private def offsetsStartAtZero(vector: FieldVector): Boolean = { + val firstOffset = vector match { + case v: BaseVariableWidthVector if v.getValueCount > 0 => v.getOffsetBuffer.getInt(0) + case v: BaseRepeatedValueVector if v.getValueCount > 0 => v.getOffsetBuffer.getInt(0) + case _ => 0 + } + firstOffset == 0 && vector.getChildrenFromFields.asScala.forall(offsetsStartAtZero) + } + /** * The classes that carry the output of [[serializeBatches]] out of Comet, for Kryo registration * by [[org.apache.comet.CometKryoRegistrator]]. @@ -380,6 +408,7 @@ object Utils extends CometTypeShim with Logging { targetRoot.allocateNew() } try { + // Appending needs offsets that start at zero, which serializeBatches writes. VectorSchemaRootAppender.append(targetRoot, sourceRoot) } catch { case e: IllegalArgumentException => diff --git a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala index 36e7c5c4f4a..d96d4f7b8e7 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -19,6 +19,8 @@ package org.apache.comet.exec +import java.nio.charset.StandardCharsets + import org.scalactic.source.Position import org.scalatest.Tag @@ -1716,6 +1718,34 @@ class CometJoinSuite extends CometTestBase { } } + test("Broadcast coalescing keeps the values of sliced batches") { + // A grouped aggregate emits its groups as batch-size slices of one batch, so every slice after + // the first reaches the JVM with offsets that start past zero. Coalescing those slices used to + // prepend the bytes of the rows before each slice to its first value (#6779). + withSQLConf( + CometConf.COMET_BATCH_SIZE.key -> "100", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + val rows = (0 until 1000).map { i => + val s = if (i % 7 == 0) null else s"value-$i" + val bytes = if (s == null) null else s.getBytes(StandardCharsets.UTF_8) + val array = if (i % 11 == 0) null else Seq(s, s"element-$i") + (i, s, bytes, array, (s, i)) + } + withParquetTable(rows, "t") { + val (_, cometPlan) = checkSparkAnswerAndOperator( + sql(""" + |SELECT /*+ BROADCAST(d) */ t._1, d._2, d._3, d._4, d._5 + |FROM t JOIN (SELECT DISTINCT * FROM t) d ON t._1 = d._1 + |""".stripMargin), + Seq(classOf[CometBroadcastExchangeExec], classOf[CometBroadcastHashJoinExec])) + + val broadcast = collect(cometPlan) { case b: CometBroadcastExchangeExec => b }.head + assert(broadcast.metrics("numCoalescedRows").value == rows.length) + } + } + } + // Reproducer for SPARK-43113: full outer SMJ with a join filter that references // a nullable column should not match when the filter evaluates to NULL. test("SPARK-43113: Full outer SMJ with NULL in join filter") { diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala index 4510f9d0ac1..54a55ccde82 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala @@ -19,11 +19,20 @@ package org.apache.spark.sql.comet.util +import java.nio.channels.ClosedByInterruptException +import java.nio.charset.StandardCharsets + +import scala.jdk.CollectionConverters._ + import org.apache.arrow.c.CDataDictionaryProvider +import org.apache.arrow.memory.BufferAllocator +import org.apache.arrow.vector.VarCharVector +import org.apache.arrow.vector.ipc.message.ArrowFieldNode import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.execution.vectorized.ConstantColumnVector import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType, TimestampType} import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} +import org.apache.spark.util.io.ChunkedByteBuffer import org.apache.comet.CometArrowAllocator import org.apache.comet.vector.CometVector @@ -58,6 +67,64 @@ class UtilsSuite extends CometTestBase { assert(decoded.map(_.numRows()).sum == expected) } + test("serializeBatches rebases offsets that start past zero") { + // A slice of a native string array reaches the JVM with its original offsets, which start + // past zero, and with the bytes of the rows before the slice at the front of its data buffer. + // Coalescing such batches for a broadcast used to fold those bytes into each batch's first + // value (#6779). + val values = (0 until 24).map(i => if (i % 5 == 0) null else s"value-$i") + val skips = Seq(8, 16) + + def serialize(vectors: Seq[VarCharVector]): Seq[ChunkedByteBuffer] = { + val batches = vectors.map { v => + val column = CometVector.getVector(v, new CDataDictionaryProvider) + new ColumnarBatch(Array[ColumnVector](column), v.getValueCount) + } + try Utils.serializeBatches(batches.iterator).map(_._2).toList + finally batches.foreach(_.close()) + } + + val sliced = serialize(skips.map(slicedStrings(values, _))) + val (coalesced, batchCount, _) = Utils.coalesceBroadcastBatches(sliced.iterator) + assert(batchCount == skips.length) + val it = Utils.decodeBatches(coalesced.head, "test") + val out = it.next() + val got = (0 until out.numRows()).map { i => + if (out.column(0).isNullAt(i)) null else out.column(0).getUTF8String(i).toString + } + assert(!it.hasNext) + assert(got == skips.flatMap(values.drop)) + + // The bytes before a slice are not written either: it takes as many bytes as the same values + // in a vector of their own. + val fresh = serialize(skips.map(skip => strings(values.drop(skip)))) + assert(sliced.map(_.size) == fresh.map(_.size)) + } + + test("serializeBatches closes the slice it writes when the write fails") { + // A cancelled broadcast interrupts its task, so the next write to the interruptible channel + // serializeBatches writes through fails. The slice holds its own references on the batch's + // buffers, and they would outlive the batch unless the slice is closed on that path too. + val allocator = CometArrowAllocator.newChildAllocator("sliced-batch", 0, Long.MaxValue) + try { + val vector = slicedStrings((0 until 24).map(i => s"value-$i"), 8, allocator) + val column = CometVector.getVector(vector, new CDataDictionaryProvider) + val batch = new ColumnarBatch(Array[ColumnVector](column), vector.getValueCount) + Thread.currentThread().interrupt() + try { + intercept[ClosedByInterruptException] { + Utils.serializeBatches(Iterator(batch)).next() + } + } finally { + Thread.interrupted() + batch.close() + } + assert(allocator.getAllocatedMemory == 0) + } finally { + allocator.close() + } + } + test("serializeBatches materializes ConstantColumnVector columns") { // Spark wraps file-source partition columns and other per-batch constants in // ConstantColumnVector. When such a batch reaches Comet's serialization/export path @@ -197,4 +264,43 @@ class UtilsSuite extends CometTestBase { } } } + + private def strings( + values: Seq[String], + allocator: BufferAllocator = CometArrowAllocator): VarCharVector = { + val vector = new VarCharVector("s", allocator) + vector.allocateNew() + values.zipWithIndex.foreach { + case (null, i) => vector.setNull(i) + case (value, i) => vector.setSafe(i, value.getBytes(StandardCharsets.UTF_8)) + } + vector.setValueCount(values.length) + vector + } + + /** + * `values.drop(skip)` in the shape the C Data import gives a slice of a native string array: + * the offsets are the original ones from `skip` on, and the data buffer still starts with the + * bytes of the first `skip` values. `skip` is a multiple of 8, so the validity bitmap slices on + * a byte boundary. + */ + private def slicedStrings( + values: Seq[String], + skip: Int, + allocator: BufferAllocator = CometArrowAllocator): VarCharVector = { + val full = strings(values, allocator) + try { + val length = values.length - skip + val sliced = new VarCharVector("s", allocator) + sliced.loadFieldBuffers( + new ArrowFieldNode(length.toLong, values.drop(skip).count(_ == null).toLong), + Seq( + full.getValidityBuffer.slice(skip / 8L, (length + 7) / 8L), + full.getOffsetBuffer.slice(skip * 4L, (length + 1) * 4L), + full.getDataBuffer).asJava) + sliced + } finally { + full.close() + } + } }