This is an automated email from the ASF dual-hosted git repository. taiyang-li pushed a commit to branch fake_add_bolt_backend in repository https://gitbox.apache.org/repos/asf/gluten.git
commit ab54bf2a8167b8ebbd015d06069244d8af9559d8 Author: liyang.127 <[email protected]> AuthorDate: Fri Jun 26 17:19:40 2026 +0800 [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 --- .../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 | 40 +++++++ .../apache/gluten/substrait/rel/JoinRelNode.java | 10 ++ .../gluten/substrait/rel/ProjectRelNode.java | 6 + .../apache/gluten/substrait/rel/ReadRelNode.java | 57 ++++++++- .../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 | 132 +++++++++++++++++++++ .../ColumnarCollapseTransformStages.scala | 35 +++++- 31 files changed, 513 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..075aa0fbe4 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> childNode() { + 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..e2faa88103 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> childNode() { + 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..844c21910c 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> childNode() { + 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..7d8d591172 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> childNode() { + 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..7a7d2fa366 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> childNode() { + 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..1d1a7a58b5 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> childNode() { + 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..cf896c8a2f 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 rowSize; + + 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 getRowSize() { + return rowSize; + } + + public void setRowSize(BigInt rowSize) { + this.rowSize = rowSize; + } + + 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,21 @@ public class InputIteratorRelNode implements RelNode { LocalFilesBuilder.makeLocalFiles(ConverterUtils.ITERATOR_PREFIX() + iteratorIndex); readBuilder.setLocalFiles(iteratorIndexNode.toProtobuf()); + if (null != rowSize) { + Any inputRowSize = + Any.pack(StringValue.newBuilder().setValue("rowSize=" + rowSize.toLong() + "\n").build()); + AdvancedExtensionNode advancedExtension = + ExtensionBuilder.makeAdvancedExtension(inputRowSize, null); + readBuilder.setAdvancedExtension(advancedExtension.toProtobuf()); + } + Rel.Builder builder = Rel.newBuilder(); builder.setRead(readBuilder.build()); return builder.build(); } + + @Override + public List<RelNode> childNode() { + 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..fcf7a3fea3 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> childNode() { + 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..b5e486baaa 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> childNode() { + 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..b44a9631f8 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 rowSize; + 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 getRowSize() { + return rowSize; + } + + public void setRowSize(BigInt rowSize) { + this.rowSize = rowSize; + } + + public InputStats getInputStats() { + return inputStats; + } + + public void setInputStats(InputStats inputStats) { + this.inputStats = inputStats; + } + public void setStreamKafka(boolean streamKafka) { this.streamKafka = streamKafka; } @@ -74,11 +100,40 @@ public class ReadRelNode implements RelNode, Serializable { } if (extensionNode != null) { - readBuilder.setAdvancedExtension(extensionNode.toProtobuf()); + if (null != rowSize) { + Any optimization = extensionNode.getOptimization(); + Any enhancement = extensionNode.getEnhancement(); + try { + Any any = Any.parseFrom(optimization.toByteArray()); + String newOptimizationStr = any.getValue() + "rowSize=" + rowSize.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 != rowSize) { + Any inputRowSize = + Any.pack( + StringValue.newBuilder().setValue("rowSize=" + rowSize.toLong() + "\n").build()); + AdvancedExtensionNode advancedExtension = + ExtensionBuilder.makeAdvancedExtension(inputRowSize, null); + readBuilder.setAdvancedExtension(advancedExtension.toProtobuf()); + } } Rel.Builder builder = Rel.newBuilder(); builder.setRead(readBuilder.build()); return builder.build(); } + + @Override + public List<RelNode> childNode() { + 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..817e39cf9f 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> childNode(); } 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..eab234acb5 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> childNode() { + 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..5266a4bd8e 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> childNode() { + 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..efd4f09a9a 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> childNode() { + 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..e4a6e8b647 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> childNode() { + 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..409d143bf4 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> childNode() { + 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..af6aba7f1e 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> childNode() { + 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 053d715125..d6e8fded5f 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 @@ -325,6 +325,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) @@ -1431,6 +1433,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..bb91c3ae9b --- /dev/null +++ b/gluten-substrait/src/main/scala/org/apache/gluten/extension/ApplyStageInputStatsRule.scala @@ -0,0 +1,132 @@ +/* + * 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 { + 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) { + val totalBytes = Math.max(stats.bytesByPartitionId.sum, 1) + val estimatedRowCount = + if (stats.inputStatsKind == BroadcastStats) { + stats.rowCount.toLong + } else { + (1.0d * stats.rowCount.toLong / totalBytes * stats.bytesByPartitionId(index)).toLong + } + logInfo(s"pass input stats idx: $index with estimated row count $estimatedRowCount") + input.setRowSize(estimatedRowCount) + } + case scan: ReadRelNode => + val stats = scan.getInputStats + if (null != stats) { + val estimatedRowCount = (1.0d * stats.rowCount.toLong / partitionLength).toLong + logInfo(s"pass scan stats idx: $index with estimated row count $estimatedRowCount") + scan.setRowSize(estimatedRowCount) + } + case _ => + } + node.childNode().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 maybeCatalogTable = logicalPlan.collectFirst { + case LogicalRelation(_, _, catalogTable, _) => catalogTable + } + if ( + maybeCatalogTable.isDefined && maybeCatalogTable.get.isDefined && + maybeCatalogTable.get.get.stats.isDefined + ) { + val scanStats = maybeCatalogTable.get.get.stats.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]
