[GLUTEN-10892][VL] Use veloxPreferredBatchBytes to control the max size of memory of batches combined

backends-velox/src/main/java/org/apache/gluten/utils/VeloxBatchResizer.java

@@ -28,14 +28,18 @@
 
 public final class VeloxBatchResizer {
   public static ColumnarBatchOutIterator create(
-      int minOutputBatchSize, int maxOutputBatchSize, Iterator<ColumnarBatch> in) {
+      int minOutputBatchSize,
+      int maxOutputBatchSize,
+      long preferredBatchBytes,
+      Iterator<ColumnarBatch> in) {
     final Runtime runtime =
         Runtimes.contextInstance(BackendsApiManager.getBackendName(), "VeloxBatchResizer");
     long outHandle =
         VeloxBatchResizerJniWrapper.create(runtime)
             .create(
                 minOutputBatchSize,
                 maxOutputBatchSize,
+                preferredBatchBytes,
                 new ColumnarBatchInIterator(BackendsApiManager.getBackendName(), in));
     return new ColumnarBatchOutIterator(runtime, outHandle);
   }

backends-velox/src/main/java/org/apache/gluten/utils/VeloxBatchResizerJniWrapper.java

@@ -37,5 +37,8 @@ public long rtHandle() {
   }
 
   public native long create(
-      int minOutputBatchSize, int maxOutputBatchSize, ColumnarBatchInIterator itr);
+      int minOutputBatchSize,
+      int maxOutputBatchSize,
+      long preferredBatchBytes,
+      ColumnarBatchInIterator itr);
 }

backends-velox/src/main/scala/org/apache/gluten/execution/VeloxResizeBatchesExec.scala

