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 f9820178c25..8e20a72a65c 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 73867aace94..9ad1e8dc8e2 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 00000000000..47c23dc9c77 --- /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 f313888d4ff..b75f4e587fc 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 AdvancedExtensionNode(Any optimization, Any enhancement) { this.enhancement = enhancement; } + public Any getOptimization() { + return optimization; + } + + public Any getEnhancement() { + return enhancement; + } + public AdvancedExtension toProtobuf() { AdvancedExtension.Builder extensionBuilder = AdvancedExtension.newBuilder(); if (optimization != null) { 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 0dd4a74bbce..373508392a6 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 Plan toProtobuf() { } return planBuilder.build(); } + + public List 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 61ed751e453..31424e17747 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 java.io.Serializable; import java.util.ArrayList; +import java.util.Collections; import java.util.List; public class AggregateRelNode implements RelNode, Serializable { @@ -83,4 +84,9 @@ public Rel toProtobuf() { builder.setAggregate(aggBuilder.build()); return builder.build(); } + + @Override + public List 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 338a10b4d6c..070e95e37ec 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.RelCommon; import java.io.Serializable; +import java.util.ArrayList; +import java.util.List; public class CrossRelNode implements RelNode, Serializable { private final RelNode left; @@ -69,4 +71,12 @@ public Rel toProtobuf() { } return Rel.newBuilder().setCross(crossRelBuilder.build()).build(); } + + @Override + public List childNodes() { + List 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 25d007e5641..29aa595818e 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 java.io.Serializable; import java.util.ArrayList; +import java.util.Collections; import java.util.List; public class ExpandRelNode implements RelNode, Serializable { @@ -76,4 +77,9 @@ public Rel toProtobuf() { builder.setExpand(expandBuilder.build()); return builder.build(); } + + @Override + public List 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 3c8b991c5b5..91f5a00ce56 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.RelCommon; import java.io.Serializable; +import java.util.Collections; +import java.util.List; public class FetchRelNode implements RelNode, Serializable { private final RelNode input; @@ -66,4 +68,9 @@ public Rel toProtobuf() { relBuilder.setFetch(fetchRelBuilder.build()); return relBuilder.build(); } + + @Override + public List 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 356b12b292a..df699f8f781 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.RelCommon; import java.io.Serializable; +import java.util.Collections; +import java.util.List; public class FilterRelNode implements RelNode, Serializable { private final RelNode input; @@ -60,4 +62,9 @@ public Rel toProtobuf() { builder.setFilter(filterBuilder.build()); return builder.build(); } + + @Override + public List 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 efd47089190..a648f8f66de 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.RelCommon; import java.io.Serializable; +import java.util.Collections; import java.util.List; public class GenerateRelNode implements RelNode, Serializable { @@ -81,4 +82,9 @@ public Rel toProtobuf() { relBuilder.setGenerate(generateRelBuilder.build()); return relBuilder.build(); } + + @Override + public List 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 55b0e313be5..e31d6609008 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 types; private final List names; private final Long iteratorIndex; + private BigInt rowCount; + + private InputStats inputStats; public InputIteratorRelNode(List types, List names, Long iteratorIndex) { this.types = types; @@ -39,6 +50,22 @@ public InputIteratorRelNode(List types, List names, Long itera this.iteratorIndex = iteratorIndex; } + public BigInt getRowCount() { + return rowCount; + } + + public void setRowCount(BigInt rowCount) { + this.rowCount = rowCount; + } + + public InputStats getInputStats() { + return inputStats; + } + + public void setInputStats(InputStats inputStats) { + this.inputStats = inputStats; + } + @Override public Rel toProtobuf() { Type.Struct.Builder structBuilder = Type.Struct.newBuilder(); @@ -59,8 +86,22 @@ public Rel toProtobuf() { LocalFilesBuilder.makeLocalFiles(ConverterUtils.ITERATOR_PREFIX() + iteratorIndex); readBuilder.setLocalFiles(iteratorIndexNode.toProtobuf()); + if (null != rowCount) { + Any inputRowCount = + Any.pack( + StringValue.newBuilder().setValue("rowSize=" + rowCount.toLong() + "\n").build()); + AdvancedExtensionNode advancedExtension = + ExtensionBuilder.makeAdvancedExtension(inputRowCount, null); + readBuilder.setAdvancedExtension(advancedExtension.toProtobuf()); + } + Rel.Builder builder = Rel.newBuilder(); builder.setRead(readBuilder.build()); return builder.build(); } + + @Override + public List 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 2bd98500fee..1abb7374d95 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.RelCommon; import java.io.Serializable; +import java.util.ArrayList; +import java.util.List; public class JoinRelNode implements RelNode, Serializable { private final RelNode left; @@ -79,4 +81,12 @@ public Rel toProtobuf() { return Rel.newBuilder().setJoin(joinBuilder.build()).build(); } + + @Override + public List childNodes() { + List 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 d1cccd15958..8b76619be1d 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 java.io.Serializable; import java.util.ArrayList; +import java.util.Collections; import java.util.List; public class ProjectRelNode implements RelNode, Serializable { @@ -78,4 +79,9 @@ public Rel toProtobuf() { builder.setProject(projectBuilder.build()); return builder.build(); } + + @Override + public List 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 b82e05fd367..33332f18672 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 @@ 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 types = new ArrayList<>(); private final List 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 types, List 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 Rel toProtobuf() { } 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 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 3dacd947c0f..8cc0d2143e3 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 @@ 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 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 ddcfb1701dd..410fa59d84b 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 Rel toProtobuf() { builder.setSet(setBuilder.build()); return builder.build(); } + + @Override + public List 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 ec7368cd3ac..b34f5a85d17 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.SortRel; import java.io.Serializable; +import java.util.Collections; import java.util.List; public class SortRelNode implements RelNode, Serializable { @@ -67,4 +68,9 @@ public Rel toProtobuf() { builder.setSort(sortBuilder.build()); return builder.build(); } + + @Override + public List 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 a33d7aaca5b..dbc325674f9 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 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 Rel toProtobuf() { relBuilder.setTopN(topNBuilder.build()); return relBuilder.build(); } + + @Override + public List 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 a7a0ea62aa0..3bc448892f0 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 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 Rel toProtobuf() { builder.setWindowGroupLimit(windowBuilder.build()); return builder.build(); } + + @Override + public List 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 415bbc350fa..5e619cc6472 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 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 Rel toProtobuf() { builder.setWindow(windowBuilder.build()); return builder.build(); } + + @Override + public List 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 74ffc8282c6..9c875b1d014 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 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 Rel toProtobuf() { builder.setWrite(writeBuilder.build()); return builder.build(); } + + @Override + public List 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 a6b00535935..2cb45b97a82 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 37d8b2b3abc..fc13549decd 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) @@ -1432,6 +1434,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 3b142fecb95..6d265100d0f 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 8fea175d7ee..83023f027ac 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 4a00dbb5872..afec9cc10f6 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 8de9d407bfd..1f3ae4d753c 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 393716bb7b1..73068241f7d 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 00000000000..dfd24209e63 --- /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 c81d380a781..f35aea8fdc2 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) }