[GLUTEN-12465][CORE] Support optional stage InputStats plumbing in scan/input-iterator transformers (#12433)

* [GLUTEN][CORE] Support optional stage InputStats plumbing in scan/input-iterator transformers

Introduce an opt-in framework to pass stage input statistics (scan / shuffle /
broadcast) to the native engine as estimated row-size hints.

- Add InputStats / InputStatsKind / ShuffleStageWrapper and ApplyStageInputStatsRule.
- Add RelNode.childNode() with implementations across all rel nodes; carry
  rowSize / inputStats on ReadRelNode and InputIteratorRelNode; expose
  PlanNode.getRelNodes() and AdvancedExtensionNode optimization/enhancement getters.
- Add BasicScanExecTransformer.getInputStats hook and optional inputStats field on
  FileSourceScanExecTransformer; thread wsContext through genFirstStageIterator and
  the whole-stage RDDs; inject shuffle/broadcast stats in InputIteratorTransformer.
- Gate everything behind spark.gluten.sql.enablePassStageInputStats (default false),
  so velox / clickhouse behavior is unchanged.

Generated-by: TraeCli openrouter-3o
Co-Authored-By: Aime <aime@bytedance.com>
Change-Id: I7cd0a5abc98b6c41857230131513bfcb5ec6bf95

* change as request

---------

Co-authored-by: Aime <aime@bytedance.com>
Co-authored-by: liyang.127 <liyang.127@bytedance.com>

Author: 3 people

backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHIteratorApi.scala

@@ -277,7 +277,8 @@ class CHIteratorApi extends IteratorApi with Logging with LogLevelUtil {
       updateNativeMetrics: IMetrics => Unit,
       partitionIndex: Int,
       inputIterators: Seq[Iterator[ColumnarBatch]] = Seq(),
-      enableCudf: Boolean = false
+      enableCudf: Boolean = false,
+      wsContext: WholeStageTransformContext = null
   ): Iterator[ColumnarBatch] = {
 
     require(

backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxIteratorApi.scala

@@ -190,7 +190,8 @@ class VeloxIteratorApi extends IteratorApi with Logging {
       updateNativeMetrics: IMetrics => Unit,
       partitionIndex: Int,
       inputIterators: Seq[Iterator[ColumnarBatch]] = Seq(),
-      enableCudf: Boolean = false): Iterator[ColumnarBatch] = {
+      enableCudf: Boolean = false,
+      wsContext: WholeStageTransformContext = null): Iterator[ColumnarBatch] = {
     assert(
       inputPartition.isInstanceOf[GlutenPartition],
       "Velox backend only accept GlutenPartition.")

gluten-core/src/main/scala/org/apache/spark/sql/execution/adaptive/ShuffleStageWrapper.scala

@@ -0,0 +1,83 @@
+/*
+ * 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.adaptive
+
+import org.apache.spark.MapOutputStatistics
+import org.apache.spark.internal.Logging
+import org.apache.spark.sql.catalyst.trees.TreeNodeTag
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.exchange.ENSURE_REQUIREMENTS
+
+/**
+ * Represents stage input stats, eg scan or shuffle
+ *
+ * @param known
+ * @param sizeInBytes
+ * @param rowCount
+ * @param bytesByPartitionId
+ */
+sealed trait InputStatsKind
+case object ScanInputStats extends InputStatsKind
+case object BroadcastStats extends InputStatsKind
+case object ShuffleStats extends InputStatsKind
+case object UnknownStats extends InputStatsKind
+case class InputStats(
+    known: Boolean,
+    sizeInBytes: BigInt,
+    rowCount: BigInt,
+    bytesByPartitionId: Array[Long],
+    inputStatsKind: InputStatsKind)
+  extends Serializable {
+  override def toString: String = {
+    s"known : ($known) sizeInBytes:($sizeInBytes), rowCount: ($rowCount), " +
+      s"bytesByPartitionId:  [${bytesByPartitionId.mkString(",")}], statsKind: $inputStatsKind"
+  }
+}
+object ShuffleStageWrapper extends Logging {
+  val INPUT_STATS: TreeNodeTag[java.util.List[InputStats]] = TreeNodeTag("InputStats")
+  def unapply(plan: SparkPlan): Option[ShuffleQueryStageRuntimeStats] = plan match {
+    case s: ShuffleQueryStageExec
+        if s.isMaterialized && s.mapStats.isDefined &&
+          s.shuffle.shuffleOrigin == ENSURE_REQUIREMENTS =>
+      logInfo("hit InputIteratorTransformer ShuffleQueryStageExec with stats")
+      Some(
+        ShuffleQueryStageRuntimeStats(
+          s.getRuntimeStatistics.sizeInBytes,
+          s.getRuntimeStatistics.rowCount,
+          s.mapStats))
+    case a: AQEShuffleReadExec if a.child.isInstanceOf[ShuffleQueryStageExec] =>
+      logInfo("hit InputIteratorTransformer AQEShuffleReadExec with stats")
+      val queryStageExec = a.child.asInstanceOf[ShuffleQueryStageExec]
+      Some(
+        ShuffleQueryStageRuntimeStats(
+          queryStageExec.getRuntimeStatistics.sizeInBytes,
+          queryStageExec.getRuntimeStatistics.rowCount,
+          queryStageExec.mapStats))
+    case b: BroadcastQueryStageExec if b.isMaterialized =>
+      logInfo("hit InputIteratorTransformer BroadcastQueryStageExec with stats")
+      Some(
+        ShuffleQueryStageRuntimeStats(
+          b.getRuntimeStatistics.sizeInBytes,
+          b.getRuntimeStatistics.rowCount,
+          None))
+    case _ => None
+  }
+}
+case class ShuffleQueryStageRuntimeStats(
+    sizeInBytes: BigInt,
+    rowCount: Option[BigInt] = None,
+    mapOutputStatistics: Option[MapOutputStatistics] = None)

gluten-substrait/src/main/java/org/apache/gluten/substrait/extensions/AdvancedExtensionNode.java

@@ -40,6 +40,14 @@ public AdvancedExtensionNode(Any optimization, Any enhancement) {
     this.enhancement = enhancement;
   }
 
+  public Any getOptimization() {
+    return optimization;
+  }
+
+  public Any getEnhancement() {
+    return enhancement;
+  }
+
   public AdvancedExtension toProtobuf() {
     AdvancedExtension.Builder extensionBuilder = AdvancedExtension.newBuilder();
     if (optimization != null) {

gluten-substrait/src/main/java/org/apache/gluten/substrait/plan/PlanNode.java

@@ -77,4 +77,8 @@ public Plan toProtobuf() {
     }
     return planBuilder.build();
   }
+
+  public List<RelNode> getRelNodes() {
+    return relNodes;
+  }
 }

gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/AggregateRelNode.java

@@ -26,6 +26,7 @@
 
 import java.io.Serializable;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 
 public class AggregateRelNode implements RelNode, Serializable {
@@ -83,4 +84,9 @@ public Rel toProtobuf() {
     builder.setAggregate(aggBuilder.build());
     return builder.build();
   }
+
+  @Override
+  public List<RelNode> childNodes() {
+    return Collections.singletonList(input);
+  }
 }

gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/CrossRelNode.java

@@ -24,6 +24,8 @@
 import io.substrait.proto.RelCommon;
 
 import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.List;
 
 public class CrossRelNode implements RelNode, Serializable {
   private final RelNode left;
@@ -69,4 +71,12 @@ public Rel toProtobuf() {
     }
     return Rel.newBuilder().setCross(crossRelBuilder.build()).build();
   }
+
+  @Override
+  public List<RelNode> childNodes() {
+    List<RelNode> children = new ArrayList<>();
+    children.add(left);
+    children.add(right);
+    return children;
+  }
 }

gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/ExpandRelNode.java

@@ -25,6 +25,7 @@
 
 import java.io.Serializable;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 
 public class ExpandRelNode implements RelNode, Serializable {
@@ -76,4 +77,9 @@ public Rel toProtobuf() {
     builder.setExpand(expandBuilder.build());
     return builder.build();
   }
+
+  @Override
+  public List<RelNode> childNodes() {
+    return Collections.singletonList(input);
+  }
 }

gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/FetchRelNode.java

@@ -23,6 +23,8 @@
 import io.substrait.proto.RelCommon;
 
 import java.io.Serializable;
+import java.util.Collections;
+import java.util.List;
 
 public class FetchRelNode implements RelNode, Serializable {
   private final RelNode input;
@@ -66,4 +68,9 @@ public Rel toProtobuf() {
     relBuilder.setFetch(fetchRelBuilder.build());
     return relBuilder.build();
   }
+
+  @Override
+  public List<RelNode> childNodes() {
+    return Collections.singletonList(input);
+  }
 }

gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/FilterRelNode.java

@@ -24,6 +24,8 @@
 import io.substrait.proto.RelCommon;
 
 import java.io.Serializable;
+import java.util.Collections;
+import java.util.List;
 
 public class FilterRelNode implements RelNode, Serializable {
   private final RelNode input;
@@ -60,4 +62,9 @@ public Rel toProtobuf() {
     builder.setFilter(filterBuilder.build());
     return builder.build();
   }
+
+  @Override
+  public List<RelNode> childNodes() {
+    return Collections.singletonList(input);
+  }
 }