feat: Support CollectLimit operator #100

advancedxy · 2024-02-23T18:46:08Z

Which issue does this PR close?

Closes #37 .

Rationale for this change

Better operator coverage

What changes are included in this PR?

Add CometCollectLimitExec

How are these changes tested?

Add new test and some manual verification.

advancedxy · 2024-02-26T01:19:36Z

spark/src/main/scala/org/apache/spark/sql/comet/CometCoalesceExec.scala

-object CometCoalesceExec {
-
-  /** A simple RDD with no data, but with the given number of partitions. */
-  class EmptyRDDWithPartitions(@transient private val sc: SparkContext, numPartitions: Int)


It's moving to CometExecUtils, So it can be reused for several places.

advancedxy · 2024-02-26T04:48:26Z

All cleared. cc @sunchao @viirya

spark/src/main/scala/org/apache/spark/sql/comet/CometCollectLimitExec.scala

sunchao · 2024-02-27T07:08:31Z

spark/src/main/scala/org/apache/comet/shims/ShimCometSparkSessionExtensions.scala

@@ -32,4 +33,9 @@ trait ShimCometSparkSessionExtensions {
    .map { a => a.setAccessible(true); a }
    .flatMap(_.get(scan).asInstanceOf[Option[Aggregation]])
    .headOption
+
+  def getOffset(limit: LimitExec): Option[Int] = limit.getClass.getDeclaredFields


nit: I wonder if we can just return 0 if there is no offset in this method, so that we don't have to do getOffset(op).getOrElse(0) in a few places.

I considered it too. But it seems more nature to define a Option[Int] to access a potential non-existed field.

Let me reconsider this part.

sunchao · 2024-02-27T07:37:03Z

spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala

+  // `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] {


Hmm I'm trying to understand why this is necessary. The test passes even if I remove this rule.

Yea, I didn't add a test case for this part. Like noted in the comment, it's correct to add or remove the ColumnarToRowExec on top of a CometExec.

CollectLimitExec's executeCollect is optimized by using executeTake to take rows from child operator. Unlike CollectLimitExec.doExecute() or TakeOrderedAndProjectExec.doExecute(), which would shuffle all the data into a single partition and then get the limited data from shuffled partition, executeTake will retrieves rows directly from child's RDD without shuffle by partitions.

Take the following code for an example: sql("select * from a_very_large_table limit 100").collect(). CollectLimitExec's executeCollect will try to get the first 100 rows in the first partition, then the next 2 partitions if the previous partition doesn't contains 100 rows, then the next 4 partitions .... without shuffle.

I modeled this behavior(see https://github.com/apache/arrow-datafusion-comet/pull/100/files#diff-50c88b1d9b68e7ba24cb6fad9a4f20ea1b8fa63c3c868578db151b83182c627fR57) in CometCollectLimitExec as well. However, without this rule, an additional ColumnarToRowExec operator is wrapped on top of CometCollectLimitExec, which makes the override ineffective.

I added an assert in the test file, which should illustrate the basic idea.

// make sure the root node is CometCollectLimitExec assert(qe.executedPlan.isInstanceOf[CometCollectLimitExec])

I see, so if we have the extra ColumnarToRowExec, the code will go through its executeCollect instead which will call doExecuteÇolumnar, instead of calling the executeCollect in the CollectLimitExec itself.

I think we can probably do the same for CometTakeOrderedAndProjectExec too - Spark has an executeCollect implementation for this too. However, I don't know how useful it is since executeCollect is not often used? cc @viirya

I see, so if we have the extra ColumnarToRowExec, the code will go through its executeCollect instead which will call doExecuteÇolumnar, instead of calling the executeCollect in the CollectLimitExec itself.

Yeah, exactly.

I think we can probably do the same for CometTakeOrderedAndProjectExec too - Spark has an executeCollect implementation for this too.

I checked the implementation of TakeOrderedAndProjectExec.executeCollect when reviewing CometTakeOrderedAndProjectExec, it still shuffles all data into a single partition which is necessary to satisfy the ordering semantic. Hence it's not necessary to do the same for CometTakeOrderedAndProjectExec.

However, I don't know how useful it is since executeCollect is not often used

It's used in API/df scenarios, it's quite often for data scientists to collect and explore the data via collect with limit set. For the pure SQL and ETL scenario, I believe it's rarely used.

Thanks, makes sense. This also removes ColumnarToRowExec from the plan even if instead of executeCollect, doExecute is used, but I think it is OK since doExecute itself calls ColumnarToRowExec

advancedxy · 2024-02-27T15:09:54Z

spark/src/main/scala/org/apache/comet/shims/ShimCometSparkSessionExtensions.scala

+  /**
+   * TODO: delete after dropping Spark 3.2 and 3.3 support
+   */
+  def getOffset(limit: LimitExec): Int = getOffsetOpt(limit).getOrElse(0)


How do you like this? I think we should expose the getOffset method to accept LimitExec only and it could return Int directly.

The actual implementation could be generic.

Looks good 👍

advancedxy · 2024-02-27T15:11:07Z

spark/src/main/scala/org/apache/spark/sql/comet/CometTakeOrderedAndProjectExec.scala

-              CometExecUtils.getLimitNativePlan(output, limit).get
-            CometExec.getCometIterator(Seq(iter), limitOp)
-          }
+          CometExecUtils.toNativeLimitedPerPartition(childRDD, output, limit)


refactor to use the utility method.

If not appropriate, I can revert this.

sunchao

LGTM

sunchao · 2024-02-28T05:00:06Z

spark/src/main/scala/org/apache/comet/shims/ShimCometSparkSessionExtensions.scala

+  /**
+   * TODO: delete after dropping Spark 3.2 and 3.3 support
+   */
+  def getOffset(limit: LimitExec): Int = getOffsetOpt(limit).getOrElse(0)


Looks good 👍

sunchao · 2024-02-28T05:03:52Z

spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala

+  // `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] {


Thanks, makes sense. This also removes ColumnarToRowExec from the plan even if instead of executeCollect, doExecute is used, but I think it is OK since doExecute itself calls ColumnarToRowExec

viirya · 2024-02-28T08:08:14Z

spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala

+      plan.transform { case ColumnarToRowExec(child: CometCollectLimitExec) =>
+        child


Hmm, this looks like a bit dangerous if ColumnarToRowExec + CometCollectLimitExec is not end of the query.

I think the assumption here is the query is to collect data from ColumnarToRowExec + CometCollectLimitExec. So executeCollect is called on ColumnarToRowExec which makes ineffective of CometCollectLimitExec's executeCollect.

The more correct one maybe:

plan match { case ColumnarToRowExec(child: CometCollectLimitExec) => child case other => other }

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.

I'd like to point out that ColumnarToRowExec + CometCollectLimitExec will always be the end of the query as CollectLimitExec is the end of query.

Yes, this is usually the case. But I remember that at some special cases, users can produce a query tree that some others operators on top of CollectLimitExec. I think this is why CollectLimitExec still implements doExecute not just executeCollect.

viirya · 2024-02-28T08:12:17Z

spark/src/main/scala/org/apache/spark/sql/comet/CometExecUtils.scala

+   */
+  def createEmptyColumnarRDDWithSinglePartition(
+      sparkContext: SparkContext): RDD[ColumnarBatch] = {
+    new EmptyRDDWithPartitions(sparkContext, 1)
+  }


The method name is too long. This doesn't save the number of words. 😂

Maybe just keep original one.

spark/src/main/scala/org/apache/spark/sql/comet/CometExecUtils.scala

viirya · 2024-02-28T08:20:54Z

spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala

+      if (!plan.exists(op => planClass.isAssignableFrom(op.getClass))) {
+        assert(
+          false,
+          s"Expected plan to contain ${planClass.getSimpleName}.\n" +


Suggested change

s"Expected plan to contain ${planClass.getSimpleName}.\n" +

s"Expected plan to contain ${planClass.getSimpleName} but not.\n" +

spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala

viirya · 2024-02-28T08:23:11Z

spark/src/main/scala/org/apache/spark/sql/comet/CometCollectLimitExec.scala

+    new UnsafeRowSerializer(child.output.size, longMetric("dataSize"))
+
+  override def executeCollect(): Array[InternalRow] = {
+    ColumnarToRowExec(child).executeTake(limit)


Maybe we need to handle limit < 0 case.

when offset = 0, limit cannot be limit < 0.
See CollectLimitExec's assert.

Let's handle that case when we are adding offset support?

spark/src/main/scala/org/apache/spark/sql/comet/CometCollectLimitExec.scala

…_exec

sunchao · 2024-02-28T18:19:13Z

Merged, thanks!

advancedxy force-pushed the support_collect_limit_exec branch from 673bb7f to 2d0fbb2 Compare February 24, 2024 07:08

advancedxy marked this pull request as draft February 24, 2024 09:15

This was referenced Feb 24, 2024

Comet native shuffle in rust doesn't handle empty projection properly #102

Closed

data type Binary not supported in shuffle write #105

Closed

advancedxy force-pushed the support_collect_limit_exec branch 2 times, most recently from 009b7e2 to b07f8de Compare February 26, 2024 01:16

advancedxy marked this pull request as ready for review February 26, 2024 01:18

advancedxy commented Feb 26, 2024

View reviewed changes

sunchao reviewed Feb 27, 2024

View reviewed changes

advancedxy added 2 commits February 27, 2024 22:28

feat: Support CollectLimit operator

248cd03

Address comments

afb6513

advancedxy force-pushed the support_collect_limit_exec branch from b07f8de to afb6513 Compare February 27, 2024 15:05

advancedxy commented Feb 27, 2024

View reviewed changes

sunchao approved these changes Feb 28, 2024

View reviewed changes

viirya reviewed Feb 28, 2024

View reviewed changes

spark/src/main/scala/org/apache/spark/sql/comet/CometExecUtils.scala Outdated Show resolved Hide resolved

viirya reviewed Feb 28, 2024

View reviewed changes

spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala Outdated Show resolved Hide resolved

viirya reviewed Feb 28, 2024

View reviewed changes

spark/src/main/scala/org/apache/spark/sql/comet/CometCollectLimitExec.scala Outdated Show resolved Hide resolved

advancedxy added 3 commits February 28, 2024 19:10

address review comments

e15852f

Merge remote-tracking branch 'origin/main' into support_collect_limit…

5543837

…_exec

spotless:apply

f1c7718

viirya approved these changes Feb 28, 2024

View reviewed changes

sunchao merged commit 313111d into apache:main Feb 28, 2024
10 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: Support CollectLimit operator #100

feat: Support CollectLimit operator #100

advancedxy commented Feb 23, 2024

advancedxy Feb 26, 2024

advancedxy commented Feb 26, 2024

sunchao Feb 27, 2024

advancedxy Feb 27, 2024

sunchao Feb 27, 2024

advancedxy Feb 27, 2024 •

edited

Loading

advancedxy Feb 27, 2024

sunchao Feb 27, 2024

advancedxy Feb 28, 2024

sunchao Feb 28, 2024

advancedxy Feb 27, 2024

sunchao Feb 28, 2024

advancedxy Feb 27, 2024

sunchao left a comment

sunchao Feb 28, 2024

sunchao Feb 28, 2024

viirya Feb 28, 2024

advancedxy Feb 28, 2024

advancedxy Feb 28, 2024

viirya Feb 28, 2024

viirya Feb 28, 2024

advancedxy Feb 28, 2024

viirya Feb 28, 2024

viirya Feb 28, 2024

advancedxy Feb 28, 2024

viirya Feb 28, 2024

sunchao commented Feb 28, 2024

		plan.transform { case ColumnarToRowExec(child: CometCollectLimitExec) =>
		child

	s"Expected plan to contain ${planClass.getSimpleName}.\n" +
	s"Expected plan to contain ${planClass.getSimpleName} but not.\n" +

feat: Support CollectLimit operator #100

feat: Support CollectLimit operator #100

Conversation

advancedxy commented Feb 23, 2024

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

How are these changes tested?

Choose a reason for hiding this comment

advancedxy commented Feb 26, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

advancedxy Feb 27, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sunchao left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sunchao commented Feb 28, 2024

advancedxy Feb 27, 2024 •

edited

Loading