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

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);
+  }
 }

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

@@ -24,6 +24,7 @@
 import io.substrait.proto.RelCommon;
 
 import java.io.Serializable;
+import java.util.Collections;
 import java.util.List;
 
 public class GenerateRelNode implements RelNode, Serializable {
@@ -81,4 +82,9 @@ public Rel toProtobuf() {
     relBuilder.setGenerate(generateRelBuilder.build());
     return relBuilder.build();
   }
+
+  @Override
+  public List<RelNode> childNodes() {
+    return Collections.singletonList(input);
+  }
 }

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

@@ -17,12 +17,20 @@
 package org.apache.gluten.substrait.rel;
 
 import org.apache.gluten.expression.ConverterUtils;
+import org.apache.gluten.substrait.extensions.AdvancedExtensionNode;
+import org.apache.gluten.substrait.extensions.ExtensionBuilder;
 import org.apache.gluten.substrait.type.TypeNode;
 
+import com.google.protobuf.Any;
+import com.google.protobuf.StringValue;
 import io.substrait.proto.*;
+import org.apache.spark.sql.execution.adaptive.InputStats;
 
+import java.util.ArrayList;
 import java.util.List;
 
+import scala.math.BigInt;
+
 /**
  * The relation for input iterator, e.g., the input is ColumnarShuffleExchange or RowToColumnarExec.
  * It uses `ReadRel` as the substrait rel type.
@@ -32,13 +40,32 @@ public class InputIteratorRelNode implements RelNode {
   private final List<TypeNode> types;
   private final List<String> names;
   private final Long iteratorIndex;
+  private BigInt rowCount;
+
+  private InputStats inputStats;
 
   public InputIteratorRelNode(List<TypeNode> types, List<String> names, Long iteratorIndex) {
     this.types = types;
     this.names = names;
     this.iteratorIndex = iteratorIndex;
   }
 
+  public BigInt getRowCount() {
+    return rowCount;
+  }
+
+  public void setRowCount(BigInt rowCount) {
+    this.rowCount = rowCount;
+  }
+
+  public InputStats getInputStats() {
+    return inputStats;
+  }
+
+  public void setInputStats(InputStats inputStats) {
+    this.inputStats = inputStats;
+  }
+
   @Override
   public Rel toProtobuf() {
     Type.Struct.Builder structBuilder = Type.Struct.newBuilder();
@@ -59,8 +86,22 @@ public Rel toProtobuf() {
         LocalFilesBuilder.makeLocalFiles(ConverterUtils.ITERATOR_PREFIX() + iteratorIndex);
     readBuilder.setLocalFiles(iteratorIndexNode.toProtobuf());
 
+    if (null != rowCount) {
+      Any inputRowCount =
+          Any.pack(
+              StringValue.newBuilder().setValue("rowSize=" + rowCount.toLong() + "\n").build());
+      AdvancedExtensionNode advancedExtension =
+          ExtensionBuilder.makeAdvancedExtension(inputRowCount, null);
+      readBuilder.setAdvancedExtension(advancedExtension.toProtobuf());
+    }
+
     Rel.Builder builder = Rel.newBuilder();
     builder.setRead(readBuilder.build());
     return builder.build();
   }
+
+  @Override
+  public List<RelNode> childNodes() {
+    return new ArrayList<>();
+  }
 }

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

@@ -24,6 +24,8 @@
 import io.substrait.proto.RelCommon;
 
 import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.List;
 
 public class JoinRelNode implements RelNode, Serializable {
   private final RelNode left;
@@ -79,4 +81,12 @@ public Rel toProtobuf() {
 
     return Rel.newBuilder().setJoin(joinBuilder.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/ProjectRelNode.java

@@ -25,6 +25,7 @@
 
 import java.io.Serializable;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 
 public class ProjectRelNode implements RelNode, Serializable {
@@ -78,4 +79,9 @@ public Rel toProtobuf() {
     builder.setProject(projectBuilder.build());
     return builder.build();
   }
+
+  @Override
+  public List<RelNode> childNodes() {
+    return Collections.singletonList(input);
+  }
 }