@@ -35,11 +35,14 @@ import scala.collection.JavaConverters._
 case class VeloxResizeBatchesExec(
     override val child: SparkPlan,
     minOutputBatchSize: Int,
-    maxOutputBatchSize: Int)
+    maxOutputBatchSize: Int,
+    preferredBatchBytes: Long)
   extends ColumnarToColumnarExec(child) {
 
   override protected def mapIterator(in: Iterator[ColumnarBatch]): Iterator[ColumnarBatch] = {
-    VeloxBatchResizer.create(minOutputBatchSize, maxOutputBatchSize, in.asJava).asScala
+    VeloxBatchResizer
+      .create(minOutputBatchSize, maxOutputBatchSize, preferredBatchBytes, in.asJava)
+      .asScala
   }
 
   override protected def closeIterator(out: Iterator[ColumnarBatch]): Unit = {

backends-velox/src/main/scala/org/apache/gluten/extension/AppendBatchResizeForShuffleInputAndOutput.scala

@@ -40,19 +40,20 @@ case class AppendBatchResizeForShuffleInputAndOutput() extends Rule[SparkPlan] {
     }
 
     val range = VeloxConfig.get.veloxResizeBatchesShuffleInputOutputRange
+    val preferredBatchBytes = VeloxConfig.get.veloxPreferredBatchBytes
     plan.transformUp {
       case shuffle: ColumnarShuffleExchangeExec
           if resizeBatchesShuffleInputEnabled &&
             shuffle.shuffleWriterType.requiresResizingShuffleInput =>
         val appendBatches =
-          VeloxResizeBatchesExec(shuffle.child, range.min, range.max)
+          VeloxResizeBatchesExec(shuffle.child, range.min, range.max, preferredBatchBytes)
         shuffle.withNewChildren(Seq(appendBatches))
       case a @ AQEShuffleReadExec(
             ShuffleQueryStageExec(_, shuffle: ColumnarShuffleExchangeExec, _),
             _)
           if resizeBatchesShuffleOutputEnabled &&
             shuffle.shuffleWriterType.requiresResizingShuffleOutput =>
-        VeloxResizeBatchesExec(a, range.min, range.max)
+        VeloxResizeBatchesExec(a, range.min, range.max, preferredBatchBytes)
       case a @ AQEShuffleReadExec(
             ShuffleQueryStageExec(
               _,
@@ -61,42 +62,44 @@ case class AppendBatchResizeForShuffleInputAndOutput() extends Rule[SparkPlan] {
             _)
           if resizeBatchesShuffleOutputEnabled &&
             shuffle.shuffleWriterType.requiresResizingShuffleOutput =>
-        VeloxResizeBatchesExec(a, range.min, range.max)
+        VeloxResizeBatchesExec(a, range.min, range.max, preferredBatchBytes)
       // Since it's transformed in a bottom to up order, so we may first encounter
       // ShuffeQueryStageExec, which is transformed to VeloxResizeBatchesExec(ShuffeQueryStageExec),
       // then we see AQEShuffleReadExec
       case a @ AQEShuffleReadExec(
             VeloxResizeBatchesExec(
               s @ ShuffleQueryStageExec(_, shuffle: ColumnarShuffleExchangeExec, _),
               _,
+              _,
               _),
             _)
           if resizeBatchesShuffleOutputEnabled &&
             shuffle.shuffleWriterType.requiresResizingShuffleOutput =>
-        VeloxResizeBatchesExec(a.copy(child = s), range.min, range.max)
+        VeloxResizeBatchesExec(a.copy(child = s), range.min, range.max, preferredBatchBytes)
       case a @ AQEShuffleReadExec(
             VeloxResizeBatchesExec(
               s @ ShuffleQueryStageExec(
                 _,
                 ReusedExchangeExec(_, shuffle: ColumnarShuffleExchangeExec),
                 _),
               _,
+              _,
               _),
             _)
           if resizeBatchesShuffleOutputEnabled &&
             shuffle.shuffleWriterType.requiresResizingShuffleOutput =>
-        VeloxResizeBatchesExec(a.copy(child = s), range.min, range.max)
+        VeloxResizeBatchesExec(a.copy(child = s), range.min, range.max, preferredBatchBytes)
       case s @ ShuffleQueryStageExec(_, shuffle: ColumnarShuffleExchangeExec, _)
           if resizeBatchesShuffleOutputEnabled &&
             shuffle.shuffleWriterType.requiresResizingShuffleOutput =>
-        VeloxResizeBatchesExec(s, range.min, range.max)
+        VeloxResizeBatchesExec(s, range.min, range.max, preferredBatchBytes)
       case s @ ShuffleQueryStageExec(
             _,
             ReusedExchangeExec(_, shuffle: ColumnarShuffleExchangeExec),
             _)
           if resizeBatchesShuffleOutputEnabled &&
             shuffle.shuffleWriterType.requiresResizingShuffleOutput =>
-        VeloxResizeBatchesExec(s, range.min, range.max)
+        VeloxResizeBatchesExec(s, range.min, range.max, preferredBatchBytes)
     }
   }
 }

backends-velox/src/main/scala/org/apache/gluten/extension/CudfNodeValidationRule.scala

@@ -35,7 +35,7 @@ case class CudfNodeValidationRule(glutenConf: GlutenConfig) extends Rule[SparkPl
     val transformedPlan = plan.transformUp {
       case shuffle @ ColumnarShuffleExchangeExec(
             _,
-            VeloxResizeBatchesExec(w: WholeStageTransformer, _, _),
+            VeloxResizeBatchesExec(w: WholeStageTransformer, _, _, _),
             _,
             _,
             _) =>

backends-velox/src/main/scala/org/apache/gluten/extension/GpuBufferBatchResizeForShuffleInputOutput.scala

@@ -34,13 +34,14 @@ case class GpuBufferBatchResizeForShuffleInputOutput() extends Rule[SparkPlan] {
       return plan
     }
     val range = VeloxConfig.get.veloxResizeBatchesShuffleInputOutputRange
+    val preferredBatchBytes = VeloxConfig.get.veloxPreferredBatchBytes
     val batchSize = VeloxConfig.get.cudfBatchSize
     plan.transformUp {
       case shuffle: ColumnarShuffleExchangeExec
           if shuffle.shuffleWriterType == HashShuffleWriterType &&
             VeloxConfig.get.veloxResizeBatchesShuffleInput =>
         val appendBatches =
-          VeloxResizeBatchesExec(shuffle.child, range.min, range.max)
+          VeloxResizeBatchesExec(shuffle.child, range.min, range.max, preferredBatchBytes)
         shuffle.withNewChildren(Seq(appendBatches))
       case a @ AQEShuffleReadExec(
             ShuffleQueryStageExec(_, _: ColumnarShuffleExchangeExecBase, _),