[ARROW] Arrow serialization should not introduce extra shuffle for outermost limit

[cfmcgrady]

Why are the changes needed?

The fundamental concept is to execute a job similar to the way in which CollectLimitExec.executeCollect() operates.

select * from parquet.`parquet/tpcds/sf1000/catalog_sales` limit 20;

Before this PR:

After this PR:

How was this patch tested?

• Add some test cases that check the changes thoroughly including negative and positive cases if possible

• Add screenshots for manual tests if appropriate

• Run test locally before make a pull request

[cfmcgrady]

the AdaptiveSparkPlanExec.finalPhysicalPlan function was introduced in SPARK-41914, and it may present compatibility issues if the underlying Spark runtime lacks the corresponding patch.

[ulysses-you]

shall we reflect some related private method to workaround ? it's unacceptable if we break the compatibility.

[cfmcgrady]

have changed to reflective call function adaptiveSparkPlan.finalPhysicalPlan.

[cfmcgrady]

add offset support in the separate PR to adapt Spark-3.4.0

[cfmcgrady]

cc @yaooqinn @ulysses-you @pan3793 @turboFei @cxzl25

[ulysses-you]

can we unfiy the class name a bit more ? not see the difference between xxxutils and xxx helper... logically, most of those methods should be private.

[ulysses-you]

in general, we can infer row size by schema.defaultSize

[cfmcgrady]

sorry for the lack of documentation.
This class ArrowBatchIterator is derived from org.apache.spark.sql.execution.arrow.ArrowConverters.ArrowBatchWithSchemaIterator, with two key differences:

1. there is no requirement to write the schema at the batch header
2. iteration halts when rowCount equals limit

[cfmcgrady]

here is the diff, compare with latest spark master branch https://github.com/apache/spark/blob/3c189abd73afa998e8573cbfdaf0f72445284314/sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowConverters.scala

-  private[sql] class ArrowBatchWithSchemaIterator(
+  private[sql] class ArrowBatchIterator(
       rowIter: Iterator[InternalRow],
       schema: StructType,
       maxRecordsPerBatch: Long,
       maxEstimatedBatchSize: Long,
+      limit: Long,
       timeZoneId: String,
       context: TaskContext)
-    extends ArrowBatchIterator(
-      rowIter, schema, maxRecordsPerBatch, timeZoneId, context) {
+    extends Iterator[Array[Byte]] {
+

-    private val arrowSchemaSize = SizeEstimator.estimate(arrowSchema)
     var rowCountInLastBatch: Long = 0
+    var rowCount: Long = 0

     override def next(): Array[Byte] = {
       val out = new ByteArrayOutputStream()
       val writeChannel = new WriteChannel(Channels.newChannel(out))

       rowCountInLastBatch = 0
-      var estimatedBatchSize = arrowSchemaSize
+      var estimatedBatchSize = 0
       Utils.tryWithSafeFinally {
-        // Always write the schema.
-        MessageSerializer.serialize(writeChannel, arrowSchema)

         // Always write the first row.
         while (rowIter.hasNext && (
@@ -31,15 +30,17 @@
             estimatedBatchSize < maxEstimatedBatchSize ||
             // If the size of rows are 0 or negative, unlimit it.
             maxRecordsPerBatch <= 0 ||
-            rowCountInLastBatch < maxRecordsPerBatch)) {
+            rowCountInLastBatch < maxRecordsPerBatch ||
+            rowCount < limit)) {
           val row = rowIter.next()
           arrowWriter.write(row)
           estimatedBatchSize += (row match {
             case ur: UnsafeRow => ur.getSizeInBytes
-            // Trying to estimate the size of the current row, assuming 16 bytes per value.
-            case ir: InternalRow => ir.numFields * 16
+            // Trying to estimate the size of the current row
+            case _: InternalRow => schema.defaultSize
           })
           rowCountInLastBatch += 1
+          rowCount += 1
         }
         arrowWriter.finish()
         val batch = unloader.getRecordBatch()

[ulysses-you]

shall we reflect some related private method to workaround ? it's unacceptable if we break the compatibility.

[codecov-commenter]

Codecov Report

Merging #4662 (82c912e) into master (d9e14f2) will increase coverage by 0.25%.
The diff coverage is 85.91%.

@@             Coverage Diff              @@
##             master    #4662      +/-   ##
============================================
+ Coverage     57.60%   57.85%   +0.25%     
  Complexity       13       13              
============================================
  Files           579      580       +1     
  Lines         31951    32212     +261     
  Branches       4269     4304      +35     
============================================
+ Hits          18404    18635     +231     
+ Misses        11785    11780       -5     
- Partials       1762     1797      +35     

Impacted Files	Coverage Δ
...uubi/engine/spark/operation/ExecuteStatement.scala	82.60% <50.00%> (+1.80%)	⬆️
...g/apache/spark/sql/kyuubi/SparkDatasetHelper.scala	80.00% <82.60%> (+1.42%)	⬆️
...rk/sql/execution/arrow/KyuubiArrowConverters.scala	88.57% <88.57%> (ø)

... and 23 files with indirect coverage changes

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[cfmcgrady]

can we unfiy the class name a bit more ? not see the difference between xxxutils and xxx helper... logically, most of those methods should be private.

have completed the refactoring, please take another look when you have time. Thank you. @ulysses-you

[ulysses-you]

thank you @cfmcgrady , lgtm if tests pass

[pan3793]

The name "rootAllocator" is not suitable then, and why should we create a child allocator?

[cfmcgrady]

I think it's for debugging.

Memory was leaked by query. Memory leaked: (128)
Allocator(slice) 0/128/128/9223372036854775807 (res/actual/peak/limit)

java.lang.IllegalStateException: Memory was leaked by query. Memory leaked: (128)
Allocator(slice) 0/128/128/9223372036854775807 (res/actual/peak/limit)

	at org.apache.arrow.memory.BaseAllocator.close(BaseAllocator.java:437)
	at org.apache.spark.sql.execution.arrow.KyuubiArrowConverters$.slice(KyuubiArrowConverters.scala:91)
	at org.apache.spark.sql.kyuubi.SparkDatasetHelper$.doCollectLimit(SparkDatasetHelper.scala:170)
	at org.apache.spark.sql.kyuubi.SparkDatasetHelper$.$anonfun$executeArrowBatchCollect$1(SparkDatasetHelper.scala:51)
	at org.apache.spark.sql.kyuubi.SparkDatasetHelper$

[pan3793]

LGTM if CI pass

[turboFei]

Great work， thanks

[ulysses-you]

thanks, merging to master/branch-1.7