diff --git a/python/pyspark/sql/tests/arrow/test_arrow_map.py b/python/pyspark/sql/tests/arrow/test_arrow_map.py
index 853ba731f1df8..e08b65e759dd5 100644
--- a/python/pyspark/sql/tests/arrow/test_arrow_map.py
+++ b/python/pyspark/sql/tests/arrow/test_arrow_map.py
@@ -151,7 +151,8 @@ def func(iterator):
 
         df = self.spark.range(10)
         actual = df.repartition(1).mapInArrow(func, "a long").collect()
-        self.assertEqual(set((r.a for r in actual)), set(range(100)))
+        # Each input batch produces 0..99 once, in order, since there is one partition.
+        self.assertEqual([r.a for r in actual], list(range(100)) * max(1, len(actual) // 100))
 
     def test_other_than_recordbatch_iter(self):
         with self.quiet():
diff --git a/python/pyspark/sql/tests/arrow/test_arrow_python_udf.py b/python/pyspark/sql/tests/arrow/test_arrow_python_udf.py
index 417f1082cdc45..43622f61b18b5 100644
--- a/python/pyspark/sql/tests/arrow/test_arrow_python_udf.py
+++ b/python/pyspark/sql/tests/arrow/test_arrow_python_udf.py
@@ -385,6 +385,14 @@ def create_struct_with_interval(interval_val, name_val):
         self.assertIsNotNone(result["interval_field"])
         self.assertEqual(result["name"], "test_name")
 
+    def test_udf_with_output_batch_slicing(self):
+        # Each result row is joined with the next queued input row, so the sliced 10-row result
+        # batch must be output once and in order.
+        df = self.spark.range(0, 10, 1, 1)
+        with self.sql_conf({"spark.sql.execution.arrow.maxRecordsPerOutputBatch": "3"}):
+            result = df.select(col("id"), udf(lambda x: x + 1, "long")(col("id")))
+            assertDataFrameEqual(result, [Row(i, i + 1) for i in range(10)])
+
 
 @unittest.skipIf(
     not have_pandas or not have_pyarrow, pandas_requirement_message or pyarrow_requirement_message
diff --git a/python/pyspark/sql/tests/arrow/test_arrow_python_udf_cached.py b/python/pyspark/sql/tests/arrow/test_arrow_python_udf_cached.py
index ba14a9b740757..196c172de54e8 100644
--- a/python/pyspark/sql/tests/arrow/test_arrow_python_udf_cached.py
+++ b/python/pyspark/sql/tests/arrow/test_arrow_python_udf_cached.py
@@ -133,6 +133,22 @@ def test_udf_on_cached_numeric_column_alongside_unselected_string(self):
         finally:
             df.unpersist()
 
+    def test_udf_on_cached_input_with_output_batch_slicing(self):
+        # The columnar path pairs each result batch with one input batch, so the 10-row result
+        # of the single cached batch must not be sliced under the configs below.
+        df = self.spark.range(0, 10, 1, 1).cache()
+        try:
+            for conf in [
+                {"spark.sql.execution.arrow.maxRecordsPerOutputBatch": "3"},
+                {"spark.sql.execution.arrow.maxBytesPerOutputBatch": "3"},
+            ]:
+                with self.subTest(conf=conf), self.sql_conf(conf):
+                    result = df.select(col("id"), udf(lambda x: x + 1, "long")(col("id")))
+                    self._assert_columnar_arrow_eval(result, expected_nodes=1)
+                    assertDataFrameEqual(result, [Row(i, i + 1) for i in range(10)])
+        finally:
+            df.unpersist()
+
     def test_udf_on_cached_map_column(self):
         # Canonical map vectors from the cache pass the gate (their entry-struct children carry
         # the canonical "key"/"value" names in order) and ride the columnar path verbatim. The
diff --git a/python/pyspark/sql/tests/pandas/test_pandas_grouped_map_with_state.py b/python/pyspark/sql/tests/pandas/test_pandas_grouped_map_with_state.py
index 38f0933c35302..dd58f04064bd4 100644
--- a/python/pyspark/sql/tests/pandas/test_pandas_grouped_map_with_state.py
+++ b/python/pyspark/sql/tests/pandas/test_pandas_grouped_map_with_state.py
@@ -93,10 +93,14 @@ def prepare_test_resource():
             .start()
         )
 
-        self.assertEqual(q.name, "this_query")
-        self.assertTrue(q.isActive)
-        q.processAllAvailable()
-        self.assertTrue(q.exception() is None)
+        try:
+            self.assertEqual(q.name, "this_query")
+            self.assertTrue(q.isActive)
+            q.processAllAvailable()
+            self.assertTrue(q.exception() is None)
+        finally:
+            q.stop()
+            shutil.rmtree(input_path, ignore_errors=True)
 
     def test_apply_in_pandas_with_state_basic(self):
         def func(key, pdf_iter, state):
@@ -194,6 +198,30 @@ def check_results(batch_df, _):
 
         self._test_apply_in_pandas_with_state_basic(func, check_results)
 
+    def test_apply_in_pandas_with_state_output_batch_slicing(self):
+        # The worker packs the data rows (5 per group) and state rows (1 per group) of the groups
+        # in a partition into one output batch. Its data and state counts cover the whole batch,
+        # so it must not be sliced by the configs below.
+        def func(key, pdf_iter, state):
+            state.update((1,))
+            yield pd.DataFrame(
+                {"key": [f"{key[0]}_{i}" for i in range(5)], "countAsString": list("01234")}
+            )
+
+        expected = sorted(
+            Row(key=f"{k}_{i}", countAsString=str(i)) for k in ["hello", "this"] for i in range(5)
+        )
+
+        def check_results(batch_df, _):
+            assert sorted(batch_df.collect()) == expected
+
+        for conf in [
+            {"spark.sql.execution.arrow.maxRecordsPerOutputBatch": "3"},
+            {"spark.sql.execution.arrow.maxBytesPerOutputBatch": "3"},
+        ]:
+            with self.subTest(conf=conf), self.sql_conf(conf):
+                self._test_apply_in_pandas_with_state_basic(func, check_results)
+
     def test_apply_in_pandas_with_state_python_worker_random_failure(self):
         input_path = tempfile.mkdtemp()
         output_path = tempfile.mkdtemp()
diff --git a/python/pyspark/sql/tests/test_udtf.py b/python/pyspark/sql/tests/test_udtf.py
index 5e3df713ca3c7..24cfe5e9369e6 100644
--- a/python/pyspark/sql/tests/test_udtf.py
+++ b/python/pyspark/sql/tests/test_udtf.py
@@ -3786,6 +3786,27 @@ def eval(self):
                 ):
                     udtf(TestUDTF, returnType=ret_type)().collect()
 
+    def test_udtf_lateral_join_with_output_batch_slicing(self):
+        # Each output batch is joined with one input row, so it must not be sliced by the
+        # configs below.
+        @udtf(returnType="x: int")
+        class TestUDTF:
+            def eval(self, n: int):
+                for i in range(5):
+                    yield (n * 10 + i,)
+
+        self.spark.udtf.register("test_udtf_slicing", TestUDTF)
+        expected = [Row(id=n, x=n * 10 + i) for n in range(3) for i in range(5)]
+        for conf in [
+            {"spark.sql.execution.arrow.maxRecordsPerOutputBatch": "3"},
+            {"spark.sql.execution.arrow.maxBytesPerOutputBatch": "3"},
+        ]:
+            with self.subTest(conf=conf), self.sql_conf(conf):
+                df = self.spark.sql(
+                    "SELECT t.id, u.x FROM range(0, 3, 1, 1) t, LATERAL test_udtf_slicing(t.id) u"
+                )
+                assertDataFrameEqual(df, expected)
+
 
 class LegacyUDTFArrowTests(LegacyUDTFArrowTestsMixin, ReusedSQLTestCase):
     @classmethod
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
index 6af6dfd9cd6cf..ba1d0667f515e 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
@@ -5452,7 +5452,9 @@ object SQLConf {
         "there is no limit. Note that the complete ArrowRecordBatch is actually created but " +
         "the number of records is limited when sending it to the downstream operator. This is " +
         "used to avoid large batches being sent to the downstream operator including " +
-        "the columnar-based operator implemented by third-party libraries.")
+        "the columnar-based operator implemented by third-party libraries. " +
+        "applyInPandasWithState, Arrow Python UDTFs and Arrow Python UDFs that take " +
+        "columnar input directly ignore this config.")
       .version("4.1.0")
       .internal()
       .intConf
