This is an automated email from the ASF dual-hosted git repository.
marin-ma pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gluten.git
The following commit(s) were added to refs/heads/main by this push:
new cd994505b0 [GLUTEN-12465][CORE] Support optional stage InputStats
plumbing in scan/input-iterator transformers (#12433)
cd994505b0 is described below
commit cd994505b08c6ae37ab6813a56c0139b226f4e66
Author: 李扬 <[email protected]>
AuthorDate: Wed Jul 8 01:20:52 2026 +0800
[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 <[email protected]>
Change-Id: I7cd0a5abc98b6c41857230131513bfcb5ec6bf95
* change as request
---------
Co-authored-by: Aime <[email protected]>
Co-authored-by: liyang.127 <[email protected]>
---
.../backendsapi/clickhouse/CHIteratorApi.scala | 3 +-
.../backendsapi/velox/VeloxIteratorApi.scala | 3 +-
.../execution/adaptive/ShuffleStageWrapper.scala | 83 ++++++++++++
.../extensions/AdvancedExtensionNode.java | 8 ++
.../org/apache/gluten/substrait/plan/PlanNode.java | 4 +
.../gluten/substrait/rel/AggregateRelNode.java | 6 +
.../apache/gluten/substrait/rel/CrossRelNode.java | 10 ++
.../apache/gluten/substrait/rel/ExpandRelNode.java | 6 +
.../apache/gluten/substrait/rel/FetchRelNode.java | 7 +
.../apache/gluten/substrait/rel/FilterRelNode.java | 7 +
.../gluten/substrait/rel/GenerateRelNode.java | 6 +
.../gluten/substrait/rel/InputIteratorRelNode.java | 41 ++++++
.../apache/gluten/substrait/rel/JoinRelNode.java | 10 ++
.../gluten/substrait/rel/ProjectRelNode.java | 6 +
.../apache/gluten/substrait/rel/ReadRelNode.java | 59 +++++++-
.../org/apache/gluten/substrait/rel/RelNode.java | 3 +
.../apache/gluten/substrait/rel/SetRelNode.java | 5 +
.../apache/gluten/substrait/rel/SortRelNode.java | 6 +
.../org/apache/gluten/substrait/rel/TopNNode.java | 6 +
.../substrait/rel/WindowGroupLimitRelNode.java | 6 +
.../apache/gluten/substrait/rel/WindowRelNode.java | 6 +
.../apache/gluten/substrait/rel/WriteRelNode.java | 6 +
.../apache/gluten/backendsapi/IteratorApi.scala | 3 +-
.../org/apache/gluten/config/GlutenConfig.scala | 11 ++
.../execution/BasicScanExecTransformer.scala | 10 +-
.../execution/FileSourceScanExecTransformer.scala | 11 +-
.../execution/GlutenWholeStageColumnarRDD.scala | 14 +-
.../gluten/execution/WholeStageTransformer.scala | 6 +-
.../execution/WholeStageZippedPartitionsRDD.scala | 11 +-
.../extension/ApplyStageInputStatsRule.scala | 148 +++++++++++++++++++++
.../ColumnarCollapseTransformStages.scala | 35 ++++-
31 files changed, 532 insertions(+), 14 deletions(-)
diff --git
a/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHIteratorApi.scala
b/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHIteratorApi.scala
index f9820178c2..8e20a72a65 100644
---
a/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHIteratorApi.scala
+++
b/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(
diff --git
a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxIteratorApi.scala
b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxIteratorApi.scala
index 73867aace9..9ad1e8dc8e 100644
---
a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxIteratorApi.scala
+++
b/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.")
diff --git
a/gluten-core/src/main/scala/org/apache/spark/sql/execution/adaptive/ShuffleStageWrapper.scala
b/gluten-core/src/main/scala/org/apache/spark/sql/execution/adaptive/ShuffleStageWrapper.scala
new file mode 100644
index 0000000000..47c23dc9c7
--- /dev/null
+++
b/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)
diff --git
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/extensions/AdvancedExtensionNode.java
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/extensions/AdvancedExtensionNode.java
index f313888d4f..b75f4e587f 100644
---
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/extensions/AdvancedExtensionNode.java
+++
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/extensions/AdvancedExtensionNode.java
@@ -40,6 +40,14 @@ public class AdvancedExtensionNode implements Serializable {
this.enhancement = enhancement;
}
+ public Any getOptimization() {
+ return optimization;
+ }
+
+ public Any getEnhancement() {
+ return enhancement;
+ }
+
public AdvancedExtension toProtobuf() {
AdvancedExtension.Builder extensionBuilder =
AdvancedExtension.newBuilder();
if (optimization != null) {
diff --git
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/plan/PlanNode.java
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/plan/PlanNode.java
index 0dd4a74bbc..373508392a 100644
---
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/plan/PlanNode.java
+++
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/plan/PlanNode.java
@@ -77,4 +77,8 @@ public class PlanNode implements Serializable {
}
return planBuilder.build();
}
+
+ public List<RelNode> getRelNodes() {
+ return relNodes;
+ }
}
diff --git
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/AggregateRelNode.java
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/AggregateRelNode.java
index 61ed751e45..31424e1774 100644
---
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/AggregateRelNode.java
+++
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/AggregateRelNode.java
@@ -26,6 +26,7 @@ import io.substrait.proto.RelCommon;
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 class AggregateRelNode implements RelNode,
Serializable {
builder.setAggregate(aggBuilder.build());
return builder.build();
}
+
+ @Override
+ public List<RelNode> childNodes() {
+ return Collections.singletonList(input);
+ }
}
diff --git
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/CrossRelNode.java
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/CrossRelNode.java
index 338a10b4d6..070e95e37e 100644
---
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/CrossRelNode.java
+++
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/CrossRelNode.java
@@ -24,6 +24,8 @@ import io.substrait.proto.Rel;
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 class CrossRelNode implements RelNode, Serializable {
}
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;
+ }
}
diff --git
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/ExpandRelNode.java
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/ExpandRelNode.java
index 25d007e564..29aa595818 100644
---
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/ExpandRelNode.java
+++
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/ExpandRelNode.java
@@ -25,6 +25,7 @@ import io.substrait.proto.RelCommon;
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 class ExpandRelNode implements RelNode, Serializable {
builder.setExpand(expandBuilder.build());
return builder.build();
}
+
+ @Override
+ public List<RelNode> childNodes() {
+ return Collections.singletonList(input);
+ }
}
diff --git
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/FetchRelNode.java
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/FetchRelNode.java
index 3c8b991c5b..91f5a00ce5 100644
---
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/FetchRelNode.java
+++
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/FetchRelNode.java
@@ -23,6 +23,8 @@ import io.substrait.proto.Rel;
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 class FetchRelNode implements RelNode, Serializable {
relBuilder.setFetch(fetchRelBuilder.build());
return relBuilder.build();
}
+
+ @Override
+ public List<RelNode> childNodes() {
+ return Collections.singletonList(input);
+ }
}
diff --git
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/FilterRelNode.java
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/FilterRelNode.java
index 356b12b292..df699f8f78 100644
---
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/FilterRelNode.java
+++
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/FilterRelNode.java
@@ -24,6 +24,8 @@ import io.substrait.proto.Rel;
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 class FilterRelNode implements RelNode, Serializable {
builder.setFilter(filterBuilder.build());
return builder.build();
}
+
+ @Override
+ public List<RelNode> childNodes() {
+ return Collections.singletonList(input);
+ }
}
diff --git
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/GenerateRelNode.java
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/GenerateRelNode.java
index efd4708919..a648f8f66d 100644
---
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/GenerateRelNode.java
+++
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/GenerateRelNode.java
@@ -24,6 +24,7 @@ import io.substrait.proto.Rel;
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 class GenerateRelNode implements RelNode, Serializable
{
relBuilder.setGenerate(generateRelBuilder.build());
return relBuilder.build();
}
+
+ @Override
+ public List<RelNode> childNodes() {
+ return Collections.singletonList(input);
+ }
}
diff --git
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/InputIteratorRelNode.java
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/InputIteratorRelNode.java
index 55b0e313be..e31d660900 100644
---
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/InputIteratorRelNode.java
+++
b/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,6 +40,9 @@ 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;
@@ -39,6 +50,22 @@ public class InputIteratorRelNode implements RelNode {
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 class InputIteratorRelNode implements RelNode {
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<>();
+ }
}
diff --git
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/JoinRelNode.java
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/JoinRelNode.java
index 2bd98500fe..1abb7374d9 100644
---
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/JoinRelNode.java
+++
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/JoinRelNode.java
@@ -24,6 +24,8 @@ import io.substrait.proto.Rel;
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 class JoinRelNode implements RelNode, Serializable {
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;
+ }
}
diff --git
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/ProjectRelNode.java
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/ProjectRelNode.java
index d1cccd1595..8b76619be1 100644
---
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/ProjectRelNode.java
+++
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/ProjectRelNode.java
@@ -25,6 +25,7 @@ import io.substrait.proto.RelCommon;
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 class ProjectRelNode implements RelNode, Serializable {
builder.setProject(projectBuilder.build());
return builder.build();
}
+
+ @Override
+ public List<RelNode> childNodes() {
+ return Collections.singletonList(input);
+ }
}
diff --git
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/ReadRelNode.java
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/ReadRelNode.java
index b82e05fd36..33332f1867 100644
---
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/ReadRelNode.java
+++
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/ReadRelNode.java
@@ -18,19 +18,26 @@ package org.apache.gluten.substrait.rel;
import org.apache.gluten.substrait.expression.ExpressionNode;
import org.apache.gluten.substrait.extensions.AdvancedExtensionNode;
+import org.apache.gluten.substrait.extensions.ExtensionBuilder;
import org.apache.gluten.substrait.type.ColumnTypeNode;
import org.apache.gluten.substrait.type.TypeNode;
import org.apache.gluten.utils.SubstraitUtil;
+import com.google.protobuf.Any;
+import com.google.protobuf.InvalidProtocolBufferException;
+import com.google.protobuf.StringValue;
import io.substrait.proto.NamedStruct;
import io.substrait.proto.ReadRel;
import io.substrait.proto.Rel;
import io.substrait.proto.RelCommon;
+import org.apache.spark.sql.execution.adaptive.InputStats;
import java.io.Serializable;
import java.util.ArrayList;
import java.util.List;
+import scala.math.BigInt;
+
public class ReadRelNode implements RelNode, Serializable {
private final List<TypeNode> types = new ArrayList<>();
private final List<String> names = new ArrayList<>();
@@ -39,6 +46,9 @@ public class ReadRelNode implements RelNode, Serializable {
private final AdvancedExtensionNode extensionNode;
private boolean streamKafka = false;
+ private BigInt rowCount;
+ private InputStats inputStats;
+
ReadRelNode(
List<TypeNode> types,
List<String> names,
@@ -52,6 +62,22 @@ public class ReadRelNode implements RelNode, Serializable {
this.extensionNode = extensionNode;
}
+ 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;
+ }
+
public void setStreamKafka(boolean streamKafka) {
this.streamKafka = streamKafka;
}
@@ -74,11 +100,42 @@ public class ReadRelNode implements RelNode, Serializable {
}
if (extensionNode != null) {
- readBuilder.setAdvancedExtension(extensionNode.toProtobuf());
+ if (null != rowCount) {
+ Any optimization = extensionNode.getOptimization();
+ Any enhancement = extensionNode.getEnhancement();
+ try {
+ String existingOptimizationStr =
+ optimization == null ? "" :
optimization.unpack(StringValue.class).getValue();
+ String newOptimizationStr =
+ existingOptimizationStr + "rowSize=" + rowCount.toLong() + "\n";
+ Any newOptimization =
+
Any.pack(StringValue.newBuilder().setValue(newOptimizationStr).build());
+ readBuilder.setAdvancedExtension(
+ new AdvancedExtensionNode(newOptimization,
enhancement).toProtobuf());
+ } catch (InvalidProtocolBufferException e) {
+ throw new RuntimeException(e);
+ }
+ } else {
+ readBuilder.setAdvancedExtension(extensionNode.toProtobuf());
+ }
+ } else {
+ 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<>();
+ }
}
diff --git
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/RelNode.java
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/RelNode.java
index 3dacd947c0..8cc0d2143e 100644
---
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/RelNode.java
+++
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/RelNode.java
@@ -19,6 +19,7 @@ package org.apache.gluten.substrait.rel;
import io.substrait.proto.Rel;
import java.io.Serializable;
+import java.util.List;
/** Contains helper functions for constructing substrait relations. */
public interface RelNode extends Serializable {
@@ -28,4 +29,6 @@ public interface RelNode extends Serializable {
* @return A rel protobuf
*/
Rel toProtobuf();
+
+ List<RelNode> childNodes();
}
diff --git
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/SetRelNode.java
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/SetRelNode.java
index ddcfb1701d..410fa59d84 100644
---
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/SetRelNode.java
+++
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/SetRelNode.java
@@ -59,4 +59,9 @@ public class SetRelNode implements RelNode, Serializable {
builder.setSet(setBuilder.build());
return builder.build();
}
+
+ @Override
+ public List<RelNode> childNodes() {
+ return inputs;
+ }
}
diff --git
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/SortRelNode.java
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/SortRelNode.java
index ec7368cd3a..b34f5a85d1 100644
---
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/SortRelNode.java
+++
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/SortRelNode.java
@@ -24,6 +24,7 @@ import io.substrait.proto.SortField;
import io.substrait.proto.SortRel;
import java.io.Serializable;
+import java.util.Collections;
import java.util.List;
public class SortRelNode implements RelNode, Serializable {
@@ -67,4 +68,9 @@ public class SortRelNode implements RelNode, Serializable {
builder.setSort(sortBuilder.build());
return builder.build();
}
+
+ @Override
+ public List<RelNode> childNodes() {
+ return Collections.singletonList(input);
+ }
}
diff --git
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/TopNNode.java
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/TopNNode.java
index a33d7aaca5..dbc325674f 100644
---
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/TopNNode.java
+++
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/TopNNode.java
@@ -25,6 +25,7 @@ import io.substrait.proto.TopNRel;
import java.io.Serializable;
import java.util.ArrayList;
+import java.util.Collections;
import java.util.List;
public class TopNNode implements RelNode, Serializable {
@@ -73,4 +74,9 @@ public class TopNNode implements RelNode, Serializable {
relBuilder.setTopN(topNBuilder.build());
return relBuilder.build();
}
+
+ @Override
+ public List<RelNode> childNodes() {
+ return Collections.singletonList(input);
+ }
}
diff --git
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/WindowGroupLimitRelNode.java
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/WindowGroupLimitRelNode.java
index a7a0ea62aa..3bc448892f 100644
---
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/WindowGroupLimitRelNode.java
+++
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/WindowGroupLimitRelNode.java
@@ -26,6 +26,7 @@ import io.substrait.proto.WindowGroupLimitRel;
import java.io.Serializable;
import java.util.ArrayList;
+import java.util.Collections;
import java.util.List;
public class WindowGroupLimitRelNode implements RelNode, Serializable {
@@ -88,4 +89,9 @@ public class WindowGroupLimitRelNode implements RelNode,
Serializable {
builder.setWindowGroupLimit(windowBuilder.build());
return builder.build();
}
+
+ @Override
+ public List<RelNode> childNodes() {
+ return Collections.singletonList(input);
+ }
}
diff --git
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/WindowRelNode.java
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/WindowRelNode.java
index 415bbc350f..5e619cc647 100644
---
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/WindowRelNode.java
+++
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/WindowRelNode.java
@@ -27,6 +27,7 @@ import io.substrait.proto.WindowRel;
import java.io.Serializable;
import java.util.ArrayList;
+import java.util.Collections;
import java.util.List;
public class WindowRelNode implements RelNode, Serializable {
@@ -93,4 +94,9 @@ public class WindowRelNode implements RelNode, Serializable {
builder.setWindow(windowBuilder.build());
return builder.build();
}
+
+ @Override
+ public List<RelNode> childNodes() {
+ return Collections.singletonList(input);
+ }
}
diff --git
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/WriteRelNode.java
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/WriteRelNode.java
index 74ffc8282c..9c875b1d01 100644
---
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/WriteRelNode.java
+++
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/WriteRelNode.java
@@ -25,6 +25,7 @@ import io.substrait.proto.*;
import java.io.Serializable;
import java.util.ArrayList;
+import java.util.Collections;
import java.util.List;
public class WriteRelNode implements RelNode, Serializable {
@@ -83,4 +84,9 @@ public class WriteRelNode implements RelNode, Serializable {
builder.setWrite(writeBuilder.build());
return builder.build();
}
+
+ @Override
+ public List<RelNode> childNodes() {
+ return Collections.singletonList(input);
+ }
}
diff --git
a/gluten-substrait/src/main/scala/org/apache/gluten/backendsapi/IteratorApi.scala
b/gluten-substrait/src/main/scala/org/apache/gluten/backendsapi/IteratorApi.scala
index a6b0053593..2cb45b97a8 100644
---
a/gluten-substrait/src/main/scala/org/apache/gluten/backendsapi/IteratorApi.scala
+++
b/gluten-substrait/src/main/scala/org/apache/gluten/backendsapi/IteratorApi.scala
@@ -68,7 +68,8 @@ trait IteratorApi {
updateNativeMetrics: IMetrics => Unit,
partitionIndex: Int,
inputIterators: Seq[Iterator[ColumnarBatch]] = Seq(),
- enableCudf: Boolean = false
+ enableCudf: Boolean = false,
+ wsContext: WholeStageTransformContext = null
): Iterator[ColumnarBatch]
/**
diff --git
a/gluten-substrait/src/main/scala/org/apache/gluten/config/GlutenConfig.scala
b/gluten-substrait/src/main/scala/org/apache/gluten/config/GlutenConfig.scala
index bd5394095d..93fb3888e9 100644
---
a/gluten-substrait/src/main/scala/org/apache/gluten/config/GlutenConfig.scala
+++
b/gluten-substrait/src/main/scala/org/apache/gluten/config/GlutenConfig.scala
@@ -326,6 +326,8 @@ class GlutenConfig(conf: SQLConf) extends
GlutenCoreConfig(conf) {
def enableFallbackReport: Boolean = getConf(FALLBACK_REPORTER_ENABLED)
+ def enablePassStageInputStats: Boolean =
getConf(GLUTEN_PASS_STAGE_INPUT_STATS_ENABLED)
+
def failOnFallback: Boolean = getConf(FALLBACK_FAIL_ON_FALLBACK)
def debug: Boolean = getConf(DEBUG_ENABLED)
@@ -1444,6 +1446,15 @@ object GlutenConfig extends ConfigRegistry {
.booleanConf
.createWithDefault(true)
+ val GLUTEN_PASS_STAGE_INPUT_STATS_ENABLED =
+ buildConf("spark.gluten.sql.enablePassStageInputStats")
+ .internal()
+ .doc("When true, pass stage input stats (scan/shuffle/broadcast) to the
native engine " +
+ "as estimated row size hints. Disabled by default and no-op for
backends that do not " +
+ "consume the hints.")
+ .booleanConf
+ .createWithDefault(false)
+
val FALLBACK_FAIL_ON_FALLBACK =
buildConf("spark.gluten.sql.columnar.failOnFallback")
.internal()
diff --git
a/gluten-substrait/src/main/scala/org/apache/gluten/execution/BasicScanExecTransformer.scala
b/gluten-substrait/src/main/scala/org/apache/gluten/execution/BasicScanExecTransformer.scala
index 3b142fecb9..6d265100d0 100644
---
a/gluten-substrait/src/main/scala/org/apache/gluten/execution/BasicScanExecTransformer.scala
+++
b/gluten-substrait/src/main/scala/org/apache/gluten/execution/BasicScanExecTransformer.scala
@@ -22,11 +22,12 @@ import org.apache.gluten.expression.{ConverterUtils,
ExpressionConverter}
import org.apache.gluten.substrait.`type`.ColumnTypeNode
import org.apache.gluten.substrait.SubstraitContext
import org.apache.gluten.substrait.extensions.ExtensionBuilder
-import org.apache.gluten.substrait.rel.{RelBuilder, SplitInfo}
+import org.apache.gluten.substrait.rel.{ReadRelNode, RelBuilder, SplitInfo}
import org.apache.gluten.substrait.rel.LocalFilesNode.ReadFileFormat
import org.apache.spark.Partition
import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.execution.adaptive.InputStats
import com.google.protobuf.StringValue
import io.substrait.proto.NamedStruct
@@ -94,6 +95,8 @@ trait BasicScanExecTransformer extends LeafTransformSupport
with BaseDataSource
/** Returns the file format properties. */
def getProperties: Map[String, String] = Map.empty
+ def getInputStats: Option[InputStats] = Option.empty
+
override def getSplitInfos: Seq[SplitInfo] = {
getSplitInfosFromPartitions(getPartitionWithReadFileFormats)
}
@@ -183,6 +186,11 @@ trait BasicScanExecTransformer extends
LeafTransformSupport with BaseDataSource
extensionNode,
context,
context.nextOperatorId(this.nodeName))
+ getInputStats.foreach(
+ inputStats => {
+ logInfo(s"hit scan inputStats with $inputStats")
+ readNode.asInstanceOf[ReadRelNode].setInputStats(inputStats)
+ })
TransformContext(output, readNode)
}
}
diff --git
a/gluten-substrait/src/main/scala/org/apache/gluten/execution/FileSourceScanExecTransformer.scala
b/gluten-substrait/src/main/scala/org/apache/gluten/execution/FileSourceScanExecTransformer.scala
index 8fea175d7e..83023f027a 100644
---
a/gluten-substrait/src/main/scala/org/apache/gluten/execution/FileSourceScanExecTransformer.scala
+++
b/gluten-substrait/src/main/scala/org/apache/gluten/execution/FileSourceScanExecTransformer.scala
@@ -29,6 +29,7 @@ import org.apache.spark.sql.catalyst.plans.QueryPlan
import org.apache.spark.sql.catalyst.util.truncatedString
import org.apache.spark.sql.connector.read.streaming.SparkDataStream
import org.apache.spark.sql.execution.FileSourceScanExecShim
+import org.apache.spark.sql.execution.adaptive.InputStats
import org.apache.spark.sql.execution.datasources.HadoopFsRelation
import org.apache.spark.sql.execution.metric.SQLMetric
import org.apache.spark.sql.types.StructType
@@ -48,7 +49,8 @@ case class FileSourceScanExecTransformer(
override val dataFilters: Seq[Expression],
override val tableIdentifier: Option[TableIdentifier],
override val disableBucketedScan: Boolean = false,
- override val pushDownFilters: Option[Seq[Expression]] = None)
+ override val pushDownFilters: Option[Seq[Expression]] = None,
+ inputStats: Option[InputStats] = None)
extends FileSourceScanExecTransformerBase(
relation,
stream,
@@ -62,6 +64,10 @@ case class FileSourceScanExecTransformer(
disableBucketedScan
) {
+ override def getInputStats: Option[InputStats] = {
+ inputStats
+ }
+
override def doCanonicalize(): FileSourceScanExecTransformer = {
FileSourceScanExecTransformer(
relation,
@@ -78,7 +84,8 @@ case class FileSourceScanExecTransformer(
QueryPlan.normalizePredicates(dataFilters, output),
None,
disableBucketedScan,
- pushDownFilters.map(QueryPlan.normalizePredicates(_, output))
+ pushDownFilters.map(QueryPlan.normalizePredicates(_, output)),
+ inputStats
)
}
diff --git
a/gluten-substrait/src/main/scala/org/apache/gluten/execution/GlutenWholeStageColumnarRDD.scala
b/gluten-substrait/src/main/scala/org/apache/gluten/execution/GlutenWholeStageColumnarRDD.scala
index 4a00dbb587..afec9cc10f 100644
---
a/gluten-substrait/src/main/scala/org/apache/gluten/execution/GlutenWholeStageColumnarRDD.scala
+++
b/gluten-substrait/src/main/scala/org/apache/gluten/execution/GlutenWholeStageColumnarRDD.scala
@@ -17,6 +17,8 @@
package org.apache.gluten.execution
import org.apache.gluten.backendsapi.BackendsApiManager
+import org.apache.gluten.config.GlutenConfig
+import org.apache.gluten.extension.ApplyStageInputStatsRule
import org.apache.gluten.metrics.{GlutenTimeMetric, IMetrics}
import org.apache.gluten.substrait.rel.SplitInfo
@@ -58,12 +60,19 @@ class GlutenWholeStageColumnarRDD(
pipelineTime: SQLMetric,
updateInputMetrics: InputMetricsWrapper => Unit,
updateNativeMetrics: IMetrics => Unit,
- enableCudf: Boolean = false)
+ enableCudf: Boolean = false,
+ wsContext: WholeStageTransformContext = null)
extends RDD[ColumnarBatch](sc, rdds.getDependencies) {
override def compute(split: Partition, context: TaskContext):
Iterator[ColumnarBatch] = {
GlutenTimeMetric.millis(pipelineTime) {
_ =>
+ if (GlutenConfig.get.enablePassStageInputStats && wsContext != null) {
+ ApplyStageInputStatsRule.setStageInputStatsToInputNode(
+ wsContext,
+ split.index,
+ rdds.getPartitionLength)
+ }
val (inputPartition, inputColumnarRDDPartitions) =
castNativePartition(split)
val inputIterators = rdds.getIterators(inputColumnarRDDPartitions,
context)
BackendsApiManager.getIteratorApiInstance.genFirstStageIterator(
@@ -74,7 +83,8 @@ class GlutenWholeStageColumnarRDD(
updateNativeMetrics,
split.index,
inputIterators,
- enableCudf
+ enableCudf,
+ wsContext
)
}
}
diff --git
a/gluten-substrait/src/main/scala/org/apache/gluten/execution/WholeStageTransformer.scala
b/gluten-substrait/src/main/scala/org/apache/gluten/execution/WholeStageTransformer.scala
index 8de9d407bf..1f3ae4d753 100644
---
a/gluten-substrait/src/main/scala/org/apache/gluten/execution/WholeStageTransformer.scala
+++
b/gluten-substrait/src/main/scala/org/apache/gluten/execution/WholeStageTransformer.scala
@@ -379,7 +379,8 @@ case class WholeStageTransformer(child: SparkPlan,
materializeInput: Boolean = f
wsCtx.substraitContext.registeredJoinParams,
wsCtx.substraitContext.registeredAggregationParams
),
- wsCtx.enableCudf
+ wsCtx.enableCudf,
+ wsCtx
)
val allInputPartitions = leafTransformers.map(_.getPartitions)
@@ -434,7 +435,8 @@ case class WholeStageTransformer(child: SparkPlan,
materializeInput: Boolean = f
wsCtx.substraitContext.registeredJoinParams,
wsCtx.substraitContext.registeredAggregationParams
),
- materializeInput
+ materializeInput,
+ inputRDDs.getPartitionLength
)
}
}
diff --git
a/gluten-substrait/src/main/scala/org/apache/gluten/execution/WholeStageZippedPartitionsRDD.scala
b/gluten-substrait/src/main/scala/org/apache/gluten/execution/WholeStageZippedPartitionsRDD.scala
index 393716bb7b..73068241f7 100644
---
a/gluten-substrait/src/main/scala/org/apache/gluten/execution/WholeStageZippedPartitionsRDD.scala
+++
b/gluten-substrait/src/main/scala/org/apache/gluten/execution/WholeStageZippedPartitionsRDD.scala
@@ -17,6 +17,8 @@
package org.apache.gluten.execution
import org.apache.gluten.backendsapi.BackendsApiManager
+import org.apache.gluten.config.GlutenConfig
+import org.apache.gluten.extension.ApplyStageInputStatsRule
import org.apache.gluten.metrics.{GlutenTimeMetric, IMetrics}
import org.apache.spark.{Partition, SparkConf, SparkContext, TaskContext}
@@ -36,12 +38,19 @@ class WholeStageZippedPartitionsRDD(
resCtx: WholeStageTransformContext,
pipelineTime: SQLMetric,
updateNativeMetrics: IMetrics => Unit,
- materializeInput: Boolean)
+ materializeInput: Boolean,
+ partitionLength: Int = 0)
extends RDD[ColumnarBatch](sc, rdds.getDependencies) {
override def compute(split: Partition, context: TaskContext):
Iterator[ColumnarBatch] = {
GlutenTimeMetric.millis(pipelineTime) {
_ =>
+ if (GlutenConfig.get.enablePassStageInputStats) {
+ ApplyStageInputStatsRule.setStageInputStatsToInputNode(
+ resCtx,
+ split.index,
+ partitionLength)
+ }
val partitions =
split.asInstanceOf[ZippedPartitionsPartition].inputColumnarRDDPartitions
val inputIterators: Seq[Iterator[ColumnarBatch]] =
rdds.getIterators(partitions, context)
BackendsApiManager.getIteratorApiInstance
diff --git
a/gluten-substrait/src/main/scala/org/apache/gluten/extension/ApplyStageInputStatsRule.scala
b/gluten-substrait/src/main/scala/org/apache/gluten/extension/ApplyStageInputStatsRule.scala
new file mode 100644
index 0000000000..dfd24209e6
--- /dev/null
+++
b/gluten-substrait/src/main/scala/org/apache/gluten/extension/ApplyStageInputStatsRule.scala
@@ -0,0 +1,148 @@
+/*
+ * 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.gluten.extension
+
+import org.apache.gluten.execution.WholeStageTransformContext
+import org.apache.gluten.substrait.rel.{InputIteratorRelNode, ReadRelNode,
RelNode}
+
+import org.apache.spark.internal.Logging
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.execution.{CoalescedMapperPartitionSpec,
CoalescedPartitionSpec, PartialMapperPartitionSpec,
PartialReducerPartitionSpec, ShufflePartitionSpec}
+import org.apache.spark.sql.execution.adaptive.{BroadcastStats, InputStats,
ScanInputStats, ShuffleStats}
+import org.apache.spark.sql.execution.datasources.LogicalRelation
+
+import scala.collection.{mutable, Seq}
+
+object ApplyStageInputStatsRule extends Logging {
+ private def estimateInputIteratorRowCount(stats: InputStats, index: Int):
Option[Long] = {
+ if (stats.inputStatsKind == BroadcastStats) {
+ Some(stats.rowCount.toLong)
+ } else {
+ val totalBytes = Math.max(stats.bytesByPartitionId.sum, 1L)
+ if (index >= 0 && index < stats.bytesByPartitionId.length) {
+ val partitionBytes = stats.bytesByPartitionId(index)
+ Some((1.0d * stats.rowCount.toLong / totalBytes *
partitionBytes).toLong)
+ } else {
+ logWarning(
+ s"Skip invalid input stats index $index, bytesByPartitionId length:
" +
+ s"${stats.bytesByPartitionId.length}")
+ None
+ }
+ }
+ }
+
+ private def estimateScanRowCount(stats: InputStats, partitionLength: Int):
Long = {
+ val safePartitionLength = Math.max(partitionLength, 1)
+ (1.0d * stats.rowCount.toLong / safePartitionLength).toLong
+ }
+
+ def setStageInputStatsToInputNode(
+ ws: WholeStageTransformContext,
+ index: Int,
+ partitionLength: Int): Unit = {
+ val relNodesQueue: mutable.Queue[RelNode] = mutable.Queue()
+ val nodes = ws.root.getRelNodes
+ if (nodes != null) {
+ nodes.forEach(relNodesQueue.enqueue(_))
+ }
+ while (relNodesQueue.nonEmpty) {
+ val node = relNodesQueue.dequeue()
+ node match {
+ case input: InputIteratorRelNode =>
+ val stats = input.getInputStats
+ if (null != stats) {
+ estimateInputIteratorRowCount(stats, index).foreach {
+ estimatedRowCount =>
+ logDebug(
+ s"pass input stats idx: $index with estimated row count
$estimatedRowCount")
+ input.setRowCount(estimatedRowCount)
+ }
+ }
+ case scan: ReadRelNode =>
+ val stats = scan.getInputStats
+ if (null != stats) {
+ val estimatedRowCount = estimateScanRowCount(stats,
partitionLength)
+ logDebug(s"pass scan stats idx: $index with estimated row count
$estimatedRowCount")
+ scan.setRowCount(estimatedRowCount)
+ }
+ case _ =>
+ }
+ node.childNodes().forEach(relNodesQueue.enqueue(_))
+ }
+ }
+
+ /** Compute task stats for each partition */
+ def recomputeInputStatsForAQEShuffleReadExec(
+ original: InputStats,
+ aqeShufflePartitions: Seq[ShufflePartitionSpec]): InputStats = {
+ val originalBytesByPartition = original.bytesByPartitionId
+ val totalBytes = originalBytesByPartition.sum
+ val rowCount = original.rowCount
+ val aqeReadPartitionsStats = new mutable.ArrayBuffer[Long]()
+ aqeShufflePartitions.foreach {
+ case CoalescedPartitionSpec(startReducerIndex, endReducerIndex, _) =>
+ var newPartitionDataSize = 0L
+ for (i <- startReducerIndex until endReducerIndex) {
+ if (i < originalBytesByPartition.length) {
+ newPartitionDataSize += originalBytesByPartition(i)
+ } else {
+ newPartitionDataSize += 0L
+ }
+ }
+ aqeReadPartitionsStats += newPartitionDataSize
+ case PartialReducerPartitionSpec(reducerIndex, startMapIndex,
endMapIndex, _) =>
+ // over estimated stats
+ if (reducerIndex < originalBytesByPartition.length) {
+ aqeReadPartitionsStats += originalBytesByPartition(reducerIndex)
+ } else {
+ aqeReadPartitionsStats += 0L
+ }
+ case PartialMapperPartitionSpec(mapIndex, startReducerIndex,
endReducerIndex) =>
+ // unknown stats
+ aqeReadPartitionsStats += 0L
+ case CoalescedMapperPartitionSpec(startMapIndex, endMapIndex,
numReducers) =>
+ // unknown stats
+ aqeReadPartitionsStats += 0L
+ }
+ InputStats(known = true, totalBytes, rowCount,
aqeReadPartitionsStats.toArray, ShuffleStats)
+ }
+ def createInputStats(logicalLink: Option[LogicalPlan]): Option[InputStats] =
{
+ if (logicalLink.isEmpty) {
+ logInfo(s"no logicalLink associated")
+ return None
+ }
+ val logicalPlan = logicalLink.get
+ logInfo("scan logical link:" + logicalPlan)
+ logInfo("FileSourceScanTransformer logical plan " + logicalPlan)
+ val maybeScanStats = logicalPlan.collectFirst {
+ case relation: LogicalRelation => relation.catalogTable
+ }.flatten.flatMap(_.stats)
+ if (maybeScanStats.isDefined) {
+ val scanStats = maybeScanStats.get
+ logInfo("catalogTable " + scanStats)
+ Some(
+ InputStats(
+ known = true,
+ scanStats.sizeInBytes,
+ scanStats.rowCount.getOrElse(0),
+ Array(),
+ ScanInputStats))
+ } else {
+ None
+ }
+ }
+}
diff --git
a/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ColumnarCollapseTransformStages.scala
b/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ColumnarCollapseTransformStages.scala
index c81d380a78..f35aea8fdc 100644
---
a/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ColumnarCollapseTransformStages.scala
+++
b/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ColumnarCollapseTransformStages.scala
@@ -19,10 +19,11 @@ package org.apache.spark.sql.execution
import org.apache.gluten.backendsapi.BackendsApiManager
import org.apache.gluten.config.GlutenConfig
import org.apache.gluten.execution._
+import org.apache.gluten.extension.ApplyStageInputStatsRule
import org.apache.gluten.extension.columnar.transition.{Convention,
ConventionReq}
import org.apache.gluten.metrics.MetricsUpdater
import org.apache.gluten.substrait.SubstraitContext
-import org.apache.gluten.substrait.rel.RelBuilder
+import org.apache.gluten.substrait.rel.{InputIteratorRelNode, RelBuilder}
import org.apache.spark.broadcast.Broadcast
import org.apache.spark.rdd.RDD
@@ -31,7 +32,7 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute,
SortOrder}
import org.apache.spark.sql.catalyst.plans.physical.Partitioning
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.catalyst.util.truncatedString
-import org.apache.spark.sql.execution.adaptive.{BroadcastQueryStageExec,
ShuffleQueryStageExec}
+import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec,
BroadcastQueryStageExec, BroadcastStats, InputStats, ShuffleQueryStageExec,
ShuffleStageWrapper, ShuffleStats}
import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike,
ShuffleExchangeLike}
import org.apache.spark.sql.execution.metric.SQLMetric
import org.apache.spark.sql.vectorized.ColumnarBatch
@@ -76,6 +77,36 @@ case class InputIteratorTransformer(child: SparkPlan)
extends UnaryTransformSupp
assert(child.isInstanceOf[ColumnarInputAdapter])
val operatorId = context.nextOperatorId(nodeName)
val readRel = RelBuilder.makeReadRelForInputIterator(child.output.asJava,
context, operatorId)
+ if (GlutenConfig.get.enablePassStageInputStats) {
+ val inputStatsOpt: Option[InputStats] = {
+ val inputAdapter = child.asInstanceOf[ColumnarInputAdapter]
+ inputAdapter.child match {
+ case shuffle @ ShuffleStageWrapper(stats) =>
+ val inputStats = InputStats(
+ known = true,
+ sizeInBytes = stats.sizeInBytes,
+ rowCount = stats.rowCount.getOrElse(0),
+ bytesByPartitionId =
+
stats.mapOutputStatistics.map(_.bytesByPartitionId).getOrElse(Array()),
+ if (shuffle.isInstanceOf[BroadcastQueryStageExec])
BroadcastStats else ShuffleStats
+ )
+ shuffle match {
+ case a: AQEShuffleReadExec =>
+ Some(
+ ApplyStageInputStatsRule
+ .recomputeInputStatsForAQEShuffleReadExec(inputStats,
a.partitionSpecs))
+ case _ => Some(inputStats)
+ }
+ case other =>
+ logInfo(s"hit InputIteratorTransformer ${other.getClass.toString}
with stats")
+ None
+ }
+ }
+ if (inputStatsOpt.isDefined) {
+ logInfo(s"set input iterator stats ${inputStatsOpt.get}")
+
readRel.asInstanceOf[InputIteratorRelNode].setInputStats(inputStatsOpt.get)
+ }
+ }
TransformContext(output, readRel)
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
