fix: write a sliced batch to the in-memory cache without the rows before it

docs/source/contributor-guide/ffi.md

@@ -243,13 +243,16 @@ or sliced booleans reach the JVM misaligned
 
 `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
+rows before the slice still at the front of the data buffer, or of the child vector for a list or map. The first slice
+of an array starts at zero, but a list or map still has the elements of the later slices in its child. 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)).
+`coalesceBroadcastBatches` appends it ([#6779](https://github.com/apache/datafusion-comet/issues/6779)), and
+`CachedBatchIpc.serialize` does the same before the in-memory cache writes a batch, so that a cached slice neither
+stores nor reports the rows before it ([#6995](https://github.com/apache/datafusion-comet/issues/6995)). The Arrow
+Python runner rebases the same way before it sends a batch to the worker. All three go through
+`Utils.withRebasedOffsets`, and code that appends or writes imported vectors has to use it too.
 
 ### Ownership and Lifecycle
 

spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala

@@ -224,32 +224,37 @@ private[comet] object CachedBatchIpc {
         root.setRowCount(batch.numRows())
       }
 
-      // Unloaded plain and compressed afterwards rather than by handing the codec to the unloader;
-      // see compressed for why.
-      val fields = vectors.map(_.getField)
-      val unloader = new VectorUnloader(root, true, NoCompressionCodec.INSTANCE, true)
-      val plainBatch = unloader.getRecordBatch
-      val (sizes, recordBatch) =
+      // A grouped aggregate emits its groups as slices of one batch, which reach the JVM with
+      // offsets that start past zero. Unloaded as they are, each slice would be cached, and
+      // reported by columnSizes, with the bytes of the slices before it.
+      Utils.withRebasedOffsets(root) { written =>
+        // Unloaded plain and compressed afterwards rather than by handing the codec to the
+        // unloader; see compressed for why.
+        val fields = vectors.map(_.getField)
+        val unloader = new VectorUnloader(written, true, NoCompressionCodec.INSTANCE, true)
+        val plainBatch = unloader.getRecordBatch
+        val (sizes, recordBatch) =
+          try {
+            (columnSizes(fields, plainBatch), compressed(plainBatch, codec, allocator))
+          } finally {
+            plainBatch.close()
+          }
         try {
-          (columnSizes(fields, plainBatch), compressed(plainBatch, codec, allocator))
-        } finally {
-          plainBatch.close()
-        }
-      try {
-        // Leaves the batch in the state serializeBatches leaves one. The record batch holds its
-        // own buffers by now, so this does not touch it.
-        root.clear()
+          // Leaves the batch in the state serializeBatches leaves one. The record batch holds its
+          // own buffers by now, and so does the slice, so this does not touch either.
+          root.clear()
 
-        val out = new ChunkedByteBufferOutputStream(chunkSize, ByteBuffer.allocate)
-        try {
-          val channel = new WriteChannel(Channels.newChannel(out))
-          MessageSerializer.serialize(channel, recordBatch)
+          val out = new ChunkedByteBufferOutputStream(chunkSize, ByteBuffer.allocate)
+          try {
+            val channel = new WriteChannel(Channels.newChannel(out))
+            MessageSerializer.serialize(channel, recordBatch)
+          } finally {
+            out.close()
+          }
+          (out.toChunkedByteBuffer, sizes)
         } finally {
-          out.close()
+          recordBatch.close()
         }
-        (out.toChunkedByteBuffer, sizes)
-      } finally {
-        recordBatch.close()
       }
     } finally {
       // Only the vectors allocated for this write. The rest belong to the input batch.

spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala

@@ -277,20 +277,13 @@ object Utils extends CometTypeShim with Logging {
         val provider = batchProviderOpt.getOrElse(dictionaryProvider)
 
         // 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 {
+        // at zero.
+        withRebasedOffsets(root) { written =>
           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()
         }
       } catch {
         case NonFatal(e) =>
@@ -309,20 +302,44 @@ object Utils extends CometTypeShim with Logging {
   }
 
   /**
-   * Whether every offset buffer in `vector`, its children included, starts at zero.
+   * Runs `write` on `root` itself when none of its vectors is a slice of a larger one, and
+   * otherwise on a slice of it that rebases the offsets to zero and leaves out the bytes or
+   * elements outside them, which Arrow Java would write. The slice shares the other buffers and
+   * holds its own references on them, so it is closed whether `write` returns or fails, as it
+   * does when a cancelled broadcast interrupts the task.
+   */
+  private[sql] def withRebasedOffsets[T](root: VectorSchemaRoot)(
+      write: VectorSchemaRoot => T): T = {
+    val written =
+      if (root.getFieldVectors.asScala.forall(holdsOnlyItsOwnValues)) root
+      else root.slice(0, root.getRowCount)
+    try write(written)
+    finally if (written ne root) written.close()
+  }
+
+  /**
+   * Whether every vector in `vector`, its children included, holds only the bytes or elements its
+   * own offsets cover.
    *
    * 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.
+   * its child vector for a list or map. The first slice of an array starts at zero instead, and a
+   * list or map still has the elements of the later slices in its child, which is longer than the
+   * last offset. A string or binary vector stops writing its data at the last offset, so only its
+   * first offset can show a slice.
    */
-  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
+  private def holdsOnlyItsOwnValues(vector: FieldVector): Boolean = {
+    val ownsItsValues = vector match {
+      case v: BaseVariableWidthVector if v.getValueCount > 0 =>
+        v.getOffsetBuffer.getInt(0) == 0
+      case v: BaseRepeatedValueVector if v.getValueCount > 0 =>
+        v.getOffsetBuffer.getInt(0) == 0 &&
+        v.getOffsetBuffer.getInt(v.getValueCount.toLong * BaseRepeatedValueVector.OFFSET_WIDTH) ==
+          v.getDataVector.getValueCount
+      case _ => true
     }
-    firstOffset == 0 && vector.getChildrenFromFields.asScala.forall(offsetsStartAtZero)
+    ownsItsValues && vector.getChildrenFromFields.asScala.forall(holdsOnlyItsOwnValues)
   }
 
   /**

spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala

@@ -730,6 +730,17 @@ private[python] object CometArrowPythonRunnerBase {
       allocator: BufferAllocator): Unit = {
     val sourceRoot =
       new VectorSchemaRoot(sourceVectors.map(_.getField).asJava, sourceVectors.asJava, numRows)
+    // A slice of a native array reaches the JVM with offsets that start past zero, and the bytes
+    // before the first of them would be written with the batch. The slice this makes holds its
+    // own references, so it is closed here and the source root still is not.
+    Utils.withRebasedOffsets(sourceRoot)(serializeRoot(writeChannel, _, numRows, allocator))
+  }
+
+  private def serializeRoot(
+      writeChannel: WriteChannel,
+      sourceRoot: VectorSchemaRoot,
+      numRows: Int,
+      allocator: BufferAllocator): Unit = {
     val sourceBatch = new VectorUnloader(sourceRoot).getRecordBatch
     try {
       val validityBytes = (numRows.toLong + 7L) / 8L

spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala

@@ -2993,6 +2993,63 @@ class CometInMemoryCacheSuite extends CometTestBase {
     assert(zstdSize == plainSize, s"the relation size should not depend on the codec: $sizes")
   }
 
+  test("Comet in-memory cache does not store the bytes before a sliced aggregate output") {
+    // A grouped aggregate emits its groups as batch-size slices of one batch, and a slice of a
+    // string array reaches the JVM with offsets that start past zero and the rows before it still
+    // at the front of its data buffer. Written as it stands, the k-th slice is cached with the
+    // bytes of the k - 1 slices before it, and the relation reports them in its size.
+    val rows = 1000L
+    withTempPath { dir =>
+      withSQLConf(
+        SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+        SQLConf.SHUFFLE_PARTITIONS.key -> "1",
+        CometConf.COMET_BATCH_SIZE.key -> "100",
+        CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") {
+        spark.catalog.clearCache()
+        spark
+          .range(0, rows, 1, 1)
+          .selectExpr("concat('key-', lpad(cast(id AS string), 12, '0')) AS k")
+          .write
+          .parquet(dir.getCanonicalPath)
+
+        def cachedSize(name: String, df: DataFrame): (Long, Long) = {
+          df.createOrReplaceTempView(name)
+          spark.catalog.cacheTable(name)
+          assert(spark.table(name).count() == rows)
+          val relation = spark.sharedState.cacheManager
+            .lookupCachedData(spark.table(name))
+            .get
+            .cachedRepresentation
+          val batches = relation.cacheBuilder.cachedColumnBuffers
+          val classes = batches.map(_.getClass.getSimpleName).distinct().collect()
+          assert(classes.sameElements(Array("CometCachedBatch")))
+          (relation.computeStats().sizeInBytes.toLong, batches.count())
+        }
+
+        try {
+          val scanned = spark.read.parquet(dir.getCanonicalPath)
+          val distinct = scanned.distinct()
+          assert(
+            distinct.queryExecution.executedPlan.exists(_.isInstanceOf[CometHashAggregateExec]))
+
+          // The aggregate is cached first, so that it reads the Parquet scan and not a cached copy
+          // of it. It has to emit more than one slice for the bytes before a slice to matter.
+          val (distinctSize, distinctBatches) = cachedSize("distinct_cache", distinct)
+          assert(distinctBatches > 1, s"the aggregate emitted $distinctBatches batches")
+          val (scannedSize, _) = cachedSize("scanned_cache", scanned)
+          // The same rows, so the same bytes of string data. A slice that carried the bytes before
+          // it would make this about 5 times larger, as the 10 slices average 5.5 slices of data.
+          assert(
+            distinctSize < scannedSize * 1.1,
+            s"the aggregate's cache is $distinctSize bytes, a scan's of the same rows $scannedSize")
+          checkAnswer(spark.table("distinct_cache"), scanned)
+        } finally {
+          spark.catalog.clearCache()
+        }
+      }
+    }
+  }
+
   test("Comet in-memory cache scans no columns for a row-count-only query") {
     // SELECT count(*) selects no columns, and the scan must keep it that way. Widening it -- to
     // the whole cache schema, or to a single placeholder column -- makes the emitted batches

spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala

@@ -26,8 +26,8 @@ import scala.jdk.CollectionConverters._
 
 import org.apache.arrow.c.CDataDictionaryProvider
 import org.apache.arrow.memory.BufferAllocator
-import org.apache.arrow.vector.{BaseLargeVariableWidthVector, BitVectorHelper, IntVector, LargeVarBinaryVector, LargeVarCharVector, ValueVector, VarBinaryVector, VarCharVector, VariableWidthFieldVector}
-import org.apache.arrow.vector.complex.StructVector
+import org.apache.arrow.vector.{BaseLargeVariableWidthVector, BitVectorHelper, FieldVector, IntVector, LargeVarBinaryVector, LargeVarCharVector, ValueVector, VarBinaryVector, VarCharVector, VariableWidthFieldVector, VectorSchemaRoot, VectorUnloader}
+import org.apache.arrow.vector.complex.{ListVector, StructVector}
 import org.apache.arrow.vector.dictionary.{Dictionary, DictionaryProvider}
 import org.apache.arrow.vector.ipc.message.ArrowFieldNode
 import org.apache.arrow.vector.types.pojo.{ArrowType, DictionaryEncoding, Field, FieldType}
@@ -105,6 +105,41 @@ class UtilsSuite extends CometTestBase {
     assert(sliced.map(_.size) == fresh.map(_.size))
   }
 
+  test("withRebasedOffsets writes a sliced list as its own values") {
+    // A native slice of a list array has the offsets of the list and a child that both still hold
+    // the rows around the slice. The first slice starts at zero, so only the child shows that it
+    // is a slice: it holds the elements of every later one. Written as it stands, a slice carries
+    // those elements too.
+    val values: Seq[Seq[String]] = (0 until 24).map { i =>
+      if (i % 5 == 0) null else (0 until i % 4).map(j => s"element-$i-$j")
+    }
+    val allocator = CometArrowAllocator.newChildAllocator("sliced-lists", 0, Long.MaxValue)
+    try {
+      Seq((0, 8), (8, 8), (8, 16), (16, 8)).foreach { case (skip, take) =>
+        val sliced = slicedLists(values, skip, take, allocator)
+        val fresh = lists(values.slice(skip, skip + take), allocator)
+        try {
+          val offsets = sliced.getOffsetBuffer
+          assert(
+            offsets.getInt(0) > 0 || offsets.getInt(
+              take * 4L) < sliced.getDataVector.getValueCount,
+            s"the test needs a slice: skip=$skip, take=$take")
+          val (slicedBytes, slicedValues) = writtenBytesAndValues(sliced)
+          val (freshBytes, freshValues) = writtenBytesAndValues(fresh)
+          assert(slicedValues == freshValues, s"values of skip=$skip, take=$take")
+          assert(slicedBytes == freshBytes, s"bytes of skip=$skip, take=$take")
+        } finally {
+          fresh.close()
+          sliced.close()
+        }
+      }
+      // The slice withRebasedOffsets made is released with the rest.
+      assert(allocator.getAllocatedMemory == 0)
+    } finally {
+      allocator.close()
+    }
+  }
+
   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
@@ -563,6 +598,80 @@ class UtilsSuite extends CometTestBase {
     vector
   }
 
+  private def lists(values: Seq[Seq[String]], allocator: BufferAllocator): ListVector = {
+    val list = ListVector.empty("l", allocator)
+    val child = list
+      .addOrGetVector[VarCharVector](FieldType.nullable(ArrowType.Utf8.INSTANCE))
+      .getVector
+    child.allocateNew()
+    var position = 0
+    values.zipWithIndex.foreach { case (elements, i) =>
+      if (elements == null) {
+        list.setNull(i)
+      } else {
+        list.startNewValue(i)
+        elements.foreach { e =>
+          child.setSafe(position, e.getBytes(StandardCharsets.UTF_8))
+          position += 1
+        }
+        list.endValue(i, elements.length)
+      }
+    }
+    child.setValueCount(position)
+    list.setValueCount(values.length)
+    list
+  }
+
+  /**
+   * `values.slice(skip, skip + take)` in the shape the C Data import gives a slice of a native
+   * list array: the offsets are the original ones from `skip` on, over a child that still holds
+   * the elements of every list. `skip` is a multiple of 8, so the validity bitmap slices on a
+   * byte boundary.
+   */
+  private def slicedLists(
+      values: Seq[Seq[String]],
+      skip: Int,
+      take: Int,
+      allocator: BufferAllocator): ListVector = {
+    val full = lists(values, allocator)
+    try {
+      val sliced =
+        new ListVector("l", allocator, FieldType.nullable(ArrowType.List.INSTANCE), null)
+      sliced.initializeChildrenFromFields(full.getField.getChildren)
+      sliced.loadFieldBuffers(
+        new ArrowFieldNode(take.toLong, values.slice(skip, skip + take).count(_ == null).toLong),
+        Seq(
+          full.getValidityBuffer.slice(skip / 8L, (take + 7) / 8L),
+          full.getOffsetBuffer.slice(skip * 4L, (take + 1) * 4L)).asJava)
+      val child = full.getDataVector
+      sliced.getDataVector.loadFieldBuffers(
+        new ArrowFieldNode(child.getValueCount.toLong, child.getNullCount.toLong),
+        child.getFieldBuffers)
+      sliced
+    } finally {
+      full.close()
+    }
+  }
+
+  /**
+   * The size of the body `vector` is written with after `withRebasedOffsets`, and the values that
+   * body holds.
+   */
+  private def writtenBytesAndValues(vector: FieldVector): (Long, Seq[String]) = {
+    val root =
+      new VectorSchemaRoot(Seq(vector.getField).asJava, Seq(vector).asJava, vector.getValueCount)
+    Utils.withRebasedOffsets(root) { written =>
+      val batch = new VectorUnloader(written).getRecordBatch
+      try {
+        val column = written.getVector(0)
+        val values = (0 until column.getValueCount).map(i => String.valueOf(column.getObject(i)))
+        (batch.computeBodyLength(), values)
+      } finally {
+        batch.close()
+      }
+    }
+  }
+
   /**
    * `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