advancedxy commented on code in PR #100: URL: https://github.com/apache/arrow-datafusion-comet/pull/100#discussion_r1505642037
########## spark/src/main/scala/org/apache/spark/sql/comet/CometCollectLimitExec.scala: ########## @@ -0,0 +1,112 @@ +/* + * 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.comet + +import java.util.Objects + +import org.apache.spark.rdd.RDD +import org.apache.spark.serializer.Serializer +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.comet.execution.shuffle.{CometShuffledBatchRDD, CometShuffleExchangeExec} +import org.apache.spark.sql.execution.{ColumnarToRowExec, SparkPlan, UnaryExecNode, UnsafeRowSerializer} +import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics, SQLShuffleReadMetricsReporter, SQLShuffleWriteMetricsReporter} +import org.apache.spark.sql.vectorized.ColumnarBatch + +/** + * Comet physical plan node for Spark `CollectLimitExec`. + * + * Similar to `CometTakeOrderedAndProjectExec`, it contains two native executions seperated by a + * Comet shuffle. + * + * TODO: support offset semantics + */ +case class CometCollectLimitExec( + override val originalPlan: SparkPlan, + limit: Int, + offset: Int, + child: SparkPlan) + extends CometExec + with UnaryExecNode { + + private lazy val writeMetrics = + SQLShuffleWriteMetricsReporter.createShuffleWriteMetrics(sparkContext) + private lazy val readMetrics = + SQLShuffleReadMetricsReporter.createShuffleReadMetrics(sparkContext) + override lazy val metrics: Map[String, SQLMetric] = Map( + "dataSize" -> SQLMetrics.createSizeMetric(sparkContext, "data size"), + "shuffleReadElapsedCompute" -> + SQLMetrics.createNanoTimingMetric(sparkContext, "shuffle read elapsed compute at native"), + "numPartitions" -> SQLMetrics.createMetric( + sparkContext, + "number of partitions")) ++ readMetrics ++ writeMetrics + + private lazy val serializer: Serializer = + new UnsafeRowSerializer(child.output.size, longMetric("dataSize")) + + override def executeCollect(): Array[InternalRow] = { + ColumnarToRowExec(child).executeTake(limit) Review Comment: when offset = 0, limit cannot be `limit < 0`. See `CollectLimitExec`'s assert. <img width="803" alt="image" src="https://github.com/apache/arrow-datafusion-comet/assets/807537/6477d247-a3c6-42d9-9b00-51a98a171440";> Let's handle that case when we are adding offset support? ########## spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala: ########## @@ -451,6 +468,23 @@ class CometSparkSessionExtensions } } } + + // CometExec already wraps a `ColumnarToRowExec` for row-based operators. Therefore, + // `ColumnarToRowExec` is redundant and can be eliminated. + // + // It was added during ApplyColumnarRulesAndInsertTransitions' insertTransitions phase when Spark + // requests row-based output such as `collect` call. It's correct to add a redundant + // `ColumnarToRowExec` for `CometExec`. However, for certain operators such as + // `CometCollectLimitExec` which overrides `executeCollect`, the redundant `ColumnarToRowExec` + // makes the override ineffective. The purpose of this rule is to eliminate the redundant + // `ColumnarToRowExec` for such operators. + case class EliminateRedundantColumnarToRow(session: SparkSession) extends Rule[SparkPlan] { + override def apply(plan: SparkPlan): SparkPlan = { + plan.transform { case ColumnarToRowExec(child: CometCollectLimitExec) => + child Review Comment: Thanks, the suggest one is better. > Hmm, this looks like a bit dangerous if ColumnarToRowExec + CometCollectLimitExec is not end of the query. I'd like to point out that `ColumnarToRowExec + CometCollectLimitExec` will always be the end of the query as `CollectLimitExec` is the end of query. You can see the `SpecialLimits` rule which only translate the end of query to a `CollectLimitExec`. ########## spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala: ########## @@ -1073,6 +1073,34 @@ class CometExecSuite extends CometTestBase { } }) } + + test("collect limit") { + Seq("true", "false").foreach(aqe => { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe) { + withParquetTable((0 until 5).map(i => (i, i + 1)), "tbl") { + val df = sql("SELECT _1 as id, _2 as value FROM tbl limit 2") + assert(df.queryExecution.executedPlan.execute().getNumPartitions === 1) + checkSparkAnswerAndOperator(df, Seq(classOf[CometCollectLimitExec])) + assert(df.collect().length === 2) + + // checks CometCollectExec.doExecuteColumnar is indirectly called Review Comment: Yea, let me refine this one. ########## spark/src/main/scala/org/apache/spark/sql/comet/CometExecUtils.scala: ########## @@ -21,15 +21,40 @@ package org.apache.spark.sql.comet import scala.collection.JavaConverters.asJavaIterableConverter +import org.apache.spark.{Partition, SparkContext, TaskContext} +import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.expressions.{Attribute, NamedExpression, SortOrder} import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.comet.serde.OperatorOuterClass import org.apache.comet.serde.OperatorOuterClass.Operator import org.apache.comet.serde.QueryPlanSerde.{exprToProto, serializeDataType} object CometExecUtils { + /** + * Create an empty ColumnarBatch RDD with a single partition. + */ + def createEmptyColumnarRDDWithSinglePartition( + sparkContext: SparkContext): RDD[ColumnarBatch] = { + new EmptyRDDWithPartitions(sparkContext, 1) + } + + /** + * Transform the given RDD into a new RDD that takes the first `limit` elements of each + * partition. The limit operation is performed on the native side. + */ + def toNativeLimitedPerPartition( Review Comment: OK, thanks for your suggestion. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
