fix: keep the values of sliced batches when coalescing a broadcast

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:

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 =>

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") {

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()
+    }
+  }
 }