@@ -5468,7 +5470,8 @@ object SQLConf {
         "the columnar-based operator implemented by third-party libraries. Spark will try to " +
         "create batches with the size equal or less than this value. Normally it should not " +
         "happen, but if in extreme case that even one record is still larger than this value, " +
-        "Spark will create a batch with one record.")
+        "Spark will create a batch with one record. applyInPandasWithState, Arrow Python " +
+        "UDTFs and Arrow Python UDFs that take columnar input directly ignore this config.")
       .version("4.1.0")
       .internal()
       .bytesConf(ByteUnit.BYTE)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonUDTFRunner.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonUDTFRunner.scala
index deceed02abb9a..f63a28effb7c5 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonUDTFRunner.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonUDTFRunner.scala
@@ -51,6 +51,11 @@ class ArrowPythonUDTFRunner(
   with BasicPythonArrowOutput {
   ArrowUtils.failDuplicatedFieldNames(schema)
 
+  // `EvalPythonUDTFExec` joins each output batch with one input row, so the batches returned
+  // from Python must not be sliced.
+  override protected def arrowMaxRecordsPerOutputBatch: Int = -1
+  override protected def arrowMaxBytesPerOutputBatch: Int = -1
+
   override protected def runnerConf: Map[String, String] = super.runnerConf ++ pythonRunnerConf
 
   override protected def evalConf: Map[String, String] = {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonRunner.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonRunner.scala
index de19b43d32c1d..81f78399784bf 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonRunner.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonRunner.scala
@@ -52,6 +52,11 @@ private[python] class ColumnarArrowPythonWithNamedArgumentRunner(
   with ColumnarArrowPythonInput
   with BasicPythonArrowOutput {
 
+  // `ColumnarArrowEvalPythonEvaluatorFactory.evalArrowColumnar` pairs each output batch with
+  // one input batch, so the batches returned from Python must not be sliced.
+  override protected def arrowMaxRecordsPerOutputBatch: Int = -1
+  override protected def arrowMaxBytesPerOutputBatch: Int = -1
+
   override protected def runnerConf: Map[String, String] = super.runnerConf ++ pythonRunnerConf
 
   // The input schema (and per-UDF nesting depth) travel in evalConf, exactly like
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala
index 1cf2664f67b35..0fbf3240bf1bb 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala
@@ -44,6 +44,9 @@ private[python] trait PythonArrowOutput[OUT <: AnyRef] { self: BasePythonRunner[
 
   protected def deserializeColumnarBatch(batch: ColumnarBatch, schema: StructType): OUT
 
+  // The limits for slicing the batches returned from Python. A runner whose consumer cannot take
+  // a batch split into several, e.g. one that pairs each batch with one input row, overrides both
+  // to -1, as `ArrowPythonUDTFRunner` does.
   protected def arrowMaxRecordsPerOutputBatch: Int = SQLConf.get.arrowMaxRecordsPerOutputBatch
 
   protected def arrowMaxBytesPerOutputBatch: Int = SQLConf.get.arrowMaxBytesPerOutputBatch.toInt
@@ -210,8 +213,7 @@ abstract class BaseSliceArrowOutputProcessor(
   extends ArrowOutputProcessorImpl(reader, pythonMetrics) {
 
   protected var currentRowIdx = -1
-  protected var prevRoot: VectorSchemaRoot = null
-  protected var prevVectors: Array[ColumnVector] = _
+  protected var prevRoot: VectorSchemaRoot = root
 
   override def loadBatch(): Boolean = {
     if (rowCount > 0 && currentRowIdx < rowCount) {
@@ -225,34 +227,34 @@ abstract class BaseSliceArrowOutputProcessor(
 
   override def produceBatch(): ColumnarBatch = {
     val batchRoot = getRoot
+    closePrevSlice()
+    prevRoot = batchRoot
 
-    if (batchRoot != prevRoot) {
-      if (prevRoot != null) {
-        prevVectors.foreach(_.close())
-        prevRoot.close()
-      }
-      prevRoot = batchRoot
-    }
-
-    val vectors = getVectors(batchRoot)
-    prevVectors = vectors
-
-    val batch = new ColumnarBatch(vectors)
+    val batch = new ColumnarBatch(getVectors(batchRoot))
     batch.setNumRows(batchRoot.getRowCount)
     batch
   }
 
   protected override def getVectors(root: VectorSchemaRoot): Array[ColumnVector] = {
-    root.getFieldVectors.asScala.map { vector =>
-      new ArrowColumnVector(vector)
-    }.toArray[ColumnVector]
+    if (root eq this.root) {
+      super.getVectors(root)
+    } else {
+      root.getFieldVectors.asScala.map { vector =>
+        new ArrowColumnVector(vector)
+      }.toArray[ColumnVector]
+    }
   }
 
-  override def close(): Unit = {
-    if (prevRoot != null) {
-      prevVectors.foreach(_.close())
+  // Releases the previous output batch if it is a slice. The reader loads every batch from Python
+  // into the same root, so it must not be closed here; `ArrowOutputProcessorImpl.close` closes it.
+  private def closePrevSlice(): Unit = {
+    if (prevRoot ne root) {
       prevRoot.close()
     }
+  }
+
+  override def close(): Unit = {
+    closePrevSlice()
     super.close()
   }
 }
@@ -264,7 +266,8 @@ class SliceBytesArrowOutputProcessorImpl(
   extends BaseSliceArrowOutputProcessor(reader, pythonMetrics) {
 
   protected override def getRoot: VectorSchemaRoot = {
-    if (getBatchBytes(root) < arrowMaxBytesPerOutputBatch) {
+    // Output the whole root only if no row of this batch has been output yet.
+    if (currentRowIdx == 0 && getBatchBytes(root) < arrowMaxBytesPerOutputBatch) {
       currentRowIdx = currentRowIdx + root.getRowCount
       root
     } else {
@@ -313,14 +316,11 @@ class SliceRecordsArrowOutputProcessorImpl(
   extends BaseSliceArrowOutputProcessor(reader, pythonMetrics) {
 
   protected override def getRoot: VectorSchemaRoot = {
-    val remainingRows = rowCount - currentRowIdx
-    val rootSlice = if (remainingRows > arrowMaxRecordsPerOutputBatch) {
-      root.slice(currentRowIdx, arrowMaxRecordsPerOutputBatch)
-    } else {
-      root
-    }
+    val numRows = math.min(rowCount - currentRowIdx, arrowMaxRecordsPerOutputBatch)
+    // Output the whole root only if this chunk is the whole batch.
+    val rootSlice = if (numRows == rowCount) root else root.slice(currentRowIdx, numRows)
 
-    currentRowIdx = currentRowIdx + rootSlice.getRowCount
+    currentRowIdx = currentRowIdx + numRows
 
     rootSlice
   }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/ApplyInPandasWithStatePythonRunner.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/ApplyInPandasWithStatePythonRunner.scala
index b512c99f7b6be..f58138c4b68e0 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/ApplyInPandasWithStatePythonRunner.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/ApplyInPandasWithStatePythonRunner.scala
@@ -72,6 +72,11 @@ class ApplyInPandasWithStatePythonRunner(
   with PythonArrowOutput[OutType] {
   ArrowUtils.failDuplicatedFieldNames(inputSchema)
 
+  // The data and state counts of each batch returned from Python cover the whole batch and are
+  // read from its first row, so the batch must not be sliced.
+  override protected def arrowMaxRecordsPerOutputBatch: Int = -1
+  override protected def arrowMaxBytesPerOutputBatch: Int = -1
+
   // Installed at worker launch. This runner is constructed per task, so each launch reads the
   // configuration that task carries. A running streaming query carries a snapshot:
   // `StreamExecution` runs its batches on `sparkSession.cloneSession()`, whose conf is copied at
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala
new file mode 100644
index 0000000000000..6ee8fcba9db97
--- /dev/null
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala
@@ -0,0 +1,151 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.python
+
+import java.io.{ByteArrayInputStream, ByteArrayOutputStream}
+
+import scala.collection.mutable.ArrayBuffer
+import scala.jdk.CollectionConverters._
+
+import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter}
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.execution.arrow.ArrowWriter
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.types._
+import org.apache.spark.sql.util.ArrowUtils
+import org.apache.spark.sql.vectorized.ArrowColumnVector
+import org.apache.spark.unsafe.types.UTF8String
+
+class PythonArrowOutputSuite extends SparkFunSuite {
+
+  private val schema = new StructType()
+    .add("id", LongType)
+    .add("s", StringType)
+    .add("st", new StructType().add("a", IntegerType).add("b", StringType))
+
+  private def row(i: Long): InternalRow = {
+    val s = if (i % 3 == 0) null else UTF8String.fromString(s"str_$i")
+    val st = if (i % 5 == 0) {
+      null
+    } else {
+      InternalRow(i.toInt * 2, if (i % 2 == 0) null else UTF8String.fromString(s"b_$i"))
+    }
+    InternalRow(i, s, st)
+  }
+
+  /** Writes an Arrow IPC stream with one record batch per entry of `batchSizes`. */
+  private def writeStream(batchSizes: Seq[Int]): Array[Byte] = {
+    val writer = ArrowWriter.create(schema, "UTC")
+    val out = new ByteArrayOutputStream()
+    try {
+      val streamWriter = new ArrowStreamWriter(writer.root, null, out)
+      streamWriter.start()
+      var next = 0L
+      batchSizes.foreach { size =>
+        (0 until size).foreach { _ =>
+          writer.write(row(next))
+          next += 1
+        }
+        writer.finish()
+        streamWriter.writeBatch()
+        writer.reset()
+      }
+      streamWriter.end()
+    } finally {
+      writer.root.close()
+    }
+    out.toByteArray
+  }
+
+  /**
+   * Drives `processor` the way `PythonArrowOutput.newReaderIterator.read()` does, and returns
+   * the rows and the buffer size in bytes of every output batch. Rows are copied before the
+   * next batch is produced, since the processor releases the previous output batch then.
+   */
+  private def readAll(
+      batchSizes: Seq[Int],
+      createProcessor: (ArrowStreamReader, Map[String, SQLMetric]) => ArrowOutputProcessor)
+    : Seq[(Seq[InternalRow], Long)] = {
+    val allocator = ArrowUtils.rootAllocator.newChildAllocator("reader", 0, Long.MaxValue)
+    val reader = new ArrowStreamReader(new ByteArrayInputStream(writeStream(batchSizes)), allocator)
+    val metrics = Map(
+      "pythonNumRowsReceived" -> new SQLMetric("sum"),
+      "pythonDataReceived" -> new SQLMetric("size"))
+    val processor = createProcessor(reader, metrics)
+    val output = ArrayBuffer.empty[(Seq[InternalRow], Long)]
+    while (processor.loadBatch()) {
+      val batch = processor.produceBatch()
+      val bytes = (0 until batch.numCols()).map { i =>
+        batch.column(i).asInstanceOf[ArrowColumnVector].getValueVector.getBufferSize.toLong
+      }.sum
+      output += ((batch.rowIterator().asScala.map(_.copy()).toSeq, bytes))
+    }
+    processor.close()
+    reader.close(false)
+    // Fails if any slice made by the processor was leaked.
+    allocator.close()
+    // `PythonArrowOutput` closes the processor again in its task completion listener.
+    processor.close()
+    assert(metrics("pythonNumRowsReceived").value === batchSizes.sum)
+    output.toSeq
+  }
+
+  private def checkAllRowsOnceInOrder(
+      batchSizes: Seq[Int],
+      output: Seq[(Seq[InternalRow], Long)]): Unit = {
+    assert(output.flatMap(_._1) === (0L until batchSizes.sum).map(row))
+  }
+
+  Seq(
+    (Seq(3, 3, 3), 2),
+    (Seq(4, 4, 4), 2),
+    (Seq(2, 5, 2, 7), 2),
+    (Seq(1, 0, 6, 1, 3), 3),
+    (Seq(10), 3)
+  ).foreach { case (batchSizes, maxRecords) =>
+    test(s"slice by records: batches $batchSizes, max $maxRecords records per batch") {
+      val output = readAll(batchSizes, (reader, metrics) =>
+        new SliceRecordsArrowOutputProcessorImpl(reader, metrics, maxRecords))
+      checkAllRowsOnceInOrder(batchSizes, output)
+      val expectedSizes = batchSizes.flatMap { n =>
+        val sizes = Seq.fill(n / maxRecords)(maxRecords) ++ Seq(n % maxRecords).filter(_ > 0)
+        if (sizes.isEmpty) Seq(0) else sizes
+      }
+      assert(output.map(_._1.size) === expectedSizes)
+    }
+  }
+
+  Seq(
+    (Seq(1, 100, 1), 512),
+    (Seq(100, 1, 100), 512),
+    (Seq(1, 0, 57, 2, 33), 512),
+    // Even a single row is larger than the limit, so every non-empty output batch has one row.
+    (Seq(2, 0, 5, 1), 3)
+  ).foreach { case (batchSizes, maxBytes) =>
+    test(s"slice by bytes: batches $batchSizes, max $maxBytes bytes per batch") {
+      val output = readAll(batchSizes, (reader, metrics) =>
+        new SliceBytesArrowOutputProcessorImpl(reader, metrics, maxBytes))
+      checkAllRowsOnceInOrder(batchSizes, output)
+      // Guards the choice of cases: each has a batch that does not fit, which must be sliced.
+      assert(output.size > batchSizes.size)
+      assert(output.forall { case (rows, bytes) => rows.size <= 1 || bytes < maxBytes })
+    }
+  }
+}